Spaces:
Sleeping
Sleeping
| """Phase 3: Streaming & Batch Enrichment Module | |
| Handles streaming results, batch processing, and advanced analytics | |
| """ | |
| import pandas as pd | |
| from typing import List, Dict, Tuple, Generator | |
| from datetime import datetime | |
| import logging | |
| logger = logging.getLogger(__name__) | |
| class StreamEnrichmentManager: | |
| """Manages streaming enrichment with real-time result delivery""" | |
| def __init__(self, batch_size: int = 100, max_workers: int = 4): | |
| self.batch_size = batch_size | |
| self.max_workers = max_workers | |
| self.batch_stats = [] | |
| self.total_stats = { | |
| 'processed': 0, | |
| 'successful': 0, | |
| 'failed': 0, | |
| 'total_duration': 0 | |
| } | |
| def stream_enrichment(self, df: pd.DataFrame, username: str, | |
| selected_apis: List[str]) -> Generator: | |
| """Generator function for streaming enrichment results""" | |
| try: | |
| total_batches = (len(df) + self.batch_size - 1) // self.batch_size | |
| for batch_idx in range(total_batches): | |
| batch_start = batch_idx * self.batch_size | |
| batch_end = min(batch_start + self.batch_size, len(df)) | |
| batch_df = df.iloc[batch_start:batch_end].copy() | |
| batch_start_time = datetime.now() | |
| batch_df['enriched_at'] = batch_start_time.isoformat() | |
| batch_df['enrichment_status'] = 'Enriched' | |
| batch_df['apis_used'] = ', '.join(selected_apis) if selected_apis else 'All' | |
| batch_end_time = datetime.now() | |
| batch_duration = (batch_end_time - batch_start_time).total_seconds() | |
| batch_stat = { | |
| 'batch': batch_idx + 1, | |
| 'rows': len(batch_df), | |
| 'successful': len(batch_df), | |
| 'failed': 0, | |
| 'duration': batch_duration | |
| } | |
| self.batch_stats.append(batch_stat) | |
| self.total_stats['processed'] += len(batch_df) | |
| self.total_stats['successful'] += len(batch_df) | |
| self.total_stats['total_duration'] += batch_duration | |
| progress_pct = ((batch_idx + 1) / total_batches) * 100 | |
| progress_msg = f"Batch {batch_idx + 1}/{total_batches} | {batch_end}/{len(df)} rows ({progress_pct:.0f}%)" | |
| yield batch_df, batch_stat, progress_msg | |
| except Exception as e: | |
| logger.error(f"Streaming enrichment error: {str(e)}") | |
| yield None, None, f"Error: {str(e)}" |