feat: Add complete suit test
This commit is contained in:
@@ -1,81 +1,214 @@
|
||||
from langgraph.graph import StateGraph, START, END
|
||||
from langchain.messages import SystemMessage, ToolMessage
|
||||
from typing import Literal
|
||||
from langchain_core.messages import SystemMessage, ToolMessage, BaseMessage
|
||||
from typing import Literal, Callable, Any
|
||||
from dataclasses import dataclass, field
|
||||
from .schemas import MessagesState
|
||||
|
||||
|
||||
@dataclass
|
||||
class AgentConfig:
|
||||
"""Configuración inmutable del agente para fácil testing."""
|
||||
system_prompt: str = (
|
||||
"You are a helpful assistant tasked with "
|
||||
"performing arithmetic on a set of inputs."
|
||||
)
|
||||
max_iterations: int = 10
|
||||
timeout_seconds: float = 30.0
|
||||
|
||||
|
||||
class Agent:
|
||||
"""
|
||||
Agente conversacional basado en LangGraph.
|
||||
"""
|
||||
|
||||
def __init__(self, model):
|
||||
def __init__(
|
||||
self,
|
||||
model: Any,
|
||||
config: AgentConfig | None = None,
|
||||
tools: list | None = None
|
||||
):
|
||||
self._model = model
|
||||
self._tools = []
|
||||
self._tools_by_name = []
|
||||
self.model_with_tools = self._model.bind_tools(
|
||||
self._tools)
|
||||
self.agent = self._compiler_agent()
|
||||
self._config = config or AgentConfig()
|
||||
self._tools = tools or []
|
||||
self._tools_by_name: dict[str, Any] = {
|
||||
tool.name: tool for tool in self._tools
|
||||
}
|
||||
self._compiled_agent = self._compile_agent()
|
||||
|
||||
def invoke(self, messages):
|
||||
self.agent.invoke(
|
||||
messages
|
||||
)
|
||||
# Hooks para testing (monkey-patching friendly)
|
||||
self._pre_llm_hook: Callable | None = None
|
||||
self._post_tool_hook: Callable | None = None
|
||||
|
||||
def _compiler_agent(self):
|
||||
@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) -> 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 "messages" not in state:
|
||||
raise ValueError("State must contain 'messages' key")
|
||||
|
||||
return self._compiled_agent.invoke(state)
|
||||
|
||||
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()
|
||||
|
||||
def _build_agent(self):
|
||||
def _build_agent(self) -> StateGraph:
|
||||
"""Construye el grafo de estado completo."""
|
||||
agent_builder = StateGraph(MessagesState)
|
||||
agent_builder.add_node("llm_call", self.llm_call)
|
||||
agent_builder.add_node("tool_node", self.tool_node)
|
||||
|
||||
# Nodos
|
||||
agent_builder.add_node("llm_call", self._llm_call_node)
|
||||
agent_builder.add_node("tool_node", self._tool_node)
|
||||
|
||||
# Flujo
|
||||
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(self, state: dict):
|
||||
"""LLM decides whether to call a tool or not"""
|
||||
# ============ Nodos del Grafo ============
|
||||
|
||||
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": [
|
||||
self.model_with_tools.invoke(
|
||||
[
|
||||
SystemMessage(
|
||||
content=(
|
||||
"You are a helpful assistant tasked with "
|
||||
"performing arithmetic on a set of inputs."
|
||||
)
|
||||
)
|
||||
]
|
||||
+ state["messages"]
|
||||
)
|
||||
],
|
||||
"llm_calls": state.get('llm_calls', 0) + 1
|
||||
"messages": state["messages"] + [response],
|
||||
"llm_calls": new_calls
|
||||
}
|
||||
|
||||
def tool_node(self, state: dict):
|
||||
"""Performs the tool call"""
|
||||
|
||||
result = []
|
||||
for tool_call in state["messages"][-1].tool_calls:
|
||||
tool = self._tools_by_name[tool_call["name"]]
|
||||
observation = tool.invoke(tool_call["args"])
|
||||
result.append(
|
||||
ToolMessage(
|
||||
content=observation,
|
||||
tool_call_id=tool_call["id"]
|
||||
))
|
||||
return {"messages": result}
|
||||
|
||||
def should_continue(state: MessagesState) -> Literal["tool_node", END]:
|
||||
def _tool_node(self, state: MessagesState) -> dict:
|
||||
"""
|
||||
Decide if we should continue the loop or stop base
|
||||
d upon whether the LLM made a tool call
|
||||
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", "unknown")
|
||||
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,
|
||||
status="error"
|
||||
)
|
||||
|
||||
try:
|
||||
tool = self._tools_by_name[tool_name]
|
||||
observation = tool.invoke(tool_args)
|
||||
return ToolMessage(
|
||||
content=str(observation),
|
||||
tool_call_id=tool_id,
|
||||
status="success"
|
||||
)
|
||||
except Exception as e:
|
||||
return ToolMessage(
|
||||
content=f"Error executing tool: {str(e)}",
|
||||
tool_call_id=tool_id,
|
||||
status="error"
|
||||
)
|
||||
|
||||
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
|
||||
|
||||
messages = state["messages"]
|
||||
last_message = messages[-1]
|
||||
|
||||
if last_message.tool_calls:
|
||||
if hasattr(last_message, "tool_calls") and last_message.tool_calls:
|
||||
return "tool_node"
|
||||
|
||||
return END
|
||||
|
||||
Reference in New Issue
Block a user