diff --git a/.gitignore b/.gitignore index 93a9d64..0e6d194 100644 --- a/.gitignore +++ b/.gitignore @@ -24,3 +24,4 @@ venv/ .sidekick/ .sidekick-*/ .financebench/ +.medical/ diff --git a/BENCHMARKS.md b/BENCHMARKS.md index ba1d6cf..74f8d33 100644 --- a/BENCHMARKS.md +++ b/BENCHMARKS.md @@ -32,6 +32,12 @@ abstention that v1 lacks, and a ~62% depth cut with precision rising to 100%._ ## Retrieval over heterogeneous personal data (financial × medical) +> **Now live** as the **"Context Runtime · Mixed"** model on chat.redevops.io — the same +> coverage-routed sharding, served from `CR_TENANTS=…mixed=shards(finance:/corpus,medical:/medical)`. +> The live demo's medical shard is the public **PubMedQA** corpus (`deploy/medical/`, ~3.3k +> passages); the measurement below used the smaller **16-note curated collision set** from +> `examples/heterogeneous_shards.py` (built to maximize the vocabulary overlap the result turns on). + **`examples/heterogeneous_shards.py`** — the interesting, non-obvious one. A real user's local files are a mix of very different data. This runs against **real diff --git a/benchmarks.html b/benchmarks.html new file mode 100644 index 0000000..1cc5692 --- /dev/null +++ b/benchmarks.html @@ -0,0 +1,21 @@ +Context Runtime — v1 vs v2 benchmarks + +
+

Context Runtime — v1 → v2, measured in both runtimes

+

The same seeded, ground-truth retrieval simulation, run in the Python source-of-truth and the Go port. v2 = calibrated relevance-in-reward + abstention + the DSpark load-aware sizer.

+
MetricPy v1Py v2ΔGo v1Go v2Δ
Learned-policy precisionThe served passages that are actually relevant, after the policy converges. v2's reward finally sees calibrated relevance.67.6%82.2%▲ +14.6 pts84.6%95.9%▲ +11.3 pts
Abstention recall (unanswerable caught)Share of truly-unanswerable queries v2 declines to answer. v1 has no abstention at all.0.0%100.0%▲ +100.0 pts0.0%100.0%▲ +100.0 pts
False-abstain rate (answerable dropped)Answerable queries v2 wrongly declined — the cost of abstention. Lower is better.0.0%0.0%—0.0%0.0%—
Expensive-stage depth (passages)Passages sent to the costly rerank/synthesis stage from a deep k=8 arm. The sizer prunes the low-relevance tail.8.003.00▼ −62%8.003.00▼ −63%
Precision after the sizerPrecision of what survives the sizer's gate — pruning the tail raises it.37.5%100.0%▲ +62.5 pts37.5%100.0%▲ +62.5 pts
+

40-seed average · precision headlined at β=0.9 (the calibration-trust knob; shipped default 0.5) · Go is an independent re-implementation on identical methodology — directional parity across languages.

+
\ No newline at end of file diff --git a/context_runtime/abstention.py b/context_runtime/abstention.py new file mode 100644 index 0000000..9950958 --- /dev/null +++ b/context_runtime/abstention.py @@ -0,0 +1,68 @@ +"""Honest abstention — Generation 5 (Whitepaper v3, Trust-Aware Execution). + +Production AI systems fail when operators stop trusting them, and the fastest way to burn trust is a +confident, wrong answer. So a trust-aware planner must be willing to **decline**: when the best available +plan's calibrated confidence is below the bar, abstain (or escalate to a stronger tier) rather than serve. +Abstaining beats guessing — an honest "I don't know" protects trust, and the TrustLedger credits it. + +This gate is optimizer-agnostic: it reads a ``PlanScore`` (the estimator's plan-time confidence, optionally +run through a calibration map so the number means P(correct), not a raw score) and returns a verdict. The +online optimizer records the verdict in ``Plan.extra["abstention"]``; a consumer serves, escalates, or +declines accordingly. +""" +from __future__ import annotations + +from dataclasses import dataclass +from typing import Callable + +from .types import Goal, PlanScore + + +@dataclass(frozen=True) +class AbstentionVerdict: + action: str # "serve" | "escalate" | "abstain" + confidence: float # calibrated confidence of the evaluated plan, in [0, 1] + reason: str + + @property + def abstained(self) -> bool: + return self.action == "abstain" + + +class AbstentionGate: + """Decide whether a plan is confident enough to serve. + + ``min_confidence`` — the bar; below it we do not serve. + ``calibrate`` — optional map from the raw plan-time confidence to a calibrated P(correct) + (e.g. the DSpark ``CalibrationMap``); identity if omitted. + ``can_escalate`` — optional predicate ``(score, goal) -> bool``: when confidence is below the bar + but a stronger option exists, prefer escalation over a flat abstention. + """ + + def __init__( + self, + min_confidence: float = 0.5, + *, + calibrate: Callable[[float], float] | None = None, + can_escalate: Callable[[PlanScore, Goal | None], bool] | None = None, + ): + self.min_confidence = min_confidence + self.calibrate = calibrate + self.can_escalate = can_escalate + + def confidence(self, score: PlanScore) -> float: + """Plan-time confidence = the estimator's expected accuracy, calibrated if a map is provided.""" + raw = score.expected_accuracy + return self.calibrate(raw) if self.calibrate else raw + + def evaluate(self, score: PlanScore, goal: Goal | None = None) -> AbstentionVerdict: + c = self.confidence(score) + if c >= self.min_confidence: + return AbstentionVerdict("serve", c, f"confidence {c:.2f} ≥ {self.min_confidence:.2f}") + if self.can_escalate is not None and self.can_escalate(score, goal): + return AbstentionVerdict( + "escalate", c, f"confidence {c:.2f} < {self.min_confidence:.2f} — escalate to a stronger tier" + ) + return AbstentionVerdict( + "abstain", c, f"confidence {c:.2f} < {self.min_confidence:.2f} — abstain (honest, protects trust)" + ) diff --git a/context_runtime/adapters/model_openai.py b/context_runtime/adapters/model_openai.py new file mode 100644 index 0000000..32b6d2f --- /dev/null +++ b/context_runtime/adapters/model_openai.py @@ -0,0 +1,109 @@ +"""OpenAICompatibleModel — a dependency-light ``ModelPlugin`` (SPEC §4.3). + +Speaks the OpenAI ``/chat/completions`` wire format over the *stdlib* (urllib), so it +drives OpenAI gpt-5.x, DeepSeek, vLLM, Ollama — any compatible endpoint — with the same +``ModelRequest → ModelResult`` contract and cost-tiered ``Tier`` routing as +``LiteLLMModel``, but **without the litellm dependency**. That matters because the slim +agent containers install only the base ``context-runtime`` (no ``[litellm]`` extra): this +adapter lets those apps run their conversational agent *on the Context Runtime model +plane* rather than reaching around it to a raw provider SDK. Degrade to ``StubModel`` +when no key is configured (``from_env`` returns ``None``). +""" +from __future__ import annotations + +import json +import os +import urllib.error +import urllib.request + +from ..types import ModelCapabilities, ModelRequest, ModelResult, PluginInfo +from .model_litellm import Tier # dataclass only; importing it does not pull in litellm + + +class OpenAICompatibleModel: + """ModelPlugin over an OpenAI-compatible endpoint using only the stdlib.""" + + def __init__(self, tiers: list[Tier], default_tier: str = "chat", timeout: float = 40.0): + self.tiers = {t.name: t for t in tiers} + self.default_tier = default_tier + self.timeout = timeout + + @classmethod + def from_env( + cls, + *, + model_env: str = "AGENT_LLM_MODEL", + default_model: str = "grok-4.5", + key_envs: tuple[str, ...] = ("AGENT_LLM_KEY", "OPENAI_API_KEY"), + base_envs: tuple[str, ...] = ("AGENT_LLM_BASE_URL", "OPENAI_BASE_URL"), + cost_per_1k: float = 0.0, + ) -> "OpenAICompatibleModel | None": + """Build from environment, or ``None`` when nothing is configured (offline → StubModel). + + Priority: an explicit OpenAI/agent key, else the self-hosted OpenAI-compatible endpoint + the agentic-os apps already point at (``REDEVOPS_LLM_BASE_URL`` / ``REDEVOPS_LLM_MODEL`` — + typically a keyless on-prem DeepSeek). That keeps the agent on our own model plane without + spreading a provider key across every container. + """ + key = next((os.environ[k] for k in key_envs if os.environ.get(k)), None) + if key: + base = next((os.environ[b] for b in base_envs if os.environ.get(b)), "https://api.openai.com/v1") + model = os.environ.get(model_env, default_model) + return cls([Tier(name="chat", model=model, base_url=base.rstrip("/"), api_key=key, cost_per_1k=cost_per_1k)]) + rbase = os.environ.get("REDEVOPS_LLM_BASE_URL") + if rbase: + rmodel = os.environ.get("REDEVOPS_LLM_MODEL", "DeepSeek-V4-Flash") + rkey = os.environ.get("REDEVOPS_LLM_KEY") or "sk-noauth" # vLLM ignores it; header must exist + return cls([Tier(name="chat", model=rmodel, base_url=rbase.rstrip("/"), api_key=rkey, cost_per_1k=cost_per_1k)]) + return None + + def _tier_for(self, capability: str) -> Tier: + return self.tiers.get(self.default_tier) or next(iter(self.tiers.values())) + + def complete(self, req: ModelRequest) -> ModelResult: + tier = self._tier_for(req.capability) + messages: list[dict] = [] + if req.system: + messages.append({"role": "system", "content": req.system}) + messages.extend(dict(m) for m in req.messages) + # gpt-5.x wants max_completion_tokens; self-hosted vLLM/DeepSeek want max_tokens. Branch on host. + tok_key = "max_completion_tokens" if "openai.com" in (tier.base_url or "") else "max_tokens" + payload: dict = {"model": tier.model, "messages": messages, tok_key: req.max_tokens} + if req.tools: + payload["tools"] = list(req.tools) + request = urllib.request.Request( + tier.base_url + "/chat/completions", + data=json.dumps(payload).encode(), + headers={"Authorization": "Bearer " + (tier.api_key or ""), "Content-Type": "application/json"}, + method="POST", + ) + try: + with urllib.request.urlopen(request, timeout=self.timeout) as resp: + d = json.load(resp) + except urllib.error.HTTPError as e: # surface provider errors so the caller can fall back + detail = e.read()[:300].decode("utf-8", "ignore") + raise RuntimeError(f"model http {e.code}: {detail}") from e + choice = (d.get("choices") or [{}])[0] + text = (choice.get("message") or {}).get("content") or "" + usage = d.get("usage") or {} + ptoks = int(usage.get("prompt_tokens") or self.count_tokens(json.dumps(messages), tier.model)) + ctoks = int(usage.get("completion_tokens") or self.count_tokens(text, tier.model)) + cost = (ptoks + ctoks) / 1000.0 * tier.cost_per_1k + return ModelResult( + text=text.strip(), + model=tier.model, + tier=tier.name, + prompt_tokens=ptoks, + completion_tokens=ctoks, + est_cost_usd=round(cost, 6), + models_used=(tier.model,), + ) + + def capabilities(self, model: str) -> ModelCapabilities: + return ModelCapabilities(max_context_tokens=128000, tool_calling=True, structured_outputs=True) + + def count_tokens(self, text: str, model: str) -> int: + return max(1, len(text) // 4) # ~4 chars/token, matches StubModel's estimate + + def info(self) -> PluginInfo: + return PluginInfo(name="openai_compatible", kind="model", capabilities=frozenset(self.tiers)) diff --git a/context_runtime/adapters/store_hipporag.py b/context_runtime/adapters/store_hipporag.py index f2cb50a..e11ccd3 100644 --- a/context_runtime/adapters/store_hipporag.py +++ b/context_runtime/adapters/store_hipporag.py @@ -97,7 +97,7 @@ class HippoRAGRetriever: """The real graph retriever — wraps redevops-io/HippoRAG. Lazy-imports so the core package and the offline path don't need its (heavy) deps.""" - def __init__(self, save_dir: str = ".context_runtime/hipporag", llm_model_name: str = "gpt-4o-mini", + def __init__(self, save_dir: str = ".context_runtime/hipporag", llm_model_name: str = "gpt-5-mini", embedding_model_name: str = "nvidia/NV-Embed-v2", source: str = "graph"): self.save_dir = save_dir self.llm_model_name = llm_model_name diff --git a/context_runtime/adapters/store_router.py b/context_runtime/adapters/store_router.py index 56d053e..4b41ab9 100644 --- a/context_runtime/adapters/store_router.py +++ b/context_runtime/adapters/store_router.py @@ -1,9 +1,15 @@ """HopRouterRetriever — dispatch single-hop vs multi-hop by retrieval method (SPEC §4.5). -The control plane's per-query decision made physical: ``method="graph"`` → the graph -retriever (HippoRAG / SimGraph); everything else → the single-hop retriever -(redevops-rag / in-memory). Itself a RetrieverPlugin, so the runtime holds ONE -retriever and the planner's method choice routes transparently underneath. +The control plane's per-query decision made physical, routing by knowledge representation: +``method="graph"`` → the graph retriever (HippoRAG / SimGraph); ``method="temporal"`` → the +bi-temporal store (Graphiti / in-memory TemporalStore); ``method="community"`` → the community +retriever; everything else → the single-hop retriever (redevops-rag / in-memory). Itself a +RetrieverPlugin, so the runtime holds ONE retriever and the planner's method choice routes +transparently underneath. + +The graph/community/temporal slots are OPTIONAL: an unwired slot (or, for temporal, a store +with no matching fact) falls back to single-hop retrieval, so adding a representation never +regresses default behavior — it only *adds* a route the planner can take once a backend is bound. """ from __future__ import annotations @@ -11,32 +17,40 @@ _GRAPH_METHODS = {"graph"} _COMMUNITY_METHODS = {"community"} +_TEMPORAL_METHODS = {"temporal"} class HopRouterRetriever: - def __init__(self, single_hop, graph, community=None): - self.single_hop = single_hop # RetrieverPlugin: bm25/vector/hybrid + def __init__(self, single_hop, graph, community=None, temporal=None): + self.single_hop = single_hop # RetrieverPlugin: bm25/vector/hybrid (document) self.graph = graph # RetrieverPlugin: graph/multi-hop self.community = community # RetrieverPlugin: community/global (optional) + self.temporal = temporal # RetrieverPlugin: bi-temporal facts (optional) def search(self, query: str, k: int, method: Retrieval = "hybrid") -> list[Hit]: if method in _COMMUNITY_METHODS and self.community is not None: return self.community.search(query, k, method) if method in _GRAPH_METHODS: return self.graph.search(query, k, method) + if method in _TEMPORAL_METHODS and self.temporal is not None: + hits = self.temporal.search(query, k, method) + if hits: # a populated bi-temporal store answered + return hits + # unpopulated store / no temporal fact matched → fall back to single-hop return self.single_hop.search(query, k, method) def index(self, path: str) -> dict: a = self.single_hop.index(path) if hasattr(self.single_hop, "index") else {} b = self.graph.index(path) if hasattr(self.graph, "index") else {} out = {"single_hop": a, "graph": b} - if self.community is not None and hasattr(self.community, "index"): - out["community"] = self.community.index(path) + for name, r in (("community", self.community), ("temporal", self.temporal)): + if r is not None and hasattr(r, "index"): + out[name] = r.index(path) return out def info(self) -> PluginInfo: caps = set() - for r in (self.single_hop, self.graph, self.community): + for r in (self.single_hop, self.graph, self.community, self.temporal): if r is not None and hasattr(r, "info"): caps |= set(r.info().capabilities) return PluginInfo(name="hop_router", kind="retriever", capabilities=frozenset(caps)) diff --git a/context_runtime/adapters/store_temporal.py b/context_runtime/adapters/store_temporal.py new file mode 100644 index 0000000..f7da8d8 --- /dev/null +++ b/context_runtime/adapters/store_temporal.py @@ -0,0 +1,125 @@ +"""Temporal / bi-temporal retrieval — "what changed, and when?" (Whitepaper v3, forthcoming retrieval). + +The next retrieval method: a **bi-temporal** store (Graphiti/Zep lineage) where every fact carries two +time axes — + + • **valid time** (``valid_from`` / ``valid_to``): when the fact is true in the world; + • **transaction time** (``recorded_at``): when the system learned it. + +That lets the planner answer point-in-time questions the other methods can't: the state *as of* a past +date, and *as it was known* at a past date (so a late-arriving correction doesn't rewrite history). It is +densest in conversational corpora (chat/tickets) where facts are constantly revised. + +Dependency-free (stdlib + the ``RetrieverPlugin`` seam), so it drops into the planner as one more +routable capability (``method="temporal"``); the full Graphiti engine can back it as an optional binding. +Times are compared as plain ISO strings (``"2026-01-15"``): ``""`` = the beginning of time, ``None`` on +``valid_to`` = still valid. +""" +from __future__ import annotations + +import re +from dataclasses import dataclass, field + +from ..types import Hit, PluginInfo, Retrieval + +_STOP = {"the", "a", "an", "of", "to", "in", "for", "is", "was", "did", "who", "what", "when", "how", + "owns", "own"} +# keep hyphen/underscore inside a token so entity ids ("auth-service") stay whole and don't collide +# with a shared fragment ("service") on unrelated subjects +_WORD = re.compile(r"[a-z0-9][a-z0-9_-]*") + + +def _tokens(text: str) -> set[str]: + return {t for t in _WORD.findall(text.lower()) if len(t) > 1 and t not in _STOP} + + +@dataclass +class TemporalFact: + subject: str + predicate: str + obj: str + valid_from: str = "" # inclusive; "" = -inf + valid_to: str | None = None # exclusive; None = still valid (+inf) + recorded_at: str = "" # transaction time; "" = -inf + meta: dict = field(default_factory=dict) + + def text(self) -> str: + return f"{self.subject} {self.predicate} {self.obj}" + + def valid_at(self, at: str | None) -> bool: + if at is None: # "current" = not yet superseded + return self.valid_to is None + return self.valid_from <= at and (self.valid_to is None or at < self.valid_to) + + def known_at(self, known: str | None) -> bool: + return known is None or self.recorded_at <= known + + +class TemporalStore: + """An in-memory bi-temporal fact store exposing the RetrieverPlugin seam.""" + + def __init__(self, facts: list[TemporalFact] | None = None): + self._facts: list[TemporalFact] = list(facts or []) + + def add(self, subject: str, predicate: str, obj: str, *, valid_from: str = "", + valid_to: str | None = None, recorded_at: str = "", meta: dict | None = None) -> "TemporalStore": + self._facts.append(TemporalFact(subject, predicate, obj, valid_from, valid_to, + recorded_at, meta or {})) + return self + + def _hit(self, f: TemporalFact, score: float) -> Hit: + return Hit( + chunk_id=f"{f.subject}:{f.predicate}:{f.valid_from or '-'}", + filename="temporal", + text=f.text(), + score=score, + created_at=f.valid_from or None, + source="temporal", + meta={"subject": f.subject, "predicate": f.predicate, "object": f.obj, + "valid_from": f.valid_from, "valid_to": f.valid_to, "recorded_at": f.recorded_at, + **f.meta}, + ) + + def _ranked(self, query: str, candidates: list[TemporalFact], k: int) -> list[Hit]: + q = _tokens(query) + scored: list[tuple[float, str, TemporalFact]] = [] + for f in candidates: + overlap = len(q & _tokens(f.text())) + if q and overlap == 0: + continue + score = overlap / (len(q) or 1) + scored.append((score, f.valid_from, f)) + # best match first; ties → most-recently-valid first + scored.sort(key=lambda t: (t[0], t[1]), reverse=True) + return [self._hit(f, s) for s, _, f in scored[:k]] + + # ── the RetrieverPlugin seam ── + def search(self, query: str, k: int = 5, method: Retrieval = "temporal") -> list[Hit]: + """Current state: facts not yet superseded (``valid_to is None``), ranked by query match.""" + return self._ranked(query, [f for f in self._facts if f.valid_at(None)], k) + + def as_of(self, query: str, k: int = 5, *, at: str, known_at: str | None = None) -> list[Hit]: + """Point-in-time query: facts valid at ``at`` (world time) and known at ``known_at`` (transaction + time). ``known_at`` lets you ask "what did we believe on date X", ignoring later corrections.""" + return self._ranked(query, [f for f in self._facts + if f.valid_at(at) and f.known_at(known_at)], k) + + def changes(self, query: str = "", *, since: str, until: str, k: int = 20) -> list[dict]: + """"What changed, and when?" — facts that began or ended validity in ``[since, until)``.""" + q = _tokens(query) + out: list[dict] = [] + for f in self._facts: + if q and not (q & _tokens(f.text())): + continue + if since <= f.valid_from < until: + out.append({"at": f.valid_from, "change": "began", "fact": f.text(), + "subject": f.subject, "predicate": f.predicate, "object": f.obj}) + if f.valid_to is not None and since <= f.valid_to < until: + out.append({"at": f.valid_to, "change": "ended", "fact": f.text(), + "subject": f.subject, "predicate": f.predicate, "object": f.obj}) + out.sort(key=lambda c: c["at"]) + return out[:k] + + def info(self) -> PluginInfo: + return PluginInfo(name="temporal", kind="retriever", + capabilities=frozenset({"temporal", "retrieval"})) diff --git a/context_runtime/control_plane/app.py b/context_runtime/control_plane/app.py index 67ad273..cb6e31d 100644 --- a/context_runtime/control_plane/app.py +++ b/context_runtime/control_plane/app.py @@ -1033,17 +1033,31 @@ def _already_in_script(text: str, lang: str) -> bool: return bool(rng and re.search(f"[{rng}]", text)) -def _query_expander(text: str) -> str: - if not _QUERY_LANGS: - return text - parts = [text] - for lang in _QUERY_LANGS: - if _already_in_script(text, lang): - continue # already in this language — no cross-language expansion needed - t = _translate_query(text, lang) - if t and t.lower() != text.lower(): - parts.append(t) - return " ".join(parts) +def _make_query_expander(langs: list[str]): + """A cross-language query expander bound to a specific target-language set: the query is + expanded with translated terms into each `lang` before retrieval (so e.g. a Spanish question + matches an English corpus). Returns None when `langs` is empty (no expansion). Per-tenant, so + a multi-corpus deployment can bridge different languages per model (CR_TENANT_LANGS).""" + langs = [l for l in (langs or []) if l] + if not langs: + return None + + def _expand(text: str) -> str: + parts = [text] + for lang in langs: + if _already_in_script(text, lang): + continue # already in this language — no cross-language expansion needed + t = _translate_query(text, lang) + if t and t.lower() != text.lower(): + parts.append(t) + return " ".join(parts) + + return _expand + + +# Backward-compatible global expander (used by the legacy single tenant): CR_QUERY_LANGS applies +# to every tenant unless a tenant overrides it via CR_TENANT_LANGS (parsed in the multi-tenant block). +_query_expander = _make_query_expander(_QUERY_LANGS) def _flag(name: str) -> bool: @@ -1102,6 +1116,158 @@ def _flag(name: str) -> bool: pass +# ── Multi-corpus tenants (opt-in via CR_TENANTS) ────────────────────────────────────────── +# One control plane can expose several corpus-scoped models. Each model id in the chat dropdown +# is its OWN Context Runtime tenant: own retriever, own learned policy/quality/calibration files, +# own corpus (or heterogeneous shard mix). The chat shim (/v1) and the Query Board (/librechat/*) +# both route by model id, so the transparency panel always reflects the corpus the visitor picked. +# +# CR_TENANTS="context-runtime-finance=/corpus | context-runtime-medical=/med +# | context-runtime-mixed=shards(finance:/corpus,medical:/med)" +# +# '|'-separated `id=spec`; a bare path ⇒ single-corpus tenant; `shards(a:/p1,b:/p2)` ⇒ a +# coverage-routed heterogeneous mix (the whitepaper's cross-domain-noise-22→0 win). Unset ⇒ the +# single legacy tenant `librechat` (id = CR_SHIM_MODEL_ID) built above, so existing deploys/tests +# are unaffected. + +def _parse_tenant_langs(spec: str) -> dict[str, list[str]]: + """CR_TENANT_LANGS → {tenant_id: [langs]} — per-tenant cross-language expansion. A query to + that model is expanded with translated terms into its langs before retrieval, so e.g. a Spanish + speaker can search an English corpus. Format: 'context-runtime-finance=es | other=ru,en' + ('|'-separated id=comma-langs). Overrides the global CR_QUERY_LANGS for the listed tenants.""" + out: dict[str, list[str]] = {} + for entry in spec.split("|"): + entry = entry.strip() + if not entry: + continue + tid, sep, langs = entry.partition("=") + tid = tid.strip() + if not sep or not tid: + continue + out[tid] = [l.strip() for l in langs.split(",") if l.strip()] + return out + + +_TENANT_LANGS = _parse_tenant_langs(os.getenv("CR_TENANT_LANGS", "")) + + +def _parse_tenants_spec(spec: str) -> list[tuple[str, str, str]]: + """CR_TENANTS → [(tenant_id, corpus_dir, shards_spec)]. Exactly one of corpus_dir/shards_spec + is set per entry (shards_spec is the text inside `shards(...)`).""" + out: list[tuple[str, str, str]] = [] + for entry in spec.split("|"): + entry = entry.strip() + if not entry: + continue + tid, sep, body = entry.partition("=") + tid, body = tid.strip(), body.strip() + if not sep or not tid: + continue + low = body.lower() + if low.startswith("shards(") and body.endswith(")"): + out.append((tid, "", body[len("shards("):-1].strip())) + else: + out.append((tid, body, "")) + return out + + +def _build_lc_retriever(shards_spec: str = ""): + """A fresh LibreChat retriever built the same way the single-tenant path builds `_lc_retriever` + (hybrid→HopRouter base, optional coverage-routed shards, optional multimodal wrap) — one per + tenant so each corpus is indexed independently.""" + retr = None + if os.getenv("CR_EMBEDDINGS") == "1": + try: + from ..adapters.store_semantic import HybridRetriever, embeddings_available + if embeddings_available(): + retr = HybridRetriever() + except Exception: + retr = None + if retr is not None: + try: + from ..adapters.store_community import CommunityRetriever + from ..adapters.store_hipporag import SimGraphRetriever + from ..adapters.store_router import HopRouterRetriever + retr = HopRouterRetriever(single_hop=retr, graph=SimGraphRetriever([]), + community=CommunityRetriever([])) + except Exception: + pass + if shards_spec: + from ..adapters.store_inmemory import InMemoryStore + from ..scheduler.parallel_fusion import ShardedRetriever + shards = [] + for _spec in shards_spec.split(","): + _spec = _spec.strip() + if not _spec: + continue + _name, _sep, _path = _spec.partition(":") + if not _sep: # bare dir → derive a source name from the folder + _name, _path = (os.path.basename(_name.rstrip("/")) or "shard"), _name + if os.path.isdir(_path): + try: + _sh = InMemoryStore([], source=_name) + _sh.index(_path) + shards.append(_sh) + except Exception: + pass + if shards: + retr = ShardedRetriever(shards, router=os.getenv("CR_SHARD_ROUTER", "coverage"), + route_margin=float(os.getenv("CR_ROUTE_MARGIN", "0.15"))) + if _flag("CR_MULTIMODAL"): + from ..adapters.store_image import ImageRetriever + from ..adapters.store_inmemory import InMemoryStore as _InMem + from ..adapters.store_multimodal import MultimodalRetriever + retr = MultimodalRetriever(text=retr or _InMem([]), image=ImageRetriever()) + return retr + + +def _make_tenant(retriever, *, tenant_id: str) -> LibreChatTenant: + """A LibreChatTenant sharing this deployment's cost/load config but with its OWN policy, + quality-ledger and calibration files (each corpus learns independently) and its OWN + cross-language config (CR_TENANT_LANGS, falling back to the global CR_QUERY_LANGS).""" + sfx = f"_{tenant_id}" + calib_path = str(fleet.home / f"librechat_calibration{sfx}.json") + calib_map = CalibrationMap.load(calib_path) + calib_log = CalibrationLog(str(fleet.home / f"librechat_calib_log{sfx}.jsonl")) if _flag("CR_CALIBRATE") else None + q_ledger = (QualityLedger(path=str(fleet.home / f"librechat_quality{sfx}.json")) + if _quality_routing or _flag("CR_QUALITY") else None) + langs = _TENANT_LANGS.get(tenant_id, _QUERY_LANGS) # per-tenant cross-language, else the global default + return LibreChatTenant( + runtime=fleet.runtime, retriever=retriever, strategies=_lc_strategies, + query_expander=_make_query_expander(langs), + persist_path=str(fleet.home / f"librechat_bandit{sfx}.json"), + calibration=calib_map, calib_log=calib_log, load_meter=_load_meter, + cost_profile=_cost_profile, load_aware=_flag("CR_LOAD_AWARE"), + abstain_threshold=_abstain_threshold, + reward_beta=float(os.getenv("CR_REWARD_BETA", "0.5")), + quality_ledger=q_ledger, quality_routing=_quality_routing) + + +_TENANTS: dict[str, LibreChatTenant] = {} +_DEFAULT_TENANT_ID = os.getenv("CR_SHIM_MODEL_ID", "context-runtime") +_tenants_spec = os.getenv("CR_TENANTS", "").strip() +if _tenants_spec: + _parsed = _parse_tenants_spec(_tenants_spec) + for _tid, _cdir, _sspec in _parsed: + _t = _make_tenant(_build_lc_retriever(_sspec), tenant_id=_tid) + if _cdir and os.path.isdir(_cdir): + try: + _t.ingest(_cdir) # single-corpus tenants index at startup; shards self-index + except Exception: + pass + _TENANTS[_tid] = _t + if _parsed: + _DEFAULT_TENANT_ID = _parsed[0][0] + librechat = _TENANTS[_DEFAULT_TENANT_ID] # default tenant backs the legacy global + judge wiring + tests +if not _TENANTS: # CR_TENANTS unset ⇒ legacy single tenant + _TENANTS[_DEFAULT_TENANT_ID] = librechat + + +def _tenant_for(model_id: str | None) -> LibreChatTenant: + """Pick the tenant a model id maps to; fall back to the default (legacy) tenant.""" + return _TENANTS.get((model_id or "").strip()) or _TENANTS.get(_DEFAULT_TENANT_ID) or librechat + + def _hits_json(hits) -> list[dict]: return [{"chunk_id": h.chunk_id, "filename": h.filename, "score": round(h.score, 4), "text": h.text[:600]} for h in hits] @@ -1109,38 +1275,43 @@ def _hits_json(hits) -> list[dict]: class LcIngestReq(BaseModel): path: str + model: str | None = None # which corpus tenant to index into (default: the default tenant) class LcRetrieveReq(BaseModel): request: str + model: str | None = None # which corpus tenant to route to (multi-corpus demos) class LcJudgeReq(BaseModel): request: str score: float + model: str | None = None @app.post("/librechat/ingest", dependencies=[Depends(require_api_key)]) def librechat_ingest(req: LcIngestReq) -> dict: - """Index a (pre-built, normalized) corpus directory into LibreChat's retriever.""" + """Index a (pre-built, normalized) corpus directory into a tenant's retriever.""" if not os.path.isdir(req.path): raise HTTPException(400, f"not a directory: {req.path}") - return {"ingested": librechat.ingest(req.path)} + return {"ingested": _tenant_for(req.model).ingest(req.path)} @app.post("/librechat/retrieve") def librechat_retrieve(req: LcRetrieveReq) -> dict: """Plan + retrieve context for a request using the learned strategy (no judging yet).""" - ctx = librechat.retrieve(req.request) + tenant = _tenant_for(req.model) + ctx = tenant.retrieve(req.request) return {"request": ctx.request, "strategy": ctx.strategy.key, "method": ctx.strategy.method, "final_k": ctx.strategy.final_k, "hits": _hits_json(ctx.hits), "context": ctx.context, "plan_id": ctx.plan.id, - "suggestion": librechat.suggest(req.request)} + "suggestion": tenant.suggest(req.request)} class LcCompareReq(BaseModel): request: str k: int = 5 + model: str | None = None # the corpus the Query Board is inspecting (per selected chat model) @app.post("/librechat/compare") @@ -1148,7 +1319,7 @@ def librechat_compare(req: LcCompareReq) -> dict: """Retrieval transparency: run every method (bm25/vector/hybrid/community/graph) side by side and report what the learned policy chose + served. Powers the LibreChat comparison panel that shows users WHY Context Runtime's answer is better than any single method.""" - return librechat.compare(req.request, k=max(1, min(req.k, 10))) + return _tenant_for(req.model).compare(req.request, k=max(1, min(req.k, 10))) @app.post("/librechat/explain") @@ -1157,16 +1328,16 @@ def librechat_explain(req: LcCompareReq) -> dict: candidate arm with its learned value + quality/cost decomposition, the per-method retrieval trace with calibrated P(relevant), what was served, the abstention decision, and how the reward is computed. Read-only, so inspecting never perturbs the learned policy.""" - return librechat.explain(req.request, k=max(1, min(req.k, 10))) + return _tenant_for(req.model).explain(req.request, k=max(1, min(req.k, 10))) @app.get("/librechat/image") -def librechat_image(chunk_id: str): +def librechat_image(chunk_id: str, model: str | None = None): """Serve an indexed image by chunk_id so the LibreQB panel can render a thumbnail for cross-modal `image` hits. Only images the retriever has indexed are servable (no arbitrary file read).""" from fastapi.responses import FileResponse - path = librechat.image_path(chunk_id) + path = _tenant_for(model).image_path(chunk_id) if not path or not os.path.isfile(path): raise HTTPException(404, "image not found") return FileResponse(path) @@ -1176,13 +1347,15 @@ def librechat_image(chunk_id: str): def librechat_judge(req: LcJudgeReq) -> dict: """Close the loop: an external LLM judge posts the retrieval-quality score (0..1) for the last retrieve of this request; the policy learns which strategy the judge rewards.""" - reward = librechat.record_judgment(req.request, max(0.0, min(1.0, req.score))) - return {"reward": reward, "suggestion": librechat.suggest(req.request), "policy": librechat.policy()} + tenant = _tenant_for(req.model) + reward = tenant.record_judgment(req.request, max(0.0, min(1.0, req.score))) + return {"reward": reward, "suggestion": tenant.suggest(req.request), "policy": tenant.policy()} class LcFeedbackReq(BaseModel): request: str signal: str # thumbs_up | kept | copied | regenerated | thumbs_down | edited | … + model: str | None = None @app.post("/librechat/feedback", dependencies=[Depends(require_api_key)]) @@ -1190,24 +1363,26 @@ def librechat_feedback(req: LcFeedbackReq) -> dict: """The app's NATIVE reward: map a user action (kept / regenerated / thumbs / …) to a retrieval-quality signal so the policy learns from REAL usage — no LLM in the loop. This is the primary online learning path; /librechat/judge is a cold-start bootstrap.""" - reward = librechat.record_feedback(req.request, req.signal) + tenant = _tenant_for(req.model) + reward = tenant.record_feedback(req.request, req.signal) return {"signal": req.signal, "reward": reward, - "suggestion": librechat.suggest(req.request), "policy": librechat.policy()} + "suggestion": tenant.suggest(req.request), "policy": tenant.policy()} @app.post("/librechat/chat", dependencies=[Depends(require_api_key)]) def librechat_chat(req: LcRetrieveReq) -> dict: """Self-contained loop: retrieve → (offline heuristic) judge → learn, in one call.""" - ctx, score, reward = librechat.handle(req.request) + tenant = _tenant_for(req.model) + ctx, score, reward = tenant.handle(req.request) return {"request": ctx.request, "strategy": ctx.strategy.key, "score": score, "reward": reward, "hits": _hits_json(ctx.hits), "context": ctx.context, - "suggestion": librechat.suggest(req.request)} + "suggestion": tenant.suggest(req.request)} @app.get("/librechat/policy") -def librechat_policy() -> dict: - """The retrieval strategy LibreChat has learned per request-type.""" - return {"policy": librechat.policy()} +def librechat_policy(model: str | None = None) -> dict: + """The retrieval strategy a tenant has learned per request-type.""" + return {"policy": _tenant_for(model).policy()} # ──────────────────────────── OpenAI-compatible shim (LibreChat self-learning bridge) ──────────────────────────── @@ -1259,16 +1434,24 @@ def _nullctx(): def _build_judge_model(): - """A ModelPlugin pointed at the upstream, for the LLM judges — or None if unavailable - (→ the heuristic judge stays in place). Guarded so a missing litellm/upstream degrades.""" - if not _UPSTREAM_BASE: + """A ModelPlugin for the LLM judges. Prefers a DEDICATED judge endpoint — CR_JUDGE_BASE_URL, + e.g. Grok 4.5 on https://api.x.ai/v1 — a strong but cheap judge kept SEPARATE from the chat model + (the chat model shouldn't grade its own retrieval, and Grok 4.5 is ~5x cheaper than GPT-5.5). Falls + back to the chat upstream if no dedicated judge is set. None if neither → the heuristic judge stays.""" + jbase = os.getenv("CR_JUDGE_BASE_URL", "").rstrip("/") + if jbase: # dedicated judge (default: Grok 4.5) + base, model, key = jbase, os.getenv("CR_JUDGE_MODEL", "grok-4.5"), os.getenv("CR_JUDGE_API_KEY", "") + elif _UPSTREAM_BASE: # else reuse the chat upstream as judge + base = _UPSTREAM_BASE + model = os.getenv("CR_JUDGE_MODEL") or _UPSTREAM_MODEL or _SHIM_MODEL_ID + key = _UPSTREAM_KEY + else: return None try: - from ..adapters.model_litellm import LiteLLMModel, Tier - model = os.getenv("CR_JUDGE_MODEL") or _UPSTREAM_MODEL or _SHIM_MODEL_ID - return LiteLLMModel(tiers=[Tier(name="judge", model=model, - base_url=_UPSTREAM_BASE, api_key=_UPSTREAM_KEY)], - default_tier="judge") + from ..adapters.model_openai import OpenAICompatibleModel # xAI is OpenAI-compatible; no litellm dep + from ..adapters.model_litellm import Tier + return OpenAICompatibleModel([Tier(name="judge", model=model, base_url=base, api_key=key or "sk-noauth")], + default_tier="judge") except Exception: return None @@ -1280,8 +1463,10 @@ def _build_judge_model(): if _jm is not None: from ..integrations.librechat import llm_judge as _mk_llm_judge from ..integrations.librechat import llm_passage_judge as _mk_passage_judge - librechat.judge = _mk_llm_judge(_jm) - librechat.passage_judge = _mk_passage_judge(_jm) + _lj, _pj = _mk_llm_judge(_jm), _mk_passage_judge(_jm) + for _t in _TENANTS.values(): # every corpus tenant shares the dedicated judge + _t.judge = _lj + _t.passage_judge = _pj class ChatMessage(BaseModel): @@ -1299,9 +1484,12 @@ class ChatCompletionReq(BaseModel): @app.get("/v1/models") def v1_models() -> dict: - """OpenAI-compatible model list — LibreChat populates its model dropdown from this.""" + """OpenAI-compatible model list — LibreChat populates its model dropdown from this. One entry + per corpus tenant (finance / medical / mixed …), so each selectable model is a scoped corpus.""" + created = int(_time.time()) return {"object": "list", "data": [ - {"id": _SHIM_MODEL_ID, "object": "model", "created": int(_time.time()), "owned_by": "context-runtime"}]} + {"id": tid, "object": "model", "created": created, "owned_by": "context-runtime"} + for tid in _TENANTS]} def _forward_to_upstream(messages: list[dict], req: ChatCompletionReq) -> tuple[str, dict, str]: @@ -1313,7 +1501,7 @@ def _forward_to_upstream(messages: list[dict], req: ChatCompletionReq) -> tuple[ body = ctx_preamble.split("RETRIEVED CONTEXT:", 1)[-1].strip() answer = ("Based on the retrieved context:\n\n" + body[:1200]) if body else \ "No relevant context was retrieved for this request." - return answer, {"prompt_tokens": 0, "completion_tokens": 0, "total_tokens": 0}, _SHIM_MODEL_ID + return answer, {"prompt_tokens": 0, "completion_tokens": 0, "total_tokens": 0}, (req.model or _SHIM_MODEL_ID) payload = {"model": _UPSTREAM_MODEL or _SHIM_MODEL_ID, "messages": messages, "stream": False} # Reasoning upstreams (e.g. kimi-k2.6) need a temperature the model accepts and enough # max_tokens to finish reasoning AND write the answer — LibreChat's own params are often @@ -1339,7 +1527,7 @@ def _forward_to_upstream(messages: list[dict], req: ChatCompletionReq) -> tuple[ _shim_seen: dict[int, bool] = {} -def _infer_feedback(req: "ChatCompletionReq", request_text: str) -> None: +def _infer_feedback(tenant, req: "ChatCompletionReq", request_text: str) -> None: """Turn LibreChat's real behaviour into the app's NATIVE reward — no UI patch needed: • the user asked a NEW question after a prior answer → that prior turn was KEPT; • the exact conversation was re-sent (Regenerate) → that answer was REGENERATED. @@ -1349,43 +1537,46 @@ def _infer_feedback(req: "ChatCompletionReq", request_text: str) -> None: if len(user_msgs) >= 2 and assistant_msgs: # a completed prior turn the user moved on from prior = user_msgs[-2].content if prior and prior != request_text: - librechat.record_feedback(prior, "kept") - key = hash(tuple((m.role, m.content) for m in req.messages)) + tenant.record_feedback(prior, "kept") + key = hash((req.model or "", tuple((m.role, m.content) for m in req.messages))) with _shim_seen_lock: # runs on the learn pool → guard the dict regenerated = _shim_seen.get(key) if len(_shim_seen) > 512: _shim_seen.clear() _shim_seen[key] = True if regenerated: # identical conversation re-sent → Regenerate - librechat.record_feedback(request_text, "regenerated") + tenant.record_feedback(request_text, "regenerated") -def _learn_in_background(req: ChatCompletionReq, request_text: str, ctx) -> None: +def _learn_in_background(tenant, req: ChatCompletionReq, request_text: str, ctx) -> None: """The learning half of a chat turn — off the response path. Infers implicit feedback, runs the (possibly LLM) judge, and records the reward (bandit update + persist).""" try: - _infer_feedback(req, request_text) # native implicit reward (primary) - librechat.annotate_relevance(request_text, ctx.hits) # per-passage calibration labels (opt-in) - score = librechat.judge(request_text, ctx.context, ctx.hits) - librechat.record_judgment(request_text, score) + _infer_feedback(tenant, req, request_text) # native implicit reward (primary) + tenant.annotate_relevance(request_text, ctx.hits) # per-passage calibration labels (opt-in) + score = tenant.judge(request_text, ctx.context, ctx.hits) + tenant.record_judgment(request_text, score) except Exception: pass # learning must never affect the served response def _chat_core(req: ChatCompletionReq) -> dict: - """The self-learning chat turn: retrieve → augment → answer (learning is deferred).""" + """The self-learning chat turn: retrieve → augment → answer (learning is deferred). + Routes to the corpus tenant named by req.model (finance / medical / mixed).""" + tenant = _tenant_for(req.model) user_msgs = [m for m in req.messages if m.role == "user"] request_text = user_msgs[-1].content if user_msgs else "" # Count this turn as in-flight so the load-aware sizer/bandit see current concurrency. tracker = _load_meter.track() if _load_meter is not None else _nullctx() with tracker: - ctx = librechat.retrieve(request_text) # learned retrieval strategy + ctx = tenant.retrieve(request_text) # learned retrieval strategy if ctx.abstain: # Best passage's calibrated P(relevant) is below the floor — answering would be # ungrounded. Say so instead of burning an upstream call, and still learn. answer = ("I don't have enough relevant context to answer that reliably. " "Try rephrasing or adding detail.") - usage, model_name = {"prompt_tokens": 0, "completion_tokens": 0, "total_tokens": 0}, _SHIM_MODEL_ID + usage = {"prompt_tokens": 0, "completion_tokens": 0, "total_tokens": 0} + model_name = req.model or _SHIM_MODEL_ID else: system = ("You are a helpful assistant. Answer using the RETRIEVED CONTEXT below when " "relevant, citing [n]. If it is insufficient, say so.\n\n" @@ -1396,9 +1587,9 @@ def _chat_core(req: ChatCompletionReq) -> dict: # judge + bandit update run on the learn pool so they never add to response latency. score = _heuristic_judge(request_text, ctx.context, ctx.hits) rel = (sum(ctx.probs) / len(ctx.probs)) if ctx.probs else None - reward = _reward_from_judgment(score, ctx.strategy, librechat.strategies, - rel_signal=rel, beta=librechat.reward_beta) - _submit_learn(_learn_in_background, req, request_text, ctx) + reward = _reward_from_judgment(score, ctx.strategy, tenant.strategies, + rel_signal=rel, beta=tenant.reward_beta) + _submit_learn(_learn_in_background, tenant, req, request_text, ctx) return { "id": f"chatcmpl-{ctx.plan.id}", "object": "chat.completion", "created": int(_time.time()), "model": model_name, @@ -1406,7 +1597,7 @@ def _chat_core(req: ChatCompletionReq) -> dict: "usage": usage or {"prompt_tokens": 0, "completion_tokens": 0, "total_tokens": 0}, "context_runtime": { "strategy": ctx.strategy.key, "retrieval_score": round(score, 4), "reward": reward, - "citations": [h.chunk_id for h in ctx.hits], "suggestion": librechat.suggest(request_text), + "citations": [h.chunk_id for h in ctx.hits], "suggestion": tenant.suggest(request_text), "confidence": round(ctx.max_p_rel, 4), "abstained": ctx.abstain, }, } diff --git a/context_runtime/costmodel/estimators.py b/context_runtime/costmodel/estimators.py index 8b4d501..857042b 100644 --- a/context_runtime/costmodel/estimators.py +++ b/context_runtime/costmodel/estimators.py @@ -21,7 +21,8 @@ TIER_ACCURACY = {"local": 0.62, "cheap": 0.78, "premium": 0.9} TIER_HALLUCINATION = {"local": 0.18, "cheap": 0.1, "premium": 0.05} -METHOD_RECALL = {"bm25": 0.6, "vector": 0.7, "hybrid": 0.85, "code": 0.8, "graph": 0.75} +METHOD_RECALL = {"bm25": 0.6, "vector": 0.7, "hybrid": 0.85, "code": 0.8, "graph": 0.75, + "temporal": 0.7} class HeuristicEstimator: @@ -52,6 +53,11 @@ def estimate(self, candidate: Candidate, goal: Goal) -> PlanScore: bucket, _risk = rules.classify(goal.text) if bucket == "multi_hop": recall = 0.93 if method == "graph" else recall * 0.55 + elif bucket == "temporal": + # temporal questions (point-in-time state, what-changed, provenance) need the + # bi-temporal store; single-hop retrieval returns the latest chunk and misses the + # time dimension. This is how the planner routes Graphiti vs document retrieval. + recall = 0.9 if method == "temporal" else recall * 0.6 is_graph = method == "graph" base_acc = TIER_ACCURACY.get(tier, 0.7) diff --git a/context_runtime/integrations/agent_console.py b/context_runtime/integrations/agent_console.py new file mode 100644 index 0000000..c1cbd8a --- /dev/null +++ b/context_runtime/integrations/agent_console.py @@ -0,0 +1,462 @@ +"""AgentConsole — a reusable conversational agent that runs entirely on the Context +Runtime stack. + +Every agentic-os app wraps a complex OSS core (ERPNext, Lago, OpenSCAP, CrowdSec, …); the +dashboards show state but leave users to learn the tool and hunt for the right action. +``AgentConsole`` is the one chat surface that fixes both: ask *"how do I…?"* or *"explain +this"* and it answers grounded in the app's guide; ask it to *do* something and it runs a +tool — always showing its work. + +It is deliberately built on our own three planes, not a raw provider SDK: + +* **context-runtime** — generation goes through a Context Runtime ``ModelPlugin`` + (``OpenAICompatibleModel`` when a key is present, else the offline ``StubModel``), so the + same cost-tiered ``ModelRequest → ModelResult`` contract and accounting apply. +* **redevops-rag** — the *"how do I / explain"* path is grounded by retrieval over the app's + primer, honouring redevops-rag's ``hybrid_search`` contract. The dependency-free + ``PrimerIndex`` here is the offline path; the ``[rag]`` extra swaps in the real reranked + index without changing callers. +* **agent-harness** — actions are dispatched through the harness ``ToolRegistry`` + + ``ApprovalPolicy``, so a side-effecting tool is gated and every decision is audited. + +One console per app: give it a ``tenant`` name, a ``primer`` (the product knowledge), and a +list of ``tools`` (each bound to that app's core API). ``respond()`` returns a transparent +payload the chat panel renders; ``panel_html()`` returns the panel itself. +""" +from __future__ import annotations + +import json +import math +import re +from collections import Counter +from dataclasses import dataclass +from typing import Any, Callable + +from ..adapters.model_openai import OpenAICompatibleModel +from ..adapters.model_stub import StubModel +from ..tools.base import ApprovalPolicy, ToolRegistry, ToolResult, ToolSpec +from ..types import ModelRequest + +_WORD = re.compile(r"[a-z0-9]+") + + +def _tokens(text: str) -> list[str]: + return _WORD.findall(text.lower()) + + +# ──────────────────────────── grounding (redevops-rag contract) ──────────────────────────── + + +@dataclass +class Chunk: + idx: int + text: str + score: float = 0.0 + + +class PrimerIndex: + """Dependency-free ranked retrieval over the app primer — the offline redevops-rag path. + + Splits the primer into blank-line-separated passages and ranks them by TF·IDF overlap + with the query. Same ``search(query, k) -> ranked passages`` shape as redevops-rag's + ``hybrid_search``; the ``[rag]`` extra replaces this with the reranked vector path. + """ + + def __init__(self, primer: str): + passages = [p.strip() for p in re.split(r"\n\s*\n", primer.strip()) if p.strip()] + self.passages = passages + self._toks = [_tokens(p) for p in passages] + df: Counter = Counter() + for toks in self._toks: + for t in set(toks): + df[t] += 1 + n = max(1, len(passages)) + self._idf = {t: math.log(1 + n / c) for t, c in df.items()} + + def search(self, query: str, k: int = 4) -> list[Chunk]: + q = set(_tokens(query)) + if not q or not self.passages: + return [] + scored: list[Chunk] = [] + for i, toks in enumerate(self._toks): + tf = Counter(toks) + score = sum(tf[t] * self._idf.get(t, 0.0) for t in q) + if score > 0: + scored.append(Chunk(idx=i, text=self.passages[i], score=round(score, 3))) + scored.sort(key=lambda c: c.score, reverse=True) + return scored[:k] + + +# ──────────────────────────── tools (agent-harness) ──────────────────────────── + + +ToolFn = Callable[[dict], Any] + + +class _Tool: + """Adapt a plain callable into an agent-harness ``ToolPlugin``.""" + + def __init__(self, spec: ToolSpec, fn: ToolFn): + self._spec = spec + self._fn = fn + + def spec(self) -> ToolSpec: + return self._spec + + def run(self, args: dict) -> ToolResult: + try: + out = self._fn(args or {}) + except Exception as e: # noqa: BLE001 — never let a tool crash the console + return ToolResult(ok=False, error=str(e), text=f"[tool error] {e}") + if isinstance(out, ToolResult): + return out + if isinstance(out, dict) and ("text" in out or "data" in out): + return ToolResult(ok=True, data=out.get("data"), text=str(out.get("text", ""))) + return ToolResult(ok=True, data=out, text="" if out is None else str(out)) + + +def tool( + name: str, + description: str, + fn: ToolFn, + *, + parameters: dict | None = None, + side_effecting: bool = False, +) -> _Tool: + """Declare an app capability. ``fn`` takes an args dict and returns a dict/str/ToolResult.""" + return _Tool( + ToolSpec( + name=name, + description=description, + parameters=parameters or {"type": "object", "properties": {}}, + side_effecting=side_effecting, + ), + fn, + ) + + +# ──────────────────────────── the console ──────────────────────────── + + +_CLASSIFY_SYS = ( + "You route a user's message for the assistant of a business app. Reply with ONLY a JSON " + 'object: {"mode":"tool"|"help","tool":,"args":{...},"reason":}. ' + 'Use "tool" when the message asks to DO or SHOW something a listed tool covers; use "help" ' + "for how-to / explain / what-is questions. Never invent a tool name." +) + +# Interrogative/how-to leads that the deterministic fallback routes to help (not a tool). Kept as +# explicit phrases so action queries like "what changed this week?" / "how's my traffic?" still route. +_HELP_MARKERS = ( + "how do i", "how do you", "how can i", "how to ", "how does ", "how would i", + "what is ", "what are ", "what does ", "what's a ", "explain", "tell me about", "why ", +) + + +def _policy_items(decisions) -> list[dict]: + """Compact policy list for the panel (§10.1): only non-allow decisions become chips; allows fold.""" + out = [] + for d, phase in decisions: + if d.action == "allow": + continue + out.append({"decision": d.action, "phase": phase, "reason": d.reason, "rule_id": d.rule_id, + "scope": d.scope, "provider": d.provider, "summary": f"{d.action} · {phase}"}) + return out + + +class AgentConsole: + def __init__( + self, + tenant: str, + primer: str, + tools: list[_Tool] | tuple[_Tool, ...] = (), + *, + suggestions: list[str] | tuple[str, ...] = (), + subtitle: str = "", + model: Any = None, + allow_side_effects: list[str] | None = None, + authorizer=None, + commands=None, + policy=None, + app: str = "", + ): + self.tenant = tenant + self.subtitle = subtitle or f"Ask about {tenant} — how things work, or get something done." + self.primer = primer.strip() + self.index = PrimerIndex(self.primer) + self.suggestions = list(suggestions) + self.model = model if model is not None else (OpenAICompatibleModel.from_env() or StubModel()) + # authorizer (optional, enterprise): gates every tool by the request principal — falls back to + # the process-wide default (set_default_authorizer), so one install() call governs all consoles. + self.registry = ToolRegistry(ApprovalPolicy(mode="allowlist", allow=list(allow_side_effects or [])), + authorizer=authorizer) + # Policy Runtime: /commands are dispatched here (dual-path); `policy` enforces input/output/tool + # phases and emits PolicyDecisions. Both optional and fall back to the process default → one + # install governs the fleet. app = the policy scope slug (e.g. "market-radar"). + self.commands = commands + self._policy = policy + self.app = app + self._tools: dict[str, _Tool] = {} + for t in tools: + self.registry.register(t) + self._tools[t.spec().name] = t + + def mount_mcp(self, client, *, prefix: str | None = None, allow: list[str] | None = None) -> list[str]: + """Mount an external MCP tool server's tools into this console. + + Registers each tool in the agent-harness registry AND the classify/dispatch catalog, so + the assistant can pick and run them like any native tool. Read-only MCP tools + (``readOnlyHint``) run ungated; side-effecting ones stay gated unless named in ``allow``. + Returns the mounted tool names. + """ + from ..tools.mcp import mount_mcp as _mount_mcp + + names = _mount_mcp(self.registry, client, prefix=prefix) + for n in names: + self._tools[n] = self.registry.get(n) # classify() sees it, respond() dispatches it + if allow: + self.registry.policy.allow.update(allow) + return names + + # ---- helpers ------------------------------------------------------------- + + @property + def _live(self) -> bool: + return not isinstance(self.model, StubModel) + + def _tool_catalog(self) -> str: + return "\n".join(f"- {n}: {t.spec().description}" for n, t in self._tools.items()) or "(no tools)" + + def _keyword_route(self, message: str) -> dict: + """Deterministic fallback when the model can't return JSON (offline / parse fail).""" + # How-to / explain questions are help, not actions — mirror the LLM classify rule so a + # single incidental word ("...part of a page?" vs a "monitor a page" tool) can't misroute + # an interrogative into a side-effecting tool call. + low = (message or "").strip().lower() + if any(low.startswith(mk) or (" " + mk) in low for mk in _HELP_MARKERS): + return {"mode": "help", "tool": None, "args": {}, "reason": "how-to → help"} + toks = set(_tokens(message)) + best, best_score = None, 0 + for name, t in self._tools.items(): + vocab = set(_tokens(name.replace("_", " ") + " " + t.spec().description)) + score = len(toks & vocab) + if score > best_score: + best, best_score = name, score + if best and best_score >= 1: + return {"mode": "tool", "tool": best, "args": {}, "reason": "keyword match"} + return {"mode": "help", "tool": None, "args": {}, "reason": "default to help"} + + def classify(self, message: str) -> dict: + if not self._tools or not self._live: + return self._keyword_route(message) + prompt = f"Tools:\n{self._tool_catalog()}\n\nMessage: {message}\n\nJSON:" + try: + res = self.model.complete( + ModelRequest(messages=({"role": "user", "content": prompt},), system=_CLASSIFY_SYS, max_tokens=200) + ) + m = re.search(r"\{.*\}", res.text, re.S) + data = json.loads(m.group(0)) if m else {} + mode = data.get("mode") + tool_name = data.get("tool") + if mode == "tool" and tool_name in self._tools: + return {"mode": "tool", "tool": tool_name, "args": data.get("args") or {}, "reason": data.get("reason", "")} + if mode == "help": + return {"mode": "help", "tool": None, "args": {}, "reason": data.get("reason", "")} + except Exception: # noqa: BLE001 — any failure drops to the deterministic router + pass + return self._keyword_route(message) + + def _answer_help(self, message: str) -> dict: + hits = self.index.search(message, k=4) + evidence = [{"n": i + 1, "text": h.text, "score": h.score} for i, h in enumerate(hits)] + if not hits: + body = "I don't have that in the guide yet. Try the dashboard, or ask about a specific action." + return {"intent": "help", "text": body, "evidence": [], "model": self.model.info().name} + context = "\n\n".join(f"[{e['n']}] {e['text']}" for e in evidence) + if not self._live: + # extractive offline answer: lead with the top passage, cite it + body = f"{hits[0].text}\n\n(Grounded in the {self.tenant} guide [1].)" + return {"intent": "help", "text": body, "evidence": evidence, "model": self.model.info().name} + system = ( + f"You are the assistant for {self.tenant}. Answer the user's question using ONLY the numbered " + "guide passages, citing them inline like [1]. For a 'what is' / explain question, define the term " + "AND explain how it fits this system from the passage context — keep the relevant specifics, don't " + "collapse it to a bare one-line definition. For a how-to question, give the concrete steps. Keep it " + "tight but complete; if the passages don't cover it, say so plainly." + ) + prompt = f"Guide:\n{context}\n\nQuestion: {message}" + try: + res = self.model.complete( + ModelRequest(messages=({"role": "user", "content": prompt},), system=system, max_tokens=600) + ) + return {"intent": "help", "text": res.text or hits[0].text, "evidence": evidence, + "model": res.model, "cost_usd": res.est_cost_usd} + except Exception: # noqa: BLE001 — model outage → fall back to the top grounded passage + return {"intent": "help", "text": f"{hits[0].text}\n\n(Grounded in the {self.tenant} guide [1].)", + "evidence": evidence, "model": self.model.info().name} + + def _effective_policy(self): + from ..policy import current_policy + return self._policy or current_policy() + + def _answer_tool(self, name: str, args: dict, principal=None) -> dict: + # tool-phase policy: an approval rule pauses an irreversible action; a deny rule blocks it. + policy = self._effective_policy() + if policy is not None: + td = policy.check(principal, "tool", (name, args), app=getattr(principal, "app", "") or self.app) + if td.action in ("require_approval", "deny"): + verb = "needs approval before it runs" if td.action == "require_approval" else "is blocked" + return {"intent": "action", "tool": name, "text": f"This action {verb}: {td.reason}", + "evidence": [{"tool": name, "gate": td.action, "ok": False}], "data": None, "approved": False, + "policy": [{"decision": td.action, "phase": "tool", "reason": td.reason, "rule_id": td.rule_id, + "scope": td.scope, "provider": td.provider, "summary": f"{name} → {td.action}"}]} + result = self.registry.run(name, args, principal=principal) + gate = self.registry.audit[-1] if self.registry.audit else {"decision": "read-only"} + evidence = [{"tool": name, "args": args, "gate": gate.get("reason") or gate.get("decision", ""), "ok": result.ok}] + summary = result.text or (json.dumps(result.data)[:800] if result.data is not None else "") + if not result.ok: + body = summary or "That action needs confirmation before it can run." + return {"intent": "action", "tool": name, "text": body, "evidence": evidence, "data": result.data, "approved": False} + if self._live and summary: + system = ( + f"You are the assistant for {self.tenant}. The user asked something and a tool returned the data " + "below. Answer them directly and concisely from that data. Do not invent numbers." + ) + prompt = f"Tool `{name}` returned:\n{summary}\n\nAnswer the user." + try: + res = self.model.complete( + ModelRequest(messages=({"role": "user", "content": prompt},), system=system, max_tokens=500) + ) + body = res.text or summary + model_name = res.model + except Exception: # noqa: BLE001 — model outage → return the raw tool summary + body = summary + model_name = self.model.info().name + else: + body = summary or "Done." + model_name = self.model.info().name + return {"intent": "action", "tool": name, "text": body, "evidence": evidence, "data": result.data, "approved": True, "model": model_name} + + def respond(self, message: str, principal=None) -> dict: + message = (message or "").strip() + # 1) command path — deterministic, no LLM, permission-gated (dual-path dispatch) + if self.commands is not None and self.commands.is_command(message): + return self.commands.dispatch(message, principal) + if not message: + return {"intent": "help", "text": "Ask me anything about " + self.tenant + ".", "evidence": []} + policy = self._effective_policy() + app = getattr(principal, "app", "") or self.app + decisions: list[tuple] = [] + # 2) input policy (guardrails) + if policy is not None: + di = policy.check(principal, "input", message, app=app) + decisions.append((di, "input")) + if di.action == "deny": + return {"intent": "blocked", "text": f"Blocked by policy: {di.reason}", "evidence": [], + "policy": _policy_items(decisions)} + if di.action == "redact": + message = policy.redact(message, di) + # 3) route (help / tool). tool-phase policy is enforced inside _answer_tool. + route = self.classify(message) + if route["mode"] == "tool" and route["tool"] in self._tools: + out = self._answer_tool(route["tool"], route.get("args") or {}, principal=principal) + else: + out = self._answer_help(message) + out["route"] = route.get("reason", "") + # 4) output policy (guardrails) + if policy is not None: + do = policy.check(principal, "output", out.get("text", ""), app=app) + decisions.append((do, "output")) + if do.action == "deny": + return {"intent": "blocked", "text": f"Blocked by policy: {do.reason}", "evidence": [], + "policy": _policy_items(decisions)} + if do.action == "redact": + out["text"] = policy.redact(out["text"], do) + # 5) attach compact policy summary (progressive disclosure — §10.1). Merge any tool-phase items. + # Only when a policy is active, so consoles without policy are byte-for-byte unchanged. + if policy is not None: + items = _policy_items(decisions) + [p for p in (out.get("policy") or []) if p.get("decision") != "allow"] + out["policy"] = items + out["policy_checks"] = len(decisions) + (1 if out.get("tool") else 0) + return out + + # ---- the panel ----------------------------------------------------------- + + def panel_html(self, mount: str = "agent", endpoint: str = "api/agent") -> str: + chips = "".join( + f'' for s in self.suggestions + ) + return _PANEL_TMPL.format( + mount=mount, endpoint=endpoint, tenant=_esc(self.tenant), subtitle=_esc(self.subtitle), chips=chips + ) + + +def _esc(s: str) -> str: + return ( + str(s).replace("&", "&").replace("<", "<").replace(">", ">").replace('"', """) + ) + + +_PANEL_TMPL = """ +
+
{tenant} · Assistant
+
{subtitle}
+
+
{chips}
+
+ + +
+
+ + +""" diff --git a/context_runtime/integrations/bandit.py b/context_runtime/integrations/bandit.py index 7d280db..15788d5 100644 --- a/context_runtime/integrations/bandit.py +++ b/context_runtime/integrations/bandit.py @@ -27,11 +27,15 @@ class EpsilonGreedyBandit: """ def __init__(self, arms: tuple, epsilon: float = 0.15, optimistic: float = 1.0, - seed: int = 0x9E3779B9, persist_path: str | None = None): + seed: int = 0x9E3779B9, persist_path: str | None = None, discount: float = 0.0): import threading self.arms = arms self.epsilon = epsilon self.optimistic = optimistic + # discount ∈ (0,1] = constant step size (exponential recency weighting) so stale evidence + # fades and the policy tracks a drifting best arm — the Whitepaper-v3 non-stationarity property. + # 0 (default) = sample-average (1/n), the stationary estimator; behavior is unchanged. + self.discount = discount self.stats: dict[str, dict[str, list[float]]] = {} # ctx → arm.key → [n, mean] self._rng = seed & 0xFFFFFFFF self.persist_path = persist_path # learned policy survives restarts if set @@ -73,7 +77,8 @@ def update(self, ctx: str, arm, reward: float) -> None: arms = self._ctx(ctx) n, mean = arms[arm.key] n += 1 - arms[arm.key] = [n, mean + (reward - mean) / n] + alpha = self.discount if self.discount > 0.0 else 1.0 / n + arms[arm.key] = [n, mean + alpha * (reward - mean)] # serialize a consistent snapshot under the lock; write it to disk outside so # the fsync/rename never blocks concurrent select/update. snapshot = json.dumps(self.stats) if self.persist_path else None diff --git a/context_runtime/integrations/job_leads.py b/context_runtime/integrations/job_leads.py new file mode 100644 index 0000000..1761c70 --- /dev/null +++ b/context_runtime/integrations/job_leads.py @@ -0,0 +1,457 @@ +"""Job-lead outreach — turn AI-engineering job listings into tailored pitches (Market Radar lead-gen). + +A company hiring an *AI Data Engineer / AI Developer* internally is signalling real, unmet AI-adoption +pain — exactly who an execution-layer runtime helps. This module: + + • **sources** such listings (via a pluggable search — Market Radar's web_search MCP, or any callable); + • **filters** to genuine internal-build roles, dropping consultancies / forward-deployed / agency / + staffing shops (they build AI *for* clients, not for themselves — no internal pain to solve); + • **writes** a listing-tailored pitch from an EDITABLE template (agent-editable, shared by all users); + • **dedupes** so a listing that reappears months later is never pitched twice; + • **drafts** distribution (default: an outbox draft — sending is an approval-gated deployment step). + +Dependency-light; the LLM tailoring and live search are optional/pluggable, so the logic is fully +testable offline. +""" +from __future__ import annotations + +import json +import os +import re +import time +from dataclasses import dataclass, field +from pathlib import Path +from typing import Callable + +# ── signals ─────────────────────────────────────────────────────────────────────────────────────── +# A company hiring for these builds AI in-house → a lead. +_AI_ROLE = re.compile( + r"\b(ai|ml|genai|llm)\b.*\b(engineer|developer|scientist|architect)\b" + r"|\b(machine learning|applied ai|ai platform|ai data|prompt|rag|retrieval[- ]augmented|mlops|" + r"llmops|foundation model|generative ai)\b", re.I) +# A company doing THIS builds AI for clients → skip (no internal pain). +_CONSULTING = re.compile( + r"\b(consult\w*|advisory|agency|staffing|system integrator|systems integration|outsourc\w*|" + r"forward[- ]deployed|client engagement|client delivery|professional services|body shop|" + r"managed services|for our clients|on behalf of clients)\b", re.I) +# Extra evidence the pain is internal (a scorer, not a gate). +_INTERNAL = re.compile( + r"\b(in[- ]house|internal|our (own )?(platform|product|data|stack|customers)|build(ing)? our|" + r"production ai|our ai (platform|systems)|self[- ]host)\b", re.I) + + +@dataclass(frozen=True) +class JobListing: + company: str + title: str + url: str = "" + description: str = "" + location: str = "" + posted: str = "" + source: str = "" + + def key(self) -> str: + """Stable dedup key: normalized company + role (so a reappearing listing collides).""" + def norm(s: str) -> str: + return re.sub(r"[^a-z0-9]+", " ", (s or "").lower()).strip() + return f"{norm(self.company)}|{norm(self.title)}" + + def _text(self) -> str: + return f"{self.title}\n{self.company}\n{self.description}" + + +@dataclass(frozen=True) +class Classification: + is_lead: bool + kind: str # "internal" | "consulting" | "other" + reason: str + score: float # 0..1 confidence it's a real internal-AI-pain lead + + +def classify(listing: JobListing) -> Classification: + """Is this listing a real internal-AI-adoption lead (vs. a consultancy hiring, or unrelated)?""" + text = listing._text() + ai = bool(_AI_ROLE.search(text)) + consulting = bool(_CONSULTING.search(text)) + if consulting: + return Classification(False, "consulting", "consultancy / forward-deployed — builds AI for clients", 0.0) + if not ai: + return Classification(False, "other", "not an AI-engineering role", 0.0) + internal = bool(_INTERNAL.search(text)) + score = 0.7 + (0.3 if internal else 0.0) + reason = "internal AI-engineering hire" + (" (explicit in-house signals)" if internal else "") + return Classification(True, "internal", reason, score) + + +# ── per-user lead profiles (each user targets their own roles, with their own template) ──────────── +# The default excludes apply to everyone (consultancies build AI *for* clients — no internal pain). +_DEFAULT_EXCLUDE = ["consult", "advisory", "forward deployed", "forward-deployed", "agency", "staffing", + "system integrator", "professional services", "outsourc", "managed services", + "for our clients", "on behalf of clients", "body shop"] +# The owner's targets (AI-infrastructure roles) — seeded for LEAD_OWNER_USER only. +_OWNER_INCLUDE = ["AI Data Engineer", "AI Engineer", "AI Developer", "ML Engineer", + "Machine Learning Engineer", "Applied AI", "AI Platform", "LLM", "RAG", "GenAI", "MLOps"] + + +@dataclass +class LeadProfile: + """One user's outreach config: which role categories to target, what to exclude, and their template.""" + user: str + include: list[str] = field(default_factory=list) # role phrases to TARGET (empty ⇒ user must set) + exclude: list[str] = field(default_factory=lambda: list(_DEFAULT_EXCLUDE)) + template: str = "" + resume_path: str = "" + + def to_dict(self) -> dict: + return {"user": self.user, "include": self.include, "exclude": self.exclude, + "template": self.template, "resume_path": self.resume_path} + + @classmethod + def from_dict(cls, d: dict) -> "LeadProfile": + return cls(user=d.get("user", ""), include=list(d.get("include", [])), + exclude=list(d.get("exclude", _DEFAULT_EXCLUDE)), template=d.get("template", ""), + resume_path=d.get("resume_path", "")) + + +class ProfileStore: + """Per-user profiles persisted as ``/.json``. A user with no saved profile gets a default: + the owner (``LEAD_OWNER_USER``) is seeded with the AI-infrastructure targets + the Context Runtime + template; everyone else gets a blank-target profile they configure via chat.""" + + def __init__(self, dir: str | None = None, owner: str | None = None): + self.dir = dir or os.getenv("LEAD_PROFILE_DIR", "/data/market-radar/profiles") + self.owner = (owner or os.getenv("LEAD_OWNER_USER", "alex")).strip() + + def _path(self, user: str): + safe = re.sub(r"[^a-z0-9_-]+", "_", (user or "default").lower()) + return Path(self.dir) / f"{safe}.json" + + def default_for(self, user: str) -> LeadProfile: + if user and user == self.owner: + return LeadProfile(user, list(_OWNER_INCLUDE), list(_DEFAULT_EXCLUDE), + DEFAULT_TEMPLATE, os.getenv("RESUME_PATH", "")) + return LeadProfile(user or "default", [], list(_DEFAULT_EXCLUDE), GENERIC_TEMPLATE, "") + + def load(self, user: str) -> LeadProfile: + p = self._path(user) + if p.exists(): + try: + return LeadProfile.from_dict(json.loads(p.read_text(encoding="utf-8"))) + except Exception: # noqa: BLE001 + pass + return self.default_for(user) + + def save(self, profile: LeadProfile) -> None: + p = self._path(profile.user) + p.parent.mkdir(parents=True, exist_ok=True) + tmp = str(p) + ".tmp" + Path(tmp).write_text(json.dumps(profile.to_dict(), indent=2), encoding="utf-8") + os.replace(tmp, p) + + +def classify_for(listing: JobListing, profile: LeadProfile) -> Classification: + """Classify against ONE user's targets: skip on any exclude match, lead on any include match.""" + text = listing._text().lower() + for kw in profile.exclude: + if kw and kw.lower() in text: + return Classification(False, "excluded", f"matched exclusion {kw!r}", 0.0) + if not profile.include: + return Classification(False, "unset", "no target roles set — tell the agent what to look for", 0.0) + hits = [kw for kw in profile.include if kw and kw.lower() in text] + if not hits: + return Classification(False, "other", "no targeted role matched", 0.0) + internal = bool(_INTERNAL.search(text)) + score = min(1.0, 0.6 + 0.1 * len(hits) + (0.1 if internal else 0.0)) + return Classification(True, "match", "matched: " + ", ".join(hits[:3]), round(score, 2)) + + +def find_leads_for(source, ledger, profile: LeadProfile, *, queries=None, limit: int = 50, + min_score: float = 0.6) -> list[dict]: + """Search (queries default to the profile's target roles) → classify against the profile → dedupe.""" + queries = queries or (tuple(f'"{c}" hiring' for c in profile.include[:4]) or DEFAULT_QUERIES) + seen: set[str] = set() + leads: list[dict] = [] + for q in queries: + for listing in source.search(q, limit): + k = listing.key() + if k in seen: + continue + seen.add(k) + c = classify_for(listing, profile) + if not c.is_lead or c.score < min_score: + continue + if ledger is not None and ledger.already_pitched(listing): + continue + leads.append({"listing": listing, "classification": c}) + leads.sort(key=lambda d: d["classification"].score, reverse=True) + return leads[:limit] + + +# ── pitch templates ──────────────────────────────────────────────────────────────────────────────── +GENERIC_TEMPLATE = """Hi, + +I came across your {title} opening at {company}. + +{match} + +I'd welcome a short conversation about how I could help. + +Best,""" + +# The owner's template (agent-editable per user). Placeholders: {title} {company} {match}. +DEFAULT_TEMPLATE = """Hi, + +I came across your {title} opening. + +{match} + +I started with RAG. Then added reranking. Then memory, model routing, verification, permissions. +Eventually I realized every AI application was rebuilding the same execution layer. + +That became Context Runtime — an open-source runtime that plans retrieval, memory, model selection, +verification and policy before execution, measures every decision, and continuously improves from +production feedback. It's implemented in both Python and Go, includes heterogeneous retrieval +benchmarks, execution observability and explainability, and is documented in an open whitepaper. + +After reading your job description, I think there's an opportunity beyond discussing a developer role. +I'd be interested in talking about whether this architecture could simplify how your team builds and +operates enterprise AI systems. + +Resources: +• Whitepaper v2 — https://redevops.io/whitepaper-v2 +• Benchmarks — https://redevops.io/benchmarks +• Context Runtime — https://redevops.io/context-runtime/under-the-hood/ +• Planner — https://redevops.io/planner/under-the-hood/ +• Retrieval Engine — https://redevops.io/redevops-rag/under-the-hood/ + +Best, +Alex""" + +_DEFAULT_MATCH = ("Almost every responsibility listed is something I've been working on over the past " + "year while building production AI systems.") + + +class PitchTemplate: + """The outreach letter template, persisted to a shared path so every user's agent edits the same + one. Placeholders: ``{title}``, ``{company}``, ``{match}`` (the listing-tailored paragraph).""" + + def __init__(self, path: str | None = None): + self.path = path or os.getenv("PITCH_TEMPLATE_PATH", "/data/market-radar/pitch_template.txt") + + def load(self) -> str: + p = Path(self.path) + if p.exists(): + try: + return p.read_text(encoding="utf-8") + except Exception: # noqa: BLE001 + pass + return DEFAULT_TEMPLATE + + def save(self, text: str) -> None: + p = Path(self.path) + p.parent.mkdir(parents=True, exist_ok=True) + tmp = str(p) + ".tmp" + Path(tmp).write_text(text, encoding="utf-8") + os.replace(tmp, self.path) # atomic; all users see the update + + def reset(self) -> None: + self.save(DEFAULT_TEMPLATE) + + +def _needs_resume(listing: JobListing) -> bool: + return bool(re.search(r"\b(resume|cv|curriculum vitae)\b", listing.description, re.I)) + + +def write_pitch(listing: JobListing, template: str | None = None, *, model=None, + resume_path: str | None = None) -> dict: + """Render a listing-tailored pitch. ``{match}`` is written by the model from the listing when one is + given (else a sensible default). Returns {subject, body, to_hint, attach}.""" + tpl = template if template is not None else DEFAULT_TEMPLATE + match = _DEFAULT_MATCH + if model is not None and listing.description: + match = _llm_match(listing, model) or match + body = tpl.format(title=listing.title.strip() or "AI role", + company=listing.company.strip() or "your team", match=match) + subject = f"Re: your {listing.title.strip() or 'AI'} opening — an execution layer for enterprise AI" + attach = (resume_path or os.getenv("RESUME_PATH", "")) if _needs_resume(listing) else "" + return {"subject": subject, "body": body, "to_hint": listing.url, "attach": attach} + + +def _llm_match(listing: JobListing, model) -> str | None: + from ..types import ModelRequest + prompt = ( + "You are tailoring a cold outreach letter. Given the job description below, write ONE short " + "paragraph (2 sentences, first person, confident, no fluff) noting that the listed " + "responsibilities — name 2-3 specific ones actually present (e.g. RAG, model routing, " + "evaluation, retrieval, agents) — are exactly what I've built into a production AI execution " + f"layer. Do not greet or sign off.\n\nTitle: {listing.title}\nCompany: {listing.company}\n" + f"Description:\n{listing.description[:1500]}" + ) + try: + res = model.complete(ModelRequest(messages=({"role": "user", "content": prompt},), + system="Write only the paragraph.", max_tokens=160)) + text = (res.text or "").strip() + return text or None + except Exception: # noqa: BLE001 + return None + + +# ── dedup ledger (never pitch the same listing twice) ────────────────────────────────────────────── +class OutreachLedger: + """Append-only record of pitched listings. ``already_pitched`` is True for a listing whose key or + url has been pitched within ``window_days`` (default: effectively forever), so a listing that + reappears months later is not re-pitched.""" + + def __init__(self, path: str | None = None, window_days: int = 3650): + self.path = path or os.getenv("OUTREACH_LEDGER_PATH", "/data/market-radar/outreach_ledger.jsonl") + self.window_s = window_days * 86400 + self._seen: dict[str, float] = {} + self._urls: dict[str, float] = {} + self._load() + + def _load(self) -> None: + p = Path(self.path) + if not p.exists(): + return + for line in p.read_text(encoding="utf-8").splitlines(): + try: + r = json.loads(line) + except Exception: # noqa: BLE001 + continue + self._seen[r.get("key", "")] = r.get("at", 0.0) + if r.get("url"): + self._urls[r["url"]] = r.get("at", 0.0) + + def already_pitched(self, listing: JobListing, *, now: float | None = None) -> bool: + now = time.time() if now is None else now + for at in (self._seen.get(listing.key()), self._urls.get(listing.url) if listing.url else None): + if at is not None and (now - at) <= self.window_s: + return True + return False + + def record(self, listing: JobListing, channel: str = "draft", *, now: float | None = None) -> None: + at = time.time() if now is None else now + self._seen[listing.key()] = at + if listing.url: + self._urls[listing.url] = at + p = Path(self.path) + p.parent.mkdir(parents=True, exist_ok=True) + with open(self.path, "a", encoding="utf-8") as f: + f.write(json.dumps({"key": listing.key(), "url": listing.url, "company": listing.company, + "title": listing.title, "channel": channel, "at": at}) + "\n") + + +# ── sourcing (pluggable) ─────────────────────────────────────────────────────────────────────────── +class StaticJobSource: + """A fixed list of listings — for tests and for feeding pre-scraped data.""" + + def __init__(self, listings: list[JobListing]): + self._listings = list(listings) + + def search(self, query: str = "", limit: int = 50) -> list[JobListing]: + return self._listings[:limit] + + +class CallableJobSource: + """Wrap any ``search_fn(query, limit) -> list[dict|JobListing]`` (e.g. Market Radar's web_search MCP + tool, a Greenhouse/Lever/Ashby public API, or an RSS parser) into a JobSource.""" + + def __init__(self, search_fn: Callable[[str, int], list]): + self.search_fn = search_fn + + def search(self, query: str, limit: int = 50) -> list[JobListing]: + out: list[JobListing] = [] + for row in self.search_fn(query, limit) or []: + if isinstance(row, JobListing): + out.append(row) + elif isinstance(row, dict): + out.append(JobListing( + company=row.get("company", ""), title=row.get("title", ""), url=row.get("url", ""), + description=row.get("description") or row.get("snippet", ""), + location=row.get("location", ""), posted=row.get("posted", ""), + source=row.get("source", "search"))) + return out + + +class AdzunaSource: + """Live listings from the Adzuna jobs API (free tier: ``ADZUNA_APP_ID`` / ``ADZUNA_APP_KEY``; set + ``ADZUNA_COUNTRY`` e.g. us/gb). Search-based (good for "AI Data Engineer") and ToS-friendly, unlike + scraping LinkedIn. Degrades to no results when unconfigured or unreachable; inject ``client`` for tests.""" + + def __init__(self, app_id: str | None = None, app_key: str | None = None, + country: str | None = None, client=None): + self.app_id = app_id or os.getenv("ADZUNA_APP_ID", "") + self.app_key = app_key or os.getenv("ADZUNA_APP_KEY", "") + self.country = (country or os.getenv("ADZUNA_COUNTRY", "us")).lower() + self._client = client + + def search(self, query: str, limit: int = 50) -> list[JobListing]: + if not (self.app_id and self.app_key): + return [] + params = {"app_id": self.app_id, "app_key": self.app_key, "what": query, + "results_per_page": min(limit, 50), "content-type": "application/json"} + try: + client = self._client + if client is None: + import httpx + client = httpx.Client(timeout=10.0) + r = client.get(f"https://api.adzuna.com/v1/api/jobs/{self.country}/search/1", params=params) + data = r.json() + except Exception: # noqa: BLE001 + return [] + out: list[JobListing] = [] + for j in (data.get("results") or [])[:limit]: + out.append(JobListing( + company=(j.get("company") or {}).get("display_name", ""), title=j.get("title", ""), + url=j.get("redirect_url", ""), description=j.get("description", ""), + location=(j.get("location") or {}).get("display_name", ""), + posted=j.get("created", ""), source="adzuna")) + return out + + +DEFAULT_QUERIES = ( + '"AI Data Engineer" hiring', '"AI Engineer" (in-house OR platform) hiring', + '"AI Developer" build production AI', '"Machine Learning Engineer" our platform hiring', +) + + +def find_leads(source, ledger: OutreachLedger | None = None, *, queries=DEFAULT_QUERIES, + limit: int = 50, min_score: float = 0.7) -> list[dict]: + """Search → classify → keep internal-AI leads → drop already-pitched. Returns lead dicts sorted by + score (best first), each with the listing + its classification.""" + seen: set[str] = set() + leads: list[dict] = [] + for q in queries: + for listing in source.search(q, limit): + k = listing.key() + if k in seen: + continue + seen.add(k) + c = classify(listing) + if not c.is_lead or c.score < min_score: + continue + if ledger is not None and ledger.already_pitched(listing): + continue + leads.append({"listing": listing, "classification": c}) + leads.sort(key=lambda d: d["classification"].score, reverse=True) + return leads[:limit] + + +# ── distribution (default: draft; real send is approval-gated deployment config) ─────────────────── +class DraftDistributor: + """Writes each pitch to an outbox directory instead of sending — the safe default. A real + SMTP/CRM sender is a deployment concern behind explicit approval; this never sends on its own.""" + + def __init__(self, outbox: str | None = None): + self.outbox = outbox or os.getenv("OUTREACH_OUTBOX", "/data/market-radar/outbox") + + def send(self, listing: JobListing, pitch: dict, *, to: str = "") -> dict: + p = Path(self.outbox) + p.parent.mkdir(parents=True, exist_ok=True) + p.mkdir(parents=True, exist_ok=True) + fname = re.sub(r"[^a-z0-9]+", "-", listing.key().replace("|", "-")).strip("-") + ".txt" + doc = (f"To: {to or '(find via ' + (listing.url or 'company site') + ')'}\n" + f"Subject: {pitch['subject']}\n" + f"Attach: {pitch['attach'] or '(none)'}\n\n{pitch['body']}\n") + (p / fname).write_text(doc, encoding="utf-8") + return {"mode": "draft", "path": str(p / fname), "subject": pitch["subject"], "attach": pitch["attach"]} diff --git a/context_runtime/integrations/librechat.py b/context_runtime/integrations/librechat.py index 4047e9f..af72e7a 100644 --- a/context_runtime/integrations/librechat.py +++ b/context_runtime/integrations/librechat.py @@ -488,7 +488,9 @@ def explain(self, request: str, k: int = 5) -> dict: "routing can prefer a genuinely better arm at equal cost"), } return { - "request": request, "intent_bucket": bucket, "query_type": qtype, "context_key": ctx_key, + "request": request, "intent_bucket": bucket, + "knowledge_representation": getattr(plan.intent, "representation", "document"), + "query_type": qtype, "context_key": ctx_key, "decision": {"chosen": {"key": chosen.key, "method": chosen.method, "final_k": chosen.final_k, "rerank": chosen.rerank}, "candidates": candidates}, diff --git a/context_runtime/integrations/modules.py b/context_runtime/integrations/modules.py index 4ceb959..08eeb74 100644 --- a/context_runtime/integrations/modules.py +++ b/context_runtime/integrations/modules.py @@ -67,6 +67,10 @@ class ModuleSpec: ("crowdsec", "threat_intel", "edr"), "correct_verdict", ("block_ip",)), "outreach": ModuleSpec("outreach", "Twenty CRM", "pilot outreach & pipeline", ("signals", "accounts", "artifacts", "sequences"), "reply_to_pilot", ("send_sequence",)), + "guide": ModuleSpec("guide", "Context Runtime", "onboarding & help", + ("docs", "walkthroughs", "faqs", "changelog"), "answer_grounded", ()), + "growth_assistant": ModuleSpec("growth_assistant", "ERPNext", "founder growth & traction", + ("playbooks", "benchmarks", "channels"), "traction_lift", ()), # ── net-new tenants from the use-cases doc ── "incident": ModuleSpec("incident", "Kubernetes", "incident response", ("logs", "git", "metrics", "runbook"), "root_cause_found", ("rollback",)), diff --git a/context_runtime/integrations/supply_chain.py b/context_runtime/integrations/supply_chain.py new file mode 100644 index 0000000..03cecc7 --- /dev/null +++ b/context_runtime/integrations/supply_chain.py @@ -0,0 +1,246 @@ +"""SupplyChainScanner — Context Runtime's software-supply-chain inspection plane. + +"Inspect the open source we ship before clients run it." Wraps Trivy (CVEs across OS packages + +language dependencies, IaC misconfig, exposed secrets, and SBOM) and Syft (SBOM) as subprocesses, +normalizes their output, and degrades gracefully when the binaries aren't installed — the same +shell-out pattern edge-sentinel already uses for ``cscli``. + +This is the PRE-DEPLOY half of the Security & Compliance block (the runtime half is Wazuh / +CrowdSec / Falco); the Edge Sentinel AgentConsole triages and explains the findings — our +scaled analog of Project Lightwell's AI clearinghouse. + +No heavy deps: stdlib subprocess + json. Install ``trivy`` (and optionally ``syft`` / ``cosign``) +on the host to light up the real path. +""" +from __future__ import annotations + +import json +import shutil +import subprocess +from dataclasses import dataclass, field + +_SEV_ORDER = {"CRITICAL": 0, "HIGH": 1, "MEDIUM": 2, "LOW": 3, "UNKNOWN": 4} + + +@dataclass +class Finding: + id: str # CVE / advisory id + pkg: str + installed: str + fixed: str # the version that resolves it — the "patched pin" to move to + severity: str + title: str + target: str = "" + cls: str = "" # "os" (base-image system package) | "lang" (this app's own dependency) + + +@dataclass +class ScanResult: + ok: bool + target: str + scanner: str + findings: list[Finding] = field(default_factory=list) + secrets: int = 0 + misconfigs: int = 0 + note: str = "" + + def summary(self) -> dict: + by: dict[str, int] = {} + for f in self.findings: + by[f.severity] = by.get(f.severity, 0) + 1 + return { + "ok": self.ok, + "target": self.target, + "total": len(self.findings), + "by_severity": by, + "fixable": sum(1 for f in self.findings if f.fixed), + "secrets": self.secrets, + "misconfigs": self.misconfigs, + "note": self.note, + } + + +class SupplyChainScanner: + def __init__(self, timeout: float = 300.0): + self.timeout = timeout + + def available(self) -> dict: + return {t: bool(shutil.which(t)) for t in ("trivy", "syft", "cosign", "clamscan")} + + def _run(self, cmd: list[str]) -> tuple[int, str, str]: + try: + p = subprocess.run(cmd, text=True, capture_output=True, timeout=self.timeout) + return p.returncode, p.stdout, p.stderr + except FileNotFoundError: + return 127, "", "binary not found" + except subprocess.TimeoutExpired: + return 124, "", "scan timed out" + except Exception as e: # noqa: BLE001 + return 1, "", str(e) + + @staticmethod + def _parse_trivy(raw: str, target: str = "") -> ScanResult: + try: + d = json.loads(raw) + except Exception: # noqa: BLE001 + return ScanResult(False, target, "trivy", note="unparseable trivy output") + findings: list[Finding] = [] + secrets = misconfigs = 0 + _cls = {"os-pkgs": "os", "lang-pkgs": "lang"} + for res in d.get("Results", []) or []: + tgt = res.get("Target", target) + fcls = _cls.get(res.get("Class", ""), "") + for v in res.get("Vulnerabilities", []) or []: + findings.append(Finding( + id=v.get("VulnerabilityID", ""), + pkg=v.get("PkgName", ""), + installed=v.get("InstalledVersion", ""), + fixed=v.get("FixedVersion", ""), + severity=(v.get("Severity") or "UNKNOWN").upper(), + title=(v.get("Title") or v.get("Description") or "")[:200], + target=tgt, + cls=fcls, + )) + secrets += len(res.get("Secrets", []) or []) + misconfigs += len(res.get("Misconfigurations", []) or []) + findings.sort(key=lambda f: (_SEV_ORDER.get(f.severity, 9), not f.fixed)) + return ScanResult(True, target, "trivy", findings=findings, secrets=secrets, misconfigs=misconfigs) + + def scan_fs(self, path: str = ".") -> ScanResult: + """Scan a directory tree (dependencies, lockfiles, IaC, secrets).""" + if not shutil.which("trivy"): + return ScanResult(False, path, "trivy", note="trivy not installed — deploy it to enable supply-chain scanning") + rc, out, err = self._run(["trivy", "fs", "--quiet", "--format", "json", + "--scanners", "vuln,secret,misconfig", path]) + if not out: + return ScanResult(False, path, "trivy", note=f"trivy fs failed (rc={rc}): {err[:140]}") + return self._parse_trivy(out, path) + + def scan_rootfs(self, path: str = "/") -> ScanResult: + """Scan an entire root filesystem — OS packages + installed language deps. Run from inside a + container this inspects the container's FULL supply chain (the base image + everything we + installed), surfacing CVEs that a lockfile scan of the app dir misses.""" + if not shutil.which("trivy"): + return ScanResult(False, path, "trivy", note="trivy not installed") + rc, out, err = self._run(["trivy", "rootfs", "--quiet", "--format", "json", "--scanners", "vuln", path]) + if not out: + return ScanResult(False, path, "trivy", note=f"trivy rootfs failed (rc={rc}): {err[:140]}") + return self._parse_trivy(out, path) + + def scan_image(self, ref: str) -> ScanResult: + """Scan a container image reference for known CVEs.""" + if not shutil.which("trivy"): + return ScanResult(False, ref, "trivy", note="trivy not installed") + rc, out, err = self._run(["trivy", "image", "--quiet", "--format", "json", ref]) + if not out: + return ScanResult(False, ref, "trivy", note=f"trivy image failed (rc={rc}): {err[:140]}") + return self._parse_trivy(out, ref) + + def sbom(self, path: str = ".") -> dict: + """Produce a component inventory (SBOM) — Syft if present, else Trivy CycloneDX.""" + if shutil.which("syft"): + rc, out, _ = self._run(["syft", "-q", "-o", "syft-json", path]) + if out: + try: + arts = json.loads(out).get("artifacts", []) or [] + return {"ok": True, "tool": "syft", "components": len(arts), + "sample": [{"name": a.get("name"), "version": a.get("version"), "type": a.get("type")} + for a in arts[:20]]} + except Exception: # noqa: BLE001 + pass + if shutil.which("trivy"): + rc, out, _ = self._run(["trivy", "fs", "--quiet", "--format", "cyclonedx", path]) + if out: + try: + comps = json.loads(out).get("components", []) or [] + return {"ok": True, "tool": "trivy-cyclonedx", "components": len(comps), + "sample": [{"name": c.get("name"), "version": c.get("version")} for c in comps[:20]]} + except Exception: # noqa: BLE001 + pass + return {"ok": False, "note": "no SBOM tool available (install syft or trivy)"} + + def resolve_image(self, container: str) -> str: + rc, out, _ = self._run(["docker", "inspect", "-f", "{{.Config.Image}}", container]) + if rc == 0: + return out.strip() + return "" + + def scan_container(self, container: str) -> ScanResult: + image = self.resolve_image(container) + if not image: + return ScanResult(False, container, "trivy", note=f"could not resolve image for {container}") + res = self.scan_image(image) + res.target = f"{container} ({image})" + return res + + def list_scannable_containers(self, name_filter: str = "") -> list[dict]: + rc, out, _ = self._run(["docker", "ps", "--format", "{{.Names}}\t{{.Image}}"]) + if rc != 0: + return [] + rows = [] + for line in out.strip().splitlines(): + if not line: + continue + parts = line.split("\t", 1) + if len(parts) != 2: + continue + name, image = parts + if name_filter and name_filter not in name: + continue + rows.append({"name": name, "image": image}) + return rows + + def triage(self, result: ScanResult, top: int = 6) -> dict: + if not result.ok: + return {"summary": result.note, "fixes": [], "note": result.note} + ordered = sorted( + (f for f in result.findings if f.fixed), + key=lambda f: (_SEV_ORDER.get(f.severity, 99), f.id), + )[:top] + fixes = [ + { + "id": f.id, + "pkg": f.pkg, + "installed": f.installed, + "fixed": f.fixed, + "severity": f.severity, + "action": f"upgrade {f.pkg} {f.installed} → {f.fixed}", + } + for f in ordered + ] + crit = sum(1 for f in result.findings if f.severity == "CRITICAL") + high = sum(1 for f in result.findings if f.severity == "HIGH") + fixable = sum(1 for f in result.findings if f.fixed) + summary = f"{len(result.findings)} vulns ({crit} critical, {high} high); {fixable} fixable" + note = "" if fixable else result.note or "no fixable findings" + return {"summary": summary, "fixes": fixes, "note": note} + + def advise(self, result: ScanResult) -> dict: + """Turn the OS-vs-dependency split into an actionable recommendation. Most container CVEs + live in the base OS image, not the app's own code — so the highest-leverage fix is usually + hardening the base image (once, fleet-wide) rather than chasing individual CVEs.""" + if not result.ok: + return {"os": 0, "lang": 0, "recommendation": result.note} + fs = result.findings + os_n = sum(1 for f in fs if f.cls == "os") + lang_n = sum(1 for f in fs if f.cls == "lang") + fix_os = sum(1 for f in fs if f.cls == "os" and f.fixed) + fix_lang = sum(1 for f in fs if f.cls == "lang" and f.fixed) + os_pkgs = sorted({f.pkg for f in fs if f.cls == "os" and f.fixed and f.severity in ("CRITICAL", "HIGH")})[:4] + if not fs: + rec = "No known vulnerabilities — nothing to do." + elif os_n and os_n >= max(1, lang_n): + eg = (" (e.g. " + ", ".join(os_pkgs) + ")") if os_pkgs else "" + rec = (f"{os_n} of {len(fs)} findings are in the OS BASE IMAGE, not this app's own dependencies. " + f"The high-leverage fix is to HARDEN THE BASE, not chase individual CVEs: rebase onto a " + f"patched python:3.12-slim digest, or add `apt-get update && apt-get upgrade -y` for the " + f"flagged system packages{eg} in the Dockerfile — clearing ~{fix_os} of them, and fleet-wide " + f"since every agent shares the base image.") + if fix_lang: + rec += f" Separately, {fix_lang} of this app's own dependency CVE(s) are fixable by upgrading the pins." + elif lang_n: + rec = (f"The exposure is in this app's OWN dependencies ({lang_n} finding(s), {fix_lang} fixable) rather " + f"than the base image — upgrade the pinned versions listed above.") + else: + rec = "Findings have no fixed version yet — monitor upstream and re-scan after updates." + return {"os": os_n, "lang": lang_n, "fixable_os": fix_os, "fixable_lang": fix_lang, "recommendation": rec} diff --git a/context_runtime/integrations/vuln_db.py b/context_runtime/integrations/vuln_db.py new file mode 100644 index 0000000..edb8f9d --- /dev/null +++ b/context_runtime/integrations/vuln_db.py @@ -0,0 +1,129 @@ +"""VulnDB — read-only client for the shared vulnerability database (Apache Doris). + +Consumes the NVD/OSV/Snyk vuln DB that ``vuln_feeds`` populates (table ``vulns``). This is the +READ path for apps like Edge Sentinel: enrich scan findings and look up advisories/patched pins, +with the SAME row-scope + column-mask permissioning the enterprise ``DorisStore`` applies — a +principal only sees the sources it owns, and sensitive columns (``refs``) are masked for +non-privileged callers. The full policy engine lives in context-runtime-v3; the minimal +``Principal`` here mirrors it for the read side. + +pymysql is an OPTIONAL, deferred import (like the store adapters); the client degrades to +``available() == False`` when it's missing or Doris is unreachable — callers stay functional. +Connection comes from the environment (Doris FE, exposed to the host via a NodePort): + + DORIS_MYSQL_HOST / DORIS_FE_HOST (127.0.0.1) DORIS_QUERY_PORT (9030) + DORIS_USER (root) DORIS_PASSWORD ('') DORIS_DATABASE (context_runtime) +""" +from __future__ import annotations + +import os +from dataclasses import dataclass, field + +_SENSITIVE = ("refs",) # columns masked for non-privileged principals + + +@dataclass(frozen=True) +class Principal: + """Who's asking — mirrors the enterprise policy engine's Identity (read side).""" + roles: frozenset = field(default_factory=frozenset) + owns_rows_of: frozenset = field(default_factory=frozenset) # sources/owners this principal may read + + @property + def privileged(self) -> bool: + return bool({"admin", "security"} & set(self.roles)) + + +def _pymysql(): + try: + import pymysql + return pymysql + except Exception as e: # noqa: BLE001 + raise ImportError("VulnDB needs the optional `pymysql` dependency — `pip install pymysql`.") from e + + +class VulnDB: + def __init__(self, *, host: str | None = None, port: int | None = None, user: str | None = None, + password: str | None = None, database: str | None = None, timeout: float = 8.0): + self.host = host or os.getenv("DORIS_MYSQL_HOST") or os.getenv("DORIS_FE_HOST", "127.0.0.1") + self.port = int(port or os.getenv("DORIS_QUERY_PORT", "9030")) + self.user = user or os.getenv("DORIS_USER", "root") + self.password = password if password is not None else os.getenv("DORIS_PASSWORD", "") + self.database = database or os.getenv("DORIS_DATABASE", "context_runtime") + self.timeout = timeout + + # ---- connection (deferred, degrades) ------------------------------------- + def _query(self, sql: str, params: tuple = ()) -> list[dict]: + pymysql = _pymysql() + conn = pymysql.connect(host=self.host, port=self.port, user=self.user, password=self.password, + database=self.database, connect_timeout=self.timeout, + cursorclass=pymysql.cursors.DictCursor) + try: + with conn.cursor() as cur: + cur.execute(sql, params) + return list(cur.fetchall()) + finally: + conn.close() + + def available(self) -> bool: + try: + self._query("SELECT 1 AS ok") + return True + except Exception: # noqa: BLE001 + return False + + # ---- permissioning (same shape as the enterprise DorisStore) ------------- + @staticmethod + def _apply_policy(rows: list[dict], principal: Principal | None) -> list[dict]: + if principal is None or principal.privileged: + return rows + owns = principal.owns_rows_of + out = [] + for r in rows: + if not owns or not ({r.get("owner"), r.get("source")} & set(owns)): + continue # row scope: only owned sources + out.append({k: v for k, v in r.items() if k not in _SENSITIVE}) # column mask + return out + + # ---- reads --------------------------------------------------------------- + def lookup(self, *, package: str | None = None, cve: str | None = None, min_cvss: float | None = None, + principal: Principal | None = None, limit: int = 50) -> list[dict]: + sql = "SELECT * FROM vulns WHERE 1=1" + params: list = [] + if package: + sql += " AND package=%s" + params.append(package) + if cve: + sql += " AND cve_id=%s" + params.append(cve) + if min_cvss is not None: + sql += " AND cvss>=%s" + params.append(min_cvss) + sql += " ORDER BY cvss DESC LIMIT %s" + params.append(int(limit)) + return self._apply_policy(self._query(sql, tuple(params)), principal) + + def count(self, principal: Principal | None = None) -> int: + rows = self.lookup(principal=principal, limit=100000) + return len(rows) + + def enrich(self, findings, principal: Principal | None = None) -> dict: + """Given scan findings (objects/dicts with cve_id or id), return which are corroborated by our + vuln DB: {cve_id: [our records]}. Batched by the CVE ids present, then policy-filtered.""" + ids = [] + for f in findings or []: + cid = getattr(f, "id", None) or (f.get("id") if isinstance(f, dict) else None) or (f.get("cve_id") if isinstance(f, dict) else None) + if cid: + ids.append(cid) + ids = list(dict.fromkeys(ids))[:200] + if not ids: + return {} + placeholders = ",".join(["%s"] * len(ids)) + try: + rows = self._query(f"SELECT * FROM vulns WHERE cve_id IN ({placeholders})", tuple(ids)) + except Exception: # noqa: BLE001 + return {} + rows = self._apply_policy(rows, principal) + out: dict = {} + for r in rows: + out.setdefault(r["cve_id"], []).append(r) + return out diff --git a/context_runtime/integrations/vuln_feeds.py b/context_runtime/integrations/vuln_feeds.py new file mode 100644 index 0000000..1944cf6 --- /dev/null +++ b/context_runtime/integrations/vuln_feeds.py @@ -0,0 +1,369 @@ +"""vuln_feeds — vulnerability-feed ingestion for the Security & Compliance block. + +Pulls normalized vulnerability records from public feeds (NVD 2.0, OSV.dev, and — when a +token is present — Snyk) and hands them to a :class:`VulnStore` (backed by Doris in the +deployed stack). This is the *feed* half of the block; SupplyChainScanner is the pre-deploy +scanner half and Edge Sentinel triages the merged findings. + +All network I/O is factored through two private helpers (:func:`_get_json` / +:func:`_post_json`) built on stdlib ``urllib`` + ``json`` so tests can monkeypatch them and +the module carries no third-party deps. Every fetcher wraps its network call in try/except and +returns ``[]`` on any failure or missing key — a dead feed never breaks ingestion. +""" +from __future__ import annotations + +import json +import os +import urllib.parse +import urllib.request +from dataclasses import dataclass, field +from typing import Protocol, runtime_checkable + +_SEV_ORDER = {"CRITICAL": 0, "HIGH": 1, "MEDIUM": 2, "LOW": 3, "UNKNOWN": 4} + + +@dataclass(frozen=True) +class VulnRecord: + cve_id: str + source: str + package: str + ecosystem: str + severity: str + cvss: float + fixed_version: str + vulnerable_range: str + published: str + modified: str + summary: str + references: tuple = field(default_factory=tuple) + aliases: tuple = field(default_factory=tuple) + + +# --------------------------------------------------------------------------- network I/O + + +def _get_json(url: str, headers: dict | None = None) -> dict: + req = urllib.request.Request(url, headers=headers or {}, method="GET") + with urllib.request.urlopen(req, timeout=20) as resp: + return json.loads(resp.read().decode("utf-8")) + + +def _post_json(url: str, body: dict, headers: dict | None = None) -> dict: + data = json.dumps(body).encode("utf-8") + hdrs = {"Content-Type": "application/json", **(headers or {})} + req = urllib.request.Request(url, data=data, headers=hdrs, method="POST") + with urllib.request.urlopen(req, timeout=20) as resp: + return json.loads(resp.read().decode("utf-8")) + + +# --------------------------------------------------------------------------- helpers + + +def _band_from_score(score: float) -> str: + """Derive a CVSS v3 severity band from a numeric score.""" + if score >= 9.0: + return "CRITICAL" + if score >= 7.0: + return "HIGH" + if score >= 4.0: + return "MEDIUM" + if score > 0.0: + return "LOW" + return "UNKNOWN" + + +def _as_tuple(seq) -> tuple: + if not seq: + return () + return tuple(seq) + + +# --------------------------------------------------------------------------- NVD + + +def fetch_nvd(query: str, limit: int = 20) -> list[VulnRecord]: + """GET the NVD 2.0 keyword-search endpoint and normalize into VulnRecords.""" + url = ( + "https://services.nvd.nist.gov/rest/json/cves/2.0" + f"?keywordSearch={urllib.parse.quote(query)}&resultsPerPage={int(limit)}" + ) + headers: dict = {} + api_key = os.environ.get("NVD_API_KEY") + if api_key: + headers["apiKey"] = api_key + try: + data = _get_json(url, headers=headers) + out: list[VulnRecord] = [] + for item in data.get("vulnerabilities", [])[:limit]: + cve = item.get("cve", {}) + cve_id = cve.get("id", "") + if not cve_id: + continue + + summary = "" + for desc in cve.get("descriptions", []): + if desc.get("lang") == "en": + summary = desc.get("value", "") + break + + severity, cvss = "UNKNOWN", 0.0 + metrics = cve.get("metrics", {}) + for key in ("cvssMetricV31", "cvssMetricV30", "cvssMetricV2"): + entries = metrics.get(key) + if entries: + cdata = entries[0].get("cvssData", {}) + cvss = float(cdata.get("baseScore", 0.0) or 0.0) + severity = ( + cdata.get("baseSeverity") + or entries[0].get("baseSeverity") + or _band_from_score(cvss) + ) + break + + refs = tuple( + r.get("url", "") for r in cve.get("references", []) if r.get("url") + ) + out.append( + VulnRecord( + cve_id=cve_id, + source="nvd", + package="", + ecosystem="", + severity=str(severity).upper(), + cvss=cvss, + fixed_version="", + vulnerable_range="", + published=cve.get("published", ""), + modified=cve.get("lastModified", ""), + summary=summary, + references=refs, + aliases=(cve_id,), + ) + ) + return out + except Exception: + return [] + + +# --------------------------------------------------------------------------- OSV + + +def _osv_to_record(vuln: dict) -> VulnRecord | None: + vuln_id = vuln.get("id", "") + aliases = list(vuln.get("aliases", []) or []) + if not vuln_id and not aliases: + return None + + # Prefer a CVE identifier for cve_id, else fall back to the OSV id. + cve_id = vuln_id + for a in [vuln_id, *aliases]: + if a.startswith("CVE-"): + cve_id = a + break + + all_aliases = tuple(dict.fromkeys([vuln_id, *aliases])) + summary = vuln.get("summary") or vuln.get("details") or "" + + severity, cvss = "UNKNOWN", 0.0 + for sev in vuln.get("severity", []) or []: + if sev.get("type") == "CVSS_V3": + score = sev.get("score", "") + try: + cvss = float(score) + except (TypeError, ValueError): + # score is a CVSS vector string, not a number — leave 0.0 but band later. + cvss = 0.0 + severity = _band_from_score(cvss) + break + + refs = tuple( + r.get("url", "") for r in vuln.get("references", []) or [] if r.get("url") + ) + + package, ecosystem, fixed_version, vulnerable_range = "", "", "", "" + affected = vuln.get("affected", []) or [] + if affected: + first = affected[0] + pkg = first.get("package", {}) or {} + package = pkg.get("name", "") + ecosystem = pkg.get("ecosystem", "") + introduced, fixed = "", "" + for rng in first.get("ranges", []) or []: + for ev in rng.get("events", []) or []: + if "introduced" in ev and not introduced: + introduced = ev["introduced"] + if "fixed" in ev: + fixed = ev["fixed"] + fixed_version = fixed + if introduced or fixed: + lo = introduced or "0" + vulnerable_range = f">={lo}" + (f", <{fixed}" if fixed else "") + + return VulnRecord( + cve_id=cve_id, + source="osv", + package=package, + ecosystem=ecosystem, + severity=severity.upper(), + cvss=cvss, + fixed_version=fixed_version, + vulnerable_range=vulnerable_range, + published=vuln.get("published", ""), + modified=vuln.get("modified", ""), + summary=summary, + references=refs, + aliases=all_aliases, + ) + + +def fetch_osv(package: str, ecosystem: str, limit: int = 50) -> list[VulnRecord]: + """POST the OSV.dev query endpoint for a package/ecosystem pair.""" + body = {"package": {"name": package, "ecosystem": ecosystem}} + try: + data = _post_json("https://api.osv.dev/v1/query", body) + out: list[VulnRecord] = [] + for vuln in data.get("vulns", [])[:limit]: + rec = _osv_to_record(vuln) + if rec is not None: + out.append(rec) + return out + except Exception: + return [] + + +def fetch_osv_by_id(vuln_id: str) -> list[VulnRecord]: + """GET a single OSV advisory by id.""" + try: + data = _get_json(f"https://api.osv.dev/v1/vulns/{vuln_id}") + rec = _osv_to_record(data) + return [rec] if rec is not None else [] + except Exception: + return [] + + +# --------------------------------------------------------------------------- Snyk + + +def fetch_snyk(org_id: str = "", limit: int = 50) -> list[VulnRecord]: + """Best-effort GET of Snyk org issues; returns [] immediately when SNYK_TOKEN is unset.""" + token = os.environ.get("SNYK_TOKEN") + if not token: + return [] + try: + org = org_id or os.environ.get("SNYK_ORG_ID", "") + url = f"https://api.snyk.io/rest/orgs/{org}/issues?version=2024-01-01&limit={int(limit)}" + headers = {"Authorization": f"token {token}"} + data = _get_json(url, headers=headers) + out: list[VulnRecord] = [] + for item in data.get("data", [])[:limit]: + attrs = item.get("attributes", {}) or {} + cvss = 0.0 + severity = str(attrs.get("effective_severity_level", "unknown")).upper() + for sev in attrs.get("severities", []) or []: + try: + cvss = float(sev.get("score", 0.0) or 0.0) + except (TypeError, ValueError): + cvss = 0.0 + break + problems = attrs.get("problems", []) or [] + cve_id = item.get("id", "") + aliases = [] + for p in problems: + pid = p.get("id", "") + if pid: + aliases.append(pid) + if pid.startswith("CVE-"): + cve_id = pid + out.append( + VulnRecord( + cve_id=cve_id, + source="snyk", + package="", + ecosystem="", + severity=severity, + cvss=cvss, + fixed_version="", + vulnerable_range="", + published=attrs.get("created_at", ""), + modified=attrs.get("updated_at", ""), + summary=attrs.get("title", ""), + references=(), + aliases=tuple(dict.fromkeys(aliases)), + ) + ) + return out + except Exception: + return [] + + +# --------------------------------------------------------------------------- dedupe / ingest + + +def _keys(rec: VulnRecord) -> set: + """The set of identifiers a record is known by (CVE id + aliases).""" + ids = set(rec.aliases) + if rec.cve_id: + ids.add(rec.cve_id) + return {i for i in ids if i} + + +def _richer(a: VulnRecord, b: VulnRecord) -> VulnRecord: + """Prefer the record with more resolved data: a fixed_version, then higher CVSS.""" + if bool(a.fixed_version) != bool(b.fixed_version): + return a if a.fixed_version else b + if a.cvss != b.cvss: + return a if a.cvss > b.cvss else b + return a + + +def dedupe(records: list[VulnRecord]) -> list[VulnRecord]: + """Collapse records sharing a CVE id or overlapping aliases, preferring richer data.""" + kept: list[VulnRecord] = [] + kept_keys: list[set] = [] + for rec in records: + rkeys = _keys(rec) + match = -1 + for i, ks in enumerate(kept_keys): + if rkeys & ks: + match = i + break + if match == -1: + kept.append(rec) + kept_keys.append(set(rkeys)) + else: + kept[match] = _richer(kept[match], rec) + kept_keys[match] |= rkeys + return kept + + +@runtime_checkable +class VulnStore(Protocol): + def upsert_vulns(self, records: list[VulnRecord]) -> int: + ... + + def query_vulns(self, **filters) -> list[dict]: + ... + + +def ingest(records: list[VulnRecord], store: VulnStore) -> int: + """Dedupe then upsert into the store; returns the count reported by the store.""" + deduped = dedupe(records) + return store.upsert_vulns(deduped) + + +def collect( + *, + nvd_query: str | None = None, + osv_packages: list | None = None, + snyk: bool = False, +) -> list[VulnRecord]: + """Run the requested fetchers and return a single deduped, combined list.""" + records: list[VulnRecord] = [] + if nvd_query: + records.extend(fetch_nvd(nvd_query)) + for pkg in osv_packages or []: + name, ecosystem = pkg if isinstance(pkg, (tuple, list)) else (pkg, "") + records.extend(fetch_osv(name, ecosystem)) + if snyk: + records.extend(fetch_snyk()) + return dedupe(records) diff --git a/context_runtime/learning/__init__.py b/context_runtime/learning/__init__.py new file mode 100644 index 0000000..008e3ec --- /dev/null +++ b/context_runtime/learning/__init__.py @@ -0,0 +1,8 @@ +"""Async learning loop — outcome events → aggregator → learned-state snapshots (Whitepaper v3).""" +from __future__ import annotations + +from .aggregator import LearnedStateSnapshot, LearningAggregator +from .bus import EventBus, InMemoryBus +from .events import OutcomeEvent + +__all__ = ["OutcomeEvent", "EventBus", "InMemoryBus", "LearningAggregator", "LearnedStateSnapshot"] diff --git a/context_runtime/learning/aggregator.py b/context_runtime/learning/aggregator.py new file mode 100644 index 0000000..b5c575b --- /dev/null +++ b/context_runtime/learning/aggregator.py @@ -0,0 +1,106 @@ +"""Learning aggregator + state snapshot — the async policy loop (Whitepaper v3, Phase 4). + + execution → OutcomeEvent → [bus] → aggregator (off hot path) → LearnedStateSnapshot → [bus] → replicas + +The aggregator is the ONE writer of learned state: it drains outcome events, folds each into the shared +bandit (and, via optional sinks, the trust ledger / calibration), and publishes a versioned snapshot. +Stateless planner replicas never learn on the serving path — they select against a local snapshot and +reconcile when a newer one is published. "The learning loop is asynchronous — and that is what makes it +scale." +""" +from __future__ import annotations + +from dataclasses import dataclass, field +from typing import Callable + +from ..integrations.bandit import EpsilonGreedyBandit +from .bus import EventBus +from .events import OutcomeEvent + + +class _Arm: + __slots__ = ("key",) + + def __init__(self, key: str): + self.key = key + + +@dataclass +class LearnedStateSnapshot: + """A versioned, serializable snapshot of learned state that a replica reconciles against.""" + version: int + bandit_stats: dict = field(default_factory=dict) # ctx -> arm.key -> [n, mean] + + def to_dict(self) -> dict: + return {"version": self.version, "bandit_stats": self.bandit_stats} + + @classmethod + def from_dict(cls, d: dict) -> "LearnedStateSnapshot": + return cls(int(d["version"]), d.get("bandit_stats", {})) + + def apply_to(self, bandit: EpsilonGreedyBandit) -> None: + """Reconcile a replica's bandit with this snapshot (replace its learned means).""" + bandit.stats = {ctx: {k: list(v) for k, v in arms.items()} + for ctx, arms in self.bandit_stats.items()} + + +class LearningAggregator: + """Consumes OutcomeEvents off the serving path and folds them into the shared learned state.""" + + def __init__( + self, + bandit: EpsilonGreedyBandit, + *, + topic: str = "outcomes", + on_trust: Callable[[OutcomeEvent], None] | None = None, + on_calibration: Callable[[OutcomeEvent], None] | None = None, + ): + self.bandit = bandit + self.topic = topic + self.on_trust = on_trust # e.g. the enterprise TrustLedger sink + self.on_calibration = on_calibration + self.version = 0 + self.processed = 0 + self._last_seq = -1 + self._known: set[str] = {a.key for a in self.bandit.arms} + + def _ensure_arm(self, key: str) -> None: + """Register a plan shape into the shared bandit so update()/value() don't miss the key.""" + if key not in self._known: + self.bandit.arms = tuple(self.bandit.arms) + (_Arm(key),) + self._known.add(key) + + def apply(self, event: OutcomeEvent) -> bool: + """Fold one event into learned state. Ignores stale/duplicate seqs (idempotent replay).""" + if event.seq and event.seq <= self._last_seq: + return False + self._last_seq = max(self._last_seq, event.seq) + # An abstention produced no served answer to reward the arm with — it still carries trust signal. + if not event.abstained: + self._ensure_arm(event.arm) + self.bandit.update(event.context, _Arm(event.arm), event.reward) + if self.on_trust is not None: + self.on_trust(event) + if self.on_calibration is not None: + self.on_calibration(event) + self.processed += 1 + return True + + def drain(self, bus: EventBus) -> int: + """Consume all pending outcome events (off the hot path). Bumps the version if anything applied.""" + applied = 0 + for ev in bus.poll(self.topic): + if self.apply(ev): + applied += 1 + if applied: + self.version += 1 + return applied + + def snapshot(self) -> LearnedStateSnapshot: + stats = {ctx: {k: list(v) for k, v in arms.items()} for ctx, arms in self.bandit.stats.items()} + return LearnedStateSnapshot(self.version, stats) + + def publish(self, bus: EventBus, topic: str = "snapshots") -> LearnedStateSnapshot: + snap = self.snapshot() + bus.publish(topic, snap) + return snap diff --git a/context_runtime/learning/bus.py b/context_runtime/learning/bus.py new file mode 100644 index 0000000..f5845eb --- /dev/null +++ b/context_runtime/learning/bus.py @@ -0,0 +1,46 @@ +"""Event bus — the transport for the async learning loop (Whitepaper v3, Phase 4). + +The planner never learns on the hot path. Instead it publishes ``OutcomeEvent``s to a bus; an +aggregator consumes them off the serving path and republishes learned-state snapshots. The engine +depends only on this small ``EventBus`` Protocol, so the default in-process ``InMemoryBus`` (for a +single node / tests) and a Kafka binding (the paper's "nervous system", for a fleet) are +interchangeable — the two streams never touch the model's context window. +""" +from __future__ import annotations + +import threading +from collections import defaultdict, deque +from typing import Protocol, runtime_checkable + + +@runtime_checkable +class EventBus(Protocol): + def publish(self, topic: str, event) -> None: ... + def poll(self, topic: str) -> list: ... + + +class InMemoryBus: + """A thread-safe, in-process bus. ``poll`` drains all pending events for a topic (at-most-once, + FIFO). Sufficient for a single node and for deterministic tests; swap in a KafkaBus (same + publish/poll seam) to distribute the loop across stateless replicas.""" + + def __init__(self): + self._topics: dict[str, deque] = defaultdict(deque) + self._lock = threading.Lock() + + def publish(self, topic: str, event) -> None: + with self._lock: + self._topics[topic].append(event) + + def poll(self, topic: str) -> list: + with self._lock: + q = self._topics.get(topic) + if not q: + return [] + items = list(q) + q.clear() + return items + + def pending(self, topic: str) -> int: + with self._lock: + return len(self._topics.get(topic, ())) diff --git a/context_runtime/learning/events.py b/context_runtime/learning/events.py new file mode 100644 index 0000000..ec90add --- /dev/null +++ b/context_runtime/learning/events.py @@ -0,0 +1,55 @@ +"""Outcome events — what execution emits for the async learning loop (Whitepaper v3, Phase 4). + +Every served (or abstained) plan produces one ``OutcomeEvent``: the selection context + arm, the +measured reward, the selection propensity (for off-policy), and any operator/verification signals. The +event is published to a bus and folded into learned state OFF the serving path — it never touches the +model's context window. This is the "policy loop" half of the paper's Kafka nervous system. +""" +from __future__ import annotations + +from dataclasses import asdict, dataclass + + +@dataclass(frozen=True) +class OutcomeEvent: + context: str # the selection context (intent bucket) + arm: str # the chosen plan shape (method:tier) + reward: float # measured reward for this execution + representation: str = "" # v4: knowledge representation the plan routed to (attribution) + seq: int = 0 # monotonic sequence number (ordering / idempotency) + mode: str = "" # exploit | explore | shadow (from the bandit metadata) + p: float = 1.0 # selection propensity, for off-policy evaluation + abstained: bool = False # the planner declined to serve + confidence: float | None = None # calibrated confidence of the served plan + # optional operator/verification signals, for a trust sink + accepted: bool | None = None + overridden: bool = False + regenerated: bool = False + verified: bool | None = None + + def to_dict(self) -> dict: + return asdict(self) + + @classmethod + def from_dict(cls, d: dict) -> "OutcomeEvent": + return cls(**{k: v for k, v in d.items() if k in cls.__dataclass_fields__}) + + @classmethod + def from_plan(cls, plan, reward: float, *, seq: int = 0, **signals) -> "OutcomeEvent": + """Build an event from a served Plan (reading the optimizer's ``extra['bandit']`` / ``extra + ['abstention']``) plus the measured reward and any operator signals.""" + extra = getattr(plan, "extra", None) or {} + b = extra.get("bandit") or {} + ab = extra.get("abstention") or {} + return cls( + context=b.get("context", ""), + arm=b.get("arm", ""), + reward=reward, + representation=getattr(getattr(plan, "intent", None), "representation", "") or "", + seq=seq, + mode=b.get("mode", ""), + p=float(b.get("p", 1.0)), + abstained=ab.get("action") == "abstain", + confidence=ab.get("confidence"), + **signals, + ) diff --git a/context_runtime/learning/kafka_bus.py b/context_runtime/learning/kafka_bus.py new file mode 100644 index 0000000..1e30328 --- /dev/null +++ b/context_runtime/learning/kafka_bus.py @@ -0,0 +1,99 @@ +"""Kafka binding for the learning loop's EventBus (Whitepaper v3, Phase 4 — the "nervous system"). + +``KafkaEventBus`` satisfies the same ``publish`` / ``poll`` seam as ``InMemoryBus``, so the async policy +loop (and the context-freshness stream) run unchanged across stateless replicas over a real broker. The +engine stays dependency-free: ``kafka-python`` is imported lazily, only when this binding is constructed. + + bus = KafkaEventBus("broker:9092", codec_out=OutcomeEvent.to_dict, codec_in=OutcomeEvent.from_dict) + bus.publish("outcomes", event) # producer → topic + for ev in bus.poll("outcomes"): ... # consumer drains what's available + +Serialization is JSON by default; pass ``codec_out``/``codec_in`` to (de)serialize a specific event type +(e.g. ``OutcomeEvent``). Payloads that are already dicts round-trip as-is (handy for the CDC/vuln stream). +""" +from __future__ import annotations + +import json +from typing import Callable + + +class KafkaEventBus: + def __init__( + self, + bootstrap_servers: str, + *, + group_id: str = "context-runtime", + codec_out: Callable[[object], dict] | None = None, + codec_in: Callable[[dict], object] | None = None, + poll_timeout_ms: int = 800, + auto_offset_reset: str = "earliest", + client=None, # inject a fake (must expose .send/.flush and a consumer factory) for tests + ): + self.bootstrap_servers = bootstrap_servers + self.group_id = group_id + self.codec_out = codec_out or (lambda e: e.to_dict() if hasattr(e, "to_dict") else e) + self.codec_in = codec_in or (lambda d: d) + self.poll_timeout_ms = poll_timeout_ms + self.auto_offset_reset = auto_offset_reset + self._injected = client + self._producer = None + self._consumers: dict[str, object] = {} # one consumer per polled topic + if client is None: + try: + import kafka # noqa: F401 + except ImportError as e: # pragma: no cover - environment dependent + raise ImportError( + "KafkaEventBus needs the 'kafka-python' package (pip install kafka-python), " + "or inject a fake client for tests." + ) from e + + # ── producer ── + def _get_producer(self): + if self._injected is not None: + return self._injected + if self._producer is None: + from kafka import KafkaProducer + self._producer = KafkaProducer( + bootstrap_servers=self.bootstrap_servers, + value_serializer=lambda v: json.dumps(v).encode("utf-8"), + ) + return self._producer + + def publish(self, topic: str, event) -> None: + self._get_producer().send(topic, self.codec_out(event)) + + def flush(self) -> None: + p = self._get_producer() + if hasattr(p, "flush"): + p.flush() + + # ── consumer ── + def _get_consumer(self, topic: str): + if self._injected is not None: + return self._injected.consumer(topic) # fake factory in tests + if topic not in self._consumers: + from kafka import KafkaConsumer + self._consumers[topic] = KafkaConsumer( + topic, + bootstrap_servers=self.bootstrap_servers, + group_id=self.group_id, + auto_offset_reset=self.auto_offset_reset, + enable_auto_commit=True, + consumer_timeout_ms=self.poll_timeout_ms, + value_deserializer=lambda b: json.loads(b.decode("utf-8")), + ) + return self._consumers[topic] + + def poll(self, topic: str) -> list: + """Drain the records available within one poll window and decode them (at-least-once).""" + consumer = self._get_consumer(topic) + out = [] + for record in consumer: # bounded by consumer_timeout_ms + out.append(self.codec_in(record.value)) + return out + + def close(self) -> None: + if self._producer is not None: + self._producer.close() + for c in self._consumers.values(): + c.close() diff --git a/context_runtime/optimizer/knapsack.py b/context_runtime/optimizer/knapsack.py index 856a897..a2daa4d 100644 --- a/context_runtime/optimizer/knapsack.py +++ b/context_runtime/optimizer/knapsack.py @@ -4,36 +4,73 @@ PlanScore. The "knapsack" proper is the token-budget packing inside context assembly (``runtime``); here selection is greedy-by-utility, which is the right move when the candidate set is small. CP-SAT replaces this in v0.2 when constraints interact. + +Governance seam: an optional ``policy`` (PolicyProvider) narrows the feasible space *before* +cost ranking — a policy-rejected candidate is infeasible regardless of its score, and its reason +is recorded in ``Plan.rejected`` (observable in EXPLAIN). An optional ``trust`` (TrustProvider) +breaks ties between equally cost-ranked feasible plans. Both default to None (OSS runs standalone); +the enterprise layer injects concrete PolicyEngine / TrustLedger adapters. """ from __future__ import annotations from ..constraints.hard import feasible -from ..plugins.base import CostEstimator +from ..plugins.base import CostEstimator, PolicyProvider, TrustProvider from ..types import Candidate, Goal, Plan, PlanScore from dataclasses import replace class KnapsackOptimizer: - def __init__(self, estimator: CostEstimator): + def __init__( + self, + estimator: CostEstimator, + policy: PolicyProvider | None = None, + trust: TrustProvider | None = None, + ): self.estimator = estimator + self.policy = policy + self.trust = trust + + def _policy_reject(self, candidate: Candidate, goal: Goal, score: PlanScore) -> str | None: + """None if policy-feasible (or no policy injected); else the rejection reason.""" + if self.policy is None: + return None + return self.policy.feasible(candidate, goal, score) def score(self, candidate: Candidate, goal: Goal) -> PlanScore: s = self.estimator.estimate(candidate, goal) ok, _ = feasible(candidate, s, goal.constraints) + if ok and self._policy_reject(candidate, goal, s) is not None: + ok = False return replace(s, feasible=ok) - def select(self, scored: list[tuple[Candidate, PlanScore]], goal: Goal) -> Plan: + def _utility(self, goal: Goal): + """Rank key: cost/quality total first, trust score as the tie-breaker. Trust defaults to a + constant when no provider is injected, so selection is identical to the pre-seam behavior.""" + trust = self.trust + def key(cs: tuple[Candidate, PlanScore]) -> tuple[float, float]: + cand, sc = cs + t = trust.score(cand, goal) if trust is not None else 0.0 + return (sc.total, t) + return key + + def select(self, scored: list[tuple[Candidate, PlanScore]], goal: Goal, context: str = "") -> Plan: + # ``context`` is accepted for interface parity with the online optimizer; the static + # cost optimizer does not use it (selection is a pure function of score + feasibility). rejected: list[tuple[Candidate, str]] = [] feasible_set: list[tuple[Candidate, PlanScore]] = [] for cand, sc in scored: ok, reason = feasible(cand, sc, goal.constraints) + if ok: + preason = self._policy_reject(cand, goal, sc) + if preason is not None: + ok, reason = False, f"policy: {preason}" if ok: feasible_set.append((cand, sc)) else: rejected.append((cand, reason or "infeasible")) pool = feasible_set or scored # never fail to produce a plan; mark infeasible - chosen, chosen_score = max(pool, key=lambda cs: cs[1].total) + chosen, chosen_score = max(pool, key=self._utility(goal)) for cand, sc in pool: if cand is not chosen: rejected.append((cand, f"lower score {sc.total:.3f} < {chosen_score.total:.3f}")) diff --git a/context_runtime/optimizer/online.py b/context_runtime/optimizer/online.py new file mode 100644 index 0000000..8d08d27 --- /dev/null +++ b/context_runtime/optimizer/online.py @@ -0,0 +1,208 @@ +"""Online-learning optimizer — Generation 4 (Whitepaper v3). + +Generations 1-2 pick the plan with the best *estimated* cost/quality. But estimates go stale, so a +Gen-4 planner must **explore**: it treats plan selection as a contextual bandit keyed by intent bucket, +seeded with the cost model's estimate as an optimistic prior and refined by measured reward. Selection +balances exploitation (the best learned plan) against exploration (trying alternatives), and every +choice records its selection probability so alternatives can be evaluated off-policy — the cheapest +exploration tier. + +Layered exploration (paper's "cheapest form that answers the question"): + • Tier 1 — off-policy evaluation: score alternative plans from executions already observed, using the + logged selection probabilities (``offpolicy_values``). Zero user-facing latency. + • Tier 2 — shadow execution: run a candidate off the serving path (``select(..., shadow=True)`` marks + a shadow pick; the caller executes it asynchronously and feeds the reward back via ``learn``). + • Tier 3 — live exploration: the ε-greedy branch actually serves an exploratory plan. + +Reuses the fleet's ``EpsilonGreedyBandit`` as the shared per-(context, arm) value store, and composes +with the Phase-0 governance seam: a PolicyProvider still filters the feasible space *before* any +exploration (you never explore an infeasible plan), and a TrustProvider breaks ties. +""" +from __future__ import annotations + +from dataclasses import replace + +from ..constraints.hard import feasible as _hard_feasible +from ..integrations.bandit import EpsilonGreedyBandit +from ..plugins.base import CostEstimator, PolicyProvider, TrustProvider +from ..types import Candidate, Goal, Intent, Plan, PlanScore + + +class _Arm: + """Minimal Arm (exposes ``.key``) so plan shapes can be registered into the shared bandit.""" + __slots__ = ("key",) + + def __init__(self, key: str): + self.key = key + + +def plan_key(candidate: Candidate) -> str: + """The bandit arm for a candidate: its retrieval method × model tier — the decision that matters.""" + method = next((s.params.get("method", "") for s in candidate.steps if s.type == "retrieve"), "") + return f"{method}:{candidate.model_tier}" + + +class BanditOptimizer: + """A CostOptimizer that selects online. ``score`` is identical to the cost optimizer (estimate + + hard/policy feasibility); ``select`` runs a subset-aware ε-greedy over the *feasible* candidates.""" + + def __init__( + self, + estimator: CostEstimator, + *, + bandit: EpsilonGreedyBandit | None = None, + epsilon: float = 0.15, + discount: float = 0.0, + policy: PolicyProvider | None = None, + trust: TrustProvider | None = None, + trust_weight: float = 0.0, + abstain_gate=None, + context_of=None, + seed: int = 0x1234567, + ): + self.estimator = estimator + # discount > 0 → recency-weighted learning (tracks a drifting best arm); see EpsilonGreedyBandit. + self.bandit = bandit or EpsilonGreedyBandit(arms=(), epsilon=epsilon, discount=discount) + self.policy = policy + self.trust = trust + # Gen-5: when > 0, trust is a weighted OBJECTIVE term in ranking, not just a tie-breaker. + self.trust_weight = trust_weight + # Gen-5: an optional AbstentionGate — records serve/escalate/abstain in Plan.extra. + self.abstain_gate = abstain_gate + self.context_of = context_of # goal -> context str; else the caller passes context= + self._rng = seed & 0xFFFFFFFF + self._known: set[str] = {a.key for a in self.bandit.arms} + + # ── deterministic rng (xorshift), independent of the bandit's own ── + def _rand(self) -> float: + x = self._rng + x ^= (x << 13) & 0xFFFFFFFF + x ^= x >> 17 + x ^= (x << 5) & 0xFFFFFFFF + self._rng = x & 0xFFFFFFFF + return self._rng / 0x100000000 + + def _ensure_arm(self, key: str) -> None: + """Register a newly-seen plan shape into the shared bandit so value()/update() don't KeyError.""" + if key not in self._known: + self.bandit.arms = tuple(self.bandit.arms) + (_Arm(key),) + self._known.add(key) + + def _context(self, goal: Goal, context: str) -> str: + if context: + return context + if self.context_of: + return self.context_of(goal) + return "default" + + def _policy_reject(self, cand: Candidate, goal: Goal, sc: PlanScore) -> str | None: + return self.policy.feasible(cand, goal, sc) if self.policy is not None else None + + def score(self, candidate: Candidate, goal: Goal) -> PlanScore: + s = self.estimator.estimate(candidate, goal) + ok, _ = _hard_feasible(candidate, s, goal.constraints) + if ok and self._policy_reject(candidate, goal, s) is not None: + ok = False + return replace(s, feasible=ok) + + def _arm_value(self, ctx: str, cand: Candidate, sc: PlanScore, goal: Goal) -> tuple[float, float]: + """Value to rank by: the learned mean once observed, else the cost-model total as an optimistic + prior. With ``trust_weight`` (Gen 5) trust is folded into the objective; it is also the secondary + tie-breaker (the Phase-0 seam shape).""" + key = plan_key(cand) + self._ensure_arm(key) + n, mean = self.bandit.value(ctx, key) + base = mean if n > 0 else sc.total + t = self.trust.score(cand, goal) if self.trust is not None else 0.0 + if self.trust_weight > 0.0 and self.trust is not None: + base = (1.0 - self.trust_weight) * base + self.trust_weight * t + return (base, t) + + def select(self, scored, goal: Goal, context: str = "", *, shadow: bool = False) -> Plan: + ctx = self._context(goal, context) + + rejected: list[tuple[Candidate, str]] = [] + pool: list[tuple[Candidate, PlanScore]] = [] + for cand, sc in scored: + ok, reason = _hard_feasible(cand, sc, goal.constraints) + if ok: + pr = self._policy_reject(cand, goal, sc) + if pr is not None: + ok, reason = False, f"policy: {pr}" + (pool if ok else rejected).append((cand, sc) if ok else (cand, reason or "infeasible")) + + eff_pool = pool or [(c, s) for c, s in scored] # never fail to produce a plan + m = len(eff_pool) + + # greedy (exploit) reference — argmax learned/estimated value + greedy_i = max(range(m), key=lambda i: self._arm_value(ctx, eff_pool[i][0], eff_pool[i][1], goal)) + + explore = m > 1 and self._rand() < self.bandit.epsilon + if explore: + # Tier 3 (or Tier 2 shadow): pick a NON-greedy feasible arm uniformly. + others = [i for i in range(m) if i != greedy_i] + chosen_i = others[int(self._rand() * len(others)) % len(others)] + mode = "shadow" if shadow else "explore" + else: + chosen_i, mode = greedy_i, "exploit" + + chosen, chosen_sc = eff_pool[chosen_i] + + # selection probability (propensity) under this ε-greedy policy — logged for off-policy eval + eps = self.bandit.epsilon + p = (eps / m) + ((1.0 - eps) if chosen_i == greedy_i else 0.0) + + for i, (cand, sc) in enumerate(eff_pool): + if i != chosen_i: + rejected.append((cand, f"not selected (value rank; mode={mode})")) + + extra = {"bandit": {"context": ctx, "arm": plan_key(chosen), "mode": mode, + "p": round(p, 6), "epsilon": eps}} + # Gen-5 honest abstention: is the chosen plan confident enough to serve? + if self.abstain_gate is not None: + v = self.abstain_gate.evaluate(chosen_sc, goal) + extra["abstention"] = {"action": v.action, "confidence": round(v.confidence, 6), "reason": v.reason} + + return Plan( + intent=Intent(bucket="unknown"), # runtime attaches the real intent + chosen=chosen, + score=chosen_sc, + rejected=tuple(rejected), + extra=extra, + ) + + # ── learning ── + def learn(self, context: str, candidate_or_key, reward: float) -> None: + """Fold a measured reward into the bandit for (context, arm). Accepts a Candidate or an arm key.""" + key = candidate_or_key if isinstance(candidate_or_key, str) else plan_key(candidate_or_key) + self._ensure_arm(key) + self.bandit.update(context, _Arm(key), reward) + + def learn_from_plan(self, plan: Plan, reward: float) -> None: + """Convenience: learn using the (context, arm) recorded in ``plan.extra['bandit']``.""" + b = (plan.extra or {}).get("bandit") + if b: + self.learn(b["context"], b["arm"], reward) + + +def offpolicy_values(logs) -> dict[str, dict[str, tuple[float, float]]]: + """Tier-1 off-policy evaluation: estimate each arm's value per context from already-observed + executions, WITHOUT running anything. Self-normalized inverse-propensity weighting over the logged + selection probabilities corrects for the fact that exploratory data isn't uniformly sampled. + + ``logs``: iterable of dicts ``{"context", "arm", "reward", "p"}`` (as recorded in Plan.extra). + Returns ``{context: {arm: (effective_weight, estimated_value)}}`` — higher value = better plan, + letting the planner rank alternatives it never served on the live path. + """ + acc: dict[str, dict[str, list[float]]] = {} + for row in logs: + ctx, arm = row["context"], row["arm"] + p = max(float(row.get("p", 1.0)), 1e-6) # guard divide-by-zero + w = 1.0 / p + cell = acc.setdefault(ctx, {}).setdefault(arm, [0.0, 0.0]) # [sum_w, sum_w*reward] + cell[0] += w + cell[1] += w * float(row["reward"]) + out: dict[str, dict[str, tuple[float, float]]] = {} + for ctx, arms in acc.items(): + out[ctx] = {a: (sw, (swr / sw if sw else 0.0)) for a, (sw, swr) in arms.items()} + return out diff --git a/context_runtime/planner/candidates.py b/context_runtime/planner/candidates.py index e31f02e..0935905 100644 --- a/context_runtime/planner/candidates.py +++ b/context_runtime/planner/candidates.py @@ -7,7 +7,11 @@ from __future__ import annotations from ..types import Candidate, Goal, Intent, PluginInfo, StepSpec -from . import rules +from . import representations, rules + +# below this intent confidence the generator WIDENS beyond the chosen representation so the +# bandit can explore other representations and learn which one wins for this context. +EXPLORE_CONFIDENCE = 0.5 class RuleCandidateGenerator: @@ -16,8 +20,26 @@ def __init__(self, default_top_k: int = 50, final_k: int = 8, target_tokens: int self.final_k = final_k self.target_tokens = target_tokens + def _methods_for_intent(self, intent: Intent, bucket_methods: tuple) -> tuple: + """v4: constrain to the chosen knowledge representation instead of the flat bucket table. + + Representation-first, with two safety valves: a document (`hybrid`) fallback so a missing or + infeasible representation engine degrades gracefully (the cost model prunes the infeasible + one), and confidence-gated widening so uncertain intents still explore across representations + (that exploration is how the bandit LEARNS representation selection).""" + rep = getattr(intent, "representation", "document") + # the bucket's own methods that already belong to the chosen representation, in bucket order… + primary = tuple(m for m in bucket_methods if representations.representation_for(m) == rep) + if not primary: # …else the representation's own methods + primary = representations.methods_for(rep) + fallback = () if rep == "document" else ("hybrid",) + widen = bucket_methods if intent.confidence < EXPLORE_CONFIDENCE else () + methods = tuple(dict.fromkeys(primary + fallback + widen)) # de-dupe, representation-first + return methods or bucket_methods + def generate(self, intent: Intent, goal: Goal) -> list[Candidate]: - methods, strategy, want_verify = rules.BUCKET_DEFAULTS[intent.bucket] + bucket_methods, strategy, want_verify = rules.BUCKET_DEFAULTS[intent.bucket] + methods = self._methods_for_intent(intent, bucket_methods) tiers = rules.BUCKET_TIERS[intent.bucket] c = goal.constraints # citations are checked by the verify step, so requiring them implies verify diff --git a/context_runtime/planner/candidates_registry.py b/context_runtime/planner/candidates_registry.py new file mode 100644 index 0000000..28a6867 --- /dev/null +++ b/context_runtime/planner/candidates_registry.py @@ -0,0 +1,68 @@ +"""Registry-backed candidate generator — the same plan shapes as RuleCandidateGenerator, but the +retrieval methods and model tiers are drawn from a CapabilityRegistry instead of the hard-coded rule +tables. Adding a capability (a new retrieval method, a new tier) becomes a ``registry.register()`` call +that immediately widens the planner's search space — with no edit to rules.py and nothing added to the +model's context window (Whitepaper v3: "capabilities live in the planner, not the prompt"). + +The reasoning strategy and the verify requirement remain per-intent planner policy (from rules); this +generator governs which *capabilities* are composed, which is the part that scales. +""" +from __future__ import annotations + +from ..registry import CapabilityRegistry +from ..types import Candidate, Goal, Intent, PluginInfo, StepSpec +from . import rules + + +class RegistryCandidateGenerator: + def __init__( + self, + registry: CapabilityRegistry | None = None, + default_top_k: int = 50, + final_k: int = 8, + target_tokens: int = 3000, + ): + self.registry = registry or CapabilityRegistry.from_rules() + self.default_top_k = default_top_k + self.final_k = final_k + self.target_tokens = target_tokens + + def generate(self, intent: Intent, goal: Goal) -> list[Candidate]: + _methods, strategy, want_verify = rules.BUCKET_DEFAULTS.get( + intent.bucket, rules.BUCKET_DEFAULTS["unknown"] + ) + methods = self.registry.values_for(intent.bucket, "retrieval") or ["hybrid"] + tiers = self.registry.values_for(intent.bucket, "model_tier") or ["local"] + + c = goal.constraints + require_verify = want_verify or c.require_verification or c.require_citations + + out: list[Candidate] = [] + for method in methods: + for tier in tiers: + steps: list[StepSpec] = [ + StepSpec("retrieve", {"method": method, "top_k": self.default_top_k}), + ] + if method in ("hybrid", "vector", "code"): + steps.append(StepSpec("rerank", {"final_k": self.final_k})) + steps.append(StepSpec("compress", {"target_tokens": self.target_tokens})) + steps.append(StepSpec("route", {"tier": tier})) + steps.append(StepSpec("reason", {"strategy": strategy, "capability": "synthesis"})) + if require_verify: + steps.append(StepSpec("verify", {"method": "citation"})) + out.append(Candidate(steps=tuple(steps), model_tier=tier)) + return out + + def prune(self, candidates: list[Candidate], goal: Goal) -> list[Candidate]: + c = goal.constraints + kept: list[Candidate] = [] + for cand in candidates: + if c.sensitivity == "restricted" and cand.model_tier != "local": + continue + if c.require_citations and not any(s.type == "verify" for s in cand.steps): + continue + kept.append(cand) + return kept or candidates[:1] + + def info(self) -> PluginInfo: + return PluginInfo(name="registry_candidates", kind="planner", capabilities=frozenset({"candidates"})) diff --git a/context_runtime/planner/intent.py b/context_runtime/planner/intent.py index bde6d52..faea097 100644 --- a/context_runtime/planner/intent.py +++ b/context_runtime/planner/intent.py @@ -6,7 +6,7 @@ from __future__ import annotations from ..types import Goal, Intent, PluginInfo -from . import rules +from . import representations, rules class RuleIntentAnalyzer: @@ -22,12 +22,15 @@ def analyze(self, goal: Goal) -> Intent: confidence = 0.8 if bucket != "unknown" else 0.3 if entities: confidence = min(1.0, confidence + 0.1) + # v4: the first decision — which knowledge representation is this question even about? + representation = representations.classify(bucket, goal.text, entities) return Intent( bucket=bucket, entities=entities, risk=risk, # type: ignore[arg-type] normalized=rules.normalize(goal.text), confidence=confidence, + representation=representation, ) def info(self) -> PluginInfo: diff --git a/context_runtime/planner/representations.py b/context_runtime/planner/representations.py new file mode 100644 index 0000000..5cf841a --- /dev/null +++ b/context_runtime/planner/representations.py @@ -0,0 +1,79 @@ +"""Knowledge representations — the planner's first-class decision axis (Whitepaper v4). + +Context Runtime does not assume every problem is document retrieval. A request is first +mapped to the *knowledge representation* that can answer it — a document, a graph of +relationships, a bi-temporal fact history, an analytical/OLAP cube, code, or media — and +only then to a concrete ``Retrieval`` method within that representation. Retrieval is one +specialization; the planner routes across representations without the application changing. + +This module holds the (deliberately small, inspectable) mapping from retrieval method → +representation, so plans, traces and the cost model can reason about *which representation* +was chosen — not just which method. The heavy engines behind each representation (a hosted +Graphiti temporal store, an OLAP connector, a real HippoRAG graph) bind at construction time +via the ``RetrieverPlugin`` seam; this table is only the taxonomy. +""" +from __future__ import annotations + +import re + +from ..types import IntentBucket, KnowledgeRepresentation, Retrieval + +# method → the knowledge representation it operates over +REPRESENTATION_OF: dict[Retrieval, KnowledgeRepresentation] = { + "vector": "document", "bm25": "document", "hybrid": "document", "file": "document", + "graph": "graph", + "community": "community", + "temporal": "temporal", + "code": "code", + "sql": "analytical", "logs": "analytical", "api": "analytical", + "image": "multimodal", "colpali": "multimodal", "video": "multimodal", +} + + +def representation_for(method: Retrieval | str) -> KnowledgeRepresentation: + """The knowledge representation a retrieval method operates over (default: document).""" + return REPRESENTATION_OF.get(method, "document") # type: ignore[arg-type] + + +def methods_for(representation: KnowledgeRepresentation | str) -> tuple[Retrieval, ...]: + """The retrieval methods that specialize a given representation.""" + return tuple(m for m, r in REPRESENTATION_OF.items() if r == representation) + + +# ─── the classify head: intent → representation (v4's first-class decision) ─── +# A bucket carries a *default* representation (the buckets are already representation-leaning: +# multi_hop→graph, temporal→temporal, code→code). Content HINTS override it — most importantly +# they reach representations no bucket produces on its own (analytical OLAP, multimodal). +BUCKET_REPRESENTATION: dict[IntentBucket, KnowledgeRepresentation] = { + "multi_hop": "graph", + "temporal": "temporal", + "code_reasoning": "code", + # everything else defaults to document unless a hint says otherwise +} + +# Ordered: first hit wins. These are the signals a bucket alone misses. +HINT_RULES: list[tuple[re.Pattern, KnowledgeRepresentation]] = [ + (re.compile(r"\b(how\s+many|number\s+of|count\s+of|\btotal\b|\bsum\b|average|avg|median|" + r"per\s+(day|week|month|quarter|user|account)|over\s+the\s+(last|past)\s+\w+|" + r"top\s+\d+|rank(ed|ing)?|group\s+by|distribution\s+of|trend|breakdown|" + r"month[- ]over[- ]month|year[- ]over[- ]year|\bmrr\b|\barr\b|conversion\s+rate)\b", + re.I), "analytical"), + (re.compile(r"\b(screenshot|screen\s?shot|image|photo|picture|diagram|figure|chart\s+image|" + r"scanned|scan\s+of|receipt|invoice\s+image|whiteboard|slide)\b", re.I), "multimodal"), + (re.compile(r"\b(as\s+of|point[- ]in[- ]time|what\s+changed|history\s+of|superseded|" + r"no\s+longer|used\s+to\s+be|previously|valid\s+(from|until))\b", re.I), "temporal"), + (re.compile(r"\b(related\s+to|connected\s+to|depend(s|ency|encies)?\s+on|graph\s+of|" + r"network\s+of|linked\s+to|relationship\s+between|traverse|multi[- ]hop)\b", + re.I), "graph"), +] + + +def classify(bucket: IntentBucket | str, text: str, entities: tuple[str, ...] = ()) -> KnowledgeRepresentation: + """Map an analysed intent to the knowledge representation that can answer it. + + Content hints win over the bucket default, so 'how many incidents last week' routes to the + analytical representation even though its bucket is 'incident'. Default: document.""" + for pat, rep in HINT_RULES: + if pat.search(text or ""): + return rep + return BUCKET_REPRESENTATION.get(bucket, "document") # type: ignore[arg-type] diff --git a/context_runtime/planner/rules.py b/context_runtime/planner/rules.py index fe0d0b2..617c8bb 100644 --- a/context_runtime/planner/rules.py +++ b/context_runtime/planner/rules.py @@ -14,6 +14,16 @@ INTENT_RULES: list[tuple[re.Pattern, IntentBucket, str]] = [ (re.compile(r"\b(error|exception|stack ?trace|code|status)\s*[:#]?\s*\w*\d", re.I), "exact_lookup", "low"), (re.compile(r"\b[A-Z]{2,}-\d+\b"), "exact_lookup", "low"), # JIRA-123, ERR-500 + # temporal: the answer depends on WHEN — point-in-time state, what changed, provenance, + # supersession. Placed before multi_hop so a temporal cue wins over an incidental + # "between X and Y". The bi-temporal store answers these; semantic retrieval cannot. + (re.compile(r"\b(as\s+of|at\s+the\s+time|over\s+time|point[- ]in[- ]time|as\s+at|" + r"when\s+(did|was|were|is|does)|was\s+(true|active|valid|the\s+case)|" + r"no\s+longer|used\s+to|previously|originally|currently|still\b|" + r"supersed\w+|replac\w+|revis\w+|deprecat\w+|(? list[Candidate]: ... @runtime_checkable class CostOptimizer(Protocol): def score(self, candidate: Candidate, goal: Goal) -> PlanScore: ... - def select(self, scored: list[tuple[Candidate, PlanScore]], goal: Goal) -> Plan: ... + # ``context`` (the intent bucket) is optional: static optimizers ignore it; an online + # (Gen-4) optimizer uses it as the contextual-bandit key. Kept keyword-defaulted for + # backward compatibility with callers that pass only (scored, goal). + def select(self, scored: list[tuple[Candidate, PlanScore]], goal: Goal, context: str = "") -> Plan: ... + + +# ──────────────────────────── governance seam (enterprise open-core) ──────────────────────────── +# The optimizer optionally consults two injected providers. Both are OPTIONAL and default to None, +# so the OSS engine runs standalone; the commercial layer (context-runtime-v3) supplies concrete +# implementations that adapt its PolicyEngine / TrustLedger to these Protocols. The engine never +# imports the enterprise types — it depends only on the interface, exactly like every other plugin. + + +@runtime_checkable +class PolicyProvider(Protocol): + """Policy-Constrained Planning (Whitepaper v3): filter the feasible execution space BEFORE cost + ranking. Return ``None`` if the candidate is policy-feasible, else a short human-readable reason + (surfaced in ``Plan.rejected`` → EXPLAIN's "rejected alternatives + why"). The score is provided + so cost/budget policies can bind to the real estimate. A policy-violating plan is infeasible + regardless of its estimated quality or cost — authorization is a first-class planning constraint.""" + + def feasible(self, candidate: Candidate, goal: Goal, score: PlanScore) -> str | None: ... + + +@runtime_checkable +class TrustProvider(Protocol): + """Trust-Aware Execution (Whitepaper v3, Generation 5): a [0, 1] trust score for a candidate in + the context of this goal. The optimizer uses it to break ties between equally cost-ranked feasible + plans (and, in the full Gen-5 build, folds it into the objective). Unobserved ⇒ a neutral prior.""" + + def score(self, candidate: Candidate, goal: Goal) -> float: ... # ──────────────────────────── §3.1 cost estimator ──────────────────────────── diff --git a/context_runtime/policy/__init__.py b/context_runtime/policy/__init__.py new file mode 100644 index 0000000..8f3854a --- /dev/null +++ b/context_runtime/policy/__init__.py @@ -0,0 +1,22 @@ +"""Policy Runtime — the feasible-execution-space plane (commands, rules, enforcement). + +See docs/policy-runtime.md. Policy defines the feasible execution space before planning and enforces +across phases; commands read/write rules (long-term memory); every decision emits a PolicyDecision. +""" +from __future__ import annotations + +from .commands import Command, CommandRegistry, parse_args, role_gate +from .plane import ( + ACTIONS, ALLOW, ApprovalProvider, Decision, DecisionSink, GuardrailProvider, Policy, PolicyDecision, + current_policy, set_default_policy, +) +from .store import Rule, RuleStore, rule_id, scopes_for +from .factories import global_policy_commands, policy_commands + +__all__ = [ + "Rule", "RuleStore", "rule_id", "scopes_for", + "Decision", "ALLOW", "ACTIONS", "PolicyDecision", "DecisionSink", + "GuardrailProvider", "ApprovalProvider", "Policy", "set_default_policy", "current_policy", + "Command", "CommandRegistry", "parse_args", "role_gate", + "policy_commands", "global_policy_commands", +] diff --git a/context_runtime/policy/commands.py b/context_runtime/policy/commands.py new file mode 100644 index 0000000..365bc2d --- /dev/null +++ b/context_runtime/policy/commands.py @@ -0,0 +1,127 @@ +"""Slash-command framework (docs/policy-runtime.md §7). Dual-path dispatch at the input boundary: +input starting with ``/`` → a deterministic, no-LLM, permission-gated handler; otherwise → the agent loop. +Commands read/write the RuleStore (long-term memory) and reply with the resulting state. The permission +check (``can``) is injected — the enterprise ``command_gate`` satisfies it from the PermissionsPlane. +""" +from __future__ import annotations + +import shlex +from dataclasses import dataclass, field +from typing import Callable + + +@dataclass(frozen=True) +class Command: + name: str + run: Callable[[dict, object], dict] # (args, principal) -> {"text", "data"?, "ok"?} + description: str + usage: str = "" + requires: str = "" # capability/role; "" = any (unrestricted) + aliases: tuple[str, ...] = () + + +def parse_args(argstr: str) -> dict: + """`--key value` / `--flag` → args[key]; the rest → args['text'] and comma-split args['items'].""" + args: dict = {"_raw": argstr} + try: + toks = shlex.split(argstr) + except ValueError: + toks = argstr.split() + rest: list[str] = [] + i = 0 + while i < len(toks): + t = toks[i] + if t.startswith("--"): + key = t[2:] + if i + 1 < len(toks) and not toks[i + 1].startswith("--"): + args[key] = toks[i + 1] + i += 2 + else: + args[key] = True + i += 1 + else: + rest.append(t) + i += 1 + text = " ".join(rest).strip() + args["text"] = text + args["items"] = [s.strip() for s in text.split(",") if s.strip()] + return args + + +def role_gate(*, admin_roles=frozenset({"admin", "security", "policy-admin"})): + """A simple role-based ``can(principal, requires)`` for deployments without the enterprise plane: + ``""`` = anyone · ``self`` = any authenticated user · ``policy-admin`` = an admin role · + ``role:X`` / ``X`` = the principal carries role X.""" + def can(principal, requires: str) -> bool: + roles = getattr(principal, "roles", frozenset()) or frozenset() + user = getattr(principal, "user", "") or "" + if not requires: + return True + if requires == "self": + return bool(user) + if requires == "policy-admin": + return bool(admin_roles & set(roles)) + if requires.startswith("role:"): + return requires.split(":", 1)[1] in roles + return requires in roles + return can + + +class CommandRegistry: + def __init__(self, *, can: Callable[[object, str], bool] | None = None): + self._cmds: dict[str, Command] = {} + self._alias: dict[str, str] = {} + self.can = can or (lambda principal, requires: not requires) + self.audit: list[dict] = [] + + def register(self, cmd: Command) -> Command: + self._cmds[cmd.name] = cmd + for a in cmd.aliases: + self._alias[a] = cmd.name + return cmd + + def register_all(self, cmds) -> None: + for c in cmds: + self.register(c) + + def is_command(self, text: str) -> bool: + return (text or "").lstrip().startswith("/") + + def parse(self, text: str) -> tuple[str, str]: + s = (text or "").lstrip() + if s.startswith("/"): + s = s[1:] + head, _, rest = s.partition(" ") + return head.strip().lower(), rest.strip() + + def _resolve(self, name: str) -> Command | None: + return self._cmds.get(name) or self._cmds.get(self._alias.get(name, "")) + + def dispatch(self, text: str, principal=None) -> dict: + name, argstr = self.parse(text) + if name in ("help", ""): + return self._help(principal) + cmd = self._resolve(name) + if cmd is None: + return {"text": f"Unknown command /{name}. Try /help.", "ok": False} + if cmd.requires and not self.can(principal, cmd.requires): + self.audit.append({"cmd": name, "user": getattr(principal, "user", None), "allowed": False, + "requires": cmd.requires}) + return {"text": f"You don't have permission to run /{name} (requires: {cmd.requires}).", "ok": False, + "policy": [{"decision": "deny", "phase": "command", "reason": f"requires {cmd.requires}", + "summary": f"/{name} → denied"}]} + args = parse_args(argstr) + self.audit.append({"cmd": name, "user": getattr(principal, "user", None), "allowed": True, "args": args}) + out = cmd.run(args, principal) + out.setdefault("ok", True) + return out + + def _help(self, principal) -> dict: + visible = [c for c in self._cmds.values() if not c.requires or self.can(principal, c.requires)] + visible.sort(key=lambda c: c.name) + lines = [f"/{c.name}{(' ' + c.usage) if c.usage else ''} — {c.description}" for c in visible] + return {"text": "Commands you can run:\n" + "\n".join(lines) if lines else "No commands available.", + "data": {"commands": [c.name for c in visible]}, "ok": True} + + def names(self) -> list[str]: + return list(self._cmds) diff --git a/context_runtime/policy/factories.py b/context_runtime/policy/factories.py new file mode 100644 index 0000000..1b4395c --- /dev/null +++ b/context_runtime/policy/factories.py @@ -0,0 +1,64 @@ +"""Ready-made command sets over a RuleStore (docs/policy-runtime.md §8). Generic + reusable: the +permission `requires` strings are declared here; the injected `can` (enterprise `command_gate`) enforces +them. `policy_commands` manages a rule tier (global by default); `rule_commands` is the app-rule variant. +""" +from __future__ import annotations + +from .commands import Command +from .store import RuleStore + + +def _fmt(rules) -> str: + return "\n".join(f"{r.id} [{r.kind}·{r.action}] {r.text}" for r in rules) or "(none)" + + +def policy_commands(store: RuleStore, *, scope: str = "global", requires: str = "policy-admin", + default_kind: str = "guardrail", prefix: str = "policy", + read_aliases: tuple = (), write_aliases: tuple = ()) -> list[Command]: + """`/show` (read, open) + `/add` `/remove` `/modify` (write, gated).""" + + def _show(args, p): + rules = store.list(scope=scope, kind=args.get("kind")) + return {"text": "Policy rules:\n" + _fmt(rules), "data": {"rules": [r.to_dict() for r in rules]}} + + def _add(args, p): + text = args.get("text", "").strip() + if not text: + return {"text": f"Usage: /add{prefix} [--kind guardrail|approval] " + "[--action deny|redact|require_approval] [--phase input|output] [--tool ] [--regex]", + "ok": False} + match = {} + for k in ("phase", "tool"): + if args.get(k): + match[k] = args[k] + if args.get("regex"): + match["regex"] = True + r = store.add(scope, args.get("kind", default_kind), text, action=args.get("action", "deny"), + match=match, created_by=getattr(p, "user", "") or "") + return {"text": f"Added {r.id} ({r.kind}·{r.action}): {r.text}", "data": {"id": r.id}} + + def _remove(args, p): + rid = args.get("text", "").strip() + ok = store.remove(rid, scope=scope) + return {"text": f"Removed {rid}." if ok else f"No rule {rid} in this scope.", "ok": ok} + + def _modify(args, p): + parts = args.get("text", "").split(" ", 1) + if len(parts) < 2: + return {"text": f"Usage: /modify{prefix} ", "ok": False} + r = store.modify(parts[0].strip(), scope=scope, text=parts[1].strip()) + return {"text": f"Updated {parts[0]}." if r else f"No rule {parts[0]}.", "ok": bool(r)} + + return [ + Command(f"show{prefix}", _show, f"show the {scope} policy rules", requires="", aliases=read_aliases), + Command(f"add{prefix}", _add, f"add a {scope} policy rule", usage=" [--kind …] [--action …]", + requires=requires, aliases=write_aliases), + Command(f"remove{prefix}", _remove, f"remove a {scope} policy rule by id", usage="", requires=requires), + Command(f"modify{prefix}", _modify, f"modify a {scope} policy rule", usage=" ", requires=requires), + ] + + +def global_policy_commands(store: RuleStore) -> list[Command]: + """The universal Global-Policy set every app registers (privileged to write; read is open).""" + return policy_commands(store, scope="global", requires="policy-admin", + read_aliases=("showguardrails",), write_aliases=("addguardrail",)) diff --git a/context_runtime/policy/plane.py b/context_runtime/policy/plane.py new file mode 100644 index 0000000..fc6ea95 --- /dev/null +++ b/context_runtime/policy/plane.py @@ -0,0 +1,182 @@ +"""The Policy Runtime plane — providers, decisions, and the mandatory PolicyDecision audit. + +Thesis (docs/policy-runtime.md §0): policy is not post-processing; it defines the feasible execution +space before planning, and enforces across phases (planning · command · input · tool · retrieval · output). +This module carries the runtime-phase providers (input/output guardrails, tool approval) and the composed +``Policy`` that evaluates them and emits a ``PolicyDecision`` for EVERY decision so policy is visible in +EXPLAIN, not hidden in enforcement. The planning-phase check lives in the optimizer (PolicyEngine.feasible). +""" +from __future__ import annotations + +import itertools +import re +import time +from collections import deque +from dataclasses import dataclass, field + +from .store import RuleStore, scopes_for + +ACTIONS = ("allow", "deny", "redact", "flag", "require_approval") + + +@dataclass(frozen=True) +class Decision: + action: str = "allow" + reason: str = "" + rule_id: str = "" + scope: str = "" + provider: str = "" + replacement: str | None = None # for redact + + @property + def ok(self) -> bool: + return self.action == "allow" + + @property + def blocked(self) -> bool: + return self.action == "deny" + + +ALLOW = Decision("allow") + +_ids = itertools.count(1) + + +@dataclass(frozen=True) +class PolicyDecision: + """The audit/EXPLAIN event emitted for every policy decision (docs §4).""" + policy_decision_id: str + principal: str + app: str + scope: str + rule_id: str + decision: str + reason: str + phase: str # planning | command | input | tool | retrieval | output + provider: str = "" + at: float = 0.0 + + def to_dict(self) -> dict: + return {"policy_decision_id": self.policy_decision_id, "principal": self.principal, "app": self.app, + "scope": self.scope, "rule_id": self.rule_id, "decision": self.decision, "reason": self.reason, + "phase": self.phase, "provider": self.provider, "at": self.at} + + def summary(self) -> str: + return f"{self.decision} · {self.phase}" + (f" · {self.reason}" if self.reason else "") + + +class DecisionSink: + """Collects PolicyDecisions for audit/EXPLAIN, with an optional forward callback (e.g. a trace/OLAP).""" + + def __init__(self, forward=None, cap: int = 2000): + self.events: deque = deque(maxlen=cap) + self.forward = forward + + def emit(self, ev: PolicyDecision) -> PolicyDecision: + self.events.append(ev) + if self.forward is not None: + try: + self.forward(ev) + except Exception: # noqa: BLE001 + pass + return ev + + def recent(self, n: int = 20, *, phase: str | None = None, decision: str | None = None) -> list[PolicyDecision]: + out = [e for e in self.events + if (phase is None or e.phase == phase) and (decision is None or e.decision == decision)] + return out[-n:] + + +def _matches(rule, text: str) -> bool: + pat = rule.match.get("pattern") or rule.text + if rule.match.get("regex"): + try: + return re.search(pat, text or "", re.I) is not None + except re.error: + return False + return pat.lower() in (text or "").lower() + + +class GuardrailProvider: + """Input/output content safety over ``guardrail``-kind rules (global + app scope). Pattern-based v1; + a rule with ``meta.semantic`` is a v2 LLM-judged extension behind the same seam.""" + + def __init__(self, store: RuleStore): + self.store = store + + def evaluate(self, principal, phase: str, text: str, *, app: str = "") -> Decision: + for scope in scopes_for(principal, app): + for r in self.store.list(scope=scope, kind="guardrail"): + if r.match.get("phase", "both") in (phase, "both") and _matches(r, text): + return Decision(r.action, f"guardrail: {r.text}", r.id, r.scope, "guardrail", + replacement=r.meta.get("replacement", "[redacted]")) + return ALLOW + + +class ApprovalProvider: + """Marks a tool call ``require_approval`` when an ``approval``-kind rule matches the tool (irreversible + ops: send email, delete, charge, CRM/financial update, block IP).""" + + def __init__(self, store: RuleStore): + self.store = store + + def evaluate(self, principal, tool: str, args: dict, *, app: str = "") -> Decision: + for scope in scopes_for(principal, app): + for r in self.store.list(scope=scope, kind="approval"): + want = r.match.get("tool") + if want in (None, "", tool) or (r.match.get("pattern") and _matches(r, tool)): + return Decision("require_approval", f"approval required: {r.text}", r.id, r.scope, "approval") + return ALLOW + + +class Policy: + """Composes the providers, evaluates the right one per phase, and emits a PolicyDecision for each.""" + + def __init__(self, *, store: RuleStore | None = None, guardrail: GuardrailProvider | None = None, + approval: ApprovalProvider | None = None, sink: DecisionSink | None = None, app: str = ""): + self.store = store + self.guardrail = guardrail or (GuardrailProvider(store) if store else None) + self.approval = approval or (ApprovalProvider(store) if store else None) + self.sink = sink or DecisionSink() + self.app = app + + def _emit(self, principal, phase: str, d: Decision) -> PolicyDecision: + return self.sink.emit(PolicyDecision( + policy_decision_id=f"pd-{next(_ids)}", principal=getattr(principal, "user", "") or "", + app=self.app, scope=d.scope, rule_id=d.rule_id, decision=d.action, reason=d.reason, + phase=phase, provider=d.provider, at=time.time())) + + def check(self, principal, phase: str, payload, *, app: str | None = None) -> Decision: + """``payload`` is the text (input/output) or a ``(tool, args)`` tuple (tool phase).""" + a = self.app if app is None else app + d = ALLOW + if phase in ("input", "output") and self.guardrail is not None: + d = self.guardrail.evaluate(principal, phase, payload if isinstance(payload, str) else "", app=a) + elif phase == "tool" and self.approval is not None: + tool, args = (payload if isinstance(payload, (tuple, list)) and len(payload) == 2 else (str(payload), {})) + d = self.approval.evaluate(principal, tool, args or {}, app=a) + self._emit(principal, phase, d) # every decision is audited (docs §4) + return d + + def redact(self, text: str, decision: Decision) -> str: + if decision.action == "redact" and self.store is not None: + for scope in scopes_for(None, self.app): + for r in self.store.list(scope=scope, kind="guardrail"): + if r.action == "redact" and _matches(r, text): + pat = r.match.get("pattern") or re.escape(r.text) + text = re.sub(pat, r.meta.get("replacement", "[redacted]"), text, + flags=re.I if not r.match.get("regex") else re.I) + return text + + +# ── fleet-wide install (like set_default_authorizer) ── +_default_policy = None + + +def set_default_policy(policy) -> None: + global _default_policy + _default_policy = policy + + +def current_policy(): + return _default_policy diff --git a/context_runtime/policy/store.py b/context_runtime/policy/store.py new file mode 100644 index 0000000..95daabd --- /dev/null +++ b/context_runtime/policy/store.py @@ -0,0 +1,160 @@ +"""Rules = the long-term memory of the Policy Runtime (see docs/policy-runtime.md §5). + +A ``Rule`` is a persisted, scoped, hard fact — distinct from chat history. Three tiers by scope: +``global`` (universal policy), ```` (app rules), ``:`` (user rules). Backed by JSONL files, +atomic writes, so a change is visible to the next request. This is the source of truth for v1; it can be +projected into the OLAP / snapshot layers later. +""" +from __future__ import annotations + +import hashlib +import json +import os +import time +from dataclasses import dataclass, field +from pathlib import Path + + +@dataclass +class Rule: + id: str # stable content id (r-…) — the handle for modify/remove + scope: str # "global" | "" | ":" + kind: str # "guardrail" | "approval" | "budget" | "target" | "rule" | … + text: str # the rule content (phrase/regex/instruction/limit) + action: str = "deny" # allow | deny | redact | flag | require_approval + match: dict = field(default_factory=dict) # {"regex":bool,"phase":..,"tool":..,"pattern":..} + meta: dict = field(default_factory=dict) + created_by: str = "" + at: float = 0.0 + + @property + def tier(self) -> str: + return "global" if self.scope == "global" else ("user" if ":" in self.scope else "app") + + def to_dict(self) -> dict: + return {"id": self.id, "scope": self.scope, "kind": self.kind, "text": self.text, + "action": self.action, "match": self.match, "meta": self.meta, + "created_by": self.created_by, "at": self.at} + + @classmethod + def from_dict(cls, d: dict) -> "Rule": + return cls(id=d["id"], scope=d["scope"], kind=d.get("kind", "rule"), text=d.get("text", ""), + action=d.get("action", "deny"), match=d.get("match", {}), meta=d.get("meta", {}), + created_by=d.get("created_by", ""), at=d.get("at", 0.0)) + + +def rule_id(scope: str, kind: str, text: str) -> str: + return "r-" + hashlib.sha1(f"{scope}|{kind}|{text}".encode("utf-8")).hexdigest()[:6] + + +class RuleStore: + """Per-scope JSONL rule files under ``dir``. ``global.jsonl`` / ``/app.jsonl`` / + ``/users/.jsonl``.""" + + def __init__(self, dir: str | None = None): + self.dir = dir or os.getenv("POLICY_DIR", "/data/context-runtime/policy") + + # ── scope → file ── + def _path(self, scope: str) -> Path: + base = Path(self.dir) + if scope == "global": + return base / "global.jsonl" + if ":" in scope: + app, user = scope.split(":", 1) + return base / _safe(app) / "users" / f"{_safe(user)}.jsonl" + return base / _safe(scope) / "app.jsonl" + + def _read(self, scope: str) -> list[Rule]: + p = self._path(scope) + if not p.exists(): + return [] + out: list[Rule] = [] + for line in p.read_text(encoding="utf-8").splitlines(): + if line.strip(): + try: + out.append(Rule.from_dict(json.loads(line))) + except Exception: # noqa: BLE001 + pass + return out + + def _write(self, scope: str, rules: list[Rule]) -> None: + p = self._path(scope) + p.parent.mkdir(parents=True, exist_ok=True) + tmp = str(p) + ".tmp" + Path(tmp).write_text("".join(json.dumps(r.to_dict()) + "\n" for r in rules), encoding="utf-8") + os.replace(tmp, p) + + # ── CRUD ── + def add(self, scope: str, kind: str, text: str, *, action: str = "deny", match: dict | None = None, + meta: dict | None = None, created_by: str = "") -> Rule: + rules = self._read(scope) + rid = rule_id(scope, kind, text) + rules = [r for r in rules if r.id != rid] # idempotent by content + rule = Rule(id=rid, scope=scope, kind=kind, text=text, action=action, match=match or {}, + meta=meta or {}, created_by=created_by, at=time.time()) + rules.append(rule) + self._write(scope, rules) + return rule + + def list(self, *, scope: str | None = None, kind: str | None = None) -> list[Rule]: + rules = self._read(scope) if scope is not None else self._all() + return [r for r in rules if kind is None or r.kind == kind] + + def get(self, rule_id: str, *, scope: str | None = None) -> Rule | None: + return next((r for r in self.list(scope=scope) if r.id == rule_id), None) + + def remove(self, rule_id: str, *, scope: str) -> bool: + rules = self._read(scope) + kept = [r for r in rules if r.id != rule_id] + if len(kept) == len(rules): + return False + self._write(scope, kept) + return True + + def modify(self, rule_id: str, *, scope: str, text: str | None = None, action: str | None = None, + match: dict | None = None) -> Rule | None: + rules = self._read(scope) + for i, r in enumerate(rules): + if r.id == rule_id: + if text is not None: + r.text = text + if action is not None: + r.action = action + if match is not None: + r.match = match + rules[i] = r + self._write(scope, rules) + return r + return None + + def _all(self) -> list[Rule]: + base = Path(self.dir) + out: list[Rule] = [] + if not base.exists(): + return out + for f in base.rglob("*.jsonl"): + if f.name.endswith(".tmp"): + continue + for line in f.read_text(encoding="utf-8").splitlines(): + if line.strip(): + try: + out.append(Rule.from_dict(json.loads(line))) + except Exception: # noqa: BLE001 + pass + return out + + +def _safe(s: str) -> str: + import re + return re.sub(r"[^a-z0-9_-]+", "_", (s or "default").lower()) + + +def scopes_for(principal, app: str) -> list[str]: + """The rule scopes that apply to a request: global always, then the app, then the app:user tier.""" + out = ["global"] + if app: + out.append(app) + user = getattr(principal, "user", "") or "" + if user: + out.append(f"{app}:{user}") + return out diff --git a/context_runtime/registry/__init__.py b/context_runtime/registry/__init__.py new file mode 100644 index 0000000..1321c9a --- /dev/null +++ b/context_runtime/registry/__init__.py @@ -0,0 +1,6 @@ +"""Capability registry — the out-of-band catalog the planner selects from (Whitepaper v3).""" +from __future__ import annotations + +from .capabilities import Capability, CapabilityRegistry + +__all__ = ["Capability", "CapabilityRegistry"] diff --git a/context_runtime/registry/capabilities.py b/context_runtime/registry/capabilities.py new file mode 100644 index 0000000..a19ad44 --- /dev/null +++ b/context_runtime/registry/capabilities.py @@ -0,0 +1,118 @@ +"""Capability Registry — capabilities live in the PLANNER, not the prompt (Whitepaper v3). + +The paper's pivotal scaling property: an assistant's abilities (retrieval methods, model tiers, +reasoners, tools) are registered *out of band*, each with cost/quality priors and the intent buckets +it serves. The planner selects among them using intent, policy, and cost — the model never sees a tool +catalog in its context window. "Capability count scales in a registry — which is cheap — rather than +in the context window, which is not." + +This module is the catalog. ``RegistryCandidateGenerator`` (planner/candidates_registry.py) consults it +to enumerate candidate plans, so adding a capability is a ``register()`` call, not an edit to the rule +tables or the prompt. ``CapabilityRegistry.from_rules()`` lifts the existing v0.1 rule tables into +registry form, so the default registry is behavior-equivalent to the hand-written planner. +""" +from __future__ import annotations + +from dataclasses import dataclass, field + + +@dataclass(frozen=True) +class Capability: + """One registered ability the planner may compose into a plan. + + ``value`` is the concrete token that lands in a StepSpec (a retrieval method like ``"hybrid"``, a + model tier like ``"premium"``). ``buckets`` are the intent buckets it serves (``"*"`` = all). The + priors are relative [0, 1] hints the cost model / planner can use before real statistics exist. + ``local`` marks an in-house resource (a policy hint — e.g. a local model tier vs. a hosted one). + """ + key: str # unique id, e.g. "retrieval:hybrid" | "model_tier:local" + kind: str # "retrieval" | "model_tier" | "reasoner" | "tool" + value: str # the token used in a StepSpec / model_tier + buckets: frozenset[str] = frozenset({"*"}) # intent buckets served + cost_prior: float = 0.5 # relative $ prior [0,1] (higher = pricier) + quality_prior: float = 0.5 # relative quality prior [0,1] (higher = better) + local: bool = False # served by an in-house / local resource + tags: frozenset[str] = frozenset() + meta: dict = field(default_factory=dict) + + def serves(self, bucket: str) -> bool: + return "*" in self.buckets or bucket in self.buckets + + +class CapabilityRegistry: + """An ordered catalog of Capabilities the planner selects from, keyed by ``Capability.key``.""" + + def __init__(self, capabilities: tuple[Capability, ...] | list[Capability] = ()): + self._by_key: dict[str, Capability] = {} + for c in capabilities: + self.register(c) + + def register(self, cap: Capability) -> "CapabilityRegistry": + """Add (or replace) a capability. Idempotent by key — re-registering updates in place.""" + self._by_key[cap.key] = cap + return self + + def get(self, key: str) -> Capability | None: + return self._by_key.get(key) + + def list(self, kind: str | None = None) -> list[Capability]: + return [c for c in self._by_key.values() if kind is None or c.kind == kind] + + def for_intent(self, bucket: str, kind: str) -> list[Capability]: + """Capabilities of ``kind`` that serve ``bucket``, best-quality first.""" + return sorted( + (c for c in self.list(kind) if c.serves(bucket)), + key=lambda c: (-c.quality_prior, c.cost_prior, c.key), + ) + + def values_for(self, bucket: str, kind: str) -> list[str]: + return [c.value for c in self.for_intent(bucket, kind)] + + # ──────────────────────────── defaults ──────────────────────────── + + @classmethod + def from_rules(cls) -> "CapabilityRegistry": + """Lift the v0.1 rule tables (planner/rules.py) into registry form. + + Retrieval methods and model tiers become Capabilities whose ``buckets`` are exactly the buckets + that reference them in ``BUCKET_DEFAULTS`` / ``BUCKET_TIERS``. This keeps a single source of truth + (the rules) while exposing it as a registry a caller can extend with ``register()``. + """ + from ..planner import rules + + # relative priors (used before measured statistics exist) + _METHOD_PRIOR = { # method: (cost, quality) + "bm25": (0.10, 0.50), "vector": (0.30, 0.70), "hybrid": (0.40, 0.80), + "code": (0.40, 0.75), "graph": (0.60, 0.85), "community": (0.70, 0.80), + } + _TIER_PRIOR = { # tier: (cost, quality, local) + "local": (0.10, 0.55, True), "cheap": (0.30, 0.72, False), "premium": (0.90, 0.95, False), + } + + reg = cls() + # retrieval capabilities: a method serves every bucket whose defaults list it + method_buckets: dict[str, set[str]] = {} + for bucket, (methods, _strategy, _verify) in rules.BUCKET_DEFAULTS.items(): + for m in methods: + method_buckets.setdefault(m, set()).add(bucket) + for method, buckets in method_buckets.items(): + cost, quality = _METHOD_PRIOR.get(method, (0.5, 0.6)) + reg.register(Capability( + key=f"retrieval:{method}", kind="retrieval", value=method, + buckets=frozenset(buckets), cost_prior=cost, quality_prior=quality, local=True, + tags=frozenset({"retrieval"}), + )) + + # model-tier capabilities: a tier serves every bucket whose tier list includes it + tier_buckets: dict[str, set[str]] = {} + for bucket, tiers in rules.BUCKET_TIERS.items(): + for t in tiers: + tier_buckets.setdefault(t, set()).add(bucket) + for tier, buckets in tier_buckets.items(): + cost, quality, local = _TIER_PRIOR.get(tier, (0.5, 0.6, False)) + reg.register(Capability( + key=f"model_tier:{tier}", kind="model_tier", value=tier, + buckets=frozenset(buckets), cost_prior=cost, quality_prior=quality, local=local, + tags=frozenset({"model"}), + )) + return reg diff --git a/context_runtime/runtime/runtime.py b/context_runtime/runtime/runtime.py index 669aafb..cbdf798 100644 --- a/context_runtime/runtime/runtime.py +++ b/context_runtime/runtime/runtime.py @@ -109,7 +109,9 @@ def _make_plan(self, goal: Goal): intent = self.intent.analyze(goal) cands = self.candidates.prune(self.candidates.generate(intent, goal), goal) scored: list[tuple[Candidate, PlanScore]] = [(c, self.optimizer.score(c, goal)) for c in cands] - plan = self.optimizer.select(scored, goal) + # pass the intent bucket as the selection context — an online (Gen-4) optimizer keys its + # contextual bandit on it; the static optimizer ignores it. + plan = self.optimizer.select(scored, goal, context=intent.bucket) plan = replace(plan, intent=intent) # attach the real intent return plan, scored, intent diff --git a/context_runtime/tools/base.py b/context_runtime/tools/base.py index a34f16a..ecd7157 100644 --- a/context_runtime/tools/base.py +++ b/context_runtime/tools/base.py @@ -80,12 +80,40 @@ def decide(self, spec: ToolSpec, args: dict) -> tuple[bool, str]: return False, "side-effecting tool requires approval" +# ── data-access authorization seam (enterprise open-core) ───────────────────────────────────────── +# An optional authorizer runs on every tool call BEFORE the tool executes: given the caller's principal +# (opaque to the engine), the ToolSpec, and the args, it returns None to allow or a short deny reason. +# The engine imports no enterprise code — it calls this callable, exactly like the PolicyProvider seam. +# ``set_default_authorizer`` installs one fleet-wide, so every AgentConsole app is gated by a single call. +Authorizer = "Callable[[object, ToolSpec, dict], str | None]" + +_default_authorizer = None + + +def set_default_authorizer(fn) -> None: + """Install a process-wide authorizer used by every ToolRegistry that wasn't given its own.""" + global _default_authorizer + _default_authorizer = fn + + +import contextvars as _contextvars + +_principal_ctx = _contextvars.ContextVar("cr_current_principal", default=None) + + +def current_principal(): + """The principal for the in-flight tool call (set by ``ToolRegistry.run``), or None. Lets a tool + scope its behavior to the caller — e.g. per-user config — without changing the ToolPlugin interface.""" + return _principal_ctx.get() + + class ToolRegistry: """Register tools; run them through the approval gate; describe them to a model.""" - def __init__(self, policy: ApprovalPolicy | None = None): + def __init__(self, policy: ApprovalPolicy | None = None, authorizer=None): self._tools: dict[str, ToolPlugin] = {} self.policy = policy or ApprovalPolicy() + self.authorizer = authorizer # data-access gate; falls back to the process default self.audit: list[dict] = [] # append-only record of every gated decision def register(self, tool: ToolPlugin) -> ToolPlugin: @@ -104,22 +132,34 @@ def specs(self) -> list[dict]: """OpenAI-style tool specs for a tool-calling model.""" return [t.spec().openai() for t in self._tools.values()] - def run(self, name: str, args: dict | None = None) -> ToolResult: + def run(self, name: str, args: dict | None = None, principal=None) -> ToolResult: try: tool = self.get(name) except ToolError as e: return ToolResult(ok=False, error=str(e), text=f"[error] {e}") spec = tool.spec() args = args or {} - ok, reason = self.policy.decide(spec, args) - if spec.side_effecting or spec.approval_required: - self.audit.append({"tool": name, "args": args, "allowed": ok, "reason": reason}) - if not ok: - return ToolResult(ok=False, error=f"denied: {reason}", text=f"[blocked] {name}: {reason}") + # expose the caller to the tool for the duration of the call (per-user scoping), reset after. + _token = _principal_ctx.set(principal) try: + # data-access authorization (enterprise): gate the tool by who is asking, before side-effect + # approval. None principal + no authorizer ⇒ unchanged behavior. + authorizer = self.authorizer or _default_authorizer + if authorizer is not None: + deny = authorizer(principal, spec, args) + if deny: + self.audit.append({"tool": name, "args": args, "allowed": False, "reason": deny}) + return ToolResult(ok=False, error=f"denied: {deny}", text=f"[blocked] {name}: {deny}") + ok, reason = self.policy.decide(spec, args) + if spec.side_effecting or spec.approval_required: + self.audit.append({"tool": name, "args": args, "allowed": ok, "reason": reason}) + if not ok: + return ToolResult(ok=False, error=f"denied: {reason}", text=f"[blocked] {name}: {reason}") return tool.run(args) except Exception as e: # a tool crash is a failed result, not a runtime crash return ToolResult(ok=False, error=f"{type(e).__name__}: {e}", text=f"[error] {name}: {e}") + finally: + _principal_ctx.reset(_token) def function_tool(name: str, fn: Callable[[dict], ToolResult], description: str = "", diff --git a/context_runtime/tools/mcp.py b/context_runtime/tools/mcp.py new file mode 100644 index 0000000..2bac820 --- /dev/null +++ b/context_runtime/tools/mcp.py @@ -0,0 +1,259 @@ +"""MCP adapter: mount external MCP tool servers into the agent-harness registry gated by ApprovalPolicy. + +Minimal JSON-RPC 2.0 client for the Model Context Protocol. Supports stdio (subprocess) +and HTTP (httpx POST) transports. Tools are wrapped as ToolPlugins so side-effects are +routed through the registry's ApprovalPolicy. Do not list tools at import time. +""" + +from __future__ import annotations + +import json +import os +import subprocess +from dataclasses import dataclass +from typing import Any + +import httpx + +from .base import ToolPlugin, ToolRegistry, ToolResult, ToolSpec + + +class MCPClient: + """Minimal JSON-RPC 2.0 MCP client. + + Create via classmethods: + MCPClient.stdio(command, env=None) + MCPClient.http(base_url, headers=None) + """ + + def __init__(self) -> None: + self._req_id: int = 0 + self._stdio_proc: subprocess.Popen[str] | None = None + self._http_base: str | None = None + self._http_headers: dict[str, str] = {} + self._initialized: bool = False + self._closed: bool = False + + @classmethod + def stdio(cls, command: list[str], env: dict[str, str] | None = None) -> MCPClient: + self = cls() + env_dict = os.environ.copy() + if env: + env_dict.update({k: str(v) for k, v in env.items()}) + self._stdio_proc = subprocess.Popen( + command, + stdin=subprocess.PIPE, + stdout=subprocess.PIPE, + stderr=subprocess.PIPE, + text=True, + bufsize=1, + env=env_dict, + ) + return self + + @classmethod + def http(cls, base_url: str, headers: dict[str, str] | None = None) -> MCPClient: + self = cls() + self._http_base = base_url + self._http_headers = dict(headers or {}) + return self + + def _next_id(self) -> int: + self._req_id += 1 + return self._req_id + + def _exchange(self, req: dict[str, Any]) -> dict[str, Any]: + """Send req and return the JSON-RPC response dict (may contain 'error').""" + if self._closed: + return {"jsonrpc": "2.0", "id": req.get("id"), "error": {"code": -32000, "message": "client closed"}} + if self._stdio_proc is not None: + try: + line = json.dumps(req, separators=(",", ":")) + "\n" + assert self._stdio_proc.stdin is not None + self._stdio_proc.stdin.write(line) + self._stdio_proc.stdin.flush() + while True: + outline = self._stdio_proc.stdout.readline() + if not outline: + try: + err = self._stdio_proc.stderr.read(256) or "" + except Exception: + err = "" + return {"jsonrpc": "2.0", "id": req.get("id"), "error": {"code": -32000, "message": f"stdio closed: {err}"}} + try: + resp = json.loads(outline.strip()) + except json.JSONDecodeError: + continue + rid = resp.get("id") + if rid == req.get("id") or "id" not in resp: + # accept matching or (rare) no-id but we treat as notif only on dedicated path + if "id" in resp: + return resp + continue # notification while waiting, keep reading for our id + # other id? ignore + except Exception as e: + return {"jsonrpc": "2.0", "id": req.get("id"), "error": {"code": -32603, "message": f"stdio transport: {e}"}} + elif self._http_base is not None: + try: + headers = {"Content-Type": "application/json", **self._http_headers} + r = httpx.post(self._http_base, json=req, headers=headers, timeout=30.0) + r.raise_for_status() + return r.json() + except Exception as e: + return {"jsonrpc": "2.0", "id": req.get("id"), "error": {"code": -32603, "message": f"http transport: {e}"}} + return {"jsonrpc": "2.0", "id": req.get("id"), "error": {"code": -32603, "message": "no transport configured"}} + + def _send_notif(self, notif: dict[str, Any]) -> None: + if self._closed: + return + if self._stdio_proc is not None: + try: + line = json.dumps(notif, separators=(",", ":")) + "\n" + assert self._stdio_proc.stdin is not None + self._stdio_proc.stdin.write(line) + self._stdio_proc.stdin.flush() + except Exception: + pass + elif self._http_base is not None: + try: + headers = {"Content-Type": "application/json", **self._http_headers} + httpx.post(self._http_base, json=notif, headers=headers, timeout=10.0) + except Exception: + pass + + def _rpc(self, method: str, params: dict[str, Any] | None = None) -> dict[str, Any]: + rid = self._next_id() + req = {"jsonrpc": "2.0", "id": rid, "method": method, "params": params or {}} + resp = self._exchange(req) + if "error" in resp and resp.get("error"): + err = resp["error"] + code = err.get("code", -1) + msg = err.get("message", str(err)) + data = err.get("data") + raise RuntimeError(f"MCP {method} error {code}: {msg}{f' {data}' if data else ''}") + return resp.get("result") or {} + + def initialize(self) -> None: + if self._initialized: + return + params: dict[str, Any] = { + "protocolVersion": "2024-11-05", + "clientInfo": {"name": "context-runtime", "version": "0.1"}, + } + self._rpc("initialize", params) + notif = {"jsonrpc": "2.0", "method": "notifications/initialized"} + self._send_notif(notif) + self._initialized = True + + def list_tools(self) -> list[dict[str, Any]]: + if not self._initialized: + self.initialize() + result = self._rpc("tools/list") + if isinstance(result, dict): + return result.get("tools", []) + return result if isinstance(result, list) else [] + + def call_tool(self, name: str, arguments: dict[str, Any] | None = None) -> dict[str, Any]: + return self._rpc("tools/call", {"name": name, "arguments": arguments or {}}) + + def close(self) -> None: + self._closed = True + if self._stdio_proc is not None: + try: + assert self._stdio_proc.stdin is not None + self._stdio_proc.stdin.close() + self._stdio_proc.terminate() + self._stdio_proc.wait(timeout=0.5) + except Exception: + pass + self._stdio_proc = None + self._http_base = None + + +@dataclass +class MCPToolPlugin: + """ToolPlugin wrapper for a single MCP tool descriptor + shared live MCPClient.""" + + _desc: dict[str, Any] + _client: MCPClient + _name: str + _default_side: bool = True + + def __init__(self, tool_desc: dict[str, Any], client: MCPClient, *, name: str | None = None, default_side_effecting: bool = True) -> None: + # support dataclass init + override for manual + object.__setattr__(self, "_desc", tool_desc) + object.__setattr__(self, "_client", client) + object.__setattr__(self, "_name", name or tool_desc.get("name", "unnamed")) + object.__setattr__(self, "_default_side", default_side_effecting) + + def spec(self) -> ToolSpec: + mcp_name = self._desc.get("name") or self._name + desc = self._desc.get("description") or "" + schema = self._desc.get("inputSchema") or self._desc.get("input_schema") or {"type": "object", "properties": {}} + ann = self._desc.get("annotations") or {} + read_only = bool(ann.get("readOnlyHint")) if "readOnlyHint" in ann else False + side = (not read_only) if "readOnlyHint" in ann else bool(self._default_side) + return ToolSpec( + name=self._name, + description=desc, + parameters=schema, + side_effecting=side, + ) + + def run(self, args: dict) -> ToolResult: + try: + res = self._client.call_tool(self._desc.get("name") or self._name, args) + except Exception as e: + return ToolResult(ok=False, error=str(e), text=f"[error] {self._name}: {e}") + is_err = bool(res.get("isError", False)) + content_list = res.get("content") or [] + texts: list[str] = [] + for block in content_list: + if isinstance(block, dict) and block.get("type") == "text": + t = block.get("text") or "" + if t: + texts.append(t) + elif block: + texts.append(str(block)) + text = "\n".join(texts) + if "structuredContent" in res and res.get("structuredContent") is not None: + data = res["structuredContent"] + else: + data = res.get("content", res) + err = None + if is_err: + err = res.get("error") or (text or "mcp tool returned isError") + return ToolResult(ok=not is_err, data=data, text=text, error=err) + + +def mount_mcp( + registry: ToolRegistry, + client: MCPClient, + *, + prefix: str | None = None, + default_side_effecting: bool = True, +) -> list[str]: + """Initialize client if needed, list tools, register each wrapped MCPToolPlugin. + + Registered name = f'{prefix}.{name}' if prefix else name. + Returns the list of registered tool names. + """ + try: + if not client._initialized: + client.initialize() + except Exception: + pass # list_tools may still work or caller will see [] + registered: list[str] = [] + try: + tools = client.list_tools() + except Exception: + return registered + for tdesc in tools: + if not isinstance(tdesc, dict) or not tdesc.get("name"): + continue + raw = tdesc["name"] + reg_name = f"{prefix}.{raw}" if prefix else raw + plugin = MCPToolPlugin(tdesc, client, name=reg_name, default_side_effecting=default_side_effecting) + registry.register(plugin) + registered.append(reg_name) + return registered diff --git a/context_runtime/tools/mcp_servers/__init__.py b/context_runtime/tools/mcp_servers/__init__.py new file mode 100644 index 0000000..8e5b6d5 --- /dev/null +++ b/context_runtime/tools/mcp_servers/__init__.py @@ -0,0 +1,6 @@ +"""Small, self-contained MCP tool servers shipped with Context Runtime. + +These are reference MCP servers (stdio transport) that any AgentConsole app can mount via +``context_runtime.tools.mcp.MCPClient.stdio`` — proving the agent-harness ↔ MCP bridge with +real, keyless tools. +""" diff --git a/context_runtime/tools/mcp_servers/web_search.py b/context_runtime/tools/mcp_servers/web_search.py new file mode 100644 index 0000000..0a55117 --- /dev/null +++ b/context_runtime/tools/mcp_servers/web_search.py @@ -0,0 +1,135 @@ +"""web_search — a keyless, dependency-free MCP web-search server (stdio transport). + +Exposes ONE read-only tool, ``web_search``, that searches the open web for a topic, company, +competitor, or trend and returns titles + URLs. Sources are keyless JSON APIs so it runs in a +slim container with no API key: + + * Wikipedia opensearch — reference/knowledge results + * Hacker News (Algolia) — live discussion / launch / news results + +Speaks newline-delimited JSON-RPC 2.0 (the MCP stdio transport) — one JSON message per line — +matching ``context_runtime.tools.mcp.MCPClient.stdio``. The tool is annotated ``readOnlyHint`` +so, once mounted, the agent-harness ApprovalPolicy lets it run without a gate. + +Run: python -m context_runtime.tools.mcp_servers.web_search +""" +from __future__ import annotations + +import json +import sys +import urllib.parse +import urllib.request + +_UA = {"User-Agent": "context-runtime-web-search/0.1"} +_TIMEOUT = 12.0 + + +def _get(url: str) -> str: + return urllib.request.urlopen(urllib.request.Request(url, headers=_UA), timeout=_TIMEOUT).read().decode("utf-8", "ignore") + + +def _wikipedia(query: str, k: int) -> list[dict]: + try: + d = json.loads(_get("https://en.wikipedia.org/w/api.php?" + urllib.parse.urlencode( + {"action": "opensearch", "search": query, "limit": k, "format": "json"}))) + return [{"title": t, "url": u, "source": "Wikipedia"} for t, u in zip(d[1], d[3])] + except Exception: # noqa: BLE001 + return [] + + +def _hackernews(query: str, k: int) -> list[dict]: + try: + d = json.loads(_get("https://hn.algolia.com/api/v1/search?" + urllib.parse.urlencode( + {"query": query, "hitsPerPage": k}))) + out: list[dict] = [] + for h in d.get("hits", [])[:k]: + title = h.get("title") or h.get("story_title") or "" + url = h.get("url") or h.get("story_url") or f"https://news.ycombinator.com/item?id={h.get('objectID', '')}" + if title: + out.append({"title": title, "url": url, "source": "HackerNews", "points": h.get("points")}) + return out + except Exception: # noqa: BLE001 + return [] + + +def web_search(query: str, max_results: int = 6) -> str: + query = (query or "").strip() + if not query: + return "Give me something to search for." + k = max(1, min(int(max_results or 6), 10)) + results = (_wikipedia(query, 2) + _hackernews(query, k))[:k] + if not results: + return f"No web results found for: {query}" + lines = [f'Web search — "{query}" ({len(results)} results):'] + for i, r in enumerate(results, 1): + extra = f" · {r['points']} pts" if r.get("points") else "" + lines.append(f"{i}. {r['title']} [{r['source']}{extra}]\n {r['url']}") + return "\n".join(lines) + + +_TOOLS = [{ + "name": "web_search", + "description": ("Search the open web (Wikipedia + Hacker News) for information, discussions, or news " + "about a topic, company, competitor, or trend. Returns titles and URLs."), + "inputSchema": { + "type": "object", + "properties": { + "query": {"type": "string", "description": "the search query"}, + "max_results": {"type": "integer", "description": "max results to return (default 6)"}, + }, + "required": ["query"], + }, + "annotations": {"readOnlyHint": True, "openWorldHint": True}, +}] + + +def _result(rid, result=None, error=None) -> dict: + m: dict = {"jsonrpc": "2.0", "id": rid} + if error is not None: + m["error"] = error + else: + m["result"] = result + return m + + +def _handle(req: dict) -> dict | None: + method = req.get("method") + rid = req.get("id") + params = req.get("params") or {} + if method == "initialize": + return _result(rid, {"protocolVersion": "2024-11-05", "capabilities": {"tools": {}}, + "serverInfo": {"name": "web-search", "version": "0.1.0"}}) + if method == "tools/list": + return _result(rid, {"tools": _TOOLS}) + if method == "tools/call": + name = params.get("name") + args = params.get("arguments") or {} + if name != "web_search": + return _result(rid, {"content": [{"type": "text", "text": f"unknown tool: {name}"}], "isError": True}) + try: + text = web_search(args.get("query", ""), args.get("max_results", 6)) + return _result(rid, {"content": [{"type": "text", "text": text}], "isError": False}) + except Exception as e: # noqa: BLE001 + return _result(rid, {"content": [{"type": "text", "text": f"search error: {e}"}], "isError": True}) + if rid is not None: # unknown request (notifications have no id → ignored) + return _result(rid, error={"code": -32601, "message": f"method not found: {method}"}) + return None + + +def main() -> None: + for line in sys.stdin: + line = line.strip() + if not line: + continue + try: + req = json.loads(line) + except Exception: # noqa: BLE001 + continue + resp = _handle(req) + if resp is not None: + sys.stdout.write(json.dumps(resp) + "\n") + sys.stdout.flush() + + +if __name__ == "__main__": + main() diff --git a/context_runtime/types.py b/context_runtime/types.py index 30bafc6..b93e9b3 100644 --- a/context_runtime/types.py +++ b/context_runtime/types.py @@ -78,7 +78,7 @@ class Goal: IntentBucket = Literal[ "exact_lookup", "conceptual", "incident", "code_reasoning", - "synthesis", "high_risk", "sensitive", "multi_hop", "unknown", + "synthesis", "high_risk", "sensitive", "multi_hop", "temporal", "unknown", ] @@ -89,6 +89,10 @@ class Intent: risk: Literal["low", "medium", "high"] = "low" normalized: str = "" # deterministic canonical form → cache key confidence: float = 0.0 + # v4: the knowledge representation this request is about — the decision engine's first + # axis, chosen before any retrieval method. Candidate generation is constrained to the + # methods that specialize this representation (see planner/representations.py). + representation: "KnowledgeRepresentation" = "document" # ──────────────────────────── §2.3 candidate / plan ──────────────────────────── @@ -141,7 +145,16 @@ class Plan: Retrieval = Literal[ "vector", "bm25", "hybrid", "graph", "community", "image", "colpali", "video", - "sql", "api", "logs", "file", "code", + "sql", "api", "logs", "file", "code", "temporal", +] + +# The knowledge *representation* a retrieval method operates over. The planner's first +# decision is which representation best answers the request (a document, a graph of +# relationships, a bi-temporal fact history, an analytical/OLAP cube, code, or media); +# the concrete `Retrieval` method is a specialization within the chosen representation. +# Retrieval is therefore one family of knowledge access, not the whole of it. +KnowledgeRepresentation = Literal[ + "document", "graph", "temporal", "analytical", "community", "code", "multimodal", ] diff --git a/deploy/medical/README.md b/deploy/medical/README.md new file mode 100644 index 0000000..86b9b76 --- /dev/null +++ b/deploy/medical/README.md @@ -0,0 +1,22 @@ +# MedicalBench corpus — PubMedQA + +The medical half of the public chat demo (`chat.redevops.io`, the "Context Runtime · +MedicalBench" and "· Mixed" models). The medical analogue of `deploy/financebench`: real +domain documents + expert Q&A, from a public dataset. + +**Source — PubMedQA** (Jin et al., EMNLP 2019, https://github.com/pubmedqa/pubmedqa, MIT): +1 000 expert-labeled biomedical questions over real PubMed abstract passages, each with gold +contexts and a long-form answer. Its clinical vocabulary (discharge, statement, balance, +chronic/acute, …) deliberately collides with financial 10-K language — which is what makes +the **Mixed** model's coverage-routing win (cross-domain noise → 0) visible. + +```bash +./download.sh # → ../../.medical/pqal.json (curl, no HF datasets dep) +python3 build_corpus.py # → ../../.medical/corpus/*.txt (~3.3k passages) + qa.jsonl +``` + +`build_corpus.py` writes one normalized `.txt` passage per (PMID, abstract section) — the same +pre-normalized shape the FinanceBench corpus uses, so the control plane ingests it directly +(mounted RO at `/medical`, tenant `context-runtime-medical`). No PDF/OCR step. + +Data lands in `../../.medical/` (gitignored). We ship the downloader, not the data. diff --git a/deploy/medical/build_corpus.py b/deploy/medical/build_corpus.py new file mode 100755 index 0000000..d048449 --- /dev/null +++ b/deploy/medical/build_corpus.py @@ -0,0 +1,67 @@ +#!/usr/bin/env python3 +"""Build the MedicalBench corpus: PubMedQA abstracts → normalized passages. + + ./download.sh && python3 deploy/medical/build_corpus.py + +PubMedQA's `ori_pqal.json` is a dict keyed by PMID; each entry has QUESTION, CONTEXTS (the +abstract's labeled sections), LABELS, LONG_ANSWER and final_decision. We emit one normalized +`.txt` passage per (PMID, section) into `.medical/corpus/` — the same pre-normalized shape the +FinanceBench corpus uses (CR_CORPUS_DIR ingests it directly, no PDF step) — plus `qa.jsonl` +(question / gold long answer / decision / gold PMID) so the corpus is probe-able. + +The passages carry real clinical vocabulary (discharge, chronic/acute, balance, statement, …) +that collides with 10-K language — which is the point: in the Mixed model, coverage routing +keeps a medical query on the clinical shard instead of surfacing a 10-K's "discharge of +liability". Keeps the whitepaper's cross-domain-noise-22→0 result on the live demo path. +""" +import json +import os +import re +import sys + +ROOT = os.path.abspath(os.path.join(os.path.dirname(__file__), "..", "..")) +MED = os.path.join(ROOT, ".medical") +SRC = os.path.join(MED, "pqal.json") +CORPUS = os.path.join(MED, "corpus") + + +def _slug(text: str, n: int = 48) -> str: + return re.sub(r"[^a-z0-9]+", "-", text.lower()).strip("-")[:n] or "section" + + +def main() -> None: + if not os.path.isfile(SRC): + sys.exit(f"missing {SRC} — run deploy/medical/download.sh first") + data = json.load(open(SRC, encoding="utf-8")) + os.makedirs(CORPUS, exist_ok=True) + n_docs = n_passages = 0 + with open(os.path.join(MED, "qa.jsonl"), "w", encoding="utf-8") as qa: + for pmid, rec in data.items(): + contexts = rec.get("CONTEXTS") or [] + labels = rec.get("LABELS") or [] + if not contexts: + continue + n_docs += 1 + for i, ctx in enumerate(contexts): + label = labels[i] if i < len(labels) else "" + # one passage file per abstract section; header carries the section label so a + # lexical index keeps the clinical framing (RESULTS / METHODS / CONCLUSIONS …). + header = f"PubMed {pmid} — {label}".strip(" —") + text = f"{header}\n\n{ctx.strip()}\n" + fn = f"pmid{pmid}_{i:02d}_{_slug(label or 'abstract')}.txt" + with open(os.path.join(CORPUS, fn), "w", encoding="utf-8") as fh: + fh.write(text) + n_passages += 1 + qa.write(json.dumps({ + "pmid": pmid, + "question": rec.get("QUESTION", ""), + "long_answer": rec.get("LONG_ANSWER", ""), + "decision": rec.get("final_decision", ""), + }, ensure_ascii=False) + "\n") + print(f"=== MedicalBench corpus built ===") + print(f" {n_docs} abstracts → {n_passages} passages → {CORPUS}") + print(f" qa.jsonl → {os.path.join(MED, 'qa.jsonl')}") + + +if __name__ == "__main__": + main() diff --git a/deploy/medical/download.sh b/deploy/medical/download.sh new file mode 100755 index 0000000..686a9dd --- /dev/null +++ b/deploy/medical/download.sh @@ -0,0 +1,26 @@ +#!/usr/bin/env bash +# Download the MedicalBench demo data — the public PubMedQA expert-labeled set. +# +# PubMedQA (Jin et al., EMNLP 2019, https://github.com/pubmedqa/pubmedqa) — 1 000 expert- +# annotated biomedical questions over real PubMed abstract passages, each with gold contexts +# and a long-form answer. It's the medical analogue of the FinanceBench half of this demo: +# real domain documents + expert Q&A, and a vocabulary (discharge, statement, balance, +# chronic/acute, …) that deliberately collides with financial 10-K language — which is what +# makes the heterogeneous-shards / coverage-routing win visible in the Mixed model. +# +# Data lands in ../../.medical/ (gitignored). Public, MIT-licensed; we ship the downloader, +# not the data — same policy as deploy/financebench/download.sh. +set -euo pipefail +DEST="$(cd "$(dirname "$0")/../.." && pwd)/.medical" +RAW="https://raw.githubusercontent.com/pubmedqa/pubmedqa/master/data" +mkdir -p "$DEST" + +echo "▸ PubMedQA expert-labeled Q&A (ori_pqal.json, 1000 items)" +curl -sfL "$RAW/ori_pqal.json" -o "$DEST/pqal.json" +python3 - "$DEST/pqal.json" <<'PY' +import json, sys +d = json.load(open(sys.argv[1])) +assert isinstance(d, dict) and d, "unexpected pqal.json shape" +print(f" ✓ {len(d)} labeled questions") +PY +echo "done → $DEST/pqal.json (now run build_corpus.py to build the passage corpus)" diff --git a/deploy/proxmox-demo/.env.example b/deploy/proxmox-demo/.env.example index 7e0feae..cfdbb66 100644 --- a/deploy/proxmox-demo/.env.example +++ b/deploy/proxmox-demo/.env.example @@ -4,12 +4,26 @@ # Public hostname (must match the cloudflared ingress rule + the manual CNAME). CHAT_HOST=chat.redevops.io -# Where the contextos repo + FinanceBench corpus live on the host (rsync'd there). +# Where the contextos repo + the two corpora live on the host (rsync'd there by the caddy role). CONTEXTOS_DIR=/projects/contextos CORPUS_DIR=/projects/contextos/.financebench/corpus +MEDICAL_CORPUS_DIR=/projects/contextos/.medical/corpus -# The LibreQB fork image (built from the LibreChat fork, LibreQB branch — see README). -LIBRECHAT_IMAGE=librechat-api:latest +# The three corpus tenants the control plane serves (one selectable model each). Bare path = +# single-corpus tenant; shards(...) = coverage-routed heterogeneous mix. CR_CORPUS_DIR stays +# EMPTY so the finance tenant below is the default (no duplicate legacy ingest). +CR_CORPUS_DIR= +CR_TENANTS=context-runtime-finance=/corpus | context-runtime-medical=/medical | context-runtime-mixed=shards(finance:/corpus,medical:/medical) + +# Per-tenant cross-language retrieval (opt-in; empty = off). Lists the CORPUS language(s) to +# translate an incoming query INTO before retrieval — these corpora are English, so `=en` lets a +# Spanish-speaking (Miami) customer query them. Adds one cached translation call per new query. +# CR_TENANT_LANGS=context-runtime-finance=en | context-runtime-medical=en | context-runtime-mixed=en +CR_TENANT_LANGS= + +# Optional local-build override for the LibreQB fork image. The registry deploy pulls +# localhost:5000/chat-librechat:${IMAGE_TAG} and ignores this; only set it for a local build. +# LIBRECHAT_IMAGE=librechat-api:latest # LibreChat secrets — Vault: openssl rand -hex 32 for JWT_*/CREDS_KEY, 16 bytes hex for CREDS_IV. JWT_SECRET= @@ -18,7 +32,16 @@ CREDS_KEY= CREDS_IV= CONTEXT_RUNTIME_API_KEY=context-runtime -# Upstream model for answer synthesis (Kimi / any OpenAI-compatible). -KIMI_BASE_URL=https://api.moonshot.ai/v1 -KIMI_API_KEY= -KIMI_MODEL=kimi-k2.6 +# ── Chat upstream (answer synthesis). Default: the shared Qwen3.6-35B-A3B via the host shim +# (Kimi fallback lives INSIDE the shim). Override any of these to point elsewhere for A/B. ── +CR_MODEL_URL=http://host.docker.internal:8000/v1 +CR_MODEL_KEY=sk-noauth +CR_MODEL_NAME=Qwen3.6-35B-A3B + +# ── LLM-as-judge for the self-learning reward (dedicated, kept separate from the chat model). +# Default: Grok 4.5 (xAI, OpenAI-compatible). Set XAI_API_KEY (from Vault) to enable; absent +# → the runtime degrades to the heuristic judge. Override CR_JUDGE_* to use a different judge. ── +CR_LLM_JUDGE=1 +CR_JUDGE_URL=https://api.x.ai/v1 +CR_JUDGE_MODEL=grok-4.5 +XAI_API_KEY= diff --git a/deploy/proxmox-demo/README.md b/deploy/proxmox-demo/README.md index b313b12..570725d 100644 --- a/deploy/proxmox-demo/README.md +++ b/deploy/proxmox-demo/README.md @@ -1,9 +1,20 @@ -# Public demo deploy — LibreQB × FinanceBench on Proxmox +# Public demo deploy — LibreQB × Context Runtime on Proxmox Ships the crown-jewel demo to `chat.redevops.io`: the LibreChat fork (with the **Libre Query Board** panel) in **no-login guest mode**, plus the Context Runtime control plane -over the FinanceBench corpus — all fronted by one Caddy so the browser talks to a single -origin. +serving **three corpus-scoped models** from one process (multi-tenant, `CR_TENANTS`) — all +fronted by one Caddy so the browser talks to a single origin. + +| Model (dropdown) | Corpus | Public source | +|---|---|---| +| Context Runtime · FinanceBench | real SEC 10-K pages (`/corpus`) | Patronus AI FinanceBench | +| Context Runtime · MedicalBench | biomedical abstracts (`/medical`) | PubMedQA | +| Context Runtime · Mixed | a heterogeneous shard mix of both, **coverage-routed** | — | + +The Mixed model is the whitepaper's heterogeneous-source win, live: a finance question is +answered only from finance docs and a medical question only from medical docs (cross-domain +noise → 0). The chat shim (`/v1`) and the Query Board (`/librechat/compare`) both route by +the selected model id, so the transparency panel reflects the corpus you picked. ``` Cloudflare edge → cloudflared (host, systemd) → caddy :8092 (host) ─┬─ /librechat/* /v1/* → contextruntime:8092 @@ -20,7 +31,7 @@ subdomain. Guests are per-browser (isolated chats, 1-week TTL) via `DEMO_MODE` i | `chat-demo.compose.yml` | caddy + contextruntime + librechat + mongo (isolated network) | | `Caddyfile` | one-origin routing | | `Dockerfile.contextruntime` | control-plane image (Python) | -| `librechat.yaml` | one FinanceBench endpoint → `contextruntime:8092/v1` | +| `librechat.yaml` | one endpoint → `contextruntime:8092/v1`, three corpus modelSpecs | | `cloudflared-ingress.snippet.yml` | the `chat.redevops.io` tunnel rule | | `.env.example` | secrets template (render from Vault) | @@ -34,12 +45,16 @@ This can't be scripted here (redevops.io is a different CF account than the tunn templated ingress in `ffmpeg-mcp-aws/ansible/local/deploy-cloudflared.yml` (next to the demo.redevops.io rule), then re-run `deploy-cloudflared.yml`. -**3. Corpus + repo onto the host.** Get the FinanceBench data + the contextos repo to the -host (e.g. under `/projects/contextos`): +**3. Corpora onto the host.** Build BOTH corpora on the control node, then get them + the +contextos repo to the host (e.g. under `/projects/contextos`). The ansible `caddy` role does +the rsync; to do it by hand: ```bash cd deploy/financebench && ./download.sh && uv run --with pdfplumber build_corpus.py # → .financebench/corpus -# then rsync the contextos repo (incl. .financebench/corpus + deploy/proxmox-demo) to the host +cd ../medical && ./download.sh && python3 build_corpus.py # → .medical/corpus (PubMedQA) +# then rsync .financebench/corpus + .medical/corpus + deploy/proxmox-demo to the host ``` +The ansible path builds both automatically (`build-push.yml`) and mounts them RO at `/corpus` +and `/medical`; `CR_TENANTS` (in the rendered `.env`) wires the three models. **4. LibreQB fork image.** Build the fork (LibreQB branch → the panel + DEMO_MODE + the auto-detecting compare URL are baked in) and tag it `librechat-api:latest` on the host: @@ -79,7 +94,17 @@ curl -s localhost:8092/librechat/compare -H 'content-type: application/json' \ - English corpus → do **not** set `CR_QUERY_LANGS`. - The control plane image installs `fastembed` (CPU MiniLM, no torch); ~a few GB RAM total with LibreChat + mongo. Fits alongside the agentic-os stack. -- To also expose the medical (Russian) demo, run a second control plane with that corpus + - `CR_QUERY_LANGS=ru CR_QUERY_XLATE_MODEL=moonshot-v1-8k` and a second endpoint. +- The three models are served by ONE control plane via `CR_TENANTS` (no second container). + Each corpus is its own tenant with its own learned policy + Query Board. Add/remove a model + by editing `CR_TENANTS` (`id=/path` for a single corpus, `id=shards(a:/p1,b:/p2)` for a + coverage-routed mix) — no code change. +- Mixed-model routing cleanliness is tunable via `CR_ROUTE_MARGIN` (default `0.15`, coverage + router). Lower = stricter single-domain routing if a borderline query leaks. +- **Per-tenant cross-language** (`CR_TENANT_LANGS`, opt-in): lists the CORPUS language(s) to + translate an incoming query INTO before retrieval, per model. These corpora are English, so + `context-runtime-finance=en | context-runtime-medical=en | context-runtime-mixed=en` lets a + Spanish-speaking (e.g. Miami) customer query them. Empty ⇒ off. Each tenant can bridge a + different language set (falls back to the global `CR_QUERY_LANGS`). Adds one cached translation + call per new query. - This stack is intentionally self-contained (own network + mongo) so it doesn't touch the agentic-os control plane; it only shares the host + the cloudflared tunnel. diff --git a/deploy/proxmox-demo/chat-demo.compose.yml b/deploy/proxmox-demo/chat-demo.compose.yml index c7eb611..662fa7f 100644 --- a/deploy/proxmox-demo/chat-demo.compose.yml +++ b/deploy/proxmox-demo/chat-demo.compose.yml @@ -1,12 +1,22 @@ -# Public LibreQB × FinanceBench demo — LibreChat (no-login guest mode) + the Context +# Public LibreQB × Context Runtime demo — LibreChat (no-login guest mode) + the Context # Runtime control plane, fronted by one Caddy so the browser talks to a single origin # (chat.redevops.io). cloudflared routes chat.redevops.io → caddy:80. # # caddy :80 /librechat/* /v1/* → contextruntime:8092 (Query Board + chat shim) # everything else → librechat:3080 (the SPA) # -# Deploy: rsync this dir + ../../.financebench/corpus + the contextos repo to the host, -# render .env (secrets), then `docker compose -f chat-demo.compose.yml up -d --build`. +# THREE corpus-scoped models from ONE control plane (multi-tenant, CR_TENANTS): +# • Context Runtime · FinanceBench — real SEC 10-K pages (/corpus) +# • Context Runtime · MedicalBench — PubMedQA biomedical abstracts (/medical) +# • Context Runtime · Mixed — a heterogeneous shard mix of both, coverage-routed +# (the whitepaper's cross-domain-noise-22→0 win) +# Each model is its own tenant with its own learned policy + Query Board; the panel and the +# chat shim both route by the selected model id, so the transparency view reflects the corpus +# the visitor picked. +# +# Deploy: rsync this dir + ../../.financebench/corpus + ../../.medical/corpus + the contextos +# repo to the host, render .env (secrets + corpus paths + CR_TENANTS), then +# `docker compose -f chat-demo.compose.yml up -d`. The ansible `caddy` role does all of this. # See README.md for the full runbook (incl. the manual cross-account CNAME). services: caddy: @@ -24,38 +34,64 @@ services: depends_on: [librechat, contextruntime] contextruntime: - # Context Runtime control plane over the FinanceBench corpus — all 5 methods - # (bm25/vector/hybrid/graph/community) + /librechat/compare + the /v1 shim. - build: - context: ${CONTEXTOS_DIR:-../..} - dockerfile: deploy/proxmox-demo/Dockerfile.contextruntime + # Context Runtime control plane — all 5 methods (bm25/vector/hybrid/graph/community) + + # /librechat/compare + the /v1 shim, serving THREE corpus tenants (finance/medical/mixed) + # from one process via CR_TENANTS. Registry deploy: built on the control node + # (build-push.yml, chat-contextruntime) and PULLED here. IMAGE_TAG defaults to :stable; + # pin a sha for rollback (-e image_tag=). + image: localhost:5000/chat-contextruntime:${IMAGE_TAG:-stable} restart: unless-stopped # Confined under docker-contextos (docker-default + abi/3.0 pin) — the host's # apparmor 4.1 breaks AF_UNIX create under docker-default, killing uvicorn's # socket.socketpair(). Load the profile first: see apparmor-docker-contextos.profile. security_opt: - apparmor=docker-contextos + # reach the host-network model shim (Qwen primary + Kimi fallback) at host.docker.internal:8000 + extra_hosts: + - "host.docker.internal:host-gateway" environment: - - CR_CORPUS_DIR=/corpus + # Multi-tenant: one model per corpus. Bare path = single-corpus tenant; shards(...) = + # coverage-routed heterogeneous mix. CR_CORPUS_DIR is left EMPTY so the legacy single + # tenant is skipped and the finance tenant below is the default. + - CR_CORPUS_DIR= + - CR_TENANTS=context-runtime-finance=/corpus | context-runtime-medical=/medical | context-runtime-mixed=shards(finance:/corpus,medical:/medical) + # Per-tenant cross-language retrieval (opt-in). Lists the CORPUS language(s) to translate an + # incoming query INTO before retrieval — these corpora are English, so `=en` lets a non-English + # (e.g. Spanish, Miami) speaker query them. Empty ⇒ off (no per-query translation latency). + # e.g. CR_TENANT_LANGS="context-runtime-finance=en | context-runtime-medical=en | context-runtime-mixed=en" + - CR_TENANT_LANGS=${CR_TENANT_LANGS:-} + - CR_SHARD_ROUTER=coverage - CR_EMBEDDINGS=1 - CR_MULTIMODAL=1 - CR_QUALITY_ROUTING=1 # the redevops.io EXPLAIN page fetches /librechat/explain cross-origin, so allow it too - CR_CORS_ORIGINS=https://${CHAT_HOST:-chat.redevops.io},https://redevops.io,https://www.redevops.io - - CR_UPSTREAM_BASE_URL=${KIMI_BASE_URL:-} - - CR_UPSTREAM_API_KEY=${KIMI_API_KEY:-} - - CR_UPSTREAM_MODEL=${KIMI_MODEL:-} + # Model plane: the shared Qwen3.6-35B-A3B via the host shim (Kimi fallback lives in the shim). + # Override CR_MODEL_* to point elsewhere for A/B testing. + - CR_UPSTREAM_BASE_URL=${CR_MODEL_URL:-http://host.docker.internal:8000/v1} + - CR_UPSTREAM_API_KEY=${CR_MODEL_KEY:-sk-noauth} + - CR_UPSTREAM_MODEL=${CR_MODEL_NAME:-Qwen3.6-35B-A3B} - CR_UPSTREAM_MAX_TOKENS=4096 - CR_UPSTREAM_TEMPERATURE=1 + # LLM-as-judge for the self-learning RAG reward (shared across all three corpus tenants). A + # DEDICATED judge = Grok 4.5 (xAI, OpenAI-compatible) — cheap and kept separate from the chat + # model so it doesn't grade its own retrieval. Needs XAI_API_KEY (from Vault); absent → + # degrades to the heuristic judge. + - CR_LLM_JUDGE=${CR_LLM_JUDGE:-1} + - CR_JUDGE_BASE_URL=${CR_JUDGE_URL:-https://api.x.ai/v1} + - CR_JUDGE_MODEL=${CR_JUDGE_MODEL:-grok-4.5} + - CR_JUDGE_API_KEY=${XAI_API_KEY:-} - CONTEXT_RUNTIME_HOME=/data volumes: - ${CORPUS_DIR:-../../.financebench/corpus}:/corpus:ro + - ${MEDICAL_CORPUS_DIR:-../../.medical/corpus}:/medical:ro - cr_data:/data networks: [chatdemo] librechat: # The LibreQB fork. DEMO_MODE=true → per-visitor guest, no login/registration. - image: ${LIBRECHAT_IMAGE:-librechat-api:latest} + # Registry deploy: built on the control node (build-push.yml, chat-librechat) and pulled here. + image: localhost:5000/chat-librechat:${IMAGE_TAG:-stable} restart: unless-stopped depends_on: [chat-mongo] # Confined under docker-contextos (see contextruntime + the .profile file). diff --git a/deploy/proxmox-demo/librechat.yaml b/deploy/proxmox-demo/librechat.yaml index 078abfc..120f4a5 100644 --- a/deploy/proxmox-demo/librechat.yaml +++ b/deploy/proxmox-demo/librechat.yaml @@ -1,13 +1,13 @@ -# LibreChat config for the public FinanceBench demo — one Context Runtime endpoint, -# pointed at the control plane over the internal compose network (contextruntime:8092). -# The Libre Query Board renders under the chat box (baked into the fork's frontend). +# LibreChat config for the public Context Runtime demo — ONE custom endpoint pointed at the +# control plane (contextruntime:8092/v1) that serves THREE corpus-scoped models, each its own +# self-learning Context Runtime tenant. The Libre Query Board (baked into the fork) renders +# under the chat box and inspects the corpus of whichever model is selected. version: 1.2.1 cache: true -# Guests land STRAIGHT on Context Runtime — no "select an Agent" step. modelSpecs with -# prioritize+enforce makes the single FinanceBench spec the default (and only) model, so -# a new chat opens ready to type. interface.agents:false removes the "My Agents" builder; -# endpointsMenu/presets off keep the guest UI to just the one thing this demo is about. +# Guests land straight on Context Runtime — no "select an Agent" step. modelSpecs.enforce keeps +# the dropdown to just these three corpus models; prioritize + FinanceBench default opens a new +# chat ready to type. interface.agents:false removes the "My Agents" builder. interface: agents: false endpointsMenu: false @@ -17,24 +17,39 @@ modelSpecs: enforce: true prioritize: true list: - - name: "context-runtime" + - name: "context-runtime-finance" label: "Context Runtime · FinanceBench" default: true - description: "Self-learning RAG over the FinanceBench corpus — Context Runtime decides what the model sees before it answers." + description: "Self-learning RAG over the FinanceBench corpus (real SEC 10-K filings) — Context Runtime decides what the model sees before it answers." preset: - endpoint: "FinanceBench (Context Runtime)" - model: "context-runtime" + endpoint: "Context Runtime" + model: "context-runtime-finance" + - name: "context-runtime-medical" + label: "Context Runtime · MedicalBench" + description: "Self-learning RAG over the PubMedQA corpus (expert-labeled biomedical abstracts) — the same runtime, a different domain." + preset: + endpoint: "Context Runtime" + model: "context-runtime-medical" + - name: "context-runtime-mixed" + label: "Context Runtime · Mixed (finance × medical)" + description: "A heterogeneous mix of the finance AND medical corpora, sharded + coverage-routed — a query is answered from the domain it belongs to, cutting cross-domain noise to zero. The whitepaper's heterogeneous-source win, live." + preset: + endpoint: "Context Runtime" + model: "context-runtime-mixed" endpoints: custom: - - name: "FinanceBench (Context Runtime)" + - name: "Context Runtime" apiKey: "${CONTEXT_RUNTIME_API_KEY}" baseURL: "${CONTEXT_RUNTIME_URL}/v1" + # the control plane's /v1/models advertises one id per corpus tenant; fetch picks them up. models: default: - - "context-runtime" + - "context-runtime-finance" + - "context-runtime-medical" + - "context-runtime-mixed" fetch: true titleConvo: true - titleModel: "context-runtime" + titleModel: "context-runtime-finance" summarize: false - modelDisplayLabel: "Context Runtime · FinanceBench (self-learning RAG)" + modelDisplayLabel: "Context Runtime (self-learning RAG)" diff --git a/docs/BENCHMARK-REPORT.md b/docs/BENCHMARK-REPORT.md new file mode 100644 index 0000000..a6b29ca --- /dev/null +++ b/docs/BENCHMARK-REPORT.md @@ -0,0 +1,69 @@ +# Context Runtime — Final Benchmark Report + +_Full `examples/` suite run end-to-end on the v3 engine. Every number is produced by a runnable +script (`PYTHONPATH=. python examples/.py`) — no invented figures._ + +**Run status: 30 / 30 benchmarks passed** (seeded, deterministic simulations; the consolidated +headline is cross-checked in the independent Go re-implementation). + +**Judge change in this cycle:** the LLM-as-judge default moved from **OpenAI GPT-5.5 → Grok 4.5** +(xAI, OpenAI-compatible, ~5× cheaper), now a *dedicated* judge endpoint kept separate from the chat +model (`CR_JUDGE_BASE_URL` / `CR_JUDGE_MODEL` / `CR_JUDGE_API_KEY`, default `grok-4.5`). Note: the +offline benchmarks below use a *seeded coverage judge* (no live LLM), so their numbers are unaffected +by the judge model — the judge swap changes the **live** self-learning RAG's cost, not these results. + +--- + +## 1. Headline — v1 → v2, measured in both runtimes (`consolidated_benchmark`) + +| Metric | Py v1 | Py v2 | Δ | Go v1 | Go v2 | Δ | +| --- | --- | --- | --- | --- | --- | --- | +| Learned-policy precision | 67.6% | 82.2% | ▲ +14.6 pts | 84.6% | 95.9% | ▲ +11.3 pts | +| Abstention recall (unanswerable caught) | 0.0% | 100.0% | ▲ +100 pts | 0.0% | 100.0% | ▲ +100 pts | +| False-abstain rate (answerable dropped) | 0.0% | 0.0% | — | 0.0% | 0.0% | — | +| Expensive-stage depth (passages) | 8.00 | 3.00 | ▼ −62% | 8.00 | 3.00 | ▼ −63% | +| Precision after the sizer | 37.5% | 100.0% | ▲ +62.5 pts | 37.5% | 100.0% | ▲ +62.5 pts | + +_40-seed average; precision headlined at β=0.9 (calibration-trust knob; shipped default 0.5). Go is an +independent re-implementation on identical methodology — directional parity across languages._ + +## 2. Calibration sweep (`dspark_calibration_bench`) + +Served true-precision, 40 seeds, coverage-biased judge: +- v1 judge-only (β=0.0): **67.6%** → v2 judge + calibrated relevance: β0.5 **68.9%** (+1.3), β0.7 **74.3%** (+6.7), β0.9 **82.2%** (+14.6). +- Abstention (P(rel)≥0.5 floor, v1 cannot abstain): sizer off → 8 passages @ 38% precision; sizer on → **3 passages @ 100%** (depth −62%, the pruned tail was the low-relevance one). + +## 3. Online adaptation under drift (Gen-4) + +- **`online_vs_static_bench`** — best plan drifts A→C at t=200. Post-drift served reward: static **0.20** · online-plain **0.33** · **online-discounted 0.67** (oracle 0.80). Discounting fades stale evidence so the planner tracks the shift. +- **`online_learning`** — priors say hybrid=0.80/graph=0.55 (static always serves hybrid); production reward returns graph=0.9/hybrid=0.3 → after learning the planner serves **graph**, adapting past the stale estimate (surfaced without re-running it, via off-policy value). +- **`learning_loop`** — async learning drained 8 events off the serving path → snapshot v1; replica reconciles and switches hybrid→graph with **zero learning on the hot path**. + +## 4. Retrieval methods & routing + +- **`hop_routing`** — intent-routed: conceptual → hybrid (acc 0.78, $0.06); **multi-hop → graph** (acc 0.88, $0.43), surfacing a bridge doc via the ATP→α-synuclein hop. +- **`parallel_fusion`** — fan-out wall-clock **5.9× faster** (50.9 ms parallel vs 300.8 ms sequential), then rerank-before-model. +- **`heterogeneous_shards`** — community/coverage routing cut cross-domain noise **23 → 0 docs** across 8 medical queries while keeping recall **8/8**. +- **`multimodal_image_search`** — cross-modal text→image: each text query retrieves the right image as top hit (no OCR, no shared terms). +- **`temporal_retrieval`** — bi-temporal as-of filtering (a record filed 2026-07-15 is correctly invisible to a 2026-06-15 as-of query). +- **`rag_tuning`** — per-intent optimal config recovered (exact_lookup / incident / conceptual all match the latent best) after 40 observed retrievals. + +## 5. Trust & governance (Gen-5) + +- **`trust_aware`** — trust folded into plan selection (trust_weight 0.6 → the relied-upon local plan wins); abstention gate: expected accuracy 0.80 clears the 0.7 threshold → **SERVE** (else abstain). +- **`capability_registry`, `explain`, `chat_memory`** — capability gating, EXPLAIN-ANALYZE traces, and conversational memory all green. + +## 6. Fleet / tenants (`fleet_tenants` + per-app demos) + +**17 tenants**, one pattern: each module = one goal, one metric, a learned *cheapest-sufficient source* +policy — replacing hand-wired controllers with one data-driven Context Runtime tenant pattern. The +per-app tenant demos all ran green: `agentic_billing, agentic_books, agentic_compliance, +agentic_support, control_tower, market_radar, growth_engine, social_autopilot, outreach_engine, +outreach_pipeline, incident_review, soc_triage, sidekick_learning, vibexgen_learning`. + +## Full run table + +All 30 scripts exited 0. Slowest: `consolidated_benchmark` (14s, runs the Go twin), `dspark_calibration_bench` (12s). One pre-existing example bug fixed this cycle: `fleet_tenants` was missing an `outreach` probe (added). + +_Reproduce any line: `PYTHONPATH=. python examples/.py`. Headline card: +`PYTHONPATH=. python examples/consolidated_benchmark.py --html out.html`._ diff --git a/docs/BENCHMARKS-v3-preliminary.md b/docs/BENCHMARKS-v3-preliminary.md new file mode 100644 index 0000000..e049aa5 --- /dev/null +++ b/docs/BENCHMARKS-v3-preliminary.md @@ -0,0 +1,19 @@ +# Benchmarks — v3 (preliminary) + +> **Preliminary.** This is a forward-looking axis distinct from — and not a replacement for — the +> shipped v1→v2 calibration results in [`BENCHMARKS.md`](../BENCHMARKS.md). It is kept in its own +> file until promoted. Reproduce with `PYTHONPATH=. python examples/consolidated_benchmark.py --v3-doc docs/BENCHMARKS-v3-preliminary.md`. + +## Online optimization under drift (Generation 4) + +The best plan drifts mid-run (a model upgrade, a corpus shift). A **static** v1/v2 planner is pinned +to the now-stale plan; the **v3 online** planner re-explores, and recency-weighted (discounted) +learning lets it track the shift. We report the post-drift average served-plan reward, seed-averaged. + +| Metric | v2 (static) | v3 (online) | Δ | +| --- | --- | --- | --- | +| Post-drift served-plan reward | 0.20 | 0.67 | ▲ +0.47 | +| Post-drift reward, online w/o discounting | 0.20 | 0.34 | +0.14 | + +_Seeded drift simulation, 24-seed average; post-drift oracle = 0.80. Discounting is what converts online learning from a marginal gain into recovery of most of the +oracle reward after the drift. This axis is orthogonal to (and preserves) the v2 calibration gains._ diff --git a/examples/capability_registry.py b/examples/capability_registry.py new file mode 100644 index 0000000..2826bc1 --- /dev/null +++ b/examples/capability_registry.py @@ -0,0 +1,50 @@ +"""Capabilities live in the planner, not the prompt (Whitepaper v3). + +Shows the CapabilityRegistry as the out-of-band catalog the planner selects from: the default registry +is lifted from the rule tables (so behavior is unchanged), and registering ONE new capability widens +the planner's search space immediately — no edit to the rules, nothing added to the model's context. + + python examples/capability_registry.py +""" +from __future__ import annotations + +from context_runtime.registry import Capability, CapabilityRegistry +from context_runtime.planner.candidates_registry import RegistryCandidateGenerator +from context_runtime.runtime.runtime import ContextRuntime + +QUERY = "how does the auth service relate to the billing outage across systems" # → multi_hop intent + + +def methods(exp) -> list[str]: + return sorted({next(s.params.get("method") for s in c.steps if s.type == "retrieve") + for c, _ in exp.candidates}) + + +def main() -> None: + registry = CapabilityRegistry.from_rules() + runtime = ContextRuntime.default( + docs=[{"id": "d1", "text": "The auth service change preceded the billing outage."}], + candidates=RegistryCandidateGenerator(registry), + ) + + exp = runtime.explain(QUERY, analyze=False) + print("Query:", QUERY) + print("Intent bucket:", exp.chosen.intent.bucket) + print("\nRegistry today —", len(registry.list("retrieval")), "retrieval capabilities,", + len(registry.list("model_tier")), "model tiers.") + print("Candidate retrieval methods for this intent:", methods(exp)) + + # A new retrieval capability comes online (e.g. a temporal/bi-temporal graph index). Registering it + # is one call — the planner can now select it wherever it serves; the prompt is untouched. + registry.register(Capability( + key="retrieval:temporal", kind="retrieval", value="temporal", + buckets=frozenset({"multi_hop", "incident"}), cost_prior=0.55, quality_prior=0.9, local=True, + )) + exp2 = runtime.explain(QUERY, analyze=False) + print("\nAfter registering 'temporal' (no rules edit, no prompt change):") + print("Candidate retrieval methods for this intent:", methods(exp2)) + print("\nCapability count scaled in the registry — which is cheap — not in the context window.") + + +if __name__ == "__main__": + main() diff --git a/examples/consolidated_benchmark.py b/examples/consolidated_benchmark.py index a6c969a..994550e 100644 --- a/examples/consolidated_benchmark.py +++ b/examples/consolidated_benchmark.py @@ -141,6 +141,69 @@ def to_html(py: dict, go: dict | None) -> str: """ +def v3_results(seeds: int = 24) -> dict: + """v3 (preliminary) — online optimization under drift (Generation 4). The best plan drifts mid-run + (a model upgrade / corpus shift). A static v1/v2 planner is pinned to the now-stale plan; the v3 online + planner re-explores, and recency-weighted (discounted) learning lets it track the shift. We report the + post-drift average served-plan reward, seed-averaged.""" + from context_runtime.optimizer.online import BanditOptimizer + from context_runtime.types import Candidate, Goal, PlanScore, StepSpec + + STEPS, DRIFT = 400, 200 + PRIOR = {"A": 0.80, "C": 0.20} + PRE = {"A": 0.80, "C": 0.20} + POST = {"A": 0.20, "C": 0.80} # best plan flips A→C at t=DRIFT + + def cand(a): + return Candidate(steps=(StepSpec(type="retrieve", params={"method": a}),), model_tier="cheap") + + def scored(): + return [(cand(a), PlanScore(total=PRIOR[a], feasible=True)) for a in PRIOR] + + def run(discount, seed): + opt = BanditOptimizer(None, epsilon=0.2, discount=discount, seed=seed) + post = [] + for t in range(STEPS): + plan = opt.select(scored(), Goal(text="q"), context="c") + arm = next(s.params["method"] for s in plan.chosen.steps if s.type == "retrieve") + r = (PRE if t < DRIFT else POST)[arm] + opt.learn_from_plan(plan, r) + if t >= DRIFT: + post.append(r) + return sum(post) / len(post) + + static = POST["A"] # v1/v2 never adapt → pinned to the stale best (A) + plain = sum(run(0.0, 0x1000 + s) for s in range(seeds)) / seeds + disc = sum(run(0.2, 0x1000 + s) for s in range(seeds)) / seeds + return {"seeds": seeds, "static": static, "online_plain": plain, "online_discount": disc, + "oracle": max(POST.values())} + + +def v3_document(v3: dict) -> str: + """A self-contained standalone markdown doc — kept OUT of the shipped BENCHMARKS.md because the v3 + online-optimization line is preliminary (a different axis from the v1→v2 calibration gains).""" + r = v3 + return ( + "# Benchmarks — v3 (preliminary)\n\n" + "> **Preliminary.** This is a forward-looking axis distinct from — and not a replacement for — the\n" + "> shipped v1→v2 calibration results in [`BENCHMARKS.md`](../BENCHMARKS.md). It is kept in its own\n" + "> file until promoted. Reproduce with " + "`PYTHONPATH=. python examples/consolidated_benchmark.py --v3-doc docs/BENCHMARKS-v3-preliminary.md`.\n\n" + "## Online optimization under drift (Generation 4)\n\n" + "The best plan drifts mid-run (a model upgrade, a corpus shift). A **static** v1/v2 planner is pinned\n" + "to the now-stale plan; the **v3 online** planner re-explores, and recency-weighted (discounted)\n" + "learning lets it track the shift. We report the post-drift average served-plan reward, seed-averaged.\n\n" + "| Metric | v2 (static) | v3 (online) | Δ |\n| --- | --- | --- | --- |\n" + f"| Post-drift served-plan reward | {r['static']:.2f} | {r['online_discount']:.2f} | " + f"▲ +{r['online_discount'] - r['static']:.2f} |\n" + f"| Post-drift reward, online w/o discounting | {r['static']:.2f} | {r['online_plain']:.2f} | " + f"+{r['online_plain'] - r['static']:.2f} |\n\n" + f"_Seeded drift simulation, {r['seeds']}-seed average; post-drift oracle = {r['oracle']:.2f}. " + "Discounting is what converts online learning from a marginal gain into recovery of most of the\n" + "oracle reward after the drift. This axis is orthogonal to (and preserves) the v2 calibration gains._\n" + ) + + def main() -> int: py = python_results() go = go_results() @@ -150,6 +213,12 @@ def main() -> int: path = sys.argv[sys.argv.index("--html") + 1] pathlib.Path(path).write_text(to_html(py, go), encoding="utf-8") print(f"\n[wrote HTML → {path}]", file=sys.stderr) + if "--v3-doc" in sys.argv: + path = sys.argv[sys.argv.index("--v3-doc") + 1] + p = pathlib.Path(path) + p.parent.mkdir(parents=True, exist_ok=True) + p.write_text(v3_document(v3_results()), encoding="utf-8") + print(f"\n[wrote v3 preliminary doc → {path}]", file=sys.stderr) if "--json" in sys.argv: print(json.dumps({"python": py, "go": go}, indent=2), file=sys.stderr) return 0 diff --git a/examples/fleet_tenants.py b/examples/fleet_tenants.py index 4b824a8..ba06ef4 100644 --- a/examples/fleet_tenants.py +++ b/examples/fleet_tenants.py @@ -23,10 +23,13 @@ "crm": ("which deals are most likely to close?", "deals"), "market_radar": ("did a competitor change pricing?", "pricing"), "growth_engine": ("which channel drove signups?", "attribution"), + "outreach": ("which prospects should we reach out to first?", "leads"), "social": ("what content is trending for us?", "trends"), "lifecycle": ("why did deliverability drop?", "segments"), "privacy": ("can we fulfill this DSAR in time?", "requests"), "edge_sentinel": ("is this IP malicious?", "threat_intel"), + "guide": ("how do I set up permissions for my app?", "docs"), + "growth_assistant": ("which channel should a first-time founder try first?", "playbooks"), "incident": ("why did the deploy fail?", "logs"), "research": ("how does mitochondrial dysfunction relate to Parkinson?", "citations"), "finance": ("should we invest given the latest filing?", "filings"), diff --git a/examples/learning_loop.py b/examples/learning_loop.py new file mode 100644 index 0000000..7801c1b --- /dev/null +++ b/examples/learning_loop.py @@ -0,0 +1,62 @@ +"""Phase 4 — the asynchronous learning loop that makes the planner scale. + +A stateless planner replica selects fast against a LOCAL snapshot and never learns on the serving path. +Executions publish outcome events to a bus; a single aggregator folds them into the shared bandit off +the hot path and republishes a versioned snapshot; replicas reconcile. Same seam works in-process +(shown here) or over Kafka across a fleet — neither stream touches the model's context window. + + python examples/learning_loop.py +""" +from __future__ import annotations + +from context_runtime.integrations.bandit import EpsilonGreedyBandit +from context_runtime.learning import InMemoryBus, LearningAggregator, OutcomeEvent +from context_runtime.optimizer.online import BanditOptimizer +from context_runtime.types import Candidate, Goal, PlanScore, StepSpec + + +def cand(method): + return Candidate(steps=(StepSpec(type="retrieve", params={"method": method}),), model_tier="cheap") + + +def score(total): + return PlanScore(total=total, expected_accuracy=0.9, feasible=True) + + +def main(): + bus = InMemoryBus() + goal = Goal(text="how does the auth change relate to the billing outage") + candidates = [(cand("hybrid"), score(0.80)), (cand("graph"), score(0.55))] # cost model prefers hybrid + + # A serving replica: selects against its own (initially empty) bandit; ε=0 so it just exploits. + replica_bandit = EpsilonGreedyBandit(arms=()) + replica = BanditOptimizer(None, bandit=replica_bandit, epsilon=0.0) + + def served(opt): + return next(s.params["method"] for s in opt.select(candidates, goal, context="multi_hop").chosen.steps + if s.type == "retrieve") + + print("Replica serves (before any learning):", served(replica), "— the cost-model favorite\n") + + # Executions happen (somewhere in the fleet) and publish outcome events. graph earns more reward. + print("Fleet executions publish outcome events to the bus (graph=0.9, hybrid=0.3)...") + seq = 0 + for _ in range(4): + for arm, r in (("graph:cheap", 0.9), ("hybrid:cheap", 0.3)): + seq += 1 + bus.publish("outcomes", OutcomeEvent(context="multi_hop", arm=arm, reward=r, seq=seq)) + + # The aggregator (one writer, off the hot path) folds them in and publishes a snapshot. + aggregator = LearningAggregator(EpsilonGreedyBandit(arms=())) + n = aggregator.drain(bus) + snap = aggregator.publish(bus, "snapshots") + print(f"Aggregator drained {n} events off the serving path → snapshot v{snap.version}\n") + + # The replica reconciles from the published snapshot — it never processed an event itself. + bus.poll("snapshots")[-1].apply_to(replica_bandit) + print("Replica serves (after reconciling the snapshot):", served(replica), + "— adapted past the stale estimate, with zero learning on the hot path") + + +if __name__ == "__main__": + main() diff --git a/examples/online_learning.py b/examples/online_learning.py new file mode 100644 index 0000000..3e7c260 --- /dev/null +++ b/examples/online_learning.py @@ -0,0 +1,70 @@ +"""Generation 4 — online learning: the planner explores, measures, and adapts. + +A static planner always serves the plan with the best *estimate*. A Gen-4 planner treats plan +selection as a contextual bandit: it mostly exploits the best-known plan, occasionally explores an +alternative, learns from the measured reward, and can evaluate plans it never served using the logged +selection probabilities (off-policy). This script shows all three, deterministically. + + python examples/online_learning.py +""" +from __future__ import annotations + +from context_runtime.optimizer.online import BanditOptimizer, offpolicy_values +from context_runtime.types import Candidate, Goal, PlanScore, StepSpec + + +def cand(tier, method): + return Candidate(steps=(StepSpec(type="retrieve", params={"method": method}),), model_tier=tier) + + +def score(total): + return PlanScore(total=total, feasible=True) + + +class _Est: + def estimate(self, c, g): return score(0.5) + def statistics(self, b=None): raise NotImplementedError + def observe(self, p, t): raise NotImplementedError + + +def main(): + goal = Goal(text="how does the auth change relate to the billing outage") + graph = cand("cheap", "graph") # cost model under-rates this one (prior 0.55) + hybrid = cand("cheap", "hybrid") # cost model prefers it (prior 0.80) + scored = [(hybrid, score(0.80)), (graph, score(0.55))] + + print("Estimated priors: hybrid=0.80 graph=0.55 → a static planner always serves hybrid.\n") + + # 1) Exploitation only (ε=0): the estimate wins. + opt = BanditOptimizer(_Est(), epsilon=0.0) + served = opt.select(scored, goal, context="multi_hop").chosen + print(f"ε=0 (exploit): serves {served.model_tier}:{_m(served)}") + + # 2) Reality disagrees — graph consistently earns higher reward. Feed it back. + print("\nProduction reward comes back: graph=0.9, hybrid=0.3 (repeatedly). Learning...") + for _ in range(4): + opt.learn("multi_hop", graph, 0.9) + opt.learn("multi_hop", hybrid, 0.3) + served = opt.select(scored, goal, context="multi_hop").chosen + print(f"After learning (exploit): serves {served.model_tier}:{_m(served)} " + f"← the planner adapted past its stale estimate") + + # 3) Off-policy (Tier 1): evaluate arms purely from logged executions, no live run. + logs = [ + {"context": "multi_hop", "arm": "hybrid:cheap", "reward": 0.3, "p": 0.85}, + {"context": "multi_hop", "arm": "hybrid:cheap", "reward": 0.3, "p": 0.85}, + {"context": "multi_hop", "arm": "graph:cheap", "reward": 0.9, "p": 0.15}, # rarely explored + ] + vals = offpolicy_values(logs)["multi_hop"] + print("\nOff-policy value estimate (IPS over logged selection probabilities):") + for arm, (w, v) in sorted(vals.items(), key=lambda kv: -kv[1][1]): + print(f" {arm:14} value≈{v:.3f}") + print(" → graph scores highest even though it was served rarely — surfaced without running it.") + + +def _m(c): + return next(s.params.get("method") for s in c.steps if s.type == "retrieve") + + +if __name__ == "__main__": + main() diff --git a/examples/online_vs_static_bench.py b/examples/online_vs_static_bench.py new file mode 100644 index 0000000..4680635 --- /dev/null +++ b/examples/online_vs_static_bench.py @@ -0,0 +1,76 @@ +#!/usr/bin/env python3 +"""Online (Gen-4) vs static planner under drift — "remaining correct in a non-stationary world". + +The static planner (v2 cost model) commits to the plan it estimated best and never revisits it. The +online planner (BanditOptimizer) learns from measured reward. We run a seeded environment where plan A +is best — until it drifts and C becomes best (a model upgrade, a corpus shift). We compare three: + + • static — always the highest-prior plan (A); never adapts. + • online, plain — sample-average learning; 200 stale samples of A drown the new signal. + • online, discount — recency-weighted learning (the whitepaper's non-stationarity property); + stale evidence fades, so the planner tracks the drift to C. + + PYTHONPATH=. python examples/online_vs_static_bench.py +""" +from __future__ import annotations + +from context_runtime.optimizer.online import BanditOptimizer +from context_runtime.types import Candidate, Goal, PlanScore, StepSpec + +CTX, GOAL = "synthesis", Goal(text="q") +STEPS, DRIFT = 400, 200 +PRIOR = {"A": 0.80, "C": 0.20} # cost model's fixed estimate (calibrated to t=200)", DRIFT, STEPS), + ("overall", 0, STEPS)]: + print(f"| {name:24} | {avg(s[lo:hi]):.3f} | {avg(plain[lo:hi]):.3f} " + f"| {avg(disc[lo:hi]):.3f} |") + print(f"\nPost-drift oracle (best possible) = {oracle:.2f}. Static is pinned to the now-stale A " + f"({avg(s[DRIFT:]):.2f}). Plain online adapts poorly ({avg(plain[DRIFT:]):.2f}) — the sample " + f"average is dominated by 200 pre-drift observations. With discounting, stale evidence fades " + f"and the planner tracks the drift to C ({avg(disc[DRIFT:]):.2f}).") + + +if __name__ == "__main__": + main() diff --git a/examples/temporal_retrieval.py b/examples/temporal_retrieval.py new file mode 100644 index 0000000..a663f29 --- /dev/null +++ b/examples/temporal_retrieval.py @@ -0,0 +1,42 @@ +"""Temporal / bi-temporal retrieval — "what changed, and when?" (Whitepaper v3, forthcoming retrieval). + +Every fact carries valid time (when it's true in the world) and transaction time (when we learned it), +so the planner can answer point-in-time questions the other retrieval methods can't. + + python examples/temporal_retrieval.py +""" +from __future__ import annotations + +from context_runtime.adapters.store_temporal import TemporalStore + + +def owner(hits): + return hits[0].meta["object"] if hits else "(none on record)" + + +def main(): + s = TemporalStore() + # the on-call owner of the auth service, revised over time — and one correction filed late + s.add("auth-service", "owner", "Alice", valid_from="2026-01-01", valid_to="2026-06-01", recorded_at="2026-01-01") + s.add("auth-service", "owner", "Bob", valid_from="2026-06-01", valid_to=None, recorded_at="2026-07-15") + + print("Current owner of auth-service:", owner(s.search("auth-service owner"))) + print() + print("Valid-time travel (what was true in the world, using the corrected record):") + for at in ("2026-03-01", "2026-06-15", "2026-09-01"): + print(f" as of {at}: {owner(s.as_of('auth-service', at=at))}") + + print() + print("Transaction-time travel (what we BELIEVED on a past date):") + print(f" as of 2026-06-15, as known then: {owner(s.as_of('auth-service', at='2026-06-15', known_at='2026-06-15'))}") + print(" → Bob's record was filed 2026-07-15, so on 2026-06-15 we hadn't recorded it yet — a late") + print(" correction doesn't silently rewrite what the system knew at the time.") + + print() + print("What changed, and when? (2026):") + for c in s.changes("auth-service", since="2026-01-01", until="2027-01-01"): + print(f" {c['at']} {c['change']:6} {c['fact']}") + + +if __name__ == "__main__": + main() diff --git a/examples/trust_aware.py b/examples/trust_aware.py new file mode 100644 index 0000000..867daef --- /dev/null +++ b/examples/trust_aware.py @@ -0,0 +1,61 @@ +"""Generation 5 — Trust-Aware Execution: optimize for trust, and abstain rather than guess. + +Two behaviors a trust-aware planner adds on top of cost/quality: + 1. trust is a weighted OBJECTIVE — a more-relied-upon plan can win over a nominally cheaper one; + 2. honest abstention — when the best plan's calibrated confidence is below the bar, decline (or + escalate) instead of serving a confident-wrong answer, which is what actually burns operator trust. + + python examples/trust_aware.py +""" +from __future__ import annotations + +from context_runtime.abstention import AbstentionGate +from context_runtime.optimizer.online import BanditOptimizer +from context_runtime.types import Candidate, Goal, PlanScore, StepSpec + + +def cand(tier, method): + return Candidate(steps=(StepSpec(type="retrieve", params={"method": method}),), model_tier=tier) + + +def score(total, acc): + return PlanScore(total=total, expected_accuracy=acc, feasible=True) + + +class Trust: + def __init__(self, scores): + self.scores = scores + + def score(self, c, g): + return self.scores.get(c.model_tier, 0.5) + + +def main(): + goal = Goal(text="summarize the incident and recommend a fix") + premium, local = cand("premium", "hybrid"), cand("local", "graph") + scored = [(premium, score(0.90, acc=0.55)), (local, score(0.60, acc=0.80))] + # operators have relied on the local graph plan; the premium plan they keep overriding + trust = Trust({"premium": 0.2, "local": 0.9}) + + print("Plans: premium (cost-total 0.90, trust 0.2) local (cost-total 0.60, trust 0.9)\n") + + tw0 = BanditOptimizer(None, epsilon=0.0, trust=trust, trust_weight=0.0) + print(f"trust_weight=0.0 (cost only): serves {tw0.select(scored, goal, context='incident').chosen.model_tier}") + + tw = BanditOptimizer(None, epsilon=0.0, trust=trust, trust_weight=0.6) + print(f"trust_weight=0.6 (trust folded in): serves {tw.select(scored, goal, context='incident').chosen.model_tier}" + f" ← the relied-upon plan wins") + + print("\nHonest abstention — the best plan's calibrated confidence must clear the bar (0.7):") + gate = AbstentionGate(min_confidence=0.7) + opt = BanditOptimizer(None, epsilon=0.0, trust=trust, trust_weight=0.6, abstain_gate=gate) + plan = opt.select(scored, goal, context="incident") + ab = plan.extra["abstention"] + print(f" chosen={plan.chosen.model_tier} confidence={ab['confidence']} → {ab['action'].upper()}") + print(f" reason: {ab['reason']}") + print("\n The chosen plan's expected accuracy (0.80) clears 0.7 → SERVE. Had it not, the planner") + print(" would abstain (or escalate) — an honest 'not sure' instead of a confident-wrong answer.") + + +if __name__ == "__main__": + main() diff --git a/tests/test_agent_console_policy.py b/tests/test_agent_console_policy.py new file mode 100644 index 0000000..c05a621 --- /dev/null +++ b/tests/test_agent_console_policy.py @@ -0,0 +1,62 @@ +"""AgentConsole ⇄ Policy Runtime: /command dispatch, input/output guardrails, tool approval, policy[].""" +from __future__ import annotations + +from types import SimpleNamespace + +from context_runtime.integrations.agent_console import AgentConsole +from context_runtime.policy import Command, CommandRegistry, Policy, RuleStore +from context_runtime.tools.base import ToolResult, function_tool + + +def _who(user="", app="market-radar"): + return SimpleNamespace(user=user, app=app, roles=frozenset()) + + +def test_slash_command_is_dispatched_without_the_model(): + reg = CommandRegistry() + reg.register(Command("ping", lambda a, p: {"text": "pong " + a["text"]}, "ping back")) + c = AgentConsole("T", "primer", commands=reg) + assert c.respond("/ping there")["text"] == "pong there" + + +def test_keyword_fallback_routes_howto_to_help_not_a_tool(): + """A how-to question must not be keyword-matched into a side-effecting tool just because it + shares an incidental word ('...part of a page?' vs the add_watch 'monitor a page' description).""" + c = AgentConsole("Market Radar", "radar", tools=[ + function_tool("add_watch", lambda a: ToolResult(ok=True, text="ok"), + description="start monitoring a new page (non-destructive)")]) + howto = c._keyword_route("How do I track just part of a page?") + assert howto["mode"] == "help" and howto["reason"] == "how-to → help" + # explicit actions still keyword-match to the tool + assert c._keyword_route("add a watch on this page")["mode"] == "tool" + + +def test_input_guardrail_blocks_and_reports(tmp_path): + store = RuleStore(dir=str(tmp_path)) + store.add("global", "guardrail", "forbidden phrase", action="deny", match={"phase": "input"}) + c = AgentConsole("T", "primer", policy=Policy(store=store, app="market-radar"), app="market-radar") + r = c.respond("here is the forbidden phrase", principal=_who("bob")) + assert r["intent"] == "blocked" and r["policy"][0]["decision"] == "deny" and r["policy"][0]["phase"] == "input" + + +def test_tool_requires_approval_and_is_not_run(tmp_path): + store = RuleStore(dir=str(tmp_path)) + store.add("global", "approval", "confirm irreversible sends", match={"tool": "send"}, action="require_approval") + ran = [] + c = AgentConsole("T", "primer", + tools=[function_tool("send", lambda a: ran.append(1) or ToolResult(ok=True, text="sent"), + side_effecting=True)], + policy=Policy(store=store, app="market-radar"), app="market-radar", allow_side_effects=["send"]) + out = c._answer_tool("send", {"to": "x"}, principal=_who("alex")) + assert out["approved"] is False and out["policy"][0]["decision"] == "require_approval" and not ran # not executed + + +def test_policy_summary_folds_allows(tmp_path): + c = AgentConsole("T", "how things work", policy=Policy(store=RuleStore(dir=str(tmp_path)), app="x"), app="x") + out = c.respond("how do I do the thing?", principal=_who("bob")) + assert out["policy"] == [] and out["policy_checks"] >= 2 # input+output allow → folded, no chips + + +def test_backward_compatible_without_policy_or_commands(): + out = AgentConsole("T", "primer").respond("hello") # no policy/commands ⇒ unchanged path + assert "text" in out and "policy" not in out diff --git a/tests/test_bandit_discount.py b/tests/test_bandit_discount.py new file mode 100644 index 0000000..f7258be --- /dev/null +++ b/tests/test_bandit_discount.py @@ -0,0 +1,33 @@ +"""Discounted (recency-weighted) bandit updates — the non-stationarity property (Whitepaper v3): +a constant step size fades stale evidence so the learned value tracks a drifting reward.""" +from __future__ import annotations + +from context_runtime.integrations.bandit import EpsilonGreedyBandit + + +class _Arm: + def __init__(self, key): + self.key = key + + +def _after_flip(discount: float) -> float: + b = EpsilonGreedyBandit(arms=(_Arm("x"),), discount=discount) + for _ in range(50): + b.update("c", _Arm("x"), 0.9) # long history at 0.9 + for _ in range(15): + b.update("c", _Arm("x"), 0.1) # world flips to 0.1 + return b.value("c", "x")[1] + + +def test_sample_average_is_the_default_and_barely_moves(): + # 0 discount = 1/n sample average: 50 old samples dominate 15 new ones + assert _after_flip(0.0) > 0.6 + + +def test_discount_tracks_the_new_regime(): + # constant step size weights recent rewards → tracks the flip toward 0.1 + assert _after_flip(0.2) < 0.2 + + +def test_discount_beats_sample_average_under_drift(): + assert _after_flip(0.2) < _after_flip(0.0) diff --git a/tests/test_capability_registry.py b/tests/test_capability_registry.py new file mode 100644 index 0000000..7862b03 --- /dev/null +++ b/tests/test_capability_registry.py @@ -0,0 +1,70 @@ +"""Capability registry + registry-backed candidate generation. + +Two claims: (1) the default registry (lifted from the rule tables) is behavior-equivalent to the +hand-written RuleCandidateGenerator, and (2) a new capability registered out of band immediately +widens the planner's candidate set — with no change to rules.py and nothing added to any prompt. +""" +from __future__ import annotations + +from context_runtime.planner.candidates import RuleCandidateGenerator +from context_runtime.planner.candidates_registry import RegistryCandidateGenerator +from context_runtime.registry import Capability, CapabilityRegistry +from context_runtime.types import Goal, Intent + + +def _sig(cands): + """Order-independent signature of a candidate set: (model_tier, retrieval method) pairs.""" + out = set() + for c in cands: + method = next((s.params.get("method") for s in c.steps if s.type == "retrieve"), None) + out.add((c.model_tier, method)) + return out + + +_BUCKETS = [ + "exact_lookup", "conceptual", "incident", "code_reasoning", + "synthesis", "high_risk", "sensitive", "multi_hop", "unknown", +] + + +def test_default_registry_matches_rule_generator_across_buckets(): + rule = RuleCandidateGenerator() + reg = RegistryCandidateGenerator() # defaults to CapabilityRegistry.from_rules() + goal = Goal(text="q") + for bucket in _BUCKETS: + intent = Intent(bucket=bucket) + assert _sig(reg.generate(intent, goal)) == _sig(rule.generate(intent, goal)), bucket + + +def test_from_rules_registers_expected_capabilities(): + reg = CapabilityRegistry.from_rules() + methods = {c.value for c in reg.list("retrieval")} + tiers = {c.value for c in reg.list("model_tier")} + assert {"bm25", "vector", "hybrid", "code", "graph"} <= methods + assert {"local", "cheap", "premium"} == tiers + # graph serves multi_hop; bm25 does not + assert reg.get("retrieval:graph").serves("multi_hop") + assert not reg.get("retrieval:bm25").serves("multi_hop") + + +def test_registering_a_new_capability_widens_the_candidate_set(): + reg = CapabilityRegistry.from_rules() + gen = RegistryCandidateGenerator(reg) + goal, intent = Goal(text="q"), Intent(bucket="multi_hop") + before = _sig(gen.generate(intent, goal)) + assert ("cheap", "temporal") not in before + + # Register a brand-new retrieval capability serving multi_hop — no edit to rules.py, no prompt. + reg.register(Capability( + key="retrieval:temporal", kind="retrieval", value="temporal", + buckets=frozenset({"multi_hop"}), cost_prior=0.5, quality_prior=0.88, local=True, + )) + after = _sig(gen.generate(intent, goal)) + assert ("cheap", "temporal") in after and ("premium", "temporal") in after + assert before < after # strictly more candidates, everything prior still present + + +def test_for_intent_orders_by_quality_prior(): + reg = CapabilityRegistry.from_rules() + tiers = reg.values_for("multi_hop", "model_tier") # cheap + premium serve multi_hop + assert tiers == ["premium", "cheap"] # best quality first diff --git a/tests/test_governed_optimizer.py b/tests/test_governed_optimizer.py new file mode 100644 index 0000000..9653f98 --- /dev/null +++ b/tests/test_governed_optimizer.py @@ -0,0 +1,87 @@ +"""The optimizer's governance seam — an injected PolicyProvider narrows the feasible space before +cost ranking, and a TrustProvider breaks ties. Uses in-file fakes (no enterprise dependency): the +same Protocols the commercial PolicyEngine / TrustLedger adapters satisfy. +""" +from __future__ import annotations + +from context_runtime.optimizer.knapsack import KnapsackOptimizer +from context_runtime.types import Candidate, Goal, PlanScore, StepSpec + + +def _cand(tier: str, method: str = "bm25") -> Candidate: + return Candidate(steps=(StepSpec(type="retrieve", params={"method": method}),), model_tier=tier) + + +def _score(total: float) -> PlanScore: + return PlanScore(total=total, feasible=True) + + +class _LocalOnlyPolicy: + """Rejects any candidate that isn't a local model — the canonical 'no external egress' policy.""" + + def feasible(self, candidate: Candidate, goal: Goal, score: PlanScore) -> str | None: + if candidate.model_tier != "local": + return "local models only" + return None + + +class _PreferMethod: + """Trust score keyed on the retrieval method — used only to break cost ties.""" + + def __init__(self, favored: str): + self.favored = favored + + def score(self, candidate: Candidate, goal: Goal) -> float: + method = candidate.steps[0].params.get("method") + return 1.0 if method == self.favored else 0.5 + + +def _estimator(): # minimal CostEstimator returning a fixed score per candidate tier + class _E: + def estimate(self, candidate: Candidate, goal: Goal) -> PlanScore: + return _score(0.9 if candidate.model_tier != "local" else 0.4) + def statistics(self, bucket=None): # unused by these tests + raise NotImplementedError + def observe(self, plan, trace): + raise NotImplementedError + return _E() + + +def test_policy_narrows_feasible_space_before_cost(): + """The external plan scores higher, but policy makes it infeasible → the local plan is chosen, + and the external plan appears in rejected with the policy reason (observable in EXPLAIN).""" + goal = Goal(text="q") + ext, loc = _cand("frontier"), _cand("local") + scored = [(ext, _score(0.9)), (loc, _score(0.4))] + opt = KnapsackOptimizer(_estimator(), policy=_LocalOnlyPolicy()) + + plan = opt.select(scored, goal) + + assert plan.chosen is loc, "policy-infeasible high scorer must not win" + ext_reasons = [r for c, r in plan.rejected if c is ext] + assert ext_reasons and "policy: local models only" in ext_reasons[0] + + +def test_policy_marks_score_infeasible(): + goal = Goal(text="q") + opt = KnapsackOptimizer(_estimator(), policy=_LocalOnlyPolicy()) + assert opt.score(_cand("local"), goal).feasible is True + assert opt.score(_cand("frontier"), goal).feasible is False + + +def test_trust_breaks_ties(): + """Two feasible candidates with equal cost total; trust prefers the graph one → it wins.""" + goal = Goal(text="q") + a = _cand("local", method="bm25") + b = _cand("local", method="graph") + scored = [(a, _score(0.5)), (b, _score(0.5))] + opt = KnapsackOptimizer(_estimator(), trust=_PreferMethod("graph")) + assert opt.select(scored, goal).chosen is b + + +def test_no_providers_is_unchanged_behavior(): + """Without policy/trust the optimizer is the pre-seam engine: highest total wins, ties keep order.""" + goal = Goal(text="q") + hi, lo = _cand("frontier"), _cand("local") + plan = KnapsackOptimizer(_estimator()).select([(lo, _score(0.4)), (hi, _score(0.9))], goal) + assert plan.chosen is hi diff --git a/tests/test_job_leads.py b/tests/test_job_leads.py new file mode 100644 index 0000000..6bee87d --- /dev/null +++ b/tests/test_job_leads.py @@ -0,0 +1,127 @@ +"""Job-lead outreach: internal-vs-consulting classification, tailored pitch, dedup, and lead-finding.""" +from __future__ import annotations + +from context_runtime.integrations.job_leads import ( + DEFAULT_TEMPLATE, JobListing, OutreachLedger, PitchTemplate, StaticJobSource, classify, + find_leads, write_pitch, +) + + +def _l(company, title, desc="", url="", **kw): + return JobListing(company=company, title=title, description=desc, url=url, **kw) + + +def test_classify_internal_ai_role_is_a_lead(): + c = classify(_l("Acme Retail", "AI Data Engineer", "Build our in-house RAG platform for production.")) + assert c.is_lead and c.kind == "internal" and c.score >= 0.9 # explicit in-house signals + + +def test_classify_consultancy_is_skipped(): + c = classify(_l("BigConsulting", "AI Engineer", "Deliver forward-deployed AI for our clients.")) + assert not c.is_lead and c.kind == "consulting" + + +def test_classify_non_ai_role_is_other(): + assert classify(_l("Acme", "Frontend Developer", "React and CSS.")).kind == "other" + + +def test_write_pitch_substitutes_and_flags_resume(): + listing = _l("Acme", "AI Developer", "You will build RAG. Attach your CV/resume.") + p = write_pitch(listing, DEFAULT_TEMPLATE, resume_path="/docs/resume.pdf") + assert "your AI Developer opening" in p["body"] + assert p["attach"] == "/docs/resume.pdf" # 'resume' in the description → attach + assert p["subject"].startswith("Re: your AI Developer opening") + # a listing that doesn't ask for a resume → no attachment + assert write_pitch(_l("Acme", "AI Engineer", "Build stuff."), DEFAULT_TEMPLATE, + resume_path="/docs/resume.pdf")["attach"] == "" + + +def test_ledger_dedupes_reappearing_listings(tmp_path): + ledger = OutreachLedger(path=str(tmp_path / "led.jsonl")) + listing = _l("Acme", "AI Data Engineer", url="https://acme/jobs/1") + assert not ledger.already_pitched(listing) + ledger.record(listing) + assert ledger.already_pitched(listing) # same key won't be pitched again + # a fresh ledger loading the same file still sees it (persistent) + assert OutreachLedger(path=str(tmp_path / "led.jsonl")).already_pitched(listing) + + +def test_find_leads_filters_dedupes_and_ranks(tmp_path): + source = StaticJobSource([ + _l("Acme Retail", "AI Data Engineer", "Build our in-house AI platform.", url="u1"), # lead (high) + _l("Startup", "AI Engineer", "Ship AI features.", url="u2"), # lead (base) + _l("BigConsulting", "AI Engineer", "Forward-deployed AI for clients.", url="u3"), # skip + _l("Shop", "Store Manager", "Retail ops.", url="u4"), # skip + ]) + ledger = OutreachLedger(path=str(tmp_path / "led.jsonl")) + leads = find_leads(source, ledger, queries=("x",)) + companies = [d["listing"].company for d in leads] + assert companies == ["Acme Retail", "Startup"] # consultancy + non-AI dropped, best first + # after pitching Acme, it's excluded next time + ledger.record(leads[0]["listing"]) + assert [d["listing"].company for d in find_leads(source, ledger, queries=("x",))] == ["Startup"] + + +def test_adzuna_source_parses_and_degrades(): + from context_runtime.integrations.job_leads import AdzunaSource + + assert AdzunaSource(app_id="", app_key="").search("ai") == [] # unconfigured → empty, no crash + + class _Resp: + def json(self): + return {"results": [{"title": "AI Data Engineer", "company": {"display_name": "Acme"}, + "redirect_url": "https://x/1", "description": "Build our in-house AI platform.", + "location": {"display_name": "NYC"}, "created": "2026-07-01"}]} + + class _Client: + def get(self, url, params=None): + return _Resp() + + got = AdzunaSource(app_id="a", app_key="k", client=_Client()).search("AI Data Engineer", 10) + assert len(got) == 1 and got[0].company == "Acme" and got[0].source == "adzuna" + assert classify(got[0]).is_lead + + +def test_profile_store_owner_vs_generic_and_isolated(tmp_path): + from context_runtime.integrations.job_leads import DEFAULT_TEMPLATE, GENERIC_TEMPLATE, ProfileStore + + store = ProfileStore(dir=str(tmp_path), owner="alex") + owner = store.load("alex") + assert "AI Data Engineer" in owner.include and owner.template == DEFAULT_TEMPLATE # seeded owner config + other = store.load("bob") + assert other.include == [] and other.template == GENERIC_TEMPLATE # blank for others + + other.include = ["DevOps Engineer"] + other.template = "Hi {company}, about {title}. {match}" + store.save(other) + assert store.load("bob").include == ["DevOps Engineer"] # persisted per user + assert store.load("alex").include != store.load("bob").include # isolated + + +def test_classify_for_uses_profile_targets(): + from context_runtime.integrations.job_leads import LeadProfile, classify_for + + bob = LeadProfile("bob", include=["DevOps"], exclude=["consult"]) + assert classify_for(_l("Acme", "Senior DevOps Engineer", "Run our k8s."), bob).is_lead + assert not classify_for(_l("Acme", "AI Data Engineer", "Build AI."), bob).is_lead # not bob's target + assert classify_for(_l("BigCo", "DevOps consultant", "for our clients"), bob).kind == "excluded" + assert classify_for(_l("x", "AI Engineer", ""), LeadProfile("new")).kind == "unset" # no targets set + + +def test_find_leads_for_scopes_to_profile(tmp_path): + from context_runtime.integrations.job_leads import LeadProfile, OutreachLedger, StaticJobSource, find_leads_for + + src = StaticJobSource([_l("A", "DevOps Engineer", "run our infra", url="u1"), + _l("B", "AI Data Engineer", "build our AI", url="u2")]) + leads = find_leads_for(src, OutreachLedger(path=str(tmp_path / "l.jsonl")), + LeadProfile("bob", include=["DevOps"]), queries=("x",)) + assert [d["listing"].company for d in leads] == ["A"] # only bob's targeted role + + +def test_pitch_template_is_shared_and_editable(tmp_path): + tpl = PitchTemplate(path=str(tmp_path / "tpl.txt")) + assert tpl.load() == DEFAULT_TEMPLATE # default until edited + tpl.save("Hi, {company} — about your {title} role. {match}") + assert PitchTemplate(path=str(tmp_path / "tpl.txt")).load().startswith("Hi, {company}") # all users see it + tpl.reset() + assert tpl.load() == DEFAULT_TEMPLATE diff --git a/tests/test_kafka_bus.py b/tests/test_kafka_bus.py new file mode 100644 index 0000000..de11b11 --- /dev/null +++ b/tests/test_kafka_bus.py @@ -0,0 +1,58 @@ +"""KafkaEventBus satisfies the EventBus seam — tested against an injected fake client (no broker), +so the aggregator/loop code is identical whether the bus is in-process or Kafka.""" +from __future__ import annotations + +from context_runtime.integrations.bandit import EpsilonGreedyBandit +from context_runtime.learning import LearningAggregator, OutcomeEvent +from context_runtime.learning.kafka_bus import KafkaEventBus + + +class _Rec: + def __init__(self, value): + self.value = value + + +class _FakeKafka: + """Stand-in for a broker: send() appends; consumer(topic) drains what's queued.""" + + def __init__(self): + self.topics: dict[str, list] = {} + + def send(self, topic, value): + self.topics.setdefault(topic, []).append(value) + + def flush(self): + pass + + def consumer(self, topic): + recs = [_Rec(v) for v in self.topics.get(topic, [])] + self.topics[topic] = [] + return recs + + +def test_event_roundtrips_through_the_kafka_seam(): + bus = KafkaEventBus("x:9092", client=_FakeKafka(), + codec_out=OutcomeEvent.to_dict, codec_in=OutcomeEvent.from_dict) + ev = OutcomeEvent(context="mh", arm="graph:cheap", reward=0.9, seq=1, accepted=True) + bus.publish("outcomes", ev) + got = bus.poll("outcomes") + assert got == [ev] # serialized out, reconstructed in + assert bus.poll("outcomes") == [] # drained + + +def test_plain_dicts_roundtrip_for_the_cdc_stream(): + """No codec → payloads pass through as-is (the vuln / context-freshness stream ships dict rows).""" + bus = KafkaEventBus("x:9092", client=_FakeKafka()) + row = {"cve_id": "CVE-2026-1", "package": "openssl", "cvss": 9.1} + bus.publish("vuln-freshness", row) + assert bus.poll("vuln-freshness") == [row] + + +def test_aggregator_drains_from_the_kafka_bus(): + bus = KafkaEventBus("x:9092", client=_FakeKafka(), + codec_out=OutcomeEvent.to_dict, codec_in=OutcomeEvent.from_dict) + for i, (arm, r) in enumerate([("graph:cheap", 0.9), ("graph:cheap", 0.9), ("hybrid:cheap", 0.2)]): + bus.publish("outcomes", OutcomeEvent(context="mh", arm=arm, reward=r, seq=i + 1)) + agg = LearningAggregator(EpsilonGreedyBandit(arms=())) + assert agg.drain(bus) == 3 + assert agg.bandit.value("mh", "graph:cheap")[1] > agg.bandit.value("mh", "hybrid:cheap")[1] diff --git a/tests/test_learning_loop.py b/tests/test_learning_loop.py new file mode 100644 index 0000000..4143b75 --- /dev/null +++ b/tests/test_learning_loop.py @@ -0,0 +1,76 @@ +"""Phase 4 — the async learning loop: execution → outcome events → aggregator → snapshot → replicas. +Learning happens off the serving path; snapshots distribute it; replays are idempotent.""" +from __future__ import annotations + +from context_runtime.integrations.bandit import EpsilonGreedyBandit +from context_runtime.learning import ( + InMemoryBus, LearnedStateSnapshot, LearningAggregator, OutcomeEvent, +) +from context_runtime.optimizer.online import BanditOptimizer +from context_runtime.types import Candidate, Goal, PlanScore, StepSpec + + +def _cand(tier, method): + return Candidate(steps=(StepSpec(type="retrieve", params={"method": method}),), model_tier=tier) + + +def _score(total, acc=0.9): + return PlanScore(total=total, expected_accuracy=acc, feasible=True) + + +def test_event_from_plan_reads_optimizer_metadata(): + opt = BanditOptimizer(None, epsilon=0.0) + plan = opt.select([(_cand("cheap", "graph"), _score(0.7))], Goal(text="q"), context="multi_hop") + ev = OutcomeEvent.from_plan(plan, reward=0.9, seq=1, accepted=True) + assert ev.context == "multi_hop" and ev.arm == "graph:cheap" and ev.reward == 0.9 + assert ev.accepted is True and OutcomeEvent.from_dict(ev.to_dict()) == ev + + +def test_aggregator_folds_events_off_the_hot_path(): + bus, bandit = InMemoryBus(), EpsilonGreedyBandit(arms=()) + agg = LearningAggregator(bandit) + for i, (arm, r) in enumerate([("graph:cheap", 0.9), ("graph:cheap", 0.9), ("hybrid:cheap", 0.2)]): + bus.publish("outcomes", OutcomeEvent(context="mh", arm=arm, reward=r, seq=i + 1)) + assert agg.drain(bus) == 3 and agg.version == 1 + assert bandit.value("mh", "graph:cheap")[1] > bandit.value("mh", "hybrid:cheap")[1] + + +def test_snapshot_distributes_learning_to_a_replica(): + bus = InMemoryBus() + canonical = EpsilonGreedyBandit(arms=()) + agg = LearningAggregator(canonical) + for i in range(3): + bus.publish("outcomes", OutcomeEvent(context="mh", arm="graph:cheap", reward=0.9, seq=i + 1)) + agg.drain(bus) + agg.publish(bus, "snapshots") + + # a stateless replica reconciles from the published snapshot — it never processed an event itself + replica = EpsilonGreedyBandit(arms=()) + snap = bus.poll("snapshots")[-1] + LearnedStateSnapshot.from_dict(snap.to_dict()).apply_to(replica) # survives serialization + assert replica.value("mh", "graph:cheap") == canonical.value("mh", "graph:cheap") + + # and a replica optimizer now exploits what the aggregator learned, without learning locally + ropt = BanditOptimizer(None, bandit=replica, epsilon=0.0) + goal = Goal(text="q") + plan = ropt.select([(_cand("cheap", "graph"), _score(0.1)), (_cand("cheap", "hybrid"), _score(0.9))], + goal, context="mh") + assert plan.chosen.model_tier == "cheap" and plan.chosen.steps[0].params["method"] == "graph" + + +def test_replay_is_idempotent_by_seq(): + bandit = EpsilonGreedyBandit(arms=()) + agg = LearningAggregator(bandit) + ev = OutcomeEvent(context="mh", arm="graph:cheap", reward=0.9, seq=5) + assert agg.apply(ev) is True + assert agg.apply(ev) is False # same seq → ignored + assert bandit.value("mh", "graph:cheap")[0] == 1 # counted exactly once + + +def test_trust_sink_receives_every_event_and_abstentions_skip_the_arm(): + seen = [] + bandit = EpsilonGreedyBandit(arms=()) + agg = LearningAggregator(bandit, on_trust=seen.append) + agg.apply(OutcomeEvent(context="mh", arm="graph:cheap", reward=0.0, seq=1, abstained=True)) + assert len(seen) == 1 and seen[0].abstained + assert "graph:cheap" not in bandit.stats.get("mh", {}) # abstention rewards no arm; trust still hears it diff --git a/tests/test_mcp.py b/tests/test_mcp.py new file mode 100644 index 0000000..917c530 --- /dev/null +++ b/tests/test_mcp.py @@ -0,0 +1,156 @@ +"""Hermetic MCP adapter tests: real stdio transport + ApprovalPolicy gating.""" +from __future__ import annotations + +import json +import sys +import tempfile +from context_runtime.tools import ApprovalPolicy, ToolRegistry, ToolResult +from context_runtime.tools.mcp import MCPClient, mount_mcp + + +FAKE_SERVER = r'''import sys, json + +def main(): + for line in sys.stdin: + line = line.strip() + if not line: + continue + try: + req = json.loads(line) + except Exception: + continue + rid = req.get("id") + meth = req.get("method") + params = req.get("params") or {} + if meth == "initialize": + res = { + "jsonrpc": "2.0", + "id": rid, + "result": { + "protocolVersion": "2024-11-05", + "capabilities": {"tools": {}}, + "serverInfo": {"name": "fake-mcp", "version": "0"}, + }, + } + print(json.dumps(res, separators=(",", ":")), flush=True) + continue + if rid is None: + # notification (e.g. initialized), no response + continue + if meth == "tools/list": + tools = [ + { + "name": "web_search", + "description": "search the web", + "inputSchema": {"type": "object", "properties": {"query": {"type": "string"}}}, + }, + { + "name": "echo", + "description": "echo text", + "inputSchema": {"type": "object", "properties": {"text": {"type": "string"}}}, + }, + ] + reply = {"jsonrpc": "2.0", "id": rid, "result": {"tools": tools}} + print(json.dumps(reply, separators=(",", ":")), flush=True) + continue + if meth == "tools/call": + nm = (params or {}).get("name", "") + args = (params or {}).get("arguments") or {} + if nm == "echo": + txt = args.get("text", "") + if txt == "ERR": + res = { + "isError": True, + "content": [{"type": "text", "text": "failed"}], + "error": "tool error", + } + else: + res = {"content": [{"type": "text", "text": "echoed:" + txt}]} + print(json.dumps({"jsonrpc": "2.0", "id": rid, "result": res}, separators=(",", ":")), flush=True) + continue + # web_search or default + q = args.get("query", "") + res = {"content": [{"type": "text", "text": "web:" + q}]} + print(json.dumps({"jsonrpc": "2.0", "id": rid, "result": res}, separators=(",", ":")), flush=True) + continue + if rid is not None: + err = {"jsonrpc": "2.0", "id": rid, "error": {"code": -32601, "message": "method not found"}} + print(json.dumps(err, separators=(",", ":")), flush=True) + +if __name__ == "__main__": + main() +''' + + +def _run_with_client_and_registry(policy: ApprovalPolicy | None = None): + reg = ToolRegistry(policy=policy) + with tempfile.TemporaryDirectory() as td: + fake_path = f"{td}/fake_mcp_server.py" + with open(fake_path, "w") as f: + f.write(FAKE_SERVER) + client = MCPClient.stdio([sys.executable, fake_path]) + try: + names = mount_mcp(reg, client) + return reg, names, client + except Exception: + client.close() + raise + + +def test_mcp_mount_registers_tools_via_real_stdio(): + reg, names, client = _run_with_client_and_registry(policy=ApprovalPolicy(mode="bypass")) + try: + assert set(names) == {"web_search", "echo"} + assert set(reg.list()) == {"web_search", "echo"} + finally: + client.close() + + +def test_mcp_echo_runs_under_bypass_and_returns_text(): + reg, names, client = _run_with_client_and_registry(policy=ApprovalPolicy(mode="bypass")) + try: + assert "echo" in reg.list() + res = reg.run("echo", {"text": "hi"}) + assert isinstance(res, ToolResult) + assert res.ok is True + assert "hi" in res.text + assert "echoed" in res.text + finally: + client.close() + + +def test_mcp_tools_are_gated_by_approval_policy(): + # default policy (deny_side_effects) blocks side-effecting MCP tool (default_side=True) + reg, names, client = _run_with_client_and_registry() + try: + res = reg.run("echo", {"text": "hi"}) + assert res.ok is False + assert "denied" in (res.error or "") + finally: + client.close() + + # with allowlist it passes + reg2 = ToolRegistry(ApprovalPolicy(mode="allowlist", allow=["echo", "web_search"])) + with tempfile.TemporaryDirectory() as td: + fake_path = f"{td}/fake_mcp_server.py" + with open(fake_path, "w") as f: + f.write(FAKE_SERVER) + client2 = MCPClient.stdio([sys.executable, fake_path]) + try: + mount_mcp(reg2, client2) + res2 = reg2.run("echo", {"text": "allow"}) + assert res2.ok is True + assert "allow" in res2.text + finally: + client2.close() + + +def test_mcp_iserror_maps_to_toolresult_failure(): + reg, names, client = _run_with_client_and_registry(policy=ApprovalPolicy(mode="bypass")) + try: + res = reg.run("echo", {"text": "ERR"}) + assert res.ok is False + assert res.error is not None and "error" in res.error.lower() + assert "failed" in (res.text or "") + finally: + client.close() diff --git a/tests/test_online_optimizer.py b/tests/test_online_optimizer.py new file mode 100644 index 0000000..9505202 --- /dev/null +++ b/tests/test_online_optimizer.py @@ -0,0 +1,95 @@ +"""Generation-4 online-learning optimizer: contextual-bandit selection with exploration, learning +from reward, off-policy evaluation, and composition with the Phase-0 governance seam. +""" +from __future__ import annotations + +from context_runtime.optimizer.online import BanditOptimizer, offpolicy_values, plan_key +from context_runtime.types import Candidate, Goal, PlanScore, StepSpec + + +def _cand(tier: str, method: str = "bm25") -> Candidate: + return Candidate(steps=(StepSpec(type="retrieve", params={"method": method}),), model_tier=tier) + + +def _score(total: float) -> PlanScore: + return PlanScore(total=total, feasible=True) + + +class _Est: # unused by select(); present so BanditOptimizer(estimator=...) is realistic + def estimate(self, c, g): return _score(0.5) + def statistics(self, bucket=None): raise NotImplementedError + def observe(self, p, t): raise NotImplementedError + + +def test_exploit_picks_highest_prior_when_epsilon_zero(): + opt = BanditOptimizer(_Est(), epsilon=0.0) + goal = Goal(text="q") + hi, lo = _cand("premium", "hybrid"), _cand("local", "bm25") + plan = opt.select([(lo, _score(0.4)), (hi, _score(0.9))], goal, context="synthesis") + assert plan.chosen is hi + assert plan.extra["bandit"]["mode"] == "exploit" + assert plan.extra["bandit"]["p"] == 1.0 # greedy under ε=0 has propensity 1 + + +def test_exploration_serves_a_non_greedy_arm(): + opt = BanditOptimizer(_Est(), epsilon=1.0) # always explore + goal = Goal(text="q") + hi, lo = _cand("premium", "hybrid"), _cand("local", "bm25") + plan = opt.select([(hi, _score(0.9)), (lo, _score(0.4))], goal, context="synthesis") + assert plan.chosen is lo # the non-greedy arm + assert plan.extra["bandit"]["mode"] == "explore" + assert plan.extra["bandit"]["p"] == 0.5 # ε/m with m=2 arms + + +def test_learning_shifts_the_exploit_choice(): + opt = BanditOptimizer(_Est(), epsilon=0.0) + goal = Goal(text="q") + a, b = _cand("premium", "hybrid"), _cand("local", "graph") # a prior 0.9 > b prior 0.4 + scored = [(a, _score(0.9)), (b, _score(0.4))] + assert opt.select(scored, goal, context="multi_hop").chosen is a # prior wins first + + for _ in range(3): # b turns out to be worth 1.0 in practice + opt.learn("multi_hop", b, 1.0) + assert opt.select(scored, goal, context="multi_hop").chosen is b # learned value beats a's prior + + +def test_learn_from_plan_uses_recorded_context_and_arm(): + opt = BanditOptimizer(_Est(), epsilon=0.0) + goal = Goal(text="q") + b = _cand("local", "graph") + plan = opt.select([(_cand("premium", "hybrid"), _score(0.9)), (b, _score(0.4))], goal, context="mh") + # feed reward back for the *served* arm via the plan's recorded (context, arm) + for _ in range(3): + opt.learn_from_plan( + opt.select([(b, _score(0.4))], goal, context="mh"), 1.0 + ) + n, mean = opt.bandit.value("mh", plan_key(b)) + assert n == 3 and mean == 1.0 + + +class _LocalOnly: + def feasible(self, cand, goal, score): + return None if cand.model_tier == "local" else "local models only" + + +def test_policy_filters_before_any_exploration(): + """Even with ε=1 (always explore), a policy-infeasible arm is never in the pool → never served.""" + opt = BanditOptimizer(_Est(), epsilon=1.0, policy=_LocalOnly()) + goal = Goal(text="q") + ext, loc = _cand("premium", "hybrid"), _cand("local", "bm25") + plan = opt.select([(ext, _score(0.9)), (loc, _score(0.4))], goal, context="synthesis") + assert plan.chosen is loc + assert any(c is ext and r.startswith("policy:") for c, r in plan.rejected) + + +def test_offpolicy_evaluation_surfaces_a_rare_but_better_arm(): + """Tier 1: a high-reward arm that was rarely served (low p) still scores highest, because IPS + up-weights it — so the planner can prefer it without ever running it on the live path.""" + logs = [ + {"context": "mh", "arm": "hybrid:cheap", "reward": 0.2, "p": 0.9}, # served often, mediocre + {"context": "mh", "arm": "hybrid:cheap", "reward": 0.3, "p": 0.9}, + {"context": "mh", "arm": "graph:cheap", "reward": 0.9, "p": 0.1}, # served rarely, great + ] + vals = offpolicy_values(logs)["mh"] + assert vals["graph:cheap"][1] > vals["hybrid:cheap"][1] + assert abs(vals["hybrid:cheap"][1] - 0.25) < 1e-9 # self-normalized IPS = mean reward here diff --git a/tests/test_policy_runtime.py b/tests/test_policy_runtime.py new file mode 100644 index 0000000..373da81 --- /dev/null +++ b/tests/test_policy_runtime.py @@ -0,0 +1,116 @@ +"""Policy Runtime core — RuleStore (long-term memory), providers, decisions+audit, command framework.""" +from __future__ import annotations + +from types import SimpleNamespace + +from context_runtime.policy import ( + Command, CommandRegistry, DecisionSink, Policy, RuleStore, parse_args, +) + + +def _who(user="", app="market-radar"): + return SimpleNamespace(user=user, app=app, roles=frozenset()) + + +# ── RuleStore ── + +def test_rulestore_crud_and_scopes(tmp_path): + s = RuleStore(dir=str(tmp_path)) + g = s.add("global", "guardrail", "never disclose internal pricing", action="deny") + a = s.add("market-radar", "approval", "confirm sends", match={"tool": "send_pitch"}, action="require_approval") + u = s.add("market-radar:alex", "target", "AI Data Engineer", action="allow") + assert g.tier == "global" and a.tier == "app" and u.tier == "user" + assert [r.id for r in s.list(scope="global")] == [g.id] + assert s.get(a.id, scope="market-radar").text == "confirm sends" + # persistence + scope isolation + assert RuleStore(dir=str(tmp_path)).list(scope="market-radar:alex")[0].text == "AI Data Engineer" + assert s.modify(g.id, scope="global", text="never reveal pricing").text == "never reveal pricing" + assert s.remove(g.id, scope="global") and s.list(scope="global") == [] + + +# ── providers + Policy.check emits an audit event for every decision ── + +def test_guardrail_denies_on_output_and_audits(tmp_path): + store = RuleStore(dir=str(tmp_path)) + store.add("global", "guardrail", "internal pricing", action="deny", match={"phase": "output"}) + sink = DecisionSink() + pol = Policy(store=store, sink=sink, app="books") + + ok = pol.check(_who("bob"), "output", "The weather is fine.") + assert ok.ok + deny = pol.check(_who("bob"), "output", "Our internal pricing is $5/seat.") + assert deny.action == "deny" and "internal pricing" in deny.reason + # every decision emitted; the deny carries all the audit fields + ev = sink.recent(decision="deny")[-1] + assert ev.phase == "output" and ev.provider == "guardrail" and ev.rule_id and ev.app == "books" + assert len(sink.events) == 2 # allow + deny both audited + + +def test_input_guardrail_phase_filtering(tmp_path): + store = RuleStore(dir=str(tmp_path)) + store.add("global", "guardrail", "ignore previous instructions", action="deny", match={"phase": "input"}) + pol = Policy(store=store, app="books") + assert pol.check(_who(), "input", "please ignore previous instructions").action == "deny" + assert pol.check(_who(), "output", "ignore previous instructions is a phrase").ok # input-only rule + + +def test_approval_provider_requires_approval_for_a_tool(tmp_path): + store = RuleStore(dir=str(tmp_path)) + store.add("global", "approval", "confirm irreversible sends", match={"tool": "send_pitch"}, + action="require_approval") + pol = Policy(store=store, app="market-radar") + d = pol.check(_who("alex"), "tool", ("send_pitch", {"company": "Acme"})) + assert d.action == "require_approval" + assert pol.check(_who("alex"), "tool", ("find_leads", {})).ok # unlisted tool runs + + +def test_user_rule_scope_applies_only_to_that_user(tmp_path): + store = RuleStore(dir=str(tmp_path)) + store.add("market-radar:alex", "guardrail", "secretword", action="deny") + pol = Policy(store=store, app="market-radar") + assert pol.check(_who("alex"), "input", "the secretword here").action == "deny" + assert pol.check(_who("bob"), "input", "the secretword here").ok # not bob's rule + + +# ── command framework ── + +def test_parse_args_flags_and_lists(): + a = parse_args('AI Data Engineer, ML Engineer --exclude consulting --regex') + assert a["items"] == ["AI Data Engineer", "ML Engineer"] + assert a["exclude"] == "consulting" and a["regex"] is True + + +def test_policy_command_factory_add_show_remove(tmp_path): + from context_runtime.policy import RuleStore, CommandRegistry, global_policy_commands + store = RuleStore(dir=str(tmp_path)) + reg = CommandRegistry(can=lambda p, req: req == "" or "admin" in getattr(p, "roles", ())) + reg.register_all(global_policy_commands(store)) + admin = SimpleNamespace(user="root", app="x", roles=frozenset({"admin"})) + + r = reg.dispatch("/addpolicy no internal pricing --phase output", admin) + assert r["ok"] and store.list(scope="global", kind="guardrail")[0].text == "no internal pricing" + assert "no internal pricing" in reg.dispatch("/showpolicy", _who("bob"))["text"] # read is open + assert reg.dispatch(f"/removepolicy {r['data']['id']}", admin)["ok"] and store.list(scope="global") == [] + assert not reg.dispatch("/addpolicy x", _who("bob"))["ok"] # non-admin blocked + assert reg.dispatch("/addguardrail block this", admin)["ok"] # alias works + + +def test_command_dispatch_permission_and_help(): + calls = [] + reg = CommandRegistry(can=lambda p, req: req == "self" or (req == "policy-admin" and "admin" in getattr(p, "roles", ()))) + reg.register(Command("addtarget", lambda args, p: (calls.append(args["items"]), {"text": "ok"})[1], + "add target roles", requires="self")) + reg.register(Command("addpolicy", lambda args, p: {"text": "added"}, "add global policy", + requires="policy-admin", aliases=("addguardrail",))) + + assert reg.is_command("/addtarget X") and not reg.is_command("hello") + assert reg.dispatch("/addtarget AI, ML", _who("bob"))["ok"] and calls[-1] == ["AI", "ML"] + # non-admin denied on a policy-admin command (with a command-phase policy note) + denied = reg.dispatch("/addpolicy no pricing", _who("bob")) + assert not denied["ok"] and denied["policy"][0]["decision"] == "deny" + # alias resolves + admin allowed + admin = SimpleNamespace(user="alex", app="x", roles=frozenset({"admin"})) + assert reg.dispatch("/addguardrail no pricing", admin)["text"] == "added" + # help lists only permitted commands + assert "addtarget" in reg.dispatch("/help", _who("bob"))["data"]["commands"] + assert "addpolicy" not in reg.dispatch("/help", _who("bob"))["data"]["commands"] diff --git a/tests/test_representation_planning.py b/tests/test_representation_planning.py new file mode 100644 index 0000000..ea90f38 --- /dev/null +++ b/tests/test_representation_planning.py @@ -0,0 +1,86 @@ +"""v4 knowledge-aware planning — intent → representation → constrained candidates → learning. + +Proves the wiring the whitepaper-v4 claims: the decision engine classifies a knowledge +representation, candidate generation is constrained to that representation's methods (with a +document fallback + confidence-gated exploration), and the choice is recorded on the plan/event. +""" +from __future__ import annotations + +from context_runtime.planner import representations +from context_runtime.planner.candidates import RuleCandidateGenerator +from context_runtime.planner.intent import RuleIntentAnalyzer +from context_runtime.learning.events import OutcomeEvent +from context_runtime.types import Goal, Intent, Candidate, StepSpec, PlanScore, Plan + + +def _methods(cands): + out = [] + for c in cands: + for s in c.steps: + if s.type == "retrieve": + out.append(s.params["method"]) + return set(out) + + +def _plan_for(text: str): + intent = RuleIntentAnalyzer().analyze(Goal(text=text)) + cands = RuleCandidateGenerator().generate(intent, Goal(text=text)) + return intent, cands + + +# ─── classify head ─────────────────────────────────────────────────────────── +def test_classify_maps_intent_to_representation(): + assert representations.classify("multi_hop", "how does X relate to Y") == "graph" + assert representations.classify("temporal", "what was the plan as of March") == "temporal" + assert representations.classify("code_reasoning", "refactor this function") == "code" + assert representations.classify("conceptual", "what is idempotency") == "document" + # HINTS reach representations no bucket produces on its own: + assert representations.classify("incident", "how many incidents last week") == "analytical" + assert representations.classify("conceptual", "what's in this screenshot") == "multimodal" + + +def test_analyzer_sets_representation(): + assert RuleIntentAnalyzer().analyze(Goal(text="how does A depend on B")).representation == "graph" + assert RuleIntentAnalyzer().analyze(Goal(text="explain caching")).representation == "document" + + +# ─── candidates constrained to the representation ──────────────────────────── +def test_graph_query_constrains_to_graph_plus_document_fallback(): + intent, cands = _plan_for("what is the dependency chain between service A and service B") + assert intent.representation == "graph" + methods = _methods(cands) + assert "graph" in methods and "hybrid" in methods # representation-first + fallback + assert "vector" not in methods and "bm25" not in methods # not the flat document table + + +def test_analytical_hint_generates_olap_candidates(): + intent, cands = _plan_for("how many deploys happened per week over the last month") + assert intent.representation == "analytical" + methods = _methods(cands) + assert methods & {"sql", "logs", "api"} # OLAP candidates now reachable + assert "hybrid" in methods # with a document fallback + + +def test_plain_document_query_unchanged(): + intent, cands = _plan_for("what is eventual consistency") # conceptual, no graph/temporal cue + assert intent.representation == "document" + assert _methods(cands) <= {"vector", "bm25", "hybrid", "file"} + + +def test_bandit_arms_span_representations_for_learning(): + # multi_hop routes to graph but keeps the document fallback -> the bandit's arms span TWO + # representations, so choosing among them IS representation selection (what it learns online). + _, cands = _plan_for("trace how the outage propagated across the payment and auth services") + reps = {representations.representation_for(m) for m in _methods(cands)} + assert {"graph", "document"} <= reps + + +# ─── representation recorded on the outcome event ──────────────────────────── +def test_representation_recorded_on_outcome_event(): + intent = Intent(bucket="multi_hop", representation="graph") + plan = Plan(intent=intent, + chosen=Candidate(steps=(StepSpec("retrieve", {"method": "graph"}),)), + score=PlanScore(), + extra={"bandit": {"context": "multi_hop", "arm": "graph:cheap"}}) + ev = OutcomeEvent.from_plan(plan, reward=1.0) + assert ev.representation == "graph" and ev.context == "multi_hop" diff --git a/tests/test_supply_chain.py b/tests/test_supply_chain.py new file mode 100644 index 0000000..98fb9f5 --- /dev/null +++ b/tests/test_supply_chain.py @@ -0,0 +1,48 @@ +"""Hermetic tests for the supply-chain scanner — parse canned Trivy output; graceful degrade.""" +from __future__ import annotations + +from context_runtime.integrations.supply_chain import ScanResult, SupplyChainScanner + +_TRIVY_JSON = """ +{"Results":[ + {"Target":"requirements.txt","Vulnerabilities":[ + {"VulnerabilityID":"CVE-2024-0001","PkgName":"requests","InstalledVersion":"2.20.0","FixedVersion":"2.31.0","Severity":"HIGH","Title":"requests TLS bypass"}, + {"VulnerabilityID":"CVE-2023-9999","PkgName":"jinja2","InstalledVersion":"2.11.0","FixedVersion":"","Severity":"CRITICAL","Title":"jinja2 sandbox escape"}, + {"VulnerabilityID":"CVE-2022-1234","PkgName":"idna","InstalledVersion":"2.8","FixedVersion":"3.3","Severity":"LOW","Title":"idna dos"}], + "Secrets":[{"RuleID":"aws-access-key"}], + "Misconfigurations":[{"ID":"DS002"}]} +]} +""" + + +def test_parse_trivy_normalizes_and_sorts_by_severity(): + r = SupplyChainScanner._parse_trivy(_TRIVY_JSON, "requirements.txt") + assert r.ok and r.scanner == "trivy" + assert len(r.findings) == 3 + # CRITICAL first, then HIGH, then LOW + assert [f.severity for f in r.findings] == ["CRITICAL", "HIGH", "LOW"] + assert r.secrets == 1 and r.misconfigs == 1 + # the fixed/"patched pin" is carried through + high = next(f for f in r.findings if f.id == "CVE-2024-0001") + assert high.pkg == "requests" and high.fixed == "2.31.0" + + +def test_summary_counts(): + r = SupplyChainScanner._parse_trivy(_TRIVY_JSON, "requirements.txt") + s = r.summary() + assert s["total"] == 3 + assert s["by_severity"] == {"CRITICAL": 1, "HIGH": 1, "LOW": 1} + assert s["fixable"] == 2 # jinja2 has no fix yet + assert s["secrets"] == 1 and s["misconfigs"] == 1 + + +def test_graceful_degrade_when_trivy_absent(monkeypatch): + monkeypatch.setattr("shutil.which", lambda _name: None) + r = SupplyChainScanner().scan_fs("/whatever") + assert isinstance(r, ScanResult) and r.ok is False + assert "not installed" in r.note + + +def test_bad_output_is_handled(): + r = SupplyChainScanner._parse_trivy("not json", "x") + assert r.ok is False and "unparseable" in r.note diff --git a/tests/test_supply_chain_scan.py b/tests/test_supply_chain_scan.py new file mode 100644 index 0000000..362f9d9 --- /dev/null +++ b/tests/test_supply_chain_scan.py @@ -0,0 +1,90 @@ +"""Hermetic tests for container scan + triage — NO docker, NO network.""" +from __future__ import annotations + +import pytest + +from context_runtime.integrations.supply_chain import ScanResult, SupplyChainScanner + +_TRIVY_JSON = """ +{"Results":[ + {"Target":"requirements.txt","Vulnerabilities":[ + {"VulnerabilityID":"CVE-2024-0001","PkgName":"requests","InstalledVersion":"2.20.0","FixedVersion":"2.31.0","Severity":"HIGH","Title":"requests TLS bypass"}, + {"VulnerabilityID":"CVE-2023-9999","PkgName":"jinja2","InstalledVersion":"2.11.0","FixedVersion":"","Severity":"CRITICAL","Title":"jinja2 sandbox escape"}, + {"VulnerabilityID":"CVE-2022-1234","PkgName":"idna","InstalledVersion":"2.8","FixedVersion":"3.3","Severity":"LOW","Title":"idna dos"}], + "Secrets":[{"RuleID":"aws-access-key"}], + "Misconfigurations":[{"ID":"DS002"}]} +]} +""" + + +def test_resolve_image_parses_docker_inspect(monkeypatch): + sc = SupplyChainScanner() + monkeypatch.setattr(sc, "_run", lambda cmd: (0, "sha256:abc123\n", "")) + assert sc.resolve_image("foo") == "sha256:abc123" + monkeypatch.setattr(sc, "_run", lambda cmd: (1, "", "boom")) + assert sc.resolve_image("bar") == "" + + +def test_scan_container_happy_and_not_resolvable(monkeypatch): + sc = SupplyChainScanner() + monkeypatch.setattr("shutil.which", lambda n: True) + monkeypatch.setattr(sc, "resolve_image", lambda c: "img:1") + monkeypatch.setattr(sc, "scan_image", lambda img: SupplyChainScanner._parse_trivy(_TRIVY_JSON, img)) + r = sc.scan_container("c1") + assert r.ok and "c1" in r.target and "img:1" in r.target + + monkeypatch.setattr(sc, "resolve_image", lambda c: "") + r2 = sc.scan_container("c2") + assert r2.ok is False and "could not resolve image" in r2.note + + +def test_list_scannable_containers_parses_and_filters(monkeypatch): + sc = SupplyChainScanner() + out = "c1\timg:1\nc2\timg:2\n" + monkeypatch.setattr(sc, "_run", lambda cmd: (0, out, "")) + rows = sc.list_scannable_containers() + assert rows == [{"name": "c1", "image": "img:1"}, {"name": "c2", "image": "img:2"}] + assert sc.list_scannable_containers("2") == [{"name": "c2", "image": "img:2"}] + monkeypatch.setattr(sc, "_run", lambda cmd: (1, "", "")) + assert sc.list_scannable_containers() == [] + + +def test_triage_orders_fixable_and_caps_and_handles_bad_result(): + sc = SupplyChainScanner() + bad = ScanResult(ok=False, target="x", scanner="trivy", note="fail") + t = sc.triage(bad) + assert t == {"summary": "fail", "fixes": [], "note": "fail"} + + findings = [ + type("F", (), {"id": "C1", "pkg": "p", "installed": "1", "fixed": "2", "severity": "CRITICAL"})(), + type("F", (), {"id": "H1", "pkg": "q", "installed": "1", "fixed": "", "severity": "HIGH"})(), + type("F", (), {"id": "L1", "pkg": "r", "installed": "1", "fixed": "3", "severity": "LOW"})(), + ] + good = ScanResult(ok=True, target="t", scanner="trivy", findings=findings) + out = sc.triage(good, top=1) + assert len(out["fixes"]) == 1 + assert out["fixes"][0]["action"] == "upgrade p 1 → 2" + assert "critical" in out["summary"] + + nofix = ScanResult(ok=True, target="t", scanner="trivy", findings=[findings[1]]) + out2 = sc.triage(nofix) + assert out2["note"] == "no fixable findings" + +def test_advise_recommends_base_hardening_when_os_dominates(): + from context_runtime.integrations.supply_chain import SupplyChainScanner as S + raw = ('{"Results":[{"Class":"os-pkgs","Vulnerabilities":[' + '{"VulnerabilityID":"CVE-1","PkgName":"openssl","InstalledVersion":"3.5.5","FixedVersion":"3.5.6","Severity":"HIGH"},' + '{"VulnerabilityID":"CVE-2","PkgName":"libcap2","InstalledVersion":"1","FixedVersion":"2","Severity":"HIGH"}]},' + '{"Class":"lang-pkgs","Vulnerabilities":[' + '{"VulnerabilityID":"CVE-3","PkgName":"pip","InstalledVersion":"25.0","FixedVersion":"25.3","Severity":"MEDIUM"}]}]}') + a = S().advise(S._parse_trivy(raw)) + assert a["os"] == 2 and a["lang"] == 1 + assert "BASE IMAGE" in a["recommendation"] and "python:3.12-slim" in a["recommendation"] + + +def test_advise_flags_app_deps_when_lang_dominates(): + from context_runtime.integrations.supply_chain import SupplyChainScanner as S + raw = ('{"Results":[{"Class":"lang-pkgs","Vulnerabilities":[' + '{"VulnerabilityID":"CVE-9","PkgName":"jinja2","InstalledVersion":"2.1","FixedVersion":"3.1.4","Severity":"HIGH"}]}]}') + a = S().advise(S._parse_trivy(raw)) + assert a["lang"] == 1 and a["os"] == 0 and "OWN dependencies" in a["recommendation"] diff --git a/tests/test_temporal_routing.py b/tests/test_temporal_routing.py new file mode 100644 index 0000000..576bc65 --- /dev/null +++ b/tests/test_temporal_routing.py @@ -0,0 +1,112 @@ +"""Temporal (bi-temporal) routing: intent → method, the router slot + fallback, cost edge, +the point-in-time / what-changed semantics, and the representation taxonomy (Whitepaper v4).""" +from __future__ import annotations + +from context_runtime import ContextRuntime, Goal +from context_runtime.adapters.store_inmemory import InMemoryStore +from context_runtime.adapters.store_router import HopRouterRetriever +from context_runtime.adapters.store_temporal import TemporalStore +from context_runtime.planner import representations +from context_runtime.plugins import base + +# A tiny bi-temporal history: the client's vitamin-D dose was revised on 2026-03-01. +DOC = [{"chunk_id": "d1", "filename": "d1", "text": "The client takes vitamin D daily.", "created_at": None}] + + +def _history() -> TemporalStore: + return (TemporalStore() + .add("client", "takes", "vitamin D 5000 IU", + valid_from="2026-01-01", valid_to="2026-03-01", recorded_at="2026-01-01") + .add("client", "takes", "vitamin D 2000 IU", + valid_from="2026-03-01", recorded_at="2026-03-01")) # current (valid_to=None) + + +TEMPORAL_QS = [ + "What was the client's dose as of February 2026?", + "What changed in the client's supplements?", + "Which supplement did the client previously take?", + "Is the client still on the higher dose?", +] + + +def test_planner_routes_temporal_queries_to_temporal_method(): + rt = ContextRuntime.default([]) + for q in TEMPORAL_QS: + p = rt.plan(q) + assert p.intent.bucket == "temporal", f"{q!r} → {p.intent.bucket}" + method = next(s.params.get("method") for s in p.chosen.steps if s.type == "retrieve") + assert method == "temporal", f"{q!r} chose {method}" + + +def test_non_temporal_queries_do_not_route_to_temporal(): + rt = ContextRuntime.default([]) + for q in ("What is reciprocal rank fusion?", "Find ERR-500 in the logs"): + p = rt.plan(q) + method = next(s.params.get("method") for s in p.chosen.steps if s.type == "retrieve") + assert method != "temporal", f"{q!r} should not use temporal" + + +def test_temporal_candidate_beats_hybrid_on_a_temporal_query(): + # inverse of the graph case: for a temporal question, temporal must score HIGHER than hybrid + rt = ContextRuntime.default([]) + g = Goal(text="What changed in the client's supplements?") + from context_runtime.types import Candidate, StepSpec + temporal = Candidate(steps=(StepSpec("retrieve", {"method": "temporal"}),), model_tier="cheap") + hybrid = Candidate(steps=(StepSpec("retrieve", {"method": "hybrid"}),), model_tier="cheap") + assert rt.optimizer.score(temporal, g).total > rt.optimizer.score(hybrid, g).total + + +def test_router_dispatches_temporal_to_the_bitemporal_store(): + router = HopRouterRetriever(single_hop=InMemoryStore(list(DOC)), + graph=InMemoryStore(list(DOC)), temporal=_history()) + assert isinstance(router, base.RetrieverPlugin) + hits = router.search("client vitamin dose", k=5, method="temporal") + assert hits and all(h.source == "temporal" for h in hits) + # "current" search returns only the not-yet-superseded fact + assert any("2000 IU" in h.text for h in hits) + assert not any("5000 IU" in h.text for h in hits) + + +def test_router_falls_back_when_temporal_unwired_or_empty(): + # (a) no temporal slot → temporal method falls through to single-hop, no crash + r1 = HopRouterRetriever(single_hop=InMemoryStore(list(DOC)), graph=InMemoryStore(list(DOC))) + h1 = r1.search("client vitamin", k=3, method="temporal") + assert all(h.source != "temporal" for h in h1) + # (b) wired but EMPTY store → also falls back (never returns nothing when documents exist) + r2 = HopRouterRetriever(single_hop=InMemoryStore(list(DOC)), + graph=InMemoryStore(list(DOC)), temporal=TemporalStore()) + h2 = r2.search("client vitamin", k=3, method="temporal") + assert h2 and all(h.source != "temporal" for h in h2) + + +def test_bitemporal_point_in_time_and_changes(): + hist = _history() + now = hist.search("client vitamin", k=5) # current state + assert any("2000 IU" in h.text for h in now) + past = hist.as_of("client vitamin", at="2026-02-01") # as it was in February + assert any("5000 IU" in h.text for h in past) + assert not any("2000 IU" in h.text for h in past) + changes = hist.changes("vitamin", since="2026-01-01", until="2026-06-01") + kinds = {(c["at"], c["change"]) for c in changes} + assert ("2026-03-01", "began") in kinds and ("2026-03-01", "ended") in kinds + + +def test_representation_taxonomy(): + assert representations.representation_for("temporal") == "temporal" + assert representations.representation_for("graph") == "graph" + assert representations.representation_for("hybrid") == "document" + assert representations.representation_for("sql") == "analytical" + assert representations.representation_for("video") == "multimodal" + assert representations.representation_for("unknown-method") == "document" # safe default + assert "hybrid" in representations.methods_for("document") + + +def test_end_to_end_temporal_run_uses_the_history(): + router = HopRouterRetriever(single_hop=InMemoryStore(list(DOC)), + graph=InMemoryStore(list(DOC)), temporal=_history()) + from context_runtime.adapters.model_stub import StubModel + rt = ContextRuntime(models={t: StubModel(tier=t) for t in ("local", "cheap", "premium")}, + retriever=router) + res = rt.run("What did the client previously take, and what is it now?") + # the planner routed to temporal; the assembled context carries the bi-temporal facts + assert any("client:takes" in c for c in res.citations) diff --git a/tests/test_temporal_store.py b/tests/test_temporal_store.py new file mode 100644 index 0000000..8fd6650 --- /dev/null +++ b/tests/test_temporal_store.py @@ -0,0 +1,61 @@ +"""Bi-temporal retrieval: point-in-time queries over valid time AND transaction time, and a +"what changed, and when?" scan. Dependency-free; satisfies the RetrieverPlugin seam.""" +from __future__ import annotations + +from context_runtime.adapters.store_temporal import TemporalStore +from context_runtime.plugins.base import RetrieverPlugin + + +def _store() -> TemporalStore: + s = TemporalStore() + # the auth service's owner is revised over time (valid-time history) + s.add("auth-service", "owner", "Alice", valid_from="2026-01-01", valid_to="2026-06-01", + recorded_at="2026-01-01") + s.add("auth-service", "owner", "Bob", valid_from="2026-06-01", valid_to=None, + recorded_at="2026-07-15") # recorded LATE (a correction filed in July) + s.add("billing-service", "owner", "Carol", valid_from="2026-02-01", valid_to=None, + recorded_at="2026-02-01") + return s + + +def test_satisfies_retriever_plugin_seam(): + assert isinstance(_store(), RetrieverPlugin) + info = _store().info() + assert "temporal" in info.capabilities and info.name == "temporal" + + +def test_current_state_returns_the_latest_valid_fact(): + hits = _store().search("who owns auth-service", k=5) + owners = [h.meta["object"] for h in hits if h.meta["subject"] == "auth-service"] + assert owners == ["Bob"] # Alice's record was superseded (valid_to set) + + +def test_as_of_valid_time_travels_the_world_history(): + s = _store() + assert s.as_of("auth-service", at="2026-03-01")[0].meta["object"] == "Alice" + assert s.as_of("auth-service", at="2026-09-01")[0].meta["object"] == "Bob" + # boundary: valid_from inclusive, valid_to exclusive + assert s.as_of("auth-service", at="2026-06-01")[0].meta["object"] == "Bob" + + +def test_as_of_transaction_time_ignores_later_corrections(): + """As known on 2026-06-15, Bob's record (filed 2026-07-15) didn't exist yet — so even at a world + time inside Bob's validity, the state we *believed then* still shows Alice (whose interval we knew).""" + s = _store() + seen = s.as_of("auth-service", at="2026-06-15", known_at="2026-06-15") + # Bob is filtered out (recorded_at 2026-07-15 > known_at); Alice's interval ended 2026-06-01, so + # nothing valid AND known → empty (we honestly had no recorded owner for that instant yet) + assert [h.meta["object"] for h in seen] == [] + # without the transaction-time bound, the corrected history shows Bob + assert s.as_of("auth-service", at="2026-06-15")[0].meta["object"] == "Bob" + + +def test_changes_reports_what_changed_and_when(): + s = _store() + changes = s.changes("auth-service", since="2026-01-01", until="2027-01-01") + # Alice began (01-01) + ended (06-01), Bob began (06-01) + assert [(c["at"], c["change"], c["object"]) for c in changes] == [ + ("2026-01-01", "began", "Alice"), + ("2026-06-01", "ended", "Alice"), + ("2026-06-01", "began", "Bob"), + ] diff --git a/tests/test_tool_authorizer.py b/tests/test_tool_authorizer.py new file mode 100644 index 0000000..d4c2e0c --- /dev/null +++ b/tests/test_tool_authorizer.py @@ -0,0 +1,38 @@ +"""The tool-access authorization seam — an optional authorizer gates every tool call by principal, +before the side-effect approval gate. The enterprise permissions plane plugs in here (open-core).""" +from __future__ import annotations + +from context_runtime.tools.base import ToolRegistry, ToolResult, function_tool, set_default_authorizer + + +def _reg(authorizer=None) -> ToolRegistry: + r = ToolRegistry(authorizer=authorizer) + r.register(function_tool("read_vulns", lambda a: ToolResult(ok=True, text="rows"))) + return r + + +def test_no_authorizer_is_unchanged_behavior(): + assert _reg().run("read_vulns", {}).ok + + +def test_authorizer_denies_by_principal(): + # deny 'read_vulns' unless the principal carries the 'security' role + def authz(principal, spec, args): + roles = (principal or {}).get("roles", set()) + return None if spec.name != "read_vulns" or "security" in roles else "not permitted" + + reg = _reg(authz) + assert reg.run("read_vulns", {}, principal={"roles": {"security"}}).ok + denied = reg.run("read_vulns", {}, principal={"roles": {"guest"}}) + assert not denied.ok and "not permitted" in denied.error + assert reg.audit[-1]["allowed"] is False # the denial is audited + + +def test_default_authorizer_governs_all_registries_then_clears(): + reg = _reg() + set_default_authorizer(lambda p, spec, args: "blocked fleet-wide") + try: + assert not reg.run("read_vulns", {}).ok # a registry with no authorizer inherits the default + finally: + set_default_authorizer(None) + assert reg.run("read_vulns", {}).ok diff --git a/tests/test_trust_aware.py b/tests/test_trust_aware.py new file mode 100644 index 0000000..7870c2a --- /dev/null +++ b/tests/test_trust_aware.py @@ -0,0 +1,83 @@ +"""Generation 5 — Trust-Aware Execution: trust as a ranking objective (not just a tie-breaker) and a +calibrated abstention gate (serve / escalate / abstain). +""" +from __future__ import annotations + +from context_runtime.abstention import AbstentionGate +from context_runtime.optimizer.online import BanditOptimizer +from context_runtime.types import Candidate, Goal, PlanScore, StepSpec + + +def _cand(tier: str, method: str = "bm25") -> Candidate: + return Candidate(steps=(StepSpec(type="retrieve", params={"method": method}),), model_tier=tier) + + +def _score(total: float, acc: float = 0.9) -> PlanScore: + return PlanScore(total=total, expected_accuracy=acc, feasible=True) + + +class _Trust: + def __init__(self, scores): + self.scores = scores + + def score(self, cand: Candidate, goal: Goal) -> float: + return self.scores.get(cand.model_tier, 0.5) + + +# ── trust as an objective ── + +def test_trust_weight_zero_is_pure_cost_ranking(): + trust = _Trust({"premium": 0.1, "local": 0.9}) + opt = BanditOptimizer(None, epsilon=0.0, trust=trust, trust_weight=0.0) + hi, lo = _cand("premium"), _cand("local") + plan = opt.select([(hi, _score(0.9)), (lo, _score(0.5))], Goal(text="q"), context="synthesis") + assert plan.chosen is hi # trust only tie-breaks; cost total still wins + + +def test_trust_weight_folds_trust_into_the_objective(): + """A less costly-attractive but more trusted plan wins once trust is weighted into the objective.""" + trust = _Trust({"premium": 0.1, "local": 0.9}) + opt = BanditOptimizer(None, epsilon=0.0, trust=trust, trust_weight=0.7) + hi, lo = _cand("premium"), _cand("local") + # premium: 0.3*0.9 + 0.7*0.1 = 0.34 ; local: 0.3*0.5 + 0.7*0.9 = 0.78 → local wins + plan = opt.select([(hi, _score(0.9)), (lo, _score(0.5))], Goal(text="q"), context="synthesis") + assert plan.chosen is lo + + +def test_trust_on_bandit_no_longer_errors(): + """Regression: _arm_value now receives the goal, so a TrustProvider on the bandit works.""" + opt = BanditOptimizer(None, epsilon=0.0, trust=_Trust({"local": 0.9})) + plan = opt.select([(_cand("local"), _score(0.5))], Goal(text="q"), context="c") + assert plan.chosen.model_tier == "local" + + +# ── abstention gate ── + +def test_gate_serves_when_confident(): + v = AbstentionGate(min_confidence=0.6).evaluate(_score(0.5, acc=0.8)) + assert v.action == "serve" and not v.abstained + + +def test_gate_abstains_when_below_bar_and_no_escalation(): + v = AbstentionGate(min_confidence=0.6).evaluate(_score(0.5, acc=0.3)) + assert v.action == "abstain" and v.abstained and "abstain" in v.reason + + +def test_gate_escalates_when_a_stronger_option_exists(): + gate = AbstentionGate(min_confidence=0.6, can_escalate=lambda s, g: True) + assert gate.evaluate(_score(0.5, acc=0.3)).action == "escalate" + + +def test_gate_uses_calibration_map(): + """A raw score above the bar can be calibrated below it → abstain (calibrated P(correct)).""" + gate = AbstentionGate(min_confidence=0.6, calibrate=lambda c: c * 0.5) + v = gate.evaluate(_score(0.5, acc=0.9)) # raw 0.9 → calibrated 0.45 < 0.6 + assert v.action == "abstain" and abs(v.confidence - 0.45) < 1e-9 + + +def test_optimizer_records_abstention_verdict_in_plan_extra(): + gate = AbstentionGate(min_confidence=0.7) + opt = BanditOptimizer(None, epsilon=0.0, abstain_gate=gate) + plan = opt.select([(_cand("local"), _score(0.5, acc=0.3))], Goal(text="q"), context="c") + ab = plan.extra["abstention"] + assert ab["action"] == "abstain" and ab["confidence"] == 0.3 diff --git a/tests/test_vuln_db.py b/tests/test_vuln_db.py new file mode 100644 index 0000000..b22b149 --- /dev/null +++ b/tests/test_vuln_db.py @@ -0,0 +1,48 @@ +"""Hermetic tests for the vuln-DB read client — no Doris, no network.""" +from __future__ import annotations + +from context_runtime.integrations.vuln_db import Principal, VulnDB + +_ROWS = [ + {"cve_id": "CVE-A", "package": "requests", "source": "osv", "owner": "osv", "cvss": 9.1, "refs": "u1", "summary": "a"}, + {"cve_id": "CVE-B", "package": "jinja2", "source": "nvd", "owner": "nvd", "cvss": 7.0, "refs": "u2", "summary": "b"}, +] + + +def _db(monkeypatch, rows=None): + db = VulnDB(host="x") + seen = {} + monkeypatch.setattr(db, "_query", lambda sql, params=(): (seen.update(sql=sql, params=params), rows if rows is not None else _ROWS)[1]) + return db, seen + + +def test_lookup_builds_filtered_sql(monkeypatch): + db, seen = _db(monkeypatch) + db.lookup(package="requests", min_cvss=7.0, limit=5) + assert "AND package=%s" in seen["sql"] and "AND cvss>=%s" in seen["sql"] + assert seen["params"] == ("requests", 7.0, 5) + + +def test_policy_admin_sees_all_columns(monkeypatch): + db, _ = _db(monkeypatch) + out = db.lookup(principal=Principal(roles=frozenset({"admin"}))) + assert len(out) == 2 and all("refs" in r for r in out) + + +def test_policy_row_scope_and_column_mask(monkeypatch): + db, _ = _db(monkeypatch) + out = db.lookup(principal=Principal(roles=frozenset({"analyst"}), owns_rows_of=frozenset({"osv"}))) + assert [r["cve_id"] for r in out] == ["CVE-A"] # only the owned source + assert all("refs" not in r for r in out) # sensitive column masked + + +def test_policy_no_scope_sees_nothing(monkeypatch): + db, _ = _db(monkeypatch) + assert db.lookup(principal=Principal(roles=frozenset({"guest"}))) == [] + + +def test_enrich_groups_by_cve_and_filters(monkeypatch): + db, _ = _db(monkeypatch) + findings = [{"id": "CVE-A"}, {"id": "CVE-B"}, {"id": "CVE-Z"}] + got = db.enrich(findings, principal=Principal(roles=frozenset({"security"}))) + assert set(got) == {"CVE-A", "CVE-B"} and got["CVE-A"][0]["package"] == "requests" diff --git a/tests/test_vuln_feeds.py b/tests/test_vuln_feeds.py new file mode 100644 index 0000000..7afa0dc --- /dev/null +++ b/tests/test_vuln_feeds.py @@ -0,0 +1,157 @@ +"""Hermetic tests for vuln_feeds — canned NVD/OSV payloads, no network I/O. + +`_get_json` / `_post_json` are monkeypatched to return parsed canned payloads so the +fetchers exercise their normalization paths without touching the network. +""" +from __future__ import annotations + +import json + +from context_runtime.integrations import vuln_feeds +from context_runtime.integrations.vuln_feeds import ( + VulnRecord, + dedupe, + fetch_nvd, + fetch_osv, + fetch_snyk, + ingest, +) + +# --------------------------------------------------------------------------- canned data + +_NVD_JSON = """ +{"vulnerabilities":[ + {"cve":{ + "id":"CVE-2024-5555", + "published":"2024-05-01T00:00:00.000", + "lastModified":"2024-05-02T00:00:00.000", + "descriptions":[ + {"lang":"es","value":"desbordamiento de bufer"}, + {"lang":"en","value":"A buffer overflow in acme-lib allows RCE."} + ], + "metrics":{"cvssMetricV31":[ + {"cvssData":{"baseScore":9.8,"baseSeverity":"critical"}} + ]}, + "references":[ + {"url":"https://example.com/advisory/5555"}, + {"url":"https://nvd.nist.gov/vuln/detail/CVE-2024-5555"} + ] + }} +]} +""" + +_OSV_JSON = """ +{"vulns":[ + { + "id":"GHSA-aaaa-bbbb-cccc", + "aliases":["CVE-2024-5555"], + "summary":"acme-lib buffer overflow", + "published":"2024-05-01T00:00:00Z", + "modified":"2024-05-03T00:00:00Z", + "affected":[ + {"package":{"name":"acme-lib","ecosystem":"PyPI"}, + "ranges":[{"type":"ECOSYSTEM","events":[ + {"introduced":"0"}, + {"fixed":"2.31.0"} + ]}]} + ], + "references":[{"url":"https://example.com/osv/5555"}] + } +]} +""" + +_NVD_PARSED = json.loads(_NVD_JSON) +_OSV_PARSED = json.loads(_OSV_JSON) + + +class _FakeStore: + """In-memory VulnStore that records what it was asked to upsert.""" + + def __init__(self): + self.records: list[VulnRecord] = [] + + def upsert_vulns(self, records: list[VulnRecord]) -> int: + self.records = list(records) + return len(self.records) + + def query_vulns(self, **filters) -> list[dict]: + return [] + + +# --------------------------------------------------------------------------- NVD + + +def test_nvd_normalization(monkeypatch): + monkeypatch.setattr(vuln_feeds, "_get_json", lambda url, headers=None: _NVD_PARSED) + recs = fetch_nvd("acme-lib") + assert len(recs) == 1 + r = recs[0] + assert r.cve_id == "CVE-2024-5555" + assert r.source == "nvd" + # severity uppercased from the canned lowercase "critical" + assert r.severity == "CRITICAL" + assert isinstance(r.cvss, float) and r.cvss == 9.8 + # english description picked, not the Spanish one + assert r.summary == "A buffer overflow in acme-lib allows RCE." + assert "https://example.com/advisory/5555" in r.references + + +# --------------------------------------------------------------------------- OSV + + +def test_osv_normalization(monkeypatch): + monkeypatch.setattr( + vuln_feeds, "_post_json", lambda url, body, headers=None: _OSV_PARSED + ) + recs = fetch_osv("acme-lib", "PyPI") + assert len(recs) == 1 + r = recs[0] + assert r.package == "acme-lib" + assert r.ecosystem == "PyPI" + # fixed_version pulled from the ranges 'fixed' event + assert r.fixed_version == "2.31.0" + # cve_id resolved to the CVE alias, not the GHSA id + assert r.cve_id == "CVE-2024-5555" + assert r.source == "osv" + + +# --------------------------------------------------------------------------- dedupe + + +def test_dedupe_collapses_shared_cve(monkeypatch): + monkeypatch.setattr(vuln_feeds, "_get_json", lambda url, headers=None: _NVD_PARSED) + monkeypatch.setattr( + vuln_feeds, "_post_json", lambda url, body, headers=None: _OSV_PARSED + ) + combined = fetch_nvd("acme-lib") + fetch_osv("acme-lib", "PyPI") + assert len(combined) == 2 + collapsed = dedupe(combined) + assert len(collapsed) == 1 + r = collapsed[0] + assert r.cve_id == "CVE-2024-5555" + # richer record wins: OSV carried the fixed_version + assert r.fixed_version == "2.31.0" + + +# --------------------------------------------------------------------------- Snyk + + +def test_snyk_returns_empty_without_token(monkeypatch): + monkeypatch.delenv("SNYK_TOKEN", raising=False) + assert fetch_snyk() == [] + + +# --------------------------------------------------------------------------- ingest + + +def test_ingest_dedupes_and_upserts(monkeypatch): + monkeypatch.setattr(vuln_feeds, "_get_json", lambda url, headers=None: _NVD_PARSED) + monkeypatch.setattr( + vuln_feeds, "_post_json", lambda url, body, headers=None: _OSV_PARSED + ) + combined = fetch_nvd("acme-lib") + fetch_osv("acme-lib", "PyPI") + store = _FakeStore() + count = ingest(combined, store) + assert count == 1 + assert len(store.records) == 1 + assert store.records[0].cve_id == "CVE-2024-5555" diff --git a/tests/test_web_search_mcp.py b/tests/test_web_search_mcp.py new file mode 100644 index 0000000..6d9a3e7 --- /dev/null +++ b/tests/test_web_search_mcp.py @@ -0,0 +1,42 @@ +"""The bundled web_search MCP server + AgentConsole.mount_mcp — the end-to-end proof that an +app can mount a real MCP tool through the agent-harness registry. + +The mount/registration test is hermetic (tools/list is static, no network). The actual search +test hits keyless public APIs and skips if the network is unavailable. +""" +from __future__ import annotations + +import sys + +import pytest + +from context_runtime.integrations.agent_console import AgentConsole, tool +from context_runtime.tools.mcp import MCPClient + + +def _client() -> MCPClient: + return MCPClient.stdio([sys.executable, "-m", "context_runtime.tools.mcp_servers.web_search"]) + + +def test_web_search_mounts_as_readonly_tool(): + console = AgentConsole("X", "A watch monitors a URL.", tools=[tool("noop", "noop", lambda a: {"text": "ok"})]) + client = _client() + try: + names = console.mount_mcp(client) + assert "web_search" in names + assert "web_search" in console._tools # visible to classify + dispatch + spec = console.registry.get("web_search").spec() + assert spec.side_effecting is False # readOnlyHint → ungated → runs from chat + assert "query" in spec.parameters.get("properties", {}) + finally: + client.close() + + +def test_web_search_returns_results(): + from context_runtime.tools.mcp_servers.web_search import web_search + + out = web_search("python programming language", 4) + if "no web results" in out.lower(): + pytest.skip("network unavailable / providers returned nothing") + assert "results" in out.lower() + assert "http" in out.lower()