t6_mem0_v2/api/models.py

"""
Pydantic models for T6 Mem0 v2 REST API
"""

from typing import Optional, List, Dict, Any
from datetime import datetime
from pydantic import BaseModel, Field, ConfigDict
from uuid import UUID


class Message(BaseModel):
    """Chat message for memory extraction"""
    role: str = Field(..., description="Message role (user, assistant, system)")
    content: str = Field(..., description="Message content")

    model_config = ConfigDict(
        json_schema_extra={
            "example": {
                "role": "user",
                "content": "I love pizza with extra cheese"
            }
        }
    )


class AddMemoryRequest(BaseModel):
    """Request to add new memory"""
    messages: List[Message] = Field(..., description="Conversation messages")
    user_id: Optional[str] = Field(None, description="User identifier")
    agent_id: Optional[str] = Field(None, description="Agent identifier")
    run_id: Optional[str] = Field(None, description="Run identifier")
    metadata: Optional[Dict[str, Any]] = Field(default_factory=dict, description="Additional metadata")

    model_config = ConfigDict(
        json_schema_extra={
            "example": {
                "messages": [
                    {"role": "user", "content": "I love pizza"},
                    {"role": "assistant", "content": "Great! What toppings?"},
                    {"role": "user", "content": "Extra cheese please"}
                ],
                "user_id": "alice",
                "metadata": {"session": "chat_123", "source": "web"}
            }
        }
    )


class SearchMemoryRequest(BaseModel):
    """Request to search memories"""
    query: str = Field(..., description="Search query")
    user_id: Optional[str] = Field(None, description="User identifier filter")
    agent_id: Optional[str] = Field(None, description="Agent identifier filter")
    run_id: Optional[str] = Field(None, description="Run identifier filter")
    limit: int = Field(default=10, ge=1, le=100, description="Maximum results")

    model_config = ConfigDict(
        json_schema_extra={
            "example": {
                "query": "What food does the user like?",
                "user_id": "alice",
                "limit": 5
            }
        }
    )


class UpdateMemoryRequest(BaseModel):
    """Request to update existing memory"""
    memory_text: Optional[str] = Field(None, description="Updated memory text")
    metadata: Optional[Dict[str, Any]] = Field(None, description="Updated metadata")

    model_config = ConfigDict(
        json_schema_extra={
            "example": {
                "memory_text": "User loves pizza with mushrooms and olives",
                "metadata": {"verified": True, "updated_by": "admin"}
            }
        }
    )


class MemoryResponse(BaseModel):
    """Memory response model"""
    id: str = Field(..., description="Memory unique identifier")
    memory: str = Field(..., description="Memory text content")
    user_id: Optional[str] = Field(None, description="User identifier")
    agent_id: Optional[str] = Field(None, description="Agent identifier")
    run_id: Optional[str] = Field(None, description="Run identifier")
    metadata: Dict[str, Any] = Field(default_factory=dict, description="Memory metadata")
    created_at: Optional[datetime] = Field(None, description="Creation timestamp")
    updated_at: Optional[datetime] = Field(None, description="Last update timestamp")
    score: Optional[float] = Field(None, description="Relevance score (for search results)")

    model_config = ConfigDict(
        json_schema_extra={
            "example": {
                "id": "mem_abc123",
                "memory": "User loves pizza with extra cheese",
                "user_id": "alice",
                "metadata": {"session": "chat_123"},
                "created_at": "2025-10-13T12:00:00Z",
                "score": 0.95
            }
        }
    )


class AddMemoryResponse(BaseModel):
    """Response from adding memory"""
    status: str = Field(..., description="Operation status")
    memories: List[MemoryResponse] = Field(..., description="Created memories")
    message: str = Field(..., description="Result message")

    model_config = ConfigDict(
        json_schema_extra={
            "example": {
                "status": "success",
                "memories": [{
                    "id": "mem_abc123",
                    "memory": "User loves pizza with extra cheese",
                    "user_id": "alice"
                }],
                "message": "Successfully added 1 memory"
            }
        }
    )


class SearchMemoryResponse(BaseModel):
    """Response from searching memories"""
    status: str = Field(..., description="Operation status")
    memories: List[MemoryResponse] = Field(..., description="Matching memories")
    count: int = Field(..., description="Number of results")

    model_config = ConfigDict(
        json_schema_extra={
            "example": {
                "status": "success",
                "memories": [{
                    "id": "mem_abc123",
                    "memory": "User loves pizza",
                    "user_id": "alice",
                    "score": 0.95
                }],
                "count": 1
            }
        }
    )


class DeleteMemoryResponse(BaseModel):
    """Response from deleting memory"""
    status: str = Field(..., description="Operation status")
    message: str = Field(..., description="Result message")

    model_config = ConfigDict(
        json_schema_extra={
            "example": {
                "status": "success",
                "message": "Memory deleted successfully"
            }
        }
    )


class HealthResponse(BaseModel):
    """Health check response"""
    status: str = Field(..., description="Service status")
    version: str = Field(..., description="API version")
    timestamp: datetime = Field(..., description="Check timestamp")
    dependencies: Dict[str, str] = Field(..., description="Dependency status")

    model_config = ConfigDict(
        json_schema_extra={
            "example": {
                "status": "healthy",
                "version": "0.1.0",
                "timestamp": "2025-10-13T12:00:00Z",
                "dependencies": {
                    "supabase": "connected",
                    "neo4j": "connected",
                    "openai": "available"
                }
            }
        }
    )


class StatsResponse(BaseModel):
    """Memory statistics response"""
    total_memories: int = Field(..., description="Total memory count")
    total_users: int = Field(..., description="Unique user count")
    total_agents: int = Field(..., description="Unique agent count")
    avg_memories_per_user: float = Field(..., description="Average memories per user")
    oldest_memory: Optional[datetime] = Field(None, description="Oldest memory timestamp")
    newest_memory: Optional[datetime] = Field(None, description="Newest memory timestamp")

    model_config = ConfigDict(
        json_schema_extra={
            "example": {
                "total_memories": 1523,
                "total_users": 42,
                "total_agents": 5,
                "avg_memories_per_user": 36.26,
                "oldest_memory": "2025-01-01T00:00:00Z",
                "newest_memory": "2025-10-13T12:00:00Z"
            }
        }
    )


class ErrorResponse(BaseModel):
    """Error response model"""
    status: str = Field(default="error", description="Error status")
    error: str = Field(..., description="Error message")
    detail: Optional[str] = Field(None, description="Detailed error information")

    model_config = ConfigDict(
        json_schema_extra={
            "example": {
                "status": "error",
                "error": "Authentication failed",
                "detail": "Invalid or missing API key"
            }
        }
    )