Compare commits
3
Commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
f085460201 | ||
|
|
839352b132 | ||
|
|
68ee4389cd |
@@ -1,79 +0,0 @@
|
|||||||
---
|
|
||||||
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,
|
|
||||||
}
|
|
||||||
}
|
|
||||||
}
|
|
||||||
```
|
|
||||||
@@ -33,7 +33,6 @@ See the list of supported vector databases below.
|
|||||||
<Card title="LangChain" href="/components/vectordbs/dbs/langchain"></Card>
|
<Card title="LangChain" href="/components/vectordbs/dbs/langchain"></Card>
|
||||||
<Card title="Amazon S3 Vectors" href="/components/vectordbs/dbs/s3_vectors"></Card>
|
<Card title="Amazon S3 Vectors" href="/components/vectordbs/dbs/s3_vectors"></Card>
|
||||||
<Card title="Databricks" href="/components/vectordbs/dbs/databricks"></Card>
|
<Card title="Databricks" href="/components/vectordbs/dbs/databricks"></Card>
|
||||||
<Card title="Turbopuffer" href="/components/vectordbs/dbs/turbopuffer"></Card>
|
|
||||||
</CardGroup>
|
</CardGroup>
|
||||||
|
|
||||||
## Usage
|
## Usage
|
||||||
|
|||||||
+1
-2
@@ -231,8 +231,7 @@
|
|||||||
"components/vectordbs/dbs/cassandra",
|
"components/vectordbs/dbs/cassandra",
|
||||||
"components/vectordbs/dbs/s3_vectors",
|
"components/vectordbs/dbs/s3_vectors",
|
||||||
"components/vectordbs/dbs/databricks",
|
"components/vectordbs/dbs/databricks",
|
||||||
"components/vectordbs/dbs/neptune_analytics",
|
"components/vectordbs/dbs/neptune_analytics"
|
||||||
"components/vectordbs/dbs/turbopuffer"
|
|
||||||
]
|
]
|
||||||
}
|
}
|
||||||
]
|
]
|
||||||
|
|||||||
@@ -35,7 +35,7 @@ graph LR
|
|||||||
Mem0’s extraction LLM identifies entities, relationships, and timestamps from the conversation payload you send to `memory.add`.
|
Mem0’s extraction LLM identifies entities, relationships, and timestamps from the conversation payload you send to `memory.add`.
|
||||||
</Step>
|
</Step>
|
||||||
<Step title="Store vectors and edges together">
|
<Step title="Store vectors and edges together">
|
||||||
Embeddings land in your configured vector database while nodes and edges flow into a graph backend (Neo4j, Memgraph, Neptune, Kuzu, or Apache AGE).
|
Embeddings land in your configured vector database while nodes and edges flow into a Bolt-compatible graph backend (Neo4j, Memgraph, Neptune, or Kuzu).
|
||||||
</Step>
|
</Step>
|
||||||
<Step title="Expose graph context at search time">
|
<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.
|
`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
|
## Decision Points
|
||||||
|
|
||||||
- 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).
|
- Select the graph store that fits your deployment (managed Aura vs. self-hosted Neo4j vs. AWS Neptune vs. local Kuzu).
|
||||||
- Decide when to enable graph writes per request; routine conversations may stay vector-only to save latency.
|
- 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.
|
- Set a policy for pruning stale relationships so your graph stays fast and affordable.
|
||||||
|
|
||||||
@@ -381,42 +381,6 @@ config = {
|
|||||||
|
|
||||||
Kuzu will clear its state when using `:memory:` once the process exits. See the [Kuzu documentation](https://kuzudb.com/docs/) for advanced settings.
|
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>
|
||||||
<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>
|
</AccordionGroup>
|
||||||
|
|
||||||
<CardGroup cols={2}>
|
<CardGroup cols={2}>
|
||||||
|
|||||||
@@ -14,7 +14,7 @@ The Mem0 REST API server exposes every OSS memory operation over HTTP. Run it al
|
|||||||
</Info>
|
</Info>
|
||||||
|
|
||||||
<Warning>
|
<Warning>
|
||||||
Enable API key authentication (see below) and HTTPS before exposing the server to anything beyond your internal network.
|
Add your own authentication and HTTPS before exposing the server to anything beyond your internal network. The default image does not include auth.
|
||||||
</Warning>
|
</Warning>
|
||||||
|
|
||||||
---
|
---
|
||||||
@@ -22,7 +22,6 @@ The Mem0 REST API server exposes every OSS memory operation over HTTP. Run it al
|
|||||||
## Feature
|
## Feature
|
||||||
|
|
||||||
- **CRUD endpoints:** Create, retrieve, search, update, delete, and reset memories by `user_id`, `agent_id`, or `run_id`.
|
- **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.
|
- **Status health check:** Access base routes to confirm the server is online.
|
||||||
- **OpenAPI explorer:** Visit `/docs` for interactive testing and schema reference.
|
- **OpenAPI explorer:** Visit `/docs` for interactive testing and schema reference.
|
||||||
|
|
||||||
@@ -92,41 +91,6 @@ 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
|
## See it in action
|
||||||
|
|
||||||
### Create and search memories via HTTP
|
### Create and search memories via HTTP
|
||||||
@@ -173,7 +137,7 @@ curl "http://localhost:8000/memories/search?user_id=alice&query=vegetable"
|
|||||||
|
|
||||||
## Best practices
|
## Best practices
|
||||||
|
|
||||||
1. **Enable authentication:** Set `ADMIN_API_KEY` to secure all endpoints, or use an API gateway for more advanced schemes.
|
1. **Add authentication:** Protect endpoints with API gateways, proxies, or custom FastAPI middleware.
|
||||||
2. **Use HTTPS:** Terminate TLS at your load balancer or reverse proxy.
|
2. **Use HTTPS:** Terminate TLS at your load balancer or reverse proxy.
|
||||||
3. **Monitor uptime:** Track request rates, latency, and error codes per endpoint.
|
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.
|
4. **Version configs:** Keep environment files and Docker Compose definitions in source control.
|
||||||
|
|||||||
@@ -47,7 +47,6 @@ 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` |
|
| 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` |
|
| 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` |
|
| 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
|
### What This Does
|
||||||
|
|
||||||
|
|||||||
@@ -183,11 +183,6 @@ 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
|
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.
|
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.
|
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.
|
||||||
|
|||||||
Vendored
+2
-2
@@ -166,13 +166,13 @@ Call out the most common mistake or edge case for this layer.
|
|||||||
title="[Related cookbook / deep dive]"
|
title="[Related cookbook / deep dive]"
|
||||||
description="[Why this pairs well with the current guide]"
|
description="[Why this pairs well with the current guide]"
|
||||||
icon="arrow-right"
|
icon="arrow-right"
|
||||||
href="#related-link"
|
href="/[related-link]"
|
||||||
/>
|
/>
|
||||||
<Card
|
<Card
|
||||||
title="[Next cookbook in journey]"
|
title="[Next cookbook in journey]"
|
||||||
description="[Set expectation for the next step]"
|
description="[Set expectation for the next step]"
|
||||||
icon="rocket"
|
icon="rocket"
|
||||||
href="#next-link"
|
href="/[next-link]"
|
||||||
/>
|
/>
|
||||||
</CardGroup>
|
</CardGroup>
|
||||||
```
|
```
|
||||||
|
|||||||
+2
-2
@@ -145,13 +145,13 @@ npm install mem0ai@[version]
|
|||||||
title="[Deep dive reference]"
|
title="[Deep dive reference]"
|
||||||
description="[Why this reference matters post-migration]"
|
description="[Why this reference matters post-migration]"
|
||||||
icon="book"
|
icon="book"
|
||||||
href="#reference-link"
|
href="/[reference-link]"
|
||||||
/>
|
/>
|
||||||
<Card
|
<Card
|
||||||
title="[Applied example or next step]"
|
title="[Applied example or next step]"
|
||||||
description="[What readers can build now]"
|
description="[What readers can build now]"
|
||||||
icon="rocket"
|
icon="rocket"
|
||||||
href="#example-link"
|
href="/[example-link]"
|
||||||
/>
|
/>
|
||||||
</CardGroup>
|
</CardGroup>
|
||||||
```
|
```
|
||||||
|
|||||||
@@ -1,419 +0,0 @@
|
|||||||
#!/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 ""
|
|
||||||
@@ -1,272 +0,0 @@
|
|||||||
# 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)
|
|
||||||
@@ -1,864 +0,0 @@
|
|||||||
#!/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.
@@ -34,6 +34,12 @@ export const DEFAULT_MEMORY_CONFIG: MemoryConfig = {
|
|||||||
username: process.env.NEO4J_USERNAME || "neo4j",
|
username: process.env.NEO4J_USERNAME || "neo4j",
|
||||||
password: process.env.NEO4J_PASSWORD || "password",
|
password: process.env.NEO4J_PASSWORD || "password",
|
||||||
},
|
},
|
||||||
|
llm: {
|
||||||
|
provider: "openai",
|
||||||
|
config: {
|
||||||
|
model: "gpt-4-turbo-preview",
|
||||||
|
},
|
||||||
|
},
|
||||||
},
|
},
|
||||||
historyStore: {
|
historyStore: {
|
||||||
provider: "sqlite",
|
provider: "sqlite",
|
||||||
|
|||||||
@@ -26,5 +26,4 @@ export * from "./vector_stores/supabase";
|
|||||||
export * from "./vector_stores/langchain";
|
export * from "./vector_stores/langchain";
|
||||||
export * from "./vector_stores/vectorize";
|
export * from "./vector_stores/vectorize";
|
||||||
export * from "./vector_stores/azure_ai_search";
|
export * from "./vector_stores/azure_ai_search";
|
||||||
export * from "./vector_stores/pgvector";
|
|
||||||
export * from "./utils/factory";
|
export * from "./utils/factory";
|
||||||
|
|||||||
@@ -80,18 +80,18 @@ export class MemoryGraph {
|
|||||||
);
|
);
|
||||||
|
|
||||||
this.llmProvider = "openai";
|
this.llmProvider = "openai";
|
||||||
let llmConfig = this.config.llm.config;
|
|
||||||
|
|
||||||
if (this.config.llm?.provider) {
|
if (this.config.llm?.provider) {
|
||||||
this.llmProvider = this.config.llm.provider;
|
this.llmProvider = this.config.llm.provider;
|
||||||
}
|
}
|
||||||
if (this.config.graphStore?.llm?.provider) {
|
if (this.config.graphStore?.llm?.provider) {
|
||||||
this.llmProvider = 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, llmConfig);
|
this.llm = LLMFactory.create(this.llmProvider, this.config.llm.config);
|
||||||
this.structuredLlm = LLMFactory.create(this.llmProvider, llmConfig);
|
this.structuredLlm = LLMFactory.create(
|
||||||
|
this.llmProvider,
|
||||||
|
this.config.llm.config,
|
||||||
|
);
|
||||||
this.threshold = 0.7;
|
this.threshold = 0.7;
|
||||||
}
|
}
|
||||||
|
|
||||||
|
|||||||
@@ -31,7 +31,6 @@ export const MemoryUpdateSchema = z.object({
|
|||||||
old_memory: z
|
old_memory: z
|
||||||
.string()
|
.string()
|
||||||
.optional()
|
.optional()
|
||||||
.nullable()
|
|
||||||
.describe(
|
.describe(
|
||||||
"The previous content of the memory item if the event was UPDATE.",
|
"The previous content of the memory item if the event was UPDATE.",
|
||||||
),
|
),
|
||||||
|
|||||||
@@ -34,7 +34,6 @@ import { LangchainLLM } from "../llms/langchain";
|
|||||||
import { LangchainEmbedder } from "../embeddings/langchain";
|
import { LangchainEmbedder } from "../embeddings/langchain";
|
||||||
import { LangchainVectorStore } from "../vector_stores/langchain";
|
import { LangchainVectorStore } from "../vector_stores/langchain";
|
||||||
import { AzureAISearch } from "../vector_stores/azure_ai_search";
|
import { AzureAISearch } from "../vector_stores/azure_ai_search";
|
||||||
import { PGVector } from "../vector_stores/pgvector";
|
|
||||||
|
|
||||||
export class EmbedderFactory {
|
export class EmbedderFactory {
|
||||||
static create(provider: string, config: EmbeddingConfig): Embedder {
|
static create(provider: string, config: EmbeddingConfig): Embedder {
|
||||||
@@ -105,8 +104,6 @@ export class VectorStoreFactory {
|
|||||||
return new VectorizeDB(config as any);
|
return new VectorizeDB(config as any);
|
||||||
case "azure-ai-search":
|
case "azure-ai-search":
|
||||||
return new AzureAISearch(config as any);
|
return new AzureAISearch(config as any);
|
||||||
case "pgvector":
|
|
||||||
return new PGVector(config as any);
|
|
||||||
default:
|
default:
|
||||||
throw new Error(`Unsupported vector store provider: ${provider}`);
|
throw new Error(`Unsupported vector store provider: ${provider}`);
|
||||||
}
|
}
|
||||||
|
|||||||
@@ -522,7 +522,7 @@ export class RedisDB implements VectorStore {
|
|||||||
|
|
||||||
return {
|
return {
|
||||||
id: vectorId,
|
id: vectorId,
|
||||||
payload: toCamelCase(payload),
|
payload,
|
||||||
};
|
};
|
||||||
} catch (error) {
|
} catch (error) {
|
||||||
console.error("Error getting vector:", error);
|
console.error("Error getting vector:", error);
|
||||||
|
|||||||
@@ -349,78 +349,6 @@ 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)
|
// Memory class – LM Studio end-to-end flow (mocked factories)
|
||||||
// ─────────────────────────────────────────────────────────────────────────
|
// ─────────────────────────────────────────────────────────────────────────
|
||||||
|
|||||||
@@ -511,137 +511,6 @@ describe("Prompt construction — all json_object sites include 'json'", () => {
|
|||||||
// 5. Edge cases – malformed entity fields in _removeSpacesFromEntities
|
// 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)", () => {
|
describe("_removeSpacesFromEntities (via _establishNodesRelationsFromData)", () => {
|
||||||
it("normalises spaces and case in entity source/relationship/destination", async () => {
|
it("normalises spaces and case in entity source/relationship/destination", async () => {
|
||||||
mockGenerateResponse.mockResolvedValueOnce({
|
mockGenerateResponse.mockResolvedValueOnce({
|
||||||
|
|||||||
@@ -1,45 +0,0 @@
|
|||||||
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)
|
|
||||||
+2
-24
@@ -77,32 +77,12 @@ class KuzuConfig(BaseModel):
|
|||||||
db: Optional[str] = Field(":memory:", description="Path to a Kuzu database file")
|
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):
|
class GraphStoreConfig(BaseModel):
|
||||||
provider: str = Field(
|
provider: str = Field(
|
||||||
description="Provider of the data store (e.g., 'neo4j', 'memgraph', 'neptune', 'kuzu', 'apache_age')",
|
description="Provider of the data store (e.g., 'neo4j', 'memgraph', 'neptune', 'kuzu')",
|
||||||
default="neo4j",
|
default="neo4j",
|
||||||
)
|
)
|
||||||
config: Union[Neo4jConfig, MemgraphConfig, NeptuneConfig, KuzuConfig, ApacheAgeConfig] = Field(
|
config: Union[Neo4jConfig, MemgraphConfig, NeptuneConfig, KuzuConfig] = Field(
|
||||||
description="Configuration for the specific data store", default=None
|
description="Configuration for the specific data store", default=None
|
||||||
)
|
)
|
||||||
llm: Optional[LlmConfig] = Field(description="LLM configuration for querying the graph store", default=None)
|
llm: Optional[LlmConfig] = Field(description="LLM configuration for querying the graph store", default=None)
|
||||||
@@ -130,7 +110,5 @@ class GraphStoreConfig(BaseModel):
|
|||||||
return NeptuneConfig(**v.model_dump())
|
return NeptuneConfig(**v.model_dump())
|
||||||
elif provider == "kuzu":
|
elif provider == "kuzu":
|
||||||
return KuzuConfig(**v.model_dump())
|
return KuzuConfig(**v.model_dump())
|
||||||
elif provider == "apache_age":
|
|
||||||
return ApacheAgeConfig(**v.model_dump())
|
|
||||||
else:
|
else:
|
||||||
raise ValueError(f"Unsupported graph store provider: {provider}")
|
raise ValueError(f"Unsupported graph store provider: {provider}")
|
||||||
|
|||||||
@@ -1,4 +1,3 @@
|
|||||||
import json
|
|
||||||
import os
|
import os
|
||||||
from typing import Dict, List, Optional
|
from typing import Dict, List, Optional
|
||||||
|
|
||||||
@@ -7,7 +6,6 @@ from openai import AzureOpenAI
|
|||||||
|
|
||||||
from mem0.configs.llms.base import BaseLlmConfig
|
from mem0.configs.llms.base import BaseLlmConfig
|
||||||
from mem0.llms.base import LLMBase
|
from mem0.llms.base import LLMBase
|
||||||
from mem0.memory.utils import extract_json
|
|
||||||
|
|
||||||
SCOPE = "https://cognitiveservices.azure.com/.default"
|
SCOPE = "https://cognitiveservices.azure.com/.default"
|
||||||
|
|
||||||
@@ -85,35 +83,9 @@ class AzureOpenAIStructuredLLM(LLMBase):
|
|||||||
params["tools"] = tools
|
params["tools"] = tools
|
||||||
params["tool_choice"] = tool_choice
|
params["tool_choice"] = tool_choice
|
||||||
|
|
||||||
|
if tools:
|
||||||
|
params["tools"] = tools
|
||||||
|
params["tool_choice"] = tool_choice
|
||||||
|
|
||||||
response = self.client.chat.completions.create(**params)
|
response = self.client.chat.completions.create(**params)
|
||||||
return self._parse_response(response, tools)
|
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
|
|
||||||
|
|||||||
+8
-11
@@ -32,25 +32,22 @@ class GeminiLLM(LLMBase):
|
|||||||
Returns:
|
Returns:
|
||||||
str or dict: The processed response.
|
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:
|
if tools:
|
||||||
processed_response = {
|
processed_response = {
|
||||||
"content": None,
|
"content": None,
|
||||||
"tool_calls": [],
|
"tool_calls": [],
|
||||||
}
|
}
|
||||||
|
|
||||||
if parts:
|
# Extract content from the first candidate
|
||||||
# Extract content from the first candidate
|
if response.candidates and response.candidates[0].content.parts:
|
||||||
for part in parts:
|
for part in response.candidates[0].content.parts:
|
||||||
if hasattr(part, "text") and part.text:
|
if hasattr(part, "text") and part.text:
|
||||||
processed_response["content"] = part.text
|
processed_response["content"] = part.text
|
||||||
break
|
break
|
||||||
|
|
||||||
# Extract function calls
|
# Extract function calls
|
||||||
for part in parts:
|
if response.candidates and response.candidates[0].content.parts:
|
||||||
|
for part in response.candidates[0].content.parts:
|
||||||
if hasattr(part, "function_call") and part.function_call:
|
if hasattr(part, "function_call") and part.function_call:
|
||||||
fn = part.function_call
|
fn = part.function_call
|
||||||
processed_response["tool_calls"].append(
|
processed_response["tool_calls"].append(
|
||||||
@@ -62,8 +59,8 @@ class GeminiLLM(LLMBase):
|
|||||||
|
|
||||||
return processed_response
|
return processed_response
|
||||||
else:
|
else:
|
||||||
if parts:
|
if response.candidates and response.candidates[0].content.parts:
|
||||||
for part in parts:
|
for part in response.candidates[0].content.parts:
|
||||||
if hasattr(part, "text") and part.text:
|
if hasattr(part, "text") and part.text:
|
||||||
return part.text
|
return part.text
|
||||||
return ""
|
return ""
|
||||||
|
|||||||
@@ -1,573 +0,0 @@
|
|||||||
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_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()
|
|
||||||
+18
-46
@@ -174,7 +174,6 @@ class MemoryGraph:
|
|||||||
|
|
||||||
query = f"""
|
query = f"""
|
||||||
MATCH (n {self.node_label} {{{node_props_str}}})-[r]->(m {self.node_label} {{{node_props_str}}})
|
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
|
RETURN n.name AS source, type(r) AS relationship, m.name AS target
|
||||||
LIMIT $limit
|
LIMIT $limit
|
||||||
"""
|
"""
|
||||||
@@ -292,12 +291,10 @@ class MemoryGraph:
|
|||||||
CALL {{
|
CALL {{
|
||||||
WITH n
|
WITH n
|
||||||
MATCH (n)-[r]->(m {self.node_label} {{{node_props_str}}})
|
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
|
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
|
UNION
|
||||||
WITH n
|
WITH n
|
||||||
MATCH (n)<-[r]-(m {self.node_label} {{{node_props_str}}})
|
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
|
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
|
WITH distinct source, source_id, relationship, relation_id, destination, destination_id, similarity
|
||||||
@@ -395,15 +392,13 @@ class MemoryGraph:
|
|||||||
source_props_str = ", ".join(source_props)
|
source_props_str = ", ".join(source_props)
|
||||||
dest_props_str = ", ".join(dest_props)
|
dest_props_str = ", ".join(dest_props)
|
||||||
|
|
||||||
# Soft-delete: mark relationship as invalid instead of removing it,
|
# Delete the specific relationship between nodes
|
||||||
# enabling temporal reasoning over historical graph state.
|
|
||||||
# See: https://github.com/mem0ai/mem0/issues/4187
|
|
||||||
cypher = f"""
|
cypher = f"""
|
||||||
MATCH (n {self.node_label} {{{source_props_str}}})
|
MATCH (n {self.node_label} {{{source_props_str}}})
|
||||||
-[r:{relationship}]->
|
-[r:{relationship}]->
|
||||||
(m {self.node_label} {{{dest_props_str}}})
|
(m {self.node_label} {{{dest_props_str}}})
|
||||||
WHERE r.valid IS NULL OR r.valid = true
|
|
||||||
SET r.valid = false, r.invalidated_at = datetime()
|
DELETE r
|
||||||
RETURN
|
RETURN
|
||||||
n.name AS source,
|
n.name AS source,
|
||||||
m.name AS target,
|
m.name AS target,
|
||||||
@@ -469,16 +464,11 @@ class MemoryGraph:
|
|||||||
CALL db.create.setNodeVectorProperty(destination, 'embedding', $destination_embedding)
|
CALL db.create.setNodeVectorProperty(destination, 'embedding', $destination_embedding)
|
||||||
WITH source, destination
|
WITH source, destination
|
||||||
MERGE (source)-[r:{relationship}]->(destination)
|
MERGE (source)-[r:{relationship}]->(destination)
|
||||||
ON CREATE SET
|
ON CREATE SET
|
||||||
r.created_at = timestamp(),
|
r.created = timestamp(),
|
||||||
r.updated_at = timestamp(),
|
r.mentions = 1
|
||||||
r.mentions = 1,
|
|
||||||
r.valid = true
|
|
||||||
ON MATCH SET
|
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
|
RETURN source.name AS source, type(r) AS relationship, destination.name AS target
|
||||||
"""
|
"""
|
||||||
|
|
||||||
@@ -518,16 +508,11 @@ class MemoryGraph:
|
|||||||
CALL db.create.setNodeVectorProperty(source, 'embedding', $source_embedding)
|
CALL db.create.setNodeVectorProperty(source, 'embedding', $source_embedding)
|
||||||
WITH source, destination
|
WITH source, destination
|
||||||
MERGE (source)-[r:{relationship}]->(destination)
|
MERGE (source)-[r:{relationship}]->(destination)
|
||||||
ON CREATE SET
|
ON CREATE SET
|
||||||
r.created_at = timestamp(),
|
r.created = timestamp(),
|
||||||
r.updated_at = timestamp(),
|
r.mentions = 1
|
||||||
r.mentions = 1,
|
|
||||||
r.valid = true
|
|
||||||
ON MATCH SET
|
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
|
RETURN source.name AS source, type(r) AS relationship, destination.name AS target
|
||||||
"""
|
"""
|
||||||
|
|
||||||
@@ -552,16 +537,11 @@ class MemoryGraph:
|
|||||||
WHERE elementId(destination) = $destination_id
|
WHERE elementId(destination) = $destination_id
|
||||||
SET destination.mentions = coalesce(destination.mentions, 0) + 1
|
SET destination.mentions = coalesce(destination.mentions, 0) + 1
|
||||||
MERGE (source)-[r:{relationship}]->(destination)
|
MERGE (source)-[r:{relationship}]->(destination)
|
||||||
ON CREATE SET
|
ON CREATE SET
|
||||||
r.created_at = timestamp(),
|
r.created_at = timestamp(),
|
||||||
r.updated_at = timestamp(),
|
r.updated_at = timestamp(),
|
||||||
r.mentions = 1,
|
r.mentions = 1
|
||||||
r.valid = true
|
ON MATCH SET r.mentions = coalesce(r.mentions, 0) + 1
|
||||||
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
|
RETURN source.name AS source, type(r) AS relationship, destination.name AS target
|
||||||
"""
|
"""
|
||||||
|
|
||||||
@@ -605,18 +585,10 @@ class MemoryGraph:
|
|||||||
WITH source, destination
|
WITH source, destination
|
||||||
CALL db.create.setNodeVectorProperty(destination, 'embedding', $dest_embedding)
|
CALL db.create.setNodeVectorProperty(destination, 'embedding', $dest_embedding)
|
||||||
WITH source, destination
|
WITH source, destination
|
||||||
MERGE (source)-[r:{relationship}]->(destination)
|
MERGE (source)-[rel:{relationship}]->(destination)
|
||||||
ON CREATE SET
|
ON CREATE SET rel.created = timestamp(), rel.mentions = 1
|
||||||
r.created_at = timestamp(),
|
ON MATCH SET rel.mentions = coalesce(rel.mentions, 0) + 1
|
||||||
r.updated_at = timestamp(),
|
RETURN source.name AS source, type(rel) AS relationship, destination.name AS target
|
||||||
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 = {
|
params = {
|
||||||
|
|||||||
+11
-55
@@ -9,7 +9,7 @@ import uuid
|
|||||||
import warnings
|
import warnings
|
||||||
from copy import deepcopy
|
from copy import deepcopy
|
||||||
from datetime import datetime, timezone
|
from datetime import datetime, timezone
|
||||||
from typing import Any, Dict, List, Optional, Union
|
from typing import Any, Dict, Optional
|
||||||
|
|
||||||
from pydantic import ValidationError
|
from pydantic import ValidationError
|
||||||
|
|
||||||
@@ -479,8 +479,7 @@ class Memory(MemoryBase):
|
|||||||
|
|
||||||
msg_content = message_dict["content"]
|
msg_content = message_dict["content"]
|
||||||
msg_embeddings = self.embedding_model.embed(msg_content, "add")
|
msg_embeddings = self.embedding_model.embed(msg_content, "add")
|
||||||
# Pass embeddings as a dict so _create_memory can reuse the cached embedding
|
mem_id = self._create_memory(msg_content, msg_embeddings, per_msg_meta)
|
||||||
mem_id = self._create_memory(msg_content, {msg_content: msg_embeddings}, per_msg_meta)
|
|
||||||
|
|
||||||
returned_memories.append(
|
returned_memories.append(
|
||||||
{
|
{
|
||||||
@@ -609,9 +608,6 @@ class Memory(MemoryBase):
|
|||||||
|
|
||||||
event_type = resp.get("event")
|
event_type = resp.get("event")
|
||||||
if event_type == "ADD":
|
if event_type == "ADD":
|
||||||
# Ensure action_text has an embedding cached to avoid redundant API calls
|
|
||||||
if action_text not in new_message_embeddings:
|
|
||||||
new_message_embeddings[action_text] = self.embedding_model.embed(action_text, "add")
|
|
||||||
memory_id = self._create_memory(
|
memory_id = self._create_memory(
|
||||||
data=action_text,
|
data=action_text,
|
||||||
existing_embeddings=new_message_embeddings,
|
existing_embeddings=new_message_embeddings,
|
||||||
@@ -619,9 +615,6 @@ class Memory(MemoryBase):
|
|||||||
)
|
)
|
||||||
returned_memories.append({"id": memory_id, "memory": action_text, "event": event_type})
|
returned_memories.append({"id": memory_id, "memory": action_text, "event": event_type})
|
||||||
elif event_type == "UPDATE":
|
elif event_type == "UPDATE":
|
||||||
# Ensure action_text has an embedding cached to avoid redundant API calls
|
|
||||||
if action_text not in new_message_embeddings:
|
|
||||||
new_message_embeddings[action_text] = self.embedding_model.embed(action_text, "update")
|
|
||||||
self._update_memory(
|
self._update_memory(
|
||||||
memory_id=temp_uuid_mapping[resp.get("id")],
|
memory_id=temp_uuid_mapping[resp.get("id")],
|
||||||
data=action_text,
|
data=action_text,
|
||||||
@@ -651,9 +644,6 @@ class Memory(MemoryBase):
|
|||||||
if memory_id and (metadata.get("agent_id") or metadata.get("run_id")):
|
if memory_id and (metadata.get("agent_id") or metadata.get("run_id")):
|
||||||
# Update only the session identifiers, keep content the same
|
# Update only the session identifiers, keep content the same
|
||||||
existing_memory = self.vector_store.get(vector_id=memory_id)
|
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)
|
updated_metadata = deepcopy(existing_memory.payload)
|
||||||
if metadata.get("agent_id"):
|
if metadata.get("agent_id"):
|
||||||
updated_metadata["agent_id"] = metadata["agent_id"]
|
updated_metadata["agent_id"] = metadata["agent_id"]
|
||||||
@@ -1160,13 +1150,10 @@ class Memory(MemoryBase):
|
|||||||
capture_event("mem0.history", self, {"memory_id": memory_id, "sync_type": "sync"})
|
capture_event("mem0.history", self, {"memory_id": memory_id, "sync_type": "sync"})
|
||||||
return self.db.get_history(memory_id)
|
return self.db.get_history(memory_id)
|
||||||
|
|
||||||
def _create_memory(self, data: str, existing_embeddings: Union[Dict[str, List[float]], List[float]], metadata=None):
|
def _create_memory(self, data, existing_embeddings, metadata=None):
|
||||||
logger.debug(f"Creating memory with {data=}")
|
logger.debug(f"Creating memory with {data=}")
|
||||||
# existing_embeddings may be a dict (preferred) or a precomputed vector
|
if data in existing_embeddings:
|
||||||
if isinstance(existing_embeddings, dict) and data in existing_embeddings:
|
|
||||||
embeddings = existing_embeddings[data]
|
embeddings = existing_embeddings[data]
|
||||||
elif not isinstance(existing_embeddings, dict):
|
|
||||||
embeddings = existing_embeddings
|
|
||||||
else:
|
else:
|
||||||
embeddings = self.embedding_model.embed(data, memory_action="add")
|
embeddings = self.embedding_model.embed(data, memory_action="add")
|
||||||
memory_id = str(uuid.uuid4())
|
memory_id = str(uuid.uuid4())
|
||||||
@@ -1230,7 +1217,7 @@ class Memory(MemoryBase):
|
|||||||
|
|
||||||
return result
|
return result
|
||||||
|
|
||||||
def _update_memory(self, memory_id, data: str, existing_embeddings: Union[Dict[str, List[float]], List[float]], metadata=None):
|
def _update_memory(self, memory_id, data, existing_embeddings, metadata=None):
|
||||||
logger.info(f"Updating memory with {data=}")
|
logger.info(f"Updating memory with {data=}")
|
||||||
|
|
||||||
try:
|
try:
|
||||||
@@ -1239,9 +1226,6 @@ class Memory(MemoryBase):
|
|||||||
logger.error(f"Error getting memory with ID {memory_id} during update.")
|
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'")
|
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")
|
prev_value = existing_memory.payload.get("data")
|
||||||
|
|
||||||
new_metadata = deepcopy(metadata) if metadata is not None else {}
|
new_metadata = deepcopy(metadata) if metadata is not None else {}
|
||||||
@@ -1263,10 +1247,8 @@ class Memory(MemoryBase):
|
|||||||
if "role" not in new_metadata and "role" in existing_memory.payload:
|
if "role" not in new_metadata and "role" in existing_memory.payload:
|
||||||
new_metadata["role"] = existing_memory.payload["role"]
|
new_metadata["role"] = existing_memory.payload["role"]
|
||||||
|
|
||||||
if isinstance(existing_embeddings, dict) and data in existing_embeddings:
|
if data in existing_embeddings:
|
||||||
embeddings = existing_embeddings[data]
|
embeddings = existing_embeddings[data]
|
||||||
elif not isinstance(existing_embeddings, dict):
|
|
||||||
embeddings = existing_embeddings
|
|
||||||
else:
|
else:
|
||||||
embeddings = self.embedding_model.embed(data, "update")
|
embeddings = self.embedding_model.embed(data, "update")
|
||||||
|
|
||||||
@@ -1292,8 +1274,6 @@ class Memory(MemoryBase):
|
|||||||
def _delete_memory(self, memory_id):
|
def _delete_memory(self, memory_id):
|
||||||
logger.info(f"Deleting memory with {memory_id=}")
|
logger.info(f"Deleting memory with {memory_id=}")
|
||||||
existing_memory = self.vector_store.get(vector_id=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")
|
|
||||||
prev_value = existing_memory.payload.get("data", "")
|
prev_value = existing_memory.payload.get("data", "")
|
||||||
self.vector_store.delete(vector_id=memory_id)
|
self.vector_store.delete(vector_id=memory_id)
|
||||||
self.db.add_history(
|
self.db.add_history(
|
||||||
@@ -1535,8 +1515,7 @@ class AsyncMemory(MemoryBase):
|
|||||||
|
|
||||||
msg_content = message_dict["content"]
|
msg_content = message_dict["content"]
|
||||||
msg_embeddings = await asyncio.to_thread(self.embedding_model.embed, msg_content, "add")
|
msg_embeddings = await asyncio.to_thread(self.embedding_model.embed, msg_content, "add")
|
||||||
# Pass embeddings as a dict so _create_memory can reuse the cached embedding
|
mem_id = await self._create_memory(msg_content, msg_embeddings, per_msg_meta)
|
||||||
mem_id = await self._create_memory(msg_content, {msg_content: msg_embeddings}, per_msg_meta)
|
|
||||||
|
|
||||||
returned_memories.append(
|
returned_memories.append(
|
||||||
{
|
{
|
||||||
@@ -1664,11 +1643,6 @@ class AsyncMemory(MemoryBase):
|
|||||||
event_type = resp.get("event")
|
event_type = resp.get("event")
|
||||||
|
|
||||||
if event_type == "ADD":
|
if event_type == "ADD":
|
||||||
# Ensure action_text has an embedding cached to avoid redundant API calls
|
|
||||||
if action_text not in new_message_embeddings:
|
|
||||||
new_message_embeddings[action_text] = await asyncio.to_thread(
|
|
||||||
self.embedding_model.embed, action_text, "add"
|
|
||||||
)
|
|
||||||
task = asyncio.create_task(
|
task = asyncio.create_task(
|
||||||
self._create_memory(
|
self._create_memory(
|
||||||
data=action_text,
|
data=action_text,
|
||||||
@@ -1678,11 +1652,6 @@ class AsyncMemory(MemoryBase):
|
|||||||
)
|
)
|
||||||
memory_tasks.append((task, resp, "ADD", None))
|
memory_tasks.append((task, resp, "ADD", None))
|
||||||
elif event_type == "UPDATE":
|
elif event_type == "UPDATE":
|
||||||
# Ensure action_text has an embedding cached to avoid redundant API calls
|
|
||||||
if action_text not in new_message_embeddings:
|
|
||||||
new_message_embeddings[action_text] = await asyncio.to_thread(
|
|
||||||
self.embedding_model.embed, action_text, "update"
|
|
||||||
)
|
|
||||||
task = asyncio.create_task(
|
task = asyncio.create_task(
|
||||||
self._update_memory(
|
self._update_memory(
|
||||||
memory_id=temp_uuid_mapping[resp["id"]],
|
memory_id=temp_uuid_mapping[resp["id"]],
|
||||||
@@ -1702,9 +1671,6 @@ class AsyncMemory(MemoryBase):
|
|||||||
# Create async task to update only the session identifiers
|
# Create async task to update only the session identifiers
|
||||||
async def update_session_ids(mem_id, meta):
|
async def update_session_ids(mem_id, meta):
|
||||||
existing_memory = await asyncio.to_thread(self.vector_store.get, vector_id=mem_id)
|
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)
|
updated_metadata = deepcopy(existing_memory.payload)
|
||||||
if meta.get("agent_id"):
|
if meta.get("agent_id"):
|
||||||
updated_metadata["agent_id"] = meta["agent_id"]
|
updated_metadata["agent_id"] = meta["agent_id"]
|
||||||
@@ -2252,13 +2218,10 @@ class AsyncMemory(MemoryBase):
|
|||||||
capture_event("mem0.history", self, {"memory_id": memory_id, "sync_type": "async"})
|
capture_event("mem0.history", self, {"memory_id": memory_id, "sync_type": "async"})
|
||||||
return await asyncio.to_thread(self.db.get_history, memory_id)
|
return await asyncio.to_thread(self.db.get_history, memory_id)
|
||||||
|
|
||||||
async def _create_memory(self, data: str, existing_embeddings: Union[Dict[str, List[float]], List[float]], metadata=None):
|
async def _create_memory(self, data, existing_embeddings, metadata=None):
|
||||||
logger.debug(f"Creating memory with {data=}")
|
logger.debug(f"Creating memory with {data=}")
|
||||||
# existing_embeddings may be a dict (preferred) or a precomputed vector
|
if data in existing_embeddings:
|
||||||
if isinstance(existing_embeddings, dict) and data in existing_embeddings:
|
|
||||||
embeddings = existing_embeddings[data]
|
embeddings = existing_embeddings[data]
|
||||||
elif not isinstance(existing_embeddings, dict):
|
|
||||||
embeddings = existing_embeddings
|
|
||||||
else:
|
else:
|
||||||
embeddings = await asyncio.to_thread(self.embedding_model.embed, data, memory_action="add")
|
embeddings = await asyncio.to_thread(self.embedding_model.embed, data, memory_action="add")
|
||||||
|
|
||||||
@@ -2341,7 +2304,7 @@ class AsyncMemory(MemoryBase):
|
|||||||
|
|
||||||
return result
|
return result
|
||||||
|
|
||||||
async def _update_memory(self, memory_id, data: str, existing_embeddings: Union[Dict[str, List[float]], List[float]], metadata=None):
|
async def _update_memory(self, memory_id, data, existing_embeddings, metadata=None):
|
||||||
logger.info(f"Updating memory with {data=}")
|
logger.info(f"Updating memory with {data=}")
|
||||||
|
|
||||||
try:
|
try:
|
||||||
@@ -2350,9 +2313,6 @@ class AsyncMemory(MemoryBase):
|
|||||||
logger.error(f"Error getting memory with ID {memory_id} during update.")
|
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'")
|
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")
|
prev_value = existing_memory.payload.get("data")
|
||||||
|
|
||||||
new_metadata = deepcopy(metadata) if metadata is not None else {}
|
new_metadata = deepcopy(metadata) if metadata is not None else {}
|
||||||
@@ -2375,10 +2335,8 @@ class AsyncMemory(MemoryBase):
|
|||||||
if "role" not in new_metadata and "role" in existing_memory.payload:
|
if "role" not in new_metadata and "role" in existing_memory.payload:
|
||||||
new_metadata["role"] = existing_memory.payload["role"]
|
new_metadata["role"] = existing_memory.payload["role"]
|
||||||
|
|
||||||
if isinstance(existing_embeddings, dict) and data in existing_embeddings:
|
if data in existing_embeddings:
|
||||||
embeddings = existing_embeddings[data]
|
embeddings = existing_embeddings[data]
|
||||||
elif not isinstance(existing_embeddings, dict):
|
|
||||||
embeddings = existing_embeddings
|
|
||||||
else:
|
else:
|
||||||
embeddings = await asyncio.to_thread(self.embedding_model.embed, data, "update")
|
embeddings = await asyncio.to_thread(self.embedding_model.embed, data, "update")
|
||||||
|
|
||||||
@@ -2406,8 +2364,6 @@ class AsyncMemory(MemoryBase):
|
|||||||
async def _delete_memory(self, memory_id):
|
async def _delete_memory(self, memory_id):
|
||||||
logger.info(f"Deleting memory with {memory_id=}")
|
logger.info(f"Deleting memory with {memory_id=}")
|
||||||
existing_memory = await asyncio.to_thread(self.vector_store.get, vector_id=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")
|
|
||||||
prev_value = existing_memory.payload.get("data", "")
|
prev_value = existing_memory.payload.get("data", "")
|
||||||
|
|
||||||
await asyncio.to_thread(self.vector_store.delete, vector_id=memory_id)
|
await asyncio.to_thread(self.vector_store.delete, vector_id=memory_id)
|
||||||
|
|||||||
@@ -188,7 +188,6 @@ class VectorStoreFactory:
|
|||||||
"baidu": "mem0.vector_stores.baidu.BaiduDB",
|
"baidu": "mem0.vector_stores.baidu.BaiduDB",
|
||||||
"cassandra": "mem0.vector_stores.cassandra.CassandraDB",
|
"cassandra": "mem0.vector_stores.cassandra.CassandraDB",
|
||||||
"neptune": "mem0.vector_stores.neptune_analytics.NeptuneAnalyticsVector",
|
"neptune": "mem0.vector_stores.neptune_analytics.NeptuneAnalyticsVector",
|
||||||
"turbopuffer": "mem0.vector_stores.turbopuffer.TurbopufferDB",
|
|
||||||
}
|
}
|
||||||
|
|
||||||
@classmethod
|
@classmethod
|
||||||
@@ -219,7 +218,6 @@ class GraphStoreFactory:
|
|||||||
"neptune": "mem0.graphs.neptune.neptunegraph.MemoryGraph",
|
"neptune": "mem0.graphs.neptune.neptunegraph.MemoryGraph",
|
||||||
"neptunedb": "mem0.graphs.neptune.neptunedb.MemoryGraph",
|
"neptunedb": "mem0.graphs.neptune.neptunedb.MemoryGraph",
|
||||||
"kuzu": "mem0.memory.kuzu_memory.MemoryGraph",
|
"kuzu": "mem0.memory.kuzu_memory.MemoryGraph",
|
||||||
"apache_age": "mem0.memory.apache_age_memory.MemoryGraph",
|
|
||||||
"default": "mem0.memory.graph_memory.MemoryGraph",
|
"default": "mem0.memory.graph_memory.MemoryGraph",
|
||||||
}
|
}
|
||||||
|
|
||||||
|
|||||||
@@ -34,7 +34,6 @@ class VectorStoreConfig(BaseModel):
|
|||||||
"faiss": "FAISSConfig",
|
"faiss": "FAISSConfig",
|
||||||
"langchain": "LangchainConfig",
|
"langchain": "LangchainConfig",
|
||||||
"s3_vectors": "S3VectorsConfig",
|
"s3_vectors": "S3VectorsConfig",
|
||||||
"turbopuffer": "TurbopufferConfig",
|
|
||||||
}
|
}
|
||||||
|
|
||||||
@model_validator(mode="after")
|
@model_validator(mode="after")
|
||||||
|
|||||||
@@ -26,7 +26,7 @@ class OutputData(BaseModel):
|
|||||||
|
|
||||||
|
|
||||||
class MongoDB(VectorStoreBase):
|
class MongoDB(VectorStoreBase):
|
||||||
VECTOR_TYPE = "vector"
|
VECTOR_TYPE = "knnVector"
|
||||||
SIMILARITY_METRIC = "cosine"
|
SIMILARITY_METRIC = "cosine"
|
||||||
|
|
||||||
def __init__(self, db_name: str, collection_name: str, embedding_model_dims: int, mongo_uri: str):
|
def __init__(self, db_name: str, collection_name: str, embedding_model_dims: int, mongo_uri: str):
|
||||||
@@ -69,16 +69,17 @@ class MongoDB(VectorStoreBase):
|
|||||||
else:
|
else:
|
||||||
search_index_model = SearchIndexModel(
|
search_index_model = SearchIndexModel(
|
||||||
name=self.index_name,
|
name=self.index_name,
|
||||||
type="vectorSearch",
|
|
||||||
definition={
|
definition={
|
||||||
"fields": [
|
"mappings": {
|
||||||
{
|
"dynamic": False,
|
||||||
"type": self.VECTOR_TYPE,
|
"fields": {
|
||||||
"path": "embedding",
|
"embedding": {
|
||||||
"numDimensions": self.embedding_model_dims,
|
"type": self.VECTOR_TYPE,
|
||||||
"similarity": self.SIMILARITY_METRIC,
|
"dimensions": self.embedding_model_dims,
|
||||||
}
|
"similarity": self.SIMILARITY_METRIC,
|
||||||
]
|
}
|
||||||
|
},
|
||||||
|
}
|
||||||
},
|
},
|
||||||
)
|
)
|
||||||
collection.create_search_index(search_index_model)
|
collection.create_search_index(search_index_model)
|
||||||
@@ -140,7 +141,7 @@ class MongoDB(VectorStoreBase):
|
|||||||
"$vectorSearch": {
|
"$vectorSearch": {
|
||||||
"index": self.index_name,
|
"index": self.index_name,
|
||||||
"limit": limit,
|
"limit": limit,
|
||||||
"numCandidates": min(limit * 20, 10000),
|
"numCandidates": limit,
|
||||||
"queryVector": vectors,
|
"queryVector": vectors,
|
||||||
"path": "embedding",
|
"path": "embedding",
|
||||||
}
|
}
|
||||||
|
|||||||
@@ -113,25 +113,6 @@ class OpenSearchDB(VectorStoreBase):
|
|||||||
if payloads is None:
|
if payloads is None:
|
||||||
payloads = [{} for _ in range(len(vectors))]
|
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 = []
|
results = []
|
||||||
for i, (vec, id_) in enumerate(zip(vectors, ids)):
|
for i, (vec, id_) in enumerate(zip(vectors, ids)):
|
||||||
body = {
|
body = {
|
||||||
@@ -143,16 +124,14 @@ class OpenSearchDB(VectorStoreBase):
|
|||||||
self.client.index(index=self.collection_name, body=body)
|
self.client.index(index=self.collection_name, body=body)
|
||||||
# Force refresh to make documents immediately searchable for tests
|
# Force refresh to make documents immediately searchable for tests
|
||||||
self.client.indices.refresh(index=self.collection_name)
|
self.client.indices.refresh(index=self.collection_name)
|
||||||
|
|
||||||
results.append(
|
results.append(OutputData(
|
||||||
OutputData(
|
id=id_,
|
||||||
id=id_,
|
score=1.0, # No score for inserts
|
||||||
score=1.0, # No score for inserts
|
payload=payloads[i]
|
||||||
payload=payloads[i],
|
))
|
||||||
)
|
|
||||||
)
|
|
||||||
except Exception as e:
|
except Exception as e:
|
||||||
logger.error(f"Error inserting vector {id_}: {e}", exc_info=True)
|
logger.error(f"Error inserting vector {id_}: {e}")
|
||||||
raise
|
raise
|
||||||
|
|
||||||
return results
|
return results
|
||||||
@@ -200,7 +179,7 @@ class OpenSearchDB(VectorStoreBase):
|
|||||||
]
|
]
|
||||||
return results
|
return results
|
||||||
except Exception as e:
|
except Exception as e:
|
||||||
logger.error(f"Error during search: {e}", exc_info=True)
|
logger.error(f"Error during search: {e}")
|
||||||
return []
|
return []
|
||||||
|
|
||||||
def delete(self, vector_id: str) -> None:
|
def delete(self, vector_id: str) -> None:
|
||||||
@@ -221,15 +200,6 @@ class OpenSearchDB(VectorStoreBase):
|
|||||||
|
|
||||||
def update(self, vector_id: str, vector: Optional[List[float]] = None, payload: Optional[Dict] = None) -> None:
|
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."""
|
"""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
|
# First, find the document by custom ID
|
||||||
search_query = {"query": {"term": {"id": vector_id}}}
|
search_query = {"query": {"term": {"id": vector_id}}}
|
||||||
@@ -252,9 +222,8 @@ class OpenSearchDB(VectorStoreBase):
|
|||||||
if doc:
|
if doc:
|
||||||
try:
|
try:
|
||||||
response = self.client.update(index=self.collection_name, id=opensearch_id, body={"doc": doc})
|
response = self.client.update(index=self.collection_name, id=opensearch_id, body={"doc": doc})
|
||||||
except Exception as e:
|
except Exception:
|
||||||
logger.error(f"Error updating vector {vector_id}: {e}", exc_info=True)
|
pass
|
||||||
raise
|
|
||||||
|
|
||||||
def get(self, vector_id: str) -> Optional[OutputData]:
|
def get(self, vector_id: str) -> Optional[OutputData]:
|
||||||
"""Retrieve a vector by ID."""
|
"""Retrieve a vector by ID."""
|
||||||
@@ -269,7 +238,7 @@ class OpenSearchDB(VectorStoreBase):
|
|||||||
|
|
||||||
return OutputData(id=hits[0]["_source"].get("id"), score=1.0, payload=hits[0]["_source"].get("payload", {}))
|
return OutputData(id=hits[0]["_source"].get("id"), score=1.0, payload=hits[0]["_source"].get("payload", {}))
|
||||||
except Exception as e:
|
except Exception as e:
|
||||||
logger.error(f"Error retrieving vector {vector_id}: {str(e)}", exc_info=True)
|
logger.error(f"Error retrieving vector {vector_id}: {str(e)}")
|
||||||
return None
|
return None
|
||||||
|
|
||||||
def list_cols(self) -> List[str]:
|
def list_cols(self) -> List[str]:
|
||||||
@@ -312,8 +281,9 @@ class OpenSearchDB(VectorStoreBase):
|
|||||||
]
|
]
|
||||||
return [results] # VectorStore expects tuple/list format
|
return [results] # VectorStore expects tuple/list format
|
||||||
except Exception as e:
|
except Exception as e:
|
||||||
logger.error(f"Error listing vectors: {e}", exc_info=True)
|
logger.error(f"Error listing vectors: {e}")
|
||||||
return []
|
return []
|
||||||
|
|
||||||
|
|
||||||
def reset(self):
|
def reset(self):
|
||||||
"""Reset the index by deleting and recreating it."""
|
"""Reset the index by deleting and recreating it."""
|
||||||
|
|||||||
@@ -1,16 +1,12 @@
|
|||||||
import logging
|
import logging
|
||||||
import os
|
import os
|
||||||
import shutil
|
import shutil
|
||||||
from typing import Optional
|
|
||||||
|
|
||||||
from qdrant_client import QdrantClient
|
from qdrant_client import QdrantClient
|
||||||
from qdrant_client.models import (
|
from qdrant_client.models import (
|
||||||
Distance,
|
Distance,
|
||||||
FieldCondition,
|
FieldCondition,
|
||||||
Filter,
|
Filter,
|
||||||
MatchAny,
|
|
||||||
MatchExcept,
|
|
||||||
MatchText,
|
|
||||||
MatchValue,
|
MatchValue,
|
||||||
PointIdsList,
|
PointIdsList,
|
||||||
PointStruct,
|
PointStruct,
|
||||||
@@ -142,148 +138,26 @@ class Qdrant(VectorStoreBase):
|
|||||||
]
|
]
|
||||||
self.client.upsert(collection_name=self.collection_name, points=points)
|
self.client.upsert(collection_name=self.collection_name, points=points)
|
||||||
|
|
||||||
def _build_field_condition(self, key: str, value) -> Optional[FieldCondition]:
|
def _create_filter(self, filters: dict) -> Filter:
|
||||||
"""
|
|
||||||
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.
|
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:
|
Args:
|
||||||
filters (dict): Filters to apply.
|
filters (dict): Filters to apply.
|
||||||
|
|
||||||
Returns:
|
Returns:
|
||||||
Filter: The created Filter object, or None if filters is empty.
|
Filter: The created Filter object.
|
||||||
"""
|
"""
|
||||||
if not filters:
|
if not filters:
|
||||||
return None
|
return None
|
||||||
|
|
||||||
# Normalize $or/$not/$and → OR/NOT/AND and deduplicate.
|
conditions = []
|
||||||
# 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():
|
for key, value in filters.items():
|
||||||
norm_key = key_map.get(key, key)
|
if isinstance(value, dict) and "gte" in value and "lte" in value:
|
||||||
if norm_key not in normalized:
|
conditions.append(FieldCondition(key=key, range=Range(gte=value["gte"], lte=value["lte"])))
|
||||||
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:
|
else:
|
||||||
condition = self._build_field_condition(key, value)
|
conditions.append(FieldCondition(key=key, match=MatchValue(value=value)))
|
||||||
if condition is not None:
|
return Filter(must=conditions) if conditions else 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:
|
def search(self, query: str, vectors: list, limit: int = 5, filters: dict = None) -> list:
|
||||||
"""
|
"""
|
||||||
|
|||||||
@@ -1,337 +0,0 @@
|
|||||||
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()
|
|
||||||
@@ -184,7 +184,7 @@ async def search_memory(query: str) -> str:
|
|||||||
for h in hits:
|
for h in hits:
|
||||||
# All vector db search functions return OutputData class
|
# All vector db search functions return OutputData class
|
||||||
id, score, payload = h.id, h.score, h.payload
|
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
|
continue
|
||||||
|
|
||||||
results.append({
|
results.append({
|
||||||
|
|||||||
@@ -31,7 +31,6 @@ graph = [
|
|||||||
"neo4j>=5.23.1",
|
"neo4j>=5.23.1",
|
||||||
"rank-bm25>=0.2.2",
|
"rank-bm25>=0.2.2",
|
||||||
"kuzu>=0.11.0",
|
"kuzu>=0.11.0",
|
||||||
"apache-age-python>=0.0.6",
|
|
||||||
]
|
]
|
||||||
vector_stores = [
|
vector_stores = [
|
||||||
"vecs>=0.4.0",
|
"vecs>=0.4.0",
|
||||||
|
|||||||
@@ -10,6 +10,3 @@ POSTGRES_DB=
|
|||||||
POSTGRES_USER=
|
POSTGRES_USER=
|
||||||
POSTGRES_PASSWORD=
|
POSTGRES_PASSWORD=
|
||||||
POSTGRES_COLLECTION_NAME=
|
POSTGRES_COLLECTION_NAME=
|
||||||
|
|
||||||
# Optional: set to enable API key authentication on all endpoints
|
|
||||||
ADMIN_API_KEY=
|
|
||||||
|
|||||||
+10
-55
@@ -1,12 +1,10 @@
|
|||||||
import logging
|
import logging
|
||||||
import os
|
import os
|
||||||
import secrets
|
|
||||||
from typing import Any, Dict, List, Optional
|
from typing import Any, Dict, List, Optional
|
||||||
|
|
||||||
from dotenv import load_dotenv
|
from dotenv import load_dotenv
|
||||||
from fastapi import Depends, FastAPI, HTTPException
|
from fastapi import FastAPI, HTTPException
|
||||||
from fastapi.responses import JSONResponse, RedirectResponse
|
from fastapi.responses import JSONResponse, RedirectResponse
|
||||||
from fastapi.security import APIKeyHeader
|
|
||||||
from pydantic import BaseModel, Field
|
from pydantic import BaseModel, Field
|
||||||
|
|
||||||
from mem0 import Memory
|
from mem0 import Memory
|
||||||
@@ -16,22 +14,6 @@ logging.basicConfig(level=logging.INFO, format="%(asctime)s - %(levelname)s - %(
|
|||||||
# Load environment variables
|
# Load environment variables
|
||||||
load_dotenv()
|
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_HOST = os.environ.get("POSTGRES_HOST", "postgres")
|
||||||
POSTGRES_PORT = os.environ.get("POSTGRES_PORT", "5432")
|
POSTGRES_PORT = os.environ.get("POSTGRES_PORT", "5432")
|
||||||
@@ -78,35 +60,10 @@ MEMORY_INSTANCE = Memory.from_config(DEFAULT_CONFIG)
|
|||||||
|
|
||||||
app = FastAPI(
|
app = FastAPI(
|
||||||
title="Mem0 REST APIs",
|
title="Mem0 REST APIs",
|
||||||
description=(
|
description="A REST API for managing and searching memories for your AI Agents and Apps.",
|
||||||
"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",
|
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):
|
class Message(BaseModel):
|
||||||
role: str = Field(..., description="Role of the message (user or assistant).")
|
role: str = Field(..., description="Role of the message (user or assistant).")
|
||||||
@@ -130,7 +87,7 @@ class SearchRequest(BaseModel):
|
|||||||
|
|
||||||
|
|
||||||
@app.post("/configure", summary="Configure Mem0")
|
@app.post("/configure", summary="Configure Mem0")
|
||||||
def set_config(config: Dict[str, Any], _api_key: Optional[str] = Depends(verify_api_key)):
|
def set_config(config: Dict[str, Any]):
|
||||||
"""Set memory configuration."""
|
"""Set memory configuration."""
|
||||||
global MEMORY_INSTANCE
|
global MEMORY_INSTANCE
|
||||||
MEMORY_INSTANCE = Memory.from_config(config)
|
MEMORY_INSTANCE = Memory.from_config(config)
|
||||||
@@ -138,7 +95,7 @@ def set_config(config: Dict[str, Any], _api_key: Optional[str] = Depends(verify_
|
|||||||
|
|
||||||
|
|
||||||
@app.post("/memories", summary="Create memories")
|
@app.post("/memories", summary="Create memories")
|
||||||
def add_memory(memory_create: MemoryCreate, _api_key: Optional[str] = Depends(verify_api_key)):
|
def add_memory(memory_create: MemoryCreate):
|
||||||
"""Store new memories."""
|
"""Store new memories."""
|
||||||
if not any([memory_create.user_id, memory_create.agent_id, memory_create.run_id]):
|
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.")
|
raise HTTPException(status_code=400, detail="At least one identifier (user_id, agent_id, run_id) is required.")
|
||||||
@@ -157,7 +114,6 @@ def get_all_memories(
|
|||||||
user_id: Optional[str] = None,
|
user_id: Optional[str] = None,
|
||||||
run_id: Optional[str] = None,
|
run_id: Optional[str] = None,
|
||||||
agent_id: Optional[str] = None,
|
agent_id: Optional[str] = None,
|
||||||
_api_key: Optional[str] = Depends(verify_api_key),
|
|
||||||
):
|
):
|
||||||
"""Retrieve stored memories."""
|
"""Retrieve stored memories."""
|
||||||
if not any([user_id, run_id, agent_id]):
|
if not any([user_id, run_id, agent_id]):
|
||||||
@@ -173,7 +129,7 @@ def get_all_memories(
|
|||||||
|
|
||||||
|
|
||||||
@app.get("/memories/{memory_id}", summary="Get a memory")
|
@app.get("/memories/{memory_id}", summary="Get a memory")
|
||||||
def get_memory(memory_id: str, _api_key: Optional[str] = Depends(verify_api_key)):
|
def get_memory(memory_id: str):
|
||||||
"""Retrieve a specific memory by ID."""
|
"""Retrieve a specific memory by ID."""
|
||||||
try:
|
try:
|
||||||
return MEMORY_INSTANCE.get(memory_id)
|
return MEMORY_INSTANCE.get(memory_id)
|
||||||
@@ -183,7 +139,7 @@ def get_memory(memory_id: str, _api_key: Optional[str] = Depends(verify_api_key)
|
|||||||
|
|
||||||
|
|
||||||
@app.post("/search", summary="Search memories")
|
@app.post("/search", summary="Search memories")
|
||||||
def search_memories(search_req: SearchRequest, _api_key: Optional[str] = Depends(verify_api_key)):
|
def search_memories(search_req: SearchRequest):
|
||||||
"""Search for memories based on a query."""
|
"""Search for memories based on a query."""
|
||||||
try:
|
try:
|
||||||
params = {k: v for k, v in search_req.model_dump().items() if v is not None and k != "query"}
|
params = {k: v for k, v in search_req.model_dump().items() if v is not None and k != "query"}
|
||||||
@@ -194,7 +150,7 @@ def search_memories(search_req: SearchRequest, _api_key: Optional[str] = Depends
|
|||||||
|
|
||||||
|
|
||||||
@app.put("/memories/{memory_id}", summary="Update a memory")
|
@app.put("/memories/{memory_id}", summary="Update a memory")
|
||||||
def update_memory(memory_id: str, updated_memory: Dict[str, Any], _api_key: Optional[str] = Depends(verify_api_key)):
|
def update_memory(memory_id: str, updated_memory: Dict[str, Any]):
|
||||||
"""Update an existing memory with new content.
|
"""Update an existing memory with new content.
|
||||||
|
|
||||||
Args:
|
Args:
|
||||||
@@ -212,7 +168,7 @@ def update_memory(memory_id: str, updated_memory: Dict[str, Any], _api_key: Opti
|
|||||||
|
|
||||||
|
|
||||||
@app.get("/memories/{memory_id}/history", summary="Get memory history")
|
@app.get("/memories/{memory_id}/history", summary="Get memory history")
|
||||||
def memory_history(memory_id: str, _api_key: Optional[str] = Depends(verify_api_key)):
|
def memory_history(memory_id: str):
|
||||||
"""Retrieve memory history."""
|
"""Retrieve memory history."""
|
||||||
try:
|
try:
|
||||||
return MEMORY_INSTANCE.history(memory_id=memory_id)
|
return MEMORY_INSTANCE.history(memory_id=memory_id)
|
||||||
@@ -222,7 +178,7 @@ def memory_history(memory_id: str, _api_key: Optional[str] = Depends(verify_api_
|
|||||||
|
|
||||||
|
|
||||||
@app.delete("/memories/{memory_id}", summary="Delete a memory")
|
@app.delete("/memories/{memory_id}", summary="Delete a memory")
|
||||||
def delete_memory(memory_id: str, _api_key: Optional[str] = Depends(verify_api_key)):
|
def delete_memory(memory_id: str):
|
||||||
"""Delete a specific memory by ID."""
|
"""Delete a specific memory by ID."""
|
||||||
try:
|
try:
|
||||||
MEMORY_INSTANCE.delete(memory_id=memory_id)
|
MEMORY_INSTANCE.delete(memory_id=memory_id)
|
||||||
@@ -237,7 +193,6 @@ def delete_all_memories(
|
|||||||
user_id: Optional[str] = None,
|
user_id: Optional[str] = None,
|
||||||
run_id: Optional[str] = None,
|
run_id: Optional[str] = None,
|
||||||
agent_id: Optional[str] = None,
|
agent_id: Optional[str] = None,
|
||||||
_api_key: Optional[str] = Depends(verify_api_key),
|
|
||||||
):
|
):
|
||||||
"""Delete all memories for a given identifier."""
|
"""Delete all memories for a given identifier."""
|
||||||
if not any([user_id, run_id, agent_id]):
|
if not any([user_id, run_id, agent_id]):
|
||||||
@@ -254,7 +209,7 @@ def delete_all_memories(
|
|||||||
|
|
||||||
|
|
||||||
@app.post("/reset", summary="Reset all memories")
|
@app.post("/reset", summary="Reset all memories")
|
||||||
def reset_memory(_api_key: Optional[str] = Depends(verify_api_key)):
|
def reset_memory():
|
||||||
"""Completely reset stored memories."""
|
"""Completely reset stored memories."""
|
||||||
try:
|
try:
|
||||||
MEMORY_INSTANCE.reset()
|
MEMORY_INSTANCE.reset()
|
||||||
|
|||||||
@@ -1,5 +1,4 @@
|
|||||||
from unittest import mock
|
from unittest import mock
|
||||||
from unittest.mock import Mock
|
|
||||||
|
|
||||||
from mem0.llms.azure_openai_structured import SCOPE, AzureOpenAIStructuredLLM
|
from mem0.llms.azure_openai_structured import SCOPE, AzureOpenAIStructuredLLM
|
||||||
|
|
||||||
@@ -99,75 +98,3 @@ def test_init_with_placeholder_api_key_uses_default_credential(
|
|||||||
args, kwargs = mock_azure_openai.call_args
|
args, kwargs = mock_azure_openai.call_args
|
||||||
assert kwargs["api_key"] is None
|
assert kwargs["api_key"] is None
|
||||||
assert kwargs["azure_ad_token_provider"] == "token-provider"
|
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"] == []
|
|
||||||
|
|||||||
@@ -117,73 +117,3 @@ def test_generate_response_with_tools(mock_gemini_client: Mock):
|
|||||||
assert len(response["tool_calls"]) == 1
|
assert len(response["tool_calls"]) == 1
|
||||||
assert response["tool_calls"][0]["name"] == "add_memory"
|
assert response["tool_calls"][0]["name"] == "add_memory"
|
||||||
assert response["tool_calls"][0]["arguments"] == {"data": "Today is a sunny day."}
|
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": []}
|
|
||||||
|
|||||||
File diff suppressed because it is too large
Load Diff
@@ -1,226 +0,0 @@
|
|||||||
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 == []
|
|
||||||
@@ -1,315 +0,0 @@
|
|||||||
"""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
|
|
||||||
@@ -320,255 +320,3 @@ def test_normalize_facts_mixed():
|
|||||||
|
|
||||||
def test_normalize_facts_filters_empty_strings():
|
def test_normalize_facts_filters_empty_strings():
|
||||||
assert normalize_facts(["", "valid", ""]) == ["valid"]
|
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()
|
|
||||||
|
|
||||||
|
|
||||||
@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_add_infer_false_embeds_once(mock_sqlite, mock_llm_factory, mock_vector_factory, mock_embedder_factory):
|
|
||||||
"""
|
|
||||||
Regression test for issue #3723: adding with infer=False should not trigger duplicate embedding calls.
|
|
||||||
|
|
||||||
Root cause: _create_memory expected a dict for existing_embeddings but received a raw list[float],
|
|
||||||
causing the cache check `data in existing_embeddings` to always fail and trigger a redundant embed.
|
|
||||||
"""
|
|
||||||
embedder = MagicMock()
|
|
||||||
embedder.embed.return_value = [0.1, 0.2, 0.3]
|
|
||||||
embedder.config = MagicMock(embedding_dims=3)
|
|
||||||
mock_embedder_factory.return_value = embedder
|
|
||||||
|
|
||||||
mock_vector_store = MagicMock()
|
|
||||||
mock_vector_store.search.return_value = []
|
|
||||||
mock_vector_store.insert.return_value = None
|
|
||||||
mock_vector_store.get.return_value = None
|
|
||||||
telemetry_vector_store = MagicMock()
|
|
||||||
mock_vector_factory.side_effect = [mock_vector_store, telemetry_vector_store]
|
|
||||||
|
|
||||||
mock_llm_factory.return_value = MagicMock()
|
|
||||||
mock_sqlite.return_value = MagicMock()
|
|
||||||
|
|
||||||
from mem0.memory.main import Memory as MemoryClass
|
|
||||||
memory = MemoryClass(MemoryConfig())
|
|
||||||
|
|
||||||
memory.add("foo", user_id="test_user", infer=False)
|
|
||||||
|
|
||||||
assert embedder.embed.call_count == 1
|
|
||||||
mock_vector_store.insert.assert_called_once()
|
|
||||||
|
|
||||||
|
|
||||||
@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_add_infer_true_caches_embedding_on_llm_rewrite(mock_sqlite, mock_llm_factory, mock_vector_factory, mock_embedder_factory):
|
|
||||||
"""
|
|
||||||
Regression test for issue #3723 (infer=True path): when the LLM rewrites a fact during the
|
|
||||||
ADD action, the embedding should be computed once and cached, not computed again inside _create_memory.
|
|
||||||
"""
|
|
||||||
embedder = MagicMock()
|
|
||||||
embedder.embed.return_value = [0.1, 0.2, 0.3]
|
|
||||||
embedder.config = MagicMock(embedding_dims=3)
|
|
||||||
mock_embedder_factory.return_value = embedder
|
|
||||||
|
|
||||||
mock_vector_store = MagicMock()
|
|
||||||
mock_vector_store.search.return_value = []
|
|
||||||
mock_vector_store.insert.return_value = None
|
|
||||||
mock_vector_store.get.return_value = None
|
|
||||||
telemetry_vector_store = MagicMock()
|
|
||||||
mock_vector_factory.side_effect = [mock_vector_store, telemetry_vector_store]
|
|
||||||
|
|
||||||
# LLM extracts fact "User likes Python", then ADD action rewrites to "The user enjoys Python"
|
|
||||||
mock_llm = MagicMock()
|
|
||||||
mock_llm.generate_response.side_effect = [
|
|
||||||
json.dumps({"facts": ["User likes Python"]}),
|
|
||||||
json.dumps({"memory": [{"id": "0", "text": "The user enjoys Python", "event": "ADD", "old_memory": None}]}),
|
|
||||||
]
|
|
||||||
mock_llm_factory.return_value = mock_llm
|
|
||||||
|
|
||||||
mock_sqlite.return_value = MagicMock()
|
|
||||||
|
|
||||||
from mem0.memory.main import Memory as MemoryClass
|
|
||||||
memory = MemoryClass(MemoryConfig())
|
|
||||||
|
|
||||||
memory.add("I like Python", user_id="test_user", infer=True)
|
|
||||||
|
|
||||||
# embed should be called exactly twice:
|
|
||||||
# 1. For the extracted fact "User likes Python" (search)
|
|
||||||
# 2. For the rewritten text "The user enjoys Python" (pre-cached before _create_memory)
|
|
||||||
# It should NOT be called a 3rd time inside _create_memory
|
|
||||||
assert embedder.embed.call_count == 2
|
|
||||||
mock_vector_store.insert.assert_called_once()
|
|
||||||
|
|
||||||
|
|
||||||
@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_infer_true_caches_embedding_on_llm_rewrite(mock_sqlite, mock_llm_factory, mock_vector_factory, mock_embedder_factory):
|
|
||||||
"""
|
|
||||||
Regression test for issue #3723 (infer=True UPDATE path): when the LLM rewrites a fact during
|
|
||||||
an UPDATE action, the embedding should be computed once and cached, not computed again inside _update_memory.
|
|
||||||
"""
|
|
||||||
embedder = MagicMock()
|
|
||||||
embedder.embed.return_value = [0.1, 0.2, 0.3]
|
|
||||||
embedder.config = MagicMock(embedding_dims=3)
|
|
||||||
mock_embedder_factory.return_value = embedder
|
|
||||||
|
|
||||||
# Existing memory that will be matched for update
|
|
||||||
existing_memory = MockVectorMemory(
|
|
||||||
memory_id="existing-mem-id",
|
|
||||||
payload={
|
|
||||||
"data": "User likes Python",
|
|
||||||
"hash": "abc123",
|
|
||||||
"created_at": "2025-01-01T00:00:00+00:00",
|
|
||||||
},
|
|
||||||
)
|
|
||||||
|
|
||||||
mock_vector_store = MagicMock()
|
|
||||||
mock_vector_store.search.return_value = [existing_memory]
|
|
||||||
mock_vector_store.get.return_value = existing_memory
|
|
||||||
mock_vector_store.insert.return_value = None
|
|
||||||
mock_vector_store.update.return_value = None
|
|
||||||
telemetry_vector_store = MagicMock()
|
|
||||||
mock_vector_factory.side_effect = [mock_vector_store, telemetry_vector_store]
|
|
||||||
|
|
||||||
# LLM extracts fact "User loves Python now", then UPDATE action rewrites to "The user loves Python"
|
|
||||||
mock_llm = MagicMock()
|
|
||||||
mock_llm.generate_response.side_effect = [
|
|
||||||
json.dumps({"facts": ["User loves Python now"]}),
|
|
||||||
json.dumps({"memory": [{"id": "0", "text": "The user loves Python", "event": "UPDATE", "old_memory": "User likes Python"}]}),
|
|
||||||
]
|
|
||||||
mock_llm_factory.return_value = mock_llm
|
|
||||||
|
|
||||||
mock_sqlite.return_value = MagicMock()
|
|
||||||
|
|
||||||
from mem0.memory.main import Memory as MemoryClass
|
|
||||||
memory = MemoryClass(MemoryConfig())
|
|
||||||
|
|
||||||
memory.add("I love Python now", user_id="test_user", infer=True)
|
|
||||||
|
|
||||||
# embed should be called exactly twice:
|
|
||||||
# 1. For the extracted fact "User loves Python now" (search)
|
|
||||||
# 2. For the rewritten text "The user loves Python" (pre-cached before _update_memory)
|
|
||||||
# It should NOT be called a 3rd time inside _update_memory
|
|
||||||
assert embedder.embed.call_count == 2
|
|
||||||
mock_vector_store.update.assert_called_once()
|
|
||||||
|
|||||||
@@ -1,496 +0,0 @@
|
|||||||
"""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)
|
|
||||||
@@ -48,16 +48,17 @@ def test_initalize_create_col(mongo_vector_fixture):
|
|||||||
search_index_model = args[0].document
|
search_index_model = args[0].document
|
||||||
assert search_index_model == {
|
assert search_index_model == {
|
||||||
"name": "test_collection_vector_index",
|
"name": "test_collection_vector_index",
|
||||||
"type": "vectorSearch",
|
|
||||||
"definition": {
|
"definition": {
|
||||||
"fields": [
|
"mappings": {
|
||||||
{
|
"dynamic": False,
|
||||||
"type": "vector",
|
"fields": {
|
||||||
"path": "embedding",
|
"embedding": {
|
||||||
"numDimensions": 1536,
|
"type": "knnVector",
|
||||||
"similarity": "cosine",
|
"dimensions": 1536,
|
||||||
}
|
"similarity": "cosine",
|
||||||
]
|
}
|
||||||
|
},
|
||||||
|
}
|
||||||
},
|
},
|
||||||
}
|
}
|
||||||
assert mongo_vector.collection == mock_collection
|
assert mongo_vector.collection == mock_collection
|
||||||
@@ -94,7 +95,7 @@ def test_search(mongo_vector_fixture):
|
|||||||
"$vectorSearch": {
|
"$vectorSearch": {
|
||||||
"index": "test_collection_vector_index",
|
"index": "test_collection_vector_index",
|
||||||
"limit": 2,
|
"limit": 2,
|
||||||
"numCandidates": 40,
|
"numCandidates": 2,
|
||||||
"queryVector": query_vector,
|
"queryVector": query_vector,
|
||||||
"path": "embedding",
|
"path": "embedding",
|
||||||
},
|
},
|
||||||
|
|||||||
@@ -18,25 +18,25 @@ from mem0.vector_stores.opensearch import OpenSearchDB
|
|||||||
# Mock classes for testing OpenSearch with AWS authentication
|
# Mock classes for testing OpenSearch with AWS authentication
|
||||||
class MockFieldInfo:
|
class MockFieldInfo:
|
||||||
"""Mock pydantic field info."""
|
"""Mock pydantic field info."""
|
||||||
|
|
||||||
def __init__(self, default=None):
|
def __init__(self, default=None):
|
||||||
self.default = default
|
self.default = default
|
||||||
|
|
||||||
|
|
||||||
class MockOpenSearchConfig:
|
class MockOpenSearchConfig:
|
||||||
|
|
||||||
model_fields = {
|
model_fields = {
|
||||||
"collection_name": MockFieldInfo(default="default_collection"),
|
'collection_name': MockFieldInfo(default="default_collection"),
|
||||||
"host": MockFieldInfo(default="localhost"),
|
'host': MockFieldInfo(default="localhost"),
|
||||||
"port": MockFieldInfo(default=9200),
|
'port': MockFieldInfo(default=9200),
|
||||||
"embedding_model_dims": MockFieldInfo(default=1536),
|
'embedding_model_dims': MockFieldInfo(default=1536),
|
||||||
"http_auth": MockFieldInfo(default=None),
|
'http_auth': MockFieldInfo(default=None),
|
||||||
"auth": MockFieldInfo(default=None),
|
'auth': MockFieldInfo(default=None),
|
||||||
"credentials": MockFieldInfo(default=None),
|
'credentials': MockFieldInfo(default=None),
|
||||||
"connection_class": MockFieldInfo(default=None),
|
'connection_class': MockFieldInfo(default=None),
|
||||||
"use_ssl": MockFieldInfo(default=False),
|
'use_ssl': MockFieldInfo(default=False),
|
||||||
"verify_certs": MockFieldInfo(default=False),
|
'verify_certs': MockFieldInfo(default=False),
|
||||||
}
|
}
|
||||||
|
|
||||||
def __init__(self, collection_name="test_collection", include_auth=True, **kwargs):
|
def __init__(self, collection_name="test_collection", include_auth=True, **kwargs):
|
||||||
self.collection_name = collection_name
|
self.collection_name = collection_name
|
||||||
self.host = kwargs.get("host", "localhost")
|
self.host = kwargs.get("host", "localhost")
|
||||||
@@ -44,7 +44,7 @@ class MockOpenSearchConfig:
|
|||||||
self.embedding_model_dims = kwargs.get("embedding_model_dims", 1536)
|
self.embedding_model_dims = kwargs.get("embedding_model_dims", 1536)
|
||||||
self.use_ssl = kwargs.get("use_ssl", True)
|
self.use_ssl = kwargs.get("use_ssl", True)
|
||||||
self.verify_certs = kwargs.get("verify_certs", True)
|
self.verify_certs = kwargs.get("verify_certs", True)
|
||||||
|
|
||||||
if any(field in kwargs for field in ["http_auth", "auth", "credentials", "connection_class"]):
|
if any(field in kwargs for field in ["http_auth", "auth", "credentials", "connection_class"]):
|
||||||
self.http_auth = kwargs.get("http_auth")
|
self.http_auth = kwargs.get("http_auth")
|
||||||
self.auth = kwargs.get("auth")
|
self.auth = kwargs.get("auth")
|
||||||
@@ -63,18 +63,20 @@ class MockOpenSearchConfig:
|
|||||||
|
|
||||||
|
|
||||||
class MockAWSAuth:
|
class MockAWSAuth:
|
||||||
|
|
||||||
def __init__(self):
|
def __init__(self):
|
||||||
self._lock = threading.Lock()
|
self._lock = threading.Lock()
|
||||||
self.region = "us-east-1"
|
self.region = "us-east-1"
|
||||||
|
|
||||||
def __deepcopy__(self, memo):
|
def __deepcopy__(self, memo):
|
||||||
raise TypeError("cannot pickle '_thread.lock' object")
|
raise TypeError("cannot pickle '_thread.lock' object")
|
||||||
|
|
||||||
|
|
||||||
class MockConnectionClass:
|
class MockConnectionClass:
|
||||||
|
|
||||||
def __init__(self):
|
def __init__(self):
|
||||||
self._state = {"connected": False}
|
self._state = {"connected": False}
|
||||||
|
|
||||||
def __deepcopy__(self, memo):
|
def __deepcopy__(self, memo):
|
||||||
raise TypeError("cannot pickle connection state")
|
raise TypeError("cannot pickle connection state")
|
||||||
|
|
||||||
@@ -253,104 +255,6 @@ class TestOpenSearchDB(unittest.TestCase):
|
|||||||
self.os_db.delete_col()
|
self.os_db.delete_col()
|
||||||
self.client_mock.indices.delete.assert_called_once_with(index="test_collection")
|
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):
|
def test_init_with_http_auth(self):
|
||||||
mock_credentials = MagicMock()
|
mock_credentials = MagicMock()
|
||||||
mock_signer = AWSV4SignerAuth(mock_credentials, "us-east-1", "es")
|
mock_signer = AWSV4SignerAuth(mock_credentials, "us-east-1", "es")
|
||||||
@@ -378,13 +282,11 @@ class TestOpenSearchDB(unittest.TestCase):
|
|||||||
|
|
||||||
|
|
||||||
# Tests for OpenSearch config deepcopy with AWS authentication (Issue #3464)
|
# Tests for OpenSearch config deepcopy with AWS authentication (Issue #3464)
|
||||||
@patch("mem0.utils.factory.EmbedderFactory.create")
|
@patch('mem0.utils.factory.EmbedderFactory.create')
|
||||||
@patch("mem0.utils.factory.VectorStoreFactory.create")
|
@patch('mem0.utils.factory.VectorStoreFactory.create')
|
||||||
@patch("mem0.utils.factory.LlmFactory.create")
|
@patch('mem0.utils.factory.LlmFactory.create')
|
||||||
@patch("mem0.memory.storage.SQLiteManager")
|
@patch('mem0.memory.storage.SQLiteManager')
|
||||||
def test_safe_deepcopy_config_handles_opensearch_auth(
|
def test_safe_deepcopy_config_handles_opensearch_auth(mock_sqlite, mock_llm_factory, mock_vector_factory, mock_embedder_factory):
|
||||||
mock_sqlite, mock_llm_factory, mock_vector_factory, mock_embedder_factory
|
|
||||||
):
|
|
||||||
"""Test that _safe_deepcopy_config handles OpenSearch configs with AWS auth objects gracefully."""
|
"""Test that _safe_deepcopy_config handles OpenSearch configs with AWS auth objects gracefully."""
|
||||||
mock_embedder_factory.return_value = MagicMock()
|
mock_embedder_factory.return_value = MagicMock()
|
||||||
mock_vector_store = MagicMock()
|
mock_vector_store = MagicMock()
|
||||||
@@ -393,9 +295,9 @@ def test_safe_deepcopy_config_handles_opensearch_auth(
|
|||||||
mock_sqlite.return_value = MagicMock()
|
mock_sqlite.return_value = MagicMock()
|
||||||
|
|
||||||
from mem0.memory.main import _safe_deepcopy_config
|
from mem0.memory.main import _safe_deepcopy_config
|
||||||
|
|
||||||
config_with_auth = MockOpenSearchConfig(collection_name="opensearch_test", include_auth=True)
|
config_with_auth = MockOpenSearchConfig(collection_name="opensearch_test", include_auth=True)
|
||||||
|
|
||||||
safe_config = _safe_deepcopy_config(config_with_auth)
|
safe_config = _safe_deepcopy_config(config_with_auth)
|
||||||
|
|
||||||
# Runtime auth objects must be preserved (Issue #3580)
|
# Runtime auth objects must be preserved (Issue #3580)
|
||||||
@@ -404,7 +306,7 @@ def test_safe_deepcopy_config_handles_opensearch_auth(
|
|||||||
assert safe_config.connection_class is not None
|
assert safe_config.connection_class is not None
|
||||||
# Credentials dict is a sensitive secret and should be redacted
|
# Credentials dict is a sensitive secret and should be redacted
|
||||||
assert safe_config.credentials is None
|
assert safe_config.credentials is None
|
||||||
|
|
||||||
assert safe_config.collection_name == "opensearch_test"
|
assert safe_config.collection_name == "opensearch_test"
|
||||||
assert safe_config.host == "localhost"
|
assert safe_config.host == "localhost"
|
||||||
assert safe_config.port == 9200
|
assert safe_config.port == 9200
|
||||||
@@ -413,10 +315,10 @@ def test_safe_deepcopy_config_handles_opensearch_auth(
|
|||||||
assert safe_config.verify_certs is True
|
assert safe_config.verify_certs is True
|
||||||
|
|
||||||
|
|
||||||
@patch("mem0.utils.factory.EmbedderFactory.create")
|
@patch('mem0.utils.factory.EmbedderFactory.create')
|
||||||
@patch("mem0.utils.factory.VectorStoreFactory.create")
|
@patch('mem0.utils.factory.VectorStoreFactory.create')
|
||||||
@patch("mem0.utils.factory.LlmFactory.create")
|
@patch('mem0.utils.factory.LlmFactory.create')
|
||||||
@patch("mem0.memory.storage.SQLiteManager")
|
@patch('mem0.memory.storage.SQLiteManager')
|
||||||
def test_safe_deepcopy_config_normal_configs(mock_sqlite, mock_llm_factory, mock_vector_factory, mock_embedder_factory):
|
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."""
|
"""Test that _safe_deepcopy_config handles normal OpenSearch configs without auth."""
|
||||||
mock_embedder_factory.return_value = MagicMock()
|
mock_embedder_factory.return_value = MagicMock()
|
||||||
@@ -426,12 +328,12 @@ def test_safe_deepcopy_config_normal_configs(mock_sqlite, mock_llm_factory, mock
|
|||||||
mock_sqlite.return_value = MagicMock()
|
mock_sqlite.return_value = MagicMock()
|
||||||
|
|
||||||
from mem0.memory.main import _safe_deepcopy_config
|
from mem0.memory.main import _safe_deepcopy_config
|
||||||
|
|
||||||
config_without_auth = MockOpenSearchConfig(collection_name="normal_test", include_auth=False)
|
config_without_auth = MockOpenSearchConfig(collection_name="normal_test", include_auth=False)
|
||||||
|
|
||||||
safe_config = _safe_deepcopy_config(config_without_auth)
|
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.host == "localhost"
|
||||||
assert safe_config.port == 9200
|
assert safe_config.port == 9200
|
||||||
assert safe_config.embedding_model_dims == 1536
|
assert safe_config.embedding_model_dims == 1536
|
||||||
@@ -439,15 +341,13 @@ def test_safe_deepcopy_config_normal_configs(mock_sqlite, mock_llm_factory, mock
|
|||||||
assert safe_config.verify_certs is True
|
assert safe_config.verify_certs is True
|
||||||
|
|
||||||
|
|
||||||
@patch("mem0.utils.factory.EmbedderFactory.create")
|
@patch('mem0.utils.factory.EmbedderFactory.create')
|
||||||
@patch("mem0.utils.factory.VectorStoreFactory.create")
|
@patch('mem0.utils.factory.VectorStoreFactory.create')
|
||||||
@patch("mem0.utils.factory.LlmFactory.create")
|
@patch('mem0.utils.factory.LlmFactory.create')
|
||||||
@patch("mem0.memory.storage.SQLiteManager")
|
@patch('mem0.memory.storage.SQLiteManager')
|
||||||
def test_memory_initialization_opensearch_aws_auth(
|
def test_memory_initialization_opensearch_aws_auth(mock_sqlite, mock_llm_factory, mock_vector_factory, mock_embedder_factory):
|
||||||
mock_sqlite, mock_llm_factory, mock_vector_factory, mock_embedder_factory
|
|
||||||
):
|
|
||||||
"""Test that Memory initialization works with OpenSearch configs containing AWS auth."""
|
"""Test that Memory initialization works with OpenSearch configs containing AWS auth."""
|
||||||
|
|
||||||
mock_embedder_factory.return_value = MagicMock()
|
mock_embedder_factory.return_value = MagicMock()
|
||||||
mock_vector_store = MagicMock()
|
mock_vector_store = MagicMock()
|
||||||
mock_vector_factory.return_value = mock_vector_store
|
mock_vector_factory.return_value = mock_vector_store
|
||||||
|
|||||||
@@ -5,15 +5,9 @@ from unittest.mock import MagicMock
|
|||||||
from qdrant_client import QdrantClient
|
from qdrant_client import QdrantClient
|
||||||
from qdrant_client.models import (
|
from qdrant_client.models import (
|
||||||
Distance,
|
Distance,
|
||||||
FieldCondition,
|
|
||||||
Filter,
|
Filter,
|
||||||
MatchAny,
|
|
||||||
MatchExcept,
|
|
||||||
MatchText,
|
|
||||||
MatchValue,
|
|
||||||
PointIdsList,
|
PointIdsList,
|
||||||
PointStruct,
|
PointStruct,
|
||||||
Range,
|
|
||||||
VectorParams,
|
VectorParams,
|
||||||
)
|
)
|
||||||
|
|
||||||
@@ -304,456 +298,3 @@ class TestQdrant(unittest.TestCase):
|
|||||||
|
|
||||||
def tearDown(self):
|
def tearDown(self):
|
||||||
del self.qdrant
|
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)
|
|
||||||
|
|||||||
@@ -1,725 +0,0 @@
|
|||||||
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
|
|
||||||
Reference in New Issue
Block a user