diff --git a/litellm/llms/opencode/chat/transformation.py b/litellm/llms/opencode/chat/transformation.py index c6d288b6d84..7b6340c93a5 100644 --- a/litellm/llms/opencode/chat/transformation.py +++ b/litellm/llms/opencode/chat/transformation.py @@ -6,7 +6,9 @@ from pydantic import TypeAdapter import litellm from litellm.litellm_core_utils.url_utils import ensure_api_route -from litellm.llms.anthropic.chat.transformation import AnthropicConfig +from litellm.llms.anthropic.chat.handler import ModelResponseIterator as AnthropicModelResponseIterator +from litellm.llms.anthropic.chat.transformation import AnthropicConfig, JsonModeRequest +from litellm.llms.base_llm.base_model_iterator import BaseModelResponseIterator from litellm.llms.base_llm.chat.transformation import BaseConfig from litellm.llms.gemini.chat.transformation import GoogleAIStudioGeminiConfig from litellm.llms.openai_like.chat.transformation import OpenAILikeChatConfig @@ -20,7 +22,7 @@ from litellm.llms.opencode.common_utils import ( ) from litellm.secret_managers.main import get_secret_str from litellm.types.llms.openai import AllMessageValues -from litellm.types.utils import ModelResponse +from litellm.types.utils import ModelResponse, Usage if TYPE_CHECKING: from litellm.litellm_core_utils.litellm_logging import Logging @@ -107,7 +109,41 @@ class OpenCodeGoChatConfig(_GoReportedCostMixin, OpenAILikeChatConfig): return [f"opencode_go/{entry['id']}" for entry in models if isinstance(entry.get("id"), str)] +class _GoMessagesResponseIterator(AnthropicModelResponseIterator, BaseModelResponseIterator): + http_response: httpx.Response | None = None + + def _handle_usage(self, anthropic_usage_chunk: Mapping[str, object]) -> Usage: + usage: Final = super()._handle_usage(anthropic_usage_chunk) + reported_usage: Final = {key: value for key, value in self._cumulative_usage.items() if key == "opencode"} + return usage.model_copy(update=reported_usage) + + def close(self) -> None: + BaseModelResponseIterator.close(self) + + async def aclose(self) -> None: + if self.sync_stream: + self.close() + return + await BaseModelResponseIterator.aclose(self) + + class OpenCodeGoMessagesChatConfig(_GoReportedCostMixin, AnthropicConfig): + def get_model_response_iterator_for_request( + self, + streaming_response: Iterator[str] | AsyncIterator[str] | ModelResponse, + sync_stream: bool, + json_mode: bool | None, + optional_params: Mapping[str, object], + logging_obj: "Logging | None" = None, + ) -> object: + iterator: Final = _GoMessagesResponseIterator( + streaming_response=streaming_response, + sync_stream=sync_stream, + json_mode=json_mode, + response_format_schema=JsonModeRequest.model_validate(optional_params).tool_schema() if json_mode else None, + ) + return iter(iterator) if sync_stream else iterator + @property def custom_llm_provider(self) -> str: return "opencode_go" diff --git a/tests/integration/pricing/test_opencode_go_reported_cost_wire.py b/tests/integration/pricing/test_opencode_go_reported_cost_wire.py index f81b2c78f21..b2157d94361 100644 --- a/tests/integration/pricing/test_opencode_go_reported_cost_wire.py +++ b/tests/integration/pricing/test_opencode_go_reported_cost_wire.py @@ -148,6 +148,26 @@ def _respond(request: Request) -> Reply: assert request.method == "POST", request.method streaming: Final = _JSON_OBJECT.validate_json(request.body).get("stream") is True tag: Final = uuid.uuid4().hex + if ":streamGenerateContent" in request.target: + body: Final = _JSON_OBJECT.validate_json(request.body) + assert "contents" in body and "messages" not in body, body + return Reply( + content_type="text/event-stream", + chunks=( + _data_frame( + { + "candidates": [ + {"content": {"role": "model", "parts": [{"text": "ok"}]}, "finishReason": "STOP"} + ], + "usageMetadata": { + "promptTokenCount": 1200, + "candidatesTokenCount": 300, + "totalTokenCount": 1500, + }, + } + ), + ), + ) if request.target.endswith("/chat/completions"): return ( Reply(content_type="text/event-stream", chunks=_chat_stream(tag)) @@ -159,6 +179,8 @@ def _respond(request: Request) -> Reply: return Reply(content_type="text/event-stream", chunks=_responses_stream(tag)) return Reply(body=json.dumps(_responses_response(tag)).encode()) assert request.target.endswith("/v1/messages"), request.target + messages_body: Final = _JSON_OBJECT.validate_json(request.body) + assert "messages" in messages_body and "contents" not in messages_body, messages_body if streaming: return Reply(content_type="text/event-stream", chunks=_anthropic_stream(tag)) return Reply(body=json.dumps(_anthropic_message(tag)).encode()) @@ -198,3 +220,46 @@ def test_request_is_billed_at_the_reported_list_cost(gateway: Gateway, path: str if not stream: assert float(response.headers["x-litellm-response-cost"]) == _approx(_LIST_COST) assert _logged_spend(response.headers["x-litellm-call-id"]) == _approx(_LIST_COST) + + +@pytest.mark.parametrize( + ("model", "stream"), + ( + ("opencode_go/minimax-m3", True), + ("opencode_go/minimax-m3", False), + ("opencode_go/minimax-m3.1-flash", True), + ("opencode_go/gemini-4-argon", True), + ), + ids=("messages-stream", "messages-non-stream", "openai-stream", "gemini-stream"), +) +def test_chat_completion_uses_go_model_stream_protocol(gateway: Gateway, model: str, stream: bool) -> None: + with wire_server(_respond) as wire, gateway.scenario() as scenario: + deployment: Final = scenario.model(model=model, api_base=wire.url, api_key="synthetic-go-key") + response: Final = gateway.request( + "POST", "/v1/chat/completions", _body("/v1/chat/completions", deployment, stream) + ) + assert response.status_code == 200, response.text + if not stream: + body: Final = _JSON_OBJECT.validate_json(response.content) + nonstream_choices: Final = TypeAdapter(tuple[dict[str, JsonValue], ...]).validate_python(body["choices"]) + message: Final = _JSON_OBJECT.validate_python(nonstream_choices[0]["message"]) + assert message["content"] == "ok", response.text + assert nonstream_choices[0]["finish_reason"] == "stop", response.text + return + events: Final = tuple( + _JSON_OBJECT.validate_json(line.removeprefix("data: ")) + for line in response.text.splitlines() + if line.startswith("data: {") + ) + choice_groups: Final = tuple( + TypeAdapter(tuple[dict[str, JsonValue], ...]).validate_python(event["choices"]) for event in events + ) + choices: Final = tuple(group[0] for group in choice_groups if group) + deltas: Final = tuple(_JSON_OBJECT.validate_python(choice["delta"]) for choice in choices) + assert "".join(str(delta.get("content") or "") for delta in deltas) == "ok", response.text + assert tuple(choice["finish_reason"] for choice in choices if choice.get("finish_reason")) == ("stop",), ( + response.text + ) + assert response.text.rstrip().endswith("data: [DONE]"), response.text + if model == "opencode_go/minimax-m3": + assert _logged_spend(response.headers["x-litellm-call-id"]) == _approx(_LIST_COST) diff --git a/tests/unit/llms/opencode/chat/test_transformation.py b/tests/unit/llms/opencode/chat/test_transformation.py index 08c80db294c..576f2201130 100644 --- a/tests/unit/llms/opencode/chat/test_transformation.py +++ b/tests/unit/llms/opencode/chat/test_transformation.py @@ -1,8 +1,8 @@ import asyncio import json -from collections.abc import Mapping +from collections.abc import AsyncIterator, Iterator, Mapping from datetime import datetime -from typing import Final +from typing import Final, cast import httpx import pytest @@ -16,6 +16,151 @@ from litellm.litellm_core_utils.litellm_logging import Logging from litellm.llms.custom_httpx.llm_http_handler import BaseLLMHTTPHandler from litellm.llms.opencode.chat.transformation import OpenCodeGoMessagesChatConfig from litellm.llms.opencode.common_utils import go_headers +from litellm.types.utils import ModelResponseStream + + +def _messages_sse( + tool: bool, + usage: Mapping[str, object] | None = None, + usage_at_start: bool = False, + tool_name: str = "weather", + arguments: str = '{"city":"Paris"}', +) -> tuple[str, ...]: + block: Final = ( + {"type": "tool_use", "id": "call_test", "name": tool_name, "input": {}} + if tool + else {"type": "text", "text": ""} + ) + delta: Final = ( + {"type": "input_json_delta", "partial_json": arguments} + if tool + else {"type": "text_delta", "text": "hello"} + ) + events: Final = ( + { + "type": "message_start", + "message": { + "id": "msg_test", + "type": "message", + "role": "assistant", + "model": "fixture-messages", + "content": [], + "usage": {"input_tokens": 10, "output_tokens": 0, **((usage or {}) if usage_at_start else {})}, + }, + }, + {"type": "content_block_start", "index": 0, "content_block": block}, + {"type": "content_block_delta", "index": 0, "delta": delta}, + {"type": "content_block_stop", "index": 0}, + { + "type": "message_delta", + "delta": {"stop_reason": "tool_use" if tool else "end_turn"}, + "usage": {"output_tokens": 5, **((usage or {}) if not usage_at_start else {})}, + }, + {"type": "message_stop"}, + ) + return tuple(f"data: {json.dumps(event)}" for event in events) + + +async def _async_lines(lines: tuple[str, ...]) -> AsyncIterator[str]: + for line in lines: + yield line + + +@pytest.mark.asyncio +@pytest.mark.parametrize("sync", (True, False), ids=("sync", "async")) +@pytest.mark.parametrize("tool", (False, True), ids=("text", "tool")) +async def test_messages_stream_translates_deltas_and_finish_reason(sync: bool, tool: bool) -> None: + lines: Final = _messages_sse(tool) + stream: Final = OpenCodeGoMessagesChatConfig().get_model_response_iterator_for_request( + iter(lines) if sync else _async_lines(lines), sync, False, {} + ) + chunks: Final = ( + tuple(cast(Iterator[ModelResponseStream], stream)) + if sync + else tuple([chunk async for chunk in cast(AsyncIterator[ModelResponseStream], stream)]) + ) + choices: Final = tuple(chunk.choices[0] for chunk in chunks if chunk.choices) + assert tuple(choice.finish_reason for choice in choices if choice.finish_reason) == ( + "tool_calls" if tool else "stop", + ) + if not tool: + assert "".join(choice.delta.content or "" for choice in choices) == "hello" + return + calls: Final = tuple(choice.delta.tool_calls for choice in choices if choice.delta.tool_calls) + assert calls[0][0].id == "call_test" + assert calls[0][0].function.name == "weather" + assert "".join(call[0].function.arguments or "" for call in calls) == '{"city":"Paris"}' + + +@pytest.mark.asyncio +@pytest.mark.parametrize("sync", (True, False), ids=("sync", "async")) +@pytest.mark.parametrize("valid_schema", (True, False, None), ids=("schema", "invalid-tools", "no-tools")) +async def test_messages_json_stream_uses_request_schema(sync: bool, valid_schema: bool | None) -> None: + answer: Final = {"city": "Paris"} + arguments: Final = json.dumps({"name": "json_tool_call", "parameters": answer}) if valid_schema else json.dumps(answer) + lines: Final = _messages_sse(True, tool_name="json_tool_call", arguments=arguments) + schema: Final = {"type": "object", "properties": {"city": {"type": "string"}}, "required": ["city"]} + tools: Final = [{"name": "json_tool_call", "input_schema": schema}] if valid_schema else "invalid" + stream: Final = OpenCodeGoMessagesChatConfig().get_model_response_iterator_for_request( + iter(lines) if sync else _async_lines(lines), sync, True, {"tools": tools} if valid_schema is not None else {} + ) + chunks: Final = ( + tuple(cast(Iterator[ModelResponseStream], stream)) + if sync + else tuple([chunk async for chunk in cast(AsyncIterator[ModelResponseStream], stream)]) + ) + choices: Final = tuple(chunk.choices[0] for chunk in chunks if chunk.choices) + assert json.loads("".join(choice.delta.content or "" for choice in choices)) == answer + assert not any(choice.delta.tool_calls for choice in choices) + assert tuple(choice.finish_reason for choice in choices if choice.finish_reason) == ("stop",) + + +class _MessagesStream(httpx.SyncByteStream, httpx.AsyncByteStream): + def __init__(self, body: bytes) -> None: + self.body = body + self.closed = False + + def __iter__(self) -> Iterator[bytes]: + yield self.body + + async def __aiter__(self) -> AsyncIterator[bytes]: + yield self.body + + def close(self) -> None: + self.closed = True + + async def aclose(self) -> None: + self.closed = True + + +@pytest.mark.asyncio +@pytest.mark.parametrize("mode", ("sync", "async", "async-close-sync")) +async def test_closing_messages_stream_releases_upstream_response(go_catalog: None, mode: str) -> None: + upstream: Final = _MessagesStream(("\n\n".join(_messages_sse(False)) + "\n\n").encode()) + with respx.mock(using="httpx") as transport: + transport.post().mock( + return_value=httpx.Response(200, headers={"content-type": "text/event-stream"}, stream=upstream) + ) + kwargs: Final = { + "model": "opencode_go/fixture-messages", + "api_key": "fake-go-key", + "messages": [{"role": "user", "content": "hello"}], + "stream": True, + } + if mode != "async": + stream: Final = litellm.completion(**kwargs) + next(stream) + assert not upstream.closed + if mode == "sync": + stream.close() + else: + await stream.aclose() + else: + async_stream: Final = await litellm.acompletion(**kwargs) + await async_stream.__anext__() + assert not upstream.closed + await async_stream.aclose() + assert upstream.closed @pytest.fixture @@ -440,6 +585,11 @@ class CostRecorder(CustomLogger): async def async_log_success_event(self, kwargs, response_obj, start_time, end_time) -> None: # noqa: ANN001 self.costs.append(kwargs.get("response_cost")) + def log_success_event( + self, kwargs: Mapping[str, object], response_obj: object, start_time: datetime, end_time: datetime + ) -> None: + self.costs.append(TypeAdapter(float | None).validate_python(kwargs.get("response_cost"))) + async def logged_cost(recorder: CostRecorder) -> float | None: for _ in range(50): @@ -496,6 +646,47 @@ def sse(*events: Mapping[str, object]) -> str: return "".join(f"data: {json.dumps(event)}\n\n" for event in events) + "data: [DONE]\n\n" +@pytest.mark.asyncio +@pytest.mark.parametrize("sync", (True, False), ids=("sync", "async")) +@pytest.mark.parametrize("usage_at_start", (False, True), ids=("delta-usage", "start-usage")) +@pytest.mark.parametrize( + ("usage", "expected"), + ( + (GO_REPORTED_USAGE, 0.00412), + ({"opencode": {"list_cost": 0}}, 0), + ({}, ROW_PRICE_FOR_FIXTURE_USAGE), + ({"opencode": {"list_cost": "invalid"}}, ROW_PRICE_FOR_FIXTURE_USAGE), + ), +) +async def test_messages_stream_bills_reported_list_cost( + go_catalog: None, + cost_recorder: CostRecorder, + sync: bool, + usage_at_start: bool, + usage: Mapping[str, object], + expected: float, +) -> None: + lines: Final = _messages_sse(False, usage, usage_at_start) + with respx.mock(using="httpx") as transport: + transport.post().mock( + return_value=httpx.Response( + 200, headers={"content-type": "text/event-stream"}, text="\n\n".join(lines) + "\n\n" + ) + ) + kwargs: Final = { + "model": "opencode_go/fixture-messages", + "api_key": "fake-go-key", + "messages": [{"role": "user", "content": "hello"}], + "stream": True, + "stream_options": {"include_usage": True}, + } + stream: Final = litellm.completion(**kwargs) if sync else await litellm.acompletion(**kwargs) + chunks: Final = tuple(stream) if sync else tuple([chunk async for chunk in stream]) + assert "".join(chunk.choices[0].delta.content or "" for chunk in chunks if chunk.choices) == "hello" + assert chunks[-1].usage.total_tokens == 15 + assert await logged_cost(cost_recorder) == pytest.approx(expected) + + @pytest.mark.asyncio async def test_chat_stream_bills_the_final_usage_chunks_list_cost( go_catalog: None, cost_recorder: CostRecorder @@ -595,7 +786,9 @@ def test_stream_chunk_builder_reads_list_cost_from_dict_chunks(): @pytest.mark.parametrize(("model", "cache_marks"), [("qwen3.8-flash", 0), ("minimax-m3", 0), ("claude-haiku-5-5", 4)]) -def test_go_messages_route_carries_prompt_cache_breakpoint_to_cache_control_only_for_claude(model: str, cache_marks: int): +def test_go_messages_route_carries_prompt_cache_breakpoint_to_cache_control_only_for_claude( + model: str, cache_marks: int +): marked_notes: Final = [ {"type": "text", "text": f"Note {n}", "prompt_cache_breakpoint": {"mode": "explicit"}} for n in range(4) ]