461 lines
16 KiB
Python
461 lines
16 KiB
Python
"""Recover DeepSeek DSML tool calls leaked through assistant text.
|
||
|
||
Some intranet OpenAI-compatible gateways render DeepSeek V4/V4.1 prompts
|
||
correctly but do not run the matching output parser. The model consequently
|
||
returns its native DSML envelope in ``message.content`` and leaves
|
||
``message.tool_calls`` empty. This module is a deliberately opt-in fallback
|
||
that converts only complete, schema-advertised calls back to OpenAI tool-call
|
||
chunks.
|
||
|
||
The parser is fail-closed:
|
||
|
||
* only tools advertised on the current provider request are accepted;
|
||
* names are matched exactly (misspellings are never corrected);
|
||
* an incomplete/malformed DSML envelope raises instead of executing a partial
|
||
call or exposing the protocol as ordinary assistant text;
|
||
* ordinary text is streamed without buffering except for the short prefix
|
||
needed to detect a marker split across provider chunks.
|
||
"""
|
||
|
||
from __future__ import annotations
|
||
|
||
import json
|
||
import re
|
||
import uuid
|
||
from collections.abc import Mapping
|
||
from dataclasses import dataclass, field
|
||
from typing import Any
|
||
|
||
from langchain_core.messages import AIMessageChunk
|
||
from langchain_core.outputs import ChatGenerationChunk
|
||
|
||
|
||
class DSMLToolCallProtocolError(RuntimeError):
|
||
"""A configured DSML provider returned an unsafe or incomplete call."""
|
||
|
||
|
||
@dataclass(frozen=True, slots=True)
|
||
class RecoveredToolCall:
|
||
name: str
|
||
arguments: dict[str, Any]
|
||
call_id: str
|
||
|
||
|
||
@dataclass(frozen=True, slots=True)
|
||
class ParsedDSMLContent:
|
||
content: str
|
||
calls: tuple[RecoveredToolCall, ...]
|
||
found: bool
|
||
|
||
|
||
_DSML_TOKEN = r"(?:｜DSML｜|\|DSML\|)"
|
||
_OUTER_NAMES = r"(?:tool_calls|toolcalls|function_calls|calls|tool)"
|
||
_START_RE = re.compile(
|
||
rf"<\s*{_DSML_TOKEN}\s*{_OUTER_NAMES}\s*>",
|
||
re.IGNORECASE,
|
||
)
|
||
_END_RE = re.compile(
|
||
rf"</\s*{_DSML_TOKEN}\s*{_OUTER_NAMES}\s*>",
|
||
re.IGNORECASE,
|
||
)
|
||
_ANY_DSML_RE = re.compile(rf"<\s*/?\s*{_DSML_TOKEN}", re.IGNORECASE)
|
||
_INVOKE_RE = re.compile(
|
||
rf"<\s*{_DSML_TOKEN}\s*invoke\s+name\s*=\s*"
|
||
rf"(?P<quote>[\"'])(?P<name>[^\"']+)(?P=quote)\s*>"
|
||
rf"(?P<body>.*?)"
|
||
rf"</\s*{_DSML_TOKEN}\s*invoke\s*>",
|
||
re.IGNORECASE | re.DOTALL,
|
||
)
|
||
_PARAM_RE = re.compile(
|
||
rf"<\s*{_DSML_TOKEN}\s*parameter\s+name\s*=\s*"
|
||
rf"(?P<quote>[\"'])(?P<name>[^\"']+)(?P=quote)"
|
||
rf"(?:\s+string\s*=\s*(?P<squote>[\"'])(?P<string>true|false)(?P=squote))?\s*>"
|
||
rf"(?P<value>.*?)"
|
||
rf"</\s*{_DSML_TOKEN}\s*parameter\s*>",
|
||
re.IGNORECASE | re.DOTALL,
|
||
)
|
||
|
||
|
||
def _marker_candidates() -> tuple[str, ...]:
|
||
names = ("tool_calls", "toolcalls", "function_calls", "calls", "tool")
|
||
tokens = ("｜DSML｜", "|DSML|")
|
||
return tuple(
|
||
f"<{token}{space}{name}>"
|
||
for token in tokens
|
||
for space in ("", " ")
|
||
for name in names
|
||
)
|
||
|
||
|
||
_START_CANDIDATES = _marker_candidates()
|
||
|
||
|
||
def _longest_possible_start_suffix(text: str) -> int:
|
||
"""Length of the suffix that may be a split DSML opening marker."""
|
||
max_len = min(len(text), max(map(len, _START_CANDIDATES)))
|
||
for size in range(max_len, 0, -1):
|
||
suffix = text[-size:]
|
||
if any(candidate.startswith(suffix) for candidate in _START_CANDIDATES):
|
||
return size
|
||
return 0
|
||
|
||
|
||
def _parse_parameter_value(raw: str, string_flag: str | None) -> Any:
|
||
if (string_flag or "").lower() == "true":
|
||
return raw
|
||
try:
|
||
return json.loads(raw)
|
||
except (TypeError, ValueError) as exc:
|
||
if string_flag is None:
|
||
# DeepSeek occasionally omits the flag for an otherwise ordinary
|
||
# string parameter. Match the provider parser's safe literal
|
||
# fallback, while keeping explicit string=false strict.
|
||
return raw
|
||
raise DSMLToolCallProtocolError(
|
||
"DeepSeek DSML contained invalid JSON for a non-string parameter"
|
||
) from exc
|
||
|
||
|
||
def _parse_invoke_body(body: str) -> dict[str, Any]:
|
||
params: dict[str, Any] = {}
|
||
spans: list[tuple[int, int]] = []
|
||
for match in _PARAM_RE.finditer(body):
|
||
name = match.group("name").strip()
|
||
if not name or name in params:
|
||
raise DSMLToolCallProtocolError(
|
||
"DeepSeek DSML contained an empty or duplicate parameter name"
|
||
)
|
||
params[name] = _parse_parameter_value(
|
||
match.group("value"),
|
||
match.group("string"),
|
||
)
|
||
spans.append(match.span())
|
||
|
||
if spans:
|
||
remainder_parts: list[str] = []
|
||
cursor = 0
|
||
for start, end in spans:
|
||
remainder_parts.append(body[cursor:start])
|
||
cursor = end
|
||
remainder_parts.append(body[cursor:])
|
||
if "".join(remainder_parts).strip():
|
||
raise DSMLToolCallProtocolError(
|
||
"DeepSeek DSML invoke mixed parameter tags with unexpected text"
|
||
)
|
||
return params
|
||
|
||
stripped = body.strip()
|
||
if not stripped:
|
||
return {}
|
||
try:
|
||
direct = json.loads(stripped)
|
||
except (TypeError, ValueError) as exc:
|
||
raise DSMLToolCallProtocolError(
|
||
"DeepSeek DSML invoke contained neither parameters nor a JSON object"
|
||
) from exc
|
||
if not isinstance(direct, dict):
|
||
raise DSMLToolCallProtocolError(
|
||
"DeepSeek DSML direct arguments must be a JSON object"
|
||
)
|
||
return direct
|
||
|
||
|
||
def parse_dsml_content(
|
||
text: str,
|
||
*,
|
||
allowed_tools: set[str],
|
||
) -> ParsedDSMLContent:
|
||
"""Parse complete V4/V4.1 DSML blocks and remove them from content."""
|
||
if not text or not _ANY_DSML_RE.search(text):
|
||
return ParsedDSMLContent(content=text, calls=(), found=False)
|
||
|
||
calls: list[RecoveredToolCall] = []
|
||
visible_parts: list[str] = []
|
||
cursor = 0
|
||
found = False
|
||
|
||
while True:
|
||
start = _START_RE.search(text, cursor)
|
||
if start is None:
|
||
visible_parts.append(text[cursor:])
|
||
break
|
||
found = True
|
||
visible_parts.append(text[cursor : start.start()])
|
||
end = _END_RE.search(text, start.end())
|
||
if end is None:
|
||
raise DSMLToolCallProtocolError(
|
||
"DeepSeek DSML tool-call block was not closed"
|
||
)
|
||
inner = text[start.end() : end.start()]
|
||
invoke_spans: list[tuple[int, int]] = []
|
||
for invoke in _INVOKE_RE.finditer(inner):
|
||
name = invoke.group("name").strip()
|
||
if name not in allowed_tools:
|
||
raise DSMLToolCallProtocolError(
|
||
f"DeepSeek DSML requested an unadvertised tool: {name!r}"
|
||
)
|
||
calls.append(
|
||
RecoveredToolCall(
|
||
name=name,
|
||
arguments=_parse_invoke_body(invoke.group("body")),
|
||
call_id=f"call_deerflow_dsml_{uuid.uuid4().hex}",
|
||
)
|
||
)
|
||
invoke_spans.append(invoke.span())
|
||
if not invoke_spans:
|
||
raise DSMLToolCallProtocolError(
|
||
"DeepSeek DSML tool-call block contained no complete invoke"
|
||
)
|
||
remainder_parts: list[str] = []
|
||
inner_cursor = 0
|
||
for invoke_start, invoke_end in invoke_spans:
|
||
remainder_parts.append(inner[inner_cursor:invoke_start])
|
||
inner_cursor = invoke_end
|
||
remainder_parts.append(inner[inner_cursor:])
|
||
if "".join(remainder_parts).strip():
|
||
raise DSMLToolCallProtocolError(
|
||
"DeepSeek DSML tool-call block contained unexpected text"
|
||
)
|
||
cursor = end.end()
|
||
|
||
if not found:
|
||
# A DSML sentinel was present but did not match any supported complete
|
||
# envelope. Do not leak it as chat text or guess a repair.
|
||
raise DSMLToolCallProtocolError(
|
||
"DeepSeek DSML marker used an unsupported tool-call envelope"
|
||
)
|
||
return ParsedDSMLContent(
|
||
content="".join(visible_parts),
|
||
calls=tuple(calls),
|
||
found=True,
|
||
)
|
||
|
||
|
||
def _payload_tool_names(payload: Mapping[str, Any]) -> set[str]:
|
||
names: set[str] = set()
|
||
tools = payload.get("tools")
|
||
if not isinstance(tools, list):
|
||
return names
|
||
for tool in tools:
|
||
function = tool.get("function") if isinstance(tool, Mapping) else None
|
||
name = function.get("name") if isinstance(function, Mapping) else None
|
||
if isinstance(name, str) and name:
|
||
names.add(name)
|
||
return names
|
||
|
||
|
||
@dataclass(slots=True)
|
||
class DSMLToolCallRecovery:
|
||
"""Per-provider-invocation DSML compatibility state."""
|
||
|
||
enabled: bool
|
||
allowed_tools: set[str] = field(default_factory=set)
|
||
pending: str = ""
|
||
capture: str = ""
|
||
capture_template: ChatGenerationChunk | None = None
|
||
recovered: bool = False
|
||
|
||
def configure_from_payload(self, payload: Mapping[str, Any]) -> None:
|
||
if self.enabled:
|
||
self.allowed_tools = _payload_tool_names(payload)
|
||
|
||
def recover_nonstream_response(self, response: dict[str, Any]) -> None:
|
||
if not self.enabled or not self.allowed_tools:
|
||
return
|
||
choices = response.get("choices")
|
||
if not isinstance(choices, list):
|
||
return
|
||
for choice in choices:
|
||
if not isinstance(choice, dict):
|
||
continue
|
||
message = choice.get("message")
|
||
if not isinstance(message, dict) or message.get("tool_calls"):
|
||
continue
|
||
content = message.get("content")
|
||
if not isinstance(content, str):
|
||
continue
|
||
parsed = parse_dsml_content(content, allowed_tools=self.allowed_tools)
|
||
if not parsed.found:
|
||
continue
|
||
message["content"] = parsed.content
|
||
message["tool_calls"] = [
|
||
{
|
||
"id": call.call_id,
|
||
"type": "function",
|
||
"function": {
|
||
"name": call.name,
|
||
"arguments": json.dumps(
|
||
call.arguments,
|
||
ensure_ascii=False,
|
||
separators=(",", ":"),
|
||
),
|
||
},
|
||
}
|
||
for call in parsed.calls
|
||
]
|
||
choice["finish_reason"] = "tool_calls"
|
||
self.recovered = True
|
||
|
||
def feed(self, chunk: ChatGenerationChunk) -> list[ChatGenerationChunk]:
|
||
if not self.enabled or not self.allowed_tools:
|
||
return [chunk]
|
||
message = chunk.message
|
||
if not isinstance(message, AIMessageChunk):
|
||
return [chunk]
|
||
if message.tool_call_chunks or message.tool_calls:
|
||
if self.capture:
|
||
raise DSMLToolCallProtocolError(
|
||
"Provider mixed native tool calls with an incomplete DSML block"
|
||
)
|
||
return [chunk]
|
||
|
||
content = message.content
|
||
if not isinstance(content, str):
|
||
return [chunk]
|
||
|
||
if self.recovered:
|
||
if _ANY_DSML_RE.search(content):
|
||
raise DSMLToolCallProtocolError(
|
||
"Provider emitted another DSML block after a recovered tool call"
|
||
)
|
||
return [self._force_tool_finish(chunk)]
|
||
|
||
if self.capture:
|
||
self.capture += content
|
||
self.capture_template = chunk
|
||
return self._complete_capture_if_ready()
|
||
|
||
combined = self.pending + content
|
||
had_pending = bool(self.pending)
|
||
self.pending = ""
|
||
start = _START_RE.search(combined)
|
||
if start is not None:
|
||
before = combined[: start.start()]
|
||
self.capture = combined[start.start() :]
|
||
self.capture_template = chunk
|
||
output: list[ChatGenerationChunk] = []
|
||
if before or self._has_noncontent_signal(message):
|
||
output.append(self._replace_content(chunk, before))
|
||
output.extend(self._complete_capture_if_ready())
|
||
return output
|
||
|
||
hold = _longest_possible_start_suffix(combined)
|
||
if hold:
|
||
safe = combined[:-hold]
|
||
self.pending = combined[-hold:]
|
||
self.capture_template = chunk
|
||
if safe or self._has_noncontent_signal(message):
|
||
return [self._replace_content(chunk, safe)]
|
||
return []
|
||
return [self._replace_content(chunk, combined)] if had_pending else [chunk]
|
||
|
||
def finish(self) -> list[ChatGenerationChunk]:
|
||
if not self.enabled:
|
||
return []
|
||
if self.capture:
|
||
raise DSMLToolCallProtocolError(
|
||
"Provider stream ended inside a DeepSeek DSML tool-call block"
|
||
)
|
||
if not self.pending:
|
||
return []
|
||
pending = self.pending
|
||
self.pending = ""
|
||
if self.capture_template is None:
|
||
return []
|
||
return [self._replace_content(self.capture_template, pending)]
|
||
|
||
def _complete_capture_if_ready(self) -> list[ChatGenerationChunk]:
|
||
if _END_RE.search(self.capture) is None:
|
||
return []
|
||
parsed = parse_dsml_content(
|
||
self.capture,
|
||
allowed_tools=self.allowed_tools,
|
||
)
|
||
if not parsed.calls or self.capture_template is None:
|
||
raise DSMLToolCallProtocolError(
|
||
"DeepSeek DSML recovery produced no complete tool call"
|
||
)
|
||
template = self.capture_template
|
||
self.capture = ""
|
||
self.recovered = True
|
||
message = template.message
|
||
assert isinstance(message, AIMessageChunk)
|
||
tool_chunks = [
|
||
{
|
||
"name": call.name,
|
||
"args": json.dumps(
|
||
call.arguments,
|
||
ensure_ascii=False,
|
||
separators=(",", ":"),
|
||
),
|
||
"id": call.call_id,
|
||
"index": index,
|
||
"type": "tool_call_chunk",
|
||
}
|
||
for index, call in enumerate(parsed.calls)
|
||
]
|
||
response_metadata = dict(message.response_metadata)
|
||
if response_metadata.get("finish_reason"):
|
||
response_metadata["finish_reason"] = "tool_calls"
|
||
recovered_message = AIMessageChunk(
|
||
content=parsed.content,
|
||
additional_kwargs=dict(message.additional_kwargs),
|
||
response_metadata=response_metadata,
|
||
id=message.id,
|
||
name=message.name,
|
||
usage_metadata=message.usage_metadata,
|
||
tool_call_chunks=tool_chunks, # type: ignore[arg-type]
|
||
)
|
||
generation_info = dict(template.generation_info or {})
|
||
if generation_info.get("finish_reason"):
|
||
generation_info["finish_reason"] = "tool_calls"
|
||
return [
|
||
ChatGenerationChunk(
|
||
message=recovered_message,
|
||
generation_info=generation_info,
|
||
)
|
||
]
|
||
|
||
@staticmethod
|
||
def _replace_content(
|
||
chunk: ChatGenerationChunk,
|
||
content: str,
|
||
) -> ChatGenerationChunk:
|
||
return ChatGenerationChunk(
|
||
message=chunk.message.model_copy(update={"content": content}),
|
||
generation_info=chunk.generation_info,
|
||
)
|
||
|
||
@staticmethod
|
||
def _has_noncontent_signal(message: AIMessageChunk) -> bool:
|
||
return bool(
|
||
message.additional_kwargs
|
||
or message.response_metadata
|
||
or message.usage_metadata
|
||
)
|
||
|
||
@staticmethod
|
||
def _force_tool_finish(chunk: ChatGenerationChunk) -> ChatGenerationChunk:
|
||
message = chunk.message
|
||
if not isinstance(message, AIMessageChunk):
|
||
return chunk
|
||
response_metadata = dict(message.response_metadata)
|
||
if response_metadata.get("finish_reason"):
|
||
response_metadata["finish_reason"] = "tool_calls"
|
||
generation_info = dict(chunk.generation_info or {})
|
||
if generation_info.get("finish_reason"):
|
||
generation_info["finish_reason"] = "tool_calls"
|
||
return ChatGenerationChunk(
|
||
message=message.model_copy(update={"response_metadata": response_metadata}),
|
||
generation_info=generation_info or chunk.generation_info,
|
||
)
|
||
|
||
|
||
__all__ = [
|
||
"DSMLToolCallProtocolError",
|
||
"DSMLToolCallRecovery",
|
||
"ParsedDSMLContent",
|
||
"RecoveredToolCall",
|
||
"parse_dsml_content",
|
||
]
|