|
|
@@ -0,0 +1,190 @@
|
|
|
+"""Optional semantic assistance for Context Broker ranking.
|
|
|
+
|
|
|
+Correctness never depends on this component: it may only reorder handles that
|
|
|
+the deterministic Broker has already authorized.
|
|
|
+"""
|
|
|
+
|
|
|
+from __future__ import annotations
|
|
|
+
|
|
|
+import asyncio
|
|
|
+import json
|
|
|
+import os
|
|
|
+from collections.abc import Mapping, Sequence
|
|
|
+from dataclasses import dataclass
|
|
|
+from hashlib import sha256
|
|
|
+from pathlib import Path
|
|
|
+from tempfile import NamedTemporaryFile
|
|
|
+from typing import Any
|
|
|
+
|
|
|
+from script_build_host.agents.prompts import CONTEXT_BROKER_PROMPT
|
|
|
+from script_build_host.domain.context_broker import (
|
|
|
+ ContextCard,
|
|
|
+ canonical_digest,
|
|
|
+ estimate_context_tokens,
|
|
|
+)
|
|
|
+
|
|
|
+_SYSTEM_PROMPT = CONTEXT_BROKER_PROMPT
|
|
|
+
|
|
|
+
|
|
|
+@dataclass(frozen=True, slots=True)
|
|
|
+class SemanticSelection:
|
|
|
+ handles: tuple[str, ...]
|
|
|
+ cache_hit: bool = False
|
|
|
+ calls: int = 0
|
|
|
+ fallback_reason: str | None = None
|
|
|
+
|
|
|
+
|
|
|
+class AdaptiveSemanticSelector:
|
|
|
+ """Bounded, cached, fail-open reranking of an authorized candidate set."""
|
|
|
+
|
|
|
+ def __init__(
|
|
|
+ self,
|
|
|
+ *,
|
|
|
+ llm_call: Any | None,
|
|
|
+ cache_root: Path | None,
|
|
|
+ mission_call_limit: int = 24,
|
|
|
+ timeout_seconds: float = 15.0,
|
|
|
+ ) -> None:
|
|
|
+ self._llm_call = llm_call
|
|
|
+ self._cache_root = cache_root
|
|
|
+ self._mission_call_limit = mission_call_limit
|
|
|
+ self._timeout_seconds = timeout_seconds
|
|
|
+ self._calls_by_root: dict[str, int] = {}
|
|
|
+ self._limit_lock = asyncio.Lock()
|
|
|
+
|
|
|
+ async def select(
|
|
|
+ self,
|
|
|
+ *,
|
|
|
+ root_trace_id: str,
|
|
|
+ query: str,
|
|
|
+ cards: Sequence[ContextCard],
|
|
|
+ model_config: Mapping[str, Any],
|
|
|
+ ) -> SemanticSelection:
|
|
|
+ eligible = tuple(card.handle for card in cards[:20])
|
|
|
+ if not query.strip() or len(eligible) < 2:
|
|
|
+ return SemanticSelection(eligible, fallback_reason="semantic_not_needed")
|
|
|
+ if self._llm_call is None:
|
|
|
+ return SemanticSelection(eligible, fallback_reason="semantic_provider_unavailable")
|
|
|
+ request: dict[str, Any] = {"query": query[:2_000], "candidates": []}
|
|
|
+ for card in cards[:20]:
|
|
|
+ candidate = {
|
|
|
+ "handle": card.handle,
|
|
|
+ "source_type": card.source_type,
|
|
|
+ "summary": card.summary[:300],
|
|
|
+ "excerpt": card.excerpt[:300],
|
|
|
+ }
|
|
|
+ trial = {**request, "candidates": [*request["candidates"], candidate]}
|
|
|
+ if request["candidates"] and estimate_context_tokens(trial) > 7_000:
|
|
|
+ break
|
|
|
+ request = trial
|
|
|
+ allowed = tuple(item["handle"] for item in request["candidates"])
|
|
|
+ key = canonical_digest(
|
|
|
+ {
|
|
|
+ "request": request,
|
|
|
+ "prompt": sha256(_SYSTEM_PROMPT.encode()).hexdigest(),
|
|
|
+ "model": dict(model_config),
|
|
|
+ }
|
|
|
+ ).removeprefix("sha256:")
|
|
|
+ async with self._limit_lock:
|
|
|
+ cached = await self._read_cache(key)
|
|
|
+ if cached is not None and set(cached).issubset(allowed):
|
|
|
+ return SemanticSelection(tuple(cached), cache_hit=True)
|
|
|
+ if not await self._claim_mission_call(root_trace_id):
|
|
|
+ return SemanticSelection(allowed, fallback_reason="semantic_call_limit")
|
|
|
+ try:
|
|
|
+ result = await asyncio.wait_for(
|
|
|
+ self._llm_call(
|
|
|
+ messages=[
|
|
|
+ {"role": "system", "content": _SYSTEM_PROMPT},
|
|
|
+ {"role": "user", "content": json.dumps(request, ensure_ascii=False)},
|
|
|
+ ],
|
|
|
+ model=str(model_config.get("model") or ""),
|
|
|
+ tools=[],
|
|
|
+ temperature=0.0,
|
|
|
+ max_tokens=1_000,
|
|
|
+ ),
|
|
|
+ timeout=self._timeout_seconds,
|
|
|
+ )
|
|
|
+ parsed = json.loads(str(result.get("content") or ""))
|
|
|
+ handles = parsed.get("handles") if isinstance(parsed, Mapping) else None
|
|
|
+ if (
|
|
|
+ not isinstance(handles, list)
|
|
|
+ or not handles
|
|
|
+ or any(not isinstance(item, str) or item not in allowed for item in handles)
|
|
|
+ ):
|
|
|
+ raise ValueError("semantic selector returned unauthorized handles")
|
|
|
+ selected = tuple(dict.fromkeys(handles))
|
|
|
+ selected += tuple(item for item in allowed if item not in selected)
|
|
|
+ await self._write_cache(key, selected)
|
|
|
+ return SemanticSelection(selected, calls=1)
|
|
|
+ except Exception as exc:
|
|
|
+ return SemanticSelection(
|
|
|
+ allowed,
|
|
|
+ calls=1,
|
|
|
+ fallback_reason=f"{type(exc).__name__}:semantic_fallback"[:120],
|
|
|
+ )
|
|
|
+
|
|
|
+ async def _claim_mission_call(self, root_trace_id: str) -> bool:
|
|
|
+ if self._cache_root is None:
|
|
|
+ used = self._calls_by_root.get(root_trace_id, 0)
|
|
|
+ if used >= self._mission_call_limit:
|
|
|
+ return False
|
|
|
+ self._calls_by_root[root_trace_id] = used + 1
|
|
|
+ return True
|
|
|
+ return await asyncio.to_thread(self._claim_persistent_call_sync, root_trace_id)
|
|
|
+
|
|
|
+ def _claim_persistent_call_sync(self, root_trace_id: str) -> bool:
|
|
|
+ assert self._cache_root is not None
|
|
|
+ mission_key = sha256(root_trace_id.encode()).hexdigest()
|
|
|
+ quota_root = self._cache_root / "mission-call-quota" / mission_key
|
|
|
+ quota_root.mkdir(parents=True, exist_ok=True)
|
|
|
+ for index in range(self._mission_call_limit):
|
|
|
+ target = quota_root / f"{index:04d}.claim"
|
|
|
+ try:
|
|
|
+ descriptor = os.open(target, os.O_CREAT | os.O_EXCL | os.O_WRONLY, 0o600)
|
|
|
+ except FileExistsError:
|
|
|
+ continue
|
|
|
+ with os.fdopen(descriptor, "w", encoding="utf-8") as stream:
|
|
|
+ stream.write("claimed\n")
|
|
|
+ stream.flush()
|
|
|
+ os.fsync(stream.fileno())
|
|
|
+ return True
|
|
|
+ return False
|
|
|
+
|
|
|
+ async def _read_cache(self, key: str) -> tuple[str, ...] | None:
|
|
|
+ if self._cache_root is None:
|
|
|
+ return None
|
|
|
+ target = self._cache_root / f"{key}.json"
|
|
|
+ if not target.is_file():
|
|
|
+ return None
|
|
|
+ try:
|
|
|
+ raw = await asyncio.to_thread(target.read_text, encoding="utf-8")
|
|
|
+ value = json.loads(raw)
|
|
|
+ handles = value.get("handles") if isinstance(value, Mapping) else None
|
|
|
+ if isinstance(handles, list) and all(isinstance(item, str) for item in handles):
|
|
|
+ return tuple(handles)
|
|
|
+ except (OSError, json.JSONDecodeError):
|
|
|
+ return None
|
|
|
+ return None
|
|
|
+
|
|
|
+ async def _write_cache(self, key: str, handles: Sequence[str]) -> None:
|
|
|
+ if self._cache_root is None:
|
|
|
+ return
|
|
|
+ await asyncio.to_thread(self._write_cache_sync, key, handles)
|
|
|
+
|
|
|
+ def _write_cache_sync(self, key: str, handles: Sequence[str]) -> None:
|
|
|
+ assert self._cache_root is not None
|
|
|
+ self._cache_root.mkdir(parents=True, exist_ok=True)
|
|
|
+ target = self._cache_root / f"{key}.json"
|
|
|
+ with NamedTemporaryFile("w", dir=self._cache_root, encoding="utf-8", delete=False) as tmp:
|
|
|
+ json.dump({"handles": list(handles)}, tmp, ensure_ascii=False, sort_keys=True)
|
|
|
+ tmp.flush()
|
|
|
+ os.fsync(tmp.fileno())
|
|
|
+ temporary = Path(tmp.name)
|
|
|
+ try:
|
|
|
+ os.replace(temporary, target)
|
|
|
+ finally:
|
|
|
+ temporary.unlink(missing_ok=True)
|
|
|
+
|
|
|
+
|
|
|
+__all__ = ["AdaptiveSemanticSelector", "SemanticSelection"]
|
SamLee