Files
docs/build/snippets/python/code-samples/long-term-memory-write-tool-postgres-py.mdx
T
2026-07-29 10:28:19 +00:00

44 lines
1.1 KiB
Plaintext

```python
from dataclasses import dataclass
from langchain.agents import create_agent
from langchain.tools import ToolRuntime, tool
from langchain_core.runnables import Runnable
from langgraph.store.postgres import PostgresStore # type: ignore[import-not-found]
from typing_extensions import TypedDict
@dataclass
class Context:
user_id: str
class UserInfo(TypedDict):
name: str
@tool
def save_user_info(user_info: UserInfo, runtime: ToolRuntime[Context]) -> str:
"""Save user info."""
assert runtime.store is not None
runtime.store.put(("users",), runtime.context.user_id, dict(user_info))
return "Successfully saved user info."
DB_URI = "postgresql://postgres:postgres@localhost:5432/postgres?sslmode=disable"
with PostgresStore.from_conn_string(DB_URI) as store:
store.setup()
agent: Runnable = create_agent(
"claude-sonnet-4-6",
tools=[save_user_info],
store=store,
context_schema=Context,
)
agent.invoke(
{"messages": [{"role": "user", "content": "My name is John Smith"}]},
context=Context(user_id="user_123"),
)
```