Skip to content

Commit 72d42dd

Browse files
parshvadaftariparshvadaftari
committed
Revamp mem0 integration to tool base approach
1 parent 84f367f commit 72d42dd

File tree

15 files changed

+1101
-1887
lines changed

15 files changed

+1101
-1887
lines changed

‎docs/mem0.md‎

Lines changed: 350 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,350 @@
1+
# Mem0 Memory Integration
2+
3+
PydanticAI provides a lightweight integration with [Mem0](https://mem0.ai) through the [`Mem0Toolset`][pydantic_ai.toolsets.Mem0Toolset]. This toolset adds memory capabilities to your agents, allowing them to save and search through conversation memories.
4+
5+
## Overview
6+
7+
The [`Mem0Toolset`][pydantic_ai.toolsets.Mem0Toolset] is a simple toolset that provides two memory tools:
8+
9+
- **`_search_memory_impl`**: Search through stored memories
10+
- **`_save_memory_impl`**: Save information to memory
11+
12+
This integration follows the same pattern as other third-party integrations like [LangChain tools](../third-party-tools.md#langchain-tools).
13+
14+
## Installation
15+
16+
To use the Mem0Toolset, install PydanticAI with the `mem0` extra:
17+
18+
```bash
19+
pip install 'pydantic-ai[mem0]'
20+
```
21+
22+
Or install Mem0 separately:
23+
24+
```bash
25+
pip install mem0ai
26+
```
27+
28+
## Quick Start
29+
30+
Here's a simple example of using the Mem0Toolset:
31+
32+
```python test="skip"
33+
import os
34+
35+
from pydantic_ai import Agent, Mem0Toolset
36+
37+
# Create Mem0 toolset
38+
mem0_toolset = Mem0Toolset(api_key=os.getenv('MEM0_API_KEY'))
39+
40+
# Create agent with memory capabilities
41+
agent = Agent(
42+
'openai:gpt-4o',
43+
toolsets=[mem0_toolset],
44+
instructions='You are a helpful assistant with memory capabilities.',
45+
)
46+
47+
48+
async def main():
49+
# Use the agent - it can now save and search memories
50+
result = await agent.run(
51+
'My name is Alice and I love Python. Please remember this.',
52+
deps='user_alice',
53+
)
54+
print(result.output)
55+
```
56+
57+
## How It Works
58+
59+
### User Identification
60+
61+
The [`Mem0Toolset`][pydantic_ai.toolsets.Mem0Toolset] requires a user identifier to scope memories. It extracts the `user_id` from the agent's `deps` in three ways:
62+
63+
1. **String deps**: If `deps` is a string, it's used directly as the `user_id`
64+
```python test="skip"
65+
import os
66+
67+
from pydantic_ai import Agent, Mem0Toolset
68+
69+
mem0_toolset = Mem0Toolset(api_key=os.getenv('MEM0_API_KEY'))
70+
agent = Agent('openai:gpt-4o', toolsets=[mem0_toolset])
71+
72+
73+
async def example():
74+
await agent.run('Remember this', deps='user_123')
75+
```
76+
77+
2. **Object with `user_id` attribute**: If `deps` has a `user_id` attribute
78+
```python test="skip"
79+
import os
80+
from dataclasses import dataclass
81+
82+
from pydantic_ai import Agent, Mem0Toolset
83+
84+
85+
@dataclass
86+
class UserSession:
87+
user_id: str
88+
session_id: str
89+
90+
91+
mem0_toolset = Mem0Toolset(api_key=os.getenv('MEM0_API_KEY'))
92+
agent = Agent('openai:gpt-4o', toolsets=[mem0_toolset])
93+
94+
95+
async def example():
96+
await agent.run(
97+
'Remember this', deps=UserSession(user_id='user_123', session_id='session_1')
98+
)
99+
```
100+
101+
3. **Object with `get_user_id()` method**: If `deps` has a `get_user_id()` method
102+
```python test="skip"
103+
import os
104+
105+
from pydantic_ai import Agent, Mem0Toolset
106+
107+
108+
class UserContext:
109+
def get_user_id(self) -> str:
110+
return 'user_123'
111+
112+
113+
mem0_toolset = Mem0Toolset(api_key=os.getenv('MEM0_API_KEY'))
114+
agent = Agent('openai:gpt-4o', toolsets=[mem0_toolset])
115+
116+
117+
async def example():
118+
await agent.run('Remember this', deps=UserContext())
119+
```
120+
121+
### Memory Tools
122+
123+
The agent automatically uses the memory tools when appropriate:
124+
125+
```python test="skip"
126+
import os
127+
128+
from pydantic_ai import Agent, Mem0Toolset
129+
130+
mem0_toolset = Mem0Toolset(api_key=os.getenv('MEM0_API_KEY'))
131+
agent = Agent('openai:gpt-4o', toolsets=[mem0_toolset])
132+
133+
134+
async def example():
135+
# Agent can save memories
136+
await agent.run(
137+
'I prefer concise responses and dark mode.',
138+
deps='user_alice',
139+
)
140+
141+
# Agent can search memories
142+
await agent.run(
143+
'What are my preferences?',
144+
deps='user_alice',
145+
)
146+
```
147+
148+
## Configuration
149+
150+
### API Key
151+
152+
Provide your Mem0 API key in one of two ways:
153+
154+
1. **Environment variable** (recommended):
155+
```bash
156+
export MEM0_API_KEY=your-api-key
157+
```
158+
159+
2. **Constructor parameter**:
160+
```python test="skip"
161+
from pydantic_ai.toolsets import Mem0Toolset
162+
163+
mem0_toolset = Mem0Toolset(api_key='your-api-key')
164+
```
165+
166+
### Custom Configuration
167+
168+
You can customize the toolset behavior:
169+
170+
```python test="skip"
171+
import os
172+
173+
from pydantic_ai.toolsets import Mem0Toolset
174+
175+
mem0_toolset = Mem0Toolset(
176+
api_key=os.getenv('MEM0_API_KEY'),
177+
limit=10, # Return top 10 memories (default: 5)
178+
host='https://api.mem0.ai', # Custom host
179+
org_id='your-org-id', # Organization ID
180+
project_id='your-project-id', # Project ID
181+
id='my-mem0-toolset', # Custom toolset ID
182+
)
183+
```
184+
185+
### Using a Custom Client
186+
187+
You can also provide a pre-configured Mem0 client:
188+
189+
```python test="skip"
190+
import os
191+
192+
from mem0 import AsyncMemoryClient
193+
194+
from pydantic_ai.toolsets import Mem0Toolset
195+
196+
client = AsyncMemoryClient(api_key=os.getenv('MEM0_API_KEY'))
197+
mem0_toolset = Mem0Toolset(client=client)
198+
```
199+
200+
## Examples
201+
202+
### Basic Memory Usage
203+
204+
```python test="skip"
205+
import os
206+
207+
from pydantic_ai import Agent, Mem0Toolset
208+
209+
mem0_toolset = Mem0Toolset(api_key=os.getenv('MEM0_API_KEY'))
210+
agent = Agent('openai:gpt-4o', toolsets=[mem0_toolset])
211+
212+
213+
async def main():
214+
# Save information
215+
await agent.run(
216+
'My name is Alice and I love Python.',
217+
deps='user_alice',
218+
)
219+
220+
# Recall information
221+
await agent.run(
222+
'What programming language do I like?',
223+
deps='user_alice',
224+
)
225+
```
226+
227+
### Multi-User Isolation
228+
229+
Memories are automatically isolated per user:
230+
231+
```python test="skip"
232+
import os
233+
234+
from pydantic_ai import Agent, Mem0Toolset
235+
236+
mem0_toolset = Mem0Toolset(api_key=os.getenv('MEM0_API_KEY'))
237+
agent = Agent('openai:gpt-4o', toolsets=[mem0_toolset])
238+
239+
240+
async def main():
241+
# Alice's memories
242+
await agent.run('My favorite color is blue.', deps='user_alice')
243+
244+
# Bob's memories
245+
await agent.run('My favorite color is red.', deps='user_bob')
246+
247+
# Each user gets their own memories back
248+
await agent.run('What is my favorite color?', deps='user_alice')
249+
# Output: "Your favorite color is blue."
250+
251+
await agent.run('What is my favorite color?', deps='user_bob')
252+
# Output: "Your favorite color is red."
253+
```
254+
255+
### Using with Dataclass Deps
256+
257+
```python test="skip"
258+
import os
259+
from dataclasses import dataclass
260+
261+
from pydantic_ai import Agent, Mem0Toolset
262+
263+
264+
@dataclass
265+
class UserSession:
266+
user_id: str
267+
session_id: str
268+
269+
270+
mem0_toolset = Mem0Toolset(api_key=os.getenv('MEM0_API_KEY'))
271+
agent = Agent('openai:gpt-4o', toolsets=[mem0_toolset])
272+
273+
274+
async def main():
275+
session = UserSession(user_id='user_charlie', session_id='session_123')
276+
277+
await agent.run(
278+
'I work as a data scientist.',
279+
deps=session,
280+
)
281+
```
282+
283+
### Personalized Assistant
284+
285+
```python test="skip"
286+
import os
287+
288+
from pydantic_ai import Agent, Mem0Toolset
289+
290+
mem0_toolset = Mem0Toolset(api_key=os.getenv('MEM0_API_KEY'))
291+
agent = Agent(
292+
'openai:gpt-4o',
293+
toolsets=[mem0_toolset],
294+
instructions=(
295+
'You are a helpful assistant with memory. '
296+
'Remember user preferences and provide personalized assistance.'
297+
),
298+
)
299+
300+
301+
async def main():
302+
# First conversation - learn preferences
303+
await agent.run(
304+
'I prefer concise responses with Python code examples.',
305+
deps='user_diana',
306+
)
307+
308+
# Later conversation - agent recalls preferences
309+
result = await agent.run(
310+
'How do I read a CSV file?',
311+
deps='user_diana',
312+
)
313+
print(result.output)
314+
# Agent will provide concise response with Python examples
315+
```
316+
317+
## Comparison with Other Approaches
318+
319+
The [`Mem0Toolset`][pydantic_ai.toolsets.Mem0Toolset] is designed to be lightweight and follow the same pattern as other third-party tool integrations like LangChain tools.
320+
321+
### Similar to LangChain Tools
322+
323+
Just like you can use [LangChain tools](../third-party-tools.md#langchain-tools) with PydanticAI:
324+
325+
```python test="skip"
326+
from langchain_community.tools import WikipediaQueryRun
327+
328+
from pydantic_ai import Agent
329+
from pydantic_ai.ext.langchain import LangChainToolset
330+
331+
toolset = LangChainToolset([WikipediaQueryRun()])
332+
agent = Agent('openai:gpt-4o', toolsets=[toolset])
333+
```
334+
335+
You can use Mem0 tools the same way:
336+
337+
```python test="skip"
338+
from pydantic_ai import Agent, Mem0Toolset
339+
340+
toolset = Mem0Toolset(api_key='your-api-key')
341+
agent = Agent('openai:gpt-4o', toolsets=[toolset])
342+
```
343+
344+
## API Reference
345+
346+
For detailed API documentation, see [`Mem0Toolset`][pydantic_ai.toolsets.Mem0Toolset].
347+
348+
## Complete Example
349+
350+
See the [complete example](../examples/mem0-toolset.md) for a full demonstration of the Mem0Toolset capabilities.

0 commit comments

Comments
 (0)