src/core/response_models.py

"""
Standardized response models for consistent API responses.

Ensures all tools and API endpoints return consistent, validated responses.
"""

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


class ResponseStatus(str, Enum):
    """Response status enum"""
    SUCCESS = "success"
    ERROR = "error"
    PARTIAL = "partial"


class SearchResultItem(BaseModel):
    """Single search result"""
    rank: int = Field(description="Result rank (1-based)")
    score: float = Field(description="Similarity score (0-1)")
    content: str = Field(description="Document content")
    source: Optional[str] = Field(default=None, description="Document source")
    metadata: Optional[Dict[str, Any]] = Field(default=None, description="Additional metadata")


class SearchResponse(BaseModel):
    """Standard search response"""
    status: ResponseStatus = Field(description="Response status")
    query: str = Field(description="Original query")
    result_count: int = Field(description="Number of results")
    results: List[SearchResultItem] = Field(description="Search results")
    elapsed_ms: float = Field(description="Query execution time in ms")
    timestamp: datetime = Field(default_factory=datetime.now, description="Response timestamp")
    error: Optional[str] = Field(default=None, description="Error message if failed")


class QueryResponse(BaseModel):
    """Standard query/QA response"""
    status: ResponseStatus = Field(description="Response status")
    question: str = Field(description="Original question")
    answer: str = Field(description="Generated answer")
    source_count: int = Field(description="Number of sources used")
    confidence: float = Field(description="Confidence score (0-1)")
    elapsed_ms: float = Field(description="Query execution time in ms")
    timestamp: datetime = Field(default_factory=datetime.now, description="Response timestamp")
    error: Optional[str] = Field(default=None, description="Error message if failed")


class ProductAnalysisResponse(BaseModel):
    """Standard product analysis response"""
    status: ResponseStatus = Field(description="Response status")
    product: str = Field(description="Product analyzed")
    analysis: str = Field(description="Analysis result")
    related_products: Optional[List[str]] = Field(default=None, description="Related products found")
    timestamp: datetime = Field(default_factory=datetime.now, description="Response timestamp")
    error: Optional[str] = Field(default=None, description="Error message if failed")


class BatchSearchResponse(BaseModel):
    """Batch search response"""
    status: ResponseStatus = Field(description="Response status")
    batch_id: str = Field(description="Batch ID")
    query_count: int = Field(description="Number of queries")
    successful: int = Field(description="Successful queries")
    failed: int = Field(description="Failed queries")
    results: List[SearchResponse] = Field(description="Individual query results")
    total_elapsed_ms: float = Field(description="Total execution time")
    timestamp: datetime = Field(default_factory=datetime.now, description="Response timestamp")


class HealthResponse(BaseModel):
    """Health check response"""
    status: str = Field(description="Health status")
    timestamp: datetime = Field(default_factory=datetime.now, description="Response timestamp")
    components: Dict[str, str] = Field(description="Component status")
    uptime_seconds: float = Field(description="Uptime in seconds")


class ErrorResponse(BaseModel):
    """Standard error response"""
    status: ResponseStatus = ResponseStatus.ERROR
    error: str = Field(description="Error message")
    code: str = Field(description="Error code")
    timestamp: datetime = Field(default_factory=datetime.now, description="Response timestamp")
    details: Optional[Dict[str, Any]] = Field(default=None, description="Additional error details")


def success_response(data: Dict[str, Any], status: ResponseStatus = ResponseStatus.SUCCESS) -> Dict[str, Any]:
    """Wrap successful response"""
    return {
        **data,
        "status": status.value,
        "timestamp": datetime.now().isoformat()
    }


def error_response(error: str, code: str = "UNKNOWN_ERROR", details: Optional[Dict] = None) -> Dict[str, Any]:
    """Wrap error response"""
    return {
        "status": ResponseStatus.ERROR.value,
        "error": error,
        "code": code,
        "details": details,
        "timestamp": datetime.now().isoformat()
    }