Glama
mem0 Memory System

"""
API endpoints for the mem0 MCP server
"""
from fastapi import FastAPI, HTTPException, Depends
from fastapi.middleware.cors import CORSMiddleware
from fastapi.responses import HTMLResponse
from typing import Dict, Any, Optional

from memory_provider import MemoryProvider
from auto_memory import AutoMemory
from models import (
    MemoryConfig, MemoryContent, SearchQuery, ChatMessage, 
    MemoryID, MemoryResponse, SearchResponse, ChatResponse, 
    StatusResponse, MemoryData
)

# Create FastAPI app
app = FastAPI(
    title="mem0 MCP Server",
    description="Memory Capabilities Provider (MCP) server for mem0",
    version="1.0.0",
)

# Add CORS middleware
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

# Global memory provider instance
memory_provider = None

# Global auto memory instance
auto_memory = None

@app.post("/configure", response_model=StatusResponse)
async def configure_memory(config: MemoryConfig):
    """Configure the memory provider"""
    global memory_provider, auto_memory
    
    try:
        # Initialize memory provider
        memory_provider = MemoryProvider(
            provider=config.provider,
            embedding_provider=config.embedding_provider,
            api_key=config.api_key,
            model=config.model,
            embedding_model=config.embedding_model,
            data_dir=config.data_dir,
        )
        
        # Initialize auto memory
        auto_memory = AutoMemory(memory_provider)
        
        return StatusResponse(
            success=True,
            message=f"Memory provider configured with {config.provider}"
        )
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))

@app.post("/memory/add", response_model=MemoryResponse)
async def add_memory(memory: MemoryContent):
    """Add a memory"""
    global memory_provider
    
    if memory_provider is None:
        raise HTTPException(status_code=500, detail="Memory provider not configured")
    
    try:
        # Add memory
        memory_id = memory_provider.store(
            content=memory.content,
            metadata=memory.metadata
        )
        
        return MemoryResponse(memory_id=memory_id)
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))

@app.post("/memory/search", response_model=SearchResponse)
async def search_memories(query: SearchQuery):
    """Search for memories"""
    global memory_provider
    
    if memory_provider is None:
        raise HTTPException(status_code=500, detail="Memory provider not configured")
    
    try:
        # Search memories
        results = memory_provider.retrieve(
            query=query.query,
            limit=query.limit
        )
        
        # Format results
        formatted_results = []
        for result in results:
            formatted_results.append({
                "memory_id": result.get("id", "unknown"),
                "content": result.get("content", ""),
                "metadata": result.get("metadata", {}),
                "score": result.get("score", 0.0)
            })
        
        return SearchResponse(
            success=True,
            results=formatted_results,
            memory_id=None
        )
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))

@app.get("/memory/{memory_id}", response_model=MemoryData)
async def get_memory(
    memory_id: str,
):
    """Get a memory by ID"""
    global memory_provider
    
    if memory_provider is None:
        raise HTTPException(status_code=500, detail="Memory provider not configured")
    
    try:
        memory = memory_provider.get(memory_id)
        if not memory:
            raise HTTPException(status_code=404, detail=f"Memory {memory_id} not found")
        
        # Ensure we have the required fields
        if "content" not in memory:
            memory["content"] = ""
        if "metadata" not in memory:
            memory["metadata"] = {}
        
        return MemoryData(
            id=memory_id,
            content=memory.get("content", ""),
            metadata=memory.get("metadata", {})
        )
    except HTTPException:
        raise
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))

@app.delete("/memory/{memory_id}", response_model=StatusResponse)
async def delete_memory(
    memory_id: str,
):
    """Delete a memory by ID"""
    global memory_provider
    
    if memory_provider is None:
        raise HTTPException(status_code=500, detail="Memory provider not configured")
    
    try:
        memory_provider.delete(memory_id)
        
        return StatusResponse(
            success=True,
            message=f"Memory {memory_id} deleted"
        )
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))

@app.delete("/memories", response_model=StatusResponse)
async def clear_memories():
    """Clear all memories"""
    global memory_provider
    
    if memory_provider is None:
        raise HTTPException(status_code=500, detail="Memory provider not configured")
    
    try:
        memory_provider.clear()
        
        return StatusResponse(
            success=True,
            message="All memories cleared"
        )
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))

@app.get("/health", response_model=StatusResponse)
async def health_check():
    """Health check endpoint"""
    return StatusResponse(
        success=True,
        message="mem0 MCP server is running"
    )

@app.get("/providers", response_model=Dict[str, Any])
async def list_providers():
    """List available providers"""
    return {
        "llm_providers": MemoryProvider.PROVIDERS,
        "embedding_providers": MemoryProvider.EMBEDDING_PROVIDERS
    }

HTML_TEMPLATE = """
<!DOCTYPE html>
<html>
<head>
    <title>mem0 MCP Server</title>
    <style>
        body {{
            font-family: system-ui, -apple-system, BlinkMacSystemFont, 'Segoe UI', Roboto, Oxygen, Ubuntu, Cantarell, 'Open Sans', 'Helvetica Neue', sans-serif;
            line-height: 1.6;
            color: #333;
            max-width: 800px;
            margin: 0 auto;
            padding: 20px;
        }}
        h1 {{
            color: #2563eb;
            border-bottom: 2px solid #e5e7eb;
            padding-bottom: 10px;
        }}
        h2 {{
            color: #4b5563;
            margin-top: 30px;
        }}
        .endpoint {{
            background-color: #f9fafb;
            border-left: 4px solid #2563eb;
            padding: 15px;
            margin: 15px 0;
            border-radius: 0 5px 5px 0;
        }}
        .method {{
            font-weight: bold;
            color: #2563eb;
        }}
        .path {{
            font-family: monospace;
            background-color: #e5e7eb;
            padding: 2px 5px;
            border-radius: 3px;
        }}
        .description {{
            margin-top: 10px;
        }}
        pre {{
            background-color: #f3f4f6;
            padding: 10px;
            border-radius: 5px;
            overflow-x: auto;
        }}
        code {{
            font-family: monospace;
        }}
        .status {{
            display: inline-block;
            padding: 5px 10px;
            border-radius: 15px;
            font-size: 0.8em;
            margin-top: 10px;
        }}
        .status.running {{
            background-color: #dcfce7;
            color: #166534;
        }}
        .status.not-configured {{
            background-color: #fee2e2;
            color: #991b1b;
        }}
        footer {{
            margin-top: 50px;
            text-align: center;
            color: #6b7280;
            font-size: 0.9em;
        }}
    </style>
</head>
<body>
    <h1>✨ mem0 MCP Server</h1>
    
    <div class="status {status_class}">
        Status: {status}
    </div>
    
    <h2>API Endpoints</h2>
    
    <div class="endpoint">
        <div><span class="method">POST</span> <span class="path">/configure</span></div>
        <div class="description">Configure the memory provider with specified settings.</div>
        <pre><code>{{
  "provider": "openai",
  "embedding_provider": "openai",
  "api_key": "your-api-key",
  "model": "gpt-4o",
  "embedding_model": "text-embedding-3-small",
  "data_dir": "./memory_data"
}}</code></pre>
    </div>
    
    <div class="endpoint">
        <div><span class="method">POST</span> <span class="path">/memory/add</span></div>
        <div class="description">Store a memory with optional metadata.</div>
        <pre><code>{{
  "content": "This is a memory to be stored",
  "metadata": {{
    "source": "user",
    "importance": "high"
  }}
}}</code></pre>
    </div>
    
    <div class="endpoint">
        <div><span class="method">POST</span> <span class="path">/memory/search</span></div>
        <div class="description">Retrieve memories based on a query.</div>
        <pre><code>{{
  "query": "What do you remember about...",
  "limit": 5
}}</code></pre>
    </div>
    
    <div class="endpoint">
        <div><span class="method">GET</span> <span class="path">/memory/{{memory_id}}</span></div>
        <div class="description">Get a memory by ID.</div>
    </div>
    
    <div class="endpoint">
        <div><span class="method">DELETE</span> <span class="path">/memory/{{memory_id}}</span></div>
        <div class="description">Delete a memory by ID.</div>
    </div>
    
    <div class="endpoint">
        <div><span class="method">DELETE</span> <span class="path">/memories</span></div>
        <div class="description">Clear all memories for a user or all users.</div>
    </div>
    
    <h2>Environment Configuration</h2>
    <p>The server can be configured using environment variables:</p>
    <ul>
        <li><code>MEM0_PROVIDER</code> - Default LLM provider (openai, anthropic, google, ollama, etc.)</li>
        <li><code>MEM0_EMBEDDING_PROVIDER</code> - Default embedding provider (openai, huggingface, ollama)</li>
        <li><code>MEM0_DATA_DIR</code> - Directory to store memory data</li>
        <li><code>OPENAI_API_KEY</code> - API key for OpenAI</li>
        <li><code>ANTHROPIC_API_KEY</code> - API key for Anthropic</li>
        <li><code>GOOGLE_API_KEY</code> - API key for Google</li>
        <li><code>OLLAMA_BASE_URL</code> - Base URL for Ollama (default: http://localhost:11434)</li>
    </ul>
    
    <footer>
        Made with ❤️ by Pink Pixel
    </footer>
</body>
</html>
"""

@app.get("/", response_class=HTMLResponse)
async def root():
    """Root endpoint that displays API documentation"""
    status = "Running" if memory_provider is not None else "Not Configured"
    status_class = "running" if memory_provider is not None else "not-configured"
    
    return HTML_TEMPLATE.format(
        status=status,
        status_class=status_class
    )