src/agents.py

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",
        }