ciris-agent 1.7.7__py3-none-any.whl

ciris_engine/logic/services/graph/tsdb_consolidation/consolidators/task.py

@@ -0,0 +1,257 @@
+"""
+Task consolidation for completed tasks and their outcomes.
+
+Creates TaskSummaryNode with final task results and handler selections.
+"""
+
+import logging
+from collections import defaultdict
+from datetime import datetime
+from typing import Any, Dict, List, Optional, Tuple, cast
+
+from ciris_engine.logic.buses.memory_bus import MemoryBus
+from ciris_engine.schemas.services.graph.consolidation import TaskCorrelationData
+from ciris_engine.schemas.services.graph_core import GraphNode, GraphScope, NodeType
+from ciris_engine.schemas.services.operations import MemoryOpStatus
+from ciris_engine.schemas.types import JSONDict
+
+logger = logging.getLogger(__name__)
+
+
+class TaskConsolidator:
+    """Consolidates task outcomes and thought processes."""
+
+    def __init__(self, memory_bus: Optional[MemoryBus] = None):
+        """
+        Initialize task consolidator.
+
+        Args:
+            memory_bus: Memory bus for storing results
+        """
+        self._memory_bus = memory_bus
+
+    async def consolidate(
+        self, period_start: datetime, period_end: datetime, period_label: str, tasks: List[TaskCorrelationData]
+    ) -> Optional[GraphNode]:
+        """
+        Consolidate tasks into a summary showing outcomes and patterns.
+
+        Args:
+            period_start: Start of consolidation period
+            period_end: End of consolidation period
+            period_label: Human-readable period label
+            tasks: List of TaskCorrelationData objects
+
+        Returns:
+            TaskSummaryNode as GraphNode if successful, None otherwise
+        """
+        if not tasks:
+            logger.info(f"No tasks found for period {period_start} to {period_end} - creating empty summary")
+
+        logger.info(f"Consolidating {len(tasks)} tasks for period {period_start}")
+
+        # Aggregate task data
+        tasks_by_status: Dict[str, int] = defaultdict(int)
+        tasks_by_channel: Dict[str, int] = defaultdict(int)
+        tasks_by_user: Dict[str, int] = defaultdict(int)
+        unique_users = set()
+        handler_usage: Dict[str, int] = defaultdict(int)
+        task_durations = []
+        task_summaries = {}
+        total_thoughts = 0
+        thoughts_per_task = []
+        retry_stats: Dict[str, int] = defaultdict(int)
+
+        for task in tasks:
+            task_id = task.task_id
+            status = task.status
+            channel = task.channel_id or "unknown"
+            user_id = task.user_id
+
+            # Track unique users and count tasks per user
+            if user_id:
+                unique_users.add(user_id)
+                tasks_by_user[user_id] += 1
+
+            # Count by status and channel
+            tasks_by_status[status] += 1
+            tasks_by_channel[channel] += 1
+
+            # Track retries
+            retry_count = task.retry_count
+            if retry_count > 0:
+                retry_stats[f"retries_{retry_count}"] += 1
+
+            # Use duration from schema or calculate
+            duration_ms = task.duration_ms
+            if duration_ms == 0:
+                try:
+                    created = task.created_at
+                    updated = task.updated_at
+                    duration_ms = (updated - created).total_seconds() * 1000
+                except (AttributeError, TypeError):
+                    duration_ms = 0
+            task_durations.append(duration_ms)
+
+            # Process thoughts and handlers
+            thoughts = task.thoughts
+            total_thoughts += len(thoughts)
+            thoughts_per_task.append(len(thoughts))
+
+            # Use handlers from schema
+            handlers_selected = task.handlers_used
+            for handler in handlers_selected:
+                handler_usage[handler] += 1
+
+            # Add final handler if present
+            if task.final_handler:
+                handler_usage[task.final_handler] += 1
+
+            # Create task summary
+            task_summaries[task_id] = {
+                "task_id": task_id,
+                "description": task.result_summary or "",
+                "status": status,
+                "channel": channel,
+                "user_id": user_id,  # OBSERVER filtering: preserve user attribution
+                "duration_ms": duration_ms,
+                "thought_count": len(thoughts),
+                "handlers_selected": handlers_selected,
+                "retry_count": retry_count,
+                "outcome": "success" if task.success else task.error_message,
+            }
+
+        # Calculate statistics
+        avg_duration = sum(task_durations) / len(task_durations) if task_durations else 0
+        avg_thoughts = sum(thoughts_per_task) / len(thoughts_per_task) if thoughts_per_task else 0
+
+        # Sort task durations for percentiles
+        if task_durations:
+            sorted_durations = sorted(task_durations)
+            p50_duration = sorted_durations[len(sorted_durations) // 2]
+            p95_duration = sorted_durations[int(len(sorted_durations) * 0.95)]
+            p99_duration = sorted_durations[int(len(sorted_durations) * 0.99)]
+        else:
+            p50_duration = p95_duration = p99_duration = 0
+
+        # Calculate completion rate
+        completed_tasks = tasks_by_status.get("completed", 0) + tasks_by_status.get("success", 0)
+        total_tasks = len(tasks)
+        completion_rate = completed_tasks / total_tasks if total_tasks > 0 else 0
+
+        # Create task summary node
+        summary_data = {
+            "id": f"task_summary_{period_start.strftime('%Y%m%d_%H')}",
+            "period_start": period_start.isoformat(),
+            "period_end": period_end.isoformat(),
+            "period_label": period_label,
+            "total_tasks": total_tasks,
+            "tasks_by_status": dict(tasks_by_status),
+            "tasks_by_channel": dict(tasks_by_channel),
+            "tasks_by_user": dict(tasks_by_user),  # OBSERVER filtering: task count per user
+            "user_list": list(unique_users),  # OBSERVER filtering: all users involved
+            "unique_users": len(unique_users),
+            "completion_rate": completion_rate,
+            "total_thoughts": total_thoughts,
+            "avg_thoughts_per_task": avg_thoughts,
+            "handler_usage": dict(handler_usage),
+            "avg_duration_ms": avg_duration,
+            "p50_duration_ms": p50_duration,
+            "p95_duration_ms": p95_duration,
+            "p99_duration_ms": p99_duration,
+            "retry_stats": dict(retry_stats),
+            "task_summaries": task_summaries,
+            "created_at": period_end.isoformat(),
+            "updated_at": period_end.isoformat(),
+        }
+
+        # Create GraphNode
+        summary_node = GraphNode(
+            id=str(summary_data["id"]),
+            type=NodeType.TASK_SUMMARY,
+            scope=GraphScope.LOCAL,
+            attributes=summary_data,
+            updated_by="tsdb_consolidation",
+            updated_at=period_end,  # Use period end as timestamp
+        )
+
+        # Store summary
+        if self._memory_bus:
+            result = await self._memory_bus.memorize(node=summary_node)
+            if result.status != MemoryOpStatus.OK:
+                logger.error(f"Failed to store task summary: {result.error}")
+                return None
+        else:
+            logger.error("Memory bus not available")
+            return None
+
+        return summary_node
+
+    def get_edges(
+        self, summary_node: GraphNode, tasks: List[TaskCorrelationData]
+    ) -> List[Tuple[GraphNode, GraphNode, str, JSONDict]]:
+        """
+        Get edges to create for task summary.
+
+        Returns edges from summary to:
+        - Failed tasks
+        - Tasks with retries
+        - Long-running tasks
+        """
+        edges = []
+
+        # Process tasks for edge creation
+        failed_count = 0
+        retry_count = 0
+        long_running_count = 0
+
+        for task in tasks:
+            task_id = task.task_id
+
+            # Failed tasks
+            if not task.success and failed_count < 5:
+                edges.append(
+                    (
+                        summary_node,
+                        summary_node,  # Self-reference with task data
+                        "FAILED_TASK",
+                        {
+                            "task_id": task_id,
+                            "failure_reason": task.error_message or "unknown",
+                            "channel_id": task.channel_id,
+                        },
+                    )
+                )
+                failed_count += 1
+
+            # Tasks with retries
+            if task.retry_count > 0 and retry_count < 5:
+                edges.append(
+                    (
+                        summary_node,
+                        summary_node,  # Self-reference with task data
+                        "RETRIED_TASK",
+                        {"task_id": task_id, "retry_count": str(task.retry_count), "final_status": task.status},
+                    )
+                )
+                retry_count += 1
+
+            # Long-running tasks (> 1 minute)
+            duration_seconds = task.duration_ms / 1000.0
+            if duration_seconds > 60 and long_running_count < 5:
+                task_attrs: Dict[str, Optional[str]] = {
+                    "task_id": task_id,
+                    "duration_seconds": str(duration_seconds),
+                    "status": task.status,
+                }
+                edges.append(
+                    (
+                        summary_node,
+                        summary_node,  # Self-reference with task data
+                        "LONG_RUNNING_TASK",
+                        task_attrs,
+                    )
+                )
+                long_running_count += 1
+
+        return cast(List[Tuple[GraphNode, GraphNode, str, JSONDict]], edges)

ciris_engine/logic/services/graph/tsdb_consolidation/consolidators/trace.py

@@ -0,0 +1,363 @@
+"""
+Trace consolidation for trace spans and task processing.
+
+Consolidates TRACE_SPAN correlations into TraceSummaryNode.
+"""
+
+import logging
+from collections import defaultdict
+from datetime import datetime
+from typing import Any, Dict, List, Optional, Set, Tuple, TypedDict, Union
+
+from ciris_engine.logic.buses.memory_bus import MemoryBus
+from ciris_engine.logic.utils.jsondict_helpers import get_dict, get_float, get_int, get_list, get_str
+from ciris_engine.schemas.services.graph.consolidation import TraceSpanData
+from ciris_engine.schemas.services.graph_core import GraphNode, GraphScope, NodeType
+from ciris_engine.schemas.services.operations import MemoryOpStatus
+from ciris_engine.schemas.types import JSONDict
+
+logger = logging.getLogger(__name__)
+
+
+class ThoughtInfo(TypedDict):
+    """Information about a thought in a task."""
+
+    thought_id: str
+    handler: str
+    timestamp: Optional[str]
+
+
+class TaskSummaryData(TypedDict, total=False):
+    """Summary data for a task (using total=False for optional fields).
+
+    Note: This TypedDict is used for documentation purposes.
+    Actual runtime uses JSONDict to allow datetime/Set objects
+    that are later serialized to JSON-compatible types.
+    """
+
+    task_id: str
+    status: str
+    thoughts: List[JSONDict]  # List of thought info dicts
+    start_time: datetime
+    end_time: datetime
+    handlers_selected: List[str]
+    trace_ids: Set[str]
+    duration_ms: float
+
+
+class TraceConsolidator:
+    """Consolidates trace span data into summaries."""
+
+    def __init__(self, memory_bus: Optional[MemoryBus] = None):
+        """
+        Initialize trace consolidator.
+
+        Args:
+            memory_bus: Memory bus for storing results
+        """
+        self._memory_bus = memory_bus
+
+    async def consolidate(
+        self, period_start: datetime, period_end: datetime, period_label: str, trace_spans: List[TraceSpanData]
+    ) -> Optional[GraphNode]:
+        """
+        Consolidate trace spans into a summary showing task processing patterns.
+
+        Args:
+            period_start: Start of consolidation period
+            period_end: End of consolidation period
+            period_label: Human-readable period label
+            trace_spans: List of TraceSpanData objects
+
+        Returns:
+            TraceSummaryNode as GraphNode if successful
+        """
+        if not trace_spans:
+            logger.info(f"No trace spans found for period {period_start} - creating empty summary")
+
+        logger.info(f"Consolidating {len(trace_spans)} trace spans")
+
+        # Initialize tracking structures
+        # Note: Using JSONDict because task_summaries contains datetime and Set objects
+        # that are later serialized to JSON-compatible types before storage
+        task_summaries: JSONDict = {}  # task_id -> summary data
+        unique_tasks: Set[str] = set()
+        unique_thoughts: Set[str] = set()
+        tasks_by_status: Dict[str, int] = defaultdict(int)
+        thoughts_by_type: Dict[str, int] = defaultdict(int)
+        component_calls: Dict[str, int] = defaultdict(int)
+        component_failures: Dict[str, int] = defaultdict(int)
+        component_latencies: Dict[str, List[float]] = defaultdict(list)
+        handler_actions: Dict[str, int] = defaultdict(int)
+        errors_by_component: Dict[str, int] = defaultdict(int)
+        total_errors = 0
+        guardrail_violations: Dict[str, int] = defaultdict(int)
+        dma_decisions: Dict[str, int] = defaultdict(int)
+
+        for span in trace_spans:
+            # Extract key identifiers from typed schema
+            trace_id = span.trace_id
+            span_id = span.span_id
+            parent_span_id = span.parent_span_id
+            timestamp = span.timestamp
+
+            # Extract from tags
+            tags = span.tags
+            task_id = span.task_id
+            thought_id = span.thought_id
+            component_type = span.component_type or "unknown"
+
+            # Track unique entities
+            if task_id:
+                unique_tasks.add(task_id)
+
+                # Initialize task summary if needed
+                if task_id not in task_summaries:
+                    task_summaries[task_id] = {
+                        "task_id": task_id,
+                        "status": "processing",
+                        "thoughts": [],
+                        "start_time": timestamp,
+                        "end_time": timestamp,
+                        "handlers_selected": [],
+                        "trace_ids": set(),
+                    }
+
+                # Type narrowing for nested access
+                task_summary_val = task_summaries.get(task_id)
+                if isinstance(task_summary_val, dict):
+                    trace_ids_val = task_summary_val.get("trace_ids")
+                    if isinstance(trace_ids_val, set):
+                        trace_ids_val.add(trace_id)
+                    task_summary_val["end_time"] = timestamp
+
+            if thought_id:
+                unique_thoughts.add(thought_id)
+
+                # Track thought type
+                thought_type = "unknown"
+                if tags and hasattr(tags, "additional_tags"):
+                    thought_type_val = tags.additional_tags.get("thought_type", "unknown")
+                    thought_type = str(thought_type_val) if thought_type_val else "unknown"
+                thoughts_by_type[thought_type] += 1
+
+                # Track handler selection
+                if component_type == "handler" and task_id:
+                    action_type = "unknown"
+                    if tags and hasattr(tags, "additional_tags"):
+                        action_type_val = tags.additional_tags.get("action_type", "unknown")
+                        action_type = str(action_type_val) if action_type_val else "unknown"
+                    handler_actions[action_type] += 1
+
+                    task_summary_handler = task_summaries.get(task_id)
+                    if isinstance(task_summary_handler, dict):
+                        handlers_sel = task_summary_handler.get("handlers_selected")
+                        if isinstance(handlers_sel, list):
+                            handlers_sel.append(action_type)
+
+                        thoughts_list = task_summary_handler.get("thoughts")
+                        if isinstance(thoughts_list, list):
+                            thoughts_list.append(
+                                {
+                                    "thought_id": thought_id,
+                                    "handler": action_type,
+                                    "timestamp": timestamp.isoformat() if timestamp else None,
+                                }
+                            )
+
+            # Track task completion
+            if task_id and tags and hasattr(tags, "additional_tags") and tags.additional_tags.get("task_status"):
+                status_val = tags.additional_tags["task_status"]
+                status = str(status_val) if isinstance(status_val, (str, int, float)) else "unknown"
+                tasks_by_status[status] += 1
+                task_summary_status = task_summaries.get(task_id)
+                if isinstance(task_summary_status, dict):
+                    task_summary_status["status"] = status
+
+            # Component tracking
+            component_calls[component_type] += 1
+
+            # Process error information
+            if span.error:
+                component_failures[component_type] += 1
+                errors_by_component[component_type] += 1
+                total_errors += 1
+
+            # Track latency
+            if span.latency_ms is not None:
+                component_latencies[component_type].append(span.latency_ms)
+            elif span.duration_ms > 0:
+                component_latencies[component_type].append(span.duration_ms)
+
+            # Track guardrail violations
+            if component_type == "guardrail":
+                guardrail_type = "unknown"
+                violation = False
+                if tags and hasattr(tags, "additional_tags"):
+                    guardrail_type_val = tags.additional_tags.get("guardrail_type", "unknown")
+                    guardrail_type = str(guardrail_type_val) if guardrail_type_val else "unknown"
+                    violation = tags.additional_tags.get("violation") == "true"
+                if violation:
+                    guardrail_violations[guardrail_type] += 1
+
+            # Track DMA decisions
+            if component_type == "dma":
+                dma_type = "unknown"
+                if tags and hasattr(tags, "additional_tags"):
+                    dma_type_val = tags.additional_tags.get("dma_type", "unknown")
+                    dma_type = str(dma_type_val) if dma_type_val else "unknown"
+                dma_decisions[dma_type] += 1
+
+        # Calculate latency statistics
+        component_latency_stats = {}
+        for component, latencies in component_latencies.items():
+            if latencies:
+                sorted_latencies = sorted(latencies)
+                component_latency_stats[component] = {
+                    "avg": sum(latencies) / len(latencies),
+                    "p50": sorted_latencies[len(sorted_latencies) // 2],
+                    "p95": sorted_latencies[int(len(sorted_latencies) * 0.95)],
+                    "p99": sorted_latencies[int(len(sorted_latencies) * 0.99)],
+                }
+
+        # Calculate task processing times
+        task_processing_times = []
+        for task_id, summary_val in task_summaries.items():
+            if not isinstance(summary_val, dict):
+                continue
+            summary = summary_val
+
+            start_time = summary.get("start_time")
+            end_time = summary.get("end_time")
+            if isinstance(start_time, datetime) and isinstance(end_time, datetime):
+                duration_ms = (end_time - start_time).total_seconds() * 1000
+                task_processing_times.append(duration_ms)
+                summary["duration_ms"] = duration_ms
+
+            # Convert sets to lists for JSON serialization
+            trace_ids = summary.get("trace_ids")
+            if isinstance(trace_ids, set):
+                summary["trace_ids"] = list(trace_ids)
+
+        # Calculate task time percentiles
+        avg_task_time = 0.0
+        p50_task_time = 0.0
+        p95_task_time = 0.0
+        p99_task_time = 0.0
+
+        if task_processing_times:
+            sorted_times = sorted(task_processing_times)
+            avg_task_time = sum(task_processing_times) / len(task_processing_times)
+            p50_task_time = sorted_times[len(sorted_times) // 2]
+            p95_task_time = sorted_times[int(len(sorted_times) * 0.95)]
+            p99_task_time = sorted_times[int(len(sorted_times) * 0.99)]
+
+        # Calculate trace depth metrics
+        trace_depths = []
+        for s in task_summaries.values():
+            if isinstance(s, dict):
+                thoughts = get_list(s, "thoughts", [])
+                trace_depths.append(len(thoughts))
+        max_trace_depth = max(trace_depths) if trace_depths else 0
+        avg_trace_depth = sum(trace_depths) / len(trace_depths) if trace_depths else 0.0
+
+        # Calculate error rate
+        total_calls = sum(component_calls.values())
+        error_rate = total_errors / total_calls if total_calls > 0 else 0.0
+
+        # Calculate avg thoughts per task
+        avg_thoughts_per_task = len(unique_thoughts) / len(unique_tasks) if unique_tasks else 0.0
+
+        # Create summary data
+        summary_data = {
+            "id": f"trace_summary_{period_start.strftime('%Y%m%d_%H')}",
+            "period_start": period_start.isoformat(),
+            "period_end": period_end.isoformat(),
+            "period_label": period_label,
+            "total_tasks_processed": len(unique_tasks),
+            "tasks_by_status": dict(tasks_by_status),
+            "unique_task_ids": list(unique_tasks),
+            "task_summaries": task_summaries,
+            "total_thoughts_processed": len(unique_thoughts),
+            "thoughts_by_type": dict(thoughts_by_type),
+            "avg_thoughts_per_task": avg_thoughts_per_task,
+            "component_calls": dict(component_calls),
+            "component_failures": dict(component_failures),
+            "component_latency_ms": component_latency_stats,
+            "dma_decisions": dict(dma_decisions),
+            "guardrail_violations": dict(guardrail_violations),
+            "handler_actions": dict(handler_actions),
+            "avg_task_processing_time_ms": avg_task_time,
+            "p50_task_processing_time_ms": p50_task_time,
+            "p95_task_processing_time_ms": p95_task_time,
+            "p99_task_processing_time_ms": p99_task_time,
+            "total_processing_time_ms": sum(task_processing_times) if task_processing_times else 0.0,
+            "total_errors": total_errors,
+            "errors_by_component": dict(errors_by_component),
+            "error_rate": error_rate,
+            "max_trace_depth": max_trace_depth,
+            "avg_trace_depth": avg_trace_depth,
+            "source_correlation_count": len(trace_spans),
+            "created_at": period_end.isoformat(),
+            "updated_at": period_end.isoformat(),
+        }
+
+        # Create GraphNode
+        summary_node = GraphNode(
+            id=str(summary_data["id"]),
+            type=NodeType.TRACE_SUMMARY,
+            scope=GraphScope.LOCAL,
+            attributes=summary_data,
+            updated_by="tsdb_consolidation",
+            updated_at=period_end,  # Use period end as timestamp
+        )
+
+        # Store summary
+        if self._memory_bus:
+            result = await self._memory_bus.memorize(node=summary_node)
+            if result.status != MemoryOpStatus.OK:
+                logger.error(f"Failed to store trace summary: {result.error}")
+                return None
+        else:
+            logger.warning("No memory bus available - summary not stored")
+
+        return summary_node
+
+    def get_edges(
+        self, summary_node: GraphNode, trace_spans: List[TraceSpanData]
+    ) -> List[Tuple[GraphNode, GraphNode, str, JSONDict]]:
+        """
+        Get edges to create for trace summary.
+
+        Returns edges from summary to:
+        - Tasks with high latency
+        - Components with errors
+        """
+        edges: List[Tuple[GraphNode, GraphNode, str, JSONDict]] = []
+
+        # Find unique tasks
+        tasks_with_errors = set()
+        high_latency_tasks = set()
+
+        for span in trace_spans:
+            task_id = span.trace_id
+            if task_id:
+                # Check for errors
+                if span.error:
+                    tasks_with_errors.add(task_id)
+
+                # Check for high latency (> 5 seconds)
+                latency = span.latency_ms or span.duration_ms
+                if latency and latency > 5000:
+                    high_latency_tasks.add(task_id)
+
+        # Create edges to problematic tasks (limit to 10 each)
+        for i, task_id in enumerate(list(tasks_with_errors)[:10]):
+            edge_attrs: JSONDict = {"task_id": task_id, "error_type": "trace_error"}
+            edges.append((summary_node, summary_node, "ERROR_TASK", edge_attrs))
+
+        for i, task_id in enumerate(list(high_latency_tasks)[:10]):
+            edge_attrs_latency: JSONDict = {"task_id": task_id, "latency_category": "high"}
+            edges.append((summary_node, summary_node, "HIGH_LATENCY_TASK", edge_attrs_latency))
+
+        return edges