NaliiaBot/src/naliiabot/bot/agent/agent.py

from langgraph.graph import StateGraph, START, END
from langchain_core.messages import SystemMessage, ToolMessage, BaseMessage
from langgraph.checkpoint.memory import InMemorySaver
from langchain_core.runnables import RunnableConfig
from typing import Literal, Callable, Any
from .schemas import MessagesState
from dataclasses import dataclass
import logging



@dataclass
class AgentConfig:
    """Configuración inmutable del agente para fácil testing."""
    system_prompt: str = (
        "helpful assistant that can call tools when needed. Always respond with a message. "
    )
    max_iterations: int = 10
    timeout_seconds: float = 30.0


class Agent:
    """
    Agente conversacional basado en LangGraph.
    
    Soporta tools pasadas como lista de instancias.
    """

    def __init__(
        self,
        model: Any,
        config: AgentConfig | None = None,
        checkpointer: InMemorySaver | None = None, 
        tools: list | None = None
    ):
        """
        Inicializa el agente.
        
        Args:
            model: Modelo LLM a usar
            config: Configuración del agente (AgentConfig)
            tools: Lista de instancias de herramientas o None
                - Si es None: sin herramientas
                - Si es list: lista de herramientas
        """

        self._model = model
        self._config = config or AgentConfig()
        self._tools = tools or []
        self._checkpointer = checkpointer or InMemorySaver()
            
        self._tools_by_name: dict[str, Any] = {
            tool.name: tool for tool in self._tools
        }
        self._compiled_agent = self._compile_agent()

        self._pre_llm_hook: Callable | None = None
        self._post_tool_hook: Callable | None = None

    @property
    def tools(self) -> list:
        """Retorna copia de herramientas registradas."""
        return self._tools.copy()

    @property
    def config(self) -> AgentConfig:
        """Retorna configuración actual."""
        return self._config

    @property
    def graph(self) -> StateGraph:
        """Expone el grafo para inspección en tests."""
        return self._build_agent()

    def invoke(self, state: dict, config: RunnableConfig | None = None) -> dict:
        """
        Ejecuta el agente con el estado inicial.

        Args:
            state: Dict con key "messages" y opcionalmente "llm_calls"

        Returns:
            Estado final después de la ejecución
        """

        if not config:
            config : RunnableConfig = {"configurable": {"thread_id": "1"}}

        if "messages" not in state:
            raise ValueError("State must contain 'messages' key")

        return self._compiled_agent.invoke(state, config=config)

    def stream(self, state: dict):
        """Streaming de la ejecución para debugging."""
        yield from self._compiled_agent.stream(state)

    # ============ Hooks para Testing ============

    def set_pre_llm_hook(self, hook: Callable[[dict], None]):
        """Registra hook pre-llm para testing."""

        self._pre_llm_hook = hook

    def set_post_tool_hook(self, hook: Callable[[dict, Any], None]):
        """Registra hook post-tool para testing."""

        self._post_tool_hook = hook

    def _compile_agent(self):
        """Compila el grafo una sola vez."""

        agent = self._build_agent()

        return agent.compile(checkpointer=self._checkpointer)

    def _build_agent(self) -> StateGraph:
        """Construye el grafo de estado completo."""
        agent_builder = StateGraph(MessagesState)

        agent_builder.add_node("llm_call", self._llm_call_node)
        agent_builder.add_node("tool_node", self._tool_node)

        agent_builder.add_edge(START, "llm_call")
        agent_builder.add_conditional_edges(
            "llm_call",
            self._should_continue,
            {
                "tool_node": "tool_node",
                END: END
            }
        )
        agent_builder.add_edge("tool_node", "llm_call")

        return agent_builder

    def _llm_call_node(self, state: MessagesState) -> dict:
        """
        Nodo LLM: Decide siguiente acción o responde.

        State transitions:
        - Si tool_calls → tool_node
        - Si no → END
        """
        if self._pre_llm_hook:
            self._pre_llm_hook(state)

        messages = [
            SystemMessage(content=self._config.system_prompt)
        ] + state["messages"]

        model_with_tools = self._model.bind_tools(self._tools)

        response = model_with_tools.invoke(messages)

        current_calls = state.get("llm_calls", 0)
        new_calls = current_calls + 1

        if new_calls > self._config.max_iterations:
            raise RuntimeError(
                f"Max iterations ({self._config.max_iterations}) exceeded"
            )

        return {
            "messages": state["messages"] + [response],
            "llm_calls": new_calls
        }

    def _tool_node(self, state: MessagesState) -> dict:
        """
        Nodo Tools: Ejecuta herramientas solicitadas por el LLM.
        """
        last_message = state["messages"][-1]

        if not hasattr(last_message, "tool_calls"):
            return {"messages": state["messages"]}

        tool_results: list[BaseMessage] = []

        for tool_call in last_message.tool_calls:
            result = self._execute_tool(tool_call)
            tool_results.append(result)

            if self._post_tool_hook:
                self._post_tool_hook(state, result)

        return {
            "messages": state["messages"] + tool_results
        }

    def _execute_tool(self, tool_call: dict) -> ToolMessage:
        """
        Ejecuta una herramienta individual con manejo de errores.
        """
        tool_name = tool_call.get("name")
        tool_id = tool_call.get("id")
        tool_args = tool_call.get("args", {})

        if tool_name not in self._tools_by_name:
            return ToolMessage(
                content=f"Error: Tool '{tool_name}' not found",
                tool_call_id=tool_id,
                name=tool_name,
            )

        try:
            tool = self._tools_by_name[tool_name]
            observation = tool.invoke(tool_args)

            return ToolMessage(
                content=str(observation),
                tool_call_id=tool_id,
                name=tool_name)

        except Exception as e:
            return ToolMessage(
                content=f"Error executing tool: {str(e)}",
                tool_call_id=tool_id,
                name=tool_name,
            )

    def _should_continue(
            self, state: MessagesState
    ) -> Literal["tool_node", END]:
        """
        Condición de transición: ¿El LLM solicitó herramientas?
        """
        messages = state.get("messages", [])
        if not messages:
            return END

        last_message = messages[-1]

        if hasattr(last_message, "tool_calls") and last_message.tool_calls:
            return "tool_node"

        return END