Compare commits

...
Author SHA1 Message Date
mintlify[bot]andGitHub 09494c2acc Fix broken placeholder links in template files
Generated-By: mintlify-agent
2026-03-21 00:10:06 +00:00
30661ab427 Fix: add pgvector support to NodeJS OSS VectorStoreFactory (fixes #3491) (#3997)
Co-authored-by: kartik-mem0 <kartik.labhshetwar@mem0.ai>
2026-03-20 20:47:04 +05:30
Failfail2603andGitHub 7ad5d6f442 fix: use toCamelCase in redis get method for the payload (#3172) 2026-03-20 20:28:41 +05:30
305ce7b6b3 feat: add Apache AGE graph store support (#4448)
Co-authored-by: utkarsh240799 <utkarsh240799@users.noreply.github.com>
Co-authored-by: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
2026-03-20 20:27:48 +05:30
4437c3e8a8 fix: add missing _parse_response to AzureOpenAIStructuredLLM (#4434)
Co-authored-by: Matt Van Horn <455140+mvanhorn@users.noreply.github.com>
Co-authored-by: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
2026-03-20 20:14:47 +05:30
54bdbde6e6 feat: add MiniMax LLM provider (#4132) (#4431)
Co-authored-by: kartik-mem0 <kartik.labhshetwar@mem0.ai>
2026-03-20 19:18:36 +05:30
KartikandGitHub 2b9558335b fix: raise ValueError when deleting nonexistent memory (#4455) 2026-03-20 18:28:58 +05:30
2520edb404 feat: add optional API key authentication to REST API server (#4442)
Co-authored-by: utkarsh240799 <utkarsh240799@users.noreply.github.com>
Co-authored-by: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
2026-03-20 16:15:25 +05:30
28 changed files with 3235 additions and 26 deletions
+2 -1
View File
@@ -273,7 +273,7 @@ config = MemoryConfig(
### Supported Providers
#### LLM Providers (19 supported)
#### LLM Providers (20 supported)
- **openai** - OpenAI GPT models (default)
- **anthropic** - Claude models
- **gemini** - Google Gemini
@@ -284,6 +284,7 @@ config = MemoryConfig(
- **azure_openai** - Azure OpenAI
- **litellm** - LiteLLM proxy
- **deepseek** - DeepSeek models
- **minimax** - MiniMax models
- **xai** - xAI models
- **sarvam** - Sarvam AI
- **lmstudio** - LM Studio local server
+56
View File
@@ -0,0 +1,56 @@
---
title: MiniMax
description: "Configure MiniMax as an LLM provider in Mem0 with API key setup and optional custom endpoint configuration."
---
To use MiniMax LLM models, you have to set the `MINIMAX_API_KEY` environment variable. You can also optionally set `MINIMAX_API_BASE` if you need to use a different API endpoint (defaults to "https://api.minimax.io/v1").
## Usage
```python
import os
from mem0 import Memory
os.environ["MINIMAX_API_KEY"] = "your-api-key"
os.environ["OPENAI_API_KEY"] = "your-api-key" # for embedder model
config = {
"llm": {
"provider": "minimax",
"config": {
"model": "MiniMax-M2.7", # default model
"temperature": 0.2,
"max_tokens": 2000,
"top_p": 1.0
}
}
}
m = Memory.from_config(config)
messages = [
{"role": "user", "content": "I'm planning to watch a movie tonight. Any recommendations?"},
{"role": "assistant", "content": "How about thriller movies? They can be quite engaging."},
{"role": "user", "content": "I'm not a big fan of thriller movies but I love sci-fi movies."},
{"role": "assistant", "content": "Got it! I'll avoid thriller recommendations and suggest sci-fi movies in the future."}
]
m.add(messages, user_id="alice", metadata={"category": "movies"})
```
You can also configure the API base URL in the config:
```python
config = {
"llm": {
"provider": "minimax",
"config": {
"model": "MiniMax-M2.7",
"minimax_base_url": "https://your-custom-endpoint.com",
"api_key": "your-api-key" # alternatively to using environment variable
}
}
}
```
## Config
All available parameters for the `minimax` config are present in [Master List of All Params in Config](../config).
+1
View File
@@ -31,6 +31,7 @@ See the list of supported LLMs below.
<Card title="Google AI" href="/components/llms/models/google_AI" />
<Card title="AWS bedrock" href="/components/llms/models/aws_bedrock" />
<Card title="DeepSeek" href="/components/llms/models/deepseek" />
<Card title="MiniMax" href="/components/llms/models/minimax" />
<Card title="xAI" href="/components/llms/models/xAI" />
<Card title="Sarvam AI" href="/components/llms/models/sarvam" />
<Card title="LM Studio" href="/components/llms/models/lmstudio" />
+1
View File
@@ -188,6 +188,7 @@
"components/llms/models/google_AI",
"components/llms/models/aws_bedrock",
"components/llms/models/deepseek",
"components/llms/models/minimax",
"components/llms/models/xAI",
"components/llms/models/sarvam",
"components/llms/models/lmstudio",
+38 -2
View File
@@ -35,7 +35,7 @@ graph LR
Mem0’s extraction LLM identifies entities, relationships, and timestamps from the conversation payload you send to `memory.add`.
</Step>
<Step title="Store vectors and edges together">
Embeddings land in your configured vector database while nodes and edges flow into a Bolt-compatible graph backend (Neo4j, Memgraph, Neptune, or Kuzu).
Embeddings land in your configured vector database while nodes and edges flow into a graph backend (Neo4j, Memgraph, Neptune, Kuzu, or Apache AGE).
</Step>
<Step title="Expose graph context at search time">
`memory.search` performs vector similarity (optionally reranked by your configured reranker) and returns the results list. Graph Memory runs in parallel and adds related entities in the `relations` array—it does not reorder the vector hits automatically.
@@ -264,7 +264,7 @@ Monitor graph growth, especially on free tiers, by periodically cleaning dormant
## Decision Points
- Select the graph store that fits your deployment (managed Aura vs. self-hosted Neo4j vs. AWS Neptune vs. local Kuzu).
- Select the graph store that fits your deployment (managed Aura vs. self-hosted Neo4j vs. AWS Neptune vs. local Kuzu vs. Apache AGE on PostgreSQL).
- Decide when to enable graph writes per request; routine conversations may stay vector-only to save latency.
- Set a policy for pruning stale relationships so your graph stays fast and affordable.
@@ -381,6 +381,42 @@ config = {
Kuzu will clear its state when using `:memory:` once the process exits. See the [Kuzu documentation](https://kuzudb.com/docs/) for advanced settings.
</Accordion>
<Accordion title="Apache AGE (PostgreSQL extension)">
[Apache AGE](https://age.apache.org/) adds graph database capabilities to PostgreSQL, letting you run Cypher queries alongside SQL on the same server. Start AGE via Docker, then configure Mem0:
```bash
docker run --name age-postgres \
-e POSTGRES_DB=mem0_db \
-e POSTGRES_USER=mem0_user \
-e POSTGRES_PASSWORD=mem0_pass \
-p 5432:5432 \
-d apache/age
```
```python
from mem0 import Memory
config = {
"graph_store": {
"provider": "apache_age",
"config": {
"host": "localhost",
"port": 5432,
"database": "mem0_db",
"username": "mem0_user",
"password": "mem0_pass",
"graph_name": "mem0_graph",
},
},
}
m = Memory.from_config(config_dict=config)
```
Apache AGE does not have a built-in vector index, so similarity search is computed client-side. This works well for moderate graph sizes; for very large graphs consider pairing AGE with pgvector for the vector store.
Reference: [Apache AGE documentation](https://age.apache.org/age-manual/master/index.html).
</Accordion>
</AccordionGroup>
<CardGroup cols={2}>
+38 -2
View File
@@ -14,7 +14,7 @@ The Mem0 REST API server exposes every OSS memory operation over HTTP. Run it al
</Info>
<Warning>
Add your own authentication and HTTPS before exposing the server to anything beyond your internal network. The default image does not include auth.
Enable API key authentication (see below) and HTTPS before exposing the server to anything beyond your internal network.
</Warning>
---
@@ -22,6 +22,7 @@ The Mem0 REST API server exposes every OSS memory operation over HTTP. Run it al
## Feature
- **CRUD endpoints:** Create, retrieve, search, update, delete, and reset memories by `user_id`, `agent_id`, or `run_id`.
- **API key authentication:** Optionally secure all endpoints with a shared API key via the `X-API-Key` header.
- **Status health check:** Access base routes to confirm the server is online.
- **OpenAPI explorer:** Visit `/docs` for interactive testing and schema reference.
@@ -91,6 +92,41 @@ uvicorn main:app --reload
---
## Authentication
The server supports optional API key authentication. When the `ADMIN_API_KEY` environment variable is set, every endpoint requires a valid `X-API-Key` header. The `/` redirect, `/docs`, and `/openapi.json` routes remain open so you can always reach the interactive API explorer.
| `ADMIN_API_KEY` value | Behavior |
|---|---|
| Not set / empty | All endpoints are open (no auth) |
| Any non-empty string | Requests must include `X-API-Key: <your-key>` |
### Enable authentication
Add the key to your `.env` file:
```bash
ADMIN_API_KEY=your-secret-api-key
```
Then include the header in every request:
```bash
curl -X POST http://localhost:8000/memories \
-H "Content-Type: application/json" \
-H "X-API-Key: your-secret-api-key" \
-d '{
"messages": [{"role": "user", "content": "I love pizza."}],
"user_id": "alice"
}'
```
<Warning>
The server logs a warning at startup when `ADMIN_API_KEY` is not set. Always set it in production.
</Warning>
---
## See it in action
### Create and search memories via HTTP
@@ -137,7 +173,7 @@ curl "http://localhost:8000/memories/search?user_id=alice&query=vegetable"
## Best practices
1. **Add authentication:** Protect endpoints with API gateways, proxies, or custom FastAPI middleware.
1. **Enable authentication:** Set `ADMIN_API_KEY` to secure all endpoints, or use an API gateway for more advanced schemes.
2. **Use HTTPS:** Terminate TLS at your load balancer or reverse proxy.
3. **Monitor uptime:** Track request rates, latency, and error codes per endpoint.
4. **Version configs:** Keep environment files and Docker Compose definitions in source control.
+2 -2
View File
@@ -166,13 +166,13 @@ Call out the most common mistake or edge case for this layer.
title="[Related cookbook / deep dive]"
description="[Why this pairs well with the current guide]"
icon="arrow-right"
href="/[related-link]"
href="#"
/>
<Card
title="[Next cookbook in journey]"
description="[Set expectation for the next step]"
icon="rocket"
href="/[next-link]"
href="#"
/>
</CardGroup>
```
+2 -2
View File
@@ -145,13 +145,13 @@ npm install mem0ai@[version]
title="[Deep dive reference]"
description="[Why this reference matters post-migration]"
icon="book"
href="/[reference-link]"
href="#"
/>
<Card
title="[Applied example or next step]"
description="[What readers can build now]"
icon="rocket"
href="/[example-link]"
href="#"
/>
</CardGroup>
```
+1
View File
@@ -26,4 +26,5 @@ export * from "./vector_stores/supabase";
export * from "./vector_stores/langchain";
export * from "./vector_stores/vectorize";
export * from "./vector_stores/azure_ai_search";
export * from "./vector_stores/pgvector";
export * from "./utils/factory";
+3
View File
@@ -34,6 +34,7 @@ import { LangchainLLM } from "../llms/langchain";
import { LangchainEmbedder } from "../embeddings/langchain";
import { LangchainVectorStore } from "../vector_stores/langchain";
import { AzureAISearch } from "../vector_stores/azure_ai_search";
import { PGVector } from "../vector_stores/pgvector";
export class EmbedderFactory {
static create(provider: string, config: EmbeddingConfig): Embedder {
@@ -104,6 +105,8 @@ export class VectorStoreFactory {
return new VectorizeDB(config as any);
case "azure-ai-search":
return new AzureAISearch(config as any);
case "pgvector":
return new PGVector(config as any);
default:
throw new Error(`Unsupported vector store provider: ${provider}`);
}
+1 -1
View File
@@ -522,7 +522,7 @@ export class RedisDB implements VectorStore {
return {
id: vectorId,
payload,
payload: toCamelCase(payload),
};
} catch (error) {
console.error("Error getting vector:", error);
+56
View File
@@ -0,0 +1,56 @@
from typing import Optional
from mem0.configs.llms.base import BaseLlmConfig
class MinimaxConfig(BaseLlmConfig):
"""
Configuration class for MiniMax-specific parameters.
Inherits from BaseLlmConfig and adds MiniMax-specific settings.
"""
def __init__(
self,
# Base parameters
model: Optional[str] = None,
temperature: float = 0.1,
api_key: Optional[str] = None,
max_tokens: int = 2000,
top_p: float = 0.1,
top_k: int = 1,
enable_vision: bool = False,
vision_details: Optional[str] = "auto",
http_client_proxies: Optional[dict] = None,
# MiniMax-specific parameters
minimax_base_url: Optional[str] = None,
):
"""
Initialize MiniMax configuration.
Args:
model: MiniMax model to use, defaults to None
temperature: Controls randomness, defaults to 0.1
api_key: MiniMax API key, defaults to None
max_tokens: Maximum tokens to generate, defaults to 2000
top_p: Nucleus sampling parameter, defaults to 0.1
top_k: Top-k sampling parameter, defaults to 1
enable_vision: Enable vision capabilities, defaults to False
vision_details: Vision detail level, defaults to "auto"
http_client_proxies: HTTP client proxy settings, defaults to None
minimax_base_url: MiniMax API base URL, defaults to None
"""
# Initialize base parameters
super().__init__(
model=model,
temperature=temperature,
api_key=api_key,
max_tokens=max_tokens,
top_p=top_p,
top_k=top_k,
enable_vision=enable_vision,
vision_details=vision_details,
http_client_proxies=http_client_proxies,
)
# MiniMax-specific parameters
self.minimax_base_url = minimax_base_url
+24 -2
View File
@@ -77,12 +77,32 @@ class KuzuConfig(BaseModel):
db: Optional[str] = Field(":memory:", description="Path to a Kuzu database file")
class ApacheAgeConfig(BaseModel):
host: Optional[str] = Field("localhost", description="PostgreSQL server hostname")
port: Optional[int] = Field(5432, description="PostgreSQL server port")
database: Optional[str] = Field(None, description="PostgreSQL database name")
username: Optional[str] = Field(None, description="PostgreSQL username")
password: Optional[str] = Field(None, description="PostgreSQL password")
graph_name: Optional[str] = Field("mem0_graph", description="Name of the Apache AGE graph")
@model_validator(mode="before")
def check_required_fields(cls, values):
database, username, password = (
values.get("database"),
values.get("username"),
values.get("password"),
)
if not database or not username or not password:
raise ValueError("Please provide 'database', 'username' and 'password'.")
return values
class GraphStoreConfig(BaseModel):
provider: str = Field(
description="Provider of the data store (e.g., 'neo4j', 'memgraph', 'neptune', 'kuzu')",
description="Provider of the data store (e.g., 'neo4j', 'memgraph', 'neptune', 'kuzu', 'apache_age')",
default="neo4j",
)
config: Union[Neo4jConfig, MemgraphConfig, NeptuneConfig, KuzuConfig] = Field(
config: Union[Neo4jConfig, MemgraphConfig, NeptuneConfig, KuzuConfig, ApacheAgeConfig] = Field(
description="Configuration for the specific data store", default=None
)
llm: Optional[LlmConfig] = Field(description="LLM configuration for querying the graph store", default=None)
@@ -110,5 +130,7 @@ class GraphStoreConfig(BaseModel):
return NeptuneConfig(**v.model_dump())
elif provider == "kuzu":
return KuzuConfig(**v.model_dump())
elif provider == "apache_age":
return ApacheAgeConfig(**v.model_dump())
else:
raise ValueError(f"Unsupported graph store provider: {provider}")
+32 -4
View File
@@ -1,3 +1,4 @@
import json
import os
from typing import Dict, List, Optional
@@ -6,6 +7,7 @@ from openai import AzureOpenAI
from mem0.configs.llms.base import BaseLlmConfig
from mem0.llms.base import LLMBase
from mem0.memory.utils import extract_json
SCOPE = "https://cognitiveservices.azure.com/.default"
@@ -83,9 +85,35 @@ class AzureOpenAIStructuredLLM(LLMBase):
params["tools"] = tools
params["tool_choice"] = tool_choice
if tools:
params["tools"] = tools
params["tool_choice"] = tool_choice
response = self.client.chat.completions.create(**params)
return self._parse_response(response, tools)
def _parse_response(self, response, tools):
"""
Process the response based on whether tools are used or not.
Args:
response: The raw response from API.
tools: The list of tools provided in the request.
Returns:
str or dict: The processed response.
"""
if tools:
processed_response = {
"content": response.choices[0].message.content,
"tool_calls": [],
}
if response.choices[0].message.tool_calls:
for tool_call in response.choices[0].message.tool_calls:
processed_response["tool_calls"].append(
{
"name": tool_call.function.name,
"arguments": json.loads(extract_json(tool_call.function.arguments)),
}
)
return processed_response
else:
return response.choices[0].message.content
+1
View File
@@ -23,6 +23,7 @@ class LlmConfig(BaseModel):
"azure_openai_structured",
"gemini",
"deepseek",
"minimax",
"xai",
"sarvam",
"lmstudio",
+114
View File
@@ -0,0 +1,114 @@
import json
import os
from typing import Dict, List, Optional, Union
from openai import OpenAI
from mem0.configs.llms.base import BaseLlmConfig
from mem0.configs.llms.minimax import MinimaxConfig
from mem0.llms.base import LLMBase
from mem0.memory.utils import extract_json
class MiniMaxLLM(LLMBase):
def __init__(self, config: Optional[Union[BaseLlmConfig, MinimaxConfig, Dict]] = None):
# Convert to MinimaxConfig if needed
if config is None:
config = MinimaxConfig()
elif isinstance(config, dict):
config = MinimaxConfig(**config)
elif isinstance(config, BaseLlmConfig) and not isinstance(config, MinimaxConfig):
# Convert BaseLlmConfig to MinimaxConfig
config = MinimaxConfig(
model=config.model,
temperature=config.temperature,
api_key=config.api_key,
max_tokens=config.max_tokens,
top_p=config.top_p,
top_k=config.top_k,
enable_vision=config.enable_vision,
vision_details=config.vision_details,
http_client_proxies=config.http_client,
)
super().__init__(config)
if not self.config.model:
self.config.model = "MiniMax-M2.7"
api_key = self.config.api_key or os.getenv("MINIMAX_API_KEY")
base_url = (
self.config.minimax_base_url
or os.getenv("MINIMAX_API_BASE")
or "https://api.minimax.io/v1"
)
self.client = OpenAI(api_key=api_key, base_url=base_url)
def _parse_response(self, response, tools):
"""
Process the response based on whether tools are used or not.
Args:
response: The raw response from API.
tools: The list of tools provided in the request.
Returns:
str or dict: The processed response.
"""
if tools:
processed_response = {
"content": response.choices[0].message.content,
"tool_calls": [],
}
if response.choices[0].message.tool_calls:
for tool_call in response.choices[0].message.tool_calls:
processed_response["tool_calls"].append(
{
"name": tool_call.function.name,
"arguments": json.loads(extract_json(tool_call.function.arguments)),
}
)
return processed_response
else:
return response.choices[0].message.content
def generate_response(
self,
messages: List[Dict[str, str]],
response_format=None,
tools: Optional[List[Dict]] = None,
tool_choice: str = "auto",
**kwargs,
):
"""
Generate a response based on the given messages using MiniMax.
Args:
messages (list): List of message dicts containing 'role' and 'content'.
response_format (str or object, optional): Format of the response. Defaults to None.
tools (list, optional): List of tools that the model can call. Defaults to None.
tool_choice (str, optional): Tool choice method. Defaults to "auto".
**kwargs: Additional MiniMax-specific parameters.
Returns:
str: The generated response.
"""
params = self._get_supported_params(messages=messages, **kwargs)
params.update(
{
"model": self.config.model,
"messages": messages,
}
)
if response_format:
params["response_format"] = response_format
if tools:
params["tools"] = tools
params["tool_choice"] = tool_choice
response = self.client.chat.completions.create(**params)
return self._parse_response(response, tools)
+573
View File
@@ -0,0 +1,573 @@
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()
+16
View File
@@ -644,6 +644,9 @@ class Memory(MemoryBase):
if memory_id and (metadata.get("agent_id") or metadata.get("run_id")):
# Update only the session identifiers, keep content the same
existing_memory = self.vector_store.get(vector_id=memory_id)
if existing_memory is None:
logger.warning(f"Memory {memory_id} not found for session ID update, skipping")
continue
updated_metadata = deepcopy(existing_memory.payload)
if metadata.get("agent_id"):
updated_metadata["agent_id"] = metadata["agent_id"]
@@ -1226,6 +1229,9 @@ class Memory(MemoryBase):
logger.error(f"Error getting memory with ID {memory_id} during update.")
raise ValueError(f"Error getting memory with ID {memory_id}. Please provide a valid 'memory_id'")
if existing_memory is None:
raise ValueError(f"Memory with id {memory_id} not found. Please provide a valid 'memory_id'")
prev_value = existing_memory.payload.get("data")
new_metadata = deepcopy(metadata) if metadata is not None else {}
@@ -1274,6 +1280,8 @@ class Memory(MemoryBase):
def _delete_memory(self, memory_id):
logger.info(f"Deleting memory with {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", "")
self.vector_store.delete(vector_id=memory_id)
self.db.add_history(
@@ -1671,6 +1679,9 @@ class AsyncMemory(MemoryBase):
# Create async task to update only the session identifiers
async def update_session_ids(mem_id, meta):
existing_memory = await asyncio.to_thread(self.vector_store.get, vector_id=mem_id)
if existing_memory is None:
logger.warning(f"Memory {mem_id} not found for session ID update, skipping")
return
updated_metadata = deepcopy(existing_memory.payload)
if meta.get("agent_id"):
updated_metadata["agent_id"] = meta["agent_id"]
@@ -2313,6 +2324,9 @@ class AsyncMemory(MemoryBase):
logger.error(f"Error getting memory with ID {memory_id} during update.")
raise ValueError(f"Error getting memory with ID {memory_id}. Please provide a valid 'memory_id'")
if existing_memory is None:
raise ValueError(f"Memory with id {memory_id} not found. Please provide a valid 'memory_id'")
prev_value = existing_memory.payload.get("data")
new_metadata = deepcopy(metadata) if metadata is not None else {}
@@ -2364,6 +2378,8 @@ class AsyncMemory(MemoryBase):
async def _delete_memory(self, memory_id):
logger.info(f"Deleting memory with {memory_id=}")
existing_memory = await asyncio.to_thread(self.vector_store.get, vector_id=memory_id)
if existing_memory is None:
raise ValueError(f"Memory with id {memory_id} not found")
prev_value = existing_memory.payload.get("data", "")
await asyncio.to_thread(self.vector_store.delete, vector_id=memory_id)
+3
View File
@@ -6,6 +6,7 @@ from mem0.configs.llms.anthropic import AnthropicConfig
from mem0.configs.llms.azure import AzureOpenAIConfig
from mem0.configs.llms.base import BaseLlmConfig
from mem0.configs.llms.deepseek import DeepSeekConfig
from mem0.configs.llms.minimax import MinimaxConfig
from mem0.configs.llms.lmstudio import LMStudioConfig
from mem0.configs.llms.ollama import OllamaConfig
from mem0.configs.llms.openai import OpenAIConfig
@@ -45,6 +46,7 @@ class LlmFactory:
"azure_openai_structured": ("mem0.llms.azure_openai_structured.AzureOpenAIStructuredLLM", AzureOpenAIConfig),
"gemini": ("mem0.llms.gemini.GeminiLLM", BaseLlmConfig),
"deepseek": ("mem0.llms.deepseek.DeepSeekLLM", DeepSeekConfig),
"minimax": ("mem0.llms.minimax.MiniMaxLLM", MinimaxConfig),
"xai": ("mem0.llms.xai.XAILLM", BaseLlmConfig),
"sarvam": ("mem0.llms.sarvam.SarvamLLM", BaseLlmConfig),
"lmstudio": ("mem0.llms.lmstudio.LMStudioLLM", LMStudioConfig),
@@ -216,6 +218,7 @@ class GraphStoreFactory:
"neptune": "mem0.graphs.neptune.neptunegraph.MemoryGraph",
"neptunedb": "mem0.graphs.neptune.neptunedb.MemoryGraph",
"kuzu": "mem0.memory.kuzu_memory.MemoryGraph",
"apache_age": "mem0.memory.apache_age_memory.MemoryGraph",
"default": "mem0.memory.graph_memory.MemoryGraph",
}
+1
View File
@@ -31,6 +31,7 @@ graph = [
"neo4j>=5.23.1",
"rank-bm25>=0.2.2",
"kuzu>=0.11.0",
"apache-age-python>=0.0.6",
]
vector_stores = [
"vecs>=0.4.0",
+3
View File
@@ -10,3 +10,6 @@ POSTGRES_DB=
POSTGRES_USER=
POSTGRES_PASSWORD=
POSTGRES_COLLECTION_NAME=
# Optional: set to enable API key authentication on all endpoints
ADMIN_API_KEY=
+55 -10
View File
@@ -1,10 +1,12 @@
import logging
import os
import secrets
from typing import Any, Dict, List, Optional
from dotenv import load_dotenv
from fastapi import FastAPI, HTTPException
from fastapi import Depends, FastAPI, HTTPException
from fastapi.responses import JSONResponse, RedirectResponse
from fastapi.security import APIKeyHeader
from pydantic import BaseModel, Field
from mem0 import Memory
@@ -14,6 +16,22 @@ logging.basicConfig(level=logging.INFO, format="%(asctime)s - %(levelname)s - %(
# Load environment variables
load_dotenv()
ADMIN_API_KEY = os.environ.get("ADMIN_API_KEY", "")
MIN_KEY_LENGTH = 16
if not ADMIN_API_KEY:
logging.warning(
"ADMIN_API_KEY not set - API endpoints are UNSECURED! "
"Set ADMIN_API_KEY environment variable for production use."
)
else:
if len(ADMIN_API_KEY) < MIN_KEY_LENGTH:
logging.warning(
"ADMIN_API_KEY is shorter than %d characters - consider using a longer key for production.",
MIN_KEY_LENGTH,
)
logging.info("API key authentication enabled")
POSTGRES_HOST = os.environ.get("POSTGRES_HOST", "postgres")
POSTGRES_PORT = os.environ.get("POSTGRES_PORT", "5432")
@@ -60,10 +78,35 @@ MEMORY_INSTANCE = Memory.from_config(DEFAULT_CONFIG)
app = FastAPI(
title="Mem0 REST APIs",
description="A REST API for managing and searching memories for your AI Agents and Apps.",
description=(
"A REST API for managing and searching memories for your AI Agents and Apps.\n\n"
"## Authentication\n"
"When the ADMIN_API_KEY environment variable is set, all endpoints require "
"the `X-API-Key` header for authentication."
),
version="1.0.0",
)
api_key_header = APIKeyHeader(name="X-API-Key", auto_error=False)
async def verify_api_key(api_key: Optional[str] = Depends(api_key_header)):
"""Validate the API key when ADMIN_API_KEY is configured. No-op otherwise."""
if ADMIN_API_KEY:
if api_key is None:
raise HTTPException(
status_code=401,
detail="X-API-Key header is required.",
headers={"WWW-Authenticate": "ApiKey"},
)
if not secrets.compare_digest(api_key, ADMIN_API_KEY):
raise HTTPException(
status_code=401,
detail="Invalid API key.",
headers={"WWW-Authenticate": "ApiKey"},
)
return api_key
class Message(BaseModel):
role: str = Field(..., description="Role of the message (user or assistant).")
@@ -87,7 +130,7 @@ class SearchRequest(BaseModel):
@app.post("/configure", summary="Configure Mem0")
def set_config(config: Dict[str, Any]):
def set_config(config: Dict[str, Any], _api_key: Optional[str] = Depends(verify_api_key)):
"""Set memory configuration."""
global MEMORY_INSTANCE
MEMORY_INSTANCE = Memory.from_config(config)
@@ -95,7 +138,7 @@ def set_config(config: Dict[str, Any]):
@app.post("/memories", summary="Create memories")
def add_memory(memory_create: MemoryCreate):
def add_memory(memory_create: MemoryCreate, _api_key: Optional[str] = Depends(verify_api_key)):
"""Store new memories."""
if not any([memory_create.user_id, memory_create.agent_id, memory_create.run_id]):
raise HTTPException(status_code=400, detail="At least one identifier (user_id, agent_id, run_id) is required.")
@@ -114,6 +157,7 @@ def get_all_memories(
user_id: Optional[str] = None,
run_id: Optional[str] = None,
agent_id: Optional[str] = None,
_api_key: Optional[str] = Depends(verify_api_key),
):
"""Retrieve stored memories."""
if not any([user_id, run_id, agent_id]):
@@ -129,7 +173,7 @@ def get_all_memories(
@app.get("/memories/{memory_id}", summary="Get a memory")
def get_memory(memory_id: str):
def get_memory(memory_id: str, _api_key: Optional[str] = Depends(verify_api_key)):
"""Retrieve a specific memory by ID."""
try:
return MEMORY_INSTANCE.get(memory_id)
@@ -139,7 +183,7 @@ def get_memory(memory_id: str):
@app.post("/search", summary="Search memories")
def search_memories(search_req: SearchRequest):
def search_memories(search_req: SearchRequest, _api_key: Optional[str] = Depends(verify_api_key)):
"""Search for memories based on a query."""
try:
params = {k: v for k, v in search_req.model_dump().items() if v is not None and k != "query"}
@@ -150,7 +194,7 @@ def search_memories(search_req: SearchRequest):
@app.put("/memories/{memory_id}", summary="Update a memory")
def update_memory(memory_id: str, updated_memory: Dict[str, Any]):
def update_memory(memory_id: str, updated_memory: Dict[str, Any], _api_key: Optional[str] = Depends(verify_api_key)):
"""Update an existing memory with new content.
Args:
@@ -168,7 +212,7 @@ def update_memory(memory_id: str, updated_memory: Dict[str, Any]):
@app.get("/memories/{memory_id}/history", summary="Get memory history")
def memory_history(memory_id: str):
def memory_history(memory_id: str, _api_key: Optional[str] = Depends(verify_api_key)):
"""Retrieve memory history."""
try:
return MEMORY_INSTANCE.history(memory_id=memory_id)
@@ -178,7 +222,7 @@ def memory_history(memory_id: str):
@app.delete("/memories/{memory_id}", summary="Delete a memory")
def delete_memory(memory_id: str):
def delete_memory(memory_id: str, _api_key: Optional[str] = Depends(verify_api_key)):
"""Delete a specific memory by ID."""
try:
MEMORY_INSTANCE.delete(memory_id=memory_id)
@@ -193,6 +237,7 @@ def delete_all_memories(
user_id: Optional[str] = None,
run_id: Optional[str] = None,
agent_id: Optional[str] = None,
_api_key: Optional[str] = Depends(verify_api_key),
):
"""Delete all memories for a given identifier."""
if not any([user_id, run_id, agent_id]):
@@ -209,7 +254,7 @@ def delete_all_memories(
@app.post("/reset", summary="Reset all memories")
def reset_memory():
def reset_memory(_api_key: Optional[str] = Depends(verify_api_key)):
"""Completely reset stored memories."""
try:
MEMORY_INSTANCE.reset()
@@ -1,4 +1,5 @@
from unittest import mock
from unittest.mock import Mock
from mem0.llms.azure_openai_structured import SCOPE, AzureOpenAIStructuredLLM
@@ -98,3 +99,75 @@ def test_init_with_placeholder_api_key_uses_default_credential(
args, kwargs = mock_azure_openai.call_args
assert kwargs["api_key"] is None
assert kwargs["azure_ad_token_provider"] == "token-provider"
@mock.patch("mem0.llms.azure_openai_structured.AzureOpenAI")
def test_generate_response_without_tools(mock_azure_openai):
mock_client = Mock()
mock_azure_openai.return_value = mock_client
config = DummyConfig(model="test-model", azure_kwargs=DummyAzureKwargs(api_key="real-key"))
llm = AzureOpenAIStructuredLLM(config)
mock_response = Mock()
mock_response.choices = [Mock(message=Mock(content="Hello there!"))]
mock_client.chat.completions.create.return_value = mock_response
messages = [{"role": "user", "content": "Hi"}]
response = llm.generate_response(messages)
assert response == "Hello there!"
@mock.patch("mem0.llms.azure_openai_structured.AzureOpenAI")
def test_generate_response_with_tools(mock_azure_openai):
mock_client = Mock()
mock_azure_openai.return_value = mock_client
config = DummyConfig(model="test-model", azure_kwargs=DummyAzureKwargs(api_key="real-key"))
llm = AzureOpenAIStructuredLLM(config)
mock_tool_call = Mock()
mock_tool_call.function.name = "add_memory"
mock_tool_call.function.arguments = '{"data": "sunny day"}'
mock_message = Mock()
mock_message.content = "I've added the memory."
mock_message.tool_calls = [mock_tool_call]
mock_response = Mock()
mock_response.choices = [Mock(message=mock_message)]
mock_client.chat.completions.create.return_value = mock_response
tools = [{"type": "function", "function": {"name": "add_memory"}}]
messages = [{"role": "user", "content": "Remember sunny day"}]
response = llm.generate_response(messages, tools=tools)
assert response["content"] == "I've added the memory."
assert len(response["tool_calls"]) == 1
assert response["tool_calls"][0]["name"] == "add_memory"
assert response["tool_calls"][0]["arguments"] == {"data": "sunny day"}
@mock.patch("mem0.llms.azure_openai_structured.AzureOpenAI")
def test_generate_response_with_tools_no_tool_calls(mock_azure_openai):
mock_client = Mock()
mock_azure_openai.return_value = mock_client
config = DummyConfig(model="test-model", azure_kwargs=DummyAzureKwargs(api_key="real-key"))
llm = AzureOpenAIStructuredLLM(config)
mock_message = Mock()
mock_message.content = "No tools needed."
mock_message.tool_calls = None
mock_response = Mock()
mock_response.choices = [Mock(message=mock_message)]
mock_client.chat.completions.create.return_value = mock_response
tools = [{"type": "function", "function": {"name": "add_memory"}}]
messages = [{"role": "user", "content": "Hello"}]
response = llm.generate_response(messages, tools=tools)
assert response["content"] == "No tools needed."
assert response["tool_calls"] == []
+194
View File
@@ -0,0 +1,194 @@
import os
from unittest.mock import Mock, patch
import pytest
from mem0.configs.llms.base import BaseLlmConfig
from mem0.configs.llms.minimax import MinimaxConfig
from mem0.llms.minimax import MiniMaxLLM
from mem0.utils.factory import LlmFactory
@pytest.fixture
def mock_minimax_client():
with patch("mem0.llms.minimax.OpenAI") as mock_openai:
mock_client = Mock()
mock_openai.return_value = mock_client
yield mock_client
def test_minimax_llm_default_base_url():
"""Default config uses MiniMax official base URL."""
config = BaseLlmConfig(
model="MiniMax-M2.7", temperature=0.7, max_tokens=100, top_p=1.0, api_key="api_key"
)
llm = MiniMaxLLM(config)
# OpenAI client may normalize URL with trailing slash
assert str(llm.client.base_url).rstrip("/") == "https://api.minimax.io/v1"
def test_minimax_llm_env_base_url():
"""Config uses MINIMAX_API_BASE env variable when set."""
provider_base_url = "https://api.provider.com/v1/"
os.environ["MINIMAX_API_BASE"] = provider_base_url
try:
config = MinimaxConfig(
model="MiniMax-M2.7",
temperature=0.7,
max_tokens=100,
top_p=1.0,
api_key="api_key",
)
llm = MiniMaxLLM(config)
assert str(llm.client.base_url).rstrip("/") == provider_base_url.rstrip("/")
finally:
os.environ.pop("MINIMAX_API_BASE", None)
def test_minimax_llm_config_base_url():
"""Config uses minimax_base_url when provided."""
config_base_url = "https://api.config.com/v1/"
config = MinimaxConfig(
model="MiniMax-M2.7",
temperature=0.7,
max_tokens=100,
top_p=1.0,
api_key="api_key",
minimax_base_url=config_base_url,
)
llm = MiniMaxLLM(config)
assert str(llm.client.base_url).rstrip("/") == config_base_url.rstrip("/")
def test_minimax_llm_default_model(mock_minimax_client):
"""Default model is MiniMax-M2.7 when not specified."""
config = MinimaxConfig(temperature=0.7, max_tokens=100, api_key="api_key")
llm = MiniMaxLLM(config)
assert llm.config.model == "MiniMax-M2.7"
def test_minimax_llm_env_api_key():
"""Uses MINIMAX_API_KEY env when api_key not in config."""
os.environ["MINIMAX_API_KEY"] = "env-api-key"
try:
with patch("mem0.llms.minimax.OpenAI") as mock_openai:
mock_client = Mock()
mock_openai.return_value = mock_client
config = MinimaxConfig(model="MiniMax-M2.7", api_key=None)
MiniMaxLLM(config)
mock_openai.assert_called_once_with(
api_key="env-api-key",
base_url="https://api.minimax.io/v1",
)
finally:
os.environ.pop("MINIMAX_API_KEY", None)
def test_generate_response_without_tools(mock_minimax_client):
"""generate_response returns text when no tools provided."""
config = BaseLlmConfig(
model="MiniMax-M2.7", temperature=0.7, max_tokens=100, top_p=1.0, api_key="api_key"
)
llm = MiniMaxLLM(config)
messages = [
{"role": "system", "content": "You are a helpful assistant."},
{"role": "user", "content": "Hello, how are you?"},
]
mock_response = Mock()
mock_response.choices = [Mock(message=Mock(content="I'm doing well, thank you for asking!"))]
mock_minimax_client.chat.completions.create.return_value = mock_response
response = llm.generate_response(messages)
mock_minimax_client.chat.completions.create.assert_called_once_with(
model="MiniMax-M2.7", messages=messages, temperature=0.7, max_tokens=100, top_p=1.0
)
assert response == "I'm doing well, thank you for asking!"
def test_generate_response_with_tools(mock_minimax_client):
"""generate_response returns tool_calls when tools provided."""
config = BaseLlmConfig(
model="MiniMax-M2.7", temperature=0.7, max_tokens=100, top_p=1.0, api_key="api_key"
)
llm = MiniMaxLLM(config)
messages = [
{"role": "system", "content": "You are a helpful assistant."},
{"role": "user", "content": "Add a new memory: Today is a sunny day."},
]
tools = [
{
"type": "function",
"function": {
"name": "add_memory",
"description": "Add a memory",
"parameters": {
"type": "object",
"properties": {"data": {"type": "string", "description": "Data to add to memory"}},
"required": ["data"],
},
},
}
]
mock_response = Mock()
mock_message = Mock()
mock_message.content = "I've added the memory for you."
mock_tool_call = Mock()
mock_tool_call.function.name = "add_memory"
mock_tool_call.function.arguments = '{"data": "Today is a sunny day."}'
mock_message.tool_calls = [mock_tool_call]
mock_response.choices = [Mock(message=mock_message)]
mock_minimax_client.chat.completions.create.return_value = mock_response
response = llm.generate_response(messages, tools=tools)
mock_minimax_client.chat.completions.create.assert_called_once_with(
model="MiniMax-M2.7",
messages=messages,
temperature=0.7,
max_tokens=100,
top_p=1.0,
tools=tools,
tool_choice="auto",
)
assert response["content"] == "I've added the memory for you."
assert len(response["tool_calls"]) == 1
assert response["tool_calls"][0]["name"] == "add_memory"
assert response["tool_calls"][0]["arguments"] == {"data": "Today is a sunny day."}
def test_generate_response_with_response_format(mock_minimax_client):
"""generate_response passes response_format to the API."""
config = BaseLlmConfig(
model="MiniMax-M2.7", temperature=0.7, max_tokens=100, top_p=1.0, api_key="api_key"
)
llm = MiniMaxLLM(config)
messages = [{"role": "user", "content": "Return JSON."}]
response_format = {"type": "json_object"}
mock_response = Mock()
mock_response.choices = [Mock(message=Mock(content='{"key": "value"}'))]
mock_minimax_client.chat.completions.create.return_value = mock_response
llm.generate_response(messages, response_format=response_format)
mock_minimax_client.chat.completions.create.assert_called_once_with(
model="MiniMax-M2.7",
messages=messages,
temperature=0.7,
max_tokens=100,
top_p=1.0,
response_format={"type": "json_object"},
)
def test_factory_creates_minimax_llm(mock_minimax_client):
"""LlmFactory.create returns MiniMaxLLM for provider 'minimax'."""
llm = LlmFactory.create("minimax", {"model": "MiniMax-M2.7", "api_key": "test-key"})
assert isinstance(llm, MiniMaxLLM)
assert llm.config.model == "MiniMax-M2.7"
File diff suppressed because it is too large Load Diff
+226
View File
@@ -0,0 +1,226 @@
from unittest.mock import MagicMock, Mock, patch
# age and rank_bm25 are optional deps — mock them so tests run without install
_age_mock = Mock()
patch.dict("sys.modules", {
"age": _age_mock,
"age.models": Mock(),
"rank_bm25": Mock(),
}).start()
from mem0.memory.apache_age_memory import MemoryGraph, _cosine_similarity # noqa: E402
def _make_instance():
with patch.object(MemoryGraph, "__init__", return_value=None):
instance = MemoryGraph.__new__(MemoryGraph)
instance.llm_provider = "openai"
instance.llm = MagicMock()
instance.embedding_model = MagicMock()
instance.config = MagicMock()
instance.config.graph_store.custom_prompt = None
instance.ag = MagicMock()
instance.graph_name = "test_graph"
instance.threshold = 0.7
return instance
class TestCosineSimilarity:
"""Tests for the _cosine_similarity helper."""
def test_identical_vectors(self):
assert abs(_cosine_similarity([1, 0, 0], [1, 0, 0]) - 1.0) < 1e-6
def test_orthogonal_vectors(self):
assert abs(_cosine_similarity([1, 0, 0], [0, 1, 0])) < 1e-6
def test_zero_vector(self):
assert _cosine_similarity([0, 0, 0], [1, 2, 3]) == 0.0
class TestRetrieveNodesFromData:
"""Tests for _retrieve_nodes_from_data in Apache AGE MemoryGraph."""
def test_normal_entities_extracted(self):
instance = _make_instance()
instance.llm.generate_response.return_value = {
"tool_calls": [{"name": "extract_entities", "arguments": {"entities": [
{"entity": "Alice", "entity_type": "person"},
{"entity": "hiking", "entity_type": "activity"},
]}}]
}
result = instance._retrieve_nodes_from_data("Alice loves hiking", {"user_id": "u1"})
assert result == {"alice": "person", "hiking": "activity"}
def test_malformed_entity_missing_entity_type_is_skipped(self):
instance = _make_instance()
instance.llm.generate_response.return_value = {
"tool_calls": [{"name": "extract_entities", "arguments": {"entities": [
{"entity": "matrix multiplication", "entity_type": "task"},
{"entity": "task"},
{"entity": "ReLU", "entity_type": "task"},
]}}]
}
result = instance._retrieve_nodes_from_data("some text", {"user_id": "u1"})
assert "matrix_multiplication" in result
assert "relu" in result
assert "task" not in result
def test_missing_entities_key_returns_empty(self):
instance = _make_instance()
instance.llm.generate_response.return_value = {
"tool_calls": [{"name": "extract_entities", "arguments": {"text": "Hello."}}]
}
result = instance._retrieve_nodes_from_data("Hello.", {"user_id": "u1"})
assert result == {}
def test_none_tool_calls_returns_empty(self):
instance = _make_instance()
instance.llm.generate_response.return_value = {"tool_calls": None}
result = instance._retrieve_nodes_from_data("hello world", {"user_id": "u1"})
assert result == {}
class TestEstablishNodesRelationsFromData:
"""Tests for _establish_nodes_relations_from_data in Apache AGE MemoryGraph."""
def test_none_response_does_not_crash(self):
instance = _make_instance()
instance.llm.generate_response.return_value = None
result = instance._establish_nodes_relations_from_data(
"Hello world", {"user_id": "u1"}, {}
)
assert result == []
def test_empty_tool_calls_returns_empty(self):
instance = _make_instance()
instance.llm.generate_response.return_value = {"tool_calls": []}
result = instance._establish_nodes_relations_from_data(
"Hello world", {"user_id": "u1"}, {}
)
assert result == []
def test_valid_entities_returned(self):
instance = _make_instance()
instance.llm.generate_response.return_value = {
"tool_calls": [{"name": "add_entities", "arguments": {"entities": [
{"source": "alice", "relationship": "loves", "destination": "hiking"}
]}}]
}
result = instance._establish_nodes_relations_from_data(
"Alice loves hiking", {"user_id": "u1"}, {"alice": "person"}
)
assert len(result) == 1
assert result[0]["source"] == "alice"
class TestRemoveSpacesFromEntities:
"""Tests for _remove_spaces_from_entities."""
def test_spaces_and_case(self):
instance = _make_instance()
entities = [{"source": "Alice Smith", "relationship": "Works At", "destination": "Big Corp"}]
result = instance._remove_spaces_from_entities(entities)
assert result[0]["source"] == "alice_smith"
assert result[0]["relationship"] == "works_at"
assert result[0]["destination"] == "big_corp"
class TestFindSimilarNode:
"""Tests for _find_similar_node."""
def test_returns_none_when_no_nodes(self):
instance = _make_instance()
instance._exec_cypher = MagicMock(return_value=[])
result = instance._find_similar_node([1.0, 0.0], {"user_id": "u1"}, threshold=0.9)
assert result is None
def test_returns_best_match_above_threshold(self):
instance = _make_instance()
instance._exec_cypher = MagicMock(return_value=[
{"name": "alice", "embedding": [1.0, 0.0], "user_id": "u1"},
{"name": "bob", "embedding": [0.0, 1.0], "user_id": "u1"},
])
result = instance._find_similar_node([1.0, 0.0], {"user_id": "u1"}, threshold=0.9)
assert result["name"] == "alice"
def test_filters_by_agent_id(self):
instance = _make_instance()
instance._exec_cypher = MagicMock(return_value=[
{"name": "alice", "embedding": [1.0, 0.0], "user_id": "u1", "agent_id": "a2"},
])
result = instance._find_similar_node(
[1.0, 0.0], {"user_id": "u1", "agent_id": "a1"}, threshold=0.9
)
assert result is None
class TestDeleteAll:
"""Tests for delete_all."""
def test_calls_exec_cypher_and_commits(self):
instance = _make_instance()
instance._exec_cypher = MagicMock(return_value=[])
instance.delete_all({"user_id": "u1"})
instance._exec_cypher.assert_called_once()
instance.ag.commit.assert_called_once()
class TestGetAll:
"""Tests for get_all."""
def test_returns_formatted_results(self):
instance = _make_instance()
instance._exec_cypher = MagicMock(return_value=[
{"source": "alice", "relationship": "KNOWS", "target": "bob"},
{"source": "alice", "relationship": "LIKES", "target": "hiking"},
])
results = instance.get_all({"user_id": "u1"}, limit=10)
assert len(results) == 2
assert results[0]["source"] == "alice"
assert results[0]["relationship"] == "KNOWS"
assert results[0]["target"] == "bob"
def test_passes_limit_to_cypher(self):
"""Limit is enforced via LIMIT in the Cypher query, not Python slicing."""
instance = _make_instance()
instance._exec_cypher = MagicMock(return_value=[
{"source": "n0", "relationship": "R", "target": "m0"},
])
instance.get_all({"user_id": "u1"}, limit=3)
# Verify limit was passed as a parameter to the query
cypher_stmt = instance._exec_cypher.call_args[0][0]
assert "LIMIT %s" in cypher_stmt
params = instance._exec_cypher.call_args[1].get("params") or instance._exec_cypher.call_args[0][2]
assert 3 in params
class TestAdd:
"""Tests for the add orchestration method."""
def test_add_returns_added_and_deleted(self):
instance = _make_instance()
instance._retrieve_nodes_from_data = MagicMock(return_value={"alice": "person"})
instance._establish_nodes_relations_from_data = MagicMock(return_value=[
{"source": "alice", "relationship": "knows", "destination": "bob"}
])
instance._search_graph_db = MagicMock(return_value=[])
instance._get_delete_entities_from_search_output = MagicMock(return_value=[])
instance._delete_entities = MagicMock(return_value=[])
instance._add_entities = MagicMock(return_value=["added"])
result = instance.add("Alice knows Bob", {"user_id": "u1"})
assert "deleted_entities" in result
assert "added_entities" in result
assert result["added_entities"] == ["added"]
class TestSearch:
"""Tests for the search method."""
def test_returns_empty_when_no_search_output(self):
instance = _make_instance()
instance._retrieve_nodes_from_data = MagicMock(return_value={"alice": "person"})
instance._search_graph_db = MagicMock(return_value=[])
result = instance.search("Who is Alice?", {"user_id": "u1"})
assert result == []
+118
View File
@@ -320,3 +320,121 @@ def test_normalize_facts_mixed():
def test_normalize_facts_filters_empty_strings():
assert normalize_facts(["", "valid", ""]) == ["valid"]
@patch('mem0.utils.factory.EmbedderFactory.create')
@patch('mem0.utils.factory.VectorStoreFactory.create')
@patch('mem0.utils.factory.LlmFactory.create')
@patch('mem0.memory.storage.SQLiteManager')
def test_delete_nonexistent_memory_raises_error(mock_sqlite, mock_llm_factory, mock_vector_factory, mock_embedder_factory):
"""
Test that delete() raises ValueError when memory_id does not exist
and does not attempt to delete from the vector store.
Issue #3849: memory.delete() fails with AttributeError when memory not found.
Should raise a clear ValueError instead.
"""
mock_embedder_factory.return_value = MagicMock()
mock_vector_store = MagicMock()
mock_vector_factory.return_value = mock_vector_store
mock_llm_factory.return_value = MagicMock()
mock_sqlite.return_value = MagicMock()
from mem0.memory.main import Memory as MemoryClass
config = MemoryConfig()
memory = MemoryClass(config)
mock_vector_store.get.return_value = None
with pytest.raises(ValueError, match="Memory with id non-existent-id not found"):
memory.delete("non-existent-id")
mock_vector_store.delete.assert_not_called()
@pytest.mark.asyncio
@patch('mem0.utils.factory.EmbedderFactory.create')
@patch('mem0.utils.factory.VectorStoreFactory.create')
@patch('mem0.utils.factory.LlmFactory.create')
@patch('mem0.memory.storage.SQLiteManager')
async def test_async_delete_nonexistent_memory_raises_error(mock_sqlite, mock_llm_factory, mock_vector_factory, mock_embedder_factory):
"""
Test that async delete() raises ValueError when memory_id does not exist
and does not attempt to delete from the vector store.
Issue #3849: memory.delete() fails with AttributeError when memory not found.
Should raise a clear ValueError instead.
"""
mock_embedder_factory.return_value = MagicMock()
mock_vector_store = MagicMock()
mock_vector_store.get.return_value = None
mock_vector_factory.return_value = mock_vector_store
mock_llm_factory.return_value = MagicMock()
mock_sqlite.return_value = MagicMock()
from mem0.memory.main import AsyncMemory
config = MemoryConfig()
memory = AsyncMemory(config)
with pytest.raises(ValueError, match="Memory with id non-existent-id not found"):
await memory.delete("non-existent-id")
mock_vector_store.delete.assert_not_called()
@patch('mem0.utils.factory.EmbedderFactory.create')
@patch('mem0.utils.factory.VectorStoreFactory.create')
@patch('mem0.utils.factory.LlmFactory.create')
@patch('mem0.memory.storage.SQLiteManager')
def test_update_nonexistent_memory_raises_error(mock_sqlite, mock_llm_factory, mock_vector_factory, mock_embedder_factory):
"""
Test that _update_memory() raises ValueError when memory_id does not exist.
Same class of bug as #3849 — vector_store.get() returns None and code
accesses .payload without a null check.
"""
mock_embedder_factory.return_value = MagicMock()
mock_vector_store = MagicMock()
mock_vector_factory.return_value = mock_vector_store
mock_llm_factory.return_value = MagicMock()
mock_sqlite.return_value = MagicMock()
from mem0.memory.main import Memory as MemoryClass
config = MemoryConfig()
memory = MemoryClass(config)
mock_vector_store.get.return_value = None
with pytest.raises(ValueError, match="Memory with id non-existent-id not found"):
memory._update_memory("non-existent-id", "new data", {"new data": [0.1, 0.2]})
mock_vector_store.update.assert_not_called()
@pytest.mark.asyncio
@patch('mem0.utils.factory.EmbedderFactory.create')
@patch('mem0.utils.factory.VectorStoreFactory.create')
@patch('mem0.utils.factory.LlmFactory.create')
@patch('mem0.memory.storage.SQLiteManager')
async def test_async_update_nonexistent_memory_raises_error(mock_sqlite, mock_llm_factory, mock_vector_factory, mock_embedder_factory):
"""
Test that async _update_memory() raises ValueError when memory_id does not exist.
Same class of bug as #3849 — vector_store.get() returns None and code
accesses .payload without a null check.
"""
mock_embedder_factory.return_value = MagicMock()
mock_vector_store = MagicMock()
mock_vector_store.get.return_value = None
mock_vector_factory.return_value = mock_vector_store
mock_llm_factory.return_value = MagicMock()
mock_sqlite.return_value = MagicMock()
from mem0.memory.main import AsyncMemory
config = MemoryConfig()
memory = AsyncMemory(config)
with pytest.raises(ValueError, match="Memory with id non-existent-id not found"):
await memory._update_memory("non-existent-id", "new data", {"new data": [0.1, 0.2]})
mock_vector_store.update.assert_not_called()
+496
View File
@@ -0,0 +1,496 @@
"""Comprehensive E2E tests for REST API server authentication.
Tests the actual server/main.py app through FastAPI's TestClient (full ASGI
round-trip) covering:
- Auth disabled mode (ADMIN_API_KEY unset)
- Auth enabled mode (ADMIN_API_KEY set)
- Edge cases: empty keys, near-miss keys, timing-safe comparison, header
casing, response headers, startup logging, and full CRUD flows through auth.
"""
import importlib
import logging
import os
from unittest.mock import MagicMock, patch
import pytest
pytest.importorskip("fastapi", reason="fastapi not installed")
from fastapi.testclient import TestClient
# ---------------------------------------------------------------------------
# Fixtures
# ---------------------------------------------------------------------------
@pytest.fixture
def _mock_memory():
"""Patch Memory.from_config so the server imports without a real backend."""
mock_instance = MagicMock()
# Set up return values so CRUD endpoints return realistic responses
mock_instance.get.return_value = {"id": "mem-1", "memory": "test memory", "user_id": "alice"}
mock_instance.get_all.return_value = [
{"id": "mem-1", "memory": "test memory", "user_id": "alice"},
]
mock_instance.add.return_value = {"results": [{"id": "mem-1", "event": "ADD", "memory": "test"}]}
mock_instance.search.return_value = [{"id": "mem-1", "memory": "test", "score": 0.9}]
mock_instance.update.return_value = {"message": "Memory updated"}
mock_instance.history.return_value = [{"id": "mem-1", "old_memory": "a", "new_memory": "b"}]
mock_instance.delete.return_value = None
mock_instance.delete_all.return_value = {"message": "Memories deleted successfully!"}
mock_instance.reset.return_value = None
with patch.dict(os.environ, {"OPENAI_API_KEY": "fake-key"}):
with patch("mem0.Memory.from_config", return_value=mock_instance):
yield mock_instance
def _load_app(env_overrides: dict):
"""Reload server/main.py with the given environment and return the FastAPI app."""
import server.main as server_main
with patch.dict(os.environ, env_overrides, clear=False):
importlib.reload(server_main)
return server_main.app
# ---------------------------------------------------------------------------
# Auth disabled (ADMIN_API_KEY not set)
# ---------------------------------------------------------------------------
class TestAuthDisabled:
"""All endpoints should be freely accessible when ADMIN_API_KEY is empty."""
@pytest.fixture(autouse=True)
def _setup(self, _mock_memory):
self.app = _load_app({"ADMIN_API_KEY": ""})
self.client = TestClient(self.app)
self.mock = _mock_memory
def test_root_redirects_to_docs(self):
resp = self.client.get("/", follow_redirects=False)
assert resp.status_code == 307
assert "/docs" in resp.headers["location"]
def test_get_memory_without_key(self):
resp = self.client.get("/memories/mem-1")
assert resp.status_code == 200
assert resp.json()["id"] == "mem-1"
def test_get_all_memories_without_key(self):
resp = self.client.get("/memories", params={"user_id": "alice"})
assert resp.status_code == 200
def test_create_memory_without_key(self):
resp = self.client.post("/memories", json={
"messages": [{"role": "user", "content": "I like pizza"}],
"user_id": "alice",
})
assert resp.status_code == 200
def test_search_without_key(self):
resp = self.client.post("/search", json={"query": "pizza", "user_id": "alice"})
assert resp.status_code == 200
def test_update_memory_without_key(self):
resp = self.client.put("/memories/mem-1", json={"data": "updated"})
assert resp.status_code == 200
def test_history_without_key(self):
resp = self.client.get("/memories/mem-1/history")
assert resp.status_code == 200
def test_delete_memory_without_key(self):
resp = self.client.delete("/memories/mem-1")
assert resp.status_code == 200
def test_delete_all_without_key(self):
resp = self.client.delete("/memories", params={"user_id": "alice"})
assert resp.status_code == 200
def test_reset_without_key(self):
resp = self.client.post("/reset")
assert resp.status_code == 200
def test_configure_without_key(self):
self.mock.from_config = MagicMock()
resp = self.client.post("/configure", json={"version": "v1.1"})
assert resp.status_code == 200
def test_supplying_key_still_works_when_auth_disabled(self):
"""A client that sends X-API-Key should not be penalized when auth is off."""
resp = self.client.get(
"/memories/mem-1", headers={"X-API-Key": "some-random-key"}
)
assert resp.status_code == 200
@pytest.mark.parametrize(
"method,path",
[
("POST", "/configure"),
("POST", "/memories"),
("GET", "/memories"),
("GET", "/memories/test-id"),
("POST", "/search"),
("PUT", "/memories/test-id"),
("GET", "/memories/test-id/history"),
("DELETE", "/memories/test-id"),
("DELETE", "/memories"),
("POST", "/reset"),
],
)
def test_no_endpoint_returns_401_when_auth_disabled(self, method, path):
resp = self.client.request(method, path)
assert resp.status_code != 401, f"{method} {path} should not require auth"
# ---------------------------------------------------------------------------
# Auth enabled (ADMIN_API_KEY set)
# ---------------------------------------------------------------------------
class TestAuthEnabled:
"""All protected endpoints must enforce the API key."""
API_KEY = "test-secret-key-12345"
@pytest.fixture(autouse=True)
def _setup(self, _mock_memory):
self.app = _load_app({"ADMIN_API_KEY": self.API_KEY})
self.client = TestClient(self.app)
self.mock = _mock_memory
# --- Rejection cases ---
def test_missing_key_returns_401(self):
resp = self.client.get("/memories/mem-1")
assert resp.status_code == 401
def test_missing_key_detail_mentions_header(self):
resp = self.client.get("/memories/mem-1")
assert "X-API-Key" in resp.json()["detail"]
def test_wrong_key_returns_401(self):
resp = self.client.get("/memories/mem-1", headers={"X-API-Key": "wrong"})
assert resp.status_code == 401
def test_wrong_key_detail_says_invalid(self):
resp = self.client.get("/memories/mem-1", headers={"X-API-Key": "wrong"})
assert "Invalid" in resp.json()["detail"]
def test_empty_string_key_returns_401(self):
resp = self.client.get("/memories/mem-1", headers={"X-API-Key": ""})
assert resp.status_code == 401
def test_401_includes_www_authenticate_header(self):
resp = self.client.get("/memories/mem-1")
assert resp.headers.get("www-authenticate") == "ApiKey"
def test_near_miss_key_rejected(self):
"""Key that differs by one character should be rejected."""
near_miss = self.API_KEY[:-1] + ("6" if self.API_KEY[-1] != "6" else "7")
resp = self.client.get("/memories/mem-1", headers={"X-API-Key": near_miss})
assert resp.status_code == 401
def test_key_with_extra_whitespace_rejected(self):
resp = self.client.get("/memories/mem-1", headers={"X-API-Key": f" {self.API_KEY} "})
assert resp.status_code == 401
def test_key_prefix_rejected(self):
resp = self.client.get("/memories/mem-1", headers={"X-API-Key": self.API_KEY[:5]})
assert resp.status_code == 401
def test_key_with_different_case_rejected(self):
resp = self.client.get("/memories/mem-1", headers={"X-API-Key": self.API_KEY.upper()})
assert resp.status_code == 401
@pytest.mark.parametrize(
"method,path",
[
("POST", "/configure"),
("POST", "/memories"),
("GET", "/memories"),
("GET", "/memories/test-id"),
("POST", "/search"),
("PUT", "/memories/test-id"),
("GET", "/memories/test-id/history"),
("DELETE", "/memories/test-id"),
("DELETE", "/memories"),
("POST", "/reset"),
],
)
def test_all_endpoints_reject_without_key(self, method, path):
resp = self.client.request(method, path)
assert resp.status_code == 401, f"{method} {path} should require auth"
@pytest.mark.parametrize(
"method,path",
[
("POST", "/configure"),
("POST", "/memories"),
("GET", "/memories"),
("GET", "/memories/test-id"),
("POST", "/search"),
("PUT", "/memories/test-id"),
("GET", "/memories/test-id/history"),
("DELETE", "/memories/test-id"),
("DELETE", "/memories"),
("POST", "/reset"),
],
)
def test_all_endpoints_reject_wrong_key(self, method, path):
resp = self.client.request(method, path, headers={"X-API-Key": "wrong-key"})
assert resp.status_code == 401, f"{method} {path} should reject wrong key"
# --- Acceptance cases ---
def test_root_does_not_require_key(self):
resp = self.client.get("/", follow_redirects=False)
assert resp.status_code == 307
def _authed(self, method, path, **kwargs):
headers = kwargs.pop("headers", {})
headers["X-API-Key"] = self.API_KEY
return self.client.request(method, path, headers=headers, **kwargs)
def test_get_memory_with_key(self):
resp = self._authed("GET", "/memories/mem-1")
assert resp.status_code == 200
assert resp.json()["id"] == "mem-1"
def test_get_all_memories_with_key(self):
resp = self._authed("GET", "/memories", params={"user_id": "alice"})
assert resp.status_code == 200
def test_create_memory_with_key(self):
resp = self._authed("POST", "/memories", json={
"messages": [{"role": "user", "content": "I like pizza"}],
"user_id": "alice",
})
assert resp.status_code == 200
data = resp.json()
assert "results" in data
def test_search_with_key(self):
resp = self._authed("POST", "/search", json={"query": "pizza", "user_id": "alice"})
assert resp.status_code == 200
def test_update_memory_with_key(self):
resp = self._authed("PUT", "/memories/mem-1", json={"data": "updated"})
assert resp.status_code == 200
def test_history_with_key(self):
resp = self._authed("GET", "/memories/mem-1/history")
assert resp.status_code == 200
def test_delete_memory_with_key(self):
resp = self._authed("DELETE", "/memories/mem-1")
assert resp.status_code == 200
def test_delete_all_with_key(self):
resp = self._authed("DELETE", "/memories", params={"user_id": "alice"})
assert resp.status_code == 200
def test_reset_with_key(self):
resp = self._authed("POST", "/reset")
assert resp.status_code == 200
def test_configure_with_key(self):
resp = self._authed("POST", "/configure", json={"version": "v1.1"})
assert resp.status_code == 200
# ---------------------------------------------------------------------------
# Full CRUD flow through auth
# ---------------------------------------------------------------------------
class TestAuthenticatedCRUDFlow:
"""Verify a complete create → read → search → update → history → delete
cycle works end-to-end through the auth layer."""
API_KEY = "flow-test-key-99"
@pytest.fixture(autouse=True)
def _setup(self, _mock_memory):
self.app = _load_app({"ADMIN_API_KEY": self.API_KEY})
self.client = TestClient(self.app)
self.mock = _mock_memory
def _authed(self, method, path, **kwargs):
headers = kwargs.pop("headers", {})
headers["X-API-Key"] = self.API_KEY
return self.client.request(method, path, headers=headers, **kwargs)
def test_full_crud_cycle(self):
# 1. Create
resp = self._authed("POST", "/memories", json={
"messages": [{"role": "user", "content": "I love fresh vegetable pizza"}],
"user_id": "alice",
})
assert resp.status_code == 200
self.mock.add.assert_called_once()
# 2. Read single
resp = self._authed("GET", "/memories/mem-1")
assert resp.status_code == 200
self.mock.get.assert_called_once_with("mem-1")
# 3. Read all
resp = self._authed("GET", "/memories", params={"user_id": "alice"})
assert resp.status_code == 200
self.mock.get_all.assert_called_once_with(user_id="alice")
# 4. Search
resp = self._authed("POST", "/search", json={"query": "pizza", "user_id": "alice"})
assert resp.status_code == 200
self.mock.search.assert_called_once()
# 5. Update
resp = self._authed("PUT", "/memories/mem-1", json={"data": "updated content"})
assert resp.status_code == 200
self.mock.update.assert_called_once()
# 6. History
resp = self._authed("GET", "/memories/mem-1/history")
assert resp.status_code == 200
self.mock.history.assert_called_once_with(memory_id="mem-1")
# 7. Delete single
resp = self._authed("DELETE", "/memories/mem-1")
assert resp.status_code == 200
self.mock.delete.assert_called_once_with(memory_id="mem-1")
# 8. Delete all
resp = self._authed("DELETE", "/memories", params={"user_id": "alice"})
assert resp.status_code == 200
self.mock.delete_all.assert_called_once()
def test_crud_flow_blocked_without_auth(self):
"""Same flow should fail at every step without the key."""
endpoints = [
("POST", "/memories", {"json": {
"messages": [{"role": "user", "content": "test"}], "user_id": "alice"
}}),
("GET", "/memories/mem-1", {}),
("GET", "/memories", {"params": {"user_id": "alice"}}),
("POST", "/search", {"json": {"query": "pizza", "user_id": "alice"}}),
("PUT", "/memories/mem-1", {"json": {"data": "x"}}),
("GET", "/memories/mem-1/history", {}),
("DELETE", "/memories/mem-1", {}),
("DELETE", "/memories", {"params": {"user_id": "alice"}}),
("POST", "/reset", {}),
]
for method, path, kwargs in endpoints:
resp = self.client.request(method, path, **kwargs)
assert resp.status_code == 401, f"Unauthenticated {method} {path} should be 401"
# Verify the mock was NOT called (auth blocked before reaching handler)
self.mock.add.assert_not_called()
self.mock.get.assert_not_called()
self.mock.search.assert_not_called()
self.mock.update.assert_not_called()
self.mock.history.assert_not_called()
self.mock.delete.assert_not_called()
self.mock.delete_all.assert_not_called()
self.mock.reset.assert_not_called()
# ---------------------------------------------------------------------------
# Edge cases
# ---------------------------------------------------------------------------
class TestAuthEdgeCases:
"""Boundary conditions and unusual inputs."""
@pytest.fixture(autouse=True)
def _setup(self, _mock_memory):
self.mock = _mock_memory
def test_very_long_api_key(self):
"""Server should handle a very long key without crashing."""
long_key = "k" * 4096
app = _load_app({"ADMIN_API_KEY": long_key})
client = TestClient(app)
resp = client.get("/memories/mem-1", headers={"X-API-Key": long_key})
assert resp.status_code == 200
def test_special_characters_in_api_key(self):
"""Keys with special ASCII characters should work."""
special_key = "sk-!@#$%^&*()_+-=[]{}|;:',.<>?/~`"
app = _load_app({"ADMIN_API_KEY": special_key})
client = TestClient(app)
resp = client.get("/memories/mem-1", headers={"X-API-Key": special_key})
assert resp.status_code == 200
resp = client.get("/memories/mem-1", headers={"X-API-Key": "wrong"})
assert resp.status_code == 401
def test_key_env_var_not_present_at_all(self):
"""When the env var is completely absent, auth should be disabled."""
import server.main as server_main
env = os.environ.copy()
env.pop("ADMIN_API_KEY", None)
with patch.dict(os.environ, env, clear=True):
importlib.reload(server_main)
client = TestClient(server_main.app)
resp = client.get("/memories/mem-1")
assert resp.status_code != 401
def test_switching_from_enabled_to_disabled(self):
"""Simulates a server restart with auth toggled off."""
# First: auth enabled
app1 = _load_app({"ADMIN_API_KEY": "secret"})
c1 = TestClient(app1)
assert c1.get("/memories/mem-1").status_code == 401
# Then: auth disabled
app2 = _load_app({"ADMIN_API_KEY": ""})
c2 = TestClient(app2)
assert c2.get("/memories/mem-1").status_code != 401
def test_openapi_schema_accessible_without_key(self):
"""The /docs and /openapi.json endpoints should always be reachable."""
app = _load_app({"ADMIN_API_KEY": "secret"})
client = TestClient(app)
resp = client.get("/openapi.json")
assert resp.status_code == 200
schema = resp.json()
assert "paths" in schema
resp = client.get("/docs")
assert resp.status_code == 200
def test_openapi_schema_documents_auth(self):
"""The OpenAPI schema should mention authentication."""
app = _load_app({"ADMIN_API_KEY": "secret"})
client = TestClient(app)
schema = client.get("/openapi.json").json()
assert "Authentication" in schema.get("info", {}).get("description", "")
# ---------------------------------------------------------------------------
# Startup logging
# ---------------------------------------------------------------------------
class TestStartupLogging:
"""Verify the server emits the correct log messages at import time."""
@pytest.fixture(autouse=True)
def _setup(self, _mock_memory):
pass
def test_warning_when_auth_disabled(self, caplog):
with caplog.at_level(logging.WARNING):
_load_app({"ADMIN_API_KEY": ""})
assert any("UNSECURED" in r.message for r in caplog.records)
def test_info_when_auth_enabled(self, caplog):
with caplog.at_level(logging.INFO):
_load_app({"ADMIN_API_KEY": "a-long-enough-secret-key"})
assert any("authentication enabled" in r.message for r in caplog.records)
def test_warning_when_key_too_short(self, caplog):
with caplog.at_level(logging.WARNING):
_load_app({"ADMIN_API_KEY": "short"})
assert any("shorter than" in r.message for r in caplog.records)