mirror of
https://github.com/andrewyng/openworker.git
synced 2026-11-04 13:27:11 +00:00
Devika Verma
08e6ab2c7d
OpenRouter names the host that served a request in a top-level `provider` field on responses and stream chunks. Keep it on AssistantTurn.served_by and persist it as the `served_by` sidecar (event + message, stripped before provider calls), so a run pinned to one host (Kimi K3 via OpenRouter -> Together, because the sandbox hosts cannot reach Together directly) can prove the pin held on every reply.
93 lines
3.6 KiB
Python
93 lines
3.6 KiB
Python
"""OpenAI-compatible routers (OpenRouter) name the upstream host that served a reply in a
|
|
top-level `provider` field. Keep it on the turn and persist it as the `served_by` sidecar,
|
|
so a run pinned to one host can prove the pin held on every reply."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import asyncio
|
|
from types import SimpleNamespace
|
|
|
|
import aisuite as ai
|
|
from coworker.engine import TurnEngine, _assistant_message
|
|
from coworker.events import EventType
|
|
from coworker.permissions import PermissionEngine
|
|
from coworker.providers import AssistantTurn, ModelCapabilities, OpenAIProvider, ProviderClient
|
|
from coworker.tools import ToolRegistry
|
|
|
|
|
|
class _FakeOpenAI:
|
|
def __init__(self, *, response=None, stream_chunks=None):
|
|
self.calls: list[dict] = []
|
|
|
|
def create(**kwargs):
|
|
self.calls.append(kwargs)
|
|
return stream_chunks if kwargs.get("stream") else response
|
|
|
|
self.chat = SimpleNamespace(completions=SimpleNamespace(create=create))
|
|
|
|
|
|
def test_openai_complete_keeps_the_serving_host():
|
|
response = SimpleNamespace(
|
|
choices=[SimpleNamespace(message=SimpleNamespace(content="hi", tool_calls=None), finish_reason="stop")],
|
|
usage=None,
|
|
provider="Together",
|
|
)
|
|
turn = OpenAIProvider(client=_FakeOpenAI(response=response)).complete(
|
|
model="moonshotai/kimi-k3", messages=[{"role": "user", "content": "x"}]
|
|
)
|
|
assert turn.served_by == "Together"
|
|
|
|
|
|
def test_openai_stream_keeps_the_serving_host_from_any_chunk():
|
|
chunks = [
|
|
SimpleNamespace(choices=[SimpleNamespace(delta=SimpleNamespace(content="hi", tool_calls=None), finish_reason=None)], usage=None),
|
|
SimpleNamespace(choices=[SimpleNamespace(delta=SimpleNamespace(content="!", tool_calls=None), finish_reason="stop")], usage=None, provider="Together"),
|
|
]
|
|
chunks_out = list(
|
|
OpenAIProvider(client=_FakeOpenAI(stream_chunks=chunks)).stream(
|
|
model="moonshotai/kimi-k3", messages=[{"role": "user", "content": "x"}]
|
|
)
|
|
)
|
|
assert chunks_out[-1].turn.served_by == "Together"
|
|
|
|
|
|
def test_plain_endpoints_leave_it_unset():
|
|
response = SimpleNamespace(
|
|
choices=[SimpleNamespace(message=SimpleNamespace(content="hi", tool_calls=None), finish_reason="stop")],
|
|
usage=None,
|
|
)
|
|
turn = OpenAIProvider(client=_FakeOpenAI(response=response)).complete(
|
|
model="m", messages=[{"role": "user", "content": "x"}]
|
|
)
|
|
assert turn.served_by is None
|
|
assert "served_by" not in _assistant_message(turn)
|
|
|
|
|
|
class _ServedProvider(ProviderClient):
|
|
def complete(self, *, model, messages, tools=None, **settings):
|
|
return AssistantTurn(text="done", finish_reason="stop", served_by="Together")
|
|
|
|
def capabilities(self, model):
|
|
return ModelCapabilities()
|
|
|
|
|
|
def test_engine_persists_served_by_and_strips_it_outbound(tmp_path):
|
|
registry = ToolRegistry()
|
|
registry.register_all(ai.toolkits.files(root=str(tmp_path), allow_write=True))
|
|
engine = TurnEngine(
|
|
provider=_ServedProvider(),
|
|
registry=registry,
|
|
permissions=PermissionEngine(workspace_root=tmp_path),
|
|
model="openrouter:moonshotai/kimi-k3",
|
|
)
|
|
|
|
async def _collect():
|
|
return [ev async for ev in engine.run("hello")]
|
|
|
|
events = asyncio.run(_collect())
|
|
assistant = next(ev for ev in events if ev.type == EventType.ASSISTANT_MESSAGE)
|
|
assert assistant.data["served_by"] == "Together"
|
|
persisted = next(m for m in engine.messages if m.get("role") == "assistant")
|
|
assert persisted["served_by"] == "Together"
|
|
assert all("served_by" not in m for m in engine._outbound_messages())
|