This commit is contained in:
@@ -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).
|
||||
@@ -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",
|
||||
]
|
||||
@@ -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]
|
||||
@@ -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"]
|
||||
@@ -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"]
|
||||
@@ -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",
|
||||
]
|
||||
@@ -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... }, ... }
|
||||
@@ -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"]
|
||||
@@ -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",
|
||||
]
|
||||
@@ -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"]
|
||||
@@ -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"]
|
||||
@@ -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"]
|
||||
@@ -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"]
|
||||
@@ -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"]
|
||||
@@ -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"]
|
||||
Reference in New Issue
Block a user