# python/src/server/services/lean/self_healing_service.py import os import re from typing import Any from ...config.logfire_config import get_logger from ...utils import get_supabase_client from ..llm.clients import get_llm_client from ..prompt_service import prompt_service from ..storage.code_storage_service import add_code_examples_to_supabase from .compiler_service import lean_compiler_service logger = get_logger(__name__) class LeanSelfHealingService: """Orchestrates 2-stage proof repair (local and cloud escalation) & database seeding.""" def __init__(self, compiler_service=None, db_client=None): self.compiler_service = compiler_service or lean_compiler_service self.db_client = db_client or get_supabase_client() async def heal_proof(self, file_path: str, proof_context: str) -> dict[str, Any]: """Orchestrates compilation, capture errors, repair loop, and database seeding.""" # 1. Run initial build logger.info(f"Initial compilation check for {file_path}") result = self.compiler_service.run_build() if result["success"]: logger.info("Compilation passed initially. No healing needed.") return {"success": True, "stage": "none", "output": result} if "error" in result and not result.get("errors"): logger.error(f"Compiler runner error: {result['error']}. Aborting self-healing.") return {"success": False, "error": f"Compiler unavailable: {result['error']}"} original_code = "" if os.path.exists(file_path): with open(file_path, encoding="utf-8") as f: original_code = f.read() errors = result.get("errors", []) logger.warning(f"Compilation failed with {len(errors)} errors. Initiating repair loop.") # --- Stage 1: Local Ollama Repair --- logger.info("Starting Stage 1: Local Ollama Repair Loop (K=2)") last_error_context = str(errors) for attempt in range(1, 3): logger.info(f"Local Repair Attempt {attempt}/2") repaired_code = await self._request_llm_repair( provider="ollama", original_code=original_code, error_info=last_error_context, proof_context=proof_context, retry_count=attempt ) if repaired_code: # Write and compile self._write_code_to_file(file_path, repaired_code) build_res = self.compiler_service.run_build() if build_res["success"]: logger.info("Stage 1 (Local Ollama) repair succeeded!") await self._seed_successful_proof(file_path, repaired_code, "local_ollama") return {"success": True, "stage": "local_ollama", "attempts": attempt} else: last_error_context = str(build_res.get("errors", [])) logger.warning(f"Attempt {attempt} failed to compile: {last_error_context}") # --- Stage 2: Cloud Escalation --- logger.info("Stage 1 failed. Initiating Stage 2: Cloud Escalation") repaired_code_cloud = await self._request_llm_repair( provider="google", # Default cloud provider original_code=original_code, error_info=last_error_context, proof_context=proof_context, retry_count=2 ) if repaired_code_cloud: self._write_code_to_file(file_path, repaired_code_cloud) build_res = self.compiler_service.run_build() if build_res["success"]: logger.info("Stage 2 (Cloud Escalation) repair succeeded!") await self._seed_successful_proof(file_path, repaired_code_cloud, "cloud_escalation") return {"success": True, "stage": "cloud_escalation"} else: logger.error(f"Stage 2 failed to compile: {build_res.get('errors')}") return {"success": False, "error": "Self-healing failed to resolve compilation errors"} async def _request_llm_repair( self, provider: str, original_code: str, error_info: str, proof_context: str, retry_count: int ) -> str | None: """Requests repair from the specified LLM provider.""" system_prompt = prompt_service.get_prompt( "LEAN_4_DEVELOPER_ASSISTANT", "You are a Lean 4 logic programming copilot. Your objective is to help write correct formal proofs, choose mathematical tactics (like intro, rfl, simp, induction), and repair compile failures. Ensure output is syntactically valid Lean 4 code." ) user_prompt = ( f"We have a Lean 4 compile failure.\n\n" f"Here is the compiler error information:\n{error_info}\n\n" f"Original proof context:\n{proof_context}\n\n" f"Current file code:\n{original_code}\n\n" f"Please repair the code. Return only the complete corrected Lean 4 code wrapped in standard ```lean ... ``` blocks." ) try: from ...config.model_ssot import SYSTEM_MODELS cloud_model = SYSTEM_MODELS["DEFAULT_PRO"].split("/")[-1] async with get_llm_client(provider=provider) as client: response = await client.chat.completions.create( model="gemma3" if provider == "ollama" else cloud_model, messages=[ {"role": "system", "content": system_prompt}, {"role": "user", "content": user_prompt} ], extra_body={"retry_count": retry_count} ) raw_content = response.choices[0].message.content return self._extract_lean_code(raw_content) except Exception as e: logger.error(f"LLM call to provider {provider} failed: {e}") return None def _extract_lean_code(self, raw_content: str) -> str: """Extracts code blocks within ```lean ... ```.""" match = re.search(r"```lean\n(.*?)\n```", raw_content, re.DOTALL | re.IGNORECASE) if match: return match.group(1).strip() # Fallback to general code blocks if specific language is missing match_general = re.search(r"```\n(.*?)\n```", raw_content, re.DOTALL) if match_general: return match_general.group(1).strip() return raw_content.strip() def _write_code_to_file(self, file_path: str, code: str) -> None: """Writes the repaired code to file.""" os.makedirs(os.path.dirname(file_path), exist_ok=True) with open(file_path, "w", encoding="utf-8") as f: f.write(code) async def _seed_successful_proof(self, file_path: str, code: str, method: str) -> None: """Seeds successful repairs into the code_examples database for RAG.""" try: logger.info(f"Seeding successful proof into code_examples for {file_path}") url = f"file://{os.path.abspath(file_path)}" meta = {"source_id": "lean_self_healing", "language": "lean", "method": method} await add_code_examples_to_supabase( client=self.db_client, urls=[url], chunk_numbers=[0], code_examples=[code], summaries=[f"Successfully healed Lean 4 proof for {os.path.basename(file_path)} via {method}"], metadatas=[meta] ) except Exception as e: logger.error(f"Failed to seed proof to database: {e}") self_healing_service = LeanSelfHealingService()