app/domain/entities/query.py

"""
Domain Layer - Query Entity

Represents a user query and its result.
"""
from dataclasses import dataclass, field
from datetime import datetime
from enum import Enum
from typing import List, Optional
from uuid import UUID, uuid4


class QueryStatus(str, Enum):
    """Query processing status"""

    PENDING = "pending"
    PROCESSING = "processing"
    COMPLETED = "completed"
    FAILED = "failed"


@dataclass
class Source:
    """Retrieved source/citation for a query answer"""

    title: str
    content: str
    relevance_score: float
    document_id: UUID
    chunk_index: int
    metadata: dict = field(default_factory=dict)

    def to_dict(self) -> dict:
        """Convert to dictionary"""
        return {
            "title": self.title,
            "content": self.content,
            "relevance_score": self.relevance_score,
            "document_id": str(self.document_id),
            "chunk_index": self.chunk_index,
            "metadata": self.metadata,
        }


@dataclass
class Query:
    """Query entity - represents user question"""

    query_text: str
    department: str
    user_id: Optional[str] = None
    session_id: Optional[str] = None
    id: UUID = field(default_factory=uuid4)
    status: QueryStatus = QueryStatus.PENDING
    answer: Optional[str] = None
    sources: List[Source] = field(default_factory=list)
    confidence: float = 0.0
    duration_ms: int = 0
    tokens_used: int = 0
    model: Optional[str] = None
    created_at: datetime = field(default_factory=datetime.utcnow)
    completed_at: Optional[datetime] = None

    def mark_as_processing(self) -> None:
        """Mark query as being processed"""
        self.status = QueryStatus.PROCESSING

    def mark_as_completed(
        self,
        answer: str,
        sources: List[Source],
        confidence: float,
        duration_ms: int,
        tokens_used: int,
        model: str,
    ) -> None:
        """Mark query as completed with results"""
        self.status = QueryStatus.COMPLETED
        self.answer = answer
        self.sources = sources
        self.confidence = confidence
        self.duration_ms = duration_ms
        self.tokens_used = tokens_used
        self.model = model
        self.completed_at = datetime.utcnow()

    def mark_as_failed(self) -> None:
        """Mark query as failed"""
        self.status = QueryStatus.FAILED
        self.completed_at = datetime.utcnow()


@dataclass
class QueryRequest:
    """Query request from user - value object"""

    query_text: str
    department: str
    user_id: Optional[str] = None
    session_id: Optional[str] = None
    top_k: int = 10
    temperature: float = 0.7
    max_tokens: int = 2048
    filters: dict = field(default_factory=dict)

    def __post_init__(self) -> None:
        """Validate query request"""
        if not self.query_text or len(self.query_text.strip()) == 0:
            raise ValueError("query_text cannot be empty")
        if self.top_k < 1 or self.top_k > 50:
            raise ValueError("top_k must be between 1 and 50")
        if self.temperature < 0 or self.temperature > 1:
            raise ValueError("temperature must be between 0 and 1")