JDSS/mem0
2
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

feat: add optional API key authentication to REST API server (#4442)

Browse Source 
Co-authored-by: utkarsh240799 <utkarsh240799@users.noreply.github.com>
Co-authored-by: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
This commit is contained in:
Utkarsh
2026-03-20 16:15:25 +05:30
committed by GitHub
co-authored by utkarsh240799 Claude Opus 4.6
parent f05e50d940
commit 2520edb404
4 changed files with 592 additions and 12 deletions
Download Patch File Download Diff File Expand all files Collapse all files
server/.env.example
+3
View File
@@ -10,3 +10,6 @@ POSTGRES_DB=
POSTGRES_USER=
POSTGRES_PASSWORD=
POSTGRES_COLLECTION_NAME=
# Optional: set to enable API key authentication on all endpoints
ADMIN_API_KEY=
server/main.py
+55 -10
View File
@@ -1,10 +1,12 @@
import logging
import os
import secrets
from typing import Any, Dict, List, Optional
from dotenv import load_dotenv
from fastapi import FastAPI, HTTPException
from fastapi import Depends, FastAPI, HTTPException
from fastapi.responses import JSONResponse, RedirectResponse
from fastapi.security import APIKeyHeader
from pydantic import BaseModel, Field
from mem0 import Memory
@@ -14,6 +16,22 @@ logging.basicConfig(level=logging.INFO, format="%(asctime)s - %(levelname)s - %(
# Load environment variables
load_dotenv()
ADMIN_API_KEY = os.environ.get("ADMIN_API_KEY", "")
MIN_KEY_LENGTH = 16
if not ADMIN_API_KEY:
logging.warning(
"ADMIN_API_KEY not set - API endpoints are UNSECURED! "
"Set ADMIN_API_KEY environment variable for production use."
)
else:
if len(ADMIN_API_KEY) < MIN_KEY_LENGTH:
logging.warning(
"ADMIN_API_KEY is shorter than %d characters - consider using a longer key for production.",
MIN_KEY_LENGTH,
)
logging.info("API key authentication enabled")
POSTGRES_HOST = os.environ.get("POSTGRES_HOST", "postgres")
POSTGRES_PORT = os.environ.get("POSTGRES_PORT", "5432")
@@ -60,10 +78,35 @@ MEMORY_INSTANCE = Memory.from_config(DEFAULT_CONFIG)
app = FastAPI(
title="Mem0 REST APIs",
description="A REST API for managing and searching memories for your AI Agents and Apps.",
description=(
"A REST API for managing and searching memories for your AI Agents and Apps.\n\n"
"## Authentication\n"
"When the ADMIN_API_KEY environment variable is set, all endpoints require "
"the `X-API-Key` header for authentication."
),
version="1.0.0",
)
api_key_header = APIKeyHeader(name="X-API-Key", auto_error=False)
async def verify_api_key(api_key: Optional[str] = Depends(api_key_header)):
"""Validate the API key when ADMIN_API_KEY is configured. No-op otherwise."""
if ADMIN_API_KEY:
if api_key is None:
raise HTTPException(
status_code=401,
detail="X-API-Key header is required.",
headers={"WWW-Authenticate": "ApiKey"},
)
if not secrets.compare_digest(api_key, ADMIN_API_KEY):
raise HTTPException(
status_code=401,
detail="Invalid API key.",
headers={"WWW-Authenticate": "ApiKey"},
)
return api_key
class Message(BaseModel):
role: str = Field(..., description="Role of the message (user or assistant).")
@@ -87,7 +130,7 @@ class SearchRequest(BaseModel):
@app.post("/configure", summary="Configure Mem0")
def set_config(config: Dict[str, Any]):
def set_config(config: Dict[str, Any], _api_key: Optional[str] = Depends(verify_api_key)):
"""Set memory configuration."""
global MEMORY_INSTANCE
MEMORY_INSTANCE = Memory.from_config(config)
@@ -95,7 +138,7 @@ def set_config(config: Dict[str, Any]):
@app.post("/memories", summary="Create memories")
def add_memory(memory_create: MemoryCreate):
def add_memory(memory_create: MemoryCreate, _api_key: Optional[str] = Depends(verify_api_key)):
"""Store new memories."""
if not any([memory_create.user_id, memory_create.agent_id, memory_create.run_id]):
raise HTTPException(status_code=400, detail="At least one identifier (user_id, agent_id, run_id) is required.")
@@ -114,6 +157,7 @@ def get_all_memories(
user_id: Optional[str] = None,
run_id: Optional[str] = None,
agent_id: Optional[str] = None,
_api_key: Optional[str] = Depends(verify_api_key),
):
"""Retrieve stored memories."""
if not any([user_id, run_id, agent_id]):
@@ -129,7 +173,7 @@ def get_all_memories(
@app.get("/memories/{memory_id}", summary="Get a memory")
def get_memory(memory_id: str):
def get_memory(memory_id: str, _api_key: Optional[str] = Depends(verify_api_key)):
"""Retrieve a specific memory by ID."""
try:
return MEMORY_INSTANCE.get(memory_id)
@@ -139,7 +183,7 @@ def get_memory(memory_id: str):
@app.post("/search", summary="Search memories")
def search_memories(search_req: SearchRequest):
def search_memories(search_req: SearchRequest, _api_key: Optional[str] = Depends(verify_api_key)):
"""Search for memories based on a query."""
try:
params = {k: v for k, v in search_req.model_dump().items() if v is not None and k != "query"}
@@ -150,7 +194,7 @@ def search_memories(search_req: SearchRequest):
@app.put("/memories/{memory_id}", summary="Update a memory")
def update_memory(memory_id: str, updated_memory: Dict[str, Any]):
def update_memory(memory_id: str, updated_memory: Dict[str, Any], _api_key: Optional[str] = Depends(verify_api_key)):
"""Update an existing memory with new content.
Args:
@@ -168,7 +212,7 @@ def update_memory(memory_id: str, updated_memory: Dict[str, Any]):
@app.get("/memories/{memory_id}/history", summary="Get memory history")
def memory_history(memory_id: str):
def memory_history(memory_id: str, _api_key: Optional[str] = Depends(verify_api_key)):
"""Retrieve memory history."""
try:
return MEMORY_INSTANCE.history(memory_id=memory_id)
@@ -178,7 +222,7 @@ def memory_history(memory_id: str):
@app.delete("/memories/{memory_id}", summary="Delete a memory")
def delete_memory(memory_id: str):
def delete_memory(memory_id: str, _api_key: Optional[str] = Depends(verify_api_key)):
"""Delete a specific memory by ID."""
try:
MEMORY_INSTANCE.delete(memory_id=memory_id)
@@ -193,6 +237,7 @@ def delete_all_memories(
user_id: Optional[str] = None,
run_id: Optional[str] = None,
agent_id: Optional[str] = None,
_api_key: Optional[str] = Depends(verify_api_key),
):
"""Delete all memories for a given identifier."""
if not any([user_id, run_id, agent_id]):
@@ -209,7 +254,7 @@ def delete_all_memories(
@app.post("/reset", summary="Reset all memories")
def reset_memory():
def reset_memory(_api_key: Optional[str] = Depends(verify_api_key)):
"""Completely reset stored memories."""
try:
MEMORY_INSTANCE.reset()