Spaces:
Sleeping
Sleeping
File size: 2,897 Bytes
ce45eb0 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 | """Two-tier context router — the centerpiece (cf. MODE / ExpertRAG / MixRAG).
Tier 1: fast centroid gate. Each expert has a running centroid of its items'
embeddings blended with a description embedding; score = cosine(query, centroid).
Tier 2: ambiguity fallback. On low confidence or an indecisive margin, WIDEN
selection rather than guess. This is the exact seam where a v1 LLM classifier
plugs in (see llm_gate.py). Every decision is returned with scores + reason so
`inspect()` can explain why an expert fired.
"""
from __future__ import annotations
from dataclasses import dataclass
from typing import Dict, List, Optional
import numpy as np
from ..embedding.base import Embedder, cosine
from ..schema.item import ContextItem
from .experts import EXPERT_DESCRIPTIONS
@dataclass
class RoutingDecision:
selected: List[str]
scores: Dict[str, float]
widened: bool
reason: str
class ContextRouter:
def __init__(self, embedder: Embedder,
confident: float = 0.18, decisive_gap: float = 0.05):
self.embedder = embedder
self.confident = confident
self.decisive_gap = decisive_gap
self._desc = {e: embedder.encode(d) for e, d in EXPERT_DESCRIPTIONS.items()}
def _centroid(self, expert: str, items: List[ContextItem]) -> Optional[np.ndarray]:
embs = [it.embedding for it in items
if it.expert == expert and it.embedding is not None]
base = self._desc.get(expert)
if not embs:
return base
c = np.mean(embs, axis=0)
if base is not None:
c = 0.7 * c + 0.3 * base
n = np.linalg.norm(c)
return c / n if n > 0 else c
def route(self, query: str, items: List[ContextItem],
top_experts: int = 3) -> RoutingDecision:
q = self.embedder.encode(query)
experts = sorted(set(list(EXPERT_DESCRIPTIONS) + [it.expert for it in items]))
scores = {e: (cosine(q, c) if (c := self._centroid(e, items)) is not None else 0.0)
for e in experts}
ranked = sorted(scores.items(), key=lambda x: x[1], reverse=True)
top = ranked[:top_experts]
top_score = top[0][1] if top else 0.0
gap = (top[0][1] - top[1][1]) if len(top) > 1 else 1.0
if top_score < self.confident or gap < self.decisive_gap:
widened = ranked[:top_experts + 2]
return RoutingDecision(
selected=[e for e, _ in widened], scores=scores, widened=True,
reason=(f"ambiguous (top={top_score:.3f} < {self.confident} or "
f"gap={gap:.3f} < {self.decisive_gap}); widened — v1 "
f"would invoke the LLM gate here"))
return RoutingDecision(
selected=[e for e, _ in top], scores=scores, widened=False,
reason=f"confident: top={top_score:.3f}, gap={gap:.3f}")
|