src/agents/coder_agent.py

"""
Coder Agent: Implements the plan as executable Python code.

This agent generates Python code that implements all steps in the current plan.
"""

from langchain_core.messages import AIMessage

from ..utils.code_execution import execute_with_debug
from ..utils.formatters import extract_code, format_data_descriptions, format_plan
from ..utils.state import DSStarState


def coder_node(state: DSStarState) -> dict:
    """
    Coder Agent Node: Generates and executes Python code for the plan.

    On first call: Generates code implementing all plan steps
    On subsequent calls: Updates code to include new plan steps

    Args:
        state: Current DSStarState

    Returns:
        Dictionary with updated state fields:
        - current_code: Generated Python code
        - execution_result: Output from code execution
        - messages: Agent communication messages
        - next: Next node to visit ("verifier")
    """
    print("=" * 60)
    print("CODER AGENT STARTING...")
    print("=" * 60)

    data_context = format_data_descriptions(state["data_descriptions"])
    plan_text = format_plan(state["plan"])

    is_initial = state["current_code"] == ""

    if is_initial:
        print("Generating INITIAL code implementation...")
        prompt = f"""You are an expert Python developer for data science.

Available Data Files:
{data_context}

Plan to Implement:
{plan_text}

Task: Write a Python script that implements ALL steps in the plan.

Requirements:
- Use pandas for data manipulation
- All files are in 'data/' directory
- Print intermediate results for each step
- No try-except blocks
- Clean, readable code

Provide ONLY the Python code in a markdown code block."""
    else:
        print(f"Updating code to implement {len(state['plan'])} steps...")
        prompt = f"""You are an expert Python developer for data science.

Available Data Files:
{data_context}

Complete Plan:
{plan_text}

Previous Code:
{state["current_code"]}

Task: Update the code to implement the COMPLETE current plan.
Build upon the previous code, extending it to include all plan steps.

Requirements:
- Use pandas for data manipulation
- All files are in 'data/' directory
- Print intermediate and final results
- No try-except blocks
- Clean, readable code

Provide ONLY the updated Python code in a markdown code block."""

    try:
        # Get LLM response
        response = state["llm"].invoke(prompt)

        # Handle different response formats
        if hasattr(response, "content") and isinstance(response.content, list):
            from ..utils.formatters import gemini_text

            response_text = gemini_text(response)
        elif hasattr(response, "content"):
            response_text = response.content
        else:
            response_text = str(response)

        code = extract_code(response_text)

        print("\nGenerated Code:")
        print("-" * 60)
        print(code[:200] + "..." if len(code) > 200 else code)
        print("-" * 60)

        print("\nExecuting code...")

        # Execute with debugging
        result = execute_with_debug(
            code, state["llm"], is_analysis=False, data_context=data_context
        )

        print("\nExecution Result:")
        print("-" * 60)
        print(result[:200] + "..." if len(result) > 200 else result)
        print("-" * 60)
        print("=" * 60)

        return {
            "current_code": code,
            "execution_result": result,
            "messages": [AIMessage(content="Code executed")],
            "next": "verifier",
        }

    except Exception as e:
        print(f"\n✗ Coder error: {str(e)}")
        return {
            "messages": [AIMessage(content=f"Coder error: {str(e)}")],
            "next": "__end__",
        }


# Standalone test function
def test_coder(llm, query: str, data_descriptions: dict, plan: list):
    """
    Test the coder agent independently.

    Args:
        llm: LLM instance
        query: User query
        data_descriptions: Dict of filename -> description
        plan: List of plan steps to implement

    Returns:
        Dictionary with coder results
    """
    # Create minimal test state
    test_state = {
        "llm": llm,
        "query": query,
        "data_descriptions": data_descriptions,
        "plan": plan,
        "current_code": "",
        "execution_result": "",
        "is_sufficient": False,
        "router_decision": "",
        "iteration": 0,
        "max_iterations": 20,
        "messages": [],
        "next": "coder",
    }

    result = coder_node(test_state)

    print("\n" + "=" * 60)
    print("CODER TEST RESULTS")
    print("=" * 60)
    print("Code:")
    print(result.get("current_code", "No code generated"))
    print("\nExecution Result:")
    print(result.get("execution_result", "No result"))

    return result