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:
@@ -0,0 +1,7 @@
|
||||
"""Infrastructure layer - adapters to the outside world.
|
||||
|
||||
Concrete implementations of what the inner layers only describe: HTTP calls to
|
||||
model gateways, the OS keyring, the filesystem, subprocesses, telemetry sinks.
|
||||
May import ``domain/`` (to speak its types) and third-party libraries, but never
|
||||
``presentation/``/``ui/``.
|
||||
"""
|
||||
@@ -0,0 +1,5 @@
|
||||
"""Provider adapters and the central provider catalogue (EPIC R03)."""
|
||||
|
||||
from .provider_registry import ProviderRegistry, default_registry
|
||||
|
||||
__all__ = ["ProviderRegistry", "default_registry"]
|
||||
@@ -0,0 +1,207 @@
|
||||
"""ProviderRegistry - the one place a provider is declared (R03-T02).
|
||||
|
||||
Replaces the three-way split between ``providers/factory.py::_REGISTRY``,
|
||||
``config.py::DEFAULT_CONFIG["providers"]`` and ``config.py::PROVIDER_LABELS``
|
||||
with a single catalogue of :class:`ProviderDescriptor` objects plus the
|
||||
implementation class each one maps to.
|
||||
|
||||
Adding a provider is now one entry in :data:`BUILT_IN_PROVIDERS` (declarative
|
||||
facts) and one line in :data:`_IMPLEMENTATIONS` (which class speaks that
|
||||
protocol) - see ``docs/governance/contributor-recipes.md`` (R10-T04).
|
||||
|
||||
Migration note (strangler fig, ADR-001 section 4): this registry does not
|
||||
re-implement any provider. It builds the SAME classes ``providers/factory.py``
|
||||
builds, so both entry points stay behaviourally identical while call sites move
|
||||
over one at a time.
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import Any, Dict, Iterable, List, Mapping, Optional
|
||||
|
||||
from cowork_local.domain.models.provider_descriptor import (
|
||||
ProviderCapability,
|
||||
ProviderDescriptor,
|
||||
)
|
||||
from cowork_local.providers.base import Provider, ProviderError
|
||||
|
||||
_CAP = ProviderCapability
|
||||
|
||||
# Every provider the app ships with, described once.
|
||||
#
|
||||
# The capability sets are deliberately conservative: a capability listed here is
|
||||
# one the adapter genuinely implements today. Claiming VISION for a provider
|
||||
# whose chat() cannot translate an image block would route an image turn into a
|
||||
# guaranteed failure, so an unimplemented capability must stay off the list.
|
||||
BUILT_IN_PROVIDERS: tuple = (
|
||||
ProviderDescriptor(
|
||||
id="openai_compat",
|
||||
label="OpenAI-compatible (Internal Gateway)",
|
||||
protocol="openai_compat",
|
||||
default_model="gpt-4o-mini",
|
||||
capabilities=frozenset({_CAP.STREAMING, _CAP.TOOLS, _CAP.VISION,
|
||||
_CAP.REASONING, _CAP.MODEL_LISTING}),
|
||||
notes="Any endpoint speaking the OpenAI Chat Completions protocol.",
|
||||
),
|
||||
ProviderDescriptor(
|
||||
id="anthropic",
|
||||
label="Anthropic Claude",
|
||||
protocol="anthropic",
|
||||
default_model="claude-sonnet-4-6",
|
||||
capabilities=frozenset({_CAP.STREAMING, _CAP.TOOLS, _CAP.VISION,
|
||||
_CAP.MODEL_LISTING}),
|
||||
),
|
||||
ProviderDescriptor(
|
||||
id="ollama",
|
||||
label="Ollama (local models)",
|
||||
protocol="openai_compat",
|
||||
default_model="llama3.1",
|
||||
capabilities=frozenset({_CAP.STREAMING, _CAP.TOOLS, _CAP.REASONING,
|
||||
_CAP.MODEL_LISTING}),
|
||||
# Ollama ignores the key, but the OpenAI client layer requires a value,
|
||||
# so the default config ships a placeholder rather than an empty string.
|
||||
requires_api_key=False,
|
||||
local=True,
|
||||
notes="Runs on this machine - no data leaves the device, no token cost.",
|
||||
),
|
||||
ProviderDescriptor(
|
||||
id="github_copilot",
|
||||
label="GitHub Copilot",
|
||||
protocol="openai_compat",
|
||||
default_model="gpt-4o",
|
||||
capabilities=frozenset({_CAP.STREAMING, _CAP.TOOLS, _CAP.MODEL_LISTING}),
|
||||
notes="Paste a Copilot token as the API key.",
|
||||
),
|
||||
ProviderDescriptor(
|
||||
id="codex",
|
||||
label="OpenAI (Codex / GPT)",
|
||||
protocol="openai_compat",
|
||||
default_model="gpt-4o-mini",
|
||||
capabilities=frozenset({_CAP.STREAMING, _CAP.TOOLS, _CAP.VISION,
|
||||
_CAP.REASONING, _CAP.MODEL_LISTING}),
|
||||
),
|
||||
)
|
||||
|
||||
|
||||
def _implementations() -> Dict[str, type]:
|
||||
"""Protocol -> adapter class.
|
||||
|
||||
Imported lazily inside the function because ``providers/anthropic.py`` and
|
||||
``providers/openai_compat.py`` pull in ``requests`` at import time; keeping
|
||||
that out of module import means a test that only inspects descriptors pays
|
||||
no import cost at all.
|
||||
"""
|
||||
from cowork_local.providers.anthropic import AnthropicProvider
|
||||
from cowork_local.providers.openai_compat import OpenAICompatProvider
|
||||
|
||||
return {
|
||||
"openai_compat": OpenAICompatProvider,
|
||||
"anthropic": AnthropicProvider,
|
||||
}
|
||||
|
||||
|
||||
class ProviderRegistry:
|
||||
"""Catalogue of known providers + the factory that instantiates them.
|
||||
|
||||
Intentionally holds no config and no app context: it is a pure lookup table
|
||||
plus a build step, so it can be constructed in a test with a custom
|
||||
descriptor list and no application running.
|
||||
"""
|
||||
|
||||
def __init__(self, descriptors: Optional[Iterable[ProviderDescriptor]] = None) -> None:
|
||||
# Dict preserves declaration order (Python 3.7+), which is the order
|
||||
# Settings lists providers in - so the catalogue order is data, not luck.
|
||||
self._by_id: Dict[str, ProviderDescriptor] = {
|
||||
d.id: d for d in (descriptors if descriptors is not None else BUILT_IN_PROVIDERS)
|
||||
}
|
||||
|
||||
# -- catalogue queries ------------------------------------------------ #
|
||||
def ids(self) -> List[str]:
|
||||
"""Known provider ids, in declaration order."""
|
||||
return list(self._by_id)
|
||||
|
||||
def all(self) -> List[ProviderDescriptor]:
|
||||
"""Every descriptor, in declaration order."""
|
||||
return list(self._by_id.values())
|
||||
|
||||
def get(self, provider_id: str) -> Optional[ProviderDescriptor]:
|
||||
"""The descriptor for ``provider_id``, or None when unknown.
|
||||
|
||||
Returns None rather than raising because the caller is often reacting to
|
||||
a config file that may name a provider from a newer version; the UI
|
||||
should be able to skip it, not crash.
|
||||
"""
|
||||
return self._by_id.get(provider_id)
|
||||
|
||||
def require(self, provider_id: str) -> ProviderDescriptor:
|
||||
"""Like :meth:`get` but raises :class:`ProviderError` when unknown.
|
||||
|
||||
Same error type ``providers/factory.py::build_provider`` already raises,
|
||||
so callers that migrate to the registry keep their existing except clause.
|
||||
"""
|
||||
descriptor = self._by_id.get(provider_id)
|
||||
if descriptor is None:
|
||||
known = ", ".join(self._by_id) or "(none)"
|
||||
raise ProviderError(f"Unsupported provider: {provider_id} (known: {known})")
|
||||
return descriptor
|
||||
|
||||
def labels(self) -> Dict[str, str]:
|
||||
"""``{id: label}`` - the drop-in replacement for ``config.PROVIDER_LABELS``."""
|
||||
return {d.id: d.label for d in self._by_id.values()}
|
||||
|
||||
def supporting(self, capability: ProviderCapability) -> List[ProviderDescriptor]:
|
||||
"""Every descriptor advertising ``capability`` - used to answer "which
|
||||
providers could serve this turn?" before any of them is built."""
|
||||
return [d for d in self._by_id.values() if d.supports(capability)]
|
||||
|
||||
def configured(self, providers_conf: Mapping[str, Mapping[str, Any]]
|
||||
) -> List[ProviderDescriptor]:
|
||||
"""Descriptors whose config section is complete enough to actually call.
|
||||
|
||||
``providers_conf`` is ``AppConfig.data["providers"]``. Passing the raw
|
||||
mapping (not the AppConfig object) keeps this layer independent of the
|
||||
config implementation, which EPIC R02 is rewriting in parallel.
|
||||
"""
|
||||
return [d for d in self._by_id.values()
|
||||
if d.is_configured(providers_conf.get(d.id, {}) or {})]
|
||||
|
||||
# -- construction ----------------------------------------------------- #
|
||||
def build(self, provider_id: str, conf: Mapping[str, Any],
|
||||
model: str = "") -> Provider:
|
||||
"""Instantiate the adapter for ``provider_id``.
|
||||
|
||||
``model`` overrides the configured model for this instance only - that is
|
||||
how the routing layer runs one turn on a different model without mutating
|
||||
the user's saved settings.
|
||||
"""
|
||||
descriptor = self.require(provider_id)
|
||||
impl = _implementations().get(descriptor.protocol)
|
||||
if impl is None: # pragma: no cover - only reachable via a bad descriptor
|
||||
raise ProviderError(
|
||||
f"Provider '{provider_id}' declares unknown protocol "
|
||||
f"'{descriptor.protocol}'."
|
||||
)
|
||||
# Copy before mutating: conf is the caller's live config dict, and
|
||||
# writing the routed model into it would silently change the user's
|
||||
# saved default for every later turn.
|
||||
resolved = dict(conf or {})
|
||||
resolved["model"] = descriptor.resolve_model(conf, model)
|
||||
instance = impl(resolved)
|
||||
# The adapter class is shared by several ids (three of them are
|
||||
# OpenAI-compatible), so its class-level `name` cannot identify which
|
||||
# provider this is. Stamping the instance keeps usage records, audit
|
||||
# entries and routing candidate keys attributed to the right provider.
|
||||
instance.name = descriptor.id
|
||||
return instance
|
||||
|
||||
def describe(self, provider_id: str, conf: Optional[Mapping[str, Any]] = None) -> str:
|
||||
"""One-line description used in logs and error messages."""
|
||||
return self.require(provider_id).describe(conf)
|
||||
|
||||
|
||||
# Shared default instance. Callers that need the built-in catalogue use this
|
||||
# instead of constructing a registry each time; tests build their own with an
|
||||
# explicit descriptor list.
|
||||
default_registry = ProviderRegistry()
|
||||
|
||||
|
||||
__all__ = ["ProviderRegistry", "BUILT_IN_PROVIDERS", "default_registry"]
|
||||
@@ -0,0 +1,21 @@
|
||||
"""Telemetry sinks: where token usage and turn metrics are recorded (EPIC R03)."""
|
||||
|
||||
from .usage_sink import (
|
||||
NullUsageSink,
|
||||
RecordingUsageSink,
|
||||
UsageEvent,
|
||||
UsageEventSink,
|
||||
UsageTrackerSink,
|
||||
default_sink,
|
||||
set_default_sink,
|
||||
)
|
||||
|
||||
__all__ = [
|
||||
"UsageEvent",
|
||||
"UsageEventSink",
|
||||
"UsageTrackerSink",
|
||||
"NullUsageSink",
|
||||
"RecordingUsageSink",
|
||||
"default_sink",
|
||||
"set_default_sink",
|
||||
]
|
||||
@@ -0,0 +1,229 @@
|
||||
"""UsageEventSink - where a turn's token usage goes (R03-T06).
|
||||
|
||||
Today each provider records its own usage inline, in the middle of the streaming
|
||||
loop::
|
||||
|
||||
# providers/openai_compat.py
|
||||
def _record_usage(self, messages, text_parts, tool_acc, usage_seen):
|
||||
from ..core import usage_tracker as ut
|
||||
...
|
||||
ut.record(self.name, self.model, ...)
|
||||
|
||||
Three problems with that shape:
|
||||
|
||||
1. **Hidden side effect.** ``chat()`` looks like a pure request/response call but
|
||||
also writes to the Dashboard's store, so a test of a provider silently
|
||||
appends rows to the developer's real usage history.
|
||||
2. **Duplicated estimation.** The "no usage block from the server, so estimate
|
||||
at ~4 chars/token" fallback is copy-pasted per provider and can drift.
|
||||
3. **One hard-wired destination.** Usage can only ever go to
|
||||
``core.usage_tracker``; a run that wants to bill a workflow, or a test that
|
||||
wants to assert on token counts, has nowhere to plug in.
|
||||
|
||||
This module introduces the seam: providers build a :class:`UsageEvent` and hand
|
||||
it to a :class:`UsageEventSink`. Production wires :class:`UsageTrackerSink`
|
||||
(same destination, same numbers as before); tests wire
|
||||
:class:`RecordingUsageSink` or :class:`NullUsageSink`.
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
import logging
|
||||
from dataclasses import dataclass
|
||||
from typing import Any, Dict, List, Optional, Protocol, Sequence
|
||||
|
||||
logger = logging.getLogger("cowork_local.telemetry")
|
||||
|
||||
# Rough characters-per-token ratio used when the gateway sends no usage block.
|
||||
# Matches the constant behaviour of ``core.usage_tracker.estimate_tokens`` so
|
||||
# moving the estimation here does not change a single recorded number.
|
||||
_CHARS_PER_TOKEN = 4
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class UsageEvent:
|
||||
"""Token usage for exactly one provider round trip.
|
||||
|
||||
``estimated`` marks a record derived from text length rather than reported by
|
||||
the server. The Dashboard shows the two differently, and conflating them
|
||||
would make cost figures look more precise than they are.
|
||||
"""
|
||||
|
||||
provider: str
|
||||
model: str
|
||||
input_tokens: int = 0
|
||||
output_tokens: int = 0
|
||||
cached_tokens: int = 0
|
||||
estimated: bool = False
|
||||
|
||||
@property
|
||||
def total_tokens(self) -> int:
|
||||
"""Input + output. Cached tokens are a subset of input, not an addition,
|
||||
so adding them here would double-count a cache hit."""
|
||||
return self.input_tokens + self.output_tokens
|
||||
|
||||
def to_dict(self) -> Dict[str, Any]:
|
||||
"""JSON-safe projection for logs and for sinks that persist raw events."""
|
||||
return {
|
||||
"provider": self.provider,
|
||||
"model": self.model,
|
||||
"input_tokens": self.input_tokens,
|
||||
"output_tokens": self.output_tokens,
|
||||
"cached_tokens": self.cached_tokens,
|
||||
"estimated": self.estimated,
|
||||
}
|
||||
|
||||
|
||||
class UsageEventSink(Protocol):
|
||||
"""Anything that can absorb a :class:`UsageEvent`.
|
||||
|
||||
Implementations MUST NOT raise: telemetry is observability, and a failure to
|
||||
record usage must never abort the turn that produced it.
|
||||
"""
|
||||
|
||||
def record(self, event: UsageEvent) -> None:
|
||||
"""Absorb one usage event."""
|
||||
|
||||
|
||||
class NullUsageSink:
|
||||
"""Discards everything. The default for tests and headless tooling, so a
|
||||
unit test never writes into the developer's real usage history."""
|
||||
|
||||
def record(self, event: UsageEvent) -> None: # noqa: D102 - see protocol
|
||||
return None
|
||||
|
||||
|
||||
class RecordingUsageSink:
|
||||
"""Keeps events in memory so a test can assert on what was recorded."""
|
||||
|
||||
def __init__(self) -> None:
|
||||
self.events: List[UsageEvent] = []
|
||||
|
||||
def record(self, event: UsageEvent) -> None: # noqa: D102 - see protocol
|
||||
self.events.append(event)
|
||||
|
||||
@property
|
||||
def total_tokens(self) -> int:
|
||||
"""Sum across every recorded event."""
|
||||
return sum(e.total_tokens for e in self.events)
|
||||
|
||||
|
||||
class UsageTrackerSink:
|
||||
"""Forwards to ``core.usage_tracker`` - the Dashboard's store.
|
||||
|
||||
This is the production sink and the only place that still knows about the
|
||||
legacy tracker module, which is what lets EPIC R10 replace the storage
|
||||
without touching a single provider.
|
||||
"""
|
||||
|
||||
def __init__(self, tracker: Optional[Any] = None) -> None:
|
||||
# Injectable for tests; imported lazily otherwise because the tracker
|
||||
# touches the config directory at import time.
|
||||
self._tracker = tracker
|
||||
|
||||
def _resolve(self) -> Any:
|
||||
if self._tracker is None:
|
||||
from cowork_local.core import usage_tracker
|
||||
|
||||
self._tracker = usage_tracker
|
||||
return self._tracker
|
||||
|
||||
def record(self, event: UsageEvent) -> None:
|
||||
"""Write the event to the usage tracker, swallowing any failure.
|
||||
|
||||
The bare except mirrors the behaviour this replaces (each provider
|
||||
already wrapped its ``ut.record`` call in ``try/except: pass``) but logs
|
||||
at debug level instead of discarding the reason entirely, so a broken
|
||||
Dashboard store can at least be diagnosed.
|
||||
"""
|
||||
try:
|
||||
self._resolve().record(
|
||||
event.provider, event.model,
|
||||
event.input_tokens, event.output_tokens, event.cached_tokens,
|
||||
estimated=event.estimated,
|
||||
)
|
||||
except Exception: # noqa: BLE001 - telemetry must never break a turn
|
||||
logger.debug("usage sink: failed to record %s", event.to_dict(), exc_info=True)
|
||||
|
||||
|
||||
def estimate_tokens(text: str) -> int:
|
||||
"""Approximate token count for ``text`` (~4 characters per token).
|
||||
|
||||
Deliberately identical to ``core.usage_tracker.estimate_tokens`` so that
|
||||
moving estimation into this layer changes no recorded number. Duplicated
|
||||
rather than imported to keep this module free of the legacy dependency;
|
||||
:class:`UsageTrackerSink` is the only bridge back to it.
|
||||
"""
|
||||
return max(0, len(text or "") // _CHARS_PER_TOKEN)
|
||||
|
||||
|
||||
def estimated_event(provider: str, model: str, sent: str, received: str) -> UsageEvent:
|
||||
"""Build an estimated :class:`UsageEvent` from the raw text of a round trip.
|
||||
|
||||
Used when the gateway sends no usage block - most self-hosted OpenAI-compatible
|
||||
servers and Ollama do not.
|
||||
"""
|
||||
return UsageEvent(
|
||||
provider=provider, model=model,
|
||||
input_tokens=estimate_tokens(sent),
|
||||
output_tokens=estimate_tokens(received),
|
||||
cached_tokens=0,
|
||||
estimated=True,
|
||||
)
|
||||
|
||||
|
||||
def openai_usage_event(provider: str, model: str, usage: Dict[str, Any]) -> UsageEvent:
|
||||
"""Build a reported :class:`UsageEvent` from an OpenAI-style usage block."""
|
||||
details = usage.get("prompt_tokens_details") or {}
|
||||
return UsageEvent(
|
||||
provider=provider, model=model,
|
||||
input_tokens=int(usage.get("prompt_tokens", 0) or 0),
|
||||
output_tokens=int(usage.get("completion_tokens", 0) or 0),
|
||||
cached_tokens=int(details.get("cached_tokens", 0) or 0),
|
||||
estimated=False,
|
||||
)
|
||||
|
||||
|
||||
def anthropic_usage_event(provider: str, model: str, usage: Dict[str, Any]) -> UsageEvent:
|
||||
"""Build a reported :class:`UsageEvent` from Anthropic's usage accumulator.
|
||||
|
||||
Anthropic reports input tokens on ``message_start`` and output tokens on
|
||||
``message_delta``, so ``providers/anthropic.py`` accumulates them into a dict
|
||||
keyed ``in``/``out``/``cache`` - this reads that shape.
|
||||
"""
|
||||
return UsageEvent(
|
||||
provider=provider, model=model,
|
||||
input_tokens=int(usage.get("in", 0) or 0),
|
||||
output_tokens=int(usage.get("out", 0) or 0),
|
||||
cached_tokens=int(usage.get("cache", 0) or 0),
|
||||
estimated=False,
|
||||
)
|
||||
|
||||
|
||||
# The sink providers use unless one is injected. A module-level default keeps
|
||||
# the change to the provider classes to a single attribute, and lets a test swap
|
||||
# the destination process-wide with one monkeypatch.
|
||||
default_sink: UsageEventSink = UsageTrackerSink()
|
||||
|
||||
|
||||
def set_default_sink(sink: UsageEventSink) -> UsageEventSink:
|
||||
"""Replace the process-wide default sink; returns the previous one so a
|
||||
caller (or fixture) can restore it."""
|
||||
global default_sink
|
||||
previous = default_sink
|
||||
default_sink = sink
|
||||
return previous
|
||||
|
||||
|
||||
__all__ = [
|
||||
"UsageEvent",
|
||||
"UsageEventSink",
|
||||
"UsageTrackerSink",
|
||||
"NullUsageSink",
|
||||
"RecordingUsageSink",
|
||||
"estimate_tokens",
|
||||
"estimated_event",
|
||||
"openai_usage_event",
|
||||
"anthropic_usage_event",
|
||||
"default_sink",
|
||||
"set_default_sink",
|
||||
]
|
||||
Reference in New Issue
Block a user