Compare commits

...
Author SHA1 Message Date
utkarsh240799andClaude Opus 4.6 6b4d2b780e fix: prevent in-place mutation of metadata dicts in _create_memory and _create_procedural_memory
Addresses #2648. The metadata dict passed by callers was mutated in-place
inside _create_memory and _create_procedural_memory (both sync and async),
causing unsafe behavior under concurrent/async usage.

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
2026-03-24 17:52:18 +05:30
KartikandGitHub 5431badfd4 fix: preserve custom metadata when updating memory (#4495) 2026-03-24 10:58:21 +05:30
bda5b726bd fix: avoid sending both temperature and top_p to Anthropic API (#4471)
Co-authored-by: utkarsh240799 <utkarsh240799@users.noreply.github.com>
Co-authored-by: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
2026-03-23 21:00:41 +05:30
KartikandGitHub 16bcc91716 chore: remove benchmark submission issue template and label matcher (#4514) 2026-03-23 20:32:55 +05:30
KartikandGitHub ba63ea4528 docs: add Vibecoding guide with Mem0 integration (#4511) 2026-03-23 19:36:21 +05:30
5332741961 fix: clean up graph store data on Memory.delete() (#4505)
Co-authored-by: utkarsh240799 <utkarsh240799@users.noreply.github.com>
Co-authored-by: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
2026-03-23 19:27:38 +05:30
d8a6960b4a fix: align Databricks docs with config and fix query mode selection (#4477)
Co-authored-by: utkarsh240799 <utkarsh240799@users.noreply.github.com>
Co-authored-by: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
2026-03-23 19:20:11 +05:30
HimanshuandGitHub 316dc67a0a fix(ts-oss): register pgvector in VectorStoreFactory (#3367) (#4502) 2026-03-23 18:52:19 +05:30
KartikandGitHub 65156d5176 chore: update issue templates and workflows for improved labeling and formatting (#4501) 2026-03-23 16:34:51 +05:30
KartikandGitHub c5e8216362 docs: add issue templates for bug, feature, benchmark, documentation, and update contact links (#4500) 2026-03-23 15:37:23 +05:30
HimanshuandGitHub ecedbc11d9 fix(qdrant): do not remove local path on init (#4473) (#4475) 2026-03-23 14:58:06 +05:30
9aadfa3221 fix: add missing limit, threshold, infer, memory_type, and prompt params to REST API (#4496)
Co-authored-by: utkarsh240799 <utkarsh240799@users.noreply.github.com>
Co-authored-by: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
2026-03-23 14:48:13 +05:30
Br1anandGitHub cc45561abd fix: accept default /tmp/chroma path in ChromaDbConfig validator (#4179) 2026-03-23 14:12:49 +05:30
7cebaba0a2 fix: upgrade MongoDB vector store from deprecated knnVector to GA vectorSearch (#3995)
Co-authored-by: kartik-mem0 <kartik.labhshetwar@mem0.ai>
Co-authored-by: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
2026-03-23 12:20:17 +05:30
mintlify[bot]andGitHub 5dabf24809 Fix 4 broken placeholder links in template files (#4478) 2026-03-21 17:29:24 -07:00
HimanshuandGitHub ec326f0f92 fix(mcp): operator precedence in search_memory filter (#4470) (#4474) 2026-03-21 21:27:35 +05:30
KartikandGitHub eb780f4880 refactor: add vector validation to OpenSearchDB to ensure non‑null, non‑empty, and correct‑dimension vectors (#4472) 2026-03-21 20:36:26 +05:30
Aditya PaulandGitHub c39d5ada4d fix: Bug: Zod Schema Incompatible with OpenAI Structured Outputs API (#3462) 2026-03-21 19:48:31 +05:30
06c25eb00b fix: use root LLM config as fallback for graph store instead of hardcoded OpenAI default (#4466)
Co-authored-by: utkarsh240799 <utkarsh240799@users.noreply.github.com>
Co-authored-by: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
2026-03-21 19:38:49 +05:30
7a09663156 fix(qdrant): implement enhanced metadata filtering operators (#4127)
Co-authored-by: kartik-mem0 <kartik.labhshetwar@mem0.ai>
Co-authored-by: Claude Sonnet 4.6 <noreply@anthropic.com>
2026-03-21 19:37:26 +05:30
mintlify[bot]GitHubmintlify[bot] <109931778+mintlify[bot]@users.noreply.github.com>
267bcf2931 (docs): add missing SEO metadata to turbopuffer page (#4468)
Co-authored-by: mintlify[bot] <109931778+mintlify[bot]@users.noreply.github.com>
2026-03-21 19:31:53 +05:30
bf9a5703b1 feat: integrate turbopuffer as vector database provider (#4428)
Co-authored-by: utkarsh240799 <utkarsh240799@users.noreply.github.com>
Co-authored-by: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
2026-03-21 19:28:25 +05:30
884e740b53 fix(graph): soft-delete graph relationships instead of hard DELETE (#4188)
Signed-off-by: sxu75374 <imshuaixu@gmail.com>
Co-authored-by: kartik-mem0 <kartik.labhshetwar@mem0.ai>
2026-03-21 17:32:59 +05:30
824032a81d feat: add NemoClaw + Mem0 plugin setup scripts and quickstart (#4464)
Co-authored-by: utkarsh240799 <utkarsh240799@users.noreply.github.com>
Co-authored-by: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
2026-03-21 13:52:42 +05:30
KartikandGitHub abdb07c204 fix: handle None content and empty candidates in GeminiLLM parsing (#4462) 2026-03-21 13:52:13 +05:30
7b26df728d docs: add Claude Code setup instructions for OpenMemory (#4430)
Co-authored-by: utkarsh240799 <utkarsh240799@users.noreply.github.com>
Co-authored-by: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
2026-03-21 13:51:47 +05:30
30661ab427 Fix: add pgvector support to NodeJS OSS VectorStoreFactory (fixes #3491) (#3997)
Co-authored-by: kartik-mem0 <kartik.labhshetwar@mem0.ai>
2026-03-20 20:47:04 +05:30
Failfail2603andGitHub 7ad5d6f442 fix: use toCamelCase in redis get method for the payload (#3172) 2026-03-20 20:28:41 +05:30
305ce7b6b3 feat: add Apache AGE graph store support (#4448)
Co-authored-by: utkarsh240799 <utkarsh240799@users.noreply.github.com>
Co-authored-by: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
2026-03-20 20:27:48 +05:30
4437c3e8a8 fix: add missing _parse_response to AzureOpenAIStructuredLLM (#4434)
Co-authored-by: Matt Van Horn <455140+mvanhorn@users.noreply.github.com>
Co-authored-by: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
2026-03-20 20:14:47 +05:30
54bdbde6e6 feat: add MiniMax LLM provider (#4132) (#4431)
Co-authored-by: kartik-mem0 <kartik.labhshetwar@mem0.ai>
2026-03-20 19:18:36 +05:30
KartikandGitHub 2b9558335b fix: raise ValueError when deleting nonexistent memory (#4455) 2026-03-20 18:28:58 +05:30
2520edb404 feat: add optional API key authentication to REST API server (#4442)
Co-authored-by: utkarsh240799 <utkarsh240799@users.noreply.github.com>
Co-authored-by: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
2026-03-20 16:15:25 +05:30
86 changed files with 11639 additions and 350 deletions
+47 -32
View File
@@ -1,41 +1,56 @@
name: 🐛 Bug Report
description: Create a report to help us reproduce and fix the bug
name: Bug Report
description: Report a bug in mem0
labels: ["bug"]
body:
- type: markdown
attributes:
value: >
#### Before submitting a bug, please make sure the issue hasn't been already addressed by searching through [the existing and past issues](https://github.com/embedchain/embedchain/issues?q=is%3Aissue+sort%3Acreated-desc+).
- type: textarea
attributes:
label: 🐛 Describe the bug
description: |
Please provide a clear and concise description of what the bug is.
- type: dropdown
id: component
attributes:
label: Component
description: Which part of mem0 is affected?
options:
- Core / Python SDK
- TypeScript SDK
- Vector Store (Qdrant, PGVector, Redis, Chroma, etc.)
- Graph Memory (Neo4j, Memgraph, etc.)
- Ollama / Local Models
- OpenMemory (MCP Server)
- OpenClaw
- REST API
- Other
validations:
required: true
If relevant, add a minimal example so that we can reproduce the error by running the code. It is very important for the snippet to be as succinct (minimal) as possible, so please take time to trim down any irrelevant code to help us debug efficiently. We are going to copy-paste your code and we expect to get the same result as you did: avoid any external data, and include the relevant imports, etc. For example:
- type: textarea
id: description
attributes:
label: Description
value: |
### Summary
```python
# All necessary imports at the beginning
import embedchain as ec
# Your code goes here
A clear summary of the bug.
### Steps to Reproduce
```
```python
from mem0 import Memory
Please also paste or describe the results you observe instead of the expected results. If you observe an error, please paste the error message including the **full** traceback of the exception. It may be relevant to wrap error messages in ```` ```triple quotes blocks``` ````.
placeholder: |
A clear and concise description of what the bug is.
m = Memory()
# Your code here...
```
```python
Sample code to reproduce the problem
```
### Expected Behavior
```
The error message you got, with the full traceback.
````
validations:
required: true
- type: markdown
attributes:
value: >
Thanks for contributing 🎉!
What you expected to happen.
### Actual Behavior
What actually happened. Paste the full error traceback if applicable.
### Environment
- mem0 version:
- Python/Node version:
- OS:
validations:
required: true
+5 -5
View File
@@ -1,8 +1,8 @@
blank_issues_enabled: true
contact_links:
- name: 1-on-1 Session
url: https://cal.com/taranjeetio/ec
about: Speak directly with Taranjeet, the founder, to discuss issues, share feedback, or explore improvements for Embedchain
- name: Discord
- name: Discord Community
url: https://discord.gg/6PzXDgEjG5
about: General community discussions
about: Ask questions and discuss with the community
- name: Documentation
url: https://docs.mem0.ai
about: Read the official mem0 documentation
+21 -9
View File
@@ -1,11 +1,23 @@
name: Documentation
description: Report an issue related to the Embedchain docs.
title: "DOC: <Please write a comprehensive title after the 'DOC: ' prefix>"
name: Documentation Issue
description: Report an issue or suggest an improvement to the mem0 docs
labels: ["documentation"]
body:
- type: textarea
attributes:
label: "Issue with current documentation:"
description: >
Please make sure to leave a reference to the document/code you're
referring to.
- type: textarea
id: description
attributes:
label: Description
value: |
### Page
Link to the docs page: https://docs.mem0.ai/...
### What's Wrong or Missing
Describe what's incorrect, unclear, or missing.
### Suggested Fix
How should the docs be improved?
validations:
required: true
+40 -21
View File
@@ -1,23 +1,42 @@
name: 🚀 Feature request
description: Submit a proposal/request for a new Embedchain feature
name: Feature Request
description: Suggest a new feature or improvement for mem0
labels: ["enhancement"]
body:
- type: textarea
id: feature-request
attributes:
label: 🚀 The feature
description: >
A clear and concise description of the feature proposal
validations:
required: true
- type: textarea
attributes:
label: Motivation, pitch
description: >
Please outline the motivation for the proposal. Is your feature request related to a specific problem? e.g., *"I'm working on X and would like Y to be possible"*. If this is related to another GitHub issue, please link here too.
validations:
required: true
- type: markdown
attributes:
value: >
Thanks for contributing 🎉!
- type: dropdown
id: component
attributes:
label: Component
description: Which part of mem0 does this relate to?
options:
- Core / Python SDK
- TypeScript SDK
- Vector Store (Qdrant, PGVector, Redis, Chroma, etc.)
- Graph Memory (Neo4j, Memgraph, etc.)
- Ollama / Local Models
- OpenMemory (MCP Server)
- OpenClaw
- REST API
- Benchmarks / Evals
- Other
validations:
required: true
- type: textarea
id: description
attributes:
label: Description
value: |
### Use Case
What problem are you trying to solve?
### Proposed Solution
How should this work? Include API examples or pseudocode if helpful.
### Alternatives Considered
Any workarounds you've tried or other approaches considered.
validations:
required: true
+25 -28
View File
@@ -1,41 +1,38 @@
## Linked Issue
Closes #<!-- issue number -->
## Description
Please include a summary of the change and which issue is fixed. Please also include relevant motivation and context. List any dependencies that are required for this change.
<!-- What does this PR do? Why is it needed? -->
Fixes # (issue)
## Type of Change
## Type of change
Please delete options that are not relevant.
- [ ] Bug fix (non-breaking change which fixes an issue)
- [ ] New feature (non-breaking change which adds functionality)
- [ ] Breaking change (fix or feature that would cause existing functionality to not work as expected)
- [ ] Refactor (does not change functionality, e.g. code style improvements, linting)
- [ ] Bug fix (non-breaking change that fixes an issue)
- [ ] New feature (non-breaking change that adds functionality)
- [ ] Breaking change (fix or feature that would cause existing functionality to change)
- [ ] Refactor (no functional changes)
- [ ] Documentation update
## How Has This Been Tested?
## Breaking Changes
Please describe the tests that you ran to verify your changes. Provide instructions so we can reproduce. Please also list any relevant details for your test configuration
<!-- If this is a breaking change, describe what breaks and the migration path. Delete this section if not applicable. -->
Please delete options that are not relevant.
N/A
- [ ] Unit Test
- [ ] Test Script (please provide)
## Test Coverage
## Checklist:
- [ ] I added/updated unit tests
- [ ] I added/updated integration tests
- [ ] I tested manually (describe below)
- [ ] No tests needed (explain why)
- [ ] My code follows the style guidelines of this project
- [ ] I have performed a self-review of my own code
- [ ] I have commented my code, particularly in hard-to-understand areas
- [ ] I have made corresponding changes to the documentation
- [ ] My changes generate no new warnings
- [ ] I have added tests that prove my fix is effective or that my feature works
- [ ] New and existing unit tests pass locally with my changes
- [ ] Any dependent changes have been merged and published in downstream modules
- [ ] I have checked my code and corrected any misspellings
<!-- Describe how you tested this, or link to CI results. -->
## Maintainer Checklist
## Checklist
- [ ] closes #xxxx (Replace xxxx with the GitHub issue number)
- [ ] Made sure Checks passed
- [ ] My code follows the project's style guidelines
- [ ] I have performed a self-review of my code
- [ ] I have added tests that prove my fix/feature works
- [ ] New and existing tests pass locally
- [ ] I have updated documentation if needed
+20
View File
@@ -0,0 +1,20 @@
# Maps dropdown selections to GitHub labels
# Used by the advanced-issue-labeler GitHub Action
component:
- label: "sdk-python"
matcher: "Core / Python SDK"
- label: "sdk-typescript"
matcher: "TypeScript SDK"
- label: "vector-store"
matcher: "Vector Store"
- label: "graph-memory"
matcher: "Graph Memory"
- label: "ollama"
matcher: "Ollama"
- label: "openmemory"
matcher: "OpenMemory"
- label: "openclaw"
matcher: "OpenClaw"
- label: "rest-api"
matcher: "REST API"
+39
View File
@@ -0,0 +1,39 @@
name: Auto-label issues
on:
issues:
types: [opened]
permissions:
contents: read
issues: write
jobs:
label:
runs-on: ubuntu-latest
steps:
- uses: stefanbuck/github-issue-parser@v3
id: issue-parser
with:
template-path: .github/ISSUE_TEMPLATE/bug_report.yml
- uses: redhat-plumbers-in-action/advanced-issue-labeler@v3
with:
issue-form: ${{ steps.issue-parser.outputs.jsonString }}
section: component
token: ${{ secrets.GITHUB_TOKEN }}
config-path: .github/advanced-issue-labeler.yml
- uses: stefanbuck/github-issue-parser@v3
id: feature-parser
if: contains(github.event.issue.labels.*.name, 'enhancement')
with:
template-path: .github/ISSUE_TEMPLATE/feature_request.yml
- uses: redhat-plumbers-in-action/advanced-issue-labeler@v3
if: contains(github.event.issue.labels.*.name, 'enhancement')
with:
issue-form: ${{ steps.feature-parser.outputs.jsonString }}
section: component
token: ${{ secrets.GITHUB_TOKEN }}
config-path: .github/advanced-issue-labeler.yml
+48
View File
@@ -0,0 +1,48 @@
name: Close stale issues
on:
schedule:
- cron: '0 0 * * *'
workflow_dispatch:
permissions:
issues: write
pull-requests: write
jobs:
stale:
runs-on: ubuntu-latest
steps:
- uses: actions/stale@v9
with:
# Issue settings
days-before-issue-stale: 90
days-before-issue-close: 14
stale-issue-label: 'stale'
stale-issue-message: >
This issue has been automatically marked as stale because it has not
had any activity in 90 days. It will be closed in 14 days if no
further activity occurs. If this is still relevant, please leave a
comment or remove the `stale` label.
close-issue-message: >
This issue has been closed due to inactivity. If this is still
relevant, feel free to reopen it or create a new issue.
# PR settings — mark stale but never auto-close
days-before-pr-stale: 90
days-before-pr-close: -1
stale-pr-label: 'stale'
stale-pr-message: >
This pull request has been automatically marked as stale because it
has not had any activity in 90 days. Please update your branch and
address any review comments, or it may be closed in the future.
# Exempt these labels from stale processing
exempt-issue-labels: 'P0-critical,P1-high,good first issue,security'
exempt-pr-labels: 'P0-critical,P1-high'
# Remove stale label when there is new activity
remove-stale-when-updated: true
# Process up to 100 issues per run to stay within API limits
operations-per-run: 100
+2 -1
View File
@@ -273,7 +273,7 @@ config = MemoryConfig(
### Supported Providers
#### LLM Providers (19 supported)
#### LLM Providers (20 supported)
- **openai** - OpenAI GPT models (default)
- **anthropic** - Claude models
- **gemini** - Google Gemini
@@ -284,6 +284,7 @@ config = MemoryConfig(
- **azure_openai** - Azure OpenAI
- **litellm** - LiteLLM proxy
- **deepseek** - DeepSeek models
- **minimax** - MiniMax models
- **xai** - xAI models
- **sarvam** - Sarvam AI
- **lmstudio** - LM Studio local server
+56
View File
@@ -0,0 +1,56 @@
---
title: MiniMax
description: "Configure MiniMax as an LLM provider in Mem0 with API key setup and optional custom endpoint configuration."
---
To use MiniMax LLM models, you have to set the `MINIMAX_API_KEY` environment variable. You can also optionally set `MINIMAX_API_BASE` if you need to use a different API endpoint (defaults to "https://api.minimax.io/v1").
## Usage
```python
import os
from mem0 import Memory
os.environ["MINIMAX_API_KEY"] = "your-api-key"
os.environ["OPENAI_API_KEY"] = "your-api-key" # for embedder model
config = {
"llm": {
"provider": "minimax",
"config": {
"model": "MiniMax-M2.7", # default model
"temperature": 0.2,
"max_tokens": 2000,
"top_p": 1.0
}
}
}
m = Memory.from_config(config)
messages = [
{"role": "user", "content": "I'm planning to watch a movie tonight. Any recommendations?"},
{"role": "assistant", "content": "How about thriller movies? They can be quite engaging."},
{"role": "user", "content": "I'm not a big fan of thriller movies but I love sci-fi movies."},
{"role": "assistant", "content": "Got it! I'll avoid thriller recommendations and suggest sci-fi movies in the future."}
]
m.add(messages, user_id="alice", metadata={"category": "movies"})
```
You can also configure the API base URL in the config:
```python
config = {
"llm": {
"provider": "minimax",
"config": {
"model": "MiniMax-M2.7",
"minimax_base_url": "https://your-custom-endpoint.com",
"api_key": "your-api-key" # alternatively to using environment variable
}
}
}
```
## Config
All available parameters for the `minimax` config are present in [Master List of All Params in Config](../config).
+1
View File
@@ -31,6 +31,7 @@ See the list of supported LLMs below.
<Card title="Google AI" href="/components/llms/models/google_AI" />
<Card title="AWS bedrock" href="/components/llms/models/aws_bedrock" />
<Card title="DeepSeek" href="/components/llms/models/deepseek" />
<Card title="MiniMax" href="/components/llms/models/minimax" />
<Card title="xAI" href="/components/llms/models/xAI" />
<Card title="Sarvam AI" href="/components/llms/models/sarvam" />
<Card title="LM Studio" href="/components/llms/models/lmstudio" />
+29 -20
View File
@@ -17,8 +17,10 @@ config = {
"workspace_url": "https://your-workspace.databricks.com",
"access_token": "your-access-token",
"endpoint_name": "your-vector-search-endpoint",
"index_name": "catalog.schema.index_name",
"source_table_name": "catalog.schema.source_table",
"catalog": "your_catalog",
"schema": "your_schema",
"table_name": "your_table",
"collection_name": "your_index_name",
"embedding_dimension": 1536
}
}
@@ -42,17 +44,22 @@ Here are the parameters available for configuring Databricks Vector Search:
| --- | --- | --- |
| `workspace_url` | The URL of your Databricks workspace | **Required** |
| `access_token` | Personal Access Token for authentication | `None` |
| `service_principal_client_id` | Service principal client ID (alternative to access_token) | `None` |
| `service_principal_client_secret` | Service principal client secret (required with client_id) | `None` |
| `client_id` | Service principal client ID (alternative to access_token) | `None` |
| `client_secret` | Service principal client secret (required with client_id) | `None` |
| `azure_client_id` | Azure AD application client ID (for Azure Databricks) | `None` |
| `azure_client_secret` | Azure AD application client secret (for Azure Databricks) | `None` |
| `endpoint_name` | Name of the Vector Search endpoint | **Required** |
| `index_name` | Name of the vector index (Unity Catalog format: catalog.schema.index) | **Required** |
| `source_table_name` | Name of the source Delta table (Unity Catalog format: catalog.schema.table) | **Required** |
| `embedding_dimension` | Dimension of self-managed embeddings | `1536` |
| `embedding_source_column` | Column name for text when using Databricks-computed embeddings | `None` |
| `catalog` | Unity Catalog catalog name | **Required** |
| `schema` | Unity Catalog schema name | **Required** |
| `table_name` | Source Delta table name | **Required** |
| `collection_name` | Vector search index name | `mem0` |
| `index_type` | Index type: `DELTA_SYNC` or `DIRECT_ACCESS` | `DELTA_SYNC` |
| `embedding_model_endpoint_name` | Databricks serving endpoint for embeddings | `None` |
| `embedding_vector_column` | Column name for self-managed embedding vectors | `embedding` |
| `embedding_dimension` | Dimension of self-managed embeddings | `1536` |
| `endpoint_type` | Type of endpoint (`STANDARD` or `STORAGE_OPTIMIZED`) | `STANDARD` |
| `sync_computed_embeddings` | Whether to sync computed embeddings automatically | `True` |
| `pipeline_type` | Sync pipeline type: `TRIGGERED` or `CONTINUOUS` | `TRIGGERED` |
| `warehouse_name` | Databricks SQL warehouse name (if using SQL warehouse) | `None` |
| `query_type` | Query type: `ANN` or `HYBRID` | `ANN` |
### Authentication
@@ -65,11 +72,13 @@ config = {
"provider": "databricks",
"config": {
"workspace_url": "https://your-workspace.databricks.com",
"service_principal_client_id": "your-service-principal-id",
"service_principal_client_secret": "your-service-principal-secret",
"client_id": "your-service-principal-id",
"client_secret": "your-service-principal-secret",
"endpoint_name": "your-endpoint",
"index_name": "catalog.schema.index_name",
"source_table_name": "catalog.schema.source_table"
"catalog": "your_catalog",
"schema": "your_schema",
"table_name": "your_table",
"collection_name": "your_index_name",
}
}
}
@@ -84,8 +93,10 @@ config = {
"workspace_url": "https://your-workspace.databricks.com",
"access_token": "your-personal-access-token",
"endpoint_name": "your-endpoint",
"index_name": "catalog.schema.index_name",
"source_table_name": "catalog.schema.source_table"
"catalog": "your_catalog",
"schema": "your_schema",
"table_name": "your_table",
"collection_name": "your_index_name",
}
}
}
@@ -103,7 +114,6 @@ config = {
"config": {
# ... authentication config ...
"embedding_dimension": 768, # Match your embedding model
"embedding_vector_column": "embedding"
}
}
}
@@ -118,7 +128,6 @@ config = {
"provider": "databricks",
"config": {
# ... authentication config ...
"embedding_source_column": "text",
"embedding_model_endpoint_name": "e5-small-v2"
}
}
@@ -127,8 +136,8 @@ config = {
### Important Notes
- **Delta Sync Index**: This implementation uses Delta Sync Index, which automatically syncs with your source Delta table. Direct vector insertion/deletion/update operations will log warnings as they're not supported with Delta Sync.
- **Unity Catalog**: Both the source table and index must be in Unity Catalog format (`catalog.schema.table_name`).
- **Index Types**: This implementation supports both `DELTA_SYNC` (auto-syncs with source Delta table) and `DIRECT_ACCESS` (manage vectors directly) index types.
- **Unity Catalog**: The source table and index are created under the specified `catalog.schema` namespace.
- **Endpoint Auto-Creation**: If the specified endpoint doesn't exist, it will be created automatically.
- **Index Auto-Creation**: If the specified index doesn't exist, it will be created automatically with the provided configuration.
- **Filter Support**: Supports filtering by metadata fields, with different syntax for STANDARD vs STORAGE_OPTIMIZED endpoints.
+3 -1
View File
@@ -48,7 +48,7 @@ const config = {
password: '123',
host: '127.0.0.1',
port: 5432,
dbname: 'vector_store', // Optional, defaults to 'postgres'
dbname: 'vector_store', // Optional; TypeScript OSS defaults to `vector_store` when omitted
diskann: false, // Optional, requires pgvectorscale extension
hnsw: false, // Optional, for HNSW indexing
},
@@ -85,6 +85,8 @@ Here are the parameters available for configuring pgvector:
| `connection_string` | PostgreSQL connection string (overrides individual connection parameters) | `None` |
| `connection_pool` | psycopg2 connection pool object (overrides connection string and individual parameters) | `None` |
**Note (TypeScript OSS):** If you omit `dbname`, the TypeScript client uses the database name `vector_store`. Python defaults to `postgres` for `dbname`, as in the table above.
**Note**: The connection parameters have the following priority:
1. `connection_pool` (highest priority)
2. `connection_string`
@@ -0,0 +1,79 @@
---
title: "Turbopuffer"
description: "Use Turbopuffer as a serverless vector database in Mem0 for low-latency search at scale with native metadata filtering."
---
[Turbopuffer](https://turbopuffer.com) is a serverless vector database optimized for low-latency search at scale. It offers cost-effective vector storage with native metadata filtering.
### Usage
```python
import os
from mem0 import Memory
os.environ["OPENAI_API_KEY"] = "sk-xx"
os.environ["TURBOPUFFER_API_KEY"] = "tpuf_xxxxxxxxxxxx"
config = {
"vector_store": {
"provider": "turbopuffer",
"config": {
"collection_name": "movie_preferences",
"embedding_model_dims": 1536,
"region": "gcp-us-central1",
}
}
}
m = Memory.from_config(config)
messages = [
{"role": "user", "content": "I'm planning to watch a movie tonight. Any recommendations?"},
{"role": "assistant", "content": "How about thriller movies? They can be quite engaging."},
{"role": "user", "content": "I'm not a big fan of thrillers but I love sci-fi."},
{"role": "assistant", "content": "Got it! I'll suggest sci-fi movies instead."}
]
m.add(messages, user_id="alice", metadata={"category": "movies"})
# Search memories
results = m.search(query="sci-fi recommendations", user_id="alice")
```
### Config
Here are the parameters available for configuring Turbopuffer:
| Parameter | Description | Default Value |
| --- | --- | --- |
| `collection_name` | Name of the namespace/collection | `mem0` |
| `embedding_model_dims` | Dimensions of the embedding model (must match your chosen embedding model) | `1536` |
| `api_key` | Turbopuffer API key | Environment variable: `TURBOPUFFER_API_KEY` |
| `region` | Turbopuffer region | `gcp-us-central1` |
| `distance_metric` | Distance metric for vector similarity (`cosine_distance` or `euclidean_squared`) | `cosine_distance` |
| `batch_size` | Batch size for bulk operations | `100` |
| `extra_params` | Additional parameters for the Turbopuffer client | `None` |
### Regions
| Region | Location |
| --- | --- |
| `gcp-us-central1` | Iowa, USA (Default) |
| `aws-us-west-2` | Oregon, USA |
### Config Example
```python
config = {
"vector_store": {
"provider": "turbopuffer",
"config": {
"collection_name": "my_memories",
"embedding_model_dims": 1536,
"api_key": "tpuf_xxxxxxxxxxxx",
"region": "aws-us-west-2",
"distance_metric": "cosine_distance",
"batch_size": 200,
}
}
}
```
+1
View File
@@ -33,6 +33,7 @@ See the list of supported vector databases below.
<Card title="LangChain" href="/components/vectordbs/dbs/langchain"></Card>
<Card title="Amazon S3 Vectors" href="/components/vectordbs/dbs/s3_vectors"></Card>
<Card title="Databricks" href="/components/vectordbs/dbs/databricks"></Card>
<Card title="Turbopuffer" href="/components/vectordbs/dbs/turbopuffer"></Card>
</CardGroup>
## Usage
+6 -1
View File
@@ -40,6 +40,7 @@
"icon": "rocket",
"pages": [
"platform/overview",
"vibecoding",
"platform/mem0-mcp",
"platform/platform-vs-oss",
"platform/quickstart"
@@ -142,6 +143,7 @@
"icon": "rocket",
"pages": [
"open-source/overview",
"vibecoding",
"open-source/python-quickstart",
"open-source/node-quickstart"
]
@@ -188,6 +190,7 @@
"components/llms/models/google_AI",
"components/llms/models/aws_bedrock",
"components/llms/models/deepseek",
"components/llms/models/minimax",
"components/llms/models/xAI",
"components/llms/models/sarvam",
"components/llms/models/lmstudio",
@@ -230,7 +233,8 @@
"components/vectordbs/dbs/cassandra",
"components/vectordbs/dbs/s3_vectors",
"components/vectordbs/dbs/databricks",
"components/vectordbs/dbs/neptune_analytics"
"components/vectordbs/dbs/neptune_analytics",
"components/vectordbs/dbs/turbopuffer"
]
}
]
@@ -300,6 +304,7 @@
"icon": "square-terminal",
"pages": [
"openmemory/overview",
"vibecoding",
"openmemory/quickstart",
"openmemory/integrations"
]
+38 -2
View File
@@ -35,7 +35,7 @@ graph LR
Mem0’s extraction LLM identifies entities, relationships, and timestamps from the conversation payload you send to `memory.add`.
</Step>
<Step title="Store vectors and edges together">
Embeddings land in your configured vector database while nodes and edges flow into a Bolt-compatible graph backend (Neo4j, Memgraph, Neptune, or Kuzu).
Embeddings land in your configured vector database while nodes and edges flow into a graph backend (Neo4j, Memgraph, Neptune, Kuzu, or Apache AGE).
</Step>
<Step title="Expose graph context at search time">
`memory.search` performs vector similarity (optionally reranked by your configured reranker) and returns the results list. Graph Memory runs in parallel and adds related entities in the `relations` array—it does not reorder the vector hits automatically.
@@ -264,7 +264,7 @@ Monitor graph growth, especially on free tiers, by periodically cleaning dormant
## Decision Points
- Select the graph store that fits your deployment (managed Aura vs. self-hosted Neo4j vs. AWS Neptune vs. local Kuzu).
- Select the graph store that fits your deployment (managed Aura vs. self-hosted Neo4j vs. AWS Neptune vs. local Kuzu vs. Apache AGE on PostgreSQL).
- Decide when to enable graph writes per request; routine conversations may stay vector-only to save latency.
- Set a policy for pruning stale relationships so your graph stays fast and affordable.
@@ -381,6 +381,42 @@ config = {
Kuzu will clear its state when using `:memory:` once the process exits. See the [Kuzu documentation](https://kuzudb.com/docs/) for advanced settings.
</Accordion>
<Accordion title="Apache AGE (PostgreSQL extension)">
[Apache AGE](https://age.apache.org/) adds graph database capabilities to PostgreSQL, letting you run Cypher queries alongside SQL on the same server. Start AGE via Docker, then configure Mem0:
```bash
docker run --name age-postgres \
-e POSTGRES_DB=mem0_db \
-e POSTGRES_USER=mem0_user \
-e POSTGRES_PASSWORD=mem0_pass \
-p 5432:5432 \
-d apache/age
```
```python
from mem0 import Memory
config = {
"graph_store": {
"provider": "apache_age",
"config": {
"host": "localhost",
"port": 5432,
"database": "mem0_db",
"username": "mem0_user",
"password": "mem0_pass",
"graph_name": "mem0_graph",
},
},
}
m = Memory.from_config(config_dict=config)
```
Apache AGE does not have a built-in vector index, so similarity search is computed client-side. This works well for moderate graph sizes; for very large graphs consider pairing AGE with pgvector for the vector store.
Reference: [Apache AGE documentation](https://age.apache.org/age-manual/master/index.html).
</Accordion>
</AccordionGroup>
<CardGroup cols={2}>
+38 -2
View File
@@ -14,7 +14,7 @@ The Mem0 REST API server exposes every OSS memory operation over HTTP. Run it al
</Info>
<Warning>
Add your own authentication and HTTPS before exposing the server to anything beyond your internal network. The default image does not include auth.
Enable API key authentication (see below) and HTTPS before exposing the server to anything beyond your internal network.
</Warning>
---
@@ -22,6 +22,7 @@ The Mem0 REST API server exposes every OSS memory operation over HTTP. Run it al
## Feature
- **CRUD endpoints:** Create, retrieve, search, update, delete, and reset memories by `user_id`, `agent_id`, or `run_id`.
- **API key authentication:** Optionally secure all endpoints with a shared API key via the `X-API-Key` header.
- **Status health check:** Access base routes to confirm the server is online.
- **OpenAPI explorer:** Visit `/docs` for interactive testing and schema reference.
@@ -91,6 +92,41 @@ uvicorn main:app --reload
---
## Authentication
The server supports optional API key authentication. When the `ADMIN_API_KEY` environment variable is set, every endpoint requires a valid `X-API-Key` header. The `/` redirect, `/docs`, and `/openapi.json` routes remain open so you can always reach the interactive API explorer.
| `ADMIN_API_KEY` value | Behavior |
|---|---|
| Not set / empty | All endpoints are open (no auth) |
| Any non-empty string | Requests must include `X-API-Key: <your-key>` |
### Enable authentication
Add the key to your `.env` file:
```bash
ADMIN_API_KEY=your-secret-api-key
```
Then include the header in every request:
```bash
curl -X POST http://localhost:8000/memories \
-H "Content-Type: application/json" \
-H "X-API-Key: your-secret-api-key" \
-d '{
"messages": [{"role": "user", "content": "I love pizza."}],
"user_id": "alice"
}'
```
<Warning>
The server logs a warning at startup when `ADMIN_API_KEY` is not set. Always set it in production.
</Warning>
---
## See it in action
### Create and search memories via HTTP
@@ -137,7 +173,7 @@ curl "http://localhost:8000/memories/search?user_id=alice&query=vegetable"
## Best practices
1. **Add authentication:** Protect endpoints with API gateways, proxies, or custom FastAPI middleware.
1. **Enable authentication:** Set `ADMIN_API_KEY` to secure all endpoints, or use an API gateway for more advanced schemes.
2. **Use HTTPS:** Terminate TLS at your load balancer or reverse proxy.
3. **Monitor uptime:** Track request rates, latency, and error codes per endpoint.
4. **Version configs:** Keep environment files and Docker Compose definitions in source control.
+1
View File
@@ -47,6 +47,7 @@ Replace `<client-name>` with the desired client name and `<user-id>` with the va
| Witsy | `npx install-mcp http://localhost:8765/mcp/witsy/sse/<user-id> --client witsy` |
| Enconvo | `npx install-mcp http://localhost:8765/mcp/enconvo/sse/<user-id> --client enconvo` |
| Augment | `npx install-mcp http://localhost:8765/mcp/augment/sse/<user-id> --client augment` |
| Claude Code | `claude mcp add -s user -t sse OpenMemory http://localhost:8765/mcp/openmemory/sse/<user-id>` |
### What This Does
+5
View File
@@ -183,6 +183,11 @@ You can configure the MCP client using the following command (replace `username`
npx @openmemory/install local "http://localhost:8765/mcp/cursor/sse/username" --client cursor
```
**For Claude Code (CLI):**
```bash
claude mcp add -s user -t sse OpenMemory "http://localhost:8765/mcp/openmemory/sse/username"
```
The OpenMemory dashboard will be available at http://localhost:3000. From here, you can view and manage your memories and check connection status with your MCP clients.
Once set up, OpenMemory runs locally on your machine, ensuring all your AI memories remain private and secure while being accessible across any compatible MCP client.
+2 -2
View File
@@ -166,13 +166,13 @@ Call out the most common mistake or edge case for this layer.
title="[Related cookbook / deep dive]"
description="[Why this pairs well with the current guide]"
icon="arrow-right"
href="/[related-link]"
href="#related-link"
/>
<Card
title="[Next cookbook in journey]"
description="[Set expectation for the next step]"
icon="rocket"
href="/[next-link]"
href="#next-link"
/>
</CardGroup>
```
+2 -2
View File
@@ -145,13 +145,13 @@ npm install mem0ai@[version]
title="[Deep dive reference]"
description="[Why this reference matters post-migration]"
icon="book"
href="/[reference-link]"
href="#reference-link"
/>
<Card
title="[Applied example or next step]"
description="[What readers can build now]"
icon="rocket"
href="/[example-link]"
href="#example-link"
/>
</CardGroup>
```
+181
View File
@@ -0,0 +1,181 @@
---
title: "Vibecoding with Mem0"
sidebarTitle: "Vibecoding"
description: "Agent skills, starter prompts, and setup for building with Mem0 using AI coding tools."
icon: "wand-magic-sparkles"
---
These docs are designed to be easily consumable by LLMs. Each page has a button that lets you copy the page as Markdown or paste directly into ChatGPT, Claude, or any AI coding tool.
We follow the llms.txt standard:
- [llms.txt](https://docs.mem0.ai/llms.txt)
<CardGroup cols={2}>
<Card title="Get an API Key" icon="key" href="https://app.mem0.ai">
Sign up for Mem0 Platform and start building
</Card>
<Card title="Quickstart" icon="rocket" href="/platform/quickstart">
Store your first memory in under 5 minutes
</Card>
</CardGroup>
## Agent Skills
Teach your coding assistant how to build with Mem0:
```bash
npx skills add https://github.com/mem0ai/mem0 --skill mem0
```
Works with Claude Code, Cursor, Windsurf, and any assistant that supports skills. Once installed, your assistant understands Mem0's full API, framework integrations, and common patterns.
## Claude Code Plugin
The [OpenMemory plugin](https://github.com/mem0ai/claude-code-plugin) gives Claude Code **persistent memory across sessions, projects, and teams** — automatically.
<Steps>
<Step title="Get your API key">
Sign up at [app.openmemory.dev](https://app.openmemory.dev).
</Step>
<Step title="Install the plugin">
```bash
/plugin add mem0ai/claude-code-plugin
```
</Step>
<Step title="Set your environment variable">
```bash
export OPENMEMORY_API_KEY="your-key-here"
```
</Step>
<Step title="Start coding">
The plugin activates automatically. It captures decisions at session end, preserves context during compaction, and retrieves relevant memories at session start.
</Step>
</Steps>
## MCP Server Setup
Connect Cursor, Windsurf, Claude Desktop, or any MCP-compatible client to Mem0.
<Tabs>
<Tab title="OpenMemory Hosted">
Sign up at [app.openmemory.dev](https://app.openmemory.dev), then pick your client:
<CodeGroup>
```bash Claude Desktop
npx @openmemory/install --client claude --env OPENMEMORY_API_KEY=your-key
```
```bash Cursor
npx @openmemory/install --client cursor --env OPENMEMORY_API_KEY=your-key
```
```bash Windsurf
npx @openmemory/install --client windsurf --env OPENMEMORY_API_KEY=your-key
```
</CodeGroup>
For full setup options, see [OpenMemory Quickstart](/openmemory/quickstart).
</Tab>
<Tab title="Mem0 Platform MCP">
Get your API key from [app.mem0.ai](https://app.mem0.ai), then add to your MCP config:
```json
{
"mcpServers": {
"mem0": {
"command": "uvx",
"args": ["mem0-mcp-server"],
"env": {
"MEM0_API_KEY": "m0-...",
"MEM0_DEFAULT_USER_ID": "your-handle"
}
}
}
}
```
For Docker, Smithery, and advanced options, see [Mem0 MCP Setup](/platform/mem0-mcp).
</Tab>
</Tabs>
## Universal Starter Prompt
Copy this into any AI tool to start building with Mem0:
```text
I want to start building with Mem0 — a self-improving memory layer for LLM
applications that gives agents persistent context across sessions.
## Mem0 Resources
**Documentation:**
- Main docs: https://docs.mem0.ai
- Platform Quickstart: https://docs.mem0.ai/platform/quickstart
- OSS Python Quickstart: https://docs.mem0.ai/open-source/python-quickstart
- OSS Node.js Quickstart: https://docs.mem0.ai/open-source/node-quickstart
- API Reference: https://docs.mem0.ai/api-reference
- Full LLM-friendly docs: https://docs.mem0.ai/llms.txt
**Code & Examples:**
- Core repo: https://github.com/mem0ai/mem0
- Python SDK: pip install mem0ai
- TypeScript SDK: npm install mem0ai
- Cookbooks: https://docs.mem0.ai/cookbooks/overview
**What Mem0 Does:**
Mem0 is a memory layer for AI apps — managed (Mem0 Platform) or self-hosted
(Open Source). It stores, retrieves, and manages user memories so agents
remember preferences, learn from interactions, and personalize over time.
Sub-50ms retrieval. Dual storage: vector embeddings + graph databases.
**Architecture Overview:**
- Memory is scoped by user_id, agent_id, or run_id
- Core operations: add, search, update, delete
- Memory types: factual (preferences, facts), episodic (past interactions),
semantic (concept relationships), working (session state)
- Integration pattern: retrieve relevant memories → generate response → store
new memories
**Quick Usage (Python Platform):**
from mem0 import MemoryClient
client = MemoryClient(api_key="m0-xxx")
client.add("I prefer dark mode and use VS Code.", user_id="user1")
results = client.search("What editor do they use?", user_id="user1")
**Quick Usage (JavaScript Platform):**
import MemoryClient from 'mem0ai';
const client = new MemoryClient({ apiKey: 'm0-xxx' });
await client.add([{ role: "user", content: "I prefer dark mode." }], { user_id: "user1" });
const results = await client.search("What editor?", { user_id: "user1" });
**Quick Usage (Python Open Source):**
from mem0 import Memory
m = Memory()
m.add("I prefer dark mode and use VS Code.", user_id="user1")
results = m.search("What editor do they use?", user_id="user1")
Help me integrate Mem0 into my project. Start by asking what I'm building,
what language/framework I'm using, and whether I want managed or self-hosted.
```
## Go Deeper
<CardGroup cols={2}>
<Card title="Platform Quickstart" icon="cloud" href="/platform/quickstart">
Get started with the managed API
</Card>
<Card title="Open Source" icon="code-branch" href="/open-source/overview">
Self-host with full control
</Card>
<Card title="Cookbooks" icon="book" href="/cookbooks/overview">
Production-ready tutorials and examples
</Card>
<Card title="API Reference" icon="code" href="/api-reference">
Explore every REST endpoint
</Card>
</CardGroup>
+419
View File
@@ -0,0 +1,419 @@
#!/usr/bin/env bash
# ─────────────────────────────────────────────────────────────────────────────
# install-mem0-plugin.sh
#
# Installs and configures the @mem0/openclaw-mem0 plugin for an existing
# NemoClaw sandbox. Assumes NemoClaw is already installed and onboarded.
#
# Usage:
# chmod +x install-mem0-plugin.sh && ./install-mem0-plugin.sh
#
# Requirements:
# - NemoClaw installed and onboarded (sandbox in Ready state)
# - Mem0 API key (from app.mem0.ai)
# ─────────────────────────────────────────────────────────────────────────────
set -euo pipefail
# ── Colors and formatting ────────────────────────────────────────────────────
RED='\033[0;31m'
GREEN='\033[0;32m'
YELLOW='\033[1;33m'
BLUE='\033[0;34m'
CYAN='\033[0;36m'
BOLD='\033[1m'
DIM='\033[2m'
NC='\033[0m'
info() { echo -e "${BLUE}[INFO]${NC} $*"; }
success() { echo -e "${GREEN}[OK]${NC} $*"; }
warn() { echo -e "${YELLOW}[WARN]${NC} $*"; }
error() { echo -e "${RED}[ERROR]${NC} $*"; }
step() { echo -e "\n${BOLD}${CYAN}── $* ──${NC}\n"; }
ask() { echo -en "${BOLD}$*${NC}"; }
die() {
error "$*"
exit 1
}
# ── Defaults ─────────────────────────────────────────────────────────────────
MEM0_USER_ID="${MEM0_USER_ID:-default}"
PLUGIN_PKG="@mem0/openclaw-mem0"
CONTAINER_NAME="nemoclaw-dev"
# ── Detect platform ──────────────────────────────────────────────────────────
OS_TYPE="$(uname -s)"
IS_MACOS=false
IS_LINUX=false
case "$OS_TYPE" in
Darwin) IS_MACOS=true ;;
Linux) IS_LINUX=true ;;
*) die "Unsupported OS: $OS_TYPE" ;;
esac
# ── Helper functions ─────────────────────────────────────────────────────────
check_command() {
command -v "$1" &>/dev/null
}
ensure_nvm() {
export NVM_DIR="${NVM_DIR:-$HOME/.nvm}"
if [[ -s "$NVM_DIR/nvm.sh" ]]; then
source "$NVM_DIR/nvm.sh"
fi
}
ensure_path() {
for p in "$HOME/.local/bin" "$HOME/.nvm/versions/node/"*/bin; do
if [[ -d "$p" ]] && [[ ":$PATH:" != *":$p:"* ]]; then
export PATH="$p:$PATH"
fi
done
}
# Wrapper to run a command natively (Linux) or in the container (macOS)
run_cmd() {
if $IS_MACOS; then
docker exec "$CONTAINER_NAME" bash -c "export NVM_DIR=/root/.nvm && source /root/.nvm/nvm.sh && $*"
else
eval "$@"
fi
}
# ── Banner ───────────────────────────────────────────────────────────────────
echo ""
echo -e "${BOLD}${CYAN}╔══════════════════════════════════════════════════════════════╗${NC}"
echo -e "${BOLD}${CYAN}║ Mem0 Plugin Installer for NemoClaw ║${NC}"
echo -e "${BOLD}${CYAN}║ Long-term memory for your OpenClaw agent ║${NC}"
echo -e "${BOLD}${CYAN}╚══════════════════════════════════════════════════════════════╝${NC}"
echo ""
if $IS_MACOS; then
info "Platform: macOS (using Docker container '$CONTAINER_NAME')"
else
info "Platform: Linux"
fi
# ═════════════════════════════════════════════════════════════════════════════
# PRE-CHECK: Verify NemoClaw is installed and sandbox is ready
# ═════════════════════════════════════════════════════════════════════════════
step "Pre-check: Verifying NemoClaw setup"
if $IS_LINUX; then
ensure_nvm
ensure_path
fi
# Verify nemoclaw/openshell are available
if $IS_MACOS; then
if ! docker ps --format '{{.Names}}' | grep -q "^${CONTAINER_NAME}$"; then
die "Docker container '$CONTAINER_NAME' is not running. Start it with: docker start $CONTAINER_NAME"
fi
if ! run_cmd "command -v nemoclaw" &>/dev/null; then
die "NemoClaw is not installed in container '$CONTAINER_NAME'. Run the full setup script first."
fi
else
if ! check_command openshell; then
die "openshell not found. Is NemoClaw installed and onboarded? Try: source ~/.bashrc"
fi
fi
# Detect sandbox
SANDBOX_NAME=$(run_cmd "openshell sandbox list 2>/dev/null" | awk 'NR>1 && $1!="" {print $1; exit}' || true)
if [[ -z "$SANDBOX_NAME" ]]; then
die "No sandbox found. Run 'nemoclaw onboard' first."
fi
# Verify sandbox is ready
SANDBOX_PHASE=$(run_cmd "openshell sandbox get '$SANDBOX_NAME' 2>/dev/null" | grep -i "phase" | awk '{print $NF}' || true)
if [[ "$SANDBOX_PHASE" != "Ready" ]]; then
warn "Sandbox '$SANDBOX_NAME' is not in Ready state (current: ${SANDBOX_PHASE:-unknown})."
warn "Waiting up to 2 minutes..."
WAIT_OK=false
for i in $(seq 1 24); do
SANDBOX_PHASE=$(run_cmd "openshell sandbox get '$SANDBOX_NAME' 2>/dev/null" | grep -i "phase" | awk '{print $NF}' || true)
if [[ "$SANDBOX_PHASE" == "Ready" ]]; then
WAIT_OK=true
break
fi
sleep 5
done
if ! $WAIT_OK; then
die "Sandbox '$SANDBOX_NAME' did not become Ready. Run: openshell sandbox list"
fi
fi
success "NemoClaw installed"
success "Sandbox '$SANDBOX_NAME' is ready"
# ═════════════════════════════════════════════════════════════════════════════
# STEP 1: Install Mem0 Plugin
# ═════════════════════════════════════════════════════════════════════════════
step "Step 1: Installing Mem0 plugin ($PLUGIN_PKG)"
# Helper: run a command inside the sandbox non-interactively via piped stdin
sandbox_exec() {
local cmd="$1"
if $IS_MACOS; then
printf '%s\nexit\n' "$cmd" | docker exec -i "$CONTAINER_NAME" bash -c \
"export NVM_DIR=/root/.nvm && source /root/.nvm/nvm.sh && openshell sandbox connect '$SANDBOX_NAME'" 2>&1
else
printf '%s\nexit\n' "$cmd" | openshell sandbox connect "$SANDBOX_NAME" 2>&1
fi
}
# Check if plugin is already installed
PLUGIN_EXISTS=false
if run_cmd "openshell sandbox download '$SANDBOX_NAME' /sandbox/.openclaw/extensions/openclaw-mem0/package.json /tmp/_mem0_plugin_check.json" &>/dev/null; then
if [[ -f /tmp/_mem0_plugin_check.json ]] || run_cmd "test -f /tmp/_mem0_plugin_check.json" &>/dev/null; then
PLUGIN_EXISTS=true
fi
fi
rm -f /tmp/_mem0_plugin_check.json 2>/dev/null || true
run_cmd "rm -f /tmp/_mem0_plugin_check.json" 2>/dev/null || true
if $PLUGIN_EXISTS; then
success "Mem0 plugin already installed in sandbox"
else
info "Downloading $PLUGIN_PKG..."
# Download and build outside the sandbox (on host or in container)
run_cmd "cd /tmp && rm -rf openclaw-mem0-full mem0-openclaw-mem0-*.tgz openclaw-mem0-full.tgz"
if ! run_cmd "cd /tmp && npm pack '$PLUGIN_PKG' 2>/dev/null"; then
die "Failed to download $PLUGIN_PKG from npm. Check your internet connection."
fi
success "Downloaded plugin"
info "Installing plugin dependencies..."
run_cmd "mkdir -p /tmp/openclaw-mem0-full && cd /tmp/openclaw-mem0-full && tar xzf /tmp/mem0-openclaw-mem0-*.tgz --strip-components=1 && npm install --omit=dev 2>&1 | tail -3"
success "Dependencies installed"
info "Uploading plugin to sandbox..."
run_cmd "cd /tmp && tar czf openclaw-mem0-full.tgz -C openclaw-mem0-full ."
if ! run_cmd "openshell sandbox upload '$SANDBOX_NAME' /tmp/openclaw-mem0-full.tgz /sandbox/openclaw-mem0-full.tgz 2>&1"; then
die "Failed to upload plugin to sandbox. Check: openshell sandbox list"
fi
success "Plugin uploaded"
info "Extracting plugin inside sandbox..."
sandbox_exec "mkdir -p ~/.openclaw/extensions/openclaw-mem0 && tar xzf /sandbox/openclaw-mem0-full.tgz/openclaw-mem0-full.tgz -C ~/.openclaw/extensions/openclaw-mem0 2>/dev/null || tar xzf /sandbox/openclaw-mem0-full.tgz -C ~/.openclaw/extensions/openclaw-mem0 2>/dev/null && echo EXTRACT_OK" >/dev/null 2>&1 || true
# Verify
VERIFY_OK=false
if run_cmd "openshell sandbox download '$SANDBOX_NAME' /sandbox/.openclaw/extensions/openclaw-mem0/package.json /tmp/_mem0_verify.json" &>/dev/null; then
VERIFY_OK=true
fi
rm -f /tmp/_mem0_verify.json 2>/dev/null || true
run_cmd "rm -f /tmp/_mem0_verify.json" 2>/dev/null || true
if $VERIFY_OK; then
success "Plugin extracted inside sandbox"
else
warn "Could not verify plugin extraction. You may need to extract manually."
if $IS_MACOS; then
echo " docker exec -it $CONTAINER_NAME bash"
echo " export NVM_DIR=/root/.nvm && source /root/.nvm/nvm.sh"
fi
echo " nemoclaw $SANDBOX_NAME connect"
echo " mkdir -p ~/.openclaw/extensions/openclaw-mem0"
echo " tar xzf /sandbox/openclaw-mem0-full.tgz/openclaw-mem0-full.tgz -C ~/.openclaw/extensions/openclaw-mem0"
fi
# Clean up
run_cmd "rm -rf /tmp/openclaw-mem0-full /tmp/mem0-openclaw-mem0-*.tgz /tmp/openclaw-mem0-full.tgz" 2>/dev/null || true
fi
# ═════════════════════════════════════════════════════════════════════════════
# STEP 2: Update Network Policy
# ═════════════════════════════════════════════════════════════════════════════
step "Step 2: Updating network policy to allow api.mem0.ai and telemetry"
# Find baseline policy
BASELINE_POLICY=$(run_cmd "find / -path '*/nemoclaw-blueprint/policies/openclaw-sandbox.yaml' 2>/dev/null | head -1" || true)
if [[ -z "$BASELINE_POLICY" ]]; then
die "Cannot find NemoClaw baseline policy file. Is NemoClaw installed?"
fi
info "Baseline policy: $BASELINE_POLICY"
# Check if mem0_api already exists
HAS_MEM0=$(run_cmd "grep -c mem0_api '$BASELINE_POLICY' 2>/dev/null" || echo "0")
if [[ "$HAS_MEM0" != "0" ]]; then
success "mem0_api already in baseline policy"
else
info "Adding api.mem0.ai to network policy..."
fi
# Create custom policy with mem0_api + telemetry
run_cmd "node -e \"
const fs = require('fs');
let c = fs.readFileSync('$BASELINE_POLICY', 'utf8');
const mem0Block = '\\n mem0_api:\\n name: mem0_api\\n endpoints:\\n - host: api.mem0.ai\\n port: 443\\n access: full\\n binaries:\\n - { path: /usr/local/bin/node }\\n - { path: /usr/local/bin/openclaw }\\n';
const telemetryBlock = '\\n mem0_telemetry:\\n name: mem0_telemetry\\n endpoints:\\n - host: us.i.posthog.com\\n port: 443\\n access: full\\n binaries:\\n - { path: /usr/local/bin/node }\\n - { path: /usr/local/bin/openclaw }\\n';
if (!c.includes('mem0_api')) {
if (c.includes('# ── Messaging')) {
c = c.replace(' # ── Messaging', mem0Block + '\\n # ── Messaging');
} else {
c += mem0Block;
}
}
if (!c.includes('mem0_telemetry')) {
if (c.includes('mem0_api:')) {
c = c.replace(' mem0_api:', telemetryBlock + '\\n mem0_api:');
} else {
c += telemetryBlock;
}
}
fs.writeFileSync('/tmp/nemoclaw-mem0-policy.yaml', c);
console.log('ok');
\""
success "Custom policy file created"
# Apply the policy
info "Applying network policy..."
if ! run_cmd "openshell policy set '$SANDBOX_NAME' --policy /tmp/nemoclaw-mem0-policy.yaml --wait 2>&1"; then
error "Failed to apply network policy."
echo ""
echo " If you see 'sandbox not found', re-run: nemoclaw onboard"
echo " Then re-run this script."
echo ""
die "Network policy update failed."
fi
success "Network policy applied — api.mem0.ai and telemetry allowed"
# ═════════════════════════════════════════════════════════════════════════════
# STEP 3: Configure Plugin
# ═════════════════════════════════════════════════════════════════════════════
step "Step 3: Configuring Mem0 plugin"
echo ""
echo -e " ${DIM}Get your Mem0 API key from: https://app.mem0.ai${NC}"
echo -e " ${DIM}The key starts with 'm0-'${NC}"
echo ""
ask "Enter your Mem0 API key: "
read -r MEM0_API_KEY
if [[ -z "$MEM0_API_KEY" ]]; then
die "Mem0 API key is required."
fi
if [[ ! "$MEM0_API_KEY" =~ ^m0- ]]; then
warn "Key doesn't start with 'm0-'. Make sure this is correct."
fi
echo ""
echo -e " ${DIM}The user ID scopes all memories. Pick any unique identifier.${NC}"
echo -e " ${DIM}Examples: alice, user_123, your-email@example.com${NC}"
echo ""
ask "Enter user ID [$MEM0_USER_ID]: "
read -r custom_user_id
MEM0_USER_ID="${custom_user_id:-$MEM0_USER_ID}"
info "Configuring plugin inside sandbox..."
CONFIG_SCRIPT="openclaw config set plugins.slots.memory openclaw-mem0 2>&1 | tail -1 && \
openclaw config set plugins.entries.openclaw-mem0.enabled true 2>&1 | tail -1 && \
openclaw config set plugins.entries.openclaw-mem0.config.apiKey '$MEM0_API_KEY' 2>&1 | tail -1 && \
openclaw config set plugins.entries.openclaw-mem0.config.userId '$MEM0_USER_ID' 2>&1 | tail -1 && \
echo SETUP_DONE"
CONFIG_OUTPUT=$(sandbox_exec "$CONFIG_SCRIPT" || true)
if echo "$CONFIG_OUTPUT" | grep -q "SETUP_DONE"; then
success "Plugin configured (mode: platform, user: $MEM0_USER_ID)"
else
warn "Could not verify config. You may need to configure manually:"
echo ""
if $IS_MACOS; then
echo " docker exec -it $CONTAINER_NAME bash"
echo " export NVM_DIR=/root/.nvm && source /root/.nvm/nvm.sh"
fi
echo " nemoclaw $SANDBOX_NAME connect"
echo " openclaw config set plugins.slots.memory openclaw-mem0"
echo " openclaw config set plugins.entries.openclaw-mem0.enabled true"
echo " openclaw config set plugins.entries.openclaw-mem0.config.apiKey \"$MEM0_API_KEY\""
echo " openclaw config set plugins.entries.openclaw-mem0.config.userId \"$MEM0_USER_ID\""
echo ""
fi
# ═════════════════════════════════════════════════════════════════════════════
# Done
# ═════════════════════════════════════════════════════════════════════════════
step "Verification"
echo ""
echo -e "${BOLD}${GREEN}╔══════════════════════════════════════════════════════════════╗${NC}"
echo -e "${BOLD}${GREEN}║ Setup Complete! ║${NC}"
echo -e "${BOLD}${GREEN}╚══════════════════════════════════════════════════════════════╝${NC}"
echo ""
echo -e " ${BOLD}Sandbox:${NC} $SANDBOX_NAME"
echo -e " ${BOLD}Plugin:${NC} @mem0/openclaw-mem0 (platform mode)"
echo -e " ${BOLD}User ID:${NC} $MEM0_USER_ID"
if $IS_MACOS; then
echo -e " ${BOLD}Container:${NC} $CONTAINER_NAME"
fi
echo ""
echo -e " ${BOLD}${CYAN}Next steps:${NC}"
echo ""
if $IS_MACOS; then
echo -e " 1. Open a shell in the container:"
echo ""
echo -e " ${DIM}docker exec -it $CONTAINER_NAME bash${NC}"
echo -e " ${DIM}export NVM_DIR=/root/.nvm && source /root/.nvm/nvm.sh${NC}"
echo ""
echo -e " 2. Connect to the sandbox and start the gateway:"
echo ""
echo -e " ${DIM}nemoclaw $SANDBOX_NAME connect${NC}"
echo -e " ${DIM}nemoclaw-start${NC}"
else
echo -e " 1. Connect to the sandbox and start the gateway:"
echo ""
echo -e " ${DIM}source ~/.bashrc${NC}"
echo -e " ${DIM}nemoclaw $SANDBOX_NAME connect${NC}"
echo -e " ${DIM}nemoclaw-start${NC}"
fi
echo ""
echo -e " Then verify the plugin loaded (look for 'openclaw-mem0: registered'):"
echo ""
echo -e " ${DIM}openclaw plugins list${NC}"
echo ""
echo -e " Test auto-capture (storing memories):"
echo ""
echo -e " ${DIM}openclaw agent --agent main --local -m \"My name is Alice\" --session-id test1${NC}"
echo ""
echo -e " Test auto-recall (new session, memories should appear):"
echo ""
echo -e " ${DIM}openclaw agent --agent main --local -m \"What do you know about me?\" --session-id test2${NC}"
echo ""
echo -e " Or use the interactive TUI:"
echo ""
echo -e " ${DIM}openclaw tui${NC}"
echo ""
echo -e " ${YELLOW}Note:${NC} You may see 'Telemetry event capture failed' errors."
echo -e " These are harmless and do not affect memory functionality."
echo ""
echo -e " ${BOLD}Documentation:${NC} https://docs.mem0.ai"
echo -e " ${BOLD}Plugin source:${NC} https://www.npmjs.com/package/@mem0/openclaw-mem0"
echo ""
+272
View File
@@ -0,0 +1,272 @@
# Mem0 Plugin for NemoClaw — Quickstart
Add persistent long-term memory to your [NemoClaw](https://docs.nvidia.com/nemoclaw/latest/get-started/quickstart.html) OpenClaw agent using the `@mem0/openclaw-mem0` plugin.
> **Note:** This plugin requires **Mem0 Platform mode** (i.e., a Mem0 API key from [app.mem0.ai](https://app.mem0.ai)). Open-source mode is not supported in NemoClaw sandboxes because the sandbox proxy blocks `/v1/embeddings` requests required by the open-source backend. See [Known Limitations](#known-limitations) for details.
## Prerequisites
| Resource | Recommended | Minimum |
|----------|-------------|---------|
| CPU | 4+ vCPU | 2 vCPU |
| RAM | 16 GB | 8 GB |
| Disk | 40 GB free | 20 GB free |
**Accounts required:**
- **NVIDIA** — sign up at [build.nvidia.com](https://build.nvidia.com), generate an API key at [build.nvidia.com/settings/api-keys](https://build.nvidia.com/settings/api-keys) (starts with `nvapi-`)
- **Mem0** — sign up at [app.mem0.ai](https://app.mem0.ai), generate an API key from the dashboard (starts with `m0-`)
**Supported platforms:** Ubuntu 22.04+, macOS (via Docker), Windows (WSL 2 + Docker)
---
## Choose Your Path
### Option A: Full Setup (NemoClaw + Mem0 Plugin)
Use this if you **don't have NemoClaw installed yet**. The script handles everything: Docker, Node.js, NemoClaw installation, onboarding, Mem0 plugin installation, network policy, and configuration.
```bash
# Download
curl -fsSL https://raw.githubusercontent.com/mem0ai/mem0/main/examples/nemoclaw/setup-mem0-nemoclaw.sh -o setup-mem0-nemoclaw.sh
# Run
chmod +x setup-mem0-nemoclaw.sh
./setup-mem0-nemoclaw.sh
```
The script runs through 7 phases:
| Phase | What it does | User input |
|-------|-------------|------------|
| 1 | Prerequisites (Docker, RAM, disk) | None (automatic) |
| 2 | Install NemoClaw | None (automatic) |
| 3 | NemoClaw onboarding (sandbox + k3s) | Sandbox name, NVIDIA API key |
| 4 | Install Mem0 plugin into sandbox | None (automatic) |
| 5 | Update network policy for `api.mem0.ai` | None (automatic) |
| 6 | Configure plugin | Mem0 API key, user ID |
| 7 | Verification | None |
### Option B: Plugin Only (NemoClaw Already Installed)
Use this if you **already have NemoClaw installed and onboarded** with a sandbox in `Ready` state.
```bash
# Download
curl -fsSL https://raw.githubusercontent.com/mem0ai/mem0/main/examples/nemoclaw/install-mem0-plugin.sh -o install-mem0-plugin.sh
# Run
chmod +x install-mem0-plugin.sh
./install-mem0-plugin.sh
```
The script auto-detects your sandbox and runs 3 steps:
| Step | What it does | User input |
|------|-------------|------------|
| 1 | Install Mem0 plugin into sandbox | None (automatic) |
| 2 | Update network policy for `api.mem0.ai` | None (automatic) |
| 3 | Configure plugin | Mem0 API key, user ID |
---
## Verify the Installation
After either script completes, connect to the sandbox and start the gateway:
```bash
source ~/.bashrc
nemoclaw <sandbox-name> connect
nemoclaw-start
```
Look for this line in the startup output:
```
openclaw-mem0: registered (mode: platform, user: your-user-id, graph: false, autoRecall: true, autoCapture: true)
```
You can also verify with:
```bash
openclaw plugins list
```
The `Memory (Mem0)` plugin should show status **loaded**.
## Test the Integration
All test commands run **inside the sandbox**.
### Test 1: Auto-capture (storing memories)
```bash
openclaw agent --agent main --local \
-m "My name is Alice and I work on distributed systems" \
--session-id test1
```
Look for: `openclaw-mem0: auto-captured 1 memories`
### Test 2: Auto-recall (retrieving memories across sessions)
Start a **new session** (different `--session-id`):
```bash
openclaw agent --agent main --local \
-m "What do you know about me?" \
--session-id test2
```
Look for: `openclaw-mem0: injecting 1 memories into context (1 long-term, 0 session)`
The agent should respond with information from the previous session ("Alice", "distributed systems").
### Test 3: Interactive TUI
```bash
openclaw tui
```
Send messages and the plugin will automatically capture and recall memories in the background.
---
## Plugin Configuration Reference
All options are set inside the sandbox via:
```bash
openclaw config set plugins.entries.openclaw-mem0.config.<key> <value>
```
| Key | Type | Default | Description |
|-----|------|---------|-------------|
| `mode` | `"platform"` \| `"open-source"` | `"platform"` | Backend mode |
| `apiKey` | string | — | Mem0 API key (starts with `m0-`) |
| `userId` | string | `"default"` | Unique identifier for the user |
| `autoRecall` | boolean | `true` | Inject memories before each agent turn |
| `autoCapture` | boolean | `true` | Store facts after each agent turn |
| `topK` | number | `5` | Max memories per recall |
| `searchThreshold` | number | `0.3` | Min similarity score (0-1) |
| `orgId` | string | — | Mem0 organization ID |
| `projectId` | string | — | Mem0 project ID |
| `enableGraph` | boolean | `false` | Enable entity graph for relationships |
| `customInstructions` | string | — | Rules for what Mem0 should store/exclude |
## Agent Memory Tools
Once the plugin is active, the agent can use these tools during conversations:
| Tool | Description |
|------|-------------|
| `memory_search` | Search memories by natural language |
| `memory_list` | List all stored memories for a user |
| `memory_store` | Explicitly save a fact |
| `memory_get` | Retrieve a memory by ID |
| `memory_forget` | Delete by ID or by query |
## CLI Commands
Run these inside the sandbox:
```bash
# Search all memories (long-term + session)
openclaw mem0 search "what languages does the user know"
# Search only long-term memories
openclaw mem0 search "user preferences" --scope long-term
# Search only session memories
openclaw mem0 search "current task" --scope session
# Memory stats
openclaw mem0 stats
```
---
## Troubleshooting
For detailed troubleshooting steps, see the [troubleshooting guide](troubleshooting-guide.pdf) included in this directory.
### Common Issues
#### `npm tar TAR_ENTRY_ERROR ENOENT` during NemoClaw installation
This is a known npm bug where concurrent tar extraction races cause `ENOENT` errors on deeply nested packages. The `setup-mem0-nemoclaw.sh` script works around this by cloning the NemoClaw repo and installing with `--maxsockets=1` to serialize downloads. If you installed NemoClaw manually and hit this error:
```bash
git clone --depth 1 https://github.com/NVIDIA/NemoClaw.git ~/.nemoclaw-src
cd ~/.nemoclaw-src
npm install --maxsockets=1
npm link
```
#### `sandbox not found` during onboarding step 7
The gateway restarted during onboarding and lost the sandbox state. Re-run `nemoclaw onboard`. When prompted that the sandbox already exists, choose `y` to recreate it.
#### `npm error 403 Forbidden` when installing plugin inside sandbox
The OpenShell gateway's TLS proxy blocks scoped npm packages (the `%2f` in the URL). Use the manual installation method (Method 2 in the scripts) which downloads outside the sandbox and uploads the tarball.
#### `capture failed: Connection error` or `recall failed: Connection error`
The network policy is not applied or missing the `mem0_api` entry. Re-run the plugin install script or manually apply the policy:
```bash
openshell policy set <sandbox-name> --policy /tmp/nemoclaw-mem0-policy.yaml --wait
```
#### `Telemetry event capture failed: TypeError: fetch failed`
This is harmless. The Mem0 SDK's telemetry endpoint (`us.i.posthog.com`) is blocked by the sandbox proxy. It does not affect memory functionality.
#### Plugin shows `disabled` in `openclaw plugins list`
The memory slot is not set to `openclaw-mem0`. Inside the sandbox, run:
```bash
openclaw config set plugins.slots.memory openclaw-mem0
```
#### `K8s namespace not ready` on Ubuntu 24.04
Ubuntu 24.04 defaults to cgroup v2 which causes k3s (used by NemoClaw) to fail. Apply the cgroup fix:
```bash
sudo python3 -c "
import json, os
p = '/etc/docker/daemon.json'
c = json.load(open(p)) if os.path.exists(p) else {}
c['default-cgroupns-mode'] = 'host'
json.dump(c, open(p, 'w'), indent=2)
"
sudo systemctl restart docker
```
Then re-run `nemoclaw onboard`.
### Known Limitations
**Open-source mode is not supported in NemoClaw sandboxes.** The Mem0 open-source mode requires calling `/v1/embeddings` on an external LLM provider. NemoClaw's sandbox proxy intercepts all OpenAI-compatible API requests but only allows `/v1/chat/completions` through. Use **platform mode** instead.
---
## Files in This Directory
| File | Description |
|------|-------------|
| `setup-mem0-nemoclaw.sh` | Full setup script (NemoClaw + Mem0 plugin) |
| `install-mem0-plugin.sh` | Plugin-only install script (NemoClaw already set up) |
| `troubleshooting-guide.pdf` | Detailed setup and troubleshooting guide |
| `quickstart.md` | This file |
## Links
- [Mem0 Documentation](https://docs.mem0.ai)
- [NemoClaw Documentation](https://docs.nvidia.com/nemoclaw/latest/get-started/quickstart.html)
- [`@mem0/openclaw-mem0` on npm](https://www.npmjs.com/package/@mem0/openclaw-mem0)
- [Mem0 Dashboard](https://app.mem0.ai)
+864
View File
@@ -0,0 +1,864 @@
#!/usr/bin/env bash
# ─────────────────────────────────────────────────────────────────────────────
# setup-mem0-nemoclaw.sh
#
# One-script setup for NemoClaw + Mem0 OpenClaw plugin.
# Supports Ubuntu servers (native) and macOS (via Docker container).
#
# Usage:
# curl -fsSL https://raw.githubusercontent.com/mem0ai/mem0/main/scripts/setup-mem0-nemoclaw.sh | bash
# # or
# chmod +x setup-mem0-nemoclaw.sh && ./setup-mem0-nemoclaw.sh
#
# Requirements:
# - Ubuntu 22.04+ OR macOS with Docker Desktop OR Windows with WSL 2 + Docker Desktop
# - 8 GB RAM minimum (16 GB recommended)
# - 40 GB free disk (20 GB minimum)
# - NVIDIA API key (from build.nvidia.com)
# - Mem0 API key (from app.mem0.ai)
# ─────────────────────────────────────────────────────────────────────────────
set -euo pipefail
# ── Colors and formatting ────────────────────────────────────────────────────
RED='\033[0;31m'
GREEN='\033[0;32m'
YELLOW='\033[1;33m'
BLUE='\033[0;34m'
CYAN='\033[0;36m'
BOLD='\033[1m'
DIM='\033[2m'
NC='\033[0m'
info() { echo -e "${BLUE}[INFO]${NC} $*"; }
success() { echo -e "${GREEN}[OK]${NC} $*"; }
warn() { echo -e "${YELLOW}[WARN]${NC} $*"; }
error() { echo -e "${RED}[ERROR]${NC} $*"; }
step() { echo -e "\n${BOLD}${CYAN}── $* ──${NC}\n"; }
ask() { echo -en "${BOLD}$*${NC}"; }
die() {
error "$*"
exit 1
}
# ── Defaults ─────────────────────────────────────────────────────────────────
SANDBOX_NAME="${SANDBOX_NAME:-my-assistant}"
MEM0_USER_ID="${MEM0_USER_ID:-default}"
MIN_RAM_MB=6000
MIN_DISK_MB=15000
PLUGIN_PKG="@mem0/openclaw-mem0"
CONTAINER_NAME="nemoclaw-dev"
# ── Detect platform ─────────────────────────────────────────────────────────
OS_TYPE="$(uname -s)"
IS_MACOS=false
IS_LINUX=false
IS_WSL=false
case "$OS_TYPE" in
Darwin) IS_MACOS=true ;;
Linux)
IS_LINUX=true
# Detect WSL (Windows Subsystem for Linux)
if grep -qi "microsoft\|wsl" /proc/version 2>/dev/null; then
IS_WSL=true
fi
;;
MINGW*|MSYS*|CYGWIN*)
error "This script cannot run in Git Bash, MSYS2, or Cygwin."
echo ""
echo " NemoClaw requires a full Linux environment. On Windows, use WSL 2:"
echo ""
echo " 1. Open PowerShell as Administrator and run:"
echo " wsl --install -d Ubuntu-24.04"
echo ""
echo " 2. Restart your computer when prompted"
echo ""
echo " 3. Open 'Ubuntu' from the Start menu (this opens a WSL 2 shell)"
echo ""
echo " 4. Install Docker Desktop for Windows:"
echo " https://www.docker.com/products/docker-desktop/"
echo " Enable 'Use the WSL 2 based engine' in Docker Desktop settings"
echo " Enable 'Ubuntu-24.04' under Resources → WSL Integration"
echo ""
echo " 5. In the Ubuntu WSL 2 shell, re-run this script:"
echo " curl -fsSL https://raw.githubusercontent.com/mem0ai/mem0/main/scripts/setup-mem0-nemoclaw.sh | bash"
echo ""
die "Please use WSL 2 instead."
;;
*)
die "Unsupported OS: $OS_TYPE. This script supports Linux (Ubuntu), macOS, and Windows (WSL 2)."
;;
esac
# ── Helper functions ─────────────────────────────────────────────────────────
check_command() {
command -v "$1" &>/dev/null
}
get_ram_mb() {
if $IS_MACOS; then
sysctl -n hw.memsize 2>/dev/null | awk '{printf "%.0f", $1/1024/1024}' || echo "0"
else
free -m 2>/dev/null | awk '/^Mem:/ {print $2}' || echo "0"
fi
}
get_disk_mb() {
if $IS_MACOS; then
df -m / 2>/dev/null | awk 'NR==2 {print $4}' || echo "0"
else
df -m / 2>/dev/null | awk 'NR==2 {print $4}' || echo "0"
fi
}
ensure_nvm() {
export NVM_DIR="${NVM_DIR:-$HOME/.nvm}"
if [[ -s "$NVM_DIR/nvm.sh" ]]; then
source "$NVM_DIR/nvm.sh"
fi
}
ensure_path() {
for p in "$HOME/.local/bin" "$HOME/.nvm/versions/node/"*/bin; do
if [[ -d "$p" ]] && [[ ":$PATH:" != *":$p:"* ]]; then
export PATH="$p:$PATH"
fi
done
}
wait_for_sandbox() {
local name="$1"
local timeout=120
local elapsed=0
while (( elapsed < timeout )); do
local phase
phase=$($RUN_CMD openshell sandbox get "$name" 2>/dev/null | grep -i "phase" | awk '{print $NF}' || true)
if [[ "$phase" == "Ready" ]]; then
return 0
fi
sleep 5
elapsed=$((elapsed + 5))
done
return 1
}
# ── macOS: run command inside the container ──────────────────────────────────
# On macOS, NemoClaw runs inside a Docker container. All nemoclaw/openshell
# commands must be exec'd into the container. On Linux, they run natively.
setup_run_cmd() {
if $IS_MACOS; then
RUN_CMD="docker exec $CONTAINER_NAME bash -c 'export NVM_DIR=/root/.nvm && source /root/.nvm/nvm.sh && "
RUN_CMD_SUFFIX="'"
RUN_CMD_IT="docker exec -it $CONTAINER_NAME bash -c 'export NVM_DIR=/root/.nvm && source /root/.nvm/nvm.sh && "
RUN_CMD_IT_SUFFIX="'"
else
RUN_CMD=""
RUN_CMD_SUFFIX=""
RUN_CMD_IT=""
RUN_CMD_IT_SUFFIX=""
fi
}
# Wrapper to run a command natively (Linux) or in the container (macOS)
run_cmd() {
if $IS_MACOS; then
docker exec "$CONTAINER_NAME" bash -c "export NVM_DIR=/root/.nvm && source /root/.nvm/nvm.sh && $*"
else
eval "$@"
fi
}
run_cmd_it() {
if $IS_MACOS; then
docker exec -it "$CONTAINER_NAME" bash -c "export NVM_DIR=/root/.nvm && source /root/.nvm/nvm.sh && $*"
else
eval "$@"
fi
}
# ── Banner ───────────────────────────────────────────────────────────────────
echo ""
echo -e "${BOLD}${CYAN}╔══════════════════════════════════════════════════════════════╗${NC}"
echo -e "${BOLD}${CYAN}║ Mem0 Plugin Setup for NemoClaw ║${NC}"
echo -e "${BOLD}${CYAN}║ Long-term memory for your OpenClaw agent ║${NC}"
echo -e "${BOLD}${CYAN}╚══════════════════════════════════════════════════════════════╝${NC}"
echo ""
if $IS_MACOS; then
info "Platform: macOS (will use Docker container)"
elif $IS_WSL; then
info "Platform: Windows (WSL 2)"
else
info "Platform: Linux"
fi
# ═════════════════════════════════════════════════════════════════════════════
# PHASE 1: Prerequisites
# ═════════════════════════════════════════════════════════════════════════════
step "Phase 1: Checking prerequisites"
# ── OS check ─────────────────────────────────────────────────────────────────
if $IS_LINUX; then
if [[ -f /etc/os-release ]]; then
source /etc/os-release
if [[ "$ID" != "ubuntu" ]]; then
warn "Detected OS: $PRETTY_NAME (not Ubuntu). This script is tested on Ubuntu 22.04+."
ask "Continue anyway? [y/N]: "
read -r cont
[[ "$cont" =~ ^[Yy]$ ]] || exit 0
else
success "OS: $PRETTY_NAME"
fi
else
warn "Cannot detect Linux distribution. Proceeding anyway."
fi
elif $IS_MACOS; then
MACOS_VERSION=$(sw_vers -productVersion 2>/dev/null || echo "unknown")
success "OS: macOS $MACOS_VERSION"
fi
# ── RAM check ────────────────────────────────────────────────────────────────
RAM_MB=$(get_ram_mb)
if (( RAM_MB < MIN_RAM_MB )); then
error "Insufficient RAM: ${RAM_MB} MB available, ${MIN_RAM_MB} MB required."
echo ""
echo " NemoClaw needs at least 8 GB RAM (16 GB recommended)."
if $IS_WSL; then
echo ""
echo " WSL 2 may have limited memory. Increase it:"
echo " 1. Create/edit %USERPROFILE%\\.wslconfig in Windows"
echo " 2. Add:"
echo " [wsl2]"
echo " memory=8GB"
echo " 3. Restart WSL: wsl --shutdown (from PowerShell)"
elif $IS_LINUX; then
echo ""
echo " If running on AWS EC2:"
echo " 1. Stop the instance"
echo " 2. Change instance type to t3.large (8 GB) or t3.xlarge (16 GB)"
echo " 3. Start the instance and re-run this script"
fi
echo ""
die "Aborting due to insufficient RAM."
else
success "RAM: ${RAM_MB} MB available"
fi
# ── Disk check ───────────────────────────────────────────────────────────────
DISK_MB=$(get_disk_mb)
if (( DISK_MB < MIN_DISK_MB )); then
error "Insufficient disk space: ${DISK_MB} MB available, ${MIN_DISK_MB} MB required."
echo ""
echo " NemoClaw needs at least 20 GB free disk (40 GB recommended)."
echo ""
echo " Quick fix:"
echo " docker system prune -a -f # Remove unused Docker data"
if $IS_LINUX; then
echo ""
echo " If running on AWS EC2, expand the EBS volume:"
echo " 1. Go to AWS Console → EC2 → Volumes"
echo " 2. Select the volume, Actions → Modify Volume → increase size"
echo " 3. Then run:"
echo " sudo growpart /dev/xvda 1"
echo " sudo resize2fs /dev/xvda1"
fi
echo ""
die "Aborting due to insufficient disk space."
else
success "Disk: ${DISK_MB} MB available"
fi
# ── Docker check ─────────────────────────────────────────────────────────────
if ! check_command docker; then
if $IS_MACOS; then
error "Docker not found."
echo ""
echo " Install Docker Desktop for macOS:"
echo " https://www.docker.com/products/docker-desktop/"
echo ""
echo " After installing, start Docker Desktop and re-run this script."
echo ""
die "Docker Desktop is required on macOS."
elif $IS_WSL; then
error "Docker not found inside WSL."
echo ""
echo " Docker Desktop must be installed on Windows with WSL 2 integration enabled:"
echo ""
echo " 1. Install Docker Desktop for Windows:"
echo " https://www.docker.com/products/docker-desktop/"
echo ""
echo " 2. Open Docker Desktop → Settings → General:"
echo " ✓ Enable 'Use the WSL 2 based engine'"
echo ""
echo " 3. Open Docker Desktop → Settings → Resources → WSL Integration:"
echo " ✓ Enable integration with your Ubuntu distribution"
echo ""
echo " 4. Click 'Apply & restart', then re-run this script in WSL."
echo ""
die "Docker Desktop WSL 2 integration is required."
else
info "Docker not found. Installing Docker..."
sudo apt-get update -qq
sudo apt-get install -y -qq docker.io >/dev/null 2>&1
sudo systemctl enable --now docker
sudo usermod -aG docker "$USER"
success "Docker installed"
fi
fi
if $IS_LINUX; then
# Ensure Docker daemon is running
if ! sudo docker info &>/dev/null; then
info "Starting Docker daemon..."
sudo systemctl start docker
sleep 2
fi
# Handle Docker group permissions
if ! docker info &>/dev/null; then
if id -nG "$USER" | grep -qw docker || grep -q "^docker:.*\b${USER}\b" /etc/group; then
info "Activating docker group for current session..."
exec sg docker -c "$0 $*"
else
info "Adding $USER to docker group..."
sudo usermod -aG docker "$USER"
info "Activating docker group for current session..."
exec sg docker -c "$0 $*"
fi
fi
else
# macOS: just verify Docker is responding
if ! docker info &>/dev/null; then
error "Docker is not running."
echo ""
echo " Start Docker Desktop and wait for it to be ready, then re-run this script."
echo ""
die "Docker Desktop is not running."
fi
fi
success "Docker: running ($(docker --version | awk '{print $3}' | tr -d ','))"
# ── Linux-only: cgroup v2 fix for Ubuntu 24.04 ──────────────────────────────
if $IS_LINUX && [[ "${VERSION_ID:-}" == "24.04" ]]; then
DAEMON_JSON="/etc/docker/daemon.json"
NEEDS_CGROUP_FIX=false
if [[ ! -f "$DAEMON_JSON" ]]; then
NEEDS_CGROUP_FIX=true
elif ! grep -q '"default-cgroupns-mode"' "$DAEMON_JSON" 2>/dev/null; then
NEEDS_CGROUP_FIX=true
fi
if $NEEDS_CGROUP_FIX; then
warn "Ubuntu 24.04 detected — applying cgroup v2 fix for Docker."
warn "This prevents 'K8s namespace not ready' errors during onboarding."
sudo python3 -c "
import json, os
p = '$DAEMON_JSON'
c = json.load(open(p)) if os.path.exists(p) else {}
c['default-cgroupns-mode'] = 'host'
json.dump(c, open(p, 'w'), indent=2)
"
sudo systemctl restart docker
success "Docker cgroup v2 fix applied"
else
success "Docker cgroup v2: already configured"
fi
fi
# ── Linux-only: Swap check ──────────────────────────────────────────────────
if $IS_LINUX; then
SWAP_MB=$(free -m 2>/dev/null | awk '/^Swap:/ {print $2}' || echo "0")
if (( RAM_MB < 12000 && SWAP_MB < 2000 )); then
warn "Low RAM (${RAM_MB} MB) and low swap. Adding 4 GB swap to prevent OOM kills."
if [[ ! -f /swapfile ]]; then
sudo fallocate -l 4G /swapfile
sudo chmod 600 /swapfile
sudo mkswap /swapfile >/dev/null
sudo swapon /swapfile
success "4 GB swap enabled"
else
success "Swap file already exists"
fi
fi
fi
# ── macOS-only: Create Docker container ──────────────────────────────────────
if $IS_MACOS; then
step "Phase 1b: Setting up NemoClaw Docker container"
if docker ps --format '{{.Names}}' | grep -q "^${CONTAINER_NAME}$"; then
success "Container '$CONTAINER_NAME' is already running"
elif docker ps -a --format '{{.Names}}' | grep -q "^${CONTAINER_NAME}$"; then
info "Starting existing container '$CONTAINER_NAME'..."
docker start "$CONTAINER_NAME"
success "Container started"
else
info "Creating Ubuntu container '$CONTAINER_NAME' for NemoClaw..."
echo -e " ${DIM}This container runs NemoClaw with --network host and Docker socket access.${NC}"
docker run -d \
--name "$CONTAINER_NAME" \
--privileged \
--network host \
-v /var/run/docker.sock:/var/run/docker.sock \
ubuntu:24.04 sleep infinity
success "Container '$CONTAINER_NAME' created"
info "Installing dependencies inside container..."
docker exec "$CONTAINER_NAME" bash -c "apt-get update -qq && apt-get install -y -qq curl git docker.io >/dev/null 2>&1"
success "Dependencies installed"
fi
# Check if NemoClaw is installed in the container
NEMO_IN_CONTAINER=$(docker exec "$CONTAINER_NAME" bash -c "export NVM_DIR=/root/.nvm && source /root/.nvm/nvm.sh 2>/dev/null && command -v nemoclaw" 2>/dev/null || true)
if [[ -z "$NEMO_IN_CONTAINER" ]]; then
info "Installing NemoClaw inside container (this takes a few minutes)..."
# The NVIDIA installer triggers npm tar race conditions. Workaround: clone
# and install with --maxsockets=1 to serialize downloads.
docker exec -it "$CONTAINER_NAME" bash -c "
export NVM_DIR=/root/.nvm &&
if [ ! -s \"\$NVM_DIR/nvm.sh\" ]; then
curl -fsSL https://raw.githubusercontent.com/nvm-sh/nvm/v0.40.3/install.sh | bash &&
source \"\$NVM_DIR/nvm.sh\" &&
nvm install 22
else
source \"\$NVM_DIR/nvm.sh\"
fi &&
git clone --depth 1 https://github.com/NVIDIA/NemoClaw.git /root/.nemoclaw-src &&
cd /root/.nemoclaw-src &&
npm install --maxsockets=1 &&
npm link
"
success "NemoClaw installed in container"
else
success "NemoClaw already installed in container"
fi
fi
# ═════════════════════════════════════════════════════════════════════════════
# PHASE 2: Install NemoClaw (Linux-only, macOS handled above)
# ═════════════════════════════════════════════════════════════════════════════
if $IS_LINUX; then
step "Phase 2: Installing NemoClaw"
ensure_nvm
ensure_path
if check_command nemoclaw; then
success "NemoClaw already installed: $(nemoclaw --version 2>/dev/null || echo 'found')"
else
info "Installing NemoClaw (this takes a few minutes)..."
# Install Node.js via nvm if not available
if ! check_command node; then
info "Node.js not found — installing via nvm..."
export NVM_DIR="${NVM_DIR:-$HOME/.nvm}"
curl -fsSL https://raw.githubusercontent.com/nvm-sh/nvm/v0.40.3/install.sh | bash
source "$NVM_DIR/nvm.sh"
nvm install 22
ensure_nvm
ensure_path
fi
# The NVIDIA installer (curl | bash) triggers npm tar race conditions
# causing ENOENT errors on deeply nested packages. Workaround: clone the
# repo and install with --maxsockets=1 to serialize downloads.
NEMOCLAW_DIR="$HOME/.nemoclaw-src"
rm -rf "$NEMOCLAW_DIR"
info "Cloning NemoClaw from GitHub..."
git clone --depth 1 https://github.com/NVIDIA/NemoClaw.git "$NEMOCLAW_DIR"
cd "$NEMOCLAW_DIR"
info "Installing dependencies (serialized to avoid tar race)..."
npm install --maxsockets=1
npm link
cd - >/dev/null
ensure_nvm
ensure_path
source "$HOME/.bashrc" 2>/dev/null || true
hash -r 2>/dev/null || true
if ! check_command nemoclaw; then
# npm link may place the binary outside the current PATH; find and add it
NEMOCLAW_BIN=$(find "$HOME/.nvm" -name nemoclaw \( -type f -o -type l \) -path "*/bin/*" 2>/dev/null | head -1)
if [[ -n "$NEMOCLAW_BIN" ]]; then
export PATH="$(dirname "$NEMOCLAW_BIN"):$PATH"
fi
fi
if ! check_command nemoclaw; then
die "NemoClaw installation failed. Try: source ~/.bashrc && nemoclaw --help"
fi
success "NemoClaw installed"
fi
ensure_path
if ! check_command openshell; then
warn "openshell not on PATH yet — it will be installed during onboarding."
fi
fi
# ═════════════════════════════════════════════════════════════════════════════
# PHASE 3: NemoClaw Onboarding
# ═════════════════════════════════════════════════════════════════════════════
step "Phase 3: NemoClaw Onboarding"
# Check if a sandbox already exists
EXISTING_SANDBOX=""
if $IS_MACOS; then
EXISTING_SANDBOX=$(run_cmd "openshell sandbox list 2>/dev/null" | awk 'NR>1 && $1!="" {print $1; exit}' || true)
else
ensure_path
if check_command openshell; then
EXISTING_SANDBOX=$(openshell sandbox list 2>/dev/null | awk 'NR>1 && $1!="" {print $1; exit}' || true)
fi
fi
if [[ -n "$EXISTING_SANDBOX" ]]; then
success "Sandbox already exists: $EXISTING_SANDBOX"
SANDBOX_NAME="$EXISTING_SANDBOX"
ask "Use existing sandbox '$SANDBOX_NAME'? [Y/n]: "
read -r use_existing
if [[ "$use_existing" =~ ^[Nn]$ ]]; then
ask "Enter sandbox name [my-assistant]: "
read -r custom_name
SANDBOX_NAME="${custom_name:-my-assistant}"
info "Running NemoClaw onboarding..."
echo -e " ${DIM}You'll need your NVIDIA API key (nvapi-...) from build.nvidia.com${NC}"
echo ""
run_cmd_it "nemoclaw onboard"
if $IS_LINUX; then ensure_path; fi
fi
else
info "No existing sandbox found. Running NemoClaw onboarding..."
echo ""
echo -e " ${DIM}The onboarding wizard will guide you through 7 steps:${NC}"
echo -e " ${DIM} 1. Preflight checks (automatic)${NC}"
echo -e " ${DIM} 2. Start gateway (automatic, takes 1-2 min)${NC}"
echo -e " ${DIM} 3. Sandbox name — enter a name (e.g. my-assistant)${NC}"
echo -e " ${DIM} 4. NVIDIA API key — paste your nvapi-... key${NC}"
echo -e " ${DIM} 5. Inference provider (automatic)${NC}"
echo -e " ${DIM} 6. OpenClaw setup (automatic)${NC}"
echo -e " ${DIM} 7. Policy presets — type Y to apply pypi and npm${NC}"
echo ""
ask "Press Enter to start onboarding..."
read -r
run_cmd_it "nemoclaw onboard"
if $IS_LINUX; then
ensure_nvm
ensure_path
fi
# Detect sandbox name
SANDBOX_NAME=$(run_cmd "openshell sandbox list 2>/dev/null" | awk 'NR>1 && $1!="" {print $1; exit}' || echo "$SANDBOX_NAME")
fi
# Verify sandbox exists
SANDBOX_PHASE=$(run_cmd "openshell sandbox get '$SANDBOX_NAME' 2>/dev/null" | grep -i "phase" | awk '{print $NF}' || true)
if [[ "$SANDBOX_PHASE" != "Ready" ]]; then
warn "Sandbox '$SANDBOX_NAME' is not in Ready state (current: ${SANDBOX_PHASE:-unknown})."
warn "Waiting up to 2 minutes..."
WAIT_OK=false
for i in $(seq 1 24); do
SANDBOX_PHASE=$(run_cmd "openshell sandbox get '$SANDBOX_NAME' 2>/dev/null" | grep -i "phase" | awk '{print $NF}' || true)
if [[ "$SANDBOX_PHASE" == "Ready" ]]; then
WAIT_OK=true
break
fi
sleep 5
done
if ! $WAIT_OK; then
die "Sandbox '$SANDBOX_NAME' did not become Ready. Run: openshell sandbox list"
fi
fi
success "Sandbox '$SANDBOX_NAME' is ready"
# ═════════════════════════════════════════════════════════════════════════════
# PHASE 4: Install Mem0 Plugin
# ═════════════════════════════════════════════════════════════════════════════
step "Phase 4: Installing Mem0 plugin ($PLUGIN_PKG)"
# Helper: run a command inside the sandbox non-interactively via piped stdin
sandbox_exec() {
local cmd="$1"
if $IS_MACOS; then
printf '%s\nexit\n' "$cmd" | docker exec -i "$CONTAINER_NAME" bash -c \
"export NVM_DIR=/root/.nvm && source /root/.nvm/nvm.sh && openshell sandbox connect '$SANDBOX_NAME'" 2>&1
else
printf '%s\nexit\n' "$cmd" | openshell sandbox connect "$SANDBOX_NAME" 2>&1
fi
}
# Check if plugin is already installed by trying to download a known file from the sandbox.
# We avoid sandbox_exec for the check because piped stdin to openshell sandbox connect
# produces shell prompt noise that causes false positives with grep.
PLUGIN_EXISTS=false
if run_cmd "openshell sandbox download '$SANDBOX_NAME' /sandbox/.openclaw/extensions/openclaw-mem0/package.json /tmp/_mem0_plugin_check.json" &>/dev/null; then
if [[ -f /tmp/_mem0_plugin_check.json ]] || run_cmd "test -f /tmp/_mem0_plugin_check.json" &>/dev/null; then
PLUGIN_EXISTS=true
fi
fi
rm -f /tmp/_mem0_plugin_check.json 2>/dev/null || true
run_cmd "rm -f /tmp/_mem0_plugin_check.json" 2>/dev/null || true
if $PLUGIN_EXISTS; then
success "Mem0 plugin already installed in sandbox"
else
info "Downloading $PLUGIN_PKG..."
# Download and build outside the sandbox (on host or in container)
run_cmd "cd /tmp && rm -rf openclaw-mem0-full mem0-openclaw-mem0-*.tgz openclaw-mem0-full.tgz"
if ! run_cmd "cd /tmp && npm pack '$PLUGIN_PKG' 2>/dev/null"; then
die "Failed to download $PLUGIN_PKG from npm. Check your internet connection."
fi
success "Downloaded plugin"
info "Installing plugin dependencies..."
run_cmd "mkdir -p /tmp/openclaw-mem0-full && cd /tmp/openclaw-mem0-full && tar xzf /tmp/mem0-openclaw-mem0-*.tgz --strip-components=1 && npm install --omit=dev 2>&1 | tail -3"
success "Dependencies installed"
info "Uploading plugin to sandbox..."
run_cmd "cd /tmp && tar czf openclaw-mem0-full.tgz -C openclaw-mem0-full ."
if ! run_cmd "openshell sandbox upload '$SANDBOX_NAME' /tmp/openclaw-mem0-full.tgz /sandbox/openclaw-mem0-full.tgz 2>&1"; then
die "Failed to upload plugin to sandbox. Check: openshell sandbox list"
fi
success "Plugin uploaded"
info "Extracting plugin inside sandbox..."
sandbox_exec "mkdir -p ~/.openclaw/extensions/openclaw-mem0 && tar xzf /sandbox/openclaw-mem0-full.tgz/openclaw-mem0-full.tgz -C ~/.openclaw/extensions/openclaw-mem0 2>/dev/null || tar xzf /sandbox/openclaw-mem0-full.tgz -C ~/.openclaw/extensions/openclaw-mem0 2>/dev/null && echo EXTRACT_OK" >/dev/null 2>&1 || true
# Verify
VERIFY_OK=false
if run_cmd "openshell sandbox download '$SANDBOX_NAME' /sandbox/.openclaw/extensions/openclaw-mem0/package.json /tmp/_mem0_verify.json" &>/dev/null; then
VERIFY_OK=true
fi
rm -f /tmp/_mem0_verify.json 2>/dev/null || true
run_cmd "rm -f /tmp/_mem0_verify.json" 2>/dev/null || true
if $VERIFY_OK; then
success "Plugin extracted inside sandbox"
else
warn "Could not verify plugin extraction. You may need to extract manually."
if $IS_MACOS; then
echo " docker exec -it $CONTAINER_NAME bash"
echo " export NVM_DIR=/root/.nvm && source /root/.nvm/nvm.sh"
fi
echo " nemoclaw $SANDBOX_NAME connect"
echo " mkdir -p ~/.openclaw/extensions/openclaw-mem0"
echo " tar xzf /sandbox/openclaw-mem0-full.tgz/openclaw-mem0-full.tgz -C ~/.openclaw/extensions/openclaw-mem0"
fi
# Clean up
run_cmd "rm -rf /tmp/openclaw-mem0-full /tmp/mem0-openclaw-mem0-*.tgz /tmp/openclaw-mem0-full.tgz" 2>/dev/null || true
fi
# ═════════════════════════════════════════════════════════════════════════════
# PHASE 5: Update Network Policy
# ═════════════════════════════════════════════════════════════════════════════
step "Phase 5: Updating network policy to allow api.mem0.ai and telemetry"
# Find baseline policy
BASELINE_POLICY=$(run_cmd "find / -path '*/nemoclaw-blueprint/policies/openclaw-sandbox.yaml' 2>/dev/null | head -1" || true)
if [[ -z "$BASELINE_POLICY" ]]; then
die "Cannot find NemoClaw baseline policy file. Is NemoClaw installed?"
fi
info "Baseline policy: $BASELINE_POLICY"
# Check if mem0_api already exists
HAS_MEM0=$(run_cmd "grep -c mem0_api '$BASELINE_POLICY' 2>/dev/null" || echo "0")
if [[ "$HAS_MEM0" != "0" ]]; then
success "mem0_api already in baseline policy"
else
info "Adding api.mem0.ai to network policy..."
fi
# Create custom policy with mem0_api + telemetry — use node for reliable cross-platform YAML editing
run_cmd "node -e \"
const fs = require('fs');
let c = fs.readFileSync('$BASELINE_POLICY', 'utf8');
const mem0Block = '\\n mem0_api:\\n name: mem0_api\\n endpoints:\\n - host: api.mem0.ai\\n port: 443\\n access: full\\n binaries:\\n - { path: /usr/local/bin/node }\\n - { path: /usr/local/bin/openclaw }\\n';
const telemetryBlock = '\\n mem0_telemetry:\\n name: mem0_telemetry\\n endpoints:\\n - host: us.i.posthog.com\\n port: 443\\n access: full\\n binaries:\\n - { path: /usr/local/bin/node }\\n - { path: /usr/local/bin/openclaw }\\n';
if (!c.includes('mem0_api')) {
if (c.includes('# ── Messaging')) {
c = c.replace(' # ── Messaging', mem0Block + '\\n # ── Messaging');
} else {
c += mem0Block;
}
}
if (!c.includes('mem0_telemetry')) {
if (c.includes('mem0_api:')) {
c = c.replace(' mem0_api:', telemetryBlock + '\\n mem0_api:');
} else {
c += telemetryBlock;
}
}
fs.writeFileSync('/tmp/nemoclaw-mem0-policy.yaml', c);
console.log('ok');
\""
success "Custom policy file created"
# Apply the policy
info "Applying network policy..."
if ! run_cmd "openshell policy set '$SANDBOX_NAME' --policy /tmp/nemoclaw-mem0-policy.yaml --wait 2>&1"; then
error "Failed to apply network policy."
echo ""
echo " If you see 'sandbox not found', re-run: nemoclaw onboard"
echo " Then re-run this script."
echo ""
die "Network policy update failed."
fi
success "Network policy applied — api.mem0.ai and telemetry allowed"
# ═════════════════════════════════════════════════════════════════════════════
# PHASE 6: Configure Plugin
# ═════════════════════════════════════════════════════════════════════════════
step "Phase 6: Configuring Mem0 plugin"
echo ""
echo -e " ${DIM}Get your Mem0 API key from: https://app.mem0.ai${NC}"
echo -e " ${DIM}The key starts with 'm0-'${NC}"
echo ""
ask "Enter your Mem0 API key: "
read -r MEM0_API_KEY
if [[ -z "$MEM0_API_KEY" ]]; then
die "Mem0 API key is required."
fi
if [[ ! "$MEM0_API_KEY" =~ ^m0- ]]; then
warn "Key doesn't start with 'm0-'. Make sure this is correct."
fi
echo ""
echo -e " ${DIM}The user ID scopes all memories. Pick any unique identifier.${NC}"
echo -e " ${DIM}Examples: alice, user_123, your-email@example.com${NC}"
echo ""
ask "Enter user ID [$MEM0_USER_ID]: "
read -r custom_user_id
MEM0_USER_ID="${custom_user_id:-$MEM0_USER_ID}"
info "Configuring plugin inside sandbox..."
CONFIG_SCRIPT="openclaw config set plugins.slots.memory openclaw-mem0 2>&1 | tail -1 && \
openclaw config set plugins.entries.openclaw-mem0.enabled true 2>&1 | tail -1 && \
openclaw config set plugins.entries.openclaw-mem0.config.apiKey '$MEM0_API_KEY' 2>&1 | tail -1 && \
openclaw config set plugins.entries.openclaw-mem0.config.userId '$MEM0_USER_ID' 2>&1 | tail -1 && \
echo SETUP_DONE"
CONFIG_OUTPUT=$(sandbox_exec "$CONFIG_SCRIPT" || true)
if echo "$CONFIG_OUTPUT" | grep -q "SETUP_DONE"; then
success "Plugin configured (mode: platform, user: $MEM0_USER_ID)"
else
warn "Could not verify config. You may need to configure manually:"
echo ""
if $IS_MACOS; then
echo " docker exec -it $CONTAINER_NAME bash"
echo " export NVM_DIR=/root/.nvm && source /root/.nvm/nvm.sh"
fi
echo " nemoclaw $SANDBOX_NAME connect"
echo " openclaw config set plugins.slots.memory openclaw-mem0"
echo " openclaw config set plugins.entries.openclaw-mem0.enabled true"
echo " openclaw config set plugins.entries.openclaw-mem0.config.apiKey \"$MEM0_API_KEY\""
echo " openclaw config set plugins.entries.openclaw-mem0.config.userId \"$MEM0_USER_ID\""
echo ""
fi
# ═════════════════════════════════════════════════════════════════════════════
# PHASE 7: Verify & Test
# ═════════════════════════════════════════════════════════════════════════════
step "Phase 7: Verification"
echo ""
echo -e "${BOLD}${GREEN}╔══════════════════════════════════════════════════════════════╗${NC}"
echo -e "${BOLD}${GREEN}║ Setup Complete! ║${NC}"
echo -e "${BOLD}${GREEN}╚══════════════════════════════════════════════════════════════╝${NC}"
echo ""
echo -e " ${BOLD}Sandbox:${NC} $SANDBOX_NAME"
echo -e " ${BOLD}Plugin:${NC} @mem0/openclaw-mem0 (platform mode)"
echo -e " ${BOLD}User ID:${NC} $MEM0_USER_ID"
if $IS_MACOS; then
echo -e " ${BOLD}Container:${NC} $CONTAINER_NAME"
fi
echo ""
echo -e " ${BOLD}${CYAN}Next steps:${NC}"
echo ""
if $IS_MACOS; then
echo -e " 1. Open a shell in the container:"
echo ""
echo -e " ${DIM}docker exec -it $CONTAINER_NAME bash${NC}"
echo -e " ${DIM}export NVM_DIR=/root/.nvm && source /root/.nvm/nvm.sh${NC}"
echo ""
echo -e " 2. Connect to the sandbox and start the gateway:"
echo ""
echo -e " ${DIM}nemoclaw $SANDBOX_NAME connect${NC}"
echo -e " ${DIM}nemoclaw-start${NC}"
else
echo -e " 1. Connect to the sandbox and start the gateway:"
echo ""
echo -e " ${DIM}source ~/.bashrc${NC}"
echo -e " ${DIM}nemoclaw $SANDBOX_NAME connect${NC}"
echo -e " ${DIM}nemoclaw-start${NC}"
fi
echo ""
echo -e " Then verify the plugin loaded (look for 'openclaw-mem0: registered'):"
echo ""
echo -e " ${DIM}openclaw plugins list${NC}"
echo ""
echo -e " Test auto-capture (storing memories):"
echo ""
echo -e " ${DIM}openclaw agent --agent main --local -m \"My name is Alice\" --session-id test1${NC}"
echo ""
echo -e " Test auto-recall (new session, memories should appear):"
echo ""
echo -e " ${DIM}openclaw agent --agent main --local -m \"What do you know about me?\" --session-id test2${NC}"
echo ""
echo -e " Or use the interactive TUI:"
echo ""
echo -e " ${DIM}openclaw tui${NC}"
echo ""
echo -e " ${YELLOW}Note:${NC} You may see 'Telemetry event capture failed' errors."
echo -e " These are harmless and do not affect memory functionality."
echo ""
echo -e " ${BOLD}Documentation:${NC} https://docs.mem0.ai"
echo -e " ${BOLD}Plugin source:${NC} https://www.npmjs.com/package/@mem0/openclaw-mem0"
echo ""
Binary file not shown.
-6
View File
@@ -34,12 +34,6 @@ export const DEFAULT_MEMORY_CONFIG: MemoryConfig = {
username: process.env.NEO4J_USERNAME || "neo4j",
password: process.env.NEO4J_PASSWORD || "password",
},
llm: {
provider: "openai",
config: {
model: "gpt-4-turbo-preview",
},
},
},
historyStore: {
provider: "sqlite",
+1
View File
@@ -26,4 +26,5 @@ export * from "./vector_stores/supabase";
export * from "./vector_stores/langchain";
export * from "./vector_stores/vectorize";
export * from "./vector_stores/azure_ai_search";
export * from "./vector_stores/pgvector";
export * from "./utils/factory";
+5 -5
View File
@@ -80,18 +80,18 @@ export class MemoryGraph {
);
this.llmProvider = "openai";
let llmConfig = this.config.llm.config;
if (this.config.llm?.provider) {
this.llmProvider = this.config.llm.provider;
}
if (this.config.graphStore?.llm?.provider) {
this.llmProvider = this.config.graphStore.llm.provider;
llmConfig = this.config.graphStore.llm.config ?? llmConfig;
}
this.llm = LLMFactory.create(this.llmProvider, this.config.llm.config);
this.structuredLlm = LLMFactory.create(
this.llmProvider,
this.config.llm.config,
);
this.llm = LLMFactory.create(this.llmProvider, llmConfig);
this.structuredLlm = LLMFactory.create(this.llmProvider, llmConfig);
this.threshold = 0.7;
}
+1
View File
@@ -31,6 +31,7 @@ export const MemoryUpdateSchema = z.object({
old_memory: z
.string()
.optional()
.nullable()
.describe(
"The previous content of the memory item if the event was UPDATE.",
),
+3
View File
@@ -34,6 +34,7 @@ import { LangchainLLM } from "../llms/langchain";
import { LangchainEmbedder } from "../embeddings/langchain";
import { LangchainVectorStore } from "../vector_stores/langchain";
import { AzureAISearch } from "../vector_stores/azure_ai_search";
import { PGVector } from "../vector_stores/pgvector";
export class EmbedderFactory {
static create(provider: string, config: EmbeddingConfig): Embedder {
@@ -104,6 +105,8 @@ export class VectorStoreFactory {
return new VectorizeDB(config as any);
case "azure-ai-search":
return new AzureAISearch(config as any);
case "pgvector":
return new PGVector(config as any);
default:
throw new Error(`Unsupported vector store provider: ${provider}`);
}
@@ -1,4 +1,4 @@
import { Client, Pool } from "pg";
import { Client } from "pg";
import { VectorStore } from "./base";
import { SearchFilters, VectorStoreConfig, VectorStoreResult } from "../types";
@@ -35,6 +35,7 @@ export class PGVector implements VectorStore {
host: config.host,
port: config.port,
});
this.initialize().catch(console.error);
}
async initialize(): Promise<void> {
+1 -1
View File
@@ -522,7 +522,7 @@ export class RedisDB implements VectorStore {
return {
id: vectorId,
payload,
payload: toCamelCase(payload),
};
} catch (error) {
console.error("Error getting vector:", error);
@@ -349,6 +349,78 @@ describe("ConfigManager", () => {
});
});
// ─────────────────────────────────────────────────────────────────────────
// Graph store LLM config propagation (issue #3425)
// ─────────────────────────────────────────────────────────────────────────
describe("mergeConfig - graph store LLM config (issue #3425)", () => {
const baseEmbedder = {
provider: "openai",
config: { apiKey: "test-key" },
};
const baseVectorStore = {
provider: "memory",
config: { collectionName: "test" },
};
const graphStoreNeo4j = {
provider: "neo4j",
config: {
url: "neo4j://localhost:7687",
username: "neo4j",
password: "password",
},
};
it("should NOT have a default graphStore.llm — root llm should be the fallback", () => {
const config = ConfigManager.mergeConfig({
embedder: baseEmbedder,
vectorStore: baseVectorStore,
llm: {
provider: "anthropic",
config: { model: "claude-sonnet-4-20250514" },
},
graphStore: graphStoreNeo4j,
});
// graphStore should NOT have its own llm after merge
expect(config.graphStore?.llm).toBeUndefined();
// root llm should be anthropic
expect(config.llm.provider).toBe("anthropic");
expect(config.llm.config.model).toBe("claude-sonnet-4-20250514");
});
it("should preserve explicit graphStore.llm when user provides it", () => {
const config = ConfigManager.mergeConfig({
embedder: baseEmbedder,
vectorStore: baseVectorStore,
llm: {
provider: "anthropic",
config: { model: "claude-sonnet-4-20250514" },
},
graphStore: {
...graphStoreNeo4j,
llm: { provider: "openai", config: { model: "gpt-4o" } },
},
});
// graphStore should have its own llm
expect(config.graphStore?.llm?.provider).toBe("openai");
expect(config.graphStore?.llm?.config).toEqual({ model: "gpt-4o" });
// root llm should still be anthropic
expect(config.llm.provider).toBe("anthropic");
});
it("should not have graphStore.llm when user does not provide one", () => {
const config = ConfigManager.mergeConfig({
embedder: baseEmbedder,
vectorStore: baseVectorStore,
llm: { provider: "openai", config: { model: "gpt-4o" } },
});
// Default graphStore should not have llm
expect(config.graphStore?.llm).toBeUndefined();
});
});
// ─────────────────────────────────────────────────────────────────────────
// Memory class – LM Studio end-to-end flow (mocked factories)
// ─────────────────────────────────────────────────────────────────────────
@@ -118,6 +118,11 @@ jest.mock("../src/vector_stores/azure_ai_search", () => ({
.fn()
.mockImplementation((config) => ({ type: "azure-ai-search", config })),
}));
jest.mock("../src/vector_stores/pgvector", () => ({
PGVector: jest
.fn()
.mockImplementation((config) => ({ type: "pgvector", config })),
}));
jest.mock("../src/storage/SupabaseHistoryManager", () => ({
SupabaseHistoryManager: jest
.fn()
@@ -236,6 +241,7 @@ describe("VectorStoreFactory", () => {
["langchain"],
["vectorize"],
["azure-ai-search"],
["pgvector"],
])("creates vector store for provider '%s'", (provider) => {
expect(() =>
VectorStoreFactory.create(provider, dummyVSConfig),
@@ -511,6 +511,137 @@ describe("Prompt construction — all json_object sites include 'json'", () => {
// 5. Edge cases – malformed entity fields in _removeSpacesFromEntities
// ═══════════════════════════════════════════════════════════════════════════
// ═══════════════════════════════════════════════════════════════════════════
// 5a. LLM config propagation — graph store uses correct provider & config
// Regression test for https://github.com/mem0ai/mem0/issues/3425
// ═══════════════════════════════════════════════════════════════════════════
describe("LLM config propagation to graph store (issue #3425)", () => {
const { LLMFactory } = require("../src/utils/factory");
beforeEach(() => {
(LLMFactory.create as jest.Mock).mockClear();
});
it("uses root llm config when no graphStore.llm is provided", () => {
const config = {
graphStore: {
config: {
url: "bolt://localhost:7687",
username: "neo4j",
password: "test",
},
},
embedder: { provider: "openai", config: {} },
llm: {
provider: "anthropic",
config: { model: "claude-sonnet-4-20250514", apiKey: "sk-ant-test" },
},
} as any;
new MemoryGraph(config);
expect(LLMFactory.create).toHaveBeenCalledWith("anthropic", {
model: "claude-sonnet-4-20250514",
apiKey: "sk-ant-test",
});
// Both llm and structuredLlm should use the same config
expect(LLMFactory.create).toHaveBeenCalledTimes(2);
expect(LLMFactory.create).toHaveBeenNthCalledWith(1, "anthropic", {
model: "claude-sonnet-4-20250514",
apiKey: "sk-ant-test",
});
expect(LLMFactory.create).toHaveBeenNthCalledWith(2, "anthropic", {
model: "claude-sonnet-4-20250514",
apiKey: "sk-ant-test",
});
});
it("uses graphStore.llm config when provided, overriding root llm", () => {
const config = {
graphStore: {
config: {
url: "bolt://localhost:7687",
username: "neo4j",
password: "test",
},
llm: {
provider: "openai",
config: { model: "gpt-4o", apiKey: "sk-openai-test" },
},
},
embedder: { provider: "openai", config: {} },
llm: {
provider: "anthropic",
config: { model: "claude-sonnet-4-20250514", apiKey: "sk-ant-test" },
},
} as any;
new MemoryGraph(config);
// Should use graphStore.llm, NOT root llm
expect(LLMFactory.create).toHaveBeenNthCalledWith(1, "openai", {
model: "gpt-4o",
apiKey: "sk-openai-test",
});
expect(LLMFactory.create).toHaveBeenNthCalledWith(2, "openai", {
model: "gpt-4o",
apiKey: "sk-openai-test",
});
});
it("falls back to root llm config when graphStore.llm.config is undefined", () => {
// Note: in practice, Zod schema requires config when graphStore.llm is
// present. This tests the defensive fallback in MemoryGraph itself.
const config = {
graphStore: {
config: {
url: "bolt://localhost:7687",
username: "neo4j",
password: "test",
},
llm: {
provider: "openai",
// config explicitly undefined
config: undefined,
},
},
embedder: { provider: "openai", config: {} },
llm: {
provider: "anthropic",
config: { model: "claude-sonnet-4-20250514" },
},
} as any;
new MemoryGraph(config);
// Provider from graphStore.llm, but config falls back to root llm.config
expect(LLMFactory.create).toHaveBeenNthCalledWith(1, "openai", {
model: "claude-sonnet-4-20250514",
});
});
it("defaults to openai when neither root nor graphStore llm provider is set", () => {
const config = {
graphStore: {
config: {
url: "bolt://localhost:7687",
username: "neo4j",
password: "test",
},
},
embedder: { provider: "openai", config: {} },
llm: { config: { model: "gpt-4" } },
} as any;
new MemoryGraph(config);
expect(LLMFactory.create).toHaveBeenNthCalledWith(1, "openai", {
model: "gpt-4",
});
});
});
describe("_removeSpacesFromEntities (via _establishNodesRelationsFromData)", () => {
it("normalises spaces and case in entity source/relationship/destination", async () => {
mockGenerateResponse.mockResolvedValueOnce({
+2 -2
View File
@@ -16,7 +16,7 @@ class AnthropicConfig(BaseLlmConfig):
temperature: float = 0.1,
api_key: Optional[str] = None,
max_tokens: int = 2000,
top_p: float = 0.1,
top_p: Optional[float] = None,
top_k: int = 1,
enable_vision: bool = False,
vision_details: Optional[str] = "auto",
@@ -32,7 +32,7 @@ class AnthropicConfig(BaseLlmConfig):
temperature: Controls randomness, defaults to 0.1
api_key: Anthropic API key, defaults to None
max_tokens: Maximum tokens to generate, defaults to 2000
top_p: Nucleus sampling parameter, defaults to 0.1
top_p: Nucleus sampling parameter, defaults to None (omitted to avoid conflict with temperature)
top_k: Top-k sampling parameter, defaults to 1
enable_vision: Enable vision capabilities, defaults to False
vision_details: Vision detail level, defaults to "auto"
+56
View File
@@ -0,0 +1,56 @@
from typing import Optional
from mem0.configs.llms.base import BaseLlmConfig
class MinimaxConfig(BaseLlmConfig):
"""
Configuration class for MiniMax-specific parameters.
Inherits from BaseLlmConfig and adds MiniMax-specific settings.
"""
def __init__(
self,
# Base parameters
model: Optional[str] = None,
temperature: float = 0.1,
api_key: Optional[str] = None,
max_tokens: int = 2000,
top_p: float = 0.1,
top_k: int = 1,
enable_vision: bool = False,
vision_details: Optional[str] = "auto",
http_client_proxies: Optional[dict] = None,
# MiniMax-specific parameters
minimax_base_url: Optional[str] = None,
):
"""
Initialize MiniMax configuration.
Args:
model: MiniMax model to use, defaults to None
temperature: Controls randomness, defaults to 0.1
api_key: MiniMax API key, defaults to None
max_tokens: Maximum tokens to generate, defaults to 2000
top_p: Nucleus sampling parameter, defaults to 0.1
top_k: Top-k sampling parameter, defaults to 1
enable_vision: Enable vision capabilities, defaults to False
vision_details: Vision detail level, defaults to "auto"
http_client_proxies: HTTP client proxy settings, defaults to None
minimax_base_url: MiniMax API base URL, defaults to None
"""
# Initialize base parameters
super().__init__(
model=model,
temperature=temperature,
api_key=api_key,
max_tokens=max_tokens,
top_p=top_p,
top_k=top_k,
enable_vision=enable_vision,
vision_details=vision_details,
http_client_proxies=http_client_proxies,
)
# MiniMax-specific parameters
self.minimax_base_url = minimax_base_url
+2 -2
View File
@@ -32,8 +32,8 @@ class ChromaDbConfig(BaseModel):
values.pop("path", None)
return values
# Check if local/server configuration is provided (excluding default tmp path for cloud config)
local_config = bool(path and path != "/tmp/chroma") or bool(host and port)
# Check if local/server configuration is provided
local_config = bool(path) or bool(host and port)
if not cloud_config and not local_config:
raise ValueError("Either ChromaDB Cloud configuration (api_key, tenant) or local configuration (path or host/port) must be provided.")
+1 -1
View File
@@ -16,7 +16,7 @@ class QdrantConfig(BaseModel):
path: Optional[str] = Field("/tmp/qdrant", description="Path for local Qdrant database")
url: Optional[str] = Field(None, description="Full URL for Qdrant server")
api_key: Optional[str] = Field(None, description="API key for Qdrant server")
on_disk: Optional[bool] = Field(False, description="Enables persistent storage")
on_disk: Optional[bool] = Field(False,description="Enables persistent storage. Vectors are kept on disk (True) or in memory (False). Does not delete the local database path.")
@model_validator(mode="before")
@classmethod
+45
View File
@@ -0,0 +1,45 @@
import os
from typing import Any, Dict, Optional
from pydantic import BaseModel, ConfigDict, Field, model_validator
class TurbopufferConfig(BaseModel):
collection_name: str = Field("mem0", description="Name of the namespace/collection")
embedding_model_dims: int = Field(1536, description="Dimensions of the embedding model")
api_key: Optional[str] = Field(None, description="API key for Turbopuffer")
region: str = Field("gcp-us-central1", description="Turbopuffer region (e.g., 'gcp-us-central1', 'aws-us-west-2')")
distance_metric: str = Field(
"cosine_distance",
description="Distance metric for vector similarity ('cosine_distance' or 'euclidean_squared')",
)
batch_size: int = Field(100, description="Batch size for bulk operations")
extra_params: Optional[Dict[str, Any]] = Field(
None,
description="Additional parameters for Turbopuffer client",
)
@model_validator(mode="before")
@classmethod
def check_api_key(cls, values: Dict[str, Any]) -> Dict[str, Any]:
api_key = values.get("api_key")
if not api_key and "TURBOPUFFER_API_KEY" not in os.environ:
raise ValueError(
"Either 'api_key' must be provided or TURBOPUFFER_API_KEY environment variable must be set."
)
return values
@model_validator(mode="before")
@classmethod
def validate_extra_fields(cls, values: Dict[str, Any]) -> Dict[str, Any]:
allowed_fields = set(cls.model_fields.keys())
input_fields = set(values.keys())
extra_fields = input_fields - allowed_fields
if extra_fields:
raise ValueError(
f"Extra fields not allowed: {', '.join(extra_fields)}. "
f"Please input only the following fields: {', '.join(allowed_fields)}"
)
return values
model_config = ConfigDict(arbitrary_types_allowed=True)
+24 -2
View File
@@ -77,12 +77,32 @@ class KuzuConfig(BaseModel):
db: Optional[str] = Field(":memory:", description="Path to a Kuzu database file")
class ApacheAgeConfig(BaseModel):
host: Optional[str] = Field("localhost", description="PostgreSQL server hostname")
port: Optional[int] = Field(5432, description="PostgreSQL server port")
database: Optional[str] = Field(None, description="PostgreSQL database name")
username: Optional[str] = Field(None, description="PostgreSQL username")
password: Optional[str] = Field(None, description="PostgreSQL password")
graph_name: Optional[str] = Field("mem0_graph", description="Name of the Apache AGE graph")
@model_validator(mode="before")
def check_required_fields(cls, values):
database, username, password = (
values.get("database"),
values.get("username"),
values.get("password"),
)
if not database or not username or not password:
raise ValueError("Please provide 'database', 'username' and 'password'.")
return values
class GraphStoreConfig(BaseModel):
provider: str = Field(
description="Provider of the data store (e.g., 'neo4j', 'memgraph', 'neptune', 'kuzu')",
description="Provider of the data store (e.g., 'neo4j', 'memgraph', 'neptune', 'kuzu', 'apache_age')",
default="neo4j",
)
config: Union[Neo4jConfig, MemgraphConfig, NeptuneConfig, KuzuConfig] = Field(
config: Union[Neo4jConfig, MemgraphConfig, NeptuneConfig, KuzuConfig, ApacheAgeConfig] = Field(
description="Configuration for the specific data store", default=None
)
llm: Optional[LlmConfig] = Field(description="LLM configuration for querying the graph store", default=None)
@@ -110,5 +130,7 @@ class GraphStoreConfig(BaseModel):
return NeptuneConfig(**v.model_dump())
elif provider == "kuzu":
return KuzuConfig(**v.model_dump())
elif provider == "apache_age":
return ApacheAgeConfig(**v.model_dump())
else:
raise ValueError(f"Unsupported graph store provider: {provider}")
+22
View File
@@ -409,6 +409,28 @@ class NeptuneBase(ABC):
"""
pass
def delete(self, data, filters):
"""
Delete graph entities associated with the given memory text.
Extracts entities and relationships from the memory text using the same
pipeline as add(), then deletes the matching relationships in the graph.
Args:
data (str): The memory text whose graph entities should be removed.
filters (dict): Scope filters (user_id, agent_id, run_id).
"""
try:
entity_type_map = self._retrieve_nodes_from_data(data, filters)
if not entity_type_map:
logger.debug("No entities found in memory text, skipping graph cleanup")
return
to_be_deleted = self._establish_nodes_relations_from_data(data, filters, entity_type_map)
if to_be_deleted:
self._delete_entities(to_be_deleted, filters["user_id"])
except Exception as e:
logger.error(f"Error during graph cleanup for memory delete: {e}")
def delete_all(self, filters):
cypher, params = self._delete_all_cypher(filters)
self.graph.query(cypher, params=params)
+25
View File
@@ -40,6 +40,31 @@ class AnthropicLLM(LLMBase):
api_key = self.config.api_key or os.getenv("ANTHROPIC_API_KEY")
self.client = anthropic.Anthropic(api_key=api_key)
def _get_common_params(self, **kwargs) -> Dict:
"""Get common parameters, avoiding sending both temperature and top_p together.
Anthropic rejects requests that include both temperature and top_p.
When both are set, we keep temperature and drop top_p.
"""
params = {}
if self.config.max_tokens is not None:
params["max_tokens"] = self.config.max_tokens
has_temperature = self.config.temperature is not None
has_top_p = self.config.top_p is not None
if has_temperature and has_top_p:
# Anthropic forbids both; prefer temperature
params["temperature"] = self.config.temperature
elif has_temperature:
params["temperature"] = self.config.temperature
elif has_top_p:
params["top_p"] = self.config.top_p
params.update(kwargs)
return params
def generate_response(
self,
messages: List[Dict[str, str]],
+32 -4
View File
@@ -1,3 +1,4 @@
import json
import os
from typing import Dict, List, Optional
@@ -6,6 +7,7 @@ from openai import AzureOpenAI
from mem0.configs.llms.base import BaseLlmConfig
from mem0.llms.base import LLMBase
from mem0.memory.utils import extract_json
SCOPE = "https://cognitiveservices.azure.com/.default"
@@ -83,9 +85,35 @@ class AzureOpenAIStructuredLLM(LLMBase):
params["tools"] = tools
params["tool_choice"] = tool_choice
if tools:
params["tools"] = tools
params["tool_choice"] = tool_choice
response = self.client.chat.completions.create(**params)
return self._parse_response(response, tools)
def _parse_response(self, response, tools):
"""
Process the response based on whether tools are used or not.
Args:
response: The raw response from API.
tools: The list of tools provided in the request.
Returns:
str or dict: The processed response.
"""
if tools:
processed_response = {
"content": response.choices[0].message.content,
"tool_calls": [],
}
if response.choices[0].message.tool_calls:
for tool_call in response.choices[0].message.tool_calls:
processed_response["tool_calls"].append(
{
"name": tool_call.function.name,
"arguments": json.loads(extract_json(tool_call.function.arguments)),
}
)
return processed_response
else:
return response.choices[0].message.content
+1
View File
@@ -23,6 +23,7 @@ class LlmConfig(BaseModel):
"azure_openai_structured",
"gemini",
"deepseek",
"minimax",
"xai",
"sarvam",
"lmstudio",
+11 -8
View File
@@ -32,22 +32,25 @@ class GeminiLLM(LLMBase):
Returns:
str or dict: The processed response.
"""
# Get parts safely — content can be None when Gemini blocks the response
candidate = response.candidates[0] if response.candidates else None
parts = candidate.content.parts if candidate and candidate.content else None
if tools:
processed_response = {
"content": None,
"tool_calls": [],
}
# Extract content from the first candidate
if response.candidates and response.candidates[0].content.parts:
for part in response.candidates[0].content.parts:
if parts:
# Extract content from the first candidate
for part in parts:
if hasattr(part, "text") and part.text:
processed_response["content"] = part.text
break
# Extract function calls
if response.candidates and response.candidates[0].content.parts:
for part in response.candidates[0].content.parts:
# Extract function calls
for part in parts:
if hasattr(part, "function_call") and part.function_call:
fn = part.function_call
processed_response["tool_calls"].append(
@@ -59,8 +62,8 @@ class GeminiLLM(LLMBase):
return processed_response
else:
if response.candidates and response.candidates[0].content.parts:
for part in response.candidates[0].content.parts:
if parts:
for part in parts:
if hasattr(part, "text") and part.text:
return part.text
return ""
+114
View File
@@ -0,0 +1,114 @@
import json
import os
from typing import Dict, List, Optional, Union
from openai import OpenAI
from mem0.configs.llms.base import BaseLlmConfig
from mem0.configs.llms.minimax import MinimaxConfig
from mem0.llms.base import LLMBase
from mem0.memory.utils import extract_json
class MiniMaxLLM(LLMBase):
def __init__(self, config: Optional[Union[BaseLlmConfig, MinimaxConfig, Dict]] = None):
# Convert to MinimaxConfig if needed
if config is None:
config = MinimaxConfig()
elif isinstance(config, dict):
config = MinimaxConfig(**config)
elif isinstance(config, BaseLlmConfig) and not isinstance(config, MinimaxConfig):
# Convert BaseLlmConfig to MinimaxConfig
config = MinimaxConfig(
model=config.model,
temperature=config.temperature,
api_key=config.api_key,
max_tokens=config.max_tokens,
top_p=config.top_p,
top_k=config.top_k,
enable_vision=config.enable_vision,
vision_details=config.vision_details,
http_client_proxies=config.http_client,
)
super().__init__(config)
if not self.config.model:
self.config.model = "MiniMax-M2.7"
api_key = self.config.api_key or os.getenv("MINIMAX_API_KEY")
base_url = (
self.config.minimax_base_url
or os.getenv("MINIMAX_API_BASE")
or "https://api.minimax.io/v1"
)
self.client = OpenAI(api_key=api_key, base_url=base_url)
def _parse_response(self, response, tools):
"""
Process the response based on whether tools are used or not.
Args:
response: The raw response from API.
tools: The list of tools provided in the request.
Returns:
str or dict: The processed response.
"""
if tools:
processed_response = {
"content": response.choices[0].message.content,
"tool_calls": [],
}
if response.choices[0].message.tool_calls:
for tool_call in response.choices[0].message.tool_calls:
processed_response["tool_calls"].append(
{
"name": tool_call.function.name,
"arguments": json.loads(extract_json(tool_call.function.arguments)),
}
)
return processed_response
else:
return response.choices[0].message.content
def generate_response(
self,
messages: List[Dict[str, str]],
response_format=None,
tools: Optional[List[Dict]] = None,
tool_choice: str = "auto",
**kwargs,
):
"""
Generate a response based on the given messages using MiniMax.
Args:
messages (list): List of message dicts containing 'role' and 'content'.
response_format (str or object, optional): Format of the response. Defaults to None.
tools (list, optional): List of tools that the model can call. Defaults to None.
tool_choice (str, optional): Tool choice method. Defaults to "auto".
**kwargs: Additional MiniMax-specific parameters.
Returns:
str: The generated response.
"""
params = self._get_supported_params(messages=messages, **kwargs)
params.update(
{
"model": self.config.model,
"messages": messages,
}
)
if response_format:
params["response_format"] = response_format
if tools:
params["tools"] = tools
params["tool_choice"] = tool_choice
response = self.client.chat.completions.create(**params)
return self._parse_response(response, tools)
+595
View File
@@ -0,0 +1,595 @@
import json
import logging
import time
from mem0.memory.utils import format_entities, sanitize_relationship_for_cypher
try:
import age
except ImportError:
raise ImportError("apache-age-python is not installed. Please install it using pip install apache-age-python")
try:
from rank_bm25 import BM25Okapi
except ImportError:
raise ImportError("rank_bm25 is not installed. Please install it using pip install rank-bm25")
from mem0.graphs.tools import (
DELETE_MEMORY_STRUCT_TOOL_GRAPH,
DELETE_MEMORY_TOOL_GRAPH,
EXTRACT_ENTITIES_STRUCT_TOOL,
EXTRACT_ENTITIES_TOOL,
RELATIONS_STRUCT_TOOL,
RELATIONS_TOOL,
)
from mem0.graphs.utils import EXTRACT_RELATIONS_PROMPT, get_delete_messages
from mem0.utils.factory import EmbedderFactory, LlmFactory
logger = logging.getLogger(__name__)
def _cosine_similarity(vec1, vec2):
"""Compute cosine similarity between two vectors without numpy."""
dot = sum(a * b for a, b in zip(vec1, vec2))
norm1 = sum(a * a for a in vec1) ** 0.5
norm2 = sum(b * b for b in vec2) ** 0.5
if norm1 == 0 or norm2 == 0:
return 0.0
return dot / (norm1 * norm2)
def _get_similar_nodes(nodes, query_embedding, filters, threshold):
"""Find nodes above the similarity threshold from a fetched node list.
Shared by ``_find_similar_node`` and ``_search_graph_db`` to avoid
duplicating the client-side cosine similarity logic.
"""
matches = []
for node in nodes:
props = node if isinstance(node, dict) else {}
stored_emb = props.get("embedding")
if not stored_emb:
continue
if isinstance(stored_emb, str):
stored_emb = json.loads(stored_emb)
if filters.get("agent_id") and props.get("agent_id") != filters["agent_id"]:
continue
if filters.get("run_id") and props.get("run_id") != filters["run_id"]:
continue
sim = _cosine_similarity(query_embedding, stored_emb)
if sim >= threshold:
matches.append({"name": props.get("name"), "similarity": sim, "props": props})
matches.sort(key=lambda x: x["similarity"], reverse=True)
return matches
class MemoryGraph:
def __init__(self, config):
self.config = config
graph_cfg = self.config.graph_store.config
self.graph_name = graph_cfg.graph_name
# Connect using the Apache AGE Python driver (psycopg2-based)
self.ag = age.connect(
graph=self.graph_name,
host=graph_cfg.host,
port=graph_cfg.port,
dbname=graph_cfg.database,
user=graph_cfg.username,
password=graph_cfg.password,
)
self.embedding_model = EmbedderFactory.create(
self.config.embedder.provider, self.config.embedder.config, self.config.vector_store.config
)
# Default to openai if no specific provider is configured
self.llm_provider = "openai"
if self.config.llm and self.config.llm.provider:
self.llm_provider = self.config.llm.provider
if self.config.graph_store and self.config.graph_store.llm and self.config.graph_store.llm.provider:
self.llm_provider = self.config.graph_store.llm.provider
# Get LLM config with proper null checks
llm_config = None
if self.config.graph_store and self.config.graph_store.llm and hasattr(self.config.graph_store.llm, "config"):
llm_config = self.config.graph_store.llm.config
elif hasattr(self.config.llm, "config"):
llm_config = self.config.llm.config
self.llm = LlmFactory.create(self.llm_provider, llm_config)
self.user_id = None
self.threshold = self.config.graph_store.threshold if hasattr(self.config.graph_store, "threshold") else 0.7
# -- helpers ---------------------------------------------------------------
def _exec_cypher(self, cypher_stmt, cols=None, params=None):
"""Execute a Cypher query via the AGE driver and return results.
Uses ``ag.execCypher`` which delegates to psycopg2's safe parameter
substitution (``%s`` placeholders). The *cols* argument specifies the
column names in the ``AS (…)`` clause — when ``None`` the driver
defaults to a single ``v agtype`` column.
When *cols* are provided, returns a list of dicts keyed by column name.
When *cols* is ``None``, returns vertex/edge property dicts or raw values.
"""
cursor = self.ag.execCypher(cypher_stmt, cols=cols, params=params)
rows = cursor.fetchall()
if not rows:
return []
col_names = [desc[0] for desc in cursor.description] if cursor.description else None
results = []
for row in rows:
if col_names and len(col_names) > 1:
record = {}
for i, col_name in enumerate(col_names):
val = row[i]
if hasattr(val, "properties"):
record[col_name] = val.properties
else:
record[col_name] = val
results.append(record)
else:
val = row[0] if len(row) == 1 else row
if hasattr(val, "properties"):
results.append(val.properties)
else:
results.append(val)
return results
def _fetch_user_nodes_with_embeddings(self, user_id):
"""Fetch all nodes with embeddings for a given user_id."""
return self._exec_cypher(
"MATCH (n {user_id: %s}) WHERE n.embedding IS NOT NULL RETURN n",
params=(user_id,),
)
def _find_similar_node(self, embedding, filters, threshold=0.9):
"""Find the most similar existing node by cosine similarity.
Apache AGE does not have a built-in vector index, so we fetch all
node embeddings matching the filters and compute cosine similarity
on the client side. This is adequate for moderate graph sizes; for
very large graphs consider pairing AGE with pgvector.
"""
nodes = self._fetch_user_nodes_with_embeddings(filters["user_id"])
matches = _get_similar_nodes(nodes, embedding, filters, threshold)
return matches[0]["props"] if matches else None
def _merge_node(self, user_id, name, embedding, agent_id=None, run_id=None):
"""Create a node if it doesn't exist, or update mentions if it does.
Apache AGE does not support ``ON CREATE SET`` / ``ON MATCH SET``, so
we use ``MERGE … SET`` which always applies the SET clause. Embeddings
and optional filter properties are set in a single query.
"""
set_parts = [
"n.embedding = %s",
"n.mentions = coalesce(n.mentions, 0) + 1",
"n.created = coalesce(n.created, %s)",
]
params = [user_id, name, json.dumps(embedding), int(time.time() * 1000)]
if agent_id:
set_parts.append("n.agent_id = %s")
params.append(agent_id)
if run_id:
set_parts.append("n.run_id = %s")
params.append(run_id)
set_clause = ", ".join(set_parts)
self._exec_cypher(
f"MERGE (n {{user_id: %s, name: %s}}) SET {set_clause}",
params=tuple(params),
)
def close(self):
"""Close the underlying database connection."""
if self.ag:
self.ag.close()
# -- public API ------------------------------------------------------------
def add(self, data, filters):
"""
Adds data to the graph.
Args:
data (str): The data to add to the graph.
filters (dict): A dictionary containing filters to be applied during the addition.
"""
entity_type_map = self._retrieve_nodes_from_data(data, filters)
to_be_added = self._establish_nodes_relations_from_data(data, filters, entity_type_map)
search_output = self._search_graph_db(node_list=list(entity_type_map.keys()), filters=filters)
to_be_deleted = self._get_delete_entities_from_search_output(search_output, data, filters)
deleted_entities = self._delete_entities(to_be_deleted, filters)
added_entities = self._add_entities(to_be_added, filters, entity_type_map)
return {"deleted_entities": deleted_entities, "added_entities": added_entities}
def search(self, query, filters, limit=100):
"""
Search for memories and related graph data.
Args:
query (str): Query to search for.
filters (dict): A dictionary containing filters to be applied during the search.
limit (int): The maximum number of nodes and relationships to retrieve. Defaults to 100.
Returns:
list: A list of dicts with keys "source", "relationship", "destination".
"""
entity_type_map = self._retrieve_nodes_from_data(query, filters)
search_output = self._search_graph_db(node_list=list(entity_type_map.keys()), filters=filters)
if not search_output:
return []
search_outputs_sequence = [
[item["source"], item["relationship"], item["destination"]] for item in search_output
]
bm25 = BM25Okapi(search_outputs_sequence)
tokenized_query = query.split(" ")
reranked_results = bm25.get_top_n(tokenized_query, search_outputs_sequence, n=5)
search_results = []
for item in reranked_results:
search_results.append({"source": item[0], "relationship": item[1], "destination": item[2]})
logger.info(f"Returned {len(search_results)} search results")
return search_results
def delete(self, data, filters):
"""
Delete graph entities associated with the given memory text.
Extracts entities and relationships from the memory text using the same
pipeline as add(), then deletes the matching relationships in the graph.
Args:
data (str): The memory text whose graph entities should be removed.
filters (dict): Scope filters (user_id, agent_id, run_id).
"""
try:
entity_type_map = self._retrieve_nodes_from_data(data, filters)
if not entity_type_map:
logger.debug("No entities found in memory text, skipping graph cleanup")
return
to_be_deleted = self._establish_nodes_relations_from_data(data, filters, entity_type_map)
if to_be_deleted:
self._delete_entities(to_be_deleted, filters)
except Exception as e:
logger.error(f"Error during graph cleanup for memory delete: {e}")
def delete_all(self, filters):
"""Delete all nodes and relationships for a user or specific agent."""
where_parts = ["n.user_id = %s"]
params = [filters["user_id"]]
if filters.get("agent_id"):
where_parts.append("n.agent_id = %s")
params.append(filters["agent_id"])
if filters.get("run_id"):
where_parts.append("n.run_id = %s")
params.append(filters["run_id"])
where_clause = " AND ".join(where_parts)
self._exec_cypher(
f"MATCH (n) WHERE {where_clause} DETACH DELETE n",
params=tuple(params),
)
self.ag.commit()
def get_all(self, filters, limit=100):
"""
Retrieves all nodes and relationships from the graph database based on optional filtering criteria.
Args:
filters (dict): A dictionary containing filters to be applied during the retrieval.
limit (int): The maximum number of nodes and relationships to retrieve. Defaults to 100.
Returns:
list: A list of dictionaries, each containing:
- 'source': The source node name.
- 'relationship': The relationship type.
- 'target': The target node name.
"""
where_parts = ["n.user_id = %s", "m.user_id = %s"]
params = [filters["user_id"], filters["user_id"]]
if filters.get("agent_id"):
where_parts.extend(["n.agent_id = %s", "m.agent_id = %s"])
params.extend([filters["agent_id"], filters["agent_id"]])
if filters.get("run_id"):
where_parts.extend(["n.run_id = %s", "m.run_id = %s"])
params.extend([filters["run_id"], filters["run_id"]])
where_clause = " AND ".join(where_parts)
params.append(limit)
results = self._exec_cypher(
f"MATCH (n)-[r]->(m) WHERE {where_clause} "
f"RETURN n.name, type(r), m.name LIMIT %s",
cols=["source", "relationship", "target"],
params=tuple(params),
)
final_results = []
for result in results:
final_results.append(
{
"source": result["source"],
"relationship": result["relationship"],
"target": result["target"],
}
)
logger.info(f"Retrieved {len(final_results)} relationships")
return final_results
# -- LLM-driven extraction -------------------------------------------------
def _retrieve_nodes_from_data(self, data, filters):
"""Extracts all the entities mentioned in the query."""
_tools = [EXTRACT_ENTITIES_TOOL]
if self.llm_provider in ["azure_openai_structured", "openai_structured"]:
_tools = [EXTRACT_ENTITIES_STRUCT_TOOL]
search_results = self.llm.generate_response(
messages=[
{
"role": "system",
"content": f"You are a smart assistant who understands entities and their types in a given text. If user message contains self reference such as 'I', 'me', 'my' etc. then use {filters['user_id']} as the source entity. Extract all the entities from the text. ***DO NOT*** answer the question itself if the given text is a question.",
},
{"role": "user", "content": data},
],
tools=_tools,
)
entity_type_map = {}
try:
for tool_call in search_results["tool_calls"]:
if tool_call["name"] != "extract_entities":
continue
for item in tool_call.get("arguments", {}).get("entities", []):
if "entity" in item and "entity_type" in item:
entity_type_map[item["entity"]] = item["entity_type"]
except Exception as e:
logger.exception(
f"Error in search tool: {e}, llm_provider={self.llm_provider}, search_results={search_results}"
)
entity_type_map = {k.lower().replace(" ", "_"): v.lower().replace(" ", "_") for k, v in entity_type_map.items()}
logger.debug(f"Entity type map: {entity_type_map}\n search_results={search_results}")
return entity_type_map
def _establish_nodes_relations_from_data(self, data, filters, entity_type_map):
"""Establish relations among the extracted nodes."""
user_identity = f"user_id: {filters['user_id']}"
if filters.get("agent_id"):
user_identity += f", agent_id: {filters['agent_id']}"
if filters.get("run_id"):
user_identity += f", run_id: {filters['run_id']}"
if self.config.graph_store.custom_prompt:
system_content = EXTRACT_RELATIONS_PROMPT.replace("USER_ID", user_identity)
system_content = system_content.replace("CUSTOM_PROMPT", f"4. {self.config.graph_store.custom_prompt}")
messages = [
{"role": "system", "content": system_content},
{"role": "user", "content": data},
]
else:
system_content = EXTRACT_RELATIONS_PROMPT.replace("USER_ID", user_identity)
messages = [
{"role": "system", "content": system_content},
{"role": "user", "content": f"List of entities: {list(entity_type_map.keys())}. \n\nText: {data}"},
]
_tools = [RELATIONS_TOOL]
if self.llm_provider in ["azure_openai_structured", "openai_structured"]:
_tools = [RELATIONS_STRUCT_TOOL]
extracted_entities = self.llm.generate_response(
messages=messages,
tools=_tools,
)
entities = []
if extracted_entities and extracted_entities.get("tool_calls"):
entities = extracted_entities["tool_calls"][0].get("arguments", {}).get("entities", [])
entities = self._remove_spaces_from_entities(entities)
logger.debug(f"Extracted entities: {entities}")
return entities
# -- graph DB operations ---------------------------------------------------
def _search_graph_db(self, node_list, filters, limit=100):
"""Search similar nodes and their respective incoming and outgoing relations."""
result_relations = []
for node in node_list:
n_embedding = self.embedding_model.embed(node)
nodes = self._fetch_user_nodes_with_embeddings(filters["user_id"])
similar_nodes = _get_similar_nodes(nodes, n_embedding, filters, self.threshold)
# Build WHERE clause for relationship target filtering
rel_where_parts = ["m.user_id = %s"]
rel_params_suffix = [filters["user_id"]]
if filters.get("agent_id"):
rel_where_parts.append("m.agent_id = %s")
rel_params_suffix.append(filters["agent_id"])
if filters.get("run_id"):
rel_where_parts.append("m.run_id = %s")
rel_params_suffix.append(filters["run_id"])
rel_where = " AND ".join(rel_where_parts)
# For each similar node, fetch its relationships
for sn in similar_nodes[:limit]:
node_name = sn["name"]
similarity = sn["similarity"]
out_params = (filters["user_id"], node_name) + tuple(rel_params_suffix)
out_results = self._exec_cypher(
f"MATCH (n {{user_id: %s, name: %s}})-[r]->(m) "
f"WHERE {rel_where} "
f"RETURN n.name, type(r), m.name",
cols=["source", "relationship", "destination"],
params=out_params,
)
in_params = (filters["user_id"], node_name) + tuple(rel_params_suffix)
in_results = self._exec_cypher(
f"MATCH (n {{user_id: %s, name: %s}})<-[r]-(m) "
f"WHERE {rel_where} "
f"RETURN m.name, type(r), n.name",
cols=["source", "relationship", "destination"],
params=in_params,
)
for rel in out_results + in_results:
rel["similarity"] = similarity
result_relations.append(rel)
return result_relations
def _get_delete_entities_from_search_output(self, search_output, data, filters):
"""Get the entities to be deleted from the search output."""
search_output_string = format_entities(search_output)
user_identity = f"user_id: {filters['user_id']}"
if filters.get("agent_id"):
user_identity += f", agent_id: {filters['agent_id']}"
if filters.get("run_id"):
user_identity += f", run_id: {filters['run_id']}"
system_prompt, user_prompt = get_delete_messages(search_output_string, data, user_identity)
_tools = [DELETE_MEMORY_TOOL_GRAPH]
if self.llm_provider in ["azure_openai_structured", "openai_structured"]:
_tools = [DELETE_MEMORY_STRUCT_TOOL_GRAPH]
memory_updates = self.llm.generate_response(
messages=[
{"role": "system", "content": system_prompt},
{"role": "user", "content": user_prompt},
],
tools=_tools,
)
to_be_deleted = []
for item in memory_updates.get("tool_calls", []):
if item.get("name") == "delete_graph_memory":
to_be_deleted.append(item.get("arguments"))
to_be_deleted = self._remove_spaces_from_entities(to_be_deleted)
logger.debug(f"Deleted relationships: {to_be_deleted}")
return to_be_deleted
def _delete_entities(self, to_be_deleted, filters):
"""Delete the entities from the graph."""
user_id = filters["user_id"]
agent_id = filters.get("agent_id")
run_id = filters.get("run_id")
results = []
try:
for item in to_be_deleted:
source = item["source"]
destination = item["destination"]
relationship = item["relationship"]
where_parts = [
"n.user_id = %s", "n.name = %s",
"m.user_id = %s", "m.name = %s",
]
params = [user_id, source, user_id, destination]
if agent_id:
where_parts.extend(["n.agent_id = %s", "m.agent_id = %s"])
params.extend([agent_id, agent_id])
if run_id:
where_parts.extend(["n.run_id = %s", "m.run_id = %s"])
params.extend([run_id, run_id])
where_clause = " AND ".join(where_parts)
result = self._exec_cypher(
f"MATCH (n)-[r:{relationship}]->(m) "
f"WHERE {where_clause} "
f"DELETE r "
f"RETURN n.name, type(r), m.name",
cols=["source", "relationship", "target"],
params=tuple(params),
)
results.append(result)
self.ag.commit()
except Exception:
self.ag.rollback()
raise
return results
def _add_entities(self, to_be_added, filters, entity_type_map):
"""Add new entities to the graph. Merge nodes if they already exist.
Apache AGE does not support ``ON CREATE SET`` / ``ON MATCH SET``, so we
use ``MERGE … SET`` which always applies. The ``coalesce`` pattern
ensures ``created`` is only set on the first merge.
"""
user_id = filters["user_id"]
agent_id = filters.get("agent_id")
run_id = filters.get("run_id")
results = []
try:
for item in to_be_added:
source = item["source"]
destination = item["destination"]
relationship = item["relationship"]
source_embedding = self.embedding_model.embed(source)
dest_embedding = self.embedding_model.embed(destination)
source_match = self._find_similar_node(source_embedding, filters, threshold=self.threshold)
dest_match = self._find_similar_node(dest_embedding, filters, threshold=self.threshold)
effective_source = source_match["name"] if source_match else source
effective_dest = dest_match["name"] if dest_match else destination
# Merge source and destination nodes
self._merge_node(user_id, effective_source, source_embedding, agent_id, run_id)
self._merge_node(user_id, effective_dest, dest_embedding, agent_id, run_id)
# Merge relationship
result = self._exec_cypher(
f"MATCH (s {{user_id: %s, name: %s}}), (d {{user_id: %s, name: %s}}) "
f"MERGE (s)-[r:{relationship}]->(d) "
f"RETURN s.name, type(r), d.name",
cols=["source", "relationship", "target"],
params=(user_id, effective_source, user_id, effective_dest),
)
results.append(result)
self.ag.commit()
except Exception:
self.ag.rollback()
raise
return results
def _remove_spaces_from_entities(self, entity_list):
for item in entity_list:
item["source"] = item["source"].lower().replace(" ", "_")
item["relationship"] = sanitize_relationship_for_cypher(item["relationship"].lower().replace(" ", "_"))
item["destination"] = item["destination"].lower().replace(" ", "_")
return entity_list
def reset(self):
"""Reset the graph by clearing all nodes and relationships."""
logger.warning("Clearing graph...")
self._exec_cypher("MATCH (n) DETACH DELETE n")
self.ag.commit()
+68 -18
View File
@@ -129,6 +129,28 @@ class MemoryGraph:
return search_results
def delete(self, data, filters):
"""
Delete graph entities associated with the given memory text.
Extracts entities and relationships from the memory text using the same
pipeline as add(), then soft-deletes the matching relationships in the graph.
Args:
data (str): The memory text whose graph entities should be removed.
filters (dict): Scope filters (user_id, agent_id, run_id).
"""
try:
entity_type_map = self._retrieve_nodes_from_data(data, filters)
if not entity_type_map:
logger.debug("No entities found in memory text, skipping graph cleanup")
return
to_be_deleted = self._establish_nodes_relations_from_data(data, filters, entity_type_map)
if to_be_deleted:
self._delete_entities(to_be_deleted, filters)
except Exception as e:
logger.error(f"Error during graph cleanup for memory delete: {e}")
def delete_all(self, filters):
# Build node properties for filtering
node_props = ["user_id: $user_id"]
@@ -174,6 +196,7 @@ class MemoryGraph:
query = f"""
MATCH (n {self.node_label} {{{node_props_str}}})-[r]->(m {self.node_label} {{{node_props_str}}})
WHERE r.valid IS NULL OR r.valid = true
RETURN n.name AS source, type(r) AS relationship, m.name AS target
LIMIT $limit
"""
@@ -291,10 +314,12 @@ class MemoryGraph:
CALL {{
WITH n
MATCH (n)-[r]->(m {self.node_label} {{{node_props_str}}})
WHERE r.valid IS NULL OR r.valid = true
RETURN n.name AS source, elementId(n) AS source_id, type(r) AS relationship, elementId(r) AS relation_id, m.name AS destination, elementId(m) AS destination_id
UNION
WITH n
MATCH (n)<-[r]-(m {self.node_label} {{{node_props_str}}})
WHERE r.valid IS NULL OR r.valid = true
RETURN m.name AS source, elementId(m) AS source_id, type(r) AS relationship, elementId(r) AS relation_id, n.name AS destination, elementId(n) AS destination_id
}}
WITH distinct source, source_id, relationship, relation_id, destination, destination_id, similarity
@@ -392,13 +417,15 @@ class MemoryGraph:
source_props_str = ", ".join(source_props)
dest_props_str = ", ".join(dest_props)
# Delete the specific relationship between nodes
# Soft-delete: mark relationship as invalid instead of removing it,
# enabling temporal reasoning over historical graph state.
# See: https://github.com/mem0ai/mem0/issues/4187
cypher = f"""
MATCH (n {self.node_label} {{{source_props_str}}})
-[r:{relationship}]->
(m {self.node_label} {{{dest_props_str}}})
DELETE r
WHERE r.valid IS NULL OR r.valid = true
SET r.valid = false, r.invalidated_at = datetime()
RETURN
n.name AS source,
m.name AS target,
@@ -464,11 +491,16 @@ class MemoryGraph:
CALL db.create.setNodeVectorProperty(destination, 'embedding', $destination_embedding)
WITH source, destination
MERGE (source)-[r:{relationship}]->(destination)
ON CREATE SET
r.created = timestamp(),
r.mentions = 1
ON CREATE SET
r.created_at = timestamp(),
r.updated_at = timestamp(),
r.mentions = 1,
r.valid = true
ON MATCH SET
r.mentions = coalesce(r.mentions, 0) + 1
r.mentions = coalesce(r.mentions, 0) + 1,
r.valid = true,
r.updated_at = timestamp(),
r.invalidated_at = null
RETURN source.name AS source, type(r) AS relationship, destination.name AS target
"""
@@ -508,11 +540,16 @@ class MemoryGraph:
CALL db.create.setNodeVectorProperty(source, 'embedding', $source_embedding)
WITH source, destination
MERGE (source)-[r:{relationship}]->(destination)
ON CREATE SET
r.created = timestamp(),
r.mentions = 1
ON CREATE SET
r.created_at = timestamp(),
r.updated_at = timestamp(),
r.mentions = 1,
r.valid = true
ON MATCH SET
r.mentions = coalesce(r.mentions, 0) + 1
r.mentions = coalesce(r.mentions, 0) + 1,
r.valid = true,
r.updated_at = timestamp(),
r.invalidated_at = null
RETURN source.name AS source, type(r) AS relationship, destination.name AS target
"""
@@ -537,11 +574,16 @@ class MemoryGraph:
WHERE elementId(destination) = $destination_id
SET destination.mentions = coalesce(destination.mentions, 0) + 1
MERGE (source)-[r:{relationship}]->(destination)
ON CREATE SET
ON CREATE SET
r.created_at = timestamp(),
r.updated_at = timestamp(),
r.mentions = 1
ON MATCH SET r.mentions = coalesce(r.mentions, 0) + 1
r.mentions = 1,
r.valid = true
ON MATCH SET
r.mentions = coalesce(r.mentions, 0) + 1,
r.valid = true,
r.updated_at = timestamp(),
r.invalidated_at = null
RETURN source.name AS source, type(r) AS relationship, destination.name AS target
"""
@@ -585,10 +627,18 @@ class MemoryGraph:
WITH source, destination
CALL db.create.setNodeVectorProperty(destination, 'embedding', $dest_embedding)
WITH source, destination
MERGE (source)-[rel:{relationship}]->(destination)
ON CREATE SET rel.created = timestamp(), rel.mentions = 1
ON MATCH SET rel.mentions = coalesce(rel.mentions, 0) + 1
RETURN source.name AS source, type(rel) AS relationship, destination.name AS target
MERGE (source)-[r:{relationship}]->(destination)
ON CREATE SET
r.created_at = timestamp(),
r.updated_at = timestamp(),
r.mentions = 1,
r.valid = true
ON MATCH SET
r.mentions = coalesce(r.mentions, 0) + 1,
r.valid = true,
r.updated_at = timestamp(),
r.invalidated_at = null
RETURN source.name AS source, type(r) AS relationship, destination.name AS target
"""
params = {
+22
View File
@@ -149,6 +149,28 @@ class MemoryGraph:
return search_results
def delete(self, data, filters):
"""
Delete graph entities associated with the given memory text.
Extracts entities and relationships from the memory text using the same
pipeline as add(), then deletes the matching relationships in the graph.
Args:
data (str): The memory text whose graph entities should be removed.
filters (dict): Scope filters (user_id, agent_id, run_id).
"""
try:
entity_type_map = self._retrieve_nodes_from_data(data, filters)
if not entity_type_map:
logger.debug("No entities found in memory text, skipping graph cleanup")
return
to_be_deleted = self._establish_nodes_relations_from_data(data, filters, entity_type_map)
if to_be_deleted:
self._delete_entities(to_be_deleted, filters)
except Exception as e:
logger.error(f"Error during graph cleanup for memory delete: {e}")
def delete_all(self, filters):
# Build node properties for filtering
node_props = ["user_id: $user_id"]
+86 -26
View File
@@ -644,6 +644,9 @@ class Memory(MemoryBase):
if memory_id and (metadata.get("agent_id") or metadata.get("run_id")):
# Update only the session identifiers, keep content the same
existing_memory = self.vector_store.get(vector_id=memory_id)
if existing_memory is None:
logger.warning(f"Memory {memory_id} not found for session ID update, skipping")
continue
updated_metadata = deepcopy(existing_memory.payload)
if metadata.get("agent_id"):
updated_metadata["agent_id"] = metadata["agent_id"]
@@ -1098,7 +1101,27 @@ class Memory(MemoryBase):
memory_id (str): ID of the memory to delete.
"""
capture_event("mem0.delete", self, {"memory_id": memory_id, "sync_type": "sync"})
self._delete_memory(memory_id)
existing_memory = self.vector_store.get(vector_id=memory_id)
if existing_memory is None:
raise ValueError(f"Memory with id {memory_id} not found")
# Clean up graph entities before deleting from vector store
if self.enable_graph:
try:
memory_text = existing_memory.payload.get("data", "")
if memory_text:
filters = {}
for key in ("user_id", "agent_id", "run_id"):
val = existing_memory.payload.get(key)
if val:
filters[key] = val
if filters.get("user_id"):
self.graph.delete(memory_text, filters)
except Exception as e:
logger.error(f"Error cleaning up graph for memory {memory_id}: {e}")
self._delete_memory(memory_id, existing_memory)
return {"message": "Memory deleted successfully!"}
def delete_all(self, user_id: Optional[str] = None, agent_id: Optional[str] = None, run_id: Optional[str] = None):
@@ -1157,24 +1180,24 @@ class Memory(MemoryBase):
else:
embeddings = self.embedding_model.embed(data, memory_action="add")
memory_id = str(uuid.uuid4())
metadata = metadata or {}
metadata["data"] = data
metadata["hash"] = hashlib.md5(data.encode()).hexdigest()
metadata["created_at"] = datetime.now(timezone.utc).isoformat()
new_metadata = deepcopy(metadata) if metadata is not None else {}
new_metadata["data"] = data
new_metadata["hash"] = hashlib.md5(data.encode()).hexdigest()
new_metadata["created_at"] = datetime.now(timezone.utc).isoformat()
self.vector_store.insert(
vectors=[embeddings],
ids=[memory_id],
payloads=[metadata],
payloads=[new_metadata],
)
self.db.add_history(
memory_id,
None,
data,
"ADD",
created_at=metadata.get("created_at"),
actor_id=metadata.get("actor_id"),
role=metadata.get("role"),
created_at=new_metadata.get("created_at"),
actor_id=new_metadata.get("actor_id"),
role=new_metadata.get("role"),
)
return memory_id
@@ -1208,9 +1231,10 @@ class Memory(MemoryBase):
if metadata is None:
raise ValueError("Metadata cannot be done for procedural memory.")
metadata["memory_type"] = MemoryType.PROCEDURAL.value
new_metadata = deepcopy(metadata)
new_metadata["memory_type"] = MemoryType.PROCEDURAL.value
embeddings = self.embedding_model.embed(procedural_memory, memory_action="add")
memory_id = self._create_memory(procedural_memory, {procedural_memory: embeddings}, metadata=metadata)
memory_id = self._create_memory(procedural_memory, {procedural_memory: embeddings}, metadata=new_metadata)
capture_event("mem0._create_procedural_memory", self, {"memory_id": memory_id, "sync_type": "sync"})
result = {"results": [{"id": memory_id, "memory": procedural_memory, "event": "ADD"}]}
@@ -1226,6 +1250,9 @@ class Memory(MemoryBase):
logger.error(f"Error getting memory with ID {memory_id} during update.")
raise ValueError(f"Error getting memory with ID {memory_id}. Please provide a valid 'memory_id'")
if existing_memory is None:
raise ValueError(f"Memory with id {memory_id} not found. Please provide a valid 'memory_id'")
prev_value = existing_memory.payload.get("data")
new_metadata = deepcopy(metadata) if metadata is not None else {}
@@ -1271,9 +1298,12 @@ class Memory(MemoryBase):
)
return memory_id
def _delete_memory(self, memory_id):
def _delete_memory(self, memory_id, existing_memory=None):
logger.info(f"Deleting memory with {memory_id=}")
existing_memory = self.vector_store.get(vector_id=memory_id)
if existing_memory is None:
existing_memory = self.vector_store.get(vector_id=memory_id)
if existing_memory is None:
raise ValueError(f"Memory with id {memory_id} not found")
prev_value = existing_memory.payload.get("data", "")
self.vector_store.delete(vector_id=memory_id)
self.db.add_history(
@@ -1671,6 +1701,9 @@ class AsyncMemory(MemoryBase):
# Create async task to update only the session identifiers
async def update_session_ids(mem_id, meta):
existing_memory = await asyncio.to_thread(self.vector_store.get, vector_id=mem_id)
if existing_memory is None:
logger.warning(f"Memory {mem_id} not found for session ID update, skipping")
return
updated_metadata = deepcopy(existing_memory.payload)
if meta.get("agent_id"):
updated_metadata["agent_id"] = meta["agent_id"]
@@ -2163,7 +2196,27 @@ class AsyncMemory(MemoryBase):
memory_id (str): ID of the memory to delete.
"""
capture_event("mem0.delete", self, {"memory_id": memory_id, "sync_type": "async"})
await self._delete_memory(memory_id)
existing_memory = await asyncio.to_thread(self.vector_store.get, vector_id=memory_id)
if existing_memory is None:
raise ValueError(f"Memory with id {memory_id} not found")
# Clean up graph entities before deleting from vector store
if self.enable_graph:
try:
memory_text = existing_memory.payload.get("data", "")
if memory_text:
filters = {}
for key in ("user_id", "agent_id", "run_id"):
val = existing_memory.payload.get(key)
if val:
filters[key] = val
if filters.get("user_id"):
await asyncio.to_thread(self.graph.delete, memory_text, filters)
except Exception as e:
logger.error(f"Error cleaning up graph for memory {memory_id}: {e}")
await self._delete_memory(memory_id, existing_memory)
return {"message": "Memory deleted successfully!"}
async def delete_all(self, user_id=None, agent_id=None, run_id=None):
@@ -2226,16 +2279,16 @@ class AsyncMemory(MemoryBase):
embeddings = await asyncio.to_thread(self.embedding_model.embed, data, memory_action="add")
memory_id = str(uuid.uuid4())
metadata = metadata or {}
metadata["data"] = data
metadata["hash"] = hashlib.md5(data.encode()).hexdigest()
metadata["created_at"] = datetime.now(timezone.utc).isoformat()
new_metadata = deepcopy(metadata) if metadata is not None else {}
new_metadata["data"] = data
new_metadata["hash"] = hashlib.md5(data.encode()).hexdigest()
new_metadata["created_at"] = datetime.now(timezone.utc).isoformat()
await asyncio.to_thread(
self.vector_store.insert,
vectors=[embeddings],
ids=[memory_id],
payloads=[metadata],
payloads=[new_metadata],
)
await asyncio.to_thread(
@@ -2244,9 +2297,9 @@ class AsyncMemory(MemoryBase):
None,
data,
"ADD",
created_at=metadata.get("created_at"),
actor_id=metadata.get("actor_id"),
role=metadata.get("role"),
created_at=new_metadata.get("created_at"),
actor_id=new_metadata.get("actor_id"),
role=new_metadata.get("role"),
)
return memory_id
@@ -2295,9 +2348,10 @@ class AsyncMemory(MemoryBase):
if metadata is None:
raise ValueError("Metadata cannot be done for procedural memory.")
metadata["memory_type"] = MemoryType.PROCEDURAL.value
new_metadata = deepcopy(metadata)
new_metadata["memory_type"] = MemoryType.PROCEDURAL.value
embeddings = await asyncio.to_thread(self.embedding_model.embed, procedural_memory, memory_action="add")
memory_id = await self._create_memory(procedural_memory, {procedural_memory: embeddings}, metadata=metadata)
memory_id = await self._create_memory(procedural_memory, {procedural_memory: embeddings}, metadata=new_metadata)
capture_event("mem0._create_procedural_memory", self, {"memory_id": memory_id, "sync_type": "async"})
result = {"results": [{"id": memory_id, "memory": procedural_memory, "event": "ADD"}]}
@@ -2313,6 +2367,9 @@ class AsyncMemory(MemoryBase):
logger.error(f"Error getting memory with ID {memory_id} during update.")
raise ValueError(f"Error getting memory with ID {memory_id}. Please provide a valid 'memory_id'")
if existing_memory is None:
raise ValueError(f"Memory with id {memory_id} not found. Please provide a valid 'memory_id'")
prev_value = existing_memory.payload.get("data")
new_metadata = deepcopy(metadata) if metadata is not None else {}
@@ -2361,9 +2418,12 @@ class AsyncMemory(MemoryBase):
)
return memory_id
async def _delete_memory(self, memory_id):
async def _delete_memory(self, memory_id, existing_memory=None):
logger.info(f"Deleting memory with {memory_id=}")
existing_memory = await asyncio.to_thread(self.vector_store.get, vector_id=memory_id)
if existing_memory is None:
existing_memory = await asyncio.to_thread(self.vector_store.get, vector_id=memory_id)
if existing_memory is None:
raise ValueError(f"Memory with id {memory_id} not found")
prev_value = existing_memory.payload.get("data", "")
await asyncio.to_thread(self.vector_store.delete, vector_id=memory_id)
+22
View File
@@ -134,6 +134,28 @@ class MemoryGraph:
return search_results
def delete(self, data, filters):
"""
Delete graph entities associated with the given memory text.
Extracts entities and relationships from the memory text using the same
pipeline as add(), then deletes the matching relationships in the graph.
Args:
data (str): The memory text whose graph entities should be removed.
filters (dict): Scope filters (user_id, agent_id).
"""
try:
entity_type_map = self._retrieve_nodes_from_data(data, filters)
if not entity_type_map:
logger.debug("No entities found in memory text, skipping graph cleanup")
return
to_be_deleted = self._establish_nodes_relations_from_data(data, filters, entity_type_map)
if to_be_deleted:
self._delete_entities(to_be_deleted, filters)
except Exception as e:
logger.error(f"Error during graph cleanup for memory delete: {e}")
def delete_all(self, filters):
"""Delete all nodes and relationships for a user or specific agent."""
if filters.get("agent_id"):
+4
View File
@@ -6,6 +6,7 @@ from mem0.configs.llms.anthropic import AnthropicConfig
from mem0.configs.llms.azure import AzureOpenAIConfig
from mem0.configs.llms.base import BaseLlmConfig
from mem0.configs.llms.deepseek import DeepSeekConfig
from mem0.configs.llms.minimax import MinimaxConfig
from mem0.configs.llms.lmstudio import LMStudioConfig
from mem0.configs.llms.ollama import OllamaConfig
from mem0.configs.llms.openai import OpenAIConfig
@@ -45,6 +46,7 @@ class LlmFactory:
"azure_openai_structured": ("mem0.llms.azure_openai_structured.AzureOpenAIStructuredLLM", AzureOpenAIConfig),
"gemini": ("mem0.llms.gemini.GeminiLLM", BaseLlmConfig),
"deepseek": ("mem0.llms.deepseek.DeepSeekLLM", DeepSeekConfig),
"minimax": ("mem0.llms.minimax.MiniMaxLLM", MinimaxConfig),
"xai": ("mem0.llms.xai.XAILLM", BaseLlmConfig),
"sarvam": ("mem0.llms.sarvam.SarvamLLM", BaseLlmConfig),
"lmstudio": ("mem0.llms.lmstudio.LMStudioLLM", LMStudioConfig),
@@ -186,6 +188,7 @@ class VectorStoreFactory:
"baidu": "mem0.vector_stores.baidu.BaiduDB",
"cassandra": "mem0.vector_stores.cassandra.CassandraDB",
"neptune": "mem0.vector_stores.neptune_analytics.NeptuneAnalyticsVector",
"turbopuffer": "mem0.vector_stores.turbopuffer.TurbopufferDB",
}
@classmethod
@@ -216,6 +219,7 @@ class GraphStoreFactory:
"neptune": "mem0.graphs.neptune.neptunegraph.MemoryGraph",
"neptunedb": "mem0.graphs.neptune.neptunedb.MemoryGraph",
"kuzu": "mem0.memory.kuzu_memory.MemoryGraph",
"apache_age": "mem0.memory.apache_age_memory.MemoryGraph",
"default": "mem0.memory.graph_memory.MemoryGraph",
}
+1
View File
@@ -34,6 +34,7 @@ class VectorStoreConfig(BaseModel):
"faiss": "FAISSConfig",
"langchain": "LangchainConfig",
"s3_vectors": "S3VectorsConfig",
"turbopuffer": "TurbopufferConfig",
}
@model_validator(mode="after")
+60 -42
View File
@@ -65,7 +65,7 @@ class Databricks(VectorStoreBase):
catalog (str): Unity Catalog catalog name.
schema (str): Unity Catalog schema name.
table_name (str): Source Delta table name.
index_name (str, optional): Vector search index name (default: "mem0").
collection_name (str, optional): Vector search index name (default: "mem0").
index_type (str, optional): Index type, either "DELTA_SYNC" or "DIRECT_ACCESS" (default: "DELTA_SYNC").
embedding_model_endpoint_name (str, optional): Embedding model endpoint for Databricks-computed embeddings.
embedding_dimension (int, optional): Vector embedding dimensions (default: 1536).
@@ -85,7 +85,7 @@ class Databricks(VectorStoreBase):
self.fully_qualified_index_name = f"{self.catalog}.{self.schema}.{self.index_name}"
# Configuration
self.index_type = index_type
self.index_type = VectorIndexType(index_type) if isinstance(index_type, str) else index_type
self.embedding_model_endpoint_name = embedding_model_endpoint_name
self.embedding_dimension = embedding_dimension
self.endpoint_type = endpoint_type
@@ -261,11 +261,11 @@ class Databricks(VectorStoreBase):
)
logger.info(f"Successfully created source table '{self.fully_qualified_table_name}'")
self.client.table_constraints.create(
full_name_arg="logistics_dev.ai.dev_memory",
full_name_arg=self.fully_qualified_table_name,
constraint=TableConstraint(
primary_key_constraint=PrimaryKeyConstraint(
name="pk_dev_memory", # Name of the primary key constraint
child_columns=["memory_id"], # Columns that make up the primary key
name=f"pk_{self.table_name}",
child_columns=["memory_id"],
)
),
)
@@ -439,29 +439,29 @@ class Databricks(VectorStoreBase):
try:
filters_json = json.dumps(filters) if filters else None
# Choose query type
if self.index_type == VectorIndexType.DELTA_SYNC and query:
# Text-based search
sdk_results = self.client.vector_search_indexes.query_index(
index_name=self.fully_qualified_index_name,
columns=self.column_names,
query_text=query,
num_results=limit,
query_type=self.query_type,
filters_json=filters_json,
)
elif self.index_type == VectorIndexType.DIRECT_ACCESS and vectors:
# Vector-based search
sdk_results = self.client.vector_search_indexes.query_index(
index_name=self.fully_qualified_index_name,
columns=self.column_names,
query_vector=vectors,
num_results=limit,
query_type=self.query_type,
filters_json=filters_json,
)
# Choose query mode per Databricks SDK contract:
# - query_text: for Delta Sync Index with model endpoint
# - query_vector: for Direct Access Index and Delta Sync Index with self-managed vectors
query_kwargs = {
"index_name": self.fully_qualified_index_name,
"columns": self.column_names,
"num_results": limit,
"query_type": self.query_type,
"filters_json": filters_json,
}
uses_model_endpoint = (
self.index_type == VectorIndexType.DELTA_SYNC and self.embedding_model_endpoint_name
)
if uses_model_endpoint:
if not query:
raise ValueError("Query text is required for Delta Sync Index with model endpoint.")
query_kwargs["query_text"] = query
elif vectors:
query_kwargs["query_vector"] = vectors
else:
raise ValueError("Must provide query text for DELTA_SYNC or vectors for DIRECT_ACCESS.")
raise ValueError("Must provide vectors for search.")
sdk_results = self.client.vector_search_indexes.query_index(**query_kwargs)
# Parse results
result_data = sdk_results.result if hasattr(sdk_results, "result") else sdk_results
@@ -572,14 +572,23 @@ class Databricks(VectorStoreBase):
filters = {"memory_id": vector_id}
filters_json = json.dumps(filters)
results = self.client.vector_search_indexes.query_index(
index_name=self.fully_qualified_index_name,
columns=self.column_names,
query_text=" ", # Empty query, rely on filters
num_results=1,
query_type=self.query_type,
filters_json=filters_json,
# Use query_text for Delta Sync with model endpoint, query_vector otherwise
query_kwargs = {
"index_name": self.fully_qualified_index_name,
"columns": self.column_names,
"num_results": 1,
"query_type": self.query_type,
"filters_json": filters_json,
}
uses_model_endpoint = (
self.index_type == VectorIndexType.DELTA_SYNC and self.embedding_model_endpoint_name
)
if uses_model_endpoint:
query_kwargs["query_text"] = " "
else:
query_kwargs["query_vector"] = [0.0] * self.embedding_dimension
results = self.client.vector_search_indexes.query_index(**query_kwargs)
# Process results
result_data = results.result if hasattr(results, "result") else results
@@ -589,7 +598,7 @@ class Databricks(VectorStoreBase):
raise KeyError(f"Vector with ID {vector_id} not found")
result = data_array[0]
columns = columns = [col.name for col in results.manifest.columns] if results.manifest and results.manifest.columns else []
columns = [col.name for col in results.manifest.columns] if results.manifest and results.manifest.columns else []
row_data = dict(zip(columns, result))
# Build payload following the standard schema
@@ -686,14 +695,23 @@ class Databricks(VectorStoreBase):
filters_json = json.dumps(filters) if filters else None
num_results = limit or 100
columns = self.column_names
sdk_results = self.client.vector_search_indexes.query_index(
index_name=self.fully_qualified_index_name,
columns=columns,
query_text=" ",
num_results=num_results,
query_type=self.query_type,
filters_json=filters_json,
# Use query_text for Delta Sync with model endpoint, query_vector otherwise
query_kwargs = {
"index_name": self.fully_qualified_index_name,
"columns": columns,
"num_results": num_results,
"query_type": self.query_type,
"filters_json": filters_json,
}
uses_model_endpoint = (
self.index_type == VectorIndexType.DELTA_SYNC and self.embedding_model_endpoint_name
)
if uses_model_endpoint:
query_kwargs["query_text"] = " "
else:
query_kwargs["query_vector"] = [0.0] * self.embedding_dimension
sdk_results = self.client.vector_search_indexes.query_index(**query_kwargs)
result_data = sdk_results.result if hasattr(sdk_results, "result") else sdk_results
data_array = result_data.data_array if hasattr(result_data, "data_array") else []
+13 -13
View File
@@ -26,7 +26,7 @@ class OutputData(BaseModel):
class MongoDB(VectorStoreBase):
VECTOR_TYPE = "knnVector"
VECTOR_TYPE = "vector"
SIMILARITY_METRIC = "cosine"
def __init__(self, db_name: str, collection_name: str, embedding_model_dims: int, mongo_uri: str):
@@ -69,17 +69,16 @@ class MongoDB(VectorStoreBase):
else:
search_index_model = SearchIndexModel(
name=self.index_name,
type="vectorSearch",
definition={
"mappings": {
"dynamic": False,
"fields": {
"embedding": {
"type": self.VECTOR_TYPE,
"dimensions": self.embedding_model_dims,
"similarity": self.SIMILARITY_METRIC,
}
},
}
"fields": [
{
"type": self.VECTOR_TYPE,
"path": "embedding",
"numDimensions": self.embedding_model_dims,
"similarity": self.SIMILARITY_METRIC,
}
]
},
)
collection.create_search_index(search_index_model)
@@ -141,7 +140,7 @@ class MongoDB(VectorStoreBase):
"$vectorSearch": {
"index": self.index_name,
"limit": limit,
"numCandidates": limit,
"numCandidates": min(limit * 20, 10000),
"queryVector": vectors,
"path": "embedding",
}
@@ -198,7 +197,8 @@ class MongoDB(VectorStoreBase):
if vector is not None:
update_fields["embedding"] = vector
if payload is not None:
update_fields["payload"] = payload
for key, value in payload.items():
update_fields[f"payload.{key}"] = value
if update_fields:
try:
+43 -13
View File
@@ -113,6 +113,25 @@ class OpenSearchDB(VectorStoreBase):
if payloads is None:
payloads = [{} for _ in range(len(vectors))]
for idx, vec in enumerate(vectors):
if vec is None:
raise ValueError(
f"Vector at index {idx} is null. "
f"This usually means the embedding model failed to generate an embedding. "
f"Check that your embedding model is configured correctly and returning valid vectors."
)
if len(vec) == 0:
raise ValueError(
f"Vector at index {idx} is empty. "
f"Expected a vector of dimension {self.embedding_model_dims}, got an empty vector."
)
if len(vec) != self.embedding_model_dims:
raise ValueError(
f"Vector at index {idx} has dimension {len(vec)}, "
f"but the index '{self.collection_name}' expects dimension {self.embedding_model_dims}. "
f"Ensure your embedding model's output dimensions match the vector store configuration."
)
results = []
for i, (vec, id_) in enumerate(zip(vectors, ids)):
body = {
@@ -124,14 +143,16 @@ class OpenSearchDB(VectorStoreBase):
self.client.index(index=self.collection_name, body=body)
# Force refresh to make documents immediately searchable for tests
self.client.indices.refresh(index=self.collection_name)
results.append(OutputData(
id=id_,
score=1.0, # No score for inserts
payload=payloads[i]
))
results.append(
OutputData(
id=id_,
score=1.0, # No score for inserts
payload=payloads[i],
)
)
except Exception as e:
logger.error(f"Error inserting vector {id_}: {e}")
logger.error(f"Error inserting vector {id_}: {e}", exc_info=True)
raise
return results
@@ -179,7 +200,7 @@ class OpenSearchDB(VectorStoreBase):
]
return results
except Exception as e:
logger.error(f"Error during search: {e}")
logger.error(f"Error during search: {e}", exc_info=True)
return []
def delete(self, vector_id: str) -> None:
@@ -200,6 +221,15 @@ class OpenSearchDB(VectorStoreBase):
def update(self, vector_id: str, vector: Optional[List[float]] = None, payload: Optional[Dict] = None) -> None:
"""Update a vector and its payload using the custom 'id' field."""
if vector is not None:
if len(vector) == 0:
raise ValueError("Cannot update with an empty vector.")
if len(vector) != self.embedding_model_dims:
raise ValueError(
f"Update vector has dimension {len(vector)}, "
f"but the index '{self.collection_name}' expects dimension {self.embedding_model_dims}. "
f"Ensure your embedding model's output dimensions match the vector store configuration."
)
# First, find the document by custom ID
search_query = {"query": {"term": {"id": vector_id}}}
@@ -222,8 +252,9 @@ class OpenSearchDB(VectorStoreBase):
if doc:
try:
response = self.client.update(index=self.collection_name, id=opensearch_id, body={"doc": doc})
except Exception:
pass
except Exception as e:
logger.error(f"Error updating vector {vector_id}: {e}", exc_info=True)
raise
def get(self, vector_id: str) -> Optional[OutputData]:
"""Retrieve a vector by ID."""
@@ -238,7 +269,7 @@ class OpenSearchDB(VectorStoreBase):
return OutputData(id=hits[0]["_source"].get("id"), score=1.0, payload=hits[0]["_source"].get("payload", {}))
except Exception as e:
logger.error(f"Error retrieving vector {vector_id}: {str(e)}")
logger.error(f"Error retrieving vector {vector_id}: {str(e)}", exc_info=True)
return None
def list_cols(self) -> List[str]:
@@ -281,9 +312,8 @@ class OpenSearchDB(VectorStoreBase):
]
return [results] # VectorStore expects tuple/list format
except Exception as e:
logger.error(f"Error listing vectors: {e}")
logger.error(f"Error listing vectors: {e}", exc_info=True)
return []
def reset(self):
"""Reset the index by deleting and recreating it."""
+136 -14
View File
@@ -1,12 +1,14 @@
import logging
import os
import shutil
from typing import Optional
from qdrant_client import QdrantClient
from qdrant_client.models import (
Distance,
FieldCondition,
Filter,
MatchAny,
MatchExcept,
MatchText,
MatchValue,
PointIdsList,
PointStruct,
@@ -44,7 +46,8 @@ class Qdrant(VectorStoreBase):
path (str, optional): Path for local Qdrant database. Defaults to None.
url (str, optional): Full URL for Qdrant server. Defaults to None.
api_key (str, optional): API key for Qdrant server. Defaults to None.
on_disk (bool, optional): Enables persistent storage. Defaults to False.
on_disk (bool, optional): Enables persistent storage. Vectors are stored on disk (True) or in memory (False).
Does not delete the local database path. Defaults to False.
"""
if client:
self.client = client
@@ -62,9 +65,6 @@ class Qdrant(VectorStoreBase):
if not params:
params["path"] = path
self.is_local = True
if not on_disk:
if os.path.exists(path) and os.path.isdir(path):
shutil.rmtree(path)
else:
self.is_local = False
@@ -138,26 +138,148 @@ class Qdrant(VectorStoreBase):
]
self.client.upsert(collection_name=self.collection_name, points=points)
def _create_filter(self, filters: dict) -> Filter:
def _build_field_condition(self, key: str, value) -> Optional[FieldCondition]:
"""
Build a single FieldCondition from a key-value filter pair.
Supports the enhanced filter syntax documented at
https://docs.mem0.ai/open-source/features/metadata-filtering
Args:
key (str): The payload field name.
value: A scalar for simple equality, or a dict with one operator key.
Returns:
Optional[FieldCondition]: The Qdrant field condition, or None if the
value is the wildcard '*' (match any / field exists — skip filter).
"""
if not isinstance(value, dict):
if value == "*":
# Wildcard: match any value. Qdrant has no direct "field exists"
# condition via FieldCondition, so we skip this filter (match all).
return None
if isinstance(value, list):
# List shorthand: {"field": ["a", "b"]} treated as in-operator.
return FieldCondition(key=key, match=MatchAny(any=value))
# Simple equality: {"field": "value"}
return FieldCondition(key=key, match=MatchValue(value=value))
ops = set(value.keys())
range_ops = {"gt", "gte", "lt", "lte"}
non_range_ops = ops - range_ops
if ops & range_ops:
if non_range_ops:
raise ValueError(
f"Cannot mix range operators ({ops & range_ops}) with "
f"non-range operators ({non_range_ops}) for field '{key}'. "
f"Use AND to combine them as separate conditions."
)
range_kwargs = {op: value[op] for op in range_ops if op in value}
return FieldCondition(key=key, range=Range(**range_kwargs))
elif "eq" in value:
return FieldCondition(key=key, match=MatchValue(value=value["eq"]))
elif "ne" in value:
return FieldCondition(key=key, match=MatchExcept(**{"except": [value["ne"]]}))
elif "in" in value:
return FieldCondition(key=key, match=MatchAny(any=value["in"]))
elif "nin" in value:
return FieldCondition(key=key, match=MatchExcept(**{"except": value["nin"]}))
elif "contains" in value or "icontains" in value:
# MatchText: with a full-text index, tokenized matching (all words must appear).
# Without a full-text index, exact substring match.
op = "icontains" if "icontains" in value else "contains"
text = value[op]
if op == "icontains":
logger.debug(
"icontains on field '%s': Qdrant MatchText case sensitivity depends on "
"full-text index configuration. Without a full-text index this behaves "
"as a case-sensitive substring match (same as 'contains').",
key,
)
return FieldCondition(key=key, match=MatchText(text=text))
else:
supported = {"eq", "ne", "gt", "gte", "lt", "lte", "in", "nin", "contains", "icontains"}
raise ValueError(
f"Unsupported filter operator(s) for field '{key}': {ops}. "
f"Supported operators: {supported}"
)
def _create_filter(self, filters: dict) -> Optional[Filter]:
"""
Create a Filter object from the provided filters.
Supports the enhanced filter syntax with comparison operators (eq, ne,
gt, gte, lt, lte), list operators (in, nin), string operators (contains,
icontains), and logical operators (AND, OR, NOT).
Args:
filters (dict): Filters to apply.
Returns:
Filter: The created Filter object.
Filter: The created Filter object, or None if filters is empty.
"""
if not filters:
return None
conditions = []
# Normalize $or/$not/$and → OR/NOT/AND and deduplicate.
# Memory._process_metadata_filters() renames OR→$or and NOT→$not,
# but effective_filters retains the original OR/NOT keys from
# deepcopy(input_filters). Without dedup the same sub-conditions
# would be evaluated twice.
key_map = {"$or": "OR", "$not": "NOT", "$and": "AND"}
normalized = {}
for key, value in filters.items():
if isinstance(value, dict) and "gte" in value and "lte" in value:
conditions.append(FieldCondition(key=key, range=Range(gte=value["gte"], lte=value["lte"])))
norm_key = key_map.get(key, key)
if norm_key not in normalized:
normalized[norm_key] = value
must = []
should = []
must_not = []
for key, value in normalized.items():
if key in ("AND", "OR", "NOT"):
if not isinstance(value, list):
raise ValueError(
f"{key} filter value must be a list of filter dicts, "
f"got {type(value).__name__}"
)
for i, item in enumerate(value):
if not isinstance(item, dict):
raise ValueError(
f"{key} filter list item at index {i} must be a dict, "
f"got {type(item).__name__}: {item!r}"
)
if key == "AND":
for sub in value:
built = self._create_filter(sub)
if built:
must.append(built)
elif key == "OR":
for sub in value:
built = self._create_filter(sub)
if built:
should.append(built)
elif key == "NOT":
for sub in value:
built = self._create_filter(sub)
if built:
must_not.append(built)
else:
conditions.append(FieldCondition(key=key, match=MatchValue(value=value)))
return Filter(must=conditions) if conditions else None
condition = self._build_field_condition(key, value)
if condition is not None:
must.append(condition)
if not any([must, should, must_not]):
return None
return Filter(
must=must or None,
should=should or None,
must_not=must_not or None,
)
def search(self, query: str, vectors: list, limit: int = 5, filters: dict = None) -> list:
"""
+337
View File
@@ -0,0 +1,337 @@
import logging
import os
from typing import Any, Dict, List, Optional, Union
try:
from turbopuffer import Turbopuffer as TurbopufferClient
except ImportError:
raise ImportError(
"Turbopuffer requires extra dependencies. Install with `pip install turbopuffer`"
) from None
from pydantic import BaseModel
from mem0.vector_stores.base import VectorStoreBase
logger = logging.getLogger(__name__)
class OutputData(BaseModel):
id: Optional[str]
score: Optional[float]
payload: Optional[Dict]
class TurbopufferDB(VectorStoreBase):
def __init__(
self,
collection_name: str,
embedding_model_dims: int,
api_key: Optional[str] = None,
region: str = "gcp-us-central1",
distance_metric: str = "cosine_distance",
batch_size: int = 100,
extra_params: Optional[Dict[str, Any]] = None,
):
"""
Initialize the Turbopuffer vector store.
Args:
collection_name (str): Name of the namespace/collection.
embedding_model_dims (int): Dimensions of the embedding model.
api_key (str, optional): API key for Turbopuffer. Defaults to None.
region (str, optional): Turbopuffer region. Defaults to "gcp-us-central1".
distance_metric (str, optional): Distance metric for vector similarity.
Options: "cosine_distance" or "euclidean_squared". Defaults to "cosine_distance".
batch_size (int, optional): Batch size for operations. Defaults to 100.
extra_params (Dict, optional): Additional parameters for Turbopuffer client. Defaults to None.
"""
api_key = api_key or os.environ.get("TURBOPUFFER_API_KEY")
if not api_key:
raise ValueError(
"Turbopuffer API key must be provided either as a parameter or via TURBOPUFFER_API_KEY environment variable"
)
params = extra_params or {}
params["region"] = region
self.client = TurbopufferClient(api_key=api_key, **params)
self.collection_name = collection_name
self.embedding_model_dims = embedding_model_dims
self.distance_metric = distance_metric
self.batch_size = batch_size
self.namespace = self.client.namespace(self.collection_name)
def create_col(self, name=None, vector_size=None, distance=None):
"""
Create a new namespace in Turbopuffer.
Namespaces are created implicitly on first upsert, so this is a no-op.
"""
pass
def insert(
self,
vectors: List[List[float]],
payloads: Optional[List[Dict]] = None,
ids: Optional[List[Union[str, int]]] = None,
):
"""
Insert vectors into the namespace.
Args:
vectors (list): List of vectors to insert.
payloads (list, optional): List of payloads corresponding to vectors. Defaults to None.
ids (list, optional): List of IDs corresponding to vectors. Defaults to None.
"""
logger.info(f"Inserting {len(vectors)} vectors into namespace {self.collection_name}")
if ids is None:
ids = [str(i) for i in range(len(vectors))]
for i in range(0, len(vectors), self.batch_size):
batch_end = i + self.batch_size
rows = []
for j in range(i, min(batch_end, len(vectors))):
row = {}
if payloads and payloads[j]:
row.update(payloads[j])
row["id"] = str(ids[j])
row["vector"] = vectors[j]
rows.append(row)
self.namespace.write(
upsert_rows=rows,
distance_metric=self.distance_metric,
)
def _parse_output(self, rows) -> List[OutputData]:
"""
Parse the output data from Turbopuffer query results.
Args:
rows: List of Row objects from Turbopuffer query.
Returns:
List[OutputData]: Parsed output data.
"""
results = []
for row in rows:
row_dict = row.model_dump()
row_id = str(row_dict.pop("id"))
dist = row_dict.pop("$dist", None)
row_dict.pop("vector", None)
score = 1 - dist if dist is not None else None
results.append(OutputData(
id=row_id,
score=score,
payload=row_dict,
))
return results
def _convert_filters(self, filters: Optional[Dict]):
"""
Convert mem0 filters to Turbopuffer filter format.
Turbopuffer filters use tuple format: ("And", (("field", "Op", value), ...))
"""
if not filters:
return None
conditions = []
for key, value in filters.items():
if isinstance(value, dict):
if "gte" in value:
conditions.append((key, "Gte", value["gte"]))
if "lte" in value:
conditions.append((key, "Lte", value["lte"]))
else:
conditions.append((key, "Eq", value))
if not conditions:
return None
if len(conditions) == 1:
return conditions[0]
return ("And", tuple(conditions))
def search(
self, query: str, vectors: List[float], limit: int = 5, filters: Optional[Dict] = None
) -> List[OutputData]:
"""
Search for similar vectors.
Args:
query (str): Query text (unused in vector search, kept for interface consistency).
vectors (list): Query vector to search with.
limit (int, optional): Number of results to return. Defaults to 5.
filters (dict, optional): Filters to apply to the search. Defaults to None.
Returns:
list: Search results.
"""
query_params = {
"rank_by": ("vector", "ANN", vectors),
"top_k": limit,
"include_attributes": True,
}
tpuf_filters = self._convert_filters(filters)
if tpuf_filters is not None:
query_params["filters"] = tpuf_filters
response = self.namespace.query(**query_params)
return self._parse_output(response.rows or [])
def delete(self, vector_id: Union[str, int]):
"""
Delete a vector by ID.
Args:
vector_id (Union[str, int]): ID of the vector to delete.
"""
self.namespace.write(deletes=[str(vector_id)])
def update(
self,
vector_id: Union[str, int],
vector: Optional[List[float]] = None,
payload: Optional[Dict] = None,
):
"""
Update a vector and its payload.
Args:
vector_id (Union[str, int]): ID of the vector to update.
vector (list, optional): Updated vector. Defaults to None.
payload (dict, optional): Updated payload. Defaults to None.
"""
if vector is not None:
row = {}
if payload:
row.update(payload)
row["id"] = str(vector_id)
row["vector"] = vector
self.namespace.write(
upsert_rows=[row],
distance_metric=self.distance_metric,
)
elif payload is not None:
row = dict(payload)
row["id"] = str(vector_id)
self.namespace.write(patch_rows=[row])
def get(self, vector_id: Union[str, int]) -> Optional[OutputData]:
"""
Retrieve a vector by ID.
Args:
vector_id (Union[str, int]): ID of the vector to retrieve.
Returns:
OutputData: Retrieved vector data, or None if not found.
"""
try:
response = self.namespace.query(
top_k=1,
rank_by=("vector", "ANN", [0.0] * self.embedding_model_dims),
filters=("id", "Eq", str(vector_id)),
include_attributes=True,
)
rows = response.rows or []
if rows:
return self._parse_output(rows)[0]
return None
except Exception as e:
logger.error(f"Error retrieving vector {vector_id}: {e}")
return None
def list_cols(self) -> list:
"""
List all namespaces.
Returns:
list: List of namespace summaries.
"""
try:
result = []
for ns in self.client.namespaces():
result.append(ns)
return result
except Exception as e:
logger.error(f"Error listing namespaces: {e}")
return []
def delete_col(self):
"""Delete the entire namespace."""
try:
self.namespace.delete_all()
logger.info(f"Namespace {self.collection_name} deleted successfully")
except Exception as e:
logger.error(f"Error deleting namespace {self.collection_name}: {e}")
def col_info(self) -> Dict:
"""
Get information about the namespace.
Returns:
dict: Namespace metadata.
"""
try:
metadata = self.namespace.metadata()
return {
"name": self.collection_name,
"approx_row_count": metadata.approx_row_count,
"approx_logical_bytes": metadata.approx_logical_bytes,
"created_at": str(metadata.created_at),
"updated_at": str(metadata.updated_at),
}
except Exception:
return {"name": self.collection_name}
def list(self, filters: Optional[Dict] = None, limit: int = 100) -> list:
"""
List vectors in the namespace with optional filtering.
Args:
filters (dict, optional): Filters to apply. Defaults to None.
limit (int, optional): Number of vectors to return. Defaults to 100.
Returns:
list: Wrapped list of OutputData objects ([[results]]).
"""
query_params = {
"rank_by": ("vector", "ANN", [0.0] * self.embedding_model_dims),
"top_k": limit,
"include_attributes": True,
}
tpuf_filters = self._convert_filters(filters)
if tpuf_filters is not None:
query_params["filters"] = tpuf_filters
try:
response = self.namespace.query(**query_params)
results = self._parse_output(response.rows or [])
except Exception as e:
logger.error(f"Error listing vectors: {e}")
results = []
return [results]
def count(self) -> int:
"""
Get approximate count of vectors in the namespace.
Returns:
int: Approximate number of vectors.
"""
try:
metadata = self.namespace.metadata()
return metadata.approx_row_count
except Exception:
return 0
def reset(self):
"""Reset the namespace by deleting all vectors."""
self.delete_col()
+1 -1
View File
@@ -184,7 +184,7 @@ async def search_memory(query: str) -> str:
for h in hits:
# All vector db search functions return OutputData class
id, score, payload = h.id, h.score, h.payload
if allowed and h.id is None or h.id not in allowed:
if allowed and (h.id is None or h.id not in allowed):
continue
results.append({
+1
View File
@@ -31,6 +31,7 @@ graph = [
"neo4j>=5.23.1",
"rank-bm25>=0.2.2",
"kuzu>=0.11.0",
"apache-age-python>=0.0.6",
]
vector_stores = [
"vecs>=0.4.0",
+3
View File
@@ -10,3 +10,6 @@ POSTGRES_DB=
POSTGRES_USER=
POSTGRES_PASSWORD=
POSTGRES_COLLECTION_NAME=
# Optional: set to enable API key authentication on all endpoints
ADMIN_API_KEY=
+60 -10
View File
@@ -1,10 +1,12 @@
import logging
import os
import secrets
from typing import Any, Dict, List, Optional
from dotenv import load_dotenv
from fastapi import FastAPI, HTTPException
from fastapi import Depends, FastAPI, HTTPException
from fastapi.responses import JSONResponse, RedirectResponse
from fastapi.security import APIKeyHeader
from pydantic import BaseModel, Field
from mem0 import Memory
@@ -14,6 +16,22 @@ logging.basicConfig(level=logging.INFO, format="%(asctime)s - %(levelname)s - %(
# Load environment variables
load_dotenv()
ADMIN_API_KEY = os.environ.get("ADMIN_API_KEY", "")
MIN_KEY_LENGTH = 16
if not ADMIN_API_KEY:
logging.warning(
"ADMIN_API_KEY not set - API endpoints are UNSECURED! "
"Set ADMIN_API_KEY environment variable for production use."
)
else:
if len(ADMIN_API_KEY) < MIN_KEY_LENGTH:
logging.warning(
"ADMIN_API_KEY is shorter than %d characters - consider using a longer key for production.",
MIN_KEY_LENGTH,
)
logging.info("API key authentication enabled")
POSTGRES_HOST = os.environ.get("POSTGRES_HOST", "postgres")
POSTGRES_PORT = os.environ.get("POSTGRES_PORT", "5432")
@@ -60,10 +78,35 @@ MEMORY_INSTANCE = Memory.from_config(DEFAULT_CONFIG)
app = FastAPI(
title="Mem0 REST APIs",
description="A REST API for managing and searching memories for your AI Agents and Apps.",
description=(
"A REST API for managing and searching memories for your AI Agents and Apps.\n\n"
"## Authentication\n"
"When the ADMIN_API_KEY environment variable is set, all endpoints require "
"the `X-API-Key` header for authentication."
),
version="1.0.0",
)
api_key_header = APIKeyHeader(name="X-API-Key", auto_error=False)
async def verify_api_key(api_key: Optional[str] = Depends(api_key_header)):
"""Validate the API key when ADMIN_API_KEY is configured. No-op otherwise."""
if ADMIN_API_KEY:
if api_key is None:
raise HTTPException(
status_code=401,
detail="X-API-Key header is required.",
headers={"WWW-Authenticate": "ApiKey"},
)
if not secrets.compare_digest(api_key, ADMIN_API_KEY):
raise HTTPException(
status_code=401,
detail="Invalid API key.",
headers={"WWW-Authenticate": "ApiKey"},
)
return api_key
class Message(BaseModel):
role: str = Field(..., description="Role of the message (user or assistant).")
@@ -76,6 +119,9 @@ class MemoryCreate(BaseModel):
agent_id: Optional[str] = None
run_id: Optional[str] = None
metadata: Optional[Dict[str, Any]] = None
infer: Optional[bool] = Field(None, description="Whether to extract facts from messages. Defaults to True.")
memory_type: Optional[str] = Field(None, description="Type of memory to store (e.g. 'core').")
prompt: Optional[str] = Field(None, description="Custom prompt to use for fact extraction.")
class SearchRequest(BaseModel):
@@ -84,10 +130,12 @@ class SearchRequest(BaseModel):
run_id: Optional[str] = None
agent_id: Optional[str] = None
filters: Optional[Dict[str, Any]] = None
limit: Optional[int] = Field(None, description="Maximum number of results to return.")
threshold: Optional[float] = Field(None, description="Minimum similarity score for results.")
@app.post("/configure", summary="Configure Mem0")
def set_config(config: Dict[str, Any]):
def set_config(config: Dict[str, Any], _api_key: Optional[str] = Depends(verify_api_key)):
"""Set memory configuration."""
global MEMORY_INSTANCE
MEMORY_INSTANCE = Memory.from_config(config)
@@ -95,7 +143,7 @@ def set_config(config: Dict[str, Any]):
@app.post("/memories", summary="Create memories")
def add_memory(memory_create: MemoryCreate):
def add_memory(memory_create: MemoryCreate, _api_key: Optional[str] = Depends(verify_api_key)):
"""Store new memories."""
if not any([memory_create.user_id, memory_create.agent_id, memory_create.run_id]):
raise HTTPException(status_code=400, detail="At least one identifier (user_id, agent_id, run_id) is required.")
@@ -114,6 +162,7 @@ def get_all_memories(
user_id: Optional[str] = None,
run_id: Optional[str] = None,
agent_id: Optional[str] = None,
_api_key: Optional[str] = Depends(verify_api_key),
):
"""Retrieve stored memories."""
if not any([user_id, run_id, agent_id]):
@@ -129,7 +178,7 @@ def get_all_memories(
@app.get("/memories/{memory_id}", summary="Get a memory")
def get_memory(memory_id: str):
def get_memory(memory_id: str, _api_key: Optional[str] = Depends(verify_api_key)):
"""Retrieve a specific memory by ID."""
try:
return MEMORY_INSTANCE.get(memory_id)
@@ -139,7 +188,7 @@ def get_memory(memory_id: str):
@app.post("/search", summary="Search memories")
def search_memories(search_req: SearchRequest):
def search_memories(search_req: SearchRequest, _api_key: Optional[str] = Depends(verify_api_key)):
"""Search for memories based on a query."""
try:
params = {k: v for k, v in search_req.model_dump().items() if v is not None and k != "query"}
@@ -150,7 +199,7 @@ def search_memories(search_req: SearchRequest):
@app.put("/memories/{memory_id}", summary="Update a memory")
def update_memory(memory_id: str, updated_memory: Dict[str, Any]):
def update_memory(memory_id: str, updated_memory: Dict[str, Any], _api_key: Optional[str] = Depends(verify_api_key)):
"""Update an existing memory with new content.
Args:
@@ -168,7 +217,7 @@ def update_memory(memory_id: str, updated_memory: Dict[str, Any]):
@app.get("/memories/{memory_id}/history", summary="Get memory history")
def memory_history(memory_id: str):
def memory_history(memory_id: str, _api_key: Optional[str] = Depends(verify_api_key)):
"""Retrieve memory history."""
try:
return MEMORY_INSTANCE.history(memory_id=memory_id)
@@ -178,7 +227,7 @@ def memory_history(memory_id: str):
@app.delete("/memories/{memory_id}", summary="Delete a memory")
def delete_memory(memory_id: str):
def delete_memory(memory_id: str, _api_key: Optional[str] = Depends(verify_api_key)):
"""Delete a specific memory by ID."""
try:
MEMORY_INSTANCE.delete(memory_id=memory_id)
@@ -193,6 +242,7 @@ def delete_all_memories(
user_id: Optional[str] = None,
run_id: Optional[str] = None,
agent_id: Optional[str] = None,
_api_key: Optional[str] = Depends(verify_api_key),
):
"""Delete all memories for a given identifier."""
if not any([user_id, run_id, agent_id]):
@@ -209,7 +259,7 @@ def delete_all_memories(
@app.post("/reset", summary="Reset all memories")
def reset_memory():
def reset_memory(_api_key: Optional[str] = Depends(verify_api_key)):
"""Completely reset stored memories."""
try:
MEMORY_INSTANCE.reset()
+100
View File
@@ -0,0 +1,100 @@
from unittest.mock import Mock, patch
import pytest
pytest.importorskip("anthropic", reason="anthropic package not installed")
from mem0.configs.llms.anthropic import AnthropicConfig
from mem0.configs.llms.base import BaseLlmConfig
from mem0.llms.anthropic import AnthropicLLM
@pytest.fixture
def mock_anthropic_client():
with patch("mem0.llms.anthropic.anthropic") as mock_anthropic:
mock_client = Mock()
mock_anthropic.Anthropic.return_value = mock_client
yield mock_client
def test_default_config_omits_top_p(mock_anthropic_client):
"""Default AnthropicConfig should not set top_p to avoid conflict with temperature."""
config = AnthropicConfig(model="claude-3-5-sonnet-20240620", api_key="test-key")
assert config.top_p is None
assert config.temperature == 0.1
def test_generate_response_does_not_send_top_p_by_default(mock_anthropic_client):
"""Anthropic API rejects temperature and top_p together; top_p must be omitted by default."""
config = AnthropicConfig(model="claude-3-5-sonnet-20240620", api_key="test-key")
llm = AnthropicLLM(config)
mock_response = Mock()
mock_response.content = [Mock(text="Hello!")]
mock_anthropic_client.messages.create.return_value = mock_response
messages = [
{"role": "system", "content": "You are a helpful assistant."},
{"role": "user", "content": "Hi"},
]
llm.generate_response(messages)
call_kwargs = mock_anthropic_client.messages.create.call_args[1]
assert "top_p" not in call_kwargs
assert call_kwargs["temperature"] == 0.1
def test_generate_response_sends_top_p_alone_when_no_temperature(mock_anthropic_client):
"""When user sets only top_p (no temperature), top_p should be sent."""
config = AnthropicConfig(model="claude-3-5-sonnet-20240620", api_key="test-key", top_p=0.9, temperature=None)
llm = AnthropicLLM(config)
mock_response = Mock()
mock_response.content = [Mock(text="Hello!")]
mock_anthropic_client.messages.create.return_value = mock_response
messages = [
{"role": "system", "content": "You are a helpful assistant."},
{"role": "user", "content": "Hi"},
]
llm.generate_response(messages)
call_kwargs = mock_anthropic_client.messages.create.call_args[1]
assert call_kwargs["top_p"] == 0.9
assert "temperature" not in call_kwargs
def test_both_set_prefers_temperature_over_top_p(mock_anthropic_client):
"""When both temperature and top_p are set, temperature wins and top_p is dropped."""
config = AnthropicConfig(model="claude-3-5-sonnet-20240620", api_key="test-key", top_p=0.9, temperature=0.5)
llm = AnthropicLLM(config)
mock_response = Mock()
mock_response.content = [Mock(text="Hello!")]
mock_anthropic_client.messages.create.return_value = mock_response
messages = [{"role": "user", "content": "Hi"}]
llm.generate_response(messages)
call_kwargs = mock_anthropic_client.messages.create.call_args[1]
assert call_kwargs["temperature"] == 0.5
assert "top_p" not in call_kwargs
def test_base_config_conversion_does_not_send_both(mock_anthropic_client):
"""BaseLlmConfig defaults both temperature=0.1 and top_p=0.1; Anthropic must not send both."""
base_config = BaseLlmConfig(model="claude-3-5-sonnet-20240620", api_key="test-key")
llm = AnthropicLLM(base_config)
mock_response = Mock()
mock_response.content = [Mock(text="Hello!")]
mock_anthropic_client.messages.create.return_value = mock_response
messages = [{"role": "user", "content": "Hi"}]
llm.generate_response(messages)
call_kwargs = mock_anthropic_client.messages.create.call_args[1]
assert "temperature" in call_kwargs
assert "top_p" not in call_kwargs
@@ -1,4 +1,5 @@
from unittest import mock
from unittest.mock import Mock
from mem0.llms.azure_openai_structured import SCOPE, AzureOpenAIStructuredLLM
@@ -98,3 +99,75 @@ def test_init_with_placeholder_api_key_uses_default_credential(
args, kwargs = mock_azure_openai.call_args
assert kwargs["api_key"] is None
assert kwargs["azure_ad_token_provider"] == "token-provider"
@mock.patch("mem0.llms.azure_openai_structured.AzureOpenAI")
def test_generate_response_without_tools(mock_azure_openai):
mock_client = Mock()
mock_azure_openai.return_value = mock_client
config = DummyConfig(model="test-model", azure_kwargs=DummyAzureKwargs(api_key="real-key"))
llm = AzureOpenAIStructuredLLM(config)
mock_response = Mock()
mock_response.choices = [Mock(message=Mock(content="Hello there!"))]
mock_client.chat.completions.create.return_value = mock_response
messages = [{"role": "user", "content": "Hi"}]
response = llm.generate_response(messages)
assert response == "Hello there!"
@mock.patch("mem0.llms.azure_openai_structured.AzureOpenAI")
def test_generate_response_with_tools(mock_azure_openai):
mock_client = Mock()
mock_azure_openai.return_value = mock_client
config = DummyConfig(model="test-model", azure_kwargs=DummyAzureKwargs(api_key="real-key"))
llm = AzureOpenAIStructuredLLM(config)
mock_tool_call = Mock()
mock_tool_call.function.name = "add_memory"
mock_tool_call.function.arguments = '{"data": "sunny day"}'
mock_message = Mock()
mock_message.content = "I've added the memory."
mock_message.tool_calls = [mock_tool_call]
mock_response = Mock()
mock_response.choices = [Mock(message=mock_message)]
mock_client.chat.completions.create.return_value = mock_response
tools = [{"type": "function", "function": {"name": "add_memory"}}]
messages = [{"role": "user", "content": "Remember sunny day"}]
response = llm.generate_response(messages, tools=tools)
assert response["content"] == "I've added the memory."
assert len(response["tool_calls"]) == 1
assert response["tool_calls"][0]["name"] == "add_memory"
assert response["tool_calls"][0]["arguments"] == {"data": "sunny day"}
@mock.patch("mem0.llms.azure_openai_structured.AzureOpenAI")
def test_generate_response_with_tools_no_tool_calls(mock_azure_openai):
mock_client = Mock()
mock_azure_openai.return_value = mock_client
config = DummyConfig(model="test-model", azure_kwargs=DummyAzureKwargs(api_key="real-key"))
llm = AzureOpenAIStructuredLLM(config)
mock_message = Mock()
mock_message.content = "No tools needed."
mock_message.tool_calls = None
mock_response = Mock()
mock_response.choices = [Mock(message=mock_message)]
mock_client.chat.completions.create.return_value = mock_response
tools = [{"type": "function", "function": {"name": "add_memory"}}]
messages = [{"role": "user", "content": "Hello"}]
response = llm.generate_response(messages, tools=tools)
assert response["content"] == "No tools needed."
assert response["tool_calls"] == []
+70
View File
@@ -117,3 +117,73 @@ def test_generate_response_with_tools(mock_gemini_client: Mock):
assert len(response["tool_calls"]) == 1
assert response["tool_calls"][0]["name"] == "add_memory"
assert response["tool_calls"][0]["arguments"] == {"data": "Today is a sunny day."}
def test_parse_response_none_content_no_tools(mock_gemini_client: Mock):
"""Gemini can return content=None when response is blocked by safety filters."""
config = BaseLlmConfig(model="gemini-2.0-flash", temperature=0.7, max_tokens=100, top_p=1.0)
llm = GeminiLLM(config)
mock_candidate = Mock(content=None)
mock_response = Mock(candidates=[mock_candidate])
result = llm._parse_response(mock_response, tools=None)
assert result == ""
def test_parse_response_none_content_with_tools(mock_gemini_client: Mock):
"""Gemini can return content=None when response is blocked by safety filters (tools path)."""
config = BaseLlmConfig(model="gemini-2.0-flash", temperature=0.7, max_tokens=100, top_p=1.0)
llm = GeminiLLM(config)
mock_candidate = Mock(content=None)
mock_response = Mock(candidates=[mock_candidate])
result = llm._parse_response(mock_response, tools=[{"function": {"name": "test"}}])
assert result == {"content": None, "tool_calls": []}
def test_parse_response_empty_candidates(mock_gemini_client: Mock):
"""Gemini can return an empty candidates list."""
config = BaseLlmConfig(model="gemini-2.0-flash", temperature=0.7, max_tokens=100, top_p=1.0)
llm = GeminiLLM(config)
mock_response = Mock(candidates=[])
result = llm._parse_response(mock_response, tools=None)
assert result == ""
def test_parse_response_none_candidates(mock_gemini_client: Mock):
"""Gemini can return candidates=None."""
config = BaseLlmConfig(model="gemini-2.0-flash", temperature=0.7, max_tokens=100, top_p=1.0)
llm = GeminiLLM(config)
mock_response = Mock(candidates=None)
result = llm._parse_response(mock_response, tools=None)
assert result == ""
def test_parse_response_empty_parts_no_tools(mock_gemini_client: Mock):
"""Gemini can return content with an empty parts list."""
config = BaseLlmConfig(model="gemini-2.0-flash", temperature=0.7, max_tokens=100, top_p=1.0)
llm = GeminiLLM(config)
mock_content = Mock(parts=[])
mock_candidate = Mock(content=mock_content)
mock_response = Mock(candidates=[mock_candidate])
result = llm._parse_response(mock_response, tools=None)
assert result == ""
def test_parse_response_empty_parts_with_tools(mock_gemini_client: Mock):
"""Gemini can return content with an empty parts list (tools path)."""
config = BaseLlmConfig(model="gemini-2.0-flash", temperature=0.7, max_tokens=100, top_p=1.0)
llm = GeminiLLM(config)
mock_content = Mock(parts=[])
mock_candidate = Mock(content=mock_content)
mock_response = Mock(candidates=[mock_candidate])
result = llm._parse_response(mock_response, tools=[{"function": {"name": "test"}}])
assert result == {"content": None, "tool_calls": []}
+194
View File
@@ -0,0 +1,194 @@
import os
from unittest.mock import Mock, patch
import pytest
from mem0.configs.llms.base import BaseLlmConfig
from mem0.configs.llms.minimax import MinimaxConfig
from mem0.llms.minimax import MiniMaxLLM
from mem0.utils.factory import LlmFactory
@pytest.fixture
def mock_minimax_client():
with patch("mem0.llms.minimax.OpenAI") as mock_openai:
mock_client = Mock()
mock_openai.return_value = mock_client
yield mock_client
def test_minimax_llm_default_base_url():
"""Default config uses MiniMax official base URL."""
config = BaseLlmConfig(
model="MiniMax-M2.7", temperature=0.7, max_tokens=100, top_p=1.0, api_key="api_key"
)
llm = MiniMaxLLM(config)
# OpenAI client may normalize URL with trailing slash
assert str(llm.client.base_url).rstrip("/") == "https://api.minimax.io/v1"
def test_minimax_llm_env_base_url():
"""Config uses MINIMAX_API_BASE env variable when set."""
provider_base_url = "https://api.provider.com/v1/"
os.environ["MINIMAX_API_BASE"] = provider_base_url
try:
config = MinimaxConfig(
model="MiniMax-M2.7",
temperature=0.7,
max_tokens=100,
top_p=1.0,
api_key="api_key",
)
llm = MiniMaxLLM(config)
assert str(llm.client.base_url).rstrip("/") == provider_base_url.rstrip("/")
finally:
os.environ.pop("MINIMAX_API_BASE", None)
def test_minimax_llm_config_base_url():
"""Config uses minimax_base_url when provided."""
config_base_url = "https://api.config.com/v1/"
config = MinimaxConfig(
model="MiniMax-M2.7",
temperature=0.7,
max_tokens=100,
top_p=1.0,
api_key="api_key",
minimax_base_url=config_base_url,
)
llm = MiniMaxLLM(config)
assert str(llm.client.base_url).rstrip("/") == config_base_url.rstrip("/")
def test_minimax_llm_default_model(mock_minimax_client):
"""Default model is MiniMax-M2.7 when not specified."""
config = MinimaxConfig(temperature=0.7, max_tokens=100, api_key="api_key")
llm = MiniMaxLLM(config)
assert llm.config.model == "MiniMax-M2.7"
def test_minimax_llm_env_api_key():
"""Uses MINIMAX_API_KEY env when api_key not in config."""
os.environ["MINIMAX_API_KEY"] = "env-api-key"
try:
with patch("mem0.llms.minimax.OpenAI") as mock_openai:
mock_client = Mock()
mock_openai.return_value = mock_client
config = MinimaxConfig(model="MiniMax-M2.7", api_key=None)
MiniMaxLLM(config)
mock_openai.assert_called_once_with(
api_key="env-api-key",
base_url="https://api.minimax.io/v1",
)
finally:
os.environ.pop("MINIMAX_API_KEY", None)
def test_generate_response_without_tools(mock_minimax_client):
"""generate_response returns text when no tools provided."""
config = BaseLlmConfig(
model="MiniMax-M2.7", temperature=0.7, max_tokens=100, top_p=1.0, api_key="api_key"
)
llm = MiniMaxLLM(config)
messages = [
{"role": "system", "content": "You are a helpful assistant."},
{"role": "user", "content": "Hello, how are you?"},
]
mock_response = Mock()
mock_response.choices = [Mock(message=Mock(content="I'm doing well, thank you for asking!"))]
mock_minimax_client.chat.completions.create.return_value = mock_response
response = llm.generate_response(messages)
mock_minimax_client.chat.completions.create.assert_called_once_with(
model="MiniMax-M2.7", messages=messages, temperature=0.7, max_tokens=100, top_p=1.0
)
assert response == "I'm doing well, thank you for asking!"
def test_generate_response_with_tools(mock_minimax_client):
"""generate_response returns tool_calls when tools provided."""
config = BaseLlmConfig(
model="MiniMax-M2.7", temperature=0.7, max_tokens=100, top_p=1.0, api_key="api_key"
)
llm = MiniMaxLLM(config)
messages = [
{"role": "system", "content": "You are a helpful assistant."},
{"role": "user", "content": "Add a new memory: Today is a sunny day."},
]
tools = [
{
"type": "function",
"function": {
"name": "add_memory",
"description": "Add a memory",
"parameters": {
"type": "object",
"properties": {"data": {"type": "string", "description": "Data to add to memory"}},
"required": ["data"],
},
},
}
]
mock_response = Mock()
mock_message = Mock()
mock_message.content = "I've added the memory for you."
mock_tool_call = Mock()
mock_tool_call.function.name = "add_memory"
mock_tool_call.function.arguments = '{"data": "Today is a sunny day."}'
mock_message.tool_calls = [mock_tool_call]
mock_response.choices = [Mock(message=mock_message)]
mock_minimax_client.chat.completions.create.return_value = mock_response
response = llm.generate_response(messages, tools=tools)
mock_minimax_client.chat.completions.create.assert_called_once_with(
model="MiniMax-M2.7",
messages=messages,
temperature=0.7,
max_tokens=100,
top_p=1.0,
tools=tools,
tool_choice="auto",
)
assert response["content"] == "I've added the memory for you."
assert len(response["tool_calls"]) == 1
assert response["tool_calls"][0]["name"] == "add_memory"
assert response["tool_calls"][0]["arguments"] == {"data": "Today is a sunny day."}
def test_generate_response_with_response_format(mock_minimax_client):
"""generate_response passes response_format to the API."""
config = BaseLlmConfig(
model="MiniMax-M2.7", temperature=0.7, max_tokens=100, top_p=1.0, api_key="api_key"
)
llm = MiniMaxLLM(config)
messages = [{"role": "user", "content": "Return JSON."}]
response_format = {"type": "json_object"}
mock_response = Mock()
mock_response.choices = [Mock(message=Mock(content='{"key": "value"}'))]
mock_minimax_client.chat.completions.create.return_value = mock_response
llm.generate_response(messages, response_format=response_format)
mock_minimax_client.chat.completions.create.assert_called_once_with(
model="MiniMax-M2.7",
messages=messages,
temperature=0.7,
max_tokens=100,
top_p=1.0,
response_format={"type": "json_object"},
)
def test_factory_creates_minimax_llm(mock_minimax_client):
"""LlmFactory.create returns MiniMaxLLM for provider 'minimax'."""
llm = LlmFactory.create("minimax", {"model": "MiniMax-M2.7", "api_key": "test-key"})
assert isinstance(llm, MiniMaxLLM)
assert llm.config.model == "MiniMax-M2.7"
File diff suppressed because it is too large Load Diff
+226
View File
@@ -0,0 +1,226 @@
from unittest.mock import MagicMock, Mock, patch
# age and rank_bm25 are optional deps — mock them so tests run without install
_age_mock = Mock()
patch.dict("sys.modules", {
"age": _age_mock,
"age.models": Mock(),
"rank_bm25": Mock(),
}).start()
from mem0.memory.apache_age_memory import MemoryGraph, _cosine_similarity # noqa: E402
def _make_instance():
with patch.object(MemoryGraph, "__init__", return_value=None):
instance = MemoryGraph.__new__(MemoryGraph)
instance.llm_provider = "openai"
instance.llm = MagicMock()
instance.embedding_model = MagicMock()
instance.config = MagicMock()
instance.config.graph_store.custom_prompt = None
instance.ag = MagicMock()
instance.graph_name = "test_graph"
instance.threshold = 0.7
return instance
class TestCosineSimilarity:
"""Tests for the _cosine_similarity helper."""
def test_identical_vectors(self):
assert abs(_cosine_similarity([1, 0, 0], [1, 0, 0]) - 1.0) < 1e-6
def test_orthogonal_vectors(self):
assert abs(_cosine_similarity([1, 0, 0], [0, 1, 0])) < 1e-6
def test_zero_vector(self):
assert _cosine_similarity([0, 0, 0], [1, 2, 3]) == 0.0
class TestRetrieveNodesFromData:
"""Tests for _retrieve_nodes_from_data in Apache AGE MemoryGraph."""
def test_normal_entities_extracted(self):
instance = _make_instance()
instance.llm.generate_response.return_value = {
"tool_calls": [{"name": "extract_entities", "arguments": {"entities": [
{"entity": "Alice", "entity_type": "person"},
{"entity": "hiking", "entity_type": "activity"},
]}}]
}
result = instance._retrieve_nodes_from_data("Alice loves hiking", {"user_id": "u1"})
assert result == {"alice": "person", "hiking": "activity"}
def test_malformed_entity_missing_entity_type_is_skipped(self):
instance = _make_instance()
instance.llm.generate_response.return_value = {
"tool_calls": [{"name": "extract_entities", "arguments": {"entities": [
{"entity": "matrix multiplication", "entity_type": "task"},
{"entity": "task"},
{"entity": "ReLU", "entity_type": "task"},
]}}]
}
result = instance._retrieve_nodes_from_data("some text", {"user_id": "u1"})
assert "matrix_multiplication" in result
assert "relu" in result
assert "task" not in result
def test_missing_entities_key_returns_empty(self):
instance = _make_instance()
instance.llm.generate_response.return_value = {
"tool_calls": [{"name": "extract_entities", "arguments": {"text": "Hello."}}]
}
result = instance._retrieve_nodes_from_data("Hello.", {"user_id": "u1"})
assert result == {}
def test_none_tool_calls_returns_empty(self):
instance = _make_instance()
instance.llm.generate_response.return_value = {"tool_calls": None}
result = instance._retrieve_nodes_from_data("hello world", {"user_id": "u1"})
assert result == {}
class TestEstablishNodesRelationsFromData:
"""Tests for _establish_nodes_relations_from_data in Apache AGE MemoryGraph."""
def test_none_response_does_not_crash(self):
instance = _make_instance()
instance.llm.generate_response.return_value = None
result = instance._establish_nodes_relations_from_data(
"Hello world", {"user_id": "u1"}, {}
)
assert result == []
def test_empty_tool_calls_returns_empty(self):
instance = _make_instance()
instance.llm.generate_response.return_value = {"tool_calls": []}
result = instance._establish_nodes_relations_from_data(
"Hello world", {"user_id": "u1"}, {}
)
assert result == []
def test_valid_entities_returned(self):
instance = _make_instance()
instance.llm.generate_response.return_value = {
"tool_calls": [{"name": "add_entities", "arguments": {"entities": [
{"source": "alice", "relationship": "loves", "destination": "hiking"}
]}}]
}
result = instance._establish_nodes_relations_from_data(
"Alice loves hiking", {"user_id": "u1"}, {"alice": "person"}
)
assert len(result) == 1
assert result[0]["source"] == "alice"
class TestRemoveSpacesFromEntities:
"""Tests for _remove_spaces_from_entities."""
def test_spaces_and_case(self):
instance = _make_instance()
entities = [{"source": "Alice Smith", "relationship": "Works At", "destination": "Big Corp"}]
result = instance._remove_spaces_from_entities(entities)
assert result[0]["source"] == "alice_smith"
assert result[0]["relationship"] == "works_at"
assert result[0]["destination"] == "big_corp"
class TestFindSimilarNode:
"""Tests for _find_similar_node."""
def test_returns_none_when_no_nodes(self):
instance = _make_instance()
instance._exec_cypher = MagicMock(return_value=[])
result = instance._find_similar_node([1.0, 0.0], {"user_id": "u1"}, threshold=0.9)
assert result is None
def test_returns_best_match_above_threshold(self):
instance = _make_instance()
instance._exec_cypher = MagicMock(return_value=[
{"name": "alice", "embedding": [1.0, 0.0], "user_id": "u1"},
{"name": "bob", "embedding": [0.0, 1.0], "user_id": "u1"},
])
result = instance._find_similar_node([1.0, 0.0], {"user_id": "u1"}, threshold=0.9)
assert result["name"] == "alice"
def test_filters_by_agent_id(self):
instance = _make_instance()
instance._exec_cypher = MagicMock(return_value=[
{"name": "alice", "embedding": [1.0, 0.0], "user_id": "u1", "agent_id": "a2"},
])
result = instance._find_similar_node(
[1.0, 0.0], {"user_id": "u1", "agent_id": "a1"}, threshold=0.9
)
assert result is None
class TestDeleteAll:
"""Tests for delete_all."""
def test_calls_exec_cypher_and_commits(self):
instance = _make_instance()
instance._exec_cypher = MagicMock(return_value=[])
instance.delete_all({"user_id": "u1"})
instance._exec_cypher.assert_called_once()
instance.ag.commit.assert_called_once()
class TestGetAll:
"""Tests for get_all."""
def test_returns_formatted_results(self):
instance = _make_instance()
instance._exec_cypher = MagicMock(return_value=[
{"source": "alice", "relationship": "KNOWS", "target": "bob"},
{"source": "alice", "relationship": "LIKES", "target": "hiking"},
])
results = instance.get_all({"user_id": "u1"}, limit=10)
assert len(results) == 2
assert results[0]["source"] == "alice"
assert results[0]["relationship"] == "KNOWS"
assert results[0]["target"] == "bob"
def test_passes_limit_to_cypher(self):
"""Limit is enforced via LIMIT in the Cypher query, not Python slicing."""
instance = _make_instance()
instance._exec_cypher = MagicMock(return_value=[
{"source": "n0", "relationship": "R", "target": "m0"},
])
instance.get_all({"user_id": "u1"}, limit=3)
# Verify limit was passed as a parameter to the query
cypher_stmt = instance._exec_cypher.call_args[0][0]
assert "LIMIT %s" in cypher_stmt
params = instance._exec_cypher.call_args[1].get("params") or instance._exec_cypher.call_args[0][2]
assert 3 in params
class TestAdd:
"""Tests for the add orchestration method."""
def test_add_returns_added_and_deleted(self):
instance = _make_instance()
instance._retrieve_nodes_from_data = MagicMock(return_value={"alice": "person"})
instance._establish_nodes_relations_from_data = MagicMock(return_value=[
{"source": "alice", "relationship": "knows", "destination": "bob"}
])
instance._search_graph_db = MagicMock(return_value=[])
instance._get_delete_entities_from_search_output = MagicMock(return_value=[])
instance._delete_entities = MagicMock(return_value=[])
instance._add_entities = MagicMock(return_value=["added"])
result = instance.add("Alice knows Bob", {"user_id": "u1"})
assert "deleted_entities" in result
assert "added_entities" in result
assert result["added_entities"] == ["added"]
class TestSearch:
"""Tests for the search method."""
def test_returns_empty_when_no_search_output(self):
instance = _make_instance()
instance._retrieve_nodes_from_data = MagicMock(return_value={"alice": "person"})
instance._search_graph_db = MagicMock(return_value=[])
result = instance.search("Who is Alice?", {"user_id": "u1"})
assert result == []
@@ -0,0 +1,315 @@
"""Tests for graph memory soft-delete behavior.
Verifies that _delete_entities marks relationships as invalid (soft-delete)
rather than permanently removing them, and that search/retrieval queries
exclude soft-deleted relationships by default.
See: https://github.com/mem0ai/mem0/issues/4187
"""
from unittest.mock import Mock, patch
# Mock optional deps at module level so the import works across all Python
# versions without triggering transitive C-extension reloads (numpy via
# qdrant_client). This matches the pattern in test_memgraph_memory.py.
_neo4j_mock = Mock()
patch.dict("sys.modules", {
"langchain_neo4j": _neo4j_mock,
"rank_bm25": Mock(),
}).start()
from mem0.memory.graph_memory import MemoryGraph # noqa: E402
def _create_graph_memory():
"""Create a MemoryGraph instance with mocked dependencies."""
with patch.object(MemoryGraph, "__init__", lambda self, *a, **kw: None):
mg = MemoryGraph.__new__(MemoryGraph)
mg.graph = Mock()
mg.graph.query = Mock(return_value=[])
mg.embedding_model = Mock()
mg.embedding_model.embed = Mock(return_value=[0.1] * 128)
mg.llm = Mock()
mg.node_label = ":Entity"
mg.threshold = 0.7
mg.llm_provider = "openai"
return mg
class TestSoftDelete:
"""Verify _delete_entities uses SET r.valid = false, not DELETE r."""
def test_delete_entities_sends_soft_delete_cypher(self):
mg = _create_graph_memory()
mg.graph.query.return_value = [
{"source": "Alice", "target": "Bob", "relationship": "KNOWS"}
]
mg._delete_entities(
[{"source": "alice", "destination": "bob", "relationship": "KNOWS"}],
{"user_id": "user1"},
)
cypher = mg.graph.query.call_args[0][0]
assert "SET r.valid = false" in cypher
assert "r.invalidated_at = datetime()" in cypher
assert "DELETE r" not in cypher
def test_delete_entities_only_targets_valid_edges(self):
mg = _create_graph_memory()
mg._delete_entities(
[{"source": "alice", "destination": "bob", "relationship": "KNOWS"}],
{"user_id": "user1"},
)
cypher = mg.graph.query.call_args[0][0]
assert "r.valid IS NULL OR r.valid = true" in cypher
def test_delete_entities_is_idempotent(self):
mg = _create_graph_memory()
item = [{"source": "alice", "destination": "bob", "relationship": "KNOWS"}]
filters = {"user_id": "user1"}
mg.graph.query.return_value = [
{"source": "Alice", "target": "Bob", "relationship": "KNOWS"}
]
mg._delete_entities(item, filters)
mg.graph.query.return_value = []
mg._delete_entities(item, filters)
# Both calls should have the same WHERE filter
for c in mg.graph.query.call_args_list:
assert "r.valid IS NULL OR r.valid = true" in c[0][0]
class TestSearchExcludesSoftDeleted:
"""Verify search and get_all filter out soft-deleted relationships."""
def test_get_all_filters_soft_deleted(self):
mg = _create_graph_memory()
mg.get_all(filters={"user_id": "user1"}, limit=10)
cypher = mg.graph.query.call_args[0][0]
assert "r.valid IS NULL OR r.valid = true" in cypher
def test_search_graph_db_filters_both_directions(self):
"""_search_graph_db must filter soft-deleted edges in both outgoing and incoming queries."""
mg = _create_graph_memory()
mg.graph.query.return_value = []
mg._search_graph_db(node_list=["alice"], filters={"user_id": "user1"})
cypher = mg.graph.query.call_args[0][0]
# The UNION query has two MATCH branches — both must filter
occurrences = cypher.count("r.valid IS NULL OR r.valid = true")
assert occurrences >= 2, (
f"_search_graph_db has {occurrences} valid-filter(s) but needs >= 2 "
"(one for outgoing, one for incoming relationships)"
)
def test_delete_all_still_hard_deletes(self):
mg = _create_graph_memory()
mg.delete_all(filters={"user_id": "user1"})
cypher = mg.graph.query.call_args[0][0]
assert "DETACH DELETE" in cypher
class TestMergeResetsValidFlag:
"""Verify MERGE in _add_entities sets r.valid = true.
Critical: after soft-delete, a MERGE that matches the existing
(invalidated) edge must reset valid=true, or the edge becomes
a zombie -- exists but invisible to queries.
"""
def _run_add_entities(self, source_found, dest_found):
"""Helper: call _add_entities with configurable node search results."""
mg = _create_graph_memory()
source_result = (
[{"elementId(source_candidate)": "src_id_1"}] if source_found else []
)
dest_result = (
[{"elementId(destination_candidate)": "dst_id_1"}] if dest_found else []
)
mg._search_source_node = Mock(return_value=source_result)
mg._search_destination_node = Mock(return_value=dest_result)
mg._add_entities(
[{"source": "alice", "destination": "bob", "relationship": "KNOWS"}],
{"user_id": "user1"},
entity_type_map={},
)
cypher = mg.graph.query.call_args[0][0]
return cypher
def test_merge_sets_valid_true_when_source_found(self):
cypher = self._run_add_entities(source_found=True, dest_found=False)
assert "r.valid = true" in cypher
def test_merge_sets_valid_true_when_dest_found(self):
cypher = self._run_add_entities(source_found=False, dest_found=True)
assert "r.valid = true" in cypher
def test_merge_sets_valid_true_when_both_found(self):
cypher = self._run_add_entities(source_found=True, dest_found=True)
assert "r.valid = true" in cypher
def test_merge_sets_valid_true_when_neither_found(self):
cypher = self._run_add_entities(source_found=False, dest_found=False)
assert "r.valid = true" in cypher
def test_merge_clears_invalidated_at_on_resurrection(self):
"""When a soft-deleted edge is resurrected via MERGE, invalidated_at must be cleared.
Without this, a resurrected edge (valid=true) still carries stale
invalidated_at metadata, which corrupts temporal reasoning queries.
"""
for label, src, dst in [
("source found", True, False),
("dest found", False, True),
("both found", True, True),
("neither found", False, False),
]:
cypher = self._run_add_entities(source_found=src, dest_found=dst)
assert "r.invalidated_at = null" in cypher, (
f"MERGE path '{label}': ON MATCH SET does not clear r.invalidated_at. "
"Resurrected edges will have stale invalidation timestamps."
)
class TestCypherConsistency:
"""Verify all MERGE blocks use consistent property names and variable aliases."""
def _get_merge_cypher(self, source_found, dest_found):
mg = _create_graph_memory()
mg._search_source_node = Mock(
return_value=[{"elementId(source_candidate)": "id1"}] if source_found else []
)
mg._search_destination_node = Mock(
return_value=[{"elementId(destination_candidate)": "id2"}] if dest_found else []
)
mg._add_entities(
[{"source": "alice", "destination": "bob", "relationship": "KNOWS"}],
{"user_id": "user1"},
entity_type_map={},
)
return mg.graph.query.call_args[0][0]
def test_all_blocks_use_created_at_not_created(self):
"""All MERGE blocks must use r.created_at, not r.created."""
for label, src, dst in [
("source found", True, False),
("dest found", False, True),
("both found", True, True),
("neither found", False, False),
]:
cypher = self._get_merge_cypher(src, dst)
assert "r.created_at" in cypher, (
f"MERGE path '{label}': uses r.created instead of r.created_at"
)
def test_all_blocks_use_r_not_rel(self):
"""All MERGE blocks must use 'r' as the relationship variable, not 'rel'."""
for label, src, dst in [
("source found", True, False),
("dest found", False, True),
("both found", True, True),
("neither found", False, False),
]:
cypher = self._get_merge_cypher(src, dst)
assert "rel." not in cypher, (
f"MERGE path '{label}': uses 'rel' variable instead of 'r'"
)
def test_all_blocks_set_updated_at_on_create(self):
"""All MERGE blocks must set r.updated_at on CREATE for consistent timestamps."""
for label, src, dst in [
("source found", True, False),
("dest found", False, True),
("both found", True, True),
("neither found", False, False),
]:
cypher = self._get_merge_cypher(src, dst)
assert "r.updated_at = timestamp()" in cypher, (
f"MERGE path '{label}': missing r.updated_at on CREATE SET"
)
class TestSoftDeleteWithFilters:
"""Verify soft-delete works correctly with agent_id and run_id filters."""
def test_delete_entities_with_agent_id(self):
mg = _create_graph_memory()
mg._delete_entities(
[{"source": "alice", "destination": "bob", "relationship": "KNOWS"}],
{"user_id": "user1", "agent_id": "agent1"},
)
cypher = mg.graph.query.call_args[0][0]
params = mg.graph.query.call_args[1]["params"]
assert "SET r.valid = false" in cypher
assert "agent_id: $agent_id" in cypher
assert params["agent_id"] == "agent1"
def test_delete_entities_with_run_id(self):
mg = _create_graph_memory()
mg._delete_entities(
[{"source": "alice", "destination": "bob", "relationship": "KNOWS"}],
{"user_id": "user1", "run_id": "run1"},
)
cypher = mg.graph.query.call_args[0][0]
params = mg.graph.query.call_args[1]["params"]
assert "SET r.valid = false" in cypher
assert "run_id: $run_id" in cypher
assert params["run_id"] == "run1"
def test_get_all_with_agent_id_filters_soft_deleted(self):
mg = _create_graph_memory()
mg.get_all(filters={"user_id": "user1", "agent_id": "agent1"}, limit=10)
cypher = mg.graph.query.call_args[0][0]
assert "r.valid IS NULL OR r.valid = true" in cypher
assert "agent_id: $agent_id" in cypher
def test_merge_with_agent_id_sets_valid_true(self):
mg = _create_graph_memory()
mg._search_source_node = Mock(
return_value=[{"elementId(source_candidate)": "id1"}]
)
mg._search_destination_node = Mock(return_value=[])
mg._add_entities(
[{"source": "alice", "destination": "bob", "relationship": "KNOWS"}],
{"user_id": "user1", "agent_id": "agent1"},
entity_type_map={},
)
cypher = mg.graph.query.call_args[0][0]
assert "r.valid = true" in cypher
assert "agent_id: $agent_id" in cypher
class TestResetAndCleanup:
"""Verify reset and delete_all use hard-delete (DETACH DELETE)."""
def test_reset_uses_detach_delete(self):
mg = _create_graph_memory()
mg.reset()
cypher = mg.graph.query.call_args[0][0]
assert "DETACH DELETE" in cypher
assert "valid" not in cypher.lower()
def test_delete_all_does_not_soft_delete(self):
mg = _create_graph_memory()
mg.delete_all(filters={"user_id": "user1"})
cypher = mg.graph.query.call_args[0][0]
assert "DETACH DELETE" in cypher
assert "r.valid = false" not in cypher
+195
View File
@@ -188,6 +188,201 @@ async def test_async_update_memory_uses_utc_timestamps(mocker):
_assert_utc_timestamp(payload["updated_at"])
class TestMetadataNotMutated:
"""Tests that metadata dicts passed to memory methods are not mutated in-place (issue #2648)."""
def test_create_memory_does_not_mutate_metadata(self, mocker):
memory = _build_memory_instance(mocker, Memory)
original_metadata = {"user_id": "test_user", "category": "sports"}
metadata_copy = original_metadata.copy()
memory._create_memory("test data", {"test data": [0.1, 0.2, 0.3]}, metadata=original_metadata)
assert original_metadata == metadata_copy, (
f"_create_memory mutated the caller's metadata dict: {original_metadata} != {metadata_copy}"
)
def test_create_memory_stores_correct_payload(self, mocker):
memory = _build_memory_instance(mocker, Memory)
metadata = {"user_id": "test_user", "category": "sports"}
memory._create_memory("test data", {"test data": [0.1, 0.2, 0.3]}, metadata=metadata)
payload = memory.vector_store.insert.call_args.kwargs["payloads"][0]
assert payload["data"] == "test data"
assert payload["user_id"] == "test_user"
assert payload["category"] == "sports"
assert "hash" in payload
assert "created_at" in payload
def test_create_memory_with_none_metadata(self, mocker):
memory = _build_memory_instance(mocker, Memory)
memory._create_memory("test data", {"test data": [0.1, 0.2, 0.3]}, metadata=None)
payload = memory.vector_store.insert.call_args.kwargs["payloads"][0]
assert payload["data"] == "test data"
assert "hash" in payload
def test_create_memory_shared_metadata_across_calls(self, mocker):
"""Verify that sharing a metadata dict between multiple _create_memory calls is safe."""
memory = _build_memory_instance(mocker, Memory)
shared_metadata = {"user_id": "test_user"}
memory._create_memory("first memory", {"first memory": [0.1, 0.2, 0.3]}, metadata=shared_metadata)
memory._create_memory("second memory", {"second memory": [0.4, 0.5, 0.6]}, metadata=shared_metadata)
assert shared_metadata == {"user_id": "test_user"}, "shared metadata was mutated across calls"
# Verify each call got the correct data
first_payload = memory.vector_store.insert.call_args_list[0].kwargs["payloads"][0]
second_payload = memory.vector_store.insert.call_args_list[1].kwargs["payloads"][0]
assert first_payload["data"] == "first memory"
assert second_payload["data"] == "second memory"
def test_create_memory_preserves_role_and_actor_id_in_history(self, mocker):
"""Verify that role and actor_id from metadata flow through to add_history after deepcopy."""
memory = _build_memory_instance(mocker, Memory)
metadata = {"user_id": "test_user", "role": "assistant", "actor_id": "bot-1"}
memory._create_memory("test data", {"test data": [0.1, 0.2, 0.3]}, metadata=metadata)
# Verify the payload stored in vector store has all fields
payload = memory.vector_store.insert.call_args.kwargs["payloads"][0]
assert payload["role"] == "assistant"
assert payload["actor_id"] == "bot-1"
assert payload["user_id"] == "test_user"
assert payload["data"] == "test data"
# Verify add_history received the correct role and actor_id
history_call = memory.db.add_history.call_args
assert history_call.kwargs["role"] == "assistant"
assert history_call.kwargs["actor_id"] == "bot-1"
# And the original metadata is still untouched
assert metadata == {"user_id": "test_user", "role": "assistant", "actor_id": "bot-1"}
def test_create_memory_with_nested_metadata_not_mutated(self, mocker):
"""Verify deepcopy protects nested structures in metadata."""
memory = _build_memory_instance(mocker, Memory)
metadata = {"user_id": "test_user", "tags": ["important", "urgent"], "config": {"key": "val"}}
import copy
metadata_snapshot = copy.deepcopy(metadata)
memory._create_memory("test data", {"test data": [0.1, 0.2, 0.3]}, metadata=metadata)
assert metadata == metadata_snapshot, "Nested metadata structures were mutated"
def test_update_memory_does_not_mutate_metadata(self, mocker):
memory = _build_memory_instance(mocker, Memory)
memory.vector_store.get.return_value = MagicMock(
payload={"data": "old data", "user_id": "test_user", "created_at": "2026-01-01T00:00:00+00:00"}
)
original_metadata = {"category": "updated"}
metadata_copy = original_metadata.copy()
memory._update_memory("mem-id", "new data", {"new data": [0.1, 0.2, 0.3]}, metadata=original_metadata)
assert original_metadata == metadata_copy, (
f"_update_memory mutated the caller's metadata dict: {original_metadata} != {metadata_copy}"
)
def test_add_to_vector_store_no_infer_does_not_mutate_metadata(self, mocker):
"""Verify _add_to_vector_store with infer=False doesn't leak metadata between messages."""
memory = _build_memory_instance(mocker, Memory)
memory.embedding_model.embed.return_value = [0.1, 0.2, 0.3]
original_metadata = {"user_id": "test_user"}
metadata_copy = original_metadata.copy()
messages = [
{"role": "user", "content": "Hello"},
{"role": "assistant", "content": "Hi there", "name": "bot-1"},
]
result = memory._add_to_vector_store(messages, original_metadata, filters={}, infer=False)
# Metadata should not be mutated
assert original_metadata == metadata_copy, (
f"_add_to_vector_store mutated the caller's metadata: {original_metadata}"
)
# Should have created 2 memories
assert len(result) == 2
assert result[0]["role"] == "user"
assert result[1]["role"] == "assistant"
assert result[1]["actor_id"] == "bot-1"
# Verify each insert got distinct payloads with correct roles
insert_calls = memory.vector_store.insert.call_args_list
first_payload = insert_calls[0].kwargs["payloads"][0]
second_payload = insert_calls[1].kwargs["payloads"][0]
assert first_payload["role"] == "user"
assert "actor_id" not in first_payload # user message has no name
assert second_payload["role"] == "assistant"
assert second_payload["actor_id"] == "bot-1"
@pytest.mark.asyncio
async def test_async_create_memory_does_not_mutate_metadata(self, mocker):
memory = _build_memory_instance(mocker, AsyncMemory)
original_metadata = {"user_id": "test_user", "category": "sports"}
metadata_copy = original_metadata.copy()
await memory._create_memory("test data", {"test data": [0.1, 0.2, 0.3]}, metadata=original_metadata)
assert original_metadata == metadata_copy, (
f"async _create_memory mutated the caller's metadata dict: {original_metadata} != {metadata_copy}"
)
@pytest.mark.asyncio
async def test_async_create_memory_shared_metadata_across_calls(self, mocker):
memory = _build_memory_instance(mocker, AsyncMemory)
shared_metadata = {"user_id": "test_user"}
await memory._create_memory("first memory", {"first memory": [0.1, 0.2, 0.3]}, metadata=shared_metadata)
await memory._create_memory("second memory", {"second memory": [0.4, 0.5, 0.6]}, metadata=shared_metadata)
assert shared_metadata == {"user_id": "test_user"}, "shared metadata was mutated across async calls"
@pytest.mark.asyncio
async def test_async_add_to_vector_store_no_infer_does_not_mutate_metadata(self, mocker):
"""Verify async _add_to_vector_store with infer=False doesn't leak metadata between messages."""
memory = _build_memory_instance(mocker, AsyncMemory)
memory.embedding_model.embed.return_value = [0.1, 0.2, 0.3]
original_metadata = {"user_id": "test_user"}
metadata_copy = original_metadata.copy()
messages = [
{"role": "user", "content": "Hello"},
{"role": "assistant", "content": "Hi there", "name": "bot-1"},
]
result = await memory._add_to_vector_store(messages, original_metadata, effective_filters={}, infer=False)
assert original_metadata == metadata_copy, (
f"async _add_to_vector_store mutated the caller's metadata: {original_metadata}"
)
assert len(result) == 2
assert result[0]["role"] == "user"
assert result[1]["role"] == "assistant"
@pytest.mark.asyncio
async def test_async_update_memory_does_not_mutate_metadata(self, mocker):
memory = _build_memory_instance(mocker, AsyncMemory)
memory.vector_store.get.return_value = MagicMock(
payload={"data": "old data", "user_id": "test_user", "created_at": "2026-01-01T00:00:00+00:00"}
)
original_metadata = {"category": "updated"}
metadata_copy = original_metadata.copy()
await memory._update_memory("mem-id", "new data", {"new data": [0.1, 0.2, 0.3]}, metadata=original_metadata)
assert original_metadata == metadata_copy, (
f"async _update_memory mutated the caller's metadata dict: {original_metadata} != {metadata_copy}"
)
def test_normalize_iso_timestamp_to_utc_preserves_naive_values():
assert _normalize_iso_timestamp_to_utc("2026-03-18T00:00:00") == "2026-03-18T00:00:00"
+517
View File
@@ -0,0 +1,517 @@
"""Tests for graph cleanup on memory deletion (issue #3245)."""
from unittest.mock import MagicMock, patch
import pytest
from mem0.configs.base import MemoryConfig
class MockVectorMemory:
def __init__(self, memory_id, payload, score=0.8):
self.id = memory_id
self.payload = payload
self.score = score
@patch("mem0.utils.factory.EmbedderFactory.create")
@patch("mem0.utils.factory.VectorStoreFactory.create")
@patch("mem0.utils.factory.LlmFactory.create")
@patch("mem0.memory.storage.SQLiteManager")
def test_delete_calls_graph_cleanup_when_graph_enabled(
mock_sqlite, mock_llm_factory, mock_vector_factory, mock_embedder_factory
):
"""When graph is enabled, delete() should call graph.delete() with memory text and filters."""
mock_embedder_factory.return_value = MagicMock()
mock_vector_store = MagicMock()
mock_vector_factory.return_value = mock_vector_store
mock_llm_factory.return_value = MagicMock()
mock_sqlite.return_value = MagicMock()
from mem0.memory.main import Memory
config = MemoryConfig()
memory = Memory(config)
# Enable graph with a mock
memory.enable_graph = True
memory.graph = MagicMock()
# Set up vector store to return a memory with graph-relevant data
mock_vector_store.get.return_value = MockVectorMemory(
"mem-1",
{
"data": "Alice likes Bob",
"user_id": "user-1",
"agent_id": "agent-1",
"hash": "abc",
},
)
memory.delete("mem-1")
# graph.delete should have been called with the memory text and filters
memory.graph.delete.assert_called_once_with(
"Alice likes Bob", {"user_id": "user-1", "agent_id": "agent-1"}
)
# _delete_memory should still have been called (vector store + history cleanup)
mock_vector_store.delete.assert_called_once_with(vector_id="mem-1")
@patch("mem0.utils.factory.EmbedderFactory.create")
@patch("mem0.utils.factory.VectorStoreFactory.create")
@patch("mem0.utils.factory.LlmFactory.create")
@patch("mem0.memory.storage.SQLiteManager")
def test_delete_skips_graph_when_not_enabled(
mock_sqlite, mock_llm_factory, mock_vector_factory, mock_embedder_factory
):
"""When graph is not enabled, delete() should not attempt graph cleanup."""
mock_embedder_factory.return_value = MagicMock()
mock_vector_store = MagicMock()
mock_vector_factory.return_value = mock_vector_store
mock_llm_factory.return_value = MagicMock()
mock_sqlite.return_value = MagicMock()
from mem0.memory.main import Memory
config = MemoryConfig()
memory = Memory(config)
assert memory.enable_graph is False
mock_vector_store.get.return_value = MockVectorMemory(
"mem-1", {"data": "Alice likes Bob", "user_id": "user-1", "hash": "abc"}
)
result = memory.delete("mem-1")
assert result == {"message": "Memory deleted successfully!"}
mock_vector_store.delete.assert_called_once_with(vector_id="mem-1")
@patch("mem0.utils.factory.EmbedderFactory.create")
@patch("mem0.utils.factory.VectorStoreFactory.create")
@patch("mem0.utils.factory.LlmFactory.create")
@patch("mem0.memory.storage.SQLiteManager")
def test_delete_continues_if_graph_cleanup_fails(
mock_sqlite, mock_llm_factory, mock_vector_factory, mock_embedder_factory
):
"""If graph cleanup raises an exception, delete() should still succeed."""
mock_embedder_factory.return_value = MagicMock()
mock_vector_store = MagicMock()
mock_vector_factory.return_value = mock_vector_store
mock_llm_factory.return_value = MagicMock()
mock_sqlite.return_value = MagicMock()
from mem0.memory.main import Memory
config = MemoryConfig()
memory = Memory(config)
memory.enable_graph = True
memory.graph = MagicMock()
memory.graph.delete.side_effect = RuntimeError("Neo4j connection lost")
mock_vector_store.get.return_value = MockVectorMemory(
"mem-1", {"data": "Alice likes Bob", "user_id": "user-1", "hash": "abc"}
)
# Should not raise
result = memory.delete("mem-1")
assert result == {"message": "Memory deleted successfully!"}
# Vector store deletion should still proceed
mock_vector_store.delete.assert_called_once_with(vector_id="mem-1")
@patch("mem0.utils.factory.EmbedderFactory.create")
@patch("mem0.utils.factory.VectorStoreFactory.create")
@patch("mem0.utils.factory.LlmFactory.create")
@patch("mem0.memory.storage.SQLiteManager")
def test_delete_skips_graph_when_no_user_id(
mock_sqlite, mock_llm_factory, mock_vector_factory, mock_embedder_factory
):
"""Graph cleanup should be skipped if the memory has no user_id."""
mock_embedder_factory.return_value = MagicMock()
mock_vector_store = MagicMock()
mock_vector_factory.return_value = mock_vector_store
mock_llm_factory.return_value = MagicMock()
mock_sqlite.return_value = MagicMock()
from mem0.memory.main import Memory
config = MemoryConfig()
memory = Memory(config)
memory.enable_graph = True
memory.graph = MagicMock()
# Memory with no user_id
mock_vector_store.get.return_value = MockVectorMemory(
"mem-1", {"data": "Some data", "hash": "abc"}
)
memory.delete("mem-1")
# graph.delete should NOT have been called since there's no user_id
memory.graph.delete.assert_not_called()
mock_vector_store.delete.assert_called_once_with(vector_id="mem-1")
@patch("mem0.utils.factory.EmbedderFactory.create")
@patch("mem0.utils.factory.VectorStoreFactory.create")
@patch("mem0.utils.factory.LlmFactory.create")
@patch("mem0.memory.storage.SQLiteManager")
def test_delete_skips_graph_when_no_memory_text(
mock_sqlite, mock_llm_factory, mock_vector_factory, mock_embedder_factory
):
"""Graph cleanup should be skipped if the memory has no text data."""
mock_embedder_factory.return_value = MagicMock()
mock_vector_store = MagicMock()
mock_vector_factory.return_value = mock_vector_store
mock_llm_factory.return_value = MagicMock()
mock_sqlite.return_value = MagicMock()
from mem0.memory.main import Memory
config = MemoryConfig()
memory = Memory(config)
memory.enable_graph = True
memory.graph = MagicMock()
mock_vector_store.get.return_value = MockVectorMemory(
"mem-1", {"user_id": "user-1", "hash": "abc"}
)
memory.delete("mem-1")
memory.graph.delete.assert_not_called()
mock_vector_store.delete.assert_called_once_with(vector_id="mem-1")
@patch("mem0.utils.factory.EmbedderFactory.create")
@patch("mem0.utils.factory.VectorStoreFactory.create")
@patch("mem0.utils.factory.LlmFactory.create")
@patch("mem0.memory.storage.SQLiteManager")
def test_delete_passes_all_filters_to_graph(
mock_sqlite, mock_llm_factory, mock_vector_factory, mock_embedder_factory
):
"""Graph cleanup should include all available filters (user_id, agent_id, run_id)."""
mock_embedder_factory.return_value = MagicMock()
mock_vector_store = MagicMock()
mock_vector_factory.return_value = mock_vector_store
mock_llm_factory.return_value = MagicMock()
mock_sqlite.return_value = MagicMock()
from mem0.memory.main import Memory
config = MemoryConfig()
memory = Memory(config)
memory.enable_graph = True
memory.graph = MagicMock()
mock_vector_store.get.return_value = MockVectorMemory(
"mem-1",
{
"data": "Alice likes Bob",
"user_id": "user-1",
"agent_id": "agent-1",
"run_id": "run-1",
"hash": "abc",
},
)
memory.delete("mem-1")
memory.graph.delete.assert_called_once_with(
"Alice likes Bob",
{"user_id": "user-1", "agent_id": "agent-1", "run_id": "run-1"},
)
@pytest.mark.asyncio
@patch("mem0.utils.factory.EmbedderFactory.create")
@patch("mem0.utils.factory.VectorStoreFactory.create")
@patch("mem0.utils.factory.LlmFactory.create")
@patch("mem0.memory.storage.SQLiteManager")
async def test_async_delete_calls_graph_cleanup(
mock_sqlite, mock_llm_factory, mock_vector_factory, mock_embedder_factory
):
"""Async delete() should also perform graph cleanup."""
mock_embedder_factory.return_value = MagicMock()
mock_vector_store = MagicMock()
mock_vector_factory.return_value = mock_vector_store
mock_llm_factory.return_value = MagicMock()
mock_sqlite.return_value = MagicMock()
from mem0.memory.main import AsyncMemory
config = MemoryConfig()
memory = AsyncMemory(config)
memory.enable_graph = True
memory.graph = MagicMock()
mock_vector_store.get.return_value = MockVectorMemory(
"mem-1",
{
"data": "Alice likes Bob",
"user_id": "user-1",
"hash": "abc",
},
)
result = await memory.delete("mem-1")
assert result == {"message": "Memory deleted successfully!"}
memory.graph.delete.assert_called_once_with("Alice likes Bob", {"user_id": "user-1"})
mock_vector_store.delete.assert_called_once_with(vector_id="mem-1")
@pytest.mark.asyncio
@patch("mem0.utils.factory.EmbedderFactory.create")
@patch("mem0.utils.factory.VectorStoreFactory.create")
@patch("mem0.utils.factory.LlmFactory.create")
@patch("mem0.memory.storage.SQLiteManager")
async def test_async_delete_continues_if_graph_cleanup_fails(
mock_sqlite, mock_llm_factory, mock_vector_factory, mock_embedder_factory
):
"""Async delete() should continue even if graph cleanup fails."""
mock_embedder_factory.return_value = MagicMock()
mock_vector_store = MagicMock()
mock_vector_factory.return_value = mock_vector_store
mock_llm_factory.return_value = MagicMock()
mock_sqlite.return_value = MagicMock()
from mem0.memory.main import AsyncMemory
config = MemoryConfig()
memory = AsyncMemory(config)
memory.enable_graph = True
memory.graph = MagicMock()
memory.graph.delete.side_effect = RuntimeError("Graph error")
mock_vector_store.get.return_value = MockVectorMemory(
"mem-1", {"data": "Alice likes Bob", "user_id": "user-1", "hash": "abc"}
)
result = await memory.delete("mem-1")
assert result == {"message": "Memory deleted successfully!"}
mock_vector_store.delete.assert_called_once_with(vector_id="mem-1")
@patch("mem0.utils.factory.EmbedderFactory.create")
@patch("mem0.utils.factory.VectorStoreFactory.create")
@patch("mem0.utils.factory.LlmFactory.create")
@patch("mem0.memory.storage.SQLiteManager")
def test_delete_raises_for_nonexistent_memory_with_graph_enabled(
mock_sqlite, mock_llm_factory, mock_vector_factory, mock_embedder_factory
):
"""delete() should raise ValueError for non-existent memory even with graph enabled."""
mock_embedder_factory.return_value = MagicMock()
mock_vector_store = MagicMock()
mock_vector_factory.return_value = mock_vector_store
mock_llm_factory.return_value = MagicMock()
mock_sqlite.return_value = MagicMock()
from mem0.memory.main import Memory
config = MemoryConfig()
memory = Memory(config)
memory.enable_graph = True
memory.graph = MagicMock()
mock_vector_store.get.return_value = None
with pytest.raises(ValueError, match="Memory with id non-existent not found"):
memory.delete("non-existent")
memory.graph.delete.assert_not_called()
mock_vector_store.delete.assert_not_called()
@pytest.mark.asyncio
@patch("mem0.utils.factory.EmbedderFactory.create")
@patch("mem0.utils.factory.VectorStoreFactory.create")
@patch("mem0.utils.factory.LlmFactory.create")
@patch("mem0.memory.storage.SQLiteManager")
async def test_async_delete_raises_for_nonexistent_memory_with_graph_enabled(
mock_sqlite, mock_llm_factory, mock_vector_factory, mock_embedder_factory
):
"""Async delete() should raise ValueError for non-existent memory even with graph enabled."""
mock_embedder_factory.return_value = MagicMock()
mock_vector_store = MagicMock()
mock_vector_store.get.return_value = None
mock_vector_factory.return_value = mock_vector_store
mock_llm_factory.return_value = MagicMock()
mock_sqlite.return_value = MagicMock()
from mem0.memory.main import AsyncMemory
config = MemoryConfig()
memory = AsyncMemory(config)
memory.enable_graph = True
memory.graph = MagicMock()
with pytest.raises(ValueError, match="Memory with id non-existent not found"):
await memory.delete("non-existent")
memory.graph.delete.assert_not_called()
mock_vector_store.delete.assert_not_called()
@patch("mem0.utils.factory.EmbedderFactory.create")
@patch("mem0.utils.factory.VectorStoreFactory.create")
@patch("mem0.utils.factory.LlmFactory.create")
@patch("mem0.memory.storage.SQLiteManager")
def test_delete_all_does_not_trigger_per_memory_graph_cleanup(
mock_sqlite, mock_llm_factory, mock_vector_factory, mock_embedder_factory
):
"""delete_all() should use graph.delete_all(), not per-memory graph.delete()."""
mock_embedder_factory.return_value = MagicMock()
mock_vector_store = MagicMock()
mock_vector_factory.return_value = mock_vector_store
mock_llm_factory.return_value = MagicMock()
mock_sqlite.return_value = MagicMock()
from mem0.memory.main import Memory
config = MemoryConfig()
memory = Memory(config)
memory.enable_graph = True
memory.graph = MagicMock()
mem1 = MockVectorMemory("mem-1", {"data": "Alice likes Bob", "user_id": "user-1"})
mem2 = MockVectorMemory("mem-2", {"data": "Bob likes Charlie", "user_id": "user-1"})
mock_vector_store.list.return_value = ([mem1, mem2], 2)
mock_vector_store.get.return_value = MockVectorMemory(
"mem-1", {"data": "Alice likes Bob", "user_id": "user-1"}
)
memory.delete_all(user_id="user-1")
# graph.delete (per-memory) should NOT be called
memory.graph.delete.assert_not_called()
# graph.delete_all (bulk) SHOULD be called
memory.graph.delete_all.assert_called_once_with({"user_id": "user-1"})
@patch("mem0.utils.factory.EmbedderFactory.create")
@patch("mem0.utils.factory.VectorStoreFactory.create")
@patch("mem0.utils.factory.LlmFactory.create")
@patch("mem0.memory.storage.SQLiteManager")
def test_internal_delete_memory_does_not_trigger_graph_cleanup(
mock_sqlite, mock_llm_factory, mock_vector_factory, mock_embedder_factory
):
"""_delete_memory() should NOT call graph.delete() — only the public delete() does.
This ensures that the DELETE branch inside _add_to_vector_store() (which calls
_delete_memory directly) does not interfere with the parallel graph pipeline
running in _add_to_graph().
"""
mock_embedder_factory.return_value = MagicMock()
mock_vector_store = MagicMock()
mock_vector_factory.return_value = mock_vector_store
mock_llm_factory.return_value = MagicMock()
mock_sqlite.return_value = MagicMock()
from mem0.memory.main import Memory
config = MemoryConfig()
memory = Memory(config)
memory.enable_graph = True
memory.graph = MagicMock()
mock_vector_store.get.return_value = MockVectorMemory(
"mem-1", {"data": "Alice likes Bob", "user_id": "user-1", "hash": "abc"}
)
# Call _delete_memory directly (as _add_to_vector_store does for DELETE events)
memory._delete_memory("mem-1")
# graph.delete should NOT have been called — graph cleanup is only in delete()
memory.graph.delete.assert_not_called()
# But vector store deletion should proceed
mock_vector_store.delete.assert_called_once_with(vector_id="mem-1")
def test_graph_memory_delete_calls_internal_methods():
"""Test that MemoryGraph.delete() calls the expected internal pipeline methods."""
from unittest.mock import patch as _patch
# We need to mock the Neo4j import
with _patch.dict("sys.modules", {"langchain_neo4j": MagicMock(), "rank_bm25": MagicMock()}):
from mem0.memory.graph_memory import MemoryGraph
with _patch.object(MemoryGraph, "__init__", return_value=None):
graph = MemoryGraph.__new__(MemoryGraph)
# Mock the internal methods
graph._retrieve_nodes_from_data = MagicMock(
return_value={"alice": "person", "bob": "person"}
)
graph._establish_nodes_relations_from_data = MagicMock(
return_value=[
{"source": "alice", "destination": "bob", "relationship": "likes"}
]
)
graph._delete_entities = MagicMock(return_value=[])
filters = {"user_id": "user-1"}
graph.delete("Alice likes Bob", filters)
graph._retrieve_nodes_from_data.assert_called_once_with("Alice likes Bob", filters)
graph._establish_nodes_relations_from_data.assert_called_once_with(
"Alice likes Bob", filters, {"alice": "person", "bob": "person"}
)
graph._delete_entities.assert_called_once_with(
[{"source": "alice", "destination": "bob", "relationship": "likes"}],
filters,
)
def test_graph_memory_delete_skips_when_no_entities():
"""Test that MemoryGraph.delete() does nothing when no entities are extracted."""
from unittest.mock import patch as _patch
with _patch.dict("sys.modules", {"langchain_neo4j": MagicMock(), "rank_bm25": MagicMock()}):
from mem0.memory.graph_memory import MemoryGraph
with _patch.object(MemoryGraph, "__init__", return_value=None):
graph = MemoryGraph.__new__(MemoryGraph)
graph._retrieve_nodes_from_data = MagicMock(return_value={})
graph._establish_nodes_relations_from_data = MagicMock()
graph._delete_entities = MagicMock()
graph.delete("Some text", {"user_id": "user-1"})
graph._retrieve_nodes_from_data.assert_called_once()
graph._establish_nodes_relations_from_data.assert_not_called()
graph._delete_entities.assert_not_called()
def test_graph_memory_delete_handles_exception():
"""Test that MemoryGraph.delete() catches exceptions without raising."""
from unittest.mock import patch as _patch
with _patch.dict("sys.modules", {"langchain_neo4j": MagicMock(), "rank_bm25": MagicMock()}):
from mem0.memory.graph_memory import MemoryGraph
with _patch.object(MemoryGraph, "__init__", return_value=None):
graph = MemoryGraph.__new__(MemoryGraph)
graph._retrieve_nodes_from_data = MagicMock(
side_effect=RuntimeError("LLM error")
)
# Should not raise
graph.delete("Some text", {"user_id": "user-1"})
+821
View File
@@ -0,0 +1,821 @@
"""
End-to-end tests for graph cleanup on memory deletion against real
Neo4j, Memgraph, and Apache AGE instances running in Docker.
Requires:
docker run -d --name mem0-neo4j-test -p 7687:7687 -e NEO4J_AUTH=neo4j/testpassword neo4j:5.23
docker run -d --name mem0-memgraph-test -p 7688:7687 memgraph/memgraph:latest
docker run -d --name mem0-age-test -p 5432:5432 -e POSTGRES_USER=postgres -e POSTGRES_PASSWORD=testpassword -e POSTGRES_DB=testdb apache/age:latest
Tests are skipped automatically if the databases or required Python
packages are not available.
"""
import hashlib
import sys
import warnings
from unittest.mock import MagicMock
import pytest
warnings.filterwarnings("ignore")
EMBEDDING_DIMS = 64
# ---------------------------------------------------------------------------
# Deterministic embedding helper (shared across backends)
# ---------------------------------------------------------------------------
def _make_deterministic_embedder():
cache = {}
counter = [0]
def embed(text, *args, **kwargs):
t = text.lower().strip()
if t not in cache:
vec = [0.0] * EMBEDDING_DIMS
idx = counter[0] % EMBEDDING_DIMS
vec[idx] = 1.0
h = hashlib.sha256(t.encode()).digest()
for i in range(EMBEDDING_DIMS):
vec[i] += float(h[i % len(h)]) / 25500.0
norm = sum(v * v for v in vec) ** 0.5
cache[t] = [v / norm for v in vec]
counter[0] += 1
return cache[t]
mock = MagicMock()
mock.embed.side_effect = embed
mock.config.embedding_dims = EMBEDDING_DIMS
return mock
def _make_mock_llm(entities, relations):
"""Create an LLM mock that returns specific entities and relations."""
mock = MagicMock()
def generate_response(messages, tools):
tool_names = []
for t in tools:
if isinstance(t, dict):
fn = t.get("function", t)
tool_names.append(fn.get("name", ""))
else:
tool_names.append(getattr(t, "name", str(t)))
if any("extract_entities" in n for n in tool_names):
return {
"tool_calls": [
{"name": "extract_entities", "arguments": {"entities": entities}}
]
}
elif any("establish" in n or "relation" in n for n in tool_names):
return {
"tool_calls": [
{"name": "establish_nodes_relations", "arguments": {"entities": relations}}
]
}
elif any("delete" in n for n in tool_names):
return {"tool_calls": []}
return {"tool_calls": []}
mock.generate_response.side_effect = generate_response
return mock
# ===========================================================================
# NEO4J
# ===========================================================================
def _port_open(host, port, timeout=1):
"""Quick TCP check — avoids slow driver-level timeouts."""
import socket
try:
with socket.create_connection((host, port), timeout=timeout):
return True
except OSError:
return False
def _neo4j_available():
if not _port_open("localhost", 7687):
return False
try:
from langchain_neo4j import Neo4jGraph
g = Neo4jGraph(
url="bolt://localhost:7687",
username="neo4j",
password="testpassword",
refresh_schema=False,
driver_config={"notifications_min_severity": "OFF"},
)
g.query("RETURN 1")
return True
except Exception:
return False
requires_neo4j = pytest.mark.skipif(not _neo4j_available(), reason="Neo4j not available")
@pytest.fixture
def neo4j_graph():
"""Create a Neo4j-backed MemoryGraph with mocked LLM/embedder."""
from langchain_neo4j import Neo4jGraph
from mem0.memory.graph_memory import MemoryGraph
mg = MemoryGraph.__new__(MemoryGraph)
mg.graph = Neo4jGraph(
url="bolt://localhost:7687",
username="neo4j",
password="testpassword",
refresh_schema=False,
driver_config={"notifications_min_severity": "OFF"},
)
mg.graph.query("MATCH (n) DETACH DELETE n")
mg.node_label = ":`__Entity__`"
mg.llm_provider = "openai"
mg.user_id = None
mg.threshold = 0.99
mg.embedding_model = _make_deterministic_embedder()
mg.llm = MagicMock()
mg.config = MagicMock()
mg.config.graph_store.custom_prompt = None
mg.config.graph_store.config.base_label = True
yield mg
mg.graph.query("MATCH (n) DETACH DELETE n")
@requires_neo4j
class TestNeo4jDeleteE2E:
def _node_count(self, mg):
return mg.graph.query("MATCH (n) RETURN count(n) AS cnt")[0]["cnt"]
def _valid_edge_count(self, mg):
return mg.graph.query(
"MATCH ()-[r]->() WHERE r.valid IS NULL OR r.valid = true RETURN count(r) AS cnt"
)[0]["cnt"]
def _invalid_edge_count(self, mg):
return mg.graph.query(
"MATCH ()-[r]->() WHERE r.valid = false RETURN count(r) AS cnt"
)[0]["cnt"]
def test_add_creates_graph_data(self, neo4j_graph):
mg = neo4j_graph
mg.llm = _make_mock_llm(
[{"entity": "Alice", "entity_type": "person"}, {"entity": "Bob", "entity_type": "person"}],
[{"source": "Alice", "destination": "Bob", "relationship": "likes"}],
)
mg.add("Alice likes Bob", {"user_id": "u1"})
assert self._node_count(mg) == 2
assert self._valid_edge_count(mg) == 1
def test_delete_soft_deletes_relationships(self, neo4j_graph):
"""Neo4j delete() should set r.valid=false (soft-delete), not hard-delete."""
mg = neo4j_graph
mg.llm = _make_mock_llm(
[{"entity": "Alice", "entity_type": "person"}, {"entity": "Bob", "entity_type": "person"}],
[{"source": "Alice", "destination": "Bob", "relationship": "likes"}],
)
mg.add("Alice likes Bob", {"user_id": "u1"})
assert self._valid_edge_count(mg) == 1
assert self._invalid_edge_count(mg) == 0
mg.delete("Alice likes Bob", {"user_id": "u1"})
assert self._valid_edge_count(mg) == 0
assert self._invalid_edge_count(mg) == 1 # soft-deleted, not removed
assert self._node_count(mg) == 2 # nodes preserved
def test_delete_preserves_other_relationships(self, neo4j_graph):
mg = neo4j_graph
mg.llm = _make_mock_llm(
[{"entity": "Alice", "entity_type": "person"}, {"entity": "Bob", "entity_type": "person"}],
[{"source": "Alice", "destination": "Bob", "relationship": "likes"}],
)
mg.add("Alice likes Bob", {"user_id": "u1"})
mg.llm = _make_mock_llm(
[{"entity": "Alice", "entity_type": "person"}, {"entity": "Charlie", "entity_type": "person"}],
[{"source": "Alice", "destination": "Charlie", "relationship": "knows"}],
)
mg.add("Alice knows Charlie", {"user_id": "u1"})
assert self._valid_edge_count(mg) == 2
mg.llm = _make_mock_llm(
[{"entity": "Alice", "entity_type": "person"}, {"entity": "Bob", "entity_type": "person"}],
[{"source": "Alice", "destination": "Bob", "relationship": "likes"}],
)
mg.delete("Alice likes Bob", {"user_id": "u1"})
assert self._valid_edge_count(mg) == 1
assert self._invalid_edge_count(mg) == 1
def test_delete_user_isolation(self, neo4j_graph):
mg = neo4j_graph
mg.llm = _make_mock_llm(
[{"entity": "Alice", "entity_type": "person"}, {"entity": "Bob", "entity_type": "person"}],
[{"source": "Alice", "destination": "Bob", "relationship": "likes"}],
)
mg.add("Alice likes Bob", {"user_id": "u1"})
mg.add("Alice likes Bob", {"user_id": "u2"})
assert self._valid_edge_count(mg) == 2
mg.delete("Alice likes Bob", {"user_id": "u1"})
assert self._valid_edge_count(mg) == 1
def test_delete_all_hard_deletes(self, neo4j_graph):
mg = neo4j_graph
mg.llm = _make_mock_llm(
[{"entity": "Alice", "entity_type": "person"}, {"entity": "Bob", "entity_type": "person"}],
[{"source": "Alice", "destination": "Bob", "relationship": "likes"}],
)
mg.add("Alice likes Bob", {"user_id": "u1"})
assert self._node_count(mg) == 2
mg.delete_all({"user_id": "u1"})
assert self._node_count(mg) == 0
def test_add_delete_add_cycle(self, neo4j_graph):
mg = neo4j_graph
mg.llm = _make_mock_llm(
[{"entity": "Alice", "entity_type": "person"}, {"entity": "Bob", "entity_type": "person"}],
[{"source": "Alice", "destination": "Bob", "relationship": "likes"}],
)
mg.add("Alice likes Bob", {"user_id": "u1"})
assert self._valid_edge_count(mg) == 1
mg.delete("Alice likes Bob", {"user_id": "u1"})
assert self._valid_edge_count(mg) == 0
mg.add("Alice likes Bob", {"user_id": "u1"})
assert self._valid_edge_count(mg) == 1
# ===========================================================================
# MEMGRAPH
# ===========================================================================
def _memgraph_available():
if not _port_open("localhost", 7688):
return False
try:
from langchain_memgraph.graphs.memgraph import Memgraph
g = Memgraph("bolt://localhost:7688", "memgraph", "memgraph")
g.query("RETURN 1")
return True
except Exception:
return False
requires_memgraph = pytest.mark.skipif(
not _memgraph_available(), reason="Memgraph not available"
)
@pytest.fixture
def memgraph_graph():
"""Create a Memgraph-backed MemoryGraph with mocked LLM/embedder."""
from langchain_memgraph.graphs.memgraph import Memgraph
from mem0.memory.memgraph_memory import MemoryGraph
mg = MemoryGraph.__new__(MemoryGraph)
mg.graph = Memgraph("bolt://localhost:7688", "memgraph", "memgraph")
mg.graph.query("MATCH (n) DETACH DELETE n")
try:
mg.graph.query("DROP VECTOR INDEX memzero;")
except Exception:
pass
mg.graph.query(
f"CREATE VECTOR INDEX memzero ON :Entity(embedding) "
f"WITH CONFIG {{'dimension': {EMBEDDING_DIMS}, 'capacity': 1000, 'metric': 'cos'}};"
)
try:
mg.graph.query("CREATE INDEX ON :Entity(user_id);")
except Exception:
pass
try:
mg.graph.query("CREATE INDEX ON :Entity;")
except Exception:
pass
mg.llm_provider = "openai"
mg.user_id = None
mg.threshold = 0.99
mg.embedding_model = _make_deterministic_embedder()
mg.llm = MagicMock()
mg.config = MagicMock()
mg.config.graph_store.custom_prompt = None
mg.config.embedder.config = {"embedding_dims": EMBEDDING_DIMS}
yield mg
mg.graph.query("MATCH (n) DETACH DELETE n")
@requires_memgraph
class TestMemgraphDeleteE2E:
def _node_count(self, mg):
return mg.graph.query("MATCH (n:Entity) RETURN count(n) AS cnt")[0]["cnt"]
def _edge_count(self, mg):
return mg.graph.query("MATCH ()-[r]->() RETURN count(r) AS cnt")[0]["cnt"]
def test_add_creates_graph_data(self, memgraph_graph):
mg = memgraph_graph
mg.llm = _make_mock_llm(
[{"entity": "Alice", "entity_type": "person"}, {"entity": "Bob", "entity_type": "person"}],
[{"source": "Alice", "destination": "Bob", "relationship": "likes"}],
)
mg.add("Alice likes Bob", {"user_id": "u1"})
assert self._node_count(mg) == 2
assert self._edge_count(mg) == 1
def test_delete_hard_deletes_relationships(self, memgraph_graph):
"""Memgraph delete() should hard-delete the relationship (DELETE r)."""
mg = memgraph_graph
mg.llm = _make_mock_llm(
[{"entity": "Alice", "entity_type": "person"}, {"entity": "Bob", "entity_type": "person"}],
[{"source": "Alice", "destination": "Bob", "relationship": "likes"}],
)
mg.add("Alice likes Bob", {"user_id": "u1"})
assert self._edge_count(mg) == 1
mg.delete("Alice likes Bob", {"user_id": "u1"})
assert self._edge_count(mg) == 0
assert self._node_count(mg) == 2
def test_delete_preserves_other_relationships(self, memgraph_graph):
mg = memgraph_graph
mg.llm = _make_mock_llm(
[{"entity": "Alice", "entity_type": "person"}, {"entity": "Bob", "entity_type": "person"}],
[{"source": "Alice", "destination": "Bob", "relationship": "likes"}],
)
mg.add("Alice likes Bob", {"user_id": "u1"})
mg.llm = _make_mock_llm(
[{"entity": "Alice", "entity_type": "person"}, {"entity": "Charlie", "entity_type": "person"}],
[{"source": "Alice", "destination": "Charlie", "relationship": "knows"}],
)
mg.add("Alice knows Charlie", {"user_id": "u1"})
assert self._edge_count(mg) == 2
mg.llm = _make_mock_llm(
[{"entity": "Alice", "entity_type": "person"}, {"entity": "Bob", "entity_type": "person"}],
[{"source": "Alice", "destination": "Bob", "relationship": "likes"}],
)
mg.delete("Alice likes Bob", {"user_id": "u1"})
assert self._edge_count(mg) == 1
def test_delete_user_isolation(self, memgraph_graph):
mg = memgraph_graph
mg.llm = _make_mock_llm(
[{"entity": "Alice", "entity_type": "person"}, {"entity": "Bob", "entity_type": "person"}],
[{"source": "Alice", "destination": "Bob", "relationship": "likes"}],
)
mg.add("Alice likes Bob", {"user_id": "u1"})
mg.add("Alice likes Bob", {"user_id": "u2"})
assert self._edge_count(mg) == 2
mg.delete("Alice likes Bob", {"user_id": "u1"})
assert self._edge_count(mg) == 1
def test_delete_all_hard_deletes(self, memgraph_graph):
mg = memgraph_graph
mg.llm = _make_mock_llm(
[{"entity": "Alice", "entity_type": "person"}, {"entity": "Bob", "entity_type": "person"}],
[{"source": "Alice", "destination": "Bob", "relationship": "likes"}],
)
mg.add("Alice likes Bob", {"user_id": "u1"})
assert self._node_count(mg) == 2
mg.delete_all({"user_id": "u1"})
assert self._node_count(mg) == 0
assert self._edge_count(mg) == 0
def test_add_delete_add_cycle(self, memgraph_graph):
mg = memgraph_graph
mg.llm = _make_mock_llm(
[{"entity": "Alice", "entity_type": "person"}, {"entity": "Bob", "entity_type": "person"}],
[{"source": "Alice", "destination": "Bob", "relationship": "likes"}],
)
mg.add("Alice likes Bob", {"user_id": "u1"})
assert self._edge_count(mg) == 1
mg.delete("Alice likes Bob", {"user_id": "u1"})
assert self._edge_count(mg) == 0
mg.add("Alice likes Bob", {"user_id": "u1"})
assert self._edge_count(mg) == 1
# ===========================================================================
# APACHE AGE
# ===========================================================================
def _age_available():
if not _port_open("localhost", 5432):
return False
try:
import age
ag = age.connect(
host="localhost",
port=5432,
dbname="testdb",
user="postgres",
password="testpassword",
)
with ag.connection.cursor() as cur:
cur.execute("CREATE EXTENSION IF NOT EXISTS age;")
cur.execute("SET search_path = ag_catalog, '$user', public;")
ag.connection.commit()
ag.close()
return True
except Exception:
return False
requires_age = pytest.mark.skipif(not _age_available(), reason="Apache AGE not available")
@pytest.fixture
def age_graph():
"""Create an Apache AGE-backed MemoryGraph with mocked LLM/embedder."""
import age
from mem0.memory.apache_age_memory import MemoryGraph
graph_name = "mem0_test_delete"
ag = age.connect(
graph=graph_name,
host="localhost",
port=5432,
dbname="testdb",
user="postgres",
password="testpassword",
)
with ag.connection.cursor() as cur:
cur.execute("CREATE EXTENSION IF NOT EXISTS age;")
cur.execute("SET search_path = ag_catalog, '$user', public;")
ag.connection.commit()
age.setUpAge(ag.connection, graph_name)
ag.connection.commit()
mg = MemoryGraph.__new__(MemoryGraph)
mg.ag = ag
mg.graph_name = graph_name
mg.llm_provider = "openai"
mg.user_id = None
mg.threshold = 0.99
mg.embedding_model = _make_deterministic_embedder()
mg.llm = MagicMock()
mg.config = MagicMock()
mg.config.graph_store.custom_prompt = None
try:
ag.execCypher("MATCH (n) DETACH DELETE n")
ag.commit()
except Exception:
ag.rollback()
yield mg
try:
ag.execCypher("MATCH (n) DETACH DELETE n")
ag.commit()
except Exception:
ag.rollback()
ag.close()
def _age_node_count(mg):
cursor = mg.ag.execCypher("MATCH (n) RETURN count(n)", cols=["cnt"])
rows = cursor.fetchall()
return rows[0][0] if rows else 0
def _age_edge_count(mg):
cursor = mg.ag.execCypher("MATCH ()-[r]->() RETURN count(r)", cols=["cnt"])
rows = cursor.fetchall()
return rows[0][0] if rows else 0
@requires_age
class TestApacheAgeDeleteE2E:
def test_add_creates_graph_data(self, age_graph):
mg = age_graph
mg.llm = _make_mock_llm(
[{"entity": "Alice", "entity_type": "person"}, {"entity": "Bob", "entity_type": "person"}],
[{"source": "Alice", "destination": "Bob", "relationship": "likes"}],
)
mg.add("Alice likes Bob", {"user_id": "u1"})
assert _age_node_count(mg) == 2
assert _age_edge_count(mg) == 1
def test_delete_hard_deletes_relationships(self, age_graph):
"""Apache AGE delete() should hard-delete the relationship."""
mg = age_graph
mg.llm = _make_mock_llm(
[{"entity": "Alice", "entity_type": "person"}, {"entity": "Bob", "entity_type": "person"}],
[{"source": "Alice", "destination": "Bob", "relationship": "likes"}],
)
mg.add("Alice likes Bob", {"user_id": "u1"})
assert _age_edge_count(mg) == 1
mg.delete("Alice likes Bob", {"user_id": "u1"})
assert _age_edge_count(mg) == 0
assert _age_node_count(mg) == 2
def test_delete_preserves_other_relationships(self, age_graph):
mg = age_graph
mg.llm = _make_mock_llm(
[{"entity": "Alice", "entity_type": "person"}, {"entity": "Bob", "entity_type": "person"}],
[{"source": "Alice", "destination": "Bob", "relationship": "likes"}],
)
mg.add("Alice likes Bob", {"user_id": "u1"})
mg.llm = _make_mock_llm(
[{"entity": "Alice", "entity_type": "person"}, {"entity": "Charlie", "entity_type": "person"}],
[{"source": "Alice", "destination": "Charlie", "relationship": "knows"}],
)
mg.add("Alice knows Charlie", {"user_id": "u1"})
assert _age_edge_count(mg) == 2
mg.llm = _make_mock_llm(
[{"entity": "Alice", "entity_type": "person"}, {"entity": "Bob", "entity_type": "person"}],
[{"source": "Alice", "destination": "Bob", "relationship": "likes"}],
)
mg.delete("Alice likes Bob", {"user_id": "u1"})
assert _age_edge_count(mg) == 1
def test_delete_user_isolation(self, age_graph):
mg = age_graph
mg.llm = _make_mock_llm(
[{"entity": "Alice", "entity_type": "person"}, {"entity": "Bob", "entity_type": "person"}],
[{"source": "Alice", "destination": "Bob", "relationship": "likes"}],
)
mg.add("Alice likes Bob", {"user_id": "u1"})
mg.add("Alice likes Bob", {"user_id": "u2"})
assert _age_edge_count(mg) == 2
mg.delete("Alice likes Bob", {"user_id": "u1"})
assert _age_edge_count(mg) == 1
def test_delete_all_hard_deletes(self, age_graph):
mg = age_graph
mg.llm = _make_mock_llm(
[{"entity": "Alice", "entity_type": "person"}, {"entity": "Bob", "entity_type": "person"}],
[{"source": "Alice", "destination": "Bob", "relationship": "likes"}],
)
mg.add("Alice likes Bob", {"user_id": "u1"})
assert _age_node_count(mg) == 2
mg.delete_all({"user_id": "u1"})
assert _age_node_count(mg) == 0
assert _age_edge_count(mg) == 0
def test_add_delete_add_cycle(self, age_graph):
mg = age_graph
mg.llm = _make_mock_llm(
[{"entity": "Alice", "entity_type": "person"}, {"entity": "Bob", "entity_type": "person"}],
[{"source": "Alice", "destination": "Bob", "relationship": "likes"}],
)
mg.add("Alice likes Bob", {"user_id": "u1"})
assert _age_edge_count(mg) == 1
mg.delete("Alice likes Bob", {"user_id": "u1"})
assert _age_edge_count(mg) == 0
mg.add("Alice likes Bob", {"user_id": "u1"})
assert _age_edge_count(mg) == 1
# ===========================================================================
# NEPTUNE (tested via Neo4j OpenCypher — same query language)
# ===========================================================================
def _neptune_test_available():
"""Neptune uses OpenCypher — we test NeptuneBase.delete() against Neo4j."""
if not _port_open("localhost", 7687):
return False
try:
# Mock langchain_aws so NeptuneBase can be imported without AWS deps
sys.modules.setdefault("langchain_aws", MagicMock())
sys.modules.setdefault("botocore", MagicMock())
sys.modules.setdefault("botocore.config", MagicMock())
from mem0.graphs.neptune.base import NeptuneBase # noqa: F401
from langchain_neo4j import Neo4jGraph
g = Neo4jGraph(
url="bolt://localhost:7687",
username="neo4j",
password="testpassword",
refresh_schema=False,
driver_config={"notifications_min_severity": "OFF"},
)
g.query("RETURN 1")
return True
except Exception:
return False
requires_neptune_test = pytest.mark.skipif(
not _neptune_test_available(),
reason="Neo4j not available (used as OpenCypher backend for Neptune tests)",
)
def _make_concrete_neptune_subclass():
"""Create a concrete NeptuneBase subclass for testing, backed by Neo4j."""
# Ensure mocks are in place for import
sys.modules.setdefault("langchain_aws", MagicMock())
sys.modules.setdefault("botocore", MagicMock())
sys.modules.setdefault("botocore.config", MagicMock())
from mem0.graphs.neptune.base import NeptuneBase
class TestableNeptune(NeptuneBase):
def __init__(self):
pass
def _delete_entities_cypher(self, source, destination, relationship, user_id):
cypher = f"""
MATCH (n:`__Entity__` {{name: $source_name, user_id: $user_id}})
-[r:{relationship}]->
(m:`__Entity__` {{name: $dest_name, user_id: $user_id}})
DELETE r
RETURN n.name AS source, m.name AS target, type(r) AS relationship
"""
return cypher, {"source_name": source, "dest_name": destination, "user_id": user_id}
def _delete_all_cypher(self, filters):
return (
"MATCH (n:`__Entity__` {user_id: $user_id}) DETACH DELETE n",
{"user_id": filters["user_id"]},
)
# Stubs for abstract methods not used in delete path
def _add_entities_by_source_cypher(self, *a, **kw): pass
def _add_entities_by_destination_cypher(self, *a, **kw): pass
def _add_relationship_entities_cypher(self, *a, **kw): pass
def _add_new_entities_cypher(self, *a, **kw): pass
def _search_source_node_cypher(self, *a, **kw): pass
def _search_destination_node_cypher(self, *a, **kw): pass
def _get_all_cypher(self, *a, **kw): pass
def _search_graph_db_cypher(self, *a, **kw): pass
return TestableNeptune
@pytest.fixture
def neptune_graph():
"""NeptuneBase subclass backed by a real Neo4j container."""
from langchain_neo4j import Neo4jGraph
cls = _make_concrete_neptune_subclass()
mg = cls()
mg.graph = Neo4jGraph(
url="bolt://localhost:7687",
username="neo4j",
password="testpassword",
refresh_schema=False,
driver_config={"notifications_min_severity": "OFF"},
)
mg.graph.query("MATCH (n) DETACH DELETE n")
mg.node_label = ":`__Entity__`"
mg.llm_provider = "openai"
mg.user_id = None
mg.threshold = 0.99
mg.embedding_model = _make_deterministic_embedder()
mg.llm = MagicMock()
mg.config = MagicMock()
mg.config.graph_store.custom_prompt = None
yield mg
mg.graph.query("MATCH (n) DETACH DELETE n")
def _neptune_node_count(mg):
return mg.graph.query("MATCH (n) RETURN count(n) AS cnt")[0]["cnt"]
def _neptune_edge_count(mg):
return mg.graph.query("MATCH ()-[r]->() RETURN count(r) AS cnt")[0]["cnt"]
def _neptune_create_entities(mg, user_id):
"""Create test entities directly via Cypher."""
mg.graph.query(f"""
CREATE (a:`__Entity__` {{name: 'alice', user_id: '{user_id}'}})
CREATE (b:`__Entity__` {{name: 'bob', user_id: '{user_id}'}})
CREATE (a)-[:likes]->(b)
""")
@requires_neptune_test
class TestNeptuneDeleteE2E:
"""Test NeptuneBase.delete() using Neo4j as the OpenCypher backend.
Neptune uses standard OpenCypher, the same query language as Neo4j.
This validates that:
- NeptuneBase.delete() correctly calls _delete_entities(to_be_deleted, user_id) with a string
- The generated Cypher from _delete_entities_cypher runs correctly
- User isolation works
"""
def test_delete_removes_relationship(self, neptune_graph):
mg = neptune_graph
mg.llm = _make_mock_llm(
[{"entity": "Alice", "entity_type": "person"}, {"entity": "Bob", "entity_type": "person"}],
[{"source": "Alice", "destination": "Bob", "relationship": "likes"}],
)
_neptune_create_entities(mg, "u1")
assert _neptune_edge_count(mg) == 1
mg.delete("Alice likes Bob", {"user_id": "u1"})
assert _neptune_edge_count(mg) == 0
assert _neptune_node_count(mg) == 2
def test_delete_user_isolation(self, neptune_graph):
mg = neptune_graph
mg.llm = _make_mock_llm(
[{"entity": "Alice", "entity_type": "person"}, {"entity": "Bob", "entity_type": "person"}],
[{"source": "Alice", "destination": "Bob", "relationship": "likes"}],
)
_neptune_create_entities(mg, "u1")
_neptune_create_entities(mg, "u2")
assert _neptune_edge_count(mg) == 2
mg.delete("Alice likes Bob", {"user_id": "u1"})
assert _neptune_edge_count(mg) == 1
def test_delete_passes_user_id_string_not_dict(self, neptune_graph):
"""Verify NeptuneBase.delete() passes filters['user_id'] (string) to _delete_entities."""
mg = neptune_graph
mg.llm = _make_mock_llm(
[{"entity": "Alice", "entity_type": "person"}, {"entity": "Bob", "entity_type": "person"}],
[{"source": "Alice", "destination": "Bob", "relationship": "likes"}],
)
original = mg._delete_entities
call_args = []
def spy(to_be_deleted, user_id):
call_args.append(("to_be_deleted", to_be_deleted, "user_id", user_id))
return original(to_be_deleted, user_id)
mg._delete_entities = spy
_neptune_create_entities(mg, "u1")
mg.delete("Alice likes Bob", {"user_id": "u1"})
assert len(call_args) == 1
assert call_args[0][3] == "u1"
assert isinstance(call_args[0][3], str)
def test_delete_all(self, neptune_graph):
mg = neptune_graph
_neptune_create_entities(mg, "u1")
assert _neptune_node_count(mg) == 2
mg.delete_all({"user_id": "u1"})
assert _neptune_node_count(mg) == 0
def test_add_delete_add_cycle(self, neptune_graph):
mg = neptune_graph
mg.llm = _make_mock_llm(
[{"entity": "Alice", "entity_type": "person"}, {"entity": "Bob", "entity_type": "person"}],
[{"source": "Alice", "destination": "Bob", "relationship": "likes"}],
)
_neptune_create_entities(mg, "u1")
assert _neptune_edge_count(mg) == 1
mg.delete("Alice likes Bob", {"user_id": "u1"})
assert _neptune_edge_count(mg) == 0
_neptune_create_entities(mg, "u1")
assert _neptune_edge_count(mg) == 1
+736
View File
@@ -0,0 +1,736 @@
"""
End-to-end tests for graph cleanup on memory deletion (issue #3245).
Uses a real Kuzu embedded database to verify that graph entities are
correctly cleaned up when memories are deleted. LLM and embedding calls
are mocked to provide deterministic entity extraction.
Tests are skipped automatically if kuzu is not installed.
"""
import shutil
import tempfile
from unittest.mock import MagicMock, patch
import pytest
from mem0.configs.base import MemoryConfig
try:
import kuzu # noqa: F401
_kuzu_available = True
except ImportError:
_kuzu_available = False
requires_kuzu = pytest.mark.skipif(not _kuzu_available, reason="kuzu is not installed")
# ---------------------------------------------------------------------------
# Helpers
# ---------------------------------------------------------------------------
def _node_count(kuzu_graph):
"""Return total node count in the Kuzu graph."""
result = kuzu_graph.execute("MATCH (n:Entity) RETURN count(n) AS cnt")
rows = list(result.rows_as_dict())
return int(rows[0]["cnt"])
def _edge_count(kuzu_graph):
"""Return total edge count in the Kuzu graph."""
result = kuzu_graph.execute("MATCH ()-[r:CONNECTED_TO]->() RETURN count(r) AS cnt")
rows = list(result.rows_as_dict())
return int(rows[0]["cnt"])
def _get_edges(kuzu_graph):
"""Return all edges as list of (source, relationship, destination) tuples."""
result = kuzu_graph.execute(
"MATCH (s:Entity)-[r:CONNECTED_TO]->(d:Entity) "
"RETURN s.name AS src, r.name AS rel, d.name AS dst"
)
return [(row["src"], row["rel"], row["dst"]) for row in result.rows_as_dict()]
def _get_nodes(kuzu_graph):
"""Return all node names."""
result = kuzu_graph.execute("MATCH (n:Entity) RETURN n.name AS name, n.user_id AS uid")
return [(row["name"], row["uid"]) for row in result.rows_as_dict()]
# ---------------------------------------------------------------------------
# Fixtures
# ---------------------------------------------------------------------------
class MockVectorMemory:
"""Mimics the object returned by vector_store.get()."""
def __init__(self, memory_id, payload, score=0.8):
self.id = memory_id
self.payload = payload
self.score = score
@pytest.fixture
def kuzu_graph_memory():
"""
Create a real Kuzu-backed MemoryGraph with mocked LLM and embedder.
Yields (graph_memory_instance, kuzu_connection) then cleans up.
"""
import os
import kuzu
tmpdir = tempfile.mkdtemp()
db_path = os.path.join(tmpdir, "test.kuzu")
db = kuzu.Database(db_path)
conn = kuzu.Connection(db)
# We'll construct the MemoryGraph by bypassing __init__ and setting up manually
from mem0.memory.kuzu_memory import MemoryGraph
mg = MemoryGraph.__new__(MemoryGraph)
# Real Kuzu connection
mg.db = db
mg.graph = conn
mg.node_label = ":Entity"
mg.rel_label = ":CONNECTED_TO"
mg.kuzu_create_schema()
# Deterministic embedding: use one-hot-style vectors per entity name
# to avoid accidental cosine similarity matches between different entities
embedding_dims = 64
mg.embedding_dims = embedding_dims
_embed_cache = {}
_embed_counter = [0]
def deterministic_embed(text):
"""Generate a deterministic, near-orthogonal embedding for each unique text."""
text_lower = text.lower().strip()
if text_lower not in _embed_cache:
# Create a sparse vector — set a unique dimension to 1.0
vec = [0.0] * embedding_dims
idx = _embed_counter[0] % embedding_dims
vec[idx] = 1.0
# Add small noise to other dims so it's not exactly zero
import hashlib
h = hashlib.sha256(text_lower.encode()).digest()
for i in range(embedding_dims):
vec[i] += float(h[i % len(h)]) / 25500.0 # tiny noise
norm = sum(v * v for v in vec) ** 0.5
_embed_cache[text_lower] = [v / norm for v in vec]
_embed_counter[0] += 1
return _embed_cache[text_lower]
mock_embedder = MagicMock()
mock_embedder.embed.side_effect = deterministic_embed
mock_embedder.config.embedding_dims = embedding_dims
mg.embedding_model = mock_embedder
# Mock LLM — configured per-test via mock_embedder
mg.llm = MagicMock()
mg.llm_provider = "openai"
mg.user_id = None
# High threshold so only identical entity names merge, not similar ones
mg.threshold = 0.99
mg.config = MagicMock()
mg.config.graph_store.custom_prompt = None
yield mg, conn
# Cleanup
conn.close()
shutil.rmtree(tmpdir, ignore_errors=True)
def _setup_llm_for_entities(mg, entities, relations):
"""
Configure the mock LLM to return specific entities and relations.
entities: list of {"entity": str, "entity_type": str}
relations: list of {"source": str, "destination": str, "relationship": str}
"""
def generate_response(messages, tools):
# Detect which tool is being called based on tool definition names
tool_names = []
for t in tools:
if isinstance(t, dict):
fn = t.get("function", t)
tool_names.append(fn.get("name", ""))
else:
tool_names.append(getattr(t, "name", str(t)))
if any("extract_entities" in n for n in tool_names):
return {
"tool_calls": [
{
"name": "extract_entities",
"arguments": {"entities": entities},
}
]
}
elif any("establish" in n or "relation" in n for n in tool_names):
return {
"tool_calls": [
{
"name": "establish_nodes_relations",
"arguments": {"entities": relations},
}
]
}
elif any("delete" in n for n in tool_names):
# For _get_delete_entities_from_search_output during add() — return nothing to delete
return {"tool_calls": []}
return {"tool_calls": []}
mg.llm.generate_response.side_effect = generate_response
# ---------------------------------------------------------------------------
# End-to-end tests
# ---------------------------------------------------------------------------
@requires_kuzu
class TestKuzuGraphDeleteE2E:
"""End-to-end tests using a real Kuzu database."""
def test_add_creates_nodes_and_edges(self, kuzu_graph_memory):
"""Baseline: verify add() actually creates graph data."""
mg, conn = kuzu_graph_memory
_setup_llm_for_entities(
mg,
entities=[
{"entity": "Alice", "entity_type": "person"},
{"entity": "Bob", "entity_type": "person"},
],
relations=[
{"source": "Alice", "destination": "Bob", "relationship": "likes"},
],
)
filters = {"user_id": "test_user"}
mg.add("Alice likes Bob", filters)
assert _node_count(conn) == 2
assert _edge_count(conn) == 1
edges = _get_edges(conn)
assert ("alice", "likes", "bob") in edges
def test_delete_removes_edges_created_by_add(self, kuzu_graph_memory):
"""Core test: delete() should remove the relationships that add() created."""
mg, conn = kuzu_graph_memory
_setup_llm_for_entities(
mg,
entities=[
{"entity": "Alice", "entity_type": "person"},
{"entity": "Bob", "entity_type": "person"},
],
relations=[
{"source": "Alice", "destination": "Bob", "relationship": "likes"},
],
)
filters = {"user_id": "test_user"}
mg.add("Alice likes Bob", filters)
assert _edge_count(conn) == 1
# Now delete using the same text — should remove the relationship
mg.delete("Alice likes Bob", filters)
assert _edge_count(conn) == 0
# Nodes remain (we don't delete nodes on single memory delete)
assert _node_count(conn) == 2
def test_delete_only_removes_matching_edges(self, kuzu_graph_memory):
"""delete() should only remove edges matching the extracted relationships."""
mg, conn = kuzu_graph_memory
# First add: Alice likes Bob
_setup_llm_for_entities(
mg,
entities=[
{"entity": "Alice", "entity_type": "person"},
{"entity": "Bob", "entity_type": "person"},
],
relations=[
{"source": "Alice", "destination": "Bob", "relationship": "likes"},
],
)
filters = {"user_id": "test_user"}
mg.add("Alice likes Bob", filters)
# Second add: Alice knows Charlie
_setup_llm_for_entities(
mg,
entities=[
{"entity": "Alice", "entity_type": "person"},
{"entity": "Charlie", "entity_type": "person"},
],
relations=[
{"source": "Alice", "destination": "Charlie", "relationship": "knows"},
],
)
mg.add("Alice knows Charlie", filters)
assert _edge_count(conn) == 2
# Delete only the "Alice likes Bob" memory
_setup_llm_for_entities(
mg,
entities=[
{"entity": "Alice", "entity_type": "person"},
{"entity": "Bob", "entity_type": "person"},
],
relations=[
{"source": "Alice", "destination": "Bob", "relationship": "likes"},
],
)
mg.delete("Alice likes Bob", filters)
assert _edge_count(conn) == 1
edges = _get_edges(conn)
assert ("alice", "knows", "charlie") in edges
assert ("alice", "likes", "bob") not in edges
def test_delete_with_different_user_id_does_not_affect_other_users(self, kuzu_graph_memory):
"""delete() scoped to user_id should not touch another user's graph data."""
mg, conn = kuzu_graph_memory
_setup_llm_for_entities(
mg,
entities=[
{"entity": "Alice", "entity_type": "person"},
{"entity": "Bob", "entity_type": "person"},
],
relations=[
{"source": "Alice", "destination": "Bob", "relationship": "likes"},
],
)
# Add for user1
mg.add("Alice likes Bob", {"user_id": "user1"})
# Add same data for user2
mg.add("Alice likes Bob", {"user_id": "user2"})
assert _edge_count(conn) == 2
# Delete only user1's data
mg.delete("Alice likes Bob", {"user_id": "user1"})
assert _edge_count(conn) == 1
# Remaining edge belongs to user2
nodes = _get_nodes(conn)
user2_nodes = [n for n in nodes if n[1] == "user2"]
assert len(user2_nodes) == 2
def test_delete_nonexistent_relationship_is_safe(self, kuzu_graph_memory):
"""delete() on data that doesn't exist in the graph should be a no-op."""
mg, conn = kuzu_graph_memory
_setup_llm_for_entities(
mg,
entities=[
{"entity": "Alice", "entity_type": "person"},
{"entity": "Bob", "entity_type": "person"},
],
relations=[
{"source": "Alice", "destination": "Bob", "relationship": "hates"},
],
)
filters = {"user_id": "test_user"}
# Nothing in the graph yet
assert _edge_count(conn) == 0
assert _node_count(conn) == 0
# Should not raise
mg.delete("Alice hates Bob", filters)
assert _edge_count(conn) == 0
assert _node_count(conn) == 0
def test_delete_with_llm_failure_does_not_raise(self, kuzu_graph_memory):
"""If LLM fails during entity extraction, delete() should not raise."""
mg, conn = kuzu_graph_memory
# Make LLM raise
mg.llm.generate_response.side_effect = RuntimeError("LLM service down")
filters = {"user_id": "test_user"}
# Should not raise
mg.delete("Alice likes Bob", filters)
def test_delete_with_empty_entity_extraction(self, kuzu_graph_memory):
"""If LLM returns no entities, delete() should be a no-op."""
mg, conn = kuzu_graph_memory
# Add real data
_setup_llm_for_entities(
mg,
entities=[
{"entity": "Alice", "entity_type": "person"},
{"entity": "Bob", "entity_type": "person"},
],
relations=[
{"source": "Alice", "destination": "Bob", "relationship": "likes"},
],
)
filters = {"user_id": "test_user"}
mg.add("Alice likes Bob", filters)
assert _edge_count(conn) == 1
# Now delete but LLM returns no entities
_setup_llm_for_entities(mg, entities=[], relations=[])
mg.delete("some text", filters)
# Data should still be there
assert _edge_count(conn) == 1
def test_delete_all_removes_everything_for_user(self, kuzu_graph_memory):
"""delete_all() should remove all nodes/edges for a user (baseline behavior)."""
mg, conn = kuzu_graph_memory
_setup_llm_for_entities(
mg,
entities=[
{"entity": "Alice", "entity_type": "person"},
{"entity": "Bob", "entity_type": "person"},
],
relations=[
{"source": "Alice", "destination": "Bob", "relationship": "likes"},
],
)
filters = {"user_id": "test_user"}
mg.add("Alice likes Bob", filters)
_setup_llm_for_entities(
mg,
entities=[
{"entity": "Bob", "entity_type": "person"},
{"entity": "Charlie", "entity_type": "person"},
],
relations=[
{"source": "Bob", "destination": "Charlie", "relationship": "knows"},
],
)
mg.add("Bob knows Charlie", filters)
assert _node_count(conn) >= 3
assert _edge_count(conn) == 2
mg.delete_all(filters)
assert _node_count(conn) == 0
assert _edge_count(conn) == 0
def test_add_delete_add_cycle(self, kuzu_graph_memory):
"""Verify that add → delete → re-add works correctly."""
mg, conn = kuzu_graph_memory
_setup_llm_for_entities(
mg,
entities=[
{"entity": "Alice", "entity_type": "person"},
{"entity": "Bob", "entity_type": "person"},
],
relations=[
{"source": "Alice", "destination": "Bob", "relationship": "likes"},
],
)
filters = {"user_id": "test_user"}
# Add
mg.add("Alice likes Bob", filters)
assert _edge_count(conn) == 1
# Delete
mg.delete("Alice likes Bob", filters)
assert _edge_count(conn) == 0
# Re-add
mg.add("Alice likes Bob", filters)
assert _edge_count(conn) == 1
edges = _get_edges(conn)
assert ("alice", "likes", "bob") in edges
@requires_kuzu
class TestMemoryDeleteWithGraphE2E:
"""
End-to-end tests for Memory.delete() with graph enabled.
Uses a real Kuzu database for the graph store and mocks for
the vector store, LLM, and embedder.
"""
@pytest.fixture
def memory_with_graph(self):
"""Create a Memory instance with a real Kuzu graph backend."""
import os
import kuzu
tmpdir = tempfile.mkdtemp()
with (
patch("mem0.utils.factory.EmbedderFactory.create") as mock_embedder_factory,
patch("mem0.utils.factory.VectorStoreFactory.create") as mock_vector_factory,
patch("mem0.utils.factory.LlmFactory.create") as mock_llm_factory,
patch("mem0.memory.storage.SQLiteManager") as mock_sqlite,
):
_mem_embed_cache = {}
_mem_embed_counter = [0]
def _mem_deterministic_embed(text, *args, **kwargs):
text_lower = text.lower().strip()
if text_lower not in _mem_embed_cache:
import hashlib
vec = [0.0] * 64
idx = _mem_embed_counter[0] % 64
vec[idx] = 1.0
h = hashlib.sha256(text_lower.encode()).digest()
for i in range(64):
vec[i] += float(h[i % len(h)]) / 25500.0
norm = sum(v * v for v in vec) ** 0.5
_mem_embed_cache[text_lower] = [v / norm for v in vec]
_mem_embed_counter[0] += 1
return _mem_embed_cache[text_lower]
mock_embedder = MagicMock()
mock_embedder.embed.side_effect = _mem_deterministic_embed
mock_embedder.config.embedding_dims = 64
mock_embedder_factory.return_value = mock_embedder
mock_vector_store = MagicMock()
mock_vector_factory.return_value = mock_vector_store
mock_llm = MagicMock()
mock_llm_factory.return_value = mock_llm
mock_sqlite.return_value = MagicMock()
from mem0.memory.main import Memory
config = MemoryConfig()
memory = Memory(config)
# Now wire up a real Kuzu graph
db_path = os.path.join(tmpdir, "test.kuzu")
db = kuzu.Database(db_path)
conn = kuzu.Connection(db)
from mem0.memory.kuzu_memory import MemoryGraph as KuzuMemoryGraph
graph = KuzuMemoryGraph.__new__(KuzuMemoryGraph)
graph.db = db
graph.graph = conn
graph.node_label = ":Entity"
graph.rel_label = ":CONNECTED_TO"
graph.kuzu_create_schema()
graph.embedding_dims = 64
graph.embedding_model = mock_embedder
graph.llm = mock_llm
graph.llm_provider = "openai"
graph.user_id = None
graph.threshold = 0.99
graph.config = MagicMock()
graph.config.graph_store.custom_prompt = None
memory.graph = graph
memory.enable_graph = True
yield memory, mock_vector_store, mock_llm, conn
conn.close()
shutil.rmtree(tmpdir, ignore_errors=True)
def test_memory_delete_triggers_graph_cleanup(self, memory_with_graph):
"""
Full integration: Memory.delete() should clean up both vector store and graph.
"""
memory, mock_vs, mock_llm, conn = memory_with_graph
# 1. Manually add entities to the graph (simulating what add() would do)
_setup_llm_for_memory_graph(
mock_llm,
entities=[
{"entity": "Alice", "entity_type": "person"},
{"entity": "Bob", "entity_type": "person"},
],
relations=[
{"source": "Alice", "destination": "Bob", "relationship": "likes"},
],
)
memory.graph.add("Alice likes Bob", {"user_id": "user-1"})
assert _edge_count(conn) == 1
# 2. Set up mock vector store to return this memory
mock_vs.get.return_value = MockVectorMemory(
"mem-1",
{"data": "Alice likes Bob", "user_id": "user-1", "hash": "abc"},
)
# 3. Delete the memory
result = memory.delete("mem-1")
assert result == {"message": "Memory deleted successfully!"}
# 4. Verify graph was cleaned up
assert _edge_count(conn) == 0
# 5. Verify vector store was also cleaned up
mock_vs.delete.assert_called_once_with(vector_id="mem-1")
def test_memory_delete_with_graph_preserves_other_users_data(self, memory_with_graph):
"""Deleting user1's memory should not affect user2's graph data."""
memory, mock_vs, mock_llm, conn = memory_with_graph
_setup_llm_for_memory_graph(
mock_llm,
entities=[
{"entity": "Alice", "entity_type": "person"},
{"entity": "Bob", "entity_type": "person"},
],
relations=[
{"source": "Alice", "destination": "Bob", "relationship": "likes"},
],
)
# Add data for two users
memory.graph.add("Alice likes Bob", {"user_id": "user-1"})
memory.graph.add("Alice likes Bob", {"user_id": "user-2"})
assert _edge_count(conn) == 2
# Delete only user-1's memory
mock_vs.get.return_value = MockVectorMemory(
"mem-1",
{"data": "Alice likes Bob", "user_id": "user-1", "hash": "abc"},
)
memory.delete("mem-1")
# user-2's data should be intact
assert _edge_count(conn) == 1
nodes = _get_nodes(conn)
remaining_user_ids = set(uid for _, uid in nodes)
assert "user-2" in remaining_user_ids
def test_memory_delete_graph_failure_still_deletes_vector(self, memory_with_graph):
"""If graph cleanup fails, vector store deletion should still proceed."""
memory, mock_vs, mock_llm, conn = memory_with_graph
# Make LLM raise during entity extraction (graph cleanup will fail)
mock_llm.generate_response.side_effect = RuntimeError("LLM exploded")
mock_vs.get.return_value = MockVectorMemory(
"mem-1",
{"data": "Alice likes Bob", "user_id": "user-1", "hash": "abc"},
)
result = memory.delete("mem-1")
assert result == {"message": "Memory deleted successfully!"}
mock_vs.delete.assert_called_once_with(vector_id="mem-1")
def test_memory_delete_all_uses_bulk_not_per_memory(self, memory_with_graph):
"""delete_all() should use delete_all() on graph, not per-memory delete()."""
memory, mock_vs, mock_llm, conn = memory_with_graph
_setup_llm_for_memory_graph(
mock_llm,
entities=[
{"entity": "Alice", "entity_type": "person"},
{"entity": "Bob", "entity_type": "person"},
],
relations=[
{"source": "Alice", "destination": "Bob", "relationship": "likes"},
],
)
memory.graph.add("Alice likes Bob", {"user_id": "user-1"})
assert _edge_count(conn) == 1
# Set up vector store to return memories for deletion
mem1 = MockVectorMemory("mem-1", {"data": "Alice likes Bob", "user_id": "user-1"})
mock_vs.list.return_value = ([mem1], 1)
mock_vs.get.return_value = mem1
memory.delete_all(user_id="user-1")
# After delete_all, graph should be empty (via graph.delete_all)
assert _edge_count(conn) == 0
assert _node_count(conn) == 0
def test_memory_delete_nonexistent_raises_without_graph_side_effects(self, memory_with_graph):
"""Deleting a non-existent memory should raise ValueError without touching graph."""
memory, mock_vs, mock_llm, conn = memory_with_graph
# Add some graph data that should NOT be affected
_setup_llm_for_memory_graph(
mock_llm,
entities=[
{"entity": "Alice", "entity_type": "person"},
{"entity": "Bob", "entity_type": "person"},
],
relations=[
{"source": "Alice", "destination": "Bob", "relationship": "likes"},
],
)
memory.graph.add("Alice likes Bob", {"user_id": "user-1"})
assert _edge_count(conn) == 1
# Memory doesn't exist in vector store
mock_vs.get.return_value = None
with pytest.raises(ValueError, match="Memory with id non-existent not found"):
memory.delete("non-existent")
# Graph data should be untouched
assert _edge_count(conn) == 1
def _setup_llm_for_memory_graph(mock_llm, entities, relations):
"""Configure mock LLM for the Memory-level graph operations."""
def generate_response(messages, tools):
tool_names = []
for t in tools:
if isinstance(t, dict):
fn = t.get("function", t)
tool_names.append(fn.get("name", ""))
else:
tool_names.append(getattr(t, "name", str(t)))
if any("extract_entities" in n for n in tool_names):
return {
"tool_calls": [
{
"name": "extract_entities",
"arguments": {"entities": entities},
}
]
}
elif any("establish" in n or "relation" in n for n in tool_names):
return {
"tool_calls": [
{
"name": "establish_nodes_relations",
"arguments": {"entities": relations},
}
]
}
elif any("delete" in n for n in tool_names):
return {"tool_calls": []}
return {"tool_calls": []}
mock_llm.generate_response.side_effect = generate_response
+3 -1
View File
@@ -186,7 +186,9 @@ def test_delete(memory_instance):
result = memory_instance.delete("test_id")
memory_instance._delete_memory.assert_called_once_with("test_id")
# delete() now fetches the memory first and passes it to _delete_memory
existing_memory = memory_instance.vector_store.get.return_value
memory_instance._delete_memory.assert_called_once_with("test_id", existing_memory)
assert result["message"] == "Memory deleted successfully!"
+118
View File
@@ -320,3 +320,121 @@ def test_normalize_facts_mixed():
def test_normalize_facts_filters_empty_strings():
assert normalize_facts(["", "valid", ""]) == ["valid"]
@patch('mem0.utils.factory.EmbedderFactory.create')
@patch('mem0.utils.factory.VectorStoreFactory.create')
@patch('mem0.utils.factory.LlmFactory.create')
@patch('mem0.memory.storage.SQLiteManager')
def test_delete_nonexistent_memory_raises_error(mock_sqlite, mock_llm_factory, mock_vector_factory, mock_embedder_factory):
"""
Test that delete() raises ValueError when memory_id does not exist
and does not attempt to delete from the vector store.
Issue #3849: memory.delete() fails with AttributeError when memory not found.
Should raise a clear ValueError instead.
"""
mock_embedder_factory.return_value = MagicMock()
mock_vector_store = MagicMock()
mock_vector_factory.return_value = mock_vector_store
mock_llm_factory.return_value = MagicMock()
mock_sqlite.return_value = MagicMock()
from mem0.memory.main import Memory as MemoryClass
config = MemoryConfig()
memory = MemoryClass(config)
mock_vector_store.get.return_value = None
with pytest.raises(ValueError, match="Memory with id non-existent-id not found"):
memory.delete("non-existent-id")
mock_vector_store.delete.assert_not_called()
@pytest.mark.asyncio
@patch('mem0.utils.factory.EmbedderFactory.create')
@patch('mem0.utils.factory.VectorStoreFactory.create')
@patch('mem0.utils.factory.LlmFactory.create')
@patch('mem0.memory.storage.SQLiteManager')
async def test_async_delete_nonexistent_memory_raises_error(mock_sqlite, mock_llm_factory, mock_vector_factory, mock_embedder_factory):
"""
Test that async delete() raises ValueError when memory_id does not exist
and does not attempt to delete from the vector store.
Issue #3849: memory.delete() fails with AttributeError when memory not found.
Should raise a clear ValueError instead.
"""
mock_embedder_factory.return_value = MagicMock()
mock_vector_store = MagicMock()
mock_vector_store.get.return_value = None
mock_vector_factory.return_value = mock_vector_store
mock_llm_factory.return_value = MagicMock()
mock_sqlite.return_value = MagicMock()
from mem0.memory.main import AsyncMemory
config = MemoryConfig()
memory = AsyncMemory(config)
with pytest.raises(ValueError, match="Memory with id non-existent-id not found"):
await memory.delete("non-existent-id")
mock_vector_store.delete.assert_not_called()
@patch('mem0.utils.factory.EmbedderFactory.create')
@patch('mem0.utils.factory.VectorStoreFactory.create')
@patch('mem0.utils.factory.LlmFactory.create')
@patch('mem0.memory.storage.SQLiteManager')
def test_update_nonexistent_memory_raises_error(mock_sqlite, mock_llm_factory, mock_vector_factory, mock_embedder_factory):
"""
Test that _update_memory() raises ValueError when memory_id does not exist.
Same class of bug as #3849 — vector_store.get() returns None and code
accesses .payload without a null check.
"""
mock_embedder_factory.return_value = MagicMock()
mock_vector_store = MagicMock()
mock_vector_factory.return_value = mock_vector_store
mock_llm_factory.return_value = MagicMock()
mock_sqlite.return_value = MagicMock()
from mem0.memory.main import Memory as MemoryClass
config = MemoryConfig()
memory = MemoryClass(config)
mock_vector_store.get.return_value = None
with pytest.raises(ValueError, match="Memory with id non-existent-id not found"):
memory._update_memory("non-existent-id", "new data", {"new data": [0.1, 0.2]})
mock_vector_store.update.assert_not_called()
@pytest.mark.asyncio
@patch('mem0.utils.factory.EmbedderFactory.create')
@patch('mem0.utils.factory.VectorStoreFactory.create')
@patch('mem0.utils.factory.LlmFactory.create')
@patch('mem0.memory.storage.SQLiteManager')
async def test_async_update_nonexistent_memory_raises_error(mock_sqlite, mock_llm_factory, mock_vector_factory, mock_embedder_factory):
"""
Test that async _update_memory() raises ValueError when memory_id does not exist.
Same class of bug as #3849 — vector_store.get() returns None and code
accesses .payload without a null check.
"""
mock_embedder_factory.return_value = MagicMock()
mock_vector_store = MagicMock()
mock_vector_store.get.return_value = None
mock_vector_factory.return_value = mock_vector_store
mock_llm_factory.return_value = MagicMock()
mock_sqlite.return_value = MagicMock()
from mem0.memory.main import AsyncMemory
config = MemoryConfig()
memory = AsyncMemory(config)
with pytest.raises(ValueError, match="Memory with id non-existent-id not found"):
await memory._update_memory("non-existent-id", "new data", {"new data": [0.1, 0.2]})
mock_vector_store.update.assert_not_called()
+496
View File
@@ -0,0 +1,496 @@
"""Comprehensive E2E tests for REST API server authentication.
Tests the actual server/main.py app through FastAPI's TestClient (full ASGI
round-trip) covering:
- Auth disabled mode (ADMIN_API_KEY unset)
- Auth enabled mode (ADMIN_API_KEY set)
- Edge cases: empty keys, near-miss keys, timing-safe comparison, header
casing, response headers, startup logging, and full CRUD flows through auth.
"""
import importlib
import logging
import os
from unittest.mock import MagicMock, patch
import pytest
pytest.importorskip("fastapi", reason="fastapi not installed")
from fastapi.testclient import TestClient
# ---------------------------------------------------------------------------
# Fixtures
# ---------------------------------------------------------------------------
@pytest.fixture
def _mock_memory():
"""Patch Memory.from_config so the server imports without a real backend."""
mock_instance = MagicMock()
# Set up return values so CRUD endpoints return realistic responses
mock_instance.get.return_value = {"id": "mem-1", "memory": "test memory", "user_id": "alice"}
mock_instance.get_all.return_value = [
{"id": "mem-1", "memory": "test memory", "user_id": "alice"},
]
mock_instance.add.return_value = {"results": [{"id": "mem-1", "event": "ADD", "memory": "test"}]}
mock_instance.search.return_value = [{"id": "mem-1", "memory": "test", "score": 0.9}]
mock_instance.update.return_value = {"message": "Memory updated"}
mock_instance.history.return_value = [{"id": "mem-1", "old_memory": "a", "new_memory": "b"}]
mock_instance.delete.return_value = None
mock_instance.delete_all.return_value = {"message": "Memories deleted successfully!"}
mock_instance.reset.return_value = None
with patch.dict(os.environ, {"OPENAI_API_KEY": "fake-key"}):
with patch("mem0.Memory.from_config", return_value=mock_instance):
yield mock_instance
def _load_app(env_overrides: dict):
"""Reload server/main.py with the given environment and return the FastAPI app."""
import server.main as server_main
with patch.dict(os.environ, env_overrides, clear=False):
importlib.reload(server_main)
return server_main.app
# ---------------------------------------------------------------------------
# Auth disabled (ADMIN_API_KEY not set)
# ---------------------------------------------------------------------------
class TestAuthDisabled:
"""All endpoints should be freely accessible when ADMIN_API_KEY is empty."""
@pytest.fixture(autouse=True)
def _setup(self, _mock_memory):
self.app = _load_app({"ADMIN_API_KEY": ""})
self.client = TestClient(self.app)
self.mock = _mock_memory
def test_root_redirects_to_docs(self):
resp = self.client.get("/", follow_redirects=False)
assert resp.status_code == 307
assert "/docs" in resp.headers["location"]
def test_get_memory_without_key(self):
resp = self.client.get("/memories/mem-1")
assert resp.status_code == 200
assert resp.json()["id"] == "mem-1"
def test_get_all_memories_without_key(self):
resp = self.client.get("/memories", params={"user_id": "alice"})
assert resp.status_code == 200
def test_create_memory_without_key(self):
resp = self.client.post("/memories", json={
"messages": [{"role": "user", "content": "I like pizza"}],
"user_id": "alice",
})
assert resp.status_code == 200
def test_search_without_key(self):
resp = self.client.post("/search", json={"query": "pizza", "user_id": "alice"})
assert resp.status_code == 200
def test_update_memory_without_key(self):
resp = self.client.put("/memories/mem-1", json={"data": "updated"})
assert resp.status_code == 200
def test_history_without_key(self):
resp = self.client.get("/memories/mem-1/history")
assert resp.status_code == 200
def test_delete_memory_without_key(self):
resp = self.client.delete("/memories/mem-1")
assert resp.status_code == 200
def test_delete_all_without_key(self):
resp = self.client.delete("/memories", params={"user_id": "alice"})
assert resp.status_code == 200
def test_reset_without_key(self):
resp = self.client.post("/reset")
assert resp.status_code == 200
def test_configure_without_key(self):
self.mock.from_config = MagicMock()
resp = self.client.post("/configure", json={"version": "v1.1"})
assert resp.status_code == 200
def test_supplying_key_still_works_when_auth_disabled(self):
"""A client that sends X-API-Key should not be penalized when auth is off."""
resp = self.client.get(
"/memories/mem-1", headers={"X-API-Key": "some-random-key"}
)
assert resp.status_code == 200
@pytest.mark.parametrize(
"method,path",
[
("POST", "/configure"),
("POST", "/memories"),
("GET", "/memories"),
("GET", "/memories/test-id"),
("POST", "/search"),
("PUT", "/memories/test-id"),
("GET", "/memories/test-id/history"),
("DELETE", "/memories/test-id"),
("DELETE", "/memories"),
("POST", "/reset"),
],
)
def test_no_endpoint_returns_401_when_auth_disabled(self, method, path):
resp = self.client.request(method, path)
assert resp.status_code != 401, f"{method} {path} should not require auth"
# ---------------------------------------------------------------------------
# Auth enabled (ADMIN_API_KEY set)
# ---------------------------------------------------------------------------
class TestAuthEnabled:
"""All protected endpoints must enforce the API key."""
API_KEY = "test-secret-key-12345"
@pytest.fixture(autouse=True)
def _setup(self, _mock_memory):
self.app = _load_app({"ADMIN_API_KEY": self.API_KEY})
self.client = TestClient(self.app)
self.mock = _mock_memory
# --- Rejection cases ---
def test_missing_key_returns_401(self):
resp = self.client.get("/memories/mem-1")
assert resp.status_code == 401
def test_missing_key_detail_mentions_header(self):
resp = self.client.get("/memories/mem-1")
assert "X-API-Key" in resp.json()["detail"]
def test_wrong_key_returns_401(self):
resp = self.client.get("/memories/mem-1", headers={"X-API-Key": "wrong"})
assert resp.status_code == 401
def test_wrong_key_detail_says_invalid(self):
resp = self.client.get("/memories/mem-1", headers={"X-API-Key": "wrong"})
assert "Invalid" in resp.json()["detail"]
def test_empty_string_key_returns_401(self):
resp = self.client.get("/memories/mem-1", headers={"X-API-Key": ""})
assert resp.status_code == 401
def test_401_includes_www_authenticate_header(self):
resp = self.client.get("/memories/mem-1")
assert resp.headers.get("www-authenticate") == "ApiKey"
def test_near_miss_key_rejected(self):
"""Key that differs by one character should be rejected."""
near_miss = self.API_KEY[:-1] + ("6" if self.API_KEY[-1] != "6" else "7")
resp = self.client.get("/memories/mem-1", headers={"X-API-Key": near_miss})
assert resp.status_code == 401
def test_key_with_extra_whitespace_rejected(self):
resp = self.client.get("/memories/mem-1", headers={"X-API-Key": f" {self.API_KEY} "})
assert resp.status_code == 401
def test_key_prefix_rejected(self):
resp = self.client.get("/memories/mem-1", headers={"X-API-Key": self.API_KEY[:5]})
assert resp.status_code == 401
def test_key_with_different_case_rejected(self):
resp = self.client.get("/memories/mem-1", headers={"X-API-Key": self.API_KEY.upper()})
assert resp.status_code == 401
@pytest.mark.parametrize(
"method,path",
[
("POST", "/configure"),
("POST", "/memories"),
("GET", "/memories"),
("GET", "/memories/test-id"),
("POST", "/search"),
("PUT", "/memories/test-id"),
("GET", "/memories/test-id/history"),
("DELETE", "/memories/test-id"),
("DELETE", "/memories"),
("POST", "/reset"),
],
)
def test_all_endpoints_reject_without_key(self, method, path):
resp = self.client.request(method, path)
assert resp.status_code == 401, f"{method} {path} should require auth"
@pytest.mark.parametrize(
"method,path",
[
("POST", "/configure"),
("POST", "/memories"),
("GET", "/memories"),
("GET", "/memories/test-id"),
("POST", "/search"),
("PUT", "/memories/test-id"),
("GET", "/memories/test-id/history"),
("DELETE", "/memories/test-id"),
("DELETE", "/memories"),
("POST", "/reset"),
],
)
def test_all_endpoints_reject_wrong_key(self, method, path):
resp = self.client.request(method, path, headers={"X-API-Key": "wrong-key"})
assert resp.status_code == 401, f"{method} {path} should reject wrong key"
# --- Acceptance cases ---
def test_root_does_not_require_key(self):
resp = self.client.get("/", follow_redirects=False)
assert resp.status_code == 307
def _authed(self, method, path, **kwargs):
headers = kwargs.pop("headers", {})
headers["X-API-Key"] = self.API_KEY
return self.client.request(method, path, headers=headers, **kwargs)
def test_get_memory_with_key(self):
resp = self._authed("GET", "/memories/mem-1")
assert resp.status_code == 200
assert resp.json()["id"] == "mem-1"
def test_get_all_memories_with_key(self):
resp = self._authed("GET", "/memories", params={"user_id": "alice"})
assert resp.status_code == 200
def test_create_memory_with_key(self):
resp = self._authed("POST", "/memories", json={
"messages": [{"role": "user", "content": "I like pizza"}],
"user_id": "alice",
})
assert resp.status_code == 200
data = resp.json()
assert "results" in data
def test_search_with_key(self):
resp = self._authed("POST", "/search", json={"query": "pizza", "user_id": "alice"})
assert resp.status_code == 200
def test_update_memory_with_key(self):
resp = self._authed("PUT", "/memories/mem-1", json={"data": "updated"})
assert resp.status_code == 200
def test_history_with_key(self):
resp = self._authed("GET", "/memories/mem-1/history")
assert resp.status_code == 200
def test_delete_memory_with_key(self):
resp = self._authed("DELETE", "/memories/mem-1")
assert resp.status_code == 200
def test_delete_all_with_key(self):
resp = self._authed("DELETE", "/memories", params={"user_id": "alice"})
assert resp.status_code == 200
def test_reset_with_key(self):
resp = self._authed("POST", "/reset")
assert resp.status_code == 200
def test_configure_with_key(self):
resp = self._authed("POST", "/configure", json={"version": "v1.1"})
assert resp.status_code == 200
# ---------------------------------------------------------------------------
# Full CRUD flow through auth
# ---------------------------------------------------------------------------
class TestAuthenticatedCRUDFlow:
"""Verify a complete create → read → search → update → history → delete
cycle works end-to-end through the auth layer."""
API_KEY = "flow-test-key-99"
@pytest.fixture(autouse=True)
def _setup(self, _mock_memory):
self.app = _load_app({"ADMIN_API_KEY": self.API_KEY})
self.client = TestClient(self.app)
self.mock = _mock_memory
def _authed(self, method, path, **kwargs):
headers = kwargs.pop("headers", {})
headers["X-API-Key"] = self.API_KEY
return self.client.request(method, path, headers=headers, **kwargs)
def test_full_crud_cycle(self):
# 1. Create
resp = self._authed("POST", "/memories", json={
"messages": [{"role": "user", "content": "I love fresh vegetable pizza"}],
"user_id": "alice",
})
assert resp.status_code == 200
self.mock.add.assert_called_once()
# 2. Read single
resp = self._authed("GET", "/memories/mem-1")
assert resp.status_code == 200
self.mock.get.assert_called_once_with("mem-1")
# 3. Read all
resp = self._authed("GET", "/memories", params={"user_id": "alice"})
assert resp.status_code == 200
self.mock.get_all.assert_called_once_with(user_id="alice")
# 4. Search
resp = self._authed("POST", "/search", json={"query": "pizza", "user_id": "alice"})
assert resp.status_code == 200
self.mock.search.assert_called_once()
# 5. Update
resp = self._authed("PUT", "/memories/mem-1", json={"data": "updated content"})
assert resp.status_code == 200
self.mock.update.assert_called_once()
# 6. History
resp = self._authed("GET", "/memories/mem-1/history")
assert resp.status_code == 200
self.mock.history.assert_called_once_with(memory_id="mem-1")
# 7. Delete single
resp = self._authed("DELETE", "/memories/mem-1")
assert resp.status_code == 200
self.mock.delete.assert_called_once_with(memory_id="mem-1")
# 8. Delete all
resp = self._authed("DELETE", "/memories", params={"user_id": "alice"})
assert resp.status_code == 200
self.mock.delete_all.assert_called_once()
def test_crud_flow_blocked_without_auth(self):
"""Same flow should fail at every step without the key."""
endpoints = [
("POST", "/memories", {"json": {
"messages": [{"role": "user", "content": "test"}], "user_id": "alice"
}}),
("GET", "/memories/mem-1", {}),
("GET", "/memories", {"params": {"user_id": "alice"}}),
("POST", "/search", {"json": {"query": "pizza", "user_id": "alice"}}),
("PUT", "/memories/mem-1", {"json": {"data": "x"}}),
("GET", "/memories/mem-1/history", {}),
("DELETE", "/memories/mem-1", {}),
("DELETE", "/memories", {"params": {"user_id": "alice"}}),
("POST", "/reset", {}),
]
for method, path, kwargs in endpoints:
resp = self.client.request(method, path, **kwargs)
assert resp.status_code == 401, f"Unauthenticated {method} {path} should be 401"
# Verify the mock was NOT called (auth blocked before reaching handler)
self.mock.add.assert_not_called()
self.mock.get.assert_not_called()
self.mock.search.assert_not_called()
self.mock.update.assert_not_called()
self.mock.history.assert_not_called()
self.mock.delete.assert_not_called()
self.mock.delete_all.assert_not_called()
self.mock.reset.assert_not_called()
# ---------------------------------------------------------------------------
# Edge cases
# ---------------------------------------------------------------------------
class TestAuthEdgeCases:
"""Boundary conditions and unusual inputs."""
@pytest.fixture(autouse=True)
def _setup(self, _mock_memory):
self.mock = _mock_memory
def test_very_long_api_key(self):
"""Server should handle a very long key without crashing."""
long_key = "k" * 4096
app = _load_app({"ADMIN_API_KEY": long_key})
client = TestClient(app)
resp = client.get("/memories/mem-1", headers={"X-API-Key": long_key})
assert resp.status_code == 200
def test_special_characters_in_api_key(self):
"""Keys with special ASCII characters should work."""
special_key = "sk-!@#$%^&*()_+-=[]{}|;:',.<>?/~`"
app = _load_app({"ADMIN_API_KEY": special_key})
client = TestClient(app)
resp = client.get("/memories/mem-1", headers={"X-API-Key": special_key})
assert resp.status_code == 200
resp = client.get("/memories/mem-1", headers={"X-API-Key": "wrong"})
assert resp.status_code == 401
def test_key_env_var_not_present_at_all(self):
"""When the env var is completely absent, auth should be disabled."""
import server.main as server_main
env = os.environ.copy()
env.pop("ADMIN_API_KEY", None)
with patch.dict(os.environ, env, clear=True):
importlib.reload(server_main)
client = TestClient(server_main.app)
resp = client.get("/memories/mem-1")
assert resp.status_code != 401
def test_switching_from_enabled_to_disabled(self):
"""Simulates a server restart with auth toggled off."""
# First: auth enabled
app1 = _load_app({"ADMIN_API_KEY": "secret"})
c1 = TestClient(app1)
assert c1.get("/memories/mem-1").status_code == 401
# Then: auth disabled
app2 = _load_app({"ADMIN_API_KEY": ""})
c2 = TestClient(app2)
assert c2.get("/memories/mem-1").status_code != 401
def test_openapi_schema_accessible_without_key(self):
"""The /docs and /openapi.json endpoints should always be reachable."""
app = _load_app({"ADMIN_API_KEY": "secret"})
client = TestClient(app)
resp = client.get("/openapi.json")
assert resp.status_code == 200
schema = resp.json()
assert "paths" in schema
resp = client.get("/docs")
assert resp.status_code == 200
def test_openapi_schema_documents_auth(self):
"""The OpenAPI schema should mention authentication."""
app = _load_app({"ADMIN_API_KEY": "secret"})
client = TestClient(app)
schema = client.get("/openapi.json").json()
assert "Authentication" in schema.get("info", {}).get("description", "")
# ---------------------------------------------------------------------------
# Startup logging
# ---------------------------------------------------------------------------
class TestStartupLogging:
"""Verify the server emits the correct log messages at import time."""
@pytest.fixture(autouse=True)
def _setup(self, _mock_memory):
pass
def test_warning_when_auth_disabled(self, caplog):
with caplog.at_level(logging.WARNING):
_load_app({"ADMIN_API_KEY": ""})
assert any("UNSECURED" in r.message for r in caplog.records)
def test_info_when_auth_enabled(self, caplog):
with caplog.at_level(logging.INFO):
_load_app({"ADMIN_API_KEY": "a-long-enough-secret-key"})
assert any("authentication enabled" in r.message for r in caplog.records)
def test_warning_when_key_too_short(self, caplog):
with caplog.at_level(logging.WARNING):
_load_app({"ADMIN_API_KEY": "short"})
assert any("shorter than" in r.message for r in caplog.records)
+507
View File
@@ -0,0 +1,507 @@
"""Tests for REST API parameter forwarding.
Verifies that the Pydantic request models in server/main.py correctly accept
and forward all parameters supported by the underlying Memory class methods,
including limit, threshold, infer, memory_type, and prompt — which were
previously silently dropped by Pydantic v2's default extra='ignore' behavior.
"""
import importlib
import os
from unittest.mock import MagicMock, patch
import pytest
pytest.importorskip("fastapi", reason="fastapi not installed")
from fastapi.testclient import TestClient
# ---------------------------------------------------------------------------
# Fixtures
# ---------------------------------------------------------------------------
@pytest.fixture
def _mock_memory():
"""Patch Memory.from_config so the server imports without a real backend."""
mock_instance = MagicMock()
mock_instance.add.return_value = {"results": [{"id": "mem-1", "event": "ADD", "memory": "test"}]}
mock_instance.search.return_value = [{"id": "mem-1", "memory": "test", "score": 0.9}]
mock_instance.get.return_value = {"id": "mem-1", "memory": "test memory"}
mock_instance.get_all.return_value = [{"id": "mem-1", "memory": "test memory"}]
mock_instance.update.return_value = {"message": "Memory updated"}
mock_instance.history.return_value = [{"id": "mem-1", "old_memory": "a", "new_memory": "b"}]
mock_instance.delete.return_value = None
mock_instance.delete_all.return_value = {"message": "Memories deleted"}
mock_instance.reset.return_value = None
with patch.dict(os.environ, {"OPENAI_API_KEY": "fake-key", "ADMIN_API_KEY": ""}):
with patch("mem0.Memory.from_config", return_value=mock_instance):
yield mock_instance
@pytest.fixture
def client(_mock_memory):
"""Return a TestClient wired to the server app with mocked Memory."""
import server.main as server_main
with patch.dict(os.environ, {"ADMIN_API_KEY": ""}):
importlib.reload(server_main)
return TestClient(server_main.app)
@pytest.fixture
def mock_memory(_mock_memory):
return _mock_memory
# ===========================================================================
# SearchRequest: limit parameter
# ===========================================================================
class TestSearchLimit:
"""Verify that the limit parameter is accepted and forwarded to Memory.search()."""
def test_limit_forwarded(self, client, mock_memory):
resp = client.post("/search", json={"query": "food", "user_id": "u1", "limit": 5})
assert resp.status_code == 200
_, kwargs = mock_memory.search.call_args
assert kwargs["limit"] == 5
def test_limit_one(self, client, mock_memory):
resp = client.post("/search", json={"query": "food", "user_id": "u1", "limit": 1})
assert resp.status_code == 200
_, kwargs = mock_memory.search.call_args
assert kwargs["limit"] == 1
def test_limit_omitted_uses_memory_default(self, client, mock_memory):
"""When limit is not sent, it should not appear in the kwargs,
allowing Memory.search() to use its own default (100)."""
resp = client.post("/search", json={"query": "food", "user_id": "u1"})
assert resp.status_code == 200
_, kwargs = mock_memory.search.call_args
assert "limit" not in kwargs
# ===========================================================================
# SearchRequest: threshold parameter
# ===========================================================================
class TestSearchThreshold:
"""Verify that the threshold parameter is accepted and forwarded."""
def test_threshold_forwarded(self, client, mock_memory):
resp = client.post("/search", json={"query": "food", "user_id": "u1", "threshold": 0.8})
assert resp.status_code == 200
_, kwargs = mock_memory.search.call_args
assert kwargs["threshold"] == 0.8
def test_threshold_zero(self, client, mock_memory):
"""threshold=0.0 is a valid falsy value that must not be filtered out."""
resp = client.post("/search", json={"query": "food", "user_id": "u1", "threshold": 0.0})
assert resp.status_code == 200
_, kwargs = mock_memory.search.call_args
assert kwargs["threshold"] == 0.0
def test_threshold_omitted_uses_memory_default(self, client, mock_memory):
resp = client.post("/search", json={"query": "food", "user_id": "u1"})
assert resp.status_code == 200
_, kwargs = mock_memory.search.call_args
assert "threshold" not in kwargs
# ===========================================================================
# SearchRequest: limit + threshold together
# ===========================================================================
class TestSearchLimitAndThreshold:
def test_both_forwarded(self, client, mock_memory):
resp = client.post("/search", json={
"query": "food", "user_id": "u1", "limit": 10, "threshold": 0.5
})
assert resp.status_code == 200
_, kwargs = mock_memory.search.call_args
assert kwargs["limit"] == 10
assert kwargs["threshold"] == 0.5
# ===========================================================================
# MemoryCreate: infer parameter
# ===========================================================================
class TestAddInfer:
"""Verify that the infer parameter is accepted and forwarded to Memory.add()."""
def test_infer_false_forwarded(self, client, mock_memory):
resp = client.post("/memories", json={
"messages": [{"role": "user", "content": "Store this exactly"}],
"user_id": "u1",
"infer": False,
})
assert resp.status_code == 200
_, kwargs = mock_memory.add.call_args
assert kwargs["infer"] is False
def test_infer_true_forwarded(self, client, mock_memory):
resp = client.post("/memories", json={
"messages": [{"role": "user", "content": "I like pizza"}],
"user_id": "u1",
"infer": True,
})
assert resp.status_code == 200
_, kwargs = mock_memory.add.call_args
assert kwargs["infer"] is True
def test_infer_omitted_uses_memory_default(self, client, mock_memory):
"""When infer is not sent, it should not appear in kwargs,
allowing Memory.add() to use its own default (True)."""
resp = client.post("/memories", json={
"messages": [{"role": "user", "content": "hello"}],
"user_id": "u1",
})
assert resp.status_code == 200
_, kwargs = mock_memory.add.call_args
assert "infer" not in kwargs
# ===========================================================================
# MemoryCreate: memory_type parameter
# ===========================================================================
class TestAddMemoryType:
"""Verify that the memory_type parameter is accepted and forwarded."""
def test_memory_type_forwarded(self, client, mock_memory):
resp = client.post("/memories", json={
"messages": [{"role": "user", "content": "I like pizza"}],
"user_id": "u1",
"memory_type": "core",
})
assert resp.status_code == 200
_, kwargs = mock_memory.add.call_args
assert kwargs["memory_type"] == "core"
def test_memory_type_omitted(self, client, mock_memory):
resp = client.post("/memories", json={
"messages": [{"role": "user", "content": "hello"}],
"user_id": "u1",
})
assert resp.status_code == 200
_, kwargs = mock_memory.add.call_args
assert "memory_type" not in kwargs
# ===========================================================================
# MemoryCreate: prompt parameter
# ===========================================================================
class TestAddPrompt:
"""Verify that the prompt parameter is accepted and forwarded."""
def test_prompt_forwarded(self, client, mock_memory):
resp = client.post("/memories", json={
"messages": [{"role": "user", "content": "I like pizza"}],
"user_id": "u1",
"prompt": "Extract food preferences only.",
})
assert resp.status_code == 200
_, kwargs = mock_memory.add.call_args
assert kwargs["prompt"] == "Extract food preferences only."
def test_prompt_omitted(self, client, mock_memory):
resp = client.post("/memories", json={
"messages": [{"role": "user", "content": "hello"}],
"user_id": "u1",
})
assert resp.status_code == 200
_, kwargs = mock_memory.add.call_args
assert "prompt" not in kwargs
# ===========================================================================
# MemoryCreate: all new params together
# ===========================================================================
class TestAddAllNewParams:
def test_infer_memory_type_and_prompt_together(self, client, mock_memory):
resp = client.post("/memories", json={
"messages": [{"role": "user", "content": "I like pizza"}],
"user_id": "u1",
"infer": False,
"memory_type": "core",
"prompt": "Custom extraction prompt.",
})
assert resp.status_code == 200
_, kwargs = mock_memory.add.call_args
assert kwargs["infer"] is False
assert kwargs["memory_type"] == "core"
assert kwargs["prompt"] == "Custom extraction prompt."
# ===========================================================================
# Edge cases: falsy-but-valid values must not be filtered out
# ===========================================================================
class TestFalsyValues:
"""The handler filters with `v is not None`. Falsy values like False, 0,
0.0, and empty string must still be forwarded."""
def test_infer_false_not_filtered(self, client, mock_memory):
resp = client.post("/memories", json={
"messages": [{"role": "user", "content": "test"}],
"user_id": "u1",
"infer": False,
})
assert resp.status_code == 200
_, kwargs = mock_memory.add.call_args
assert kwargs["infer"] is False
def test_threshold_zero_not_filtered(self, client, mock_memory):
resp = client.post("/search", json={
"query": "food", "user_id": "u1", "threshold": 0.0,
})
assert resp.status_code == 200
_, kwargs = mock_memory.search.call_args
assert kwargs["threshold"] == 0.0
# ===========================================================================
# Extra/unknown fields are still silently ignored (existing Pydantic behavior)
# ===========================================================================
class TestUnknownFieldsIgnored:
def test_unknown_search_field_ignored(self, client, mock_memory):
resp = client.post("/search", json={
"query": "food", "user_id": "u1", "bogus_field": "xyz",
})
assert resp.status_code == 200
_, kwargs = mock_memory.search.call_args
assert "bogus_field" not in kwargs
def test_unknown_add_field_ignored(self, client, mock_memory):
resp = client.post("/memories", json={
"messages": [{"role": "user", "content": "test"}],
"user_id": "u1",
"unknown_param": 42,
})
assert resp.status_code == 200
_, kwargs = mock_memory.add.call_args
assert "unknown_param" not in kwargs
# ===========================================================================
# Backward compatibility: existing params still work
# ===========================================================================
class TestExistingParamsUnchanged:
def test_search_filters_still_forwarded(self, client, mock_memory):
resp = client.post("/search", json={
"query": "food",
"user_id": "u1",
"agent_id": "a1",
"filters": {"category": "food"},
})
assert resp.status_code == 200
_, kwargs = mock_memory.search.call_args
assert kwargs["user_id"] == "u1"
assert kwargs["agent_id"] == "a1"
assert kwargs["filters"] == {"category": "food"}
def test_add_metadata_still_forwarded(self, client, mock_memory):
resp = client.post("/memories", json={
"messages": [{"role": "user", "content": "test"}],
"user_id": "u1",
"agent_id": "a1",
"metadata": {"source": "test"},
})
assert resp.status_code == 200
_, kwargs = mock_memory.add.call_args
assert kwargs["user_id"] == "u1"
assert kwargs["agent_id"] == "a1"
assert kwargs["metadata"] == {"source": "test"}
# ===========================================================================
# OpenAPI schema: new fields are documented
# ===========================================================================
class TestOpenAPISchema:
"""Verify the new fields appear in the auto-generated OpenAPI schema."""
def test_search_schema_includes_limit(self, client):
schema = client.get("/openapi.json").json()
search_props = schema["components"]["schemas"]["SearchRequest"]["properties"]
assert "limit" in search_props
assert search_props["limit"]["description"] == "Maximum number of results to return."
def test_search_schema_includes_threshold(self, client):
schema = client.get("/openapi.json").json()
search_props = schema["components"]["schemas"]["SearchRequest"]["properties"]
assert "threshold" in search_props
def test_add_schema_includes_infer(self, client):
schema = client.get("/openapi.json").json()
add_props = schema["components"]["schemas"]["MemoryCreate"]["properties"]
assert "infer" in add_props
def test_add_schema_includes_memory_type(self, client):
schema = client.get("/openapi.json").json()
add_props = schema["components"]["schemas"]["MemoryCreate"]["properties"]
assert "memory_type" in add_props
def test_add_schema_includes_prompt(self, client):
schema = client.get("/openapi.json").json()
add_props = schema["components"]["schemas"]["MemoryCreate"]["properties"]
assert "prompt" in add_props
# ===========================================================================
# Pydantic type validation: invalid types return 422
# ===========================================================================
class TestTypeValidation:
"""Verify FastAPI/Pydantic rejects invalid types with 422."""
def test_limit_string_rejected(self, client):
resp = client.post("/search", json={
"query": "food", "user_id": "u1", "limit": "not_a_number",
})
assert resp.status_code == 422
def test_threshold_string_rejected(self, client):
resp = client.post("/search", json={
"query": "food", "user_id": "u1", "threshold": "high",
})
assert resp.status_code == 422
def test_infer_string_coerced_by_pydantic(self, client, mock_memory):
"""Pydantic v2 coerces truthy strings like 'yes' to True for bool fields."""
resp = client.post("/memories", json={
"messages": [{"role": "user", "content": "test"}],
"user_id": "u1",
"infer": "yes",
})
assert resp.status_code == 200
_, kwargs = mock_memory.add.call_args
assert kwargs["infer"] is True
def test_infer_invalid_value_rejected(self, client):
"""A value that cannot be coerced to bool should be rejected."""
resp = client.post("/memories", json={
"messages": [{"role": "user", "content": "test"}],
"user_id": "u1",
"infer": [1, 2, 3],
})
assert resp.status_code == 422
def test_limit_float_rejected(self, client):
resp = client.post("/search", json={
"query": "food", "user_id": "u1", "limit": 5.7,
})
assert resp.status_code == 422
def test_memory_type_int_rejected(self, client):
resp = client.post("/memories", json={
"messages": [{"role": "user", "content": "test"}],
"user_id": "u1",
"memory_type": 123,
})
assert resp.status_code == 422
# ===========================================================================
# Explicit null values: treated as omitted (filtered by `is not None`)
# ===========================================================================
class TestExplicitNull:
"""When a client sends null for an optional field, it should be treated
as omitted — the Memory class default should be used."""
def test_limit_null_uses_memory_default(self, client, mock_memory):
resp = client.post("/search", json={
"query": "food", "user_id": "u1", "limit": None,
})
assert resp.status_code == 200
_, kwargs = mock_memory.search.call_args
assert "limit" not in kwargs
def test_infer_null_uses_memory_default(self, client, mock_memory):
resp = client.post("/memories", json={
"messages": [{"role": "user", "content": "test"}],
"user_id": "u1",
"infer": None,
})
assert resp.status_code == 200
_, kwargs = mock_memory.add.call_args
assert "infer" not in kwargs
def test_prompt_null_uses_memory_default(self, client, mock_memory):
resp = client.post("/memories", json={
"messages": [{"role": "user", "content": "test"}],
"user_id": "u1",
"prompt": None,
})
assert resp.status_code == 200
_, kwargs = mock_memory.add.call_args
assert "prompt" not in kwargs
# ===========================================================================
# Verify exact call signatures match Memory method params
# ===========================================================================
class TestCallSignatureMatch:
"""Ensure forwarded params exactly match Memory.add() and Memory.search()
keyword argument names — a typo here would cause a TypeError at runtime."""
def test_search_kwargs_are_valid(self, client, mock_memory):
"""All kwargs forwarded to Memory.search() must be in its signature."""
resp = client.post("/search", json={
"query": "food", "user_id": "u1", "agent_id": "a1",
"run_id": "r1", "filters": {"k": "v"},
"limit": 10, "threshold": 0.5,
})
assert resp.status_code == 200
# The handler passes query= as a keyword arg, so it appears in kwargs too
_, kwargs = mock_memory.search.call_args
valid_params = {"query", "user_id", "agent_id", "run_id", "limit", "filters", "threshold", "rerank"}
for key in kwargs:
assert key in valid_params, f"Unexpected kwarg '{key}' forwarded to Memory.search()"
def test_add_kwargs_are_valid(self, client, mock_memory):
"""All kwargs forwarded to Memory.add() must be in its signature."""
resp = client.post("/memories", json={
"messages": [{"role": "user", "content": "hi"}],
"user_id": "u1", "agent_id": "a1", "run_id": "r1",
"metadata": {"k": "v"},
"infer": False, "memory_type": "core", "prompt": "custom",
})
assert resp.status_code == 200
# The handler passes messages= as a keyword arg, so it appears in kwargs too
_, kwargs = mock_memory.add.call_args
valid_params = {"messages", "user_id", "agent_id", "run_id", "metadata", "infer", "memory_type", "prompt"}
for key in kwargs:
assert key in valid_params, f"Unexpected kwarg '{key}' forwarded to Memory.add()"
def test_messages_excluded_from_params_dict(self, client, mock_memory):
"""messages is passed separately via messages= kwarg, not duplicated from model_dump."""
resp = client.post("/memories", json={
"messages": [{"role": "user", "content": "hi"}],
"user_id": "u1",
})
assert resp.status_code == 200
_, kwargs = mock_memory.add.call_args
# messages should be present (passed explicitly) and be a list of dicts
assert "messages" in kwargs
assert isinstance(kwargs["messages"], list)
assert kwargs["messages"][0] == {"role": "user", "content": "hi"}
def test_query_passed_explicitly(self, client, mock_memory):
"""query is passed as an explicit keyword arg to Memory.search()."""
resp = client.post("/search", json={"query": "food", "user_id": "u1"})
assert resp.status_code == 200
_, kwargs = mock_memory.search.call_args
assert kwargs["query"] == "food"
+13
View File
@@ -3,6 +3,7 @@ from unittest.mock import Mock, patch
import pytest
from mem0.vector_stores.chroma import ChromaDB
from mem0.configs.vector_stores.chroma import ChromaDbConfig
@pytest.fixture
@@ -249,3 +250,15 @@ def test_generate_where_clause_non_string_values():
# ChromaDB accepts non-string values in filters
expected = {"$and": [{"user_id": {"$eq": "alice"}}, {"count": {"$eq": 5}}, {"active": {"$eq": True}}]}
assert result == expected
def test_chroma_config_accepts_default_tmp_path():
"""Test that ChromaDbConfig accepts the default /tmp/chroma path."""
config = ChromaDbConfig(path="/tmp/chroma")
assert config.path == "/tmp/chroma"
def test_chroma_config_rejects_no_config():
"""Test that ChromaDbConfig rejects when no connection config is provided."""
with pytest.raises(ValueError):
ChromaDbConfig()
+499 -2
View File
@@ -1,8 +1,12 @@
from types import SimpleNamespace
from unittest.mock import MagicMock, patch
import pytest
pytest.importorskip("databricks", reason="databricks-sdk package not installed")
from databricks.sdk.service.vectorsearch import VectorIndexType, QueryVectorIndexResponse, ResultManifest, ResultData, ColumnInfo
from mem0.vector_stores.databricks import Databricks
import pytest
# ---------------------- Fixtures ---------------------- #
@@ -205,9 +209,34 @@ def test_search_direct_access_vector(db_instance_direct, mock_workspace_client):
assert results[0].score == 0.77
def test_search_delta_sync_self_managed_vectors(mock_workspace_client):
"""DELTA_SYNC without embedding model endpoint should use query_vector, not query_text."""
mock_workspace_client.tables.exists.return_value = SimpleNamespace(table_exists=True)
# DELTA_SYNC without embedding_model_endpoint_name = self-managed vectors
inst = Databricks(
workspace_url="https://test",
access_token="tok",
endpoint_name="vs-endpoint",
catalog="catalog",
schema="schema",
table_name="table",
warehouse_name="test-warehouse",
index_type=VectorIndexType.DELTA_SYNC,
embedding_dimension=4,
# NOTE: no embedding_model_endpoint_name
)
mock_workspace_client.vector_search_indexes.query_index.return_value = SimpleNamespace(
result=SimpleNamespace(data_array=[])
)
inst.search(query="ignored", vectors=[0.1, 0.2, 0.3, 0.4], limit=5)
call_kwargs = mock_workspace_client.vector_search_indexes.query_index.call_args.kwargs
assert "query_vector" in call_kwargs
assert "query_text" not in call_kwargs
def test_search_missing_params_raises(db_instance_delta):
with pytest.raises(ValueError):
db_instance_delta.search(query="", vectors=[0.1, 0.2]) # DELTA_SYNC requires query text
db_instance_delta.search(query="", vectors=[0.1, 0.2]) # DELTA_SYNC with model endpoint requires query text
# ---------------------- Delete Tests ---------------------- #
@@ -275,6 +304,54 @@ def test_get_vector(db_instance_delta, mock_workspace_client):
assert res.id == "id-get"
assert res.payload["data"] == "some memory"
assert res.payload["tag"] == "x"
# DELTA_SYNC should use query_text, not query_vector
call_kwargs = mock_workspace_client.vector_search_indexes.query_index.call_args.kwargs
assert "query_text" in call_kwargs
assert "query_vector" not in call_kwargs
def test_get_vector_direct_access(db_instance_direct, mock_workspace_client):
"""get() on a DIRECT_ACCESS index must use query_vector instead of query_text."""
mock_workspace_client.vector_search_indexes.query_index.return_value = QueryVectorIndexResponse(
manifest=ResultManifest(columns=[
ColumnInfo(name="memory_id"),
ColumnInfo(name="hash"),
ColumnInfo(name="agent_id"),
ColumnInfo(name="run_id"),
ColumnInfo(name="user_id"),
ColumnInfo(name="memory"),
ColumnInfo(name="metadata"),
ColumnInfo(name="created_at"),
ColumnInfo(name="updated_at"),
ColumnInfo(name="embedding"),
ColumnInfo(name="score"),
]),
result=ResultData(
data_array=[
[
"id-get-da",
"h",
"a",
"r",
"u",
"direct access memory",
'{"tag":"da"}',
"2024-01-01T00:00:00",
"2024-01-01T00:00:00",
[0.1, 0.2, 0.3, 0.4],
"0.88",
]
]
)
)
res = db_instance_direct.get("id-get-da")
assert res.id == "id-get-da"
assert res.payload["data"] == "direct access memory"
# DIRECT_ACCESS should use query_vector, not query_text
call_kwargs = mock_workspace_client.vector_search_indexes.query_index.call_args.kwargs
assert "query_vector" in call_kwargs
assert "query_text" not in call_kwargs
assert call_kwargs["query_vector"] == [0.0] * 4 # embedding_dimension=4
# ---------------------- Collection Info / Listing Tests ---------------------- #
@@ -330,6 +407,185 @@ def test_list_memories(db_instance_delta, mock_workspace_client):
assert isinstance(res, list)
assert len(res[0]) == 1
assert res[0][0].id == "id-get"
# DELTA_SYNC should use query_text
call_kwargs = mock_workspace_client.vector_search_indexes.query_index.call_args.kwargs
assert "query_text" in call_kwargs
assert "query_vector" not in call_kwargs
def test_list_memories_direct_access(db_instance_direct, mock_workspace_client):
"""list() on a DIRECT_ACCESS index must use query_vector instead of query_text."""
mock_workspace_client.vector_search_indexes.query_index.return_value = SimpleNamespace(
result=SimpleNamespace(
data_array=[
[
"id-da-list",
"h",
"a",
"r",
"u",
"direct memory",
None,
"2024-01-01T00:00:00",
"2024-01-01T00:00:00",
[0.1, 0.2, 0.3, 0.4],
]
]
)
)
res = db_instance_direct.list(limit=5)
assert isinstance(res, list)
assert len(res[0]) == 1
assert res[0][0].id == "id-da-list"
assert res[0][0].payload["data"] == "direct memory"
# DIRECT_ACCESS should use query_vector
call_kwargs = mock_workspace_client.vector_search_indexes.query_index.call_args.kwargs
assert "query_vector" in call_kwargs
assert "query_text" not in call_kwargs
assert call_kwargs["query_vector"] == [0.0] * 4
def test_get_vector_delta_sync_self_managed(mock_workspace_client):
"""get() on DELTA_SYNC without model endpoint should use query_vector."""
mock_workspace_client.tables.exists.return_value = SimpleNamespace(table_exists=True)
inst = Databricks(
workspace_url="https://test",
access_token="tok",
endpoint_name="vs-endpoint",
catalog="catalog",
schema="schema",
table_name="table",
warehouse_name="test-warehouse",
index_type=VectorIndexType.DELTA_SYNC,
embedding_dimension=4,
# NOTE: no embedding_model_endpoint_name
)
mock_workspace_client.vector_search_indexes.query_index.return_value = QueryVectorIndexResponse(
manifest=ResultManifest(columns=[
ColumnInfo(name="memory_id"), ColumnInfo(name="hash"),
ColumnInfo(name="agent_id"), ColumnInfo(name="run_id"),
ColumnInfo(name="user_id"), ColumnInfo(name="memory"),
ColumnInfo(name="metadata"), ColumnInfo(name="created_at"),
ColumnInfo(name="updated_at"),
]),
result=ResultData(data_array=[["id-sm", "h", None, None, None, "self-managed mem", None, None, None]]),
)
res = inst.get("id-sm")
assert res.id == "id-sm"
call_kwargs = mock_workspace_client.vector_search_indexes.query_index.call_args.kwargs
assert "query_vector" in call_kwargs
assert "query_text" not in call_kwargs
assert call_kwargs["query_vector"] == [0.0] * 4
def test_list_memories_delta_sync_self_managed(mock_workspace_client):
"""list() on DELTA_SYNC without model endpoint should use query_vector."""
mock_workspace_client.tables.exists.return_value = SimpleNamespace(table_exists=True)
inst = Databricks(
workspace_url="https://test",
access_token="tok",
endpoint_name="vs-endpoint",
catalog="catalog",
schema="schema",
table_name="table",
warehouse_name="test-warehouse",
index_type=VectorIndexType.DELTA_SYNC,
embedding_dimension=4,
# NOTE: no embedding_model_endpoint_name
)
mock_workspace_client.vector_search_indexes.query_index.return_value = SimpleNamespace(
result=SimpleNamespace(data_array=[])
)
inst.list(limit=5)
call_kwargs = mock_workspace_client.vector_search_indexes.query_index.call_args.kwargs
assert "query_vector" in call_kwargs
assert "query_text" not in call_kwargs
def test_list_memories_default_limit(db_instance_delta, mock_workspace_client):
"""list() with no limit should default to 100."""
mock_workspace_client.vector_search_indexes.query_index.return_value = SimpleNamespace(
result=SimpleNamespace(data_array=[])
)
db_instance_delta.list(limit=None)
call_kwargs = mock_workspace_client.vector_search_indexes.query_index.call_args.kwargs
assert call_kwargs["num_results"] == 100
# ---------------------- Table Creation Tests ---------------------- #
def test_ensure_source_table_uses_dynamic_names(mock_workspace_client):
"""Verify _ensure_source_table_exists uses self.fully_qualified_table_name and
self.table_name for the PK constraint, not hardcoded values."""
mock_workspace_client.tables.exists.return_value = SimpleNamespace(table_exists=False)
Databricks(
workspace_url="https://test",
access_token="tok",
endpoint_name="vs-endpoint",
catalog="my_catalog",
schema="my_schema",
table_name="my_memories",
collection_name="my_index",
warehouse_name="test-warehouse",
index_type=VectorIndexType.DELTA_SYNC,
embedding_model_endpoint_name="embedding-endpoint",
)
# _ensure_source_table_exists was called during __init__ via create_col
constraint_call = mock_workspace_client.table_constraints.create.call_args
assert constraint_call.kwargs["full_name_arg"] == "my_catalog.my_schema.my_memories"
pk_name = constraint_call.kwargs["constraint"].primary_key_constraint.name
assert pk_name == "pk_my_memories"
# ---------------------- Config Validation Tests ---------------------- #
def test_config_rejects_old_doc_params():
"""Config should reject the old documentation parameter names like index_name and source_table_name."""
from mem0.configs.vector_stores.databricks import DatabricksConfig
with pytest.raises(ValueError, match="Extra fields not allowed"):
DatabricksConfig(
workspace_url="https://test",
access_token="tok",
endpoint_name="ep",
catalog="cat",
schema="sch",
table_name="tbl",
index_name="catalog.schema.index", # old param from docs
)
def test_config_rejects_source_table_name():
"""Config should reject source_table_name which was in old docs."""
from mem0.configs.vector_stores.databricks import DatabricksConfig
with pytest.raises(ValueError, match="Extra fields not allowed"):
DatabricksConfig(
workspace_url="https://test",
access_token="tok",
endpoint_name="ep",
catalog="cat",
schema="sch",
table_name="tbl",
source_table_name="catalog.schema.table", # old param from docs
)
def test_config_accepts_correct_params():
"""Config should accept all the correct parameter names."""
from mem0.configs.vector_stores.databricks import DatabricksConfig
config = DatabricksConfig(
workspace_url="https://test",
access_token="tok",
endpoint_name="ep",
catalog="cat",
schema="sch",
table_name="tbl",
collection_name="my_index",
embedding_dimension=768,
)
assert config.collection_name == "my_index"
assert config.embedding_dimension == 768
# ---------------------- Reset Tests ---------------------- #
@@ -341,3 +597,244 @@ def test_reset(db_instance_delta, mock_workspace_client):
with patch.object(db_instance_delta, "create_col", wraps=db_instance_delta.create_col) as create_spy:
db_instance_delta.reset()
assert create_spy.called
# ---------------------- End-to-End Config → Factory → CRUD Tests ---------------------- #
def test_e2e_config_to_factory_delta_sync(mock_workspace_client):
"""End-to-end: VectorStoreConfig validates docs-correct params, factory creates Databricks instance."""
from mem0.vector_stores.configs import VectorStoreConfig
from mem0.utils.factory import VectorStoreFactory
# Step 1: Config validation (simulates what Memory.from_config does)
vs_config = VectorStoreConfig(
provider="databricks",
config={
"workspace_url": "https://my-workspace.databricks.com",
"access_token": "my-token",
"endpoint_name": "my-endpoint",
"catalog": "prod_catalog",
"schema": "ai_schema",
"table_name": "memories_table",
"collection_name": "my_index",
"embedding_dimension": 768,
"warehouse_name": "test-warehouse",
},
)
assert vs_config.config.collection_name == "my_index"
assert vs_config.config.catalog == "prod_catalog"
# Step 2: Factory instantiation (same as MemoryBase.__init__)
instance = VectorStoreFactory.create("databricks", vs_config.config)
assert isinstance(instance, Databricks)
assert instance.fully_qualified_table_name == "prod_catalog.ai_schema.memories_table"
assert instance.fully_qualified_index_name == "prod_catalog.ai_schema.my_index"
assert instance.embedding_dimension == 768
def test_e2e_config_to_factory_direct_access(mock_workspace_client):
"""End-to-end: DIRECT_ACCESS via config → factory creates correct instance."""
from mem0.vector_stores.configs import VectorStoreConfig
from mem0.utils.factory import VectorStoreFactory
mock_workspace_client.tables.exists.return_value = SimpleNamespace(table_exists=True)
vs_config = VectorStoreConfig(
provider="databricks",
config={
"workspace_url": "https://my-workspace.databricks.com",
"access_token": "my-token",
"endpoint_name": "my-endpoint",
"catalog": "cat",
"schema": "sch",
"table_name": "tbl",
"index_type": "DIRECT_ACCESS",
"embedding_dimension": 4,
"warehouse_name": "test-warehouse",
},
)
instance = VectorStoreFactory.create("databricks", vs_config.config)
assert isinstance(instance, Databricks)
assert "embedding" in instance.column_names
def test_e2e_old_docs_config_rejected():
"""End-to-end: Config from old docs (with index_name, source_table_name) is rejected at validation."""
from mem0.vector_stores.configs import VectorStoreConfig
with pytest.raises(ValueError, match="Extra fields not allowed"):
VectorStoreConfig(
provider="databricks",
config={
"workspace_url": "https://my-workspace.databricks.com",
"access_token": "my-token",
"endpoint_name": "my-endpoint",
"index_name": "catalog.schema.index_name",
"source_table_name": "catalog.schema.source_table",
"embedding_dimension": 1536,
},
)
def test_e2e_crud_lifecycle_delta_sync(mock_workspace_client):
"""End-to-end CRUD lifecycle: insert → search → get → list → update → delete."""
from mem0.vector_stores.configs import VectorStoreConfig
from mem0.utils.factory import VectorStoreFactory
vs_config = VectorStoreConfig(
provider="databricks",
config={
"workspace_url": "https://test",
"access_token": "tok",
"endpoint_name": "ep",
"catalog": "cat",
"schema": "sch",
"table_name": "tbl",
"warehouse_name": "test-warehouse",
"embedding_model_endpoint_name": "emb-ep",
},
)
db = VectorStoreFactory.create("databricks", vs_config.config)
# INSERT
db.insert(
vectors=[[0.1, 0.2]],
payloads=[{"data": "test memory", "user_id": "u1", "hash": "h1"}],
ids=["mem-001"],
)
insert_sql = mock_workspace_client.statement_execution.execute_statement.call_args.kwargs["statement"]
assert "INSERT INTO cat.sch.tbl" in insert_sql
assert "mem-001" in insert_sql
# SEARCH
mock_workspace_client.vector_search_indexes.query_index.return_value = SimpleNamespace(
result=SimpleNamespace(
data_array=[["mem-001", "h1", None, None, "u1", "test memory", None, None, None, 0.95]]
)
)
results = db.search(query="test", vectors=None, limit=5)
assert len(results) == 1
assert results[0].id == "mem-001"
assert results[0].payload["data"] == "test memory"
search_kwargs = mock_workspace_client.vector_search_indexes.query_index.call_args.kwargs
assert search_kwargs["query_text"] == "test"
# GET
mock_workspace_client.vector_search_indexes.query_index.return_value = QueryVectorIndexResponse(
manifest=ResultManifest(columns=[
ColumnInfo(name="memory_id"), ColumnInfo(name="hash"),
ColumnInfo(name="agent_id"), ColumnInfo(name="run_id"),
ColumnInfo(name="user_id"), ColumnInfo(name="memory"),
ColumnInfo(name="metadata"), ColumnInfo(name="created_at"),
ColumnInfo(name="updated_at"),
]),
result=ResultData(data_array=[["mem-001", "h1", None, None, "u1", "test memory", None, None, None]]),
)
got = db.get("mem-001")
assert got.id == "mem-001"
assert got.payload["data"] == "test memory"
get_kwargs = mock_workspace_client.vector_search_indexes.query_index.call_args.kwargs
assert "query_text" in get_kwargs
assert "query_vector" not in get_kwargs
# LIST
mock_workspace_client.vector_search_indexes.query_index.return_value = SimpleNamespace(
result=SimpleNamespace(
data_array=[["mem-001", "h1", None, None, "u1", "test memory", None, None, None]]
)
)
listed = db.list(filters={"user_id": "u1"}, limit=10)
assert len(listed[0]) == 1
assert listed[0][0].id == "mem-001"
list_kwargs = mock_workspace_client.vector_search_indexes.query_index.call_args.kwargs
assert "query_text" in list_kwargs
assert list_kwargs["num_results"] == 10
# UPDATE
db.update(vector_id="mem-001", payload={"memory": "updated memory"})
update_sql = mock_workspace_client.statement_execution.execute_statement.call_args.kwargs["statement"]
assert "UPDATE cat.sch.tbl" in update_sql
assert "mem-001" in update_sql
assert "updated memory" in update_sql
# DELETE
db.delete("mem-001")
delete_sql = mock_workspace_client.statement_execution.execute_statement.call_args.kwargs["statement"]
assert "DELETE FROM cat.sch.tbl" in delete_sql
assert "mem-001" in delete_sql
def test_e2e_crud_lifecycle_direct_access(mock_workspace_client):
"""End-to-end CRUD lifecycle for DIRECT_ACCESS: insert → search → get → list."""
from mem0.vector_stores.configs import VectorStoreConfig
from mem0.utils.factory import VectorStoreFactory
mock_workspace_client.tables.exists.return_value = SimpleNamespace(table_exists=True)
vs_config = VectorStoreConfig(
provider="databricks",
config={
"workspace_url": "https://test",
"access_token": "tok",
"endpoint_name": "ep",
"catalog": "cat",
"schema": "sch",
"table_name": "tbl",
"index_type": "DIRECT_ACCESS",
"embedding_dimension": 4,
"warehouse_name": "test-warehouse",
"embedding_model_endpoint_name": "emb-ep",
},
)
db = VectorStoreFactory.create("databricks", vs_config.config)
assert "embedding" in db.column_names
# INSERT with vector
db.insert(
vectors=[[0.1, 0.2, 0.3, 0.4]],
payloads=[{"data": "direct memory", "user_id": "u1", "hash": "h1"}],
ids=["mem-da-001"],
)
insert_sql = mock_workspace_client.statement_execution.execute_statement.call_args.kwargs["statement"]
assert "array(0.1, 0.2, 0.3, 0.4)" in insert_sql
# SEARCH with vector
mock_workspace_client.vector_search_indexes.query_index.return_value = SimpleNamespace(
result=SimpleNamespace(
data_array=[["mem-da-001", "h1", None, None, "u1", "direct memory", None, None, None, [0.1, 0.2, 0.3, 0.4], 0.9]]
)
)
results = db.search(query="", vectors=[0.1, 0.2, 0.3, 0.4], limit=5)
assert len(results) == 1
search_kwargs = mock_workspace_client.vector_search_indexes.query_index.call_args.kwargs
assert "query_vector" in search_kwargs
assert "query_text" not in search_kwargs
# GET — must use query_vector for DIRECT_ACCESS
mock_workspace_client.vector_search_indexes.query_index.return_value = QueryVectorIndexResponse(
manifest=ResultManifest(columns=[
ColumnInfo(name="memory_id"), ColumnInfo(name="hash"),
ColumnInfo(name="agent_id"), ColumnInfo(name="run_id"),
ColumnInfo(name="user_id"), ColumnInfo(name="memory"),
ColumnInfo(name="metadata"), ColumnInfo(name="created_at"),
ColumnInfo(name="updated_at"), ColumnInfo(name="embedding"),
]),
result=ResultData(data_array=[["mem-da-001", "h1", None, None, "u1", "direct memory", None, None, None, [0.1, 0.2, 0.3, 0.4]]]),
)
got = db.get("mem-da-001")
assert got.id == "mem-da-001"
get_kwargs = mock_workspace_client.vector_search_indexes.query_index.call_args.kwargs
assert "query_vector" in get_kwargs
assert get_kwargs["query_vector"] == [0.0] * 4
# LIST — must use query_vector for DIRECT_ACCESS
mock_workspace_client.vector_search_indexes.query_index.return_value = SimpleNamespace(
result=SimpleNamespace(
data_array=[["mem-da-001", "h1", None, None, "u1", "direct memory", None, None, None, [0.1, 0.2, 0.3, 0.4]]]
)
)
db.list(limit=5)
list_kwargs = mock_workspace_client.vector_search_indexes.query_index.call_args.kwargs
assert "query_vector" in list_kwargs
assert "query_text" not in list_kwargs
+55 -12
View File
@@ -48,17 +48,16 @@ def test_initalize_create_col(mongo_vector_fixture):
search_index_model = args[0].document
assert search_index_model == {
"name": "test_collection_vector_index",
"type": "vectorSearch",
"definition": {
"mappings": {
"dynamic": False,
"fields": {
"embedding": {
"type": "knnVector",
"dimensions": 1536,
"similarity": "cosine",
}
},
}
"fields": [
{
"type": "vector",
"path": "embedding",
"numDimensions": 1536,
"similarity": "cosine",
}
]
},
}
assert mongo_vector.collection == mock_collection
@@ -95,7 +94,7 @@ def test_search(mongo_vector_fixture):
"$vectorSearch": {
"index": "test_collection_vector_index",
"limit": 2,
"numCandidates": 2,
"numCandidates": 40,
"queryVector": query_vector,
"path": "embedding",
},
@@ -204,6 +203,10 @@ def test_delete(mongo_vector_fixture):
def test_update(mongo_vector_fixture):
"""
Test that update() uses dot notation for payload fields instead of replacing
the entire payload document.
"""
mongo_vector, mock_collection, _ = mongo_vector_fixture
vector_id = "id1"
updated_vector = [0.3] * 1536
@@ -213,11 +216,51 @@ def test_update(mongo_vector_fixture):
mongo_vector.update(vector_id=vector_id, vector=updated_vector, payload=updated_payload)
# Should use dot notation (payload.name) instead of full replacement (payload)
mock_collection.update_one.assert_called_once_with(
{"_id": vector_id}, {"$set": {"embedding": updated_vector, "payload": updated_payload}}
{"_id": vector_id}, {"$set": {"embedding": updated_vector, "payload.name": "updated_vector"}}
)
def test_update_payload_only_uses_dot_notation(mongo_vector_fixture):
"""
Test that updating only the payload uses dot notation for each field,
preserving existing metadata fields not included in the update.
"""
mongo_vector, mock_collection, _ = mongo_vector_fixture
mock_collection.update_one.return_value = MagicMock(matched_count=1)
mongo_vector.update(
vector_id="id1",
payload={"data": "updated text", "hash": "def456", "updated_at": "2025-06-01"},
)
set_arg = mock_collection.update_one.call_args[0][1]["$set"]
# Only the specified fields should be in $set, using dot notation
assert set_arg == {
"payload.data": "updated text",
"payload.hash": "def456",
"payload.updated_at": "2025-06-01",
}
# "payload" key itself should NOT appear (that would replace the whole document)
assert "payload" not in set_arg
def test_update_vector_only_does_not_touch_payload(mongo_vector_fixture):
"""Test that updating only the vector does not touch any payload fields."""
mongo_vector, mock_collection, _ = mongo_vector_fixture
mock_collection.update_one.return_value = MagicMock(matched_count=1)
mongo_vector.update(vector_id="id1", vector=[0.7, 0.8, 0.9])
set_arg = mock_collection.update_one.call_args[0][1]["$set"]
assert set_arg == {"embedding": [0.7, 0.8, 0.9]}
assert not any(k.startswith("payload.") for k in set_arg)
def test_get(mongo_vector_fixture):
mongo_vector, mock_collection, _ = mongo_vector_fixture
vector_id = "id1"
+139 -39
View File
@@ -18,25 +18,25 @@ from mem0.vector_stores.opensearch import OpenSearchDB
# Mock classes for testing OpenSearch with AWS authentication
class MockFieldInfo:
"""Mock pydantic field info."""
def __init__(self, default=None):
self.default = default
class MockOpenSearchConfig:
model_fields = {
'collection_name': MockFieldInfo(default="default_collection"),
'host': MockFieldInfo(default="localhost"),
'port': MockFieldInfo(default=9200),
'embedding_model_dims': MockFieldInfo(default=1536),
'http_auth': MockFieldInfo(default=None),
'auth': MockFieldInfo(default=None),
'credentials': MockFieldInfo(default=None),
'connection_class': MockFieldInfo(default=None),
'use_ssl': MockFieldInfo(default=False),
'verify_certs': MockFieldInfo(default=False),
"collection_name": MockFieldInfo(default="default_collection"),
"host": MockFieldInfo(default="localhost"),
"port": MockFieldInfo(default=9200),
"embedding_model_dims": MockFieldInfo(default=1536),
"http_auth": MockFieldInfo(default=None),
"auth": MockFieldInfo(default=None),
"credentials": MockFieldInfo(default=None),
"connection_class": MockFieldInfo(default=None),
"use_ssl": MockFieldInfo(default=False),
"verify_certs": MockFieldInfo(default=False),
}
def __init__(self, collection_name="test_collection", include_auth=True, **kwargs):
self.collection_name = collection_name
self.host = kwargs.get("host", "localhost")
@@ -44,7 +44,7 @@ class MockOpenSearchConfig:
self.embedding_model_dims = kwargs.get("embedding_model_dims", 1536)
self.use_ssl = kwargs.get("use_ssl", True)
self.verify_certs = kwargs.get("verify_certs", True)
if any(field in kwargs for field in ["http_auth", "auth", "credentials", "connection_class"]):
self.http_auth = kwargs.get("http_auth")
self.auth = kwargs.get("auth")
@@ -63,20 +63,18 @@ class MockOpenSearchConfig:
class MockAWSAuth:
def __init__(self):
self._lock = threading.Lock()
self.region = "us-east-1"
def __deepcopy__(self, memo):
raise TypeError("cannot pickle '_thread.lock' object")
class MockConnectionClass:
def __init__(self):
self._state = {"connected": False}
def __deepcopy__(self, memo):
raise TypeError("cannot pickle connection state")
@@ -255,6 +253,104 @@ class TestOpenSearchDB(unittest.TestCase):
self.os_db.delete_col()
self.client_mock.indices.delete.assert_called_once_with(index="test_collection")
def test_insert_rejects_null_vectors(self):
"""Vectors that are None should raise ValueError before hitting OpenSearch."""
vectors = [None]
payloads = [{"key": "value"}]
ids = ["id1"]
with self.assertRaises(ValueError) as ctx:
self.os_db.insert(vectors=vectors, payloads=payloads, ids=ids)
self.assertIn("null", str(ctx.exception).lower())
self.client_mock.index.assert_not_called()
def test_insert_rejects_empty_vectors(self):
"""Empty vector lists should raise ValueError."""
vectors = [[]]
payloads = [{"key": "value"}]
ids = ["id1"]
with self.assertRaises(ValueError) as ctx:
self.os_db.insert(vectors=vectors, payloads=payloads, ids=ids)
self.assertIn("empty", str(ctx.exception).lower())
self.client_mock.index.assert_not_called()
def test_insert_rejects_dimension_mismatch(self):
"""Vectors with wrong dimensions should raise ValueError with a clear message."""
vectors = [[0.1] * 768]
payloads = [{"key": "value"}]
ids = ["id1"]
with self.assertRaises(ValueError) as ctx:
self.os_db.insert(vectors=vectors, payloads=payloads, ids=ids)
error_msg = str(ctx.exception)
self.assertIn("768", error_msg)
self.assertIn("1536", error_msg)
self.client_mock.index.assert_not_called()
def test_update_rejects_empty_vector(self):
"""Update with an empty vector should raise ValueError."""
with self.assertRaises(ValueError) as ctx:
self.os_db.update("id1", vector=[], payload={"key": "value"})
self.assertIn("empty", str(ctx.exception).lower())
self.client_mock.search.assert_not_called()
def test_update_rejects_dimension_mismatch(self):
"""Update with wrong vector dimensions should raise ValueError."""
vector = [0.1] * 768
payload = {"key": "value"}
with self.assertRaises(ValueError) as ctx:
self.os_db.update("id1", vector=vector, payload=payload)
error_msg = str(ctx.exception)
self.assertIn("768", error_msg)
self.assertIn("1536", error_msg)
self.client_mock.search.assert_not_called()
@patch("mem0.vector_stores.opensearch.logger")
def test_update_error_logs_with_exc_info(self, mock_logger):
"""Update errors should log with exc_info and re-raise."""
mock_search_response = {"hits": {"hits": [{"_id": "doc1", "_source": {"id": "id1"}}]}}
self.client_mock.search.return_value = mock_search_response
self.client_mock.update.side_effect = Exception("Update failed")
with self.assertRaises(Exception):
self.os_db.update("id1", vector=[0.1] * 1536, payload={"key": "value"})
mock_logger.error.assert_called_once()
call_kwargs = mock_logger.error.call_args
self.assertTrue(call_kwargs[1].get("exc_info"), "logger.error must be called with exc_info=True")
@patch("mem0.vector_stores.opensearch.logger")
def test_insert_error_logs_with_exc_info(self, mock_logger):
"""Error logging should include exc_info for full stack trace."""
vectors = [[0.1] * 1536]
payloads = [{"key": "value"}]
ids = ["id1"]
self.client_mock.index.side_effect = Exception("Connection refused")
with self.assertRaises(Exception):
self.os_db.insert(vectors=vectors, payloads=payloads, ids=ids)
mock_logger.error.assert_called_once()
call_kwargs = mock_logger.error.call_args
self.assertTrue(call_kwargs[1].get("exc_info"), "logger.error must be called with exc_info=True")
@patch("mem0.vector_stores.opensearch.logger")
def test_search_error_logs_with_exc_info(self, mock_logger):
"""Search error logging should include exc_info for full stack trace."""
self.client_mock.search.side_effect = Exception("Search failed")
results = self.os_db.search(query="", vectors=[[0.1] * 1536], limit=5)
self.assertEqual(results, [])
mock_logger.error.assert_called_once()
call_kwargs = mock_logger.error.call_args
self.assertTrue(call_kwargs[1].get("exc_info"), "logger.error must be called with exc_info=True")
def test_init_with_http_auth(self):
mock_credentials = MagicMock()
mock_signer = AWSV4SignerAuth(mock_credentials, "us-east-1", "es")
@@ -282,11 +378,13 @@ class TestOpenSearchDB(unittest.TestCase):
# Tests for OpenSearch config deepcopy with AWS authentication (Issue #3464)
@patch('mem0.utils.factory.EmbedderFactory.create')
@patch('mem0.utils.factory.VectorStoreFactory.create')
@patch('mem0.utils.factory.LlmFactory.create')
@patch('mem0.memory.storage.SQLiteManager')
def test_safe_deepcopy_config_handles_opensearch_auth(mock_sqlite, mock_llm_factory, mock_vector_factory, mock_embedder_factory):
@patch("mem0.utils.factory.EmbedderFactory.create")
@patch("mem0.utils.factory.VectorStoreFactory.create")
@patch("mem0.utils.factory.LlmFactory.create")
@patch("mem0.memory.storage.SQLiteManager")
def test_safe_deepcopy_config_handles_opensearch_auth(
mock_sqlite, mock_llm_factory, mock_vector_factory, mock_embedder_factory
):
"""Test that _safe_deepcopy_config handles OpenSearch configs with AWS auth objects gracefully."""
mock_embedder_factory.return_value = MagicMock()
mock_vector_store = MagicMock()
@@ -295,9 +393,9 @@ def test_safe_deepcopy_config_handles_opensearch_auth(mock_sqlite, mock_llm_fact
mock_sqlite.return_value = MagicMock()
from mem0.memory.main import _safe_deepcopy_config
config_with_auth = MockOpenSearchConfig(collection_name="opensearch_test", include_auth=True)
safe_config = _safe_deepcopy_config(config_with_auth)
# Runtime auth objects must be preserved (Issue #3580)
@@ -306,7 +404,7 @@ def test_safe_deepcopy_config_handles_opensearch_auth(mock_sqlite, mock_llm_fact
assert safe_config.connection_class is not None
# Credentials dict is a sensitive secret and should be redacted
assert safe_config.credentials is None
assert safe_config.collection_name == "opensearch_test"
assert safe_config.host == "localhost"
assert safe_config.port == 9200
@@ -315,10 +413,10 @@ def test_safe_deepcopy_config_handles_opensearch_auth(mock_sqlite, mock_llm_fact
assert safe_config.verify_certs is True
@patch('mem0.utils.factory.EmbedderFactory.create')
@patch('mem0.utils.factory.VectorStoreFactory.create')
@patch('mem0.utils.factory.LlmFactory.create')
@patch('mem0.memory.storage.SQLiteManager')
@patch("mem0.utils.factory.EmbedderFactory.create")
@patch("mem0.utils.factory.VectorStoreFactory.create")
@patch("mem0.utils.factory.LlmFactory.create")
@patch("mem0.memory.storage.SQLiteManager")
def test_safe_deepcopy_config_normal_configs(mock_sqlite, mock_llm_factory, mock_vector_factory, mock_embedder_factory):
"""Test that _safe_deepcopy_config handles normal OpenSearch configs without auth."""
mock_embedder_factory.return_value = MagicMock()
@@ -328,12 +426,12 @@ def test_safe_deepcopy_config_normal_configs(mock_sqlite, mock_llm_factory, mock
mock_sqlite.return_value = MagicMock()
from mem0.memory.main import _safe_deepcopy_config
config_without_auth = MockOpenSearchConfig(collection_name="normal_test", include_auth=False)
safe_config = _safe_deepcopy_config(config_without_auth)
assert safe_config.collection_name == "normal_test"
assert safe_config.collection_name == "normal_test"
assert safe_config.host == "localhost"
assert safe_config.port == 9200
assert safe_config.embedding_model_dims == 1536
@@ -341,13 +439,15 @@ def test_safe_deepcopy_config_normal_configs(mock_sqlite, mock_llm_factory, mock
assert safe_config.verify_certs is True
@patch('mem0.utils.factory.EmbedderFactory.create')
@patch('mem0.utils.factory.VectorStoreFactory.create')
@patch('mem0.utils.factory.LlmFactory.create')
@patch('mem0.memory.storage.SQLiteManager')
def test_memory_initialization_opensearch_aws_auth(mock_sqlite, mock_llm_factory, mock_vector_factory, mock_embedder_factory):
@patch("mem0.utils.factory.EmbedderFactory.create")
@patch("mem0.utils.factory.VectorStoreFactory.create")
@patch("mem0.utils.factory.LlmFactory.create")
@patch("mem0.memory.storage.SQLiteManager")
def test_memory_initialization_opensearch_aws_auth(
mock_sqlite, mock_llm_factory, mock_vector_factory, mock_embedder_factory
):
"""Test that Memory initialization works with OpenSearch configs containing AWS auth."""
mock_embedder_factory.return_value = MagicMock()
mock_vector_store = MagicMock()
mock_vector_factory.return_value = mock_vector_store
+479 -1
View File
@@ -1,13 +1,21 @@
import os
import tempfile
import unittest
import uuid
from unittest.mock import MagicMock
from unittest.mock import MagicMock, patch
from qdrant_client import QdrantClient
from qdrant_client.models import (
Distance,
FieldCondition,
Filter,
MatchAny,
MatchExcept,
MatchText,
MatchValue,
PointIdsList,
PointStruct,
Range,
VectorParams,
)
@@ -25,6 +33,23 @@ class TestQdrant(unittest.TestCase):
on_disk=True,
)
def test_local_path_on_disk_false_preserves_existing_directory(self):
"""#4473: local path must not be removed when on_disk is False."""
with tempfile.TemporaryDirectory() as tmp:
sentinel = os.path.join(tmp, "sentinel")
with open(sentinel, "w", encoding="utf-8") as f:
f.write("keep")
mock_client = MagicMock()
mock_client.get_collections.return_value = MagicMock(collections=[])
with patch("mem0.vector_stores.qdrant.QdrantClient", return_value=mock_client):
Qdrant(
collection_name="c",
embedding_model_dims=128,
path=tmp,
on_disk=False,
)
self.assertTrue(os.path.isfile(sentinel))
def test_create_col(self):
self.client_mock.get_collections.return_value = MagicMock(collections=[])
@@ -298,3 +323,456 @@ class TestQdrant(unittest.TestCase):
def tearDown(self):
del self.qdrant
class TestQdrantEnhancedFilters(unittest.TestCase):
"""Tests for enhanced metadata filtering operators (issue #3975)."""
def setUp(self):
self.client_mock = MagicMock(spec=QdrantClient)
self.qdrant = Qdrant(
collection_name="test_collection",
embedding_model_dims=128,
client=self.client_mock,
)
# ------------------------------------------------------------------ #
# _build_field_condition #
# ------------------------------------------------------------------ #
def test_simple_equality(self):
"""Plain value maps to MatchValue."""
cond = self.qdrant._build_field_condition("category", "programming")
self.assertIsInstance(cond, FieldCondition)
self.assertEqual(cond.key, "category")
self.assertIsInstance(cond.match, MatchValue)
self.assertEqual(cond.match.value, "programming")
def test_eq_operator(self):
"""{\"eq\": v} maps to MatchValue."""
cond = self.qdrant._build_field_condition("category", {"eq": "programming"})
self.assertIsInstance(cond.match, MatchValue)
self.assertEqual(cond.match.value, "programming")
def test_ne_operator(self):
"""{\"ne\": v} maps to MatchExcept with a single-element list."""
cond = self.qdrant._build_field_condition("category", {"ne": "spam"})
self.assertIsInstance(cond.match, MatchExcept)
self.assertIn("spam", getattr(cond.match, "except_", None) or getattr(cond.match, "except"))
def test_in_operator(self):
"""{\"in\": [...]} maps to MatchAny."""
cond = self.qdrant._build_field_condition("category", {"in": ["a", "b", "c"]})
self.assertIsInstance(cond.match, MatchAny)
self.assertEqual(cond.match.any, ["a", "b", "c"])
def test_nin_operator(self):
"""{\"nin\": [...]} maps to MatchExcept."""
cond = self.qdrant._build_field_condition("status", {"nin": ["deleted", "banned"]})
self.assertIsInstance(cond.match, MatchExcept)
excluded = getattr(cond.match, "except_", None) or getattr(cond.match, "except")
self.assertIn("deleted", excluded)
self.assertIn("banned", excluded)
def test_contains_operator(self):
"""{\"contains\": x} maps to MatchText."""
cond = self.qdrant._build_field_condition("bio", {"contains": "python"})
self.assertIsInstance(cond.match, MatchText)
self.assertEqual(cond.match.text, "python")
def test_icontains_operator(self):
"""{\"icontains\": x} maps to MatchText (same as contains for Qdrant)."""
cond = self.qdrant._build_field_condition("bio", {"icontains": "Python"})
self.assertIsInstance(cond.match, MatchText)
self.assertEqual(cond.match.text, "Python")
def test_gt_operator(self):
"""{\"gt\": v} maps to Range with gt only."""
cond = self.qdrant._build_field_condition("priority", {"gt": 5})
self.assertIsInstance(cond.range, Range)
self.assertEqual(cond.range.gt, 5)
self.assertIsNone(cond.range.gte)
self.assertIsNone(cond.range.lt)
self.assertIsNone(cond.range.lte)
def test_gte_operator(self):
"""{\"gte\": v} maps to Range with gte only."""
cond = self.qdrant._build_field_condition("priority", {"gte": 5})
self.assertIsInstance(cond.range, Range)
self.assertEqual(cond.range.gte, 5)
self.assertIsNone(cond.range.gt)
def test_lt_operator(self):
"""{\"lt\": v} maps to Range with lt only."""
cond = self.qdrant._build_field_condition("priority", {"lt": 10})
self.assertIsInstance(cond.range, Range)
self.assertEqual(cond.range.lt, 10)
def test_lte_operator(self):
"""{\"lte\": v} maps to Range with lte only."""
cond = self.qdrant._build_field_condition("priority", {"lte": 10})
self.assertIsInstance(cond.range, Range)
self.assertEqual(cond.range.lte, 10)
def test_range_gte_lte(self):
"""Combined {\"gte\": x, \"lte\": y} maps to Range with both bounds."""
cond = self.qdrant._build_field_condition("priority", {"gte": 5, "lte": 10})
self.assertIsInstance(cond.range, Range)
self.assertEqual(cond.range.gte, 5)
self.assertEqual(cond.range.lte, 10)
def test_range_gt_lt(self):
"""Open interval {\"gt\": x, \"lt\": y} maps to Range with both bounds."""
cond = self.qdrant._build_field_condition("score", {"gt": 0.5, "lt": 0.9})
self.assertIsInstance(cond.range, Range)
self.assertEqual(cond.range.gt, 0.5)
self.assertEqual(cond.range.lt, 0.9)
# ------------------------------------------------------------------ #
# _create_filter — comparison and list operators #
# ------------------------------------------------------------------ #
def test_create_filter_eq(self):
result = self.qdrant._create_filter({"category": {"eq": "programming"}})
self.assertIsInstance(result, Filter)
self.assertEqual(len(result.must), 1)
self.assertIsInstance(result.must[0].match, MatchValue)
def test_create_filter_ne(self):
result = self.qdrant._create_filter({"category": {"ne": "spam"}})
self.assertIsInstance(result, Filter)
self.assertIsInstance(result.must[0].match, MatchExcept)
def test_create_filter_in(self):
result = self.qdrant._create_filter({"category": {"in": ["prog", "data"]}})
self.assertIsInstance(result, Filter)
self.assertIsInstance(result.must[0].match, MatchAny)
self.assertEqual(result.must[0].match.any, ["prog", "data"])
def test_create_filter_nin(self):
result = self.qdrant._create_filter({"status": {"nin": ["deleted", "banned"]}})
self.assertIsInstance(result, Filter)
self.assertIsInstance(result.must[0].match, MatchExcept)
def test_create_filter_gt(self):
result = self.qdrant._create_filter({"priority": {"gt": 7}})
self.assertIsInstance(result, Filter)
self.assertIsInstance(result.must[0].range, Range)
self.assertEqual(result.must[0].range.gt, 7)
def test_create_filter_gte_lte_range(self):
result = self.qdrant._create_filter({"priority": {"gte": 5, "lte": 9}})
self.assertIsInstance(result, Filter)
r = result.must[0].range
self.assertEqual(r.gte, 5)
self.assertEqual(r.lte, 9)
def test_create_filter_contains(self):
result = self.qdrant._create_filter({"bio": {"contains": "python"}})
self.assertIsInstance(result, Filter)
self.assertIsInstance(result.must[0].match, MatchText)
# ------------------------------------------------------------------ #
# _create_filter — logical operators #
# ------------------------------------------------------------------ #
def test_and_operator(self):
"""AND populates Filter.must with nested Filter objects."""
filters = {"AND": [{"category": "programming"}, {"priority": 10}]}
result = self.qdrant._create_filter(filters)
self.assertIsInstance(result, Filter)
self.assertEqual(len(result.must), 2)
# Each item is a Filter that wraps a single FieldCondition
for item in result.must:
self.assertIsInstance(item, Filter)
def test_or_operator(self):
"""OR populates Filter.should with nested Filter objects."""
filters = {"OR": [{"category": "programming"}, {"category": "data"}]}
result = self.qdrant._create_filter(filters)
self.assertIsInstance(result, Filter)
self.assertEqual(len(result.should), 2)
self.assertIsNone(result.must)
def test_not_operator(self):
"""NOT populates Filter.must_not with nested Filter objects."""
filters = {"NOT": [{"category": "spam"}]}
result = self.qdrant._create_filter(filters)
self.assertIsInstance(result, Filter)
self.assertEqual(len(result.must_not), 1)
self.assertIsNone(result.must)
self.assertIsNone(result.should)
def test_mixed_field_and_and(self):
"""Top-level field conditions and AND can coexist in must."""
filters = {
"user_id": "alice",
"AND": [{"priority": {"gte": 5}}, {"category": "programming"}],
}
result = self.qdrant._create_filter(filters)
self.assertIsInstance(result, Filter)
# user_id FieldCondition + 2 nested AND Filters = 3 items in must
self.assertEqual(len(result.must), 3)
def test_nested_and_or(self):
"""AND containing an OR sub-condition produces correct nesting."""
filters = {
"AND": [
{"OR": [{"category": "prog"}, {"category": "data"}]},
{"priority": {"gt": 3}},
]
}
result = self.qdrant._create_filter(filters)
self.assertIsInstance(result, Filter)
self.assertEqual(len(result.must), 2)
# First item is a Filter with should (the OR)
or_filter = result.must[0]
self.assertIsInstance(or_filter, Filter)
self.assertEqual(len(or_filter.should), 2)
# ------------------------------------------------------------------ #
# Edge cases #
# ------------------------------------------------------------------ #
def test_empty_filter_returns_none(self):
self.assertIsNone(self.qdrant._create_filter({}))
self.assertIsNone(self.qdrant._create_filter(None))
def test_backward_compat_simple_equality(self):
"""Plain scalar equality still works as before."""
result = self.qdrant._create_filter({"user_id": "alice", "run_id": "r1"})
self.assertIsInstance(result, Filter)
keys = [c.key for c in result.must]
self.assertIn("user_id", keys)
self.assertIn("run_id", keys)
def test_backward_compat_gte_lte_range(self):
"""Original gte+lte range filter still produces a Range condition."""
result = self.qdrant._create_filter({"count": {"gte": 5, "lte": 10}})
self.assertIsInstance(result, Filter)
self.assertIsInstance(result.must[0].range, Range)
self.assertEqual(result.must[0].range.gte, 5)
self.assertEqual(result.must[0].range.lte, 10)
def test_unknown_operator_raises_error(self):
"""Unknown operator dict should raise ValueError, not ValidationError."""
with self.assertRaises(ValueError) as ctx:
self.qdrant._build_field_condition("field", {"unknown_op": "foo"})
self.assertIn("Unsupported", str(ctx.exception))
def test_mixed_range_and_non_range_raises_error(self):
"""Mixing range ops with non-range ops should raise ValueError."""
with self.assertRaises(ValueError):
self.qdrant._build_field_condition("priority", {"gte": 5, "ne": 10})
def test_mixed_range_and_eq_raises_error(self):
"""Mixing range ops with eq should raise ValueError."""
with self.assertRaises(ValueError):
self.qdrant._build_field_condition("score", {"gt": 0.5, "eq": 1.0})
def test_wildcard_returns_none(self):
"""Wildcard '*' should return None (skip filter — match any)."""
result = self.qdrant._build_field_condition("category", "*")
self.assertIsNone(result)
def test_create_filter_with_wildcard_skips_it(self):
"""Wildcard fields should be skipped in the final filter."""
result = self.qdrant._create_filter({"category": "*", "user_id": "alice"})
self.assertIsInstance(result, Filter)
self.assertEqual(len(result.must), 1)
self.assertEqual(result.must[0].key, "user_id")
def test_create_filter_only_wildcard_returns_none(self):
"""Filter with only wildcard should return None."""
result = self.qdrant._create_filter({"category": "*"})
self.assertIsNone(result)
def test_and_with_non_list_raises_error(self):
"""AND with non-list value should raise ValueError."""
with self.assertRaises(ValueError):
self.qdrant._create_filter({"AND": "not_a_list"})
def test_or_with_non_list_raises_error(self):
"""OR with non-list value should raise ValueError."""
with self.assertRaises(ValueError):
self.qdrant._create_filter({"OR": {"category": "work"}})
def test_not_with_non_list_raises_error(self):
"""NOT with non-list value should raise ValueError."""
with self.assertRaises(ValueError):
self.qdrant._create_filter({"NOT": "invalid"})
def test_empty_and_list(self):
"""AND with empty list should return None."""
result = self.qdrant._create_filter({"AND": []})
self.assertIsNone(result)
def test_empty_or_list(self):
"""OR with empty list should return None."""
result = self.qdrant._create_filter({"OR": []})
self.assertIsNone(result)
def test_empty_not_list(self):
"""NOT with empty list should return None."""
result = self.qdrant._create_filter({"NOT": []})
self.assertIsNone(result)
def test_deeply_nested_logical(self):
"""3-level nesting: AND > OR > NOT."""
filters = {
"AND": [
{
"OR": [
{"category": "work"},
{"category": "personal"}
]
},
{"priority": {"gte": 5}},
{
"NOT": [
{"status": "archived"}
]
}
]
}
result = self.qdrant._create_filter(filters)
self.assertIsInstance(result, Filter)
self.assertEqual(len(result.must), 3)
def test_boolean_equality(self):
"""Boolean values should work with MatchValue."""
cond = self.qdrant._build_field_condition("active", True)
self.assertIsInstance(cond.match, MatchValue)
self.assertEqual(cond.match.value, True)
def test_integer_equality(self):
"""Integer values should work with MatchValue."""
cond = self.qdrant._build_field_condition("count", 42)
self.assertIsInstance(cond.match, MatchValue)
self.assertEqual(cond.match.value, 42)
def test_eq_with_boolean(self):
"""eq operator with boolean should work."""
cond = self.qdrant._build_field_condition("active", {"eq": False})
self.assertIsInstance(cond.match, MatchValue)
self.assertEqual(cond.match.value, False)
def test_in_with_integers(self):
"""in operator with integer list should use MatchAny."""
cond = self.qdrant._build_field_condition("priority", {"in": [1, 2, 3]})
self.assertIsInstance(cond.match, MatchAny)
self.assertEqual(cond.match.any, [1, 2, 3])
def test_wildcard_inside_and(self):
"""Wildcard inside AND should be skipped, other conditions preserved."""
result = self.qdrant._create_filter({
"AND": [{"category": "*"}, {"user_id": "alice"}]
})
self.assertIsInstance(result, Filter)
# AND produces nested Filters; the wildcard sub-filter returns None and is skipped
# Only the user_id sub-filter remains
self.assertEqual(len(result.must), 1)
def test_list_value_treated_as_match_any(self):
"""List value shorthand should be treated as in-operator (MatchAny)."""
cond = self.qdrant._build_field_condition("tags", ["a", "b", "c"])
self.assertIsInstance(cond.match, MatchAny)
self.assertEqual(cond.match.any, ["a", "b", "c"])
def test_list_value_in_create_filter(self):
"""List value shorthand should work through _create_filter too."""
result = self.qdrant._create_filter({"tags": ["python", "rust"]})
self.assertIsInstance(result, Filter)
self.assertEqual(len(result.must), 1)
self.assertIsInstance(result.must[0].match, MatchAny)
def test_non_dict_item_in_and_raises_error(self):
"""Non-dict item inside AND list should raise clear ValueError."""
with self.assertRaises(ValueError) as ctx:
self.qdrant._create_filter({"AND": ["not_a_dict", {"field": "val"}]})
self.assertIn("index 0", str(ctx.exception))
self.assertIn("must be a dict", str(ctx.exception))
def test_non_dict_item_in_or_raises_error(self):
"""Non-dict item inside OR list should raise clear ValueError."""
with self.assertRaises(ValueError) as ctx:
self.qdrant._create_filter({"OR": [{"field": "val"}, 42]})
self.assertIn("index 1", str(ctx.exception))
def test_empty_dict_value_raises_error(self):
"""Empty dict as filter value should raise ValueError."""
with self.assertRaises(ValueError):
self.qdrant._build_field_condition("field", {})
def test_ne_alone_does_not_trigger_mixed_error(self):
"""ne as sole operator should NOT trigger mixed-operator error (regression guard)."""
cond = self.qdrant._build_field_condition("status", {"ne": "deleted"})
self.assertIsInstance(cond.match, MatchExcept)
def test_eq_with_literal_star(self):
"""eq operator with literal '*' should match the string '*', not wildcard."""
cond = self.qdrant._build_field_condition("category", {"eq": "*"})
self.assertIsInstance(cond.match, MatchValue)
self.assertEqual(cond.match.value, "*")
# ------------------------------------------------------------------ #
# $or / $not normalization (Memory._process_metadata_filters injects #
# these keys alongside the original OR/NOT) #
# ------------------------------------------------------------------ #
def test_dollar_or_handled_as_or(self):
"""$or injected by Memory middleware should be treated as OR."""
filters = {
"user_id": "alice",
"$or": [{"category": "programming"}, {"category": "data"}],
}
result = self.qdrant._create_filter(filters)
self.assertIsInstance(result, Filter)
self.assertIsNotNone(result.should)
self.assertEqual(len(result.should), 2)
def test_dollar_not_handled_as_not(self):
"""$not injected by Memory middleware should be treated as NOT."""
filters = {
"user_id": "alice",
"$not": [{"category": "spam"}],
}
result = self.qdrant._create_filter(filters)
self.assertIsInstance(result, Filter)
self.assertIsNotNone(result.must_not)
self.assertEqual(len(result.must_not), 1)
def test_memory_search_or_shape(self):
"""Simulate exact shape Memory.search() sends for OR filters.
effective_filters keeps the original OR key (via deepcopy of
input_filters) and _process_metadata_filters adds $or with the
same content. _create_filter should deduplicate so only the
first occurrence (OR) is used — exactly 2 should entries.
"""
filters = {
"OR": [{"category": "programming"}, {"category": "data"}],
"user_id": "test_user",
"$or": [{"category": "programming"}, {"category": "data"}],
}
result = self.qdrant._create_filter(filters)
self.assertIsInstance(result, Filter)
self.assertIsNotNone(result.should)
# Deduplicated: OR wins, $or is skipped — exactly 2 entries
self.assertEqual(len(result.should), 2)
def test_memory_search_not_shape(self):
"""Simulate exact shape Memory.search() sends for NOT filters.
Same deduplication as OR: NOT wins, $not is skipped.
"""
filters = {
"NOT": [{"category": "spam"}],
"user_id": "test_user",
"$not": [{"category": "spam"}],
}
result = self.qdrant._create_filter(filters)
self.assertIsInstance(result, Filter)
self.assertIsNotNone(result.must_not)
# Deduplicated: NOT wins, $not is skipped — exactly 1 entry
self.assertEqual(len(result.must_not), 1)
+725
View File
@@ -0,0 +1,725 @@
from datetime import datetime
from unittest.mock import MagicMock, patch
import pytest
pytest.importorskip("turbopuffer", reason="turbopuffer not installed")
from turbopuffer.types import Row
from mem0.vector_stores.turbopuffer import OutputData, TurbopufferDB
def _make_row(id, dist=None, vector=None, **attributes):
"""Helper to create a turbopuffer Row with extra attributes."""
data = {"id": id, "vector": vector}
if dist is not None:
data["$dist"] = dist
data.update(attributes)
return Row.model_validate(data)
@pytest.fixture
def mock_client():
with patch("mem0.vector_stores.turbopuffer.TurbopufferClient") as MockClient:
client_instance = MagicMock()
namespace_instance = MagicMock()
client_instance.namespace.return_value = namespace_instance
MockClient.return_value = client_instance
yield client_instance, namespace_instance
@pytest.fixture
def db(mock_client):
_, _ = mock_client
return TurbopufferDB(
collection_name="test_ns",
embedding_model_dims=4,
api_key="tpuf_test_key",
region="gcp-us-central1",
distance_metric="cosine_distance",
batch_size=2,
)
# ── Initialization ──────────────────────────────────────────────────
class TestInit:
def test_init_with_api_key(self, mock_client):
client_instance, namespace_instance = mock_client
db = TurbopufferDB(
collection_name="my_ns",
embedding_model_dims=128,
api_key="tpuf_key",
region="gcp-us-central1",
)
assert db.collection_name == "my_ns"
assert db.embedding_model_dims == 128
assert db.distance_metric == "cosine_distance"
assert db.batch_size == 100
client_instance.namespace.assert_called_with("my_ns")
def test_init_with_env_var(self, mock_client):
with patch.dict("os.environ", {"TURBOPUFFER_API_KEY": "tpuf_env_key"}):
db = TurbopufferDB(
collection_name="test",
embedding_model_dims=4,
region="gcp-us-central1",
)
assert db.collection_name == "test"
def test_init_without_api_key_raises(self, mock_client):
with patch.dict("os.environ", {}, clear=True):
# Remove TURBOPUFFER_API_KEY if it exists
import os
os.environ.pop("TURBOPUFFER_API_KEY", None)
with pytest.raises(ValueError, match="API key must be provided"):
TurbopufferDB(
collection_name="test",
embedding_model_dims=4,
region="gcp-us-central1",
)
def test_init_with_extra_params(self, mock_client):
client_instance, _ = mock_client
with patch("mem0.vector_stores.turbopuffer.TurbopufferClient") as MockClient:
MockClient.return_value = client_instance
TurbopufferDB(
collection_name="test",
embedding_model_dims=4,
api_key="key",
region="aws-us-west-2",
extra_params={"compression": True},
)
MockClient.assert_called_with(
api_key="key",
region="aws-us-west-2",
compression=True,
)
def test_init_default_region(self, mock_client):
client_instance, _ = mock_client
with patch("mem0.vector_stores.turbopuffer.TurbopufferClient") as MockClient:
MockClient.return_value = client_instance
TurbopufferDB(
collection_name="test",
embedding_model_dims=4,
api_key="key",
)
MockClient.assert_called_with(
api_key="key",
region="gcp-us-central1",
)
# ── create_col ───────────────────────────────────────────────────────
class TestCreateCol:
def test_create_col_is_noop(self, db):
# Should not raise or call anything
result = db.create_col()
assert result is None
def test_create_col_with_args_is_noop(self, db):
result = db.create_col(name="x", vector_size=128, distance="cosine")
assert result is None
# ── insert ───────────────────────────────────────────────────────────
class TestInsert:
def test_insert_with_ids_and_payloads(self, db):
vectors = [[0.1, 0.2, 0.3, 0.4], [0.5, 0.6, 0.7, 0.8]]
payloads = [{"data": "hello", "user_id": "u1"}, {"data": "world", "user_id": "u2"}]
ids = ["id1", "id2"]
db.insert(vectors, payloads, ids)
db.namespace.write.assert_called_once()
call_kwargs = db.namespace.write.call_args[1]
rows = call_kwargs["upsert_rows"]
assert len(rows) == 2
assert rows[0]["id"] == "id1"
assert rows[0]["vector"] == [0.1, 0.2, 0.3, 0.4]
assert rows[0]["data"] == "hello"
assert rows[0]["user_id"] == "u1"
assert rows[1]["id"] == "id2"
assert call_kwargs["distance_metric"] == "cosine_distance"
def test_insert_without_ids_generates_ids(self, db):
vectors = [[0.1, 0.2, 0.3, 0.4]]
db.insert(vectors)
db.namespace.write.assert_called_once()
call_kwargs = db.namespace.write.call_args[1]
assert call_kwargs["upsert_rows"][0]["id"] == "0"
def test_insert_without_payloads(self, db):
vectors = [[0.1, 0.2, 0.3, 0.4]]
ids = ["id1"]
db.insert(vectors, ids=ids)
call_kwargs = db.namespace.write.call_args[1]
row = call_kwargs["upsert_rows"][0]
assert row["id"] == "id1"
assert row["vector"] == [0.1, 0.2, 0.3, 0.4]
assert len(row) == 2 # only id and vector
def test_insert_payload_does_not_overwrite_id_or_vector(self, db):
"""Payload with 'id' or 'vector' keys must not overwrite the actual values."""
vectors = [[0.1, 0.2, 0.3, 0.4]]
payloads = [{"id": "fake_id", "vector": "fake_vector", "data": "hello"}]
ids = ["real_id"]
db.insert(vectors, payloads, ids)
call_kwargs = db.namespace.write.call_args[1]
row = call_kwargs["upsert_rows"][0]
assert row["id"] == "real_id"
assert row["vector"] == [0.1, 0.2, 0.3, 0.4]
assert row["data"] == "hello"
def test_insert_batching(self, db):
"""batch_size=2, so 3 vectors should produce 2 write calls."""
vectors = [[0.1] * 4, [0.2] * 4, [0.3] * 4]
ids = ["a", "b", "c"]
db.insert(vectors, ids=ids)
assert db.namespace.write.call_count == 2
first_call = db.namespace.write.call_args_list[0][1]
second_call = db.namespace.write.call_args_list[1][1]
assert len(first_call["upsert_rows"]) == 2
assert len(second_call["upsert_rows"]) == 1
# ── _parse_output ────────────────────────────────────────────────────
class TestParseOutput:
def test_parse_rows_with_dist_and_attributes(self, db):
rows = [
_make_row("id1", dist=0.1, data="hello", user_id="u1"),
_make_row("id2", dist=0.3, data="world", user_id="u2"),
]
results = db._parse_output(rows)
assert len(results) == 2
assert results[0].id == "id1"
assert results[0].score == pytest.approx(0.9)
assert results[0].payload == {"data": "hello", "user_id": "u1"}
assert results[1].id == "id2"
assert results[1].score == pytest.approx(0.7)
assert results[1].payload == {"data": "world", "user_id": "u2"}
def test_parse_rows_without_dist(self, db):
rows = [_make_row("id1", data="hello")]
results = db._parse_output(rows)
assert results[0].score is None
assert results[0].payload == {"data": "hello"}
def test_parse_rows_strips_vector_and_id(self, db):
rows = [_make_row("id1", dist=0.0, vector=[1.0, 2.0, 3.0, 4.0], data="test")]
results = db._parse_output(rows)
assert "vector" not in results[0].payload
assert "id" not in results[0].payload
assert "$dist" not in results[0].payload
assert results[0].id == "id1"
def test_parse_empty_rows(self, db):
assert db._parse_output([]) == []
# ── _convert_filters ─────────────────────────────────────────────────
class TestConvertFilters:
def test_none_filters(self, db):
assert db._convert_filters(None) is None
def test_empty_filters(self, db):
assert db._convert_filters({}) is None
def test_single_eq_filter(self, db):
result = db._convert_filters({"user_id": "u1"})
assert result == ("user_id", "Eq", "u1")
def test_multiple_eq_filters(self, db):
result = db._convert_filters({"user_id": "u1", "agent_id": "a1"})
assert result[0] == "And"
conditions = result[1]
assert ("user_id", "Eq", "u1") in conditions
assert ("agent_id", "Eq", "a1") in conditions
def test_range_filter_gte_lte(self, db):
result = db._convert_filters({"score": {"gte": 0.5, "lte": 1.0}})
assert result[0] == "And"
conditions = result[1]
assert ("score", "Gte", 0.5) in conditions
assert ("score", "Lte", 1.0) in conditions
def test_range_filter_gte_only(self, db):
result = db._convert_filters({"score": {"gte": 0.5}})
assert result == ("score", "Gte", 0.5)
def test_range_filter_lte_only(self, db):
result = db._convert_filters({"score": {"lte": 1.0}})
assert result == ("score", "Lte", 1.0)
def test_mixed_eq_and_range_filters(self, db):
result = db._convert_filters({"user_id": "u1", "score": {"gte": 0.5}})
assert result[0] == "And"
conditions = result[1]
assert ("user_id", "Eq", "u1") in conditions
assert ("score", "Gte", 0.5) in conditions
# ── search ───────────────────────────────────────────────────────────
class TestSearch:
def test_search_basic(self, db):
mock_response = MagicMock()
mock_response.rows = [
_make_row("id1", dist=0.1, data="hello"),
_make_row("id2", dist=0.2, data="world"),
]
db.namespace.query.return_value = mock_response
results = db.search("test query", [0.1, 0.2, 0.3, 0.4], limit=2)
db.namespace.query.assert_called_once_with(
rank_by=("vector", "ANN", [0.1, 0.2, 0.3, 0.4]),
top_k=2,
include_attributes=True,
)
assert len(results) == 2
assert results[0].id == "id1"
assert results[0].score == pytest.approx(0.9)
def test_search_with_filters(self, db):
mock_response = MagicMock()
mock_response.rows = [_make_row("id1", dist=0.1, data="hello")]
db.namespace.query.return_value = mock_response
results = db.search(
"query", [0.1, 0.2, 0.3, 0.4], limit=5, filters={"user_id": "u1"}
)
call_kwargs = db.namespace.query.call_args[1]
assert call_kwargs["filters"] == ("user_id", "Eq", "u1")
assert len(results) == 1
def test_search_no_results(self, db):
mock_response = MagicMock()
mock_response.rows = None
db.namespace.query.return_value = mock_response
results = db.search("query", [0.1, 0.2, 0.3, 0.4])
assert results == []
def test_search_empty_rows(self, db):
mock_response = MagicMock()
mock_response.rows = []
db.namespace.query.return_value = mock_response
results = db.search("query", [0.1, 0.2, 0.3, 0.4])
assert results == []
# ── delete ───────────────────────────────────────────────────────────
class TestDelete:
def test_delete_by_string_id(self, db):
db.delete("id1")
db.namespace.write.assert_called_once_with(deletes=["id1"])
def test_delete_by_int_id(self, db):
db.delete(123)
db.namespace.write.assert_called_once_with(deletes=["123"])
# ── update ───────────────────────────────────────────────────────────
class TestUpdate:
def test_update_with_vector_and_payload(self, db):
db.update("id1", vector=[0.5, 0.6, 0.7, 0.8], payload={"data": "updated"})
db.namespace.write.assert_called_once()
call_kwargs = db.namespace.write.call_args[1]
row = call_kwargs["upsert_rows"][0]
assert row["id"] == "id1"
assert row["vector"] == [0.5, 0.6, 0.7, 0.8]
assert row["data"] == "updated"
assert call_kwargs["distance_metric"] == "cosine_distance"
def test_update_vector_only(self, db):
db.update("id1", vector=[0.5, 0.6, 0.7, 0.8])
call_kwargs = db.namespace.write.call_args[1]
row = call_kwargs["upsert_rows"][0]
assert row["id"] == "id1"
assert row["vector"] == [0.5, 0.6, 0.7, 0.8]
assert "data" not in row
def test_update_payload_only_uses_patch(self, db):
"""Payload-only updates should use patch_rows, not upsert_rows."""
db.update("id1", vector=None, payload={"data": "patched", "user_id": "u1"})
call_kwargs = db.namespace.write.call_args[1]
row = call_kwargs["patch_rows"][0]
assert row["id"] == "id1"
assert row["data"] == "patched"
assert row["user_id"] == "u1"
assert "upsert_rows" not in call_kwargs
def test_update_payload_does_not_overwrite_id(self, db):
"""Payload with an 'id' key must not overwrite the actual vector ID."""
db.update("real_id", vector=None, payload={"id": "fake_id", "data": "test"})
call_kwargs = db.namespace.write.call_args[1]
row = call_kwargs["patch_rows"][0]
assert row["id"] == "real_id"
def test_update_nothing(self, db):
"""Neither vector nor payload: should not call write."""
db.update("id1")
db.namespace.write.assert_not_called()
# ── get ──────────────────────────────────────────────────────────────
class TestGet:
def test_get_found(self, db):
mock_response = MagicMock()
mock_response.rows = [_make_row("id1", dist=0.0, data="hello", user_id="u1")]
db.namespace.query.return_value = mock_response
result = db.get("id1")
db.namespace.query.assert_called_once_with(
top_k=1,
rank_by=("vector", "ANN", [0.0, 0.0, 0.0, 0.0]),
filters=("id", "Eq", "id1"),
include_attributes=True,
)
assert result is not None
assert result.id == "id1"
assert result.payload["data"] == "hello"
assert result.payload["user_id"] == "u1"
def test_get_not_found(self, db):
mock_response = MagicMock()
mock_response.rows = []
db.namespace.query.return_value = mock_response
result = db.get("nonexistent")
assert result is None
def test_get_none_rows(self, db):
mock_response = MagicMock()
mock_response.rows = None
db.namespace.query.return_value = mock_response
result = db.get("id1")
assert result is None
def test_get_handles_exception(self, db):
db.namespace.query.side_effect = Exception("API error")
result = db.get("id1")
assert result is None
# ── list_cols ────────────────────────────────────────────────────────
class TestListCols:
def test_list_cols(self, db):
ns1 = MagicMock()
ns1.id = "ns1"
ns2 = MagicMock()
ns2.id = "ns2"
db.client.namespaces.return_value = [ns1, ns2]
result = db.list_cols()
assert len(result) == 2
db.client.namespaces.assert_called_once()
# ── delete_col ───────────────────────────────────────────────────────
class TestDeleteCol:
def test_delete_col(self, db):
db.delete_col()
db.namespace.delete_all.assert_called_once()
def test_delete_col_handles_error(self, db):
db.namespace.delete_all.side_effect = Exception("API error")
# Should not raise
db.delete_col()
# ── col_info ─────────────────────────────────────────────────────────
class TestColInfo:
def test_col_info(self, db):
mock_metadata = MagicMock()
mock_metadata.approx_row_count = 42
mock_metadata.approx_logical_bytes = 1024
mock_metadata.created_at = datetime(2025, 1, 1)
mock_metadata.updated_at = datetime(2025, 6, 1)
db.namespace.metadata.return_value = mock_metadata
info = db.col_info()
assert info["name"] == "test_ns"
assert info["approx_row_count"] == 42
assert info["approx_logical_bytes"] == 1024
def test_col_info_handles_error(self, db):
db.namespace.metadata.side_effect = Exception("not found")
info = db.col_info()
assert info == {"name": "test_ns"}
# ── list ─────────────────────────────────────────────────────────────
class TestList:
def test_list_returns_wrapped_format(self, db):
"""list() must return [[results]] for compatibility with main.py."""
mock_response = MagicMock()
mock_response.rows = [
_make_row("id1", dist=0.1, data="hello"),
_make_row("id2", dist=0.2, data="world"),
]
db.namespace.query.return_value = mock_response
result = db.list()
# Must be wrapped: result[0] is the actual list
assert isinstance(result, list)
assert isinstance(result[0], list)
assert len(result[0]) == 2
assert result[0][0].id == "id1"
assert result[0][1].id == "id2"
def test_list_with_filters(self, db):
mock_response = MagicMock()
mock_response.rows = [_make_row("id1", dist=0.1, data="hello")]
db.namespace.query.return_value = mock_response
db.list(filters={"user_id": "u1"}, limit=50)
call_kwargs = db.namespace.query.call_args[1]
assert call_kwargs["filters"] == ("user_id", "Eq", "u1")
assert call_kwargs["top_k"] == 50
def test_list_empty(self, db):
mock_response = MagicMock()
mock_response.rows = None
db.namespace.query.return_value = mock_response
result = db.list()
assert result == [[]]
def test_list_uses_zero_vector(self, db):
mock_response = MagicMock()
mock_response.rows = []
db.namespace.query.return_value = mock_response
db.list()
call_kwargs = db.namespace.query.call_args[1]
assert call_kwargs["rank_by"] == ("vector", "ANN", [0.0, 0.0, 0.0, 0.0])
def test_list_compatible_with_main_py_get_all(self, db):
"""Simulate how main.py _get_all_from_vector_store unwraps list()."""
mock_response = MagicMock()
mock_response.rows = [_make_row("id1", dist=0.1, data="hello")]
db.namespace.query.return_value = mock_response
memories_result = db.list(filters={"user_id": "u1"})
# Reproduce main.py unwrapping logic
first_element = memories_result[0]
if isinstance(first_element, (list, tuple)):
actual_memories = first_element
else:
actual_memories = memories_result
assert isinstance(actual_memories, list)
assert len(actual_memories) == 1
assert actual_memories[0].id == "id1"
assert actual_memories[0].payload["data"] == "hello"
def test_list_compatible_with_main_py_delete_all(self, db):
"""Simulate how main.py delete_all uses list()[0]."""
mock_response = MagicMock()
mock_response.rows = [
_make_row("id1", dist=0.1, data="hello"),
_make_row("id2", dist=0.2, data="world"),
]
db.namespace.query.return_value = mock_response
result = db.list(filters={"user_id": "u1"})
memories = result[0]
assert isinstance(memories, list)
for mem in memories:
assert hasattr(mem, "id")
assert hasattr(mem, "payload")
# ── count ────────────────────────────────────────────────────────────
class TestCount:
def test_count(self, db):
mock_metadata = MagicMock()
mock_metadata.approx_row_count = 100
db.namespace.metadata.return_value = mock_metadata
assert db.count() == 100
def test_count_handles_error(self, db):
db.namespace.metadata.side_effect = Exception("error")
assert db.count() == 0
# ── reset ────────────────────────────────────────────────────────────
class TestReset:
def test_reset_calls_delete_all(self, db):
db.reset()
db.namespace.delete_all.assert_called_once()
# ── Config ───────────────────────────────────────────────────────────
class TestConfig:
def test_config_valid(self):
with patch.dict("os.environ", {"TURBOPUFFER_API_KEY": "key"}):
from mem0.configs.vector_stores.turbopuffer import TurbopufferConfig
config = TurbopufferConfig()
assert config.collection_name == "mem0"
assert config.embedding_model_dims == 1536
assert config.distance_metric == "cosine_distance"
assert config.batch_size == 100
assert config.region == "gcp-us-central1"
def test_config_custom_values(self):
from mem0.configs.vector_stores.turbopuffer import TurbopufferConfig
config = TurbopufferConfig(
collection_name="custom",
embedding_model_dims=768,
api_key="tpuf_key",
region="aws-us-west-2",
distance_metric="euclidean_squared",
batch_size=50,
)
assert config.collection_name == "custom"
assert config.embedding_model_dims == 768
assert config.api_key == "tpuf_key"
assert config.region == "aws-us-west-2"
assert config.distance_metric == "euclidean_squared"
assert config.batch_size == 50
def test_config_rejects_extra_fields(self):
from mem0.configs.vector_stores.turbopuffer import TurbopufferConfig
with pytest.raises(ValueError, match="Extra fields not allowed"):
TurbopufferConfig(api_key="key", unknown_field="value")
def test_config_requires_api_key(self):
from mem0.configs.vector_stores.turbopuffer import TurbopufferConfig
with patch.dict("os.environ", {}, clear=True):
import os
os.environ.pop("TURBOPUFFER_API_KEY", None)
with pytest.raises(ValueError, match="api_key"):
TurbopufferConfig()
def test_config_accepts_env_var(self):
from mem0.configs.vector_stores.turbopuffer import TurbopufferConfig
with patch.dict("os.environ", {"TURBOPUFFER_API_KEY": "env_key"}):
config = TurbopufferConfig()
assert config.api_key is None # not set explicitly, but env var is present
# ── Factory Registration ─────────────────────────────────────────────
class TestFactoryRegistration:
def test_vector_store_factory_has_turbopuffer(self):
from mem0.utils.factory import VectorStoreFactory
assert "turbopuffer" in VectorStoreFactory.provider_to_class
assert VectorStoreFactory.provider_to_class["turbopuffer"] == "mem0.vector_stores.turbopuffer.TurbopufferDB"
def test_vector_store_config_has_turbopuffer(self):
from mem0.vector_stores.configs import VectorStoreConfig
config = VectorStoreConfig.__private_attributes__["_provider_configs"].default
assert "turbopuffer" in config
assert config["turbopuffer"] == "TurbopufferConfig"
def test_config_validation_pipeline(self):
"""Test that VectorStoreConfig correctly resolves turbopuffer config."""
from mem0.vector_stores.configs import VectorStoreConfig
with patch.dict("os.environ", {"TURBOPUFFER_API_KEY": "key"}):
config = VectorStoreConfig(
provider="turbopuffer",
config={"collection_name": "test", "region": "gcp-us-central1"},
)
assert config.config.collection_name == "test"
assert config.config.region == "gcp-us-central1"
assert config.config.embedding_model_dims == 1536
# ── OutputData ───────────────────────────────────────────────────────
class TestOutputData:
def test_output_data_has_required_fields(self):
od = OutputData(id="test", score=0.9, payload={"data": "hello"})
assert od.id == "test"
assert od.score == 0.9
assert od.payload == {"data": "hello"}
def test_output_data_nullable_fields(self):
od = OutputData(id=None, score=None, payload=None)
assert od.id is None
assert od.score is None
assert od.payload is None
def test_output_data_payload_access_pattern(self):
"""Test the exact access pattern used by main.py."""
od = OutputData(
id="mem1",
score=0.95,
payload={
"data": "User likes sci-fi movies",
"hash": "abc123",
"created_at": "2025-01-01",
"updated_at": "2025-06-01",
"user_id": "alice",
"agent_id": "agent1",
},
)
assert od.payload.get("data", "") == "User likes sci-fi movies"
assert od.payload.get("hash") == "abc123"
assert od.payload.get("created_at") == "2025-01-01"
assert od.payload.get("user_id") == "alice"
assert od.payload.get("run_id") is None # not present, should return None