diff --git a/litellm/litellm_core_utils/prompt_templates/factory.py b/litellm/litellm_core_utils/prompt_templates/factory.py index 973a2787b6a..4580f9bd01b 100644 --- a/litellm/litellm_core_utils/prompt_templates/factory.py +++ b/litellm/litellm_core_utils/prompt_templates/factory.py @@ -5188,7 +5188,7 @@ def function_call_prompt( messages: list[dict[str, object]], functions: list[object], ) -> list[dict[str, object]]: - function_prompt = """Produce JSON OUTPUT ONLY! Adhere to this format {"name": "function_name", "arguments":{"argument_name": "argument_value"}} The following functions are available to you:""" + function_prompt = """To call a function, reply with JSON ONLY in this format {"name": "function_name", "arguments":{"argument_name": "argument_value"}}. Once a function result answers the request, reply to the user in plain text instead of calling a function again. The following functions are available to you:""" for function in functions: function_prompt += f"""\n{function}\n""" diff --git a/tests/e2e/AGENTS.md b/tests/e2e/AGENTS.md index 2e2169f8604..98a56bf1860 100644 --- a/tests/e2e/AGENTS.md +++ b/tests/e2e/AGENTS.md @@ -215,7 +215,7 @@ llm..... | rerank | images_generations | audio_speech | audio_transcriptions | moderations | realtime route : openai | azure_openai | anthropic | bedrock_converse | bedrock_invoke | vertex - | azure_foundry | cohere | together_ai + | azure_foundry | cohere | together_ai | ollama | ollama_chat (vocab varies per endpoint; messages is anthropic-format only) capability : basic | tool_use | prompt_cache_5m | vision | thinking | structured_output | service_tier | mid_conversation_system diff --git a/tests/e2e/coverage_registry/llm_conversational.yaml b/tests/e2e/coverage_registry/llm_conversational.yaml index 919884b66f0..cd41a4548ee 100644 --- a/tests/e2e/coverage_registry/llm_conversational.yaml +++ b/tests/e2e/coverage_registry/llm_conversational.yaml @@ -101,6 +101,32 @@ - {id: llm.messages.together_ai.basic.stream.works, module: llm, tier: P1, subject_endpoint: messages, route: together_ai, capability: basic, streaming: stream, assertions: [works], source: "llm_translation/test_together_ai_e2e.py", rationale: "Together over /v1/messages streaming"} - {id: llm.messages.together_ai.tool_use.nonstream.works, module: llm, tier: P1, subject_endpoint: messages, route: together_ai, capability: tool_use, streaming: nonstream, assertions: [works], source: "llm_translation/test_together_ai_e2e.py", rationale: "Together tool calls over /v1/messages"} - {id: llm.messages.together_ai.multi_turn.nonstream.works, module: llm, tier: P1, subject_endpoint: messages, route: together_ai, capability: multi_turn, streaming: nonstream, assertions: [works], source: "llm_translation/test_together_ai_e2e.py", rationale: "Together tool result round trip over /v1/messages"} +- {id: llm.chat_completions.ollama_chat.basic.nonstream.works, module: llm, tier: P1, subject_endpoint: chat_completions, route: ollama_chat, capability: basic, streaming: nonstream, assertions: [works], source: "llm_translation/test_ollama_e2e.py", rationale: "Ollama /api/chat (native tools) over /chat/completions returns text and usage"} +- {id: llm.chat_completions.ollama_chat.basic.stream.works, module: llm, tier: P1, subject_endpoint: chat_completions, route: ollama_chat, capability: basic, streaming: stream, assertions: [works], source: "llm_translation/test_ollama_e2e.py", rationale: "Ollama /api/chat (native tools) over /chat/completions streams text deltas, usage and a terminal event"} +- {id: llm.chat_completions.ollama_chat.tool_use.nonstream.works, module: llm, tier: P1, subject_endpoint: chat_completions, route: ollama_chat, capability: tool_use, streaming: nonstream, assertions: [works], source: "llm_translation/test_ollama_e2e.py", rationale: "Ollama /api/chat (native tools) over /chat/completions returns one addressable get_weather call"} +- {id: llm.chat_completions.ollama_chat.multi_turn.nonstream.works, module: llm, tier: P1, subject_endpoint: chat_completions, route: ollama_chat, capability: multi_turn, streaming: nonstream, assertions: [works], source: "llm_translation/test_ollama_e2e.py", rationale: "Ollama /api/chat (native tools) over /chat/completions tool result round trip reaches the model"} +- {id: llm.chat_completions.ollama_chat.tool_use.stream.works, module: llm, tier: P1, subject_endpoint: chat_completions, route: ollama_chat, capability: tool_use, streaming: stream, assertions: [works], source: "llm_translation/test_ollama_e2e.py", rationale: "Ollama /api/chat (native tools) streams the call as tool_call deltas ending in finish_reason tool_calls, the shape coding agents like OpenCode execute"} +- {id: llm.messages.ollama_chat.basic.nonstream.works, module: llm, tier: P1, subject_endpoint: messages, route: ollama_chat, capability: basic, streaming: nonstream, assertions: [works], source: "llm_translation/test_ollama_e2e.py", rationale: "Ollama /api/chat (native tools) over /v1/messages returns text and usage"} +- {id: llm.messages.ollama_chat.basic.stream.works, module: llm, tier: P1, subject_endpoint: messages, route: ollama_chat, capability: basic, streaming: stream, assertions: [works], source: "llm_translation/test_ollama_e2e.py", rationale: "Ollama /api/chat (native tools) over /v1/messages streams text deltas, usage and a terminal event"} +- {id: llm.messages.ollama_chat.tool_use.nonstream.works, module: llm, tier: P1, subject_endpoint: messages, route: ollama_chat, capability: tool_use, streaming: nonstream, assertions: [works], source: "llm_translation/test_ollama_e2e.py", rationale: "Ollama /api/chat (native tools) over /v1/messages returns one addressable get_weather call"} +- {id: llm.messages.ollama_chat.multi_turn.nonstream.works, module: llm, tier: P1, subject_endpoint: messages, route: ollama_chat, capability: multi_turn, streaming: nonstream, assertions: [works], source: "llm_translation/test_ollama_e2e.py", rationale: "Ollama /api/chat (native tools) over /v1/messages tool result round trip reaches the model"} +- {id: llm.responses.ollama_chat.basic.nonstream.works, module: llm, tier: P1, subject_endpoint: responses, route: ollama_chat, capability: basic, streaming: nonstream, assertions: [works], source: "llm_translation/test_ollama_e2e.py", rationale: "Ollama /api/chat (native tools) over /v1/responses returns text and usage"} +- {id: llm.responses.ollama_chat.basic.stream.works, module: llm, tier: P1, subject_endpoint: responses, route: ollama_chat, capability: basic, streaming: stream, assertions: [works], source: "llm_translation/test_ollama_e2e.py", rationale: "Ollama /api/chat (native tools) over /v1/responses streams text deltas, usage and a terminal event"} +- {id: llm.responses.ollama_chat.tool_use.nonstream.works, module: llm, tier: P1, subject_endpoint: responses, route: ollama_chat, capability: tool_use, streaming: nonstream, assertions: [works], source: "llm_translation/test_ollama_e2e.py", rationale: "Ollama /api/chat (native tools) over /v1/responses returns one addressable get_weather call"} +- {id: llm.responses.ollama_chat.multi_turn.nonstream.works, module: llm, tier: P1, subject_endpoint: responses, route: ollama_chat, capability: multi_turn, streaming: nonstream, assertions: [works], source: "llm_translation/test_ollama_e2e.py", rationale: "Ollama /api/chat (native tools) over /v1/responses tool result round trip reaches the model"} +- {id: llm.chat_completions.ollama.basic.nonstream.works, module: llm, tier: P1, subject_endpoint: chat_completions, route: ollama, capability: basic, streaming: nonstream, assertions: [works], source: "llm_translation/test_ollama_e2e.py", rationale: "Ollama /api/generate (prompt-based tools) over /chat/completions returns text and usage"} +- {id: llm.chat_completions.ollama.basic.stream.works, module: llm, tier: P1, subject_endpoint: chat_completions, route: ollama, capability: basic, streaming: stream, assertions: [works], source: "llm_translation/test_ollama_e2e.py", rationale: "Ollama /api/generate (prompt-based tools) over /chat/completions streams text deltas, usage and a terminal event"} +- {id: llm.chat_completions.ollama.tool_use.nonstream.works, module: llm, tier: P1, subject_endpoint: chat_completions, route: ollama, capability: tool_use, streaming: nonstream, assertions: [works], source: "llm_translation/test_ollama_e2e.py", rationale: "Ollama /api/generate (prompt-based tools) over /chat/completions returns one addressable get_weather call"} +- {id: llm.chat_completions.ollama.multi_turn.nonstream.works, module: llm, tier: P1, subject_endpoint: chat_completions, route: ollama, capability: multi_turn, streaming: nonstream, assertions: [works], source: "llm_translation/test_ollama_e2e.py", rationale: "Ollama /api/generate (prompt-based tools) over /chat/completions tool result round trip reaches the model"} +- {id: llm.chat_completions.ollama.tool_use.stream.works, module: llm, tier: P1, subject_endpoint: chat_completions, route: ollama, capability: tool_use, streaming: stream, assertions: [works], source: "llm_translation/test_ollama_e2e.py", rationale: "Ollama /api/generate (prompt-based tools) streams the call as tool_call deltas ending in finish_reason tool_calls, the shape coding agents like OpenCode execute; prompt-based JSON used to arrive as plain text with finish_reason stop (GitHub issue #35711)"} +- {id: llm.messages.ollama.basic.nonstream.works, module: llm, tier: P1, subject_endpoint: messages, route: ollama, capability: basic, streaming: nonstream, assertions: [works], source: "llm_translation/test_ollama_e2e.py", rationale: "Ollama /api/generate (prompt-based tools) over /v1/messages returns text and usage"} +- {id: llm.messages.ollama.basic.stream.works, module: llm, tier: P1, subject_endpoint: messages, route: ollama, capability: basic, streaming: stream, assertions: [works], source: "llm_translation/test_ollama_e2e.py", rationale: "Ollama /api/generate (prompt-based tools) over /v1/messages streams text deltas, usage and a terminal event"} +- {id: llm.messages.ollama.tool_use.nonstream.works, module: llm, tier: P1, subject_endpoint: messages, route: ollama, capability: tool_use, streaming: nonstream, assertions: [works], source: "llm_translation/test_ollama_e2e.py", rationale: "Ollama /api/generate (prompt-based tools) over /v1/messages returns one addressable get_weather call"} +- {id: llm.messages.ollama.multi_turn.nonstream.works, module: llm, tier: P1, subject_endpoint: messages, route: ollama, capability: multi_turn, streaming: nonstream, assertions: [works], source: "llm_translation/test_ollama_e2e.py", rationale: "Ollama /api/generate (prompt-based tools) over /v1/messages tool result round trip reaches the model"} +- {id: llm.responses.ollama.basic.nonstream.works, module: llm, tier: P1, subject_endpoint: responses, route: ollama, capability: basic, streaming: nonstream, assertions: [works], source: "llm_translation/test_ollama_e2e.py", rationale: "Ollama /api/generate (prompt-based tools) over /v1/responses returns text and usage"} +- {id: llm.responses.ollama.basic.stream.works, module: llm, tier: P1, subject_endpoint: responses, route: ollama, capability: basic, streaming: stream, assertions: [works], source: "llm_translation/test_ollama_e2e.py", rationale: "Ollama /api/generate (prompt-based tools) over /v1/responses streams text deltas, usage and a terminal event"} +- {id: llm.responses.ollama.tool_use.nonstream.works, module: llm, tier: P1, subject_endpoint: responses, route: ollama, capability: tool_use, streaming: nonstream, assertions: [works], source: "llm_translation/test_ollama_e2e.py", rationale: "Ollama /api/generate (prompt-based tools) over /v1/responses returns one addressable get_weather call"} +- {id: llm.responses.ollama.multi_turn.nonstream.works, module: llm, tier: P1, subject_endpoint: responses, route: ollama, capability: multi_turn, streaming: nonstream, assertions: [works], source: "llm_translation/test_ollama_e2e.py", rationale: "Ollama /api/generate (prompt-based tools) over /v1/responses tool result round trip reaches the model"} - {id: llm.chat_completions.sail.service_tier.nonstream.cost_logged, module: llm, tier: P1, subject_endpoint: chat_completions, route: sail, capability: service_tier, streaming: nonstream, assertions: [cost_logged], source: "llm_translation/test_sail_e2e.py", rationale: "service_tier balanced and auto map to Sail completion windows and bill the matching price columns; Sail serves flex only to background responses and Batch"} - {id: llm.chat_completions.sail.service_tier.nonstream.rejects_unknown_tier, module: llm, tier: P1, subject_endpoint: chat_completions, route: sail, capability: service_tier, streaming: nonstream, assertions: [rejects_unknown_tier], source: "llm_translation/test_sail_e2e.py", rationale: "A service_tier Sail has no completion window for is a 400 without drop_params"} - {id: llm.chat_completions.sail.service_tier.nonstream.drops_unknown_tier_and_bills_asap, module: llm, tier: P1, subject_endpoint: chat_completions, route: sail, capability: service_tier, streaming: nonstream, assertions: [drops_unknown_tier_and_bills_asap], source: "llm_translation/test_sail_e2e.py", rationale: "An unknown service_tier under drop_params is dropped and billed at asap in both the cost header and spend log"} diff --git a/tests/e2e/coverage_registry/schema.py b/tests/e2e/coverage_registry/schema.py index d089b9c1ed8..931085cbbc2 100644 --- a/tests/e2e/coverage_registry/schema.py +++ b/tests/e2e/coverage_registry/schema.py @@ -55,6 +55,8 @@ LlmRoute = Literal[ "cohere", "gemini", "hosted_vllm", + "ollama", + "ollama_chat", "openai", "sail", "together_ai", diff --git a/tests/e2e/llm_translation/test_ollama_e2e.py b/tests/e2e/llm_translation/test_ollama_e2e.py new file mode 100644 index 00000000000..6dea5a51562 --- /dev/null +++ b/tests/e2e/llm_translation/test_ollama_e2e.py @@ -0,0 +1,258 @@ +"""Ollama behind the proxy on /chat/completions, /v1/messages and /v1/responses. + +Ollama has two litellm routes with different tool plumbing: `ollama_chat/` calls +/api/chat and forwards native tools, while `ollama/` calls /api/generate, which +has no tools field, so litellm prompts the model for a JSON function call and +turns that JSON back into a tool call. Each route runs the same conversation +contract as the conversational matrix on every surface, through the matrix's +SDK-backed surfaces, plus a streamed tool call on chat completions, the shape +coding agents such as OpenCode consume. + +The deployments set drop_params because Ollama has no parallel_tool_calls, +which the matrix surfaces send alongside a forced tool_choice. Live only: Ollama +Cloud has no provider edge mount, and its requests are not priced in the cost +map, so there is no cost cell here. +""" + +from __future__ import annotations + +import json +from collections.abc import Iterator, Mapping +from dataclasses import dataclass +from itertools import chain, product +from types import MappingProxyType +from typing import Final, Literal, cast + +import pytest +from _pytest.mark.structures import ParameterSet +from e2e_config import unique_marker +from e2e_metadata import Capability as SubjectCapability +from e2e_metadata import Domain, Mode, Provider, Route, Subject, meta +from lifecycle import ResourceManager +from llm_translation.conversational_matrix import ( + GREETING_PROMPT, + INSTRUCTIONS, + MAX_OUTPUT_TOKENS, + SURFACES, + WEATHER_PROMPT, + WEATHER_REPORT, + WEATHER_TOOL_DESCRIPTION, + WEATHER_TOOL_NAME, + WEATHER_TOOL_SCHEMA, + Surface, + SurfaceName, + ToolCall, + WeatherArgs, + build_surfaces, +) +from llm_translation.sdk_clients import NO_PROXY_CACHE, SdkClients +from models import LiteLLMParamsBody +from openai.types.chat import ChatCompletionChunk, ChatCompletionToolParam +from proxy_client import ProxyClient + +pytestmark = pytest.mark.e2e + +OllamaRoute = Literal["ollama_chat", "ollama"] +Capability = Literal["basic", "tool_use", "multi_turn"] +Streaming = Literal["stream", "nonstream"] + +OLLAMA_API_BASE: Final = "https://ollama.com" +OLLAMA_MODEL: Final = "gemma4:31b" +ROUTES: Final[tuple[OllamaRoute, ...]] = ("ollama_chat", "ollama") + + +@dataclass(frozen=True, slots=True) +class Cell: + surface: SurfaceName + route: OllamaRoute + + @property + def id(self) -> str: + return f"{self.surface}-{self.route}" + + +def _cells(capability: Capability, streaming: Streaming) -> tuple[ParameterSet, ...]: + return tuple( + pytest.param( + Cell(surface=surface, route=route), + id=f"{surface}-{route}", + marks=pytest.mark.covers(f"llm.{surface}.{route}.{capability}.{streaming}.works"), + ) + for surface, route in product(SURFACES, ROUTES) + ) + + +def _streamed_tool_cells() -> tuple[ParameterSet, ...]: + return tuple( + pytest.param(route, id=route, marks=pytest.mark.covers(f"llm.chat_completions.{route}.tool_use.stream.works")) + for route in ROUTES + ) + + +def _register(proxy: ProxyClient, resources: ResourceManager, route: OllamaRoute) -> str: + alias: Final = f"e2e-ollama-{route}-{unique_marker()}" + model_id: Final = proxy.create_model( + alias, + LiteLLMParamsBody( + model=f"{route}/{OLLAMA_MODEL}", + api_base=OLLAMA_API_BASE, + api_key="os.environ/OLLAMA_API_KEY", + drop_params=True, + ), + ) + resources.defer(lambda: proxy.delete_model(model_id)) + return alias + + +@pytest.fixture(scope="module") +def aliases(proxy: ProxyClient) -> Iterator[Mapping[OllamaRoute, str]]: + resources: Final = ResourceManager(client=proxy) + try: + yield MappingProxyType({route: _register(proxy, resources, route) for route in ROUTES}) + finally: + resources.teardown() + + +@pytest.fixture(scope="module") +def surfaces(sdk: SdkClients) -> Mapping[SurfaceName, Surface]: + return build_surfaces(sdk) + + +def _weather_call(cell: Cell, surface: Surface, key: str, model: str) -> ToolCall: + reply: Final = surface.reply(key, model, WEATHER_PROMPT, with_tool=True) + assert len(reply.tool_calls) == 1, ( + f"{cell.id}: expected one {WEATHER_TOOL_NAME} call, got {reply.tool_calls} text={reply.text!r}" + ) + call: Final = reply.tool_calls[0] + assert call.name == WEATHER_TOOL_NAME, f"{cell.id}: called {call.name!r}, not {WEATHER_TOOL_NAME!r}" + assert call.call_id, f"{cell.id}: tool call has no id, so the caller cannot answer it: {call}" + assert "paris" in call.parsed().location.lower(), f"{cell.id}: tool arguments lost the location: {call}" + return call + + +def _weather_tool() -> ChatCompletionToolParam: + return { + "type": "function", + "function": { + "name": WEATHER_TOOL_NAME, + "description": WEATHER_TOOL_DESCRIPTION, + "parameters": dict(WEATHER_TOOL_SCHEMA), + }, + } + + +def _subject(mode: Mode, *, tools: bool, route: Route | None = None) -> Subject: + return Subject( + domain=Domain.LLM_TRANSLATION, + route=route, + providers=(Provider.OLLAMA,), + models=(OLLAMA_MODEL,), + capabilities=(SubjectCapability.FUNCTION_CALLING,) if tools else (), + mode=mode, + ) + + +class TestOllamaConversation: + @pytest.mark.parametrize("cell", _cells("basic", "nonstream")) + @meta(_subject(Mode.NONSTREAM, tools=False)) + def test_reply_carries_assistant_text_and_usage( + self, + cell: Cell, + aliases: Mapping[OllamaRoute, str], + surfaces: Mapping[SurfaceName, Surface], + resources: ResourceManager, + ) -> None: + reply: Final = surfaces[cell.surface].reply(resources.key(), aliases[cell.route], GREETING_PROMPT) + + assert reply.response_id, f"{cell.id}: response has no id" + assert reply.text.strip(), f"{cell.id}: response carried no assistant text" + assert reply.usage is not None and reply.usage.input_tokens > 0 and reply.usage.output_tokens > 0, ( + f"{cell.id}: usage missing or zero: {reply.usage}" + ) + assert reply.call_id_header, f"{cell.id}: x-litellm-call-id header missing" + + @pytest.mark.parametrize("cell", _cells("basic", "stream")) + @meta(_subject(Mode.STREAM, tools=False)) + def test_stream_delivers_text_usage_and_a_terminal_event( + self, + cell: Cell, + aliases: Mapping[OllamaRoute, str], + surfaces: Mapping[SurfaceName, Surface], + resources: ResourceManager, + ) -> None: + streamed: Final = surfaces[cell.surface].stream(resources.key(), aliases[cell.route], GREETING_PROMPT) + + assert streamed.event_count > 1, f"{cell.id}: stream arrived as {streamed.event_count} event(s)" + assert streamed.text.strip(), f"{cell.id}: stream carried no text deltas" + assert streamed.finished, f"{cell.id}: stream never sent its terminal event" + assert streamed.usage_reported, f"{cell.id}: stream never reported usage" + + @pytest.mark.parametrize("cell", _cells("tool_use", "nonstream")) + @meta(_subject(Mode.NONSTREAM, tools=True)) + def test_tool_call_is_returned_named_and_addressable( + self, + cell: Cell, + aliases: Mapping[OllamaRoute, str], + surfaces: Mapping[SurfaceName, Surface], + resources: ResourceManager, + ) -> None: + _ = _weather_call(cell, surfaces[cell.surface], resources.key(), aliases[cell.route]) + + @pytest.mark.parametrize("cell", _cells("multi_turn", "nonstream")) + @meta(_subject(Mode.NONSTREAM, tools=True)) + def test_tool_result_round_trip_reaches_the_model( + self, + cell: Cell, + aliases: Mapping[OllamaRoute, str], + surfaces: Mapping[SurfaceName, Surface], + resources: ResourceManager, + ) -> None: + key: Final = resources.key() + model: Final = aliases[cell.route] + surface: Final = surfaces[cell.surface] + call: Final = _weather_call(cell, surface, key, model) + + answer: Final = surface.reply_to_tool_result(key, model, WEATHER_PROMPT, call, WEATHER_REPORT) + assert "22" in answer.text, f"{cell.id}: the model never saw the tool result: {answer.text!r}" + + +class TestOllamaStreamedToolCall: + @pytest.mark.parametrize("route", _streamed_tool_cells()) + @meta(_subject(Mode.STREAM, tools=True, route=Route.CHAT_COMPLETIONS)) + def test_tool_call_streams_as_tool_call_deltas( + self, + route: OllamaRoute, + aliases: Mapping[OllamaRoute, str], + sdk: SdkClients, + resources: ResourceManager, + ) -> None: + chunks: Final[tuple[ChatCompletionChunk, ...]] = tuple( + sdk.openai(resources.key()).chat.completions.create( + model=aliases[route], + messages=[ + {"role": "system", "content": INSTRUCTIONS}, + {"role": "user", "content": WEATHER_PROMPT}, + ], + tools=[_weather_tool()], + max_completion_tokens=MAX_OUTPUT_TOKENS, + stream=True, + extra_body=NO_PROXY_CACHE, + ) + ) + choices: Final = tuple(chunk.choices[0] for chunk in chunks if chunk.choices) + text: Final = "".join(choice.delta.content or "" for choice in choices) + deltas: Final = tuple(chain.from_iterable(choice.delta.tool_calls or () for choice in choices)) + call_ids: Final = tuple(delta.id for delta in deltas if delta.id) + indexes: Final = frozenset(delta.index for delta in deltas) + functions: Final = tuple(delta.function for delta in deltas if delta.function is not None) + names: Final = tuple(function.name for function in functions if function.name) + arguments: Final = "".join(function.arguments or "" for function in functions) + finish_reasons: Final = tuple(choice.finish_reason for choice in choices if choice.finish_reason is not None) + + assert names == (WEATHER_TOOL_NAME,), f"{route}: streamed tool names {names}, text={text!r}" + assert len(call_ids) == 1, f"{route}: expected one streamed tool call id, got {call_ids}" + assert indexes == {0}, f"{route}: streamed tool call deltas used indexes {sorted(indexes)}" + assert WEATHER_TOOL_NAME not in text, f"{route}: the tool call leaked into assistant text: {text!r}" + location: Final = WeatherArgs.model_validate(cast(object, json.loads(arguments))).location + assert "paris" in location.lower(), f"{route}: streamed tool arguments lost the location: {arguments!r}" + assert finish_reasons[-1:] == ("tool_calls",), f"{route}: stream finished with {finish_reasons}" diff --git a/tests/e2e/models.py b/tests/e2e/models.py index 8796393f269..3ae200ff5e7 100644 --- a/tests/e2e/models.py +++ b/tests/e2e/models.py @@ -1257,6 +1257,7 @@ class LiteLLMParamsBody(BaseModel): api_version: str | None = None realtime_protocol: str | None = None allowed_openai_params: list[str] | None = None + drop_params: bool | None = None aws_access_key_id: str | None = Field(default=None, repr=False) aws_secret_access_key: str | None = Field(default=None, repr=False) aws_region_name: str | None = None diff --git a/tests/integration/providers/test_ollama_prompt_tools_chaos.py b/tests/integration/providers/test_ollama_prompt_tools_chaos.py new file mode 100644 index 00000000000..243c85b2a64 --- /dev/null +++ b/tests/integration/providers/test_ollama_prompt_tools_chaos.py @@ -0,0 +1,404 @@ +import asyncio +import json +import re +import signal +import threading +import uuid +from collections.abc import Callable +from dataclasses import dataclass +from pathlib import Path +from queue import SimpleQueue +from types import MappingProxyType +from typing import Final, Literal +from urllib.parse import urlsplit + +import httpx +import psutil +import pytest +import yaml +from integration._support.client import Gateway, eventually +from integration._support.database import read_rows +from integration._support.process import graceful_stop_seconds, owned_proxy_process +from integration._support.wire import Reply, Request, Wire, wire_server +from pydantic import JsonValue, TypeAdapter + +_BACKEND: Final = "llama3-prompt-tools-chaos" +_API_KEY: Final = "synthetic-ollama-key" +_CONFIG_MODEL: Final = "ollama-prompt-tools-chaos" +_INSTRUCTION: Final = ( + 'To call a function, reply with JSON ONLY in this format {"name": "function_name", ' + '"arguments":{"argument_name": "argument_value"}}. Once a function result answers the request, ' + "reply to the user in plain text instead of calling a function again. " + "The following functions are available to you:" +) +_CALL_ID: Final = "call_prompt_tools_chaos_1" +_ARGUMENTS: Final[dict[str, JsonValue]] = {"city": "Paris"} +_PARAMETERS: Final[dict[str, JsonValue]] = { + "type": "object", + "properties": {"city": {"type": "string"}}, + "required": ["city"], +} +_CHAT_TOOL: Final[dict[str, JsonValue]] = { + "type": "function", + "function": {"name": "get_weather", "description": "Weather for a city", "parameters": _PARAMETERS}, +} +_ANTHROPIC_TOOL: Final[dict[str, JsonValue]] = { + "name": "get_weather", + "description": "Weather for a city", + "input_schema": _PARAMETERS, +} +_RESPONSES_TOOL: Final[dict[str, JsonValue]] = { + "type": "function", + "name": "get_weather", + "description": "Weather for a city", + "parameters": _PARAMETERS, +} +_JSON_OBJECT: Final = TypeAdapter(dict[str, JsonValue]) +_MARKER: Final = re.compile(r"marker-([0-9a-f]{32})") +_STARTED_WORKER: Final = re.compile(r"Started server process \[(\d+)\]") +_OWNED_PROXY_CELL_SECONDS: Final = 2 * graceful_stop_seconds() + 120 + +Endpoint = Literal["chat", "messages", "responses"] + + +@dataclass(frozen=True, slots=True) +class _Call: + endpoint: Endpoint + stream: bool + marker: str + + +@dataclass(frozen=True, slots=True) +class _Served: + call: _Call + status: int + text: str + + +def _result(marker: str) -> str: + return f"Paris: 22 degrees Celsius marker-{marker}" + + +def _answer(marker: str) -> str: + return f"answer marker-{marker}" + + +def _path(endpoint: Endpoint) -> str: + match endpoint: + case "chat": + return "/v1/chat/completions" + case "messages": + return "/v1/messages" + case "responses": + return "/v1/responses" + + +def _body(model: str, call: _Call) -> dict[str, JsonValue]: + question: Final = "What is the weather in Paris?" + common: Final[dict[str, JsonValue]] = { + "model": model, + "stream": call.stream, + "num_retries": 0, + "cache": {"no-cache": True}, + } + match call.endpoint: + case "chat": + return { + **common, + "tools": [_CHAT_TOOL], + "messages": [ + {"role": "user", "content": question}, + { + "role": "assistant", + "content": None, + "tool_calls": [ + { + "id": _CALL_ID, + "type": "function", + "function": {"name": "get_weather", "arguments": json.dumps(_ARGUMENTS)}, + } + ], + }, + {"role": "tool", "tool_call_id": _CALL_ID, "content": _result(call.marker)}, + ], + } + case "messages": + return { + **common, + "max_tokens": 64, + "tools": [_ANTHROPIC_TOOL], + "messages": [ + {"role": "user", "content": question}, + { + "role": "assistant", + "content": [{"type": "tool_use", "id": _CALL_ID, "name": "get_weather", "input": _ARGUMENTS}], + }, + { + "role": "user", + "content": [{"type": "tool_result", "tool_use_id": _CALL_ID, "content": _result(call.marker)}], + }, + ], + } + case "responses": + return { + **common, + "store": False, + "tools": [_RESPONSES_TOOL], + "input": [ + {"role": "user", "content": question}, + { + "type": "function_call", + "call_id": _CALL_ID, + "name": "get_weather", + "arguments": json.dumps(_ARGUMENTS), + }, + {"type": "function_call_output", "call_id": _CALL_ID, "output": _result(call.marker)}, + ], + } + + +def _generate_reply(marker: str, stream: bool, drop_connection: bool = False) -> Reply: + done: Final = { + "model": _BACKEND, + "created_at": "2026-10-07T00:00:00Z", + "response": "", + "done": True, + "done_reason": "stop", + "prompt_eval_count": 30, + "eval_count": 12, + } + if not stream: + return Reply(body=json.dumps({**done, "response": _answer(marker)}).encode(), drop_connection=drop_connection) + pieces: Final = ("answer ", f"marker-{marker}") + frames: Final = ( + *( + {"model": _BACKEND, "created_at": "2026-10-07T00:00:00Z", "response": piece, "done": False} + for piece in pieces + ), + done, + ) + return Reply( + content_type="application/x-ndjson", + chunks=tuple(json.dumps(frame).encode() + b"\n" for frame in frames), + drop_connection=drop_connection, + ) + + +def _is_generate(request: Request) -> bool: + return (request.method, request.target) == ("POST", "/api/generate") + + +def _marker_of(request: Request) -> str: + found: Final = _MARKER.search(request.body.decode()) + assert found is not None, request.body + return found.group(1) + + +def _echo(request: Request) -> Reply: + if not _is_generate(request): + return Reply(body=b"{}") + stream: Final = _JSON_OBJECT.validate_json(request.body).get("stream") is True + return _generate_reply(_marker_of(request), stream) + + +def _assert_each_prompt_is_instructed_once(received: tuple[Request, ...], markers: frozenset[str]) -> None: + generates: Final = tuple(request for request in received if _is_generate(request)) + assert sorted(_marker_of(request) for request in generates) == sorted(markers) + for request in generates: + body: Final = _JSON_OBJECT.validate_json(request.body) + assert body["format"] == "json", sorted(body) + prompt: Final = body["prompt"] + assert isinstance(prompt, str) + assert prompt.count(_INSTRUCTION) == 1, prompt + assert set(_MARKER.findall(prompt)) == {_marker_of(request)}, prompt + + +def _response_id(served: _Served) -> str | None: + if served.call.endpoint == "responses": + return None + if not served.call.stream: + identity: Final = _JSON_OBJECT.validate_json(served.text)["id"] + assert isinstance(identity, str) + return identity + for line in served.text.splitlines(): + if not line.startswith("data: ") or line == "data: [DONE]": + continue + payload: Final = _JSON_OBJECT.validate_json(line.removeprefix("data: ")) + if served.call.endpoint == "chat": + first: Final = payload["id"] + assert isinstance(first, str) + return first + if payload.get("type") == "message_start": + message: Final = payload["message"] + assert isinstance(message, dict) and isinstance(message["id"], str) + return message["id"] + raise AssertionError(served.text) + + +def _spend_statuses(model: str, expected: int) -> MappingProxyType[str, str]: + rows: Final = eventually( + lambda: read_rows('SELECT request_id, status FROM "LiteLLM_SpendLogs" WHERE model_group=%s', (model,)), + lambda found: len(found) >= expected, + seconds=70, + ) + statuses: Final = MappingProxyType({str(row["request_id"]): str(row["status"]) for row in rows}) + assert len(statuses) == len(rows) == expected, rows + return statuses + + +def _successes(statuses: MappingProxyType[str, str]) -> frozenset[str]: + return frozenset(identity for identity, status in statuses.items() if status == "success") + + +async def _send(client: httpx.AsyncClient, key: str, model: str, call: _Call) -> _Served: + async with client.stream( + "POST", + _path(call.endpoint), + json=_body(model, call), + headers={"Authorization": f"Bearer {key}", "anthropic-version": "2023-06-01"}, + ) as response: + raw: Final = await response.aread() + return _Served(call=call, status=response.status_code, text=raw.decode()) + + +async def _burst( + base_url: str, key: str, model: str, calls: tuple[_Call, ...], *, tolerate_transport_errors: bool = False +) -> tuple[_Served, ...]: + async with httpx.AsyncClient(base_url=base_url, timeout=60, trust_env=False) as client: + results: Final = await asyncio.gather( + *(_send(client, key, model, call) for call in calls), return_exceptions=tolerate_transport_errors + ) + for result in results: + assert not isinstance(result, BaseException) or isinstance(result, httpx.TransportError), repr(result) + return tuple(result for result in results if isinstance(result, _Served)) + + +def _calls(count: int, endpoints: tuple[Endpoint, ...], stream: Callable[[int], bool]) -> tuple[_Call, ...]: + return tuple( + _Call(endpoint=endpoints[index % len(endpoints)], stream=stream(index), marker=uuid.uuid4().hex) + for index in range(count) + ) + + +def _assert_answered_with_its_own_marker(served: _Served) -> None: + assert served.status == 200, served.text + assert set(_MARKER.findall(served.text)) == {served.call.marker}, served.text + + +async def test_concurrent_tool_result_turns_across_endpoints_each_get_their_own_instructed_prompt( + gateway: Gateway, +) -> None: + calls: Final = _calls(24, ("chat", "messages", "responses"), lambda index: index % 2 == 0) + with wire_server(_echo) as wire, gateway.scenario() as scenario: + model: Final = scenario.model(model=f"ollama/{_BACKEND}", api_base=wire.url, api_key=_API_KEY) + served: Final = await _burst(str(gateway.client.base_url), gateway.key, model, calls) + assert len(served) == 24 + for item in served: + _assert_answered_with_its_own_marker(item) + _assert_each_prompt_is_instructed_once(wire.drain(), frozenset(call.marker for call in calls)) + known: Final = frozenset(identity for identity in map(_response_id, served) if identity is not None) + assert len(known) == 16, known + statuses: Final = _spend_statuses(model, 24) + assert _successes(statuses) == frozenset(statuses), statuses + assert known <= _successes(statuses) + + +async def test_dropped_ollama_connections_fail_their_callers_and_the_rest_keep_their_prompts(gateway: Gateway) -> None: + calls: Final = _calls(12, ("chat",), lambda index: index % 2 == 1) + dropped: Final = frozenset(call.marker for index, call in enumerate(calls) if index % 3 == 0) + + def respond(request: Request) -> Reply: + if not _is_generate(request): + return Reply(body=b"{}") + marker: Final = _marker_of(request) + stream: Final = _JSON_OBJECT.validate_json(request.body).get("stream") is True + return _generate_reply(marker, stream, drop_connection=marker in dropped) + + with wire_server(respond) as wire, gateway.scenario() as scenario: + model: Final = scenario.model(model=f"ollama/{_BACKEND}", api_base=wire.url, api_key=_API_KEY) + served: Final = await _burst(str(gateway.client.base_url), gateway.key, model, calls) + assert len(served) == 12 + for item in served: + if item.call.marker in dropped: + assert item.status == 500, item.text + assert "APIConnectionError" in item.text and "marker-" not in item.text, item.text + else: + _assert_answered_with_its_own_marker(item) + recovery: Final = _Call(endpoint="chat", stream=False, marker=uuid.uuid4().hex) + (recovered,) = await _burst(str(gateway.client.base_url), gateway.key, model, (recovery,)) + _assert_answered_with_its_own_marker(recovered) + _assert_each_prompt_is_instructed_once(wire.drain(), frozenset(call.marker for call in (*calls, recovery))) + answered: Final = tuple(item for item in served if item.call.marker not in dropped) + survivors: Final = frozenset( + identity for identity in map(_response_id, (*answered, recovered)) if identity is not None + ) + assert len(survivors) == 9, survivors + statuses: Final = _spend_statuses(model, 13) + assert _successes(statuses) == survivors, statuses + assert sum(status == "failure" for status in statuses.values()) == len(dropped), statuses + + +def _chaos_config(wire: Wire, tmp_path: Path) -> Path: + base: Final = yaml.safe_load(Path("tests/integration/proxy_config.yaml").read_text()) + deployment: Final = { + "model_name": _CONFIG_MODEL, + "litellm_params": {"model": f"ollama/{_BACKEND}", "api_base": wire.url, "api_key": _API_KEY}, + } + path: Final = tmp_path / "ollama-prompt-tools-chaos.yaml" + path.write_text(yaml.safe_dump({**base, "model_list": [deployment]})) + return path + + +def _open_upstream_connections(pid: int, upstream: str) -> int: + port: Final = urlsplit(upstream).port + return sum( + 1 + for connection in psutil.Process(pid).net_connections(kind="tcp") + if connection.status == psutil.CONN_ESTABLISHED and connection.raddr and connection.raddr.port == port + ) + + +@pytest.mark.timeout(_OWNED_PROXY_CELL_SECONDS) +async def test_worker_sigkill_mid_burst_leaves_the_sibling_instructing_ollama(gateway: Gateway, tmp_path: Path) -> None: + calls: Final = _calls(20, ("chat",), lambda _: False) + release: Final = threading.Event() + held_markers: Final[SimpleQueue[str]] = SimpleQueue() + + def held(request: Request) -> Reply: + if not _is_generate(request): + return Reply(body=b"{}") + held_markers.put(_marker_of(request)) + assert release.wait(timeout=60), "The burst was never released" + return _echo(request) + + with wire_server(held) as wire: + path: Final = _chaos_config(wire, tmp_path) + with owned_proxy_process(gateway, tmp_path, {}, config=path, workers=2) as owned: + candidate: Final = owned.gateway + workers: Final = eventually( + lambda: tuple(int(pid) for pid in _STARTED_WORKER.findall(owned.log.read_text())), + lambda pids: len(pids) == 2, + seconds=30, + ) + burst: Final = asyncio.create_task( + _burst( + str(candidate.client.base_url), candidate.key, _CONFIG_MODEL, calls, tolerate_transport_errors=True + ) + ) + await asyncio.to_thread(eventually, held_markers.qsize, lambda size: size == 20, 60) + held_by: Final = MappingProxyType({pid: _open_upstream_connections(pid, wire.url) for pid in workers}) + assert sum(held_by.values()) == 20, held_by + victim_pid, survivor_pid = sorted(workers, key=held_by.__getitem__) + victim: Final = psutil.Process(victim_pid) + victim.suspend() + victim.send_signal(signal.SIGKILL) + release.set() + served: Final = await burst + assert held_by[survivor_pid] >= 10, held_by + assert len(served) == held_by[survivor_pid], (held_by, len(served)) + for item in served: + _assert_answered_with_its_own_marker(item) + follow_up: Final = _Call(endpoint="chat", stream=False, marker=uuid.uuid4().hex) + (answered,) = await _burst(str(candidate.client.base_url), candidate.key, _CONFIG_MODEL, (follow_up,)) + _assert_answered_with_its_own_marker(answered) + _assert_each_prompt_is_instructed_once(wire.drain(), frozenset(call.marker for call in (*calls, follow_up))) diff --git a/tests/integration/providers/test_ollama_prompt_tools_wire.py b/tests/integration/providers/test_ollama_prompt_tools_wire.py new file mode 100644 index 00000000000..115fdc65e33 --- /dev/null +++ b/tests/integration/providers/test_ollama_prompt_tools_wire.py @@ -0,0 +1,796 @@ +import itertools +import json +import uuid +from collections.abc import Callable, Iterator, Sequence +from contextlib import contextmanager +from typing import Final + +import anthropic +import openai +import pytest +from openai.types.chat import ChatCompletionChunk +from openai.types.chat.chat_completion_chunk import Choice as ChunkChoice +from openai.types.chat.chat_completion_chunk import ChoiceDeltaToolCall +from integration._support.client import Gateway, eventually +from integration._support.database import read_rows +from integration._support.wire import Reply, Request, Wire, wire_server +from pydantic import JsonValue, TypeAdapter + +_BACKEND: Final = "llama3-prompt-tools" +_API_KEY: Final = "synthetic-ollama-key" +_INSTRUCTION: Final = ( + 'To call a function, reply with JSON ONLY in this format {"name": "function_name", ' + '"arguments":{"argument_name": "argument_value"}}. Once a function result answers the request, ' + "reply to the user in plain text instead of calling a function again. " + "The following functions are available to you:" +) +_QUESTION: Final = "What is the weather in Paris?" +_RESULT: Final = "Paris: 22 degrees Celsius, clear skies" +_ANSWER: Final = "Paris is 22 degrees Celsius with clear skies." +_CALL_ID: Final = "call_prompt_tools_1" +_ARGUMENTS: Final[dict[str, JsonValue]] = {"city": "Paris"} +_CALL_JSON: Final = json.dumps({"name": "get_weather", "arguments": _ARGUMENTS}) +_CALL_JSON_FIELDS: Final = frozenset({"get_weather"}) +_PARAMETERS: Final[dict[str, JsonValue]] = { + "type": "object", + "properties": {"city": {"type": "string"}}, + "required": ["city"], +} +_WEATHER_TOOL: Final[dict[str, JsonValue]] = { + "type": "function", + "function": {"name": "get_weather", "description": "Weather for a city", "parameters": _PARAMETERS}, +} +_TIME_TOOL: Final[dict[str, JsonValue]] = { + "type": "function", + "function": {"name": "get_time", "description": "Local time for a city", "parameters": _PARAMETERS}, +} +_ANTHROPIC_TOOL: Final[dict[str, JsonValue]] = { + "name": "get_weather", + "description": "Weather for a city", + "input_schema": _PARAMETERS, +} +_RESPONSES_TOOL: Final[dict[str, JsonValue]] = { + "type": "function", + "name": "get_weather", + "description": "Weather for a city", + "parameters": _PARAMETERS, +} +_NO_CACHE: Final[dict[str, JsonValue]] = {"no-cache": True} +_JSON_OBJECT: Final = TypeAdapter(dict[str, JsonValue]) +_RAW_ANTHROPIC_EVENTS: Final = frozenset( + { + "message_start", + "content_block_start", + "content_block_delta", + "content_block_stop", + "message_delta", + "message_stop", + } +) + + +def _generate_reply(text: str) -> Reply: + return Reply( + body=json.dumps( + { + "model": _BACKEND, + "created_at": "2026-10-07T00:00:00Z", + "response": text, + "done": True, + "done_reason": "stop", + "prompt_eval_count": 30, + "eval_count": 12, + } + ).encode() + ) + + +def _streamed_reply(text: str) -> Reply: + pieces: Final = tuple(text[index : index + 7] for index in range(0, len(text), 7)) + lines: Final = tuple( + json.dumps({"model": _BACKEND, "created_at": "2026-10-07T00:00:00Z", "response": piece, "done": False}).encode() + + b"\n" + for piece in pieces + ) + final: Final = ( + json.dumps( + { + "model": _BACKEND, + "created_at": "2026-10-07T00:00:00Z", + "response": "", + "done": True, + "done_reason": "stop", + "prompt_eval_count": 30, + "eval_count": 12, + } + ).encode() + + b"\n" + ) + return Reply(content_type="application/x-ndjson", chunks=(*lines, final)) + + +def _is_generate(request: Request) -> bool: + return (request.method, request.target) == ("POST", "/api/generate") + + +@contextmanager +def _ollama_server(respond: Callable[[Request], Reply]) -> Iterator[Wire]: + with wire_server(lambda request: respond(request) if _is_generate(request) else Reply(body=b"{}")) as wire: + yield wire + + +def _generate_calls(wire: Wire) -> tuple[Request, ...]: + return tuple(request for request in wire.drain() if _is_generate(request)) + + +def _only_generate(wire: Wire) -> dict[str, JsonValue]: + received: Final = _generate_calls(wire) + assert len(received) == 1, [(request.method, request.target) for request in received] + assert received[0].headers["authorization"] == f"Bearer {_API_KEY}" + return _JSON_OBJECT.validate_json(received[0].body) + + +def _prompt_of(body: dict[str, JsonValue]) -> str: + assert body["model"] == _BACKEND + assert body["format"] == "json" + assert "tools" not in body and "messages" not in body, sorted(body) + prompt: Final = body["prompt"] + assert isinstance(prompt, str) + return prompt + + +def _assert_instructed_once(prompt: str, *tool_names: str) -> None: + assert prompt.count(_INSTRUCTION) == 1, prompt + assert prompt.count("### System:") == 1, prompt + for name in tool_names: + assert f"'name': '{name}'" in prompt, prompt + + +def _assert_tool_turn(prompt: str, result: str = _RESULT) -> None: + _assert_instructed_once(prompt, "get_weather") + assert f"### User:\n{_QUESTION}\n\n### Assistant:\n{_CALL_JSON}\n\n### User:\n{result}\n\n" in prompt, prompt + + +def _spend_row(identity: str) -> dict[str, JsonValue]: + rows: Final = eventually( + lambda: read_rows( + 'SELECT model_group, status, prompt_tokens, completion_tokens FROM "LiteLLM_SpendLogs" WHERE request_id=%s', + (identity,), + ), + lambda found: len(found) == 1, + seconds=70, + ) + return rows[0] + + +def _model_spend_rows(model: str, expected: int) -> list[dict[str, JsonValue]]: + rows: Final = eventually( + lambda: read_rows( + 'SELECT request_id, status, prompt_tokens, completion_tokens FROM "LiteLLM_SpendLogs" WHERE model_group=%s', + (model,), + ), + lambda found: len(found) >= expected, + seconds=70, + ) + assert len({row["request_id"] for row in rows}) == len(rows) == expected, rows + return rows + + +def _billed(model: str) -> dict[str, JsonValue]: + return {"model_group": model, "status": "success", "prompt_tokens": 30, "completion_tokens": 12} + + +def _openai_client(gateway: Gateway) -> openai.OpenAI: + return openai.OpenAI(base_url=str(gateway.client.base_url) + "/v1", api_key=gateway.key, max_retries=0) + + +def _async_openai_client(gateway: Gateway) -> openai.AsyncOpenAI: + return openai.AsyncOpenAI(base_url=str(gateway.client.base_url) + "/v1", api_key=gateway.key, max_retries=0) + + +def _anthropic_client(gateway: Gateway) -> anthropic.Anthropic: + return anthropic.Anthropic(base_url=str(gateway.client.base_url), api_key=gateway.key, max_retries=0) + + +def _async_anthropic_client(gateway: Gateway) -> anthropic.AsyncAnthropic: + return anthropic.AsyncAnthropic(base_url=str(gateway.client.base_url), api_key=gateway.key, max_retries=0) + + +def _stream_choices(chunks: Sequence[ChatCompletionChunk]) -> Iterator[ChunkChoice]: + for chunk in chunks: + yield from chunk.choices + + +def _delta_tool_calls(choices: Sequence[ChunkChoice]) -> Iterator[ChoiceDeltaToolCall]: + for choice in choices: + yield from choice.delta.tool_calls or () + + +def _first_turn() -> list[dict[str, JsonValue]]: + return [{"role": "user", "content": _QUESTION}] + + +def _second_turn(result: JsonValue = _RESULT) -> list[dict[str, JsonValue]]: + return [ + {"role": "user", "content": _QUESTION}, + { + "role": "assistant", + "content": None, + "tool_calls": [ + { + "id": _CALL_ID, + "type": "function", + "function": {"name": "get_weather", "arguments": json.dumps(_ARGUMENTS)}, + } + ], + }, + {"role": "tool", "tool_call_id": _CALL_ID, "content": result}, + ] + + +def _anthropic_second_turn() -> list[dict[str, JsonValue]]: + return [ + {"role": "user", "content": _QUESTION}, + { + "role": "assistant", + "content": [{"type": "tool_use", "id": _CALL_ID, "name": "get_weather", "input": _ARGUMENTS}], + }, + {"role": "user", "content": [{"type": "tool_result", "tool_use_id": _CALL_ID, "content": _RESULT}]}, + ] + + +def _responses_second_turn() -> list[dict[str, JsonValue]]: + return [ + {"role": "user", "content": _QUESTION}, + {"type": "function_call", "call_id": _CALL_ID, "name": "get_weather", "arguments": json.dumps(_ARGUMENTS)}, + {"type": "function_call_output", "call_id": _CALL_ID, "output": _RESULT}, + ] + + +def _post(gateway: Gateway, path: str, body: dict[str, JsonValue], key: str | None = None) -> tuple[int, str]: + response: Final = gateway.request("POST", path, {**body, "cache": _NO_CACHE}, key=key) + return response.status_code, response.text + + +def _post_chat( + gateway: Gateway, model: str, messages: Sequence[dict[str, JsonValue]], **extra: JsonValue +) -> dict[str, JsonValue]: + code, text = _post(gateway, "/v1/chat/completions", {"model": model, "messages": list(messages), **extra}) + assert code == 200, text + return _JSON_OBJECT.validate_json(text) + + +def test_openai_sdk_tool_request_reaches_ollama_as_an_instructed_prompt(gateway: Gateway) -> None: + with _ollama_server(lambda _: _generate_reply(_CALL_JSON)) as wire, gateway.scenario() as scenario: + model: Final = scenario.model(model=f"ollama/{_BACKEND}", api_base=wire.url, api_key=_API_KEY) + completion: Final = _openai_client(gateway).chat.completions.create( + model=model, + messages=_first_turn(), # pyright: ignore[reportArgumentType] # plain JSON messages + tools=[_WEATHER_TOOL], # pyright: ignore[reportArgumentType] # plain JSON tool + extra_body={"cache": _NO_CACHE}, + ) + choice: Final = completion.choices[0] + assert choice.finish_reason == "tool_calls" + assert choice.message.tool_calls is not None and len(choice.message.tool_calls) == 1 + call: Final = choice.message.tool_calls[0] + assert call.type == "function" + assert call.function.name == "get_weather" + assert json.loads(call.function.arguments) == _ARGUMENTS + assert completion.usage is not None + assert (completion.usage.prompt_tokens, completion.usage.completion_tokens) == (30, 12) + body: Final = _only_generate(wire) + assert body["stream"] is False + prompt: Final = _prompt_of(body) + _assert_instructed_once(prompt, "get_weather") + assert f"### User:\n{_QUESTION}\n\n" in prompt, prompt + assert "Weather for a city" in prompt, prompt + assert _spend_row(completion.id) == _billed(model) + + +def test_openai_sdk_tool_result_turn_gets_a_plain_text_answer(gateway: Gateway) -> None: + with _ollama_server(lambda _: _generate_reply(_ANSWER)) as wire, gateway.scenario() as scenario: + model: Final = scenario.model(model=f"ollama/{_BACKEND}", api_base=wire.url, api_key=_API_KEY) + completion: Final = _openai_client(gateway).chat.completions.create( + model=model, + messages=_second_turn(), # pyright: ignore[reportArgumentType] # plain JSON messages + tools=[_WEATHER_TOOL], # pyright: ignore[reportArgumentType] # plain JSON tool + extra_body={"cache": _NO_CACHE}, + ) + choice: Final = completion.choices[0] + assert choice.finish_reason == "stop" + assert choice.message.content == _ANSWER + assert choice.message.tool_calls is None + _assert_tool_turn(_prompt_of(_only_generate(wire))) + assert _spend_row(completion.id) == _billed(model) + + +async def test_async_openai_sdk_stream_flushes_the_held_tool_call_once(gateway: Gateway) -> None: + with _ollama_server(lambda _: _streamed_reply(_CALL_JSON)) as wire, gateway.scenario() as scenario: + model: Final = scenario.model(model=f"ollama/{_BACKEND}", api_base=wire.url, api_key=_API_KEY) + stream: Final = await _async_openai_client(gateway).chat.completions.create( + model=model, + messages=_first_turn(), # pyright: ignore[reportArgumentType] # plain JSON messages + tools=[_WEATHER_TOOL], # pyright: ignore[reportArgumentType] # plain JSON tool + stream=True, + stream_options={"include_usage": True}, + extra_body={"cache": _NO_CACHE}, + ) + chunks: Final = [chunk async for chunk in stream] + assert {chunk.id for chunk in chunks} == {chunks[0].id} + choices: Final = tuple(_stream_choices(chunks)) + deltas: Final = tuple(_delta_tool_calls(choices)) + assert len(deltas) == 1, deltas + assert deltas[0].function is not None and deltas[0].function.name == "get_weather" + assert json.loads(deltas[0].function.arguments or "") == _ARGUMENTS + assert [choice.finish_reason for choice in choices if choice.finish_reason] == ["tool_calls"] + usages: Final = [chunk.usage for chunk in chunks if chunk.usage is not None] + assert [(usage.prompt_tokens, usage.completion_tokens) for usage in usages] == [(30, 12)] + body: Final = _only_generate(wire) + assert body["stream"] is True + _assert_instructed_once(_prompt_of(body), "get_weather") + assert _spend_row(chunks[0].id) == _billed(model) + + +async def test_async_openai_sdk_stream_answers_the_tool_result_in_plain_text(gateway: Gateway) -> None: + with _ollama_server(lambda _: _streamed_reply(_ANSWER)) as wire, gateway.scenario() as scenario: + model: Final = scenario.model(model=f"ollama/{_BACKEND}", api_base=wire.url, api_key=_API_KEY) + stream: Final = await _async_openai_client(gateway).chat.completions.create( + model=model, + messages=_second_turn(), # pyright: ignore[reportArgumentType] # plain JSON messages + tools=[_WEATHER_TOOL], # pyright: ignore[reportArgumentType] # plain JSON tool + stream=True, + stream_options={"include_usage": True}, + extra_body={"cache": _NO_CACHE}, + ) + chunks: Final = [chunk async for chunk in stream] + choices: Final = tuple(_stream_choices(chunks)) + assert "".join(choice.delta.content or "" for choice in choices) == _ANSWER + assert tuple(_delta_tool_calls(choices)) == () + assert [choice.finish_reason for choice in choices if choice.finish_reason] == ["stop"] + body: Final = _only_generate(wire) + assert body["stream"] is True + _assert_tool_turn(_prompt_of(body)) + assert _spend_row(chunks[0].id) == _billed(model) + + +def test_anthropic_sdk_tool_request_comes_back_as_tool_use(gateway: Gateway) -> None: + with _ollama_server(lambda _: _generate_reply(_CALL_JSON)) as wire, gateway.scenario() as scenario: + model: Final = scenario.model(model=f"ollama/{_BACKEND}", api_base=wire.url, api_key=_API_KEY) + message: Final = _anthropic_client(gateway).messages.create( + model=model, + max_tokens=64, + messages=_first_turn(), # pyright: ignore[reportArgumentType] # plain JSON messages + tools=[_ANTHROPIC_TOOL], # pyright: ignore[reportArgumentType] # plain JSON tool + extra_body={"cache": _NO_CACHE}, + ) + assert message.stop_reason == "tool_use" + assert [block.type for block in message.content] == ["tool_use"] + block: Final = message.content[0] + assert block.type == "tool_use" + assert block.name == "get_weather" + assert block.input == _ARGUMENTS + assert (message.usage.input_tokens, message.usage.output_tokens) == (30, 12) + prompt: Final = _prompt_of(_only_generate(wire)) + _assert_instructed_once(prompt, "get_weather") + assert f"### User:\n{_QUESTION}\n\n" in prompt, prompt + assert _spend_row(message.id) == _billed(model) + + +def test_anthropic_sdk_tool_result_turn_ends_with_text(gateway: Gateway) -> None: + with _ollama_server(lambda _: _generate_reply(_ANSWER)) as wire, gateway.scenario() as scenario: + model: Final = scenario.model(model=f"ollama/{_BACKEND}", api_base=wire.url, api_key=_API_KEY) + message: Final = _anthropic_client(gateway).messages.create( + model=model, + max_tokens=64, + messages=_anthropic_second_turn(), # pyright: ignore[reportArgumentType] # plain JSON messages + tools=[_ANTHROPIC_TOOL], # pyright: ignore[reportArgumentType] # plain JSON tool + extra_body={"cache": _NO_CACHE}, + ) + assert message.stop_reason == "end_turn" + assert [(block.type, getattr(block, "text", None)) for block in message.content] == [("text", _ANSWER)] + _assert_tool_turn(_prompt_of(_only_generate(wire))) + assert _spend_row(message.id) == _billed(model) + + +async def test_async_anthropic_sdk_stream_emits_the_tool_use_block(gateway: Gateway) -> None: + with _ollama_server(lambda _: _streamed_reply(_CALL_JSON)) as wire, gateway.scenario() as scenario: + model: Final = scenario.model(model=f"ollama/{_BACKEND}", api_base=wire.url, api_key=_API_KEY) + async with _async_anthropic_client(gateway).messages.stream( + model=model, + max_tokens=64, + messages=_first_turn(), # pyright: ignore[reportArgumentType] # plain JSON messages + tools=[_ANTHROPIC_TOOL], # pyright: ignore[reportArgumentType] # plain JSON tool + extra_body={"cache": _NO_CACHE}, + ) as stream: + events: Final = [event async for event in stream if event.type in _RAW_ANTHROPIC_EVENTS] + final: Final = await stream.get_final_message() + starts: Final = [event for event in events if event.type == "content_block_start"] + assert [event.content_block.type for event in starts] == ["tool_use"], [event.type for event in events] + assert any( + event.type == "content_block_delta" and event.delta.type == "input_json_delta" for event in events + ), [event.type for event in events] + assert final.stop_reason == "tool_use" + block: Final = final.content[0] + assert block.type == "tool_use" and block.name == "get_weather" and block.input == _ARGUMENTS + body: Final = _only_generate(wire) + assert body["stream"] is True + _assert_instructed_once(_prompt_of(body), "get_weather") + assert _spend_row(final.id) == _billed(model) + + +async def test_async_anthropic_sdk_stream_answers_the_tool_result_in_text(gateway: Gateway) -> None: + with _ollama_server(lambda _: _streamed_reply(_ANSWER)) as wire, gateway.scenario() as scenario: + model: Final = scenario.model(model=f"ollama/{_BACKEND}", api_base=wire.url, api_key=_API_KEY) + async with _async_anthropic_client(gateway).messages.stream( + model=model, + max_tokens=64, + messages=_anthropic_second_turn(), # pyright: ignore[reportArgumentType] # plain JSON messages + tools=[_ANTHROPIC_TOOL], # pyright: ignore[reportArgumentType] # plain JSON tool + extra_body={"cache": _NO_CACHE}, + ) as stream: + texts: Final = [event.text async for event in stream if event.type == "text"] + final: Final = await stream.get_final_message() + assert "".join(texts) == _ANSWER + assert final.stop_reason == "end_turn" + assert [(block.type, getattr(block, "text", None)) for block in final.content] == [("text", _ANSWER)] + body: Final = _only_generate(wire) + assert body["stream"] is True + _assert_tool_turn(_prompt_of(body)) + assert _spend_row(final.id) == _billed(model) + + +def test_openai_sdk_responses_tool_request_comes_back_as_a_function_call(gateway: Gateway) -> None: + with _ollama_server(lambda _: _generate_reply(_CALL_JSON)) as wire, gateway.scenario() as scenario: + model: Final = scenario.model(model=f"ollama/{_BACKEND}", api_base=wire.url, api_key=_API_KEY) + response: Final = _openai_client(gateway).responses.create( + model=model, + input=_QUESTION, + tools=[_RESPONSES_TOOL], # pyright: ignore[reportArgumentType] # plain JSON tool + store=False, + extra_body={"cache": _NO_CACHE}, + ) + assert [item.type for item in response.output] == ["function_call"] + item: Final = response.output[0] + assert item.type == "function_call" + assert item.name == "get_weather" + assert json.loads(item.arguments) == _ARGUMENTS + assert response.usage is not None and (response.usage.input_tokens, response.usage.output_tokens) == (30, 12) + prompt: Final = _prompt_of(_only_generate(wire)) + _assert_instructed_once(prompt, "get_weather") + assert f"### User:\n{_QUESTION}\n\n" in prompt, prompt + assert _model_spend_rows(model, 1)[0]["status"] == "success" + + +def test_openai_sdk_responses_function_output_turn_gets_a_message(gateway: Gateway) -> None: + with _ollama_server(lambda _: _generate_reply(_ANSWER)) as wire, gateway.scenario() as scenario: + model: Final = scenario.model(model=f"ollama/{_BACKEND}", api_base=wire.url, api_key=_API_KEY) + response: Final = _openai_client(gateway).responses.create( + model=model, + input=_responses_second_turn(), # pyright: ignore[reportArgumentType] # plain JSON input items + tools=[_RESPONSES_TOOL], # pyright: ignore[reportArgumentType] # plain JSON tool + store=False, + extra_body={"cache": _NO_CACHE}, + ) + assert [item.type for item in response.output] == ["message"] + assert response.output_text == _ANSWER + _assert_tool_turn(_prompt_of(_only_generate(wire))) + rows: Final = _model_spend_rows(model, 1) + assert (rows[0]["status"], rows[0]["prompt_tokens"], rows[0]["completion_tokens"]) == ("success", 30, 12) + + +async def test_async_openai_sdk_responses_stream_emits_the_function_call_item(gateway: Gateway) -> None: + with _ollama_server(lambda _: _streamed_reply(_CALL_JSON)) as wire, gateway.scenario() as scenario: + model: Final = scenario.model(model=f"ollama/{_BACKEND}", api_base=wire.url, api_key=_API_KEY) + stream: Final = await _async_openai_client(gateway).responses.create( + model=model, + input=_QUESTION, + tools=[_RESPONSES_TOOL], # pyright: ignore[reportArgumentType] # plain JSON tool + store=False, + stream=True, + extra_body={"cache": _NO_CACHE}, + ) + events: Final = [event async for event in stream] + done_items: Final = [event.item for event in events if event.type == "response.output_item.done"] + assert [item.type for item in done_items] == ["function_call"], [event.type for event in events] + item: Final = done_items[0] + assert item.type == "function_call" and item.name == "get_weather" + assert json.loads(item.arguments) == _ARGUMENTS + completed: Final = [event for event in events if event.type == "response.completed"] + assert len(completed) == 1 + final_calls: Final = [item for item in completed[0].response.output if item.type == "function_call"] + assert [(item.name, json.loads(item.arguments)) for item in final_calls] == [("get_weather", _ARGUMENTS)] + assert completed[0].response.output_text == "" + body: Final = _only_generate(wire) + assert body["stream"] is True + _assert_instructed_once(_prompt_of(body), "get_weather") + assert _model_spend_rows(model, 1)[0]["status"] == "success" + + +async def test_async_openai_sdk_responses_stream_answers_the_function_output_in_text(gateway: Gateway) -> None: + with _ollama_server(lambda _: _streamed_reply(_ANSWER)) as wire, gateway.scenario() as scenario: + model: Final = scenario.model(model=f"ollama/{_BACKEND}", api_base=wire.url, api_key=_API_KEY) + stream: Final = await _async_openai_client(gateway).responses.create( + model=model, + input=_responses_second_turn(), # pyright: ignore[reportArgumentType] # plain JSON input items + tools=[_RESPONSES_TOOL], # pyright: ignore[reportArgumentType] # plain JSON tool + store=False, + stream=True, + extra_body={"cache": _NO_CACHE}, + ) + events: Final = [event async for event in stream] + assert "".join(event.delta for event in events if event.type == "response.output_text.delta") == _ANSWER + completed: Final = [event for event in events if event.type == "response.completed"] + assert len(completed) == 1 and completed[0].response.output_text == _ANSWER + done_types: Final = [event.type for event in events if event.type == "response.output_item.done"] + assert done_types == ["response.output_item.done"] + body: Final = _only_generate(wire) + assert body["stream"] is True + _assert_tool_turn(_prompt_of(body)) + assert _model_spend_rows(model, 1)[0]["status"] == "success" + + +def test_legacy_functions_param_is_instructed_the_same_way(gateway: Gateway) -> None: + with _ollama_server(lambda _: _generate_reply(_CALL_JSON)) as wire, gateway.scenario() as scenario: + model: Final = scenario.model(model=f"ollama/{_BACKEND}", api_base=wire.url, api_key=_API_KEY) + function: Final = _WEATHER_TOOL["function"] + payload: Final = _post_chat(gateway, model, _first_turn(), functions=[function]) + choices: Final = payload["choices"] + assert isinstance(choices, list) and len(choices) == 1 + assert _CALL_JSON_FIELDS <= set(json.dumps(choices[0]).split('"')), choices + prompt: Final = _prompt_of(_only_generate(wire)) + _assert_instructed_once(prompt, "get_weather") + assert "Weather for a city" in prompt, prompt + identity: Final = payload["id"] + assert isinstance(identity, str) + assert _spend_row(identity) == _billed(model) + + +def test_string_system_message_keeps_one_system_section_with_the_instruction(gateway: Gateway) -> None: + system: Final = f"You are a terse weather bot {uuid.uuid4().hex}." + with _ollama_server(lambda _: _generate_reply(_CALL_JSON)) as wire, gateway.scenario() as scenario: + model: Final = scenario.model(model=f"ollama/{_BACKEND}", api_base=wire.url, api_key=_API_KEY) + _post_chat(gateway, model, [{"role": "system", "content": system}, *_first_turn()], tools=[_WEATHER_TOOL]) + prompt: Final = _prompt_of(_only_generate(wire)) + _assert_instructed_once(prompt, "get_weather") + assert f"### System:\n{system} {_INSTRUCTION}\n" in prompt, prompt + + +def test_list_system_message_keeps_one_system_section_with_the_instruction(gateway: Gateway) -> None: + system: Final = f"You are a terse weather bot {uuid.uuid4().hex}." + with _ollama_server(lambda _: _generate_reply(_CALL_JSON)) as wire, gateway.scenario() as scenario: + model: Final = scenario.model(model=f"ollama/{_BACKEND}", api_base=wire.url, api_key=_API_KEY) + _post_chat( + gateway, + model, + [{"role": "system", "content": [{"type": "text", "text": system}]}, *_first_turn()], + tools=[_WEATHER_TOOL], + ) + prompt: Final = _prompt_of(_only_generate(wire)) + _assert_instructed_once(prompt, "get_weather") + section: Final = prompt.split("### System:\n", 1)[1] + assert section.startswith(system), section + assert section.count(_INSTRUCTION) == 1, section + + +def test_two_tools_are_both_listed_under_one_instruction(gateway: Gateway) -> None: + with _ollama_server(lambda _: _generate_reply(_CALL_JSON)) as wire, gateway.scenario() as scenario: + model: Final = scenario.model(model=f"ollama/{_BACKEND}", api_base=wire.url, api_key=_API_KEY) + _post_chat(gateway, model, _first_turn(), tools=[_WEATHER_TOOL, _TIME_TOOL]) + prompt: Final = _prompt_of(_only_generate(wire)) + _assert_instructed_once(prompt, "get_weather", "get_time") + assert prompt.index("'name': 'get_weather'") < prompt.index("'name': 'get_time'"), prompt + + +def test_unauthenticated_tool_request_never_reaches_ollama(gateway: Gateway) -> None: + with _ollama_server(lambda _: _generate_reply(_CALL_JSON)) as wire, gateway.scenario() as scenario: + model: Final = scenario.model(model=f"ollama/{_BACKEND}", api_base=wire.url, api_key=_API_KEY) + code, text = _post( + gateway, + "/v1/chat/completions", + {"model": model, "messages": _first_turn(), "tools": [_WEATHER_TOOL]}, + key=f"sk-not-a-key-{uuid.uuid4().hex}", + ) + assert code == 401, text + assert _generate_calls(wire) == () + + +def test_ollama_model_not_found_reaches_the_caller_after_one_attempt(gateway: Gateway) -> None: + message: Final = f"model '{_BACKEND}' not found {uuid.uuid4().hex}" + reply: Final = Reply(status=404, body=json.dumps({"error": message}).encode()) + with _ollama_server(lambda _: reply) as wire, gateway.scenario() as scenario: + model: Final = scenario.model(model=f"ollama/{_BACKEND}", api_base=wire.url, api_key=_API_KEY) + code, text = _post( + gateway, "/v1/chat/completions", {"model": model, "messages": _first_turn(), "tools": [_WEATHER_TOOL]} + ) + assert code == 404, text + assert message in text, text + _assert_instructed_once(_prompt_of(_only_generate(wire)), "get_weather") + + +def test_ollama_server_error_on_the_tool_result_turn_does_not_take_the_deployment_down(gateway: Gateway) -> None: + attempts: Final = itertools.count() + failure: Final = f"internal failure {uuid.uuid4().hex}" + + def respond(_: Request) -> Reply: + if next(attempts) == 0: + return Reply(status=500, body=json.dumps({"error": failure}).encode()) + return _generate_reply(_ANSWER) + + with _ollama_server(respond) as wire, gateway.scenario() as scenario: + model: Final = scenario.model(model=f"ollama/{_BACKEND}", api_base=wire.url, api_key=_API_KEY) + code, text = _post( + gateway, "/v1/chat/completions", {"model": model, "messages": _second_turn(), "tools": [_WEATHER_TOOL]} + ) + assert code == 500, text + assert failure in text, text + payload: Final = _post_chat(gateway, model, _second_turn(), tools=[_WEATHER_TOOL]) + choices: Final = payload["choices"] + assert isinstance(choices, list) and len(choices) == 1 + assert json.dumps(choices[0]).count(_ANSWER) == 1, choices + prompts: Final = [_prompt_of(_JSON_OBJECT.validate_json(request.body)) for request in _generate_calls(wire)] + assert len(prompts) == 2, prompts + for prompt in prompts: + _assert_tool_turn(prompt) + + +@pytest.mark.parametrize( + ("result", "forwarded"), + [ + pytest.param("", None, id="empty-string-drops-the-section"), + pytest.param("r" * 5120, "r" * 5120, id="5kb-string-forwarded-intact"), + pytest.param( + [{"type": "text", "text": "Paris: 22 degrees"}, {"type": "text", "text": "clear skies"}], + "Paris: 22 degreesclear skies", + id="text-parts-joined", + ), + ], +) +def test_tool_result_content_shapes_reach_the_prompt( + gateway: Gateway, result: JsonValue, forwarded: str | None +) -> None: + with _ollama_server(lambda _: _generate_reply(_ANSWER)) as wire, gateway.scenario() as scenario: + model: Final = scenario.model(model=f"ollama/{_BACKEND}", api_base=wire.url, api_key=_API_KEY) + payload: Final = _post_chat(gateway, model, _second_turn(result), tools=[_WEATHER_TOOL]) + assert json.dumps(payload["choices"]).count(_ANSWER) == 1, payload + prompt: Final = _prompt_of(_only_generate(wire)) + _assert_instructed_once(prompt, "get_weather") + if forwarded is None: + assert f"### User:\n{_QUESTION}\n\n### Assistant:\n{_CALL_JSON}\n\n### System:\n" in prompt, prompt + else: + _assert_tool_turn(prompt, forwarded) + + +def test_the_same_tool_result_twice_is_forwarded_twice_under_one_instruction(gateway: Gateway) -> None: + with _ollama_server(lambda _: _generate_reply(_ANSWER)) as wire, gateway.scenario() as scenario: + model: Final = scenario.model(model=f"ollama/{_BACKEND}", api_base=wire.url, api_key=_API_KEY) + messages: Final = [*_second_turn(), {"role": "tool", "tool_call_id": _CALL_ID, "content": _RESULT}] + _post_chat(gateway, model, messages, tools=[_WEATHER_TOOL]) + prompt: Final = _prompt_of(_only_generate(wire)) + _assert_instructed_once(prompt, "get_weather") + assert prompt.count(_RESULT) == 2, prompt + assert prompt.count("### User:") == 2 and prompt.count("### Assistant:") == 1, prompt + + +def test_a_second_function_call_after_the_result_is_surfaced_as_tool_calls(gateway: Gateway) -> None: + with _ollama_server(lambda _: _generate_reply(_CALL_JSON)) as wire, gateway.scenario() as scenario: + model: Final = scenario.model(model=f"ollama/{_BACKEND}", api_base=wire.url, api_key=_API_KEY) + completion: Final = _openai_client(gateway).chat.completions.create( + model=model, + messages=_second_turn(), # pyright: ignore[reportArgumentType] # plain JSON messages + tools=[_WEATHER_TOOL], # pyright: ignore[reportArgumentType] # plain JSON tool + extra_body={"cache": _NO_CACHE}, + ) + choice: Final = completion.choices[0] + assert choice.finish_reason == "tool_calls" + assert choice.message.tool_calls is not None and [call.function.name for call in choice.message.tool_calls] == [ + "get_weather" + ] + _assert_tool_turn(_prompt_of(_only_generate(wire))) + assert _spend_row(completion.id) == _billed(model) + + +def test_a_non_function_json_answer_is_returned_as_text(gateway: Gateway) -> None: + answer: Final = json.dumps({"city": "Paris", "temperature_c": 22, "sky": "clear"}) + with _ollama_server(lambda _: _generate_reply(answer)) as wire, gateway.scenario() as scenario: + model: Final = scenario.model(model=f"ollama/{_BACKEND}", api_base=wire.url, api_key=_API_KEY) + completion: Final = _openai_client(gateway).chat.completions.create( + model=model, + messages=_second_turn(), # pyright: ignore[reportArgumentType] # plain JSON messages + tools=[_WEATHER_TOOL], # pyright: ignore[reportArgumentType] # plain JSON tool + extra_body={"cache": _NO_CACHE}, + ) + choice: Final = completion.choices[0] + assert choice.finish_reason == "stop" + assert choice.message.tool_calls is None + assert choice.message.content is not None and json.loads(choice.message.content) == json.loads(answer) + _assert_tool_turn(_prompt_of(_only_generate(wire))) + assert _spend_row(completion.id) == _billed(model) + + +def test_int_tool_result_content_fails_in_the_response_body_and_leaves_the_deployment_serving(gateway: Gateway) -> None: + with _ollama_server(lambda _: _generate_reply(_ANSWER)) as wire, gateway.scenario() as scenario: + model: Final = scenario.model(model=f"ollama/{_BACKEND}", api_base=wire.url, api_key=_API_KEY) + code, text = _post( + gateway, "/v1/chat/completions", {"model": model, "messages": _second_turn(22), "tools": [_WEATHER_TOOL]} + ) + assert code >= 400, text + error: Final = _JSON_OBJECT.validate_json(text)["error"] + assert isinstance(error, dict) and isinstance(error["message"], str) and error["message"], text + assert _generate_calls(wire) == () + payload: Final = _post_chat(gateway, model, _second_turn(), tools=[_WEATHER_TOOL]) + assert json.dumps(payload["choices"]).count(_ANSWER) == 1, payload + _assert_tool_turn(_prompt_of(_only_generate(wire))) + + +def test_ollama_chat_keeps_native_tools_and_gets_no_instruction(gateway: Gateway) -> None: + reply: Final = Reply( + body=json.dumps( + { + "model": _BACKEND, + "created_at": "2026-10-07T00:00:00Z", + "message": {"role": "assistant", "content": _ANSWER}, + "done": True, + "prompt_eval_count": 30, + "eval_count": 12, + } + ).encode() + ) + with wire_server(lambda _: reply) as wire, gateway.scenario() as scenario: + model: Final = scenario.model(model=f"ollama_chat/{_BACKEND}", api_base=wire.url, api_key=_API_KEY) + completion: Final = _openai_client(gateway).chat.completions.create( + model=model, + messages=_second_turn(), # pyright: ignore[reportArgumentType] # plain JSON messages + tools=[_WEATHER_TOOL], # pyright: ignore[reportArgumentType] # plain JSON tool + extra_body={"cache": _NO_CACHE}, + ) + assert completion.choices[0].message.content == _ANSWER + received: Final = tuple(request for request in wire.drain() if request.method == "POST") + assert [request.target for request in received] == ["/api/chat"] + body: Final = _JSON_OBJECT.validate_json(received[0].body) + assert "prompt" not in body and "format" not in body, sorted(body) + tools: Final = body["tools"] + assert isinstance(tools, list) and len(tools) == 1 + messages: Final = body["messages"] + assert isinstance(messages, list) and [item["role"] for item in messages if isinstance(item, dict)] == [ + "user", + "assistant", + "tool", + ] + assert "function_name" not in received[0].body.decode(), received[0].body + assert _spend_row(completion.id) == _billed(model) + + +def test_identical_uncached_tool_result_turns_are_each_forwarded_and_billed_once(gateway: Gateway) -> None: + with _ollama_server(lambda _: _generate_reply(_ANSWER)) as wire, gateway.scenario() as scenario: + model: Final = scenario.model(model=f"ollama/{_BACKEND}", api_base=wire.url, api_key=_API_KEY) + first: Final = _post_chat(gateway, model, _second_turn(), tools=[_WEATHER_TOOL]) + second: Final = _post_chat(gateway, model, _second_turn(), tools=[_WEATHER_TOOL]) + assert first["id"] != second["id"], (first["id"], second["id"]) + prompts: Final = [_prompt_of(_JSON_OBJECT.validate_json(request.body)) for request in _generate_calls(wire)] + assert len(prompts) == 2, prompts + for prompt in prompts: + _assert_tool_turn(prompt) + for payload in (first, second): + identity: Final = payload["id"] + assert isinstance(identity, str) + assert _spend_row(identity) == _billed(model) + + +def test_the_cell_deployment_is_gone_after_its_scenario(gateway: Gateway) -> None: + with _ollama_server(lambda _: _generate_reply(_CALL_JSON)) as wire: + with gateway.scenario() as scenario: + model: Final = scenario.model(model=f"ollama/{_BACKEND}", api_base=wire.url, api_key=_API_KEY) + _post_chat(gateway, model, _first_turn(), tools=[_WEATHER_TOOL]) + _only_generate(wire) + listed: Final = eventually( + lambda: [entry["model_name"] for entry in _deployments(gateway) if isinstance(entry, dict)], + lambda names: model not in names, + seconds=70, + ) + assert model not in listed + + +def _deployments(gateway: Gateway) -> list[JsonValue]: + data: Final = gateway.get("/model/info")["data"] + assert isinstance(data, list) + return data diff --git a/tests/unit/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_factory.py b/tests/unit/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_factory.py index a2987730851..61b59f7dbe6 100644 --- a/tests/unit/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_factory.py +++ b/tests/unit/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_factory.py @@ -51,6 +51,17 @@ def test_function_call_prompt_preserves_append_failure_for_non_string_content() function_call_prompt(messages, []) +def test_function_call_prompt_lets_the_model_answer_after_a_function_result() -> None: + messages: Final[list[dict[str, object]]] = [{"role": "system", "content": "Be terse."}] + + prompted: Final = function_call_prompt(messages, [{"name": "get_weather"}]) + + system: Final = str(prompted[0]["content"]) + assert "JSON OUTPUT ONLY" not in system + assert "reply to the user in plain text instead of calling a function again" in system + assert "{'name': 'get_weather'}" in system + + @pytest.mark.parametrize( ("thought_signature", "expected"), [