feat: Implement MemorySaver

This commit is contained in:
2026-02-27 00:47:16 -05:00
parent 4b2d02eca8
commit da2c2a1528
3 changed files with 144 additions and 7 deletions

View File

@@ -1,9 +1,12 @@
from langgraph.graph import StateGraph, START, END
from langchain_core.messages import SystemMessage, ToolMessage, BaseMessage
import logging
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
@@ -27,6 +30,7 @@ class Agent:
self,
model: Any,
config: AgentConfig | None = None,
checkpointer: InMemorySaver | None = None,
tools: list | None = None
):
"""
@@ -43,6 +47,7 @@ class Agent:
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
@@ -67,7 +72,7 @@ class Agent:
"""Expone el grafo para inspección en tests."""
return self._build_agent()
def invoke(self, state: dict, config: dict | None = None) -> dict:
def invoke(self, state: dict, config: RunnableConfig | None = None) -> dict:
"""
Ejecuta el agente con el estado inicial.
@@ -78,6 +83,9 @@ class Agent:
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")
@@ -104,7 +112,7 @@ class Agent:
agent = self._build_agent()
return agent.compile()
return agent.compile(checkpointer=self._checkpointer)
def _build_agent(self) -> StateGraph:
"""Construye el grafo de estado completo."""