Spaces:
Sleeping
Sleeping
Add enhanced RAG-powered analytics agent with community insights, versioning hooks, and rollback support
f088210 verified | """Enhanced RAG-Powered Analytics Agent with Community Insights. | |
| Integrates retrieval-augmented generation for context-aware analytics, | |
| versioning hooks per SCHEMA_VERSIONING.md, and rollback capabilities. | |
| Enhanced with patterns from bziiit/AGENT_ANALYSE_RAG and fahmiaziz/rag-sql-agent. | |
| """ | |
| import os | |
| import json | |
| import logging | |
| from datetime import datetime | |
| from typing import Dict, List, Optional, Any, Tuple | |
| from dataclasses import dataclass, asdict | |
| import sys | |
| # Import RAG pipeline for context retrieval | |
| from smolagents.rag_pipeline import RAGPipeline, RetrievalConfig | |
| from schema_registry import SchemaRegistry | |
| logger = logging.getLogger(__name__) | |
| class AnalyticsConfig: | |
| """Configuration for RAG-powered analytics.""" | |
| rag_enabled: bool = True | |
| top_k_docs: int = 5 | |
| similarity_threshold: float = 0.7 | |
| enable_versioning: bool = True | |
| schema_version: str = "1.0.0" | |
| rollback_enabled: bool = True | |
| audit_trail: bool = True | |
| class AnalyticsQuery: | |
| """Structured analytics query with RAG context.""" | |
| query_id: str | |
| query_text: str | |
| context_docs: List[Dict[str, Any]] | |
| filters: Dict[str, Any] | |
| aggregations: List[str] | |
| timestamp: str | |
| schema_version: str | |
| class AnalyticsResult: | |
| """Enhanced analytics result with citations and provenance.""" | |
| query_id: str | |
| data: Dict[str, Any] | |
| citations: List[Dict[str, Any]] | |
| confidence_score: float | |
| schema_version: str | |
| rollback_checkpoint: Optional[str] | |
| timestamp: str | |
| class RAGAnalyticsAgent: | |
| """Enhanced Analytics Agent with RAG, versioning, and rollback.""" | |
| def __init__( | |
| self, | |
| config: Optional[AnalyticsConfig] = None, | |
| rag_pipeline: Optional[RAGPipeline] = None, | |
| schema_registry: Optional[SchemaRegistry] = None | |
| ): | |
| self.config = config or AnalyticsConfig() | |
| self.rag_pipeline = rag_pipeline | |
| self.schema_registry = schema_registry | |
| self.checkpoints: Dict[str, Dict] = {} | |
| logger.info(f"Initialized RAGAnalyticsAgent with config: {self.config}") | |
| def create_checkpoint(self, checkpoint_id: str, state: Dict[str, Any]) -> None: | |
| """Create rollback checkpoint per SCHEMA_VERSIONING.md.""" | |
| self.checkpoints[checkpoint_id] = { | |
| "state": state, | |
| "timestamp": datetime.utcnow().isoformat(), | |
| "schema_version": self.config.schema_version | |
| } | |
| logger.info(f"Created checkpoint: {checkpoint_id}") | |
| def rollback_to_checkpoint(self, checkpoint_id: str) -> Optional[Dict[str, Any]]: | |
| """Rollback to a previous checkpoint.""" | |
| if checkpoint_id in self.checkpoints: | |
| checkpoint = self.checkpoints[checkpoint_id] | |
| logger.info(f"Rolling back to checkpoint: {checkpoint_id}") | |
| return checkpoint["state"] | |
| logger.warning(f"Checkpoint not found: {checkpoint_id}") | |
| return None | |
| def retrieve_context( | |
| self, | |
| query_text: str, | |
| top_k: Optional[int] = None | |
| ) -> List[Dict[str, Any]]: | |
| """Retrieve relevant context using RAG pipeline.""" | |
| if not self.config.rag_enabled or not self.rag_pipeline: | |
| return [] | |
| top_k = top_k or self.config.top_k_docs | |
| try: | |
| results = self.rag_pipeline.retrieve( | |
| query=query_text, | |
| top_k=top_k, | |
| threshold=self.config.similarity_threshold | |
| ) | |
| context_docs = [ | |
| { | |
| "content": r.get("content", ""), | |
| "source": r.get("metadata", {}).get("source", "unknown"), | |
| "score": r.get("score", 0.0), | |
| "metadata": r.get("metadata", {}) | |
| } | |
| for r in results | |
| ] | |
| logger.info(f"Retrieved {len(context_docs)} context documents") | |
| return context_docs | |
| except Exception as e: | |
| logger.error(f"Context retrieval failed: {e}") | |
| return [] | |
| def validate_schema( | |
| self, | |
| data: Dict[str, Any], | |
| schema_version: Optional[str] = None | |
| ) -> Tuple[bool, List[str]]: | |
| """Validate data against schema version.""" | |
| if not self.config.enable_versioning or not self.schema_registry: | |
| return True, [] | |
| schema_version = schema_version or self.config.schema_version | |
| try: | |
| is_valid = self.schema_registry.validate(data, schema_version) | |
| errors = [] if is_valid else ["Schema validation failed"] | |
| return is_valid, errors | |
| except Exception as e: | |
| logger.error(f"Schema validation error: {e}") | |
| return False, [str(e)] | |
| def analyze( | |
| self, | |
| query_text: str, | |
| data_source: Dict[str, Any], | |
| filters: Optional[Dict[str, Any]] = None, | |
| aggregations: Optional[List[str]] = None | |
| ) -> AnalyticsResult: | |
| """Perform RAG-enhanced analytics with versioning and rollback.""" | |
| query_id = f"analytics_{datetime.utcnow().timestamp()}" | |
| # Create checkpoint before analysis | |
| if self.config.rollback_enabled: | |
| self.create_checkpoint( | |
| f"{query_id}_pre", | |
| {"data_source": data_source, "filters": filters} | |
| ) | |
| # Retrieve context | |
| context_docs = self.retrieve_context(query_text) | |
| # Build query | |
| query = AnalyticsQuery( | |
| query_id=query_id, | |
| query_text=query_text, | |
| context_docs=context_docs, | |
| filters=filters or {}, | |
| aggregations=aggregations or [], | |
| timestamp=datetime.utcnow().isoformat(), | |
| schema_version=self.config.schema_version | |
| ) | |
| # Perform analysis (simplified - extend with actual analytics logic) | |
| try: | |
| analysis_data = self._execute_analysis(query, data_source) | |
| # Validate result schema | |
| is_valid, errors = self.validate_schema(analysis_data) | |
| if not is_valid: | |
| logger.error(f"Schema validation failed: {errors}") | |
| # Attempt rollback | |
| if self.config.rollback_enabled: | |
| rollback_state = self.rollback_to_checkpoint(f"{query_id}_pre") | |
| if rollback_state: | |
| logger.info("Successfully rolled back to pre-analysis state") | |
| # Build citations from context | |
| citations = [ | |
| { | |
| "source": doc["source"], | |
| "relevance_score": doc["score"], | |
| "excerpt": doc["content"][:200] | |
| } | |
| for doc in context_docs | |
| ] | |
| # Calculate confidence based on context quality | |
| avg_score = sum(d["score"] for d in context_docs) / len(context_docs) if context_docs else 0.5 | |
| result = AnalyticsResult( | |
| query_id=query_id, | |
| data=analysis_data, | |
| citations=citations, | |
| confidence_score=avg_score, | |
| schema_version=self.config.schema_version, | |
| rollback_checkpoint=f"{query_id}_pre" if self.config.rollback_enabled else None, | |
| timestamp=datetime.utcnow().isoformat() | |
| ) | |
| # Audit trail | |
| if self.config.audit_trail: | |
| self._log_audit(query, result) | |
| return result | |
| except Exception as e: | |
| logger.error(f"Analysis failed: {e}") | |
| # Return error result with rollback info | |
| return AnalyticsResult( | |
| query_id=query_id, | |
| data={"error": str(e)}, | |
| citations=[], | |
| confidence_score=0.0, | |
| schema_version=self.config.schema_version, | |
| rollback_checkpoint=f"{query_id}_pre" if self.config.rollback_enabled else None, | |
| timestamp=datetime.utcnow().isoformat() | |
| ) | |
| def _execute_analysis( | |
| self, | |
| query: AnalyticsQuery, | |
| data_source: Dict[str, Any] | |
| ) -> Dict[str, Any]: | |
| """Execute the actual analysis logic (extend as needed).""" | |
| # Simplified analytics - enhance with real logic | |
| results = { | |
| "query_id": query.query_id, | |
| "metrics": {}, | |
| "insights": [], | |
| "context_applied": len(query.context_docs) > 0 | |
| } | |
| # Apply filters | |
| filtered_data = data_source | |
| for key, value in query.filters.items(): | |
| if key in filtered_data: | |
| filtered_data = {k: v for k, v in filtered_data.items() if v == value} | |
| # Apply aggregations | |
| for agg in query.aggregations: | |
| if agg == "count": | |
| results["metrics"]["count"] = len(filtered_data) | |
| elif agg == "sum": | |
| results["metrics"]["sum"] = sum(v for v in filtered_data.values() if isinstance(v, (int, float))) | |
| # Extract insights from RAG context | |
| if query.context_docs: | |
| for doc in query.context_docs[:3]: # Top 3 docs | |
| results["insights"].append({ | |
| "source": doc["source"], | |
| "insight": doc["content"][:150] | |
| }) | |
| return results | |
| def _log_audit( | |
| self, | |
| query: AnalyticsQuery, | |
| result: AnalyticsResult | |
| ) -> None: | |
| """Log audit trail for compliance.""" | |
| audit_entry = { | |
| "timestamp": datetime.utcnow().isoformat(), | |
| "query_id": query.query_id, | |
| "query_text": query.query_text, | |
| "schema_version": query.schema_version, | |
| "confidence_score": result.confidence_score, | |
| "citations_count": len(result.citations) | |
| } | |
| logger.info(f"Audit: {json.dumps(audit_entry)}") | |
| def get_version_info(self) -> Dict[str, Any]: | |
| """Get current version and rollback information.""" | |
| return { | |
| "schema_version": self.config.schema_version, | |
| "checkpoints": list(self.checkpoints.keys()), | |
| "versioning_enabled": self.config.enable_versioning, | |
| "rollback_enabled": self.config.rollback_enabled | |
| } | |
| # Example usage | |
| if __name__ == "__main__": | |
| logging.basicConfig(level=logging.INFO) | |
| # Initialize agent | |
| config = AnalyticsConfig( | |
| rag_enabled=True, | |
| top_k_docs=5, | |
| enable_versioning=True, | |
| rollback_enabled=True | |
| ) | |
| agent = RAGAnalyticsAgent(config=config) | |
| # Example query | |
| result = agent.analyze( | |
| query_text="What are the top engagement metrics for Q4?", | |
| data_source={"engagement": 100, "reach": 500}, | |
| aggregations=["count", "sum"] | |
| ) | |
| print(f"\nAnalysis Result:") | |
| print(f"Query ID: {result.query_id}") | |
| print(f"Confidence: {result.confidence_score}") | |
| print(f"Citations: {len(result.citations)}") | |
| print(f"Rollback checkpoint: {result.rollback_checkpoint}") | |
| print(f"\nVersion Info: {json.dumps(agent.get_version_info(), indent=2)}") | |