lean-explore 0.2.2__py3-none-any.whl → 1.0.0__py3-none-any.whl

lean_explore/mcp/tools.py

@@ -1,27 +1,35 @@
-# src/lean_explore/mcp/tools.py
+"""Defines MCP tools for interacting with the Lean Explore search engine."""
 
-"""Defines MCP tools for interacting with the Lean Explore search engine.
-
-These tools provide functionalities such as searching for statement groups,
-retrieving specific groups by ID, and getting their dependencies. They
-utilize a backend service (either an API client or a local service)
-made available through the MCP application context.
-"""
-
-import asyncio  # Needed for asyncio.iscoroutinefunction
+import asyncio
 import logging
-from typing import Any, Dict, List, Optional
+from typing import TypedDict
 
 from mcp.server.fastmcp import Context as MCPContext
 
 from lean_explore.mcp.app import AppContext, BackendServiceType, mcp_app
+from lean_explore.models import SearchResponse, SearchResult
+
+
+class SearchResultDict(TypedDict, total=False):
+    """Serialized SearchResult for MCP tool responses."""
+
+    id: int
+    name: str
+    module: str
+    docstring: str | None
+    source_text: str
+    source_link: str
+    dependencies: str | None
+    informalization: str | None
+
 
-# Import Pydantic models for type hinting and for creating response dicts
-from lean_explore.shared.models.api import (
-    APICitationsResponse,
-    APISearchResponse,
-    APISearchResultItem,
-)
+class SearchResponseDict(TypedDict, total=False):
+    """Serialized SearchResponse for MCP tool responses."""
+
+    query: str
+    results: list[SearchResultDict]
+    count: int
+    processing_time_ms: int | None
 
 logger = logging.getLogger(__name__)
 
@@ -33,210 +41,95 @@ async def _get_backend_from_context(ctx: MCPContext) -> BackendServiceType:
         ctx: The MCP context provided to the tool.
 
     Returns:
-        The configured backend service (APIClient or LocalService).
-        Guaranteed to be non-None if this function returns, otherwise
-        it raises an exception.
+        The configured backend service (ApiClient or Service).
 
     Raises:
-        RuntimeError: If the backend service is not available in the context,
-                      indicating a server configuration issue.
+        RuntimeError: If the backend service is not available in the context.
     """
     app_ctx: AppContext = ctx.request_context.lifespan_context
     backend = app_ctx.backend_service
     if not backend:
-        logger.error(
-            "MCP Tool Error: Backend service is not available in lifespan_context."
-        )
+        logger.error("MCP Tool Error: Backend service is not available.")
         raise RuntimeError("Backend service not configured or available for MCP tool.")
     return backend
 
 
-def _prepare_mcp_result_item(backend_item: APISearchResultItem) -> APISearchResultItem:
-    """Prepares an APISearchResultItem for MCP response.
-
-    This helper ensures that the item sent over MCP does not include
-    the display_statement_text, as the full statement_text is preferred
-    for model consumption.
-
-    Args:
-        backend_item: The item as received from the backend service.
-
-    Returns:
-        A new APISearchResultItem instance suitable for MCP responses.
-    """
-    # Create a new instance or use .model_copy(update=...) for Pydantic v2
-    return APISearchResultItem(
-        id=backend_item.id,
-        primary_declaration=backend_item.primary_declaration.model_copy()
-        if backend_item.primary_declaration
-        else None,
-        source_file=backend_item.source_file,
-        range_start_line=backend_item.range_start_line,
-        statement_text=backend_item.statement_text,
-        docstring=backend_item.docstring,
-        informal_description=backend_item.informal_description,
-        display_statement_text=None,  # Ensure this is not sent over MCP
-    )
-
-
 @mcp_app.tool()
 async def search(
     ctx: MCPContext,
     query: str,
-    package_filters: Optional[List[str]] = None,
     limit: int = 10,
-) -> Dict[str, Any]:
-    """Searches Lean statement groups by a query string.
-
-    This tool allows for filtering by package names and limits the number
-    of results returned.
+    rerank_top: int | None = 50,
+    packages: list[str] | None = None,
+) -> SearchResponseDict:
+    """Searches Lean declarations by a query string.
 
     Args:
-        ctx: The MCP context, providing access to shared resources like the
-             backend service.
-        query: The search query string. For example, "continuous function" or
-               "prime number theorem".
-        package_filters: An optional list of package names to filter the search
-                         results by. For example, `["Mathlib.Analysis",
-                         "Mathlib.Order"]`. If None or empty, no package filter
-                         is applied.
-        limit: The maximum number of search results to return from this tool.
-               Defaults to 10. Must be a positive integer.
+        ctx: The MCP context, providing access to the backend service.
+        query: A search query string, e.g., "continuous function".
+        limit: The maximum number of search results to return. Defaults to 10.
+        rerank_top: Number of candidates to rerank with cross-encoder. Set to 0 or
+            None to skip reranking. Defaults to 50. Only used with local backend.
+        packages: Filter results to specific packages (e.g., ["Mathlib", "Std"]).
+            Defaults to None (all packages).
 
     Returns:
-        A dictionary corresponding to the APISearchResponse model, containing
-        the search results (potentially truncated by the `limit` parameter of
-        this tool), and metadata about the search operation. The
-        `display_statement_text` field within each result item is omitted.
+        A dictionary containing the search response with results.
     """
     backend = await _get_backend_from_context(ctx)
     logger.info(
-        f"MCP Tool 'search' called with query: '{query}', "
-        f"packages: {package_filters}, tool_limit: {limit}"
+        f"MCP Tool 'search' called with query: '{query}', limit: {limit}, "
+        f"rerank_top: {rerank_top}, packages: {packages}"
     )
 
     if not hasattr(backend, "search"):
         logger.error("Backend service does not have a 'search' method.")
-        # This should ideally return a structured error for MCP if possible.
-        # For now, FastMCP will convert this RuntimeError.
         raise RuntimeError("Search functionality not available on configured backend.")
 
-    tool_limit = max(1, limit)  # Ensure limit is at least 1 for slicing
-    api_response_pydantic: Optional[APISearchResponse]
-
-    # Conditionally await based on the backend's search method type
+    # Call backend search (handle both async and sync)
     if asyncio.iscoroutinefunction(backend.search):
-        api_response_pydantic = await backend.search(
-            query=query,
-            package_filters=package_filters,
-            # The backend.search method uses its own internal default for limit
-            # if None is passed, or the passed limit.
-            # The MCP tool will truncate the results later using tool_limit.
+        response: SearchResponse = await backend.search(
+            query=query, limit=limit, rerank_top=rerank_top, packages=packages
         )
     else:
-        api_response_pydantic = backend.search(
-            query=query, package_filters=package_filters
-        )
-
-    if not api_response_pydantic:
-        logger.warning("Backend search returned None, responding with empty results.")
-        empty_response = APISearchResponse(
-            query=query,
-            packages_applied=package_filters or [],
-            results=[],
-            count=0,
-            total_candidates_considered=0,
-            processing_time_ms=0,
+        response: SearchResponse = backend.search(
+            query=query, limit=limit, rerank_top=rerank_top, packages=packages
         )
-        return empty_response.model_dump(exclude_none=True)
-
-    actual_backend_results = api_response_pydantic.results
 
-    mcp_results_list = []
-    for backend_item in actual_backend_results[:tool_limit]:  # Apply MCP tool's limit
-        mcp_results_list.append(_prepare_mcp_result_item(backend_item))
-
-    final_mcp_response = APISearchResponse(
-        query=api_response_pydantic.query,
-        packages_applied=api_response_pydantic.packages_applied,
-        results=mcp_results_list,
-        count=len(mcp_results_list),  # Count is after this tool's truncation
-        total_candidates_considered=api_response_pydantic.total_candidates_considered,
-        processing_time_ms=api_response_pydantic.processing_time_ms,
-    )
-
-    return final_mcp_response.model_dump(exclude_none=True)
+    # Return as dict for MCP
+    return response.model_dump(exclude_none=True)
 
 
 @mcp_app.tool()
-async def get_by_id(ctx: MCPContext, group_id: int) -> Optional[Dict[str, Any]]:
-    """Retrieves a specific statement group by its unique identifier.
-
-    The `display_statement_text` field is omitted from the response.
+async def get_by_id(
+    ctx: MCPContext,
+    declaration_id: int,
+) -> SearchResultDict | None:
+    """Retrieves a specific declaration by its unique identifier.
 
     Args:
         ctx: The MCP context, providing access to the backend service.
-        group_id: The unique integer identifier of the statement group to retrieve.
-                  For example, `12345`.
+        declaration_id: The unique integer identifier of the declaration.
 
     Returns:
-        A dictionary corresponding to the APISearchResultItem model if a
-        statement group with the given ID is found (with
-        `display_statement_text` omitted). Returns None (which will be
-        serialized as JSON null by MCP) if no such group exists.
+        A dictionary representing the SearchResult, or None if not found.
     """
     backend = await _get_backend_from_context(ctx)
-    logger.info(f"MCP Tool 'get_by_id' called for group_id: {group_id}")
-
-    backend_item: Optional[APISearchResultItem]
-    if asyncio.iscoroutinefunction(backend.get_by_id):
-        backend_item = await backend.get_by_id(group_id=group_id)
-    else:
-        backend_item = backend.get_by_id(group_id=group_id)
-
-    if backend_item:
-        mcp_item = _prepare_mcp_result_item(backend_item)
-        return mcp_item.model_dump(exclude_none=True)
-    return None
-
-
-@mcp_app.tool()
-async def get_dependencies(ctx: MCPContext, group_id: int) -> Optional[Dict[str, Any]]:
-    """Retrieves the direct dependencies (citations) for a specific statement group.
-
-    The `display_statement_text` field within each cited item is omitted
-    from the response.
+    logger.info(f"MCP Tool 'get_by_id' called for declaration_id: {declaration_id}")
 
-    Args:
-        ctx: The MCP context, providing access to the backend service.
-        group_id: The unique integer identifier of the statement group for which
-                  to fetch its direct dependencies. For example, `12345`.
-
-    Returns:
-        A dictionary corresponding to the APICitationsResponse model, which
-        contains a list of cited statement groups (each with
-        `display_statement_text` omitted), if the source group_id
-        is found and has dependencies. Returns None (serialized as JSON null
-        by MCP) if the source group is not found or has no dependencies.
-    """
-    backend = await _get_backend_from_context(ctx)
-    logger.info(f"MCP Tool 'get_dependencies' called for group_id: {group_id}")
+    if not hasattr(backend, "get_by_id"):
+        logger.error("Backend service does not have a 'get_by_id' method.")
+        raise RuntimeError(
+            "Get by ID functionality not available on configured backend."
+        )
 
-    backend_response: Optional[APICitationsResponse]
-    if asyncio.iscoroutinefunction(backend.get_dependencies):
-        backend_response = await backend.get_dependencies(group_id=group_id)
+    # Call backend get_by_id (handle both async and sync)
+    if asyncio.iscoroutinefunction(backend.get_by_id):
+        result: SearchResult | None = await backend.get_by_id(
+            declaration_id=declaration_id
+        )
     else:
-        backend_response = backend.get_dependencies(group_id=group_id)
+        result: SearchResult | None = backend.get_by_id(declaration_id=declaration_id)
 
-    if backend_response:
-        mcp_citations_list = []
-        for backend_item in backend_response.citations:
-            mcp_citations_list.append(_prepare_mcp_result_item(backend_item))
-
-        final_mcp_response = APICitationsResponse(
-            source_group_id=backend_response.source_group_id,
-            citations=mcp_citations_list,
-            count=len(mcp_citations_list),
-        )
-        return final_mcp_response.model_dump(exclude_none=True)
-    return None
+    # Return as dict for MCP, or None
+    return result.model_dump(exclude_none=True) if result else None

lean_explore/models/__init__.py

@@ -0,0 +1,9 @@
+"""Data models for lean_explore.
+
+This package contains database models and type definitions for search results.
+"""
+
+from lean_explore.models.search_db import Base, Declaration
+from lean_explore.models.search_types import SearchResponse, SearchResult
+
+__all__ = ["Base", "Declaration", "SearchResult", "SearchResponse"]

lean_explore/models/search_db.py

@@ -0,0 +1,76 @@
+"""SQLAlchemy ORM models for Lean declaration database.
+
+Simple schema for a Lean declaration search engine.
+Uses SQLAlchemy 2.0 syntax with SQLite for storage and FAISS for vector search.
+"""
+
+import struct
+
+from sqlalchemy import Integer, LargeBinary, Text
+from sqlalchemy.orm import DeclarativeBase, Mapped, mapped_column
+from sqlalchemy.types import TypeDecorator
+
+
+class BinaryEmbedding(TypeDecorator):
+    """Custom type for storing embeddings as binary blobs.
+
+    Converts between Python list[float] and compact binary representation.
+    Uses float32 (4 bytes per dimension) for ~5x space savings over JSON.
+    """
+
+    impl = LargeBinary
+    cache_ok = True
+
+    def process_bind_param(self, value: list[float] | None, dialect) -> bytes | None:
+        """Convert list[float] to binary for storage."""
+        if value is None:
+            return None
+        return struct.pack(f"{len(value)}f", *value)
+
+    def process_result_value(self, value: bytes | None, dialect) -> list[float] | None:
+        """Convert binary back to list[float] on retrieval."""
+        if value is None:
+            return None
+        num_floats = len(value) // 4
+        return list(struct.unpack(f"{num_floats}f", value))
+
+
+class Base(DeclarativeBase):
+    """Base class for SQLAlchemy declarative models."""
+
+    pass
+
+
+class Declaration(Base):
+    """Represents a Lean declaration for search."""
+
+    __tablename__ = "declarations"
+
+    id: Mapped[int] = mapped_column(Integer, primary_key=True)
+    """Primary key identifier."""
+
+    name: Mapped[str] = mapped_column(Text, unique=True, index=True, nullable=False)
+    """Fully qualified Lean name (e.g., 'Nat.add')."""
+
+    module: Mapped[str] = mapped_column(Text, index=True, nullable=False)
+    """Module name (e.g., 'Mathlib.Data.List.Basic')."""
+
+    docstring: Mapped[str | None] = mapped_column(Text, nullable=True)
+    """Documentation string from the source code, if available."""
+
+    source_text: Mapped[str] = mapped_column(Text, nullable=False)
+    """The actual Lean source code for this declaration."""
+
+    source_link: Mapped[str] = mapped_column(Text, nullable=False)
+    """GitHub URL to the declaration source code."""
+
+    dependencies: Mapped[str | None] = mapped_column(Text, nullable=True)
+    """JSON array of declaration names this declaration depends on."""
+
+    informalization: Mapped[str | None] = mapped_column(Text, nullable=True)
+    """Natural language description of the declaration."""
+
+    informalization_embedding: Mapped[list[float] | None] = mapped_column(
+        BinaryEmbedding, nullable=True
+    )
+    """1024-dimensional embedding of the informalization text (binary float32)."""

lean_explore/models/search_types.py

@@ -0,0 +1,53 @@
+"""Type definitions for search results and related data structures."""
+
+from pydantic import BaseModel, ConfigDict
+
+
+class SearchResult(BaseModel):
+    """A search result representing a Lean declaration.
+
+    This model represents the core information returned from a search query,
+    mirroring the essential fields from the database Declaration model.
+    """
+
+    id: int
+    """Primary key identifier."""
+
+    name: str
+    """Fully qualified Lean name (e.g., 'Nat.add')."""
+
+    module: str
+    """Module name (e.g., 'Mathlib.Data.List.Basic')."""
+
+    docstring: str | None
+    """Documentation string from the source code, if available."""
+
+    source_text: str
+    """The actual Lean source code for this declaration."""
+
+    source_link: str
+    """GitHub URL to the declaration source code."""
+
+    dependencies: str | None
+    """JSON array of declaration names this declaration depends on."""
+
+    informalization: str | None
+    """Natural language description of the declaration."""
+
+    model_config = ConfigDict(from_attributes=True)
+
+
+class SearchResponse(BaseModel):
+    """Response from a search operation containing results and metadata."""
+
+    query: str
+    """The original search query string."""
+
+    results: list[SearchResult]
+    """List of search results."""
+
+    count: int
+    """Number of results returned."""
+
+    processing_time_ms: int | None = None
+    """Processing time in milliseconds, if available."""

lean_explore/search/__init__.py

@@ -0,0 +1,32 @@
+"""Search package for Lean Explore.
+
+This package provides hybrid search for Lean declarations using BM25 lexical
+matching and FAISS semantic search, combined via Reciprocal Rank Fusion.
+
+Modules:
+    engine: Core SearchEngine class with hybrid retrieval and cross-encoder reranking.
+    scoring: Score normalization and fusion algorithms (RRF, weighted fusion).
+    service: Service layer wrapper for search operations.
+    tokenization: Text tokenization utilities for Lean declaration names.
+
+Note: SearchEngine and Service are lazily imported to avoid loading FAISS at module
+import time, which helps prevent OpenMP library conflicts with torch on macOS.
+"""
+
+from lean_explore.models import SearchResponse, SearchResult
+
+
+def __getattr__(name: str):
+    """Lazy import SearchEngine and Service to avoid FAISS loading at import time."""
+    if name == "SearchEngine":
+        from lean_explore.search.engine import SearchEngine
+
+        return SearchEngine
+    if name == "Service":
+        from lean_explore.search.service import Service
+
+        return Service
+    raise AttributeError(f"module {__name__!r} has no attribute {name!r}")
+
+
+__all__ = ["SearchEngine", "Service", "SearchResponse", "SearchResult"]