aleksey
a2d238d92e
Make Zed Agent closer to Cursor: deterministic DevOps path index, live Traefik port probe before blind edits, stop-after-edit, and quieter Russian progress.
319 lines
11 KiB
Python
319 lines
11 KiB
Python
"""Lane orchestration A/B/C with Redis session context."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import hashlib
|
|
import json
|
|
import logging
|
|
import os
|
|
import re
|
|
from dataclasses import dataclass, field
|
|
from enum import Enum
|
|
from typing import Any
|
|
|
|
from classifier_llm import LlmClassifier
|
|
from rules_loader import load_model_matrix, load_orchestration, load_routing_rules
|
|
|
|
log = logging.getLogger("orchestrator")
|
|
|
|
LANE_PREFIX = {"A": "a", "B": "b", "C": "c"}
|
|
TIER_SUFFIX = {
|
|
"SIMPLE": "simple",
|
|
"MEDIUM_OPS": "medium-ops",
|
|
"MEDIUM_CODE": "medium-code",
|
|
"COMPLEX": "complex",
|
|
"REASONING": "reasoning",
|
|
"VISION_OCR": "vision-ocr",
|
|
"VISION_UI": "vision",
|
|
}
|
|
|
|
|
|
class Tier(str, Enum):
|
|
SIMPLE = "SIMPLE"
|
|
MEDIUM_OPS = "MEDIUM_OPS"
|
|
MEDIUM_CODE = "MEDIUM_CODE"
|
|
COMPLEX = "COMPLEX"
|
|
REASONING = "REASONING"
|
|
VISION_OCR = "VISION_OCR"
|
|
VISION_UI = "VISION_UI"
|
|
|
|
|
|
@dataclass
|
|
class SessionCtx:
|
|
escalation_level: int = 0
|
|
last_prompt_hash: str = ""
|
|
turn_count: int = 0
|
|
budget_pct: float = 0.0
|
|
|
|
|
|
@dataclass
|
|
class RouteDecision:
|
|
tier: Tier
|
|
lane: str
|
|
model: str
|
|
quality_mode: str
|
|
escalation_level: int
|
|
confidence: float
|
|
delegated_internal: bool = False
|
|
classifier_source: str = "heuristic"
|
|
|
|
|
|
class SessionStore:
|
|
def __init__(self, redis_url: str | None) -> None:
|
|
self._redis = None
|
|
if redis_url:
|
|
try:
|
|
import redis
|
|
|
|
self._redis = redis.from_url(redis_url, decode_responses=True)
|
|
self._redis.ping()
|
|
except Exception as exc: # noqa: BLE001
|
|
log.warning("Redis unavailable (%s), using in-memory sessions", exc)
|
|
self._memory: dict[str, SessionCtx] = {}
|
|
orch = load_orchestration()
|
|
self._prefix = orch.get("redis", {}).get("key_prefix", "ai-router:session:")
|
|
self._ttl = int(orch.get("redis", {}).get("ttl_sec", 1800))
|
|
|
|
def get(self, session_id: str) -> SessionCtx:
|
|
if not session_id:
|
|
return SessionCtx()
|
|
key = f"{self._prefix}{session_id}"
|
|
if self._redis:
|
|
raw = self._redis.get(key)
|
|
if raw:
|
|
data = json.loads(raw)
|
|
return SessionCtx(**data)
|
|
return SessionCtx()
|
|
return self._memory.get(session_id, SessionCtx())
|
|
|
|
def save(self, session_id: str, ctx: SessionCtx) -> None:
|
|
if not session_id:
|
|
return
|
|
key = f"{self._prefix}{session_id}"
|
|
payload = json.dumps(
|
|
{
|
|
"escalation_level": ctx.escalation_level,
|
|
"last_prompt_hash": ctx.last_prompt_hash,
|
|
"turn_count": ctx.turn_count,
|
|
"budget_pct": ctx.budget_pct,
|
|
}
|
|
)
|
|
if self._redis:
|
|
self._redis.setex(key, self._ttl, payload)
|
|
else:
|
|
self._memory[session_id] = ctx
|
|
|
|
|
|
class Classifier:
|
|
def __init__(self) -> None:
|
|
rules = load_routing_rules()
|
|
gw = rules.get("gateway", {})
|
|
self._simple_re = re.compile(
|
|
gw.get("simple_patterns", [r"^(hello)\\b"])[0],
|
|
re.IGNORECASE,
|
|
)
|
|
self._complex = self._kw_re(gw.get("complex_keywords", []))
|
|
self._reasoning = self._kw_re(gw.get("reasoning_keywords", []))
|
|
self._medium_ops = self._kw_re(gw.get("medium_ops_keywords", []))
|
|
self._medium_code = self._kw_re(gw.get("medium_code_keywords", []))
|
|
self._ocr = self._kw_re(gw.get("ocr_keywords", []))
|
|
self._escalation = self._kw_re(gw.get("escalation_keywords", []))
|
|
wt = gw.get("word_thresholds", {})
|
|
self._simple_max = int(wt.get("simple_max_words", 12))
|
|
self._simple_q_max = int(wt.get("simple_question_max_words", 25))
|
|
self._complex_min = int(wt.get("complex_min_words", 400))
|
|
self._low_conf = float(rules.get("gateway", {}).get("confidence", {}).get("low_threshold", 0.6))
|
|
self._code_block = re.compile(r"```[\s\S]*?```|`[^`]+`")
|
|
|
|
@staticmethod
|
|
def _kw_re(keywords: list[str]) -> re.Pattern[str]:
|
|
if not keywords:
|
|
return re.compile(r"(?!x)x")
|
|
escaped = [re.escape(k) for k in keywords]
|
|
return re.compile("|".join(escaped), re.IGNORECASE)
|
|
|
|
def wants_escalation(self, text: str, prompt_hash: str, prev_hash: str) -> bool:
|
|
if self._escalation.search(text):
|
|
return True
|
|
return bool(prompt_hash and prompt_hash == prev_hash)
|
|
|
|
def classify(
|
|
self,
|
|
messages: list[dict[str, Any]],
|
|
*,
|
|
has_image: bool,
|
|
text: str,
|
|
) -> tuple[Tier, float]:
|
|
if has_image:
|
|
if self._ocr.search(text) or len(text.split()) < 30:
|
|
return Tier.VISION_OCR, 0.95
|
|
return Tier.VISION_UI, 0.9
|
|
|
|
text = text.strip()
|
|
if not text:
|
|
return Tier.SIMPLE, 0.9
|
|
|
|
words = len(text.split())
|
|
if words <= self._simple_max and self._simple_re.search(text):
|
|
return Tier.SIMPLE, 0.92
|
|
if self._reasoning.search(text):
|
|
return Tier.REASONING, 0.88
|
|
if self._complex.search(text):
|
|
return Tier.COMPLEX, 0.88
|
|
if self._code_block.search(text) or self._medium_code.search(text):
|
|
return Tier.MEDIUM_CODE, 0.85
|
|
if self._medium_ops.search(text):
|
|
return Tier.MEDIUM_OPS, 0.82
|
|
if words <= self._simple_q_max and "?" in text and not self._code_block.search(text):
|
|
return Tier.SIMPLE, 0.75
|
|
if words >= self._complex_min:
|
|
return Tier.COMPLEX, 0.7
|
|
return Tier.MEDIUM_OPS, 0.55
|
|
|
|
|
|
class Orchestrator:
|
|
def __init__(self, session_store: SessionStore) -> None:
|
|
self.sessions = session_store
|
|
self.classifier = Classifier()
|
|
self.llm_classifier = LlmClassifier()
|
|
self._orch = load_orchestration()
|
|
self._matrix = load_model_matrix()
|
|
self._models: dict[str, dict] = self._matrix.get("models", {})
|
|
self._no_esc = {
|
|
name for name, cfg in self._models.items() if cfg.get("no_escalation")
|
|
}
|
|
self.default_quality = os.environ.get(
|
|
"DEFAULT_QUALITY_MODE",
|
|
self._orch.get("default_quality_mode", "auto"),
|
|
)
|
|
|
|
def _lane_from_mode(self, quality_mode: str) -> str | None:
|
|
if quality_mode == "auto":
|
|
return None
|
|
return self._orch.get("quality_mode_map", {}).get(quality_mode)
|
|
|
|
def _start_lane(self, tier: Tier) -> str:
|
|
return self._orch.get("start_lanes", {}).get(tier.value, "A")
|
|
|
|
def _bump_lane(self, lane: str, levels: int) -> str:
|
|
order: list[str] = self._orch.get("lane_order", ["A", "B", "C"])
|
|
try:
|
|
idx = order.index(lane)
|
|
except ValueError:
|
|
return lane
|
|
return order[min(idx + levels, len(order) - 1)]
|
|
|
|
def _cap_lane(self, lane: str, budget_pct: float) -> str:
|
|
caps = self._orch.get("budget_caps", {})
|
|
hard = float(caps.get("hard_pct", 95))
|
|
warn = float(caps.get("warn_pct", 80))
|
|
if budget_pct >= hard:
|
|
return caps.get("hard_max_lane", "A")
|
|
if budget_pct >= warn:
|
|
max_lane = caps.get("warn_max_lane", "B")
|
|
order = self._orch.get("lane_order", ["A", "B", "C"])
|
|
if order.index(lane) > order.index(max_lane):
|
|
return max_lane
|
|
return lane
|
|
|
|
def _model_name(self, lane: str, tier: Tier) -> str:
|
|
prefix = LANE_PREFIX.get(lane, "a")
|
|
suffix = TIER_SUFFIX[tier]
|
|
return f"{prefix}-{suffix}"
|
|
|
|
async def resolve(
|
|
self,
|
|
messages: list[dict[str, Any]],
|
|
*,
|
|
quality_mode: str | None = None,
|
|
session_id: str = "",
|
|
text: str = "",
|
|
has_image: bool = False,
|
|
token_estimate: int = 0,
|
|
) -> RouteDecision:
|
|
mode = quality_mode or self.default_quality
|
|
tier, confidence = self.classifier.classify(messages, has_image=has_image, text=text)
|
|
classifier_source = "heuristic"
|
|
|
|
if self.llm_classifier.should_use_llm(tier.value, confidence, has_image):
|
|
llm_result = await self.llm_classifier.classify(text)
|
|
if llm_result:
|
|
tier = Tier(llm_result[0])
|
|
confidence = llm_result[1]
|
|
classifier_source = "gigachat"
|
|
elif self.llm_classifier.mode == "llm":
|
|
classifier_source = "heuristic_fallback"
|
|
|
|
ctx = self.sessions.get(session_id)
|
|
prompt_hash = hashlib.sha256(text.encode()).hexdigest()[:16]
|
|
|
|
if mode != "auto":
|
|
lane = self._lane_from_mode(mode) or "A"
|
|
escalation = 0
|
|
else:
|
|
lane = self._start_lane(tier)
|
|
escalation = ctx.escalation_level
|
|
if tier != Tier.VISION_OCR and self.classifier.wants_escalation(
|
|
text, prompt_hash, ctx.last_prompt_hash
|
|
):
|
|
escalation += 1
|
|
min_tokens = int(
|
|
self._orch.get("escalation", {}).get("context_tokens_min_lane_b", 32000)
|
|
)
|
|
if token_estimate > min_tokens:
|
|
order = self._orch.get("lane_order", ["A", "B", "C"])
|
|
if order.index(lane) < order.index("B"):
|
|
lane = "B"
|
|
lane = self._bump_lane(lane, escalation)
|
|
lane = self._cap_lane(lane, ctx.budget_pct)
|
|
|
|
model = self._model_name(lane, tier)
|
|
if model not in self._models:
|
|
log.warning("model %s missing from matrix, fallback smart-router-internal", model)
|
|
return RouteDecision(
|
|
tier=tier,
|
|
lane=lane,
|
|
model="smart-router-internal",
|
|
quality_mode=mode,
|
|
escalation_level=escalation,
|
|
confidence=confidence,
|
|
delegated_internal=True,
|
|
classifier_source=classifier_source,
|
|
)
|
|
|
|
delegated = (
|
|
classifier_source == "heuristic"
|
|
and confidence < self.classifier._low_conf
|
|
and tier in (Tier.MEDIUM_OPS, Tier.SIMPLE)
|
|
)
|
|
|
|
return RouteDecision(
|
|
tier=tier,
|
|
lane=lane,
|
|
model="smart-router-internal" if delegated else model,
|
|
quality_mode=mode,
|
|
escalation_level=escalation,
|
|
confidence=confidence,
|
|
delegated_internal=delegated,
|
|
classifier_source=classifier_source,
|
|
)
|
|
|
|
def after_request(
|
|
self,
|
|
session_id: str,
|
|
*,
|
|
prompt_hash: str,
|
|
success: bool,
|
|
escalate: bool,
|
|
) -> None:
|
|
if not session_id:
|
|
return
|
|
ctx = self.sessions.get(session_id)
|
|
ctx.last_prompt_hash = prompt_hash
|
|
ctx.turn_count += 1
|
|
if success and not escalate:
|
|
ctx.escalation_level = 0
|
|
elif escalate:
|
|
ctx.escalation_level = min(ctx.escalation_level + 1, 2)
|
|
self.sessions.save(session_id, ctx)
|