chore(repo): initialize Cowork Local Gitea repository
CI / test (push) Canceled after 0s

This commit is contained in:
thanhnv
2026-08-09 20:12:05 +07:00
commit 414eaddca3
192 changed files with 48160 additions and 0 deletions
+125
View File
@@ -0,0 +1,125 @@
# Auto Model Assessment & Routing
Tự động **đánh giá** từng model (provider/model do user cấu hình), **chấm điểm phù hợp**
cho mỗi loại task, rồi **định tuyến** mỗi lượt chat/agent tới model phù hợp nhất — theo
3 chế độ **Off / Auto / Manual** bật ngay trên màn hình chat (Cowork, Co4E, AI-Edit).
Module này được xây dựng để **hoà vào đúng stack sẵn có** của Cowork-Local (PySide6
desktop app), thay vì dựng một service FastAPI riêng:
| Bản mô tả gốc (đề bài) | Hiện thực trong app này |
|---|---|
| Config YAML | Config JSON `~/.cowork_local/config.json` (chuẩn của app) + file assessment riêng |
| FastAPI REST endpoints | `RoutingService` (Python facade) — các method ánh xạ 1-1 với endpoint |
| `httpx` async + `asyncio.Semaphore` | Tái dùng `providers/` (requests) + `ThreadPoolExecutor` với **semaphore theo từng provider** |
| APScheduler | `QTimer` (giống `core/task_scheduler.py`) — không thêm dependency |
| `clients.py` (Anthropic/OpenAI) | `AppProbeClient` bọc `AppContext.build_provider_for` (đã có sẵn TLS-trust, retry 429, gateway) |
## Kiến trúc
```
core/routing/
models.py # Pydantic v2: ModelMetadata, ProbeResult, ModelAssessment,
# SwitchDecision, PendingSwitch, TaskType/Policy/SwitchMode
store.py # AssessmentStore: JSON, atomic write (temp+rename), history backup
metadata.py # STATIC_METADATA + enrich() (dùng lại core/model_pricing cho giá)
clients.py # AppProbeClient (bọc Provider có sẵn) + ProbeClient protocol
prober.py # BENCHMARK_TASKS, probe_model(), make_judge(), semaphore/provider
scorer.py # compute_fit_score() + POLICY_WEIGHTS
selector.py # rank_models()/best_model() — tính lại fit theo policy, không probe lại
classifier.py # classify(prompt) -> TaskType (heuristic, fallback LLM tuỳ chọn)
switch_controller.py # decide() (thuần) + PendingSwitchRegistry (TTL, idempotent)
orchestrator.py # check_and_update(): enrich -> probe -> score -> store
service.py # RoutingService — facade UI gọi
scheduler.py # RoutingScheduler (QTimer): reassess định kỳ + dọn pending hết hạn
```
## Công thức fit score
```
fit = w_quality * quality
+ w_cost * 1/(1 + cost)
+ w_latency * 1/(1 + latency_s)
```
`POLICY_WEIGHTS` (mỗi hàng cộng = 1.0):
| Policy | quality | cost | latency |
|---|---|---|---|
| `quality` | 0.80 | 0.10 | 0.10 |
| `cost` | 0.20 | 0.70 | 0.10 |
| `latency` | 0.20 | 0.10 | 0.70 |
| `balanced` | 0.50 | 0.25 | 0.25 |
- Probe **fail** → fit = 0 (model không dùng được thì không bao giờ được chọn).
- Giá **không rõ** → để `None`, đánh dấu `metadata_incomplete=True` (không đoán bừa).
## Config (trong `config.json`, mục `routing`)
```jsonc
"routing": {
"switch_mode": "off", // mặc định toàn cục: "off" | "auto" | "manual"
"policy": "balanced", // "quality" | "cost" | "latency" | "balanced"
"min_score_gain": 0.05, // chỉ chuyển nếu model mới hơn model hiện tại ≥ ngưỡng này
"confirm_timeout_sec": 60, // (manual) hết giờ chờ confirm → giữ model hiện tại
"reassess_interval_hours": 24, // lịch reassess; 0 = tắt
"per_provider_concurrency": 2, // số probe song song tối đa mỗi provider (chống rate limit)
"judge_provider": "", // provider của judge ("" → active provider)
"judge_model": "", // model chấm điểm cố định ("" → default rẻ theo provider)
"candidates": [ // model muốn đánh giá; rỗng → tự lấy model đang cấu hình
{"provider": "anthropic", "model_id": "claude-opus-4-8", "tier": "powerful"},
{"provider": "anthropic", "model_id": "claude-haiku-4-5", "tier": "fast"}
],
"auto_reassess_on_add": true, // thêm model mới → reassess ngay
"surface_modes": { // toggle Off/Auto/Manual của TỪNG màn hình ("" = theo switch_mode)
"cowork": "", "co4e": "", "ai_edit": ""
}
}
```
Kết quả assessment **KHÔNG** nằm trong `config.json` mà ở file riêng:
`~/.cowork_local/assessments.json` (+ backup lịch sử ở `assessments_history/<timestamp>.json`).
## Toggle Off / Auto / Manual (trên màn hình chat)
Mỗi màn hình chat có một toggle nhỏ cạnh ô chọn model:
- **Off** — tắt định tuyến, luôn dùng model đang chọn.
- **Auto** — tự động chuyển sang model phù hợp nhất (nếu `gain ≥ min_score_gain`),
chạy luôn, hiện dòng thông báo `↪ Auto-routed to …`.
- **Manual** — hiện hộp thoại xác nhận (có đếm ngược `confirm_timeout_sec`); user
đồng ý mới chuyển, từ chối / hết giờ thì giữ model hiện tại.
Toggle được lưu **riêng cho từng màn hình** (`surface_modes`) và ghi đè `switch_mode` toàn cục.
## Logical API (RoutingService)
| Method | Tương đương REST trong đề bài |
|---|---|
| `reassess(policy=None)` / `reassess_background()` | `POST /models/reassess` |
| `best_for(task_type, policy)` | `GET /models/best` |
| `status()` / `assessments()` | `GET /models/assessments` |
| `add_candidate(provider, model_id, tier)` | `POST /models/add` (tự trigger reassess) |
| `route(surface, prompt, provider, model)` | phần quyết định của `POST /task/execute` |
| `create_pending()` / `resolve_pending(id, approve, run)` | `POST /task/confirm-switch` (idempotent) |
| `get_routing_config()` / `update_routing_config(**)` | `GET`/`PATCH /routing/config` |
## Bảo mật & chi phí
- **API key** đọc từ env (qua `api_key_env` của provider) — không ghi key vào config/log.
- **Probe tốn tiền** → chỉ chạy theo lịch / khi thêm model / khi bấm "Reassess now".
Mỗi lần reassess ghi log số lượng API call.
- **Idempotent** — reassess ổn định (chỉ latency dao động ~µs, dưới xa `min_score_gain`);
ghi atomic nên ngắt giữa chừng không hỏng config.
- **Không gọi API thật trong test** — `clients.py`/`judge()` được mock hoàn toàn.
## Chạy test
```bash
# từ thư mục cha của package (…/cowork_local_20260722)
python -m pytest cowork_local/tests/routing/ -q
```
Bao phủ: `scorer`, `store` (atomic + history), `selector`, `switch_controller`
(Auto/Manual/Off, timeout, idempotent), `orchestrator` (mock client), `classifier`,
và `service` (end-to-end reassess → route → confirm).
+52
View File
@@ -0,0 +1,52 @@
"""Auto Model Assessment & Routing.
Reads the configured providers/models, assesses each model (static metadata +
dynamic probes judged by a fixed cheap judge model), scores them per task type
under a policy, and routes each chat/agent turn to the best-fit model — either
silently (Auto), after user confirmation (Manual), or not at all (Off).
Public entry point for the app is :class:`service.RoutingService`, wired into
``AppContext`` and driven from the Off/Auto/Manual toggle on each chat screen.
Sub-modules
-----------
* ``models`` — Pydantic data models shared by everything here.
* ``scorer`` — fit-score formula + policy weights.
* ``store`` — persist/version assessments (atomic write + history).
* ``metadata`` — static metadata table + enrich() with fallbacks.
* ``clients`` — thin adapter over the app's existing Provider layer.
* ``prober`` — benchmark prompts, probe_model(), judge().
* ``scorer``/``selector`` — score and rank candidates per task type.
* ``switch_controller`` — Auto/Manual/Off switch decisions + pending confirms.
* ``classifier`` — classify a prompt into a TaskType.
* ``orchestrator`` — check_and_update(): the full assess→score→store loop.
* ``service`` — façade the UI talks to.
* ``scheduler`` — periodic + on-model-add reassess triggers.
"""
from .models import (
ModelAssessment,
ModelMetadata,
PendingSwitch,
Policy,
ProbeResult,
SwitchDecision,
SwitchMode,
SwitchStatus,
TaskType,
candidate_key,
split_key,
)
__all__ = [
"TaskType",
"Policy",
"SwitchMode",
"SwitchStatus",
"ModelMetadata",
"ProbeResult",
"ModelAssessment",
"SwitchDecision",
"PendingSwitch",
"candidate_key",
"split_key",
]
+122
View File
@@ -0,0 +1,122 @@
"""Classify a user prompt into a :class:`TaskType`.
Routing runs on every turn, so classification must be cheap — a keyword
heuristic first, with an optional one-shot LLM fallback only when the heuristic
is unsure. The heuristic is intentionally conservative: it defaults to ``QA``
(the safest general bucket) rather than mis-routing an ambiguous prompt.
"""
from __future__ import annotations
import re
from typing import Callable, List, Optional, Tuple
from .models import TaskType
# Signal words per task type. Matched case-insensitively on word boundaries.
# Ordered by specificity when scoring ties (CODING/REASONING beat QA).
_KEYWORDS: dict[TaskType, List[str]] = {
TaskType.CODING: [
"code", "function", "class", "bug", "debug", "refactor", "compile",
"stack trace", "traceback", "python", "javascript", "typescript",
"java", "c++", "golang", "rust", "sql", "regex", "api", "endpoint",
"unit test", "pytest", "npm", "docker", "git", "implement", "algorithm",
"syntax", "exception", "import", "def ", "async", "lập trình", "hàm",
"sửa lỗi", "biên dịch",
],
TaskType.REASONING: [
"why", "prove", "explain why", "reason", "logic", "deduce", "infer",
"step by step", "step-by-step", "solve", "calculate", "how many",
"puzzle", "riddle", "strategy", "trade-off", "tradeoff", "analyze",
"compare and", "chứng minh", "suy luận", "tính toán", "phân tích",
],
TaskType.SUMMARIZATION: [
"summarize", "summary", "tl;dr", "tldr", "condense", "shorten",
"key points", "in short", "brief", "recap", "abstract of", "gist",
"tóm tắt", "rút gọn", "tóm lược",
],
TaskType.CREATIVE: [
"poem", "story", "write a", "creative", "imagine", "fiction", "lyrics",
"song", "haiku", "screenplay", "dialogue", "brainstorm", "slogan",
"tagline", "marketing copy", "viết truyện", "bài thơ", "sáng tạo",
"kịch bản",
],
TaskType.QA: [
"what is", "who is", "when did", "where is", "define", "meaning of",
"how do i", "how to", "is it", "does", "can you tell", "fact",
"là gì", "ai là", "khi nào", "ở đâu", "định nghĩa",
],
}
# Precompiled boundary regexes; ASCII \b doesn't hug Vietnamese diacritics well,
# so multi-word/diacritic phrases fall back to plain substring matching.
_COMPILED: dict[TaskType, List[Tuple[str, Optional[re.Pattern]]]] = {}
for _tt, _words in _KEYWORDS.items():
entries: List[Tuple[str, Optional[re.Pattern]]] = []
for w in _words:
if w.isascii() and " " not in w and w.strip().isalpha():
entries.append((w, re.compile(rf"\b{re.escape(w)}\b", re.IGNORECASE)))
else:
entries.append((w, None)) # substring match
_COMPILED[_tt] = entries
# Tie-break priority when multiple task types score equally.
_PRIORITY = [
TaskType.CODING,
TaskType.REASONING,
TaskType.SUMMARIZATION,
TaskType.CREATIVE,
TaskType.QA,
]
# LLM fallback: given the prompt, return a TaskType value string.
LLMClassifier = Callable[[str], str]
def _heuristic_scores(text: str) -> dict[TaskType, int]:
low = (text or "").lower()
scores: dict[TaskType, int] = {tt: 0 for tt in TaskType}
for tt, entries in _COMPILED.items():
for raw, pat in entries:
if pat is not None:
if pat.search(low):
scores[tt] += 1
elif raw in low:
scores[tt] += 1
return scores
def classify(
text: str,
*,
llm_classifier: Optional[LLMClassifier] = None,
min_confidence: int = 1,
) -> TaskType:
"""Return the most likely :class:`TaskType` for ``text``.
Uses the keyword heuristic first. If nothing scores at least
``min_confidence`` and an ``llm_classifier`` is provided, defers to it once;
otherwise defaults to :attr:`TaskType.QA`.
"""
scores = _heuristic_scores(text)
best_score = max(scores.values()) if scores else 0
if best_score >= min_confidence:
# Highest score, ties broken by _PRIORITY order.
for tt in _PRIORITY:
if scores[tt] == best_score:
return tt
if llm_classifier is not None:
try:
raw = (llm_classifier(text) or "").strip().lower()
return TaskType(raw)
except Exception: # noqa: BLE001 — bad/failed classification → default
pass
return TaskType.QA
__all__ = ["classify", "LLMClassifier", "BENCHMARK_HINT"]
# Small doc alias so callers can show which task types exist.
BENCHMARK_HINT = [tt.value for tt in TaskType]
+88
View File
@@ -0,0 +1,88 @@
"""A thin, unified calling surface over the app's existing Provider layer.
The task asks for a ``clients.py`` abstraction that talks to Anthropic / OpenAI
behind one interface. This app **already has** that — ``providers/`` with
``build_provider`` and a canonical ``chat()`` that streams text and returns the
final assistant message. Rather than duplicate it (and re-solve TLS trust,
429-retry, gateway config…), this module adapts it to the shape the prober
wants: a single blocking ``complete()`` that returns text + token estimate.
Tests inject a fake :class:`ProbeClient` so assessment never hits a real API.
"""
from __future__ import annotations
from dataclasses import dataclass
from typing import Any, Dict, List, Optional, Protocol
@dataclass
class CompletionResult:
"""Outcome of one non-streaming completion used for probing."""
text: str = ""
tokens_out: int = 0
error: Optional[str] = None
@property
def ok(self) -> bool:
return self.error is None
class ProbeClient(Protocol):
"""Minimal interface the prober/judge depend on (so they're mockable)."""
def complete(
self,
provider: str,
model_id: str,
messages: List[Dict[str, Any]],
) -> CompletionResult:
...
def _estimate_tokens(text: str) -> int:
"""Rough output-token count. Uses the app's estimator when importable
(keeps the number consistent with the usage tracker), else ~4 chars/token."""
try:
from ..usage_tracker import estimate_tokens
return int(estimate_tokens(text or ""))
except Exception: # noqa: BLE001
return max(0, len(text or "") // 4)
class AppProbeClient:
"""Real :class:`ProbeClient` backed by :class:`AppContext`.
Builds a fresh provider per call via ``ctx.build_provider_for`` — the same
path interactive chat uses — so the internal gateway, per-host TLS trust and
rate-limit retry all apply to assessment calls too.
"""
def __init__(self, ctx: Any) -> None:
self.ctx = ctx
def complete(
self,
provider: str,
model_id: str,
messages: List[Dict[str, Any]],
) -> CompletionResult:
try:
prov = self.ctx.build_provider_for(provider, model_id or None)
# Non-streaming: no on_text/on_reasoning callbacks. cancel=None.
result = prov.chat(messages, tools=None, on_text=None, cancel=None)
except Exception as exc: # noqa: BLE001 — surfaced as a failed probe
return CompletionResult(error=str(exc))
content = ""
if isinstance(result, dict):
content = result.get("content") or ""
# Strip any inline <think> block a reasoning model may have inlined.
try:
from ...providers.base import Provider
content = Provider.strip_think(content)
except Exception: # noqa: BLE001
pass
return CompletionResult(text=content, tokens_out=_estimate_tokens(content))
__all__ = ["CompletionResult", "ProbeClient", "AppProbeClient"]
+202
View File
@@ -0,0 +1,202 @@
"""Enrich a candidate model with static metadata (price / context / caps).
Order of precedence when filling in a model's facts:
1. **Existing price table** — the app already lets users maintain a per-model
USD price sheet (``core/model_pricing.py``, shown on the Monitoring
Overview). If the model is in there, its real prices win.
2. **Built-in ``STATIC_METADATA``** — a small hard-coded table for well-known
models (context window + capabilities + rough tier), since those rarely
change and the price sheet may not carry them.
3. **Provider ``/models`` discovery** — used only to confirm the model is
actually *available* on the provider right now.
4. **One-shot LLM self-report** — for a genuinely unknown model, an injected
``llm_declarer`` may be called ONCE to have the model describe its own
capabilities; the result is cached by the caller.
Crucially, when a price is genuinely unknown we leave it ``None`` and set
``metadata_incomplete=True`` rather than inventing a number (per the task's
"KHÔNG hardcode giá đoán bừa" rule).
"""
from __future__ import annotations
from typing import Any, Callable, Dict, Iterable, Optional
from .models import ModelMetadata
# Optional hook: given (provider, model_id) return a dict of self-reported
# facts (capabilities/max_context). Injected so tests never hit a real API.
LLMDeclarer = Callable[[str, str], Dict[str, Any]]
# --------------------------------------------------------------------------- #
# Built-in static table for well-known models.
#
# Keyed by a model-id PREFIX (longest match wins), so "claude-opus-4-8" is
# matched by the "claude-opus-4" entry. Prices here are deliberately absent for
# most rows — the user's own price sheet is the source of truth for cost, and a
# wrong hard-coded price is worse than a known-unknown. Context windows and
# capabilities, which are stable, are provided.
# --------------------------------------------------------------------------- #
STATIC_METADATA: Dict[str, Dict[str, Any]] = {
# Anthropic Claude
"claude-opus-4": {
"tier": "powerful", "max_context": 200000,
"capabilities": {"tools", "vision", "reasoning", "long_context"},
},
"claude-sonnet-4": {
"tier": "balanced", "max_context": 200000,
"capabilities": {"tools", "vision", "reasoning", "long_context"},
},
"claude-sonnet-5": {
"tier": "balanced", "max_context": 200000,
"capabilities": {"tools", "vision", "reasoning", "long_context"},
},
"claude-haiku-4": {
"tier": "fast", "max_context": 200000,
"capabilities": {"tools", "vision", "long_context"},
},
"claude-3-5-haiku": {
"tier": "fast", "max_context": 200000,
"capabilities": {"tools", "vision"},
},
# OpenAI / GPT
"gpt-4o-mini": {
"tier": "fast", "max_context": 128000,
"capabilities": {"tools", "vision"},
},
"gpt-4o": {
"tier": "balanced", "max_context": 128000,
"capabilities": {"tools", "vision", "reasoning"},
},
"gpt-4-turbo": {
"tier": "powerful", "max_context": 128000,
"capabilities": {"tools", "vision", "reasoning"},
},
"o1": {
"tier": "powerful", "max_context": 200000,
"capabilities": {"reasoning", "long_context"},
},
"o3": {
"tier": "powerful", "max_context": 200000,
"capabilities": {"reasoning", "long_context", "tools"},
},
# Local / open models (Ollama)
"llama3.1": {
"tier": "fast", "max_context": 128000,
"capabilities": {"tools"},
},
"qwen": {
"tier": "fast", "max_context": 32000,
"capabilities": {"tools", "reasoning"},
},
"deepseek": {
"tier": "balanced", "max_context": 64000,
"capabilities": {"reasoning", "tools"},
},
"gemma": {
"tier": "fast", "max_context": 8192,
"capabilities": set(),
},
}
def _static_for(model_id: str) -> Dict[str, Any]:
"""Longest-prefix lookup in ``STATIC_METADATA`` (empty dict if no match)."""
m = (model_id or "").lower()
best_key = ""
for key in STATIC_METADATA:
if m.startswith(key) and len(key) > len(best_key):
best_key = key
return dict(STATIC_METADATA[best_key]) if best_key else {}
def _cost_from_price_table(model_id: str, config) -> tuple[Optional[float], Optional[float]]:
"""USD cost **per 1k tokens** from the app's price sheet, or ``(None, None)``.
``model_pricing.usd_rates_for`` returns USD per **1M** tokens, so we divide
by 1000. A zero/absent entry is treated as unknown, not as free.
"""
try:
from .. import model_pricing
except Exception: # noqa: BLE001 — module optional in some contexts (tests)
return None, None
if config is None:
return None, None
rates = model_pricing.usd_rates_for(model_id, config)
if not rates:
return None, None
ci = rates.get("in")
co = rates.get("out")
ci = (ci / 1000.0) if ci else None
co = (co / 1000.0) if co else None
return ci, co
def enrich(
provider: str,
model_id: str,
*,
config: Any = None,
tier: Optional[str] = None,
available_models: Optional[Iterable[str]] = None,
llm_declarer: Optional[LLMDeclarer] = None,
) -> ModelMetadata:
"""Build a :class:`ModelMetadata` for one candidate.
Parameters
----------
provider, model_id:
Identify the candidate.
config:
The app config, used to read the user's price sheet (optional).
tier:
User-declared tier from the provider config (e.g. "fast"); overrides
any static-table tier when given.
available_models:
Model ids the provider currently lists. When provided, availability is
set from membership; when ``None`` the model is assumed available (the
prober will discover a truly-dead model via a failed probe anyway).
llm_declarer:
Optional one-shot capability self-report hook for unknown models.
"""
static = _static_for(model_id)
ci, co = _cost_from_price_table(model_id, config)
max_context = static.get("max_context")
capabilities = set(static.get("capabilities") or set())
# Unknown model + a declarer available → ask it once to describe itself.
if not static and llm_declarer is not None:
try:
declared = llm_declarer(provider, model_id) or {}
except Exception: # noqa: BLE001 — a failed self-report must not crash enrichment
declared = {}
if declared.get("max_context"):
max_context = int(declared["max_context"])
for cap in declared.get("capabilities") or []:
capabilities.add(str(cap))
available = True
if available_models is not None:
avail = {str(m) for m in available_models}
available = model_id in avail
# Price genuinely unknown → flag incomplete rather than guessing.
metadata_incomplete = ci is None or co is None
return ModelMetadata(
provider=provider,
model_id=model_id,
tier=tier or static.get("tier"),
cost_per_1k_input=ci,
cost_per_1k_output=co,
max_context=max_context,
capabilities=capabilities,
available=available,
metadata_incomplete=metadata_incomplete,
)
__all__ = ["STATIC_METADATA", "enrich", "LLMDeclarer"]
+214
View File
@@ -0,0 +1,214 @@
"""Pydantic v2 data models for Auto Model Assessment & Routing.
These are the provider-agnostic shapes shared by every routing module — the
enricher, prober, scorer, selector and switch-controller all speak in terms of
these. They serialize cleanly to/from JSON so the assessment store and the
app config (``~/.cowork_local/…``) can round-trip them.
Terminology
-----------
* A **candidate** is a ``(provider, model_id)`` pair the app can call.
* An **assessment** is what we learned about one candidate: its static
metadata, the dynamic probe results per task type, and the derived
``fit_scores`` per task type.
* A **task type** is the kind of work a message represents (qa / coding / …).
* A **policy** is how we weigh quality vs cost vs latency when scoring.
"""
from __future__ import annotations
from enum import Enum
from typing import Any, Dict, List, Optional, Set
from pydantic import BaseModel, Field
# --------------------------------------------------------------------------- #
# Enums
# --------------------------------------------------------------------------- #
class TaskType(str, Enum):
"""The kinds of work a chat/agent turn can represent.
A message is classified into exactly one of these before routing (see
``classifier.py``). ``BENCHMARK_TASKS`` in the prober has one fixed prompt
per value so every candidate model is compared on the same yardstick.
"""
QA = "qa"
CODING = "coding"
REASONING = "reasoning"
SUMMARIZATION = "summarization"
CREATIVE = "creative"
class Policy(str, Enum):
"""How to trade off quality, cost and latency when scoring a model."""
QUALITY = "quality"
COST = "cost"
LATENCY = "latency"
BALANCED = "balanced"
class SwitchMode(str, Enum):
"""Per-surface routing behaviour, driven by the Off/Auto/Manual toggle.
* ``OFF`` — routing disabled; always use the manually-selected model.
* ``AUTO`` — silently switch to the best model when it clears the gain
threshold, then run the task.
* ``MANUAL`` — propose the switch and wait for the user to confirm before
running with the new model.
"""
OFF = "off"
AUTO = "auto"
MANUAL = "manual"
class SwitchStatus(str, Enum):
"""Lifecycle of a :class:`PendingSwitch` awaiting user confirmation."""
PENDING = "pending"
CONFIRMED = "confirmed"
REJECTED = "rejected"
EXPIRED = "expired"
# --------------------------------------------------------------------------- #
# Static metadata + dynamic probe
# --------------------------------------------------------------------------- #
class ModelMetadata(BaseModel):
"""Static, mostly-price/capability facts about one candidate model.
``cost_per_1k_*`` are USD per 1,000 tokens. They are ``None`` — not a
guess — when the price is genuinely unknown; ``metadata_incomplete`` is
then set True so the scorer/UI can flag it rather than silently trusting a
fabricated number (see ``metadata.py``).
"""
provider: str
model_id: str
tier: Optional[str] = None # e.g. "fast" | "powerful" — free-form, user-supplied
cost_per_1k_input: Optional[float] = None
cost_per_1k_output: Optional[float] = None
max_context: Optional[int] = None
capabilities: Set[str] = Field(default_factory=set) # e.g. {"vision", "tools"}
available: bool = True
metadata_incomplete: bool = False
@property
def key(self) -> str:
"""Stable ``provider/model_id`` identity used as a dict key everywhere."""
return candidate_key(self.provider, self.model_id)
@property
def avg_cost_per_1k(self) -> Optional[float]:
"""Blended input/output price, or None if either side is unknown.
A rough 1:3 input:output ratio (typical chat workload) is used so a
single scalar can feed the cost term of the fit score.
"""
ci, co = self.cost_per_1k_input, self.cost_per_1k_output
if ci is None or co is None:
return None
return (ci + 3.0 * co) / 4.0
class ProbeResult(BaseModel):
"""Outcome of running one benchmark task against one model.
``success=False`` means the call itself failed (network/auth/model error);
``error`` then holds a human-readable reason and ``quality_score`` stays 0.
"""
latency_ms: float = 0.0
success: bool = False
quality_score: float = 0.0 # 0..1, from the judge model
tokens_out: int = 0
error: Optional[str] = None
class ModelAssessment(BaseModel):
"""Everything we know about one candidate after an assessment run."""
metadata: ModelMetadata
# Keyed by TaskType.value (JSON-friendly string keys).
probes: Dict[str, ProbeResult] = Field(default_factory=dict)
fit_scores: Dict[str, float] = Field(default_factory=dict)
assessed_at: Optional[str] = None # ISO-8601 UTC timestamp
@property
def key(self) -> str:
return self.metadata.key
def fit_for(self, task_type: TaskType) -> float:
"""Fit score for ``task_type`` (0.0 if this model was never scored for it)."""
return float(self.fit_scores.get(task_type.value, 0.0))
# --------------------------------------------------------------------------- #
# Switch decision + pending confirmation
# --------------------------------------------------------------------------- #
class SwitchDecision(BaseModel):
"""The verdict of comparing the current model against the selector's best.
``should_switch`` is False when routing is Off, when the best candidate IS
the current model, or when the score gain is below ``min_score_gain``.
"""
should_switch: bool
from_model: Optional[str] = None # candidate key, or None if nothing active yet
to_model: Optional[str] = None
from_score: float = 0.0
to_score: float = 0.0
score_gain: float = 0.0
reason: str = ""
mode: SwitchMode = SwitchMode.OFF
task_type: Optional[str] = None
class PendingSwitch(BaseModel):
"""A Manual-mode switch proposal held until the user confirms/rejects.
Stored in-memory with a TTL; ``result`` caches the executed task output so
a repeated confirm of the same ``request_id`` is idempotent (returns the
cached result instead of running the task twice).
"""
request_id: str
task_payload: Dict[str, Any] = Field(default_factory=dict)
decision: SwitchDecision
created_at: float # epoch seconds (monotonic wall clock at creation)
expires_at: float
status: SwitchStatus = SwitchStatus.PENDING
result: Optional[Dict[str, Any]] = None # cached task result once executed
# --------------------------------------------------------------------------- #
# Helpers
# --------------------------------------------------------------------------- #
def candidate_key(provider: str, model_id: str) -> str:
"""The canonical ``provider/model_id`` string used as a dict key."""
return f"{provider}/{model_id}"
def split_key(key: str) -> tuple[str, str]:
"""Inverse of :func:`candidate_key`. Splits on the first ``/`` only, so a
model id that itself contains ``/`` (some gateways use ``org/model``) is
preserved intact."""
provider, _, model_id = key.partition("/")
return provider, model_id
__all__ = [
"TaskType",
"Policy",
"SwitchMode",
"SwitchStatus",
"ModelMetadata",
"ProbeResult",
"ModelAssessment",
"SwitchDecision",
"PendingSwitch",
"candidate_key",
"split_key",
]
+61
View File
@@ -0,0 +1,61 @@
# Sample Auto Model Assessment & Routing config (REFERENCE / DOCUMENTATION).
#
# NOTE: The running app stores config as JSON at ~/.cowork_local/config.json
# (see config.py) — this YAML mirrors that structure only to document the
# routing schema in the shape the original spec described. Copy the values into
# the JSON "routing" section (or edit them in Settings → "Auto Model Routing").
#
# API keys are NEVER stored here — each provider reads its key from an env var
# named by `api_key_env`; the app resolves it at call time and never logs it.
providers:
- name: anthropic
api_key_env: ANTHROPIC_API_KEY
base_url: https://api.anthropic.com
models:
- id: claude-opus-4-8
tier: powerful
- id: claude-haiku-4-5-20251001
tier: fast
- name: codex # OpenAI-compatible
api_key_env: OPENAI_API_KEY
base_url: https://api.openai.com/v1
models:
- id: gpt-4o
tier: balanced
- id: gpt-4o-mini
tier: fast
# Behaviour of the router (maps to config.json → "routing").
routing:
switch_mode: manual # "off" | "auto" | "manual" (global default)
policy: balanced # "quality" | "cost" | "latency" | "balanced"
min_score_gain: 0.05 # only propose a switch if new model beats current by >= this
confirm_timeout_sec: 60 # (manual) keep current model if not confirmed in time
reassess_interval_hours: 24 # periodic reassess cadence; 0 disables it
per_provider_concurrency: 2 # max concurrent probe calls per provider (rate-limit safety)
judge_provider: anthropic # provider of the fixed judge model ("" = active provider)
judge_model: claude-haiku-4-5-20251001 # one cheap judge for ALL candidates (fair grading)
auto_reassess_on_add: true # reassess a newly-added model immediately
# Explicit candidate set to assess. Leave empty to auto-discover from each
# provider's currently-configured model.
candidates:
- {provider: anthropic, model_id: claude-opus-4-8, tier: powerful}
- {provider: anthropic, model_id: claude-haiku-4-5-20251001, tier: fast}
- {provider: codex, model_id: gpt-4o, tier: balanced}
- {provider: codex, model_id: gpt-4o-mini, tier: fast}
# Per-chat-screen Off/Auto/Manual toggle state. "" = follow switch_mode above.
surface_modes:
cowork: ""
co4e: ""
ai_edit: ""
# Assessment results are written by the system (do NOT hand-edit) — the app
# keeps them in ~/.cowork_local/assessments.json, with versioned backups under
# assessments_history/<timestamp>.json. Shown here for reference only:
assessments:
last_updated: null # ISO-8601 UTC, e.g. "2026-07-22T09:30:00+00:00"
policy: balanced
results: {} # { "anthropic/claude-opus-4-8": { ...ModelAssessment... }, ... }
+165
View File
@@ -0,0 +1,165 @@
"""Orchestrate a full assessment run: enrich → probe → score → store.
``check_and_update`` is the single entry point the API/scheduler call. It:
1. Enriches each candidate's static metadata (price / context / capabilities).
2. Probes every candidate on every task type concurrently, bounded per provider
(delegated to ``prober.probe_candidates``), grading each answer with one
fixed judge.
3. Computes fit scores per task type under the active policy.
4. Persists the results atomically, backing up the previous version to history
first (so a model that *degrades* between runs can be spotted).
Cost-awareness: probing spends real tokens, so this runs only on a schedule,
when a model is added, or on an explicit reassess — never per chat turn. The
number of API calls made is logged so the cost is visible.
"""
from __future__ import annotations
import logging
from typing import Callable, Dict, List, Optional, Tuple
from .clients import ProbeClient
from .metadata import LLMDeclarer, enrich
from .models import ModelAssessment, Policy, TaskType, candidate_key
from .prober import JudgeFn, make_judge, probe_candidates
from .scorer import compute_fit_score
from .store import AssessmentStore, utc_now_iso
logger = logging.getLogger("cowork_local.routing")
# A candidate to assess: (provider, model_id, tier|None).
Candidate = Tuple[str, str, Optional[str]]
def build_assessment(
provider: str,
model_id: str,
tier: Optional[str],
probes: Dict[str, "object"],
policy: Policy,
*,
config=None,
task_types: Optional[List[TaskType]] = None,
llm_declarer: Optional[LLMDeclarer] = None,
) -> ModelAssessment:
"""Assemble one :class:`ModelAssessment` from its probe results.
Pure except for metadata enrichment (which may read the config price
sheet). Kept separate from I/O so it's unit-testable without any network.
"""
from .models import ProbeResult
task_types = task_types or list(TaskType)
meta = enrich(provider, model_id, config=config, tier=tier, llm_declarer=llm_declarer)
# A model that failed EVERY probe is effectively unavailable this run.
typed_probes: Dict[str, ProbeResult] = {}
any_success = False
for tt in task_types:
probe = probes.get(tt.value)
if isinstance(probe, ProbeResult):
typed_probes[tt.value] = probe
any_success = any_success or probe.success
if typed_probes and not any_success:
meta.available = False
fit_scores: Dict[str, float] = {}
for tt in task_types:
probe = typed_probes.get(tt.value)
if probe is not None:
fit_scores[tt.value] = compute_fit_score(meta, probe, policy)
return ModelAssessment(
metadata=meta,
probes=typed_probes,
fit_scores=fit_scores,
assessed_at=utc_now_iso(),
)
def check_and_update(
candidates: List[Candidate],
client: ProbeClient,
*,
judge: Optional[JudgeFn] = None,
judge_provider: str = "",
judge_model: str = "",
store: Optional[AssessmentStore] = None,
config=None,
policy: Policy = Policy.BALANCED,
task_types: Optional[List[TaskType]] = None,
per_provider_concurrency: int = 2,
max_workers: int = 8,
llm_declarer: Optional[LLMDeclarer] = None,
persist: bool = True,
) -> Dict[str, ModelAssessment]:
"""Assess every candidate and (optionally) persist the results.
Provide either a ready ``judge`` callable, or ``judge_provider`` +
``judge_model`` to build the standard rubric judge from ``client``.
Returns ``{candidate_key: ModelAssessment}``. ``persist=False`` skips the
store write (used by tests / dry runs).
"""
task_types = task_types or list(TaskType)
if not candidates:
logger.info("routing.reassess: no candidates configured — nothing to do")
return {}
if judge is None:
if not (judge_provider and judge_model):
raise ValueError("check_and_update needs either `judge` or judge_provider+judge_model")
judge = make_judge(client, judge_provider, judge_model)
judge_key = candidate_key(judge_provider, judge_model) if judge_provider else None
if judge_key and any(candidate_key(p, m) == judge_key for p, m, _ in candidates):
# The judge is also a candidate — its own answers are self-graded. We
# keep it routable (it may genuinely be a fine cheap model) but flag the
# bias so it's not mistaken for an independent score.
logger.warning(
"routing.reassess: judge model %s is also a candidate — its quality "
"scores are self-judged and may be optimistic", judge_key,
)
# --- count API calls so the cost of a reassess is visible ------------- #
call_count = {"n": 0}
def _tick() -> None:
call_count["n"] += 1
pairs = [(p, m) for (p, m, _tier) in candidates]
probe_map = probe_candidates(
client, pairs, task_types, judge,
per_provider_concurrency=per_provider_concurrency,
max_workers=max_workers,
call_counter=_tick,
)
assessments: Dict[str, ModelAssessment] = {}
for (provider, model_id, tier) in candidates:
key = candidate_key(provider, model_id)
assessments[key] = build_assessment(
provider, model_id, tier,
probe_map.get(key, {}), policy,
config=config, task_types=task_types, llm_declarer=llm_declarer,
)
# ~1 probe call + 1 judge call per (candidate, task). The counter above only
# counts probe calls (judge calls happen inside probe_model), so report both.
probe_calls = call_count["n"]
logger.info(
"routing.reassess: %d candidate(s) × %d task(s) → ~%d probe calls "
"(+~%d judge calls), policy=%s",
len(candidates), len(task_types), probe_calls, probe_calls, policy.value,
)
if persist:
store = store or AssessmentStore()
store.save(assessments, policy)
store.prune_history(keep=30)
return assessments
__all__ = ["check_and_update", "build_assessment", "Candidate"]
+238
View File
@@ -0,0 +1,238 @@
"""Dynamic benchmark probing + LLM-as-judge quality scoring.
For each ``(model, task_type)`` we send a fixed benchmark prompt, measure real
latency, and score the answer's quality with a single **fixed, cheap judge
model** (configurable) using a rubric that returns strict JSON. Using the same
judge for every candidate keeps the comparison fair, and never letting a model
judge its own answer avoids self-grading bias.
Concurrency is bounded **per provider** with a semaphore (the sync analogue of
``asyncio.Semaphore``, since the app's Provider layer is ``requests``-based) so
a reassess never trips a provider's rate limit. Every probe is timed and every
exception is captured as ``success=False`` — a dead model scores 0, it never
crashes the run.
"""
from __future__ import annotations
import json
import re
import threading
import time
from concurrent.futures import ThreadPoolExecutor, as_completed
from typing import Callable, Dict, List, Optional, Tuple
from .clients import CompletionResult, ProbeClient
from .models import ProbeResult, TaskType
# One fixed prompt per task type — shared by EVERY candidate so scores are
# comparable. Kept short to keep probing cheap (probes cost real tokens).
BENCHMARK_TASKS: Dict[TaskType, str] = {
TaskType.QA: (
"Answer concisely and correctly: What is the capital of Australia, and "
"name one reason it — rather than Sydney — was chosen as the capital?"
),
TaskType.CODING: (
"Write a correct Python function `is_balanced(s: str) -> bool` that returns "
"True iff the brackets (), [], {} in `s` are balanced and properly nested. "
"Return only the function, no explanation."
),
TaskType.REASONING: (
"A bat and a ball cost $1.10 in total. The bat costs $1.00 more than the "
"ball. How much does the ball cost? Show the reasoning in one or two lines "
"and give the final numeric answer."
),
TaskType.SUMMARIZATION: (
"Summarize the following in exactly one sentence: 'Photosynthesis is the "
"process by which green plants, algae and some bacteria convert light "
"energy, usually from the sun, into chemical energy stored in glucose, "
"releasing oxygen as a by-product and forming the base of most food chains.'"
),
TaskType.CREATIVE: (
"Write a vivid two-line poem about a lighthouse at dawn. Use one concrete "
"sensory image per line."
),
}
# Rubric handed to the judge. It must return STRICT JSON: {"score": 0..1}.
_JUDGE_RUBRIC = (
"You are grading an AI assistant's answer to a {task} task on a 0.0–1.0 scale.\n"
"Judge correctness, relevance and quality only — ignore verbosity/style unless "
"it harms the answer. 0.0 = wrong/empty/off-task, 0.5 = partially correct, "
"1.0 = fully correct and high quality.\n\n"
"TASK PROMPT:\n{prompt}\n\nANSWER TO GRADE:\n{answer}\n\n"
'Respond with ONLY a JSON object, no prose: {{"score": <float 0..1>}}'
)
# JudgeFn: given (task_type, prompt, answer) → quality score in [0,1].
JudgeFn = Callable[[TaskType, str, str], float]
_SCORE_RE = re.compile(r'"score"\s*:\s*([0-9]*\.?[0-9]+)')
def _clamp01(x: float) -> float:
return min(1.0, max(0.0, float(x)))
def parse_judge_score(text: str) -> float:
"""Extract the 0..1 score from a judge reply, tolerating minor noise.
Tries strict JSON first, then a regex fallback for models that wrap the
JSON in prose despite instructions. Returns 0.0 if nothing parseable.
"""
if not text:
return 0.0
try:
obj = json.loads(text.strip())
if isinstance(obj, dict) and "score" in obj:
return _clamp01(obj["score"])
except (json.JSONDecodeError, TypeError, ValueError):
pass
m = _SCORE_RE.search(text)
if m:
try:
return _clamp01(float(m.group(1)))
except ValueError:
return 0.0
return 0.0
def make_judge(
client: ProbeClient,
judge_provider: str,
judge_model: str,
) -> JudgeFn:
"""Build a :data:`JudgeFn` bound to one fixed judge model.
The same judge grades every candidate (fair comparison). The orchestrator
is responsible for not pointing the judge at the model being graded.
"""
def judge(task_type: TaskType, prompt: str, answer: str) -> float:
rubric = _JUDGE_RUBRIC.format(
task=task_type.value, prompt=prompt, answer=(answer or "")[:4000]
)
messages = [{"role": "user", "content": rubric}]
result = client.complete(judge_provider, judge_model, messages)
if not result.ok:
return 0.0
return parse_judge_score(result.text)
return judge
def probe_model(
client: ProbeClient,
provider: str,
model_id: str,
task_type: TaskType,
judge: JudgeFn,
) -> ProbeResult:
"""Run one benchmark task against one model and score it.
Measures wall-clock latency around the completion call. Any exception or a
provider-level error becomes ``success=False`` with the error captured; the
quality score then stays 0.
"""
prompt = BENCHMARK_TASKS[task_type]
messages = [{"role": "user", "content": prompt}]
started = time.perf_counter()
try:
result: CompletionResult = client.complete(provider, model_id, messages)
except Exception as exc: # noqa: BLE001 — defensive; client should not raise
elapsed_ms = (time.perf_counter() - started) * 1000.0
return ProbeResult(latency_ms=elapsed_ms, success=False, error=str(exc))
elapsed_ms = (time.perf_counter() - started) * 1000.0
if not result.ok:
return ProbeResult(latency_ms=elapsed_ms, success=False, error=result.error)
quality = judge(task_type, prompt, result.text)
return ProbeResult(
latency_ms=elapsed_ms,
success=True,
quality_score=quality,
tokens_out=result.tokens_out,
)
class _PerProviderSemaphores:
"""Lazily-created, per-provider bounded semaphores for rate-limit safety."""
def __init__(self, limit: int) -> None:
self._limit = max(1, int(limit))
self._sems: Dict[str, threading.Semaphore] = {}
self._lock = threading.Lock()
def get(self, provider: str) -> threading.Semaphore:
with self._lock:
sem = self._sems.get(provider)
if sem is None:
sem = threading.Semaphore(self._limit)
self._sems[provider] = sem
return sem
def probe_candidates(
client: ProbeClient,
candidates: List[Tuple[str, str]],
task_types: List[TaskType],
judge: JudgeFn,
*,
per_provider_concurrency: int = 2,
max_workers: int = 8,
call_counter: Optional[Callable[[], None]] = None,
) -> Dict[str, Dict[str, ProbeResult]]:
"""Probe every ``(provider, model_id)`` on every ``task_type`` concurrently.
Concurrency is capped globally by ``max_workers`` and, more importantly,
**per provider** by ``per_provider_concurrency`` — so many models on one
provider never fire more than N calls at once at that provider.
``call_counter`` (if given) is invoked once per probe call, letting the
orchestrator log "how many API calls this reassess made".
Returns ``{candidate_key: {task_type_value: ProbeResult}}``.
"""
from .models import candidate_key
sems = _PerProviderSemaphores(per_provider_concurrency)
results: Dict[str, Dict[str, ProbeResult]] = {}
results_lock = threading.Lock()
def _one(provider: str, model_id: str, task_type: TaskType) -> None:
sem = sems.get(provider)
with sem:
if call_counter is not None:
call_counter()
probe = probe_model(client, provider, model_id, task_type, judge)
key = candidate_key(provider, model_id)
with results_lock:
results.setdefault(key, {})[task_type.value] = probe
jobs = [
(provider, model_id, task_type)
for (provider, model_id) in candidates
for task_type in task_types
]
if not jobs:
return results
with ThreadPoolExecutor(max_workers=max(1, max_workers)) as pool:
futures = [pool.submit(_one, p, m, t) for (p, m, t) in jobs]
for f in as_completed(futures):
# _one swallows its own errors into a ProbeResult; this is just to
# surface any truly-unexpected exception without killing the pool.
f.result()
return results
__all__ = [
"BENCHMARK_TASKS",
"JudgeFn",
"make_judge",
"probe_model",
"probe_candidates",
"parse_judge_score",
]
+126
View File
@@ -0,0 +1,126 @@
"""Periodic reassessment scheduler (Qt layer).
No APScheduler dependency — this mirrors the app's existing ``TaskScheduler``:
a lightweight ``QTimer`` ticks periodically and, when the configured interval
has elapsed since the last assessment, launches a background reassess on a
daemon thread (so the UI never blocks). It also expires stale Manual-mode
pending switches on each tick.
Reassessment is expensive (it spends real tokens), so the cadence is
deliberately coarse — default every 24h, configurable via
``routing.reassess_interval_hours`` (0 disables the periodic run entirely).
"""
from __future__ import annotations
import logging
from datetime import datetime, timezone
from typing import Any, Optional
from PySide6.QtCore import QObject, QTimer, Signal
logger = logging.getLogger("cowork_local.routing")
# How often the timer wakes to CHECK whether a reassess is due. The actual
# reassess cadence is governed by reassess_interval_hours; this is just the
# polling granularity (cheap — it only reads a timestamp).
_TICK_MS = 30 * 60 * 1000 # 30 minutes
class RoutingScheduler(QObject):
"""Drives periodic reassessment + pending-switch expiry for a service."""
reassess_started = Signal()
reassess_finished = Signal(int) # number of models assessed
def __init__(self, ctx: Any, service: Any, parent: Optional[QObject] = None) -> None:
super().__init__(parent)
self.ctx = ctx
self.service = service
self._timer = QTimer(self)
self._timer.setInterval(_TICK_MS)
self._timer.timeout.connect(self.tick)
# -- lifecycle ------------------------------------------------------ #
def start(self) -> None:
"""Begin periodic checks. Does NOT force an immediate reassess — the
first one happens when the interval is genuinely due (or never, if the
store is fresh), to avoid a burst of API calls at every app launch."""
self.tick()
self._timer.start()
def stop(self) -> None:
self._timer.stop()
# -- tick ----------------------------------------------------------- #
def _interval_hours(self) -> float:
try:
return float(self.ctx.config.routing.get("reassess_interval_hours", 24) or 0)
except Exception: # noqa: BLE001
return 24.0
def _hours_since_last(self) -> Optional[float]:
last = self.service.store.last_updated()
if not last:
return None # never assessed
try:
dt = datetime.fromisoformat(last)
if dt.tzinfo is None:
dt = dt.replace(tzinfo=timezone.utc)
return (datetime.now(timezone.utc) - dt).total_seconds() / 3600.0
except (ValueError, TypeError):
return None
def _routing_enabled_anywhere(self) -> bool:
"""Is routing actually in use? True if the global mode is auto/manual OR
any chat surface overrides to auto/manual. When everything is Off, the
assessment scores would never be consulted — so we don't spend tokens
probing for them (no surprise cost on a fresh install)."""
try:
routing = self.ctx.config.routing
if (routing.get("switch_mode") or "off") in ("auto", "manual"):
return True
for m in (routing.get("surface_modes") or {}).values():
if m in ("auto", "manual"):
return True
except Exception: # noqa: BLE001
pass
return False
def is_due(self) -> bool:
if not self._routing_enabled_anywhere():
return False # routing off everywhere → don't probe (would be wasted cost)
interval = self._interval_hours()
if interval <= 0:
return False # periodic reassess disabled
since = self._hours_since_last()
if since is None:
return True # never assessed → due once routing is actually enabled
return since >= interval
def tick(self) -> None:
"""Expire stale pending switches; reassess if the interval is due."""
try:
self.service.sweep_pending()
except Exception: # noqa: BLE001
logger.exception("routing.scheduler: sweep_pending failed")
if not self.is_due() or self.service.is_reassessing():
return
logger.info("routing.scheduler: reassess is due — starting background run")
self.reassess_started.emit()
def _done(result) -> None:
self.reassess_finished.emit(len(result or {}))
self.service.reassess_background(on_done=_done)
def trigger_now(self) -> None:
"""Force an out-of-band reassess (e.g. Settings' 'Reassess now' button)."""
if self.service.is_reassessing():
return
self.reassess_started.emit()
self.service.reassess_background(on_done=lambda r: self.reassess_finished.emit(len(r or {})))
__all__ = ["RoutingScheduler"]
+80
View File
@@ -0,0 +1,80 @@
"""Fit scoring: turn a model's metadata + probe result into a 0..1 score.
The score blends three normalized terms — quality (from the judge), cost
(cheaper is better) and latency (faster is better) — weighted by the active
:class:`~cowork_local.core.routing.models.Policy`::
fit = w_quality * quality
+ w_cost * 1/(1 + cost)
+ w_latency * 1/(1 + latency_s)
Each term is in ``[0, 1]`` and the weights sum to 1, so ``fit`` is in ``[0, 1]``.
A probe that failed scores 0 outright — an unusable model must never win.
"""
from __future__ import annotations
from typing import Dict
from .models import ModelMetadata, Policy, ProbeResult
# Weights per policy: (quality, cost, latency). Each row sums to 1.0.
# quality — pick the smartest model, cost/speed barely matter.
# cost — pick the cheapest usable model.
# latency — pick the fastest usable model.
# balanced — a sensible default that still leans on quality.
POLICY_WEIGHTS: Dict[Policy, tuple[float, float, float]] = {
Policy.QUALITY: (0.80, 0.10, 0.10),
Policy.COST: (0.20, 0.70, 0.10),
Policy.LATENCY: (0.20, 0.10, 0.70),
Policy.BALANCED: (0.50, 0.25, 0.25),
}
# When a model's price is unknown (metadata_incomplete), we cannot compute a
# real cost term. Rather than reward the gap (cost=0 → term=1.0, unfairly
# best) or nuke the model (term=0), we assume a neutral middling price so it
# competes on quality/latency without a fabricated cost advantage.
_UNKNOWN_COST_PER_1K = 0.01
def _cost_term(metadata: ModelMetadata) -> float:
"""Normalized cost term ``1/(1+cost)`` in ``(0, 1]`` — higher is cheaper."""
cost = metadata.avg_cost_per_1k
if cost is None:
cost = _UNKNOWN_COST_PER_1K
cost = max(0.0, float(cost))
return 1.0 / (1.0 + cost)
def _latency_term(probe: ProbeResult) -> float:
"""Normalized latency term ``1/(1+latency_s)`` in ``(0, 1]`` — higher is faster."""
latency_s = max(0.0, float(probe.latency_ms)) / 1000.0
return 1.0 / (1.0 + latency_s)
def compute_fit_score(
metadata: ModelMetadata,
probe: ProbeResult,
policy: Policy = Policy.BALANCED,
) -> float:
"""Fit score in ``[0, 1]`` for one model on one task, under ``policy``.
Returns 0.0 immediately if the probe failed or the model is unavailable —
an unusable model is never routable regardless of its price/speed.
"""
if not probe.success or not metadata.available:
return 0.0
w_quality, w_cost, w_latency = POLICY_WEIGHTS.get(
policy, POLICY_WEIGHTS[Policy.BALANCED]
)
quality = min(1.0, max(0.0, float(probe.quality_score)))
cost_term = _cost_term(metadata)
latency_term = _latency_term(probe)
score = w_quality * quality + w_cost * cost_term + w_latency * latency_term
# Clamp defensively against float drift; the math already bounds it to [0,1].
return round(min(1.0, max(0.0, score)), 6)
__all__ = ["POLICY_WEIGHTS", "compute_fit_score"]
+128
View File
@@ -0,0 +1,128 @@
"""Select the best-fit model for a task type under a policy.
The selector is deliberately *stateless and pure*: given a set of assessments,
a task type and a policy, it recomputes each candidate's fit score from its
stored probe + metadata (via :func:`scorer.compute_fit_score`) and ranks them.
Recomputing (rather than trusting the ``fit_scores`` cached at assess time)
means changing the routing **policy** — quality → cost, say — re-ranks instantly
from existing measurements, with **no** expensive re-probing. Probes are the raw
truth; fit is a pure function of ``(probe, metadata, policy)``.
"""
from __future__ import annotations
from dataclasses import dataclass, field
from typing import Dict, Iterable, List, Optional, Set
from .models import ModelAssessment, Policy, TaskType
from .scorer import compute_fit_score
@dataclass
class RankedCandidate:
"""One candidate's standing for a given task type + policy."""
assessment: ModelAssessment
score: float
@property
def key(self) -> str:
return self.assessment.key
@dataclass
class Ranking:
"""Full ordering of candidates for a task type, best first."""
task_type: TaskType
policy: Policy
ranked: List[RankedCandidate] = field(default_factory=list)
@property
def best(self) -> Optional[RankedCandidate]:
return self.ranked[0] if self.ranked else None
def score_of(self, key: str) -> float:
"""Score of a specific candidate key, or 0.0 if it isn't ranked
(unavailable / filtered out / failed probe)."""
for c in self.ranked:
if c.key == key:
return c.score
return 0.0
def as_dicts(self) -> List[Dict]:
"""JSON-friendly ranking for API responses / the UI."""
return [
{
"key": c.key,
"provider": c.assessment.metadata.provider,
"model_id": c.assessment.metadata.model_id,
"score": c.score,
"tier": c.assessment.metadata.tier,
"metadata_incomplete": c.assessment.metadata.metadata_incomplete,
}
for c in self.ranked
]
def _has_capabilities(assessment: ModelAssessment, required: Set[str]) -> bool:
return required.issubset(assessment.metadata.capabilities)
def rank_models(
assessments: Iterable[ModelAssessment],
task_type: TaskType,
policy: Policy = Policy.BALANCED,
*,
required_capabilities: Optional[Iterable[str]] = None,
) -> Ranking:
"""Rank candidates for ``task_type`` under ``policy``, best first.
A candidate is excluded when it is unavailable, lacks a probe for this task
type, fails the required-capability filter, or scores 0 (failed probe).
Ties break by lower average cost, then by model id, for stable ordering.
"""
required: Set[str] = set(required_capabilities or ())
scored: List[RankedCandidate] = []
for a in assessments:
if not a.metadata.available:
continue
if required and not _has_capabilities(a, required):
continue
probe = a.probes.get(task_type.value)
if probe is None:
continue
score = compute_fit_score(a.metadata, probe, policy)
if score <= 0.0:
continue
scored.append(RankedCandidate(assessment=a, score=score))
def _sort_key(c: RankedCandidate):
cost = c.assessment.metadata.avg_cost_per_1k
cost = cost if cost is not None else float("inf")
# score desc, then cheaper, then model id for determinism.
return (-c.score, cost, c.assessment.metadata.model_id)
scored.sort(key=_sort_key)
return Ranking(task_type=task_type, policy=policy, ranked=scored)
def best_model(
assessments: Iterable[ModelAssessment],
task_type: TaskType,
policy: Policy = Policy.BALANCED,
*,
required_capabilities: Optional[Iterable[str]] = None,
) -> Optional[RankedCandidate]:
"""The single best-fit candidate for ``task_type``, or ``None`` if none
qualify (all unavailable / filtered / failed)."""
return rank_models(
assessments,
task_type,
policy,
required_capabilities=required_capabilities,
).best
__all__ = ["Ranking", "RankedCandidate", "rank_models", "best_model"]
+357
View File
@@ -0,0 +1,357 @@
"""RoutingService — the façade the UI (and the "REST-equivalent" API) talk to.
It wires the pieces together and holds the per-app state (assessment store +
pending-switch registry). It is deliberately **Qt-free and thread-safe** so it
can run from a chat worker thread, the scheduler, or a test. The UI layer adds
the toggle widget and the Manual-mode confirm dialog on top of these methods.
Logical API surface (mirrors the task's REST endpoints):
* :meth:`reassess` ↔ ``POST /models/reassess``
* :meth:`best_for` ↔ ``GET /models/best``
* :meth:`assessments` / :meth:`status` ↔ ``GET /models/assessments``
* :meth:`add_candidate` ↔ ``POST /models/add``
* :meth:`route` ↔ the decision half of ``POST /task/execute``
* :meth:`create_pending` / :meth:`resolve_pending` ↔ ``POST /task/confirm-switch``
* :meth:`get_routing_config` / :meth:`update_routing_config` ↔ ``/routing/config``
"""
from __future__ import annotations
import logging
import threading
from dataclasses import dataclass
from typing import Any, Callable, Dict, List, Optional, Tuple
from .classifier import classify
from .clients import AppProbeClient, ProbeClient
from .models import (
ModelAssessment,
PendingSwitch,
Policy,
SwitchDecision,
SwitchMode,
TaskType,
candidate_key,
split_key,
)
from .orchestrator import Candidate, check_and_update
from .prober import make_judge
from .selector import Ranking, rank_models
from .store import AssessmentStore
from .switch_controller import Executor, PendingSwitchRegistry, decide
logger = logging.getLogger("cowork_local.routing")
# A sensible cheap judge model per known provider, used when the user hasn't
# pinned one in Settings. Falls back to the provider's own configured model.
_CHEAP_JUDGE_MODEL = {
"anthropic": "claude-haiku-4-5-20251001",
"codex": "gpt-4o-mini",
"github_copilot": "gpt-4o-mini",
"openai_compat": "", # unknown gateway → use configured model
"ollama": "", # local → use configured model
}
@dataclass
class RouteResult:
"""Outcome of routing one turn (before any execution)."""
mode: SwitchMode
task_type: TaskType
decision: SwitchDecision
ranking: Optional[Ranking] = None
@property
def should_switch(self) -> bool:
return self.decision.should_switch
@property
def needs_confirmation(self) -> bool:
"""Manual mode with a worthwhile switch → the UI must ask the user."""
return self.mode == SwitchMode.MANUAL and self.decision.should_switch
def target(self) -> Optional[Tuple[str, str]]:
"""The (provider, model_id) to switch to, or None."""
if not self.decision.to_model:
return None
return split_key(self.decision.to_model)
class RoutingService:
"""Central routing coordinator, one per :class:`AppContext`."""
def __init__(
self,
ctx: Any,
*,
store: Optional[AssessmentStore] = None,
client: Optional[ProbeClient] = None,
clock: Optional[Callable[[], float]] = None,
) -> None:
self.ctx = ctx
self.store = store or AssessmentStore()
self._client = client # None → lazily build AppProbeClient(ctx)
import time as _time
self.pending = PendingSwitchRegistry(clock=clock or _time.time)
self._reassess_lock = threading.Lock()
self._reassessing = False
# -- config helpers ------------------------------------------------- #
@property
def _routing_cfg(self) -> Dict[str, Any]:
return self.ctx.config.routing
def get_routing_config(self) -> Dict[str, Any]:
"""Current routing behaviour config (for ``GET /routing/config``)."""
return dict(self._routing_cfg)
def update_routing_config(self, **changes) -> Dict[str, Any]:
"""Patch routing config (``PATCH /routing/config``) and persist.
Only known keys are accepted; unknown keys are ignored so a typo can't
silently poison the config.
"""
cfg = self._routing_cfg
allowed = {
"switch_mode", "policy", "min_score_gain", "confirm_timeout_sec",
"reassess_interval_hours", "per_provider_concurrency",
"judge_provider", "judge_model", "auto_reassess_on_add",
}
for k, v in changes.items():
if k in allowed:
cfg[k] = v
self.ctx.config.save()
return dict(cfg)
def _policy(self) -> Policy:
raw = (self._routing_cfg.get("policy") or "balanced").lower()
try:
return Policy(raw)
except ValueError:
return Policy.BALANCED
def _client_or_build(self) -> ProbeClient:
if self._client is None:
self._client = AppProbeClient(self.ctx)
return self._client
def _resolve_judge(self) -> Tuple[str, str]:
"""Which (provider, model) grades every probe.
Uses the pinned judge from config when set, else a cheap default for
the active provider (falling back to that provider's configured model).
"""
cfg = self._routing_cfg
provider = cfg.get("judge_provider") or self.ctx.config.active_provider
model = cfg.get("judge_model") or ""
if not model:
model = _CHEAP_JUDGE_MODEL.get(provider, "")
if not model:
model = self.ctx.config.provider_conf(provider).get("model", "")
return provider, model
# -- candidates ----------------------------------------------------- #
def candidates(self) -> List[Candidate]:
"""The models to assess: explicit ``routing.candidates`` plus each
provider's currently-configured model (so the model in use is always
scored). Deduplicated, order-stable."""
out: List[Candidate] = []
seen = set()
def _add(provider: str, model_id: str, tier: Optional[str]) -> None:
if not provider or not model_id:
return
key = candidate_key(provider, model_id)
if key in seen:
return
seen.add(key)
out.append((provider, model_id, tier))
for c in self._routing_cfg.get("candidates") or []:
if isinstance(c, dict):
_add(c.get("provider", ""), c.get("model_id", ""), c.get("tier"))
# Always include each configured provider's active model.
for name, conf in (self.ctx.config.data.get("providers") or {}).items():
_add(name, conf.get("model", ""), None)
return out
def add_candidate(
self,
provider: str,
model_id: str,
tier: Optional[str] = None,
*,
reassess: Optional[bool] = None,
) -> bool:
"""Add a model to the assessed set (``POST /models/add``).
Returns True if it was newly added. When ``reassess`` (defaults to the
``auto_reassess_on_add`` config) is True, kicks off a background
reassess so the new model gets scored right away.
"""
cfg = self._routing_cfg
cand = cfg.setdefault("candidates", [])
key = candidate_key(provider, model_id)
if any(candidate_key(c.get("provider", ""), c.get("model_id", "")) == key
for c in cand if isinstance(c, dict)):
return False
cand.append({"provider": provider, "model_id": model_id, "tier": tier})
self.ctx.config.save()
do_reassess = cfg.get("auto_reassess_on_add", True) if reassess is None else reassess
if do_reassess:
self.reassess_background()
return True
# -- assessment run ------------------------------------------------- #
def reassess(
self,
policy: Optional[Policy] = None,
*,
client: Optional[ProbeClient] = None,
) -> Dict[str, ModelAssessment]:
"""Run a full assessment (blocking). Safe to call from a worker thread.
Guarded so two reassessments never run at once (a second call while one
is in flight is a no-op returning the current store)."""
with self._reassess_lock:
if self._reassessing:
logger.info("routing.reassess: already running — skipping duplicate")
return self.store.load()
self._reassessing = True
try:
policy = policy or self._policy()
judge_provider, judge_model = self._resolve_judge()
cli = client or self._client_or_build()
if not judge_model:
logger.warning("routing.reassess: no judge model resolved — aborting")
return self.store.load()
return check_and_update(
self.candidates(), cli,
judge_provider=judge_provider, judge_model=judge_model,
store=self.store, config=self.ctx.config, policy=policy,
per_provider_concurrency=int(self._routing_cfg.get("per_provider_concurrency", 2)),
)
finally:
with self._reassess_lock:
self._reassessing = False
def reassess_background(
self,
policy: Optional[Policy] = None,
on_done: Optional[Callable[[Dict[str, ModelAssessment]], None]] = None,
) -> threading.Thread:
"""Run :meth:`reassess` on a daemon thread (non-Qt, headless-safe)."""
def _run() -> None:
try:
result = self.reassess(policy)
except Exception: # noqa: BLE001 — never let a reassess crash the app
logger.exception("routing.reassess background run failed")
result = {}
if on_done is not None:
try:
on_done(result)
except Exception: # noqa: BLE001
logger.exception("routing.reassess on_done callback failed")
t = threading.Thread(target=_run, name="routing-reassess", daemon=True)
t.start()
return t
def is_reassessing(self) -> bool:
return self._reassessing
# -- query ---------------------------------------------------------- #
def assessments(self) -> Dict[str, ModelAssessment]:
return self.store.load()
def status(self) -> Dict[str, Any]:
"""``GET /models/assessments`` — last_updated + per-model summary."""
assessments = self.store.load()
return {
"last_updated": self.store.last_updated(),
"policy": self.store.policy(),
"count": len(assessments),
"models": sorted(assessments.keys()),
}
def best_for(
self,
task_type: TaskType,
policy: Optional[Policy] = None,
*,
required_capabilities: Optional[List[str]] = None,
) -> Ranking:
"""Ranking + best model for a task type (``GET /models/best``)."""
policy = policy or self._policy()
return rank_models(
self.store.load().values(), task_type, policy,
required_capabilities=required_capabilities,
)
# -- routing decision ----------------------------------------------- #
def route(
self,
surface: str,
prompt: str,
current_provider: str,
current_model: str,
*,
mode_override: Optional[str] = None,
required_capabilities: Optional[List[str]] = None,
task_type: Optional[TaskType] = None,
) -> RouteResult:
"""Decide whether/how to switch models for one turn on ``surface``.
Does NOT execute anything — returns a :class:`RouteResult` the caller
acts on (Auto → switch & run; Manual+should_switch → confirm; else run
as-is). Never raises: any internal failure yields an Off/no-switch
result so a broken assessment store can't block chatting.
"""
try:
mode = (mode_override or self.ctx.config.routing_mode_for(surface) or "off").lower()
mode_enum = SwitchMode(mode) if mode in ("off", "auto", "manual") else SwitchMode.OFF
tt = task_type or classify(prompt)
current_key = candidate_key(current_provider, current_model) if current_model else None
if mode_enum == SwitchMode.OFF:
decision = decide(current_key, rank_models([], tt), SwitchMode.OFF, 0.0, task_type=tt)
return RouteResult(mode=mode_enum, task_type=tt, decision=decision)
policy = self._policy()
ranking = rank_models(
self.store.load().values(), tt, policy,
required_capabilities=required_capabilities,
)
min_gain = float(self._routing_cfg.get("min_score_gain", 0.05) or 0.0)
decision = decide(current_key, ranking, mode_enum, min_gain, task_type=tt)
return RouteResult(mode=mode_enum, task_type=tt, decision=decision, ranking=ranking)
except Exception: # noqa: BLE001 — routing must never break a chat turn
logger.exception("routing.route failed — falling back to no-switch")
tt = task_type or TaskType.QA
current_key = candidate_key(current_provider, current_model) if current_model else None
decision = decide(current_key, rank_models([], tt), SwitchMode.OFF, 0.0, task_type=tt)
return RouteResult(mode=SwitchMode.OFF, task_type=tt, decision=decision)
# -- manual pending switches ---------------------------------------- #
def create_pending(self, decision: SwitchDecision, task_payload: Dict) -> PendingSwitch:
"""Register a Manual-mode proposal awaiting the user's confirm."""
timeout = float(self._routing_cfg.get("confirm_timeout_sec", 60) or 60)
return self.pending.create(decision, task_payload, timeout)
def resolve_pending(self, request_id: str, approve: bool, run: Executor) -> Optional[Dict]:
"""Confirm/reject a pending switch (idempotent) — ``POST /task/confirm-switch``."""
return self.pending.resolve(request_id, approve, run)
def get_pending(self, request_id: str) -> Optional[PendingSwitch]:
return self.pending.get(request_id)
def sweep_pending(self) -> List[str]:
"""Expire overdue pending switches (called periodically by the scheduler)."""
return self.pending.sweep_expired()
__all__ = ["RoutingService", "RouteResult"]
+211
View File
@@ -0,0 +1,211 @@
"""Persistence for model assessments — the routing feature's "loader/writer".
Assessments can be large and are rewritten wholesale on every reassess, so they
live in their OWN file (``~/.cowork_local/assessments.json`` by default) rather
than bloating the main ``config.json``. Two robustness guarantees the task
requires:
* **Atomic write** — results are written to a temp file in the same directory
and then ``os.replace``'d over the target, so an interrupted reassess can
never leave a half-written / corrupt store behind.
* **Versioned history** — before each overwrite, the previous store is copied
to ``assessments_history/<timestamp>.json`` so a model that *degrades*
between runs can be detected after the fact.
On-disk shape (JSON, mirrors the task's YAML ``assessments`` block)::
{
"last_updated": "2026-07-22T09:30:00+00:00",
"policy": "balanced",
"results": {
"anthropic/claude-opus-4-8": { <ModelAssessment> },
...
}
}
"""
from __future__ import annotations
import json
import os
import tempfile
from datetime import datetime, timezone
from pathlib import Path
from typing import Dict, Optional
from .models import ModelAssessment, Policy
# Default location under the app's config dir. Imported lazily so tests can
# point the store anywhere without touching the real home directory.
_DEFAULT_STORE_NAME = "assessments.json"
_HISTORY_DIR_NAME = "assessments_history"
def utc_now_iso() -> str:
"""Current UTC time as an ISO-8601 string (used for ``last_updated``)."""
return datetime.now(timezone.utc).isoformat()
def _fs_safe_stamp() -> str:
"""A filesystem-safe timestamp for history filenames (no ``:``)."""
return datetime.now(timezone.utc).strftime("%Y%m%dT%H%M%S%fZ")
class AssessmentStore:
"""Reads/writes the assessment JSON file with atomic writes + history.
Parameters
----------
store_path:
Path to the assessments JSON file. If ``None``, defaults to
``~/.cowork_local/assessments.json``.
history_dir:
Directory for pre-overwrite backups. Defaults to a sibling
``assessments_history/`` next to ``store_path``.
"""
def __init__(
self,
store_path: Optional[Path] = None,
history_dir: Optional[Path] = None,
) -> None:
if store_path is None:
from ...config import CONFIG_DIR # lazy: avoids import cost in tests
store_path = CONFIG_DIR / _DEFAULT_STORE_NAME
self.store_path = Path(store_path)
self.history_dir = Path(
history_dir or self.store_path.parent / _HISTORY_DIR_NAME
)
# -- read ----------------------------------------------------------- #
def load_raw(self) -> Dict:
"""Return the raw JSON dict, or an empty skeleton if absent/corrupt.
A corrupt store must never crash the app (same philosophy as
``AppConfig.load``) — we fall back to an empty result set so the next
reassess simply rebuilds it.
"""
if not self.store_path.exists():
return {"last_updated": None, "policy": Policy.BALANCED.value, "results": {}}
try:
return json.loads(self.store_path.read_text(encoding="utf-8"))
except (json.JSONDecodeError, OSError):
return {"last_updated": None, "policy": Policy.BALANCED.value, "results": {}}
def load(self) -> Dict[str, ModelAssessment]:
"""Return ``{candidate_key: ModelAssessment}`` parsed from disk.
Individual malformed entries are skipped rather than failing the whole
load — one bad row shouldn't hide every good assessment.
"""
raw = self.load_raw()
out: Dict[str, ModelAssessment] = {}
for key, payload in (raw.get("results") or {}).items():
try:
out[key] = ModelAssessment.model_validate(payload)
except Exception: # noqa: BLE001 — skip a single corrupt entry
continue
return out
def last_updated(self) -> Optional[str]:
return self.load_raw().get("last_updated")
def policy(self) -> str:
return self.load_raw().get("policy") or Policy.BALANCED.value
# -- write ---------------------------------------------------------- #
def save(
self,
results: Dict[str, ModelAssessment],
policy: Policy | str = Policy.BALANCED,
*,
last_updated: Optional[str] = None,
backup: bool = True,
) -> Path:
"""Atomically write ``results`` to the store, backing up the previous
version to history first.
Returns the store path. Never leaves a partially-written file: the
payload is fully serialized to a temp file and only then swapped into
place with ``os.replace`` (atomic on the same filesystem, incl. NTFS).
"""
if backup:
self._backup_existing()
policy_val = policy.value if isinstance(policy, Policy) else str(policy)
payload = {
"last_updated": last_updated or utc_now_iso(),
"policy": policy_val,
"results": {
key: assessment.model_dump(mode="json")
for key, assessment in results.items()
},
}
self.store_path.parent.mkdir(parents=True, exist_ok=True)
text = json.dumps(payload, indent=2, ensure_ascii=False)
# Temp file MUST be on the same volume as the target for os.replace to
# be atomic — so create it in the target's own directory.
fd, tmp_name = tempfile.mkstemp(
dir=str(self.store_path.parent),
prefix=".assessments-",
suffix=".tmp",
)
try:
with os.fdopen(fd, "w", encoding="utf-8") as fh:
fh.write(text)
fh.flush()
os.fsync(fh.fileno()) # durability: survive a crash right after
os.replace(tmp_name, self.store_path) # atomic swap
except BaseException:
# Clean up the temp file on any failure so we never leave litter.
try:
os.unlink(tmp_name)
except OSError:
pass
raise
return self.store_path
def _backup_existing(self) -> Optional[Path]:
"""Copy the current store into history as ``<timestamp>.json``.
No-op when there is nothing to back up. Best-effort: a failed backup
must not block the (more important) new write.
"""
if not self.store_path.exists():
return None
try:
self.history_dir.mkdir(parents=True, exist_ok=True)
# The wall clock can be coarse (Windows ~15ms), so two rapid
# backups may share a timestamp — disambiguate with a counter so a
# snapshot is never silently overwritten.
stamp = _fs_safe_stamp()
dst = self.history_dir / f"{stamp}.json"
n = 1
while dst.exists():
dst = self.history_dir / f"{stamp}-{n}.json"
n += 1
dst.write_text(
self.store_path.read_text(encoding="utf-8"), encoding="utf-8"
)
return dst
except OSError:
return None
def history_files(self) -> list[Path]:
"""All history snapshots, oldest first (for degrade detection / UI)."""
if not self.history_dir.exists():
return []
return sorted(self.history_dir.glob("*.json"))
def prune_history(self, keep: int = 30) -> None:
"""Keep only the newest ``keep`` history snapshots; delete the rest."""
files = self.history_files()
for old in files[:-keep] if keep > 0 else files:
try:
old.unlink()
except OSError:
pass
__all__ = ["AssessmentStore", "utc_now_iso"]
+280
View File
@@ -0,0 +1,280 @@
"""Decide whether to switch models, and orchestrate Auto vs Manual execution.
Two independent pieces:
* :func:`decide` — a **pure** function turning ``(current model, ranking, mode,
threshold)`` into a :class:`SwitchDecision`. No I/O, no state; trivially
testable.
* :class:`PendingSwitchRegistry` — an in-memory, TTL'd, thread-safe store of
Manual-mode switch proposals awaiting user confirmation, with an
**idempotent** ``resolve`` (confirming the same ``request_id`` twice never
runs the task twice).
Flow (from the task spec)::
task arrives → classify task_type → selector.best_model()
→ decide() compares best vs current
→ gain < min_score_gain → keep current model
→ gain ok, mode == AUTO → switch now, run task
→ gain ok, mode == MANUAL → create PendingSwitch, ask user
→ gain ok, mode == OFF → never switch (keep current)
"""
from __future__ import annotations
import threading
import time
import uuid
from typing import Callable, Dict, List, Optional, Set
from .models import (
PendingSwitch,
SwitchDecision,
SwitchMode,
SwitchStatus,
TaskType,
)
from .selector import Ranking
# Executor signature used by the registry: given the resolved model key and
# whether that represents a switch away from the original, run the task and
# return a JSON-serializable result dict.
Executor = Callable[[str, bool], Dict]
def decide(
current_key: Optional[str],
ranking: Ranking,
mode: SwitchMode,
min_score_gain: float,
*,
task_type: Optional[TaskType] = None,
) -> SwitchDecision:
"""Compare the current model against the ranking's best under ``mode``.
Returns a :class:`SwitchDecision` whose ``should_switch`` is True only when
routing is enabled, a better candidate exists, and it beats the current
model by at least ``min_score_gain``. ``reason`` always explains the call
in words (e.g. *"coding fit 0.82 > current 0.71, gain 0.11"*).
"""
tt = task_type or ranking.task_type
best = ranking.best
tt_name = tt.value if tt else "?"
base = dict(
from_model=current_key,
to_model=best.key if best else None,
mode=mode,
task_type=tt.value if tt else None,
)
# Routing disabled → never switch.
if mode == SwitchMode.OFF:
return SwitchDecision(
should_switch=False, score_gain=0.0,
reason="routing off — keeping current model", **base,
)
# Nothing assessed / nothing usable → cannot switch.
if best is None:
return SwitchDecision(
should_switch=False, score_gain=0.0,
reason="no assessed candidate available for this task", **base,
)
to_score = best.score
from_score = ranking.score_of(current_key) if current_key else 0.0
best_name = best.assessment.metadata.model_id
# No current model yet (fresh surface) → adopt the best outright.
if not current_key:
return SwitchDecision(
should_switch=to_score > 0.0,
from_score=0.0, to_score=to_score, score_gain=to_score,
reason=f"no current model — selecting best-fit {best_name} ({tt_name} fit {to_score:.2f})",
**base,
)
# Current model is already the best-fit → stay put.
if best.key == current_key:
return SwitchDecision(
should_switch=False,
from_score=from_score, to_score=to_score, score_gain=0.0,
reason=f"current model is already best-fit for {tt_name} (fit {to_score:.2f})",
**base,
)
gain = round(to_score - from_score, 6)
if gain < min_score_gain:
return SwitchDecision(
should_switch=False,
from_score=from_score, to_score=to_score, score_gain=gain,
reason=(
f"best {best_name} fit {to_score:.2f} vs current {from_score:.2f}, "
f"gain {gain:.2f} < threshold {min_score_gain:.2f} — keeping current"
),
**base,
)
return SwitchDecision(
should_switch=True,
from_score=from_score, to_score=to_score, score_gain=gain,
reason=(
f"{tt_name} fit {to_score:.2f} > current {from_score:.2f}, "
f"gain {gain:.2f} — switch to {best_name}"
),
**base,
)
class PendingSwitchRegistry:
"""Thread-safe, TTL'd store of Manual-mode switch proposals.
A proposal is created when Manual mode wants to switch; the UI shows it and
later calls :meth:`resolve` with the user's approve/reject. Idempotency:
resolving the same ``request_id`` more than once runs the task exactly once
and returns the cached result to every caller.
``clock`` is injectable so tests can drive expiry deterministically.
"""
# A short cap so a wedged executor can't hang a waiting confirm forever.
_RESOLVE_WAIT_SEC = 600.0
def __init__(self, clock: Callable[[], float] = time.time) -> None:
self._items: Dict[str, PendingSwitch] = {}
self._events: Dict[str, threading.Event] = {}
self._running: Set[str] = set()
self._lock = threading.Lock()
self._clock = clock
# -- creation ------------------------------------------------------- #
def create(
self,
decision: SwitchDecision,
task_payload: Dict,
timeout_sec: float,
) -> PendingSwitch:
"""Register a new pending switch and return it (with a fresh id)."""
rid = uuid.uuid4().hex
now = self._clock()
ps = PendingSwitch(
request_id=rid,
task_payload=task_payload,
decision=decision,
created_at=now,
expires_at=now + max(0.0, float(timeout_sec)),
status=SwitchStatus.PENDING,
)
with self._lock:
self._items[rid] = ps
self._events[rid] = threading.Event()
return ps
# -- lookup --------------------------------------------------------- #
def get(self, request_id: str) -> Optional[PendingSwitch]:
"""Fetch a pending switch, lazily marking it EXPIRED if its TTL passed."""
with self._lock:
ps = self._items.get(request_id)
if ps is not None:
self._maybe_expire_locked(ps)
return ps
def _maybe_expire_locked(self, ps: PendingSwitch) -> None:
if ps.status == SwitchStatus.PENDING and self._clock() >= ps.expires_at:
ps.status = SwitchStatus.EXPIRED
# -- resolution ----------------------------------------------------- #
def resolve(self, request_id: str, approve: bool, run: Executor) -> Optional[Dict]:
"""Confirm (``approve=True``) or reject (``approve=False``) a proposal.
On the FIRST resolution: runs ``run(model_key, switched)`` where
``model_key`` is the proposed model when approved, else the current
model; caches and returns its result. Subsequent resolutions of the
same id return the cached result **without** re-running (idempotent).
An already-EXPIRED proposal is forced down the reject path (run with the
current model) — matching "timeout → keep current model".
Returns ``None`` if ``request_id`` is unknown.
"""
with self._lock:
ps = self._items.get(request_id)
if ps is None:
return None
self._maybe_expire_locked(ps)
event = self._events[request_id]
# Already executed → idempotent replay, no matter who asks.
if ps.result is not None:
return ps.result
expired = ps.status == SwitchStatus.EXPIRED
effective_approve = bool(approve) and not expired
# First caller to arrive wins the right to execute exactly once.
i_run = request_id not in self._running
if i_run:
self._running.add(request_id)
ps.status = (
SwitchStatus.CONFIRMED if effective_approve else SwitchStatus.REJECTED
)
if not i_run:
# Another thread is executing — wait for it, then replay its result.
event.wait(timeout=self._RESOLVE_WAIT_SEC)
with self._lock:
return self._items[request_id].result
# Execute outside the lock (the network/LLM call may be slow).
decision = ps.decision
model_key = decision.to_model if effective_approve else decision.from_model
try:
result = run(model_key or "", bool(effective_approve))
finally:
with self._lock:
self._running.discard(request_id)
with self._lock:
ps.result = result
event.set()
return result
# -- maintenance ---------------------------------------------------- #
def sweep_expired(self) -> List[str]:
"""Mark all overdue PENDING proposals EXPIRED; return their ids."""
expired: List[str] = []
with self._lock:
for rid, ps in self._items.items():
if ps.status == SwitchStatus.PENDING and self._clock() >= ps.expires_at:
ps.status = SwitchStatus.EXPIRED
expired.append(rid)
return expired
def purge(self, keep_resolved: bool = False) -> int:
"""Drop resolved/expired entries to free memory. Returns count removed.
With ``keep_resolved=True``, entries that carry a cached ``result`` are
retained so their idempotent replay still works.
"""
removed = 0
with self._lock:
for rid in list(self._items):
ps = self._items[rid]
terminal = ps.status in (
SwitchStatus.CONFIRMED, SwitchStatus.REJECTED, SwitchStatus.EXPIRED
)
if terminal and not (keep_resolved and ps.result is not None):
self._items.pop(rid, None)
self._events.pop(rid, None)
self._running.discard(rid)
removed += 1
return removed
def pending_ids(self) -> List[str]:
with self._lock:
return [
rid for rid, ps in self._items.items()
if ps.status == SwitchStatus.PENDING
]
__all__ = ["decide", "PendingSwitchRegistry", "Executor"]