Skip to content

Commit ecafc65

Browse files
author
Your Name
committed
Parse XML-like structure in agent mode (especially useful for local Qwen models)
1 parent a1bcba6 commit ecafc65

2 files changed

Lines changed: 125 additions & 64 deletions

File tree

‎cecli/coders/base_coder.py‎

Lines changed: 14 additions & 64 deletions
Original file line numberDiff line numberDiff line change
@@ -32,15 +32,15 @@
3232

3333
import httpx
3434
from litellm import experimental_mcp_client
35-
from litellm.types.utils import ChatCompletionMessageToolCall, Function, ModelResponse
35+
from litellm.types.utils import ModelResponse
3636
from prompt_toolkit.patch_stdout import patch_stdout
3737
from rich.console import Console
3838

3939
import cecli.prompts.utils.system as prompts
4040
from cecli import __version__, models, urls, utils
4141
from cecli.commands import Commands, SwitchCoderSignal
4242
from cecli.exceptions import LiteLLMExceptions
43-
from cecli.helpers import command_parser, coroutines, nested
43+
from cecli.helpers import command_parser, coroutines, nested, responses
4444
from cecli.helpers.conversation import (
4545
ConversationChunks,
4646
ConversationManager,
@@ -1581,8 +1581,8 @@ async def run_one(self, user_message, preproc):
15811581
self.reflected_message = None
15821582
self.tool_reflection = False
15831583

1584-
if float(self.total_cost) > self.cost_multiplier * nested.getter(
1585-
self.args, "cost_limit", float("inf")
1584+
if float(self.total_cost) > self.cost_multiplier * (
1585+
nested.getter(self.args, "cost_limit", float("inf")) or float("inf")
15861586
):
15871587
if await self.io.confirm_ask(
15881588
"You have reached your configured cost limit. Continue?",
@@ -3311,66 +3311,16 @@ def consolidate_chunks(self):
33113311
# If no native tool calls, check if the content contains JSON tool calls
33123312
# This handles models that write JSON in text instead of using native calling
33133313
if not self.partial_response_tool_calls and self.partial_response_content:
3314-
try:
3315-
# Simple extraction of JSON-like structures that look like tool calls
3316-
# Only look for tool calls if it looks like JSON
3317-
if "{" in self.partial_response_content or "[" in self.partial_response_content:
3318-
json_chunks = utils.split_concatenated_json(self.partial_response_content)
3319-
extracted_calls = []
3320-
chunk_index = 0
3321-
3322-
for chunk in json_chunks:
3323-
chunk_index += 1
3324-
try:
3325-
json_obj = json.loads(chunk)
3326-
if (
3327-
isinstance(json_obj, dict)
3328-
and "name" in json_obj
3329-
and "arguments" in json_obj
3330-
):
3331-
# Create a Pydantic model for the tool call
3332-
function_obj = Function(
3333-
name=json_obj["name"],
3334-
arguments=(
3335-
json.dumps(json_obj["arguments"])
3336-
if isinstance(json_obj["arguments"], (dict, list))
3337-
else str(json_obj["arguments"])
3338-
),
3339-
)
3340-
tool_call_obj = ChatCompletionMessageToolCall(
3341-
type="function",
3342-
function=function_obj,
3343-
id=f"call_{len(extracted_calls)}_{int(time.time())}_{chunk_index}",
3344-
)
3345-
extracted_calls.append(tool_call_obj)
3346-
elif isinstance(json_obj, list):
3347-
for item in json_obj:
3348-
if (
3349-
isinstance(item, dict)
3350-
and "name" in item
3351-
and "arguments" in item
3352-
):
3353-
function_obj = Function(
3354-
name=item["name"],
3355-
arguments=(
3356-
json.dumps(item["arguments"])
3357-
if isinstance(item["arguments"], (dict, list))
3358-
else str(item["arguments"])
3359-
),
3360-
)
3361-
tool_call_obj = ChatCompletionMessageToolCall(
3362-
type="function",
3363-
function=function_obj,
3364-
id=f"call_{len(extracted_calls)}_{int(time.time())}_{chunk_index}",
3365-
)
3366-
extracted_calls.append(tool_call_obj)
3367-
except json.JSONDecodeError:
3368-
continue
3369-
3370-
if extracted_calls:
3371-
self.partial_response_tool_calls = extracted_calls
3372-
except Exception:
3373-
pass
3314+
extracted_calls = responses.extract_tools_from_content_json(
3315+
self.partial_response_content
3316+
)
3317+
if not extracted_calls:
3318+
extracted_calls = responses.extract_tools_from_content_xml(
3319+
self.partial_response_content
3320+
)
3321+
3322+
if extracted_calls:
3323+
self.partial_response_tool_calls = extracted_calls
33743324

33753325
return response, func_err, content_err
33763326

‎cecli/helpers/responses.py‎

Lines changed: 111 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,11 @@
1+
import json
12
import re
3+
import time
4+
from typing import List, Optional
5+
6+
from litellm.types.utils import ChatCompletionMessageToolCall, Function
7+
8+
from cecli import utils
29

310

411
def preprocess_json(response: str) -> str:
@@ -19,3 +26,107 @@ def normalize(match):
1926
return "\\\\" + suffix
2027

2128
return re.sub(pattern, normalize, response)
29+
30+
31+
def extract_tools_from_content_json(content: str) -> Optional[List[ChatCompletionMessageToolCall]]:
32+
"""
33+
Simple extraction of JSON-like structures that look like tool calls.
34+
This handles models that write JSON in text instead of using native calling.
35+
"""
36+
if not content or ("{" not in content and "[" not in content):
37+
return None
38+
39+
try:
40+
json_chunks = utils.split_concatenated_json(content)
41+
extracted_calls = []
42+
chunk_index = 0
43+
44+
for chunk in json_chunks:
45+
chunk_index += 1
46+
try:
47+
json_obj = json.loads(chunk)
48+
if isinstance(json_obj, dict) and "name" in json_obj and "arguments" in json_obj:
49+
# Create a Pydantic model for the tool call
50+
function_obj = Function(
51+
name=json_obj["name"],
52+
arguments=(
53+
json.dumps(json_obj["arguments"])
54+
if isinstance(json_obj["arguments"], (dict, list))
55+
else str(json_obj["arguments"])
56+
),
57+
)
58+
tool_call_obj = ChatCompletionMessageToolCall(
59+
type="function",
60+
function=function_obj,
61+
id=f"call_{len(extracted_calls)}_{int(time.time())}_{chunk_index}",
62+
)
63+
extracted_calls.append(tool_call_obj)
64+
elif isinstance(json_obj, list):
65+
for item in json_obj:
66+
if isinstance(item, dict) and "name" in item and "arguments" in item:
67+
function_obj = Function(
68+
name=item["name"],
69+
arguments=(
70+
json.dumps(item["arguments"])
71+
if isinstance(item["arguments"], (dict, list))
72+
else str(item["arguments"])
73+
),
74+
)
75+
tool_call_obj = ChatCompletionMessageToolCall(
76+
type="function",
77+
function=function_obj,
78+
id=f"call_{len(extracted_calls)}_{int(time.time())}_{chunk_index}",
79+
)
80+
extracted_calls.append(tool_call_obj)
81+
except json.JSONDecodeError:
82+
continue
83+
84+
return extracted_calls if extracted_calls else None
85+
except Exception:
86+
return None
87+
88+
89+
def extract_tools_from_content_xml(content: str) -> Optional[List[ChatCompletionMessageToolCall]]:
90+
"""
91+
Extraction of Qwen-style XML tool calls.
92+
Example:
93+
<function=UpdateTodoList>
94+
<parameter=tasks>
95+
[{"task": "Update task list", "done": false, "current": true}]
96+
</parameter>
97+
</function>
98+
"""
99+
if not content or "<function=" not in content:
100+
return None
101+
102+
try:
103+
extracted_calls = []
104+
# Find all blocks between <function=...> and </function>
105+
func_blocks = re.finditer(r"<function=(.*?)>(.*?)</function>", content, re.DOTALL)
106+
107+
for i, block_match in enumerate(func_blocks):
108+
func_name = block_match.group(1).strip()
109+
block_content = block_match.group(2).strip()
110+
111+
params_dict = {}
112+
param_pattern = r"<parameter=(.*?)>(.*?)</parameter>"
113+
for param_match in re.finditer(param_pattern, block_content, re.DOTALL):
114+
key = param_match.group(1).strip()
115+
value_str = param_match.group(2).strip()
116+
try:
117+
params_dict[key] = json.loads(value_str)
118+
except json.JSONDecodeError:
119+
params_dict[key] = value_str
120+
121+
function_obj = Function(name=func_name, arguments=json.dumps(params_dict))
122+
123+
tool_call_obj = ChatCompletionMessageToolCall(
124+
type="function",
125+
function=function_obj,
126+
id=f"xml_call_{i}_{int(time.time())}",
127+
)
128+
extracted_calls.append(tool_call_obj)
129+
130+
return extracted_calls if extracted_calls else None
131+
except Exception:
132+
return None

0 commit comments

Comments
 (0)