backend/models.py

"""
Pydantic models for request and response validation
"""

from typing import Optional, List, Dict, Any, Union
from pydantic import BaseModel, Field
from datetime import datetime
from enum import Enum


class KnowledgeGraphResponse(BaseModel):
    """Response model for knowledge graphs"""
    files: List[str]


class PlatformStatsResponse(BaseModel):
    """Response model for platform statistics"""
    total_graphs: int
    total_entities: int
    total_relations: int
    entity_distribution: Dict[str, int]
    relation_distribution: Dict[str, int]
    recent_graphs: List[Dict[str, Any]]


# Context Document Models
class ContextDocumentType(str, Enum):
    """Enumeration of context document types."""
    DOMAIN_KNOWLEDGE = "domain_knowledge"
    SCHEMA = "schema"
    DOCUMENTATION = "documentation"
    GUIDELINES = "guidelines"
    EXAMPLES = "examples"


class ContextDocument(BaseModel):
    """Context document model."""
    id: str
    title: str = Field(..., max_length=200, description="Document title")
    document_type: ContextDocumentType
    content: str = Field(..., max_length=100000, description="Document content")
    file_name: Optional[str] = Field(None, description="Original file name if uploaded")
    created_at: datetime
    is_active: bool = Field(default=True, description="Whether document is active")


class CreateContextRequest(BaseModel):
    """Request model for creating context documents."""
    title: str = Field(..., max_length=200, description="Document title")
    document_type: ContextDocumentType
    content: str = Field(..., max_length=100000, description="Document content")
    file_name: Optional[str] = Field(None, description="Original file name if uploaded")


class UpdateContextRequest(BaseModel):
    """Request model for updating context documents."""
    title: Optional[str] = Field(None, max_length=200, description="Document title")
    document_type: Optional[ContextDocumentType] = None
    content: Optional[str] = Field(None, max_length=100000, description="Document content")
    is_active: Optional[bool] = None


class ContextDocumentResponse(BaseModel):
    """Response model for context document operations."""
    success: bool
    message: str
    data: Optional[ContextDocument] = None


class ContextStats(BaseModel):
    """Statistics for context documents."""
    total_count: int
    active_count: int
    by_type: Dict[ContextDocumentType, int]