Crumb-Core-v.1/app/models/rag_models.py

# app/models/rag_models.py
from pydantic import BaseModel, Field
from typing import Optional, List


class IndexRequest(BaseModel):
    """Request model for indexing posts."""
    provider: str = Field(..., description="Provider to use for embeddings (openai, openrouter, claude)")
    locale: Optional[str] = Field(None, description="Locale filter (de, en, or None for all)")


class IndexSingleRequest(BaseModel):
    """Request model for indexing a single post."""
    provider: str = Field(..., description="Provider to use for embeddings")


class IndexResponse(BaseModel):
    """Response model for indexing operations."""
    status: str
    indexed: Optional[int] = None
    errors: Optional[int] = None
    unchanged: Optional[int] = None
    total: Optional[int] = None
    message: Optional[str] = None
    post_id: Optional[int] = None
    chunks: Optional[int] = None
    collection: Optional[str] = None
    provider: Optional[str] = None


class SearchRequest(BaseModel):
    """Request model for semantic search."""
    query: str = Field(..., description="Search query text")
    provider: str = Field(..., description="Provider to use for search")
    locale: str = Field(..., description="Locale to search in (de, en)")
    limit: int = Field(5, description="Maximum number of results", ge=1, le=20)


class SearchResult(BaseModel):
    """Individual search result."""
    post_id: int
    title: str
    slug: str
    content: str
    header: str
    score: float


class SearchResponse(BaseModel):
    """Response model for search operations."""
    results: List[SearchResult]
    query: str
    locale: str
    provider: str


class QueryRequest(BaseModel):
    """Request model for RAG query."""
    question: str = Field(..., description="Question to answer")
    provider: str = Field(..., description="Provider to use for completion")
    locale: str = Field(..., description="Locale to search in (de, en)")
    context_limit: int = Field(3, description="Number of context chunks to retrieve", ge=1, le=10)


class QuerySource(BaseModel):
    """Source document for RAG query."""
    post_id: int
    title: str
    slug: str
    score: float


class QueryResponse(BaseModel):
    """Response model for RAG query."""
    answer: str
    sources: List[QuerySource]
    provider: str
    model: str
    question: str


class StatusResponse(BaseModel):
    """Response model for indexing status."""
    total_posts: int
    indexed_posts: int
    collections: dict


class ProviderInfo(BaseModel):
    """Information about a provider."""
    name: str
    available: bool
    provider_name: Optional[str] = None
    model: Optional[str] = None
    dimension: Optional[int] = None
    supports_embeddings: Optional[bool] = None
    supports_completions: Optional[bool] = None
    error: Optional[str] = None


class ProvidersResponse(BaseModel):
    """Response model for provider status."""
    providers: List[ProviderInfo]


# ===== Diary-specific models =====

class DiaryIndexRequest(BaseModel):
    """Request model for indexing diary entries."""
    entry_id: int = Field(..., description="Diary entry ID")
    child_id: int = Field(..., description="Child ID (owner of diary)")
    content: str = Field(..., description="Markdown content of diary entry")
    provider: str = Field("openai", description="Provider to use for embeddings")


class DiaryIndexResponse(BaseModel):
    """Response model for diary indexing."""
    status: str
    entry_id: int
    child_id: int
    chunks: Optional[int] = None
    collection: Optional[str] = None
    provider: Optional[str] = None
    message: Optional[str] = None


class DiarySearchRequest(BaseModel):
    """Request model for searching child's diary."""
    child_id: int = Field(..., description="Child ID")
    query: str = Field(..., description="Search query")
    provider: str = Field("openai", description="Provider to use")
    limit: int = Field(5, description="Max results", ge=1, le=20)


class DiarySearchResult(BaseModel):
    """Individual diary search result."""
    entry_id: int
    content: str
    score: float
    created_at: Optional[str] = None


class DiarySearchResponse(BaseModel):
    """Response model for diary search."""
    results: List[DiarySearchResult]
    query: str
    child_id: int
    provider: str


class DiaryAskRequest(BaseModel):
    """Request model for RAG query on diary."""
    child_id: int = Field(..., description="Child ID")
    question: str = Field(..., description="Question to answer")
    provider: str = Field("openai", description="Provider to use")
    context_limit: int = Field(3, description="Number of entries to retrieve", ge=1, le=10)


class DiaryAskResponse(BaseModel):
    """Response model for diary RAG query."""
    answer: str
    question: str
    child_id: int
    sources: List[DiarySearchResult]
    provider: str
    model: str