lean-explore 0.3.0__py3-none-any.whl → 1.0.1__py3-none-any.whl

lean_explore/search/engine.py

@@ -0,0 +1,651 @@
+"""Core search engine for Lean declarations.
+
+This module provides the SearchEngine class that implements hybrid search using
+BM25 lexical matching and FAISS semantic search, combined via Reciprocal Rank
+Fusion (RRF) and cross-encoder reranking.
+
+Note: On macOS, torch and FAISS have OpenMP library conflicts. To avoid segfaults:
+- FAISS is imported lazily (not at module level)
+- When semantic search is needed, torch/embeddings are loaded FIRST, then FAISS
+"""
+
+import json
+import logging
+from pathlib import Path
+from typing import TYPE_CHECKING
+
+import bm25s
+import numpy as np
+from sqlalchemy import select
+from sqlalchemy.ext.asyncio import AsyncEngine, AsyncSession, create_async_engine
+
+from lean_explore.config import Config
+from lean_explore.models import Declaration, SearchResult
+from lean_explore.search.scoring import (
+    fuzzy_name_score,
+    normalize_dependency_counts,
+    normalize_scores,
+)
+from lean_explore.search.tokenization import (
+    is_autogenerated,
+    tokenize_raw,
+    tokenize_spaced,
+    tokenize_words,
+)
+
+if TYPE_CHECKING:
+    import faiss
+
+    from lean_explore.util import EmbeddingClient, RerankerClient
+
+logger = logging.getLogger(__name__)
+
+
+class SearchEngine:
+    """Core search engine for Lean declarations.
+
+    Uses two-stage retrieval:
+    1. FAISS semantic search on informalizations
+    2. BM25 lexical search on declaration names (independent)
+    Then merges and reranks candidates.
+    """
+
+    def __init__(
+        self,
+        db_url: str | None = None,
+        embedding_client: "EmbeddingClient | None" = None,
+        embedding_model_name: str = "Qwen/Qwen3-Embedding-0.6B",
+        reranker_client: "RerankerClient | None" = None,
+        reranker_model_name: str = "Qwen/Qwen3-Reranker-0.6B",
+        faiss_index_path: Path | None = None,
+        faiss_ids_map_path: Path | None = None,
+        use_local_data: bool = True,
+    ):
+        """Initialize the search engine.
+
+        Args:
+            db_url: Database URL. Defaults to configured URL.
+            embedding_client: Client for generating embeddings. Created lazily if None.
+            embedding_model_name: Name of the embedding model to use.
+            reranker_client: Client for reranking results. Created lazily if None.
+            reranker_model_name: Name of the reranker model to use.
+            faiss_index_path: Path to FAISS index. Defaults to config path.
+            faiss_ids_map_path: Path to FAISS ID mapping. Defaults to config path.
+            use_local_data: If True, use DATA_DIRECTORY paths. If False, use
+                CACHE_DIRECTORY paths (for downloaded remote data).
+        """
+        self._embedding_client = embedding_client
+        self._embedding_model_name = embedding_model_name
+        self._reranker_client = reranker_client
+        self._reranker_model_name = reranker_model_name
+
+        if use_local_data:
+            base_path = Config.ACTIVE_DATA_PATH
+            default_db_url = Config.EXTRACTION_DATABASE_URL
+        else:
+            base_path = Config.ACTIVE_CACHE_PATH
+            default_db_url = Config.DATABASE_URL
+
+        self.db_url = db_url or default_db_url
+        self.engine: AsyncEngine = create_async_engine(self.db_url)
+
+        self._faiss_informal_path = faiss_index_path or (
+            base_path / "informalization_faiss.index"
+        )
+        self._faiss_informal_ids_path = faiss_ids_map_path or (
+            base_path / "informalization_faiss_ids_map.json"
+        )
+        self._faiss_informal_index: faiss.Index | None = None
+        self._faiss_informal_id_map: list[int] | None = None
+
+        self._bm25_spaced_path = base_path / "bm25_name_spaced"
+        self._bm25_raw_path = base_path / "bm25_name_raw"
+        self._bm25_ids_map_path = base_path / "bm25_ids_map.json"
+        self._all_declaration_ids: list[int] | None = None
+        self._bm25_name_spaced: bm25s.BM25 | None = None
+        self._bm25_name_raw: bm25s.BM25 | None = None
+
+        self._validate_paths()
+
+    def _validate_paths(self) -> None:
+        """Validate that required data files exist."""
+        required_paths = [
+            self._faiss_informal_path,
+            self._faiss_informal_ids_path,
+            self._bm25_spaced_path,
+            self._bm25_raw_path,
+            self._bm25_ids_map_path,
+        ]
+        for path in required_paths:
+            if not path.exists():
+                raise FileNotFoundError(
+                    f"Required file not found at {path}. "
+                    "Please run 'lean-explore data fetch' to download the data."
+                )
+
+    @property
+    def embedding_client(self) -> "EmbeddingClient":
+        """Lazily create the embedding client to avoid loading torch at import time."""
+        if self._embedding_client is None:
+            from lean_explore.util import EmbeddingClient
+
+            self._embedding_client = EmbeddingClient(
+                model_name=self._embedding_model_name,
+                max_length=512,
+            )
+        return self._embedding_client
+
+    @property
+    def reranker_client(self) -> "RerankerClient":
+        """Lazily create the reranker client to avoid loading torch at import time."""
+        if self._reranker_client is None:
+            from lean_explore.util import RerankerClient
+
+            self._reranker_client = RerankerClient(
+                model_name=self._reranker_model_name,
+                max_length=256,
+            )
+        return self._reranker_client
+
+    def _ensure_faiss_loaded(self) -> None:
+        """Load the FAISS index if not already loaded."""
+        if self._faiss_informal_index is not None:
+            return
+
+        import faiss
+
+        logger.info(f"Loading FAISS index from {self._faiss_informal_path}")
+        self._faiss_informal_index = faiss.read_index(str(self._faiss_informal_path))
+        with open(self._faiss_informal_ids_path) as f:
+            self._faiss_informal_id_map = json.load(f)
+
+    @property
+    def faiss_informal_index(self) -> "faiss.Index":
+        """Get the informalization FAISS index."""
+        self._ensure_faiss_loaded()
+        return self._faiss_informal_index  # type: ignore[return-value]
+
+    @property
+    def faiss_informal_id_map(self) -> list[int]:
+        """Get the informalization FAISS ID mapping."""
+        self._ensure_faiss_loaded()
+        return self._faiss_informal_id_map  # type: ignore[return-value]
+
+    def _ensure_bm25_loaded(self) -> None:
+        """Load pre-built BM25 indices from disk."""
+        if self._bm25_name_spaced is not None:
+            return
+
+        logger.info(f"Loading BM25 indices from {self._bm25_spaced_path.parent}")
+
+        self._bm25_name_spaced = bm25s.BM25.load(str(self._bm25_spaced_path))
+        self._bm25_name_raw = bm25s.BM25.load(str(self._bm25_raw_path))
+
+        with open(self._bm25_ids_map_path) as f:
+            self._all_declaration_ids = json.load(f)
+
+        logger.info(f"BM25 indices loaded ({len(self._all_declaration_ids)} decls)")
+
+    def _retrieve_bm25_candidates(self, query: str, bm25_k: int) -> dict[int, float]:
+        """Retrieve candidates using BM25 on declaration names.
+
+        Args:
+            query: Search query string.
+            bm25_k: Number of candidates to retrieve.
+
+        Returns:
+            Map of declaration ID to BM25 score.
+        """
+        self._ensure_bm25_loaded()
+
+        query_tokens_spaced = tokenize_spaced(query)
+        query_tokens_raw = tokenize_raw(query)
+
+        results_spaced, scores_spaced = self._bm25_name_spaced.retrieve(
+            [query_tokens_spaced], k=bm25_k
+        )
+        results_raw, scores_raw = self._bm25_name_raw.retrieve(
+            [query_tokens_raw], k=bm25_k
+        )
+
+        bm25_map: dict[int, float] = {}
+        for idx, score in zip(results_spaced[0], scores_spaced[0]):
+            decl_id = self._all_declaration_ids[idx]
+            bm25_map[decl_id] = max(bm25_map.get(decl_id, 0.0), float(score))
+        for idx, score in zip(results_raw[0], scores_raw[0]):
+            decl_id = self._all_declaration_ids[idx]
+            bm25_map[decl_id] = max(bm25_map.get(decl_id, 0.0), float(score))
+
+        logger.info(f"BM25 name: {len(bm25_map)} candidates")
+        return bm25_map
+
+    async def _retrieve_semantic_candidates(
+        self, query: str, faiss_k: int
+    ) -> dict[int, float]:
+        """Retrieve candidates using semantic search on informalizations.
+
+        Args:
+            query: Search query string.
+            faiss_k: Number of candidates to retrieve from FAISS.
+
+        Returns:
+            Map of declaration ID to semantic similarity score.
+        """
+        embedding_response = await self.embedding_client.embed([query], is_query=True)
+        query_embedding = np.array([embedding_response.embeddings[0]], dtype=np.float32)
+
+        import faiss as faiss_module
+
+        faiss_module.normalize_L2(query_embedding)
+
+        informal_index = self.faiss_informal_index
+        informal_id_map = self.faiss_informal_id_map
+
+        if hasattr(informal_index, "nprobe"):
+            informal_index.nprobe = 64
+
+        distances, indices = informal_index.search(query_embedding, faiss_k)
+
+        semantic_map: dict[int, float] = {}
+        for idx, dist in zip(indices[0], distances[0]):
+            if idx == -1 or idx >= len(informal_id_map):
+                continue
+            decl_id = informal_id_map[idx]
+            similarity = float(dist)
+            semantic_map[decl_id] = max(semantic_map.get(decl_id, 0.0), similarity)
+
+        logger.info(f"FAISS informal: {len(semantic_map)} candidates")
+        return semantic_map
+
+    def _compute_rrf_scores(
+        self,
+        bm25_map: dict[int, float],
+        semantic_map: dict[int, float],
+    ) -> list[tuple[int, float]]:
+        """Compute RRF scores from BM25 and semantic retrieval signals.
+
+        Args:
+            bm25_map: Map of declaration ID to BM25 score.
+            semantic_map: Map of declaration ID to semantic similarity score.
+
+        Returns:
+            List of (declaration_id, rrf_score) sorted by score descending.
+        """
+        all_candidate_ids = set(bm25_map.keys()) | set(semantic_map.keys())
+        logger.info(f"Total merged candidates: {len(all_candidate_ids)}")
+
+        if not all_candidate_ids:
+            return []
+
+        bm25_sorted = sorted(bm25_map.items(), key=lambda x: x[1], reverse=True)
+        sem_sorted = sorted(semantic_map.items(), key=lambda x: x[1], reverse=True)
+
+        bm25_rank_map = {cid: rank + 1 for rank, (cid, _) in enumerate(bm25_sorted)}
+        sem_rank_map = {cid: rank + 1 for rank, (cid, _) in enumerate(sem_sorted)}
+
+        default_bm25_rank = len(bm25_sorted) + 1
+        default_sem_rank = len(sem_sorted) + 1
+
+        rrf_scores: list[tuple[int, float]] = []
+        for cid in all_candidate_ids:
+            name_rank = bm25_rank_map.get(cid, default_bm25_rank)
+            inf_rank = sem_rank_map.get(cid, default_sem_rank)
+            rrf_score = 1.0 / name_rank + 1.0 / inf_rank
+            rrf_scores.append((cid, rrf_score))
+
+        rrf_scores.sort(key=lambda x: x[1], reverse=True)
+        return rrf_scores
+
+    async def _apply_dependency_boost(
+        self,
+        rrf_scores: list[tuple[int, float]],
+        top_n: int = 500,
+    ) -> tuple[list[tuple[int, float]], dict[int, Declaration]]:
+        """Apply dependency-based boost to RRF scores.
+
+        Declarations that are dependencies of other top candidates get a boost.
+
+        Args:
+            rrf_scores: List of (declaration_id, rrf_score) sorted by score.
+            top_n: Number of top candidates to consider for dependency analysis.
+
+        Returns:
+            Tuple of (boosted_scores, declarations_map).
+        """
+        top_ids = [cid for cid, _ in rrf_scores[:top_n]]
+
+        async with AsyncSession(self.engine) as session:
+            stmt = select(Declaration).where(Declaration.id.in_(top_ids))
+            result = await session.execute(stmt)
+            declarations_map = {d.id: d for d in result.scalars().all()}
+
+        name_to_id = {
+            declarations_map[cid].name: cid
+            for cid in top_ids
+            if cid in declarations_map
+        }
+        dep_counts: dict[int, int] = {cid: 0 for cid in top_ids}
+
+        for cid in top_ids:
+            decl = declarations_map.get(cid)
+            if decl and decl.dependencies:
+                try:
+                    deps = json.loads(decl.dependencies)
+                    for dep_name in deps:
+                        if dep_name in name_to_id:
+                            dep_counts[name_to_id[dep_name]] += 1
+                except json.JSONDecodeError:
+                    pass
+
+        max_deps = max(dep_counts.values()) if dep_counts else 0
+        boosted_scores: list[tuple[int, float]] = []
+
+        for rank, (cid, _) in enumerate(rrf_scores[:top_n], 1):
+            dep_count = dep_counts.get(cid, 0)
+            if max_deps > 0 and dep_count > 0:
+                dep_rank = (max_deps - dep_count) + 1
+            else:
+                dep_rank = max_deps + 1 if max_deps > 0 else top_n + 1
+
+            boosted_score = 1.0 / rank + 1.0 / dep_rank
+            boosted_scores.append((cid, boosted_score))
+
+        boosted_scores.sort(key=lambda x: x[1], reverse=True)
+        logger.info(f"Applied dependency boost to top {top_n} candidates")
+        return boosted_scores, declarations_map
+
+    async def _rerank_candidates(
+        self,
+        query: str,
+        scored_results: list[tuple[Declaration, float]],
+        limit: int,
+    ) -> list[SearchResult]:
+        """Apply cross-encoder reranking with additional signals.
+
+        Args:
+            query: Search query string.
+            scored_results: List of (declaration, score) tuples.
+            limit: Maximum number of results to return.
+
+        Returns:
+            List of SearchResult objects after reranking.
+        """
+        logger.info(f"Reranking top {len(scored_results)} candidates")
+
+        documents = [
+            f"{decl.name}: {decl.informalization}"
+            if decl.informalization
+            else decl.name
+            for decl, _ in scored_results
+        ]
+
+        rerank_response = await self.reranker_client.rerank(query, documents)
+        reranker_scores = rerank_response.scores
+
+        fuzzy_scores = [
+            fuzzy_name_score(query, decl.name) for decl, _ in scored_results
+        ]
+
+        bm25_informal_scores = self._compute_bm25_on_informalizations(
+            query, scored_results
+        )
+
+        dep_counts = self._compute_candidate_dependency_counts(scored_results)
+
+        norm_reranker = normalize_scores(reranker_scores)
+        norm_fuzzy = normalize_scores(fuzzy_scores)
+        norm_bm25 = normalize_scores(bm25_informal_scores)
+        norm_dep = normalize_dependency_counts(dep_counts)
+
+        final_scores = []
+        for i, (decl, _) in enumerate(scored_results):
+            score = 1.0 * norm_reranker[i] + 0.4 * norm_bm25[i] + 0.2 * norm_dep[i]
+            if fuzzy_scores[i] >= 0.7:
+                score += 1.0 * norm_fuzzy[i]
+            final_scores.append(score)
+
+        combined = sorted(
+            zip(scored_results, final_scores),
+            key=lambda x: x[1],
+            reverse=True,
+        )
+
+        return self._filter_and_convert_results(combined, limit)
+
+    def _compute_bm25_on_informalizations(
+        self,
+        query: str,
+        scored_results: list[tuple[Declaration, float]],
+    ) -> list[float]:
+        """Compute BM25 scores on informalizations for reranking.
+
+        Args:
+            query: Search query string.
+            scored_results: List of (declaration, score) tuples.
+
+        Returns:
+            List of BM25 scores for each candidate.
+        """
+        informalizations = [
+            decl.informalization if decl.informalization else decl.name
+            for decl, _ in scored_results
+        ]
+        informal_tokens = [tokenize_words(text) for text in informalizations]
+        query_tokens = tokenize_words(query)
+
+        bm25_informal = bm25s.BM25(method="bm25+")
+        bm25_informal.index(informal_tokens)
+        results, scores = bm25_informal.retrieve([query_tokens], k=len(informal_tokens))
+
+        bm25_scores = [0.0] * len(scored_results)
+        for idx, score in zip(results[0], scores[0]):
+            if int(idx) < len(bm25_scores):
+                bm25_scores[int(idx)] = float(score)
+
+        return bm25_scores
+
+    def _compute_candidate_dependency_counts(
+        self,
+        scored_results: list[tuple[Declaration, float]],
+    ) -> list[int]:
+        """Count how many candidates depend on each declaration.
+
+        Args:
+            scored_results: List of (declaration, score) tuples.
+
+        Returns:
+            List of dependency counts for each candidate.
+        """
+        candidate_names = {decl.name for decl, _ in scored_results}
+        dep_counts_map: dict[str, int] = {name: 0 for name in candidate_names}
+
+        for decl, _ in scored_results:
+            if decl.dependencies:
+                try:
+                    deps = json.loads(decl.dependencies)
+                    for dep_name in deps:
+                        if dep_name in dep_counts_map:
+                            dep_counts_map[dep_name] += 1
+                except json.JSONDecodeError:
+                    pass
+
+        return [dep_counts_map.get(decl.name, 0) for decl, _ in scored_results]
+
+    def _filter_and_convert_results(
+        self,
+        combined: list[tuple[tuple[Declaration, float], float]],
+        limit: int,
+    ) -> list[SearchResult]:
+        """Filter auto-generated declarations and convert to SearchResult.
+
+        Args:
+            combined: List of ((declaration, old_score), final_score) tuples.
+            limit: Maximum number of results to return.
+
+        Returns:
+            List of SearchResult objects.
+        """
+        results = []
+        for (decl, _), _ in combined:
+            if not is_autogenerated(decl.name):
+                results.append(self._to_search_result(decl))
+                if len(results) >= limit:
+                    break
+        return results
+
+    def _extract_package(self, module: str) -> str:
+        """Extract package name from module path.
+
+        Args:
+            module: Full module path (e.g., "Mathlib.Algebra.Group").
+
+        Returns:
+            Package name (first component of module path).
+        """
+        return module.split(".")[0] if module else ""
+
+    def _filter_by_packages(
+        self,
+        declarations_map: dict[int, Declaration],
+        packages: list[str],
+    ) -> dict[int, Declaration]:
+        """Filter declarations to only include specified packages.
+
+        Args:
+            declarations_map: Map of declaration ID to Declaration.
+            packages: List of package names to include.
+
+        Returns:
+            Filtered declarations map.
+        """
+        if not packages:
+            return declarations_map
+
+        package_set = set(packages)
+        return {
+            cid: decl
+            for cid, decl in declarations_map.items()
+            if self._extract_package(decl.module) in package_set
+        }
+
+    async def search(
+        self,
+        query: str,
+        limit: int = 50,
+        faiss_k: int = 1000,
+        bm25_k: int = 1000,
+        rerank_top: int | None = 25,
+        packages: list[str] | None = None,
+    ) -> list[SearchResult]:
+        """Search for Lean declarations using Reciprocal Rank Fusion.
+
+        Two-signal approach:
+        1. BM25+ on declaration names (lexical match)
+        2. Semantic search on informalizations (meaning match)
+
+        Combined via RRF: score = 1/name_rank + 1/informal_rank
+
+        Optionally applies cross-encoder reranking to the top candidates.
+
+        Args:
+            query: Search query string.
+            limit: Maximum number of results to return. Defaults to 50.
+            faiss_k: Number of candidates from FAISS index. Defaults to 1000.
+            bm25_k: Number of candidates from BM25 index. Defaults to 1000.
+            rerank_top: If set, apply cross-encoder reranking to top N candidates.
+                Set to 0 or None to skip reranking.
+            packages: Optional list of package names to filter by. If provided,
+                only declarations from these packages will be returned.
+
+        Returns:
+            List of SearchResult objects, ranked by combined score.
+        """
+        if not query.strip():
+            return []
+
+        bm25_map = self._retrieve_bm25_candidates(query, bm25_k)
+        semantic_map = await self._retrieve_semantic_candidates(query, faiss_k)
+        rrf_scores = self._compute_rrf_scores(bm25_map, semantic_map)
+
+        if not rrf_scores:
+            return []
+
+        boosted_scores, declarations_map = await self._apply_dependency_boost(
+            rrf_scores
+        )
+
+        # Apply package filtering if specified
+        if packages:
+            declarations_map = self._filter_by_packages(declarations_map, packages)
+            # Filter boosted_scores to only include filtered declarations
+            boosted_scores = [
+                (cid, score) for cid, score in boosted_scores if cid in declarations_map
+            ]
+            logger.info(f"Filtered to {len(declarations_map)} in {packages}")
+
+        top_n = rerank_top if rerank_top and rerank_top > 0 else limit
+
+        scored_results: list[tuple[Declaration, float]] = [
+            (declarations_map[cid], score)
+            for cid, score in boosted_scores[:top_n]
+            if cid in declarations_map
+        ]
+
+        if rerank_top and rerank_top > 0:
+            return await self._rerank_candidates(query, scored_results, limit)
+
+        results = []
+        for decl, _ in scored_results:
+            if not is_autogenerated(decl.name):
+                results.append(self._to_search_result(decl))
+                if len(results) >= limit:
+                    break
+        return results
+
+    async def get_by_id(self, declaration_id: int) -> SearchResult | None:
+        """Retrieve a declaration by ID.
+
+        Args:
+            declaration_id: The declaration ID.
+
+        Returns:
+            SearchResult if found, None otherwise.
+        """
+        async with AsyncSession(self.engine) as session:
+            decl = await session.get(Declaration, declaration_id)
+            return self._to_search_result(decl) if decl else None
+
+    async def get_by_name(self, name: str) -> SearchResult | None:
+        """Retrieve a declaration by its exact name.
+
+        Args:
+            name: The exact declaration name (e.g., "AlgebraicGeometry.Scheme").
+
+        Returns:
+            SearchResult if found, None otherwise.
+        """
+        async with AsyncSession(self.engine) as session:
+            stmt = select(Declaration).where(Declaration.name == name)
+            result = await session.execute(stmt)
+            decl = result.scalar_one_or_none()
+            return self._to_search_result(decl) if decl else None
+
+    def _to_search_result(self, decl: Declaration) -> SearchResult:
+        """Convert Declaration ORM object to SearchResult.
+
+        Args:
+            decl: Declaration ORM object.
+
+        Returns:
+            SearchResult pydantic model.
+        """
+        return SearchResult(
+            id=decl.id,
+            name=decl.name,
+            module=decl.module,
+            docstring=decl.docstring,
+            source_text=decl.source_text,
+            source_link=decl.source_link,
+            dependencies=decl.dependencies,
+            informalization=decl.informalization,
+        )