| import asyncio |
| from collections.abc import Callable |
| from typing import Any |
|
|
| from src.server.config.logfire_config import search_logger |
|
|
|
|
| class ProgressTracker: |
| """Manages progress reporting and SSE communication for storage operations.""" |
|
|
| def __init__(self, progress_callback: Callable | None = None): |
| self.progress_callback = progress_callback |
|
|
| async def report_progress(self, message: str, progress: int, batch_info: dict[str, Any] | None = None): |
| if self.progress_callback and asyncio.iscoroutinefunction(self.progress_callback): |
| try: |
| if batch_info: |
| await self.progress_callback("document_storage", progress, message, **batch_info) |
| else: |
| await self.progress_callback("document_storage", progress, message) |
| except Exception as e: |
| search_logger.warning(f"Progress callback failed: {e}. Storage continuing...") |
|
|
| async def embedding_progress_wrapper(self, message: str, percentage: float, current_progress: int, batch_num: int): |
| |
| if self.progress_callback and "rate limit" in message.lower(): |
| try: |
| await self.progress_callback( |
| "document_storage", |
| current_progress, |
| message, |
| batch=batch_num, |
| type="rate_limit_wait", |
| ) |
| except Exception as e: |
| search_logger.warning(f"Progress callback failed during rate limiting: {e}") |
|
|
| async def report_final_progress(self, total_contents: int, total_batches: int): |
| if self.progress_callback and asyncio.iscoroutinefunction(self.progress_callback): |
| try: |
| await self.progress_callback( |
| "document_storage", |
| 100, |
| f"Document storage completed: {total_contents} chunks stored in {total_batches} batches", |
| completed_batches=total_batches, |
| total_batches=total_batches, |
| current_batch=total_batches, |
| chunks_processed=total_contents, |
| ) |
| except Exception as e: |
| search_logger.warning(f"Progress callback failed during completion: {e}. Storage still successful.") |
|
|
