mirror of
https://github.com/BerriAI/litellm.git
synced 2026-11-14 09:03:16 +00:00
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:
parent
8ef451a187
commit
752eb7c5b7
3 changed files with 299 additions and 5 deletions
|
|
@ -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"
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
]
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue