Skip to content

feat(proxy): rescue tool calls emitted inside reasoning_content (#1) #4

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
Hrt-Htk merged 1 commit into from
Jun 15, 2026
Merged

feat(proxy): rescue tool calls emitted inside reasoning_content (#1) #4

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
266 changes: 234 additions & 32 deletions proxy.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -815,13 +815,28 @@ def _stringify_message_content(content: object) -> str:


class SSEChunkLogger:
# Markers that can start a tool-call XML block
_START_MARKERS = ("<tool_call>", "<function=")
# End markers keyed by kind
_END_MARKERS = {"tool_call": "</tool_call>", "function": "</function>"}

def __init__(self, wrapped, chat_logger: ChatLogger) -> None:
self._wrapped = wrapped
self._chat_logger = chat_logger
self._buffer = b""
# --- existing logging state ---
self._current_kind: str | None = None
self._current_text = ""
self._tool_calls: dict[int, dict[str, str]] = {}
# --- rescue state machine ---
self._rescue_capturing: bool = False
self._rescue_buf: str = ""
self._rescue_kind: str | None = None # "tool_call" or "function"
self._reasoning_holdback: str = ""
self._rescued_any: bool = False
self._rescue_index: int = 0
# cache upstream chunk id for synthesised events
self._last_chunk_id: str = "rescued"

async def _flush_text(self) -> None:
if self._current_kind and self._current_text:
Expand All @@ -845,12 +860,93 @@ async def _flush_all(self) -> None:
await self._flush_text()
await self._flush_tool_calls()

# ---- rescue helpers ----

@staticmethod
def _split_safe_prefix(text: str, markers: tuple[str, ...]) -> tuple[str, str]:
"""Return (emit, holdback) where holdback is the longest suffix of *text*
that is a proper prefix of any *marker*."""
for length in range(len(text), 0, -1):
suffix = text[len(text) - length:]
for marker in markers:
if len(suffix) < len(marker) and marker.startswith(suffix):
return text[: len(text) - length], suffix
return text, ""

@staticmethod
def _parse_tool_call_xml(block: str) -> dict | None:
"""Parse a tool-call XML block. Returns {name, arguments} or None."""
# Find <function=NAME>
import re as _re
m = _re.search(r"<function=([^>\s]+)", block)
if not m:
return None
name: str = m.group(1)
# Find all <parameter=KEY>VALUE</parameter>
args: dict[str, object] = {}
for pm in _re.finditer(r"<parameter=(\S+)>(.*?)</parameter>", block, _re.DOTALL):
key = pm.group(1)
value = pm.group(2).strip()
# Coerce: try JSON parse
try:
value = json.loads(value)
except (json.JSONDecodeError, ValueError):
pass
args[key] = value
return {"name": name, "arguments": args}

def _build_synthesized_event(self, parsed: dict) -> bytes:
"""Build a synthesised tool_calls SSE event from a parsed tool call."""
import secrets as _secrets
event: dict = {
"id": self._last_chunk_id,
"object": "chat.completion.chunk",
"choices": [
{
"index": 0,
"delta": {
"tool_calls": [
{
"index": self._rescue_index,
"id": f"call_{_secrets.token_hex(4)}",
"type": "function",
"function": {
"name": parsed["name"],
"arguments": json.dumps(parsed["arguments"]),
},
}
]
},
"finish_reason": None,
}
],
}
self._rescue_index += 1
body = json.dumps(event, ensure_ascii=False)
return f"data: {body}\r\n\r\n".encode()

# ---- main read loop ----

async def readany(self) -> bytes:
# Loop so we only ever return b"" at true EOF — the consumer treats an
# empty return as end-of-stream. A single upstream chunk may not complete
# an SSE event, in which case _readany_once returns None and we read more.
while True:
out = await self._readany_once()
if out is not None:
return out

async def _readany_once(self) -> bytes | None:
data = await self._wrapped.content.readany()
if not data:
await self._flush_all()
return data
if self._buffer:
leftover = self._buffer
self._buffer = b""
return leftover
return b""
self._buffer += data
outbound: list[bytes] = []
while True:
crlf_idx = self._buffer.find(b"\r\n\r\n")
lf_idx = self._buffer.find(b"\n\n")
Expand All @@ -860,53 +956,159 @@ async def readany(self) -> bytes:
idx, sep_len = crlf_idx, 4
else:
idx, sep_len = lf_idx, 2
event = self._buffer[:idx]
raw_event = self._buffer[: idx + sep_len]
self._buffer = self._buffer[idx + sep_len:]
text = event.decode("utf-8", errors="replace").strip()
if not text:
event_text = raw_event.decode("utf-8", errors="replace").strip()
if not event_text:
outbound.append(raw_event)
continue
# Extract payload line
payload = ""
for line in text.splitlines():
for line in event_text.splitlines():
if line.startswith("data:"):
payload = line[5:].strip()
# Non-data lines, comments, [DONE] → pass through
if not payload:
outbound.append(raw_event)
continue
if payload == "[DONE]":
await self._flush_all()
await self._chat_logger.log_response("[DONE]", True)
outbound.append(raw_event)
continue
try:
obj = json.loads(payload)
except json.JSONDecodeError:
outbound.append(raw_event)
continue
# --- (a) existing logging on original delta ---
choices = obj.get("choices") or []
if not choices:
continue
delta = choices[0].get("delta") or {}
reasoning = delta.get("reasoning_content")
content = delta.get("content")
tool_calls = delta.get("tool_calls")
if reasoning:
if self._current_kind != "thinking":
await self._flush_all()
self._current_kind = "thinking"
self._current_text += reasoning
if content:
if self._current_kind != "content":
await self._flush_all()
self._current_kind = "content"
self._current_text += content
if tool_calls:
await self._flush_text()
for tc in tool_calls:
i = tc.get("index", 0)
slot = self._tool_calls.setdefault(i, {"name": "", "arguments": ""})
fn = tc.get("function") or {}
if fn.get("name"):
slot["name"] = fn["name"]
if fn.get("arguments"):
slot["arguments"] += fn["arguments"]
return data
if choices:
delta = choices[0].get("delta") or {}
reasoning = delta.get("reasoning_content")
content = delta.get("content")
tool_calls = delta.get("tool_calls")
if reasoning:
if self._current_kind != "thinking":
await self._flush_all()
self._current_kind = "thinking"
self._current_text += reasoning
if content:
if self._current_kind != "content":
await self._flush_all()
self._current_kind = "content"
self._current_text += content
if tool_calls:
await self._flush_text()
for tc in tool_calls:
i = tc.get("index", 0)
slot = self._tool_calls.setdefault(i, {"name": "", "arguments": ""})
fn = tc.get("function") or {}
if fn.get("name"):
slot["name"] = fn["name"]
if fn.get("arguments"):
slot["arguments"] += fn["arguments"]
# Track chunk id for synthesised events
cid = obj.get("id")
if cid:
self._last_chunk_id = cid
# --- (b) build outbound bytes with rescue transform ---
outbound_event = self._transform_event(obj)
outbound.append(outbound_event)
return b"".join(outbound) if outbound else None

def _transform_event(self, obj: dict) -> bytes:
"""Transform a single parsed event dict into outbound SSE bytes,
applying the rescue state machine."""
choices = obj.get("choices") or []
if not choices:
# No choices — pass through
body = json.dumps(obj, ensure_ascii=False)
return f"data: {body}\r\n\r\n".encode()

delta = choices[0].get("delta") or {}
reasoning = delta.get("reasoning_content")

# If no reasoning_content, just rewrite finish_reason if needed
if not reasoning:
if self._rescued_any and choices[0].get("finish_reason") == "stop":
choices[0]["finish_reason"] = "tool_calls"
body = json.dumps(obj, ensure_ascii=False)
return f"data: {body}\r\n\r\n".encode()

# Run rescue state machine on reasoning_content
work = self._reasoning_holdback + reasoning
self._reasoning_holdback = ""
prose_parts: list[str] = []
synthesized: list[bytes] = []

while work:
if not self._rescue_capturing:
# Look for earliest start marker
earliest_pos = len(work)
earliest_marker: str | None = None
for marker in self._START_MARKERS:
pos = work.find(marker)
if pos != -1 and pos < earliest_pos:
earliest_pos = pos
earliest_marker = marker

if earliest_marker is None:
# No marker found — apply split_safe_prefix
emit, holdback = self._split_safe_prefix(work, self._START_MARKERS)
prose_parts.append(emit)
self._reasoning_holdback = holdback
work = ""
else:
# A complete start marker is present, so everything before it
# is safe prose — no partial-marker holdback needed here.
prose_parts.append(work[:earliest_pos])
# Start capturing
self._rescue_capturing = True
self._rescue_kind = (
"tool_call" if earliest_marker == "<tool_call>" else "function"
)
self._rescue_buf = earliest_marker
work = work[earliest_pos + len(earliest_marker):]
else:
# Capturing — look for end marker
end_marker = self._END_MARKERS[self._rescue_kind]
end_pos = work.find(end_marker)
if end_pos != -1:
self._rescue_buf += work[: end_pos + len(end_marker)]
# Parse the block
parsed = self._parse_tool_call_xml(self._rescue_buf)
if parsed:
synthesized.append(self._build_synthesized_event(parsed))
self._rescued_any = True
self._rescue_capturing = False
self._rescue_buf = ""
self._rescue_kind = None
work = work[end_pos + len(end_marker):]
else:
# End marker not found — keep all of work in buffer
self._rescue_buf += work
work = ""

# Build the outbound event
forwarded_reasoning = "".join(prose_parts)
if forwarded_reasoning:
delta["reasoning_content"] = forwarded_reasoning
else:
delta.pop("reasoning_content", None)
# If delta is now empty and has no other keys, we still emit the event
# (the caller handles skipping if needed)

# Rewrite finish_reason if rescued
if self._rescued_any and choices[0].get("finish_reason") == "stop":
choices[0]["finish_reason"] = "tool_calls"

body = json.dumps(obj, ensure_ascii=False)
result = f"data: {body}\r\n\r\n".encode()
# Append any synthesised events after the reasoning event
for syn in synthesized:
result += syn
return result

def __getattr__(self, name: str) -> object:
return getattr(self._wrapped, name)
Expand Down
Empty file added tests/__init__.py
View file
Open in desktop
Empty file.
Loading
Loading