atakan commited on
Commit
53fb4c4
·
1 Parent(s): 2739ebf

fix: Fall back to end-of-string match for truncated tool-call JSON

Browse files

When generation is truncated mid-array (e.g. a repetition loop hits the
token budget before closing), the closing-brace-anchored regexes in
_extract_tool_calls fail to match, letting raw truncated JSON leak into
the chat instead of being caught by the degenerate-array guard.

Files changed (1) hide show
  1. controlai_agent/orchestrator.py +10 -0
controlai_agent/orchestrator.py CHANGED
@@ -200,6 +200,13 @@ def _extract_tool_calls(text: str) -> tuple[list[dict[str, Any]], str]:
200
  # 3. Raw JSON object containing "name" and "arguments" / "parameters"
201
  if not calls and ('"name"' in text or "'name'" in text):
202
  match = re.search(r"(\{\s*[\"']name[\"']\s*:\s*[\"'][a-zA-Z0-9_]+[\"'][\s\S]*\})", text)
 
 
 
 
 
 
 
203
  if match:
204
  obj = parse_flexible_json(match.group(1))
205
  if obj and "name" in obj:
@@ -209,6 +216,9 @@ def _extract_tool_calls(text: str) -> tuple[list[dict[str, Any]], str]:
209
  cleaned = re.sub(r"<tool_call>[\s\S]*?</tool_call>", "", text, flags=re.DOTALL)
210
  cleaned = re.sub(r"```(?:json)?\s*\{\s*[\"']name[\"']\s*:[\s\S]*?\}\s*```", "", cleaned)
211
  cleaned = re.sub(r"\{\s*[\"']name[\"']\s*:\s*[\"'][a-zA-Z0-9_]+[\"'][\s\S]*\}", "", cleaned)
 
 
 
212
  # An orphaned <tool_call> tag with no matching close (the model started a
213
  # call, then abandoned it mid-generation for plain text) survives the
214
  # paired regex above -- strip any leftover tag so it never reaches the UI.
 
200
  # 3. Raw JSON object containing "name" and "arguments" / "parameters"
201
  if not calls and ('"name"' in text or "'name'" in text):
202
  match = re.search(r"(\{\s*[\"']name[\"']\s*:\s*[\"'][a-zA-Z0-9_]+[\"'][\s\S]*\})", text)
203
+ if not match:
204
+ # Generation was truncated mid-JSON (e.g. a repetition loop hit the
205
+ # token budget before the array closed), so there's no closing "}"
206
+ # to anchor on -- fall back to matching to the end of the string so
207
+ # close_unbalanced_json can still repair it and _degenerate_array_reason
208
+ # can refuse it, instead of the raw truncated JSON leaking into the chat.
209
+ match = re.search(r"(\{\s*[\"']name[\"']\s*:\s*[\"'][a-zA-Z0-9_]+[\"'][\s\S]*)", text)
210
  if match:
211
  obj = parse_flexible_json(match.group(1))
212
  if obj and "name" in obj:
 
216
  cleaned = re.sub(r"<tool_call>[\s\S]*?</tool_call>", "", text, flags=re.DOTALL)
217
  cleaned = re.sub(r"```(?:json)?\s*\{\s*[\"']name[\"']\s*:[\s\S]*?\}\s*```", "", cleaned)
218
  cleaned = re.sub(r"\{\s*[\"']name[\"']\s*:\s*[\"'][a-zA-Z0-9_]+[\"'][\s\S]*\}", "", cleaned)
219
+ # Same truncated-JSON fallback as above, so the leftover raw JSON is
220
+ # stripped from the visible text even when it never closed.
221
+ cleaned = re.sub(r"\{\s*[\"']name[\"']\s*:\s*[\"'][a-zA-Z0-9_]+[\"'][\s\S]*", "", cleaned)
222
  # An orphaned <tool_call> tag with no matching close (the model started a
223
  # call, then abandoned it mid-generation for plain text) survives the
224
  # paired regex above -- strip any leftover tag so it never reaches the UI.