fix(opencode): stream Go Messages-protocol models through chat completions (#51715)

* test(opencode): cover Go Messages chat streaming

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

* fix(opencode): delegate Go Messages chat streaming to Anthropic

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

* fix(opencode): preserve Messages stream usage and cleanup

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

---------

Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
This commit is contained in:
moe-berri 2026-11-14 11:21:52 +08:00 • committed by GitHub
parent 8ef451a187
commit 752eb7c5b7
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
3 changed files with 299 additions and 5 deletions

View file

@ -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"

View file

@ -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)

View file

@ -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)
]