diff --git a/atom/entrypoints/openai/api_server.py b/atom/entrypoints/openai/api_server.py index 6ff565249..a276726c6 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,72 @@ 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( @@ -440,17 +535,36 @@ def do_preprocess(): raise engine.core_mgr.add_request([seq]) - while True: - item = await token_queue.get() - token_ids = item.get("token_ids") or [] - if token_ids: - if first_token_at is None: - first_token_at = item.get("ts", time.time()) - last_token_at = item.get("ts", time.time()) - all_token_ids.extend(token_ids) - if item.get("finished", False): - finish_reason = item.get("finish_reason") - break + _finished_ok = False + try: + while True: + item = await token_queue.get() + token_ids = item.get("token_ids") or [] + if token_ids: + if first_token_at is None: + first_token_at = item.get("ts", time.time()) + last_token_at = item.get("ts", time.time()) + all_token_ids.extend(token_ids) + if item.get("finished", False): + finish_reason = item.get("finish_reason") + _finished_ok = True + break + finally: + # Two responsibilities, on EVERY exit path: + # 1) If we didn't finish (client disconnected / cancelled), tell the + # engine to stop so the seq doesn't run to max_tokens and burn GPU. + # 2) Always drop the seq from io_processor.requests. The engine frees + # its own KV on finish, but this dict is only cleaned up here for + # non-stream requests -- without an unconditional pop, every + # completed non-stream request leaks a Sequence (pending grows + # forever). Streaming pops via cleanup_streaming_request instead. + if seq is not None: + if not _finished_ok: + try: + engine.core_mgr.abort_request(seq.id) + except Exception: + pass + engine.io_processor.requests.pop(seq.id, None) text = tokenizer.decode(all_token_ids, skip_special_tokens=True) num_tokens_input = ( @@ -531,17 +645,29 @@ def do_preprocess(): raise engine.core_mgr.add_request([seq]) - while True: - item = await token_queue.get() - token_ids_out = item.get("token_ids") or [] - if token_ids_out: - if first_token_at is None: - first_token_at = item.get("ts", time.time()) - last_token_at = item.get("ts", time.time()) - all_token_ids.extend(token_ids_out) - if item.get("finished", False): - finish_reason = item.get("finish_reason") - break + _finished_ok = False + try: + while True: + item = await token_queue.get() + token_ids_out = item.get("token_ids") or [] + if token_ids_out: + if first_token_at is None: + first_token_at = item.get("ts", time.time()) + last_token_at = item.get("ts", time.time()) + all_token_ids.extend(token_ids_out) + if item.get("finished", False): + finish_reason = item.get("finish_reason") + _finished_ok = True + break + finally: + # See generate_async: abort on early exit, always pop to avoid leak. + if seq is not None: + if not _finished_ok: + try: + engine.core_mgr.abort_request(seq.id) + except Exception: + pass + engine.io_processor.requests.pop(seq.id, None) text = tokenizer.decode(all_token_ids, skip_special_tokens=True) num_tokens_output = len(all_token_ids) @@ -647,19 +773,31 @@ def do_preprocess(): engine.core_mgr.add_request(seqs) num_tokens_input = seqs[0].num_prompt_tokens - while not all(finished): - idx, item = await shared_queue.get() - if finished[idx]: - continue - tokens = item.get("token_ids") or [] - if tokens: - if per_first_token_at[idx] is None: - per_first_token_at[idx] = item.get("ts", time.time()) - per_last_token_at[idx] = item.get("ts", time.time()) - per_tokens[idx].extend(tokens) - if item.get("finished", False): - per_finish_reason[idx] = item.get("finish_reason") - finished[idx] = True + _all_finished = False + try: + while not all(finished): + idx, item = await shared_queue.get() + if finished[idx]: + continue + tokens = item.get("token_ids") or [] + if tokens: + if per_first_token_at[idx] is None: + per_first_token_at[idx] = item.get("ts", time.time()) + per_last_token_at[idx] = item.get("ts", time.time()) + per_tokens[idx].extend(tokens) + if item.get("finished", False): + per_finish_reason[idx] = item.get("finish_reason") + finished[idx] = True + _all_finished = True + finally: + # Abort any sibling still running on early exit; always pop all seqs. + for _seq in seqs: + if not _all_finished: + try: + engine.core_mgr.abort_request(_seq.id) + except Exception: + pass + engine.io_processor.requests.pop(_seq.id, None) finished_at = time.time() outputs: List[Dict[str, Any]] = [] @@ -779,9 +917,113 @@ def cleanup_streaming_request(request_id: str, seq_id: int) -> None: _seq_id_to_request_id.pop(seq_id, None) _stream_loops.pop(request_id, None) _request_start_times.pop(request_id, None) + # If the stream ended early (client disconnected) the seq may still be + # generating in the engine core -> tell it to stop so it doesn't run to + # max_tokens and pile up. No-op if the seq already finished. + try: + engine.core_mgr.abort_request(seq_id) + except Exception: + pass engine.io_processor.requests.pop(seq_id, None) +class _ClientDisconnected(Exception): + """Raised when a non-streaming client hangs up mid-generation.""" + + def __init__(self, request_id: str): + super().__init__(request_id) + self.request_id = request_id + + +async def _listen_for_disconnect(request) -> None: + """Block until the client sends an ``http.disconnect`` ASGI event. + + Unlike polling ``request.is_disconnected()`` on a timer, this awaits the + disconnect event directly, so detection is immediate and costs nothing while + the client stays connected. + """ + while True: + message = await request.receive() + if message["type"] == "http.disconnect": + break + + +async def _race_disconnect(coro, raw_request, request_id): + """Race an awaitable against client disconnect (vLLM ``with_cancellation`` + style). + + Starlette does NOT cancel a *non-streaming* request handler when the client + goes away (unlike StreamingResponse, which is cancelled on http.disconnect). + Without this, an abandoned non-stream request keeps ``await``-ing the engine + until it hits ``max_tokens`` -- burning GPU on output nobody will read AND + leaking the seq(s) in ``io_processor.requests`` (their finally never fires). + + We run ``coro`` (which produces the final result) as a task alongside a task + that awaits the ASGI ``http.disconnect`` event. Whichever finishes first + wins; the loser is cancelled. On disconnect, the coro's cancellation + propagates into its ``await`` points so its own ``try/finally`` runs -> + ``abort_request`` + ``io_processor.requests.pop`` (for fan-out, this aborts + every sibling). We then raise ``_ClientDisconnected``. + + ``request.receive()`` is safe here because FastAPI has already parsed the + request body into a pydantic model before this handler runs, so there is no + unread body for ``receive()`` to race against. + """ + handler_task = asyncio.ensure_future(coro) + + # No ASGI request object (e.g. internal call) -> just await the coro. + if raw_request is None: + return await handler_task + + disconnect_task = asyncio.ensure_future(_listen_for_disconnect(raw_request)) + + done, pending = await asyncio.wait( + [handler_task, disconnect_task], + return_when=asyncio.FIRST_COMPLETED, + ) + + # Cancel the loser and let its cancellation settle (drives the coro's own + # finally -> abort_request when the handler is the loser). Only swallow the + # expected CancelledError; log anything else, and let BaseException + # (KeyboardInterrupt/SystemExit) propagate. + for task in pending: + task.cancel() + for task in pending: + try: + await task + except asyncio.CancelledError: + pass + except Exception: + logger.warning( + f"Error tearing down cancelled task for request {request_id}", + exc_info=True, + ) + + if handler_task in done: + return handler_task.result() + + logger.info(f"Client disconnected (non-stream), aborting request {request_id}") + raise _ClientDisconnected(request_id) + + +async def _run_nonstream_with_disconnect(agen, raw_request, request_id): + """Drive a non-stream ``generate_async*`` async-*generator* while watching + for client disconnect. + + Thin wrapper over :func:`_race_disconnect` that collects the generator's + last yielded output. Use :func:`_race_disconnect` directly for the fan-out + path, whose ``generate_async_fanout`` is a coroutine returning a list. + """ + + async def _collect(): + final_output = None + async for output in agen: + final_output = output + return final_output + + return await _race_disconnect(_collect(), raw_request, request_id) + + async def setup_streaming_request_fanout( prompt_or_tokens: str | List[int], sampling_params: SamplingParams, @@ -908,7 +1150,7 @@ async def general_error_handler(request: Request, exc: Exception): @app.post("/v1/chat/completions") -async def chat_completions(request: ChatCompletionRequest): +async def chat_completions(request: ChatCompletionRequest, raw_request: Request): """Handle chat completion requests (OpenAI-compatible).""" global engine, tokenizer, model_name @@ -1003,12 +1245,16 @@ async def chat_completions(request: ChatCompletionRequest): # Non-streaming if is_multimodal and effective_n > 1: - outputs = await generate_async_fanout( - token_ids, - sampling_params, + outputs = await _race_disconnect( + generate_async_fanout( + token_ids, + sampling_params, + request_id, + multimodal_data=multimodal_data, + kv_transfer_params=request.kv_transfer_params, + ), + raw_request, request_id, - multimodal_data=multimodal_data, - kv_transfer_params=request.kv_transfer_params, ) if not outputs: raise RuntimeError("No output generated") @@ -1016,14 +1262,16 @@ async def chat_completions(request: ChatCompletionRequest): request_id, model_name, outputs, tools=request.tools ) elif is_multimodal: - final_output = None - async for output in generate_async_multimodal( - token_ids, - multimodal_data, - sampling_params, + final_output = await _run_nonstream_with_disconnect( + generate_async_multimodal( + token_ids, + multimodal_data, + sampling_params, + request_id, + ), + raw_request, request_id, - ): - final_output = output + ) if final_output is None: raise RuntimeError("No output generated") resp = build_chat_response( @@ -1034,11 +1282,15 @@ async def chat_completions(request: ChatCompletionRequest): tools=request.tools, ) elif effective_n > 1: - outputs = await generate_async_fanout( - prompt, - sampling_params, + outputs = await _race_disconnect( + generate_async_fanout( + prompt, + sampling_params, + request_id, + kv_transfer_params=request.kv_transfer_params, + ), + raw_request, request_id, - kv_transfer_params=request.kv_transfer_params, ) if not outputs: raise RuntimeError("No output generated") @@ -1046,14 +1298,16 @@ async def chat_completions(request: ChatCompletionRequest): request_id, model_name, outputs, tools=request.tools ) else: - final_output = None - async for output in generate_async( - prompt, - sampling_params, + final_output = await _run_nonstream_with_disconnect( + generate_async( + prompt, + sampling_params, + request_id, + kv_transfer_params=request.kv_transfer_params, + ), + raw_request, request_id, - kv_transfer_params=request.kv_transfer_params, - ): - final_output = output + ) if final_output is None: raise RuntimeError("No output generated") resp = build_chat_response( @@ -1066,6 +1320,9 @@ async def chat_completions(request: ChatCompletionRequest): _log_request_event("response", request_id, resp.model_dump()) return resp + except _ClientDisconnected: + # Client hung up; seq already aborted + popped. Nothing to return. + return JSONResponse(status_code=499, content={"detail": "client disconnected"}) except ValueError as e: logger.error(f"Validation error in chat_completions: {e}") raise HTTPException(status_code=400, detail=str(e)) @@ -1075,7 +1332,7 @@ async def chat_completions(request: ChatCompletionRequest): @app.post("/v1/completions") -async def completions(request: CompletionRequest): +async def completions(request: CompletionRequest, raw_request: Request): """Handle text completion requests (OpenAI-compatible).""" global engine, tokenizer, model_name @@ -1138,26 +1395,32 @@ async def completions(request: CompletionRequest): # Non-streaming if effective_n > 1: - outputs = await generate_async_fanout( - request.prompt, - sampling_params, + outputs = await _race_disconnect( + generate_async_fanout( + request.prompt, + sampling_params, + request_id, + kv_transfer_params=request.kv_transfer_params, + data_parallel_rank=request.data_parallel_rank, + ), + raw_request, request_id, - kv_transfer_params=request.kv_transfer_params, - data_parallel_rank=request.data_parallel_rank, ) if not outputs: raise RuntimeError("No output generated") resp = build_completion_response_multi(request_id, model_name, outputs) else: - final_output = None - async for output in generate_async( - request.prompt, - sampling_params, + final_output = await _run_nonstream_with_disconnect( + generate_async( + request.prompt, + sampling_params, + request_id, + kv_transfer_params=request.kv_transfer_params, + data_parallel_rank=request.data_parallel_rank, + ), + raw_request, request_id, - kv_transfer_params=request.kv_transfer_params, - data_parallel_rank=request.data_parallel_rank, - ): - final_output = output + ) if final_output is None: raise RuntimeError("No output generated") @@ -1166,6 +1429,9 @@ async def completions(request: CompletionRequest): _log_request_event("response", request_id, resp.model_dump()) return resp + except _ClientDisconnected: + # Client hung up; seq already aborted + popped. Nothing to return. + return JSONResponse(status_code=499, content={"detail": "client disconnected"}) except ValueError as e: logger.error(f"Validation error in completions: {e}") raise HTTPException(status_code=400, detail=str(e)) @@ -1255,6 +1521,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 @@ -1419,15 +1686,19 @@ 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): @@ -1444,6 +1715,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( @@ -1455,6 +1729,231 @@ 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 d3fb0462c..9021ab85d 100644 --- a/atom/entrypoints/openai/chat_encoders.py +++ b/atom/entrypoints/openai/chat_encoders.py @@ -85,6 +85,50 @@ 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 +141,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 6fb8a8e00..f7cd8f0b4 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_chat.py b/atom/entrypoints/openai/serving_chat.py index 7c707e968..9c21e8534 100644 --- a/atom/entrypoints/openai/serving_chat.py +++ b/atom/entrypoints/openai/serving_chat.py @@ -61,56 +61,78 @@ async def stream_chat_response( cleanup_fn, tools=None, ) -> AsyncGenerator[str, None]: - """Generate streaming chat completion response with reasoning and tool calls. - - Yields SSE chunks with: - - reasoning_content deltas during thinking phase - - content deltas for the answer - - tool_calls deltas when model invokes tools - - ``num_prompt_tokens`` is the engine-computed prompt length (``Sequence. - num_prompt_tokens``); reusing it avoids re-tokenizing the prompt on the - event loop at stream start. - """ - num_tokens_input = num_prompt_tokens - num_tokens_output = 0 - num_cached_tokens = 0 - reasoning_filter = ReasoningFilter() - tool_parser = ToolCallStreamParser(tools=tools) - has_tool_calls = False - - # Send initial role chunk - yield create_chat_chunk(request_id, model, delta={"role": "assistant"}) - - kv_transfer_params_value = None - - while True: - chunk_data = await stream_queue.get() - new_text = chunk_data["text"] - num_tokens_output += len(chunk_data.get("token_ids", [])) - _ct = chunk_data.get("num_cached_tokens", 0) - if _ct: - num_cached_tokens = _ct - - if "kv_transfer_params" in chunk_data: - kv_transfer_params_value = chunk_data["kv_transfer_params"] - - # Phase 1: Process through reasoning filter - segments = reasoning_filter.process(new_text) - if chunk_data.get("finished", False): - segments.extend(reasoning_filter.flush()) - - # Phase 2: For content segments, check for tool calls - for field, text in segments: - if field == "reasoning_content": - if text: - yield create_chat_chunk( - request_id, model, delta={"reasoning_content": text} - ) - elif field == "content": - # Run through tool parser - events = tool_parser.process(text) - for event_type, data in events: + try: + """Generate streaming chat completion response with reasoning and tool calls. + + Yields SSE chunks with: + - reasoning_content deltas during thinking phase + - content deltas for the answer + - tool_calls deltas when model invokes tools + + ``num_prompt_tokens`` is the engine-computed prompt length (``Sequence. + num_prompt_tokens``); reusing it avoids re-tokenizing the prompt on the + event loop at stream start. + """ + num_tokens_input = num_prompt_tokens + num_tokens_output = 0 + num_cached_tokens = 0 + reasoning_filter = ReasoningFilter() + tool_parser = ToolCallStreamParser(tools=tools) + has_tool_calls = False + + # Send initial role chunk + yield create_chat_chunk(request_id, model, delta={"role": "assistant"}) + + kv_transfer_params_value = None + + while True: + chunk_data = await stream_queue.get() + new_text = chunk_data["text"] + num_tokens_output += len(chunk_data.get("token_ids", [])) + _ct = chunk_data.get("num_cached_tokens", 0) + if _ct: + num_cached_tokens = _ct + + if "kv_transfer_params" in chunk_data: + kv_transfer_params_value = chunk_data["kv_transfer_params"] + + # Phase 1: Process through reasoning filter + segments = reasoning_filter.process(new_text) + if chunk_data.get("finished", False): + segments.extend(reasoning_filter.flush()) + + # Phase 2: For content segments, check for tool calls + for field, text in segments: + if field == "reasoning_content": + if text: + yield create_chat_chunk( + request_id, model, delta={"reasoning_content": text} + ) + elif field == "content": + # Run through tool parser + events = tool_parser.process(text) + for event_type, data in events: + if event_type == "content": + yield create_chat_chunk( + request_id, model, delta={"content": data} + ) + elif event_type == "tool_call_start": + has_tool_calls = True + yield create_chat_chunk( + request_id, + model, + delta={"tool_calls": [data]}, + ) + elif event_type == "tool_call_args": + yield create_chat_chunk( + request_id, + model, + delta={"tool_calls": [data]}, + ) + + if chunk_data.get("finished", False): + # Flush tool parser + for event_type, data in tool_parser.flush(): if event_type == "content": yield create_chat_chunk( request_id, model, delta={"content": data} @@ -118,60 +140,44 @@ async def stream_chat_response( elif event_type == "tool_call_start": has_tool_calls = True yield create_chat_chunk( - request_id, - model, - delta={"tool_calls": [data]}, + request_id, model, delta={"tool_calls": [data]} ) elif event_type == "tool_call_args": yield create_chat_chunk( - request_id, - model, - delta={"tool_calls": [data]}, + request_id, model, delta={"tool_calls": [data]} ) - - if chunk_data.get("finished", False): - # Flush tool parser - for event_type, data in tool_parser.flush(): - if event_type == "content": - yield create_chat_chunk(request_id, model, delta={"content": data}) - elif event_type == "tool_call_start": - has_tool_calls = True - yield create_chat_chunk( - request_id, model, delta={"tool_calls": [data]} - ) - elif event_type == "tool_call_args": - yield create_chat_chunk( - request_id, model, delta={"tool_calls": [data]} - ) - break - - cleanup_fn(request_id, seq_id) - - # Final chunks - finish_reason = "tool_calls" if has_tool_calls else "stop" - usage = { - "prompt_tokens": num_tokens_input, - "completion_tokens": num_tokens_output, - "total_tokens": num_tokens_input + num_tokens_output, - "prompt_tokens_details": {"cached_tokens": num_cached_tokens}, - } - usage_chunk = { - "id": request_id, - "object": CHAT_COMPLETION_CHUNK_OBJECT, - "created": int(time.time()), - "model": model, - "usage": usage, - } - if kv_transfer_params_value is not None: - usage_chunk["kv_transfer_params"] = kv_transfer_params_value - # Coalesce finish + usage + [DONE] into one send: at a wave boundary many - # requests finalize at once, so collapsing 3 socket writes/req to 1 cuts - # the syscalls that saturate the API event loop. - yield ( - create_chat_chunk(request_id, model, finish_reason=finish_reason) - + f"data: {json.dumps(usage_chunk)}\n\n" - + STREAM_DONE_MESSAGE - ) + break + + cleanup_fn(request_id, seq_id) + + # Final chunks + finish_reason = "tool_calls" if has_tool_calls else "stop" + usage = { + "prompt_tokens": num_tokens_input, + "completion_tokens": num_tokens_output, + "total_tokens": num_tokens_input + num_tokens_output, + "prompt_tokens_details": {"cached_tokens": num_cached_tokens}, + } + usage_chunk = { + "id": request_id, + "object": CHAT_COMPLETION_CHUNK_OBJECT, + "created": int(time.time()), + "model": model, + "choices": [], + "usage": usage, + } + if kv_transfer_params_value is not None: + usage_chunk["kv_transfer_params"] = kv_transfer_params_value + # Coalesce finish + usage + [DONE] into one send: at a wave boundary many + # requests finalize at once, so collapsing 3 socket writes/req to 1 cuts + # the syscalls that saturate the API event loop. + yield ( + create_chat_chunk(request_id, model, finish_reason=finish_reason) + + f"data: {json.dumps(usage_chunk)}\n\n" + + STREAM_DONE_MESSAGE + ) + finally: + cleanup_fn(request_id, seq_id) def _build_chat_choice( @@ -298,59 +304,84 @@ async def stream_chat_response_fanout( cleanup_fn, tools=None, ) -> AsyncGenerator[str, None]: - """Streaming variant that multiplexes ``len(seq_ids)`` fan-out siblings - into a single SSE stream, tagging every chunk with ``choices[0].index``. - - The shared queue receives ``(sibling_index, chunk_data)`` tuples from - the engine callbacks registered in :func:`setup_streaming_request_fanout`. - Reasoning + tool-call state is kept independently per sibling. + try: + """Streaming variant that multiplexes ``len(seq_ids)`` fan-out siblings + into a single SSE stream, tagging every chunk with ``choices[0].index``. + + The shared queue receives ``(sibling_index, chunk_data)`` tuples from + the engine callbacks registered in :func:`setup_streaming_request_fanout`. + Reasoning + tool-call state is kept independently per sibling. + + ``num_prompt_tokens`` is the engine-computed prompt length shared by all + siblings (they tokenize the same prompt once); reusing it avoids + re-tokenizing on the event loop at stream start. + """ + n = len(seq_ids) + num_tokens_input = num_prompt_tokens + num_tokens_output = [0] * n + num_cached_tokens = 0 + reasoning_filters = [ReasoningFilter() for _ in range(n)] + tool_parsers = [ToolCallStreamParser(tools=tools) for _ in range(n)] + has_tool_calls = [False] * n + finished = [False] * n + kv_transfer_params_value = None + + for i in range(n): + yield create_chat_chunk( + request_id, model, delta={"role": "assistant"}, index=i + ) - ``num_prompt_tokens`` is the engine-computed prompt length shared by all - siblings (they tokenize the same prompt once); reusing it avoids - re-tokenizing on the event loop at stream start. - """ - n = len(seq_ids) - num_tokens_input = num_prompt_tokens - num_tokens_output = [0] * n - reasoning_filters = [ReasoningFilter() for _ in range(n)] - tool_parsers = [ToolCallStreamParser(tools=tools) for _ in range(n)] - has_tool_calls = [False] * n - finished = [False] * n - kv_transfer_params_value = None - num_cached_tokens = 0 - - for i in range(n): - yield create_chat_chunk(request_id, model, delta={"role": "assistant"}, index=i) - - while not all(finished): - idx, chunk_data = await shared_queue.get() - if finished[idx]: - # Defensive: should not happen, engine emits finished once per seq. - continue - new_text = chunk_data["text"] - num_tokens_output[idx] += len(chunk_data.get("token_ids", [])) - _ct = chunk_data.get("num_cached_tokens", 0) - if _ct: - num_cached_tokens = _ct - - if "kv_transfer_params" in chunk_data: - kv_transfer_params_value = chunk_data["kv_transfer_params"] - - segments = reasoning_filters[idx].process(new_text) - if chunk_data.get("finished", False): - segments.extend(reasoning_filters[idx].flush()) - - for field, text in segments: - if field == "reasoning_content": - if text: - yield create_chat_chunk( - request_id, - model, - delta={"reasoning_content": text}, - index=idx, - ) - elif field == "content": - for event_type, data in tool_parsers[idx].process(text): + while not all(finished): + idx, chunk_data = await shared_queue.get() + if finished[idx]: + # Defensive: should not happen, engine emits finished once per seq. + continue + new_text = chunk_data["text"] + num_tokens_output[idx] += len(chunk_data.get("token_ids", [])) + _ct = chunk_data.get("num_cached_tokens", 0) + if _ct: + num_cached_tokens = _ct + + if "kv_transfer_params" in chunk_data: + kv_transfer_params_value = chunk_data["kv_transfer_params"] + + segments = reasoning_filters[idx].process(new_text) + if chunk_data.get("finished", False): + segments.extend(reasoning_filters[idx].flush()) + + for field, text in segments: + if field == "reasoning_content": + if text: + yield create_chat_chunk( + request_id, + model, + delta={"reasoning_content": text}, + index=idx, + ) + elif field == "content": + for event_type, data in tool_parsers[idx].process(text): + if event_type == "content": + yield create_chat_chunk( + request_id, model, delta={"content": data}, index=idx + ) + elif event_type == "tool_call_start": + has_tool_calls[idx] = True + yield create_chat_chunk( + request_id, + model, + delta={"tool_calls": [data]}, + index=idx, + ) + elif event_type == "tool_call_args": + yield create_chat_chunk( + request_id, + model, + delta={"tool_calls": [data]}, + index=idx, + ) + + if chunk_data.get("finished", False): + for event_type, data in tool_parsers[idx].flush(): if event_type == "content": yield create_chat_chunk( request_id, model, delta={"content": data}, index=idx @@ -370,61 +401,43 @@ async def stream_chat_response_fanout( delta={"tool_calls": [data]}, index=idx, ) - - if chunk_data.get("finished", False): - for event_type, data in tool_parsers[idx].flush(): - if event_type == "content": - yield create_chat_chunk( - request_id, model, delta={"content": data}, index=idx - ) - elif event_type == "tool_call_start": - has_tool_calls[idx] = True - yield create_chat_chunk( - request_id, - model, - delta={"tool_calls": [data]}, - index=idx, - ) - elif event_type == "tool_call_args": - yield create_chat_chunk( - request_id, - model, - delta={"tool_calls": [data]}, - index=idx, - ) - finished[idx] = True - - # Clean up all sibling seq_id entries then the shared request state. - for sid in seq_ids: - cleanup_fn(request_id, sid) - - usage = { - "prompt_tokens": num_tokens_input, - "completion_tokens": sum(num_tokens_output), - "total_tokens": num_tokens_input + sum(num_tokens_output), - "num_choices": n, - "prompt_tokens_details": {"cached_tokens": num_cached_tokens}, - } - usage_chunk = { - "id": request_id, - "object": CHAT_COMPLETION_CHUNK_OBJECT, - "created": int(time.time()), - "model": model, - "usage": usage, - } - if kv_transfer_params_value is not None: - usage_chunk["kv_transfer_params"] = kv_transfer_params_value - # Coalesce the per-sibling finish chunks + usage + [DONE] into one send. - yield ( - "".join( - create_chat_chunk( - request_id, - model, - finish_reason="tool_calls" if has_tool_calls[i] else "stop", - index=i, + finished[idx] = True + + # Clean up all sibling seq_id entries then the shared request state. + for sid in seq_ids: + cleanup_fn(request_id, sid) + + usage = { + "prompt_tokens": num_tokens_input, + "completion_tokens": sum(num_tokens_output), + "total_tokens": num_tokens_input + sum(num_tokens_output), + "num_choices": n, + "prompt_tokens_details": {"cached_tokens": num_cached_tokens}, + } + usage_chunk = { + "id": request_id, + "object": CHAT_COMPLETION_CHUNK_OBJECT, + "created": int(time.time()), + "model": model, + "choices": [], + "usage": usage, + } + if kv_transfer_params_value is not None: + usage_chunk["kv_transfer_params"] = kv_transfer_params_value + # Coalesce the per-sibling finish chunks + usage + [DONE] into one send. + yield ( + "".join( + create_chat_chunk( + request_id, + model, + finish_reason="tool_calls" if has_tool_calls[i] else "stop", + index=i, + ) + for i in range(n) ) - for i in range(n) + + f"data: {json.dumps(usage_chunk)}\n\n" + + STREAM_DONE_MESSAGE ) - + f"data: {json.dumps(usage_chunk)}\n\n" - + STREAM_DONE_MESSAGE - ) + finally: + for _sid in seq_ids: + cleanup_fn(request_id, _sid) diff --git a/atom/entrypoints/openai/serving_responses.py b/atom/entrypoints/openai/serving_responses.py new file mode 100644 index 000000000..26d22af82 --- /dev/null +++ b/atom/entrypoints/openai/serving_responses.py @@ -0,0 +1,718 @@ +"""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 549277e8d..21a98f5df 100644 --- a/atom/entrypoints/openai/tool_parser.py +++ b/atom/entrypoints/openai/tool_parser.py @@ -183,21 +183,391 @@ 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 +681,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 +1025,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 = [] diff --git a/atom/model_engine/engine_core_mgr.py b/atom/model_engine/engine_core_mgr.py index b254a792b..eddebc82c 100644 --- a/atom/model_engine/engine_core_mgr.py +++ b/atom/model_engine/engine_core_mgr.py @@ -460,6 +460,18 @@ def send_utility_command(self, cmd: str, dp_rank: int = None): copy=False, ) + def abort_request(self, req_id): + """Tell the engine core(s) to drop a request (client disconnected). + + Broadcast to every DP rank (only the one holding ``req_id`` acts). The + scheduler finishes the seq at its next step via the normal stop path, + freeing its KV blocks. Fire-and-forget; safe if the seq already finished. + """ + try: + self.broadcast_utility_command("abort_request", req_id=req_id) + except Exception as e: + logger.warning(f"{self.label}: abort_request({req_id}) failed: {e}") + def broadcast_utility_command(self, cmd: str, **kwargs): payload = {"cmd": cmd, **kwargs} # Serialize once and reuse for all ranks (optimization: avoid repeated pickle.dumps) diff --git a/atom/model_engine/engine_utility.py b/atom/model_engine/engine_utility.py index 49fb9219d..c1aed18b2 100644 --- a/atom/model_engine/engine_utility.py +++ b/atom/model_engine/engine_utility.py @@ -41,6 +41,7 @@ class EngineUtilityHandler: "stop_profile": "_handle_stop_profile", "get_mtp_stats": "_handle_get_mtp_stats", "get_mtp_statistics": "_handle_get_mtp_statistics", + "abort_request": "_handle_abort_request", } def __init__( @@ -210,6 +211,19 @@ def _handle_clear_kv_cache(self, args: dict): ("UTILITY_RESPONSE", {"cmd": "clear_kv_cache", "result": result}) ) + def _handle_abort_request(self, args: dict): + """Mark a sequence aborted (client disconnected) so the scheduler finishes + it at the next step via the normal stop path (frees KV, drops it).""" + req_id = args.get("req_id") if isinstance(args, dict) else None + if req_id is None or self.scheduler is None: + return + found = False + for seq in list(self.scheduler.running) + list(self.scheduler.waiting): + if seq.id == req_id: + seq.aborted = True + found = True + logger.info(f"{self.label}: abort_request req_id={req_id} found={found}") + def _handle_configure_hidden_states(self, args: dict): """Configure hidden states extraction on all model runners (TorchSpec).""" aux_layer_ids = args.get("aux_layer_ids", []) diff --git a/atom/model_engine/scheduler.py b/atom/model_engine/scheduler.py index 295a76f3e..166ae79f2 100644 --- a/atom/model_engine/scheduler.py +++ b/atom/model_engine/scheduler.py @@ -1430,6 +1430,11 @@ def postprocess( num_tokens = seq.num_tokens - self.mtp_k - num_rejected leave_reason = None + # Client disconnected -> finish now via the normal stop path (frees + # KV blocks, emits a finished RequestOutput). A natural stop below + # may still overwrite the reason; either way the seq terminates. + if getattr(seq, "aborted", False): + leave_reason = "aborted" # MTP edge case: `rejection_sampler` does NOT inspect EOS — it # only compares draft vs target_argmax for acceptance. So when # the verified token is EOS the kernel still emits 1+ accepted diff --git a/atom/model_engine/sequence.py b/atom/model_engine/sequence.py index e925823ac..ea73665df 100644 --- a/atom/model_engine/sequence.py +++ b/atom/model_engine/sequence.py @@ -86,6 +86,9 @@ def __init__( # out-of-window SWA blocks can be freed while compressed blocks persist). # Empty / unused for non-SWA models. self.swa_block_table = [] + # Set True when the client disconnected; the scheduler finishes the seq + # at the next step via the normal stop path (frees KV, emits finished). + self.aborted = False # Per-request cache slot index (filled by BlockManager.allocate()). # -1 = unallocated. The slot indexes into the per-req cache tensors # owned by ModelRunner (e.g. mamba_k_cache for GDN).