Spaces:
Running
Running
| # Enterprise Agent State - Extended for Advanced AI Modes | |
| """ | |
| State management for the AI Business Analyst agent. | |
| Tracks: | |
| - Query information | |
| - Routing decisions | |
| - Context and sources | |
| - Confidence metrics | |
| - Traversal paths | |
| """ | |
| from dataclasses import dataclass, field | |
| from typing import List, Dict, Optional, Any | |
| class AgentState: | |
| """ | |
| State object passed through agent workflow. | |
| Attributes: | |
| company_id: User/tenant identifier | |
| question: Original user question | |
| route: Selected processing mode (rag, graph, hybrid, vision) | |
| answer: Generated response | |
| context: Additional context and metadata | |
| sources: List of sources used in response | |
| confidence: Confidence score (0-1) | |
| query_type: Classified query type | |
| reasoning_depth: Required reasoning complexity | |
| """ | |
| company_id: str | |
| question: str | |
| route: str = "" | |
| answer: str = "" | |
| context: Dict[str, Any] = field(default_factory=dict) | |
| sources: List[str] = field(default_factory=list) | |
| # Enhanced fields for enterprise modes | |
| confidence: float = 0.0 | |
| query_type: str = "" | |
| reasoning_depth: str = "" | |
| entities: List[str] = field(default_factory=list) | |
| # Graph traversal metadata | |
| traversal_paths: int = 0 | |
| visited_nodes: int = 0 | |
| max_hops: int = 0 | |
| # Hybrid fusion metadata | |
| fusion_weights: Dict[str, float] = field(default_factory=dict) | |
| primary_mode: str = "" | |
| # Vision metadata | |
| vision_extracted_text: str = "" | |
| vision_tables: List[Dict] = field(default_factory=list) | |
| vision_chart_data: Dict = field(default_factory=dict) | |
| # Performance metrics | |
| processing_time_ms: float = 0.0 | |
| token_count: int = 0 | |
| # Role Intelligence | |
| user_role: str = "analyst" # executive, manager, analyst, operator | |
| def to_dict(self) -> Dict[str, Any]: | |
| """Convert state to dictionary""" | |
| return { | |
| "company_id": self.company_id, | |
| "question": self.question, | |
| "route": self.route, | |
| "answer": self.answer, | |
| "sources": self.sources, | |
| "confidence": self.confidence, | |
| "query_type": self.query_type, | |
| "reasoning_depth": self.reasoning_depth, | |
| "entities": self.entities, | |
| "traversal_paths": self.traversal_paths, | |
| "visited_nodes": self.visited_nodes, | |
| "max_hops": self.max_hops, | |
| "fusion_weights": self.fusion_weights, | |
| "primary_mode": self.primary_mode, | |
| "processing_time_ms": self.processing_time_ms, | |
| "token_count": self.token_count | |
| } | |
| def from_dict(cls, data: Dict[str, Any]) -> "AgentState": | |
| """Create state from dictionary""" | |
| return cls( | |
| company_id=data.get("company_id", ""), | |
| question=data.get("question", ""), | |
| route=data.get("route", ""), | |
| answer=data.get("answer", ""), | |
| context=data.get("context", {}), | |
| sources=data.get("sources", []), | |
| confidence=data.get("confidence", 0.0), | |
| query_type=data.get("query_type", ""), | |
| reasoning_depth=data.get("reasoning_depth", ""), | |
| entities=data.get("entities", []), | |
| traversal_paths=data.get("traversal_paths", 0), | |
| visited_nodes=data.get("visited_nodes", 0), | |
| max_hops=data.get("max_hops", 0), | |
| fusion_weights=data.get("fusion_weights", {}), | |
| primary_mode=data.get("primary_mode", ""), | |
| processing_time_ms=data.get("processing_time_ms", 0.0), | |
| token_count=data.get("token_count", 0) | |
| ) | |
| def add_source(self, source: str) -> None: | |
| """Add a source if not already present""" | |
| if source and source not in self.sources: | |
| self.sources.append(source) | |
| def set_context(self, key: str, value: Any) -> None: | |
| """Set a context value""" | |
| self.context[key] = value | |
| def get_context(self, key: str, default: Any = None) -> Any: | |
| """Get a context value""" | |
| return self.context.get(key, default) | |
| def merge_context(self, new_context: Dict[str, Any]) -> None: | |
| """Merge new context into existing""" | |
| self.context.update(new_context) | |