Fix LLM hallucinated imports and slow thinking mode (#2)

* Fix import paths and Mathlib auto-detection for LangChain 1.x

EnsembleRetriever and CrossEncoderReranker moved to langchain_classic
in LangChain 1.x — update retriever.py imports accordingly.

MathLibCorpus now searches lean-interact's TempRequireProject cache
(.lake/packages/mathlib) for Mathlib source and prefers the Mathlib/
subdirectory to avoid import-only root files returning zero declarations.

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>

* Fix LLM hallucinated imports and slow thinking mode

- langgraph_agent.py: add _sanitize_imports() which strips all LLM-generated
import lines and replaces them with `import Mathlib`. LLMs consistently
hallucinate non-existent Lean paths (e.g. Mathlib.TacticLibrary.Peel) that
break compilation; the correct import for any Mathlib proof is always
`import Mathlib`.

- rag_chain.py: tighten system prompt with explicit rules (keep imports,
no spurious `open`, prefer simple tactics). Add num_predict=1024 cap and
think=False to disable chain-of-thought mode on qwen3/gemma3 models —
without this, models spend 20+ minutes generating reasoning tokens before
the actual answer.

Tested: simple_add.lean (n + 0 = n) solved in 2 attempts with gemma3:12b,
verified by Lean REPL. Proof: `simp`.

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>

---------

Co-authored-by: p4r5kpftnp-cmd <p4r5kpftnp@privaterelay.appleid.com>
Co-authored-by: Claude Sonnet 4.6 <noreply@anthropic.com>

src/langgraph_agent.py

@@ -45,6 +45,17 @@ def _extract_lean_code(text: str) -> str:
     return text.strip()
 
 
+def _sanitize_imports(code: str) -> str:
+    """
+    LLMs often hallucinate Lean import paths. This function strips all `import`
+    lines from the generated code and replaces them with `import Mathlib`, which
+    is the correct single import for any Mathlib-based proof.
+    """
+    lines = code.splitlines()
+    non_import_lines = [l for l in lines if not l.strip().startswith("import ")]
+    return "import Mathlib\n\n" + "\n".join(non_import_lines).lstrip()
+
+
 def make_verify_node(lean_env: LeanEnvironment):
     def verify_node(state: ProofState) -> ProofState:
         print(f"\n--- Attempt {state['attempt'] + 1} / {state['max_retries']} ---")
@@ -89,7 +100,7 @@ def make_generate_node(chain: RAGProofChain):
             errors=state["errors"],
             retrieved_lemmas=state["retrieved_lemmas"],
         )
-        new_code = _extract_lean_code(raw)
+        new_code = _sanitize_imports(_extract_lean_code(raw))
 
         if not new_code or new_code.strip() == state["lean_code"].strip():
             print("LLM produced no changes.")

src/mathlib_corpus.py

@@ -22,28 +22,36 @@ def _find_mathlib_root() -> Optional[str]:
     """
     Returns the path to the Mathlib4 source directory, searching common locations.
     """
+    # 1. .lake/packages next to the current project file
+    here = Path(__file__).resolve().parent.parent
+    lake_pkg = here / ".lake" / "packages" / "mathlib"
+    if (lake_pkg / "Mathlib").exists():
+        return str(lake_pkg)
+
+    # 2. lean-interact's TempRequireProject cache
+    try:
+        from lean_interact.config import DEFAULT_CACHE_DIR
+        tmp_root = DEFAULT_CACHE_DIR / "tmp_projects"
+        for lean_ver_dir in sorted(tmp_root.iterdir(), reverse=True):
+            for project_dir in lean_ver_dir.iterdir():
+                candidate = project_dir / ".lake" / "packages" / "mathlib"
+                if (candidate / "Mathlib").exists():
+                    return str(candidate)
+    except Exception:
+        pass
+
+    # 3. Broad filesystem search in common Lean 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:
+            if depth > 6:
                 dirnames.clear()
                 continue
             if "Mathlib" in dirnames:
@@ -105,7 +113,10 @@ class MathLibCorpus:
                 "Pass mathlib_root explicitly or run `lake exe cache get` first."
             )
 
-        pattern = os.path.join(self.mathlib_root, "**", "*.lean")
+        # Prefer Mathlib/ subdirectory if it exists (avoids index-only files at root)
+        mathlib_src = os.path.join(self.mathlib_root, "Mathlib")
+        search_root = mathlib_src if os.path.isdir(mathlib_src) else self.mathlib_root
+        pattern = os.path.join(search_root, "**", "*.lean")
         files = glob.glob(pattern, recursive=True)
         if max_files:
             files = files[:max_files]

src/rag_chain.py

@@ -9,8 +9,12 @@ 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."
+    "RULES:\n"
+    "1. Keep `import Mathlib` exactly as-is at the top. Do NOT add, remove, or change any import lines.\n"
+    "2. Do NOT add `open` statements unless they were already in the original code.\n"
+    "3. Keep the theorem signature exactly as given — do not change argument names or types.\n"
+    "4. Replace `sorry` with valid Lean 4 tactic(s). Prefer simple tactics: `simp`, `omega`, `ring`, `exact`, `apply`.\n"
+    "5. Respond ONLY with the complete corrected Lean code inside a single ```lean ... ``` block."
 )
 
 _HUMAN = """\
@@ -50,7 +54,13 @@ class RAGProofChain:
             ("system", _SYSTEM),
             ("human", _HUMAN),
         ])
-        llm = OllamaLLM(model=model_name)
+        # Disable thinking/chain-of-thought mode (qwen3, gemma3) and cap output
+        # so the agent doesn't spend minutes generating reasoning tokens.
+        llm = OllamaLLM(
+            model=model_name,
+            num_predict=1024,           # cap response length
+            options={"think": False},   # disable thinking mode (qwen3/gemma3)
+        )
         self._chain = prompt | llm | StrOutputParser()
 
     def generate(

src/retriever.py

@@ -5,9 +5,8 @@ from typing import List, Optional
 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_classic.retrievers import EnsembleRetriever, ContextualCompressionRetriever
+from langchain_classic.retrievers.document_compressors import CrossEncoderReranker
 from langchain_huggingface import HuggingFaceEmbeddings
 from langchain_core.documents import Document