Compare commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
09494c2acc | ||
|
|
30661ab427 | ||
|
|
7ad5d6f442 | ||
|
|
305ce7b6b3 | ||
|
|
4437c3e8a8 | ||
|
|
54bdbde6e6 | ||
|
|
2b9558335b | ||
|
|
2520edb404 |
@@ -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
|
||||
|
||||
@@ -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).
|
||||
@@ -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" />
|
||||
|
||||
@@ -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",
|
||||
|
||||
@@ -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}>
|
||||
|
||||
@@ -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.
|
||||
|
||||
Vendored
+2
-2
@@ -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
@@ -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>
|
||||
```
|
||||
|
||||
@@ -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";
|
||||
|
||||
@@ -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}`);
|
||||
}
|
||||
|
||||
@@ -522,7 +522,7 @@ export class RedisDB implements VectorStore {
|
||||
|
||||
return {
|
||||
id: vectorId,
|
||||
payload,
|
||||
payload: toCamelCase(payload),
|
||||
};
|
||||
} catch (error) {
|
||||
console.error("Error getting vector:", error);
|
||||
|
||||
@@ -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
@@ -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}")
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -23,6 +23,7 @@ class LlmConfig(BaseModel):
|
||||
"azure_openai_structured",
|
||||
"gemini",
|
||||
"deepseek",
|
||||
"minimax",
|
||||
"xai",
|
||||
"sarvam",
|
||||
"lmstudio",
|
||||
|
||||
@@ -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)
|
||||
@@ -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()
|
||||
@@ -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)
|
||||
|
||||
@@ -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",
|
||||
}
|
||||
|
||||
|
||||
@@ -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",
|
||||
|
||||
@@ -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
@@ -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"] == []
|
||||
|
||||
@@ -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
@@ -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 == []
|
||||
@@ -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()
|
||||
|
||||
@@ -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)
|
||||
Reference in New Issue
Block a user