sentienceapi 0.95.0__py3-none-any.whl

sentience/trace_file_manager.py

@@ -0,0 +1,197 @@
+"""
+Trace file management utilities for consistent file operations.
+
+This module provides helper functions for common trace file operations
+shared between JsonlTraceSink and CloudTraceSink.
+"""
+
+import json
+from collections.abc import Callable
+from pathlib import Path
+from typing import Any, Optional
+
+from .models import TraceStats
+
+
+class TraceFileManager:
+    """
+    Helper for common trace file operations.
+
+    Provides static methods for file operations shared across trace sinks.
+    """
+
+    @staticmethod
+    def write_event(file_handle: Any, event: dict[str, Any]) -> None:
+        """
+        Write a trace event to a file handle as JSONL.
+
+        Args:
+            file_handle: Open file handle (must be writable)
+            event: Event dictionary to write
+        """
+        json_str = json.dumps(event, ensure_ascii=False)
+        file_handle.write(json_str + "\n")
+        file_handle.flush()  # Ensure written to disk
+
+    @staticmethod
+    def ensure_directory(path: Path) -> None:
+        """
+        Ensure the parent directory of a path exists.
+
+        Args:
+            path: File path whose parent directory should exist
+        """
+        path.parent.mkdir(parents=True, exist_ok=True)
+
+    @staticmethod
+    def read_events(path: Path) -> list[dict[str, Any]]:
+        """
+        Read all events from a JSONL trace file.
+
+        Args:
+            path: Path to JSONL trace file
+
+        Returns:
+            List of event dictionaries
+
+        Raises:
+            FileNotFoundError: If file doesn't exist
+            json.JSONDecodeError: If file contains invalid JSON
+        """
+        events = []
+        with open(path, encoding="utf-8") as f:
+            for line in f:
+                line = line.strip()
+                if not line:
+                    continue
+                try:
+                    event = json.loads(line)
+                    events.append(event)
+                except json.JSONDecodeError:
+                    # Skip invalid lines but continue reading
+                    continue
+        return events
+
+    @staticmethod
+    def extract_stats(
+        events: list[dict[str, Any]],
+        infer_status_func: None | (
+            Callable[[list[dict[str, Any]], dict[str, Any] | None], str]
+        ) = None,
+    ) -> TraceStats:
+        """
+        Extract execution statistics from trace events.
+
+        This is a common operation shared between JsonlTraceSink and CloudTraceSink.
+
+        Args:
+            events: List of trace event dictionaries
+            infer_status_func: Optional function to infer final_status from events.
+                             If None, uses default inference logic.
+
+        Returns:
+            TraceStats with execution statistics
+        """
+        if not events:
+            return TraceStats(
+                total_steps=0,
+                total_events=0,
+                duration_ms=None,
+                final_status="unknown",
+                started_at=None,
+                ended_at=None,
+            )
+
+        # Find run_start and run_end events
+        run_start = next((e for e in events if e.get("type") == "run_start"), None)
+        run_end = next((e for e in events if e.get("type") == "run_end"), None)
+
+        # Extract timestamps
+        started_at: str | None = None
+        ended_at: str | None = None
+        if run_start:
+            started_at = run_start.get("ts")
+        if run_end:
+            ended_at = run_end.get("ts")
+
+        # Calculate duration
+        duration_ms: int | None = None
+        if started_at and ended_at:
+            try:
+                from datetime import datetime
+
+                start_dt = datetime.fromisoformat(started_at.replace("Z", "+00:00"))
+                end_dt = datetime.fromisoformat(ended_at.replace("Z", "+00:00"))
+                delta = end_dt - start_dt
+                duration_ms = int(delta.total_seconds() * 1000)
+            except Exception:
+                pass
+
+        # Count steps (from step_start events, only first attempt)
+        step_indices = set()
+        for event in events:
+            if event.get("type") == "step_start":
+                step_index = event.get("data", {}).get("step_index")
+                if step_index is not None:
+                    step_indices.add(step_index)
+        total_steps = len(step_indices) if step_indices else 0
+
+        # If run_end has steps count, use that (more accurate)
+        if run_end:
+            steps_from_end = run_end.get("data", {}).get("steps")
+            if steps_from_end is not None:
+                total_steps = max(total_steps, steps_from_end)
+
+        # Count total events
+        total_events = len(events)
+
+        # Infer final status
+        if infer_status_func:
+            final_status = infer_status_func(events, run_end)
+        else:
+            final_status = TraceFileManager._infer_final_status(events, run_end)
+
+        return TraceStats(
+            total_steps=total_steps,
+            total_events=total_events,
+            duration_ms=duration_ms,
+            final_status=final_status,
+            started_at=started_at,
+            ended_at=ended_at,
+        )
+
+    @staticmethod
+    def _infer_final_status(
+        events: list[dict[str, Any]],
+        run_end: dict[str, Any] | None,
+    ) -> str:
+        """
+        Infer final status from trace events.
+
+        Args:
+            events: List of trace event dictionaries
+            run_end: Optional run_end event dictionary
+
+        Returns:
+            Final status string: "success", "failure", "partial", or "unknown"
+        """
+        # Check for run_end event with status
+        if run_end:
+            status = run_end.get("data", {}).get("status")
+            if status in ("success", "failure", "partial", "unknown"):
+                return status
+
+        # Infer from error events
+        has_errors = any(e.get("type") == "error" for e in events)
+        if has_errors:
+            step_ends = [e for e in events if e.get("type") == "step_end"]
+            if step_ends:
+                return "partial"
+            else:
+                return "failure"
+        else:
+            step_ends = [e for e in events if e.get("type") == "step_end"]
+            if step_ends:
+                return "success"
+            else:
+                return "unknown"

sentience/trace_indexing/__init__.py

@@ -0,0 +1,27 @@
+"""
+Trace indexing module for Sentience SDK.
+"""
+
+from .index_schema import (
+    ActionInfo,
+    SnapshotInfo,
+    StepCounters,
+    StepIndex,
+    TraceFileInfo,
+    TraceIndex,
+    TraceSummary,
+)
+from .indexer import build_trace_index, read_step_events, write_trace_index
+
+__all__ = [
+    "build_trace_index",
+    "write_trace_index",
+    "read_step_events",
+    "TraceIndex",
+    "StepIndex",
+    "TraceSummary",
+    "TraceFileInfo",
+    "SnapshotInfo",
+    "ActionInfo",
+    "StepCounters",
+]

sentience/trace_indexing/index_schema.py

@@ -0,0 +1,199 @@
+"""
+Type definitions for trace index schema using concrete classes.
+"""
+
+from dataclasses import asdict, dataclass, field
+from typing import List, Literal, Optional
+
+
+@dataclass
+class TraceFileInfo:
+    """Metadata about the trace file."""
+
+    path: str
+    size_bytes: int
+    sha256: str
+    line_count: int | None = None  # Number of lines in the trace file
+
+    def to_dict(self) -> dict:
+        return asdict(self)
+
+
+@dataclass
+class TraceSummary:
+    """High-level summary of the trace."""
+
+    first_ts: str
+    last_ts: str
+    event_count: int
+    step_count: int
+    error_count: int
+    final_url: str | None
+    status: Literal["success", "failure", "partial", "unknown"] | None = None
+    agent_name: str | None = None  # Agent name from run_start event
+    duration_ms: int | None = None  # Calculated duration in milliseconds
+    counters: dict[str, int] | None = (
+        None  # Aggregated counters (snapshot_count, action_count, error_count)
+    )
+
+    def to_dict(self) -> dict:
+        return asdict(self)
+
+
+@dataclass
+class SnapshotInfo:
+    """Snapshot metadata for index."""
+
+    snapshot_id: str | None = None
+    digest: str | None = None
+    url: str | None = None
+
+    def to_dict(self) -> dict:
+        return asdict(self)
+
+
+@dataclass
+class ActionInfo:
+    """Action metadata for index."""
+
+    type: str | None = None
+    target_element_id: int | None = None
+    args_digest: str | None = None
+    success: bool | None = None
+
+    def to_dict(self) -> dict:
+        return asdict(self)
+
+
+@dataclass
+class StepCounters:
+    """Event counters per step."""
+
+    events: int = 0
+    snapshots: int = 0
+    actions: int = 0
+    llm_calls: int = 0
+
+    def to_dict(self) -> dict:
+        return asdict(self)
+
+
+@dataclass
+class StepIndex:
+    """Index entry for a single step."""
+
+    step_index: int
+    step_id: str
+    goal: str | None
+    status: Literal["success", "failure", "partial", "unknown"]
+    ts_start: str
+    ts_end: str
+    offset_start: int
+    offset_end: int
+    line_number: int | None = None  # Line number for byte-range fetching
+    url_before: str | None = None
+    url_after: str | None = None
+    snapshot_before: SnapshotInfo = field(default_factory=SnapshotInfo)
+    snapshot_after: SnapshotInfo = field(default_factory=SnapshotInfo)
+    action: ActionInfo = field(default_factory=ActionInfo)
+    counters: StepCounters = field(default_factory=StepCounters)
+
+    def to_dict(self) -> dict:
+        result = asdict(self)
+        return result
+
+
+@dataclass
+class TraceIndex:
+    """Complete trace index schema."""
+
+    version: int
+    run_id: str
+    created_at: str
+    trace_file: TraceFileInfo
+    summary: TraceSummary
+    steps: list[StepIndex] = field(default_factory=list)
+
+    def to_dict(self) -> dict:
+        """Convert to dictionary for JSON serialization."""
+        return asdict(self)
+
+    def to_sentience_studio_dict(self) -> dict:
+        """
+        Convert to SS-compatible format.
+
+        Maps SDK field names to frontend expectations:
+        - created_at -> generated_at
+        - first_ts -> start_time
+        - last_ts -> end_time
+        - step_index (0-based) -> step (1-based)
+        - ts_start -> timestamp
+        - Filters out "unknown" status
+        """
+        from datetime import datetime
+
+        # Calculate duration if not already set
+        duration_ms = self.summary.duration_ms
+        if duration_ms is None and self.summary.first_ts and self.summary.last_ts:
+            try:
+                start = datetime.fromisoformat(self.summary.first_ts.replace("Z", "+00:00"))
+                end = datetime.fromisoformat(self.summary.last_ts.replace("Z", "+00:00"))
+                duration_ms = int((end - start).total_seconds() * 1000)
+            except (ValueError, AttributeError):
+                duration_ms = None
+
+        # Aggregate counters if not already set
+        counters = self.summary.counters
+        if counters is None:
+            snapshot_count = sum(step.counters.snapshots for step in self.steps)
+            action_count = sum(step.counters.actions for step in self.steps)
+            counters = {
+                "snapshot_count": snapshot_count,
+                "action_count": action_count,
+                "error_count": self.summary.error_count,
+            }
+
+        return {
+            "version": self.version,
+            "run_id": self.run_id,
+            "generated_at": self.created_at,  # Renamed from created_at
+            "trace_file": {
+                "path": self.trace_file.path,
+                "size_bytes": self.trace_file.size_bytes,
+                "line_count": self.trace_file.line_count,  # Added
+            },
+            "summary": {
+                "agent_name": self.summary.agent_name,  # Added
+                "total_steps": self.summary.step_count,  # Renamed from step_count
+                "status": (
+                    self.summary.status if self.summary.status != "unknown" else None
+                ),  # Filter out unknown
+                "start_time": self.summary.first_ts,  # Renamed from first_ts
+                "end_time": self.summary.last_ts,  # Renamed from last_ts
+                "duration_ms": duration_ms,  # Added
+                "counters": counters,  # Added
+            },
+            "steps": [
+                {
+                    "step": s.step_index + 1,  # Convert 0-based to 1-based
+                    "byte_offset": s.offset_start,
+                    "line_number": s.line_number,  # Added
+                    "timestamp": s.ts_start,  # Use start time
+                    "action": {
+                        "type": s.action.type or "",
+                        "goal": s.goal,  # Move goal into action
+                        "digest": s.action.args_digest,
+                    },
+                    "snapshot": (
+                        {
+                            "url": s.snapshot_after.url,
+                            "digest": s.snapshot_after.digest,
+                        }
+                        if s.snapshot_after.url
+                        else None
+                    ),
+                    "status": s.status if s.status != "unknown" else None,  # Filter out unknown
+                }
+                for s in self.steps
+            ],
+        }