Upload arcgis_test_agent/agent.py

arcgis_test_agent/agent.py

@@ -0,0 +1,179 @@
+"""
+Core agent loop for the ArcGIS Test Script Generator.
+
+This module implements the main agent that orchestrates:
+  1. Multi-turn conversations with Azure OpenAI
+  2. Tool dispatching
+  3. Output extraction and file writing
+"""
+
+import json
+import logging
+import time
+from dataclasses import dataclass, field
+from typing import Optional
+
+from openai import AzureOpenAI
+
+from .config import AzureConfig, AgentConfig, load_config
+from .tools import TOOL_SCHEMAS, ToolHandlers
+from .prompts import SYSTEM_PROMPT
+
+logger = logging.getLogger(__name__)
+
+
+@dataclass
+class AgentResult:
+    """Result from an agent run."""
+    success: bool
+    script: Optional[str] = None        # The generated test script
+    tool_name: str = ""                  # Tool we generated tests for
+    iterations: int = 0                  # How many loop iterations
+    tool_calls_made: list = field(default_factory=list)  # Log of all tool calls
+    error: Optional[str] = None         # Error message if failed
+    duration_seconds: float = 0.0
+
+
+class TestGeneratorAgent:
+    """
+    The main agent that generates ArcGIS test scripts.
+
+    Usage:
+        agent = TestGeneratorAgent()
+        result = agent.generate("Buffer")
+        if result.success:
+            print(result.script)
+    """
+
+    def __init__(self, azure_config: Optional[AzureConfig] = None,
+                 agent_config: Optional[AgentConfig] = None):
+        if azure_config is None or agent_config is None:
+            _azure, _agent = load_config()
+            azure_config = azure_config or _azure
+            agent_config = agent_config or _agent
+
+        self.azure_config = azure_config
+        self.agent_config = agent_config
+
+        self.client = AzureOpenAI(
+            azure_endpoint=azure_config.endpoint,
+            api_key=azure_config.api_key,
+            api_version=azure_config.api_version,
+        )
+
+        self.tools = ToolHandlers(agent_config)
+
+        logger.info(
+            f"Agent initialized. Model: {azure_config.deployment_name}, "
+            f"Max iterations: {agent_config.max_iterations}"
+        )
+
+    def generate(self, tool_name: str, additional_context: str = "") -> AgentResult:
+        """Generate a test script for the given ArcGIS tool."""
+        start_time = time.time()
+        tool_calls_log = []
+
+        user_message = f"Generate a complete test script for the ArcGIS tool: {tool_name}"
+        if additional_context:
+            user_message += f"\n\nAdditional context: {additional_context}"
+
+        messages = [
+            {"role": "system", "content": SYSTEM_PROMPT},
+            {"role": "user", "content": user_message},
+        ]
+
+        logger.info(f"Starting generation for tool: {tool_name}")
+
+        for iteration in range(1, self.agent_config.max_iterations + 1):
+            logger.debug(f"Iteration {iteration}/{self.agent_config.max_iterations}")
+
+            try:
+                response = self.client.chat.completions.create(
+                    model=self.azure_config.deployment_name,
+                    messages=messages,
+                    tools=TOOL_SCHEMAS,
+                    tool_choice="auto",
+                    temperature=self.agent_config.temperature,
+                )
+            except Exception as e:
+                logger.error(f"API call failed at iteration {iteration}: {e}")
+                return AgentResult(
+                    success=False, tool_name=tool_name, iterations=iteration,
+                    tool_calls_made=tool_calls_log, error=f"API error: {str(e)}",
+                    duration_seconds=time.time() - start_time,
+                )
+
+            msg = response.choices[0].message
+            messages.append(msg)
+
+            if msg.tool_calls:
+                for tc in msg.tool_calls:
+                    func_name = tc.function.name
+                    func_args = json.loads(tc.function.arguments)
+
+                    logger.info(f"  Tool call: {func_name}({json.dumps(func_args)[:100]}...)")
+                    tool_calls_log.append({
+                        "iteration": iteration, "tool": func_name, "arguments": func_args,
+                    })
+
+                    result = self.tools.dispatch(func_name, func_args)
+                    messages.append({
+                        "role": "tool", "tool_call_id": tc.id,
+                        "content": json.dumps(result, default=str),
+                    })
+            else:
+                final_output = msg.content
+                if final_output:
+                    script = self._extract_script(final_output)
+                    duration = time.time() - start_time
+                    logger.info(
+                        f"Generation complete. {iteration} iterations, "
+                        f"{len(tool_calls_log)} tool calls, {duration:.1f}s"
+                    )
+                    return AgentResult(
+                        success=True, script=script, tool_name=tool_name,
+                        iterations=iteration, tool_calls_made=tool_calls_log,
+                        duration_seconds=duration,
+                    )
+                else:
+                    logger.warning("Empty response from model, retrying...")
+                    messages.append({
+                        "role": "user",
+                        "content": "Your response was empty. Please output the complete test script."
+                    })
+
+        duration = time.time() - start_time
+        logger.error(f"Max iterations ({self.agent_config.max_iterations}) reached without completion.")
+        return AgentResult(
+            success=False, tool_name=tool_name,
+            iterations=self.agent_config.max_iterations,
+            tool_calls_made=tool_calls_log,
+            error=f"Exceeded max iterations ({self.agent_config.max_iterations})",
+            duration_seconds=duration,
+        )
+
+    def _extract_script(self, raw_output: str) -> str:
+        """Extract clean Python code from the model's output."""
+        text = raw_output.strip()
+        if text.startswith("```python"):
+            text = text[len("```python"):].strip()
+        elif text.startswith("```"):
+            text = text[len("```"):].strip()
+        if text.endswith("```"):
+            text = text[:-3].strip()
+        return text
+
+    def generate_and_save(self, tool_name: str, output_path: Optional[str] = None,
+                          additional_context: str = "") -> AgentResult:
+        """Generate a test script and save it to disk."""
+        import os
+        result = self.generate(tool_name, additional_context)
+        if result.success and result.script:
+            if output_path is None:
+                os.makedirs(self.agent_config.output_dir, exist_ok=True)
+                safe_name = tool_name.lower().replace(" ", "_")
+                output_path = os.path.join(self.agent_config.output_dir, f"test_{safe_name}.py")
+            with open(output_path, "w") as f:
+                f.write(result.script)
+            logger.info(f"Test script saved to: {output_path}")
+        return result