283 lines
9.3 KiB
Python
283 lines
9.3 KiB
Python
from langgraph.graph import StateGraph, START, END
|
|
from langchain_core.messages import SystemMessage, ToolMessage, BaseMessage
|
|
from langgraph.checkpoint.base import BaseCheckpointSaver
|
|
from langgraph.checkpoint.memory import InMemorySaver
|
|
from langchain_core.runnables import RunnableConfig
|
|
from typing import Literal, Callable, Any, Union
|
|
from .schemas import MessagesState
|
|
from dataclasses import dataclass
|
|
import json
|
|
|
|
|
|
@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 = 200
|
|
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: Union[BaseCheckpointSaver, 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 checkpointer(self):
|
|
"""Retorna el checkpointer para acceso al estado."""
|
|
return self._checkpointer
|
|
|
|
@property
|
|
def graph(self) -> StateGraph:
|
|
"""Expone el grafo para inspección en tests."""
|
|
return self._build_agent()
|
|
|
|
async def ainvoke(self, state: dict, config: RunnableConfig | None = None) -> dict:
|
|
"""
|
|
Versión asíncrona de invoke, necesaria para checkpointers de DB.
|
|
"""
|
|
if not config:
|
|
config: RunnableConfig = {"configurable": {"thread_id": "default"}}
|
|
|
|
if "messages" not in state:
|
|
raise ValueError("State must contain 'messages' key")
|
|
|
|
# Llamamos al método ainvoke del grafo compilado
|
|
return await self._compiled_agent.ainvoke(state, config=config)
|
|
|
|
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("identify_customer", self._identify_customer_node)
|
|
agent_builder.add_node("llm_call", self._llm_call_node)
|
|
agent_builder.add_node("tool_node", self._tool_node)
|
|
|
|
agent_builder.add_edge(START, "identify_customer")
|
|
agent_builder.add_edge("identify_customer", "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 _identify_customer_node(self, state: MessagesState) -> dict:
|
|
"""
|
|
Nodo de identificación del cliente.
|
|
Busca el cliente por teléfono y actualiza el estado.
|
|
"""
|
|
from ..logger import logger
|
|
|
|
customer_phone = state.get("customer_phone", "")
|
|
if not customer_phone:
|
|
return {"customer_name": "Cliente Anónimo", "customer_id": 0}
|
|
|
|
find_customer_tool = self._tools_by_name.get("find_customer_by_identifier")
|
|
if not find_customer_tool:
|
|
logger.warning("Tool find_customer_by_identifier not found")
|
|
return {"customer_name": "Cliente Final", "customer_id": 0}
|
|
|
|
try:
|
|
result = find_customer_tool.invoke({"identifier": customer_phone})
|
|
customer_data = json.loads(result) if result else {}
|
|
|
|
if isinstance(customer_data, list):
|
|
customer_data = customer_data[0] if customer_data else {}
|
|
|
|
if customer_data and customer_data.get("id"):
|
|
customer_name = customer_data.get("party.", {}).get(
|
|
"name", "Cliente Final"
|
|
)
|
|
customer_id = customer_data.get("id", 0)
|
|
logger.info(
|
|
f"Cliente identificado: {customer_name} (ID: {customer_id})"
|
|
)
|
|
return {"customer_name": customer_name, "customer_id": customer_id}
|
|
else:
|
|
logger.info(f"Cliente no encontrado para teléfono: {customer_phone}")
|
|
return {"customer_name": "Cliente Nuevo", "customer_id": 0}
|
|
except Exception as e:
|
|
logger.error(f"Error identificando cliente: {e}")
|
|
return {"customer_name": "Cliente Final", "customer_id": 0}
|
|
|
|
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
|