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

"""Retrieval: SQLite FTS5 does the matching, RRF merges the signals that exist.

The first version pulled every atom into Python and rebuilt a BM25 index per query, so search
cost grew with the corpus instead of with the answer, and a Chinese query matched nothing
because the tokenizer had no CJK rule. Both are fixed below the score line: the index and the
query share `text_index`, and the ranking is a SQL query.
"""

from __future__ import annotations

import base64
from collections import defaultdict
from datetime import datetime
from pathlib import Path
from typing import Any

from . import cartridge_store, lifecycle_store
from .config import Settings, get_settings
from .db import Database
from .embeddings import EmbeddingService
from .schemas import EvidenceAtom, RetrievalHit, SearchFilters, SearchResult, utc_now
from .store_search import FtsPage, search_atoms
from .vector_index import VectorIndex
from .visual_index import VisualIndex

RRF_K = 60
CANDIDATE_WINDOW = 10
# Words decide most searches; a picture vector is only worth its encoding cost when words fell short.
THIN_SPARSE_PAGE = 2


class RetrievalService:
    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.embeddings = EmbeddingService(self.settings, self.db)
        self.vectors = VectorIndex(self.settings, self.db, self.embeddings)
        self.visual = VisualIndex(self.settings, self.db)

    def search(
        self,
        query: str,
        top_k: int | None = None,
        *,
        filters: SearchFilters | None = None,
        cursor: str | None = None,
        image: Path | None = None,
    ) -> SearchResult:
        """Answer one query against the memories that are open right now.

        The candidate window is wider than the page so fusion has something to fuse; the cursor
        advances one page at a time and re-reads the tail of the window it came from, which is
        why no result is ever skipped between pages.
        """
        limit = top_k or self.settings.retrieval_top_k
        if limit <= 0:
            raise ValueError(f"search needs a positive limit, got {limit}")
        if not query.strip() and image is None:
            raise ValueError("search needs words or a picture to ask with, and got neither")
        offset = decode_cursor(cursor)
        criteria = (filters or SearchFilters()).model_dump(exclude_none=True)
        scope = criteria.get("cartridge_ids")
        if scope is None:
            scope = cartridge_store.mounted_ids(self.db) or None
        criteria["cartridge_ids"] = scope
        # A picture-only question has no words for BM25, so the word ranking is skipped rather than matching an empty expression.
        page = (
            search_atoms(self.db, query, limit=limit * CANDIDATE_WINDOW, offset=offset, **criteria)
            if query.strip()
            else FtsPage(rows=[], has_more=False)
        )
        rankings: list[list[str]] = [[atom_id for atom_id, _ in page.rows]]
        scores: dict[str, dict[str, float]] = defaultdict(dict)
        for atom_id, bm25 in page.rows:
            scores[atom_id]["sparse_score"] = bm25
        if self.embeddings.enabled():
            dense = self._dense_search(query, limit * CANDIDATE_WINDOW, criteria)
            rankings.append(list(dense))
            for atom_id, cosine in dense.items():
                scores[atom_id]["dense_score"] = cosine
        visual = self._visual_search(
            query,
            image,
            limit * CANDIDATE_WINDOW,
            scope,
            thin=len(rankings[0]) < limit * THIN_SPARSE_PAGE,
        )
        if visual:
            rankings.append(list(visual))
            for atom_id, cosine in visual.items():
                scores[atom_id]["visual_score"] = cosine
        hits = self._settle_conflicts(self._rank(rankings, scores, limit * CANDIDATE_WINDOW))[:limit]
        lifecycle_store.touch_recalls(self.db, [hit.atom.atom_id for hit in hits], utc_now())
        more = page.has_more or len(rankings[0]) > limit
        return SearchResult(
            query=query,
            status="hits" if hits else "no_match",
            hits=hits,
            next_cursor=encode_cursor(offset + limit) if hits and more else None,
        )

    def _dense_search(self, query: str, limit: int, criteria: dict[str, Any]) -> dict[str, float]:
        """The vectors sit beside the atoms, so a memory written a moment ago is findable.

        Every filter is applied inside the engine, in the same statement as the nearest-neighbour
        walk, so a scoped search cannot return a memory the caller was never allowed to see. The
        index catches up inside `search`, so nothing here has to rebuild it first.
        """
        query_vector = self.embeddings.embed_query(query)
        return self.vectors.search(
            query_vector,
            limit,
            holder=criteria.get("holder"),
            scenario=criteria.get("scenario"),
            source_id=criteria.get("source_id"),
            modality=criteria.get("modality"),
            since=_text(criteria.get("since")),
            until=_text(criteria.get("until")),
            cartridge_ids=criteria.get("cartridge_ids"),
            as_holder=criteria.get("as_holder"),
        )

    def _visual_search(
        self,
        query: str,
        image: Path | None,
        limit: int,
        scope: list[str] | None,
        *,
        thin: bool,
    ) -> dict[str, float]:
        """Ask the picture index when words are not enough.

        A screenshot with no text on it has nothing for FTS to match, so the sparse page comes
        back short; that is the signal to spend the encoding. An image query always asks for it.
        """
        if image is None and not thin:
            return {}
        if not self.visual.available() or self.visual.load() is None:
            return {}
        if image is not None:
            try:
                ranked = self.visual.search_image(image, limit)
            except (OSError, ValueError) as error:
                raise ValueError(f"the query image {image} cannot be read: {error}") from error
        else:
            ranked = self.visual.search_text(query, limit)
        if scope:
            allowed = cartridge_store.atoms_of_cartridges(self.db, scope)
            ranked = [(atom_id, score) for atom_id, score in ranked if atom_id in allowed]
        return dict(ranked)

    def _rank(self, rankings: list[list[str]], scores: dict[str, dict[str, float]], limit: int) -> list[RetrievalHit]:
        """Reciprocal rank fusion: rank position is comparable across signals, raw scores are not.

        The old blend multiplied BM25 by 0.55 and cosine by 0.35, which compared a term-frequency
        count to a bounded cosine and let the larger number decide. graph_score stays 0.0 until
        evidence edges are retrieved as their own ranking.
        """
        fused: dict[str, float] = {}
        for ranking in rankings:
            for position, atom_id in enumerate(ranking):
                fused[atom_id] = fused.get(atom_id, 0.0) + 1.0 / (RRF_K + position + 1)
        order = sorted(fused, key=lambda atom_id: (-fused[atom_id], atom_id))[:limit]
        atoms = {atom.atom_id: atom for atom in self.db.get_atoms(order)}
        hits: list[RetrievalHit] = []
        for atom_id in order:
            atom = atoms.get(atom_id)
            if atom is None:
                raise ValueError(f"search returned {atom_id} but the atom is not in the store")
            signal = scores.get(atom_id, {})
            hits.append(
                RetrievalHit(
                    atom=atom,
                    sparse_score=signal.get("sparse_score", 0.0),
                    dense_score=signal.get("dense_score", 0.0),
                    visual_score=signal.get("visual_score", 0.0),
                    final_score=fused[atom_id],
                )
            )
        return hits

    def _settle_conflicts(self, hits: list[RetrievalHit]) -> list[RetrievalHit]:
        """Drop the older side of a disagreement nobody settled, so one page cannot contradict itself.

        Both memories stay in the store and this only filters the page; the conflict row names
        which side was dropped and why. Settling is the holder's call, and until it happens the
        newer statement is what a reader should be shown.
        """
        losers = lifecycle_store.losing_sides(self.db, [hit.atom.atom_id for hit in hits])
        if not losers:
            return hits
        return [hit for hit in hits if hit.atom.atom_id not in losers]

    def recall(self, query: str, top_k: int | None = None, *, filters: SearchFilters | None = None) -> list[EvidenceAtom]:
        """The atoms a search would put in front of a model, no scoring attached."""
        return [hit.atom for hit in self.search(query, top_k, filters=filters).hits]


def encode_cursor(offset: int) -> str:
    return base64.urlsafe_b64encode(str(offset).encode()).decode()


def decode_cursor(cursor: str | None) -> int:
    if cursor is None or cursor == "":
        return 0
    try:
        offset = int(base64.urlsafe_b64decode(cursor.encode()).decode())
    except (ValueError, UnicodeDecodeError) as error:
        raise ValueError(f"search cursor is unreadable: {cursor!r}") from error
    if offset < 0:
        raise ValueError(f"search cursor is unreadable: {cursor!r}")
    return offset


def _text(moment: datetime | None) -> str | None:
    """A time bound reaches SQLite as the same ISO text the atoms table stores."""
    return None if moment is None else moment.isoformat()