agents/code_agent.py

"""
Code Specialist Agent Node
==========================
Generates, reviews, refactors, and explains code in Python, JavaScript/TypeScript, SQL, and Shell.
"""
from __future__ import annotations

import structlog
from app.services.llm_gateway import llm_gateway
from agents.state import CopilotState

logger = structlog.get_logger(__name__)

CODE_SYSTEM_PROMPT = """You are a Senior Software Engineer AI Agent.
Your job is to generate clean, performant, secure, and fully documented code based on the user's request.

Rules:
1. Provide production-ready code with complete error handling.
2. Include brief explanations of key design choices.
3. Enclose all code blocks inside proper markdown triple backticks with language tags (e.g. ```python, ```typescript, ```sql).
"""


async def code_node(state: CopilotState) -> CopilotState:
    """
    Code Specialist Node.
    Generates and explains production code snippets.
    """
    query = state.get("query", "")
    logger.info("Code Agent executing", query=query)

    try:
        response = await llm_gateway.generate(
            prompt=f"User Request: {query}\nProvide solution code and explanation:",
            system_prompt=CODE_SYSTEM_PROMPT,
            temperature=0.1,
        )

        code_response = response.content.strip()

        state["retrieved_chunks"] = [{
            "document_id": "code_gen",
            "document_name": "Code Assistant",
            "text": code_response,
            "score": 1.0,
            "doc_type": "code"
        }]

        outputs = state.get("agent_outputs", [])
        outputs.append({
            "agent_name": "code",
            "content": "Generated code snippet and solution.",
        })
        state["agent_outputs"] = outputs
        state["active_agent"] = "code"

    except Exception as e:
        logger.error("Code Agent error", error=str(e))
        state["error"] = f"Code Agent Error: {str(e)}"
        state["retrieved_chunks"] = [{
            "document_name": "Code Agent Error",
            "text": f"Could not generate code: {str(e)}"
        }]

    return state