ciris-agent 1.7.7__py3-none-any.whl

ciris_engine/logic/services/graph/tsdb_consolidation/cleanup_helpers.py

@@ -0,0 +1,438 @@
+"""
+Cleanup helper functions for TSDB consolidation.
+
+Functions to validate and delete expired nodes safely.
+"""
+
+import json
+import logging
+from datetime import datetime
+from sqlite3 import Cursor
+from typing import Dict, List, Optional, Tuple
+
+logger = logging.getLogger(__name__)
+
+
+def validate_and_count_nodes(
+    cursor: Cursor,
+    node_type: str,
+    period_start: str,
+    period_end: str,
+) -> int:
+    """
+    Count nodes of a specific type within a period.
+
+    Database-agnostic: Works with both SQLite and PostgreSQL.
+
+    Args:
+        cursor: Database cursor
+        node_type: Type of nodes to count (e.g., 'tsdb_data', 'audit_entry')
+        period_start: ISO format period start
+        period_end: ISO format period end
+
+    Returns:
+        Count of nodes in the period
+
+    Examples:
+        >>> count = validate_and_count_nodes(cursor, 'tsdb_data', '2023-10-01T00:00:00+00:00', '2023-10-02T00:00:00+00:00')
+        >>> count
+        42
+    """
+    from ciris_engine.logic.persistence.db.dialect import get_adapter
+
+    adapter = get_adapter()
+
+    if adapter.is_postgresql():
+        cursor.execute(
+            """
+            SELECT COUNT(*) FROM graph_nodes
+            WHERE node_type = %s
+              AND created_at::timestamp >= %s::timestamp
+              AND created_at::timestamp < %s::timestamp
+        """,
+            (node_type, period_start, period_end),
+        )
+    else:
+        cursor.execute(
+            """
+            SELECT COUNT(*) FROM graph_nodes
+            WHERE node_type = ?
+              AND datetime(created_at) >= datetime(?)
+              AND datetime(created_at) < datetime(?)
+        """,
+            (node_type, period_start, period_end),
+        )
+
+    result = cursor.fetchone()
+    return result[0] if result else 0
+
+
+def validate_and_count_correlations(
+    cursor: Cursor,
+    period_start: str,
+    period_end: str,
+) -> int:
+    """
+    Count service correlations within a period.
+
+    Database-agnostic: Works with both SQLite and PostgreSQL.
+
+    Args:
+        cursor: Database cursor
+        period_start: ISO format period start
+        period_end: ISO format period end
+
+    Returns:
+        Count of correlations in the period
+
+    Examples:
+        >>> count = validate_and_count_correlations(cursor, '2023-10-01T00:00:00+00:00', '2023-10-02T00:00:00+00:00')
+        >>> count
+        15
+    """
+    from ciris_engine.logic.persistence.db.dialect import get_adapter
+
+    adapter = get_adapter()
+
+    if adapter.is_postgresql():
+        cursor.execute(
+            """
+            SELECT COUNT(*) FROM service_correlations
+            WHERE created_at::timestamp >= %s::timestamp
+              AND created_at::timestamp < %s::timestamp
+        """,
+            (period_start, period_end),
+        )
+    else:
+        cursor.execute(
+            """
+            SELECT COUNT(*) FROM service_correlations
+            WHERE datetime(created_at) >= datetime(?)
+              AND datetime(created_at) < datetime(?)
+        """,
+            (period_start, period_end),
+        )
+
+    result = cursor.fetchone()
+    return result[0] if result else 0
+
+
+def delete_nodes_in_period(
+    cursor: Cursor,
+    node_type: str,
+    period_start: str,
+    period_end: str,
+) -> int:
+    """
+    Delete nodes of a specific type within a period.
+
+    CRITICAL: Deletes edges referencing the nodes first to avoid FOREIGN KEY constraint violations.
+    Database-agnostic: Works with both SQLite and PostgreSQL.
+
+    Args:
+        cursor: Database cursor
+        node_type: Type of nodes to delete
+        period_start: ISO format period start
+        period_end: ISO format period end
+
+    Returns:
+        Number of nodes deleted
+
+    Examples:
+        >>> deleted = delete_nodes_in_period(cursor, 'tsdb_data', '2023-10-01T00:00:00+00:00', '2023-10-02T00:00:00+00:00')
+        >>> logger.info(f"Deleted {deleted} nodes")
+    """
+    from ciris_engine.logic.persistence.db.dialect import get_adapter
+
+    adapter = get_adapter()
+
+    # Step 1: Delete edges referencing these nodes to avoid FOREIGN KEY constraint violations
+    # Use correct column names: source_node_id and target_node_id (not source_id/target_id)
+    if adapter.is_postgresql():
+        # PostgreSQL: Use CAST for datetime comparison
+        cursor.execute(
+            """
+            DELETE FROM graph_edges
+            WHERE source_node_id IN (
+                SELECT node_id FROM graph_nodes
+                WHERE node_type = %s
+                  AND created_at::timestamp >= %s::timestamp
+                  AND created_at::timestamp < %s::timestamp
+            )
+            OR target_node_id IN (
+                SELECT node_id FROM graph_nodes
+                WHERE node_type = %s
+                  AND created_at::timestamp >= %s::timestamp
+                  AND created_at::timestamp < %s::timestamp
+            )
+        """,
+            (node_type, period_start, period_end, node_type, period_start, period_end),
+        )
+    else:
+        # SQLite: Use datetime() function
+        cursor.execute(
+            """
+            DELETE FROM graph_edges
+            WHERE source_node_id IN (
+                SELECT node_id FROM graph_nodes
+                WHERE node_type = ?
+                  AND datetime(created_at) >= datetime(?)
+                  AND datetime(created_at) < datetime(?)
+            )
+            OR target_node_id IN (
+                SELECT node_id FROM graph_nodes
+                WHERE node_type = ?
+                  AND datetime(created_at) >= datetime(?)
+                  AND datetime(created_at) < datetime(?)
+            )
+        """,
+            (node_type, period_start, period_end, node_type, period_start, period_end),
+        )
+
+    # Step 2: Delete the nodes
+    if adapter.is_postgresql():
+        cursor.execute(
+            """
+            DELETE FROM graph_nodes
+            WHERE node_type = %s
+              AND created_at::timestamp >= %s::timestamp
+              AND created_at::timestamp < %s::timestamp
+        """,
+            (node_type, period_start, period_end),
+        )
+    else:
+        cursor.execute(
+            """
+            DELETE FROM graph_nodes
+            WHERE node_type = ?
+              AND datetime(created_at) >= datetime(?)
+              AND datetime(created_at) < datetime(?)
+        """,
+            (node_type, period_start, period_end),
+        )
+
+    return cursor.rowcount
+
+
+def delete_correlations_in_period(
+    cursor: Cursor,
+    period_start: str,
+    period_end: str,
+) -> int:
+    """
+    Delete service correlations within a period.
+
+    Database-agnostic: Works with both SQLite and PostgreSQL.
+
+    Args:
+        cursor: Database cursor
+        period_start: ISO format period start
+        period_end: ISO format period end
+
+    Returns:
+        Number of correlations deleted
+
+    Examples:
+        >>> deleted = delete_correlations_in_period(cursor, '2023-10-01T00:00:00+00:00', '2023-10-02T00:00:00+00:00')
+        >>> logger.info(f"Deleted {deleted} correlations")
+    """
+    from ciris_engine.logic.persistence.db.dialect import get_adapter
+
+    adapter = get_adapter()
+
+    if adapter.is_postgresql():
+        cursor.execute(
+            """
+            DELETE FROM service_correlations
+            WHERE created_at::timestamp >= %s::timestamp
+              AND created_at::timestamp < %s::timestamp
+        """,
+            (period_start, period_end),
+        )
+    else:
+        cursor.execute(
+            """
+            DELETE FROM service_correlations
+            WHERE datetime(created_at) >= datetime(?)
+              AND datetime(created_at) < datetime(?)
+        """,
+            (period_start, period_end),
+        )
+
+    return cursor.rowcount
+
+
+def parse_summary_period(attrs_json: Optional[str]) -> Tuple[Optional[str], Optional[str]]:
+    """
+    Parse period_start and period_end from summary attributes JSON.
+
+    Args:
+        attrs_json: JSON string of summary attributes
+
+    Returns:
+        Tuple of (period_start, period_end) as ISO strings, or (None, None) if missing
+
+    Examples:
+        >>> start, end = parse_summary_period('{"period_start": "2023-10-01T00:00:00+00:00", "period_end": "2023-10-02T00:00:00+00:00"}')
+        >>> start
+        '2023-10-01T00:00:00+00:00'
+    """
+    if not attrs_json:
+        return None, None
+
+    try:
+        attrs = json.loads(attrs_json)
+        period_start = attrs.get("period_start")
+        period_end = attrs.get("period_end")
+        return period_start, period_end
+    except json.JSONDecodeError:
+        logger.warning("Failed to parse summary attributes JSON")
+        return None, None
+
+
+def should_cleanup_summary(
+    claimed_count: int,
+    actual_count: int,
+) -> bool:
+    """
+    Determine if a summary should be cleaned up based on count validation.
+
+    Only cleanup if:
+    1. Claimed count matches actual count (data integrity)
+    2. Actual count > 0 (there's data to cleanup)
+
+    Args:
+        claimed_count: Count claimed in summary attributes
+        actual_count: Actual count from database
+
+    Returns:
+        True if summary should be cleaned up, False otherwise
+
+    Examples:
+        >>> should_cleanup_summary(10, 10)
+        True
+        >>> should_cleanup_summary(10, 9)  # Mismatch
+        False
+        >>> should_cleanup_summary(10, 0)  # No data
+        False
+    """
+    return claimed_count == actual_count and actual_count > 0
+
+
+def cleanup_tsdb_summary(
+    cursor: Cursor,
+    node_id: str,
+    attrs_json: str,
+) -> int:
+    """
+    Cleanup nodes associated with a tsdb_summary.
+
+    Validates claimed count matches actual count before deletion.
+
+    Args:
+        cursor: Database cursor
+        node_id: Summary node ID (for logging)
+        attrs_json: Summary attributes JSON
+
+    Returns:
+        Number of nodes deleted
+
+    Examples:
+        >>> deleted = cleanup_tsdb_summary(cursor, 'tsdb_summary_20231001', '{"source_node_count": 10, ...}')
+    """
+    period_start, period_end = parse_summary_period(attrs_json)
+    if not period_start or not period_end:
+        return 0
+
+    attrs = json.loads(attrs_json) if attrs_json else {}
+    claimed_count = attrs.get("source_node_count", 0)
+
+    # Validate count
+    actual_count = validate_and_count_nodes(cursor, "tsdb_data", period_start, period_end)
+
+    if should_cleanup_summary(claimed_count, actual_count):
+        deleted = delete_nodes_in_period(cursor, "tsdb_data", period_start, period_end)
+        if deleted > 0:
+            logger.info(f"Deleted {deleted} tsdb_data nodes for period {node_id}")
+        return deleted
+
+    return 0
+
+
+def cleanup_audit_summary(
+    cursor: Cursor,
+    node_id: str,
+    attrs_json: str,
+) -> int:
+    """
+    Cleanup graph audit nodes associated with an audit_summary.
+
+    IMPORTANT: This only deletes graph_nodes entries, NOT the audit_log table.
+    The audit_log table is preserved forever for reputability.
+
+    Args:
+        cursor: Database cursor
+        node_id: Summary node ID (for logging)
+        attrs_json: Summary attributes JSON
+
+    Returns:
+        Number of nodes deleted
+
+    Examples:
+        >>> deleted = cleanup_audit_summary(cursor, 'audit_summary_20231001', '{"source_node_count": 5, ...}')
+    """
+    period_start, period_end = parse_summary_period(attrs_json)
+    if not period_start or not period_end:
+        return 0
+
+    attrs = json.loads(attrs_json) if attrs_json else {}
+    claimed_count = attrs.get("source_node_count", 0)
+
+    # Validate count
+    actual_count = validate_and_count_nodes(cursor, "audit_entry", period_start, period_end)
+
+    if should_cleanup_summary(claimed_count, actual_count):
+        deleted = delete_nodes_in_period(cursor, "audit_entry", period_start, period_end)
+        if deleted > 0:
+            logger.info(f"Deleted {deleted} audit_entry graph nodes for period {node_id} (audit_log table preserved)")
+        return deleted
+
+    return 0
+
+
+def cleanup_trace_summary(
+    cursor: Cursor,
+    node_id: str,
+    attrs_json: str,
+) -> int:
+    """
+    Cleanup service correlations associated with a trace_summary.
+
+    Args:
+        cursor: Database cursor
+        node_id: Summary node ID (for logging)
+        attrs_json: Summary attributes JSON
+
+    Returns:
+        Number of correlations deleted
+
+    Examples:
+        >>> deleted = cleanup_trace_summary(cursor, 'trace_summary_20231001', '{"source_correlation_count": 15, ...}')
+    """
+    period_start, period_end = parse_summary_period(attrs_json)
+    if not period_start or not period_end:
+        return 0
+
+    attrs = json.loads(attrs_json) if attrs_json else {}
+    claimed_count = attrs.get("source_correlation_count", 0)
+
+    # Validate count
+    actual_count = validate_and_count_correlations(cursor, period_start, period_end)
+
+    if should_cleanup_summary(claimed_count, actual_count):
+        deleted = delete_correlations_in_period(cursor, period_start, period_end)
+        if deleted > 0:
+            logger.info(f"Deleted {deleted} correlations for period {node_id}")
+        return deleted
+
+    return 0

ciris_engine/logic/services/graph/tsdb_consolidation/compressor.py

@@ -0,0 +1,260 @@
+"""
+Compression utilities for profound consolidation.
+
+This module handles in-place compression of daily summaries for long-term storage efficiency.
+Future versions will include multimedia compression for images, video, and telemetry data.
+"""
+
+import json
+import logging
+from typing import Dict, List, cast
+
+from ciris_engine.logic.utils.jsondict_helpers import get_int, get_str
+from ciris_engine.schemas.services.graph.tsdb_models import CompressionResult, SummaryAttributes
+from ciris_engine.schemas.types import JSONDict
+
+logger = logging.getLogger(__name__)
+
+
+class SummaryCompressor:
+    """Handles compression of summary nodes for profound consolidation."""
+
+    def __init__(self, target_mb_per_day: float):
+        """
+        Initialize the compressor.
+
+        Args:
+            target_mb_per_day: Target size in MB per day after compression
+        """
+        self.target_mb_per_day = target_mb_per_day
+
+    def compress_summary(self, attributes: SummaryAttributes) -> CompressionResult:
+        """
+        Compress a summary's attributes in-place.
+
+        Args:
+            attributes: The summary attributes to compress
+
+        Returns:
+            CompressionResult with compressed attributes and metrics
+        """
+        # Calculate original size
+        original_data = attributes.model_dump(exclude_none=True)
+        original_size = len(json.dumps(original_data, default=str))
+
+        # Create a copy for compression
+        compressed = attributes.model_copy(deep=True)
+
+        # Current compression strategies (text-based)
+        compressed = self._compress_metrics(compressed)
+        compressed = self._compress_descriptions(compressed)
+        compressed = self._remove_redundancy(compressed)
+
+        # Future compression strategies
+        # compressed = self._compress_images(compressed)
+        # compressed = self._compress_video_thumbnails(compressed)
+        # compressed = self._compress_telemetry_data(compressed)
+
+        # Calculate compressed size
+        compressed_data = compressed.model_dump(exclude_none=True)
+        compressed_size = len(json.dumps(compressed_data, default=str))
+        reduction_ratio = 1.0 - (compressed_size / original_size) if original_size > 0 else 0.0
+
+        return CompressionResult(
+            compressed_attributes=compressed,
+            original_size=original_size,
+            compressed_size=compressed_size,
+            reduction_ratio=reduction_ratio,
+            compression_method="text",
+        )
+
+    def _compress_metrics(self, attrs: SummaryAttributes) -> SummaryAttributes:
+        """
+        Compress metrics by keeping only significant patterns.
+
+        Args:
+            attrs: Attributes containing metrics
+
+        Returns:
+            Attributes with compressed metrics
+        """
+        # Compress basic metrics into a simplified format
+        compressed_metrics: Dict[str, float] = {}
+
+        # Only keep non-zero metrics (convert to float for schema compatibility)
+        if attrs.total_interactions > 0:
+            compressed_metrics["ti"] = float(attrs.total_interactions)  # Shortened key
+        if attrs.unique_services > 0:
+            compressed_metrics["us"] = float(attrs.unique_services)
+        if attrs.total_tasks > 0:
+            compressed_metrics["tt"] = float(attrs.total_tasks)
+        if attrs.total_thoughts > 0:
+            compressed_metrics["tth"] = float(attrs.total_thoughts)
+
+        # Store compressed version
+        attrs.compressed_metrics = compressed_metrics if compressed_metrics else None
+
+        # Clear original fields if we compressed them
+        if attrs.compressed_metrics:
+            attrs.total_interactions = 0
+            attrs.unique_services = 0
+            attrs.total_tasks = 0
+            attrs.total_thoughts = 0
+
+        return attrs
+
+    def _compress_descriptions(self, attrs: SummaryAttributes) -> SummaryAttributes:
+        """
+        Compress text descriptions by removing redundancy.
+
+        Args:
+            attrs: Attributes containing descriptions
+
+        Returns:
+            Attributes with compressed descriptions
+        """
+        # Compress conversation summaries
+        if attrs.messages_by_channel:
+            # Build new dict with proper typing - compress to Dict[str, int]
+            compressed_channels: Dict[str, int] = {}
+            for channel, data in attrs.messages_by_channel.items():
+                # Keep only channel ID and count
+                if isinstance(data, dict):
+                    # Extract count from dict with type narrowing
+                    count_val = data.get("count", 0)
+                    if isinstance(count_val, int):
+                        compressed_channels[channel] = count_val
+                    else:
+                        compressed_channels[channel] = 0
+                elif isinstance(data, int):
+                    # Already an int, keep it
+                    compressed_channels[channel] = data
+                else:
+                    # Unknown type, default to 0
+                    compressed_channels[channel] = 0
+
+            # Now cast the Dict[str, int] to the Union type the field expects
+            from typing import Union
+
+            attrs.messages_by_channel = cast(Dict[str, Union[int, JSONDict]], compressed_channels)
+
+        # Compress participant data
+        if attrs.participants:
+            compressed_participants: JSONDict = {}
+            for user_id, data in attrs.participants.items():
+                # Keep only essential data
+                if isinstance(data, dict):
+                    # Use get_str and get_int for type-safe access
+                    author_name = get_str(data, "author_name", "")
+                    msg_count = get_int(data, "message_count", 0)
+                    compressed_participants[user_id] = {
+                        "msg_count": msg_count,
+                        "name": author_name[:20],  # Truncate names
+                    }
+            attrs.participants = cast(Dict[str, JSONDict], compressed_participants)
+
+        # Compress patterns and events
+        if attrs.dominant_patterns and len(attrs.dominant_patterns) > 5:
+            attrs.dominant_patterns = attrs.dominant_patterns[:5]
+
+        if attrs.significant_events and len(attrs.significant_events) > 10:
+            attrs.significant_events = attrs.significant_events[:10]
+
+        return attrs
+
+    def _remove_redundancy(self, attrs: SummaryAttributes) -> SummaryAttributes:
+        """
+        Remove redundant information from attributes.
+
+        Args:
+            attrs: Attributes to clean
+
+        Returns:
+            Attributes with redundancy removed
+        """
+        # Since we're using typed models, redundancy removal is less relevant
+        # The model already has defined fields
+
+        # We can clear any extra fields if using extra="allow"
+        # But for now, just ensure compressed fields are used efficiently
+
+        # If we have compressed versions, clear originals
+        if attrs.compressed_metrics:
+            attrs.total_interactions = 0
+            attrs.unique_services = 0
+            attrs.total_tasks = 0
+            attrs.total_thoughts = 0
+
+        if attrs.compressed_descriptions:
+            # Clear long descriptions if compressed version exists
+            attrs.dominant_patterns = []
+            attrs.significant_events = []
+
+        return attrs
+
+    def estimate_daily_size(self, summaries: List[SummaryAttributes], days_in_period: int) -> float:
+        """
+        Estimate the daily storage size for a set of summaries.
+
+        Args:
+            summaries: List of summary attributes
+            days_in_period: Number of days covered by these summaries
+
+        Returns:
+            Estimated MB per day
+        """
+        total_size = sum(len(json.dumps(s.model_dump(exclude_none=True), default=str)) for s in summaries)
+        size_mb = total_size / (1024 * 1024)
+        return size_mb / days_in_period if days_in_period > 0 else 0
+
+    def needs_compression(self, summaries: List[SummaryAttributes], days_in_period: int) -> bool:
+        """
+        Check if summaries need compression based on target size.
+
+        Args:
+            summaries: List of summary attributes
+            days_in_period: Number of days covered
+
+        Returns:
+            True if compression is needed
+        """
+        current_daily_mb = self.estimate_daily_size(summaries, days_in_period)
+        return current_daily_mb > self.target_mb_per_day
+
+    # Future multimedia compression methods
+
+    def _compress_images(self, attrs: SummaryAttributes) -> SummaryAttributes:
+        """
+        Future: Compress embedded images using lossy compression.
+
+        Will handle:
+        - Converting to JPEG with quality reduction
+        - Resizing to thumbnails
+        - Extracting key frames only
+        """
+        # TODO: Implement when image support is added
+        return attrs
+
+    def _compress_video_thumbnails(self, attrs: SummaryAttributes) -> SummaryAttributes:
+        """
+        Future: Compress video data to thumbnails and metadata.
+
+        Will handle:
+        - Extracting keyframes
+        - Creating timeline thumbnails
+        - Keeping only metadata (duration, resolution, codec)
+        """
+        # TODO: Implement when video support is added
+        return attrs
+
+    def _compress_telemetry_data(self, attrs: SummaryAttributes) -> SummaryAttributes:
+        """
+        Future: Compress robotic/sensor telemetry data.
+
+        Will handle:
+        - Downsampling time series data
+        - Statistical aggregation (min/max/avg/std)
+        - Anomaly detection and preservation
+        """
+        # TODO: Implement when telemetry support is added
+        return attrs