"""Standardized progress response models for consistent API responses.""" from typing import Any, Literal, cast from pydantic import BaseModel, Field, field_validator class ProgressDetails(BaseModel): """Detailed progress information for granular tracking.""" current_chunk: int | None = Field(None, alias="currentChunk") total_chunks: int | None = Field(None, alias="totalChunks") current_batch: int | None = Field(None, alias="currentBatch") total_batches: int | None = Field(None, alias="totalBatches") current_operation: str | None = Field(None, alias="currentOperation") chunks_per_second: float | None = Field(None, alias="chunksPerSecond") estimated_time_remaining: int | None = Field(None, alias="estimatedTimeRemaining") elapsed_time: int | None = Field(None, alias="elapsedTime") pages_crawled: int | None = Field(None, alias="pagesCrawled") total_pages: int | None = Field(None, alias="totalPages") embeddings_created: int | None = Field(None, alias="embeddingsCreated") code_blocks_found: int | None = Field(None, alias="codeBlocksFound") class Config: populate_by_name = True class BaseProgressResponse(BaseModel): """Base progress response with common fields.""" progress_id: str = Field(alias="progressId") status: str progress: float = Field(ge=0, le=100, description="Progress percentage 0-100") message: str = "" error: str | None = None # Current operation details current_step: str | None = Field(None, alias="currentStep") step_message: str | None = Field(None, alias="stepMessage") logs: list[str] = Field(default_factory=list) details: ProgressDetails | None = None @field_validator("logs", mode="before") @classmethod def ensure_logs_is_list(cls, v): """Ensure logs is always a list of strings, converting from dict if necessary.""" if v is None: return [] if isinstance(v, str): return [v] if isinstance(v, list): # Convert list of dicts to list of strings if needed result = [] for item in v: if isinstance(item, str): result.append(item) elif isinstance(item, dict): # Extract the message from the log dict message = item.get("message", str(item)) result.append(message) else: result.append(str(item)) return result return [] class Config: populate_by_name = True # Accept both snake_case and camelCase class CrawlProgressResponse(BaseProgressResponse): """Progress response for crawl operations.""" status: Literal[ "starting", "analyzing", "crawling", "processing", "source_creation", "document_storage", "code_extraction", "code_storage", "finalization", "completed", "failed", "cancelled", "stopping", "error", ] # Crawl-specific fields current_url: str | None = Field(None, alias="currentUrl") total_pages: int = Field(0, alias="totalPages") processed_pages: int = Field(0, alias="processedPages") crawl_type: str | None = Field(None, alias="crawlType") # 'normal', 'sitemap', 'llms-txt', 'refresh' # Code extraction specific fields code_blocks_found: int = Field(0, alias="codeBlocksFound") code_examples_stored: int = Field(0, alias="codeExamplesStored") completed_documents: int = Field(0, alias="completedDocuments") total_documents: int = Field(0, alias="totalDocuments") completed_summaries: int = Field(0, alias="completedSummaries") total_summaries: int = Field(0, alias="totalSummaries") # Batch processing fields parallel_workers: int | None = Field(None, alias="parallelWorkers") total_jobs: int | None = Field(None, alias="totalJobs") total_batches: int | None = Field(None, alias="totalBatches") completed_batches: int = Field(0, alias="completedBatches") active_workers: int = Field(0, alias="activeWorkers") current_batch: int | None = Field(None, alias="currentBatch") chunks_in_batch: int = Field(0, alias="chunksInBatch") total_chunks_in_batch: int | None = Field(None, alias="totalChunksInBatch") # Results (when completed) chunks_stored: int | None = Field(None, alias="chunksStored") word_count: int | None = Field(None, alias="wordCount") source_id: str | None = Field(None, alias="sourceId") duration: str | None = None @field_validator("duration", mode="before") @classmethod def convert_duration_to_string(cls, v): """Convert duration to string if it's a float.""" if v is None: return None if isinstance(v, int | float): return str(v) return v class Config: populate_by_name = True # Accept both snake_case and camelCase class UploadProgressResponse(BaseProgressResponse): """Progress response for document upload operations.""" status: Literal[ "starting", "reading", "extracting", "chunking", "creating_source", "summarizing", "storing", "source_creation", "completed", "failed", "cancelled", ] # Upload-specific fields upload_type: Literal["document"] = Field("document", alias="uploadType") file_name: str | None = Field(None, alias="fileName") file_type: str | None = Field(None, alias="fileType") # Results (when completed) chunks_stored: int | None = Field(None, alias="chunksStored") word_count: int | None = Field(None, alias="wordCount") source_id: str | None = Field(None, alias="sourceId") class Config: populate_by_name = True # Accept both snake_case and camelCase class ProjectCreationProgressResponse(BaseProgressResponse): """Progress response for project creation operations.""" status: Literal[ "starting", "analyzing", "generating_prp", "creating_tasks", "organizing", "completed", "failed", "error" ] # Project creation specific project_title: str | None = Field(None, alias="projectTitle") tasks_created: int = Field(0, alias="tasksCreated") total_tasks_planned: int | None = Field(None, alias="totalTasksPlanned") class Config: populate_by_name = True # Accept both snake_case and camelCase def create_progress_response(operation_type: str, progress_data: dict[str, Any]) -> BaseProgressResponse: """ Factory function to create the appropriate progress response based on operation type. Args: operation_type: Type of operation (crawl, upload, project_creation) progress_data: Raw progress data from ProgressTracker Returns: Appropriate progress response model """ # Map operation types to response models response_models = { "crawl": CrawlProgressResponse, "upload": UploadProgressResponse, "project_creation": ProjectCreationProgressResponse, } # Get the appropriate model or default to base model_class = response_models.get(operation_type, BaseProgressResponse) # Ensure essential fields have defaults if missing if "status" not in progress_data: progress_data["status"] = "starting" if "progress" not in progress_data: progress_data["progress"] = 0 if "message" not in progress_data and "log" in progress_data: progress_data["message"] = progress_data["log"] # Build details object from various progress fields details_data = {} # Map snake_case fields to camelCase for details detail_field_mappings = { "current_chunk": "currentChunk", "total_chunks": "totalChunks", "current_batch": "currentBatch", "total_batches": "totalBatches", "completed_batches": "currentBatch", # Alternative name "current_operation": "currentOperation", "chunks_per_second": "chunksPerSecond", "estimated_time_remaining": "estimatedTimeRemaining", "elapsed_time": "elapsedTime", "pages_crawled": "pagesCrawled", "processed_pages": "pagesCrawled", # Alternative name "total_pages": "totalPages", "embeddings_created": "embeddingsCreated", "code_blocks_found": "codeBlocksFound", } for snake_field, camel_field in detail_field_mappings.items(): if snake_field in progress_data: # Use the camelCase name since ProgressDetails expects it details_data[camel_field] = progress_data[snake_field] # Also check for crawl-specific fields that might use alternative names if "pages_crawled" not in progress_data and "processed_pages" in progress_data: details_data["pagesCrawled"] = progress_data["processed_pages"] if "totalPages" not in details_data and "total_pages" in progress_data: details_data["totalPages"] = progress_data["total_pages"] # Create details object if we have any detail fields if details_data: progress_data["details"] = ProgressDetails(**details_data) # Create the response, the model will handle field mapping try: # Debug logging for code extraction fields if operation_type == "crawl" and "completed_summaries" in progress_data: from ..config.logfire_config import get_logger logger = get_logger(__name__) logger.info( f"Code extraction progress fields present: completed_summaries={progress_data.get('completed_summaries')}, total_summaries={progress_data.get('total_summaries')}" ) instance = model_class(**progress_data) return cast(BaseProgressResponse, instance) except Exception as e: # Log validation errors for debugging from ..config.logfire_config import get_logger logger = get_logger(__name__) logger.error(f"Failed to create {model_class.__name__}: {e}", exc_info=True) essential_fields = { "progress_id": progress_data.get("progress_id", "unknown"), "status": progress_data.get("status", "running"), "progress": progress_data.get("progress", 0), "message": progress_data.get("message", progress_data.get("log", "")), "error": progress_data.get("error"), } return BaseProgressResponse(**essential_fields)