feat(R03): unify provider catalogue, routing decisions and usage telemetry
EPIC R03 (Team Duy) - one provider catalogue, one routing flow, one usage seam.
R03-T01 tests/contracts/test_providers.py
29 contract tests every provider must satisfy: canonical assistant message,
streamed text == returned content, reasoning never joins the answer, parsed
tool arguments, ProviderError for every failure. Real adapters exercised
offline by stubbing Provider._request.
R03-T02 domain/models/provider_descriptor.py
infrastructure/providers/provider_registry.py
Provider facts declared once (was split across providers/factory.py,
DEFAULT_CONFIG and PROVIDER_LABELS). ProviderRegistry.build() also stamps the
descriptor id onto the instance, so ollama/github_copilot/codex usage is no
longer all attributed to "openai_compat", and never mutates the caller config.
R03-T03 application/model_routing/routing_application_service.py
Pure-Python routing policy with four modes: Off, Auto, Manual and the new
Fallback (switch only AFTER the current model fails). Depends on a RoutingPort
protocol; production wires the existing core.routing engine underneath.
R03-T04/T05 ui/chat_panel.py, ui/co4e_tab.py, ui/folder_tab.py
Three near-identical routing copies (~40 lines each) replaced by a call to
ctx.routing_application() plus a confirm callback. Mode vocabulary now lives
in one place (normalize_mode/is_valid_mode) instead of four literal tuples.
R03-T06 infrastructure/telemetry/usage_sink.py
Token usage extracted from both providers into UsageEvent + UsageEventSink.
Estimation pinned against core.usage_tracker so no recorded number changes.
Also fixes a deadlock introduced while wiring AppContext: routing_application()
held _routing_lock and called routing(), which takes the same non-reentrant lock.
Suite: 186 passed, 1.22s. check_imports: PASS. All new files < 400 LOC.
2 pre-existing failures remain in test_config_security.py (EPIC R02/Team Nam).
Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
This commit is contained in:
Claude Opus 5
@@ -0,0 +1,8 @@
|
||||
"""Contract tests: one shared behaviour suite every implementation must satisfy.
|
||||
|
||||
Unlike unit tests (which test one module in isolation) a contract test is
|
||||
parametrised over EVERY implementation of an interface, so a newly added
|
||||
provider either satisfies the same promises as the existing ones or the suite
|
||||
goes red on the day it is added - not months later, in production, on the one
|
||||
code path that assumed the promise held.
|
||||
"""
|
||||
@@ -0,0 +1,342 @@
|
||||
"""Provider contract suite (R03-T01).
|
||||
|
||||
Every provider - the two real adapters and the test double - must honour the
|
||||
same promises declared in ``providers/base.py``:
|
||||
|
||||
1. ``chat()`` returns the canonical assistant message
|
||||
``{"role": "assistant", "content": str, "tool_calls": [...]}``.
|
||||
2. Answer text is streamed through ``on_text`` and equals the returned content.
|
||||
3. Private reasoning goes to ``on_reasoning`` ONLY - it must never leak into the
|
||||
answer, or a reasoning model's chain of thought ends up persisted in history.
|
||||
4. Tool calls come back as ``{"id", "name", "arguments": dict}`` with arguments
|
||||
already parsed - callers must never have to json.loads() them.
|
||||
5. A failure raises ``ProviderError`` and nothing else, so one except clause in
|
||||
the agent loop covers every provider.
|
||||
|
||||
The real adapters are exercised WITHOUT network access by replacing
|
||||
``Provider._request`` with a canned SSE response - which is exactly the seam
|
||||
``providers/base.py`` documents for its TLS retry, so no production code needed
|
||||
changing to make this testable.
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
from typing import Any, Dict, List, Optional
|
||||
|
||||
import pytest
|
||||
|
||||
from cowork_local.domain.models.provider_descriptor import ProviderCapability
|
||||
from cowork_local.infrastructure.providers.provider_registry import (
|
||||
BUILT_IN_PROVIDERS,
|
||||
ProviderRegistry,
|
||||
)
|
||||
from cowork_local.providers.anthropic import AnthropicProvider
|
||||
from cowork_local.providers.base import Provider, ProviderError, ToolSpec
|
||||
from cowork_local.providers.openai_compat import OpenAICompatProvider
|
||||
from tests.fakes import FakeProvider, ScriptedTurn
|
||||
|
||||
|
||||
class _StubResponse:
|
||||
"""Minimal stand-in for a streamed ``requests.Response``.
|
||||
|
||||
Only the members the provider code actually touches are implemented; adding
|
||||
more would invite tests that pass against the stub but not against requests.
|
||||
"""
|
||||
|
||||
def __init__(self, lines: List[str], status_code: int = 200, text: str = "") -> None:
|
||||
self._lines = lines
|
||||
self.status_code = status_code
|
||||
self.text = text
|
||||
self.headers: Dict[str, str] = {}
|
||||
self.encoding = "utf-8"
|
||||
self.closed = False
|
||||
|
||||
def iter_lines(self, decode_unicode: bool = False):
|
||||
yield from self._lines
|
||||
|
||||
def close(self) -> None:
|
||||
self.closed = True
|
||||
|
||||
def json(self) -> Any:
|
||||
return json.loads(self.text or "{}")
|
||||
|
||||
|
||||
def _sse(*payloads: Dict[str, Any]) -> List[str]:
|
||||
"""Render payloads as SSE ``data:`` lines, the wire shape both adapters parse."""
|
||||
return [f"data: {json.dumps(p)}" for p in payloads]
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def canned(monkeypatch):
|
||||
"""Return a helper that makes every provider request answer with ``lines``."""
|
||||
|
||||
def _install(lines: List[str], status_code: int = 200, text: str = "") -> Dict[str, Any]:
|
||||
seen: Dict[str, Any] = {}
|
||||
|
||||
def fake_request(self, method, url, **kwargs):
|
||||
# Capture the outgoing payload so tests can assert on how the
|
||||
# canonical message list was translated to the provider's wire format.
|
||||
seen["method"] = method
|
||||
seen["url"] = url
|
||||
seen["json"] = kwargs.get("json")
|
||||
return _StubResponse(lines, status_code=status_code, text=text)
|
||||
|
||||
monkeypatch.setattr(Provider, "_request", fake_request, raising=True)
|
||||
return seen
|
||||
|
||||
return _install
|
||||
|
||||
|
||||
# --------------------------------------------------------------------------- #
|
||||
# Shared base-class behaviour every provider inherits
|
||||
# --------------------------------------------------------------------------- #
|
||||
def _providers_under_test() -> List[Provider]:
|
||||
"""One instance of each implementation, configured but never called."""
|
||||
conf = {"base_url": "https://example.invalid/v1", "api_key": "k", "model": "m"}
|
||||
return [
|
||||
OpenAICompatProvider(dict(conf)),
|
||||
AnthropicProvider(dict(conf)),
|
||||
FakeProvider(),
|
||||
]
|
||||
|
||||
|
||||
@pytest.mark.parametrize("provider", _providers_under_test(), ids=lambda p: type(p).__name__)
|
||||
def test_every_provider_exposes_the_base_contract(provider):
|
||||
assert isinstance(provider, Provider)
|
||||
assert callable(provider.chat)
|
||||
assert callable(provider.list_models)
|
||||
assert callable(provider.test_connection)
|
||||
# `name` identifies the provider in usage records and audit entries; an
|
||||
# implementation that forgot to set it would silently report as "base".
|
||||
assert provider.name and provider.name != "base"
|
||||
assert isinstance(provider.supports_vision, bool)
|
||||
assert provider.describe() == f"{provider.name}:{provider.model}"
|
||||
|
||||
|
||||
@pytest.mark.parametrize("provider", _providers_under_test(), ids=lambda p: type(p).__name__)
|
||||
def test_strip_think_removes_inline_reasoning_from_a_final_answer(provider):
|
||||
"""Safety net for gateways that fold reasoning into the content stream: the
|
||||
answer stored in history must never contain a <think> block."""
|
||||
assert provider.strip_think("<think>secret</think>Answer") == "Answer"
|
||||
assert provider.strip_think("Plain answer") == "Plain answer"
|
||||
|
||||
|
||||
def test_tool_spec_translates_to_both_wire_formats():
|
||||
"""One ToolSpec must render for both protocols - this is what lets the agent
|
||||
loop build its tool catalogue once and reuse it across providers."""
|
||||
spec = ToolSpec(name="save_file", description="Write a file",
|
||||
parameters={"type": "object", "properties": {}})
|
||||
|
||||
openai_shape = spec.to_openai()
|
||||
anthropic_shape = spec.to_anthropic()
|
||||
|
||||
assert openai_shape["type"] == "function"
|
||||
assert openai_shape["function"]["name"] == "save_file"
|
||||
assert openai_shape["function"]["parameters"] == spec.parameters
|
||||
# Anthropic names the same field `input_schema`; the values must stay equal,
|
||||
# otherwise the same tool would validate differently per provider.
|
||||
assert anthropic_shape["name"] == "save_file"
|
||||
assert anthropic_shape["input_schema"] == spec.parameters
|
||||
|
||||
|
||||
# --------------------------------------------------------------------------- #
|
||||
# Streaming contract - real adapters, canned transport
|
||||
# --------------------------------------------------------------------------- #
|
||||
def test_openai_compat_streams_text_and_returns_canonical_message(canned):
|
||||
canned(_sse(
|
||||
{"choices": [{"delta": {"content": "Hel"}}]},
|
||||
{"choices": [{"delta": {"content": "lo"}}]},
|
||||
) + ["data: [DONE]"])
|
||||
provider = OpenAICompatProvider({"base_url": "https://x.invalid/v1",
|
||||
"api_key": "k", "model": "m"})
|
||||
chunks: List[str] = []
|
||||
|
||||
result = provider.chat([{"role": "user", "content": "hi"}], on_text=chunks.append)
|
||||
|
||||
assert "".join(chunks) == "Hello"
|
||||
assert result["role"] == "assistant"
|
||||
assert result["content"] == "Hello"
|
||||
assert result["tool_calls"] == []
|
||||
|
||||
|
||||
def test_openai_compat_keeps_reasoning_out_of_the_answer(canned):
|
||||
canned(_sse(
|
||||
{"choices": [{"delta": {"reasoning_content": "hmm..."}}]},
|
||||
{"choices": [{"delta": {"content": "42"}}]},
|
||||
) + ["data: [DONE]"])
|
||||
provider = OpenAICompatProvider({"base_url": "https://x.invalid/v1",
|
||||
"api_key": "k", "model": "m"})
|
||||
text: List[str] = []
|
||||
reasoning: List[str] = []
|
||||
|
||||
result = provider.chat([{"role": "user", "content": "q"}],
|
||||
on_text=text.append, on_reasoning=reasoning.append)
|
||||
|
||||
assert reasoning == ["hmm..."]
|
||||
assert result["content"] == "42"
|
||||
assert "hmm" not in result["content"]
|
||||
|
||||
|
||||
def test_openai_compat_returns_tool_calls_with_parsed_arguments(canned):
|
||||
"""Arguments arrive as a JSON string split across chunks; the contract says
|
||||
the caller receives a ready-to-use dict."""
|
||||
canned(_sse(
|
||||
{"choices": [{"delta": {"tool_calls": [
|
||||
{"index": 0, "id": "call_a", "function": {"name": "save_file",
|
||||
"arguments": '{"filename":'}}]}}]},
|
||||
{"choices": [{"delta": {"tool_calls": [
|
||||
{"index": 0, "function": {"arguments": '"a.md"}'}}]}}]},
|
||||
) + ["data: [DONE]"])
|
||||
provider = OpenAICompatProvider({"base_url": "https://x.invalid/v1",
|
||||
"api_key": "k", "model": "m"})
|
||||
|
||||
result = provider.chat([{"role": "user", "content": "save it"}])
|
||||
|
||||
assert len(result["tool_calls"]) == 1
|
||||
call = result["tool_calls"][0]
|
||||
assert call["id"] == "call_a"
|
||||
assert call["name"] == "save_file"
|
||||
assert call["arguments"] == {"filename": "a.md"}
|
||||
|
||||
|
||||
def test_anthropic_streams_text_and_returns_canonical_message(canned):
|
||||
canned(_sse(
|
||||
{"type": "content_block_delta", "index": 0,
|
||||
"delta": {"type": "text_delta", "text": "Hel"}},
|
||||
{"type": "content_block_delta", "index": 0,
|
||||
"delta": {"type": "text_delta", "text": "lo"}},
|
||||
{"type": "message_stop"},
|
||||
))
|
||||
provider = AnthropicProvider({"base_url": "https://x.invalid",
|
||||
"api_key": "k", "model": "m"})
|
||||
chunks: List[str] = []
|
||||
|
||||
result = provider.chat([{"role": "user", "content": "hi"}], on_text=chunks.append)
|
||||
|
||||
assert "".join(chunks) == "Hello"
|
||||
assert result["content"] == "Hello"
|
||||
assert result["role"] == "assistant"
|
||||
|
||||
|
||||
def test_anthropic_keeps_extended_thinking_out_of_the_answer(canned):
|
||||
canned(_sse(
|
||||
{"type": "content_block_delta", "index": 0,
|
||||
"delta": {"type": "thinking_delta", "thinking": "reasoning..."}},
|
||||
{"type": "content_block_delta", "index": 0,
|
||||
"delta": {"type": "text_delta", "text": "42"}},
|
||||
{"type": "message_stop"},
|
||||
))
|
||||
provider = AnthropicProvider({"base_url": "https://x.invalid",
|
||||
"api_key": "k", "model": "m"})
|
||||
reasoning: List[str] = []
|
||||
|
||||
result = provider.chat([{"role": "user", "content": "q"}], on_reasoning=reasoning.append)
|
||||
|
||||
assert reasoning == ["reasoning..."]
|
||||
assert result["content"] == "42"
|
||||
|
||||
|
||||
def test_anthropic_returns_tool_calls_with_parsed_arguments(canned):
|
||||
canned(_sse(
|
||||
{"type": "content_block_start", "index": 0,
|
||||
"content_block": {"type": "tool_use", "id": "toolu_1", "name": "save_file"}},
|
||||
{"type": "content_block_delta", "index": 0,
|
||||
"delta": {"type": "input_json_delta", "partial_json": '{"filename":"a.md"}'}},
|
||||
{"type": "message_stop"},
|
||||
))
|
||||
provider = AnthropicProvider({"base_url": "https://x.invalid",
|
||||
"api_key": "k", "model": "m"})
|
||||
|
||||
result = provider.chat([{"role": "user", "content": "save"}])
|
||||
|
||||
assert result["tool_calls"] == [
|
||||
{"id": "toolu_1", "name": "save_file", "arguments": {"filename": "a.md"}}
|
||||
]
|
||||
|
||||
|
||||
@pytest.mark.parametrize("factory", [
|
||||
lambda: OpenAICompatProvider({"base_url": "https://x.invalid/v1", "api_key": "k", "model": "m"}),
|
||||
lambda: AnthropicProvider({"base_url": "https://x.invalid", "api_key": "k", "model": "m"}),
|
||||
], ids=["openai_compat", "anthropic"])
|
||||
def test_transport_failure_surfaces_as_provider_error(canned, factory):
|
||||
"""Every failure mode must arrive as ProviderError so the agent loop needs
|
||||
exactly one except clause, whichever provider is active."""
|
||||
canned([], status_code=500, text="boom")
|
||||
|
||||
with pytest.raises(ProviderError):
|
||||
factory().chat([{"role": "user", "content": "hi"}])
|
||||
|
||||
|
||||
def test_fake_provider_satisfies_the_same_streaming_contract():
|
||||
"""The double is only useful as a stand-in if it keeps the same promises the
|
||||
real adapters are held to above."""
|
||||
provider = FakeProvider([ScriptedTurn(text="Hello", reasoning="hmm")])
|
||||
text: List[str] = []
|
||||
reasoning: List[str] = []
|
||||
|
||||
result = provider.chat([{"role": "user", "content": "hi"}],
|
||||
on_text=text.append, on_reasoning=reasoning.append)
|
||||
|
||||
assert "".join(text) == result["content"] == "Hello"
|
||||
assert reasoning == ["hmm"]
|
||||
assert result["role"] == "assistant"
|
||||
assert result["tool_calls"] == []
|
||||
|
||||
|
||||
def test_fake_provider_raises_provider_error_like_the_real_ones():
|
||||
provider = FakeProvider([ScriptedTurn(error="gateway exploded")])
|
||||
|
||||
with pytest.raises(ProviderError):
|
||||
provider.chat([{"role": "user", "content": "hi"}])
|
||||
|
||||
|
||||
# --------------------------------------------------------------------------- #
|
||||
# Registry <-> implementation agreement
|
||||
# --------------------------------------------------------------------------- #
|
||||
@pytest.mark.parametrize("descriptor", BUILT_IN_PROVIDERS, ids=lambda d: d.id)
|
||||
def test_every_descriptor_builds_a_working_provider(descriptor):
|
||||
"""A descriptor that cannot be built is a catalogue lying to the UI: Settings
|
||||
would list the provider and selecting it would fail at the first message."""
|
||||
registry = ProviderRegistry()
|
||||
conf = {"base_url": "https://x.invalid/v1", "api_key": "k"}
|
||||
|
||||
provider = registry.build(descriptor.id, conf)
|
||||
|
||||
assert isinstance(provider, Provider)
|
||||
# The id, not the shared adapter class name: three descriptors map onto
|
||||
# OpenAICompatProvider, and usage/audit records must still tell them apart.
|
||||
assert provider.name == descriptor.id
|
||||
assert provider.model == descriptor.default_model
|
||||
|
||||
|
||||
@pytest.mark.parametrize("descriptor", BUILT_IN_PROVIDERS, ids=lambda d: d.id)
|
||||
def test_declared_vision_capability_matches_the_implementation(descriptor):
|
||||
"""``supports_vision`` decides whether an image block may be sent. A
|
||||
descriptor claiming vision for an adapter that cannot translate the block
|
||||
would route image turns into a guaranteed failure."""
|
||||
provider = ProviderRegistry().build(descriptor.id, {"base_url": "u", "api_key": "k"})
|
||||
|
||||
if descriptor.supports(ProviderCapability.VISION):
|
||||
assert provider.supports_vision is True
|
||||
|
||||
|
||||
def test_registry_build_never_mutates_the_caller_config():
|
||||
"""The routing layer runs one turn on a different model; if build() wrote
|
||||
that model back into the config dict it was handed, the override would
|
||||
silently become the user's saved default."""
|
||||
registry = ProviderRegistry()
|
||||
conf = {"base_url": "u", "api_key": "k", "model": "configured-model"}
|
||||
|
||||
provider = registry.build("openai_compat", conf, model="routed-model")
|
||||
|
||||
assert provider.model == "routed-model"
|
||||
assert conf["model"] == "configured-model"
|
||||
|
||||
|
||||
def test_registry_rejects_an_unknown_provider_with_provider_error():
|
||||
with pytest.raises(ProviderError) as excinfo:
|
||||
ProviderRegistry().build("does_not_exist", {})
|
||||
|
||||
# The message lists what IS known, so a typo in config is fixable from the
|
||||
# error alone without opening the source.
|
||||
assert "openai_compat" in str(excinfo.value)
|
||||
@@ -0,0 +1,346 @@
|
||||
"""Unit tests for :mod:`application.model_routing` (R03-T03).
|
||||
|
||||
These run against a hand-written fake router rather than ``core.routing``: the
|
||||
point of the service is the DECISION policy around the engine (mode handling,
|
||||
the manual confirm, never-raise behaviour, failure fallback), and mixing in the
|
||||
real scorer would test the wrong thing and drag the suite over its time budget.
|
||||
|
||||
No Qt, no config, no network - the whole file runs in milliseconds, which is the
|
||||
concrete payoff of moving this logic out of ``ui/chat_panel.py``.
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
from dataclasses import dataclass
|
||||
from typing import Any, List, Optional, Tuple
|
||||
|
||||
import pytest
|
||||
|
||||
from cowork_local.application.model_routing import (
|
||||
RoutingApplicationService,
|
||||
RoutingDecision,
|
||||
RoutingMode,
|
||||
)
|
||||
|
||||
|
||||
# --------------------------------------------------------------------------- #
|
||||
# Test doubles shaped like core.routing's RouteResult / SwitchDecision
|
||||
# --------------------------------------------------------------------------- #
|
||||
@dataclass
|
||||
class _TaskType:
|
||||
value: str
|
||||
|
||||
|
||||
@dataclass
|
||||
class _Decision:
|
||||
score_gain: float = 0.0
|
||||
reason: str = ""
|
||||
|
||||
|
||||
@dataclass
|
||||
class _RouteResult:
|
||||
should_switch: bool
|
||||
to: Optional[Tuple[str, str]] = None
|
||||
task_type: Any = None
|
||||
decision: Any = None
|
||||
|
||||
def target(self) -> Optional[Tuple[str, str]]:
|
||||
return self.to
|
||||
|
||||
|
||||
class _FakeRouter:
|
||||
"""Records every route() call and replays a canned result."""
|
||||
|
||||
def __init__(self, result: Any = None, raises: bool = False) -> None:
|
||||
self._result = result or _RouteResult(should_switch=False, decision=_Decision())
|
||||
self._raises = raises
|
||||
self.calls: List[dict] = []
|
||||
|
||||
def route(self, surface, prompt, current_provider, current_model, **kwargs):
|
||||
self.calls.append({"surface": surface, "prompt": prompt,
|
||||
"provider": current_provider, "model": current_model, **kwargs})
|
||||
if self._raises:
|
||||
raise RuntimeError("assessment store is corrupt")
|
||||
return self._result
|
||||
|
||||
|
||||
def _switch_to(provider: str, model: str, gain: float = 0.2, task: str = "coding") -> _RouteResult:
|
||||
return _RouteResult(
|
||||
should_switch=True, to=(provider, model), task_type=_TaskType(task),
|
||||
decision=_Decision(score_gain=gain, reason=f"{task} fit beats current by {gain}"),
|
||||
)
|
||||
|
||||
|
||||
# --------------------------------------------------------------------------- #
|
||||
# Mode parsing
|
||||
# --------------------------------------------------------------------------- #
|
||||
@pytest.mark.parametrize("raw,expected", [
|
||||
("off", RoutingMode.OFF),
|
||||
("AUTO", RoutingMode.AUTO),
|
||||
(" manual ", RoutingMode.MANUAL),
|
||||
("fallback", RoutingMode.FALLBACK),
|
||||
])
|
||||
def test_parse_accepts_the_config_spellings(raw, expected):
|
||||
assert RoutingMode.parse(raw) is expected
|
||||
|
||||
|
||||
@pytest.mark.parametrize("raw", ["", None, "nonsense", 0])
|
||||
def test_parse_degrades_unknown_values_to_off(raw):
|
||||
"""A corrupt setting must leave the user's own model alone rather than
|
||||
silently moving their work onto another model."""
|
||||
assert RoutingMode.parse(raw) is RoutingMode.OFF
|
||||
|
||||
|
||||
# --------------------------------------------------------------------------- #
|
||||
# OFF
|
||||
# --------------------------------------------------------------------------- #
|
||||
def test_off_never_consults_the_engine():
|
||||
router = _FakeRouter(_switch_to("anthropic", "claude"))
|
||||
service = RoutingApplicationService(router)
|
||||
|
||||
decision = service.route_turn("cowork", "hi", "openai_compat", "gpt-4o-mini",
|
||||
mode="off")
|
||||
|
||||
assert router.calls == [] # not even scored: OFF costs nothing
|
||||
assert decision.switched is False
|
||||
assert decision.target() == ("openai_compat", "gpt-4o-mini")
|
||||
|
||||
|
||||
def test_blank_prompt_is_never_routed():
|
||||
"""An empty message carries no signal to classify; all three legacy copies
|
||||
guarded this and the guard has to survive the move."""
|
||||
router = _FakeRouter(_switch_to("anthropic", "claude"))
|
||||
service = RoutingApplicationService(router)
|
||||
|
||||
decision = service.route_turn("cowork", " ", "openai_compat", "m", mode="auto")
|
||||
|
||||
assert router.calls == []
|
||||
assert decision.switched is False
|
||||
|
||||
|
||||
# --------------------------------------------------------------------------- #
|
||||
# AUTO
|
||||
# --------------------------------------------------------------------------- #
|
||||
def test_auto_switches_silently_and_reports_the_target():
|
||||
router = _FakeRouter(_switch_to("anthropic", "claude-sonnet-4-6", gain=0.31))
|
||||
service = RoutingApplicationService(router)
|
||||
|
||||
decision = service.route_turn("cowork", "write a function", "openai_compat", "gpt-4o-mini",
|
||||
mode="auto")
|
||||
|
||||
assert decision.switched is True
|
||||
assert decision.target() == ("anthropic", "claude-sonnet-4-6")
|
||||
assert decision.task_type == "coding"
|
||||
assert decision.score_gain == pytest.approx(0.31)
|
||||
assert decision.should_notify is True
|
||||
|
||||
|
||||
def test_auto_keeps_the_current_model_when_no_candidate_wins():
|
||||
router = _FakeRouter(_RouteResult(should_switch=False, decision=_Decision(reason="no gain")))
|
||||
service = RoutingApplicationService(router)
|
||||
|
||||
decision = service.route_turn("cowork", "hello", "openai_compat", "gpt-4o-mini", mode="auto")
|
||||
|
||||
assert decision.switched is False
|
||||
# The decision still names a model to run on, so the call site never has to
|
||||
# re-derive the fallback itself - the exact drift the three copies suffered.
|
||||
assert decision.target() == ("openai_compat", "gpt-4o-mini")
|
||||
assert decision.should_notify is False
|
||||
|
||||
|
||||
def test_auto_never_asks_for_confirmation():
|
||||
router = _FakeRouter(_switch_to("anthropic", "claude"))
|
||||
asked: List[RoutingDecision] = []
|
||||
service = RoutingApplicationService(router)
|
||||
|
||||
service.route_turn("cowork", "q", "openai_compat", "m", mode="auto",
|
||||
confirm=lambda d: asked.append(d) or True)
|
||||
|
||||
assert asked == []
|
||||
|
||||
|
||||
# --------------------------------------------------------------------------- #
|
||||
# MANUAL
|
||||
# --------------------------------------------------------------------------- #
|
||||
def test_manual_switches_only_after_the_user_approves():
|
||||
router = _FakeRouter(_switch_to("anthropic", "claude"))
|
||||
service = RoutingApplicationService(router)
|
||||
seen: List[RoutingDecision] = []
|
||||
|
||||
def confirm(proposal: RoutingDecision) -> bool:
|
||||
seen.append(proposal)
|
||||
return True
|
||||
|
||||
decision = service.route_turn("cowork", "q", "openai_compat", "m",
|
||||
mode="manual", confirm=confirm)
|
||||
|
||||
assert decision.switched is True
|
||||
assert decision.target() == ("anthropic", "claude")
|
||||
# The dialog is handed the full proposal so it can explain the trade-off.
|
||||
assert seen[0].model == "claude"
|
||||
assert seen[0].score_gain > 0
|
||||
|
||||
|
||||
def test_manual_keeps_the_current_model_when_declined():
|
||||
router = _FakeRouter(_switch_to("anthropic", "claude"))
|
||||
service = RoutingApplicationService(router)
|
||||
|
||||
decision = service.route_turn("cowork", "q", "openai_compat", "gpt-4o-mini",
|
||||
mode="manual", confirm=lambda d: False)
|
||||
|
||||
assert decision.switched is False
|
||||
assert decision.declined is True
|
||||
assert decision.target() == ("openai_compat", "gpt-4o-mini")
|
||||
|
||||
|
||||
def test_manual_without_a_confirm_callback_does_not_switch():
|
||||
"""A headless caller (scheduler) has nobody to ask, so Manual must behave as
|
||||
"not approved" rather than as "approved by default"."""
|
||||
router = _FakeRouter(_switch_to("anthropic", "claude"))
|
||||
service = RoutingApplicationService(router)
|
||||
|
||||
decision = service.route_turn("cowork", "q", "openai_compat", "m", mode="manual")
|
||||
|
||||
assert decision.switched is False
|
||||
assert decision.declined is True
|
||||
|
||||
|
||||
def test_a_confirm_dialog_that_raises_counts_as_declined():
|
||||
"""If the modal blows up (window closing mid-turn) the safe reading is that
|
||||
the user did NOT consent to running on another model."""
|
||||
router = _FakeRouter(_switch_to("anthropic", "claude"))
|
||||
service = RoutingApplicationService(router)
|
||||
|
||||
def confirm(_proposal):
|
||||
raise RuntimeError("dialog destroyed")
|
||||
|
||||
decision = service.route_turn("cowork", "q", "openai_compat", "m",
|
||||
mode="manual", confirm=confirm)
|
||||
|
||||
assert decision.switched is False
|
||||
|
||||
|
||||
# --------------------------------------------------------------------------- #
|
||||
# FALLBACK
|
||||
# --------------------------------------------------------------------------- #
|
||||
def test_fallback_does_not_switch_up_front():
|
||||
"""The whole point of the mode: honour the user's model choice until it
|
||||
actually fails."""
|
||||
router = _FakeRouter(_switch_to("anthropic", "claude"))
|
||||
service = RoutingApplicationService(router)
|
||||
|
||||
decision = service.route_turn("cowork", "q", "openai_compat", "gpt-4o-mini",
|
||||
mode="fallback")
|
||||
|
||||
assert router.calls == []
|
||||
assert decision.switched is False
|
||||
assert decision.target() == ("openai_compat", "gpt-4o-mini")
|
||||
|
||||
|
||||
def test_fallback_switches_after_a_failure():
|
||||
router = _FakeRouter(_switch_to("anthropic", "claude", gain=0.4))
|
||||
service = RoutingApplicationService(router)
|
||||
|
||||
decision = service.fallback_after_failure("cowork", "q", "openai_compat", "gpt-4o-mini",
|
||||
mode="fallback")
|
||||
|
||||
assert decision is not None
|
||||
assert decision.switched is True
|
||||
assert decision.target() == ("anthropic", "claude")
|
||||
assert "failed" in decision.reason
|
||||
|
||||
|
||||
def test_fallback_never_returns_the_model_that_just_failed():
|
||||
"""Retrying the model that just went down would spin on the outage."""
|
||||
router = _FakeRouter(_switch_to("openai_compat", "gpt-4o-mini"))
|
||||
service = RoutingApplicationService(router)
|
||||
|
||||
assert service.fallback_after_failure(
|
||||
"cowork", "q", "openai_compat", "gpt-4o-mini", mode="fallback") is None
|
||||
|
||||
|
||||
def test_fallback_returns_none_when_there_is_no_alternative():
|
||||
router = _FakeRouter(_RouteResult(should_switch=False, decision=_Decision()))
|
||||
service = RoutingApplicationService(router)
|
||||
|
||||
assert service.fallback_after_failure("cowork", "q", "openai_compat", "m",
|
||||
mode="auto") is None
|
||||
|
||||
|
||||
@pytest.mark.parametrize("mode", ["off", "manual"])
|
||||
def test_off_and_manual_do_not_auto_recover_from_a_failure(mode):
|
||||
"""Both modes exist to keep the user in control of which model runs their
|
||||
work; moving it on failure would break that promise silently."""
|
||||
router = _FakeRouter(_switch_to("anthropic", "claude"))
|
||||
service = RoutingApplicationService(router)
|
||||
|
||||
assert service.fallback_after_failure("cowork", "q", "openai_compat", "m",
|
||||
mode=mode) is None
|
||||
|
||||
|
||||
# --------------------------------------------------------------------------- #
|
||||
# Robustness - routing must never break a chat turn
|
||||
# --------------------------------------------------------------------------- #
|
||||
def test_engine_failure_degrades_to_keeping_the_current_model():
|
||||
service = RoutingApplicationService(_FakeRouter(raises=True))
|
||||
|
||||
decision = service.route_turn("cowork", "q", "openai_compat", "gpt-4o-mini", mode="auto")
|
||||
|
||||
assert decision.switched is False
|
||||
assert decision.target() == ("openai_compat", "gpt-4o-mini")
|
||||
|
||||
|
||||
def test_engine_failure_during_fallback_returns_none():
|
||||
"""A broken router must not mask the original provider error with its own."""
|
||||
service = RoutingApplicationService(_FakeRouter(raises=True))
|
||||
|
||||
assert service.fallback_after_failure("cowork", "q", "p", "m", mode="auto") is None
|
||||
|
||||
|
||||
def test_a_malformed_route_result_is_treated_as_no_switch():
|
||||
"""The engine is a legacy module still under refactor; a missing attribute
|
||||
must degrade, not raise into the middle of a turn."""
|
||||
class _Garbage:
|
||||
should_switch = True # claims a switch but exposes no target()
|
||||
|
||||
service = RoutingApplicationService(_FakeRouter(_Garbage()))
|
||||
|
||||
decision = service.route_turn("cowork", "q", "openai_compat", "m", mode="auto")
|
||||
|
||||
assert decision.switched is False
|
||||
assert decision.target() == ("openai_compat", "m")
|
||||
|
||||
|
||||
# --------------------------------------------------------------------------- #
|
||||
# Per-surface mode lookup
|
||||
# --------------------------------------------------------------------------- #
|
||||
def test_mode_is_read_per_surface_when_not_passed_explicitly():
|
||||
"""Each screen has its own Off/Auto/Manual toggle, and workspaces override
|
||||
it - so the surface, not a global setting, decides."""
|
||||
router = _FakeRouter(_switch_to("anthropic", "claude"))
|
||||
modes = {"cowork": "auto", "ai_edit": "off"}
|
||||
service = RoutingApplicationService(router, mode_reader=modes.get)
|
||||
|
||||
assert service.route_turn("cowork", "q", "p", "m").switched is True
|
||||
assert service.route_turn("ai_edit", "q", "p", "m").switched is False
|
||||
|
||||
|
||||
def test_a_failing_mode_reader_falls_back_to_off():
|
||||
def broken(_surface):
|
||||
raise KeyError("config not loaded yet")
|
||||
|
||||
service = RoutingApplicationService(_FakeRouter(_switch_to("a", "b")),
|
||||
mode_reader=broken)
|
||||
|
||||
assert service.route_turn("cowork", "q", "p", "m").switched is False
|
||||
|
||||
|
||||
def test_required_capabilities_are_passed_through_to_the_engine():
|
||||
"""An image turn must only be routed to a vision-capable model; the filter
|
||||
has to reach the scorer or the constraint is silently dropped."""
|
||||
router = _FakeRouter()
|
||||
service = RoutingApplicationService(router)
|
||||
|
||||
service.route_turn("cowork", "describe this", "p", "m", mode="auto",
|
||||
required_capabilities=["vision"])
|
||||
|
||||
assert router.calls[0]["required_capabilities"] == ["vision"]
|
||||
@@ -0,0 +1,112 @@
|
||||
"""Integration test for the AppContext routing wiring (R03-T04 / R03-T05).
|
||||
|
||||
The three chat surfaces now call ``ctx.routing_application()`` instead of each
|
||||
carrying their own copy of the routing algorithm. The unit tests cover the
|
||||
policy; this file covers the WIRING, which unit tests with a fake router cannot
|
||||
see:
|
||||
|
||||
* the service is built and memoised on the context
|
||||
* it reads the per-workspace mode through ``project_routing_mode``
|
||||
* the legacy ``core.routing.RoutingService`` is what sits underneath it
|
||||
* ``fallback`` survives a round trip through the per-workspace mode store
|
||||
|
||||
Still Qt-free: ``AppContext`` itself imports no widgets, and the config is
|
||||
written into a tmp dir so nothing touches ``~/.cowork_local``.
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
from pathlib import Path
|
||||
|
||||
import pytest
|
||||
|
||||
from cowork_local.application.model_routing import (
|
||||
RoutingApplicationService,
|
||||
RoutingMode,
|
||||
)
|
||||
from cowork_local.config import AppConfig
|
||||
from cowork_local.state import AppContext
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def ctx(tmp_path: Path) -> AppContext:
|
||||
"""An AppContext backed by a throwaway config file."""
|
||||
return AppContext(AppConfig.load(tmp_path / "config.json"))
|
||||
|
||||
|
||||
def test_routing_application_is_built_and_memoised(ctx):
|
||||
"""One instance per app: the pending-switch registry underneath it must be
|
||||
shared by every surface, so a second call has to return the same object."""
|
||||
first = ctx.routing_application()
|
||||
|
||||
assert isinstance(first, RoutingApplicationService)
|
||||
assert ctx.routing_application() is first
|
||||
|
||||
|
||||
def test_the_legacy_engine_sits_underneath_the_new_service():
|
||||
"""Strangler-fig check (ADR-001 section 4): the scoring engine is reused, not
|
||||
reimplemented. If this ever stops holding, the assessment scores the
|
||||
scheduler probes would no longer be the ones routing decisions use."""
|
||||
from cowork_local.core.routing.service import RoutingService
|
||||
|
||||
config = AppConfig.load(Path("does-not-exist.json"))
|
||||
context = AppContext(config)
|
||||
|
||||
service = context.routing_application()
|
||||
|
||||
assert isinstance(service._router, RoutingService)
|
||||
assert service._router is context.routing()
|
||||
|
||||
|
||||
def test_mode_is_read_through_the_per_workspace_lookup(ctx, monkeypatch):
|
||||
seen = []
|
||||
|
||||
def fake_mode(surface: str) -> str:
|
||||
seen.append(surface)
|
||||
return "off"
|
||||
|
||||
monkeypatch.setattr(ctx, "project_routing_mode", fake_mode)
|
||||
# Built after the patch so the service captures the patched reader.
|
||||
service = RoutingApplicationService(ctx.routing(), mode_reader=ctx.project_routing_mode)
|
||||
|
||||
decision = service.route_turn("co4e", "hello", "openai_compat", "gpt-4o-mini")
|
||||
|
||||
assert seen == ["co4e"]
|
||||
assert decision.switched is False
|
||||
|
||||
|
||||
def test_routing_off_by_default_leaves_the_selected_model_alone(ctx):
|
||||
"""Default config has routing off on every surface, so a fresh install must
|
||||
never move a turn to another model."""
|
||||
decision = ctx.routing_application().route_turn(
|
||||
"cowork", "write a function", "openai_compat", "gpt-4o-mini")
|
||||
|
||||
assert decision.mode is RoutingMode.OFF
|
||||
assert decision.switched is False
|
||||
assert decision.target() == ("openai_compat", "gpt-4o-mini")
|
||||
|
||||
|
||||
@pytest.mark.parametrize("mode", ["off", "auto", "manual", "fallback"])
|
||||
def test_every_mode_survives_a_round_trip_through_the_config(ctx, mode):
|
||||
"""``fallback`` is new (R03-T03); the per-surface store used to whitelist
|
||||
only three values and would have silently downgraded it to "off"."""
|
||||
ctx.set_project_routing_mode("cowork", mode)
|
||||
|
||||
assert ctx.project_routing_mode("cowork") == mode
|
||||
|
||||
|
||||
def test_an_unknown_mode_still_falls_back_to_off(ctx):
|
||||
ctx.set_project_routing_mode("cowork", "turbo")
|
||||
|
||||
assert ctx.project_routing_mode("cowork") == "off"
|
||||
|
||||
|
||||
def test_a_real_route_call_never_raises_without_any_assessments(ctx):
|
||||
"""The store is empty on a fresh install. Routing must degrade to "keep the
|
||||
current model" rather than raise into the middle of the first message."""
|
||||
ctx.set_project_routing_mode("cowork", "auto")
|
||||
|
||||
decision = ctx.routing_application().route_turn(
|
||||
"cowork", "hello there", "openai_compat", "gpt-4o-mini")
|
||||
|
||||
assert decision.switched is False
|
||||
assert decision.target() == ("openai_compat", "gpt-4o-mini")
|
||||
@@ -0,0 +1,249 @@
|
||||
"""Unit tests for :mod:`infrastructure.telemetry.usage_sink` (R03-T06).
|
||||
|
||||
Two things are being protected here:
|
||||
|
||||
1. The **numbers do not change**. Extracting usage recording out of the two
|
||||
providers is only safe if the events built from each wire format carry
|
||||
exactly what ``core.usage_tracker.record`` used to receive - a silent change
|
||||
would corrupt the Dashboard's cost history.
|
||||
2. The **sink can never break a turn**. Telemetry is observability; a broken
|
||||
store must be swallowed (and logged), never raised into a chat turn.
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
from typing import Any, Dict, List
|
||||
|
||||
import pytest
|
||||
|
||||
from cowork_local.infrastructure.telemetry import usage_sink as telemetry
|
||||
from cowork_local.providers.anthropic import AnthropicProvider
|
||||
from cowork_local.providers.base import Provider
|
||||
from cowork_local.providers.openai_compat import OpenAICompatProvider
|
||||
|
||||
|
||||
# --------------------------------------------------------------------------- #
|
||||
# Event construction - one per wire format
|
||||
# --------------------------------------------------------------------------- #
|
||||
def test_openai_usage_block_maps_onto_the_canonical_event():
|
||||
event = telemetry.openai_usage_event("openai_compat", "gpt-4o-mini", {
|
||||
"prompt_tokens": 120,
|
||||
"completion_tokens": 45,
|
||||
"prompt_tokens_details": {"cached_tokens": 100},
|
||||
})
|
||||
|
||||
assert (event.input_tokens, event.output_tokens, event.cached_tokens) == (120, 45, 100)
|
||||
assert event.estimated is False
|
||||
# Cached tokens are a SUBSET of input, so adding them would double-count.
|
||||
assert event.total_tokens == 165
|
||||
|
||||
|
||||
def test_anthropic_usage_accumulator_maps_onto_the_canonical_event():
|
||||
"""Anthropic reports input on message_start and output on message_delta, so
|
||||
providers/anthropic.py accumulates them into in/out/cache keys."""
|
||||
event = telemetry.anthropic_usage_event("anthropic", "claude", {
|
||||
"in": 200, "out": 80, "cache": 150,
|
||||
})
|
||||
|
||||
assert (event.input_tokens, event.output_tokens, event.cached_tokens) == (200, 80, 150)
|
||||
assert event.estimated is False
|
||||
|
||||
|
||||
def test_a_missing_usage_block_produces_an_estimated_event():
|
||||
event = telemetry.estimated_event("ollama", "llama3.1", "x" * 400, "y" * 40)
|
||||
|
||||
assert event.estimated is True
|
||||
assert event.input_tokens == 100 # ~4 characters per token
|
||||
assert event.output_tokens == 10
|
||||
assert event.cached_tokens == 0
|
||||
|
||||
|
||||
def test_estimation_matches_the_legacy_tracker_formula():
|
||||
"""The extraction must not shift a single recorded number, so the estimator
|
||||
is pinned against the one it replaced."""
|
||||
from cowork_local.core import usage_tracker
|
||||
|
||||
for text in ("", "short", "x" * 4001, "unicode - tiếng Việt"):
|
||||
assert telemetry.estimate_tokens(text) == usage_tracker.estimate_tokens(text)
|
||||
|
||||
|
||||
# --------------------------------------------------------------------------- #
|
||||
# Sinks
|
||||
# --------------------------------------------------------------------------- #
|
||||
def test_recording_sink_collects_events_for_assertions():
|
||||
sink = telemetry.RecordingUsageSink()
|
||||
|
||||
sink.record(telemetry.UsageEvent("p", "m", input_tokens=10, output_tokens=5))
|
||||
sink.record(telemetry.UsageEvent("p", "m", input_tokens=1, output_tokens=1))
|
||||
|
||||
assert len(sink.events) == 2
|
||||
assert sink.total_tokens == 17
|
||||
|
||||
|
||||
def test_null_sink_discards_without_error():
|
||||
telemetry.NullUsageSink().record(telemetry.UsageEvent("p", "m"))
|
||||
|
||||
|
||||
def test_tracker_sink_forwards_every_field_positionally():
|
||||
"""``core.usage_tracker.record`` takes positional counts plus an ``estimated``
|
||||
keyword; the adapter has to preserve that exact call shape."""
|
||||
seen: Dict[str, Any] = {}
|
||||
|
||||
class _Tracker:
|
||||
@staticmethod
|
||||
def record(provider, model, input_tokens, output_tokens, cached_tokens,
|
||||
estimated=False):
|
||||
# Fields captured explicitly rather than via locals(), which would
|
||||
# also drag in the closed-over `seen` binding itself.
|
||||
seen.update({"provider": provider, "model": model,
|
||||
"input_tokens": input_tokens, "output_tokens": output_tokens,
|
||||
"cached_tokens": cached_tokens, "estimated": estimated})
|
||||
|
||||
telemetry.UsageTrackerSink(tracker=_Tracker()).record(
|
||||
telemetry.UsageEvent("anthropic", "claude", 7, 3, 2, estimated=True))
|
||||
|
||||
assert seen == {"provider": "anthropic", "model": "claude", "input_tokens": 7,
|
||||
"output_tokens": 3, "cached_tokens": 2, "estimated": True}
|
||||
|
||||
|
||||
def test_a_failing_tracker_never_raises_into_the_turn():
|
||||
class _Broken:
|
||||
@staticmethod
|
||||
def record(*_args, **_kwargs):
|
||||
raise OSError("usage store is read-only")
|
||||
|
||||
# Must not raise - the turn that produced this event has already succeeded.
|
||||
telemetry.UsageTrackerSink(tracker=_Broken()).record(telemetry.UsageEvent("p", "m"))
|
||||
|
||||
|
||||
def test_set_default_sink_returns_the_previous_one_for_restoration():
|
||||
replacement = telemetry.RecordingUsageSink()
|
||||
|
||||
previous = telemetry.set_default_sink(replacement)
|
||||
try:
|
||||
assert telemetry.default_sink is replacement
|
||||
finally:
|
||||
telemetry.set_default_sink(previous)
|
||||
assert telemetry.default_sink is previous
|
||||
|
||||
|
||||
# --------------------------------------------------------------------------- #
|
||||
# Provider integration - the seam actually being used
|
||||
# --------------------------------------------------------------------------- #
|
||||
class _StubResponse:
|
||||
"""The few members the provider streaming loop touches."""
|
||||
|
||||
def __init__(self, lines: List[str]) -> None:
|
||||
self._lines = lines
|
||||
self.status_code = 200
|
||||
self.headers: Dict[str, str] = {}
|
||||
self.encoding = "utf-8"
|
||||
self.text = ""
|
||||
|
||||
def iter_lines(self, decode_unicode: bool = False):
|
||||
yield from self._lines
|
||||
|
||||
def close(self) -> None:
|
||||
return None
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def sink(monkeypatch):
|
||||
"""A per-instance recording sink, so nothing touches the real usage store."""
|
||||
return telemetry.RecordingUsageSink()
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def canned(monkeypatch):
|
||||
def _install(lines: List[str]):
|
||||
monkeypatch.setattr(Provider, "_request",
|
||||
lambda self, method, url, **kw: _StubResponse(lines))
|
||||
return _install
|
||||
|
||||
|
||||
def test_openai_provider_reports_server_counts_to_its_sink(canned, sink):
|
||||
canned([
|
||||
'data: ' + json.dumps({"choices": [{"delta": {"content": "hi"}}],
|
||||
"usage": {"prompt_tokens": 11, "completion_tokens": 2,
|
||||
"prompt_tokens_details": {"cached_tokens": 4}}}),
|
||||
"data: [DONE]",
|
||||
])
|
||||
provider = OpenAICompatProvider({"base_url": "https://x.invalid/v1",
|
||||
"api_key": "k", "model": "gpt-4o-mini"})
|
||||
provider.usage_sink = sink
|
||||
|
||||
provider.chat([{"role": "user", "content": "hi"}])
|
||||
|
||||
assert len(sink.events) == 1
|
||||
event = sink.events[0]
|
||||
assert (event.input_tokens, event.output_tokens, event.cached_tokens) == (11, 2, 4)
|
||||
assert event.estimated is False
|
||||
assert event.model == "gpt-4o-mini"
|
||||
|
||||
|
||||
def test_openai_provider_estimates_when_the_gateway_sends_no_usage(canned, sink):
|
||||
canned(['data: ' + json.dumps({"choices": [{"delta": {"content": "hello"}}]}),
|
||||
"data: [DONE]"])
|
||||
provider = OpenAICompatProvider({"base_url": "https://x.invalid/v1",
|
||||
"api_key": "k", "model": "m"})
|
||||
provider.usage_sink = sink
|
||||
|
||||
provider.chat([{"role": "user", "content": "hi"}])
|
||||
|
||||
assert sink.events[0].estimated is True
|
||||
assert sink.events[0].output_tokens >= 1
|
||||
|
||||
|
||||
def test_anthropic_provider_reports_stream_counts_to_its_sink(canned, sink):
|
||||
canned(['data: ' + json.dumps(p) for p in (
|
||||
{"type": "message_start", "message": {"usage": {"input_tokens": 30,
|
||||
"cache_read_input_tokens": 10}}},
|
||||
{"type": "content_block_delta", "index": 0,
|
||||
"delta": {"type": "text_delta", "text": "ok"}},
|
||||
{"type": "message_delta", "usage": {"output_tokens": 5}},
|
||||
{"type": "message_stop"},
|
||||
)])
|
||||
provider = AnthropicProvider({"base_url": "https://x.invalid", "api_key": "k",
|
||||
"model": "claude"})
|
||||
provider.usage_sink = sink
|
||||
|
||||
provider.chat([{"role": "user", "content": "hi"}])
|
||||
|
||||
event = sink.events[0]
|
||||
assert (event.input_tokens, event.output_tokens, event.cached_tokens) == (30, 5, 10)
|
||||
assert event.estimated is False
|
||||
|
||||
|
||||
def test_a_provider_without_an_explicit_sink_uses_the_process_default(canned):
|
||||
"""Existing call sites set no sink, so the default has to keep working -
|
||||
that is what makes this extraction a no-op for production behaviour."""
|
||||
canned(['data: ' + json.dumps({"choices": [{"delta": {"content": "x"}}]}),
|
||||
"data: [DONE]"])
|
||||
recorder = telemetry.RecordingUsageSink()
|
||||
previous = telemetry.set_default_sink(recorder)
|
||||
try:
|
||||
OpenAICompatProvider({"base_url": "https://x.invalid/v1", "api_key": "k",
|
||||
"model": "m"}).chat([{"role": "user", "content": "hi"}])
|
||||
finally:
|
||||
telemetry.set_default_sink(previous)
|
||||
|
||||
assert len(recorder.events) == 1
|
||||
|
||||
|
||||
def test_a_sink_that_raises_does_not_fail_the_turn(canned):
|
||||
"""The answer has already been produced by the time usage is recorded;
|
||||
losing the telemetry is strictly better than losing the answer."""
|
||||
canned(['data: ' + json.dumps({"choices": [{"delta": {"content": "x"}}]}),
|
||||
"data: [DONE]"])
|
||||
|
||||
class _Exploding:
|
||||
def record(self, _event):
|
||||
raise RuntimeError("sink is down")
|
||||
|
||||
provider = OpenAICompatProvider({"base_url": "https://x.invalid/v1",
|
||||
"api_key": "k", "model": "m"})
|
||||
provider.usage_sink = _Exploding()
|
||||
|
||||
result = provider.chat([{"role": "user", "content": "hi"}])
|
||||
|
||||
assert result["content"] == "x"
|
||||
Reference in New Issue
Block a user