import re import json from typing import Optional, List, Dict, Any from .llm import LLMClient from .compiler import CppCompiler class BaseAgent: def __init__(self, llm: LLMClient, compiler: CppCompiler): self.llm = llm self.compiler = compiler def extract_code_blocks(self, text: str, expected_count: int = None) -> List[str]: pattern = re.compile(r"```(?:cpp|c\+\+)?\s*\n(.*?)```", re.DOTALL | re.IGNORECASE) matches = pattern.findall(text) blocks = [m.strip() for m in matches if m.strip()] if expected_count is not None and len(blocks) != expected_count: pass return blocks def extract_json(self, text: str) -> Optional[Dict[str, Any]]: text = text.strip() if text.startswith("```"): text = re.sub(r"^```(?:json)?\s*\n", "", text, flags=re.IGNORECASE) text = re.sub(r"\n```$", "", text, flags=re.IGNORECASE) try: return json.loads(text) except json.JSONDecodeError: match = re.search(r"\{.*\}", text, re.DOTALL) if match: try: return json.loads(match.group(0)) except json.JSONDecodeError: return None return None class CoderAgent(BaseAgent): def run(self, problem: str, objective: str, max_retries: int = 3) -> Dict[str, Any]: prompt = self.llm.get_coder_prompt(problem, objective) code = "" response = "" compile_result = {"compiled": False, "compile_error": "No attempts made", "compile_stderr": ""} for attempt in range(1, max_retries + 1): response = self.llm.generate(prompt, max_tokens=4096, temperature=0.1) code_blocks = self.extract_code_blocks(response) if not code_blocks: code_blocks = [response] code = code_blocks[0] compile_result = self.compiler.compile("solution", code) if compile_result["compiled"]: run_result = self.compiler.run(compile_result["binary_path"]) return { "code": code, "compile_result": compile_result, "run_result": run_result, "attempts": attempt, "agent": "coder", } prompt = f"{prompt}\n\nPREVIOUS ATTEMPT FAILED WITH COMPILATION ERROR:\n{compile_result['compile_error']}\n\nFix the error and output ONLY valid C++ code." return { "code": code if code else response, "compile_result": compile_result, "run_result": None, "attempts": max_retries, "agent": "coder", "failed": True, } class ResearcherAgent(BaseAgent): def run(self, problem: str, baseline_code: str, objective: str, n: int = 3) -> List[Dict[str, Any]]: prompt = self.llm.get_researcher_prompt(problem, baseline_code, objective, n) response = self.llm.generate(prompt, max_tokens=8192, temperature=0.4) code_blocks = self.extract_code_blocks(response, expected_count=n) if not code_blocks: code_blocks = [response] algorithms = [] for i, code in enumerate(code_blocks[:n]): label = f"algo_{i+1}" compile_result = self.compiler.compile(label, code) run_result = None if compile_result["compiled"]: run_result = self.compiler.run(compile_result["binary_path"]) algorithms.append({ "index": i + 1, "label": label, "code": code, "compile_result": compile_result, "run_result": run_result, "agent": "researcher", }) while len(algorithms) < n: algorithms.append({ "index": len(algorithms) + 1, "label": f"algo_{len(algorithms)+1}", "code": "", "compile_result": {"compiled": False, "compile_error": "No code generated"}, "run_result": None, "agent": "researcher", "failed": True, }) return algorithms class ValidatorAgent(BaseAgent): def _build_metrics_table(self, algorithms: List[Dict[str, Any]], baseline: Dict[str, Any]) -> str: lines = ["| Algorithm | Compiled | Return Code | Exec Time (s) | Memory (KB) | Stdout |"] lines.append("|-----------|----------|-------------|---------------|-------------|--------|") b_run = baseline.get("run_result") or {} lines.append(f"| Baseline | {'Yes' if baseline.get('compile_result', {}).get('compiled') else 'No'} | {'0' if b_run.get('ran') else 'N/A'} | {b_run.get('execution_time_s', 'N/A')} | {b_run.get('memory_kb', 'N/A')} | (baseline) |") for algo in algorithms: compiled = "Yes" if (algo.get("compile_result") or {}).get("compiled") else "No" a_run = algo.get("run_result") or {} if a_run.get("ran"): rc = str(a_run.get("returncode")) et = str(a_run.get("execution_time_s")) mem = str(a_run.get("memory_kb")) else: rc = "N/A" et = "N/A" mem = "N/A" stdout_preview = (a_run.get("stdout", "") or "")[:30].replace("\n", " ") lines.append(f"| Algo {algo['index']} | {compiled} | {rc} | {et} | {mem} | {stdout_preview} |") return "\n".join(lines) def run(self, problem: str, objective: str, user_metric: str, algorithms: List[Dict[str, Any]], baseline: Dict[str, Any]) -> Dict[str, Any]: metrics_table = self._build_metrics_table(algorithms, baseline) prompt = self.llm.get_validator_prompt(problem, objective, user_metric, metrics_table, len(algorithms)) response = self.llm.generate(prompt, max_tokens=1024, temperature=0.0) result = self.extract_json(response) if result is None: result = {"winner": 1, "failed": False, "reason": "Default fallback selection"} return { "validation": result, "metrics_table": metrics_table, "agent": "validator", } class RealWorldAgent(BaseAgent): def run(self, problem: str, winner_code: str, baseline_code: str, objective: str, n: int = 3) -> Dict[str, Any]: prompt = self.llm.get_realworld_prompt(problem, winner_code, baseline_code, objective, n) response = self.llm.generate(prompt, max_tokens=8192, temperature=0.3) code_blocks = self.extract_code_blocks(response, expected_count=n) if not code_blocks: code_blocks = [response] scenarios = [] for i, code in enumerate(code_blocks[:n]): label = f"scenario_{i+1}" compile_result = self.compiler.compile(label, code) run_result = None if compile_result["compiled"]: run_result = self.compiler.run(compile_result["binary_path"], timeout=60) scenarios.append({ "index": i + 1, "label": label, "code": code, "compile_result": compile_result, "run_result": run_result, }) while len(scenarios) < n: scenarios.append({ "index": len(scenarios) + 1, "label": f"scenario_{len(scenarios)+1}", "code": "", "compile_result": {"compiled": False, "compile_error": "No code generated"}, "run_result": None, "failed": True, }) return { "scenarios": scenarios, "agent": "realworld", }