11"""Middleware for loading agent-specific long-term memory into the system prompt."""
22
33from collections .abc import Awaitable , Callable
4- from typing import TYPE_CHECKING , Any
5-
6- if TYPE_CHECKING :
7- from langgraph .runtime import Runtime
4+ from typing import NotRequired
85
6+ from deepagents .backends .protocol import BackendProtocol
97from langchain .agents .middleware .types import (
108 AgentMiddleware ,
119 AgentState ,
1210 ModelRequest ,
1311 ModelResponse ,
1412)
15- from typing_extensions import NotRequired , TypedDict
16-
17- from deepagents .backends .protocol import BackendProtocol
1813
1914
2015class AgentMemoryState (AgentState ):
@@ -83,6 +78,7 @@ class AgentMemoryState(AgentState):
8378</agent_memory>
8479"""
8580
81+
8682class AgentMemoryMiddleware (AgentMiddleware ):
8783 """Middleware for loading agent-specific long-term memory.
8884
@@ -105,7 +101,7 @@ class AgentMemoryMiddleware(AgentMiddleware):
105101 # Set up backend pointing to agent's directory
106102 agent_dir = Path.home() / ".deepagents" / "my-agent"
107103 backend = FilesystemBackend(root_dir=agent_dir)
108-
104+
109105 # Create middleware
110106 middleware = AgentMemoryMiddleware(backend=backend)
111107 ```
@@ -185,13 +181,17 @@ def wrap_model_call(
185181 """
186182 # Get agent memory from state
187183 agent_memory = request .state .get ("agent_memory" , "" )
188-
184+
189185 memory_section = self .system_prompt_template .format (agent_memory = agent_memory )
190186 if request .system_prompt :
191187 request .system_prompt = memory_section + "\n \n " + request .system_prompt
192188 else :
193189 request .system_prompt = memory_section
194- request .system_prompt = request .system_prompt + "\n \n " + LONGTERM_MEMORY_SYSTEM_PROMPT .format (memory_path = self .memory_path )
190+ request .system_prompt = (
191+ request .system_prompt
192+ + "\n \n "
193+ + LONGTERM_MEMORY_SYSTEM_PROMPT .format (memory_path = self .memory_path )
194+ )
195195
196196 return handler (request )
197197
@@ -211,12 +211,16 @@ async def awrap_model_call(
211211 """
212212 # Get agent memory from state
213213 agent_memory = request .state .get ("agent_memory" , "" )
214-
214+
215215 memory_section = self .system_prompt_template .format (agent_memory = agent_memory )
216216 if request .system_prompt :
217217 request .system_prompt = memory_section + "\n \n " + request .system_prompt
218218 else :
219219 request .system_prompt = memory_section
220- request .system_prompt = request .system_prompt + "\n \n " + LONGTERM_MEMORY_SYSTEM_PROMPT .format (memory_path = self .memory_path )
220+ request .system_prompt = (
221+ request .system_prompt
222+ + "\n \n "
223+ + LONGTERM_MEMORY_SYSTEM_PROMPT .format (memory_path = self .memory_path )
224+ )
221225
222226 return await handler (request )
0 commit comments