Spaces:
Sleeping
Sleeping
File size: 8,227 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 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 | """ContextManager — the small public facade that wires the whole engine.
ctx = ContextManager.create("my-agent")
ctx.remember("The user prefers local-first tools", expert="profile", importance=0.9)
pack = ctx.build_pack("How should I design this agent?", max_tokens=400)
print(pack.to_prompt())
print(ctx.inspect("How should I design this agent?"))
"""
from __future__ import annotations
from typing import Optional
from .config import Config
from .context.assembler import assemble_pack
from .embedding.base import Embedder
from .embedding.hashing import HashingEmbedder
from .retrieval.fusion import hybrid_retrieve
from .routing.router import ContextRouter
from .schema.item import ContextItem
from .schema.pack import ContextPack
from .schema.query import RecallQuery
from .store.sqlite import SqliteStore
def _item_dict(item: ContextItem) -> dict:
"""JSON-serializable view of a stored item (embedding omitted)."""
return {
"id": item.id,
"content": item.content,
"expert": item.expert,
"scope": item.scope,
"importance": item.importance,
"tags": list(item.tags),
"created_at": item.created_at,
"ttl": item.ttl,
}
def _packed_dict(packed) -> dict:
"""JSON-serializable view of a kept pack item with its score breakdown."""
d = _item_dict(packed.item)
d["final_score"] = packed.final_score
d["breakdown"] = packed.breakdown
return d
def _make_embedder(name: str) -> Embedder:
if name in ("hashing", "", None):
return HashingEmbedder()
if name in ("sentence-transformers", "st"):
from .embedding.sentence_transformers import SentenceTransformerEmbedder
return SentenceTransformerEmbedder()
raise ValueError(f"unknown embedder: {name}")
class ContextManager:
def __init__(self, store: SqliteStore, router: ContextRouter, embedder: Embedder,
config: Optional[Config] = None):
self.store, self.router, self.embedder = store, router, embedder
self.config = config or Config()
@classmethod
def create(cls, name: str = "default", path: Optional[str] = None,
embedder: Optional[Embedder] = None) -> "ContextManager":
embedder = embedder or HashingEmbedder()
store = SqliteStore(path or f"{name}.matrix-context.db", embedder)
return cls(store, ContextRouter(embedder), embedder, Config(name=name))
@classmethod
def from_env(cls) -> "ContextManager":
cfg = Config.from_env()
emb = _make_embedder(cfg.embedder)
store = SqliteStore(cfg.path or f"{cfg.name}.matrix-context.db", emb)
return cls(store, ContextRouter(emb), emb, cfg)
def remember(self, content: str, expert: str = "semantic", scope: str = "/",
importance: float = 0.5, tags=(), ttl: Optional[float] = None) -> ContextItem:
return self.store.add(ContextItem(content=content, expert=expert, scope=scope,
importance=importance, tags=tuple(tags), ttl=ttl))
# Default expert fan-out. The bake-off (embedder=sentence-transformers,
# store=memory) measured moc_rag winning at top_experts=2 — fewer distractors
# and tokens at equal recall — so 2 is the promoted engine default.
DEFAULT_TOP_EXPERTS = 2
def items(self, scope: Optional[str] = None,
expert: Optional[str] = None) -> list:
"""List stored items, optionally filtered by scope prefix and/or expert."""
out = self.store.all_items()
if scope and scope.rstrip("/"):
pref = scope.rstrip("/")
out = [it for it in out
if it.scope == scope or it.scope.startswith(pref + "/")]
if expert:
out = [it for it in out if it.expert == expert]
return out
def forget(self, item_id: str) -> bool:
"""Delete an item by id. Returns True if it existed."""
existed = self.store.get(item_id) is not None
self.store.delete(item_id)
return existed
def _route_and_score(self, query: str, scope: str, top_experts: int,
pin_experts: tuple = ()):
decision = self.router.route(query, self.store.all_items(), top_experts)
# Pinned experts are always injectable (e.g. profile), regardless of the
# routing decision — appended without disturbing the ranked order.
selected = list(decision.selected)
for e in pin_experts:
if e not in selected:
selected.append(e)
decision.selected = selected
cands = self.store.candidates(selected, scope)
by_id = {it.id: it for it in cands}
scores = hybrid_retrieve(query, self.embedder.encode(query), cands) if cands else {}
return decision, scores, by_id
def build_pack(self, query: str, scope: str = "/", top_experts: int = DEFAULT_TOP_EXPERTS,
max_tokens: int = 600, pin_experts: tuple = ()) -> ContextPack:
decision, scores, by_id = self._route_and_score(query, scope, top_experts,
pin_experts)
return assemble_pack(scores, by_id, decision.selected, decision.reason,
max_tokens=max_tokens)
def recall(self, query: RecallQuery) -> ContextPack:
return self.build_pack(query.text, scope=query.scopes[0],
top_experts=query.top_experts, max_tokens=query.max_tokens)
def build_inspection(self, query: str, scope: str = "/",
top_experts: int = DEFAULT_TOP_EXPERTS,
max_tokens: int = 600, pin_experts: tuple = ()) -> dict:
"""Structured, JSON-serializable explanation of a routed pack.
This is the single source of truth behind both the human-readable
``inspect()`` string and the REST ``POST /v1/inspect`` contract: routing
scores, selected vs. unselected experts, kept and dropped items with
their score breakdown, and the final prompt-ready pack.
"""
decision, scores, by_id = self._route_and_score(query, scope, top_experts,
pin_experts)
pack = assemble_pack(scores, by_id, decision.selected, decision.reason,
max_tokens=max_tokens)
selected = list(decision.selected)
unselected = [e for e in sorted(decision.scores) if e not in selected]
return {
"query": query,
"routing": {
"selected_experts": selected,
"unselected_experts": unselected,
"scores": {e: round(s, 4) for e, s in decision.scores.items()},
"widened": decision.widened,
"reason": decision.reason,
},
"pack": {
"tokens": pack.tokens,
"max_tokens": max_tokens,
"selected_experts": pack.selected_experts,
"routing_reason": pack.routing_reason,
"items": [_packed_dict(p) for p in pack.items],
"dropped": list(pack.dropped),
"citations": pack.citations,
"prompt": pack.to_prompt(),
},
}
def inspect(self, query: str, scope: str = "/", top_experts: int = DEFAULT_TOP_EXPERTS,
max_tokens: int = 600, pin_experts: tuple = ()) -> str:
ins = self.build_inspection(query, scope, top_experts, max_tokens, pin_experts)
r, pk = ins["routing"], ins["pack"]
lines = [f"ROUTING: {r['reason']}",
" scores: " + ", ".join(f"{e}={s:.3f}" for e, s in
sorted(r["scores"].items(), key=lambda x: -x[1])),
f" selected experts: {r['selected_experts']}",
f"PACK ({pk['tokens']} tokens, {len(pk['items'])} items):"]
for p in pk["items"]:
lines.append(f" [{p['expert']}] score={p['final_score']} {p['breakdown']} "
f":: {p['content'][:60]}")
for d in pk["dropped"]:
lines.append(f" DROPPED [{d['expert']}] {d['reason']}")
return "\n".join(lines)
|