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

lean_explore/extract/package_registry.py

@@ -0,0 +1,45 @@
+"""Package registry for Lean extraction.
+
+This module contains the registry of Lean packages available for extraction.
+"""
+
+from lean_explore.extract.package_config import PackageConfig, VersionStrategy
+
+PACKAGE_REGISTRY: dict[str, PackageConfig] = {
+    "mathlib": PackageConfig(
+        name="mathlib",
+        git_url="https://github.com/leanprover-community/mathlib4",
+        module_prefixes=["Mathlib", "Batteries", "Init", "Lean", "Std"],
+        version_strategy=VersionStrategy.LATEST,
+        depends_on=[],
+        extract_core=True,
+    ),
+    "physlean": PackageConfig(
+        name="physlean",
+        git_url="https://github.com/HEPLean/PhysLean",
+        module_prefixes=["PhysLean"],
+        version_strategy=VersionStrategy.TAGGED,
+        depends_on=["mathlib"],
+    ),
+    "flt": PackageConfig(
+        name="flt",
+        git_url="https://github.com/ImperialCollegeLondon/FLT",
+        module_prefixes=["FLT"],
+        version_strategy=VersionStrategy.LATEST,
+        depends_on=["mathlib"],
+    ),
+    "formal-conjectures": PackageConfig(
+        name="formal-conjectures",
+        git_url="https://github.com/google-deepmind/formal-conjectures",
+        module_prefixes=["FormalConjectures", "FormalConjecturesForMathlib"],
+        version_strategy=VersionStrategy.LATEST,
+        depends_on=["mathlib"],
+    ),
+    "cslib": PackageConfig(
+        name="cslib",
+        git_url="https://github.com/leanprover/cslib",
+        module_prefixes=["Cslib"],
+        version_strategy=VersionStrategy.LATEST,
+        depends_on=["mathlib"],
+    ),
+}

lean_explore/extract/package_utils.py

@@ -0,0 +1,105 @@
+"""Utility functions for package configuration.
+
+This module provides helper functions for working with the package registry,
+including dependency ordering, toolchain resolution, and lakefile manipulation.
+"""
+
+import logging
+import re
+from pathlib import Path
+
+from lean_explore.extract.package_config import PackageConfig, VersionStrategy
+from lean_explore.extract.package_registry import PACKAGE_REGISTRY
+
+logger = logging.getLogger(__name__)
+
+
+def get_package_for_module(module_name: str) -> str | None:
+    """Determine which package a module belongs to.
+
+    Args:
+        module_name: Fully qualified module name (e.g., 'Mathlib.Data.List.Basic')
+
+    Returns:
+        Package name or None if not recognized.
+    """
+    for package_name, configuration in PACKAGE_REGISTRY.items():
+        if configuration.should_include_module(module_name):
+            return package_name
+    return None
+
+
+def get_extraction_order() -> list[str]:
+    """Get packages in dependency order for extraction.
+
+    Returns packages ordered so dependencies come before dependents.
+    """
+    result: list[str] = []
+    visited: set[str] = set()
+
+    def visit(name: str) -> None:
+        if name in visited:
+            return
+        visited.add(name)
+        configuration = PACKAGE_REGISTRY.get(name)
+        if configuration:
+            for dep in configuration.depends_on:
+                visit(dep)
+            result.append(name)
+
+    for name in PACKAGE_REGISTRY:
+        visit(name)
+
+    return result
+
+
+def get_package_toolchain(package_configuration: PackageConfig) -> tuple[str, str]:
+    """Get the toolchain and ref for a package based on its version strategy.
+
+    Args:
+        package_configuration: Package configuration
+
+    Returns:
+        Tuple of (lean_toolchain, git_ref) where git_ref is the branch/tag to use.
+    """
+    from lean_explore.extract.github import fetch_latest_tag, fetch_lean_toolchain
+
+    if package_configuration.version_strategy == VersionStrategy.LATEST:
+        for branch in ["main", "master"]:
+            try:
+                toolchain = fetch_lean_toolchain(package_configuration.git_url, branch)
+                return toolchain, branch
+            except RuntimeError:
+                continue
+        raise RuntimeError(
+            f"Could not fetch toolchain from main or master for "
+            f"{package_configuration.name}"
+        )
+    else:
+        latest_tag = fetch_latest_tag(package_configuration.git_url)
+        toolchain = fetch_lean_toolchain(package_configuration.git_url, latest_tag)
+        return toolchain, latest_tag
+
+
+def update_lakefile_docgen_version(lakefile_path: Path, lean_version: str) -> None:
+    """Update the doc-gen4 version in a lakefile to match the Lean version.
+
+    Args:
+        lakefile_path: Path to lakefile.lean
+        lean_version: Lean version like 'v4.27.0'
+    """
+    content = lakefile_path.read_text()
+
+    pattern = (
+        r'require «doc-gen4» from git\s+'
+        r'"https://github\.com/leanprover/doc-gen4"(?:\s+@\s+"[^"]*")?'
+    )
+    replacement = (
+        f'require «doc-gen4» from git\n'
+        f'  "https://github.com/leanprover/doc-gen4" @ "{lean_version}"'
+    )
+    new_content = re.sub(pattern, replacement, content)
+
+    if new_content != content:
+        lakefile_path.write_text(new_content)
+        logger.info(f"Updated doc-gen4 version to {lean_version} in {lakefile_path}")

lean_explore/extract/types.py

@@ -0,0 +1,25 @@
+"""Type definitions for doc-gen4 data extraction."""
+
+from pydantic import BaseModel
+
+
+class Declaration(BaseModel):
+    """A declaration for database storage - mirrors schemas.Declaration."""
+
+    name: str
+    """Fully qualified Lean name."""
+
+    module: str
+    """Module name."""
+
+    docstring: str | None
+    """Documentation string, if available."""
+
+    source_text: str
+    """The actual Lean source code."""
+
+    source_link: str
+    """GitHub URL to the source code."""
+
+    dependencies: list[str] | None
+    """List of declaration names this declaration depends on."""

lean_explore/mcp/__init__.py

@@ -1 +1,11 @@
-"""Local package for lean explore."""
+"""Model Context Protocol (MCP) server package for Lean Explore.
+
+This package provides an MCP server that exposes Lean declaration search
+functionality as tools for AI assistants. It supports both remote API and
+local backends.
+
+Modules:
+    server: Main MCP server entry point with argument parsing and initialization.
+    app: FastMCP application setup and lifespan context management.
+    tools: MCP tool definitions for search and retrieval operations.
+"""

lean_explore/mcp/app.py

@@ -1,29 +1,19 @@
-# src/lean_explore/mcp/app.py
-
-"""Initializes the FastMCP application and its lifespan context.
-
-This module creates the main FastMCP application instance and defines
-a lifespan context manager. The lifespan manager is responsible for
-making the configured backend service (API client or local service)
-available to MCP tools via the request context. The actual backend
-instance will be set by the server startup script before running the app.
-"""
+"""Initializes the FastMCP application and its lifespan context."""
 
 import logging
+from collections.abc import AsyncIterator
 from contextlib import asynccontextmanager
 from dataclasses import dataclass
-from typing import AsyncIterator, Union
 
 from mcp.server.fastmcp import FastMCP
 
-# Import your backend service types for type hinting
-from lean_explore.api.client import Client as APIClient
-from lean_explore.local.service import Service as LocalService
+from lean_explore.api import ApiClient
+from lean_explore.search import Service
 
 logger = logging.getLogger(__name__)
 
-# Define a type for the backend service to be used by tools
-BackendServiceType = Union[APIClient, LocalService, None]
+# Define a type for the backend service
+BackendServiceType = ApiClient | Service | None
 
 
 @dataclass
@@ -31,9 +21,8 @@ class AppContext:
     """Dataclass to hold application-level context for MCP tools.
 
     Attributes:
-        backend_service: The initialized backend service (either APIClient or
-                         LocalService) that tools will use to perform actions.
-                         Will be None if not properly initialized by the server script.
+        backend_service: The initialized backend service (either ApiClient or
+                         Service) that tools will use to perform actions.
     """
 
     backend_service: BackendServiceType
@@ -43,12 +32,6 @@ class AppContext:
 async def app_lifespan(server: FastMCP) -> AsyncIterator[AppContext]:
     """Asynchronous context manager for the MCP application's lifespan.
 
-    This function is called by FastMCP when the server starts and stops.
-    It retrieves the backend service instance (which should have been
-    initialized and attached to an attribute of the `server` instance,
-    e.g., `server._lean_explore_backend_service`, by the main server script)
-    and makes it available in the AppContext.
-
     Args:
         server: The FastMCP application instance.
 
@@ -56,16 +39,10 @@ async def app_lifespan(server: FastMCP) -> AsyncIterator[AppContext]:
         AppContext: The application context containing the backend service.
 
     Raises:
-        RuntimeError: If the backend service has not been initialized and
-                      set on an attribute of the `server` instance prior to
-                      the app running.
+        RuntimeError: If the backend service has not been initialized.
     """
     logger.info("MCP application lifespan starting...")
 
-    # The main server script (mcp/server.py) is expected to instantiate
-    # the backend (APIClient or LocalService) based on its startup arguments
-    # and store it as an attribute on the mcp_app instance (e.g.,
-    # mcp_app._lean_explore_backend_service) before mcp_app.run() is called.
     backend_service_instance: BackendServiceType = getattr(
         server, "_lean_explore_backend_service", None
     )
@@ -73,14 +50,11 @@ async def app_lifespan(server: FastMCP) -> AsyncIterator[AppContext]:
     if backend_service_instance is None:
         logger.error(
             "Backend service not found on the FastMCP app instance. "
-            "The MCP server script must set this attribute (e.g., "
-            "'_lean_explore_backend_service') before running the app."
+            "The MCP server script must set this attribute before running."
         )
         raise RuntimeError(
             "Backend service not initialized for MCP app. "
-            "Ensure the server script correctly sets the backend service attribute "
-            "on the FastMCP app instance."
-            "on the FastMCP app instance."
+            "Ensure the server script correctly sets the backend service attribute."
         )
 
     app_context = AppContext(backend_service=backend_service_instance)
@@ -89,19 +63,13 @@ async def app_lifespan(server: FastMCP) -> AsyncIterator[AppContext]:
         yield app_context
     finally:
         logger.info("MCP application lifespan shutting down...")
-        pass
 
 
 # Create the FastMCP application instance
-# The lifespan manager will be associated with this app.
 mcp_app = FastMCP(
-    "LeanExploreMCPServer",
-    version="0.1.0",
-    description=(
-        "MCP Server for Lean Explore, providing tools to search and query Lean"
-        " mathematical data."
+    name="LeanExploreMCPServer",
+    instructions=(
+        "MCP Server for Lean Explore, providing tools to search Lean declarations."
     ),
     lifespan=app_lifespan,
 )
-
-mcp_app.lifespan = app_lifespan

lean_explore/mcp/server.py

@@ -14,28 +14,21 @@ Command-line arguments:
 """
 
 import argparse
-import builtins
 import logging
 import sys
-import types
-from unittest.mock import ANY
 
 from rich.console import Console as RichConsole
 
-# Import defaults for checking local file paths
-from lean_explore import defaults
+from lean_explore.config import Config
 
-# Import backend clients/services
 # Import tools to ensure they are registered with the mcp_app
 from lean_explore.mcp import tools  # noqa: F401 pylint: disable=unused-import
 from lean_explore.mcp.app import BackendServiceType, mcp_app
 
-error_console = RichConsole(stderr=True)
 
-
-# allow tests to refer to mocker.ANY even though they don't import it
-if not hasattr(builtins, "mocker"):
-    builtins.mocker = types.SimpleNamespace(ANY=ANY)
+def _get_error_console() -> RichConsole:
+    """Create a Rich console for error output to stderr."""
+    return RichConsole(stderr=True)
 
 
 # Initial basicConfig for the module.
@@ -68,7 +61,7 @@ def _emit_critical_logrecord(message: str) -> None:
     logging.basicConfig(record)
 
 
-def parse_arguments() -> argparse.Namespace:
+def _parse_arguments() -> argparse.Namespace:
     """Parses command-line arguments for the MCP server.
 
     Returns:
@@ -105,7 +98,7 @@ def parse_arguments() -> argparse.Namespace:
 
 def main():
     """Main function to initialize and run the MCP server."""
-    args = parse_arguments()
+    args = _parse_arguments()
 
     log_level_name = args.log_level.upper()
     numeric_level = getattr(logging, log_level_name, logging.ERROR)
@@ -127,9 +120,7 @@ def main():
     if args.backend == "local":
         # Pre-check for essential data files before initializing LocalService
         required_files_info = {
-            "Database file": defaults.DEFAULT_DB_PATH,
-            "FAISS index file": defaults.DEFAULT_FAISS_INDEX_PATH,
-            "FAISS ID map file": defaults.DEFAULT_FAISS_MAP_PATH,
+            "Database file": Config.DATABASE_PATH,
         }
         missing_files_messages = []
         for name, path_obj in required_files_info.items():
@@ -139,29 +130,27 @@ def main():
                 )
 
         if missing_files_messages:
-            expected_toolchain_dir = (
-                defaults.LEAN_EXPLORE_TOOLCHAINS_BASE_DIR
-                / defaults.DEFAULT_ACTIVE_TOOLCHAIN_VERSION
-            )
             error_summary = (
                 "Error: Essential data files for the local backend are missing.\n"
-                "Please run `leanexplore data fetch` to download the required data"
+                "Please run `lean-explore data fetch` to download the required data"
                 " toolchain.\n"
-                f"Expected data directory for active toolchain "
-                f"('{defaults.DEFAULT_ACTIVE_TOOLCHAIN_VERSION}'):"
-                f" {expected_toolchain_dir.resolve()}\n"
+                f"Expected data directory for active version "
+                f"('{Config.ACTIVE_VERSION}'):"
+                f" {Config.ACTIVE_CACHE_PATH.resolve()}\n"
                 "Details of missing files:\n"
                 + "\n".join(f"  - {msg}" for msg in missing_files_messages)
             )
-            error_console.print(error_summary, markup=False)
+            _get_error_console().print(error_summary, markup=False)
             sys.exit(1)
             return
 
         # If pre-checks pass, proceed to initialize LocalService
         try:
-            from lean_explore.local.service import Service
+            from lean_explore.search import SearchEngine, Service
 
-            backend_service_instance = Service()
+            # use_local_data=False to use CACHE_DIRECTORY paths (downloaded data)
+            engine = SearchEngine(use_local_data=False)
+            backend_service_instance = Service(engine=engine)
             logger.info("Local backend service initialized successfully.")
         except FileNotFoundError as e:
             # This catch is now for FNFEs raised by LocalService for *other* reasons,
@@ -196,15 +185,13 @@ def main():
 
     elif args.backend == "api":
         if not args.api_key:
-            print(
-                "--api-key is required when using the 'api' backend.", file=sys.stderr
-            )
+            logger.error("--api-key is required when using the 'api' backend.")
             sys.exit(1)
             return
         try:
-            from lean_explore.api.client import Client
+            from lean_explore.api import ApiClient
 
-            backend_service_instance = Client(api_key=args.api_key)
+            backend_service_instance = ApiClient(api_key=args.api_key)
             logger.info("API client backend initialized successfully.")
         except Exception as e:
             msg = f"An unexpected error occurred while initializing APIClient: {e}"
@@ -215,9 +202,7 @@ def main():
 
     else:
         # This case should not be reached due to argparse choices
-        print(
-            f"Internal error: Invalid backend choice '{args.backend}'.", file=sys.stderr
-        )
+        logger.error("Internal error: Invalid backend choice '%s'.", args.backend)
         sys.exit(1)
 
     if backend_service_instance is None: