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

lean_explore/config.py

@@ -0,0 +1,244 @@
+# src/lean_explore/config.py
+
+"""Centralized configuration for lean_explore.
+
+This module provides all configuration settings including paths, URLs,
+and other constants used throughout the application.
+"""
+
+import os
+import pathlib
+
+
+def _get_active_cache_version() -> str:
+    """Get the active cache version from the version file or environment.
+
+    The version is determined by (in order of priority):
+    1. LEAN_EXPLORE_VERSION environment variable
+    2. Contents of ~/.lean_explore/active_version file (set by data fetch)
+    3. Default fallback version
+
+    Returns:
+        The active version string (e.g., "2025.01.27" or "v4.24.0").
+    """
+    env_version = os.getenv("LEAN_EXPLORE_VERSION")
+    if env_version:
+        return env_version
+
+    version_file = pathlib.Path.home() / ".lean_explore" / "active_version"
+    if version_file.exists():
+        return version_file.read_text().strip()
+
+    return "v4.24.0"
+
+
+def _get_data_directory() -> pathlib.Path:
+    """Get the data directory path."""
+    return pathlib.Path(
+        os.getenv(
+            "LEAN_EXPLORE_DATA_DIR",
+            pathlib.Path(__file__).parent.parent.parent / "data",
+        )
+    )
+
+
+def _get_timestamped_directories(data_directory: pathlib.Path) -> list[pathlib.Path]:
+    """Get all timestamped extraction directories sorted by name descending."""
+    import re
+
+    if not data_directory.exists():
+        return []
+
+    timestamp_pattern = re.compile(r"^\d{8}_\d{6}$")
+    timestamped_directories = [
+        directory
+        for directory in data_directory.iterdir()
+        if directory.is_dir() and timestamp_pattern.match(directory.name)
+    ]
+
+    timestamped_directories.sort(key=lambda d: d.name, reverse=True)
+    return timestamped_directories
+
+
+def _resolve_active_data_path(
+    data_directory: pathlib.Path, active_version: str
+) -> pathlib.Path:
+    """Resolve the active data path using the best available source.
+
+    Priority:
+    1. DATA_DIRECTORY if it contains lean_explore.db directly
+    2. Most recent timestamped extraction directory (YYYYMMDD_HHMMSS)
+    3. DATA_DIRECTORY / ACTIVE_VERSION as fallback
+    """
+    if (data_directory / "lean_explore.db").exists():
+        return data_directory
+
+    timestamped_dirs = _get_timestamped_directories(data_directory)
+    if timestamped_dirs:
+        latest = timestamped_dirs[0]
+        if (latest / "lean_explore.db").exists():
+            return latest
+
+    return data_directory / active_version
+
+
+class Config:
+    """Application-wide configuration settings."""
+
+    CACHE_DIRECTORY: pathlib.Path = pathlib.Path(
+        os.getenv(
+            "LEAN_EXPLORE_CACHE_DIR",
+            pathlib.Path.home() / ".lean_explore" / "cache",
+        )
+    )
+    """Cache directory for downloaded data (used by search engine and MCP server).
+
+    Can be overridden with LEAN_EXPLORE_CACHE_DIR environment variable.
+    Default: ~/.lean_explore/cache
+    """
+
+    DATA_DIRECTORY: pathlib.Path = _get_data_directory()
+    """Local data directory for extraction pipeline output.
+
+    Can be overridden with LEAN_EXPLORE_DATA_DIR environment variable.
+    Default: <repo-root>/data
+    """
+
+    DEFAULT_LEAN_VERSION: str = "4.24.0"
+    """Lean version for database naming and dependency resolution."""
+
+    ACTIVE_VERSION: str = _get_active_cache_version()
+    """Active version identifier for cached data (e.g., "2025.01.27").
+
+    Determined by LEAN_EXPLORE_VERSION env var, ~/.lean_explore/active_version
+    file, or defaults to v4.24.0.
+    """
+
+    ACTIVE_CACHE_PATH: pathlib.Path = CACHE_DIRECTORY / ACTIVE_VERSION
+    """Directory for the active version's cached data files."""
+
+    ACTIVE_DATA_PATH: pathlib.Path = _resolve_active_data_path(
+        DATA_DIRECTORY, ACTIVE_VERSION
+    )
+    """Directory for the active version's local data files.
+
+    Resolved using (in priority order):
+    1. DATA_DIRECTORY if it contains lean_explore.db directly
+    2. Most recent timestamped extraction directory (YYYYMMDD_HHMMSS)
+    3. DATA_DIRECTORY / ACTIVE_VERSION as fallback
+    """
+
+    # =========================================================================
+    # Timestamped Extraction Directory Methods
+    # =========================================================================
+
+    @staticmethod
+    def _get_timestamped_directories() -> list[pathlib.Path]:
+        """Get all timestamped extraction directories sorted by name descending."""
+        return _get_timestamped_directories(Config.DATA_DIRECTORY)
+
+    @staticmethod
+    def get_latest_extraction_path() -> pathlib.Path | None:
+        """Get the most recent timestamped extraction directory.
+
+        Looks for directories matching YYYYMMDD_HHMMSS pattern in DATA_DIRECTORY.
+
+        Returns:
+            Path to most recent extraction directory, or None if none exist.
+        """
+        timestamped_directories = Config._get_timestamped_directories()
+        return timestamped_directories[0] if timestamped_directories else None
+
+    DATABASE_PATH: pathlib.Path = ACTIVE_CACHE_PATH / "lean_explore.db"
+    """Path to SQLite database file in cache (used by search engine)."""
+
+    FAISS_INDEX_PATH: pathlib.Path = ACTIVE_CACHE_PATH / "informalization_faiss.index"
+    """Path to FAISS index file in cache (using informalization embeddings)."""
+
+    FAISS_IDS_MAP_PATH: pathlib.Path = (
+        ACTIVE_CACHE_PATH / "informalization_faiss_ids_map.json"
+    )
+    """Path to FAISS ID mapping file in cache."""
+
+    BM25_SPACED_PATH: pathlib.Path = ACTIVE_CACHE_PATH / "bm25_name_spaced"
+    """Path to BM25 spaced tokenization index directory in cache."""
+
+    BM25_RAW_PATH: pathlib.Path = ACTIVE_CACHE_PATH / "bm25_name_raw"
+    """Path to BM25 raw tokenization index directory in cache."""
+
+    BM25_IDS_MAP_PATH: pathlib.Path = ACTIVE_CACHE_PATH / "bm25_ids_map.json"
+    """Path to BM25 ID mapping file in cache."""
+
+    DATABASE_URL: str = f"sqlite+aiosqlite:///{DATABASE_PATH}"
+    """Async SQLAlchemy database URL for SQLite (used by search engine)."""
+
+    EXTRACTION_DATABASE_PATH: pathlib.Path = ACTIVE_DATA_PATH / "lean_explore.db"
+    """Path to SQLite database file in data directory (used by extraction)."""
+
+    EXTRACTION_DATABASE_URL: str = f"sqlite+aiosqlite:///{EXTRACTION_DATABASE_PATH}"
+    """Async SQLAlchemy database URL for extraction pipeline."""
+
+    @staticmethod
+    def get_latest_database_path() -> pathlib.Path | None:
+        """Get the path to the most recent extraction database.
+
+        Returns:
+            Path to lean_explore.db in the most recent extraction, or None.
+        """
+        latest = Config.get_latest_extraction_path()
+        if latest:
+            database_path = latest / "lean_explore.db"
+            if database_path.exists():
+                return database_path
+        return None
+
+    @staticmethod
+    def create_timestamped_extraction_path() -> pathlib.Path:
+        """Create a new timestamped extraction directory.
+
+        Returns:
+            Path to the newly created directory (YYYYMMDD_HHMMSS format).
+        """
+        from datetime import datetime
+
+        timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
+        extraction_path = Config.DATA_DIRECTORY / timestamp
+        extraction_path.mkdir(parents=True, exist_ok=True)
+        return extraction_path
+
+    # =========================================================================
+    # Package Workspace Paths
+    # =========================================================================
+
+    PACKAGES_ROOT: pathlib.Path = pathlib.Path(
+        os.getenv(
+            "LEAN_EXPLORE_PACKAGES_ROOT",
+            pathlib.Path(__file__).parent.parent.parent / "lean",
+        )
+    )
+    """Root directory for per-package Lean workspaces.
+
+    Can be overridden with LEAN_EXPLORE_PACKAGES_ROOT environment variable.
+    Default: <repo-root>/lean
+    """
+
+    EXTRACT_PACKAGES: set[str] = {
+        "batteries",
+        "init",
+        "lean4",
+        "mathlib",
+        "physlean",
+        "std",
+    }
+    """Set of package names to extract from doc-gen4 output."""
+
+    MANIFEST_URL: str = (
+        "https://pub-48b75babc4664808b15520033423c765.r2.dev/manifest.json"
+    )
+    """Remote URL for the data toolchain manifest."""
+
+    R2_ASSETS_BASE_URL: str = "https://pub-48b75babc4664808b15520033423c765.r2.dev"
+    """Base URL for Cloudflare R2 asset storage."""
+
+    API_BASE_URL: str = "https://www.leanexplore.com/api/v2"
+    """Base URL for the LeanExplore remote API service."""

lean_explore/extract/__init__.py

@@ -0,0 +1,5 @@
+"""Lean declaration extraction and processing tools.
+
+This package contains modules for extracting, parsing, and enriching
+Lean mathematical declarations from documentation files.
+"""

lean_explore/extract/__main__.py

@@ -0,0 +1,368 @@
+"""Pipeline orchestration for Lean declaration extraction and enrichment.
+
+This module provides functions to coordinate the complete data extraction pipeline:
+1. Extract declarations from doc-gen4 output
+2. Generate informal natural language descriptions
+3. Generate vector embeddings for semantic search
+4. Build FAISS indices for vector similarity search
+"""
+
+import asyncio
+import logging
+import os
+from pathlib import Path
+
+import click
+from sqlalchemy.ext.asyncio import AsyncEngine, create_async_engine
+
+from lean_explore.config import Config
+from lean_explore.models import Base
+from lean_explore.util.logging import setup_logging
+
+logger = logging.getLogger(__name__)
+
+
+async def _create_database_schema(engine: AsyncEngine) -> None:
+    """Create database tables if they don't exist.
+
+    Args:
+        engine: SQLAlchemy async engine instance.
+    """
+    logger.info("Creating database schema...")
+    async with engine.begin() as connection:
+        await connection.run_sync(Base.metadata.create_all)
+    logger.info("Database schema created successfully")
+
+
+async def _run_doc_gen4_step(fresh: bool = False) -> None:
+    """Run doc-gen4 to generate documentation.
+
+    Args:
+        fresh: Clear cached dependencies to force fresh resolution.
+    """
+    from lean_explore.extract.doc_gen4 import run_doc_gen4
+
+    logger.info("Running doc-gen4...")
+    await run_doc_gen4(fresh=fresh)
+    logger.info("doc-gen4 complete")
+
+
+async def _run_extract_step(engine: AsyncEngine) -> None:
+    """Extract declarations from doc-gen4 output."""
+    from lean_explore.extract.doc_parser import extract_declarations
+
+    logger.info("Step 1: Extracting declarations from doc-gen4...")
+    await extract_declarations(engine)
+    logger.info("Declaration extraction complete")
+
+
+async def _run_informalize_step(
+    engine: AsyncEngine,
+    model: str,
+    batch_size: int,
+    max_concurrent: int,
+    limit: int | None,
+) -> None:
+    """Generate informal descriptions for declarations."""
+    from lean_explore.extract.informalize import informalize_declarations
+
+    logger.info("Step 2: Generating informal descriptions...")
+    await informalize_declarations(
+        engine,
+        model=model,
+        commit_batch_size=batch_size,
+        max_concurrent=max_concurrent,
+        limit=limit,
+    )
+    logger.info("Informalization complete")
+
+
+async def _run_embeddings_step(
+    engine: AsyncEngine,
+    model_name: str,
+    batch_size: int,
+    limit: int | None,
+    max_seq_length: int,
+) -> None:
+    """Generate embeddings for all declaration fields."""
+    from lean_explore.extract.embeddings import generate_embeddings
+
+    logger.info("Step 3: Generating embeddings...")
+    await generate_embeddings(
+        engine,
+        model_name=model_name,
+        batch_size=batch_size,
+        limit=limit,
+        max_seq_length=max_seq_length,
+    )
+    logger.info("Embedding generation complete")
+
+
+async def _run_index_step(engine: AsyncEngine, extraction_path: Path) -> None:
+    """Build search indices (FAISS and BM25).
+
+    Args:
+        engine: SQLAlchemy async engine instance.
+        extraction_path: Directory to save indices (same as database location).
+    """
+    from lean_explore.extract.index import build_bm25_indices, build_faiss_indices
+
+    logger.info("Step 4: Building search indices...")
+    await build_faiss_indices(engine, output_directory=extraction_path)
+    await build_bm25_indices(engine, output_directory=extraction_path)
+    logger.info("Index building complete")
+
+
+async def run_pipeline(
+    database_url: str,
+    extraction_path: Path,
+    run_doc_gen4: bool = False,
+    fresh: bool = False,
+    parse_docs: bool = True,
+    informalize: bool = True,
+    embeddings: bool = True,
+    index: bool = True,
+    informalize_model: str = "google/gemini-3-flash-preview",
+    informalize_batch_size: int = 1000,
+    informalize_max_concurrent: int = 100,
+    informalize_limit: int | None = None,
+    embedding_model: str = "Qwen/Qwen3-Embedding-0.6B",
+    embedding_batch_size: int = 250,
+    embedding_limit: int | None = None,
+    embedding_max_seq_length: int = 512,
+    verbose: bool = False,
+) -> None:
+    """Run the Lean declaration extraction and enrichment pipeline.
+
+    Args:
+        database_url: SQLite database URL (e.g., sqlite+aiosqlite:///path/to/db)
+        extraction_path: Directory containing the extraction (for saving indices).
+        run_doc_gen4: Run doc-gen4 to generate documentation before parsing
+        fresh: Clear cached dependencies to force fresh resolution (for nightly updates)
+        parse_docs: Run doc-gen4 parsing step
+        informalize: Run informalization step
+        embeddings: Run embeddings generation step
+        index: Run FAISS index building step
+        informalize_model: LLM model for generating informalizations
+        informalize_batch_size: Commit batch size for informalization
+        informalize_max_concurrent: Maximum concurrent informalization requests
+        informalize_limit: Limit number of declarations to informalize
+        embedding_model: Sentence transformer model for embeddings
+        embedding_batch_size: Batch size for embedding generation
+        embedding_limit: Limit number of declarations for embeddings
+        embedding_max_seq_length: Max sequence length for embeddings (lower=less mem)
+        verbose: Enable verbose logging
+    """
+    setup_logging(verbose)
+
+    # Validate OpenRouter API key if informalization is needed
+    if informalize:
+        if not os.getenv("OPENROUTER_API_KEY"):
+            logger.error(
+                "OPENROUTER_API_KEY environment variable is required for "
+                "informalization"
+            )
+            raise RuntimeError("OPENROUTER_API_KEY not set")
+
+    steps_enabled = []
+    if parse_docs:
+        steps_enabled.append("parse-docs")
+    if informalize:
+        steps_enabled.append("informalize")
+    if embeddings:
+        steps_enabled.append("embeddings")
+    if index:
+        steps_enabled.append("index")
+
+    logger.info("Starting Lean Explore extraction pipeline")
+    logger.info(f"Database URL: {database_url}")
+    logger.info(f"Steps to run: {', '.join(steps_enabled)}")
+
+    engine = create_async_engine(database_url, echo=verbose)
+
+    try:
+        await _create_database_schema(engine)
+
+        if run_doc_gen4:
+            await _run_doc_gen4_step(fresh=fresh)
+
+        if parse_docs:
+            await _run_extract_step(engine)
+
+        if informalize:
+            await _run_informalize_step(
+                engine,
+                informalize_model,
+                informalize_batch_size,
+                informalize_max_concurrent,
+                informalize_limit,
+            )
+
+        if embeddings:
+            await _run_embeddings_step(
+                engine,
+                embedding_model,
+                embedding_batch_size,
+                embedding_limit,
+                embedding_max_seq_length,
+            )
+
+        if index:
+            await _run_index_step(engine, extraction_path)
+
+        logger.info("Pipeline completed successfully!")
+
+    finally:
+        await engine.dispose()
+
+
+@click.command()
+@click.option(
+    "--run-doc-gen4",
+    is_flag=True,
+    help="Run doc-gen4 to generate documentation before parsing",
+)
+@click.option(
+    "--fresh",
+    is_flag=True,
+    help="Clear cached dependencies to fetch latest versions (use for nightly updates)",
+)
+@click.option(
+    "--parse-docs/--no-parse-docs",
+    default=None,
+    help="Run doc-gen4 parsing step (creates new timestamped directory)",
+)
+@click.option(
+    "--informalize/--no-informalize",
+    default=None,
+    help="Run informalization step (uses latest extraction)",
+)
+@click.option(
+    "--embeddings/--no-embeddings",
+    default=None,
+    help="Run embeddings generation step (uses latest extraction)",
+)
+@click.option(
+    "--index/--no-index",
+    default=None,
+    help="Run FAISS index building step (uses latest extraction)",
+)
+@click.option(
+    "--informalize-model",
+    default="google/gemini-3-flash-preview",
+    help="LLM model for generating informalizations",
+)
+@click.option(
+    "--informalize-max-concurrent",
+    type=int,
+    default=100,
+    help="Maximum concurrent informalization requests",
+)
+@click.option(
+    "--informalize-limit",
+    type=int,
+    default=None,
+    help="Limit number of declarations to informalize (for testing)",
+)
+@click.option(
+    "--embedding-model",
+    default="Qwen/Qwen3-Embedding-0.6B",
+    help="Sentence transformer model for embeddings",
+)
+@click.option(
+    "--embedding-batch-size",
+    type=int,
+    default=250,
+    help="Batch size for embedding generation (lower = less memory, default 250)",
+)
+@click.option(
+    "--embedding-limit",
+    type=int,
+    default=None,
+    help="Limit number of declarations for embeddings (for testing)",
+)
+@click.option(
+    "--embedding-max-seq-length",
+    type=int,
+    default=512,
+    help="Max sequence length for embeddings (lower = less memory, default 512)",
+)
+@click.option("--verbose", is_flag=True, help="Enable verbose logging")
+def main(
+    run_doc_gen4: bool,
+    fresh: bool,
+    parse_docs: bool | None,
+    informalize: bool | None,
+    embeddings: bool | None,
+    index: bool | None,
+    informalize_model: str,
+    informalize_max_concurrent: int,
+    informalize_limit: int | None,
+    embedding_model: str,
+    embedding_batch_size: int,
+    embedding_limit: int | None,
+    embedding_max_seq_length: int,
+    verbose: bool,
+) -> None:
+    """Run the Lean declaration extraction and enrichment pipeline.
+
+    Extraction creates timestamped directories (YYYYMMDD_HHMMSS format).
+    Subsequent steps (informalize, embeddings, index) use the latest extraction.
+    """
+    # Determine if any flags were explicitly set (including --run-doc-gen4)
+    step_flags = [run_doc_gen4, parse_docs, informalize, embeddings, index]
+    any_flag_explicitly_set = run_doc_gen4 or any(
+        flag is not None for flag in step_flags[1:]
+    )
+
+    # If no flags were explicitly set, run all pipeline steps by default
+    # Otherwise, only run what was explicitly requested
+    if not any_flag_explicitly_set:
+        parse_docs = informalize = embeddings = index = True
+    else:
+        parse_docs = parse_docs if parse_docs is not None else False
+        informalize = informalize if informalize is not None else False
+        embeddings = embeddings if embeddings is not None else False
+        index = index if index is not None else False
+
+    # Determine extraction directory
+    if parse_docs:
+        # Create new timestamped directory for fresh extraction
+        extraction_path = Config.create_timestamped_extraction_path()
+        logger.info(f"Created new extraction directory: {extraction_path}")
+    else:
+        # Use latest existing extraction for subsequent steps
+        extraction_path = Config.get_latest_extraction_path()
+        if extraction_path is None:
+            raise click.ClickException(
+                "No existing extraction found. Run with --parse-docs first."
+            )
+        logger.info(f"Using existing extraction: {extraction_path}")
+
+    database_path = extraction_path / "lean_explore.db"
+    database_url = f"sqlite+aiosqlite:///{database_path}"
+
+    asyncio.run(
+        run_pipeline(
+            database_url=database_url,
+            extraction_path=extraction_path,
+            run_doc_gen4=run_doc_gen4,
+            fresh=fresh,
+            parse_docs=parse_docs,
+            informalize=informalize,
+            embeddings=embeddings,
+            index=index,
+            informalize_model=informalize_model,
+            informalize_max_concurrent=informalize_max_concurrent,
+            informalize_limit=informalize_limit,
+            embedding_model=embedding_model,
+            embedding_batch_size=embedding_batch_size,
+            embedding_limit=embedding_limit,
+            embedding_max_seq_length=embedding_max_seq_length,
+            verbose=verbose,
+        )
+    )
+
+
+if __name__ == "__main__":
+    main()