ely-eye/backend/src/ely_eye/compiler.py

from __future__ import annotations

import hashlib
from collections.abc import Iterable

from . import trust
from .cache_fabric import SEMANTIC_LAYER, CacheFabric
from .config import Settings, get_settings
from .db import Database
from .retrieval import RetrievalService
from .schemas import (
    CompiledContext,
    ContextPlan,
    Modality,
    RetrievalHit,
    RuntimeProfile,
    SearchFilters,
)

PROFILE_TOKEN_BUDGETS = {
    RuntimeProfile.live_demo: 128_000,
    RuntimeProfile.extreme_context: 1_010_000,
    RuntimeProfile.library_100m: 262_144,
    RuntimeProfile.research_theater: 128_000,
}


class ContextCompiler:
    def __init__(self, settings: Settings | None = None, database: Database | None = None) -> None:
        self.settings = settings or get_settings()
        self.db = database or Database(self.settings)
        self.retrieval = RetrievalService(self.settings, self.db)
        self.cache = CacheFabric(self.settings, self.db)

    def compile(self, question: str, profile: RuntimeProfile, as_holder: str | None = None) -> CompiledContext:
        """Build the context one question gets, seen through one holder's eyes.

        `as_holder` is passed to retrieval and folded into the cache key: a context compiled for
        a holder carries that holder's own memories, and serving it to somebody else would leak.
        """
        cache_key = self.cache.semantic_key(
            question,
            profile.value,
            self.db.revision(),
            "the whole library" if as_holder is None else f"as holder {as_holder!r}",
        )
        cached = self.cache.read_semantic_cache(cache_key)
        if cached is not None:
            compiled = CompiledContext.model_validate(cached)
            if compiled.hits:
                self.cache.record_hit(SEMANTIC_LAYER, cache_key)
                return compiled
            # A cached context with no evidence is not an answer; counting it would make an empty entry look like a warm cache.
            self.cache.record_miss(
                SEMANTIC_LAYER, cache_key, metadata={"reason": "cached context has no hits"}
            )
        else:
            self.cache.record_miss(SEMANTIC_LAYER, cache_key)
        plan = self.plan(question, profile)
        hits = self.retrieval.search(
            plan.retrieval_query, plan.evidence_budget_atoms, filters=SearchFilters(as_holder=as_holder)
        ).hits
        citable, held = self._split_citable(hits)
        packed_text, token_total = self._pack_hits(citable, plan.token_budget, plan.retrieval_query)
        trace_id = hashlib.sha256(f"{cache_key}:{token_total}".encode()).hexdigest()[:16]
        compiled = CompiledContext(
            plan=plan,
            hits=hits,
            packed_text=packed_text,
            token_equivalent=token_total,
            cache_trace_id=trace_id,
            quarantined_excluded=[hit.atom.atom_id for hit in held],
        )
        self.cache.write_semantic_cache(cache_key, compiled.model_dump(mode="json"))
        return compiled

    def plan(self, question: str, profile: RuntimeProfile) -> ContextPlan:
        modalities = infer_modalities(question)
        token_budget = PROFILE_TOKEN_BUDGETS[profile]
        evidence_budget = 48 if profile in {RuntimeProfile.library_100m, RuntimeProfile.extreme_context} else 24
        return ContextPlan(
            question=question,
            profile=profile,
            required_modalities=modalities,
            evidence_budget_atoms=evidence_budget,
            token_budget=token_budget,
            retrieval_query=question,
            compression_strategy=self._compression_strategy(profile, modalities),
            verifier_contract=[
                "Every factual claim must cite atom ids.",
                "Citations must include page, image region, frame, or code path when available.",
                "Contradictions must be reported as contradiction_notes.",
            ],
        )

    def _split_citable(self, hits: list[RetrievalHit]) -> tuple[list[RetrievalHit], list[RetrievalHit]]:
        """A quarantined atom reaches the prompt only when a trusted source says the same thing."""

        citable: list[RetrievalHit] = []
        held: list[RetrievalHit] = []
        for hit in hits:
            (citable if trust.is_citable(hit.atom, self.db) else held).append(hit)
        return citable, held

    def _compression_strategy(self, profile: RuntimeProfile, modalities: Iterable[Modality]) -> str:
        has_visual = any(modality in {Modality.image, Modality.pdf_page, Modality.video_frame, Modality.ui_screenshot} for modality in modalities)
        if profile == RuntimeProfile.extreme_context:
            return "YaRN + V-NSA sparse blocks + DVKV visual KV compression + CPU/NVMe parked context"
        if profile == RuntimeProfile.library_100m:
            return "Context Cartridge retrieval + ColQwen-compatible dense vectors + BM25 + evidence graph pack"
        if has_visual:
            return "Visual patch cache + DVKV evidence-biased region retention"
        return "Prefix cache + sparse evidence packing"

    def _pack_hits(self, hits: list[RetrievalHit], token_budget: int, query: str) -> tuple[str, int]:
        sections: list[str] = []
        token_total = 0
        for hit in hits:
            atom = hit.atom
            evidence_text = focused_evidence_text(atom.text, query)
            atom_tokens = max(1, len(evidence_text) // 4)
            if token_total + atom_tokens > token_budget:
                break
            citation = {
                "atom_id": atom.atom_id,
                "modality": atom.modality.value,
                "source": atom.source,
                "image_ref": atom.image_ref,
                "layout": atom.layout.model_dump() if atom.layout else None,
                "score": hit.final_score,
            }
            sections.append(f"[EVIDENCE]\n{citation}\n{evidence_text}\n[/EVIDENCE]")
            token_total += atom_tokens
        return "\n\n".join(sections), token_total


def infer_modalities(question: str) -> list[Modality]:
    lowered = question.lower()
    modalities: list[Modality] = []
    if any(word in lowered for word in ("pdf", "page", "paper", "论文", "页面", "页码")):
        modalities.append(Modality.pdf_page)
    if any(word in lowered for word in ("image", "screenshot", "figma", "ui", "图", "截图", "设计")):
        modalities.extend([Modality.image, Modality.ui_screenshot])
    if any(word in lowered for word in ("video", "frame", "clip", "视频", "帧")):
        modalities.append(Modality.video_frame)
    if any(word in lowered for word in ("code", "function", "class", "代码", "函数")):
        modalities.append(Modality.code)
    if not modalities:
        modalities.append(Modality.mixed)
    return list(dict.fromkeys(modalities))


def focused_evidence_text(text: str, query: str, window: int = 1400, max_windows: int = 4) -> str:
    if len(text) <= window * max_windows:
        return text
    lowered = text.lower()
    query_terms = [term for term in query.lower().replace("/", " ").replace("-", " ").split() if len(term) > 2]
    anchors: list[int] = []
    for term in query_terms:
        index = lowered.find(term)
        if index >= 0:
            anchors.append(index)
    anchors.extend([0, len(text) // 3, (len(text) * 2) // 3])
    windows: list[tuple[int, int]] = []
    for anchor in anchors:
        start = max(0, anchor - window // 2)
        end = min(len(text), start + window)
        start = max(0, end - window)
        if all(abs(start - existing_start) > window // 2 for existing_start, _ in windows):
            windows.append((start, end))
        if len(windows) >= max_windows:
            break
    return "\n\n[...]\n\n".join(text[start:end].strip() for start, end in windows if text[start:end].strip())