Merge pull request #1 from ray5th/claude/intelligent-varahamihira-81b22c

.env.example

@@ -1 +0,0 @@
-ANTHROPIC_API_KEY=your_key_here

.gitignore

@@ -1,9 +0,0 @@
-.env
-.venv/
-venv/
-data/mathlib.index
-data/mathlib_meta.pkl
-__pycache__/
-*.pyc
-*.pyo
-.DS_Store

README.md

@@ -1,107 +1,334 @@
-# Agentic Lean 4 Theorem Prover
+## Brief Non-Technical MVP: AI Lean 4 Proof Assistant
 
-An agentic LLM system that iteratively generates and validates formal Lean 4 proofs using a Claude-powered feedback loop, grounded in Mathlib4 via FAISS RAG.
+### Project goal
 
-## Features
+Build an AI assistant that helps users write correct **Lean 4 mathematical proofs**.
 
-- 🤖 **Agentic loop** — Claude generates proofs, Lean validates them, errors are fed back for self-correction
-- 📚 **Mathlib4 RAG** — FAISS semantic search retrieves relevant verified lemmas to ground each generation
-- 🔁 **Self-correcting** — Parses compiler errors and re-prompts until successful compilation or max retries
-- 🖥️ **Interactive CLI** — Pretty terminal output with `rich`, syntax highlighting, live feedback
+The user provides a theorem. The system searches Lean’s Mathlib library, finds useful related theorems, asks an AI model to write a proof, checks the proof using Lean, and automatically retries if the proof fails.
 
-## Setup
+The key idea is:
 
-### Prerequisites
+```text
+AI suggests the proof.
+Lean verifies the proof.
+Only Lean decides if the answer is correct.
+```
+
+---
+
+## MVP flow
+
+```text
+User enters theorem
+        ↓
+Lean tooling reads the theorem and current proof goal
+        ↓
+System retrieves useful Mathlib lemmas
+        ↓
+System reranks the lemmas and keeps the best ones
+        ↓
+LangChain organizes the prompt and retrieval flow
+        ↓
+Local LLM generates a Lean proof
+        ↓
+Lean checks the proof
+        ↓
+If correct → return final proof
+If wrong → send Lean error back to the AI and retry
+```
+
+---
+
+## Simple explanation
+
+This project is like an **AI coding assistant for mathematical proofs**.
+
+But unlike a normal chatbot, it does not just guess. It uses Lean itself to check whether the proof is truly valid.
+
+The assistant works in three stages:
+
+```text
+1. Search
+Find useful existing Mathlib theorems.
+
+2. Generate
+Ask the AI model to write a proof using those theorems.
+
+3. Verify
+Run the proof through Lean and retry if Lean finds an error.
+```
+
+---
+
+## Main MVP components
+
+### 1. Lean tooling
+
+Use **LeanInteract** to connect Python with Lean 4.
+
+LeanInteract lets Python interact with Lean through the Lean REPL, which means your program can send Lean code, inspect proof goals, and receive Lean feedback programmatically. ([GitHub][1])
+
+In simple terms:
+
+```text
+LeanInteract lets the AI system talk to Lean.
+```
+
+Instead of only seeing an error like:
+
+```text
+type mismatch
+```
+
+the system can see more useful information like:
+
+```text
+Current goal:
+n : Nat
+⊢ n = n
+```
+
+That makes the retry much smarter.
+
+---
+
+### 2. Retrieve Mathlib lemmas
+
+The system searches Mathlib for useful theorems.
+
+For example, if the theorem is about even numbers, the retriever may find lemmas related to:
+
+```text
+Even
+Odd
+Nat.even_or_odd
+parity
+multiplication
+powers
+```
+
+This is the “retrieve” part.
+
+Use:
+
+```text
+FAISS for fast semantic search
+BM25 / keyword search for exact Lean names
+```
+
+FAISS is useful for quickly finding similar text from a large collection of embedded documents, while keyword search helps when exact theorem names matter.
+
+---
+
+### 3. Rerank retrieved lemmas
+
+Retrieval alone is not enough.
+
+The system may retrieve 50 possible Mathlib lemmas, but many may be weak or irrelevant. So we add a **reranker**.
 
-- Python 3.11+
-- [Lean 4 + elan](https://leanprover.github.io/lean4/doc/setup.html) installed and `lean` on your `$PATH`
-- An Anthropic API key
+The improved flow is:
 
-### Install
+```text
+Retrieve top 50 lemmas
+        ↓
+Rerank them based on the current theorem and proof goal
+        ↓
+Keep the best 8–12 lemmas
+        ↓
+Send only those to the AI model
+```
 
-```bash
-git clone https://github.com/YOUR_HANDLE/lean4-helper
-cd lean4-helper
+A Cross-Encoder reranker is useful here because it scores each query-document pair more carefully than the first-pass retriever. Sentence Transformers describes this as a common retrieve-and-rerank setup where fast retrieval finds candidates and the Cross-Encoder reranks them for better relevance. ([SentenceTransformers][2])
 
-python -m venv .venv && source .venv/bin/activate
-pip install -r requirements.txt
+In simple terms:
 
-cp .env.example .env
-# Edit .env and add your ANTHROPIC_API_KEY
+```text
+Retrieve finds possible useful lemmas.
+Rerank chooses the best ones.
 ```
 
-## Usage
+This improves the quality of the context given to the AI.
+
+---
+
+### 4. LangChain role
 
-### Step 1: Build the FAISS Index (one-time)
+Use **LangChain** for orchestration, not for proof correctness.
 
-First, extract lemmas from your local Mathlib4 source:
+LangChain can help organize:
 
-```bash
-python extract_mathlib.py --mathlib-path ~/.elan/toolchains/leanprover-lean4-v4.X.0/lib/lean/library
+```text
+retrieval
+reranking
+prompt templates
+document formatting
+LLM calls
+retry history
 ```
 
-Then build the FAISS index:
+LangChain has retrieval utilities such as contextual compression, where retrieved documents can be filtered or shortened before sending them to the model. ([LangChain][3])
 
-```bash
-python -m src.retriever
+But LangChain does **not** make Lean proofs correct by itself.
+
+The correctness still comes from:
+
+```text
+Lean compiler
+Lean proof state
+Mathlib retrieval quality
+LLM proof generation
+retry loop
 ```
 
-### Step 2: Run the Prover
+So the right positioning is:
 
-```bash
-# Interactive mode
-python -m src.cli
+```text
+LangChain manages the workflow.
+Lean verifies the truth.
+```
 
-# Inline theorem
-python -m src.cli --theorem "theorem add_comm (a b : Nat) : a + b = b + a := by sorry"
+---
+
+### 5. Local LLM first
+
+For the MVP, use a local model first.
+
+Example:
+
+```text
+Ollama + Qwen Coder
+```
 
-# From a .lean file
-python -m src.cli --file my_theorem.lean
+Why local first?
 
-# Without RAG (faster, less accurate)
-python -m src.cli --no-rag --max-retries 8
+```text
+No API cost
+Easy to test many times
+Good for debugging the pipeline
+Private and simple for development
 ```
 
-## Project Structure
+Later, once the workflow works, replace the local model with Claude.
 
+The architecture stays the same.
+
+---
+
+### 6. Lean verification and retry
+
+After the AI generates a proof, the system sends it to Lean.
+
+If Lean accepts the proof:
+
+```text
+Success → return proof
 ```
-lean4-helper/
-├── src/
-│   ├── __init__.py
-│   ├── lean_env.py      # Lean 4 compiler wrapper & error parser
-│   ├── retriever.py     # FAISS index build & retrieval
-│   ├── agent.py         # Agentic Claude feedback loop
-│   └── cli.py           # Rich interactive CLI
-├── data/                # FAISS index + metadata (git-ignored)
-├── extract_mathlib.py   # Mathlib4 lemma extractor
-├── requirements.txt
-├── .env.example
-└── README.md
+
+If Lean rejects the proof:
+
+```text
+Failure → capture Lean error → add error to prompt → retry
 ```
 
-## How It Works
+Example:
+
+```text
+Attempt 1:
+AI generates proof.
+Lean says: unknown theorem name.
+
+Attempt 2:
+AI sees the error and tries a different lemma.
+Lean says: type mismatch.
 
+Attempt 3:
+AI fixes the proof.
+Lean accepts it.
 ```
+
+This is the most important part of the MVP.
+
+---
+
+## Final MVP architecture
+
+```text
 User Theorem
      │
      ▼
-FAISS Retrieval ──► Top-K Mathlib Lemmas
+LeanInteract
+Get theorem goal and Lean feedback
+     │
+     ▼
+Hybrid Retrieval
+FAISS semantic search + keyword search
+     │
+     ▼
+Reranker
+Keep only the most useful Mathlib lemmas
      │
      ▼
-Claude Prompt (theorem + context + history)
+LangChain Prompt Flow
+Organize theorem, lemmas, history, and Lean errors
      │
      ▼
-Generated Lean 4 Proof
+Local LLM
+Generate Lean proof
      │
      ▼
-Lean Compiler
-     ├── ✅ Success ──► Return proof
-     └── ❌ Failure ──► Append error to history ──► Claude (retry)
+Lean Checker
+Verify proof
+     │
+     ├── Success → Return final proof
+     │
+     └── Failure → Retry with Lean error
+```
+
+---
+
+## Libraries to use
+
+| Purpose                      | Library                                           |
+| ---------------------------- | ------------------------------------------------- |
+| Interact with Lean           | `LeanInteract`                                    |
+| Semantic retrieval           | `sentence-transformers`                           |
+| Vector search                | `faiss-cpu`                                       |
+| Keyword search               | `rank-bm25`                                       |
+| Reranking                    | `sentence-transformers` CrossEncoder or FlashRank |
+| Workflow orchestration       | `LangChain`                                       |
+| More advanced workflow graph | `LangGraph`                                       |
+| Local model                  | `Ollama`                                          |
+| Future stronger model        | Claude API                                        |
+
+LangGraph is useful later if you want the proof process to behave like a clear state machine: retrieve → generate → verify → retry → success/failure. LangGraph is designed for stateful agent workflows, which fits this kind of multi-step proof loop. ([Emergent Mind][4])
+
+---
+
+## Non-technical value proposition
+
+This MVP creates an AI assistant that helps people write mathematically verified Lean proofs.
+
+It does not only generate text. It checks every answer with Lean, learns from errors, and keeps trying until it produces a valid proof.
+
+The project demonstrates:
+
+```text
+AI-assisted formal reasoning
+verified code generation
+search over mathematical knowledge
+automatic error correction
+LLM + compiler feedback loops
 ```
 
-## Benchmarks
+---
+
+## One-line MVP summary
+
+```text
+An AI-powered Lean 4 proof assistant that searches Mathlib, reranks useful lemmas, generates a proof with a local LLM, checks it using Lean tooling, and automatically retries using Lean errors until the proof is verified.
+```
 
-| Configuration        | First-attempt success rate |
-|----------------------|---------------------------|
-| Baseline (no RAG)    | ~18%                      |
-| With FAISS RAG       | ~38%                      |
+[1]: https://github.com/augustepoiroux/LeanInteract?utm_source=chatgpt.com "LeanInteract: A Python Interface for Lean 4"
+[2]: https://sbert.net/examples/sentence_transformer/applications/retrieve_rerank/README.html?utm_source=chatgpt.com "Retrieve & Re-Rank Pipeline"
+[3]: https://www.langchain.com/blog/improving-document-retrieval-with-contextual-compression?utm_source=chatgpt.com "Improving Document Retrieval with Contextual Compression"
+[4]: https://www.emergentmind.com/topics/langgraph?utm_source=chatgpt.com "LangGraph: Modular LLM Agent Orchestration"

extract_mathlib.py

@@ -1,82 +0,0 @@
-"""
-extract_mathlib.py
-
-Extracts lemma names, type signatures, and docstrings from Mathlib4 source files
-and writes them to data/mathlib_lemmas.jsonl for indexing.
-
-Usage:
-    python extract_mathlib.py --mathlib-path /path/to/mathlib4
-
-Each output line is a JSON object:
-    {"name": "Nat.add_comm", "type": "∀ (n m : ℕ), n + m = m + n", "doc": "..."}
-"""
-
-import argparse
-import json
-import os
-import re
-from pathlib import Path
-
-# Regex patterns for Lean 4 declarations
-DECL_PATTERN = re.compile(
-    r'(?P<doc>/--.*?-/\s*)?'          # optional docstring
-    r'(?:theorem|lemma|def|abbrev)\s+'
-    r'(?P<name>[\w\.\']+)\s*'
-    r'(?P<rest>[^:=]*?):\s*'
-    r'(?P<type>.+?)\s*(?::=|where|by)',
-    re.DOTALL
-)
-
-DOC_PATTERN = re.compile(r'/--\s*(.*?)\s*-/', re.DOTALL)
-
-
-def extract_from_file(path: Path) -> list[dict]:
-    """Extract lemma entries from a single .lean file."""
-    text = path.read_text(encoding="utf-8", errors="ignore")
-    entries = []
-    for m in DECL_PATTERN.finditer(text):
-        name = m.group("name").strip()
-        raw_type = m.group("type").strip()
-        # Clean up multi-line types
-        clean_type = " ".join(raw_type.split())
-        doc = ""
-        if m.group("doc"):
-            doc_match = DOC_PATTERN.search(m.group("doc"))
-            if doc_match:
-                doc = " ".join(doc_match.group(1).split())
-        entries.append({"name": name, "type": clean_type, "doc": doc})
-    return entries
-
-
-def extract_mathlib(mathlib_path: str, output_path: str = "data/mathlib_lemmas.jsonl"):
-    """Walk all .lean files in mathlib_path and write JSONL output."""
-    root = Path(mathlib_path)
-    os.makedirs(Path(output_path).parent, exist_ok=True)
-
-    total = 0
-    with open(output_path, "w") as out:
-        for lean_file in sorted(root.rglob("*.lean")):
-            entries = extract_from_file(lean_file)
-            for entry in entries:
-                out.write(json.dumps(entry) + "\n")
-            total += len(entries)
-            if total % 1000 == 0:
-                print(f"  Extracted {total} lemmas so far...")
-
-    print(f"Done. {total} lemmas written to {output_path}")
-
-
-if __name__ == "__main__":
-    parser = argparse.ArgumentParser(description="Extract Mathlib4 lemmas to JSONL")
-    parser.add_argument(
-        "--mathlib-path",
-        required=True,
-        help="Path to the root of the Mathlib4 source tree (e.g. ~/.elan/toolchains/...)",
-    )
-    parser.add_argument(
-        "--output",
-        default="data/mathlib_lemmas.jsonl",
-        help="Output JSONL file path.",
-    )
-    args = parser.parse_args()
-    extract_mathlib(args.mathlib_path, args.output)

problems/simple_add.lean

@@ -0,0 +1,4 @@
+import Mathlib
+
+theorem add_zero_simple (n : ℕ) : n + 0 = n := by
+  sorry

problems/test_problem.lean

@@ -0,0 +1,4 @@
+import Mathlib
+
+theorem square_root_two_irrational (n m : ℕ) (h : n^2 = 2 * m^2) (hm : m ≠ 0) : False := by
+  sorry

requirements.txt

@@ -1,6 +1,9 @@
-anthropic>=0.30.0
-faiss-cpu>=1.8.0
-sentence-transformers>=2.7.0
-rich>=13.7.1
-pexpect>=4.9.0
-python-dotenv>=1.0.1
+lean-interact
+ollama
+langchain
+langchain-community
+langchain-ollama
+langgraph
+faiss-cpu
+rank-bm25
+sentence-transformers

scripts/build_index.py

@@ -0,0 +1,40 @@
+#!/usr/bin/env python3
+"""
+One-time script to build the FAISS + BM25 index from Mathlib4 source files.
+Run this before using the LangGraph agent for the first time:
+
+    python scripts/build_index.py
+
+Optional flags:
+    --mathlib-root PATH   Path to Mathlib4 source (auto-detected if omitted)
+    --max-files N         Limit to first N .lean files (useful for quick testing)
+    --index-dir PATH      Where to save the index (default: data/mathlib_index)
+"""
+import sys
+import os
+import argparse
+
+sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), '..', 'src')))
+
+from retriever import MathLibRetriever
+
+
+def main():
+    parser = argparse.ArgumentParser(description="Build Mathlib FAISS index.")
+    parser.add_argument("--mathlib-root", default=None, help="Path to Mathlib4 source root")
+    parser.add_argument("--max-files", type=int, default=None, help="Limit number of .lean files processed")
+    parser.add_argument("--index-dir", default=None, help="Directory to save the index")
+    args = parser.parse_args()
+
+    retriever = MathLibRetriever(index_dir=args.index_dir)
+
+    if retriever.is_index_built() and not args.max_files:
+        print(f"Index already exists at {retriever.index_dir}. Delete it to rebuild.")
+        return
+
+    retriever.build(mathlib_root=args.mathlib_root, max_files=args.max_files)
+    print("Done. Index is ready for use.")
+
+
+if __name__ == "__main__":
+    main()

scripts/run_agent.py

@@ -0,0 +1,36 @@
+#!/usr/bin/env python3
+import sys
+import os
+import argparse
+
+sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), '..', 'src')))
+
+from langgraph_agent import LangGraphAgent
+
+
+def main():
+    parser = argparse.ArgumentParser(description="Run the LangGraph Lean Proof Agent on a file.")
+    parser.add_argument("file", help="Path to the .lean file to solve")
+    parser.add_argument("--model", default="qwen3-vl:4b", help="Ollama model name")
+    parser.add_argument("--retries", type=int, default=5, help="Max retries")
+    parser.add_argument("--index-dir", default=None, help="Path to pre-built FAISS index directory")
+
+    args = parser.parse_args()
+
+    print(f"Starting LangGraph Proof Agent with model: {args.model}")
+    agent = LangGraphAgent(
+        model_name=args.model,
+        max_retries=args.retries,
+        index_dir=args.index_dir,
+    )
+
+    success = agent.solve_file(args.file)
+
+    if success:
+        print("\nSuccess! The proof has been verified.")
+    else:
+        print("\nFailed to verify the proof.")
+
+
+if __name__ == "__main__":
+    main()

src/__init__.py

@@ -1 +0,0 @@
-# lean4-helper src package

src/agent.py

@@ -1,159 +0,0 @@
-"""
-agent.py
-
-The core agentic self-correcting loop.
-
-Given a Lean 4 theorem statement, this agent will:
-  1. Retrieve relevant Mathlib4 lemmas via FAISS.
-  2. Prompt Claude to generate a proof.
-  3. Validate the proof with the Lean compiler.
-  4. On failure, feed the compiler error back to Claude and retry.
-  5. Return the final proof on success or raise after max retries.
-"""
-
-import os
-from typing import Optional
-
-import anthropic
-
-from .lean_env import LeanResult, run_lean
-from .retriever import Retriever
-
-SYSTEM_PROMPT = """\
-You are an expert Lean 4 theorem prover with deep knowledge of Mathlib4.
-Your task is to write a complete, valid Lean 4 proof for the given theorem.
-
-Rules:
-- Output ONLY valid Lean 4 code. No explanation, no markdown, no fences.
-- Begin with the necessary imports (e.g., `import Mathlib`).
-- Use `theorem` or `lemma` keyword as appropriate.
-- Use the provided relevant Mathlib4 lemmas as hints.
-- If a previous attempt failed with a compiler error, fix only the reported issue.
-"""
-
-
-def _build_user_message(
-    theorem: str,
-    context_lemmas: list[dict],
-    previous_error: Optional[str] = None,
-    previous_code: Optional[str] = None,
-) -> str:
-    """
-    Construct the user message for Claude.
-    """
-    parts = []
-
-    # Relevant Mathlib4 lemmas from RAG
-    if context_lemmas:
-        parts.append("## Potentially Relevant Mathlib4 Lemmas\n")
-        for lemma in context_lemmas:
-            name = lemma.get("name", "?")
-            typ = lemma.get("type", "?")
-            doc = lemma.get("doc", "")
-            parts.append(f"- `{name} : {typ}`" + (f" — {doc}" if doc else ""))
-        parts.append("")
-
-    # The theorem to prove
-    parts.append(f"## Theorem to Prove\n```lean\n{theorem}\n```\n")
-
-    # Feedback from previous attempt
-    if previous_code and previous_error:
-        parts.append("## Previous Attempt (FAILED)\n")
-        parts.append(f"```lean\n{previous_code}\n```\n")
-        parts.append(f"**Compiler Error:**\n```\n{previous_error}\n```\n")
-        parts.append("Please fix the error above and output the corrected proof.")
-    else:
-        parts.append("Please generate a complete Lean 4 proof for the theorem above.")
-
-    return "\n".join(parts)
-
-
-class ProverAgent:
-    """
-    Agentic prover that iteratively generates and validates Lean 4 proofs.
-    """
-
-    def __init__(
-        self,
-        retriever: Optional[Retriever] = None,
-        model: str = "claude-opus-4-5",
-        max_retries: int = 5,
-        k_lemmas: int = 5,
-        lean_timeout: int = 60,
-    ):
-        """
-        Args:
-            retriever:    Optional FAISS Retriever. If None, no RAG context is used.
-            model:        Claude model ID to use.
-            max_retries:  Maximum number of proof attempts before giving up.
-            k_lemmas:     Number of Mathlib lemmas to retrieve per query.
-            lean_timeout: Timeout in seconds for the Lean compiler.
-        """
-        api_key = os.environ.get("ANTHROPIC_API_KEY")
-        if not api_key:
-            raise EnvironmentError(
-                "ANTHROPIC_API_KEY is not set. Please add it to your .env file."
-            )
-        self.client = anthropic.Anthropic(api_key=api_key)
-        self.retriever = retriever
-        self.model = model
-        self.max_retries = max_retries
-        self.k_lemmas = k_lemmas
-        self.lean_timeout = lean_timeout
-
-    def prove(self, theorem: str, on_attempt=None) -> tuple[bool, str, int]:
-        """
-        Attempt to prove the given Lean 4 theorem statement.
-
-        Args:
-            theorem:    A Lean 4 theorem/lemma declaration string.
-            on_attempt: Optional callback(attempt, code, result) for streaming UI.
-
-        Returns:
-            (success, final_code, num_attempts)
-        """
-        # RAG: retrieve relevant lemmas
-        context_lemmas = []
-        if self.retriever:
-            context_lemmas = self.retriever.retrieve(theorem, k=self.k_lemmas)
-
-        messages = []
-        previous_code = None
-        previous_error = None
-
-        for attempt in range(1, self.max_retries + 1):
-            user_content = _build_user_message(
-                theorem=theorem,
-                context_lemmas=context_lemmas,
-                previous_error=previous_error,
-                previous_code=previous_code,
-            )
-
-            messages.append({"role": "user", "content": user_content})
-
-            # Call Claude
-            response = self.client.messages.create(
-                model=self.model,
-                max_tokens=2048,
-                system=SYSTEM_PROMPT,
-                messages=messages,
-            )
-            generated_code = response.content[0].text.strip()
-
-            # Append assistant response to conversation history
-            messages.append({"role": "assistant", "content": generated_code})
-
-            # Validate with Lean compiler
-            lean_result: LeanResult = run_lean(generated_code, timeout=self.lean_timeout)
-
-            if on_attempt:
-                on_attempt(attempt, generated_code, lean_result)
-
-            if lean_result.success:
-                return True, generated_code, attempt
-
-            # Feed error back for next iteration
-            previous_code = generated_code
-            previous_error = "\n".join(lean_result.errors) or lean_result.output[:500]
-
-        return False, previous_code or "", self.max_retries

src/cli.py

@@ -1,162 +0,0 @@
-"""
-cli.py
-
-Interactive CLI for the Agentic Lean 4 Theorem Prover.
-
-Usage:
-    python -m src.cli
-    python -m src.cli --theorem "theorem add_comm (a b : Nat) : a + b = b + a := by sorry"
-    python -m src.cli --file theorems.lean
-"""
-
-import argparse
-import os
-import sys
-from pathlib import Path
-
-from dotenv import load_dotenv
-from rich.console import Console
-from rich.markdown import Markdown
-from rich.panel import Panel
-from rich.syntax import Syntax
-from rich.text import Text
-from rich import print as rprint
-from rich.rule import Rule
-from rich.spinner import Spinner
-from rich.live import Live
-from rich import box
-
-from .agent import ProverAgent
-from .lean_env import LeanResult
-
-# Load .env for ANTHROPIC_API_KEY
-load_dotenv()
-
-console = Console()
-
-BANNER = """
-[bold cyan]╔══════════════════════════════════════════════════╗[/bold cyan]
-[bold cyan]║   🧮  Agentic Lean 4 Theorem Prover              ║[/bold cyan]
-[bold cyan]║   Powered by Claude + FAISS + Mathlib4 RAG       ║[/bold cyan]
-[bold cyan]╚══════════════════════════════════════════════════╝[/bold cyan]
-"""
-
-
-def on_attempt_callback(attempt: int, code: str, result: LeanResult):
-    """Rich-formatted callback for each proof attempt."""
-    console.print(Rule(f"[bold]Attempt {attempt}[/bold]", style="dim"))
-    console.print(Syntax(code, "lean", theme="monokai", line_numbers=True))
-
-    if result.success:
-        console.print(Panel("[bold green]✅ Lean accepted this proof![/bold green]", box=box.ROUNDED))
-    else:
-        errors_text = "\n".join(result.errors) if result.errors else result.output[:300]
-        console.print(
-            Panel(
-                f"[bold red]❌ Compiler rejected proof:[/bold red]\n\n[red]{errors_text}[/red]",
-                box=box.ROUNDED,
-            )
-        )
-
-
-def run_prover(theorem: str, use_rag: bool = True, max_retries: int = 5):
-    """
-    Main entry point: set up the agent and run it on the given theorem.
-    """
-    console.print(BANNER)
-    console.print(Panel(f"[bold yellow]Theorem to prove:[/bold yellow]\n\n{theorem}", box=box.ROUNDED))
-
-    # Optionally initialize retriever
-    retriever = None
-    if use_rag:
-        index_path = "data/mathlib.index"
-        meta_path = "data/mathlib_meta.pkl"
-        if os.path.exists(index_path) and os.path.exists(meta_path):
-            from .retriever import Retriever
-            console.print("[dim]Loading FAISS index for Mathlib4 RAG...[/dim]")
-            retriever = Retriever(index_path=index_path, meta_path=meta_path)
-            console.print("[green]✓ FAISS index loaded.[/green]\n")
-        else:
-            console.print(
-                "[yellow]⚠ FAISS index not found. Running without RAG context.[/yellow]\n"
-                "[dim]Run `python -m src.retriever` to build the index first.[/dim]\n"
-            )
-
-    agent = ProverAgent(retriever=retriever, max_retries=max_retries)
-
-    console.print(Rule("[bold]Starting Proof Search[/bold]", style="cyan"))
-    success, final_code, num_attempts = agent.prove(theorem, on_attempt=on_attempt_callback)
-
-    console.print(Rule(style="cyan"))
-    if success:
-        console.print(
-            Panel(
-                f"[bold green]🎉 Proof found in {num_attempts} attempt(s)![/bold green]",
-                box=box.DOUBLE,
-            )
-        )
-        console.print("\n[bold]Final Proof:[/bold]")
-        console.print(Syntax(final_code, "lean", theme="monokai", line_numbers=True))
-    else:
-        console.print(
-            Panel(
-                f"[bold red]💀 Failed to find a proof after {num_attempts} attempts.[/bold red]\n"
-                "[dim]Try increasing --max-retries or refining the theorem statement.[/dim]",
-                box=box.DOUBLE,
-            )
-        )
-        if final_code:
-            console.print("\n[bold]Last Generated Code:[/bold]")
-            console.print(Syntax(final_code, "lean", theme="monokai", line_numbers=True))
-
-
-def main():
-    parser = argparse.ArgumentParser(
-        description="Agentic Lean 4 Theorem Prover powered by Claude + FAISS + Mathlib4"
-    )
-    group = parser.add_mutually_exclusive_group()
-    group.add_argument(
-        "--theorem",
-        type=str,
-        help="A Lean 4 theorem declaration string to prove.",
-    )
-    group.add_argument(
-        "--file",
-        type=str,
-        help="Path to a .lean file containing the theorem to prove.",
-    )
-    parser.add_argument(
-        "--no-rag",
-        action="store_true",
-        help="Disable Mathlib4 RAG retrieval.",
-    )
-    parser.add_argument(
-        "--max-retries",
-        type=int,
-        default=5,
-        help="Maximum number of proof attempts (default: 5).",
-    )
-    args = parser.parse_args()
-
-    if args.file:
-        theorem = Path(args.file).read_text().strip()
-    elif args.theorem:
-        theorem = args.theorem.strip()
-    else:
-        console.print("[bold cyan]Enter your Lean 4 theorem (end with Ctrl+D on a new line):[/bold cyan]")
-        try:
-            lines = sys.stdin.read()
-            theorem = lines.strip()
-        except KeyboardInterrupt:
-            console.print("\n[yellow]Cancelled.[/yellow]")
-            sys.exit(0)
-
-    if not theorem:
-        console.print("[red]No theorem provided. Exiting.[/red]")
-        sys.exit(1)
-
-    run_prover(theorem=theorem, use_rag=not args.no_rag, max_retries=args.max_retries)
-
-
-if __name__ == "__main__":
-    main()

src/langgraph_agent.py

@@ -0,0 +1,173 @@
+import os
+from typing import List, TypedDict
+
+from langgraph.graph import END, StateGraph
+
+from lean_verifier import LeanEnvironment
+from rag_chain import RAGProofChain
+from retriever import MathLibRetriever
+
+
+# ---------------------------------------------------------------------------
+# State
+# ---------------------------------------------------------------------------
+
+class ProofState(TypedDict):
+    file_path: str
+    lean_code: str
+    goals: List[str]
+    errors: List[str]
+    attempt: int
+    max_retries: int
+    status: str          # "pending" | "success" | "failed"
+    retrieved_lemmas: list
+
+
+# ---------------------------------------------------------------------------
+# Nodes
+# ---------------------------------------------------------------------------
+
+def _read_file(path: str) -> str:
+    with open(path, "r") as f:
+        return f.read()
+
+
+def _write_file(path: str, code: str) -> None:
+    with open(path, "w") as f:
+        f.write(code)
+
+
+def _extract_lean_code(text: str) -> str:
+    if "```lean" in text:
+        return text.split("```lean")[1].split("```")[0].strip()
+    if "```" in text:
+        return text.split("```")[1].split("```")[0].strip()
+    return text.strip()
+
+
+def make_verify_node(lean_env: LeanEnvironment):
+    def verify_node(state: ProofState) -> ProofState:
+        print(f"\n--- Attempt {state['attempt'] + 1} / {state['max_retries']} ---")
+        code = _read_file(state["file_path"])
+        result = lean_env.verify_proof(code)
+
+        new_status = "success" if result["status"] == "success" else "pending"
+        if new_status == "success":
+            print("Proof verified successfully!")
+        else:
+            print(
+                f"Verification failed. "
+                f"Errors: {len(result['errors'])}, Goals: {len(result['goals'])}"
+            )
+
+        return {
+            **state,
+            "lean_code": code,
+            "errors": result["errors"],
+            "goals": result["goals"],
+            "status": new_status,
+        }
+    return verify_node
+
+
+def make_retrieve_node(retriever: MathLibRetriever):
+    def retrieve_node(state: ProofState) -> ProofState:
+        query = " ".join(state["goals"] + state["errors"])
+        print("Retrieving relevant Mathlib lemmas…")
+        lemmas = retriever.retrieve(query)
+        print(f"  Retrieved {len(lemmas)} lemma(s).")
+        return {**state, "retrieved_lemmas": lemmas}
+    return retrieve_node
+
+
+def make_generate_node(chain: RAGProofChain):
+    def generate_node(state: ProofState) -> ProofState:
+        print("Generating proof with LLM…")
+        raw = chain.generate(
+            lean_code=state["lean_code"],
+            goals=state["goals"],
+            errors=state["errors"],
+            retrieved_lemmas=state["retrieved_lemmas"],
+        )
+        new_code = _extract_lean_code(raw)
+
+        if not new_code or new_code.strip() == state["lean_code"].strip():
+            print("LLM produced no changes.")
+            return {**state, "attempt": state["attempt"] + 1, "status": "failed"}
+
+        _write_file(state["file_path"], new_code)
+        print("File updated.")
+        return {
+            **state,
+            "lean_code": new_code,
+            "attempt": state["attempt"] + 1,
+        }
+    return generate_node
+
+
+# ---------------------------------------------------------------------------
+# Router
+# ---------------------------------------------------------------------------
+
+def should_continue(state: ProofState) -> str:
+    if state["status"] == "success":
+        return END
+    if state["attempt"] >= state["max_retries"]:
+        return END
+    return "retrieve"
+
+
+# ---------------------------------------------------------------------------
+# Graph assembly
+# ---------------------------------------------------------------------------
+
+def build_graph(lean_env: LeanEnvironment, retriever: MathLibRetriever, chain: RAGProofChain):
+    g = StateGraph(ProofState)
+
+    g.add_node("verify", make_verify_node(lean_env))
+    g.add_node("retrieve", make_retrieve_node(retriever))
+    g.add_node("generate", make_generate_node(chain))
+
+    g.set_entry_point("verify")
+    g.add_conditional_edges("verify", should_continue, {"retrieve": "retrieve", END: END})
+    g.add_edge("retrieve", "generate")
+    g.add_edge("generate", "verify")
+
+    return g.compile()
+
+
+# ---------------------------------------------------------------------------
+# Public entry point
+# ---------------------------------------------------------------------------
+
+class LangGraphAgent:
+    def __init__(
+        self,
+        model_name: str = "qwen3-vl:4b",
+        max_retries: int = 5,
+        index_dir: str | None = None,
+    ):
+        self._lean_env = LeanEnvironment(use_mathlib=True)
+        self._retriever = MathLibRetriever(index_dir=index_dir)
+        self._chain = RAGProofChain(model_name=model_name)
+        self._graph = build_graph(self._lean_env, self._retriever, self._chain)
+        self._max_retries = max_retries
+
+    def solve_file(self, file_path: str) -> bool:
+        if not os.path.exists(file_path):
+            print(f"Error: {file_path} not found.")
+            return False
+
+        initial: ProofState = {
+            "file_path": file_path,
+            "lean_code": "",
+            "goals": [],
+            "errors": [],
+            "attempt": 0,
+            "max_retries": self._max_retries,
+            "status": "pending",
+            "retrieved_lemmas": [],
+        }
+
+        final = self._graph.invoke(initial)
+        return final["status"] == "success"

src/lean_env.py

@@ -1,83 +0,0 @@
-"""
-lean_env.py
-
-Wrapper around the Lean 4 compiler. Sends Lean 4 code to the compiler,
-captures stdout/stderr, and parses the output to determine success or failure.
-"""
-
-import subprocess
-import tempfile
-import os
-import re
-from dataclasses import dataclass
-from typing import Optional
-
-
-@dataclass
-class LeanResult:
-    success: bool
-    output: str
-    errors: list[str]
-    raw: str
-
-
-def run_lean(code: str, timeout: int = 60) -> LeanResult:
-    """
-    Write `code` to a temporary .lean file, run `lean` on it, and parse
-    the compiler output.
-
-    Args:
-        code:    Full Lean 4 source code string to validate.
-        timeout: Maximum seconds to wait for the compiler.
-
-    Returns:
-        A LeanResult with success flag, cleaned output, and parsed errors.
-    """
-    with tempfile.NamedTemporaryFile(
-        mode="w", suffix=".lean", delete=False
-    ) as f:
-        f.write(code)
-        tmp_path = f.name
-
-    try:
-        result = subprocess.run(
-            ["lean", tmp_path],
-            capture_output=True,
-            text=True,
-            timeout=timeout,
-        )
-        raw = result.stdout + result.stderr
-        errors = _parse_errors(raw)
-        success = result.returncode == 0 and len(errors) == 0
-        return LeanResult(success=success, output=raw.strip(), errors=errors, raw=raw)
-    except FileNotFoundError:
-        return LeanResult(
-            success=False,
-            output="",
-            errors=["Lean 4 not found. Please install it via elan."],
-            raw="",
-        )
-    except subprocess.TimeoutExpired:
-        return LeanResult(
-            success=False,
-            output="",
-            errors=[f"Lean compiler timed out after {timeout}s."],
-            raw="",
-        )
-    finally:
-        os.unlink(tmp_path)
-
-
-def _parse_errors(output: str) -> list[str]:
-    """
-    Extract human-readable error messages from Lean 4 compiler output.
-    Lean 4 errors look like:
-        /path/to/file.lean:10:5: error: unknown identifier 'foo'
-    """
-    pattern = re.compile(r".*?:\d+:\d+:\s*(error|warning):\s*(.+)")
-    errors = []
-    for line in output.splitlines():
-        m = pattern.match(line)
-        if m and m.group(1) == "error":
-            errors.append(m.group(2).strip())
-    return errors

src/lean_verifier.py

@@ -0,0 +1,70 @@
+from typing import Dict, Any, List
+
+from lean_interact import LeanREPLConfig, LeanServer, Command, TempRequireProject, LeanRequire
+
+
+class LeanEnvironment:
+    """
+    Manages the Lean REPL environment for verifying Lean 4 proofs.
+    """
+
+    def __init__(self, use_mathlib: bool = True, lean_version: str = "v4.8.0"):
+        """
+        Initializes the Lean environment.
+        
+        Args:
+            use_mathlib (bool): If True, configures a TempRequireProject with Mathlib.
+                                This may take a while to build on the first run.
+            lean_version (str): The Lean 4 version to use. Default is v4.8.0.
+        """
+        self.lean_version = lean_version
+        self.use_mathlib = use_mathlib
+
+        if self.use_mathlib:
+            # We use TempRequireProject with mathlib as specified in lean_interact documentation
+            project = TempRequireProject(
+                lean_version=self.lean_version, 
+                require="mathlib"
+            )
+            self.config = LeanREPLConfig(project=project)
+        else:
+            self.config = LeanREPLConfig(lean_version=self.lean_version)
+
+        self.server = LeanServer(self.config)
+
+    def verify_proof(self, lean_code: str) -> Dict[str, Any]:
+        """
+        Executes a block of Lean code and verifies if it is a correct proof.
+        
+        Args:
+            lean_code (str): The full Lean 4 code string containing imports, theorem statement, and proof.
+            
+        Returns:
+            dict: A dictionary containing the status, errors (if any), and goals (if open sorries remain).
+        """
+        response = self.server.run(Command(cmd=lean_code))
+
+        errors = []
+        goals = []
+        
+        # Check for error or warning messages
+        if hasattr(response, 'messages') and response.messages:
+            for msg in response.messages:
+                if msg.severity in ['error', 'warning']:
+                    # E.g., 'declaration uses 'sorry'' is a warning, but we might want to capture it
+                    errors.append(msg.data)
+        
+        # Check for open goals (sorries)
+        if hasattr(response, 'sorries') and response.sorries:
+            for sorry in response.sorries:
+                if sorry.goal:
+                    goals.append(sorry.goal)
+                    
+        is_success = len(errors) == 0 and len(goals) == 0
+
+        return {
+            "status": "success" if is_success else "failure",
+            "errors": errors,
+            "goals": goals,
+            "env": getattr(response, "env", None)
+        }

src/lmm_client.py

@@ -0,0 +1,54 @@
+import ollama
+from typing import List, Dict, Any, Optional
+
+class LMMClient:
+    """
+    Client for interacting with local LMMs via Ollama.
+    Focuses on Qwen3-VL:4B for high-reasoning tasks.
+    """
+    
+    def __init__(self, model_name: str = "qwen3-vl:4b"):
+        self.model_name = model_name
+
+    def chat(self, prompt: str, system_prompt: Optional[str] = None) -> str:
+        """
+        Sends a chat request to the model.
+        """
+        messages = []
+        if system_prompt:
+            messages.append({'role': 'system', 'content': system_prompt})
+        
+        messages.append({'role': 'user', 'content': prompt})
+        
+        response = ollama.chat(
+            model=self.model_name,
+            messages=messages
+        )
+        return response['message']['content']
+
+    def generate_proof_steps(self, lean_code: str, goals: List[str], errors: List[str]) -> str:
+        """
+        Specific helper to generate proof steps based on current Lean state.
+        """
+        system_prompt = (
+            "You are an expert Lean 4 proof assistant. "
+            "Your goal is to complete the proof by replacing 'sorry' with valid Lean 4 code. "
+            "Use Mathlib theorems where appropriate. "
+            "Respond ONLY with the corrected Lean code block."
+        )
+        
+        prompt = f"""
+Current Lean Code:
+```lean
+{lean_code}
+```
+
+Current Proof Goals:
+{chr(10).join(goals)}
+
+Lean Errors:
+{chr(10).join(errors)}
+
+Please provide the corrected Lean code. Focus on solving the current goals and fixing the errors.
+"""
+        return self.chat(prompt, system_prompt)

src/mathlib_corpus.py

@@ -0,0 +1,117 @@
+import os
+import re
+import glob
+from pathlib import Path
+from typing import List, Optional
+
+from langchain_core.documents import Document
+
+
+# Regex to capture optional docstring + declaration line
+_DECL_PATTERN = re.compile(
+    r'(?:/--\s*(.*?)\s*-/)?\s*'               # optional /-- docstring -/
+    r'(?:@\[.*?\]\s*)*'                        # optional attributes
+    r'(theorem|lemma|def|noncomputable def)\s+'
+    r'(\S+)\s*'                                # declaration name
+    r'(.*?)\s*:=',                             # everything up to :=
+    re.DOTALL,
+)
+
+
+def _find_mathlib_root() -> Optional[str]:
+    """
+    Returns the path to the Mathlib4 source directory, searching common locations.
+    """
+    candidates = [
+        # Lake package cache (used by lean-interact's TempRequireProject)
+        os.path.expanduser("~/.elan/toolchains"),
+        os.path.expanduser("~/.cache/mathlib"),
+        # Nix / Homebrew Lean setups
+        "/usr/local/lib/lean",
+        "/opt/homebrew/lib/lean",
+    ]
+
+    # Also check for a .lake/packages directory next to the current file
+    here = Path(__file__).resolve().parent.parent
+    lake_pkg = here / ".lake" / "packages" / "mathlib" / "Mathlib"
+    if lake_pkg.exists():
+        return str(lake_pkg.parent)
+
+    for root in candidates:
+        if not os.path.isdir(root):
+            continue
+        # Walk up to 4 levels looking for a Mathlib directory
+        for dirpath, dirnames, _ in os.walk(root):
+            depth = dirpath.replace(root, "").count(os.sep)
+            if depth > 4:
+                dirnames.clear()
+                continue
+            if "Mathlib" in dirnames:
+                return dirpath
+    return None
+
+
+def _parse_lean_file(path: str) -> List[Document]:
+    """
+    Extracts theorem/lemma/def declarations from a single .lean file.
+    Returns a list of Documents with page_content = "<name> : <signature>".
+    """
+    try:
+        text = Path(path).read_text(encoding="utf-8", errors="ignore")
+    except OSError:
+        return []
+
+    docs = []
+    for match in _DECL_PATTERN.finditer(text):
+        docstring = (match.group(1) or "").strip()
+        kind = match.group(2)
+        name = match.group(3)
+        signature = re.sub(r'\s+', ' ', match.group(4)).strip()
+
+        content = f"{name} : {signature}"
+        if docstring:
+            content = f"{docstring}\n{content}"
+
+        line = text[: match.start()].count("\n") + 1
+        docs.append(Document(
+            page_content=content,
+            metadata={"kind": kind, "name": name, "file": path, "line": line},
+        ))
+    return docs
+
+
+class MathLibCorpus:
+    """
+    Extracts LangChain Documents from Mathlib4 source files on disk.
+    """
+
+    def __init__(self, mathlib_root: Optional[str] = None):
+        self.mathlib_root = mathlib_root or _find_mathlib_root()
+
+    def extract(self, max_files: Optional[int] = None) -> List[Document]:
+        """
+        Walks Mathlib source files and extracts declaration Documents.
+
+        Args:
+            max_files: If set, stop after processing this many .lean files
+                       (useful for quick tests).
+
+        Returns:
+            List of LangChain Documents, one per declaration found.
+        """
+        if not self.mathlib_root:
+            raise RuntimeError(
+                "Could not locate Mathlib4 source. "
+                "Pass mathlib_root explicitly or run `lake exe cache get` first."
+            )
+
+        pattern = os.path.join(self.mathlib_root, "**", "*.lean")
+        files = glob.glob(pattern, recursive=True)
+        if max_files:
+            files = files[:max_files]
+
+        docs: List[Document] = []
+        for path in files:
+            docs.extend(_parse_lean_file(path))
+
+        return docs

src/proof_agent.py

@@ -0,0 +1,11 @@
+from langgraph_agent import LangGraphAgent
+
+
+class ProofAgent:
+    """Thin compatibility wrapper around LangGraphAgent."""
+
+    def __init__(self, model_name: str = "qwen3-vl:4b", max_retries: int = 5):
+        self._agent = LangGraphAgent(model_name=model_name, max_retries=max_retries)
+
+    def solve_file(self, file_path: str) -> bool:
+        return self._agent.solve_file(file_path)

src/rag_chain.py

@@ -0,0 +1,71 @@
+from typing import List
+
+from langchain_core.documents import Document
+from langchain_core.output_parsers import StrOutputParser
+from langchain_core.prompts import ChatPromptTemplate
+from langchain_ollama import OllamaLLM
+
+
+_SYSTEM = (
+    "You are an expert Lean 4 proof assistant with deep knowledge of Mathlib. "
+    "Your task is to complete the proof by replacing every `sorry` with valid Lean 4 tactic code. "
+    "Use only Mathlib theorems and tactics. "
+    "Respond ONLY with the corrected Lean code inside a single ```lean ... ``` block."
+)
+
+_HUMAN = """\
+## Current Lean Code
+```lean
+{lean_code}
+```
+
+## Open Proof Goals
+{goals}
+
+## Lean Errors
+{errors}
+
+## Relevant Mathlib Lemmas
+{retrieved_lemmas}
+
+Provide the corrected Lean code that solves all goals and fixes all errors.
+"""
+
+
+def _format_docs(docs: List[Document]) -> str:
+    if not docs:
+        return "(none retrieved)"
+    return "\n".join(
+        f"- `{d.metadata.get('name', '?')}`: {d.page_content}" for d in docs
+    )
+
+
+class RAGProofChain:
+    """
+    LangChain LCEL chain: retrieved context + proof state → corrected Lean code.
+    """
+
+    def __init__(self, model_name: str = "qwen3-vl:4b"):
+        prompt = ChatPromptTemplate.from_messages([
+            ("system", _SYSTEM),
+            ("human", _HUMAN),
+        ])
+        llm = OllamaLLM(model=model_name)
+        self._chain = prompt | llm | StrOutputParser()
+
+    def generate(
+        self,
+        lean_code: str,
+        goals: List[str],
+        errors: List[str],
+        retrieved_lemmas: List[Document],
+    ) -> str:
+        """
+        Generate corrected Lean code given the current proof state and retrieved lemmas.
+        """
+        return self._chain.invoke({
+            "lean_code": lean_code,
+            "goals": "\n".join(goals) or "(none)",
+            "errors": "\n".join(errors) or "(none)",
+            "retrieved_lemmas": _format_docs(retrieved_lemmas),
+        })

src/retriever.py

@@ -1,113 +1,126 @@
-"""
-retriever.py
-
-Builds and queries a FAISS index over Mathlib4 lemma docstrings.
-
-Usage:
-    # Build index from a jsonl file of lemmas:
-    build_index("data/mathlib_lemmas.jsonl", "data/mathlib.index")
-
-    # Query at runtime:
-    retriever = Retriever("data/mathlib.index", "data/mathlib_lemmas.jsonl")
-    lemmas = retriever.retrieve("commutativity of addition", k=5)
-"""
-
-import json
 import os
-import pickle
 from pathlib import Path
-from typing import Optional
+from typing import List, Optional
 
-import faiss
-import numpy as np
-from sentence_transformers import SentenceTransformer
+from langchain_community.retrievers import BM25Retriever
+from langchain_community.vectorstores import FAISS
+from langchain_community.cross_encoders import HuggingFaceCrossEncoder
+from langchain.retrievers import EnsembleRetriever
+from langchain.retrievers.document_compressors import CrossEncoderReranker
+from langchain.retrievers import ContextualCompressionRetriever
+from langchain_huggingface import HuggingFaceEmbeddings
+from langchain_core.documents import Document
 
-EMBEDDING_MODEL = "all-MiniLM-L6-v2"
-INDEX_FILE = "data/mathlib.index"
-META_FILE = "data/mathlib_meta.pkl"
+from mathlib_corpus import MathLibCorpus
 
+_DEFAULT_INDEX_DIR = Path(__file__).resolve().parent.parent / "data" / "mathlib_index"
+_EMBED_MODEL = "all-MiniLM-L6-v2"
+_RERANK_MODEL = "cross-encoder/ms-marco-MiniLM-L-6-v2"
 
-class Retriever:
+
+class MathLibRetriever:
     """
-    Manages the FAISS index and performs semantic retrieval of Mathlib4 lemmas.
+    Hybrid FAISS + BM25 retriever with CrossEncoder reranking over Mathlib lemmas.
+
+    On first use, call build() to create and persist the index.
+    Subsequent runs load from disk automatically.
     """
 
     def __init__(
         self,
-        index_path: str = INDEX_FILE,
-        meta_path: str = META_FILE,
-        model_name: str = EMBEDDING_MODEL,
+        index_dir: Optional[str] = None,
+        top_k: int = 20,
+        rerank_top_k: int = 5,
     ):
-        self.model = SentenceTransformer(model_name)
-        self.index = faiss.read_index(index_path)
-        with open(meta_path, "rb") as f:
-            self.metadata = pickle.load(f)  # list of dicts: {name, type, doc}
+        self.index_dir = Path(index_dir) if index_dir else _DEFAULT_INDEX_DIR
+        self.top_k = top_k
+        self.rerank_top_k = rerank_top_k
+        self._retriever = None
 
-    def retrieve(self, query: str, k: int = 5) -> list[dict]:
+    # ------------------------------------------------------------------
+    # Public API
+    # ------------------------------------------------------------------
+
+    def build(self, mathlib_root: Optional[str] = None, max_files: Optional[int] = None) -> None:
+        """
+        Extract Mathlib documents, build FAISS + BM25 indices, and persist to disk.
+        Call this once (via scripts/build_index.py) before first use.
         """
-        Return the top-k most relevant Mathlib4 lemmas for a given query.
+        print("Extracting Mathlib corpus…")
+        corpus = MathLibCorpus(mathlib_root=mathlib_root)
+        docs = corpus.extract(max_files=max_files)
+        print(f"  {len(docs)} declarations extracted.")
+
+        embeddings = self._embeddings()
+
+        print("Building FAISS index…")
+        faiss_store = FAISS.from_documents(docs, embeddings)
+        self.index_dir.mkdir(parents=True, exist_ok=True)
+        faiss_store.save_local(str(self.index_dir))
+        print(f"  Index saved to {self.index_dir}")
+
+        self._retriever = self._build_retriever(faiss_store, docs)
+
+    def retrieve(self, query: str, k: Optional[int] = None) -> List[Document]:
+        """
+        Retrieve and rerank the most relevant Mathlib lemmas for a query.
 
         Args:
-            query: Natural language or Lean theorem statement.
-            k:     Number of results to return.
+            query: Natural-language or Lean-syntax query (e.g., proof goals + errors).
+            k: Number of results to return after reranking (defaults to self.rerank_top_k).
 
         Returns:
-            List of lemma dicts with keys 'name', 'type', 'doc'.
+            List of Documents ranked by relevance.
         """
-        embedding = self.model.encode([query], normalize_embeddings=True)
-        embedding = np.array(embedding, dtype=np.float32)
-        _, indices = self.index.search(embedding, k)
-        return [self.metadata[i] for i in indices[0] if i < len(self.metadata)]
-
-
-def build_index(
-    lemmas_jsonl: str = "data/mathlib_lemmas.jsonl",
-    index_path: str = INDEX_FILE,
-    meta_path: str = META_FILE,
-    model_name: str = EMBEDDING_MODEL,
-) -> None:
-    """
-    Build a FAISS index from a JSONL file of Mathlib4 lemmas.
-
-    Each line in `lemmas_jsonl` should be a JSON object with:
-        - "name": lemma name (str)
-        - "type": Lean type signature (str)
-        - "doc":  docstring or description (str, optional)
-
-    Args:
-        lemmas_jsonl: Path to the input JSONL file.
-        index_path:   Where to save the FAISS index file.
-        meta_path:    Where to save the metadata pickle file.
-        model_name:   SentenceTransformer model to use for embeddings.
-    """
-    print(f"Loading lemmas from {lemmas_jsonl}...")
-    metadata = []
-    texts = []
-    with open(lemmas_jsonl, "r") as f:
-        for line in f:
-            entry = json.loads(line.strip())
-            metadata.append(entry)
-            # Build a rich text representation for embedding
-            text = f"{entry.get('name', '')} : {entry.get('type', '')}. {entry.get('doc', '')}"
-            texts.append(text)
-
-    print(f"Loaded {len(texts)} lemmas. Generating embeddings...")
-    model = SentenceTransformer(model_name)
-    embeddings = model.encode(texts, normalize_embeddings=True, show_progress_bar=True)
-    embeddings = np.array(embeddings, dtype=np.float32)
-
-    print("Building FAISS index...")
-    dim = embeddings.shape[1]
-    index = faiss.IndexFlatIP(dim)  # Inner product = cosine similarity (normalized)
-    index.add(embeddings)
-
-    os.makedirs(Path(index_path).parent, exist_ok=True)
-    faiss.write_index(index, index_path)
-    with open(meta_path, "wb") as f:
-        pickle.dump(metadata, f)
-
-    print(f"Index saved to {index_path} ({len(texts)} vectors, dim={dim})")
-
-
-if __name__ == "__main__":
-    build_index()
+        if self._retriever is None:
+            self._load()
+        results = self._retriever.invoke(query)
+        return results[: k or self.rerank_top_k]
+
+    def is_index_built(self) -> bool:
+        return (self.index_dir / "index.faiss").exists()
+
+    # ------------------------------------------------------------------
+    # Internal helpers
+    # ------------------------------------------------------------------
+
+    def _embeddings(self) -> HuggingFaceEmbeddings:
+        return HuggingFaceEmbeddings(model_name=_EMBED_MODEL)
+
+    def _load(self) -> None:
+        if not self.is_index_built():
+            raise RuntimeError(
+                f"No FAISS index found at {self.index_dir}. "
+                "Run `python scripts/build_index.py` first."
+            )
+        print("Loading FAISS index from disk…")
+        embeddings = self._embeddings()
+        faiss_store = FAISS.load_local(
+            str(self.index_dir),
+            embeddings,
+            allow_dangerous_deserialization=True,
+        )
+
+        # Re-build BM25 from FAISS docstore
+        docs = list(faiss_store.docstore._dict.values())
+        self._retriever = self._build_retriever(faiss_store, docs)
+
+    def _build_retriever(self, faiss_store: FAISS, docs: List[Document]):
+        faiss_retriever = faiss_store.as_retriever(
+            search_kwargs={"k": self.top_k}
+        )
+        bm25_retriever = BM25Retriever.from_documents(docs)
+        bm25_retriever.k = self.top_k
+
+        ensemble = EnsembleRetriever(
+            retrievers=[faiss_retriever, bm25_retriever],
+            weights=[0.6, 0.4],
+        )
+
+        cross_encoder = HuggingFaceCrossEncoder(model_name=_RERANK_MODEL)
+        reranker = CrossEncoderReranker(model=cross_encoder, top_n=self.rerank_top_k)
+
+        return ContextualCompressionRetriever(
+            base_compressor=reranker,
+            base_retriever=ensemble,
+        )

tests/test_lean_verifier.py

@@ -0,0 +1,57 @@
+import unittest
+import sys
+import os
+
+# Add src to Python path
+sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), '..', 'src')))
+
+from lean_verifier import LeanEnvironment
+
+class TestLeanVerifier(unittest.TestCase):
+    @classmethod
+    def setUpClass(cls):
+        # We'll use Mathlib here since our goal is to verify it works with the MVP setup
+        cls.lean_env = LeanEnvironment(use_mathlib=True)
+
+    def test_correct_proof(self):
+        lean_code = """
+import Mathlib
+
+theorem add_comm_test (n m : Nat) : n + m = m + n := by
+  exact Nat.add_comm n m
+"""
+        result = self.lean_env.verify_proof(lean_code)
+        self.assertEqual(result["status"], "success")
+        self.assertEqual(len(result["errors"]), 0)
+        self.assertEqual(len(result["goals"]), 0)
+
+    def test_incorrect_proof_type_mismatch(self):
+        lean_code = """
+import Mathlib
+
+theorem add_comm_test (n m : Nat) : n + m = m + n := by
+  exact n
+"""
+        result = self.lean_env.verify_proof(lean_code)
+        self.assertEqual(result["status"], "failure")
+        self.assertTrue(any("type mismatch" in err or "application type mismatch" in err for err in result["errors"]), 
+                        f"Expected type mismatch error, got: {result['errors']}")
+
+    def test_incomplete_proof_sorry(self):
+        lean_code = """
+import Mathlib
+
+theorem my_incomplete_thm (n : Nat) : n = 5 → n = 5 := by
+  sorry
+"""
+        result = self.lean_env.verify_proof(lean_code)
+        self.assertEqual(result["status"], "failure")
+        # Ensure it has an error indicating sorry
+        self.assertTrue(any("uses 'sorry'" in err for err in result["errors"]), 
+                        f"Expected sorry warning/error, got: {result['errors']}")
+        # Ensure it outputs the goal
+        self.assertEqual(len(result["goals"]), 1)
+        self.assertIn("⊢ n = 5 → n = 5", result["goals"][0])
+
+if __name__ == '__main__':
+    unittest.main()