Glama
mem0 Memory System

"""
Data models for the mem0 MCP server
"""
from typing import Dict, Any, Optional, List, Union
from pydantic import BaseModel, Field

class MemoryConfig(BaseModel):
    """Configuration for memory operations"""
    provider: str = Field(default="openai", description="LLM provider to use")
    embedding_provider: str = Field(default="openai", description="Provider to use for embeddings")
    api_key: Optional[str] = Field(default=None, description="API key for the provider")
    model: Optional[str] = Field(default=None, description="Model to use for the provider")
    embedding_model: Optional[str] = Field(default="text-embedding-3-small", description="Model to use for embeddings")
    data_dir: str = Field(default="./memory_data", description="Directory to store memory data")

class MemoryContent(BaseModel):
    """Content for adding a memory"""
    content: Union[str, List[Dict[str, str]]] = Field(..., description="String content or list of message dicts")
    metadata: Optional[Dict[str, Any]] = Field(default=None, description="Additional metadata for the memory")

class SearchQuery(BaseModel):
    """Query for searching memories"""
    query: str = Field(..., description="Search query")
    limit: int = Field(default=5, description="Maximum number of results to return")

class ChatMessage(BaseModel):
    """Message for chatting with the AI"""
    message: str = Field(..., description="User message")
    system_prompt: Optional[str] = Field(default=None, description="Custom system prompt")
    include_memories: bool = Field(default=True, description="Whether to include memories in the prompt")

class MemoryID(BaseModel):
    """Memory ID for operations"""
    memory_id: str = Field(..., description="Memory ID")

class MemoryResponse(BaseModel):
    """Response containing a memory ID"""
    memory_id: str = Field(..., description="Memory ID")
    success: bool = Field(default=True, description="Whether the operation was successful")

class SearchResponse(BaseModel):
    """Response containing search results"""
    success: bool = Field(default=True, description="Whether the operation was successful")
    results: List[Dict[str, Any]] = Field(..., description="Search results")
    memory_id: Optional[str] = Field(default=None, description="Memory ID of the stored conversation")

class ChatResponse(BaseModel):
    """Response containing a chat message"""
    response: str = Field(..., description="AI response")
    memory_id: Optional[str] = Field(default=None, description="Memory ID of the stored conversation")

class StatusResponse(BaseModel):
    """Response containing operation status"""
    success: bool = Field(..., description="Whether the operation was successful")
    message: str = Field(..., description="Status message")

class MemoryData(BaseModel):
    """Data for a memory"""
    id: str = Field(..., description="Memory ID")
    content: Union[str, List[Dict[str, str]]] = Field(..., description="Memory content")
    metadata: Dict[str, Any] = Field(default_factory=dict, description="Memory metadata")