From e74d64a55b8065b7ed80b42ac4a58b96ba48c8e3 Mon Sep 17 00:00:00 2001 From: yihonglie Date: Fri, 10 Jul 2026 10:19:44 -0500 Subject: [PATCH 01/18] [Frontend] DeepSeek-V4 native OpenAI/Anthropic/Responses API + DSML tool parser Serve DeepSeek-V4 (Pro/Flash) so Claude Code (Anthropic /v1/messages) and Codex CLI (OpenAI Responses /v1/responses) talk to ATOM directly, with DSML tool-call parsing shared across /v1/chat/completions, /v1/messages and /v1/responses. - tool_parser: DeepSeek-V4 DSML tool-call format (<|DSML|invoke ...>) with marker-less / self-closing / direct-JSON recovery, schema-driven type coercion and key aliases; streaming + non-streaming (alongside Qwen/GLM/MiniMax). - serving_chat: reasoning-filter + tool-call streaming for /v1/chat/completions. - serving_responses (new) + /v1/responses: OpenAI Responses translation (input/tools<->chat, SSE emitter, reasoning/function_call items). Codex compat: inject a mandatory DSML tool-format instruction and normalize shell tool name/param aliases (exec/shell_exec/... -> registered exec tool; command/script -> cmd) so Codex tool calls execute instead of erroring. - /v1/messages: pass the tool schema to the parser so tool args are typed correctly (fixes "Invalid tool parameters"); streaming UTF-8 incremental detokenization (vLLM-style sliding window) so multi-byte chars (CJK, box-drawing) aren't split into U+FFFD. - reasoning / chat_encoders: DeepSeek-V4 message encoding + reasoning tags. Builds on the abort-on-disconnect path from the previous commit. Co-Authored-By: Claude Opus 4.8 (1M context) --- atom/entrypoints/openai/api_server.py | 325 ++++++++++- atom/entrypoints/openai/chat_encoders.py | 53 +- atom/entrypoints/openai/reasoning.py | 7 + atom/entrypoints/openai/serving_responses.py | 572 ++++++++++++++++++ atom/entrypoints/openai/tool_parser.py | 585 ++++++++++++++++++- 5 files changed, 1523 insertions(+), 19 deletions(-) create mode 100644 atom/entrypoints/openai/serving_responses.py diff --git a/atom/entrypoints/openai/api_server.py b/atom/entrypoints/openai/api_server.py index c8a24f2817..12640649bb 100644 --- a/atom/entrypoints/openai/api_server.py +++ b/atom/entrypoints/openai/api_server.py @@ -68,6 +68,18 @@ stream_completion_response, stream_completion_response_fanout, ) +from .serving_responses import ( + ResponsesStreamEmitter, + inject_tool_format_instruction, + shell_arg_key, + build_responses_object, + remap_tool_name, + extract_cwd, + responses_input_to_messages, + responses_tools_to_openai, + tool_name_lookup, + translate_client_tool, +) # Configure logging logger = logging.getLogger("atom") @@ -324,6 +336,24 @@ def _prepare_multimodal_inputs( return inputs["input_ids"][0].tolist(), multimodal_data +# ── Batched stream dispatch ────────────────────────────────────────────── +# Per-seq `call_soon_threadsafe` floods the API event loop at high batch size +# (one call per token). Instead the callback only buffers the raw chunk; the +# mgr flushes a whole step with a single `tokenizer.batch_decode` (one +# GIL-released call instead of one decode per seq) plus one scheduled call per +# loop (see `flush_stream_batch`). +import threading as _threading # noqa: E402 + +_stream_batch_tls = _threading.local() + +# Per-request incremental detokenization state (vLLM-style sliding window). +# Decoding each step's new tokens in isolation splits multi-byte UTF-8 chars +# (byte-BPE tokenizers like DeepSeek-V4 split one CJK char across several +# byte-tokens) into U+FFFD. Keep accumulated tokens + prefix/read offsets so +# we only emit fully-formed characters. +_stream_detok_state: Dict[str, dict] = {} + + def _send_stream_chunk_direct( request_output: RequestOutput, request_id: str, @@ -346,7 +376,70 @@ def _send_stream_chunk_direct( } if getattr(request_output, "kv_transfer_params_output", None): chunk_data["kv_transfer_params"] = request_output.kv_transfer_params_output - loop.call_soon_threadsafe(stream_queue.put_nowait, chunk_data) + + chunk_data["request_id"] = request_id + buf = getattr(_stream_batch_tls, "buf", None) + if buf is None: + buf = _stream_batch_tls.buf = [] + buf.append((loop, stream_queue, chunk_data)) + + +def _drain_batch_into_queues(items: list) -> None: + """Runs ON the event loop: push each chunk into its per-request queue. + One scheduled call handles a whole step's worth of chunks.""" + for _loop, q, chunk in items: + q.put_nowait(chunk) + + +def flush_stream_batch() -> None: + """Flush a step's buffered chunks: one ``batch_decode`` for the whole step, + then one call_soon_threadsafe per loop (normally one — all requests on a + rank share the API loop).""" + global tokenizer + + buf = getattr(_stream_batch_tls, "buf", None) + if not buf: + return + _stream_batch_tls.buf = [] + # Decode the whole step in a single call. batch_decode is element-wise + # identical to per-seq decode but acquires/releases the GIL once instead of + # once per seq, cutting GIL ping-pong against the other rank output threads + # and the API event loop at high batch size. + # Incremental per-request detokenization: correct UTF-8 at token + # boundaries (see _stream_detok_state). Emits only fully-formed chars; + # a trailing partial multi-byte char is held until the next step. + for (_loop, _q, chunk) in buf: + rid = chunk.get("request_id") + st = _stream_detok_state.get(rid) + if st is None: + st = _stream_detok_state[rid] = { + "tokens": [], "prefix_offset": 0, "read_offset": 0, + } + toks = st["tokens"] + toks.extend(chunk["token_ids"]) + prefix_text = tokenizer.decode( + toks[st["prefix_offset"]:st["read_offset"]], skip_special_tokens=True + ) + new_text = tokenizer.decode( + toks[st["prefix_offset"]:], skip_special_tokens=True + ) + if len(new_text) > len(prefix_text) and not new_text.endswith("\ufffd"): + chunk["text"] = new_text[len(prefix_text):] + st["prefix_offset"] = st["read_offset"] + st["read_offset"] = len(toks) + elif chunk["finished"]: + chunk["text"] = new_text[len(prefix_text):] + else: + chunk["text"] = "" + if chunk["finished"]: + _stream_detok_state.pop(rid, None) + # Group by loop (normally a single loop). dict preserves insertion order + # so per-request chunk ordering within the step is maintained. + by_loop: Dict[AbstractEventLoop, list] = {} + for loop, q, chunk in buf: + by_loop.setdefault(loop, []).append((loop, q, chunk)) + for loop, items in by_loop.items(): + loop.call_soon_threadsafe(_drain_batch_into_queues, items) def _send_stream_chunk_tagged( @@ -1433,6 +1526,7 @@ async def generate_anthropic_stream(): if prompt.rstrip().endswith(""): reasoning_filter.state = 1 tool_parser = ToolCallStreamParser() + tool_parser.tools = anthropic_to_openai_tools(request.tools) block_index = 0 started_text = False started_thinking = False @@ -1602,15 +1696,17 @@ async def generate_anthropic_stream(): from .reasoning import separate_reasoning from .tool_parser import parse_tool_calls - final_output = None - async for output in generate_async(prompt, sampling_params, request_id): - final_output = output + final_output = await _run_nonstream_with_disconnect( + generate_async(prompt, sampling_params, request_id), + raw_request, + request_id, + ) if final_output is None: raise RuntimeError("No output generated") raw_text = final_output["text"] reasoning_content, content_with_tools = separate_reasoning(raw_text) - content_text, tool_calls = parse_tool_calls(content_with_tools) + content_text, tool_calls = parse_tool_calls(content_with_tools, anthropic_to_openai_tools(request.tools)) output_tokens = len(tokenizer.encode(raw_text)) cache_read_input_tokens = final_output.get("num_cached_tokens", 0) if not getattr(request, "thinking", None): @@ -1627,6 +1723,9 @@ async def generate_anthropic_stream(): cache_read_input_tokens=cache_read_input_tokens, ) + except _ClientDisconnected: + # Client hung up; seq already aborted + popped. Nothing to return. + return JSONResponse(status_code=499, content={"detail": "client disconnected"}) except Exception as e: logger.error(f"Error in anthropic_messages: {e}", exc_info=True) return JSONResponse( @@ -1638,6 +1737,222 @@ async def generate_anthropic_stream(): ) +@app.post("/v1/responses") +async def responses_endpoint(raw_request: Request): + """Handle OpenAI **Responses API** requests (`/v1/responses`). + + Native support so OpenAI Codex CLI (>= 0.14x, which speaks only the Responses + API) can talk to ATOM directly — no external responses->chat proxy needed. + Reuses ATOM's proven streaming path (setup_streaming_request + ReasoningFilter + + ToolCallStreamParser), the same one /v1/messages (claude-local) uses, so it + streams correctly for reasoning models. Stateless (full input sent each turn, + as Codex does); reasoning items are dropped from visible output. + """ + global engine, tokenizer, model_name + + try: + body = await raw_request.json() + model = body.get("model") or model_name + + from .protocol import ChatMessage + from .reasoning import ReasoningFilter, separate_reasoning + from .tool_parser import ToolCallStreamParser, parse_tool_calls + + openai_tools = responses_tools_to_openai(body.get("tools")) + valid_names, shell_tool = tool_name_lookup(openai_tools) + shell_param = shell_arg_key(openai_tools, shell_tool) + req_cwd = extract_cwd(body) # Codex — used to fix hallucinated paths + openai_messages = responses_input_to_messages( + body.get("instructions"), body.get("input") + ) + if openai_tools: + openai_messages = inject_tool_format_instruction(openai_messages) + messages = [ChatMessage(**m) for m in openai_messages] + + merged_kwargs = dict(default_chat_template_kwargs) + prompt = apply_chat_template( + tokenizer, + custom_message_encoder, + [msg.to_template_dict() for msg in messages], + tools=openai_tools or None, + **merged_kwargs, + ) + + max_out = int(body.get("max_output_tokens") or 32768) + sampling_params = _build_sampling_params( + temperature=body.get("temperature") if body.get("temperature") is not None else 1.0, + max_tokens=max_out, + stop_strings=None, + ignore_eos=False, + top_k=-1, + top_p=body.get("top_p") if body.get("top_p") is not None else 1.0, + ) + + request_id = "resp_" + uuid.uuid4().hex[:24] + input_tokens = len(tokenizer.encode(prompt)) + + # Resolve max context to bound the prompt (same probes as anthropic). + max_ctx = None + for _path in ( + lambda: engine.config.max_model_len, + lambda: engine.model_config.max_model_len, + lambda: engine.scheduler.max_model_len, + lambda: getattr(engine, "max_model_len"), + ): + try: + _v = _path() + if _v: + max_ctx = int(_v) + break + except Exception: + continue + if not max_ctx: + max_ctx = 30720 + headroom = min(max_out, max(1024, max_ctx // 8)) + max_input = max_ctx - headroom + if input_tokens > max_input: + logger.warning( + f"[responses] prompt too long ({input_tokens} > {max_input}), truncating" + ) + token_ids = tokenizer.encode(prompt)[:max_input] + prompt = tokenizer.decode(token_ids, skip_special_tokens=False) + input_tokens = max_input + + if body.get("stream"): + seq_id, stream_queue, _num_prompt_tokens = await setup_streaming_request( + prompt, sampling_params, request_id + ) + + async def generate_responses_stream(): + emitter = ResponsesStreamEmitter(request_id, model) + reasoning_filter = ReasoningFilter() + if prompt.rstrip().endswith(""): + reasoning_filter.state = 1 + tool_parser = ToolCallStreamParser() + tool_parser.tools = openai_tools or None # enables schema-based + # type coercion + key-alias (command->cmd) in _parse_dsml + output_tokens = 0 + + # Buffer each tool call (name + full args) so read/grep/ls/find + # can be translated to exec_command before emitting (name AND + # args change). See translate_client_tool. + _pending = {"tc": None} + + def _flush_pending(): + tc = _pending["tc"] + if tc is None: + return [] + _pending["tc"] = None + name, args = translate_client_tool( + tc["name"], tc["args"], valid_names, shell_tool, req_cwd, shell_param + ) + out = emitter.tool_start(tc["id"], name) + if args: + out += emitter.tool_args(args) + out += emitter.tool_end() + return out + + def handle(etype, edata): + # Map ToolCallStreamParser events -> Responses SSE strings, + # buffering tool calls for client-tool translation. + if etype == "content": + out = _flush_pending() + return out + emitter.text_delta(edata) + if etype == "tool_call_start": + out = _flush_pending() + fn = edata.get("function", {}) + _pending["tc"] = { + "id": edata.get("id", ""), + "name": fn.get("name", ""), + "args": "", + } + return out + if etype == "tool_call_args": + if _pending["tc"] is not None: + _pending["tc"]["args"] += ( + edata.get("function", {}).get("arguments", "") or "" + ) + return [] + if etype == "tool_call_end": + return _flush_pending() + return [] + + try: + for s in emitter.created(): + yield s + while True: + chunk_data = await stream_queue.get() + new_text = chunk_data["text"] + output_tokens += len(chunk_data.get("token_ids", [])) + finished = chunk_data.get("finished", False) + + segments = reasoning_filter.process(new_text) + if finished: + segments.extend(reasoning_filter.flush()) + for field, text in segments: + if not text or field == "reasoning_content": + continue # drop reasoning from visible output + for etype, edata in tool_parser.process(text): + for s in handle(etype, edata): + yield s + + if finished: + for etype, edata in tool_parser.flush(): + for s in handle(etype, edata): + yield s + for s in _flush_pending(): # emit any unclosed tool call + yield s + for s in emitter.finish(input_tokens, output_tokens): + yield s + yield "data: [DONE]\n\n" + break + finally: + cleanup_streaming_request(request_id, seq_id) + + return StreamingResponse( + generate_responses_stream(), + media_type="text/event-stream", + headers={"x-request-id": request_id}, + ) + + # Non-streaming response + final_output = await _run_nonstream_with_disconnect( + generate_async(prompt, sampling_params, request_id), + raw_request, + request_id, + ) + if final_output is None: + raise RuntimeError("No output generated") + + raw_text = final_output["text"] + _reasoning, content_with_tools = separate_reasoning(raw_text) + content_text, tool_calls = parse_tool_calls(content_with_tools, openai_tools or None) + output_tokens = len(tokenizer.encode(raw_text)) + + return JSONResponse( + content=build_responses_object( + resp_id=request_id, + model=model, + content_text=content_text, + tool_calls=tool_calls, + input_tokens=input_tokens, + output_tokens=output_tokens, + valid=valid_names, + shell_tool=shell_tool, + cwd=req_cwd, + ) + ) + + except _ClientDisconnected: + return JSONResponse(status_code=499, content={"detail": "client disconnected"}) + except Exception as e: + logger.error(f"Error in responses_endpoint: {e}", exc_info=True) + return JSONResponse( + status_code=500, + content={"error": {"type": "api_error", "message": str(e)}}, + ) + + @app.get("/v1/models") async def list_models(): """List available models.""" diff --git a/atom/entrypoints/openai/chat_encoders.py b/atom/entrypoints/openai/chat_encoders.py index d3fb0462ce..863d045275 100644 --- a/atom/entrypoints/openai/chat_encoders.py +++ b/atom/entrypoints/openai/chat_encoders.py @@ -85,6 +85,46 @@ def load_custom_message_encoder(model_path: str) -> Optional[MessageEncoder]: return _load_encoder_from_dir(_resolve_model_path(model_path)) +def _content_str(c: Any) -> str: + if isinstance(c, list): + return "\n".join( + b.get("text", "") for b in c if isinstance(b, dict) and b.get("type") == "text" + ) + return c or "" + + +def _normalize_for_v4(messages: List[dict], tools: Optional[List[dict]]) -> List[dict]: + """Prepare messages for DeepSeek-V4's ``encode_messages``. + + Two things: + 1. **Hoist system messages to the front.** Clients (notably Claude Code) send + a trailing ``system``-role message (its "skills" list) AFTER the user turn. + ``encode_messages`` only appends the ``<|Assistant|>`` generation marker + after a *user*/developer message, so a trailing system message leaves the + prompt ending mid-system-text and the model just *continues* it instead of + answering. Merging all system content into one leading system message keeps + the final turn a user turn, so the assistant marker is emitted. + 2. **Attach tools** to that leading system message (``encode_messages`` reads + tool schemas from a system message's ``tools`` field). + Does not mutate the input. + """ + sys_parts, others = [], [] + for m in messages: + (sys_parts if m.get("role") == "system" else others).append(dict(m)) + + if not sys_parts and not tools: + return [dict(m) for m in messages] + + merged = "\n\n".join(s for s in (_content_str(m.get("content")) for m in sys_parts) if s) + sys_msg: dict = {"role": "system", "content": merged} + for m in sys_parts: # preserve any pre-attached tools + if m.get("tools"): + sys_msg["tools"] = m["tools"] + if tools: + sys_msg["tools"] = tools + return [sys_msg] + others + + def apply_chat_template( tokenizer: Any, custom_encoder: Optional[MessageEncoder], @@ -97,18 +137,15 @@ def apply_chat_template( Dispatches to ``custom_encoder`` if one was discovered for this model, otherwise to ``tokenizer.apply_chat_template``. Jinja-only kwargs - (``tokenize``, ``add_generation_prompt``) are stripped on the custom - path; ``tools`` are forwarded only on the Jinja path (custom encoders - don't currently have a tools API — caller is warned and tools are - dropped). + (``tokenize``, ``add_generation_prompt``) are stripped on the custom path. + ``tools`` are supported on both paths: custom encoders (e.g. DeepSeek-V4's + ``encode_messages``) read tool schemas from a system message's ``tools`` + field, so we attach them there before encoding. """ if custom_encoder is not None: for k in ("tokenize", "add_generation_prompt"): kwargs.pop(k, None) - if tools: - logger.warning( - "tools= is not supported with the custom message encoder; ignoring." - ) + messages = _normalize_for_v4(messages, tools) return custom_encoder(messages, **kwargs) kwargs["tokenize"] = False diff --git a/atom/entrypoints/openai/reasoning.py b/atom/entrypoints/openai/reasoning.py index 6fb8a8e004..f7cd8f0b41 100644 --- a/atom/entrypoints/openai/reasoning.py +++ b/atom/entrypoints/openai/reasoning.py @@ -23,6 +23,9 @@ def separate_reasoning(text: str) -> Tuple[Optional[str], str]: Tuple of (reasoning_content, content). reasoning_content is None if no thinking block was found. """ + # MiniMax M3 emits ... instead of ...; + # normalize so the shared logic below handles both. + text = text.replace("", "").replace("", "") # Check for closed thinking block: ... match = re.match(r"(.*?)\s*(.*)", text, flags=re.DOTALL) if match: @@ -75,6 +78,10 @@ def process(self, text: str) -> list: List of (field_name, text) tuples where field_name is "reasoning_content" or "content". """ + # MiniMax M3 uses /; normalize to the tags + # the state machine below keys on. These are single special tokens, so + # each arrives whole in one chunk — a plain replace is safe. + text = text.replace("", "").replace("", "") results = [] if self.state == 0: diff --git a/atom/entrypoints/openai/serving_responses.py b/atom/entrypoints/openai/serving_responses.py new file mode 100644 index 0000000000..bc14091ad0 --- /dev/null +++ b/atom/entrypoints/openai/serving_responses.py @@ -0,0 +1,572 @@ +"""OpenAI **Responses API** (`/v1/responses`) support for the ATOM server. + +OpenAI Codex CLI (>= 0.14x) dropped `wire_api = "chat"` and only speaks the +Responses API (streaming SSE). This module provides the translation between +Responses request/response shapes and ATOM's internal chat/engine machinery, +plus a streaming SSE event emitter. The `/v1/responses` route handler in +``api_server.py`` reuses ATOM's proven streaming path (``setup_streaming_request`` ++ ``ReasoningFilter`` + ``ToolCallStreamParser``) so it streams correctly for +reasoning models — the same path claude-local uses via ``/v1/messages``. + +This makes the external ``codex_responses_proxy.py`` unnecessary: Codex can point +straight at ATOM's ``:9700/v1``. +""" +import itertools +import json +import time +from typing import Any, Dict, List, Optional, Tuple + +_ids = itertools.count(1) + + +def _rid(prefix: str) -> str: + return f"{prefix}_{int(time.time() * 1000)}{next(_ids):04d}" + + +# --------------------------------------------------------------- request xlate +def _text_of(content: Any) -> str: + """Flatten Responses content (str | list of parts) to a plain string.""" + if isinstance(content, str): + return content + if isinstance(content, list): + out = [] + for p in content: + if isinstance(p, dict): + if "text" in p and isinstance(p["text"], str): + out.append(p["text"]) + elif p.get("type") in ("input_text", "output_text", "text"): + out.append(p.get("text", "")) + elif isinstance(p, str): + out.append(p) + return "".join(out) + return "" + + +def responses_input_to_messages( + instructions: Any, inp: Any +) -> List[Dict[str, Any]]: + """Translate Responses ``instructions`` + ``input`` into OpenAI chat messages. + + ``input`` may be a plain string or a list of items: ``message`` / + ``function_call`` (assistant tool call) / ``function_call_output`` (tool + result). ``reasoning`` items are dropped. + """ + messages: List[Dict[str, Any]] = [] + if instructions: + messages.append({"role": "system", "content": _text_of(instructions)}) + + if isinstance(inp, str): + messages.append({"role": "user", "content": inp}) + elif isinstance(inp, list): + for item in inp: + if not isinstance(item, dict): + continue + t = item.get("type", "message") + if t == "message": + role = item.get("role", "user") + messages.append( + {"role": role, "content": _text_of(item.get("content", ""))} + ) + elif t == "function_call": + messages.append( + { + "role": "assistant", + "content": "", + "tool_calls": [ + { + "id": item.get("call_id") + or item.get("id") + or _rid("call"), + "type": "function", + "function": { + "name": item.get("name", ""), + "arguments": item.get("arguments", "") or "", + }, + } + ], + } + ) + elif t == "function_call_output": + out = item.get("output", "") + messages.append( + { + "role": "tool", + "tool_call_id": item.get("call_id") or item.get("id") or "", + "content": out + if isinstance(out, str) + else json.dumps(out), + } + ) + elif t == "reasoning": + continue + return messages + + +_DSML_TOOL_INSTRUCTION = ( + "\n\n# Tool-call format (MANDATORY — overrides any other format instruction)\n" + "When you call a tool, output ONLY a DSML tool-call block and NOTHING else " + "in that message — no markdown, no ```json, no ///" + " tags, no prose. Use EXACTLY this syntax (the \uff5c characters " + "are U+FF5C fullwidth vertical bars, not ASCII '|'):\n" + "<\uff5cDSML\uff5ctool_calls>\n" + "<\uff5cDSML\uff5cinvoke name=\"TOOL_NAME\">\n" + "<\uff5cDSML\uff5cparameter name=\"PARAM_NAME\" string=\"true\">VALUE" + "\n" + "\n" + "\n" + "Use the exact tool and parameter names from the tools provided to you. " + "For a shell/exec tool, put the whole shell command string in its " + "command/cmd parameter. Emit one <\uff5cDSML\uff5cinvoke> per tool call." +) + + +def inject_tool_format_instruction(messages): + """Append the mandatory DSML tool-call format to the system message so the + model emits parseable DSML instead of ad-hoc /```json text. Codex + (/v1/responses) path only. Idempotent per request.""" + for m in messages: + if m.get("role") == "system": + base = m.get("content") or "" + if "\uff5cDSML\uff5ctool_calls" not in base: + m["content"] = _text_of(base) + _DSML_TOOL_INSTRUCTION + return messages + return [{"role": "system", "content": _DSML_TOOL_INSTRUCTION.strip()}] + list(messages) + + +def responses_tools_to_openai(tools: Any) -> List[Dict[str, Any]]: + """Translate Responses function tools into OpenAI chat tool defs. + + Responses puts ``name``/``description``/``parameters`` at the top level of a + ``{"type": "function", ...}`` tool (chat nests them under ``function``). + Non-function tool types (``web_search``, ``namespace``, ...) are dropped — + ATOM only executes model-emitted function/DSML tool calls; the client + (Codex) owns actual tool execution. + """ + if not tools: + return [] + ct: List[Dict[str, Any]] = [] + for tl in tools: + if not isinstance(tl, dict): + continue + if tl.get("type") == "function": + fn = tl.get("function", tl) + ct.append( + { + "type": "function", + "function": { + "name": fn.get("name"), + "description": fn.get("description", ""), + "parameters": fn.get("parameters", {}) or {}, + }, + } + ) + return ct + + +# ------------------------------------------------- tool-name normalization +# Non-codex-tuned models (e.g. DeepSeek-V4-Pro) don't reliably emit the EXACT +# tool name the client registered for shell exec. Codex 0.142.x names it +# ``exec_command`` (args ``{"cmd": "..."}``), but the model habitually calls it +# ``exec`` / ``exec_run`` / ``shell`` / ``bash``, which Codex's tool router +# rejects ("unsupported call: exec"). The ARGUMENTS are correct; only the name +# is wrong. So remap known shell-exec aliases onto whatever shell tool the +# client actually registered this turn. Pure rename; arguments untouched. +_SHELL_ALIASES = { + "exec", "exec_run", "exec_command", "execute_command", "shell", "bash", + "sh", "run", "run_command", "run_shell", "execute", "command", + "container.exec", "local_shell", "shell_command", "shell_exec", + "run_bash", "execute_shell", "bash_command", "run_terminal_cmd", + "terminal", "console", "run_shell_command", "runshell", +} +# Substrings that mark an unknown tool name as a shell/exec call (Codex's +# non-shell tools contain none of these). +_SHELL_NAME_TOKENS = ("shell", "exec", "bash", "cmd", "command", "termin", "console") +_SHELL_TOOL_PREFERENCE = ( + "exec_command", "shell", "local_shell", "bash", "container.exec", +) + + +def tool_name_lookup(openai_tools: List[Dict[str, Any]]) -> Tuple[set, Optional[str]]: + """Return (valid tool names, preferred shell tool name) for remapping.""" + valid = { + (t.get("function") or {}).get("name") + for t in (openai_tools or []) + if isinstance(t, dict) + } + valid.discard(None) + shell_tool = next((n for n in _SHELL_TOOL_PREFERENCE if n in valid), None) + return valid, shell_tool + + +def remap_tool_name(name: str, valid: set, shell_tool: Optional[str]) -> str: + """Fix a model tool_call name that doesn't match a registered tool. + + Only remaps shell-exec aliases to the registered shell tool; other + mismatches pass through unchanged (the client will surface them).""" + if name in valid: + return name + if not shell_tool: + return name + n = (name or "").lower().replace("-", "_").replace(".", "_") + if n in _SHELL_ALIASES: + return shell_tool + # Fuzzy: unknown tool whose name signals shell/exec -> the shell tool. + if any(k in n for k in _SHELL_NAME_TOKENS): + return shell_tool + return name + + +# ------------------------------------------------ Claude-tool -> shell adapter +# DeepSeek-V4-Pro is trained on Claude-Code's toolset, so under Codex it keeps +# calling read/grep/ls/find (which Codex doesn't have) instead of exec_command, +# and flails ("unsupported call: read"). When Codex registered an exec/shell +# tool, translate these read-only Claude tools into an equivalent shell command +# so the model's intent goes through. Only fires for names NOT in the registered +# set; exec_command's real calls are untouched. Codex-path only (never applied +# to /v1/messages, where read/grep ARE native tools). +def _q(s: Any) -> str: + import shlex + return shlex.quote(str(s)) + + +def _resolve_dir(path: str, cwd: Optional[str]) -> str: + """Pick a directory that actually exists: keep relative paths and paths under + cwd; otherwise fall back to cwd (the model often invents absolute prefixes + like /sglang or /sgl-workspace that don't exist here).""" + if not path or path == ".": + return cwd or "." + if not path.startswith("/"): + return path # relative — shell runs in cwd, fine + if cwd and path.startswith(cwd): + return path + return cwd or path + + +def _read_cmd(fp: str, cwd: Optional[str], start: int, end: int) -> str: + """Resilient file read: try the literal path, then cwd-prefixed, then locate + by basename under cwd — so a hallucinated absolute prefix still finds the file.""" + sed = f"sed -n {start},{end}p" + if not cwd: + return f"{sed} {_q(fp)}" + # f = literal; if missing, cwd + '/' + (f without leading /); if still + # missing, first match of `find cwd -name basename`. + return ( + f'f={_q(fp)}; [ -f "$f" ] || f={_q(cwd)}"/${{f#/}}"; ' + f'[ -f "$f" ] || f=$(find {_q(cwd)} -type f -name "$(basename {_q(fp)})" ' + f'2>/dev/null | head -1); {sed} "$f"' + ) + + +def _claude_tool_to_shell( + name: str, a: Dict[str, Any], cwd: Optional[str] = None +) -> Optional[str]: + n = (name or "").lower() + fp = (a.get("file_path") or a.get("path") or a.get("filePath") + or a.get("filename") or a.get("target_file") or a.get("file")) + pattern = a.get("pattern") or a.get("query") or a.get("regex") + path = (a.get("path") or a.get("directory") or a.get("target_directory") + or a.get("dir") or ".") + if n in ("read", "cat", "view", "view_file", "open", "read_file", + "readfile", "openfile"): + if not fp: + return None + off, lim = a.get("offset"), a.get("limit") + if off or lim: + start = int(off or 0) + 1 + end = start + int(lim or 200) - 1 + else: + start, end = 1, 400 + return _read_cmd(str(fp), cwd, start, end) + if n in ("grep", "search", "search_file", "ripgrep", "rg", "grep_search", + "codebase_search"): + if not pattern: + return None + return f"grep -rn -- {_q(pattern)} {_q(_resolve_dir(path, cwd))}" + if n in ("ls", "list", "list_dir", "list_directory", "listdir"): + return f"ls -la {_q(_resolve_dir(path, cwd))}" + if n in ("find", "glob", "glob_file_search", "file_search"): + base = _resolve_dir(path, cwd) + if pattern: + return f"find {_q(base)} -name {_q(pattern)}" + return f"find {_q(base)} -maxdepth 3" + return None + + +_SHELL_ARG_ALIASES = ( + "cmd", "command", "commandline", "command_line", "script", "bash", "sh", + "shell", "shell_command", "code", "input", "run", "cmd_string", +) + + +def shell_arg_key(openai_tools, shell_tool): + """Required (or first) param name of the registered shell tool, e.g. Codex's + exec_command -> "cmd". Used to normalize model arg keys.""" + if not shell_tool: + return None + for t in openai_tools or []: + fn = t.get("function", t) + if (fn.get("name") or t.get("name")) != shell_tool: + continue + params = fn.get("parameters") or {} + props = params.get("properties") or {} + for r in (params.get("required") or []): + if props.get(r, {}).get("type") in (None, "string"): + return r + if props: + return next(iter(props)) + return "cmd" + + +def _is_shell_name(name, shell_tool): + return name == shell_tool or name in _SHELL_ALIASES + + +def _normalize_shell_args(args_json, req): + """If the shell tool's required param `req` is absent but a known alias is + present, rename it. Keeps exec_command calls valid when the model uses + `command`/`script`/... instead of `cmd`.""" + if not req: + return args_json + try: + a = json.loads(args_json) if args_json else {} + except Exception: + return args_json + if not isinstance(a, dict) or req in a: + return args_json + for alias in _SHELL_ARG_ALIASES: + if alias != req and isinstance(a.get(alias), str): + a[req] = a.pop(alias) + return json.dumps(a) + return args_json + + +def translate_client_tool( + name: str, args_json: str, valid: set, shell_tool: Optional[str], + cwd: Optional[str] = None, shell_param: Optional[str] = None, +): + """Return (name, args_json) with Claude read-only tools rewritten to the + registered shell tool. exec_command's own calls and any already-valid tool + pass through untouched; unknown non-shell tools fall back to name-remap. + ``cwd`` (from the request's ) makes hallucinated absolute paths resolve.""" + if name in valid: + if shell_param and _is_shell_name(name, shell_tool): + args_json = _normalize_shell_args(args_json, shell_param) + return name, args_json + exec_tool = "exec_command" if "exec_command" in valid else shell_tool + if exec_tool: + try: + a = json.loads(args_json) if args_json else {} + except Exception: + a = {} + if isinstance(a, dict): + cmd = _claude_tool_to_shell(name, a, cwd) + if cmd is not None: + return exec_tool, json.dumps({(shell_param or "cmd"): cmd}) + _remapped = remap_tool_name(name, valid, shell_tool) + if shell_param and _is_shell_name(_remapped, shell_tool): + args_json = _normalize_shell_args(args_json, shell_param) + return _remapped, args_json + + +_CWD_RE = None + + +def extract_cwd(body: Dict[str, Any]) -> Optional[str]: + """Pull the working directory from Codex's ... environment_context + (sent in instructions/input), so path fix-ups target the real directory.""" + import re + global _CWD_RE + if _CWD_RE is None: + _CWD_RE = re.compile(r"\s*([^<\s]+)\s*") + blob = _text_of(body.get("instructions")) + inp = body.get("input") + if isinstance(inp, str): + blob += "\n" + inp + elif isinstance(inp, list): + for it in inp: + if isinstance(it, dict): + blob += "\n" + _text_of(it.get("content", "")) + m = _CWD_RE.search(blob or "") + return m.group(1) if m else None + + +# --------------------------------------------------------------- SSE emitter +class ResponsesStreamEmitter: + """Builds the ordered Responses SSE event stream from incremental text / + tool-call events. Each method returns a list of SSE strings to yield. + + Output-item lifecycle (Codex expects these exact event types): + message: output_item.added -> content_part.added -> output_text.delta* + -> output_text.done -> content_part.done -> output_item.done + function_call: output_item.added -> function_call_arguments.delta* + -> function_call_arguments.done -> output_item.done + end: response.completed + """ + + def __init__(self, resp_id: str, model: str): + self.resp_id = resp_id + self.model = model + self._seq = itertools.count(0) + self.out_index = 0 + self.final_output: List[Dict[str, Any]] = [] + self._open: Optional[Dict[str, Any]] = None # current open output item + + def _ev(self, ev: str, extra: Dict[str, Any]) -> str: + d = {"type": ev, "sequence_number": next(self._seq)} + d.update(extra) + return f"event: {ev}\ndata: {json.dumps(d)}\n\n" + + def _base(self, status: str, output: List[Dict[str, Any]]) -> Dict[str, Any]: + return { + "id": self.resp_id, + "object": "response", + "status": status, + "model": self.model, + "output": output, + } + + def created(self) -> List[str]: + return [ + self._ev("response.created", {"response": self._base("in_progress", [])}), + self._ev( + "response.in_progress", {"response": self._base("in_progress", [])} + ), + ] + + def _close_open(self) -> List[str]: + if self._open is None: + return [] + o = self._open + self._open = None + if o["kind"] == "message": + mid, cur, txt = o["id"], o["index"], o["text"] + item = { + "id": mid, "type": "message", "role": "assistant", + "status": "completed", + "content": [{"type": "output_text", "text": txt}], + } + self.final_output.append(item) + return [ + self._ev("response.output_text.done", { + "item_id": mid, "output_index": cur, + "content_index": 0, "text": txt}), + self._ev("response.content_part.done", { + "item_id": mid, "output_index": cur, "content_index": 0, + "part": {"type": "output_text", "text": txt}}), + self._ev("response.output_item.done", { + "output_index": cur, "item": item}), + ] + # function_call + fid, cur = o["id"], o["index"] + item = { + "id": fid, "type": "function_call", "status": "completed", + "call_id": o["call_id"], "name": o["name"], "arguments": o["args"], + } + self.final_output.append(item) + return [ + self._ev("response.function_call_arguments.done", { + "item_id": fid, "output_index": cur, "arguments": o["args"]}), + self._ev("response.output_item.done", { + "output_index": cur, "item": item}), + ] + + def text_delta(self, delta: str) -> List[str]: + out: List[str] = [] + if self._open and self._open["kind"] != "message": + out += self._close_open() + if not self._open: + mid, cur = _rid("msg"), self.out_index + self.out_index += 1 + self._open = {"kind": "message", "id": mid, "index": cur, "text": ""} + out.append(self._ev("response.output_item.added", { + "output_index": cur, + "item": {"id": mid, "type": "message", "role": "assistant", + "status": "in_progress", "content": []}})) + out.append(self._ev("response.content_part.added", { + "item_id": mid, "output_index": cur, "content_index": 0, + "part": {"type": "output_text", "text": ""}})) + self._open["text"] += delta + out.append(self._ev("response.output_text.delta", { + "item_id": self._open["id"], "output_index": self._open["index"], + "content_index": 0, "delta": delta})) + return out + + def tool_start(self, call_id: str, name: str) -> List[str]: + out = self._close_open() + fid, cur = _rid("fc"), self.out_index + self.out_index += 1 + self._open = { + "kind": "fc", "id": fid, "index": cur, + "call_id": call_id or _rid("call"), "name": name, "args": "", + } + out.append(self._ev("response.output_item.added", { + "output_index": cur, + "item": {"id": fid, "type": "function_call", "status": "in_progress", + "call_id": self._open["call_id"], "name": name, + "arguments": ""}})) + return out + + def tool_args(self, delta: str) -> List[str]: + if not self._open or self._open["kind"] != "fc" or not delta: + return [] + self._open["args"] += delta + return [self._ev("response.function_call_arguments.delta", { + "item_id": self._open["id"], "output_index": self._open["index"], + "delta": delta})] + + def tool_end(self) -> List[str]: + return self._close_open() + + def finish(self, input_tokens: int, output_tokens: int) -> List[str]: + out = self._close_open() + completed = self._base("completed", self.final_output) + completed["usage"] = { + "input_tokens": input_tokens, + "output_tokens": output_tokens, + "total_tokens": input_tokens + output_tokens, + } + out.append(self._ev("response.completed", {"response": completed})) + return out + + +# ------------------------------------------------------- non-stream response +def build_responses_object( + resp_id: str, + model: str, + content_text: str, + tool_calls: List[Any], + input_tokens: int, + output_tokens: int, + valid: set, + shell_tool: Optional[str], + cwd: Optional[str] = None, +) -> Dict[str, Any]: + """Build a full (non-streaming) Responses object from parsed output. + + ``tool_calls`` are ATOM ``ToolCall`` objects (``.id``, ``.function`` dict).""" + output: List[Dict[str, Any]] = [] + if content_text: + output.append({ + "id": _rid("msg"), "type": "message", "role": "assistant", + "status": "completed", + "content": [{"type": "output_text", "text": content_text}], + }) + for tc in tool_calls or []: + fn = getattr(tc, "function", None) or {} + name, args = translate_client_tool( + fn.get("name", ""), fn.get("arguments", "") or "", valid, shell_tool, cwd + ) + output.append({ + "id": _rid("fc"), "type": "function_call", "status": "completed", + "call_id": getattr(tc, "id", None) or _rid("call"), + "name": name, "arguments": args, + }) + return { + "id": resp_id, "object": "response", "status": "completed", + "model": model, "output": output, + "usage": { + "input_tokens": input_tokens, + "output_tokens": output_tokens, + "total_tokens": input_tokens + output_tokens, + }, + } diff --git a/atom/entrypoints/openai/tool_parser.py b/atom/entrypoints/openai/tool_parser.py index 549277e8d9..8f0b992b0e 100644 --- a/atom/entrypoints/openai/tool_parser.py +++ b/atom/entrypoints/openai/tool_parser.py @@ -183,21 +183,382 @@ def _parse_qwen_xml(text: str, tools: Optional[list]) -> Tuple[str, List[ToolCal return content.strip(), tool_calls +# --------------------------------------------------------------------------- +# DeepSeek-V4 DSML tool-call format +# --------------------------------------------------------------------------- +# +# <|DSML|tool_calls> +# <|DSML|invoke name="NAME"> +# <|DSML|parameter name="PNAME" string="true|false">VALUE +# ... +# +# +# +# string="true" -> value is a raw string; string="false" -> value is JSON. +# DeepSeek-V4-Flash occasionally malforms this (singular ``tool_call``, a missing +# ``invoke`` wrapper, or params without ``string=``); the parser recovers those +# best-effort: it infers a dropped tool name from the parameter signature vs the +# request's ``tools`` and infers a missing value type from the schema / JSON. + +_DSML = "|DSML|" +# The model often DROPS the ``|DSML|`` marker and emits bare +# ````/````/```` tags, so the marker +# is matched OPTIONALLY everywhere. +_OPT = r"(?:" + re.escape(_DSML) + r")?" # optional |DSML| prefix +_DSML_PARAM_RE = re.compile( + r"<" + _OPT + r'parameter\s+name="(.*?)"(?:\s+string="(true|false)")?\s*>' + r"(.*?)", + re.DOTALL, +) +# Long-form `...` OR self-closing `` +# (the zero-arg shape; group(2) is None for self-closing). Matches SGLang's V4 +# detector, which accepts both. +_DSML_INVOKE_RE = re.compile( + r"<" + _OPT + r'invoke\s+name="(.*?)"\s*(?:/>|>(.*?))", + re.DOTALL, +) +# Region-start markers, both marked and marker-less variants. +_DSML_STARTS = ( + "<" + _DSML + "tool_call", # marked (covers tool_call / tool_calls) + "<" + _DSML + "invoke", # marked invoke + "", # marker-less section open +) + + +def _dsml_start(text: str) -> int: + """Index of the earliest DSML tool-call marker (marked or marker-less), or -1.""" + positions = [i for i in (text.find(m) for m in _DSML_STARTS) if i != -1] + return min(positions) if positions else -1 + + +def _is_dsml(text: str) -> bool: + return _dsml_start(text) != -1 + + +def _unwrap_wrapper_args(args: Any, allowed: set) -> Any: + """Strip spurious ``{"arguments": {...}}`` / ``{"input": {...}}`` envelopes. + + Non-tuned models (DeepSeek-V4-Pro) frequently wrap the real args in an extra + ``arguments``/``input`` object — sometimes nested 2-3 deep, or stringified — + so a call meant as ``{"cmd": "ls"}`` arrives as ``{"arguments": {"cmd": + "ls"}}`` and the client (Codex) rejects it ("missing field cmd"). Recursively + unwrap while the sole key is a wrapper that is NOT itself a declared param of + the tool. Mirrors vLLM's ``_unwrap_wrapper_args`` (deepseek_v4.py).""" + for _ in range(4): # bounded against pathological nesting + if not (isinstance(args, dict) and len(args) == 1): + break + (k, v), = args.items() + if k not in ("arguments", "input"): + break + if allowed and k in allowed: + break # this tool really has a param named arguments/input + if isinstance(v, str): + try: + v = json.loads(v) + except Exception: + break + if not isinstance(v, dict): + break + args = v + return args + + +# Canonical param key -> emitted synonyms the model uses interchangeably. Only +# applied when the canonical key is in the tool's declared schema and the model +# used a synonym instead (e.g. Codex's exec_command wants `cmd`, but the model +# habitually emits `command` -> "missing field cmd" retry loop). Scoped to the +# schema so it never renames a key a tool legitimately declares. +_KEY_ALIASES: Dict[str, Tuple[str, ...]] = { + "cmd": ("command",), +} + + +def _apply_key_aliases(args: Any, allowed: set) -> Any: + if not (isinstance(args, dict) and allowed): + return args + for canon, syns in _KEY_ALIASES.items(): + if canon in allowed and canon not in args: + for s in syns: + if s in args: + args[canon] = args.pop(s) + break + return args + + +def _dsml_coerce(value: str, string_attr: Optional[str], ptype: Any) -> Any: + if string_attr == "true": + return value + if string_attr == "false": + try: + return json.loads(value) + except Exception: + return value + # attr absent -> use declared schema type if known, else infer via JSON. + if ptype is not None: + return _coerce_param_value(value, ptype) + v = value.strip() + try: + return json.loads(v) + except Exception: + return v + + +def _infer_dsml_name(arg_names: set, param_types: Dict[str, Dict[str, Any]]) -> Optional[str]: + """Pick the request tool whose parameter set best matches ``arg_names``.""" + best, best_score = None, -1e9 + for name, props in param_types.items(): + p = set(props) + if not p: + continue + score = len(p & arg_names) - 0.1 * len(p ^ arg_names) + if score > best_score: + best_score, best = score, name + return best + + +def _parse_dsml(text: str, tools: Optional[list]) -> Tuple[str, List[ToolCall]]: + """Parse DeepSeek-V4 DSML tool calls; return (leading_content, tool_calls).""" + param_types = _build_param_types(tools) + start = _dsml_start(text) + if start == -1: + return text.strip(), [] + content = text[:start] + region = text[start:] + + calls: List[Tuple[str, Dict[str, Any]]] = [] + invokes = list(_DSML_INVOKE_RE.finditer(region)) + if invokes: + for m in invokes: + name = m.group(1) + body = m.group(2) or "" # None for self-closing + types = param_types.get(name, {}) + args: Dict[str, Any] = { + pm.group(1): _dsml_coerce(pm.group(3), pm.group(2), types.get(pm.group(1))) + for pm in _DSML_PARAM_RE.finditer(body) + } + # Direct-JSON parameter body (DSML "Format 2", also accepted by + # vLLM/SGLang): ` { "k": "v" } ` with no + # tags. Falls through here with empty args; recover them. + if not args: + stripped = body.strip() + if stripped.startswith("{"): + try: + parsed = json.loads(stripped) + if isinstance(parsed, dict): + args = parsed + except Exception: + pass + args = _unwrap_wrapper_args(args, set(types)) + args = _apply_key_aliases(args, set(types)) + calls.append((name, args)) + else: + # malformed: no complete invoke wrapper -> collect params, infer tool name + raw = {pm.group(1): (pm.group(3), pm.group(2)) for pm in _DSML_PARAM_RE.finditer(region)} + if raw: + name = _infer_dsml_name(set(raw), param_types) or "unknown" + types = param_types.get(name, {}) + args = {k: _dsml_coerce(v, s, types.get(k)) for k, (v, s) in raw.items()} + args = _unwrap_wrapper_args(args, set(types)) + args = _apply_key_aliases(args, set(types)) + calls.append((name, args)) + + tool_calls = [ + ToolCall( + id=_unique_tool_call_id(), + type="function", + function={"name": name, "arguments": json.dumps(args, ensure_ascii=False)}, + ) + for name, args in calls + ] + if _DSML in content: # scrub any stray marker fragment + content = content.split("<" + _DSML, 1)[0] + return content.strip(), tool_calls + + +# --------------------------------------------------------------------------- +# GLM-4.5 / 4.6 / 5.x tool-call format +# --------------------------------------------------------------------------- +# +# NAME +# K1V1 +# K2V2 +# ... +# +# The function name follows the opening tag directly (no ``(.*?)|(.*)$", re.DOTALL +) +_GLM_ARG_RE = re.compile( + r"(.*?)\s*" + r"(.*?)(?:|(?=)|(?=)|$)", + re.DOTALL, +) + + +def _is_glm(text: str) -> bool: + """Detect the GLM ``...`` format (never Qwen/DSML).""" + if _QWEN_TOOL_PREFIX in text: # ' Qwen, not GLM + return False + return "" in text or "" in text + + +def _glm_coerce(value: str, ptype: Any) -> Any: + """Decode one GLM ````: schema type wins, else JSON, else raw.""" + v = value.strip("\n") + if ptype is not None: + return _coerce_param_value(v, ptype) + s = v.strip() + try: + return json.loads(s) + except Exception: + return v + + +def _parse_glm(text: str, tools: Optional[list] = None) -> Tuple[str, List[ToolCall]]: + """Parse GLM tool calls; return (leading_content, tool_calls).""" + param_types = _build_param_types(tools) + start = text.find("") + if start == -1: + return text.strip(), [] + content = text[:start] + tool_calls: List[ToolCall] = [] + for m in _GLM_TOOLCALL_RE.finditer(text): + body = m.group(1) if m.group(1) is not None else m.group(2) + if not body: + continue + ak = body.find("") + name = (body if ak == -1 else body[:ak]).strip() + if not name: + continue + types = param_types.get(name, {}) + args: Dict[str, Any] = {} + for pm in _GLM_ARG_RE.finditer(body): + k = pm.group(1).strip() + if k: + args[k] = _glm_coerce(pm.group(2), types.get(k)) + tool_calls.append( + ToolCall( + id=_unique_tool_call_id(), + type="function", + function={ + "name": name, + "arguments": json.dumps(args, ensure_ascii=False), + }, + ) + ) + return content.strip(), tool_calls + + +# --------------------------------------------------------------------------- +# MiniMax-M3 tool-call format +# --------------------------------------------------------------------------- +# +# Every tag is prefixed by the ns_token ``]<]minimax[>[``: +# +# ]<]minimax[>[ +# ]<]minimax[>[ +# ]<]minimax[>[value]<]minimax[>[ +# ... +# ]<]minimax[>[ +# ]<]minimax[>[ +# +# Unlike DSML, parameters are named by the TAG itself (``Paris``), +# not a ``name="..."`` attribute. Strip the ns_token first, then parse +# / pairs. Values: schema type wins, else JSON, else raw string. + +_MINIMAX_NS = "]<]minimax[>[" +_MINIMAX_INVOKE_RE = re.compile( + r'(.*?)|(.*)$', + re.DOTALL, +) +_MINIMAX_PARAM_RE = re.compile(r"<([\w-]+)>(.*?)", re.DOTALL) + + +def _is_minimax(text: str) -> bool: + """Detect the MiniMax-M3 ns_token tool-call format.""" + return _MINIMAX_NS in text + + +def _minimax_coerce(value: str, ptype: Any) -> Any: + v = value.strip("\n") + if ptype is not None: + return _coerce_param_value(v, ptype) + s = v.strip() + try: + return json.loads(s) + except Exception: + return v + + +def _parse_minimax(text: str, tools: Optional[list] = None) -> Tuple[str, List[ToolCall]]: + """Parse MiniMax-M3 tool calls; return (leading_content, tool_calls).""" + param_types = _build_param_types(tools) + clean = text.replace(_MINIMAX_NS, "") + tc = clean.find("") + content = clean[:tc] if tc > 0 else ("" if tc == 0 else clean) + tool_calls: List[ToolCall] = [] + for m in _MINIMAX_INVOKE_RE.finditer(clean): + name = m.group(1) if m.group(1) is not None else m.group(3) + body = m.group(2) if m.group(2) is not None else (m.group(4) or "") + if not name: + continue + name = name.strip() + types = param_types.get(name, {}) + args: Dict[str, Any] = {} + for pm in _MINIMAX_PARAM_RE.finditer(body): + k = pm.group(1).strip() + if k: + args[k] = _minimax_coerce(pm.group(2), types.get(k)) + tool_calls.append( + ToolCall( + id=_unique_tool_call_id(), + type="function", + function={ + "name": name, + "arguments": json.dumps(args, ensure_ascii=False), + }, + ) + ) + for mk in ("", ""): + content = content.replace(mk, "") + return content.strip(), tool_calls + + def parse_tool_calls( text: str, tools: Optional[list] = None ) -> Tuple[str, List[ToolCall]]: """Parse tool calls from model output text. Args: - text: Raw model output that may contain tool calls (Kimi token format - or Qwen3 XML format). - tools: Optional request tool definitions; used to type-coerce Qwen XML - parameter values to their declared JSON-Schema types. + text: Raw model output that may contain tool calls (DeepSeek-V4 DSML, + Kimi token format, or Qwen3 XML format). + tools: Optional request tool definitions; used to type-coerce parameter + values to their declared JSON-Schema types. Returns: Tuple of (content_text, list_of_tool_calls). ``content_text`` has the tool-call sections removed. """ + # MiniMax-M3 ns_token format (checked before DSML: both use , + # but MiniMax names params by tag and prefixes every tag with ]<]minimax[>[) + if _is_minimax(text): + return _parse_minimax(text, tools) + + # DeepSeek-V4 DSML format + if _is_dsml(text): + return _parse_dsml(text, tools) + + # GLM / format (checked before Qwen: both use + # , but GLM never emits the Qwen ' list: """Process a text chunk and return list of (event_type, data) tuples.""" if self.fmt is None: self.buf += text - if _QWEN_TOOL_PREFIX in self.buf or "" in self.buf: + if _MINIMAX_NS in self.buf: + self.fmt = "minimax" + elif _is_dsml(self.buf): + self.fmt = "dsml" + elif "" in self.buf: + self.fmt = "glm" + elif _QWEN_TOOL_PREFIX in self.buf: self.fmt = "qwen" + elif "" in self.buf: + # '' seen but neither '' (GLM) yet. A no-arg GLM call is complete once the + # closing tag arrives; otherwise wait for the sub-marker. + if "" in self.buf: + self.fmt = "glm" + else: + return [] elif "<|tool_calls_section_begin|>" in self.buf: self.fmt = "kimi" elif "<" not in self.buf and len(self.buf) > 8: @@ -297,10 +672,202 @@ def process(self, text: str) -> list: # Format decided: replay the accumulated buffer through the handler. text, self.buf = self.buf, "" + if self.fmt == "minimax": + return self._process_minimax(text) + if self.fmt == "dsml": + return self._process_dsml(text) + if self.fmt == "glm": + return self._process_glm(text) if self.fmt == "qwen": return self._process_qwen(text) return self._process_kimi(text) + # -- MiniMax-M3 ns_token ------------------------------------------------ + def _process_minimax(self, text: str) -> list: + results: list = [] + self.buf += text + if self.state == 0: + markers = [ + i + for i in (self.buf.find(_MINIMAX_NS), self.buf.find("")) + if i != -1 + ] + if markers: + m = min(markers) + before = self.buf[:m] + if before: + results.append(("content", before)) + self.buf = self.buf[m:] + self.state = 1 + else: + cut = self.buf.rfind("<") + cut = max(cut, self.buf.rfind("]")) # ns_token starts with ']' + if cut == -1: + if self.buf: + results.append(("content", self.buf)) + self.buf = "" + elif cut > 0: + results.append(("content", self.buf[:cut])) + self.buf = self.buf[cut:] + return results + + def _flush_minimax(self) -> list: + results: list = [] + if self.state == 0: + if self.buf: + results.append(("content", self.buf)) + self.buf = "" + return results + _content, tool_calls = _parse_minimax(self.buf, self.tools) + self.buf = "" + for tc in tool_calls: + results.append( + ( + "tool_call_start", + { + "index": self.current_index, + "id": tc.id, + "type": "function", + "function": {"name": tc.function["name"], "arguments": ""}, + }, + ) + ) + results.append( + ( + "tool_call_args", + { + "index": self.current_index, + "function": {"arguments": tc.function["arguments"]}, + }, + ) + ) + self.current_index += 1 + self._emitted_calls += 1 + if self._emitted_calls > 0: + results.append(("tool_call_end", None)) + return results + + # -- DeepSeek-V4 DSML --------------------------------------------------- + def _process_dsml(self, text: str) -> list: + results: list = [] + self.buf += text + if self.state == 0: + m = _dsml_start(self.buf) + if m != -1: + before = self.buf[:m] + if before: + results.append(("content", before)) + self.buf = self.buf[m:] + self.state = 1 + else: + # Emit content but hold back a possible partial '<...' marker tail. + cut = self.buf.rfind("<") + if cut == -1: + if self.buf: + results.append(("content", self.buf)) + self.buf = "" + elif cut > 0: + results.append(("content", self.buf[:cut])) + self.buf = self.buf[cut:] + return results + + def _flush_dsml(self) -> list: + results: list = [] + if self.state == 0: + if self.buf: + results.append(("content", self.buf)) + self.buf = "" + return results + _content, tool_calls = _parse_dsml(self.buf, self.tools) + self.buf = "" + for tc in tool_calls: + results.append( + ( + "tool_call_start", + { + "index": self.current_index, + "id": tc.id, + "type": "function", + "function": {"name": tc.function["name"], "arguments": ""}, + }, + ) + ) + results.append( + ( + "tool_call_args", + { + "index": self.current_index, + "function": {"arguments": tc.function["arguments"]}, + }, + ) + ) + self.current_index += 1 + self._emitted_calls += 1 + if self._emitted_calls > 0: + results.append(("tool_call_end", None)) + return results + + # -- Qwen3 XML ---------------------------------------------------------- + # -- GLM / ----------------------------------------- + def _process_glm(self, text: str) -> list: + results: list = [] + self.buf += text + if self.state == 0: + m = self.buf.find("") + if m != -1: + before = self.buf[:m] + if before: + results.append(("content", before)) + self.buf = self.buf[m:] + self.state = 1 + else: + # Emit content but hold back a possible partial '<...' marker tail. + cut = self.buf.rfind("<") + if cut == -1: + if self.buf: + results.append(("content", self.buf)) + self.buf = "" + elif cut > 0: + results.append(("content", self.buf[:cut])) + self.buf = self.buf[cut:] + return results + + def _flush_glm(self) -> list: + results: list = [] + if self.state == 0: + if self.buf: + results.append(("content", self.buf)) + self.buf = "" + return results + _content, tool_calls = _parse_glm(self.buf, self.tools) + self.buf = "" + for tc in tool_calls: + results.append( + ( + "tool_call_start", + { + "index": self.current_index, + "id": tc.id, + "type": "function", + "function": {"name": tc.function["name"], "arguments": ""}, + }, + ) + ) + results.append( + ( + "tool_call_args", + { + "index": self.current_index, + "function": {"arguments": tc.function["arguments"]}, + }, + ) + ) + self.current_index += 1 + self._emitted_calls += 1 + if self._emitted_calls > 0: + results.append(("tool_call_end", None)) + return results + # -- Qwen3 XML ---------------------------------------------------------- def _process_qwen(self, text: str) -> list: results: list = [] @@ -449,6 +1016,12 @@ def _process_buffer(self) -> list: def flush(self) -> list: """Flush remaining buffer content.""" + if self.fmt == "minimax": + return self._flush_minimax() + if self.fmt == "dsml": + return self._flush_dsml() + if self.fmt == "glm": + return self._flush_glm() if self.fmt == "qwen": return self._flush_qwen() results = [] From 7139c4300fe15ccd48c0c200bcbdb225af8fc4e1 Mon Sep 17 00:00:00 2001 From: yihonglie Date: Sun, 12 Jul 2026 22:24:17 -0500 Subject: [PATCH 02/18] style: apply black formatting Fixes "Check Code Style with Black" CI on the DSV4 API + DSML tool parser changes (serving_responses.py, tool_parser.py, api_server.py). Co-Authored-By: Claude Opus 4.8 (1M context) Signed-off-by: yihonglie --- atom/entrypoints/openai/api_server.py | 33 ++- atom/entrypoints/openai/chat_encoders.py | 10 +- atom/entrypoints/openai/serving_responses.py | 296 ++++++++++++++----- atom/entrypoints/openai/tool_parser.py | 29 +- 4 files changed, 270 insertions(+), 98 deletions(-) diff --git a/atom/entrypoints/openai/api_server.py b/atom/entrypoints/openai/api_server.py index 12640649bb..8b37ffe7dc 100644 --- a/atom/entrypoints/openai/api_server.py +++ b/atom/entrypoints/openai/api_server.py @@ -408,27 +408,29 @@ def flush_stream_batch() -> None: # Incremental per-request detokenization: correct UTF-8 at token # boundaries (see _stream_detok_state). Emits only fully-formed chars; # a trailing partial multi-byte char is held until the next step. - for (_loop, _q, chunk) in buf: + for _loop, _q, chunk in buf: rid = chunk.get("request_id") st = _stream_detok_state.get(rid) if st is None: st = _stream_detok_state[rid] = { - "tokens": [], "prefix_offset": 0, "read_offset": 0, + "tokens": [], + "prefix_offset": 0, + "read_offset": 0, } toks = st["tokens"] toks.extend(chunk["token_ids"]) prefix_text = tokenizer.decode( - toks[st["prefix_offset"]:st["read_offset"]], skip_special_tokens=True + toks[st["prefix_offset"] : st["read_offset"]], skip_special_tokens=True ) new_text = tokenizer.decode( - toks[st["prefix_offset"]:], skip_special_tokens=True + toks[st["prefix_offset"] :], skip_special_tokens=True ) if len(new_text) > len(prefix_text) and not new_text.endswith("\ufffd"): - chunk["text"] = new_text[len(prefix_text):] + chunk["text"] = new_text[len(prefix_text) :] st["prefix_offset"] = st["read_offset"] st["read_offset"] = len(toks) elif chunk["finished"]: - chunk["text"] = new_text[len(prefix_text):] + chunk["text"] = new_text[len(prefix_text) :] else: chunk["text"] = "" if chunk["finished"]: @@ -1706,7 +1708,9 @@ async def generate_anthropic_stream(): raw_text = final_output["text"] reasoning_content, content_with_tools = separate_reasoning(raw_text) - content_text, tool_calls = parse_tool_calls(content_with_tools, anthropic_to_openai_tools(request.tools)) + content_text, tool_calls = parse_tool_calls( + content_with_tools, anthropic_to_openai_tools(request.tools) + ) output_tokens = len(tokenizer.encode(raw_text)) cache_read_input_tokens = final_output.get("num_cached_tokens", 0) if not getattr(request, "thinking", None): @@ -1780,7 +1784,9 @@ async def responses_endpoint(raw_request: Request): max_out = int(body.get("max_output_tokens") or 32768) sampling_params = _build_sampling_params( - temperature=body.get("temperature") if body.get("temperature") is not None else 1.0, + temperature=( + body.get("temperature") if body.get("temperature") is not None else 1.0 + ), max_tokens=max_out, stop_strings=None, ignore_eos=False, @@ -1844,7 +1850,12 @@ def _flush_pending(): return [] _pending["tc"] = None name, args = translate_client_tool( - tc["name"], tc["args"], valid_names, shell_tool, req_cwd, shell_param + tc["name"], + tc["args"], + valid_names, + shell_tool, + req_cwd, + shell_param, ) out = emitter.tool_start(tc["id"], name) if args: @@ -1926,7 +1937,9 @@ def handle(etype, edata): raw_text = final_output["text"] _reasoning, content_with_tools = separate_reasoning(raw_text) - content_text, tool_calls = parse_tool_calls(content_with_tools, openai_tools or None) + content_text, tool_calls = parse_tool_calls( + content_with_tools, openai_tools or None + ) output_tokens = len(tokenizer.encode(raw_text)) return JSONResponse( diff --git a/atom/entrypoints/openai/chat_encoders.py b/atom/entrypoints/openai/chat_encoders.py index 863d045275..9021ab85d3 100644 --- a/atom/entrypoints/openai/chat_encoders.py +++ b/atom/entrypoints/openai/chat_encoders.py @@ -88,7 +88,9 @@ def load_custom_message_encoder(model_path: str) -> Optional[MessageEncoder]: def _content_str(c: Any) -> str: if isinstance(c, list): return "\n".join( - b.get("text", "") for b in c if isinstance(b, dict) and b.get("type") == "text" + b.get("text", "") + for b in c + if isinstance(b, dict) and b.get("type") == "text" ) return c or "" @@ -115,9 +117,11 @@ def _normalize_for_v4(messages: List[dict], tools: Optional[List[dict]]) -> List if not sys_parts and not tools: return [dict(m) for m in messages] - merged = "\n\n".join(s for s in (_content_str(m.get("content")) for m in sys_parts) if s) + merged = "\n\n".join( + s for s in (_content_str(m.get("content")) for m in sys_parts) if s + ) sys_msg: dict = {"role": "system", "content": merged} - for m in sys_parts: # preserve any pre-attached tools + for m in sys_parts: # preserve any pre-attached tools if m.get("tools"): sys_msg["tools"] = m["tools"] if tools: diff --git a/atom/entrypoints/openai/serving_responses.py b/atom/entrypoints/openai/serving_responses.py index bc14091ad0..26d22af821 100644 --- a/atom/entrypoints/openai/serving_responses.py +++ b/atom/entrypoints/openai/serving_responses.py @@ -11,6 +11,7 @@ This makes the external ``codex_responses_proxy.py`` unnecessary: Codex can point straight at ATOM's ``:9700/v1``. """ + import itertools import json import time @@ -42,9 +43,7 @@ def _text_of(content: Any) -> str: return "" -def responses_input_to_messages( - instructions: Any, inp: Any -) -> List[Dict[str, Any]]: +def responses_input_to_messages(instructions: Any, inp: Any) -> List[Dict[str, Any]]: """Translate Responses ``instructions`` + ``input`` into OpenAI chat messages. ``input`` may be a plain string or a list of items: ``message`` / @@ -92,9 +91,7 @@ def responses_input_to_messages( { "role": "tool", "tool_call_id": item.get("call_id") or item.get("id") or "", - "content": out - if isinstance(out, str) - else json.dumps(out), + "content": out if isinstance(out, str) else json.dumps(out), } ) elif t == "reasoning": @@ -109,8 +106,8 @@ def responses_input_to_messages( " tags, no prose. Use EXACTLY this syntax (the \uff5c characters " "are U+FF5C fullwidth vertical bars, not ASCII '|'):\n" "<\uff5cDSML\uff5ctool_calls>\n" - "<\uff5cDSML\uff5cinvoke name=\"TOOL_NAME\">\n" - "<\uff5cDSML\uff5cparameter name=\"PARAM_NAME\" string=\"true\">VALUE" + '<\uff5cDSML\uff5cinvoke name="TOOL_NAME">\n' + '<\uff5cDSML\uff5cparameter name="PARAM_NAME" string="true">VALUE' "\n" "\n" "\n" @@ -130,7 +127,9 @@ def inject_tool_format_instruction(messages): if "\uff5cDSML\uff5ctool_calls" not in base: m["content"] = _text_of(base) + _DSML_TOOL_INSTRUCTION return messages - return [{"role": "system", "content": _DSML_TOOL_INSTRUCTION.strip()}] + list(messages) + return [{"role": "system", "content": _DSML_TOOL_INSTRUCTION.strip()}] + list( + messages + ) def responses_tools_to_openai(tools: Any) -> List[Dict[str, Any]]: @@ -172,17 +171,40 @@ def responses_tools_to_openai(tools: Any) -> List[Dict[str, Any]]: # is wrong. So remap known shell-exec aliases onto whatever shell tool the # client actually registered this turn. Pure rename; arguments untouched. _SHELL_ALIASES = { - "exec", "exec_run", "exec_command", "execute_command", "shell", "bash", - "sh", "run", "run_command", "run_shell", "execute", "command", - "container.exec", "local_shell", "shell_command", "shell_exec", - "run_bash", "execute_shell", "bash_command", "run_terminal_cmd", - "terminal", "console", "run_shell_command", "runshell", + "exec", + "exec_run", + "exec_command", + "execute_command", + "shell", + "bash", + "sh", + "run", + "run_command", + "run_shell", + "execute", + "command", + "container.exec", + "local_shell", + "shell_command", + "shell_exec", + "run_bash", + "execute_shell", + "bash_command", + "run_terminal_cmd", + "terminal", + "console", + "run_shell_command", + "runshell", } # Substrings that mark an unknown tool name as a shell/exec call (Codex's # non-shell tools contain none of these). _SHELL_NAME_TOKENS = ("shell", "exec", "bash", "cmd", "command", "termin", "console") _SHELL_TOOL_PREFERENCE = ( - "exec_command", "shell", "local_shell", "bash", "container.exec", + "exec_command", + "shell", + "local_shell", + "bash", + "container.exec", ) @@ -226,6 +248,7 @@ def remap_tool_name(name: str, valid: set, shell_tool: Optional[str]) -> str: # to /v1/messages, where read/grep ARE native tools). def _q(s: Any) -> str: import shlex + return shlex.quote(str(s)) @@ -261,13 +284,32 @@ def _claude_tool_to_shell( name: str, a: Dict[str, Any], cwd: Optional[str] = None ) -> Optional[str]: n = (name or "").lower() - fp = (a.get("file_path") or a.get("path") or a.get("filePath") - or a.get("filename") or a.get("target_file") or a.get("file")) + fp = ( + a.get("file_path") + or a.get("path") + or a.get("filePath") + or a.get("filename") + or a.get("target_file") + or a.get("file") + ) pattern = a.get("pattern") or a.get("query") or a.get("regex") - path = (a.get("path") or a.get("directory") or a.get("target_directory") - or a.get("dir") or ".") - if n in ("read", "cat", "view", "view_file", "open", "read_file", - "readfile", "openfile"): + path = ( + a.get("path") + or a.get("directory") + or a.get("target_directory") + or a.get("dir") + or "." + ) + if n in ( + "read", + "cat", + "view", + "view_file", + "open", + "read_file", + "readfile", + "openfile", + ): if not fp: return None off, lim = a.get("offset"), a.get("limit") @@ -277,8 +319,15 @@ def _claude_tool_to_shell( else: start, end = 1, 400 return _read_cmd(str(fp), cwd, start, end) - if n in ("grep", "search", "search_file", "ripgrep", "rg", "grep_search", - "codebase_search"): + if n in ( + "grep", + "search", + "search_file", + "ripgrep", + "rg", + "grep_search", + "codebase_search", + ): if not pattern: return None return f"grep -rn -- {_q(pattern)} {_q(_resolve_dir(path, cwd))}" @@ -293,8 +342,19 @@ def _claude_tool_to_shell( _SHELL_ARG_ALIASES = ( - "cmd", "command", "commandline", "command_line", "script", "bash", "sh", - "shell", "shell_command", "code", "input", "run", "cmd_string", + "cmd", + "command", + "commandline", + "command_line", + "script", + "bash", + "sh", + "shell", + "shell_command", + "code", + "input", + "run", + "cmd_string", ) @@ -309,7 +369,7 @@ def shell_arg_key(openai_tools, shell_tool): continue params = fn.get("parameters") or {} props = params.get("properties") or {} - for r in (params.get("required") or []): + for r in params.get("required") or []: if props.get(r, {}).get("type") in (None, "string"): return r if props: @@ -341,8 +401,12 @@ def _normalize_shell_args(args_json, req): def translate_client_tool( - name: str, args_json: str, valid: set, shell_tool: Optional[str], - cwd: Optional[str] = None, shell_param: Optional[str] = None, + name: str, + args_json: str, + valid: set, + shell_tool: Optional[str], + cwd: Optional[str] = None, + shell_param: Optional[str] = None, ): """Return (name, args_json) with Claude read-only tools rewritten to the registered shell tool. exec_command's own calls and any already-valid tool @@ -375,6 +439,7 @@ def extract_cwd(body: Dict[str, Any]) -> Optional[str]: """Pull the working directory from Codex's ... environment_context (sent in instructions/input), so path fix-ups target the real directory.""" import re + global _CWD_RE if _CWD_RE is None: _CWD_RE = re.compile(r"\s*([^<\s]+)\s*") @@ -441,33 +506,53 @@ def _close_open(self) -> List[str]: if o["kind"] == "message": mid, cur, txt = o["id"], o["index"], o["text"] item = { - "id": mid, "type": "message", "role": "assistant", + "id": mid, + "type": "message", + "role": "assistant", "status": "completed", "content": [{"type": "output_text", "text": txt}], } self.final_output.append(item) return [ - self._ev("response.output_text.done", { - "item_id": mid, "output_index": cur, - "content_index": 0, "text": txt}), - self._ev("response.content_part.done", { - "item_id": mid, "output_index": cur, "content_index": 0, - "part": {"type": "output_text", "text": txt}}), - self._ev("response.output_item.done", { - "output_index": cur, "item": item}), + self._ev( + "response.output_text.done", + { + "item_id": mid, + "output_index": cur, + "content_index": 0, + "text": txt, + }, + ), + self._ev( + "response.content_part.done", + { + "item_id": mid, + "output_index": cur, + "content_index": 0, + "part": {"type": "output_text", "text": txt}, + }, + ), + self._ev( + "response.output_item.done", {"output_index": cur, "item": item} + ), ] # function_call fid, cur = o["id"], o["index"] item = { - "id": fid, "type": "function_call", "status": "completed", - "call_id": o["call_id"], "name": o["name"], "arguments": o["args"], + "id": fid, + "type": "function_call", + "status": "completed", + "call_id": o["call_id"], + "name": o["name"], + "arguments": o["args"], } self.final_output.append(item) return [ - self._ev("response.function_call_arguments.done", { - "item_id": fid, "output_index": cur, "arguments": o["args"]}), - self._ev("response.output_item.done", { - "output_index": cur, "item": item}), + self._ev( + "response.function_call_arguments.done", + {"item_id": fid, "output_index": cur, "arguments": o["args"]}, + ), + self._ev("response.output_item.done", {"output_index": cur, "item": item}), ] def text_delta(self, delta: str) -> List[str]: @@ -478,17 +563,44 @@ def text_delta(self, delta: str) -> List[str]: mid, cur = _rid("msg"), self.out_index self.out_index += 1 self._open = {"kind": "message", "id": mid, "index": cur, "text": ""} - out.append(self._ev("response.output_item.added", { - "output_index": cur, - "item": {"id": mid, "type": "message", "role": "assistant", - "status": "in_progress", "content": []}})) - out.append(self._ev("response.content_part.added", { - "item_id": mid, "output_index": cur, "content_index": 0, - "part": {"type": "output_text", "text": ""}})) + out.append( + self._ev( + "response.output_item.added", + { + "output_index": cur, + "item": { + "id": mid, + "type": "message", + "role": "assistant", + "status": "in_progress", + "content": [], + }, + }, + ) + ) + out.append( + self._ev( + "response.content_part.added", + { + "item_id": mid, + "output_index": cur, + "content_index": 0, + "part": {"type": "output_text", "text": ""}, + }, + ) + ) self._open["text"] += delta - out.append(self._ev("response.output_text.delta", { - "item_id": self._open["id"], "output_index": self._open["index"], - "content_index": 0, "delta": delta})) + out.append( + self._ev( + "response.output_text.delta", + { + "item_id": self._open["id"], + "output_index": self._open["index"], + "content_index": 0, + "delta": delta, + }, + ) + ) return out def tool_start(self, call_id: str, name: str) -> List[str]: @@ -496,23 +608,45 @@ def tool_start(self, call_id: str, name: str) -> List[str]: fid, cur = _rid("fc"), self.out_index self.out_index += 1 self._open = { - "kind": "fc", "id": fid, "index": cur, - "call_id": call_id or _rid("call"), "name": name, "args": "", + "kind": "fc", + "id": fid, + "index": cur, + "call_id": call_id or _rid("call"), + "name": name, + "args": "", } - out.append(self._ev("response.output_item.added", { - "output_index": cur, - "item": {"id": fid, "type": "function_call", "status": "in_progress", - "call_id": self._open["call_id"], "name": name, - "arguments": ""}})) + out.append( + self._ev( + "response.output_item.added", + { + "output_index": cur, + "item": { + "id": fid, + "type": "function_call", + "status": "in_progress", + "call_id": self._open["call_id"], + "name": name, + "arguments": "", + }, + }, + ) + ) return out def tool_args(self, delta: str) -> List[str]: if not self._open or self._open["kind"] != "fc" or not delta: return [] self._open["args"] += delta - return [self._ev("response.function_call_arguments.delta", { - "item_id": self._open["id"], "output_index": self._open["index"], - "delta": delta})] + return [ + self._ev( + "response.function_call_arguments.delta", + { + "item_id": self._open["id"], + "output_index": self._open["index"], + "delta": delta, + }, + ) + ] def tool_end(self) -> List[str]: return self._close_open() @@ -546,24 +680,36 @@ def build_responses_object( ``tool_calls`` are ATOM ``ToolCall`` objects (``.id``, ``.function`` dict).""" output: List[Dict[str, Any]] = [] if content_text: - output.append({ - "id": _rid("msg"), "type": "message", "role": "assistant", - "status": "completed", - "content": [{"type": "output_text", "text": content_text}], - }) + output.append( + { + "id": _rid("msg"), + "type": "message", + "role": "assistant", + "status": "completed", + "content": [{"type": "output_text", "text": content_text}], + } + ) for tc in tool_calls or []: fn = getattr(tc, "function", None) or {} name, args = translate_client_tool( fn.get("name", ""), fn.get("arguments", "") or "", valid, shell_tool, cwd ) - output.append({ - "id": _rid("fc"), "type": "function_call", "status": "completed", - "call_id": getattr(tc, "id", None) or _rid("call"), - "name": name, "arguments": args, - }) + output.append( + { + "id": _rid("fc"), + "type": "function_call", + "status": "completed", + "call_id": getattr(tc, "id", None) or _rid("call"), + "name": name, + "arguments": args, + } + ) return { - "id": resp_id, "object": "response", "status": "completed", - "model": model, "output": output, + "id": resp_id, + "object": "response", + "status": "completed", + "model": model, + "output": output, "usage": { "input_tokens": input_tokens, "output_tokens": output_tokens, diff --git a/atom/entrypoints/openai/tool_parser.py b/atom/entrypoints/openai/tool_parser.py index 8f0b992b0e..21a98f5df4 100644 --- a/atom/entrypoints/openai/tool_parser.py +++ b/atom/entrypoints/openai/tool_parser.py @@ -204,7 +204,7 @@ def _parse_qwen_xml(text: str, tools: Optional[list]) -> Tuple[str, List[ToolCal # The model often DROPS the ``|DSML|`` marker and emits bare # ````/````/```` tags, so the marker # is matched OPTIONALLY everywhere. -_OPT = r"(?:" + re.escape(_DSML) + r")?" # optional |DSML| prefix +_OPT = r"(?:" + re.escape(_DSML) + r")?" # optional |DSML| prefix _DSML_PARAM_RE = re.compile( r"<" + _OPT + r'parameter\s+name="(.*?)"(?:\s+string="(true|false)")?\s*>' r"(.*?)", @@ -219,10 +219,10 @@ def _parse_qwen_xml(text: str, tools: Optional[list]) -> Tuple[str, List[ToolCal ) # Region-start markers, both marked and marker-less variants. _DSML_STARTS = ( - "<" + _DSML + "tool_call", # marked (covers tool_call / tool_calls) - "<" + _DSML + "invoke", # marked invoke - "", # marker-less section open + "<" + _DSML + "tool_call", # marked (covers tool_call / tool_calls) + "<" + _DSML + "invoke", # marked invoke + "", # marker-less section open ) @@ -248,7 +248,7 @@ def _unwrap_wrapper_args(args: Any, allowed: set) -> Any: for _ in range(4): # bounded against pathological nesting if not (isinstance(args, dict) and len(args) == 1): break - (k, v), = args.items() + ((k, v),) = args.items() if k not in ("arguments", "input"): break if allowed and k in allowed: @@ -304,7 +304,9 @@ def _dsml_coerce(value: str, string_attr: Optional[str], ptype: Any) -> Any: return v -def _infer_dsml_name(arg_names: set, param_types: Dict[str, Dict[str, Any]]) -> Optional[str]: +def _infer_dsml_name( + arg_names: set, param_types: Dict[str, Dict[str, Any]] +) -> Optional[str]: """Pick the request tool whose parameter set best matches ``arg_names``.""" best, best_score = None, -1e9 for name, props in param_types.items(): @@ -334,7 +336,9 @@ def _parse_dsml(text: str, tools: Optional[list]) -> Tuple[str, List[ToolCall]]: body = m.group(2) or "" # None for self-closing types = param_types.get(name, {}) args: Dict[str, Any] = { - pm.group(1): _dsml_coerce(pm.group(3), pm.group(2), types.get(pm.group(1))) + pm.group(1): _dsml_coerce( + pm.group(3), pm.group(2), types.get(pm.group(1)) + ) for pm in _DSML_PARAM_RE.finditer(body) } # Direct-JSON parameter body (DSML "Format 2", also accepted by @@ -354,7 +358,10 @@ def _parse_dsml(text: str, tools: Optional[list]) -> Tuple[str, List[ToolCall]]: calls.append((name, args)) else: # malformed: no complete invoke wrapper -> collect params, infer tool name - raw = {pm.group(1): (pm.group(3), pm.group(2)) for pm in _DSML_PARAM_RE.finditer(region)} + raw = { + pm.group(1): (pm.group(3), pm.group(2)) + for pm in _DSML_PARAM_RE.finditer(region) + } if raw: name = _infer_dsml_name(set(raw), param_types) or "unknown" types = param_types.get(name, {}) @@ -496,7 +503,9 @@ def _minimax_coerce(value: str, ptype: Any) -> Any: return v -def _parse_minimax(text: str, tools: Optional[list] = None) -> Tuple[str, List[ToolCall]]: +def _parse_minimax( + text: str, tools: Optional[list] = None +) -> Tuple[str, List[ToolCall]]: """Parse MiniMax-M3 tool calls; return (leading_content, tool_calls).""" param_types = _build_param_types(tools) clean = text.replace(_MINIMAX_NS, "") From 662b1df2c33ccb3bdc2c249901dcdb468ff75c84 Mon Sep 17 00:00:00 2001 From: yihonglie Date: Mon, 27 Jul 2026 05:47:57 -0500 Subject: [PATCH 03/18] [Bugfix] Wire batched stream-flush hook so /v1/messages + /v1/chat/completions streaming drains PR #1563 added the batched stream-dispatch (per-seq stream callbacks buffer their chunks into a thread-local; flush_stream_batch() drains the whole step's buffer into the per-request asyncio queues via one call_soon_threadsafe per event loop) but left the hook unwired: engine_core_mgr._flush_stream_batch_fn was only initialized to None, never resolved, and never called in the output thread. The result: streaming /v1/chat/completions and /v1/messages emitted only the initial synthetic role:assistant chunk, then hung indefinitely as the model's content chunks piled up in the thread-local buffer and never reached the per-request queue (the client timed out and the abort-on-disconnect fired). Non-streaming requests were unaffected. Fix: (a) engine_core_mgr's output thread calls self._flush_stream_batch_fn() (if set) after each step's per-seq callback dispatch; (b) the api server resolves the hook to flush_stream_batch lazily after engine init (avoids the api_server <-> engine_core_mgr import cycle). Co-Authored-By: Claude Opus 4.7 (1M context) Signed-off-by: yihonglie (cherry picked from commit 7c2a88a07a807a7d055cef2a60351107bcd7b917) --- atom/entrypoints/openai/api_server.py | 8 ++++++++ atom/model_engine/engine_core_mgr.py | 14 ++++++++++++++ 2 files changed, 22 insertions(+) diff --git a/atom/entrypoints/openai/api_server.py b/atom/entrypoints/openai/api_server.py index 8b37ffe7dc..8c283f6466 100644 --- a/atom/entrypoints/openai/api_server.py +++ b/atom/entrypoints/openai/api_server.py @@ -2122,6 +2122,14 @@ def main(): engine_args = EngineArgs.from_cli_args(args) engine = engine_args.create_engine(tokenizer=tokenizer) + # Wire the batched stream-flush hook: per-seq stream callbacks only buffer + # their chunks into a thread-local; the engine core manager's output thread + # calls this flush after each step's callbacks to drain the buffer into the + # per-request asyncio queues (one call_soon_threadsafe per event loop). + # Registered lazily here to avoid the api_server <-> engine_core_mgr import + # cycle; the core manager leaves the hook as None until this resolves it. + engine.core_mgr._flush_stream_batch_fn = flush_stream_batch + import signal def _sigint_handler(signum, frame): diff --git a/atom/model_engine/engine_core_mgr.py b/atom/model_engine/engine_core_mgr.py index 5e058bc451..b89dcd418b 100644 --- a/atom/model_engine/engine_core_mgr.py +++ b/atom/model_engine/engine_core_mgr.py @@ -285,6 +285,20 @@ def process_outputs_socket(): logger.debug( f"{self.label}: Cleaned up callback for finished sequence {seq_id}" ) + # Batched stream dispatch: the per-seq callbacks only buffer + # their chunks into a thread-local; flush the whole step's + # buffer into the per-request asyncio queues now (one + # call_soon_threadsafe per loop). Resolved lazily by the API + # server to avoid the api_server <-> engine_core_mgr import + # cycle. No-op when no streaming request is in flight. + if self._flush_stream_batch_fn is not None: + try: + self._flush_stream_batch_fn() + except Exception as e: + logger.warning( + f"{self.label}: flush_stream_batch failed: {e}", + exc_info=True, + ) elif request_type == EngineCoreRequestType.UTILITY_RESPONSE: self.utility_response_queue.put_nowait(data) elif request_type == EngineCoreRequestType.ADD: From 34d73f76e17ef667b3b61514e8526e80b2508073 Mon Sep 17 00:00:00 2001 From: yihonglie Date: Mon, 27 Jul 2026 07:31:17 -0500 Subject: [PATCH 04/18] [Bugfix] Inject DSML tool-format instruction on /v1/messages too MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit PR #1563 only injected the mandatory DSML tool-call format instruction on the /v1/responses path (Codex), not /v1/messages (Claude Code / Anthropic). Without it, DSV4 running under Claude Code (which speaks the Anthropic /v1/messages API) does not receive the 'output ONLY a DSML tool-call block' instruction, so it emits ad-hoc/mixed tool calls instead of its native DSML format — weaker agentic behavior. Add the same inject_tool_format_instruction(openai_messages) call (gated on request.tools) that /v1/responses uses, right after anthropic_to_openai_messages. Co-Authored-By: Claude Opus 4.7 (1M context) Signed-off-by: yihonglie (cherry picked from commit 83984016fc141f14982d60cd3d0fb1637d169e4b) --- atom/entrypoints/openai/api_server.py | 8 ++++++++ 1 file changed, 8 insertions(+) diff --git a/atom/entrypoints/openai/api_server.py b/atom/entrypoints/openai/api_server.py index 8c283f6466..1184b47c1e 100644 --- a/atom/entrypoints/openai/api_server.py +++ b/atom/entrypoints/openai/api_server.py @@ -1461,6 +1461,14 @@ async def anthropic_messages(request: AnthropicMessagesRequest, raw_request: Req # Convert Anthropic messages to OpenAI format openai_messages = anthropic_to_openai_messages(request.messages, request.system) + # Inject the mandatory DSML tool-call format instruction when tools are + # present, so the model emits parseable DSML (matching the model's native + # encoding) instead of ad-hoc /```json text. Same injection the + # /v1/responses path applies; without it, Claude Code (which speaks the + # Anthropic /v1/messages API) hits malformed/weak tool calls. + if request.tools: + openai_messages = inject_tool_format_instruction(openai_messages) + # Apply chat template from .protocol import ChatMessage From 4c8a0964d8230927eddb1c9278ac802c67aef014 Mon Sep 17 00:00:00 2001 From: yihonglie Date: Mon, 27 Jul 2026 08:13:59 -0500 Subject: [PATCH 05/18] [Frontend] Inject Claude Code cwd into Bash tool description on /v1/messages MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit DSV4 (sparse indexer, index_topk=1024) cannot reliably verbatim-recall a path buried in a ~20-40k-token system prompt. Claude Code (Anthropic /v1/messages) puts the working directory in the system message but NOT in the Bash tool description, so DSV4 blind-cd's to whatever path the user named (e.g. 'cd /app/vllm-rocm && ls' on a non-existent path) instead of using the absolute path or pwd'ing first. Codex (/v1/responses) sends ... environment_context which extract_cwd rewrites against; Claude Code has no such field (Anthropic-incompatible). Empirical A/B on the same DSV4 server (temp 0, 2 runs/condition): with the cwd in the Bash tool description DSV4 ls's the absolute path / pwd's first (0/2 blind-cd); without it DSV4 blind-cd's to the user-named path (2/2). The endpoint and the #1563 DSML-format inject are NOT the discriminating variable (held constant in the A/B); wording + cwd-presence is. Add extract_cwd_from_system (pulls 'Working directory: ' from the Claude Code system message) + inject_cwd_into_bash_tool (appends 'The current working directory is ' to the Bash tool description), and wire them on the /v1/messages path. This is a prompt-content mitigation of a model-level recall cliff, not a fix to a serving bug — the /v1/messages prompt is otherwise correctly constructed (cwd is in the system message). Co-Authored-By: Claude Opus 4.7 (1M context) Signed-off-by: yihonglie (cherry picked from commit 6e0b326e5bbd88559f385973305fda40efb2ac63) --- atom/entrypoints/openai/api_server.py | 22 ++++++- atom/entrypoints/openai/serving_anthropic.py | 62 ++++++++++++++++++++ 2 files changed, 83 insertions(+), 1 deletion(-) diff --git a/atom/entrypoints/openai/api_server.py b/atom/entrypoints/openai/api_server.py index 1184b47c1e..58a46886a6 100644 --- a/atom/entrypoints/openai/api_server.py +++ b/atom/entrypoints/openai/api_server.py @@ -54,6 +54,8 @@ anthropic_to_openai_messages, anthropic_to_openai_tools, build_anthropic_response, + extract_cwd_from_system, + inject_cwd_into_bash_tool, stream_content_block_delta, stream_content_block_start, stream_content_block_stop, @@ -1469,6 +1471,24 @@ async def anthropic_messages(request: AnthropicMessagesRequest, raw_request: Req if request.tools: openai_messages = inject_tool_format_instruction(openai_messages) + # Claude Code does not send Codex's ... environment_context + # (it's Anthropic-incompatible); it puts the working directory in the + # system message. DSV4 (sparse indexer) cannot reliably recall a path + # buried in a ~20-40k-token system prompt, so without the cwd in the + # Bash tool description it blind-`cd`s to whatever path the user named. + # Extract the cwd from the system message and inject it into the Bash + # tool description so the model uses the absolute path / `pwd`s first. + # (Empirically deterministic over 2 runs/condition; this is a + # prompt-content mitigation of a model-level recall cliff, not a fix to + # a serving bug — the prompt is otherwise correctly constructed.) + tools_with_cwd = request.tools + if request.tools: + cwd = extract_cwd_from_system(request.system) + if cwd: + tools_with_cwd = inject_cwd_into_bash_tool( + [dict(t) for t in request.tools], cwd + ) + # Apply chat template from .protocol import ChatMessage @@ -1479,7 +1499,7 @@ async def anthropic_messages(request: AnthropicMessagesRequest, raw_request: Req tokenizer, custom_message_encoder, [msg.to_template_dict() for msg in messages], - tools=anthropic_to_openai_tools(request.tools), + tools=anthropic_to_openai_tools(tools_with_cwd), **merged_kwargs, ) diff --git a/atom/entrypoints/openai/serving_anthropic.py b/atom/entrypoints/openai/serving_anthropic.py index 8a1afb1f81..bca0d7b90a 100644 --- a/atom/entrypoints/openai/serving_anthropic.py +++ b/atom/entrypoints/openai/serving_anthropic.py @@ -161,6 +161,68 @@ def anthropic_to_openai_tools(tools: Optional[List[dict]]) -> Optional[List[dict return result +def _text_of(block: Any) -> str: + """Flatten an Anthropic content block (str or list of {text} blocks) to text.""" + if isinstance(block, str): + return block + if isinstance(block, list): + return "\n".join( + _text_of(b.get("text")) if isinstance(b, dict) else _text_of(b) + for b in block + if (b.get("text") if isinstance(b, dict) else b) + ) + return "" + + +def extract_cwd_from_system(system: Any) -> Optional[str]: + """Pull the working directory from Claude Code's system prompt. + + Claude Code embeds the working directory as a `Working directory: ` + line in the system message it sends (it does not use Codex's + `...` environment_context, which is Anthropic-incompatible). + DSV4 (sparse indexer, index_topk=1024) cannot reliably verbatim-recall a + path buried in a ~20-40k-token system prompt, so when the Bash tool + description omits the cwd the model blind-`cd`s to whatever path the user + named. Returning the cwd here lets the caller inject it into the Bash tool + description so the model uses the absolute path or `pwd`s instead. + """ + import re + + blob = _text_of(system) + if not blob: + return None + m = re.search(r"(?:^|\n)\s*Working directory:\s*([^\s<\n]+)", blob) + return m.group(1) if m else None + + +def inject_cwd_into_bash_tool( + tools: Optional[List[dict]], cwd: Optional[str] +) -> Optional[List[dict]]: + """Append the working directory to the Bash tool's description. + + Empirically (A/B on the same DSV4 server, temp 0, 2 runs each): with the + cwd in the Bash tool description, DSV4 stops blind-`cd`-ing to whatever + path the user named and instead `ls`'s the absolute path or `pwd`s first. + Without it, DSV4 faithfully `cd ` even when that path + does not exist. This is a prompt-content mitigation for a model-level + recall/cliff limitation (the ATOM /v1/messages prompt is otherwise + correctly constructed — cwd is in the system message), not a serving bug. + """ + if not tools or not cwd: + return tools + for tool in tools: + if tool.get("name") == "Bash": + base = tool.get("description", "") or "" + hint = ( + f" The current working directory is {cwd}." + if "current working directory" not in base.lower() + else "" + ) + if hint and cwd not in base: + tool["description"] = base + hint + return tools + + # ── Response Construction ────────────────────────────────────────────── From 1ddd0b4f3db7b1f92d87eeed4138c1df18b59544 Mon Sep 17 00:00:00 2001 From: yihonglie Date: Mon, 27 Jul 2026 11:03:05 -0500 Subject: [PATCH 06/18] [Bugfix] Abort leaked seq on streaming client disconnect (chat + responses) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Three streaming generators called cleanup_streaming_request() in their finally with the default aborted=False. On client disconnect (GeneratorExit) the engine sequence was never aborted — it kept decoding to max_tokens with no consumer, wasting GPU cycles and KV-cache blocks. cleanup_streaming_request() only calls engine.core_mgr.abort_request() when aborted=True. Restore the aborted-tracking pattern already used by serving_completion.py: assume aborted=True before the stream, flip to False only after the stream reaches its normal end, and pass aborted to cleanup in the finally. Affected: stream_chat_response, stream_chat_response_fanout (serving_chat.py) and generate_responses_stream (api_server.py). Normal completion is unchanged (aborted=False -> abort is a no-op, no extra control-path broadcast). Co-Authored-By: Claude Opus 4.8 (1M context) Signed-off-by: yihonglie (cherry picked from commit a35fbf3c7f6daf6816b9e91a96a9ae688aa881f2) --- atom/entrypoints/openai/api_server.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/atom/entrypoints/openai/api_server.py b/atom/entrypoints/openai/api_server.py index 58a46886a6..e2c4ff9ec7 100644 --- a/atom/entrypoints/openai/api_server.py +++ b/atom/entrypoints/openai/api_server.py @@ -1916,6 +1916,7 @@ def handle(etype, edata): return _flush_pending() return [] + aborted = True # abort the seq on disconnect; False on normal end try: for s in emitter.created(): yield s @@ -1944,9 +1945,10 @@ def handle(etype, edata): for s in emitter.finish(input_tokens, output_tokens): yield s yield "data: [DONE]\n\n" + aborted = False break finally: - cleanup_streaming_request(request_id, seq_id) + cleanup_streaming_request(request_id, seq_id, aborted=aborted) return StreamingResponse( generate_responses_stream(), From 4bb70af7d6c402ca65f6d7028fd6adcc8780dfe1 Mon Sep 17 00:00:00 2001 From: yihonglie Date: Mon, 27 Jul 2026 11:04:12 -0500 Subject: [PATCH 07/18] [Bugfix] Guard Responses tool translation against malformed inputs Two crash sites in serving_responses.py, both reachable from model/client input and fatal mid-stream (500 or truncated SSE): - responses_tools_to_openai(): `tl.get("function", tl)` returns None when the key is present with value null, so the following `fn.get("name")` raised AttributeError. Use `tl.get("function") or tl`. - _claude_tool_to_shell(): `int(off or 0)` / `int(lim or 200)` raised ValueError on non-numeric offset/limit (e.g. the model emits {"offset":"top"}), crashing the streaming handler. Wrap in try/except -> fall back to (1, 400). Co-Authored-By: Claude Opus 4.8 (1M context) Signed-off-by: yihonglie (cherry picked from commit ddc2897fcc830634824d838ca8b2464318ae220f) --- atom/entrypoints/openai/serving_responses.py | 9 ++++++--- 1 file changed, 6 insertions(+), 3 deletions(-) diff --git a/atom/entrypoints/openai/serving_responses.py b/atom/entrypoints/openai/serving_responses.py index 26d22af821..01cf25bb30 100644 --- a/atom/entrypoints/openai/serving_responses.py +++ b/atom/entrypoints/openai/serving_responses.py @@ -148,7 +148,7 @@ def responses_tools_to_openai(tools: Any) -> List[Dict[str, Any]]: if not isinstance(tl, dict): continue if tl.get("type") == "function": - fn = tl.get("function", tl) + fn = tl.get("function") or tl ct.append( { "type": "function", @@ -314,8 +314,11 @@ def _claude_tool_to_shell( return None off, lim = a.get("offset"), a.get("limit") if off or lim: - start = int(off or 0) + 1 - end = start + int(lim or 200) - 1 + try: + start = int(off or 0) + 1 + end = start + int(lim or 200) - 1 + except (ValueError, TypeError): + start, end = 1, 400 else: start, end = 1, 400 return _read_cmd(str(fp), cwd, start, end) From fba0e99e8642c3b4277e1fd673dc30cf2f830e03 Mon Sep 17 00:00:00 2001 From: yihonglie Date: Mon, 27 Jul 2026 11:04:12 -0500 Subject: [PATCH 08/18] [Bugfix] Preserve cwd paths containing spaces in Codex extraction extract_cwd()'s regex `\s*([^<\s]+)\s*` stopped at the first whitespace, truncating paths like "/home/user/My Projects/foo" to "/home/user/My" and corrupting all downstream path fix-ups. Use `[^<]+?` (keeps internal spaces/newlines, never spans a stray '<', non-greedy so it stops at the first ) and strip() the captured group. Co-Authored-By: Claude Opus 4.8 (1M context) Signed-off-by: yihonglie (cherry picked from commit 2be8dcb46b514b70974decd2b927f6e5661c765b) --- atom/entrypoints/openai/serving_responses.py | 7 +++++-- 1 file changed, 5 insertions(+), 2 deletions(-) diff --git a/atom/entrypoints/openai/serving_responses.py b/atom/entrypoints/openai/serving_responses.py index 01cf25bb30..662bb485a0 100644 --- a/atom/entrypoints/openai/serving_responses.py +++ b/atom/entrypoints/openai/serving_responses.py @@ -445,7 +445,10 @@ def extract_cwd(body: Dict[str, Any]) -> Optional[str]: global _CWD_RE if _CWD_RE is None: - _CWD_RE = re.compile(r"\s*([^<\s]+)\s*") + # [^<]+? keeps internal spaces/newlines but never spans a stray '<' (so a + # malformed/unclosed inner can't over-capture); non-greedy stops at + # the first . + _CWD_RE = re.compile(r"\s*([^<]+?)\s*") blob = _text_of(body.get("instructions")) inp = body.get("input") if isinstance(inp, str): @@ -455,7 +458,7 @@ def extract_cwd(body: Dict[str, Any]) -> Optional[str]: if isinstance(it, dict): blob += "\n" + _text_of(it.get("content", "")) m = _CWD_RE.search(blob or "") - return m.group(1) if m else None + return m.group(1).strip() if m else None # --------------------------------------------------------------- SSE emitter From 2bbdfb0c6c4127bfbc5e2ae0ad55d2d51619c428 Mon Sep 17 00:00:00 2001 From: yihonglie Date: Tue, 28 Jul 2026 15:01:31 +0800 Subject: [PATCH 09/18] [Frontend] Drop Claude Code cwd injection on /v1/messages Remove extract_cwd_from_system() and inject_cwd_into_bash_tool() and their /v1/messages call site. They parsed the `Working directory: ` line out of Claude Code's system prompt and appended it to the Bash tool description, as a prompt-content mitigation for DSV4's recall cliff on long system prompts. Also drops the private _text_of() helper, added by the same commit and left with no callers once the two functions above are gone. Codex's ... handling (serving_responses.extract_cwd and its _resolve_dir / _read_cmd consumers) is a separate mechanism and is unchanged. Co-Authored-By: Claude Opus 5 (1M context) Signed-off-by: yihonglie --- atom/entrypoints/openai/api_server.py | 22 +------ atom/entrypoints/openai/serving_anthropic.py | 62 -------------------- 2 files changed, 1 insertion(+), 83 deletions(-) diff --git a/atom/entrypoints/openai/api_server.py b/atom/entrypoints/openai/api_server.py index e2c4ff9ec7..165810225a 100644 --- a/atom/entrypoints/openai/api_server.py +++ b/atom/entrypoints/openai/api_server.py @@ -54,8 +54,6 @@ anthropic_to_openai_messages, anthropic_to_openai_tools, build_anthropic_response, - extract_cwd_from_system, - inject_cwd_into_bash_tool, stream_content_block_delta, stream_content_block_start, stream_content_block_stop, @@ -1471,24 +1469,6 @@ async def anthropic_messages(request: AnthropicMessagesRequest, raw_request: Req if request.tools: openai_messages = inject_tool_format_instruction(openai_messages) - # Claude Code does not send Codex's ... environment_context - # (it's Anthropic-incompatible); it puts the working directory in the - # system message. DSV4 (sparse indexer) cannot reliably recall a path - # buried in a ~20-40k-token system prompt, so without the cwd in the - # Bash tool description it blind-`cd`s to whatever path the user named. - # Extract the cwd from the system message and inject it into the Bash - # tool description so the model uses the absolute path / `pwd`s first. - # (Empirically deterministic over 2 runs/condition; this is a - # prompt-content mitigation of a model-level recall cliff, not a fix to - # a serving bug — the prompt is otherwise correctly constructed.) - tools_with_cwd = request.tools - if request.tools: - cwd = extract_cwd_from_system(request.system) - if cwd: - tools_with_cwd = inject_cwd_into_bash_tool( - [dict(t) for t in request.tools], cwd - ) - # Apply chat template from .protocol import ChatMessage @@ -1499,7 +1479,7 @@ async def anthropic_messages(request: AnthropicMessagesRequest, raw_request: Req tokenizer, custom_message_encoder, [msg.to_template_dict() for msg in messages], - tools=anthropic_to_openai_tools(tools_with_cwd), + tools=anthropic_to_openai_tools(request.tools), **merged_kwargs, ) diff --git a/atom/entrypoints/openai/serving_anthropic.py b/atom/entrypoints/openai/serving_anthropic.py index bca0d7b90a..8a1afb1f81 100644 --- a/atom/entrypoints/openai/serving_anthropic.py +++ b/atom/entrypoints/openai/serving_anthropic.py @@ -161,68 +161,6 @@ def anthropic_to_openai_tools(tools: Optional[List[dict]]) -> Optional[List[dict return result -def _text_of(block: Any) -> str: - """Flatten an Anthropic content block (str or list of {text} blocks) to text.""" - if isinstance(block, str): - return block - if isinstance(block, list): - return "\n".join( - _text_of(b.get("text")) if isinstance(b, dict) else _text_of(b) - for b in block - if (b.get("text") if isinstance(b, dict) else b) - ) - return "" - - -def extract_cwd_from_system(system: Any) -> Optional[str]: - """Pull the working directory from Claude Code's system prompt. - - Claude Code embeds the working directory as a `Working directory: ` - line in the system message it sends (it does not use Codex's - `...` environment_context, which is Anthropic-incompatible). - DSV4 (sparse indexer, index_topk=1024) cannot reliably verbatim-recall a - path buried in a ~20-40k-token system prompt, so when the Bash tool - description omits the cwd the model blind-`cd`s to whatever path the user - named. Returning the cwd here lets the caller inject it into the Bash tool - description so the model uses the absolute path or `pwd`s instead. - """ - import re - - blob = _text_of(system) - if not blob: - return None - m = re.search(r"(?:^|\n)\s*Working directory:\s*([^\s<\n]+)", blob) - return m.group(1) if m else None - - -def inject_cwd_into_bash_tool( - tools: Optional[List[dict]], cwd: Optional[str] -) -> Optional[List[dict]]: - """Append the working directory to the Bash tool's description. - - Empirically (A/B on the same DSV4 server, temp 0, 2 runs each): with the - cwd in the Bash tool description, DSV4 stops blind-`cd`-ing to whatever - path the user named and instead `ls`'s the absolute path or `pwd`s first. - Without it, DSV4 faithfully `cd ` even when that path - does not exist. This is a prompt-content mitigation for a model-level - recall/cliff limitation (the ATOM /v1/messages prompt is otherwise - correctly constructed — cwd is in the system message), not a serving bug. - """ - if not tools or not cwd: - return tools - for tool in tools: - if tool.get("name") == "Bash": - base = tool.get("description", "") or "" - hint = ( - f" The current working directory is {cwd}." - if "current working directory" not in base.lower() - else "" - ) - if hint and cwd not in base: - tool["description"] = base + hint - return tools - - # ── Response Construction ────────────────────────────────────────────── From 9e06aa11a6d67f1b1eeb41f39d6e70433a4358a7 Mon Sep 17 00:00:00 2001 From: yihonglie Date: Tue, 28 Jul 2026 15:24:22 +0800 Subject: [PATCH 10/18] refactor(openai): split tool_parser into one module per wire format MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit tool_parser.py had grown to 1048 lines holding five formats, and the streaming half was largely copy-paste: _process_/_flush_ for qwen, dsml, glm and minimax were 57 lines each differing in exactly two — the marker lookup and which _parse_* to call. A stale `# -- Qwen3 XML --` header left above _process_glm recorded where the last paste happened. All five also shared one mutable `state`/`buf` on a single dataclass, with `state`'s meaning documented for Kimi only. Split into a package, one module per format, named after vLLM's tool_parsers/ layout: tool_parser/ tool_parser.py ToolCall, ToolCallParser, BufferedMarkerParser schema.py build_param_types / coerce_param_value / coerce_json_or_raw registry.py detection order + stream sniff + parse_tool_calls stream.py ToolCallStreamParser facade kimi_tool_parser.py qwen3_tool_parser.py deepseekv4_tool_parser.py glm_tool_parser.py minimax_tool_parser.py Four of the five formats stream identically (buffer from a start marker, parse the block at flush), so that strategy is written once in BufferedMarkerParser and each format declares only START_MARKERS, HOLDBACK_CHARS and parse(). Kimi keeps its own process/flush: its token format is self-delimiting, so it emits calls incrementally rather than buffering. GLM's and MiniMax's identical value coercion is now one shared helper; DeepSeek-V4's is deliberately NOT merged with it, because on a JSON-decode miss it falls back to value.strip() where the shared one falls back to value.strip("\n"). Behavior is unchanged, including the parts that are only accidentally correct: the detection order (documented in one place in registry.py instead of scattered across if-branches), the streaming sniff being stricter than the non-streaming detect, the `len(buf) > 8` heuristic, and parse_tool_calls returning unstripped text when no format matches while every matching path strips. Public API is unchanged — parse_tool_calls, ToolCallStreamParser and ToolCall are re-exported from the package, so all five call sites, atomesh and the existing tests import exactly as before. Net +59 lines: ~150 lines of duplicated streaming logic removed, offset by per-file headers and the parser interface. Verified: existing 13 tool-parser tests pass, plus a differential harness comparing old vs new across 33 inputs (all five formats, malformed DSML variants, unclosed blocks) — 226 comparisons covering parse_tool_calls and streaming under six chunkings including one character at a time — byte-identical after normalizing random tool-call ids. Co-Authored-By: Claude Opus 5 (1M context) Signed-off-by: yihonglie --- atom/entrypoints/openai/tool_parser.py | 1048 ----------------- .../openai/tool_parser/__init__.py | 61 + .../tool_parser/deepseekv4_tool_parser.py | 211 ++++ .../openai/tool_parser/glm_tool_parser.py | 78 ++ .../openai/tool_parser/kimi_tool_parser.py | 159 +++ .../openai/tool_parser/minimax_tool_parser.py | 79 ++ .../openai/tool_parser/qwen3_tool_parser.py | 92 ++ .../openai/tool_parser/registry.py | 104 ++ atom/entrypoints/openai/tool_parser/schema.py | 87 ++ atom/entrypoints/openai/tool_parser/stream.py | 69 ++ .../openai/tool_parser/tool_parser.py | 167 +++ 11 files changed, 1107 insertions(+), 1048 deletions(-) delete mode 100644 atom/entrypoints/openai/tool_parser.py create mode 100644 atom/entrypoints/openai/tool_parser/__init__.py create mode 100644 atom/entrypoints/openai/tool_parser/deepseekv4_tool_parser.py create mode 100644 atom/entrypoints/openai/tool_parser/glm_tool_parser.py create mode 100644 atom/entrypoints/openai/tool_parser/kimi_tool_parser.py create mode 100644 atom/entrypoints/openai/tool_parser/minimax_tool_parser.py create mode 100644 atom/entrypoints/openai/tool_parser/qwen3_tool_parser.py create mode 100644 atom/entrypoints/openai/tool_parser/registry.py create mode 100644 atom/entrypoints/openai/tool_parser/schema.py create mode 100644 atom/entrypoints/openai/tool_parser/stream.py create mode 100644 atom/entrypoints/openai/tool_parser/tool_parser.py diff --git a/atom/entrypoints/openai/tool_parser.py b/atom/entrypoints/openai/tool_parser.py deleted file mode 100644 index 21a98f5df4..0000000000 --- a/atom/entrypoints/openai/tool_parser.py +++ /dev/null @@ -1,1048 +0,0 @@ -# SPDX-License-Identifier: MIT -# Copyright (C) 2024-2026, Advanced Micro Devices, Inc. All rights reserved. - -"""Tool call parser for models that output tool calls. - -Two on-the-wire formats are auto-detected and normalized into the OpenAI -``tool_calls`` structure: - -1. Kimi-K2 special-token format:: - - <|tool_calls_section_begin|> - <|tool_call_begin|>functions.NAME:INDEX<|tool_call_argument_begin|>ARGS_JSON<|tool_call_end|> - <|tool_calls_section_end|> - -2. Qwen3 (qwen3_coder / qwen3_xml) XML format:: - - - - VALUE - ... - - - -The Qwen XML carries no value types, so when the request's ``tools`` schema is -supplied each parameter is coerced to the declared JSON-Schema type (int, float, -bool, null, object, array); otherwise it is left as a string. This mirrors the -qwen3_coder/qwen3_xml parsers in vLLM and SGLang. - -OpenAI format: - {"tool_calls": [{"id": "call_0", "type": "function", - "function": {"name": "NAME", "arguments": "ARGS_JSON"}}]} -""" - -import ast -import json -import re -import uuid -from dataclasses import dataclass -from typing import Any, Dict, List, Optional, Tuple - - -def _unique_tool_call_id() -> str: - # OpenAI tool_call ids must be unique across the whole conversation, not just - # within one response. A per-response index (call_0, call_1, ...) collides - # across turns -> clients (e.g. qwen-code) dedupe by id and silently ignore - # every repeat, causing an infinite tool-call retry loop. Use a random id. - return f"call_{uuid.uuid4().hex}" - - -@dataclass -class ToolCall: - """Parsed tool call in OpenAI format.""" - - id: str - type: str - function: Dict[str, str] - - def to_dict(self) -> Dict[str, Any]: - return {"id": self.id, "type": self.type, "function": self.function} - - -# --------------------------------------------------------------------------- -# Qwen3 XML tool-call format (qwen3_coder / qwen3_xml) -# --------------------------------------------------------------------------- - -_QWEN_TOOL_PREFIX = "||(?=)|$)", - re.DOTALL, -) - - -def _is_qwen_xml(text: str) -> bool: - """Detect the Qwen3 XML tool-call format (and not the Kimi token format).""" - return _QWEN_TOOL_PREFIX in text and "<|tool_calls_section_begin|>" not in text - - -def _build_param_types(tools: Optional[list]) -> Dict[str, Dict[str, Any]]: - """Map ``function_name -> {param_name: json_schema_type}`` from request tools. - - Accepts OpenAI (``{"type": "function", "function": {...}}``) and bare - (``{"name": ..., "parameters"/"input_schema": {...}}``) tool entries. - """ - out: Dict[str, Dict[str, Any]] = {} - for tool in tools or []: - if not isinstance(tool, dict): - continue - fn = tool.get("function", tool) - if not isinstance(fn, dict): - continue - name = fn.get("name") - if not name: - continue - schema = fn.get("parameters") or fn.get("input_schema") or {} - props = schema.get("properties") if isinstance(schema, dict) else None - out[name] = { - k: (v.get("type") if isinstance(v, dict) else None) - for k, v in (props or {}).items() - } - return out - - -def _coerce_param_value(value: str, ptype: Any) -> Any: - """Coerce a string parameter value to its declared JSON-Schema type. - - No schema type (string/unknown) -> returned unchanged. Conversion failures - fall back to the original string rather than raising. - """ - v = value.strip("\n") - if ptype is None: - return v - t = str(ptype).lower() - try: - if t in ("string", "str", "text", "varchar", "char", "enum"): - return v - if t in ("null", "none"): - return None - if t.startswith(("int", "uint", "long", "short", "unsigned")): - return int(v) - if t.startswith(("num", "float", "double", "decimal")): - f = float(v) - return int(f) if f.is_integer() else f - if t.startswith(("bool", "binary")): - return v.strip().lower() == "true" - if t.startswith(("object", "dict", "map", "array", "list", "tuple")): - try: - return json.loads(v) - except Exception: - return ast.literal_eval(v) # safer for single-quoted Python literals - except Exception: - return v - return v - - -def _parse_qwen_function( - fn_text: str, param_types: Dict[str, Dict[str, Any]], index: int -) -> Optional[ToolCall]: - """Parse the inside of one ``...`` block into a ToolCall.""" - gt = fn_text.find(">") - if gt == -1: - return None - name = fn_text[:gt].strip() - if not name: - return None - body = fn_text[gt + 1 :] - types = param_types.get(name, {}) - args: Dict[str, Any] = {} - for pm in _QWEN_PARAM_RE.finditer(body): - seg = pm.group(1) - if seg is None: - continue - pgt = seg.find(">") - if pgt == -1: - continue - pname = seg[:pgt].strip() - pval = seg[pgt + 1 :] - if pname: - args[pname] = _coerce_param_value(pval, types.get(pname)) - return ToolCall( - id=_unique_tool_call_id(), - type="function", - function={"name": name, "arguments": json.dumps(args, ensure_ascii=False)}, - ) - - -def _parse_qwen_xml(text: str, tools: Optional[list]) -> Tuple[str, List[ToolCall]]: - """Parse Qwen3 XML tool calls; return (leading_content, tool_calls).""" - param_types = _build_param_types(tools) - # Content precedes the first tool marker. - markers = [ - i for i in (text.find(""), text.find(_QWEN_TOOL_PREFIX)) if i != -1 - ] - content = text[: min(markers)] if markers else text - tool_calls: List[ToolCall] = [] - for fm in _QWEN_FUNCTION_RE.finditer(text): - fn_text = fm.group(1) if fm.group(1) is not None else fm.group(2) - if not fn_text: - continue - tc = _parse_qwen_function(fn_text, param_types, len(tool_calls)) - if tc is not None: - tool_calls.append(tc) - return content.strip(), tool_calls - - -# --------------------------------------------------------------------------- -# DeepSeek-V4 DSML tool-call format -# --------------------------------------------------------------------------- -# -# <|DSML|tool_calls> -# <|DSML|invoke name="NAME"> -# <|DSML|parameter name="PNAME" string="true|false">VALUE -# ... -# -# -# -# string="true" -> value is a raw string; string="false" -> value is JSON. -# DeepSeek-V4-Flash occasionally malforms this (singular ``tool_call``, a missing -# ``invoke`` wrapper, or params without ``string=``); the parser recovers those -# best-effort: it infers a dropped tool name from the parameter signature vs the -# request's ``tools`` and infers a missing value type from the schema / JSON. - -_DSML = "|DSML|" -# The model often DROPS the ``|DSML|`` marker and emits bare -# ````/````/```` tags, so the marker -# is matched OPTIONALLY everywhere. -_OPT = r"(?:" + re.escape(_DSML) + r")?" # optional |DSML| prefix -_DSML_PARAM_RE = re.compile( - r"<" + _OPT + r'parameter\s+name="(.*?)"(?:\s+string="(true|false)")?\s*>' - r"(.*?)", - re.DOTALL, -) -# Long-form `...` OR self-closing `` -# (the zero-arg shape; group(2) is None for self-closing). Matches SGLang's V4 -# detector, which accepts both. -_DSML_INVOKE_RE = re.compile( - r"<" + _OPT + r'invoke\s+name="(.*?)"\s*(?:/>|>(.*?))", - re.DOTALL, -) -# Region-start markers, both marked and marker-less variants. -_DSML_STARTS = ( - "<" + _DSML + "tool_call", # marked (covers tool_call / tool_calls) - "<" + _DSML + "invoke", # marked invoke - "", # marker-less section open -) - - -def _dsml_start(text: str) -> int: - """Index of the earliest DSML tool-call marker (marked or marker-less), or -1.""" - positions = [i for i in (text.find(m) for m in _DSML_STARTS) if i != -1] - return min(positions) if positions else -1 - - -def _is_dsml(text: str) -> bool: - return _dsml_start(text) != -1 - - -def _unwrap_wrapper_args(args: Any, allowed: set) -> Any: - """Strip spurious ``{"arguments": {...}}`` / ``{"input": {...}}`` envelopes. - - Non-tuned models (DeepSeek-V4-Pro) frequently wrap the real args in an extra - ``arguments``/``input`` object — sometimes nested 2-3 deep, or stringified — - so a call meant as ``{"cmd": "ls"}`` arrives as ``{"arguments": {"cmd": - "ls"}}`` and the client (Codex) rejects it ("missing field cmd"). Recursively - unwrap while the sole key is a wrapper that is NOT itself a declared param of - the tool. Mirrors vLLM's ``_unwrap_wrapper_args`` (deepseek_v4.py).""" - for _ in range(4): # bounded against pathological nesting - if not (isinstance(args, dict) and len(args) == 1): - break - ((k, v),) = args.items() - if k not in ("arguments", "input"): - break - if allowed and k in allowed: - break # this tool really has a param named arguments/input - if isinstance(v, str): - try: - v = json.loads(v) - except Exception: - break - if not isinstance(v, dict): - break - args = v - return args - - -# Canonical param key -> emitted synonyms the model uses interchangeably. Only -# applied when the canonical key is in the tool's declared schema and the model -# used a synonym instead (e.g. Codex's exec_command wants `cmd`, but the model -# habitually emits `command` -> "missing field cmd" retry loop). Scoped to the -# schema so it never renames a key a tool legitimately declares. -_KEY_ALIASES: Dict[str, Tuple[str, ...]] = { - "cmd": ("command",), -} - - -def _apply_key_aliases(args: Any, allowed: set) -> Any: - if not (isinstance(args, dict) and allowed): - return args - for canon, syns in _KEY_ALIASES.items(): - if canon in allowed and canon not in args: - for s in syns: - if s in args: - args[canon] = args.pop(s) - break - return args - - -def _dsml_coerce(value: str, string_attr: Optional[str], ptype: Any) -> Any: - if string_attr == "true": - return value - if string_attr == "false": - try: - return json.loads(value) - except Exception: - return value - # attr absent -> use declared schema type if known, else infer via JSON. - if ptype is not None: - return _coerce_param_value(value, ptype) - v = value.strip() - try: - return json.loads(v) - except Exception: - return v - - -def _infer_dsml_name( - arg_names: set, param_types: Dict[str, Dict[str, Any]] -) -> Optional[str]: - """Pick the request tool whose parameter set best matches ``arg_names``.""" - best, best_score = None, -1e9 - for name, props in param_types.items(): - p = set(props) - if not p: - continue - score = len(p & arg_names) - 0.1 * len(p ^ arg_names) - if score > best_score: - best_score, best = score, name - return best - - -def _parse_dsml(text: str, tools: Optional[list]) -> Tuple[str, List[ToolCall]]: - """Parse DeepSeek-V4 DSML tool calls; return (leading_content, tool_calls).""" - param_types = _build_param_types(tools) - start = _dsml_start(text) - if start == -1: - return text.strip(), [] - content = text[:start] - region = text[start:] - - calls: List[Tuple[str, Dict[str, Any]]] = [] - invokes = list(_DSML_INVOKE_RE.finditer(region)) - if invokes: - for m in invokes: - name = m.group(1) - body = m.group(2) or "" # None for self-closing - types = param_types.get(name, {}) - args: Dict[str, Any] = { - pm.group(1): _dsml_coerce( - pm.group(3), pm.group(2), types.get(pm.group(1)) - ) - for pm in _DSML_PARAM_RE.finditer(body) - } - # Direct-JSON parameter body (DSML "Format 2", also accepted by - # vLLM/SGLang): ` { "k": "v" } ` with no - # tags. Falls through here with empty args; recover them. - if not args: - stripped = body.strip() - if stripped.startswith("{"): - try: - parsed = json.loads(stripped) - if isinstance(parsed, dict): - args = parsed - except Exception: - pass - args = _unwrap_wrapper_args(args, set(types)) - args = _apply_key_aliases(args, set(types)) - calls.append((name, args)) - else: - # malformed: no complete invoke wrapper -> collect params, infer tool name - raw = { - pm.group(1): (pm.group(3), pm.group(2)) - for pm in _DSML_PARAM_RE.finditer(region) - } - if raw: - name = _infer_dsml_name(set(raw), param_types) or "unknown" - types = param_types.get(name, {}) - args = {k: _dsml_coerce(v, s, types.get(k)) for k, (v, s) in raw.items()} - args = _unwrap_wrapper_args(args, set(types)) - args = _apply_key_aliases(args, set(types)) - calls.append((name, args)) - - tool_calls = [ - ToolCall( - id=_unique_tool_call_id(), - type="function", - function={"name": name, "arguments": json.dumps(args, ensure_ascii=False)}, - ) - for name, args in calls - ] - if _DSML in content: # scrub any stray marker fragment - content = content.split("<" + _DSML, 1)[0] - return content.strip(), tool_calls - - -# --------------------------------------------------------------------------- -# GLM-4.5 / 4.6 / 5.x tool-call format -# --------------------------------------------------------------------------- -# -# NAME -# K1V1 -# K2V2 -# ... -# -# The function name follows the opening tag directly (no ``(.*?)|(.*)$", re.DOTALL -) -_GLM_ARG_RE = re.compile( - r"(.*?)\s*" - r"(.*?)(?:|(?=)|(?=)|$)", - re.DOTALL, -) - - -def _is_glm(text: str) -> bool: - """Detect the GLM ``...`` format (never Qwen/DSML).""" - if _QWEN_TOOL_PREFIX in text: # ' Qwen, not GLM - return False - return "" in text or "" in text - - -def _glm_coerce(value: str, ptype: Any) -> Any: - """Decode one GLM ````: schema type wins, else JSON, else raw.""" - v = value.strip("\n") - if ptype is not None: - return _coerce_param_value(v, ptype) - s = v.strip() - try: - return json.loads(s) - except Exception: - return v - - -def _parse_glm(text: str, tools: Optional[list] = None) -> Tuple[str, List[ToolCall]]: - """Parse GLM tool calls; return (leading_content, tool_calls).""" - param_types = _build_param_types(tools) - start = text.find("") - if start == -1: - return text.strip(), [] - content = text[:start] - tool_calls: List[ToolCall] = [] - for m in _GLM_TOOLCALL_RE.finditer(text): - body = m.group(1) if m.group(1) is not None else m.group(2) - if not body: - continue - ak = body.find("") - name = (body if ak == -1 else body[:ak]).strip() - if not name: - continue - types = param_types.get(name, {}) - args: Dict[str, Any] = {} - for pm in _GLM_ARG_RE.finditer(body): - k = pm.group(1).strip() - if k: - args[k] = _glm_coerce(pm.group(2), types.get(k)) - tool_calls.append( - ToolCall( - id=_unique_tool_call_id(), - type="function", - function={ - "name": name, - "arguments": json.dumps(args, ensure_ascii=False), - }, - ) - ) - return content.strip(), tool_calls - - -# --------------------------------------------------------------------------- -# MiniMax-M3 tool-call format -# --------------------------------------------------------------------------- -# -# Every tag is prefixed by the ns_token ``]<]minimax[>[``: -# -# ]<]minimax[>[ -# ]<]minimax[>[ -# ]<]minimax[>[value]<]minimax[>[ -# ... -# ]<]minimax[>[ -# ]<]minimax[>[ -# -# Unlike DSML, parameters are named by the TAG itself (``Paris``), -# not a ``name="..."`` attribute. Strip the ns_token first, then parse -# / pairs. Values: schema type wins, else JSON, else raw string. - -_MINIMAX_NS = "]<]minimax[>[" -_MINIMAX_INVOKE_RE = re.compile( - r'(.*?)|(.*)$', - re.DOTALL, -) -_MINIMAX_PARAM_RE = re.compile(r"<([\w-]+)>(.*?)", re.DOTALL) - - -def _is_minimax(text: str) -> bool: - """Detect the MiniMax-M3 ns_token tool-call format.""" - return _MINIMAX_NS in text - - -def _minimax_coerce(value: str, ptype: Any) -> Any: - v = value.strip("\n") - if ptype is not None: - return _coerce_param_value(v, ptype) - s = v.strip() - try: - return json.loads(s) - except Exception: - return v - - -def _parse_minimax( - text: str, tools: Optional[list] = None -) -> Tuple[str, List[ToolCall]]: - """Parse MiniMax-M3 tool calls; return (leading_content, tool_calls).""" - param_types = _build_param_types(tools) - clean = text.replace(_MINIMAX_NS, "") - tc = clean.find("") - content = clean[:tc] if tc > 0 else ("" if tc == 0 else clean) - tool_calls: List[ToolCall] = [] - for m in _MINIMAX_INVOKE_RE.finditer(clean): - name = m.group(1) if m.group(1) is not None else m.group(3) - body = m.group(2) if m.group(2) is not None else (m.group(4) or "") - if not name: - continue - name = name.strip() - types = param_types.get(name, {}) - args: Dict[str, Any] = {} - for pm in _MINIMAX_PARAM_RE.finditer(body): - k = pm.group(1).strip() - if k: - args[k] = _minimax_coerce(pm.group(2), types.get(k)) - tool_calls.append( - ToolCall( - id=_unique_tool_call_id(), - type="function", - function={ - "name": name, - "arguments": json.dumps(args, ensure_ascii=False), - }, - ) - ) - for mk in ("", ""): - content = content.replace(mk, "") - return content.strip(), tool_calls - - -def parse_tool_calls( - text: str, tools: Optional[list] = None -) -> Tuple[str, List[ToolCall]]: - """Parse tool calls from model output text. - - Args: - text: Raw model output that may contain tool calls (DeepSeek-V4 DSML, - Kimi token format, or Qwen3 XML format). - tools: Optional request tool definitions; used to type-coerce parameter - values to their declared JSON-Schema types. - - Returns: - Tuple of (content_text, list_of_tool_calls). ``content_text`` has the - tool-call sections removed. - """ - # MiniMax-M3 ns_token format (checked before DSML: both use , - # but MiniMax names params by tag and prefixes every tag with ]<]minimax[>[) - if _is_minimax(text): - return _parse_minimax(text, tools) - - # DeepSeek-V4 DSML format - if _is_dsml(text): - return _parse_dsml(text, tools) - - # GLM / format (checked before Qwen: both use - # , but GLM never emits the Qwen '(.*?)<\|tool_calls_section_end\|>", - text, - flags=re.DOTALL, - ) - if not section_match: - # Check for unclosed section - unclosed = re.search( - r"<\|tool_calls_section_begin\|>(.*?)$", text, flags=re.DOTALL - ) - if unclosed: - content = text[: unclosed.start()] - tool_calls = _parse_tool_call_entries(unclosed.group(1)) - return content.strip(), tool_calls - return text, [] - - content = text[: section_match.start()] - tool_calls = _parse_tool_call_entries(section_match.group(1)) - - return content.strip(), tool_calls - - -def _parse_tool_call_entries(section_text: str) -> List[ToolCall]: - """Parse individual tool call entries from the section content.""" - tool_calls = [] - pattern = re.compile( - r"<\|tool_call_begin\|>" - r"functions\.(\w+):(\d+)" - r"<\|tool_call_argument_begin\|>" - r"(.*?)" - r"<\|tool_call_end\|>", - re.DOTALL, - ) - for match in pattern.finditer(section_text): - name = match.group(1) - index = match.group(2) - arguments = match.group(3).strip() - tool_id = f"functions.{name}:{index}" - tool_calls.append( - ToolCall( - id=tool_id, - type="function", - function={"name": name, "arguments": arguments}, - ) - ) - return tool_calls - - -@dataclass -class ToolCallStreamParser: - """Stateful streaming parser for tool calls (Kimi tokens or Qwen3 XML). - - Processes text chunks and emits structured events: - - ("content", text) — regular content before tool calls - - ("tool_call_start", {"index": N, "id": ..., "function": {"name": ..., "arguments": ""}}) - - ("tool_call_args", {"index": N, "function": {"arguments": chunk}}) - - ("tool_call_end", None) — all tool calls complete - - The wire format is auto-detected from the first chunks. For the Qwen3 XML - format content is streamed normally and the ```` block is buffered - and parsed when complete (robust against partial-XML streaming edge cases); - ``tools`` enables JSON-Schema type coercion of parameter values. - - Kimi states: - 0 = normal content (no tool call tokens seen) - 1 = inside tool_calls_section (buffering) - 2 = done (after tool_calls_section_end) - """ - - state: int = 0 - buf: str = "" - current_index: int = 0 - _emitted_calls: int = 0 - tools: Optional[list] = None - fmt: Optional[str] = None # None|kimi|qwen|dsml|glm|minimax - - def process(self, text: str) -> list: - """Process a text chunk and return list of (event_type, data) tuples.""" - if self.fmt is None: - self.buf += text - if _MINIMAX_NS in self.buf: - self.fmt = "minimax" - elif _is_dsml(self.buf): - self.fmt = "dsml" - elif "" in self.buf: - self.fmt = "glm" - elif _QWEN_TOOL_PREFIX in self.buf: - self.fmt = "qwen" - elif "" in self.buf: - # '' seen but neither '' (GLM) yet. A no-arg GLM call is complete once the - # closing tag arrives; otherwise wait for the sub-marker. - if "" in self.buf: - self.fmt = "glm" - else: - return [] - elif "<|tool_calls_section_begin|>" in self.buf: - self.fmt = "kimi" - elif "<" not in self.buf and len(self.buf) > 8: - # No markup possible yet; emit accumulated content and stay undecided. - out = [("content", self.buf)] - self.buf = "" - return out - else: - return [] - # Format decided: replay the accumulated buffer through the handler. - text, self.buf = self.buf, "" - - if self.fmt == "minimax": - return self._process_minimax(text) - if self.fmt == "dsml": - return self._process_dsml(text) - if self.fmt == "glm": - return self._process_glm(text) - if self.fmt == "qwen": - return self._process_qwen(text) - return self._process_kimi(text) - - # -- MiniMax-M3 ns_token ------------------------------------------------ - def _process_minimax(self, text: str) -> list: - results: list = [] - self.buf += text - if self.state == 0: - markers = [ - i - for i in (self.buf.find(_MINIMAX_NS), self.buf.find("")) - if i != -1 - ] - if markers: - m = min(markers) - before = self.buf[:m] - if before: - results.append(("content", before)) - self.buf = self.buf[m:] - self.state = 1 - else: - cut = self.buf.rfind("<") - cut = max(cut, self.buf.rfind("]")) # ns_token starts with ']' - if cut == -1: - if self.buf: - results.append(("content", self.buf)) - self.buf = "" - elif cut > 0: - results.append(("content", self.buf[:cut])) - self.buf = self.buf[cut:] - return results - - def _flush_minimax(self) -> list: - results: list = [] - if self.state == 0: - if self.buf: - results.append(("content", self.buf)) - self.buf = "" - return results - _content, tool_calls = _parse_minimax(self.buf, self.tools) - self.buf = "" - for tc in tool_calls: - results.append( - ( - "tool_call_start", - { - "index": self.current_index, - "id": tc.id, - "type": "function", - "function": {"name": tc.function["name"], "arguments": ""}, - }, - ) - ) - results.append( - ( - "tool_call_args", - { - "index": self.current_index, - "function": {"arguments": tc.function["arguments"]}, - }, - ) - ) - self.current_index += 1 - self._emitted_calls += 1 - if self._emitted_calls > 0: - results.append(("tool_call_end", None)) - return results - - # -- DeepSeek-V4 DSML --------------------------------------------------- - def _process_dsml(self, text: str) -> list: - results: list = [] - self.buf += text - if self.state == 0: - m = _dsml_start(self.buf) - if m != -1: - before = self.buf[:m] - if before: - results.append(("content", before)) - self.buf = self.buf[m:] - self.state = 1 - else: - # Emit content but hold back a possible partial '<...' marker tail. - cut = self.buf.rfind("<") - if cut == -1: - if self.buf: - results.append(("content", self.buf)) - self.buf = "" - elif cut > 0: - results.append(("content", self.buf[:cut])) - self.buf = self.buf[cut:] - return results - - def _flush_dsml(self) -> list: - results: list = [] - if self.state == 0: - if self.buf: - results.append(("content", self.buf)) - self.buf = "" - return results - _content, tool_calls = _parse_dsml(self.buf, self.tools) - self.buf = "" - for tc in tool_calls: - results.append( - ( - "tool_call_start", - { - "index": self.current_index, - "id": tc.id, - "type": "function", - "function": {"name": tc.function["name"], "arguments": ""}, - }, - ) - ) - results.append( - ( - "tool_call_args", - { - "index": self.current_index, - "function": {"arguments": tc.function["arguments"]}, - }, - ) - ) - self.current_index += 1 - self._emitted_calls += 1 - if self._emitted_calls > 0: - results.append(("tool_call_end", None)) - return results - - # -- Qwen3 XML ---------------------------------------------------------- - # -- GLM / ----------------------------------------- - def _process_glm(self, text: str) -> list: - results: list = [] - self.buf += text - if self.state == 0: - m = self.buf.find("") - if m != -1: - before = self.buf[:m] - if before: - results.append(("content", before)) - self.buf = self.buf[m:] - self.state = 1 - else: - # Emit content but hold back a possible partial '<...' marker tail. - cut = self.buf.rfind("<") - if cut == -1: - if self.buf: - results.append(("content", self.buf)) - self.buf = "" - elif cut > 0: - results.append(("content", self.buf[:cut])) - self.buf = self.buf[cut:] - return results - - def _flush_glm(self) -> list: - results: list = [] - if self.state == 0: - if self.buf: - results.append(("content", self.buf)) - self.buf = "" - return results - _content, tool_calls = _parse_glm(self.buf, self.tools) - self.buf = "" - for tc in tool_calls: - results.append( - ( - "tool_call_start", - { - "index": self.current_index, - "id": tc.id, - "type": "function", - "function": {"name": tc.function["name"], "arguments": ""}, - }, - ) - ) - results.append( - ( - "tool_call_args", - { - "index": self.current_index, - "function": {"arguments": tc.function["arguments"]}, - }, - ) - ) - self.current_index += 1 - self._emitted_calls += 1 - if self._emitted_calls > 0: - results.append(("tool_call_end", None)) - return results - - # -- Qwen3 XML ---------------------------------------------------------- - def _process_qwen(self, text: str) -> list: - results: list = [] - self.buf += text - if self.state == 0: - markers = [ - i - for i in ( - self.buf.find(""), - self.buf.find(_QWEN_TOOL_PREFIX), - ) - if i != -1 - ] - if markers: - m = min(markers) - before = self.buf[:m] - if before: - results.append(("content", before)) - self.buf = self.buf[m:] - self.state = 1 - else: - # Emit content but hold back a possible partial '<...' marker tail. - cut = self.buf.rfind("<") - if cut == -1: - if self.buf: - results.append(("content", self.buf)) - self.buf = "" - elif cut > 0: - results.append(("content", self.buf[:cut])) - self.buf = self.buf[cut:] - return results - - def _flush_qwen(self) -> list: - results: list = [] - if self.state == 0: - if self.buf: - results.append(("content", self.buf)) - self.buf = "" - return results - # state 1: parse the complete (or trailing) tool-call block. - _content, tool_calls = _parse_qwen_xml(self.buf, self.tools) - self.buf = "" - for tc in tool_calls: - tc.id = _unique_tool_call_id() - results.append( - ( - "tool_call_start", - { - "index": self.current_index, - "id": tc.id, - "type": "function", - "function": {"name": tc.function["name"], "arguments": ""}, - }, - ) - ) - results.append( - ( - "tool_call_args", - { - "index": self.current_index, - "function": {"arguments": tc.function["arguments"]}, - }, - ) - ) - self.current_index += 1 - self._emitted_calls += 1 - if self._emitted_calls > 0: - results.append(("tool_call_end", None)) - return results - - # -- Kimi tokens -------------------------------------------------------- - def _process_kimi(self, text: str) -> list: - results = [] - - if self.state == 0: - self.buf += text - if "<|tool_calls_section_begin|>" in self.buf: - before = self.buf.split("<|tool_calls_section_begin|>")[0] - if before: - results.append(("content", before)) - self.state = 1 - self.buf = self.buf.split("<|tool_calls_section_begin|>", 1)[1] - results.extend(self._process_buffer()) - elif "<|tool" not in self.buf and len(self.buf) > 30: - results.append(("content", self.buf)) - self.buf = "" - - elif self.state == 1: - self.buf += text - if "<|tool_calls_section_end|>" in self.buf: - remaining = self.buf.split("<|tool_calls_section_end|>")[0] - self.buf = remaining - results.extend(self._process_buffer()) - results.append(("tool_call_end", None)) - self.state = 2 - self.buf = "" - else: - results.extend(self._process_buffer()) - - return results - - def _process_buffer(self) -> list: - """Extract complete tool call entries from the buffer.""" - results = [] - while "<|tool_call_begin|>" in self.buf and "<|tool_call_end|>" in self.buf: - match = re.search( - r"<\|tool_call_begin\|>" - r"functions\.(\w+):(\d+)" - r"<\|tool_call_argument_begin\|>" - r"(.*?)" - r"<\|tool_call_end\|>", - self.buf, - re.DOTALL, - ) - if not match: - break - - name = match.group(1) - index = int(match.group(2)) - arguments = match.group(3).strip() - - tool_id = f"functions.{name}:{index}" - results.append( - ( - "tool_call_start", - { - "index": index, - "id": tool_id, - "type": "function", - "function": {"name": name, "arguments": ""}, - }, - ) - ) - if arguments: - results.append( - ( - "tool_call_args", - {"index": index, "function": {"arguments": arguments}}, - ) - ) - - self.buf = self.buf[match.end() :] - self._emitted_calls += 1 - - return results - - def flush(self) -> list: - """Flush remaining buffer content.""" - if self.fmt == "minimax": - return self._flush_minimax() - if self.fmt == "dsml": - return self._flush_dsml() - if self.fmt == "glm": - return self._flush_glm() - if self.fmt == "qwen": - return self._flush_qwen() - results = [] - if self.state == 0 and self.buf: - results.append(("content", self.buf)) - self.buf = "" - elif self.state == 1: - results.extend(self._process_buffer()) - if self._emitted_calls > 0: - results.append(("tool_call_end", None)) - elif self.fmt is None and self.buf: - # Undecided at EOS: no tool markers ever appeared -> plain content. - results.append(("content", self.buf)) - self.buf = "" - return results diff --git a/atom/entrypoints/openai/tool_parser/__init__.py b/atom/entrypoints/openai/tool_parser/__init__.py new file mode 100644 index 0000000000..6c7b286241 --- /dev/null +++ b/atom/entrypoints/openai/tool_parser/__init__.py @@ -0,0 +1,61 @@ +# SPDX-License-Identifier: MIT +# Copyright (C) 2024-2026, Advanced Micro Devices, Inc. All rights reserved. + +"""Tool call parsing for models that emit tool calls in their text output. + +Five on-the-wire formats are auto-detected and normalized into the OpenAI +``tool_calls`` structure: + +============================== =============================================== +Module Format +============================== =============================================== +`kimi_tool_parser` Kimi-K2 ``<|tool_call_begin|>`` special tokens +`qwen3_tool_parser` Qwen3 (qwen3_coder / qwen3_xml) ```` markup +`glm_tool_parser` GLM-4.5/4.6/5.x ````/```` +`minimax_tool_parser` MiniMax-M3 ``]<]minimax[>[``-prefixed tags +============================== =============================================== + +Formats other than Kimi carry no value types on the wire, so when the request's +``tools`` schema is supplied each parameter is coerced to its declared +JSON-Schema type; otherwise it is left as a string. + +Two entry points, both format-agnostic: + +- :func:`parse_tool_calls` — a complete output -> ``(content, [ToolCall])`` +- :class:`ToolCallStreamParser` — chunks -> ``(event_type, data)`` tuples + +OpenAI format:: + + {"tool_calls": [{"id": "call_0", "type": "function", + "function": {"name": "NAME", "arguments": "ARGS_JSON"}}]} + +To add a format: implement :class:`~.tool_parser.ToolCallParser` (or, if it +buffers from a start marker like most do, +:class:`~.tool_parser.BufferedMarkerParser`) in its own +``_tool_parser.py``, then register it in :mod:`.registry` — in both +``_DETECT_ORDER`` and ``sniff_stream``, whose ordering constraints are +documented there. +""" + +from .deepseekv4_tool_parser import DsmlParser +from .glm_tool_parser import GlmParser +from .kimi_tool_parser import KimiParser +from .minimax_tool_parser import MiniMaxParser +from .qwen3_tool_parser import QwenXmlParser +from .registry import parse_tool_calls +from .stream import ToolCallStreamParser +from .tool_parser import BufferedMarkerParser, ToolCall, ToolCallParser + +__all__ = [ + "BufferedMarkerParser", + "DsmlParser", + "GlmParser", + "KimiParser", + "MiniMaxParser", + "QwenXmlParser", + "ToolCall", + "ToolCallParser", + "ToolCallStreamParser", + "parse_tool_calls", +] diff --git a/atom/entrypoints/openai/tool_parser/deepseekv4_tool_parser.py b/atom/entrypoints/openai/tool_parser/deepseekv4_tool_parser.py new file mode 100644 index 0000000000..e263dc5eff --- /dev/null +++ b/atom/entrypoints/openai/tool_parser/deepseekv4_tool_parser.py @@ -0,0 +1,211 @@ +# SPDX-License-Identifier: MIT +# Copyright (C) 2024-2026, Advanced Micro Devices, Inc. All rights reserved. + +"""DeepSeek-V4 DSML tool-call format:: + + <|DSML|tool_calls> + <|DSML|invoke name="NAME"> + <|DSML|parameter name="PNAME" string="true|false">VALUE + ... + + + +``string="true"`` -> value is a raw string; ``string="false"`` -> value is JSON. +DeepSeek-V4-Flash occasionally malforms this (singular ``tool_call``, a missing +``invoke`` wrapper, or params without ``string=``); the parser recovers those +best-effort: it infers a dropped tool name from the parameter signature vs the +request's ``tools`` and infers a missing value type from the schema / JSON. +""" + +import json +import re +from typing import Any, ClassVar, Dict, List, Optional, Tuple + +from .schema import build_param_types, coerce_param_value +from .tool_parser import BufferedMarkerParser, ToolCall, unique_tool_call_id + +_DSML = "|DSML|" +# The model often DROPS the ``|DSML|`` marker and emits bare +# ````/````/```` tags, so the marker +# is matched OPTIONALLY everywhere. +_OPT = r"(?:" + re.escape(_DSML) + r")?" # optional |DSML| prefix +_PARAM_RE = re.compile( + r"<" + _OPT + r'parameter\s+name="(.*?)"(?:\s+string="(true|false)")?\s*>' + r"(.*?)", + re.DOTALL, +) +# Long-form `...` OR self-closing `` +# (the zero-arg shape; group(2) is None for self-closing). Matches SGLang's V4 +# detector, which accepts both. +_INVOKE_RE = re.compile( + r"<" + _OPT + r'invoke\s+name="(.*?)"\s*(?:/>|>(.*?))", + re.DOTALL, +) + +# Canonical param key -> emitted synonyms the model uses interchangeably. Only +# applied when the canonical key is in the tool's declared schema and the model +# used a synonym instead (e.g. Codex's exec_command wants `cmd`, but the model +# habitually emits `command` -> "missing field cmd" retry loop). Scoped to the +# schema so it never renames a key a tool legitimately declares. +_KEY_ALIASES: Dict[str, Tuple[str, ...]] = { + "cmd": ("command",), +} + + +def _unwrap_wrapper_args(args: Any, allowed: set) -> Any: + """Strip spurious ``{"arguments": {...}}`` / ``{"input": {...}}`` envelopes. + + Non-tuned models (DeepSeek-V4-Pro) frequently wrap the real args in an extra + ``arguments``/``input`` object — sometimes nested 2-3 deep, or stringified — + so a call meant as ``{"cmd": "ls"}`` arrives as ``{"arguments": {"cmd": + "ls"}}`` and the client (Codex) rejects it ("missing field cmd"). Recursively + unwrap while the sole key is a wrapper that is NOT itself a declared param of + the tool. Mirrors vLLM's ``_unwrap_wrapper_args`` (deepseek_v4.py).""" + for _ in range(4): # bounded against pathological nesting + if not (isinstance(args, dict) and len(args) == 1): + break + ((k, v),) = args.items() + if k not in ("arguments", "input"): + break + if allowed and k in allowed: + break # this tool really has a param named arguments/input + if isinstance(v, str): + try: + v = json.loads(v) + except Exception: + break + if not isinstance(v, dict): + break + args = v + return args + + +def _apply_key_aliases(args: Any, allowed: set) -> Any: + if not (isinstance(args, dict) and allowed): + return args + for canon, syns in _KEY_ALIASES.items(): + if canon in allowed and canon not in args: + for s in syns: + if s in args: + args[canon] = args.pop(s) + break + return args + + +def _coerce(value: str, string_attr: Optional[str], ptype: Any) -> Any: + """Decode one ```` body. + + Deliberately not :func:`~.schema.coerce_json_or_raw`: on a JSON-decode miss + this falls back to ``value.strip()`` where that one falls back to + ``value.strip("\\n")``, which differs for values with surrounding spaces. + """ + if string_attr == "true": + return value + if string_attr == "false": + try: + return json.loads(value) + except Exception: + return value + # attr absent -> use declared schema type if known, else infer via JSON. + if ptype is not None: + return coerce_param_value(value, ptype) + v = value.strip() + try: + return json.loads(v) + except Exception: + return v + + +def _infer_name( + arg_names: set, param_types: Dict[str, Dict[str, Any]] +) -> Optional[str]: + """Pick the request tool whose parameter set best matches ``arg_names``.""" + best, best_score = None, -1e9 + for name, props in param_types.items(): + p = set(props) + if not p: + continue + score = len(p & arg_names) - 0.1 * len(p ^ arg_names) + if score > best_score: + best_score, best = score, name + return best + + +class DsmlParser(BufferedMarkerParser): + NAME: ClassVar[str] = "dsml" + # Region-start markers, both marked and marker-less variants. + START_MARKERS: ClassVar[Tuple[str, ...]] = ( + "<" + _DSML + "tool_call", # marked (covers tool_call / tool_calls) + "<" + _DSML + "invoke", # marked invoke + "", # marker-less section open + ) + + # detect() is inherited: any start marker present means DSML. + + @classmethod + def parse(cls, text: str, tools: Optional[list]) -> Tuple[str, List[ToolCall]]: + """Parse DeepSeek-V4 DSML tool calls; return (leading_content, tool_calls).""" + param_types = build_param_types(tools) + start = cls.find_start(text) + if start == -1: + return text.strip(), [] + content = text[:start] + region = text[start:] + + calls: List[Tuple[str, Dict[str, Any]]] = [] + invokes = list(_INVOKE_RE.finditer(region)) + if invokes: + for m in invokes: + name = m.group(1) + body = m.group(2) or "" # None for self-closing + types = param_types.get(name, {}) + args: Dict[str, Any] = { + pm.group(1): _coerce( + pm.group(3), pm.group(2), types.get(pm.group(1)) + ) + for pm in _PARAM_RE.finditer(body) + } + # Direct-JSON parameter body (DSML "Format 2", also accepted by + # vLLM/SGLang): ` { "k": "v" } ` with no + # tags. Falls through here with empty args; recover them. + if not args: + stripped = body.strip() + if stripped.startswith("{"): + try: + parsed = json.loads(stripped) + if isinstance(parsed, dict): + args = parsed + except Exception: + pass + args = _unwrap_wrapper_args(args, set(types)) + args = _apply_key_aliases(args, set(types)) + calls.append((name, args)) + else: + # malformed: no complete invoke wrapper -> collect params, infer tool name + raw = { + pm.group(1): (pm.group(3), pm.group(2)) + for pm in _PARAM_RE.finditer(region) + } + if raw: + name = _infer_name(set(raw), param_types) or "unknown" + types = param_types.get(name, {}) + args = {k: _coerce(v, s, types.get(k)) for k, (v, s) in raw.items()} + args = _unwrap_wrapper_args(args, set(types)) + args = _apply_key_aliases(args, set(types)) + calls.append((name, args)) + + tool_calls = [ + ToolCall( + id=unique_tool_call_id(), + type="function", + function={ + "name": name, + "arguments": json.dumps(args, ensure_ascii=False), + }, + ) + for name, args in calls + ] + if _DSML in content: # scrub any stray marker fragment + content = content.split("<" + _DSML, 1)[0] + return content.strip(), tool_calls diff --git a/atom/entrypoints/openai/tool_parser/glm_tool_parser.py b/atom/entrypoints/openai/tool_parser/glm_tool_parser.py new file mode 100644 index 0000000000..7077a98ffe --- /dev/null +++ b/atom/entrypoints/openai/tool_parser/glm_tool_parser.py @@ -0,0 +1,78 @@ +# SPDX-License-Identifier: MIT +# Copyright (C) 2024-2026, Advanced Micro Devices, Inc. All rights reserved. + +"""GLM-4.5 / 4.6 / 5.x tool-call format:: + + NAME + K1V1 + K2V2 + ... + +The function name follows the opening tag directly (no ``(.*?)|(.*)$", re.DOTALL) +_ARG_RE = re.compile( + r"(.*?)\s*" + r"(.*?)(?:|(?=)|(?=)|$)", + re.DOTALL, +) + + +class GlmParser(BufferedMarkerParser): + NAME: ClassVar[str] = "glm" + START_MARKERS: ClassVar[Tuple[str, ...]] = ("",) + + @classmethod + def detect(cls, text: str) -> bool: + """Detect the GLM ``...`` format (never Qwen/DSML).""" + if QWEN_TOOL_PREFIX in text: # ' Qwen, not GLM + return False + return "" in text or "" in text + + @classmethod + def parse(cls, text: str, tools: Optional[list]) -> Tuple[str, List[ToolCall]]: + """Parse GLM tool calls; return (leading_content, tool_calls).""" + param_types = build_param_types(tools) + start = text.find("") + if start == -1: + return text.strip(), [] + content = text[:start] + tool_calls: List[ToolCall] = [] + for m in _TOOLCALL_RE.finditer(text): + body = m.group(1) if m.group(1) is not None else m.group(2) + if not body: + continue + ak = body.find("") + name = (body if ak == -1 else body[:ak]).strip() + if not name: + continue + types = param_types.get(name, {}) + args: Dict[str, Any] = {} + for pm in _ARG_RE.finditer(body): + k = pm.group(1).strip() + if k: + args[k] = coerce_json_or_raw(pm.group(2), types.get(k)) + tool_calls.append( + ToolCall( + id=unique_tool_call_id(), + type="function", + function={ + "name": name, + "arguments": json.dumps(args, ensure_ascii=False), + }, + ) + ) + return content.strip(), tool_calls diff --git a/atom/entrypoints/openai/tool_parser/kimi_tool_parser.py b/atom/entrypoints/openai/tool_parser/kimi_tool_parser.py new file mode 100644 index 0000000000..b820ba568d --- /dev/null +++ b/atom/entrypoints/openai/tool_parser/kimi_tool_parser.py @@ -0,0 +1,159 @@ +# SPDX-License-Identifier: MIT +# Copyright (C) 2024-2026, Advanced Micro Devices, Inc. All rights reserved. + +"""Kimi-K2 special-token tool-call format:: + + <|tool_calls_section_begin|> + <|tool_call_begin|>functions.NAME:INDEX<|tool_call_argument_begin|>ARGS_JSON<|tool_call_end|> + <|tool_calls_section_end|> + +Unlike the XML-ish formats this one is self-delimiting, so entries can be +emitted as soon as their ``<|tool_call_end|>`` arrives rather than buffering the +whole block. It therefore implements streaming itself instead of inheriting +:class:`~.tool_parser.BufferedMarkerParser`. + +Arguments are already JSON on the wire, so no schema coercion is applied and +``tools`` is unused. The call id is the model's own ``functions.NAME:INDEX`` +rather than a random one, and ``index`` comes from the wire too. +""" + +import re +from typing import ClassVar, List, Optional, Tuple + +from .tool_parser import ToolCall, ToolCallParser + +KIMI_SECTION_BEGIN = "<|tool_calls_section_begin|>" +KIMI_SECTION_END = "<|tool_calls_section_end|>" + +_SECTION_RE = re.compile( + re.escape(KIMI_SECTION_BEGIN) + r"(.*?)" + re.escape(KIMI_SECTION_END), + re.DOTALL, +) +_UNCLOSED_RE = re.compile(re.escape(KIMI_SECTION_BEGIN) + r"(.*?)$", re.DOTALL) +_ENTRY_RE = re.compile( + r"<\|tool_call_begin\|>" + r"functions\.(\w+):(\d+)" + r"<\|tool_call_argument_begin\|>" + r"(.*?)" + r"<\|tool_call_end\|>", + re.DOTALL, +) + + +def _parse_entries(section_text: str) -> List[ToolCall]: + """Parse individual tool call entries from the section content.""" + tool_calls = [] + for match in _ENTRY_RE.finditer(section_text): + name = match.group(1) + index = match.group(2) + arguments = match.group(3).strip() + tool_id = f"functions.{name}:{index}" + tool_calls.append( + ToolCall( + id=tool_id, + type="function", + function={"name": name, "arguments": arguments}, + ) + ) + return tool_calls + + +class KimiParser(ToolCallParser): + """States: 0 = plain content, 1 = inside section, 2 = section closed.""" + + NAME: ClassVar[str] = "kimi" + + @classmethod + def detect(cls, text: str) -> bool: + return KIMI_SECTION_BEGIN in text + + @classmethod + def parse(cls, text: str, tools: Optional[list]) -> Tuple[str, List[ToolCall]]: + section_match = _SECTION_RE.search(text) + if not section_match: + # Unclosed section: the model was cut off mid-block; salvage whatever + # complete entries it managed to emit. + unclosed = _UNCLOSED_RE.search(text) + if unclosed: + content = text[: unclosed.start()] + return content.strip(), _parse_entries(unclosed.group(1)) + return text, [] + content = text[: section_match.start()] + return content.strip(), _parse_entries(section_match.group(1)) + + def process(self, text: str) -> list: + results: list = [] + + if self.state == 0: + self.buf += text + if KIMI_SECTION_BEGIN in self.buf: + before = self.buf.split(KIMI_SECTION_BEGIN)[0] + if before: + results.append(("content", before)) + self.state = 1 + self.buf = self.buf.split(KIMI_SECTION_BEGIN, 1)[1] + results.extend(self._drain_entries()) + elif "<|tool" not in self.buf and len(self.buf) > 30: + # No partial special token possible yet -> safe to release. + results.append(("content", self.buf)) + self.buf = "" + + elif self.state == 1: + self.buf += text + if KIMI_SECTION_END in self.buf: + self.buf = self.buf.split(KIMI_SECTION_END)[0] + results.extend(self._drain_entries()) + results.append(("tool_call_end", None)) + self.state = 2 + self.buf = "" + else: + results.extend(self._drain_entries()) + + return results + + def _drain_entries(self) -> list: + """Emit every complete tool-call entry sitting in the buffer.""" + results: list = [] + while "<|tool_call_begin|>" in self.buf and "<|tool_call_end|>" in self.buf: + match = _ENTRY_RE.search(self.buf) + if not match: + break + + name = match.group(1) + index = int(match.group(2)) + arguments = match.group(3).strip() + + results.append( + ( + "tool_call_start", + { + "index": index, + "id": f"functions.{name}:{index}", + "type": "function", + "function": {"name": name, "arguments": ""}, + }, + ) + ) + if arguments: + results.append( + ( + "tool_call_args", + {"index": index, "function": {"arguments": arguments}}, + ) + ) + + self.buf = self.buf[match.end() :] + self.emitted_calls += 1 + + return results + + def flush(self) -> list: + results: list = [] + if self.state == 0 and self.buf: + results.append(("content", self.buf)) + self.buf = "" + elif self.state == 1: + results.extend(self._drain_entries()) + if self.emitted_calls > 0: + results.append(("tool_call_end", None)) + return results diff --git a/atom/entrypoints/openai/tool_parser/minimax_tool_parser.py b/atom/entrypoints/openai/tool_parser/minimax_tool_parser.py new file mode 100644 index 0000000000..81df815dbe --- /dev/null +++ b/atom/entrypoints/openai/tool_parser/minimax_tool_parser.py @@ -0,0 +1,79 @@ +# SPDX-License-Identifier: MIT +# Copyright (C) 2024-2026, Advanced Micro Devices, Inc. All rights reserved. + +"""MiniMax-M3 tool-call format. + +Every tag is prefixed by the ns_token ``]<]minimax[>[``:: + + ]<]minimax[>[ + ]<]minimax[>[ + ]<]minimax[>[value]<]minimax[>[ + ... + ]<]minimax[>[ + ]<]minimax[>[ + +Unlike DSML, parameters are named by the TAG itself (``Paris``), +not a ``name="..."`` attribute. Strip the ns_token first, then parse +/ pairs. Values: schema type wins, else JSON, else raw string. +""" + +import json +import re +from typing import Any, ClassVar, Dict, List, Optional, Tuple + +from .schema import build_param_types, coerce_json_or_raw +from .tool_parser import BufferedMarkerParser, ToolCall, unique_tool_call_id + +MINIMAX_NS = "]<]minimax[>[" + +_INVOKE_RE = re.compile( + r'(.*?)|(.*)$', + re.DOTALL, +) +_PARAM_RE = re.compile(r"<([\w-]+)>(.*?)", re.DOTALL) + + +class MiniMaxParser(BufferedMarkerParser): + NAME: ClassVar[str] = "minimax" + START_MARKERS: ClassVar[Tuple[str, ...]] = (MINIMAX_NS, "") + # The ns_token starts with ']', so a trailing ']' may be a partial marker. + HOLDBACK_CHARS: ClassVar[Tuple[str, ...]] = ("<", "]") + + @classmethod + def detect(cls, text: str) -> bool: + """Detect the MiniMax-M3 ns_token tool-call format.""" + return MINIMAX_NS in text + + @classmethod + def parse(cls, text: str, tools: Optional[list]) -> Tuple[str, List[ToolCall]]: + """Parse MiniMax-M3 tool calls; return (leading_content, tool_calls).""" + param_types = build_param_types(tools) + clean = text.replace(MINIMAX_NS, "") + tc = clean.find("") + content = clean[:tc] if tc > 0 else ("" if tc == 0 else clean) + tool_calls: List[ToolCall] = [] + for m in _INVOKE_RE.finditer(clean): + name = m.group(1) if m.group(1) is not None else m.group(3) + body = m.group(2) if m.group(2) is not None else (m.group(4) or "") + if not name: + continue + name = name.strip() + types = param_types.get(name, {}) + args: Dict[str, Any] = {} + for pm in _PARAM_RE.finditer(body): + k = pm.group(1).strip() + if k: + args[k] = coerce_json_or_raw(pm.group(2), types.get(k)) + tool_calls.append( + ToolCall( + id=unique_tool_call_id(), + type="function", + function={ + "name": name, + "arguments": json.dumps(args, ensure_ascii=False), + }, + ) + ) + for mk in ("", ""): + content = content.replace(mk, "") + return content.strip(), tool_calls diff --git a/atom/entrypoints/openai/tool_parser/qwen3_tool_parser.py b/atom/entrypoints/openai/tool_parser/qwen3_tool_parser.py new file mode 100644 index 0000000000..bfc86b3b4b --- /dev/null +++ b/atom/entrypoints/openai/tool_parser/qwen3_tool_parser.py @@ -0,0 +1,92 @@ +# SPDX-License-Identifier: MIT +# Copyright (C) 2024-2026, Advanced Micro Devices, Inc. All rights reserved. + +"""Qwen3 (qwen3_coder / qwen3_xml) XML tool-call format:: + + + + VALUE + ... + + + +The XML carries no value types, so parameters are coerced against the request's +``tools`` schema when supplied. Mirrors the qwen3_coder/qwen3_xml parsers in +vLLM and SGLang. +""" + +import json +import re +from typing import Any, ClassVar, Dict, List, Optional, Tuple + +from .kimi_tool_parser import KIMI_SECTION_BEGIN +from .schema import build_param_types, coerce_param_value +from .tool_parser import BufferedMarkerParser, ToolCall, unique_tool_call_id + +# Also read by GlmParser.detect: ' +# apart from GLM's identically-named tag. +QWEN_TOOL_PREFIX = "||(?=)|$)", + re.DOTALL, +) + + +def _parse_function( + fn_text: str, param_types: Dict[str, Dict[str, Any]] +) -> Optional[ToolCall]: + """Parse the inside of one ``...`` block into a ToolCall.""" + gt = fn_text.find(">") + if gt == -1: + return None + name = fn_text[:gt].strip() + if not name: + return None + body = fn_text[gt + 1 :] + types = param_types.get(name, {}) + args: Dict[str, Any] = {} + for pm in _PARAM_RE.finditer(body): + seg = pm.group(1) + if seg is None: + continue + pgt = seg.find(">") + if pgt == -1: + continue + pname = seg[:pgt].strip() + pval = seg[pgt + 1 :] + if pname: + args[pname] = coerce_param_value(pval, types.get(pname)) + return ToolCall( + id=unique_tool_call_id(), + type="function", + function={"name": name, "arguments": json.dumps(args, ensure_ascii=False)}, + ) + + +class QwenXmlParser(BufferedMarkerParser): + NAME: ClassVar[str] = "qwen" + START_MARKERS: ClassVar[Tuple[str, ...]] = ("", QWEN_TOOL_PREFIX) + + @classmethod + def detect(cls, text: str) -> bool: + """Detect the Qwen3 XML format (and not the Kimi token format).""" + return QWEN_TOOL_PREFIX in text and KIMI_SECTION_BEGIN not in text + + @classmethod + def parse(cls, text: str, tools: Optional[list]) -> Tuple[str, List[ToolCall]]: + """Parse Qwen3 XML tool calls; return (leading_content, tool_calls).""" + param_types = build_param_types(tools) + # Content precedes the first tool marker. + start = cls.find_start(text) + content = text[:start] if start != -1 else text + tool_calls: List[ToolCall] = [] + for fm in _FUNCTION_RE.finditer(text): + fn_text = fm.group(1) if fm.group(1) is not None else fm.group(2) + if not fn_text: + continue + tc = _parse_function(fn_text, param_types) + if tc is not None: + tool_calls.append(tc) + return content.strip(), tool_calls diff --git a/atom/entrypoints/openai/tool_parser/registry.py b/atom/entrypoints/openai/tool_parser/registry.py new file mode 100644 index 0000000000..fd4d048ffc --- /dev/null +++ b/atom/entrypoints/openai/tool_parser/registry.py @@ -0,0 +1,104 @@ +# SPDX-License-Identifier: MIT +# Copyright (C) 2024-2026, Advanced Micro Devices, Inc. All rights reserved. + +"""Format detection. + +The wire format is sniffed from the model's own output rather than configured, +so detection order is load-bearing: several formats share tags and are only +told apart by a discriminator that a later entry would also match. The order +below is the single place that ordering is expressed — do not reorder without +re-reading the notes on each entry. +""" + +from typing import List, Optional, Tuple, Type + +from .deepseekv4_tool_parser import DsmlParser +from .glm_tool_parser import GlmParser +from .kimi_tool_parser import KIMI_SECTION_BEGIN, KimiParser +from .minimax_tool_parser import MINIMAX_NS, MiniMaxParser +from .qwen3_tool_parser import QWEN_TOOL_PREFIX, QwenXmlParser +from .tool_parser import ToolCall, ToolCallParser + +# Checked in order on a COMPLETE output. Kimi is not listed: it is the terminal +# fallback, because its parse() also defines the "no tool calls at all" result. +# +# MiniMax before DSML — both use ``; MiniMax additionally +# prefixes every tag with the ns_token. +# GLM before Qwen — both use ``; GLM never emits ` Tuple[str, List[ToolCall]]: + """Parse tool calls from a complete model output. + + Args: + text: Raw model output that may contain tool calls. + tools: Optional request tool definitions; used to type-coerce parameter + values to their declared JSON-Schema types. + + Returns: + Tuple of (content_text, list_of_tool_calls). ``content_text`` has the + tool-call sections removed. + """ + for parser in _DETECT_ORDER: + if parser.detect(text): + return parser.parse(text, tools) + # Kimi is terminal: when it finds no section either, it returns the text + # unchanged. Note that path does NOT strip, unlike every format that did + # match — preserved as-is, callers rely on plain content surviving verbatim. + return KimiParser.parse(text, tools) + + +# -- streaming sniff -------------------------------------------------------- +# +# Deciding on a PREFIX is strictly harder than on a complete output: a format's +# discriminator may not have arrived yet. These two sentinels say "cannot decide +# from what I have": + +# Enough plain text to be sure no marker is starting -> release it as content +# and stay undecided. +EMIT_CONTENT = object() +# Might still become a tool call -> keep buffering, emit nothing. +WAIT = object() + + +def sniff_stream(buf: str): + """Pick a parser from a partial stream, or return EMIT_CONTENT / WAIT. + + Deliberately NOT the same rules as the per-parser ``detect``: on a prefix, + GLM is only accepted on the unambiguous ```` (a bare ```` + could still turn out to be Qwen once ``" in buf: + return GlmParser + if QWEN_TOOL_PREFIX in buf: + return QwenXmlParser + if "" in buf: + # '' seen but neither '' + # (GLM) yet. A no-arg GLM call is complete once the closing tag arrives; + # otherwise wait for the sub-marker. + if "" in buf: + return GlmParser + return WAIT + if KIMI_SECTION_BEGIN in buf: + return KimiParser + if "<" not in buf and len(buf) > 8: + # No '<' anywhere, so no tag has started: release the text. The length + # floor is an unexplained heuristic carried over verbatim from the + # original parser — it trades a little first-token latency for not + # committing on a 1-2 char buffer. Note it does not actually rule out a + # partial MiniMax ns_token, whose first char is ']'. + return EMIT_CONTENT + return WAIT diff --git a/atom/entrypoints/openai/tool_parser/schema.py b/atom/entrypoints/openai/tool_parser/schema.py new file mode 100644 index 0000000000..e3acd35101 --- /dev/null +++ b/atom/entrypoints/openai/tool_parser/schema.py @@ -0,0 +1,87 @@ +# SPDX-License-Identifier: MIT +# Copyright (C) 2024-2026, Advanced Micro Devices, Inc. All rights reserved. + +"""Request-schema lookup and value coercion shared by every wire format. + +The XML-ish tool-call formats (Qwen, DSML, GLM, MiniMax) carry no value types on +the wire, so parameter values arrive as strings. When the request supplies a +``tools`` schema each value is coerced to its declared JSON-Schema type; +otherwise it is left alone. +""" + +import ast +import json +from typing import Any, Dict, Optional + + +def build_param_types(tools: Optional[list]) -> Dict[str, Dict[str, Any]]: + """Map ``function_name -> {param_name: json_schema_type}`` from request tools. + + Accepts OpenAI (``{"type": "function", "function": {...}}``) and bare + (``{"name": ..., "parameters"/"input_schema": {...}}``) tool entries. + """ + out: Dict[str, Dict[str, Any]] = {} + for tool in tools or []: + if not isinstance(tool, dict): + continue + fn = tool.get("function", tool) + if not isinstance(fn, dict): + continue + name = fn.get("name") + if not name: + continue + schema = fn.get("parameters") or fn.get("input_schema") or {} + props = schema.get("properties") if isinstance(schema, dict) else None + out[name] = { + k: (v.get("type") if isinstance(v, dict) else None) + for k, v in (props or {}).items() + } + return out + + +def coerce_param_value(value: str, ptype: Any) -> Any: + """Coerce a string parameter value to its declared JSON-Schema type. + + No schema type (string/unknown) -> returned unchanged. Conversion failures + fall back to the original string rather than raising. + """ + v = value.strip("\n") + if ptype is None: + return v + t = str(ptype).lower() + try: + if t in ("string", "str", "text", "varchar", "char", "enum"): + return v + if t in ("null", "none"): + return None + if t.startswith(("int", "uint", "long", "short", "unsigned")): + return int(v) + if t.startswith(("num", "float", "double", "decimal")): + f = float(v) + return int(f) if f.is_integer() else f + if t.startswith(("bool", "binary")): + return v.strip().lower() == "true" + if t.startswith(("object", "dict", "map", "array", "list", "tuple")): + try: + return json.loads(v) + except Exception: + return ast.literal_eval(v) # safer for single-quoted Python literals + except Exception: + return v + return v + + +def coerce_json_or_raw(value: str, ptype: Any) -> Any: + """Decode one untyped value: schema type wins, else JSON, else raw string. + + Shared by GLM (````) and MiniMax (````), whose templates both + render non-string values with ``tojson`` and string values raw. + """ + v = value.strip("\n") + if ptype is not None: + return coerce_param_value(v, ptype) + s = v.strip() + try: + return json.loads(s) + except Exception: + return v diff --git a/atom/entrypoints/openai/tool_parser/stream.py b/atom/entrypoints/openai/tool_parser/stream.py new file mode 100644 index 0000000000..7102205d0d --- /dev/null +++ b/atom/entrypoints/openai/tool_parser/stream.py @@ -0,0 +1,69 @@ +# SPDX-License-Identifier: MIT +# Copyright (C) 2024-2026, Advanced Micro Devices, Inc. All rights reserved. + +"""Streaming facade: sniff the format once, then delegate every chunk to it.""" + +from dataclasses import dataclass, field +from typing import Optional + +from .registry import EMIT_CONTENT, WAIT, sniff_stream +from .tool_parser import ToolCallParser + + +@dataclass +class ToolCallStreamParser: + """Stateful streaming parser; format is auto-detected from the first chunks. + + Emits structured events: + - ("content", text) — regular content before tool calls + - ("tool_call_start", {"index": N, "id": ..., "function": {"name": ..., "arguments": ""}}) + - ("tool_call_args", {"index": N, "function": {"arguments": chunk}}) + - ("tool_call_end", None) — all tool calls complete + + ``tools`` enables JSON-Schema type coercion of parameter values. It may be + assigned after construction (several call sites do) and is re-read on every + delegated call, so it takes effect as long as it is set before the stream + ends. + """ + + tools: Optional[list] = None + # Pre-detection accumulator. Once a format is chosen this is handed to the + # concrete parser and never used again. + _buf: str = "" + _parser: Optional[ToolCallParser] = field(default=None, repr=False) + + @property + def fmt(self) -> Optional[str]: + """Detected format name, or None while still undecided.""" + return self._parser.NAME if self._parser is not None else None + + def process(self, text: str) -> list: + """Process a text chunk and return list of (event_type, data) tuples.""" + if self._parser is None: + self._buf += text + choice = sniff_stream(self._buf) + if choice is WAIT: + return [] + if choice is EMIT_CONTENT: + out = [("content", self._buf)] + self._buf = "" + return out + self._parser = choice(tools=self.tools) + # Replay everything accumulated while undecided. + text, self._buf = self._buf, "" + + self._parser.tools = self.tools + return self._parser.process(text) + + def flush(self) -> list: + """Flush remaining buffer content.""" + if self._parser is None: + # Undecided at EOS: no tool markers ever appeared -> plain content. + if self._buf: + out = [("content", self._buf)] + self._buf = "" + return out + return [] + + self._parser.tools = self.tools + return self._parser.flush() diff --git a/atom/entrypoints/openai/tool_parser/tool_parser.py b/atom/entrypoints/openai/tool_parser/tool_parser.py new file mode 100644 index 0000000000..66c4733f6d --- /dev/null +++ b/atom/entrypoints/openai/tool_parser/tool_parser.py @@ -0,0 +1,167 @@ +# SPDX-License-Identifier: MIT +# Copyright (C) 2024-2026, Advanced Micro Devices, Inc. All rights reserved. + +"""Parser interface and the shared buffered-marker streaming strategy. + +Every wire format implements :class:`ToolCallParser`. Four of the five +(Qwen / DSML / GLM / MiniMax) stream identically — buffer from the first start +marker, parse the whole block at flush — so that strategy lives once in +:class:`BufferedMarkerParser` and each format only declares its markers and its +``parse``. Kimi is the exception: its token format is self-delimiting, so it +emits tool calls incrementally and implements ``process``/``flush`` itself. +""" + +import uuid +from abc import ABC, abstractmethod +from dataclasses import dataclass +from typing import Any, ClassVar, Dict, List, Optional, Tuple + + +def unique_tool_call_id() -> str: + # OpenAI tool_call ids must be unique across the whole conversation, not just + # within one response. A per-response index (call_0, call_1, ...) collides + # across turns -> clients (e.g. qwen-code) dedupe by id and silently ignore + # every repeat, causing an infinite tool-call retry loop. Use a random id. + return f"call_{uuid.uuid4().hex}" + + +@dataclass +class ToolCall: + """Parsed tool call in OpenAI format.""" + + id: str + type: str + function: Dict[str, str] + + def to_dict(self) -> Dict[str, Any]: + return {"id": self.id, "type": self.type, "function": self.function} + + +class ToolCallParser(ABC): + """One on-the-wire tool-call format. + + Class side is the stateless non-streaming path (``detect`` + ``parse``); + instance side is the stateful streaming path (``process`` + ``flush``). + """ + + NAME: ClassVar[str] + + def __init__(self, tools: Optional[list] = None): + self.tools = tools + self.buf = "" + # 0 = still in plain content, 1 = inside a tool-call region. Kimi adds + # 2 = section closed; see KimiParser. + self.state = 0 + self.current_index = 0 + self.emitted_calls = 0 + + # -- non-streaming ------------------------------------------------------ + @classmethod + @abstractmethod + def detect(cls, text: str) -> bool: + """Whether a complete model output is in this format.""" + + @classmethod + @abstractmethod + def parse(cls, text: str, tools: Optional[list]) -> Tuple[str, List[ToolCall]]: + """Parse a complete output into ``(leading_content, tool_calls)``.""" + + # -- streaming ---------------------------------------------------------- + @abstractmethod + def process(self, text: str) -> list: + """Consume one chunk; return ``(event_type, data)`` tuples.""" + + @abstractmethod + def flush(self) -> list: + """Drain whatever is buffered at end of stream.""" + + def _emit_call(self, tc: ToolCall) -> list: + """Render one parsed ToolCall as start+args stream events.""" + events = [ + ( + "tool_call_start", + { + "index": self.current_index, + "id": tc.id, + "type": "function", + "function": {"name": tc.function["name"], "arguments": ""}, + }, + ), + ( + "tool_call_args", + { + "index": self.current_index, + "function": {"arguments": tc.function["arguments"]}, + }, + ), + ] + self.current_index += 1 + self.emitted_calls += 1 + return events + + +class BufferedMarkerParser(ToolCallParser): + """Formats that buffer from a start marker and parse the block at flush. + + The block is only parsed once complete because partial XML streams badly + (a half-written `` int: + """Index of the earliest start marker, or -1.""" + positions = [i for i in (text.find(m) for m in cls.START_MARKERS) if i != -1] + return min(positions) if positions else -1 + + @classmethod + def detect(cls, text: str) -> bool: + return cls.find_start(text) != -1 + + def process(self, text: str) -> list: + results: list = [] + self.buf += text + if self.state == 0: + m = self.find_start(self.buf) + if m != -1: + before = self.buf[:m] + if before: + results.append(("content", before)) + self.buf = self.buf[m:] + self.state = 1 + else: + # Emit content but hold back a possible partial marker tail. + cut = max(self.buf.rfind(c) for c in self.HOLDBACK_CHARS) + if cut == -1: + if self.buf: + results.append(("content", self.buf)) + self.buf = "" + elif cut > 0: + results.append(("content", self.buf[:cut])) + self.buf = self.buf[cut:] + return results + + def flush(self) -> list: + results: list = [] + if self.state == 0: + if self.buf: + results.append(("content", self.buf)) + self.buf = "" + return results + # state 1: parse the complete (or trailing) tool-call block. + _content, tool_calls = self.parse(self.buf, self.tools) + self.buf = "" + for tc in tool_calls: + results.extend(self._emit_call(tc)) + if self.emitted_calls > 0: + results.append(("tool_call_end", None)) + return results From b04c2f3cdccfbf079e29bcde97119d337b542fca Mon Sep 17 00:00:00 2001 From: yihonglie Date: Tue, 28 Jul 2026 15:47:01 +0800 Subject: [PATCH 11/18] [Frontend] Drop DSML tool-format instruction injection Remove inject_tool_format_instruction() and its _DSML_TOOL_INSTRUCTION prompt block, plus both call sites: /v1/messages (anthropic_messages) and /v1/responses (responses_endpoint). It appended a "MANDATORY tool-call format" section to the system message spelling out the DSML syntax, so a non-tuned model would emit parseable DSML instead of ad-hoc /```json text. Requested removed along with the other prompt-content injections on these paths. Note this is a behavior change, not dead-code cleanup: the tool parser's DSML branch only fires on output that already looks like DSML, so models that needed the instruction to produce it may now emit tool calls the parser cannot recover. Worth re-checking DSV4-Pro tool calling on both paths. Also fixes a comment left stale by the tool_parser package split: the key-alias note referenced _parse_dsml, now DsmlParser.parse. Co-Authored-By: Claude Opus 5 (1M context) Signed-off-by: yihonglie --- atom/entrypoints/openai/api_server.py | 13 +------- atom/entrypoints/openai/serving_responses.py | 33 -------------------- 2 files changed, 1 insertion(+), 45 deletions(-) diff --git a/atom/entrypoints/openai/api_server.py b/atom/entrypoints/openai/api_server.py index 165810225a..48494c02bb 100644 --- a/atom/entrypoints/openai/api_server.py +++ b/atom/entrypoints/openai/api_server.py @@ -70,7 +70,6 @@ ) from .serving_responses import ( ResponsesStreamEmitter, - inject_tool_format_instruction, shell_arg_key, build_responses_object, remap_tool_name, @@ -1461,14 +1460,6 @@ async def anthropic_messages(request: AnthropicMessagesRequest, raw_request: Req # Convert Anthropic messages to OpenAI format openai_messages = anthropic_to_openai_messages(request.messages, request.system) - # Inject the mandatory DSML tool-call format instruction when tools are - # present, so the model emits parseable DSML (matching the model's native - # encoding) instead of ad-hoc /```json text. Same injection the - # /v1/responses path applies; without it, Claude Code (which speaks the - # Anthropic /v1/messages API) hits malformed/weak tool calls. - if request.tools: - openai_messages = inject_tool_format_instruction(openai_messages) - # Apply chat template from .protocol import ChatMessage @@ -1777,8 +1768,6 @@ async def responses_endpoint(raw_request: Request): openai_messages = responses_input_to_messages( body.get("instructions"), body.get("input") ) - if openai_tools: - openai_messages = inject_tool_format_instruction(openai_messages) messages = [ChatMessage(**m) for m in openai_messages] merged_kwargs = dict(default_chat_template_kwargs) @@ -1844,7 +1833,7 @@ async def generate_responses_stream(): reasoning_filter.state = 1 tool_parser = ToolCallStreamParser() tool_parser.tools = openai_tools or None # enables schema-based - # type coercion + key-alias (command->cmd) in _parse_dsml + # type coercion + key-alias (command->cmd) in DsmlParser.parse output_tokens = 0 # Buffer each tool call (name + full args) so read/grep/ls/find diff --git a/atom/entrypoints/openai/serving_responses.py b/atom/entrypoints/openai/serving_responses.py index 662bb485a0..088f6896c8 100644 --- a/atom/entrypoints/openai/serving_responses.py +++ b/atom/entrypoints/openai/serving_responses.py @@ -99,39 +99,6 @@ def responses_input_to_messages(instructions: Any, inp: Any) -> List[Dict[str, A return messages -_DSML_TOOL_INSTRUCTION = ( - "\n\n# Tool-call format (MANDATORY — overrides any other format instruction)\n" - "When you call a tool, output ONLY a DSML tool-call block and NOTHING else " - "in that message — no markdown, no ```json, no ///" - " tags, no prose. Use EXACTLY this syntax (the \uff5c characters " - "are U+FF5C fullwidth vertical bars, not ASCII '|'):\n" - "<\uff5cDSML\uff5ctool_calls>\n" - '<\uff5cDSML\uff5cinvoke name="TOOL_NAME">\n' - '<\uff5cDSML\uff5cparameter name="PARAM_NAME" string="true">VALUE' - "\n" - "\n" - "\n" - "Use the exact tool and parameter names from the tools provided to you. " - "For a shell/exec tool, put the whole shell command string in its " - "command/cmd parameter. Emit one <\uff5cDSML\uff5cinvoke> per tool call." -) - - -def inject_tool_format_instruction(messages): - """Append the mandatory DSML tool-call format to the system message so the - model emits parseable DSML instead of ad-hoc /```json text. Codex - (/v1/responses) path only. Idempotent per request.""" - for m in messages: - if m.get("role") == "system": - base = m.get("content") or "" - if "\uff5cDSML\uff5ctool_calls" not in base: - m["content"] = _text_of(base) + _DSML_TOOL_INSTRUCTION - return messages - return [{"role": "system", "content": _DSML_TOOL_INSTRUCTION.strip()}] + list( - messages - ) - - def responses_tools_to_openai(tools: Any) -> List[Dict[str, Any]]: """Translate Responses function tools into OpenAI chat tool defs. From 5a425a7ad2f8a3bbe2773f83fc22801a242e6f31 Mon Sep 17 00:00:00 2001 From: yihonglie Date: Tue, 28 Jul 2026 15:51:25 +0800 Subject: [PATCH 12/18] [Frontend] Drop MiniMax M3 reasoning tag normalization Keep the shared reasoning parser focused on the supported standard think tags now that MiniMax M3 compatibility is out of scope. Co-authored-by: Cursor --- atom/entrypoints/openai/reasoning.py | 7 ------- 1 file changed, 7 deletions(-) diff --git a/atom/entrypoints/openai/reasoning.py b/atom/entrypoints/openai/reasoning.py index f7cd8f0b41..6fb8a8e004 100644 --- a/atom/entrypoints/openai/reasoning.py +++ b/atom/entrypoints/openai/reasoning.py @@ -23,9 +23,6 @@ def separate_reasoning(text: str) -> Tuple[Optional[str], str]: Tuple of (reasoning_content, content). reasoning_content is None if no thinking block was found. """ - # MiniMax M3 emits ... instead of ...; - # normalize so the shared logic below handles both. - text = text.replace("", "").replace("", "") # Check for closed thinking block: ... match = re.match(r"(.*?)\s*(.*)", text, flags=re.DOTALL) if match: @@ -78,10 +75,6 @@ def process(self, text: str) -> list: List of (field_name, text) tuples where field_name is "reasoning_content" or "content". """ - # MiniMax M3 uses /; normalize to the tags - # the state machine below keys on. These are single special tokens, so - # each arrives whole in one chunk — a plain replace is safe. - text = text.replace("", "").replace("", "") results = [] if self.state == 0: From 02826d65709ea1eaab75a9932a6b2446afe1c2bd Mon Sep 17 00:00:00 2001 From: yihonglie Date: Tue, 28 Jul 2026 16:19:51 +0800 Subject: [PATCH 13/18] [Frontend] Scope custom encoder preparation by model Keep the shared chat-template dispatcher model-agnostic while giving DeepSeek-V4 the same synthetic tool-system message flow used by vLLM. Co-authored-by: Cursor --- .../openai/chat_encoder_adapters.py | 67 +++++++++ atom/entrypoints/openai/chat_encoders.py | 73 +++------- tests/entrypoints/test_chat_encoders.py | 134 ++++++++++++++++++ 3 files changed, 219 insertions(+), 55 deletions(-) create mode 100644 atom/entrypoints/openai/chat_encoder_adapters.py create mode 100644 tests/entrypoints/test_chat_encoders.py diff --git a/atom/entrypoints/openai/chat_encoder_adapters.py b/atom/entrypoints/openai/chat_encoder_adapters.py new file mode 100644 index 0000000000..38254a37df --- /dev/null +++ b/atom/entrypoints/openai/chat_encoder_adapters.py @@ -0,0 +1,67 @@ +# SPDX-License-Identifier: MIT +# Copyright (C) 2024-2026, Advanced Micro Devices, Inc. All rights reserved. + +"""Model-scoped adapters for dynamically loaded chat encoders.""" + +from dataclasses import dataclass +from typing import Any, Callable, List, Optional + +MessageEncoder = Callable[..., str] +MessagePreparer = Callable[[List[dict], Optional[List[dict]]], List[dict]] + + +def _copy_messages( + messages: List[dict], _tools: Optional[List[dict]] = None +) -> List[dict]: + """Return shallow message copies without model-specific rewriting.""" + return [dict(message) for message in messages] + + +def _prepare_deepseek_v4_messages( + messages: List[dict], tools: Optional[List[dict]] +) -> List[dict]: + """Prepare the internal message shape expected by DSV4 ``encode_messages``. + + The DeepSeek-V4 reference encoder reads tool schemas from a system message's + ``tools`` field. Match vLLM's model-specific tokenizer wrapper by prepending + a synthetic tool-carrying system message, without reordering or merging the + caller's existing messages. + """ + prepared = _copy_messages(messages) + if tools: + prepared.insert(0, {"role": "system", "tools": tools}) + return prepared + + +@dataclass(frozen=True) +class MessageEncoderAdapter: + """A raw model encoder plus its model-specific message preparation.""" + + name: str + encode: MessageEncoder + prepare_messages: MessagePreparer + supports_tools: bool = False + + def __call__(self, messages: List[dict], **kwargs: Any) -> str: + """Preserve the callable behavior of the former encoder return value.""" + return self.encode(messages, **kwargs) + + +_PREPARERS: dict[str, tuple[MessagePreparer, bool]] = { + "encoding_dsv4": (_prepare_deepseek_v4_messages, True), +} + + +def build_message_encoder_adapter( + module_name: str, encoder: MessageEncoder +) -> MessageEncoderAdapter: + """Build an adapter registered for ``module_name`` or an identity adapter.""" + prepare_messages, supports_tools = _PREPARERS.get( + module_name, (_copy_messages, False) + ) + return MessageEncoderAdapter( + name=module_name, + encode=encoder, + prepare_messages=prepare_messages, + supports_tools=supports_tools, + ) diff --git a/atom/entrypoints/openai/chat_encoders.py b/atom/entrypoints/openai/chat_encoders.py index 9021ab85d3..dcf2b1c152 100644 --- a/atom/entrypoints/openai/chat_encoders.py +++ b/atom/entrypoints/openai/chat_encoders.py @@ -14,13 +14,16 @@ import importlib.util import logging import os -from typing import Any, Callable, List, Optional +from typing import Any, List, Optional from huggingface_hub import snapshot_download -logger = logging.getLogger("atom") +from .chat_encoder_adapters import ( + MessageEncoderAdapter, + build_message_encoder_adapter, +) -MessageEncoder = Callable[..., str] +logger = logging.getLogger("atom") def _resolve_model_path(model: str) -> str: @@ -32,7 +35,7 @@ def _resolve_model_path(model: str) -> str: return model -def _load_encoder_from_dir(model_path: str) -> Optional[MessageEncoder]: +def _load_encoder_from_dir(model_path: str) -> Optional[MessageEncoderAdapter]: """Look for ``/encoding/encoding_*.py`` and load ``encode_messages``. Returns ``None`` when the directory or matching file is absent (model uses @@ -72,10 +75,10 @@ def encode(messages, **kwargs): return raw(messages, **kwargs) logger.info(f"Loaded message encoder from {enc_path}") - return encode + return build_message_encoder_adapter(module_name, encode) -def load_custom_message_encoder(model_path: str) -> Optional[MessageEncoder]: +def load_custom_message_encoder(model_path: str) -> Optional[MessageEncoderAdapter]: """Probe ``model_path`` once at startup for a custom message encoder. Returns the encoder, or ``None`` when the model uses the standard Jinja @@ -85,53 +88,9 @@ def load_custom_message_encoder(model_path: str) -> Optional[MessageEncoder]: return _load_encoder_from_dir(_resolve_model_path(model_path)) -def _content_str(c: Any) -> str: - if isinstance(c, list): - return "\n".join( - b.get("text", "") - for b in c - if isinstance(b, dict) and b.get("type") == "text" - ) - return c or "" - - -def _normalize_for_v4(messages: List[dict], tools: Optional[List[dict]]) -> List[dict]: - """Prepare messages for DeepSeek-V4's ``encode_messages``. - - Two things: - 1. **Hoist system messages to the front.** Clients (notably Claude Code) send - a trailing ``system``-role message (its "skills" list) AFTER the user turn. - ``encode_messages`` only appends the ``<|Assistant|>`` generation marker - after a *user*/developer message, so a trailing system message leaves the - prompt ending mid-system-text and the model just *continues* it instead of - answering. Merging all system content into one leading system message keeps - the final turn a user turn, so the assistant marker is emitted. - 2. **Attach tools** to that leading system message (``encode_messages`` reads - tool schemas from a system message's ``tools`` field). - Does not mutate the input. - """ - sys_parts, others = [], [] - for m in messages: - (sys_parts if m.get("role") == "system" else others).append(dict(m)) - - if not sys_parts and not tools: - return [dict(m) for m in messages] - - merged = "\n\n".join( - s for s in (_content_str(m.get("content")) for m in sys_parts) if s - ) - sys_msg: dict = {"role": "system", "content": merged} - for m in sys_parts: # preserve any pre-attached tools - if m.get("tools"): - sys_msg["tools"] = m["tools"] - if tools: - sys_msg["tools"] = tools - return [sys_msg] + others - - def apply_chat_template( tokenizer: Any, - custom_encoder: Optional[MessageEncoder], + custom_encoder: Optional[MessageEncoderAdapter], messages: List[dict], *, tools: Optional[List[dict]] = None, @@ -142,14 +101,18 @@ def apply_chat_template( Dispatches to ``custom_encoder`` if one was discovered for this model, otherwise to ``tokenizer.apply_chat_template``. Jinja-only kwargs (``tokenize``, ``add_generation_prompt``) are stripped on the custom path. - ``tools`` are supported on both paths: custom encoders (e.g. DeepSeek-V4's - ``encode_messages``) read tool schemas from a system message's ``tools`` - field, so we attach them there before encoding. + Model-scoped adapters prepare tools for custom encoders that support them; + the generic path does not apply DeepSeek-V4-specific message rewriting. """ if custom_encoder is not None: for k in ("tokenize", "add_generation_prompt"): kwargs.pop(k, None) - messages = _normalize_for_v4(messages, tools) + if tools and not custom_encoder.supports_tools: + logger.warning( + "tools= is not supported by custom message encoder %s; ignoring.", + custom_encoder.name, + ) + messages = custom_encoder.prepare_messages(messages, tools) return custom_encoder(messages, **kwargs) kwargs["tokenize"] = False diff --git a/tests/entrypoints/test_chat_encoders.py b/tests/entrypoints/test_chat_encoders.py new file mode 100644 index 0000000000..b25d3c873e --- /dev/null +++ b/tests/entrypoints/test_chat_encoders.py @@ -0,0 +1,134 @@ +# SPDX-License-Identifier: MIT +# Copyright (C) 2024-2026, Advanced Micro Devices, Inc. All rights reserved. + +"""Tests for model-scoped custom chat encoder dispatch.""" + +from atom.entrypoints.openai.chat_encoder_adapters import ( + build_message_encoder_adapter, +) +from atom.entrypoints.openai.chat_encoders import ( + _load_encoder_from_dir, + apply_chat_template, +) + + +def test_loader_selects_dsv4_adapter_and_preserves_encoder_defaults(tmp_path): + encoding_dir = tmp_path / "encoding" + encoding_dir.mkdir() + (encoding_dir / "encoding_dsv4.py").write_text( + "def encode_messages(messages, **kwargs):\n" + " return repr((messages, kwargs))\n", + encoding="utf-8", + ) + + adapter = _load_encoder_from_dir(str(tmp_path)) + + assert adapter is not None + assert adapter.name == "encoding_dsv4" + assert adapter.supports_tools is True + rendered = apply_chat_template( + tokenizer=None, + custom_encoder=adapter, + messages=[{"role": "user", "content": "hello"}], + ) + assert "'thinking_mode': 'thinking'" in rendered + + +def test_dsv4_adapter_prepends_tools_without_reordering_messages(): + captured = {} + + def raw_encoder(messages, **kwargs): + captured["messages"] = messages + captured["kwargs"] = kwargs + return "rendered" + + adapter = build_message_encoder_adapter("encoding_dsv4", raw_encoder) + messages = [ + {"role": "system", "content": "policy"}, + {"role": "user", "content": "question"}, + {"role": "system", "content": "trailing context"}, + ] + original = [dict(message) for message in messages] + tools = [{"type": "function", "function": {"name": "search"}}] + + result = apply_chat_template( + tokenizer=None, + custom_encoder=adapter, + messages=messages, + tools=tools, + tokenize=True, + add_generation_prompt=True, + thinking_mode="chat", + ) + + assert result == "rendered" + assert captured["messages"] == [ + {"role": "system", "tools": tools}, + *original, + ] + assert captured["kwargs"] == {"thinking_mode": "chat"} + assert messages == original + assert captured["messages"][1:] is not messages + assert all( + prepared is not source + for prepared, source in zip(captured["messages"][1:], messages) + ) + + +def test_unknown_custom_encoder_does_not_receive_dsv4_fields(caplog): + captured = {} + + def raw_encoder(messages, **kwargs): + captured["messages"] = messages + return "rendered" + + adapter = build_message_encoder_adapter("encoding_other", raw_encoder) + messages = [{"role": "user", "content": "hello"}] + tools = [{"type": "function", "function": {"name": "search"}}] + + result = apply_chat_template( + tokenizer=None, + custom_encoder=adapter, + messages=messages, + tools=tools, + ) + + assert result == "rendered" + assert captured["messages"] == messages + assert captured["messages"] is not messages + assert captured["messages"][0] is not messages[0] + assert "tools" not in captured["messages"][0] + assert "tools= is not supported" in caplog.text + + +def test_jinja_path_forwards_tools_and_generation_kwargs(): + class Tokenizer: + def __init__(self): + self.messages = None + self.kwargs = None + + def apply_chat_template(self, messages, **kwargs): + self.messages = messages + self.kwargs = kwargs + return "jinja-rendered" + + tokenizer = Tokenizer() + messages = [{"role": "user", "content": "hello"}] + tools = [{"type": "function", "function": {"name": "search"}}] + + result = apply_chat_template( + tokenizer=tokenizer, + custom_encoder=None, + messages=messages, + tools=tools, + enable_thinking=True, + ) + + assert result == "jinja-rendered" + assert tokenizer.messages is messages + assert tokenizer.kwargs == { + "enable_thinking": True, + "tokenize": False, + "add_generation_prompt": True, + "tools": tools, + } From 68948c03570abb3e97536199e1fd44994af51e53 Mon Sep 17 00:00:00 2001 From: yihonglie Date: Tue, 28 Jul 2026 17:09:57 +0800 Subject: [PATCH 14/18] [Frontend] Unify batched stream dispatch Extract incremental UTF-8 detokenization and queue batching so direct and fan-out streams share one lifecycle-safe dispatch path. Co-authored-by: Cursor --- atom/entrypoints/openai/api_server.py | 139 +++++------------ atom/entrypoints/openai/streaming_dispatch.py | 140 ++++++++++++++++++ tests/entrypoints/test_streaming_dispatch.py | 122 +++++++++++++++ 3 files changed, 300 insertions(+), 101 deletions(-) create mode 100644 atom/entrypoints/openai/streaming_dispatch.py create mode 100644 tests/entrypoints/test_streaming_dispatch.py diff --git a/atom/entrypoints/openai/api_server.py b/atom/entrypoints/openai/api_server.py index 48494c02bb..8f9ee4332a 100644 --- a/atom/entrypoints/openai/api_server.py +++ b/atom/entrypoints/openai/api_server.py @@ -79,6 +79,7 @@ tool_name_lookup, translate_client_tool, ) +from .streaming_dispatch import StreamBatchDispatcher # Configure logging logger = logging.getLogger("atom") @@ -103,6 +104,7 @@ _stream_loops: Dict[str, AbstractEventLoop] = {} _request_start_times: Dict[str, float] = {} _request_logger: Optional[logging.Logger] = None +_stream_batch_dispatcher: Optional[StreamBatchDispatcher] = None # ============================================================================ @@ -336,36 +338,12 @@ def _prepare_multimodal_inputs( # ── Batched stream dispatch ────────────────────────────────────────────── -# Per-seq `call_soon_threadsafe` floods the API event loop at high batch size -# (one call per token). Instead the callback only buffers the raw chunk; the -# mgr flushes a whole step with a single `tokenizer.batch_decode` (one -# GIL-released call instead of one decode per seq) plus one scheduled call per -# loop (see `flush_stream_batch`). -import threading as _threading # noqa: E402 -_stream_batch_tls = _threading.local() -# Per-request incremental detokenization state (vLLM-style sliding window). -# Decoding each step's new tokens in isolation splits multi-byte UTF-8 chars -# (byte-BPE tokenizers like DeepSeek-V4 split one CJK char across several -# byte-tokens) into U+FFFD. Keep accumulated tokens + prefix/read offsets so -# we only emit fully-formed characters. -_stream_detok_state: Dict[str, dict] = {} - - -def _send_stream_chunk_direct( - request_output: RequestOutput, - request_id: str, - stream_queue: asyncio.Queue, - loop: AbstractEventLoop, -) -> None: - """Send stream chunk directly to the queue.""" - global tokenizer - - new_text = tokenizer.decode(request_output.output_tokens, skip_special_tokens=True) +def _build_stream_chunk(request_output: RequestOutput, request_id: str) -> dict: + """Build a raw chunk; detokenization happens once in the batch dispatcher.""" started_at = _request_start_times.get(request_id) chunk_data = { - "text": new_text, "token_ids": request_output.output_tokens, "finished": request_output.finished, "finish_reason": request_output.finish_reason, @@ -375,76 +353,34 @@ def _send_stream_chunk_direct( } if getattr(request_output, "kv_transfer_params_output", None): chunk_data["kv_transfer_params"] = request_output.kv_transfer_params_output + return chunk_data - chunk_data["request_id"] = request_id - buf = getattr(_stream_batch_tls, "buf", None) - if buf is None: - buf = _stream_batch_tls.buf = [] - buf.append((loop, stream_queue, chunk_data)) - -def _drain_batch_into_queues(items: list) -> None: - """Runs ON the event loop: push each chunk into its per-request queue. - One scheduled call handles a whole step's worth of chunks.""" - for _loop, q, chunk in items: - q.put_nowait(chunk) +def _send_stream_chunk_direct( + request_output: RequestOutput, + request_id: str, + stream_queue: asyncio.Queue, + loop: AbstractEventLoop, +) -> None: + """Buffer a single-request chunk for this engine step.""" + assert _stream_batch_dispatcher is not None + _stream_batch_dispatcher.enqueue( + loop=loop, + queue=stream_queue, + state_key=request_id, + chunk=_build_stream_chunk(request_output, request_id), + ) def flush_stream_batch() -> None: - """Flush a step's buffered chunks: one ``batch_decode`` for the whole step, - then one call_soon_threadsafe per loop (normally one — all requests on a - rank share the API loop).""" - global tokenizer - - buf = getattr(_stream_batch_tls, "buf", None) - if not buf: - return - _stream_batch_tls.buf = [] - # Decode the whole step in a single call. batch_decode is element-wise - # identical to per-seq decode but acquires/releases the GIL once instead of - # once per seq, cutting GIL ping-pong against the other rank output threads - # and the API event loop at high batch size. - # Incremental per-request detokenization: correct UTF-8 at token - # boundaries (see _stream_detok_state). Emits only fully-formed chars; - # a trailing partial multi-byte char is held until the next step. - for _loop, _q, chunk in buf: - rid = chunk.get("request_id") - st = _stream_detok_state.get(rid) - if st is None: - st = _stream_detok_state[rid] = { - "tokens": [], - "prefix_offset": 0, - "read_offset": 0, - } - toks = st["tokens"] - toks.extend(chunk["token_ids"]) - prefix_text = tokenizer.decode( - toks[st["prefix_offset"] : st["read_offset"]], skip_special_tokens=True - ) - new_text = tokenizer.decode( - toks[st["prefix_offset"] :], skip_special_tokens=True - ) - if len(new_text) > len(prefix_text) and not new_text.endswith("\ufffd"): - chunk["text"] = new_text[len(prefix_text) :] - st["prefix_offset"] = st["read_offset"] - st["read_offset"] = len(toks) - elif chunk["finished"]: - chunk["text"] = new_text[len(prefix_text) :] - else: - chunk["text"] = "" - if chunk["finished"]: - _stream_detok_state.pop(rid, None) - # Group by loop (normally a single loop). dict preserves insertion order - # so per-request chunk ordering within the step is maintained. - by_loop: Dict[AbstractEventLoop, list] = {} - for loop, q, chunk in buf: - by_loop.setdefault(loop, []).append((loop, q, chunk)) - for loop, items in by_loop.items(): - loop.call_soon_threadsafe(_drain_batch_into_queues, items) + """Flush this output thread's engine-step batch into asyncio queues.""" + if _stream_batch_dispatcher is not None: + _stream_batch_dispatcher.flush() def _send_stream_chunk_tagged( request_output: RequestOutput, + request_id: str, sibling_index: int, stream_queue: asyncio.Queue, loop: AbstractEventLoop, @@ -458,18 +394,14 @@ def _send_stream_chunk_tagged( This path serves ``SamplingParams.n > 1`` by tagging each sibling's chunks so the shared stream consumer can merge them in order. """ - global tokenizer - - new_text = tokenizer.decode(request_output.output_tokens, skip_special_tokens=True) - chunk_data = { - "text": new_text, - "token_ids": request_output.output_tokens, - "finished": request_output.finished, - "finish_reason": request_output.finish_reason, - } - if getattr(request_output, "kv_transfer_params_output", None): - chunk_data["kv_transfer_params"] = request_output.kv_transfer_params_output - loop.call_soon_threadsafe(stream_queue.put_nowait, (sibling_index, chunk_data)) + assert _stream_batch_dispatcher is not None + _stream_batch_dispatcher.enqueue( + loop=loop, + queue=stream_queue, + state_key=(request_id, sibling_index), + chunk=_build_stream_chunk(request_output, request_id), + tag=sibling_index, + ) async def generate_async( @@ -925,6 +857,8 @@ def cleanup_streaming_request( _seq_id_to_request_id.pop(seq_id, None) _stream_loops.pop(request_id, None) _request_start_times.pop(request_id, None) + if _stream_batch_dispatcher is not None: + _stream_batch_dispatcher.discard_request(request_id) if aborted: try: engine.core_mgr.abort_request(seq_id) @@ -1061,7 +995,9 @@ async def setup_streaming_request_fanout( def make_callback(idx: int): def _cb(request_output: RequestOutput) -> None: - _send_stream_chunk_tagged(request_output, idx, shared_queue, stream_loop) + _send_stream_chunk_tagged( + request_output, request_id, idx, shared_queue, stream_loop + ) return _cb @@ -2070,7 +2006,7 @@ async def stop_profile(): def main(): """Main entry point for the server.""" global engine, tokenizer, model_name, default_chat_template_kwargs, _request_logger - global custom_message_encoder + global custom_message_encoder, _stream_batch_dispatcher parser = FlexibleArgumentParser(description="ATOM OpenAI API Server") EngineArgs.add_cli_args(parser) @@ -2120,6 +2056,7 @@ def main(): logger.info(f"Initializing engine with model {args.model}...") engine_args = EngineArgs.from_cli_args(args) engine = engine_args.create_engine(tokenizer=tokenizer) + _stream_batch_dispatcher = StreamBatchDispatcher(tokenizer) # Wire the batched stream-flush hook: per-seq stream callbacks only buffer # their chunks into a thread-local; the engine core manager's output thread diff --git a/atom/entrypoints/openai/streaming_dispatch.py b/atom/entrypoints/openai/streaming_dispatch.py new file mode 100644 index 0000000000..10a8609797 --- /dev/null +++ b/atom/entrypoints/openai/streaming_dispatch.py @@ -0,0 +1,140 @@ +# SPDX-License-Identifier: MIT +# Copyright (C) 2024-2026, Advanced Micro Devices, Inc. All rights reserved. + +"""Batched cross-thread dispatch for streaming model output.""" + +import threading +from asyncio import AbstractEventLoop, Queue +from dataclasses import dataclass, field +from typing import Any, Hashable, Optional + + +@dataclass +class IncrementalStreamDetokenizer: + """Decode token deltas without emitting incomplete UTF-8 characters.""" + + tokenizer: Any + tokens: list[int] = field(default_factory=list) + prefix_offset: int = 0 + read_offset: int = 0 + + def update(self, token_ids: list[int], finished: bool) -> str: + self.tokens.extend(token_ids) + prefix_text = self.tokenizer.decode( + self.tokens[self.prefix_offset : self.read_offset], + skip_special_tokens=True, + ) + new_text = self.tokenizer.decode( + self.tokens[self.prefix_offset :], + skip_special_tokens=True, + ) + + if len(new_text) > len(prefix_text) and not new_text.endswith("\ufffd"): + delta = new_text[len(prefix_text) :] + self.prefix_offset = self.read_offset + self.read_offset = len(self.tokens) + return delta + if finished: + return new_text[len(prefix_text) :] + return "" + + +@dataclass +class _BufferedChunk: + loop: AbstractEventLoop + queue: Queue + state_key: Hashable + chunk: dict + tag: Optional[int] + + +class StreamBatchDispatcher: + """Collect one engine step per output thread and dispatch it by event loop.""" + + def __init__(self, tokenizer: Any): + self.tokenizer = tokenizer + self._thread_local = threading.local() + self._states: dict[Hashable, IncrementalStreamDetokenizer] = {} + self._states_lock = threading.Lock() + + def enqueue( + self, + *, + loop: AbstractEventLoop, + queue: Queue, + state_key: Hashable, + chunk: dict, + tag: Optional[int] = None, + ) -> None: + """Buffer a raw chunk until the current engine step is flushed.""" + buf = getattr(self._thread_local, "buf", None) + if buf is None: + buf = self._thread_local.buf = [] + buf.append( + _BufferedChunk( + loop=loop, + queue=queue, + state_key=state_key, + chunk=chunk, + tag=tag, + ) + ) + + def flush(self) -> None: + """Detokenize buffered chunks and schedule one drain per event loop.""" + buf = getattr(self._thread_local, "buf", None) + if not buf: + return + self._thread_local.buf = [] + + by_loop: dict[AbstractEventLoop, list[tuple[Queue, Any]]] = {} + for item in buf: + state = self._get_state(item.state_key) + item.chunk["text"] = state.update( + item.chunk.get("token_ids") or [], + bool(item.chunk.get("finished")), + ) + if item.chunk.get("finished"): + self._drop_state(item.state_key, state) + + payload = ( + item.chunk if item.tag is None else (item.tag, item.chunk) + ) + by_loop.setdefault(item.loop, []).append((item.queue, payload)) + + for loop, items in by_loop.items(): + loop.call_soon_threadsafe(self._drain_into_queues, items) + + def discard_request(self, request_id: str) -> None: + """Drop direct and fan-out detokenizer state after request cleanup.""" + with self._states_lock: + keys = [ + key + for key in self._states + if key == request_id + or (isinstance(key, tuple) and key and key[0] == request_id) + ] + for key in keys: + self._states.pop(key, None) + + def _get_state(self, state_key: Hashable) -> IncrementalStreamDetokenizer: + with self._states_lock: + state = self._states.get(state_key) + if state is None: + state = self._states[state_key] = IncrementalStreamDetokenizer( + self.tokenizer + ) + return state + + def _drop_state( + self, state_key: Hashable, state: IncrementalStreamDetokenizer + ) -> None: + with self._states_lock: + if self._states.get(state_key) is state: + self._states.pop(state_key) + + @staticmethod + def _drain_into_queues(items: list[tuple[Queue, Any]]) -> None: + """Run on the target event loop and deliver each prepared payload.""" + for queue, payload in items: + queue.put_nowait(payload) diff --git a/tests/entrypoints/test_streaming_dispatch.py b/tests/entrypoints/test_streaming_dispatch.py new file mode 100644 index 0000000000..9c561ee300 --- /dev/null +++ b/tests/entrypoints/test_streaming_dispatch.py @@ -0,0 +1,122 @@ +import asyncio + +from atom.entrypoints.openai.streaming_dispatch import ( + IncrementalStreamDetokenizer, + StreamBatchDispatcher, +) + + +class _Utf8ByteTokenizer: + def decode(self, token_ids, skip_special_tokens=True): + return bytes(token_ids).decode("utf-8", errors="replace") + + +class _ImmediateLoop: + def __init__(self): + self.calls = [] + + def call_soon_threadsafe(self, callback, *args): + self.calls.append((callback, args)) + callback(*args) + + +def test_incremental_detokenizer_holds_incomplete_utf8(): + detokenizer = IncrementalStreamDetokenizer(_Utf8ByteTokenizer()) + + assert detokenizer.update([0xE4], finished=False) == "" + assert detokenizer.update([0xBD, 0xA0], finished=False) == "你" + assert detokenizer.update([ord("!")], finished=True) == "!" + + +def test_dispatcher_batches_direct_and_tagged_chunks_per_loop(): + dispatcher = StreamBatchDispatcher(_Utf8ByteTokenizer()) + loop = _ImmediateLoop() + direct_queue = asyncio.Queue() + tagged_queue = asyncio.Queue() + + dispatcher.enqueue( + loop=loop, + queue=direct_queue, + state_key="request-1", + chunk={"token_ids": [ord("A")], "finished": True}, + ) + dispatcher.enqueue( + loop=loop, + queue=tagged_queue, + state_key=("request-2", 0), + chunk={"token_ids": [ord("B")], "finished": True}, + tag=0, + ) + dispatcher.flush() + + assert len(loop.calls) == 1 + assert direct_queue.get_nowait()["text"] == "A" + sibling_index, chunk = tagged_queue.get_nowait() + assert sibling_index == 0 + assert chunk["text"] == "B" + + +def test_dispatcher_keeps_fanout_detokenizer_state_separate(): + dispatcher = StreamBatchDispatcher(_Utf8ByteTokenizer()) + loop = _ImmediateLoop() + queue = asyncio.Queue() + + dispatcher.enqueue( + loop=loop, + queue=queue, + state_key=("request", 0), + chunk={"token_ids": [0xE4], "finished": False}, + tag=0, + ) + dispatcher.enqueue( + loop=loop, + queue=queue, + state_key=("request", 1), + chunk={"token_ids": [ord("X")], "finished": True}, + tag=1, + ) + dispatcher.flush() + + assert queue.get_nowait()[1]["text"] == "" + assert queue.get_nowait()[1]["text"] == "X" + + dispatcher.enqueue( + loop=loop, + queue=queue, + state_key=("request", 0), + chunk={"token_ids": [0xBD, 0xA0], "finished": True}, + tag=0, + ) + dispatcher.flush() + + assert queue.get_nowait()[1]["text"] == "你" + + +def test_discard_request_drops_partial_direct_and_fanout_state(): + dispatcher = StreamBatchDispatcher(_Utf8ByteTokenizer()) + loop = _ImmediateLoop() + queue = asyncio.Queue() + + for state_key in ("request", ("request", 0)): + dispatcher.enqueue( + loop=loop, + queue=queue, + state_key=state_key, + chunk={"token_ids": [0xE4], "finished": False}, + ) + dispatcher.flush() + dispatcher.discard_request("request") + + for state_key in ("request", ("request", 0)): + dispatcher.enqueue( + loop=loop, + queue=queue, + state_key=state_key, + chunk={"token_ids": [ord("A")], "finished": True}, + ) + dispatcher.flush() + + assert queue.get_nowait()["text"] == "" + assert queue.get_nowait()["text"] == "" + assert queue.get_nowait()["text"] == "A" + assert queue.get_nowait()["text"] == "A" From 4721c02e81ea16f134ef0ed9be5a3237884f23f2 Mon Sep 17 00:00:00 2001 From: yihonglie Date: Tue, 28 Jul 2026 18:36:05 +0800 Subject: [PATCH 15/18] [Frontend] Drop Python Responses API support Remove the standalone Codex Responses endpoint and its tool adaptation layer so the frontend remains focused on chat and Anthropic protocols. Co-authored-by: Cursor --- atom/entrypoints/openai/api_server.py | 236 ------ atom/entrypoints/openai/serving_responses.py | 691 ------------------ .../tool_parser/deepseekv4_tool_parser.py | 30 +- 3 files changed, 3 insertions(+), 954 deletions(-) delete mode 100644 atom/entrypoints/openai/serving_responses.py diff --git a/atom/entrypoints/openai/api_server.py b/atom/entrypoints/openai/api_server.py index 8f9ee4332a..6a4d0d1b0a 100644 --- a/atom/entrypoints/openai/api_server.py +++ b/atom/entrypoints/openai/api_server.py @@ -68,17 +68,6 @@ stream_completion_response, stream_completion_response_fanout, ) -from .serving_responses import ( - ResponsesStreamEmitter, - shell_arg_key, - build_responses_object, - remap_tool_name, - extract_cwd, - responses_input_to_messages, - responses_tools_to_openai, - tool_name_lookup, - translate_client_tool, -) from .streaming_dispatch import StreamBatchDispatcher # Configure logging @@ -1676,231 +1665,6 @@ async def generate_anthropic_stream(): ) -@app.post("/v1/responses") -async def responses_endpoint(raw_request: Request): - """Handle OpenAI **Responses API** requests (`/v1/responses`). - - Native support so OpenAI Codex CLI (>= 0.14x, which speaks only the Responses - API) can talk to ATOM directly — no external responses->chat proxy needed. - Reuses ATOM's proven streaming path (setup_streaming_request + ReasoningFilter - + ToolCallStreamParser), the same one /v1/messages (claude-local) uses, so it - streams correctly for reasoning models. Stateless (full input sent each turn, - as Codex does); reasoning items are dropped from visible output. - """ - global engine, tokenizer, model_name - - try: - body = await raw_request.json() - model = body.get("model") or model_name - - from .protocol import ChatMessage - from .reasoning import ReasoningFilter, separate_reasoning - from .tool_parser import ToolCallStreamParser, parse_tool_calls - - openai_tools = responses_tools_to_openai(body.get("tools")) - valid_names, shell_tool = tool_name_lookup(openai_tools) - shell_param = shell_arg_key(openai_tools, shell_tool) - req_cwd = extract_cwd(body) # Codex — used to fix hallucinated paths - openai_messages = responses_input_to_messages( - body.get("instructions"), body.get("input") - ) - messages = [ChatMessage(**m) for m in openai_messages] - - merged_kwargs = dict(default_chat_template_kwargs) - prompt = apply_chat_template( - tokenizer, - custom_message_encoder, - [msg.to_template_dict() for msg in messages], - tools=openai_tools or None, - **merged_kwargs, - ) - - max_out = int(body.get("max_output_tokens") or 32768) - sampling_params = _build_sampling_params( - temperature=( - body.get("temperature") if body.get("temperature") is not None else 1.0 - ), - max_tokens=max_out, - stop_strings=None, - ignore_eos=False, - top_k=-1, - top_p=body.get("top_p") if body.get("top_p") is not None else 1.0, - ) - - request_id = "resp_" + uuid.uuid4().hex[:24] - input_tokens = len(tokenizer.encode(prompt)) - - # Resolve max context to bound the prompt (same probes as anthropic). - max_ctx = None - for _path in ( - lambda: engine.config.max_model_len, - lambda: engine.model_config.max_model_len, - lambda: engine.scheduler.max_model_len, - lambda: getattr(engine, "max_model_len"), - ): - try: - _v = _path() - if _v: - max_ctx = int(_v) - break - except Exception: - continue - if not max_ctx: - max_ctx = 30720 - headroom = min(max_out, max(1024, max_ctx // 8)) - max_input = max_ctx - headroom - if input_tokens > max_input: - logger.warning( - f"[responses] prompt too long ({input_tokens} > {max_input}), truncating" - ) - token_ids = tokenizer.encode(prompt)[:max_input] - prompt = tokenizer.decode(token_ids, skip_special_tokens=False) - input_tokens = max_input - - if body.get("stream"): - seq_id, stream_queue, _num_prompt_tokens = await setup_streaming_request( - prompt, sampling_params, request_id - ) - - async def generate_responses_stream(): - emitter = ResponsesStreamEmitter(request_id, model) - reasoning_filter = ReasoningFilter() - if prompt.rstrip().endswith(""): - reasoning_filter.state = 1 - tool_parser = ToolCallStreamParser() - tool_parser.tools = openai_tools or None # enables schema-based - # type coercion + key-alias (command->cmd) in DsmlParser.parse - output_tokens = 0 - - # Buffer each tool call (name + full args) so read/grep/ls/find - # can be translated to exec_command before emitting (name AND - # args change). See translate_client_tool. - _pending = {"tc": None} - - def _flush_pending(): - tc = _pending["tc"] - if tc is None: - return [] - _pending["tc"] = None - name, args = translate_client_tool( - tc["name"], - tc["args"], - valid_names, - shell_tool, - req_cwd, - shell_param, - ) - out = emitter.tool_start(tc["id"], name) - if args: - out += emitter.tool_args(args) - out += emitter.tool_end() - return out - - def handle(etype, edata): - # Map ToolCallStreamParser events -> Responses SSE strings, - # buffering tool calls for client-tool translation. - if etype == "content": - out = _flush_pending() - return out + emitter.text_delta(edata) - if etype == "tool_call_start": - out = _flush_pending() - fn = edata.get("function", {}) - _pending["tc"] = { - "id": edata.get("id", ""), - "name": fn.get("name", ""), - "args": "", - } - return out - if etype == "tool_call_args": - if _pending["tc"] is not None: - _pending["tc"]["args"] += ( - edata.get("function", {}).get("arguments", "") or "" - ) - return [] - if etype == "tool_call_end": - return _flush_pending() - return [] - - aborted = True # abort the seq on disconnect; False on normal end - try: - for s in emitter.created(): - yield s - while True: - chunk_data = await stream_queue.get() - new_text = chunk_data["text"] - output_tokens += len(chunk_data.get("token_ids", [])) - finished = chunk_data.get("finished", False) - - segments = reasoning_filter.process(new_text) - if finished: - segments.extend(reasoning_filter.flush()) - for field, text in segments: - if not text or field == "reasoning_content": - continue # drop reasoning from visible output - for etype, edata in tool_parser.process(text): - for s in handle(etype, edata): - yield s - - if finished: - for etype, edata in tool_parser.flush(): - for s in handle(etype, edata): - yield s - for s in _flush_pending(): # emit any unclosed tool call - yield s - for s in emitter.finish(input_tokens, output_tokens): - yield s - yield "data: [DONE]\n\n" - aborted = False - break - finally: - cleanup_streaming_request(request_id, seq_id, aborted=aborted) - - return StreamingResponse( - generate_responses_stream(), - media_type="text/event-stream", - headers={"x-request-id": request_id}, - ) - - # Non-streaming response - final_output = await _run_nonstream_with_disconnect( - generate_async(prompt, sampling_params, request_id), - raw_request, - request_id, - ) - if final_output is None: - raise RuntimeError("No output generated") - - raw_text = final_output["text"] - _reasoning, content_with_tools = separate_reasoning(raw_text) - content_text, tool_calls = parse_tool_calls( - content_with_tools, openai_tools or None - ) - output_tokens = len(tokenizer.encode(raw_text)) - - return JSONResponse( - content=build_responses_object( - resp_id=request_id, - model=model, - content_text=content_text, - tool_calls=tool_calls, - input_tokens=input_tokens, - output_tokens=output_tokens, - valid=valid_names, - shell_tool=shell_tool, - cwd=req_cwd, - ) - ) - - except _ClientDisconnected: - return JSONResponse(status_code=499, content={"detail": "client disconnected"}) - except Exception as e: - logger.error(f"Error in responses_endpoint: {e}", exc_info=True) - return JSONResponse( - status_code=500, - content={"error": {"type": "api_error", "message": str(e)}}, - ) - - @app.get("/v1/models") async def list_models(): """List available models.""" diff --git a/atom/entrypoints/openai/serving_responses.py b/atom/entrypoints/openai/serving_responses.py deleted file mode 100644 index 088f6896c8..0000000000 --- a/atom/entrypoints/openai/serving_responses.py +++ /dev/null @@ -1,691 +0,0 @@ -"""OpenAI **Responses API** (`/v1/responses`) support for the ATOM server. - -OpenAI Codex CLI (>= 0.14x) dropped `wire_api = "chat"` and only speaks the -Responses API (streaming SSE). This module provides the translation between -Responses request/response shapes and ATOM's internal chat/engine machinery, -plus a streaming SSE event emitter. The `/v1/responses` route handler in -``api_server.py`` reuses ATOM's proven streaming path (``setup_streaming_request`` -+ ``ReasoningFilter`` + ``ToolCallStreamParser``) so it streams correctly for -reasoning models — the same path claude-local uses via ``/v1/messages``. - -This makes the external ``codex_responses_proxy.py`` unnecessary: Codex can point -straight at ATOM's ``:9700/v1``. -""" - -import itertools -import json -import time -from typing import Any, Dict, List, Optional, Tuple - -_ids = itertools.count(1) - - -def _rid(prefix: str) -> str: - return f"{prefix}_{int(time.time() * 1000)}{next(_ids):04d}" - - -# --------------------------------------------------------------- request xlate -def _text_of(content: Any) -> str: - """Flatten Responses content (str | list of parts) to a plain string.""" - if isinstance(content, str): - return content - if isinstance(content, list): - out = [] - for p in content: - if isinstance(p, dict): - if "text" in p and isinstance(p["text"], str): - out.append(p["text"]) - elif p.get("type") in ("input_text", "output_text", "text"): - out.append(p.get("text", "")) - elif isinstance(p, str): - out.append(p) - return "".join(out) - return "" - - -def responses_input_to_messages(instructions: Any, inp: Any) -> List[Dict[str, Any]]: - """Translate Responses ``instructions`` + ``input`` into OpenAI chat messages. - - ``input`` may be a plain string or a list of items: ``message`` / - ``function_call`` (assistant tool call) / ``function_call_output`` (tool - result). ``reasoning`` items are dropped. - """ - messages: List[Dict[str, Any]] = [] - if instructions: - messages.append({"role": "system", "content": _text_of(instructions)}) - - if isinstance(inp, str): - messages.append({"role": "user", "content": inp}) - elif isinstance(inp, list): - for item in inp: - if not isinstance(item, dict): - continue - t = item.get("type", "message") - if t == "message": - role = item.get("role", "user") - messages.append( - {"role": role, "content": _text_of(item.get("content", ""))} - ) - elif t == "function_call": - messages.append( - { - "role": "assistant", - "content": "", - "tool_calls": [ - { - "id": item.get("call_id") - or item.get("id") - or _rid("call"), - "type": "function", - "function": { - "name": item.get("name", ""), - "arguments": item.get("arguments", "") or "", - }, - } - ], - } - ) - elif t == "function_call_output": - out = item.get("output", "") - messages.append( - { - "role": "tool", - "tool_call_id": item.get("call_id") or item.get("id") or "", - "content": out if isinstance(out, str) else json.dumps(out), - } - ) - elif t == "reasoning": - continue - return messages - - -def responses_tools_to_openai(tools: Any) -> List[Dict[str, Any]]: - """Translate Responses function tools into OpenAI chat tool defs. - - Responses puts ``name``/``description``/``parameters`` at the top level of a - ``{"type": "function", ...}`` tool (chat nests them under ``function``). - Non-function tool types (``web_search``, ``namespace``, ...) are dropped — - ATOM only executes model-emitted function/DSML tool calls; the client - (Codex) owns actual tool execution. - """ - if not tools: - return [] - ct: List[Dict[str, Any]] = [] - for tl in tools: - if not isinstance(tl, dict): - continue - if tl.get("type") == "function": - fn = tl.get("function") or tl - ct.append( - { - "type": "function", - "function": { - "name": fn.get("name"), - "description": fn.get("description", ""), - "parameters": fn.get("parameters", {}) or {}, - }, - } - ) - return ct - - -# ------------------------------------------------- tool-name normalization -# Non-codex-tuned models (e.g. DeepSeek-V4-Pro) don't reliably emit the EXACT -# tool name the client registered for shell exec. Codex 0.142.x names it -# ``exec_command`` (args ``{"cmd": "..."}``), but the model habitually calls it -# ``exec`` / ``exec_run`` / ``shell`` / ``bash``, which Codex's tool router -# rejects ("unsupported call: exec"). The ARGUMENTS are correct; only the name -# is wrong. So remap known shell-exec aliases onto whatever shell tool the -# client actually registered this turn. Pure rename; arguments untouched. -_SHELL_ALIASES = { - "exec", - "exec_run", - "exec_command", - "execute_command", - "shell", - "bash", - "sh", - "run", - "run_command", - "run_shell", - "execute", - "command", - "container.exec", - "local_shell", - "shell_command", - "shell_exec", - "run_bash", - "execute_shell", - "bash_command", - "run_terminal_cmd", - "terminal", - "console", - "run_shell_command", - "runshell", -} -# Substrings that mark an unknown tool name as a shell/exec call (Codex's -# non-shell tools contain none of these). -_SHELL_NAME_TOKENS = ("shell", "exec", "bash", "cmd", "command", "termin", "console") -_SHELL_TOOL_PREFERENCE = ( - "exec_command", - "shell", - "local_shell", - "bash", - "container.exec", -) - - -def tool_name_lookup(openai_tools: List[Dict[str, Any]]) -> Tuple[set, Optional[str]]: - """Return (valid tool names, preferred shell tool name) for remapping.""" - valid = { - (t.get("function") or {}).get("name") - for t in (openai_tools or []) - if isinstance(t, dict) - } - valid.discard(None) - shell_tool = next((n for n in _SHELL_TOOL_PREFERENCE if n in valid), None) - return valid, shell_tool - - -def remap_tool_name(name: str, valid: set, shell_tool: Optional[str]) -> str: - """Fix a model tool_call name that doesn't match a registered tool. - - Only remaps shell-exec aliases to the registered shell tool; other - mismatches pass through unchanged (the client will surface them).""" - if name in valid: - return name - if not shell_tool: - return name - n = (name or "").lower().replace("-", "_").replace(".", "_") - if n in _SHELL_ALIASES: - return shell_tool - # Fuzzy: unknown tool whose name signals shell/exec -> the shell tool. - if any(k in n for k in _SHELL_NAME_TOKENS): - return shell_tool - return name - - -# ------------------------------------------------ Claude-tool -> shell adapter -# DeepSeek-V4-Pro is trained on Claude-Code's toolset, so under Codex it keeps -# calling read/grep/ls/find (which Codex doesn't have) instead of exec_command, -# and flails ("unsupported call: read"). When Codex registered an exec/shell -# tool, translate these read-only Claude tools into an equivalent shell command -# so the model's intent goes through. Only fires for names NOT in the registered -# set; exec_command's real calls are untouched. Codex-path only (never applied -# to /v1/messages, where read/grep ARE native tools). -def _q(s: Any) -> str: - import shlex - - return shlex.quote(str(s)) - - -def _resolve_dir(path: str, cwd: Optional[str]) -> str: - """Pick a directory that actually exists: keep relative paths and paths under - cwd; otherwise fall back to cwd (the model often invents absolute prefixes - like /sglang or /sgl-workspace that don't exist here).""" - if not path or path == ".": - return cwd or "." - if not path.startswith("/"): - return path # relative — shell runs in cwd, fine - if cwd and path.startswith(cwd): - return path - return cwd or path - - -def _read_cmd(fp: str, cwd: Optional[str], start: int, end: int) -> str: - """Resilient file read: try the literal path, then cwd-prefixed, then locate - by basename under cwd — so a hallucinated absolute prefix still finds the file.""" - sed = f"sed -n {start},{end}p" - if not cwd: - return f"{sed} {_q(fp)}" - # f = literal; if missing, cwd + '/' + (f without leading /); if still - # missing, first match of `find cwd -name basename`. - return ( - f'f={_q(fp)}; [ -f "$f" ] || f={_q(cwd)}"/${{f#/}}"; ' - f'[ -f "$f" ] || f=$(find {_q(cwd)} -type f -name "$(basename {_q(fp)})" ' - f'2>/dev/null | head -1); {sed} "$f"' - ) - - -def _claude_tool_to_shell( - name: str, a: Dict[str, Any], cwd: Optional[str] = None -) -> Optional[str]: - n = (name or "").lower() - fp = ( - a.get("file_path") - or a.get("path") - or a.get("filePath") - or a.get("filename") - or a.get("target_file") - or a.get("file") - ) - pattern = a.get("pattern") or a.get("query") or a.get("regex") - path = ( - a.get("path") - or a.get("directory") - or a.get("target_directory") - or a.get("dir") - or "." - ) - if n in ( - "read", - "cat", - "view", - "view_file", - "open", - "read_file", - "readfile", - "openfile", - ): - if not fp: - return None - off, lim = a.get("offset"), a.get("limit") - if off or lim: - try: - start = int(off or 0) + 1 - end = start + int(lim or 200) - 1 - except (ValueError, TypeError): - start, end = 1, 400 - else: - start, end = 1, 400 - return _read_cmd(str(fp), cwd, start, end) - if n in ( - "grep", - "search", - "search_file", - "ripgrep", - "rg", - "grep_search", - "codebase_search", - ): - if not pattern: - return None - return f"grep -rn -- {_q(pattern)} {_q(_resolve_dir(path, cwd))}" - if n in ("ls", "list", "list_dir", "list_directory", "listdir"): - return f"ls -la {_q(_resolve_dir(path, cwd))}" - if n in ("find", "glob", "glob_file_search", "file_search"): - base = _resolve_dir(path, cwd) - if pattern: - return f"find {_q(base)} -name {_q(pattern)}" - return f"find {_q(base)} -maxdepth 3" - return None - - -_SHELL_ARG_ALIASES = ( - "cmd", - "command", - "commandline", - "command_line", - "script", - "bash", - "sh", - "shell", - "shell_command", - "code", - "input", - "run", - "cmd_string", -) - - -def shell_arg_key(openai_tools, shell_tool): - """Required (or first) param name of the registered shell tool, e.g. Codex's - exec_command -> "cmd". Used to normalize model arg keys.""" - if not shell_tool: - return None - for t in openai_tools or []: - fn = t.get("function", t) - if (fn.get("name") or t.get("name")) != shell_tool: - continue - params = fn.get("parameters") or {} - props = params.get("properties") or {} - for r in params.get("required") or []: - if props.get(r, {}).get("type") in (None, "string"): - return r - if props: - return next(iter(props)) - return "cmd" - - -def _is_shell_name(name, shell_tool): - return name == shell_tool or name in _SHELL_ALIASES - - -def _normalize_shell_args(args_json, req): - """If the shell tool's required param `req` is absent but a known alias is - present, rename it. Keeps exec_command calls valid when the model uses - `command`/`script`/... instead of `cmd`.""" - if not req: - return args_json - try: - a = json.loads(args_json) if args_json else {} - except Exception: - return args_json - if not isinstance(a, dict) or req in a: - return args_json - for alias in _SHELL_ARG_ALIASES: - if alias != req and isinstance(a.get(alias), str): - a[req] = a.pop(alias) - return json.dumps(a) - return args_json - - -def translate_client_tool( - name: str, - args_json: str, - valid: set, - shell_tool: Optional[str], - cwd: Optional[str] = None, - shell_param: Optional[str] = None, -): - """Return (name, args_json) with Claude read-only tools rewritten to the - registered shell tool. exec_command's own calls and any already-valid tool - pass through untouched; unknown non-shell tools fall back to name-remap. - ``cwd`` (from the request's ) makes hallucinated absolute paths resolve.""" - if name in valid: - if shell_param and _is_shell_name(name, shell_tool): - args_json = _normalize_shell_args(args_json, shell_param) - return name, args_json - exec_tool = "exec_command" if "exec_command" in valid else shell_tool - if exec_tool: - try: - a = json.loads(args_json) if args_json else {} - except Exception: - a = {} - if isinstance(a, dict): - cmd = _claude_tool_to_shell(name, a, cwd) - if cmd is not None: - return exec_tool, json.dumps({(shell_param or "cmd"): cmd}) - _remapped = remap_tool_name(name, valid, shell_tool) - if shell_param and _is_shell_name(_remapped, shell_tool): - args_json = _normalize_shell_args(args_json, shell_param) - return _remapped, args_json - - -_CWD_RE = None - - -def extract_cwd(body: Dict[str, Any]) -> Optional[str]: - """Pull the working directory from Codex's ... environment_context - (sent in instructions/input), so path fix-ups target the real directory.""" - import re - - global _CWD_RE - if _CWD_RE is None: - # [^<]+? keeps internal spaces/newlines but never spans a stray '<' (so a - # malformed/unclosed inner can't over-capture); non-greedy stops at - # the first . - _CWD_RE = re.compile(r"\s*([^<]+?)\s*") - blob = _text_of(body.get("instructions")) - inp = body.get("input") - if isinstance(inp, str): - blob += "\n" + inp - elif isinstance(inp, list): - for it in inp: - if isinstance(it, dict): - blob += "\n" + _text_of(it.get("content", "")) - m = _CWD_RE.search(blob or "") - return m.group(1).strip() if m else None - - -# --------------------------------------------------------------- SSE emitter -class ResponsesStreamEmitter: - """Builds the ordered Responses SSE event stream from incremental text / - tool-call events. Each method returns a list of SSE strings to yield. - - Output-item lifecycle (Codex expects these exact event types): - message: output_item.added -> content_part.added -> output_text.delta* - -> output_text.done -> content_part.done -> output_item.done - function_call: output_item.added -> function_call_arguments.delta* - -> function_call_arguments.done -> output_item.done - end: response.completed - """ - - def __init__(self, resp_id: str, model: str): - self.resp_id = resp_id - self.model = model - self._seq = itertools.count(0) - self.out_index = 0 - self.final_output: List[Dict[str, Any]] = [] - self._open: Optional[Dict[str, Any]] = None # current open output item - - def _ev(self, ev: str, extra: Dict[str, Any]) -> str: - d = {"type": ev, "sequence_number": next(self._seq)} - d.update(extra) - return f"event: {ev}\ndata: {json.dumps(d)}\n\n" - - def _base(self, status: str, output: List[Dict[str, Any]]) -> Dict[str, Any]: - return { - "id": self.resp_id, - "object": "response", - "status": status, - "model": self.model, - "output": output, - } - - def created(self) -> List[str]: - return [ - self._ev("response.created", {"response": self._base("in_progress", [])}), - self._ev( - "response.in_progress", {"response": self._base("in_progress", [])} - ), - ] - - def _close_open(self) -> List[str]: - if self._open is None: - return [] - o = self._open - self._open = None - if o["kind"] == "message": - mid, cur, txt = o["id"], o["index"], o["text"] - item = { - "id": mid, - "type": "message", - "role": "assistant", - "status": "completed", - "content": [{"type": "output_text", "text": txt}], - } - self.final_output.append(item) - return [ - self._ev( - "response.output_text.done", - { - "item_id": mid, - "output_index": cur, - "content_index": 0, - "text": txt, - }, - ), - self._ev( - "response.content_part.done", - { - "item_id": mid, - "output_index": cur, - "content_index": 0, - "part": {"type": "output_text", "text": txt}, - }, - ), - self._ev( - "response.output_item.done", {"output_index": cur, "item": item} - ), - ] - # function_call - fid, cur = o["id"], o["index"] - item = { - "id": fid, - "type": "function_call", - "status": "completed", - "call_id": o["call_id"], - "name": o["name"], - "arguments": o["args"], - } - self.final_output.append(item) - return [ - self._ev( - "response.function_call_arguments.done", - {"item_id": fid, "output_index": cur, "arguments": o["args"]}, - ), - self._ev("response.output_item.done", {"output_index": cur, "item": item}), - ] - - def text_delta(self, delta: str) -> List[str]: - out: List[str] = [] - if self._open and self._open["kind"] != "message": - out += self._close_open() - if not self._open: - mid, cur = _rid("msg"), self.out_index - self.out_index += 1 - self._open = {"kind": "message", "id": mid, "index": cur, "text": ""} - out.append( - self._ev( - "response.output_item.added", - { - "output_index": cur, - "item": { - "id": mid, - "type": "message", - "role": "assistant", - "status": "in_progress", - "content": [], - }, - }, - ) - ) - out.append( - self._ev( - "response.content_part.added", - { - "item_id": mid, - "output_index": cur, - "content_index": 0, - "part": {"type": "output_text", "text": ""}, - }, - ) - ) - self._open["text"] += delta - out.append( - self._ev( - "response.output_text.delta", - { - "item_id": self._open["id"], - "output_index": self._open["index"], - "content_index": 0, - "delta": delta, - }, - ) - ) - return out - - def tool_start(self, call_id: str, name: str) -> List[str]: - out = self._close_open() - fid, cur = _rid("fc"), self.out_index - self.out_index += 1 - self._open = { - "kind": "fc", - "id": fid, - "index": cur, - "call_id": call_id or _rid("call"), - "name": name, - "args": "", - } - out.append( - self._ev( - "response.output_item.added", - { - "output_index": cur, - "item": { - "id": fid, - "type": "function_call", - "status": "in_progress", - "call_id": self._open["call_id"], - "name": name, - "arguments": "", - }, - }, - ) - ) - return out - - def tool_args(self, delta: str) -> List[str]: - if not self._open or self._open["kind"] != "fc" or not delta: - return [] - self._open["args"] += delta - return [ - self._ev( - "response.function_call_arguments.delta", - { - "item_id": self._open["id"], - "output_index": self._open["index"], - "delta": delta, - }, - ) - ] - - def tool_end(self) -> List[str]: - return self._close_open() - - def finish(self, input_tokens: int, output_tokens: int) -> List[str]: - out = self._close_open() - completed = self._base("completed", self.final_output) - completed["usage"] = { - "input_tokens": input_tokens, - "output_tokens": output_tokens, - "total_tokens": input_tokens + output_tokens, - } - out.append(self._ev("response.completed", {"response": completed})) - return out - - -# ------------------------------------------------------- non-stream response -def build_responses_object( - resp_id: str, - model: str, - content_text: str, - tool_calls: List[Any], - input_tokens: int, - output_tokens: int, - valid: set, - shell_tool: Optional[str], - cwd: Optional[str] = None, -) -> Dict[str, Any]: - """Build a full (non-streaming) Responses object from parsed output. - - ``tool_calls`` are ATOM ``ToolCall`` objects (``.id``, ``.function`` dict).""" - output: List[Dict[str, Any]] = [] - if content_text: - output.append( - { - "id": _rid("msg"), - "type": "message", - "role": "assistant", - "status": "completed", - "content": [{"type": "output_text", "text": content_text}], - } - ) - for tc in tool_calls or []: - fn = getattr(tc, "function", None) or {} - name, args = translate_client_tool( - fn.get("name", ""), fn.get("arguments", "") or "", valid, shell_tool, cwd - ) - output.append( - { - "id": _rid("fc"), - "type": "function_call", - "status": "completed", - "call_id": getattr(tc, "id", None) or _rid("call"), - "name": name, - "arguments": args, - } - ) - return { - "id": resp_id, - "object": "response", - "status": "completed", - "model": model, - "output": output, - "usage": { - "input_tokens": input_tokens, - "output_tokens": output_tokens, - "total_tokens": input_tokens + output_tokens, - }, - } diff --git a/atom/entrypoints/openai/tool_parser/deepseekv4_tool_parser.py b/atom/entrypoints/openai/tool_parser/deepseekv4_tool_parser.py index e263dc5eff..7b9c062d70 100644 --- a/atom/entrypoints/openai/tool_parser/deepseekv4_tool_parser.py +++ b/atom/entrypoints/openai/tool_parser/deepseekv4_tool_parser.py @@ -42,25 +42,15 @@ re.DOTALL, ) -# Canonical param key -> emitted synonyms the model uses interchangeably. Only -# applied when the canonical key is in the tool's declared schema and the model -# used a synonym instead (e.g. Codex's exec_command wants `cmd`, but the model -# habitually emits `command` -> "missing field cmd" retry loop). Scoped to the -# schema so it never renames a key a tool legitimately declares. -_KEY_ALIASES: Dict[str, Tuple[str, ...]] = { - "cmd": ("command",), -} - - def _unwrap_wrapper_args(args: Any, allowed: set) -> Any: """Strip spurious ``{"arguments": {...}}`` / ``{"input": {...}}`` envelopes. Non-tuned models (DeepSeek-V4-Pro) frequently wrap the real args in an extra ``arguments``/``input`` object — sometimes nested 2-3 deep, or stringified — so a call meant as ``{"cmd": "ls"}`` arrives as ``{"arguments": {"cmd": - "ls"}}`` and the client (Codex) rejects it ("missing field cmd"). Recursively - unwrap while the sole key is a wrapper that is NOT itself a declared param of - the tool. Mirrors vLLM's ``_unwrap_wrapper_args`` (deepseek_v4.py).""" + "ls"}}``. Recursively unwrap while the sole key is a wrapper that is NOT + itself a declared param of the tool. Mirrors vLLM's ``_unwrap_wrapper_args`` + (deepseek_v4.py).""" for _ in range(4): # bounded against pathological nesting if not (isinstance(args, dict) and len(args) == 1): break @@ -80,18 +70,6 @@ def _unwrap_wrapper_args(args: Any, allowed: set) -> Any: return args -def _apply_key_aliases(args: Any, allowed: set) -> Any: - if not (isinstance(args, dict) and allowed): - return args - for canon, syns in _KEY_ALIASES.items(): - if canon in allowed and canon not in args: - for s in syns: - if s in args: - args[canon] = args.pop(s) - break - return args - - def _coerce(value: str, string_attr: Optional[str], ptype: Any) -> Any: """Decode one ```` body. @@ -179,7 +157,6 @@ def parse(cls, text: str, tools: Optional[list]) -> Tuple[str, List[ToolCall]]: except Exception: pass args = _unwrap_wrapper_args(args, set(types)) - args = _apply_key_aliases(args, set(types)) calls.append((name, args)) else: # malformed: no complete invoke wrapper -> collect params, infer tool name @@ -192,7 +169,6 @@ def parse(cls, text: str, tools: Optional[list]) -> Tuple[str, List[ToolCall]]: types = param_types.get(name, {}) args = {k: _coerce(v, s, types.get(k)) for k, (v, s) in raw.items()} args = _unwrap_wrapper_args(args, set(types)) - args = _apply_key_aliases(args, set(types)) calls.append((name, args)) tool_calls = [ From e9561347eef4e40882b9c7a24ebc2c1bc02025c1 Mon Sep 17 00:00:00 2001 From: yhl-amd Date: Thu, 30 Jul 2026 07:59:47 +0000 Subject: [PATCH 16/18] style: fix Black formatting and Ruff G201 for pre-checks CI - api_server.py: logger.error(..., exc_info=True) -> logger.exception(...) (G201) - streaming_dispatch.py, deepseekv4_tool_parser.py: apply Black Co-Authored-By: Claude Opus 4.8 --- atom/entrypoints/openai/api_server.py | 2 +- atom/entrypoints/openai/streaming_dispatch.py | 4 +--- atom/entrypoints/openai/tool_parser/deepseekv4_tool_parser.py | 1 + 3 files changed, 3 insertions(+), 4 deletions(-) diff --git a/atom/entrypoints/openai/api_server.py b/atom/entrypoints/openai/api_server.py index 6a4d0d1b0a..3c5e74f21f 100644 --- a/atom/entrypoints/openai/api_server.py +++ b/atom/entrypoints/openai/api_server.py @@ -1655,7 +1655,7 @@ async def generate_anthropic_stream(): # Client hung up; seq already aborted + popped. Nothing to return. return JSONResponse(status_code=499, content={"detail": "client disconnected"}) except Exception as e: - logger.error(f"Error in anthropic_messages: {e}", exc_info=True) + logger.exception(f"Error in anthropic_messages: {e}") return JSONResponse( status_code=500, content={ diff --git a/atom/entrypoints/openai/streaming_dispatch.py b/atom/entrypoints/openai/streaming_dispatch.py index 10a8609797..2ac5ad5589 100644 --- a/atom/entrypoints/openai/streaming_dispatch.py +++ b/atom/entrypoints/openai/streaming_dispatch.py @@ -97,9 +97,7 @@ def flush(self) -> None: if item.chunk.get("finished"): self._drop_state(item.state_key, state) - payload = ( - item.chunk if item.tag is None else (item.tag, item.chunk) - ) + payload = item.chunk if item.tag is None else (item.tag, item.chunk) by_loop.setdefault(item.loop, []).append((item.queue, payload)) for loop, items in by_loop.items(): diff --git a/atom/entrypoints/openai/tool_parser/deepseekv4_tool_parser.py b/atom/entrypoints/openai/tool_parser/deepseekv4_tool_parser.py index 7b9c062d70..5b9c8bffeb 100644 --- a/atom/entrypoints/openai/tool_parser/deepseekv4_tool_parser.py +++ b/atom/entrypoints/openai/tool_parser/deepseekv4_tool_parser.py @@ -42,6 +42,7 @@ re.DOTALL, ) + def _unwrap_wrapper_args(args: Any, allowed: set) -> Any: """Strip spurious ``{"arguments": {...}}`` / ``{"input": {...}}`` envelopes. From e11037a818bfab8be17f3450fc608536da2e068f Mon Sep 17 00:00:00 2001 From: yhl-amd Date: Thu, 30 Jul 2026 08:16:32 +0000 Subject: [PATCH 17/18] style: make new frontend files pass Ruff (UP/I001/BLE/S/TRY) and Black - Modernize type hints (Dict/List/Tuple/Optional -> builtins / X|None) in the new tool_parser package, chat_encoder_adapters, streaming_dispatch, and the PR-added lines of api_server/chat_encoders (UP006/UP035/UP045). - Sort api_server imports (I001). - Narrow best-effort json.loads catches to (ValueError, TypeError) (BLE001); keep one intentional catch-all with noqa (schema). - Drop redundant exception object in logger.exception (TRY401). Behavior-preserving; only files changed by this PR are touched. Co-Authored-By: Claude Opus 4.8 --- atom/entrypoints/openai/api_server.py | 33 ++++++++++--------- .../openai/chat_encoder_adapters.py | 15 +++++---- atom/entrypoints/openai/chat_encoders.py | 12 +++---- atom/entrypoints/openai/streaming_dispatch.py | 7 ++-- .../tool_parser/deepseekv4_tool_parser.py | 24 +++++++------- .../openai/tool_parser/glm_tool_parser.py | 10 +++--- .../openai/tool_parser/kimi_tool_parser.py | 6 ++-- .../openai/tool_parser/minimax_tool_parser.py | 12 +++---- .../openai/tool_parser/qwen3_tool_parser.py | 14 ++++---- .../openai/tool_parser/registry.py | 8 ++--- atom/entrypoints/openai/tool_parser/schema.py | 12 +++---- atom/entrypoints/openai/tool_parser/stream.py | 7 ++-- .../openai/tool_parser/tool_parser.py | 14 ++++---- 13 files changed, 86 insertions(+), 88 deletions(-) diff --git a/atom/entrypoints/openai/api_server.py b/atom/entrypoints/openai/api_server.py index 3c5e74f21f..dba63bbcf5 100644 --- a/atom/entrypoints/openai/api_server.py +++ b/atom/entrypoints/openai/api_server.py @@ -26,15 +26,16 @@ from typing import Any, AsyncGenerator, Dict, List, Optional, Tuple import uvicorn +from fastapi import FastAPI, HTTPException, Request +from fastapi.responses import JSONResponse, StreamingResponse +from PIL import Image +from transformers import AutoProcessor, AutoTokenizer + from atom import SamplingParams from atom.model_engine.arg_utils import EngineArgs from atom.model_engine.llm_engine import _load_tokenizer from atom.model_engine.request import RequestOutput from atom.utils.arg_parser import FlexibleArgumentParser -from fastapi import FastAPI, HTTPException, Request -from fastapi.responses import JSONResponse, StreamingResponse -from PIL import Image -from transformers import AutoProcessor, AutoTokenizer from .chat_encoders import apply_chat_template, load_custom_message_encoder from .protocol import ( @@ -43,12 +44,6 @@ ModelCard, ModelList, ) -from .serving_chat import ( - build_chat_response, - build_chat_response_multi, - stream_chat_response, - stream_chat_response_fanout, -) from .serving_anthropic import ( AnthropicMessagesRequest, anthropic_to_openai_messages, @@ -62,6 +57,12 @@ stream_message_stop, stream_signature_delta, ) +from .serving_chat import ( + build_chat_response, + build_chat_response_multi, + stream_chat_response, + stream_chat_response_fanout, +) from .serving_completion import ( build_completion_response, build_completion_response_multi, @@ -89,11 +90,11 @@ default_chat_template_kwargs: Dict[str, Any] = {} custom_message_encoder: Optional[Any] = None _stream_queues: Dict[str, asyncio.Queue] = {} -_seq_id_to_request_id: Dict[int, str] = {} -_stream_loops: Dict[str, AbstractEventLoop] = {} -_request_start_times: Dict[str, float] = {} -_request_logger: Optional[logging.Logger] = None -_stream_batch_dispatcher: Optional[StreamBatchDispatcher] = None +_seq_id_to_request_id: dict[int, str] = {} +_stream_loops: dict[str, AbstractEventLoop] = {} +_request_start_times: dict[str, float] = {} +_request_logger: logging.Logger | None = None +_stream_batch_dispatcher: StreamBatchDispatcher | None = None # ============================================================================ @@ -1655,7 +1656,7 @@ async def generate_anthropic_stream(): # Client hung up; seq already aborted + popped. Nothing to return. return JSONResponse(status_code=499, content={"detail": "client disconnected"}) except Exception as e: - logger.exception(f"Error in anthropic_messages: {e}") + logger.exception("Error in anthropic_messages") return JSONResponse( status_code=500, content={ diff --git a/atom/entrypoints/openai/chat_encoder_adapters.py b/atom/entrypoints/openai/chat_encoder_adapters.py index 38254a37df..11773b6265 100644 --- a/atom/entrypoints/openai/chat_encoder_adapters.py +++ b/atom/entrypoints/openai/chat_encoder_adapters.py @@ -3,23 +3,24 @@ """Model-scoped adapters for dynamically loaded chat encoders.""" +from collections.abc import Callable from dataclasses import dataclass -from typing import Any, Callable, List, Optional +from typing import Any MessageEncoder = Callable[..., str] -MessagePreparer = Callable[[List[dict], Optional[List[dict]]], List[dict]] +MessagePreparer = Callable[[list[dict], list[dict] | None], list[dict]] def _copy_messages( - messages: List[dict], _tools: Optional[List[dict]] = None -) -> List[dict]: + messages: list[dict], _tools: list[dict] | None = None +) -> list[dict]: """Return shallow message copies without model-specific rewriting.""" return [dict(message) for message in messages] def _prepare_deepseek_v4_messages( - messages: List[dict], tools: Optional[List[dict]] -) -> List[dict]: + messages: list[dict], tools: list[dict] | None +) -> list[dict]: """Prepare the internal message shape expected by DSV4 ``encode_messages``. The DeepSeek-V4 reference encoder reads tool schemas from a system message's @@ -42,7 +43,7 @@ class MessageEncoderAdapter: prepare_messages: MessagePreparer supports_tools: bool = False - def __call__(self, messages: List[dict], **kwargs: Any) -> str: + def __call__(self, messages: list[dict], **kwargs: Any) -> str: """Preserve the callable behavior of the former encoder return value.""" return self.encode(messages, **kwargs) diff --git a/atom/entrypoints/openai/chat_encoders.py b/atom/entrypoints/openai/chat_encoders.py index dcf2b1c152..4791912af5 100644 --- a/atom/entrypoints/openai/chat_encoders.py +++ b/atom/entrypoints/openai/chat_encoders.py @@ -14,7 +14,7 @@ import importlib.util import logging import os -from typing import Any, List, Optional +from typing import Any from huggingface_hub import snapshot_download @@ -35,7 +35,7 @@ def _resolve_model_path(model: str) -> str: return model -def _load_encoder_from_dir(model_path: str) -> Optional[MessageEncoderAdapter]: +def _load_encoder_from_dir(model_path: str) -> MessageEncoderAdapter | None: """Look for ``/encoding/encoding_*.py`` and load ``encode_messages``. Returns ``None`` when the directory or matching file is absent (model uses @@ -78,7 +78,7 @@ def encode(messages, **kwargs): return build_message_encoder_adapter(module_name, encode) -def load_custom_message_encoder(model_path: str) -> Optional[MessageEncoderAdapter]: +def load_custom_message_encoder(model_path: str) -> MessageEncoderAdapter | None: """Probe ``model_path`` once at startup for a custom message encoder. Returns the encoder, or ``None`` when the model uses the standard Jinja @@ -90,10 +90,10 @@ def load_custom_message_encoder(model_path: str) -> Optional[MessageEncoderAdapt def apply_chat_template( tokenizer: Any, - custom_encoder: Optional[MessageEncoderAdapter], - messages: List[dict], + custom_encoder: MessageEncoderAdapter | None, + messages: list[dict], *, - tools: Optional[List[dict]] = None, + tools: list[dict] | None = None, **kwargs: Any, ) -> str: """Render ``messages`` to a prompt string. diff --git a/atom/entrypoints/openai/streaming_dispatch.py b/atom/entrypoints/openai/streaming_dispatch.py index 2ac5ad5589..71b4662855 100644 --- a/atom/entrypoints/openai/streaming_dispatch.py +++ b/atom/entrypoints/openai/streaming_dispatch.py @@ -5,8 +5,9 @@ import threading from asyncio import AbstractEventLoop, Queue +from collections.abc import Hashable from dataclasses import dataclass, field -from typing import Any, Hashable, Optional +from typing import Any @dataclass @@ -45,7 +46,7 @@ class _BufferedChunk: queue: Queue state_key: Hashable chunk: dict - tag: Optional[int] + tag: int | None class StreamBatchDispatcher: @@ -64,7 +65,7 @@ def enqueue( queue: Queue, state_key: Hashable, chunk: dict, - tag: Optional[int] = None, + tag: int | None = None, ) -> None: """Buffer a raw chunk until the current engine step is flushed.""" buf = getattr(self._thread_local, "buf", None) diff --git a/atom/entrypoints/openai/tool_parser/deepseekv4_tool_parser.py b/atom/entrypoints/openai/tool_parser/deepseekv4_tool_parser.py index 5b9c8bffeb..559338ba67 100644 --- a/atom/entrypoints/openai/tool_parser/deepseekv4_tool_parser.py +++ b/atom/entrypoints/openai/tool_parser/deepseekv4_tool_parser.py @@ -19,7 +19,7 @@ import json import re -from typing import Any, ClassVar, Dict, List, Optional, Tuple +from typing import Any, ClassVar from .schema import build_param_types, coerce_param_value from .tool_parser import BufferedMarkerParser, ToolCall, unique_tool_call_id @@ -63,7 +63,7 @@ def _unwrap_wrapper_args(args: Any, allowed: set) -> Any: if isinstance(v, str): try: v = json.loads(v) - except Exception: + except (ValueError, TypeError): break if not isinstance(v, dict): break @@ -71,7 +71,7 @@ def _unwrap_wrapper_args(args: Any, allowed: set) -> Any: return args -def _coerce(value: str, string_attr: Optional[str], ptype: Any) -> Any: +def _coerce(value: str, string_attr: str | None, ptype: Any) -> Any: """Decode one ```` body. Deliberately not :func:`~.schema.coerce_json_or_raw`: on a JSON-decode miss @@ -83,7 +83,7 @@ def _coerce(value: str, string_attr: Optional[str], ptype: Any) -> Any: if string_attr == "false": try: return json.loads(value) - except Exception: + except (ValueError, TypeError): return value # attr absent -> use declared schema type if known, else infer via JSON. if ptype is not None: @@ -91,13 +91,11 @@ def _coerce(value: str, string_attr: Optional[str], ptype: Any) -> Any: v = value.strip() try: return json.loads(v) - except Exception: + except (ValueError, TypeError): return v -def _infer_name( - arg_names: set, param_types: Dict[str, Dict[str, Any]] -) -> Optional[str]: +def _infer_name(arg_names: set, param_types: dict[str, dict[str, Any]]) -> str | None: """Pick the request tool whose parameter set best matches ``arg_names``.""" best, best_score = None, -1e9 for name, props in param_types.items(): @@ -113,7 +111,7 @@ def _infer_name( class DsmlParser(BufferedMarkerParser): NAME: ClassVar[str] = "dsml" # Region-start markers, both marked and marker-less variants. - START_MARKERS: ClassVar[Tuple[str, ...]] = ( + START_MARKERS: ClassVar[tuple[str, ...]] = ( "<" + _DSML + "tool_call", # marked (covers tool_call / tool_calls) "<" + _DSML + "invoke", # marked invoke " Tuple[str, List[ToolCall]]: content = text[:start] region = text[start:] - calls: List[Tuple[str, Dict[str, Any]]] = [] + calls: list[tuple[str, dict[str, Any]]] = [] invokes = list(_INVOKE_RE.finditer(region)) if invokes: for m in invokes: name = m.group(1) body = m.group(2) or "" # None for self-closing types = param_types.get(name, {}) - args: Dict[str, Any] = { + args: dict[str, Any] = { pm.group(1): _coerce( pm.group(3), pm.group(2), types.get(pm.group(1)) ) @@ -155,7 +153,7 @@ def parse(cls, text: str, tools: Optional[list]) -> Tuple[str, List[ToolCall]]: parsed = json.loads(stripped) if isinstance(parsed, dict): args = parsed - except Exception: + except (ValueError, TypeError): pass args = _unwrap_wrapper_args(args, set(types)) calls.append((name, args)) diff --git a/atom/entrypoints/openai/tool_parser/glm_tool_parser.py b/atom/entrypoints/openai/tool_parser/glm_tool_parser.py index 7077a98ffe..08bbdc7dc4 100644 --- a/atom/entrypoints/openai/tool_parser/glm_tool_parser.py +++ b/atom/entrypoints/openai/tool_parser/glm_tool_parser.py @@ -17,7 +17,7 @@ import json import re -from typing import Any, ClassVar, Dict, List, Optional, Tuple +from typing import Any, ClassVar from .qwen3_tool_parser import QWEN_TOOL_PREFIX from .schema import build_param_types, coerce_json_or_raw @@ -33,7 +33,7 @@ class GlmParser(BufferedMarkerParser): NAME: ClassVar[str] = "glm" - START_MARKERS: ClassVar[Tuple[str, ...]] = ("",) + START_MARKERS: ClassVar[tuple[str, ...]] = ("",) @classmethod def detect(cls, text: str) -> bool: @@ -43,14 +43,14 @@ def detect(cls, text: str) -> bool: return "" in text or "" in text @classmethod - def parse(cls, text: str, tools: Optional[list]) -> Tuple[str, List[ToolCall]]: + def parse(cls, text: str, tools: list | None) -> tuple[str, list[ToolCall]]: """Parse GLM tool calls; return (leading_content, tool_calls).""" param_types = build_param_types(tools) start = text.find("") if start == -1: return text.strip(), [] content = text[:start] - tool_calls: List[ToolCall] = [] + tool_calls: list[ToolCall] = [] for m in _TOOLCALL_RE.finditer(text): body = m.group(1) if m.group(1) is not None else m.group(2) if not body: @@ -60,7 +60,7 @@ def parse(cls, text: str, tools: Optional[list]) -> Tuple[str, List[ToolCall]]: if not name: continue types = param_types.get(name, {}) - args: Dict[str, Any] = {} + args: dict[str, Any] = {} for pm in _ARG_RE.finditer(body): k = pm.group(1).strip() if k: diff --git a/atom/entrypoints/openai/tool_parser/kimi_tool_parser.py b/atom/entrypoints/openai/tool_parser/kimi_tool_parser.py index b820ba568d..792ff9e45c 100644 --- a/atom/entrypoints/openai/tool_parser/kimi_tool_parser.py +++ b/atom/entrypoints/openai/tool_parser/kimi_tool_parser.py @@ -18,7 +18,7 @@ """ import re -from typing import ClassVar, List, Optional, Tuple +from typing import ClassVar from .tool_parser import ToolCall, ToolCallParser @@ -40,7 +40,7 @@ ) -def _parse_entries(section_text: str) -> List[ToolCall]: +def _parse_entries(section_text: str) -> list[ToolCall]: """Parse individual tool call entries from the section content.""" tool_calls = [] for match in _ENTRY_RE.finditer(section_text): @@ -68,7 +68,7 @@ def detect(cls, text: str) -> bool: return KIMI_SECTION_BEGIN in text @classmethod - def parse(cls, text: str, tools: Optional[list]) -> Tuple[str, List[ToolCall]]: + def parse(cls, text: str, tools: list | None) -> tuple[str, list[ToolCall]]: section_match = _SECTION_RE.search(text) if not section_match: # Unclosed section: the model was cut off mid-block; salvage whatever diff --git a/atom/entrypoints/openai/tool_parser/minimax_tool_parser.py b/atom/entrypoints/openai/tool_parser/minimax_tool_parser.py index 81df815dbe..a78998aa80 100644 --- a/atom/entrypoints/openai/tool_parser/minimax_tool_parser.py +++ b/atom/entrypoints/openai/tool_parser/minimax_tool_parser.py @@ -19,7 +19,7 @@ import json import re -from typing import Any, ClassVar, Dict, List, Optional, Tuple +from typing import Any, ClassVar from .schema import build_param_types, coerce_json_or_raw from .tool_parser import BufferedMarkerParser, ToolCall, unique_tool_call_id @@ -35,9 +35,9 @@ class MiniMaxParser(BufferedMarkerParser): NAME: ClassVar[str] = "minimax" - START_MARKERS: ClassVar[Tuple[str, ...]] = (MINIMAX_NS, "") + START_MARKERS: ClassVar[tuple[str, ...]] = (MINIMAX_NS, "") # The ns_token starts with ']', so a trailing ']' may be a partial marker. - HOLDBACK_CHARS: ClassVar[Tuple[str, ...]] = ("<", "]") + HOLDBACK_CHARS: ClassVar[tuple[str, ...]] = ("<", "]") @classmethod def detect(cls, text: str) -> bool: @@ -45,13 +45,13 @@ def detect(cls, text: str) -> bool: return MINIMAX_NS in text @classmethod - def parse(cls, text: str, tools: Optional[list]) -> Tuple[str, List[ToolCall]]: + def parse(cls, text: str, tools: list | None) -> tuple[str, list[ToolCall]]: """Parse MiniMax-M3 tool calls; return (leading_content, tool_calls).""" param_types = build_param_types(tools) clean = text.replace(MINIMAX_NS, "") tc = clean.find("") content = clean[:tc] if tc > 0 else ("" if tc == 0 else clean) - tool_calls: List[ToolCall] = [] + tool_calls: list[ToolCall] = [] for m in _INVOKE_RE.finditer(clean): name = m.group(1) if m.group(1) is not None else m.group(3) body = m.group(2) if m.group(2) is not None else (m.group(4) or "") @@ -59,7 +59,7 @@ def parse(cls, text: str, tools: Optional[list]) -> Tuple[str, List[ToolCall]]: continue name = name.strip() types = param_types.get(name, {}) - args: Dict[str, Any] = {} + args: dict[str, Any] = {} for pm in _PARAM_RE.finditer(body): k = pm.group(1).strip() if k: diff --git a/atom/entrypoints/openai/tool_parser/qwen3_tool_parser.py b/atom/entrypoints/openai/tool_parser/qwen3_tool_parser.py index bfc86b3b4b..55c9551b73 100644 --- a/atom/entrypoints/openai/tool_parser/qwen3_tool_parser.py +++ b/atom/entrypoints/openai/tool_parser/qwen3_tool_parser.py @@ -17,7 +17,7 @@ import json import re -from typing import Any, ClassVar, Dict, List, Optional, Tuple +from typing import Any, ClassVar from .kimi_tool_parser import KIMI_SECTION_BEGIN from .schema import build_param_types, coerce_param_value @@ -35,8 +35,8 @@ def _parse_function( - fn_text: str, param_types: Dict[str, Dict[str, Any]] -) -> Optional[ToolCall]: + fn_text: str, param_types: dict[str, dict[str, Any]] +) -> ToolCall | None: """Parse the inside of one ``...`` block into a ToolCall.""" gt = fn_text.find(">") if gt == -1: @@ -46,7 +46,7 @@ def _parse_function( return None body = fn_text[gt + 1 :] types = param_types.get(name, {}) - args: Dict[str, Any] = {} + args: dict[str, Any] = {} for pm in _PARAM_RE.finditer(body): seg = pm.group(1) if seg is None: @@ -67,7 +67,7 @@ def _parse_function( class QwenXmlParser(BufferedMarkerParser): NAME: ClassVar[str] = "qwen" - START_MARKERS: ClassVar[Tuple[str, ...]] = ("", QWEN_TOOL_PREFIX) + START_MARKERS: ClassVar[tuple[str, ...]] = ("", QWEN_TOOL_PREFIX) @classmethod def detect(cls, text: str) -> bool: @@ -75,13 +75,13 @@ def detect(cls, text: str) -> bool: return QWEN_TOOL_PREFIX in text and KIMI_SECTION_BEGIN not in text @classmethod - def parse(cls, text: str, tools: Optional[list]) -> Tuple[str, List[ToolCall]]: + def parse(cls, text: str, tools: list | None) -> tuple[str, list[ToolCall]]: """Parse Qwen3 XML tool calls; return (leading_content, tool_calls).""" param_types = build_param_types(tools) # Content precedes the first tool marker. start = cls.find_start(text) content = text[:start] if start != -1 else text - tool_calls: List[ToolCall] = [] + tool_calls: list[ToolCall] = [] for fm in _FUNCTION_RE.finditer(text): fn_text = fm.group(1) if fm.group(1) is not None else fm.group(2) if not fn_text: diff --git a/atom/entrypoints/openai/tool_parser/registry.py b/atom/entrypoints/openai/tool_parser/registry.py index fd4d048ffc..7f854eaccc 100644 --- a/atom/entrypoints/openai/tool_parser/registry.py +++ b/atom/entrypoints/openai/tool_parser/registry.py @@ -10,8 +10,6 @@ re-reading the notes on each entry. """ -from typing import List, Optional, Tuple, Type - from .deepseekv4_tool_parser import DsmlParser from .glm_tool_parser import GlmParser from .kimi_tool_parser import KIMI_SECTION_BEGIN, KimiParser @@ -26,7 +24,7 @@ # prefixes every tag with the ns_token. # GLM before Qwen — both use ``; GLM never emits ` Tuple[str, List[ToolCall]]: + text: str, tools: list | None = None +) -> tuple[str, list[ToolCall]]: """Parse tool calls from a complete model output. Args: diff --git a/atom/entrypoints/openai/tool_parser/schema.py b/atom/entrypoints/openai/tool_parser/schema.py index e3acd35101..162c574ea4 100644 --- a/atom/entrypoints/openai/tool_parser/schema.py +++ b/atom/entrypoints/openai/tool_parser/schema.py @@ -11,16 +11,16 @@ import ast import json -from typing import Any, Dict, Optional +from typing import Any -def build_param_types(tools: Optional[list]) -> Dict[str, Dict[str, Any]]: +def build_param_types(tools: list | None) -> dict[str, dict[str, Any]]: """Map ``function_name -> {param_name: json_schema_type}`` from request tools. Accepts OpenAI (``{"type": "function", "function": {...}}``) and bare (``{"name": ..., "parameters"/"input_schema": {...}}``) tool entries. """ - out: Dict[str, Dict[str, Any]] = {} + out: dict[str, dict[str, Any]] = {} for tool in tools or []: if not isinstance(tool, dict): continue @@ -64,9 +64,9 @@ def coerce_param_value(value: str, ptype: Any) -> Any: if t.startswith(("object", "dict", "map", "array", "list", "tuple")): try: return json.loads(v) - except Exception: + except (ValueError, TypeError): return ast.literal_eval(v) # safer for single-quoted Python literals - except Exception: + except Exception: # noqa: BLE001 return v return v @@ -83,5 +83,5 @@ def coerce_json_or_raw(value: str, ptype: Any) -> Any: s = v.strip() try: return json.loads(s) - except Exception: + except (ValueError, TypeError): return v diff --git a/atom/entrypoints/openai/tool_parser/stream.py b/atom/entrypoints/openai/tool_parser/stream.py index 7102205d0d..5a72ec5f29 100644 --- a/atom/entrypoints/openai/tool_parser/stream.py +++ b/atom/entrypoints/openai/tool_parser/stream.py @@ -4,7 +4,6 @@ """Streaming facade: sniff the format once, then delegate every chunk to it.""" from dataclasses import dataclass, field -from typing import Optional from .registry import EMIT_CONTENT, WAIT, sniff_stream from .tool_parser import ToolCallParser @@ -26,14 +25,14 @@ class ToolCallStreamParser: ends. """ - tools: Optional[list] = None + tools: list | None = None # Pre-detection accumulator. Once a format is chosen this is handed to the # concrete parser and never used again. _buf: str = "" - _parser: Optional[ToolCallParser] = field(default=None, repr=False) + _parser: ToolCallParser | None = field(default=None, repr=False) @property - def fmt(self) -> Optional[str]: + def fmt(self) -> str | None: """Detected format name, or None while still undecided.""" return self._parser.NAME if self._parser is not None else None diff --git a/atom/entrypoints/openai/tool_parser/tool_parser.py b/atom/entrypoints/openai/tool_parser/tool_parser.py index 66c4733f6d..a9838ad437 100644 --- a/atom/entrypoints/openai/tool_parser/tool_parser.py +++ b/atom/entrypoints/openai/tool_parser/tool_parser.py @@ -14,7 +14,7 @@ import uuid from abc import ABC, abstractmethod from dataclasses import dataclass -from typing import Any, ClassVar, Dict, List, Optional, Tuple +from typing import Any, ClassVar def unique_tool_call_id() -> str: @@ -31,9 +31,9 @@ class ToolCall: id: str type: str - function: Dict[str, str] + function: dict[str, str] - def to_dict(self) -> Dict[str, Any]: + def to_dict(self) -> dict[str, Any]: return {"id": self.id, "type": self.type, "function": self.function} @@ -46,7 +46,7 @@ class ToolCallParser(ABC): NAME: ClassVar[str] - def __init__(self, tools: Optional[list] = None): + def __init__(self, tools: list | None = None): self.tools = tools self.buf = "" # 0 = still in plain content, 1 = inside a tool-call region. Kimi adds @@ -63,7 +63,7 @@ def detect(cls, text: str) -> bool: @classmethod @abstractmethod - def parse(cls, text: str, tools: Optional[list]) -> Tuple[str, List[ToolCall]]: + def parse(cls, text: str, tools: list | None) -> tuple[str, list[ToolCall]]: """Parse a complete output into ``(leading_content, tool_calls)``.""" # -- streaming ---------------------------------------------------------- @@ -111,11 +111,11 @@ class BufferedMarkerParser(ToolCallParser): """ # Any of these opening the tool-call region; the earliest one wins. - START_MARKERS: ClassVar[Tuple[str, ...]] = () + START_MARKERS: ClassVar[tuple[str, ...]] = () # While no marker has been seen, a trailing run starting with one of these # may be the first bytes of a marker, so it is held back rather than emitted # as content (it would otherwise leak '<' into the user-visible text). - HOLDBACK_CHARS: ClassVar[Tuple[str, ...]] = ("<",) + HOLDBACK_CHARS: ClassVar[tuple[str, ...]] = ("<",) @classmethod def find_start(cls, text: str) -> int: From c1cb3c5320be9398a340517b61799a53bb9a4a11 Mon Sep 17 00:00:00 2001 From: yhl-amd Date: Thu, 30 Jul 2026 08:27:44 +0000 Subject: [PATCH 18/18] style: clear remaining Ruff findings in api_server global block Modernize the module-global type hints (Dict/Optional -> dict / X|None) and mark the shared typing import noqa: UP035 (kept for the file's many pre-existing annotations). Resolves the diff-context Ruff failures. Co-Authored-By: Claude Opus 4.8 --- atom/entrypoints/openai/api_server.py | 12 ++++++------ 1 file changed, 6 insertions(+), 6 deletions(-) diff --git a/atom/entrypoints/openai/api_server.py b/atom/entrypoints/openai/api_server.py index dba63bbcf5..0587e57dd2 100644 --- a/atom/entrypoints/openai/api_server.py +++ b/atom/entrypoints/openai/api_server.py @@ -23,7 +23,7 @@ import uuid from asyncio import AbstractEventLoop from contextlib import asynccontextmanager -from typing import Any, AsyncGenerator, Dict, List, Optional, Tuple +from typing import Any, AsyncGenerator, Dict, List, Optional, Tuple # noqa: UP035 import uvicorn from fastapi import FastAPI, HTTPException, Request @@ -84,12 +84,12 @@ # ============================================================================ engine = None -tokenizer: Optional[AutoTokenizer] = None -processor: Optional[Any] = None +tokenizer: AutoTokenizer | None = None +processor: Any | None = None model_name: str = "" -default_chat_template_kwargs: Dict[str, Any] = {} -custom_message_encoder: Optional[Any] = None -_stream_queues: Dict[str, asyncio.Queue] = {} +default_chat_template_kwargs: dict[str, Any] = {} +custom_message_encoder: Any | None = None +_stream_queues: dict[str, asyncio.Queue] = {} _seq_id_to_request_id: dict[int, str] = {} _stream_loops: dict[str, AbstractEventLoop] = {} _request_start_times: dict[str, float] = {}