File size: 6,663 Bytes
3962d8b | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 | import time
from datetime import datetime
from typing import Iterator
from .code_generator import CodeGenerator
from .code_executor import CodeExecutor
from .error_parser import parse_error
from .reflection import ReflectionModule
from .task_classifier import TaskClassifier
from .multi_language_agent import MultiLanguageAgent
class AgentLoopController:
def __init__(self, api_key: str, model: str = "llama-3.3-70b-versatile"):
self.generator = CodeGenerator(api_key, model)
self.executor = CodeExecutor()
self.reflector = ReflectionModule(api_key, model)
self.classifier = TaskClassifier()
self.multi_agent = MultiLanguageAgent(api_key, model)
def run(
self,
task: str,
test_cases: str,
max_iterations: int = 5,
timeout: int = 10,
) -> Iterator[dict]:
start_time = time.time()
yield self._evt("start", f"Agent started - max {max_iterations} iterations, {timeout}s timeout")
yield self._evt("classifying", "Detecting task type and target languages...")
classification = self.classifier.classify(task)
mode = classification["mode"]
languages = classification.get("languages", ["Python"])
yield self._evt(
"classified",
f"Mode: {mode.upper()} | Languages: {', '.join(languages)}",
mode=mode,
languages=languages,
)
if mode == "analyze":
yield from self._run_analyze(task)
else:
yield from self._run_execute(task, test_cases, max_iterations, timeout, start_time)
def _run_analyze(self, task: str) -> Iterator[dict]:
yield self._evt("analyzing", "Generating multi-language analysis with LLM")
try:
result = self.multi_agent.generate(task)
langs = [impl["language"] for impl in result.get("implementations", [])]
label = ", ".join(langs) if langs else "response ready"
yield self._evt("analysis_ready", f"Analysis complete: {label}", analysis=result)
yield self._evt(
"complete",
"Analysis finished.",
success=True,
mode="analyze",
final_code="",
total_time=0,
iterations_count=0,
iterations=[],
analysis=result,
)
except Exception as exc:
yield self._evt("failed", f"Analysis failed: {exc}", error=str(exc))
def _run_execute(
self, task: str, test_cases: str, max_iterations: int, timeout: int, start_time: float
) -> Iterator[dict]:
all_iterations: list[dict] = []
yield self._evt("generating", "Generating initial code with LLM")
try:
code, _ = self.generator.generate(task, test_cases)
except Exception as exc:
yield self._evt("failed", f"LLM call failed: {exc}", error=str(exc))
return
yield self._evt("code_ready", f"Initial code ready ({self._line_count(code)} lines)", code=code, iteration=0)
for iteration in range(1, max_iterations + 1):
yield self._evt(
"executing",
f"Executing code - iteration {iteration}/{max_iterations}",
code=code,
iteration=iteration,
)
exec_result = self.executor.execute(code, test_cases, timeout)
iter_data: dict = {
"iteration": iteration,
"code": code,
"execution": exec_result,
"timestamp": self._ts(),
}
if exec_result["success"]:
iter_data["status"] = "success"
all_iterations.append(iter_data)
yield self._evt(
"success",
f"Code passed on iteration {iteration}",
code=code,
iteration=iteration,
execution=exec_result,
iteration_data=iter_data,
)
yield self._evt(
"complete",
f"Task completed successfully in {iteration} iteration(s).",
iterations=all_iterations,
final_code=code,
success=True,
mode="execute",
total_time=round(time.time() - start_time, 2),
iterations_count=iteration,
)
return
error_info = parse_error(exec_result)
iter_data["status"] = "failed"
iter_data["error"] = error_info
all_iterations.append(iter_data)
yield self._evt(
"error",
f"Error on iteration {iteration}: {error_info['type']}: {error_info['message'][:120]}",
code=code,
iteration=iteration,
execution=exec_result,
error=error_info,
iteration_data=iter_data,
)
if iteration < max_iterations:
yield self._evt("reflecting", "Sending error to LLM for reflection and fix", iteration=iteration)
try:
new_code, reflection = self.reflector.reflect_and_refine(
task, code, exec_result, error_info
)
iter_data["reflection"] = reflection
code = new_code
yield self._evt(
"refined",
"Code refined based on error feedback",
new_code=new_code,
reflection=reflection,
iteration=iteration,
)
except Exception as exc:
yield self._evt("error", f"Reflection failed: {exc}", iteration=iteration)
yield self._evt(
"complete",
f"Reached max iterations ({max_iterations}) without full success.",
iterations=all_iterations,
final_code=code,
success=False,
mode="execute",
total_time=round(time.time() - start_time, 2),
iterations_count=max_iterations,
)
def _evt(self, event_type: str, message: str, **kwargs) -> dict:
return {"type": event_type, "message": message, "timestamp": self._ts(), **kwargs}
@staticmethod
def _ts() -> str:
return datetime.now().strftime("%H:%M:%S")
@staticmethod
def _line_count(code: str) -> int:
return len(code.strip().splitlines())
|