src/agent.py

"""Coding agent implementation with iterative LLM loop.

Main agent loop that orchestrates LLM calls, tool execution,
and conversation memory management.
"""

import json
from openai import OpenAI
from e2b_code_interpreter import Sandbox
from .tools import execute_tool


def coding_agent(
    client: OpenAI,
    query: str,
    system: str,
    tools: dict,
    tools_schemas: list,
    sbx: Sandbox,
    messages: list = None,
    max_steps: int = 5
):
    """Run coding agent with iterative tool calling loop.

    Args:
        client: OpenAI client instance
        query: User query/prompt
        system: System prompt defining agent behavior
        tools: Dict mapping tool names to functions
        tools_schemas: List of OpenAI function schemas
        sbx: E2B Sandbox instance for code execution
        messages: Optional existing message history
        max_steps: Maximum iteration steps (default 5)

    Returns:
        Tuple of (messages list, metadata dict)
        - messages: Full conversation history
        - metadata: Accumulated metadata (especially images)
    """
    if messages is None:
        messages = []
    messages.append({"role": "user", "content": query})

    metadata = {}
    steps = 0

    while steps < max_steps:
        # Call LLM with current conversation state
        response = client.responses.create(
            model="gpt-4.1-mini",
            input=[
                {"role": "developer", "content": system},
                *messages
            ],
            tools=tools_schemas
        )

        has_function_call = False

        # Process all parts of the response
        for part in response.output:
            messages.append(part.to_dict())

            if part.type == "message":
                print(part.content)

            elif part.type == "function_call":
                has_function_call = True
                name = part.name
                args = part.arguments

                # Execute the tool and get results
                result, tool_metadata = execute_tool(name, args, tools, sbx=sbx)

                # Accumulate metadata (especially images)
                if "images" in tool_metadata:
                    metadata.setdefault("images", []).extend(tool_metadata["images"])
                if "error" in tool_metadata:
                    metadata["error"] = tool_metadata["error"]

                # Append function result to conversation
                messages.append({
                    "type": "function_call_output",
                    "call_id": part.call_id,
                    "output": json.dumps(result)
                })

        # Stop if no more function calls
        if not has_function_call:
            break

        steps += 1

    return messages, metadata