Research-Stack/5-Applications/tools-scripts/llm/deepseek_v4_flash_lean_harness.py

#!/usr/bin/env python3
"""
Lean Proof Harness — provider-agnostic, optimized for DeepSeek V4 Flash.

Discovers `sorry` markers in a .lean file, sends each theorem (with context)
to an LLM, inserts generated proofs, and verifies with `lake build`.

Provider-agnostic: switch providers via `--provider` or `LLM_PROVIDER`.
Optimized for DeepSeek V4 Flash (lowest token cost, best Lean proof quality).

Usage:
  # Resolve all sorries in a file
  python3 deepseek_v4_flash_lean_harness.py resolve Semantics/E8Sidon.lean

  # Resolve a specific sorry by line number
  python3 deepseek_v4_flash_lean_harness.py resolve Semantics/E8Sidon.lean --line 950

  # Use a different provider
  python3 deepseek_v4_flash_lean_harness.py resolve --provider deepseek-api Semantics/E8Sidon.lean

  # List sorries without resolving
  python3 deepseek_v4_flash_lean_harness.py scan Semantics/E8Sidon.lean

  # Interactive mode
  python3 deepseek_v4_flash_lean_harness.py resolve --interactive Semantics/E8Sidon.lean

Providers (built-in):
  llamacpp-local   — http://100.88.57.96:30516/v1, model=deepseek-v4-flash (default)
  deepseek-api     — https://api.deepseek.com/v1, model=deepseek-chat
  ollama-local     — http://localhost:11434/v1, model=deepseek-v4-flash
  openrouter       — https://openrouter.ai/api/v1, model=deepseek/deepseek-chat

Environment:
  LLM_PROVIDER      — provider name (default: llamacpp-local)
  LLM_API_BASE      — override API base URL
  LLM_API_KEY       — override API key
  LLM_MODEL         — override model name
  LAKE_WORKDIR      — lake build working directory (default: 0-Core-Formalism/lean/Semantics)
"""

from __future__ import annotations

import argparse
import json
import os
import re
import subprocess
import sys
import time
import urllib.request
import urllib.error
from dataclasses import dataclass, field
from datetime import datetime, timezone
from pathlib import Path
from typing import Optional

# ---------------------------------------------------------------------------
# Provider registry
# ---------------------------------------------------------------------------

PROVIDERS: dict[str, dict] = {
    "llamacpp-local": {
        "api_base": "http://100.88.57.96:30516/v1",
        "api_key": "sk-local",
        "model": "deepseek-v4-flash",
        "notes": "local llama.cpp on qfox-1 RTX 4070, ~8.2 GB quant",
    },
    "deepseek-api": {
        "api_base": "https://api.deepseek.com/v1",
        "api_key": "",
        "model": "deepseek-chat",
        "notes": "DeepSeek cloud API, lowest token cost",
    },
    "ollama-local": {
        "api_base": "http://localhost:11434/v1",
        "api_key": "sk-local",
        "model": "deepseek-v4-flash",
        "notes": "local Ollama on qfox-1",
    },
    "openrouter": {
        "api_base": "https://openrouter.ai/api/v1",
        "api_key": "",
        "model": "deepseek/deepseek-v4-flash",
        "notes": "OpenRouter, deepseek-v4-flash",
    },
    "neon-deepseek-prover": {
        "api_base": "http://100.92.88.64:11434/v1",
        "api_key": "",
        "model": "hf.co/irmma/DeepSeek-Prover-V2-7B-Q4_K_M-GGUF",
        "notes": "neon-64gb, DeepSeek-Prover-V2 7B Q4_K_M, 62GB RAM CPU-only",
    },
    "neon-goedel-prover": {
        "api_base": "http://100.92.88.64:11434/v1",
        "api_key": "",
        "model": "hf.co/mradermacher/Goedel-Prover-V2-8B-GGUF",
        "notes": "neon-64gb, Goedel-Prover-V2 8B GGUF, 62GB RAM CPU-only",
    },
}

DEFAULT_PROVIDER = "llamacpp-local"

RECEIPT_DIR = Path(__file__).resolve().parents[3] / "shared-data" / "artifacts" / "deepseek_prover"

PROMPT_TEMPLATE = """You are a Lean 4 theorem prover for the Research Stack project.

Project rules:
- Use Q16_16 fixed-point (no Float in compute paths).
- No bare sorries, no tautologies.
- Use `calc`, `omega`, `native_decide`, `positivity`, `linarith`, `nlinarith`.
- Prefer explicit `calc` blocks over opaque tactic scripts.
- Follow existing patterns in the file.

Below is a Lean 4 module with one unproven theorem (marked `:= by\\n  sorry`).
The imports and surrounding definitions are shown for context.

Output ONLY the proof block — the code that replaces `:= by\\n  sorry`.
Do NOT repeat the theorem statement. Do NOT wrap in markdown fences.
Start with `:= by` and end with the closing of the proof.

---
{context}
---

The unproven theorem at line {line_no}:

{theorem_block}

Generate the proof:"""

# ---------------------------------------------------------------------------
# Data
# ---------------------------------------------------------------------------


@dataclass
class SorrySite:
    line: int
    theorem_name: str
    theorem_block: str
    context_before: str
    context_after: str
    full_context: str


@dataclass
class Provider:
    name: str
    api_base: str
    api_key: str
    model: str

    @classmethod
    def from_name(cls, name: str) -> Provider:
        spec = PROVIDERS.get(name)
        if not spec:
            available = ", ".join(PROVIDERS)
            print(f"Unknown provider '{name}'. Available: {available}")
            sys.exit(1)
        return cls(
            name=name,
            api_base=os.environ.get("LLM_API_BASE", spec["api_base"]),
            api_key=os.environ.get("LLM_API_KEY", spec["api_key"]),
            model=os.environ.get("LLM_MODEL", spec["model"]),
        )


@dataclass
class HarnessConfig:
    provider: Provider = field(default_factory=lambda: Provider.from_name(DEFAULT_PROVIDER))
    lake_workdir: str = ""
    temperature: float = 0.4
    max_tokens: int = 4096
    max_iterations: int = 5
    interactive: bool = False
    dry_run: bool = False


@dataclass
class ProofAttempt:
    sorry_site: SorrySite
    code: str = ""
    passed: bool = False
    iterations: int = 0
    compile_log: str = ""
    latency_ms: float = 0.0
    error_feedback: str = ""


# ---------------------------------------------------------------------------
# Sorry discovery
# ---------------------------------------------------------------------------


def discover_sorries(lean_path: Path) -> list[SorrySite]:
    text = lean_path.read_text()
    lines = text.split("\n")

    # Scan once to collect all theorem/lemma declarations and all sorry lines.
    # Line numbers are 1-indexed.
    theorems: list[tuple[int, str]] = []
    for i, line in enumerate(lines, 1):
        m = re.match(r"^(theorem|lemma)\s+(\w+)", line)
        if m:
            theorems.append((i, m.group(2)))

    sorry_sites: list[SorrySite] = []
    # Track block-comment depth so we don't flag `sorry` inside /- ... -/
    in_block_comment = 0

    for i, line in enumerate(lines, 1):
        stripped = line.strip()

        # Track block-comment depth
        in_block_comment += line.count("/-") - line.count("-/")
        if in_block_comment > 0:
            continue

        # Skip single-line comments and doc-comment lines
        if stripped.startswith("--") or stripped.startswith("/-") or stripped.startswith("*"):
            continue

        if re.search(r"(?<!\w)sorry(?!\w)", stripped):
            # Find the nearest theorem/lemma before this sorry
            theorem_line = 0
            theorem_name = "unknown"
            for tl, tn in theorems:
                if tl <= i:
                    theorem_line = tl
                    theorem_name = tn
                else:
                    break
            if theorem_line == 0:
                continue

            ctx_start = max(0, theorem_line - 15)
            ctx_end = min(len(lines), i + 5)
            context_before = "\n".join(lines[ctx_start - 1 : theorem_line - 1])
            theorem_block = "\n".join(lines[theorem_line - 1 : i])
            context_after = "\n".join(lines[i : ctx_end])
            full_context = "\n".join(
                lines[max(0, theorem_line - 31) : min(len(lines), i + 10)]
            )

            sorry_sites.append(SorrySite(
                line=i,
                theorem_name=theorem_name,
                theorem_block=theorem_block,
                context_before=context_before,
                context_after=context_after,
                full_context=full_context,
            ))

    return sorry_sites


# ---------------------------------------------------------------------------
# LLM API call — provider-agnostic (OpenAI-compatible)
# ---------------------------------------------------------------------------


def call_llm(prompt: str, provider: Provider, cfg: HarnessConfig) -> tuple[str, float]:
    endpoint = f"{provider.api_base.rstrip('/')}/chat/completions"
    body = json.dumps({
        "model": provider.model,
        "messages": [{"role": "user", "content": prompt}],
        "temperature": cfg.temperature,
        "max_tokens": cfg.max_tokens,
        "stream": False,
    }).encode()

    headers = {
        "Content-Type": "application/json",
        "Authorization": f"Bearer {provider.api_key}",
    }

    t0 = time.perf_counter()
    req = urllib.request.Request(endpoint, data=body, headers=headers, method="POST")
    try:
        with urllib.request.urlopen(req, timeout=180) as resp:
            data = json.loads(resp.read())
        latency = (time.perf_counter() - t0) * 1000
        return data["choices"][0]["message"]["content"], latency
    except Exception as exc:
        latency = (time.perf_counter() - t0) * 1000
        return f"ERROR: {exc}", latency


# ---------------------------------------------------------------------------
# Lake build
# ---------------------------------------------------------------------------


def run_lake_build(workdir: str, target: str = "") -> tuple[int, str]:
    cmd = ["lake", "build"]
    if target:
        cmd.append(target)
    try:
        result = subprocess.run(
            cmd, capture_output=True, text=True, timeout=240, cwd=workdir,
        )
        return result.returncode, result.stdout + "\n" + result.stderr
    except subprocess.TimeoutExpired as exc:
        return 1, f"TIMEOUT: {exc}"


def extract_errors(log: str) -> str:
    lines = log.split("\n")
    errors = [l for l in lines if "error:" in l or "sorry" in l]
    return "\n".join(errors[:15])


# ---------------------------------------------------------------------------
# Proof insertion
# ---------------------------------------------------------------------------


def extract_proof_code(response: str) -> str:
    text = re.sub(r"^\s*```(?:lean)?\s*\n?", "", response, flags=re.MULTILINE)
    text = re.sub(r"\n\s*```\s*$", "", text, flags=re.MULTILINE)
    text = text.strip()

    if text.startswith(":= by"):
        return text

    idx = text.find(":= by")
    if idx >= 0:
        return text[idx:]

    idx = text.find("\nby ")
    if idx >= 0:
        prev_newline = text.rfind("\n", 0, idx)
        return text[prev_newline + 1:].strip()

    return text


def insert_proof(lean_path: Path, sorry_line: int, proof_code: str) -> bool:
    lines = lean_path.read_text().split("\n")
    n = len(lines)

    # Scan backwards up to 30 lines to find `:= by` or `:=` + `by` split.
    insert_idx = -1
    for i in range(sorry_line - 1, max(-1, sorry_line - 31), -1):
        if i < 0 or i >= n:
            break
        if ":= by" in lines[i]:
            insert_idx = i
            break
        if ":=" in lines[i] and "sorry" not in lines[i]:
            # Check if `by` follows on the next line
            if i + 1 < n and lines[i + 1].strip() == "by":
                insert_idx = i + 1
            else:
                insert_idx = i
            break

    if insert_idx < 0:
        print(f"  No `:= by` found before line {sorry_line}")
        return False

    # Infer indentation from the existing proof body (or default to 2 spaces).
    indent = "  "
    for probe in range(insert_idx + 1, min(insert_idx + 4, sorry_line)):
        if probe < n and lines[probe].strip():
            raw = lines[probe]
            indent = raw[: len(raw) - len(raw.lstrip())] or "  "
            break

    # Strip `:= by` prefix if the LLM included it.
    code = proof_code.strip()
    if code.startswith(":= by"):
        code = code[5:].strip()

    # Build indented proof lines (preserve blank lines).
    proof_lines: list[str] = []
    for pl in code.split("\n"):
        pl_stripped = pl.strip()
        if pl_stripped:
            proof_lines.append(f"{indent}{pl_stripped}")
        else:
            proof_lines.append("")

    # Count consecutive sorry/blank lines starting at sorry_line to skip.
    skip_to = sorry_line  # 1-indexed, inclusive
    while skip_to <= n:
        line_content = lines[skip_to - 1]
        if "sorry" in line_content or line_content.strip() == "":
            skip_to += 1
        else:
            break

    # Reconstruct: header line + indented proof + lines after the removed block.
    new_lines = lines[: insert_idx + 1] + proof_lines + lines[skip_to - 1 :]

    lean_path.write_text("\n".join(new_lines))
    return True


# ---------------------------------------------------------------------------
# Receipt emission
# ---------------------------------------------------------------------------


def emit_receipt(attempt: ProofAttempt, cfg: HarnessConfig) -> Path:
    RECEIPT_DIR.mkdir(parents=True, exist_ok=True)
    ts = datetime.now(timezone.utc).strftime("%Y%m%dT%H%M%SZ")
    safe_name = attempt.sorry_site.theorem_name[:40]
    fname = f"proof_attempt_{safe_name}_{ts}.json"

    receipt = {
        "schema": "lean_proof_attempt_v1",
        "provider": cfg.provider.name,
        "model": cfg.provider.model,
        "endpoint": cfg.provider.api_base,
        "theorem": attempt.sorry_site.theorem_name,
        "line": attempt.sorry_site.line,
        "passed": attempt.passed,
        "iterations": attempt.iterations,
        "latency_ms": attempt.latency_ms,
        "timestamp": ts,
        "error_preview": extract_errors(attempt.compile_log)[:500] if not attempt.passed else "",
    }
    path = RECEIPT_DIR / fname
    path.write_text(json.dumps(receipt, indent=2) + "\n")
    return path


# ---------------------------------------------------------------------------
# Main resolve loop
# ---------------------------------------------------------------------------


def resolve_sorry(site: SorrySite, cfg: HarnessConfig, lean_path: Path) -> ProofAttempt:
    pname = cfg.provider.name
    print(f"\n{'=' * 60}")
    print(f"Theorem: {site.theorem_name} (line {site.line})  [{pname}]")
    print(f"{'=' * 60}")
    print(site.theorem_block[:200] + "..." if len(site.theorem_block) > 200 else site.theorem_block)

    if cfg.dry_run:
        print(f"\n[DRY RUN] Would send prompt ({len(site.full_context)} chars context)")
        print(f"  Provider: {cfg.provider.name} @ {cfg.provider.api_base}")
        print(f"  Model: {cfg.provider.model}")
        print(f"Theorem block:\n{site.theorem_block[:300]}...\n")
        return ProofAttempt(sorry_site=site, passed=False, iterations=0)

    if cfg.interactive:
        resp = input(f"\nSend to {cfg.provider.model} via {cfg.provider.name}? [Y/n] ").strip().lower()
        if resp == "n":
            print("  Skipping.")
            return ProofAttempt(sorry_site=site, passed=False, iterations=0)

    attempt = ProofAttempt(sorry_site=site)

    for iteration in range(1, cfg.max_iterations + 1):
        print(f"\n--- Iteration {iteration}/{cfg.max_iterations} ---")

        context = site.full_context
        error_feedback = attempt.error_feedback
        if error_feedback:
            prompt = PROMPT_TEMPLATE + f"\n\nPrevious attempt failed. Errors:\n{error_feedback}\n\nTry a different approach:"
        else:
            prompt = PROMPT_TEMPLATE.format(
                context=context,
                line_no=site.line,
                theorem_block=site.theorem_block,
            )

        # Call LLM
        response, latency = call_llm(prompt, cfg.provider, cfg)
        attempt.latency_ms += latency
        print(f"  API: {latency:.0f}ms  [{cfg.provider.name}/{cfg.provider.model}]")

        if response.startswith("ERROR:"):
            print(f"  {response}")
            if iteration < cfg.max_iterations:
                continue
            break

        proof_code = extract_proof_code(response)
        print(f"  Generated: {len(proof_code)} chars")
        if not proof_code:
            print("  Empty response, retrying...")
            continue

        if not insert_proof(lean_path, site.line, proof_code):
            print("  Failed to insert proof")
            continue

        workdir = cfg.lake_workdir or os.environ.get("LAKE_WORKDIR", "")
        rc, log = run_lake_build(workdir)
        attempt.compile_log = log
        attempt.code = proof_code
        attempt.iterations = iteration

        if rc == 0:
            print(f"  \033[32mPASSED!\033[0m (iteration {iteration})")
            attempt.passed = True
            return attempt

        errors = extract_errors(log)
        attempt.error_feedback = errors[:2000]
        print(f"  \033[31mFAILED\033[0m (return code {rc})")
        if errors:
            print(f"  Errors: {errors[:300]}...")

        if not insert_proof(lean_path, site.line, "  sorry"):
            print("  Warning: could not restore sorry marker")

    return attempt


# ---------------------------------------------------------------------------
# CLI
# ---------------------------------------------------------------------------


def cmd_scan(args):
    path = Path(args.lean_file)
    if not path.exists():
        print(f"File not found: {path}")
        sys.exit(1)

    sites = discover_sorries(path)
    if not sites:
        print("No sorries found.")
        return

    print(f"Found {len(sites)} sorry site(s) in {path}:")
    for s in sites:
        print(f"  Line {s.line:>5}: {s.theorem_name}")


def cmd_resolve(args):
    path = Path(args.lean_file)
    if not path.exists():
        print(f"File not found: {path}")
        sys.exit(1)

    provider = Provider.from_name(args.provider)

    cfg = HarnessConfig(
        provider=provider,
        lake_workdir=args.lake_workdir or os.environ.get("LAKE_WORKDIR", ""),
        temperature=args.temperature,
        max_iterations=args.max_iterations,
        interactive=args.interactive,
        dry_run=args.dry_run,
    )

    sites = discover_sorries(path)
    if args.line:
        sites = [s for s in sites if s.line == args.line]
        if not sites:
            print(f"No sorry at line {args.line}")
            sys.exit(1)

    if not sites:
        print("No sorries found.")
        return

    print(f"Found {len(sites)} sorry site(s). Provider: {provider.name} ({provider.model})")
    passed = 0
    failed = 0

    for site in sites:
        attempt = resolve_sorry(site, cfg, path)
        if attempt.passed:
            passed += 1
        else:
            failed += 1
        receipt_path = emit_receipt(attempt, cfg)
        print(f"  Receipt: {receipt_path}")

    print(f"\n{'=' * 60}")
    print(f"Results: {passed} passed, {failed} failed, {len(sites)} total")


def main():
    parser = argparse.ArgumentParser(
        description="Lean Proof Harness — provider-agnostic, optimized for DeepSeek V4 Flash",
    )
    sub = parser.add_subparsers(dest="command", required=True)

    scan_p = sub.add_parser("scan", help="List sorries in a file")
    scan_p.add_argument("lean_file", help="Path to .lean file")

    res_p = sub.add_parser("resolve", help="Resolve sorries in a file")
    res_p.add_argument("lean_file", help="Path to .lean file")
    res_p.add_argument("--provider", default=os.environ.get("LLM_PROVIDER", DEFAULT_PROVIDER),
                       help=f"Provider (default: {DEFAULT_PROVIDER}; env: LLM_PROVIDER)")
    res_p.add_argument("--line", type=int, default=0, help="Specific sorry line to resolve")
    res_p.add_argument("--interactive", "-i", action="store_true", help="Ask before each API call")
    res_p.add_argument("--dry-run", "-n", action="store_true", help="Show prompts without sending")
    res_p.add_argument("--max-iterations", type=int, default=5, help="Max generate-compile cycles per sorry")
    res_p.add_argument("--temperature", type=float, default=0.4, help="LLM temperature (default 0.4)")
    res_p.add_argument("--lake-workdir", default="", help="lake build working directory")
    res_p.add_argument("--list-providers", action="store_true", help="List available providers and exit")

    args = parser.parse_args()

    if args.command == "scan":
        cmd_scan(args)
    elif args.command == "resolve":
        if args.list_providers:
            print("Available providers:")
            for name, spec in PROVIDERS.items():
                print(f"  {name:20s}  model={spec['model']:30s}  {spec['notes']}")
            return
        cmd_resolve(args)


if __name__ == "__main__":
    main()