chore: make the MistralFunctionCallingAdapter slightly more robust at…

… parsing
zhudotexe · Dec 10, 2024 · 5b93954 · 5b93954
1 parent 252eec1
commit 5b93954
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/kani/prompts/impl/mistral.py b/kani/prompts/impl/mistral.py
@@ -197,7 +197,7 @@ def __init__(self, *args, tool_call_token="[TOOL_CALLS]", eos_token="</s>", **kw
 
     def _parse_tool_calls(self, content: str) -> tuple[str, list[ToolCall]]:
         tool_json = re.search(
-            rf"{re.escape(self.tool_call_token)}\s*(.+?)\s*{re.escape(self.eos_token)}",
+            rf"{re.escape(self.tool_call_token)}\s*(.+?)\s*({re.escape(self.eos_token)})?$",
             content,
             re.IGNORECASE | re.DOTALL,
         )