diff --git a/studio/backend/models/inference.py b/studio/backend/models/inference.py index 31c100dbe..1e7770a7f 100644 --- a/studio/backend/models/inference.py +++ b/studio/backend/models/inference.py @@ -1146,7 +1146,8 @@ class CompletionMessage(BaseModel): """The assistant's complete response message.""" role: Literal["assistant"] = "assistant" - content: str + # ``None`` on a pure tool-call turn (OpenAI content=null); string otherwise. + content: Optional[str] = None refusal: Optional[str] = None reasoning_content: Optional[str] = None tool_calls: Optional[list[dict]] = None diff --git a/studio/backend/routes/inference.py b/studio/backend/routes/inference.py index e31c03f7f..3f6f6a2a7 100644 --- a/studio/backend/routes/inference.py +++ b/studio/backend/routes/inference.py @@ -625,6 +625,66 @@ def _chat_final_chunk(completion_id, created, model_name, finish_reason) -> str: ) +def _chat_tool_calls_chunk(completion_id, created, model_name, tool_calls) -> str: + """Delta chunk carrying OpenAI tool-call deltas (sibling of ``_chat_content_chunk``).""" + return _chat_chunk_sse( + completion_id, + created, + model_name, + delta = ChoiceDelta(tool_calls = tool_calls), + finish_reason = None, + ) + + +def _sf_heal_events_to_sse( + events, + completion_id, + created, + model_name, + state, + parallel_tool_calls, + monitor_id = None, +): + """Serialize ``StreamToolCallHealer`` events into chat SSE lines. + + ``state["idx"]`` tracks the call index across ``feed``/``finalize``; + ``parallel_tool_calls is False`` caps promotion to one call (GGUF parity). + The monitor is fed from the same events the client receives, never the + healed-away markup.""" + lines = [] + for kind, value in events: + if kind == "text": + if value: + lines.append(_chat_content_chunk(completion_id, created, model_name, value)) + api_monitor.append_reply(monitor_id, value) + continue + if parallel_tool_calls is False and state["idx"] >= 1: + continue + lines.append( + _chat_tool_calls_chunk( + completion_id, + created, + model_name, + [ + { + "index": state["idx"], + "id": value["id"], + "type": "function", + "function": value["function"], + } + ], + ) + ) + _fn = value.get("function") or {} + api_monitor.append_reply( + monitor_id, + ("[tool_calls] " if state["idx"] == 0 else "; ") + + f"{_fn.get('name', '')}({_fn.get('arguments', '')})", + ) + state["idx"] += 1 + return lines + + def _rewrite_cmpl_id(raw: bytes) -> bytes: """Rewrite llama-server's chat-style ``chatcmpl-`` ids to the ``cmpl-`` prefix OpenAI's legacy /v1/completions use. Anchored on the ``"id":`` key @@ -7190,25 +7250,87 @@ async def openai_chat_completions( if payload.preserve_thinking is not None: gen_kwargs["preserve_thinking"] = payload.preserve_thinking + # ── Client-tool passthrough (safetensors + MLX) ────────────── + # Client tools (or tool-result history) without server-side tools: render + # tools into the template, generate one turn, heal text-form calls (#6801). + # supports_tools=False falls through to plain relay (GGUF gate parity). + _sf_has_tool_msgs = any(m.role == "tool" or m.tool_calls for m in payload.messages) + # Gate on _sf_use_tools (did the server-side path claim the request?), not + # raw mcp_enabled: an empty MCP registry must not silently drop client tools. + _sf_client_tools = ( + not _effective_enable_tools(payload) + and not _sf_use_tools + and image is None + and not _sf_is_gptoss + and _sf_features.get("supports_tools", False) + and ((payload.tools and len(payload.tools) > 0) or _sf_has_tool_msgs) + ) + _sf_heal = ( + heal_gate(payload.auto_heal_tool_calls, payload.tools, payload.tool_choice) + if _sf_client_tools + else None + ) + if _sf_client_tools: + # Re-derive from payload.messages so tool_calls / role="tool" history + # survives templating; fold system/developer into one leading system + # message (templates reject "developer") and clear prompt to avoid a dup. + gen_kwargs["messages"] = _set_or_prepend_system_message( + _structured_tool_history_for_local_template( + _flatten_content_parts_for_local_template(_openai_messages_for_passthrough(payload)) + ), + system_prompt, + ) + gen_kwargs["system_prompt"] = "" + # tool_choice="none": keep history templating but advertise no tools + # (heal_gate is off, markup would relay as prose). A forced function + # narrows templating to that one schema. Both mirror the GGUF path, + # where llama-server honors tool_choice itself. + _sf_tc = payload.tool_choice + _sf_forced = None + if isinstance(_sf_tc, dict) and isinstance(_sf_tc.get("function"), dict): + _sf_forced = _sf_tc["function"].get("name") + if _sf_tc == "none": + gen_kwargs["tools"] = None + elif isinstance(_sf_forced, str): + gen_kwargs["tools"] = [ + t + for t in payload.tools or [] + if isinstance(t, dict) + and isinstance(t.get("function"), dict) + and t["function"].get("name") == _sf_forced + ] or None + else: + gen_kwargs["tools"] = payload.tools + # Request-scoped usage/timings receptacle (filled at gen_done). stats_holder: dict = {} if payload.use_adapter is not None: - def generate(): + def generate(messages_override = None): + kw = ( + gen_kwargs + if messages_override is None + else {**gen_kwargs, "messages": messages_override} + ) return backend.generate_with_adapter_control( use_adapter = payload.use_adapter, cancel_event = cancel_event, stats_holder = stats_holder, - **gen_kwargs, + **kw, ) else: - def generate(): + def generate(messages_override = None): + kw = ( + gen_kwargs + if messages_override is None + else {**gen_kwargs, "messages": messages_override} + ) return backend.generate_chat_response( cancel_event = cancel_event, stats_holder = stats_holder, - **gen_kwargs, + **kw, ) # ── Streaming response ──────────────────────────────────────── @@ -7224,6 +7346,11 @@ async def openai_chat_completions( try: yield _chat_role_chunk(completion_id, created, model_name) + # Client-tool passthrough: heal text-form calls on the fly + # (None => relay verbatim). + healer = StreamToolCallHealer(_sf_heal, payload.tools) if _sf_heal else None + heal_state = {"idx": 0} + prev_text = "" # Split prefilled into reasoning_content deltas (GGUF parity); single turn, serves MLX. reasoning_extractor = _new_sf_reasoning_extractor() @@ -7255,22 +7382,76 @@ async def openai_chat_completions( prev_text = cumulative if not new_text: continue + # Split prefilled reasoning first (GGUF/MLX parity), + # then route only the visible text through the client-tool + # healer so tool markup inside a reasoning block is not promoted. reasoning_delta, visible_delta = reasoning_extractor.feed(new_text) if reasoning_delta: yield _chat_reasoning_chunk( completion_id, created, model_name, reasoning_delta ) if visible_delta: - api_monitor.append_reply(monitor_id, visible_delta) - yield _chat_content_chunk(completion_id, created, model_name, visible_delta) + if healer is None: + # Monitor mirrors the verbatim relay; with healing on, + # _sf_heal_events_to_sse records the healed events instead. + api_monitor.append_reply(monitor_id, visible_delta) + yield _chat_content_chunk( + completion_id, created, model_name, visible_delta + ) + else: + for line in _sf_heal_events_to_sse( + healer.feed(visible_delta), + completion_id, + created, + model_name, + heal_state, + payload.parallel_tool_calls, + monitor_id, + ): + yield line final_reasoning, final_visible = reasoning_extractor.finish() if final_reasoning: yield _chat_reasoning_chunk(completion_id, created, model_name, final_reasoning) if final_visible: - api_monitor.append_reply(monitor_id, final_visible) - yield _chat_content_chunk(completion_id, created, model_name, final_visible) - yield _chat_final_chunk(completion_id, created, model_name, "stop") + if healer is None: + api_monitor.append_reply(monitor_id, final_visible) + yield _chat_content_chunk(completion_id, created, model_name, final_visible) + else: + for line in _sf_heal_events_to_sse( + healer.feed(final_visible), + completion_id, + created, + model_name, + heal_state, + payload.parallel_tool_calls, + monitor_id, + ): + yield line + + # A cancelled stream must not promote buffered-but-incomplete + # markup: finalize()'s allow_incomplete heal would execute a tool + # the user just cancelled. Disconnect returns earlier; "Stop" only + # sets cancel_event, so guard on it here too. + _cancelled = cancel_event.is_set() + if healer is not None and not _cancelled: + for line in _sf_heal_events_to_sse( + healer.finalize(), + completion_id, + created, + model_name, + heal_state, + payload.parallel_tool_calls, + monitor_id, + ): + yield line + + _finish = ( + "tool_calls" + if (healer is not None and not _cancelled and healer.healed) + else "stop" + ) + yield _chat_final_chunk(completion_id, created, model_name, _finish) # Usage chunk (choices=[], usage set), same shape as the # GGUF path so the speed popover works for MLX too. # Request-scoped holder, so concurrent streams cannot @@ -7332,27 +7513,96 @@ async def openai_chat_completions( for token in generate(): full_text = token - # Split prefilled reasoning (GGUF parity); also covers MLX via the shared generate(). + # Split prefilled reasoning (GGUF parity); also covers MLX via + # the shared generate(). Client-tool healing then runs on the visible + # text so tool markup inside a reasoning block is never promoted. _reasoning_text, _visible_text = _extract_responses_reasoning( full_text, parse_think_markers = _sf_parse_think, reasoning_prefilled = _sf_reasoning_prefilled, ) - _plain_msg_kwargs = {"content": _visible_text} + # Client-tool passthrough: promote text-form calls; opt-in single + # nudge retry on unparseable tool markup. + _msg = {"role": "assistant", "content": _visible_text} if _reasoning_text: - _plain_msg_kwargs["reasoning_content"] = _reasoning_text + _msg["reasoning_content"] = _reasoning_text + _finish = "stop" + if _sf_heal: + if heal_openai_message(_msg, _sf_heal, payload.tools): + _finish = "tool_calls" + elif nudge_enabled(payload.nudge_tool_calls): + _data = { + "choices": [{"message": {"role": "assistant", "content": _visible_text}}] + } + if nudge_should_retry(_data, _sf_heal, payload.tools): + # A failed retry must not 500 the request; keep the first + # response (GGUF nudge parity). The retry's generate() + # overwrites stats_holder, so save the first attempt's stats + # and restore them if the retry is discarded. + _first_stats = stats_holder.get("stats") + try: + retry_text = "" + for token in generate( + [*gen_kwargs["messages"], *nudge_messages(_data, _sf_heal)] + ): + retry_text = token + # Re-split reasoning on the retry so its visible text is + # what heals into a call (and reaches the monitor). + _retry_reasoning, _retry_visible = _extract_responses_reasoning( + retry_text, + parse_think_markers = _sf_parse_think, + reasoning_prefilled = _sf_reasoning_prefilled, + ) + retry_msg = {"role": "assistant", "content": _retry_visible} + if _retry_reasoning: + retry_msg["reasoning_content"] = _retry_reasoning + if heal_openai_message(retry_msg, _sf_heal, payload.tools): + _visible_text, _msg, _finish = ( + _retry_visible, + retry_msg, + "tool_calls", + ) + else: + # Retry produced no healable call -> first response wins. + stats_holder["stats"] = _first_stats + except Exception as retry_exc: + logger.debug( + "Nudge retry failed; keeping first response: %s", retry_exc + ) + stats_holder["stats"] = _first_stats + # parallel_tool_calls=false: cap to one call (GGUF parity). + if payload.parallel_tool_calls is False: + _tcs = _msg.get("tool_calls") + if isinstance(_tcs, list) and len(_tcs) > 1: + _msg["tool_calls"] = _tcs[:1] + response = ChatCompletion( id = completion_id, created = created, model = model_name, choices = [ CompletionChoice( - message = CompletionMessage(**_plain_msg_kwargs), - finish_reason = "stop", + message = CompletionMessage( + content = _msg["content"], + reasoning_content = _msg.get("reasoning_content"), + tool_calls = _msg.get("tool_calls"), + ), + finish_reason = _finish, ) ], ) - api_monitor.set_reply(monitor_id, _visible_text) + _monitor_reply = _msg.get("content") or "" + if _finish == "tool_calls": + _tcs = _msg.get("tool_calls") or [] + _calls_text = "; ".join( + f"{(tc.get('function') or {}).get('name', '')}" + f"({(tc.get('function') or {}).get('arguments', '')})" + for tc in _tcs + ) + _monitor_reply = (_msg.get("content") or "") + ( + f"[tool_calls] {_calls_text}" if _calls_text else "" + ) + api_monitor.set_reply(monitor_id, _monitor_reply) _stats = stats_holder.get("stats") if _stats: _monitor_usage(monitor_id, _stats.get("usage")) @@ -11189,6 +11439,55 @@ def _openai_messages_for_passthrough(payload) -> list[dict]: return messages +def _flatten_content_parts_for_local_template(messages: list[dict]) -> list[dict]: + """Flatten OpenAI content-part lists to plain strings. + + Local text templates take string content and raise on part lists (e.g. a + remote ``image_url`` that leaves ``image is None``): keep the text parts, + drop the rest, like the plain non-GGUF path. GGUF keeps the parts.""" + out = [] + for msg in messages: + content = msg.get("content") + if isinstance(content, list): + text_parts = [ + part.get("text", "") + for part in content + if isinstance(part, dict) and part.get("type") == "text" + ] + msg = {**msg, "content": "\n".join(text_parts) if text_parts else ""} + out.append(msg) + return out + + +def _structured_tool_history_for_local_template(messages: list[dict]) -> list[dict]: + """Deserialize assistant ``tool_calls[].function.arguments`` JSON strings to + mappings for local templating. + + Clients send prior-turn arguments as JSON strings, but local templates take + mappings (some raise on strings). Only the internal messages copy is + rewritten; the HTTP response stays OpenAI-shaped and unparseable strings + are left untouched.""" + out = [] + for msg in messages: + tool_calls = msg.get("tool_calls") + if isinstance(tool_calls, list) and tool_calls: + new_calls = [] + for tc in tool_calls: + fn = tc.get("function") if isinstance(tc, dict) else None + args = fn.get("arguments") if isinstance(fn, dict) else None + if isinstance(args, str): + try: + parsed = json.loads(args) + except ValueError: + parsed = None + if isinstance(parsed, dict): + tc = {**tc, "function": {**fn, "arguments": parsed}} + new_calls.append(tc) + msg = {**msg, "tool_calls": new_calls} + out.append(msg) + return out + + def _openai_messages_for_gguf_chat(payload, is_vision: bool) -> tuple[list[dict], bool]: """Build llama-server messages for the standard GGUF chat path. diff --git a/studio/backend/tests/test_sf_client_tools_passthrough.py b/studio/backend/tests/test_sf_client_tools_passthrough.py new file mode 100644 index 000000000..01905b712 --- /dev/null +++ b/studio/backend/tests/test_sf_client_tools_passthrough.py @@ -0,0 +1,786 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Client-tools passthrough healing for the safetensors/MLX backend. + +Parity for #6801: when a NON-GGUF model is loaded and the request declares its +own ``tools`` with server-side tools OFF, text-form tool calls are promoted back +into structured ``tool_calls`` (declared tools only) via the shared healer. MLX +rides the same orchestrator path, so a single scripted backend covers both. +""" + +import asyncio +import json +from types import SimpleNamespace + +from models.inference import ChatCompletionRequest, ChatMessage +from routes.inference import openai_chat_completions +from core.inference.api_monitor import ApiMonitor + + +LOOKUP_TOOL = { + "type": "function", + "function": { + "name": "lookup", + "description": "Look something up", + "parameters": { + "type": "object", + "properties": {"q": {"type": "string"}}, + "required": ["q"], + }, + }, +} +SEARCH_TOOL = { + "type": "function", + "function": { + "name": "search", + "parameters": { + "type": "object", + "properties": {"query": {"type": "string"}}, + "required": ["query"], + }, + }, +} + +_CALL_XML = '{"name": "lookup", "arguments": {"q": "cats"}}' +_SEARCH_XML = '{"name": "search", "arguments": {"query": "dogs"}}' + + +class _Request: + state = SimpleNamespace() + url = SimpleNamespace(path = "/v1/chat/completions") + method = "POST" + scope: dict = {} + + async def is_disconnected(self): + return False + + +class _ScriptedBackend: + """Non-GGUF backend: ``generate_chat_response`` replays scripted + CUMULATIVE snapshots. ``responder(messages, tools)`` returns the snapshot + list for one generation, so nudge tests can vary output across turns.""" + + active_model_name = "sf-model" + + def __init__( + self, + responder, + *, + stats = None, + ): + self.models = { + "sf-model": { + "chat_template_info": {"template": " chatml"}, + "context_length": 2048, + } + } + self._responder = responder + self._stats = stats + self.calls: list = [] + self.reset_count = 0 + + def generate_chat_response( + self, + *, + messages, + tools = None, + stats_holder = None, + **kwargs, + ): + self.calls.append({"messages": messages, "tools": tools, **kwargs}) + snapshots = self._responder(messages, tools) + if stats_holder is not None and self._stats is not None: + stats_holder["stats"] = self._stats + for snap in snapshots: + yield snap + + def reset_generation_state(self): + self.reset_count += 1 + + +def _fixed(*snapshots): + """Responder that always replays the given cumulative snapshots.""" + return lambda messages, tools: list(snapshots) + + +def _llama_stub(): + return SimpleNamespace( + is_loaded = False, + supports_tools = False, + is_vision = False, + context_length = None, + ) + + +def _install( + monkeypatch, + backend, + *, + supports_tools = True, +): + import routes.inference as inf + from state.tool_policy import reset_tool_policy + + reset_tool_policy() + monitor = ApiMonitor(max_entries = 8) + monkeypatch.setattr(inf, "api_monitor", monitor) + monkeypatch.setattr(inf, "get_llama_cpp_backend", lambda: _llama_stub()) + monkeypatch.setattr(inf, "get_inference_backend", lambda: backend) + monkeypatch.setattr( + inf, + "_detect_safetensors_features", + lambda *a, **k: {"supports_tools": supports_tools}, + ) + return monitor + + +def _request(**kwargs): + base = dict(model = "default", messages = [ChatMessage(role = "user", content = "hi")]) + base.update(kwargs) + return ChatCompletionRequest(**base) + + +def _call(payload, monkeypatch, backend, **install_kwargs): + _install(monkeypatch, backend, **install_kwargs) + + async def _run(): + return await openai_chat_completions(payload, request = _Request(), current_subject = "u") + + return asyncio.run(_run()) + + +def _json_body(response): + return json.loads(response.body if hasattr(response, "body") else response.content) + + +def _collect_sse(response): + async def _run(): + return [c async for c in response.body_iterator] + + return asyncio.run(_run()) + + +def _sse_objects(chunks): + out = [] + for chunk in chunks: + if isinstance(chunk, bytes): + chunk = chunk.decode() + for line in str(chunk).splitlines(): + if line.startswith("data: "): + data = line.removeprefix("data: ") + if data != "[DONE]": + out.append(json.loads(data)) + return out + + +# ── Non-streaming ───────────────────────────────────────────────── + + +def test_xml_healed_to_tool_calls_non_streaming(monkeypatch): + backend = _ScriptedBackend(_fixed(_CALL_XML)) + payload = _request(tools = [LOOKUP_TOOL], stream = False) + body = _json_body(_call(payload, monkeypatch, backend)) + choice = body["choices"][0] + assert choice["finish_reason"] == "tool_calls" + assert choice["message"]["content"] is None + calls = choice["message"]["tool_calls"] + assert len(calls) == 1 + assert calls[0]["function"]["name"] == "lookup" + assert json.loads(calls[0]["function"]["arguments"]) == {"q": "cats"} + # The client tools reached the generator (template injection). + assert backend.calls[0]["tools"] == [LOOKUP_TOOL] + + +def test_undeclared_call_stays_text(monkeypatch): + xml = '{"name": "other", "arguments": {}}' + backend = _ScriptedBackend(_fixed(xml)) + payload = _request(tools = [LOOKUP_TOOL], stream = False) + body = _json_body(_call(payload, monkeypatch, backend)) + choice = body["choices"][0] + assert choice["finish_reason"] == "stop" + assert choice["message"].get("tool_calls") is None + assert choice["message"]["content"] == xml + + +def test_opt_out_relays_verbatim(monkeypatch): + backend = _ScriptedBackend(_fixed(_CALL_XML)) + payload = _request(tools = [LOOKUP_TOOL], stream = False, auto_heal_tool_calls = False) + body = _json_body(_call(payload, monkeypatch, backend)) + choice = body["choices"][0] + assert choice["finish_reason"] == "stop" + assert choice["message"].get("tool_calls") is None + assert choice["message"]["content"] == _CALL_XML + + +def test_env_kill_switch_relays_verbatim(monkeypatch): + import core.inference.passthrough_healing as ph + + monkeypatch.setattr(ph, "_HEALING_DISABLED", True) + backend = _ScriptedBackend(_fixed(_CALL_XML)) + payload = _request(tools = [LOOKUP_TOOL], stream = False) + body = _json_body(_call(payload, monkeypatch, backend)) + choice = body["choices"][0] + assert choice["finish_reason"] == "stop" + assert choice["message"].get("tool_calls") is None + assert choice["message"]["content"] == _CALL_XML + + +def test_no_tools_request_untouched(monkeypatch): + backend = _ScriptedBackend(_fixed("just a plain answer")) + payload = _request(stream = False) + body = _json_body(_call(payload, monkeypatch, backend)) + # No tools and no tool messages -> plain path, normal ChatCompletion. + choice = body["choices"][0] + assert choice["finish_reason"] == "stop" + assert choice["message"]["content"] == "just a plain answer" + assert choice["message"].get("tool_calls") is None + + +def test_prose_around_call_retained(monkeypatch): + text = "Let me look:\n" + _CALL_XML + "\ndone" + backend = _ScriptedBackend(_fixed(text)) + payload = _request(tools = [LOOKUP_TOOL], stream = False) + body = _json_body(_call(payload, monkeypatch, backend)) + choice = body["choices"][0] + assert choice["finish_reason"] == "tool_calls" + assert choice["message"]["content"] == "Let me look:\n\ndone" + assert choice["message"]["tool_calls"][0]["function"]["name"] == "lookup" + + +def test_empty_output_is_valid_stop(monkeypatch): + backend = _ScriptedBackend(_fixed("")) + payload = _request(tools = [LOOKUP_TOOL], stream = False) + body = _json_body(_call(payload, monkeypatch, backend)) + choice = body["choices"][0] + assert choice["finish_reason"] == "stop" + assert choice["message"]["content"] in ("", None) + assert choice["message"].get("tool_calls") is None + + +def test_tool_role_follow_up_turn_preserves_history(monkeypatch): + backend = _ScriptedBackend(_fixed("The weather is sunny.")) + payload = _request( + tools = [LOOKUP_TOOL], + stream = False, + messages = [ + ChatMessage(role = "user", content = "weather?"), + ChatMessage( + role = "assistant", + content = None, + tool_calls = [ + { + "id": "call_0", + "type": "function", + "function": {"name": "lookup", "arguments": '{"q": "weather"}'}, + } + ], + ), + ChatMessage(role = "tool", tool_call_id = "call_0", content = "sunny"), + ], + ) + body = _json_body(_call(payload, monkeypatch, backend)) + assert body["choices"][0]["message"]["content"] == "The weather is sunny." + # The tool history reached the generator intact (role=tool + assistant.tool_calls). + sent = backend.calls[0]["messages"] + roles = [m["role"] for m in sent] + assert "tool" in roles + assistant = next(m for m in sent if m["role"] == "assistant") + assert assistant.get("tool_calls") + + +def test_dict_arguments_history_does_not_crash(monkeypatch): + # Non-spec client: assistant tool_calls[].function.arguments as a dict. + backend = _ScriptedBackend(_fixed("ok")) + payload = _request( + tools = [LOOKUP_TOOL], + stream = False, + messages = [ + ChatMessage(role = "user", content = "hi"), + ChatMessage( + role = "assistant", + content = None, + tool_calls = [ + { + "id": "call_0", + "type": "function", + "function": {"name": "lookup", "arguments": {"q": "x"}}, + } + ], + ), + ChatMessage(role = "tool", tool_call_id = "call_0", content = "y"), + ], + ) + body = _json_body(_call(payload, monkeypatch, backend)) + assert body["choices"][0]["message"]["content"] == "ok" + + +def test_forced_tool_choice_narrows_promotion(monkeypatch): + # tool_choice forces `search`; a `lookup` text call must NOT promote. + backend = _ScriptedBackend(_fixed(_CALL_XML)) + payload = _request( + tools = [LOOKUP_TOOL, SEARCH_TOOL], + stream = False, + tool_choice = {"type": "function", "function": {"name": "search"}}, + ) + body = _json_body(_call(payload, monkeypatch, backend)) + choice = body["choices"][0] + assert choice["finish_reason"] == "stop" + assert choice["message"].get("tool_calls") is None + + +def test_parallel_cap_non_streaming(monkeypatch): + backend = _ScriptedBackend(_fixed(_CALL_XML + _SEARCH_XML)) + payload = _request(tools = [LOOKUP_TOOL, SEARCH_TOOL], stream = False, parallel_tool_calls = False) + body = _json_body(_call(payload, monkeypatch, backend)) + calls = body["choices"][0]["message"]["tool_calls"] + assert len(calls) == 1 + assert calls[0]["function"]["name"] == "lookup" + + +def test_usage_recorded_when_stats_present(monkeypatch): + stats = {"usage": {"prompt_tokens": 7, "completion_tokens": 3, "total_tokens": 10}} + backend = _ScriptedBackend(_fixed(_CALL_XML), stats = stats) + payload = _request(tools = [LOOKUP_TOOL], stream = False) + monitor = _install(monkeypatch, backend) + + async def _run(): + return await openai_chat_completions(payload, request = _Request(), current_subject = "u") + + asyncio.run(_run()) + [entry] = monitor.snapshot() + assert entry["prompt_tokens"] == 7 + assert entry["completion_tokens"] == 3 + + +# ── Nudge ───────────────────────────────────────────────────────── + + +def test_nudge_default_off_single_generation(monkeypatch): + # Signal present but unparseable; without opt-in, no retry. + truncated = '{"name": "lookup"' + backend = _ScriptedBackend(_fixed(truncated)) + payload = _request(tools = [LOOKUP_TOOL], stream = False) + _call(payload, monkeypatch, backend) + assert len(backend.calls) == 1 + + +def test_nudge_opt_in_retry_recovers(monkeypatch): + truncated = '{"name": "lookup"' + + def responder(messages, tools): + nudged = any( + "native tool-call format" in (m.get("content") or "") + for m in messages + if m.get("role") == "user" + ) + return [_CALL_XML] if nudged else [truncated] + + backend = _ScriptedBackend(responder) + payload = _request(tools = [LOOKUP_TOOL], stream = False, nudge_tool_calls = True) + body = _json_body(_call(payload, monkeypatch, backend)) + assert len(backend.calls) == 2 + choice = body["choices"][0] + assert choice["finish_reason"] == "tool_calls" + assert choice["message"]["tool_calls"][0]["function"]["name"] == "lookup" + + +def test_nudge_double_failure_relays_original(monkeypatch): + truncated = '{"name": "lookup"' + backend = _ScriptedBackend(_fixed(truncated)) + payload = _request(tools = [LOOKUP_TOOL], stream = False, nudge_tool_calls = True) + body = _json_body(_call(payload, monkeypatch, backend)) + assert len(backend.calls) == 2 # exactly one retry + choice = body["choices"][0] + assert choice["finish_reason"] == "stop" + assert choice["message"]["content"] == truncated + + +# ── Streaming ───────────────────────────────────────────────────── + + +def test_streaming_heals_split_call_into_one_delta(monkeypatch): + # Cumulative snapshots that build the call across many increments. + pieces = ["{"name": "loo', '{"name": "lookup", "argum'] + cumulative = pieces + [_CALL_XML] + backend = _ScriptedBackend(_fixed(*cumulative)) + payload = _request(tools = [LOOKUP_TOOL], stream = True) + response = _call(payload, monkeypatch, backend) + objs = _sse_objects(_collect_sse(response)) + tool_deltas = [ + tc + for o in objs + for tc in (o.get("choices", [{}])[0].get("delta", {}) or {}).get("tool_calls", []) or [] + ] + assert len(tool_deltas) == 1 + assert tool_deltas[0]["function"]["name"] == "lookup" + finishes = [ + o["choices"][0]["finish_reason"] + for o in objs + if o["choices"] and o["choices"][0].get("finish_reason") + ] + assert finishes == ["tool_calls"] + + +def test_streaming_cancel_does_not_finalize_tool_call(monkeypatch): + # A stream cancelled via the registry ("Stop") must NOT promote the + # buffered-but-unclosed tool markup at finalize, else it executes a tool + # the user just cancelled. Guarded on cancel_event at the finalize step. + import routes.inference as inf + + cancel_id = "cancel-me-6870" + # Balanced JSON but no closing -> healer HOLDS it until finalize. + held = '{"name": "lookup", "arguments": {"q": "cats"}}' + + class _CancelMidStream(_ScriptedBackend): + def __init__(self): + super().__init__(_fixed(held)) + + def generate_chat_response( + self, + *, + messages, + tools = None, + stats_holder = None, + **kwargs, + ): + self.calls.append({"messages": messages, "tools": tools, **kwargs}) + yield held # healer holds the unclosed call + inf._cancel_by_cancel_id_or_stash(cancel_id) # user hits Stop before EOF + + backend = _CancelMidStream() + payload = _request(tools = [LOOKUP_TOOL], stream = True, cancel_id = cancel_id) + response = _call(payload, monkeypatch, backend) + objs = _sse_objects(_collect_sse(response)) + tool_deltas = [ + tc + for o in objs + for tc in (o.get("choices", [{}])[0].get("delta", {}) or {}).get("tool_calls", []) or [] + ] + assert tool_deltas == [] # no tool promoted after cancel + finishes = [ + o["choices"][0]["finish_reason"] + for o in objs + if o["choices"] and o["choices"][0].get("finish_reason") + ] + assert "tool_calls" not in finishes # ends with finish_reason=stop, not tool_calls + + +def test_streaming_no_tools_verbatim(monkeypatch): + backend = _ScriptedBackend(_fixed("hello ", "hello world")) + payload = _request(stream = True) + response = _call(payload, monkeypatch, backend) + objs = _sse_objects(_collect_sse(response)) + text = "".join( + (o["choices"][0]["delta"].get("content") or "") + for o in objs + if o["choices"] and "delta" in o["choices"][0] + ) + assert text == "hello world" + finishes = [ + o["choices"][0]["finish_reason"] + for o in objs + if o["choices"] and o["choices"][0].get("finish_reason") + ] + assert finishes == ["stop"] + + +def test_streaming_repeated_snapshot_no_duplicate_call(monkeypatch): + # Repeated then shrunk cumulative snapshots must not double-heal. + backend = _ScriptedBackend(_fixed(_CALL_XML, _CALL_XML, _CALL_XML[:5], _CALL_XML)) + payload = _request(tools = [LOOKUP_TOOL], stream = True) + response = _call(payload, monkeypatch, backend) + objs = _sse_objects(_collect_sse(response)) + tool_deltas = [ + tc + for o in objs + for tc in (o.get("choices", [{}])[0].get("delta", {}) or {}).get("tool_calls", []) or [] + ] + assert len(tool_deltas) == 1 + + +def test_streaming_parallel_cap(monkeypatch): + backend = _ScriptedBackend(_fixed(_CALL_XML + _SEARCH_XML)) + payload = _request(tools = [LOOKUP_TOOL, SEARCH_TOOL], stream = True, parallel_tool_calls = False) + response = _call(payload, monkeypatch, backend) + objs = _sse_objects(_collect_sse(response)) + tool_deltas = [ + tc + for o in objs + for tc in (o.get("choices", [{}])[0].get("delta", {}) or {}).get("tool_calls", []) or [] + ] + assert len(tool_deltas) == 1 + assert tool_deltas[0]["function"]["name"] == "lookup" + + +def test_streaming_generator_error_closes_cleanly(monkeypatch): + def responder(messages, tools): + raise RuntimeError("boom /secret/path") + + backend = _ScriptedBackend(responder) + payload = _request(tools = [LOOKUP_TOOL], stream = True) + response = _call(payload, monkeypatch, backend) + chunks = _collect_sse(response) + joined = "".join(c.decode() if isinstance(c, bytes) else c for c in chunks) + assert "An internal error occurred" in joined + assert "secret/path" not in joined # CWE-209: no path leak + assert backend.reset_count >= 1 + + +def test_streaming_disconnect_resets_once(monkeypatch): + class _DisconnectRequest(_Request): + async def is_disconnected(self): + return True + + backend = _ScriptedBackend(_fixed("a", "ab", "abc")) + payload = _request(tools = [LOOKUP_TOOL], stream = True) + _install(monkeypatch, backend) + + async def _run(): + resp = await openai_chat_completions( + payload, request = _DisconnectRequest(), current_subject = "u" + ) + return [c async for c in resp.body_iterator] + + asyncio.run(_run()) + assert backend.reset_count == 1 + + +def test_mlx_uses_same_path(monkeypatch): + # MLX and safetensors share get_inference_backend(); one scripted backend covers both. + backend = _ScriptedBackend(_fixed(_CALL_XML)) + payload = _request(tools = [LOOKUP_TOOL], stream = False) + body = _json_body(_call(payload, monkeypatch, backend)) + assert body["choices"][0]["finish_reason"] == "tool_calls" + + +def test_tool_choice_none_does_not_advertise_tools(monkeypatch): + # tool_choice="none": no tools rendered into the template; history templating still applies. + backend = _ScriptedBackend(_fixed("plain answer")) + payload = _request(tools = [LOOKUP_TOOL], tool_choice = "none", stream = False) + body = _json_body(_call(payload, monkeypatch, backend)) + assert body["choices"][0]["message"]["content"] == "plain answer" + assert backend.calls[0]["tools"] is None + + +def test_developer_message_folded_into_system_prompt(monkeypatch): + # The "developer" role folds into one leading system message (local templates reject it). + backend = _ScriptedBackend(_fixed("ok")) + payload = _request( + messages = [ + ChatMessage(role = "developer", content = "always be terse"), + ChatMessage(role = "user", content = "hi"), + ], + tools = [LOOKUP_TOOL], + stream = False, + ) + _call(payload, monkeypatch, backend) + sent = backend.calls[0]["messages"] + assert sent[0]["role"] == "system" + assert "always be terse" in sent[0]["content"] + assert all(m.get("role") != "developer" for m in sent) + + +def test_failed_nudge_retry_keeps_original_response(monkeypatch): + # A raising retry must not 500; the first response is returned. + state = {"n": 0} + + def responder(messages, tools): + state["n"] += 1 + if state["n"] == 1: + return ['{"name":"lookup"'] # unhealable signal + raise RuntimeError("retry blew up") + + backend = _ScriptedBackend(responder) + payload = _request(tools = [LOOKUP_TOOL], nudge_tool_calls = True, stream = False) + body = _json_body(_call(payload, monkeypatch, backend)) + assert state["n"] == 2 + assert body["choices"][0]["finish_reason"] == "stop" + assert body["choices"][0]["message"]["content"] == '{"name":"lookup"' + + +def test_discarded_nudge_retry_reports_first_attempt_usage(monkeypatch): + # Double-failure nudge: the first response is delivered, but the retry's + # generate() overwrites stats_holder. The monitor must record the FIRST + # attempt's usage, not the discarded retry's. + first_stats = {"usage": {"prompt_tokens": 7, "completion_tokens": 3, "total_tokens": 10}} + retry_stats = {"usage": {"prompt_tokens": 99, "completion_tokens": 99, "total_tokens": 198}} + + class _PerCallStatsBackend(_ScriptedBackend): + def __init__(self): + # Unhealable truncated markup on both attempts -> retry is discarded. + super().__init__(lambda m, t: ['{"name":"lookup"']) + self._stats_seq = [first_stats, retry_stats] + + def generate_chat_response( + self, + *, + messages, + tools = None, + stats_holder = None, + **kwargs, + ): + self.calls.append({"messages": messages, "tools": tools, **kwargs}) + stats = self._stats_seq[min(len(self.calls) - 1, len(self._stats_seq) - 1)] + if stats_holder is not None: + stats_holder["stats"] = stats + for snap in self._responder(messages, tools): + yield snap + + backend = _PerCallStatsBackend() + payload = _request(tools = [LOOKUP_TOOL], nudge_tool_calls = True, stream = False) + monitor = _install(monkeypatch, backend) + + async def _run(): + return await openai_chat_completions(payload, request = _Request(), current_subject = "u") + + asyncio.run(_run()) + assert len(backend.calls) == 2 # first attempt + one discarded retry + [entry] = monitor.snapshot() + # The delivered response is the first attempt, so its usage must be reported. + assert entry["prompt_tokens"] == 7 + assert entry["completion_tokens"] == 3 + + +def test_monitor_records_healed_call_not_raw_xml(monkeypatch): + backend = _ScriptedBackend(_fixed(_CALL_XML)) + payload = _request(tools = [LOOKUP_TOOL], stream = False) + monitor = _install(monkeypatch, backend) + + async def _run(): + return await openai_chat_completions(payload, request = _Request(), current_subject = "u") + + asyncio.run(_run()) + snap = monitor.snapshot(include_details = True) + replies = json.dumps(snap) + assert "" not in replies + assert "lookup" in replies + + +def test_streaming_monitor_records_healed_call_not_raw_xml(monkeypatch): + # Monitor mirrors what the client received, never the healed-away raw markup. + backend = _ScriptedBackend( + _fixed("Sure. ", 'Sure. {"name": "loo', "Sure. " + _CALL_XML) + ) + payload = _request(tools = [LOOKUP_TOOL], stream = True) + monitor = _install(monkeypatch, backend) + + async def _run(): + return await openai_chat_completions(payload, request = _Request(), current_subject = "u") + + response = asyncio.run(_run()) + _collect_sse(response) + replies = json.dumps(monitor.snapshot(include_details = True)) + assert "" not in replies + assert "Sure. " in replies + assert "[tool_calls] lookup(" in replies + + +def test_forced_tool_choice_narrows_templated_tools(monkeypatch): + # A forced function is the only schema rendered into the template. + backend = _ScriptedBackend(_fixed(_SEARCH_XML)) + payload = _request( + tools = [LOOKUP_TOOL, SEARCH_TOOL], + stream = False, + tool_choice = {"type": "function", "function": {"name": "search"}}, + ) + body = _json_body(_call(payload, monkeypatch, backend)) + templated = backend.calls[0]["tools"] + assert [t["function"]["name"] for t in templated] == ["search"] + choice = body["choices"][0] + assert choice["finish_reason"] == "tool_calls" + assert choice["message"]["tool_calls"][0]["function"]["name"] == "search" + + +def test_multimodal_content_parts_flattened_for_local_template(monkeypatch): + # Remote image URLs leave image=None, so content arrives as a part LIST: + # text parts are kept, the image part dropped. + backend = _ScriptedBackend(_fixed(_CALL_XML)) + payload = _request( + messages = [ + ChatMessage( + role = "user", + content = [ + {"type": "text", "text": "what is this?"}, + { + "type": "image_url", + "image_url": {"url": "https://example.com/cat.png"}, + }, + ], + ) + ], + tools = [LOOKUP_TOOL], + stream = False, + ) + body = _json_body(_call(payload, monkeypatch, backend)) + templated = backend.calls[0]["messages"] + assert all(isinstance(m.get("content"), str) for m in templated) + assert any(m["content"] == "what is this?" for m in templated) + assert body["choices"][0]["finish_reason"] == "tool_calls" + + +def test_string_arguments_history_deserialized_for_template(monkeypatch): + # JSON-string tool_calls arguments become dicts in the templated copy; + # the HTTP response stays OpenAI-shaped. + backend = _ScriptedBackend(_fixed("done")) + payload = _request( + tools = [LOOKUP_TOOL], + stream = False, + messages = [ + ChatMessage(role = "user", content = "weather?"), + ChatMessage( + role = "assistant", + content = None, + tool_calls = [ + { + "id": "call_0", + "type": "function", + "function": {"name": "lookup", "arguments": '{"q": "weather"}'}, + } + ], + ), + ChatMessage(role = "tool", tool_call_id = "call_0", content = "sunny"), + ], + ) + _json_body(_call(payload, monkeypatch, backend)) + assistant = next(m for m in backend.calls[0]["messages"] if m["role"] == "assistant") + assert assistant["tool_calls"][0]["function"]["arguments"] == {"q": "weather"} + + +def test_unparseable_arguments_string_left_untouched(monkeypatch): + backend = _ScriptedBackend(_fixed("ok")) + payload = _request( + tools = [LOOKUP_TOOL], + stream = False, + messages = [ + ChatMessage(role = "user", content = "hi"), + ChatMessage( + role = "assistant", + content = None, + tool_calls = [ + { + "id": "call_0", + "type": "function", + "function": {"name": "lookup", "arguments": "not json {"}, + } + ], + ), + ChatMessage(role = "tool", tool_call_id = "call_0", content = "y"), + ], + ) + body = _json_body(_call(payload, monkeypatch, backend)) + assert body["choices"][0]["message"]["content"] == "ok" + assistant = next(m for m in backend.calls[0]["messages"] if m["role"] == "assistant") + assert assistant["tool_calls"][0]["function"]["arguments"] == "not json {" + + +def test_mcp_enabled_without_server_tools_uses_passthrough(monkeypatch): + # mcp_enabled=true with an empty registry must not silently drop the + # declared tools; the gate keys on the server-side path claiming the request. + backend = _ScriptedBackend(_fixed(_CALL_XML)) + payload = _request(tools = [LOOKUP_TOOL], stream = False, mcp_enabled = True) + body = _json_body(_call(payload, monkeypatch, backend)) + choice = body["choices"][0] + assert choice["finish_reason"] == "tool_calls" + assert choice["message"]["tool_calls"][0]["function"]["name"] == "lookup" + assert backend.calls[0]["tools"] == [LOOKUP_TOOL]