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

lean_explore/mcp/tools.py

@@ -1,27 +1,36 @@
-# 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, Union
+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
+
+
+class SearchResponseDict(TypedDict, total=False):
+    """Serialized SearchResponse for MCP tool responses."""
+
+    query: str
+    results: list[SearchResultDict]
+    count: int
+    processing_time_ms: int | None
 
-# Import Pydantic models for type hinting and for creating response dicts
-from lean_explore.shared.models.api import (
-    APICitationsResponse,
-    APISearchResponse,
-    APISearchResultItem,
-)
 
 logger = logging.getLogger(__name__)
 
@@ -33,238 +42,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: Union[str, List[str]],
-    package_filters: Optional[List[str]] = None,
+    query: str,
     limit: int = 10,
-) -> List[Dict[str, Any]]:
-    """Searches Lean statement groups by a query string or list of strings.
-
-    This tool allows for filtering by package names and limits the number
-    of results returned per query.
+    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: A single search query string or a list of query strings. For
-               example, "continuous function" or ["prime number theorem",
-               "fundamental theorem of arithmetic"].
-        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 per query.
-               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 list of dictionaries, where each dictionary corresponds to the
-        APISearchResponse model. Each response contains the search results
-        for a single query. 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/queries: '{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.")
         raise RuntimeError("Search functionality not available on configured backend.")
 
-    tool_limit = max(1, limit)
-    backend_responses: Union[APISearchResponse, List[APISearchResponse]]
-
-    # Conditionally await based on the backend's search method type
+    # Call backend search (handle both async and sync)
     if asyncio.iscoroutinefunction(backend.search):
-        backend_responses = await backend.search(
-            query=query, package_filters=package_filters
+        response: SearchResponse = await backend.search(
+            query=query, limit=limit, rerank_top=rerank_top, packages=packages
         )
     else:
-        backend_responses = backend.search(query=query, package_filters=package_filters)
-
-    # Normalize to a list for consistent processing, handling None from backend.
-    if backend_responses is None:
-        responses_list = []
-    else:
-        responses_list = (
-            [backend_responses]
-            if isinstance(backend_responses, APISearchResponse)
-            else backend_responses
-        )
-
-    final_mcp_responses = []
-
-    for response_pydantic in responses_list:
-        if not response_pydantic:
-            logger.warning("A backend search returned None; skipping this response.")
-            continue
-
-        actual_backend_results = response_pydantic.results
-        mcp_results_list = []
-        for backend_item in actual_backend_results[:tool_limit]:
-            mcp_results_list.append(_prepare_mcp_result_item(backend_item))
-
-        final_mcp_response = APISearchResponse(
-            query=response_pydantic.query,
-            packages_applied=response_pydantic.packages_applied,
-            results=mcp_results_list,
-            count=len(mcp_results_list),
-            total_candidates_considered=response_pydantic.total_candidates_considered,
-            processing_time_ms=response_pydantic.processing_time_ms,
+        response: SearchResponse = backend.search(
+            query=query, limit=limit, rerank_top=rerank_top, packages=packages
         )
-        final_mcp_responses.append(final_mcp_response.model_dump(exclude_none=True))
 
-    return final_mcp_responses
+    # Return as dict for MCP
+    return response.model_dump(exclude_none=True)
 
 
 @mcp_app.tool()
 async def get_by_id(
-    ctx: MCPContext, group_id: Union[int, List[int]]
-) -> List[Optional[Dict[str, Any]]]:
-    """Retrieves specific statement groups by their unique identifier(s).
-
-    The `display_statement_text` field is omitted from the response. This tool
-    always returns a list of results.
+    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: A single unique integer identifier or a list of identifiers
-                  of the statement group(s) to retrieve. For example, `12345` or
-                  `[12345, 67890]`.
+        declaration_id: The unique integer identifier of the declaration.
 
     Returns:
-        A list of dictionaries, where each dictionary corresponds to the
-        APISearchResultItem model. If an ID is not found, its corresponding
-        entry in the list will be None (serialized as JSON null by MCP).
+        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(s): {group_id}")
-
-    backend_items: Union[
-        Optional[APISearchResultItem], List[Optional[APISearchResultItem]]
-    ]
-    if asyncio.iscoroutinefunction(backend.get_by_id):
-        backend_items = await backend.get_by_id(group_id=group_id)
-    else:
-        backend_items = backend.get_by_id(group_id=group_id)
-
-    # Normalize to a list for consistent return type
-    items_list = (
-        [backend_items] if not isinstance(backend_items, list) else backend_items
-    )
-
-    mcp_items = []
-    for item in items_list:
-        if item:
-            mcp_item = _prepare_mcp_result_item(item)
-            mcp_items.append(mcp_item.model_dump(exclude_none=True))
-        else:
-            mcp_items.append(None)
-
-    return mcp_items
-
-
-@mcp_app.tool()
-async def get_dependencies(
-    ctx: MCPContext, group_id: Union[int, List[int]]
-) -> List[Optional[Dict[str, Any]]]:
-    """Retrieves direct dependencies (citations) for specific statement group(s).
-
-    The `display_statement_text` field within each cited item is omitted
-    from the response. This tool always returns a list of results.
+    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: A single unique integer identifier or a list of identifiers for
-                  the statement group(s) for which to fetch direct dependencies.
-                  For example, `12345` or `[12345, 67890]`.
-
-    Returns:
-        A list of dictionaries, where each dictionary corresponds to the
-        APICitationsResponse model. If a source group ID is not found or has
-        no dependencies, its corresponding entry will be None.
-    """
-    backend = await _get_backend_from_context(ctx)
-    logger.info(f"MCP Tool 'get_dependencies' called for group_id(s): {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_responses: Union[
-        Optional[APICitationsResponse], List[Optional[APICitationsResponse]]
-    ]
-    if asyncio.iscoroutinefunction(backend.get_dependencies):
-        backend_responses = 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_responses = backend.get_dependencies(group_id=group_id)
-
-    # Normalize to a list for consistent return type
-    responses_list = (
-        [backend_responses]
-        if not isinstance(backend_responses, list)
-        else backend_responses
-    )
-    final_mcp_responses = []
-
-    for response in responses_list:
-        if response:
-            mcp_citations_list = []
-            for backend_item in response.citations:
-                mcp_citations_list.append(_prepare_mcp_result_item(backend_item))
-
-            final_response = APICitationsResponse(
-                source_group_id=response.source_group_id,
-                citations=mcp_citations_list,
-                count=len(mcp_citations_list),
-            )
-            final_mcp_responses.append(final_response.model_dump(exclude_none=True))
-        else:
-            final_mcp_responses.append(None)
+        result: SearchResult | None = backend.get_by_id(declaration_id=declaration_id)
 
-    return final_mcp_responses
+    # 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"]