prela 0.1.0__py3-none-any.whl

prela/replay/loader.py

@@ -0,0 +1,231 @@
+"""Trace loading utilities for replay engine."""
+
+from __future__ import annotations
+
+import json
+import logging
+from pathlib import Path
+from typing import Any
+
+from prela.core.span import Span
+
+logger = logging.getLogger(__name__)
+
+
+class Trace:
+    """Represents a complete trace with all spans.
+
+    A trace is a collection of spans that form a complete execution tree.
+    """
+
+    def __init__(self, trace_id: str, spans: list[Span]) -> None:
+        """Initialize a trace.
+
+        Args:
+            trace_id: Unique trace identifier
+            spans: List of spans in this trace
+        """
+        self.trace_id = trace_id
+        self.spans = spans
+        self._build_tree()
+
+    def _build_tree(self) -> None:
+        """Build parent-child relationships between spans."""
+        # Create span lookup
+        self.span_map: dict[str, Span] = {s.span_id: s for s in self.spans}
+
+        # Build children mapping
+        self.children: dict[str, list[Span]] = {}
+        self.root_spans: list[Span] = []
+
+        for span in self.spans:
+            if span.parent_span_id is None:
+                self.root_spans.append(span)
+            else:
+                if span.parent_span_id not in self.children:
+                    self.children[span.parent_span_id] = []
+                self.children[span.parent_span_id].append(span)
+
+        # Sort children by start time for deterministic execution order
+        for children_list in self.children.values():
+            children_list.sort(key=lambda s: s.started_at)
+
+        self.root_spans.sort(key=lambda s: s.started_at)
+
+    def get_children(self, span_id: str) -> list[Span]:
+        """Get all child spans of a given span.
+
+        Args:
+            span_id: Parent span ID
+
+        Returns:
+            List of child spans (empty if no children)
+        """
+        return self.children.get(span_id, [])
+
+    def walk_depth_first(self) -> list[Span]:
+        """Walk the trace tree depth-first.
+
+        Returns:
+            List of spans in depth-first execution order
+        """
+        result = []
+
+        def visit(span: Span) -> None:
+            result.append(span)
+            for child in self.get_children(span.span_id):
+                visit(child)
+
+        for root in self.root_spans:
+            visit(root)
+
+        return result
+
+    def has_replay_data(self) -> bool:
+        """Check if trace has replay snapshots.
+
+        Returns:
+            True if at least one span has replay data
+        """
+        return any(s.replay_snapshot is not None for s in self.spans)
+
+    def validate_replay_completeness(self) -> tuple[bool, list[str]]:
+        """Validate that trace has complete replay data.
+
+        Returns:
+            Tuple of (is_complete, list of missing span names)
+        """
+        missing = []
+        for span in self.spans:
+            if span.replay_snapshot is None:
+                missing.append(f"{span.name} ({span.span_id})")
+
+        return len(missing) == 0, missing
+
+
+class TraceLoader:
+    """Loads traces from various sources for replay."""
+
+    @staticmethod
+    def from_file(file_path: str | Path) -> Trace:
+        """Load trace from a JSON file.
+
+        Args:
+            file_path: Path to trace JSON file
+
+        Returns:
+            Loaded trace
+
+        Raises:
+            FileNotFoundError: If file doesn't exist
+            ValueError: If file format is invalid
+        """
+        path = Path(file_path)
+        if not path.exists():
+            raise FileNotFoundError(f"Trace file not found: {file_path}")
+
+        with open(path) as f:
+            data = json.load(f)
+
+        return TraceLoader.from_dict(data)
+
+    @staticmethod
+    def from_dict(data: dict[str, Any]) -> Trace:
+        """Load trace from dictionary.
+
+        Args:
+            data: Trace data dictionary
+
+        Returns:
+            Loaded trace
+
+        Raises:
+            ValueError: If data format is invalid
+        """
+        # Handle both single span and trace array formats
+        if isinstance(data, dict) and "trace_id" in data and "spans" in data:
+            # Trace format with metadata
+            trace_id = data["trace_id"]
+            spans_data = data["spans"]
+        elif isinstance(data, dict) and "span_id" in data:
+            # Single span format
+            trace_id = data["trace_id"]
+            spans_data = [data]
+        elif isinstance(data, list):
+            # Array of spans
+            if not data:
+                raise ValueError("Empty span list")
+            trace_id = data[0]["trace_id"]
+            spans_data = data
+        else:
+            raise ValueError("Invalid trace data format")
+
+        # Deserialize spans
+        spans = [Span.from_dict(span_data) for span_data in spans_data]
+
+        return Trace(trace_id, spans)
+
+    @staticmethod
+    def from_jsonl(file_path: str | Path) -> list[Trace]:
+        """Load multiple traces from JSONL file.
+
+        Each line should be a complete trace or span.
+
+        Args:
+            file_path: Path to JSONL file
+
+        Returns:
+            List of traces
+
+        Raises:
+            FileNotFoundError: If file doesn't exist
+        """
+        path = Path(file_path)
+        if not path.exists():
+            raise FileNotFoundError(f"JSONL file not found: {file_path}")
+
+        # Group spans by trace_id
+        traces_data: dict[str, list[dict[str, Any]]] = {}
+
+        with open(path) as f:
+            for line in f:
+                line = line.strip()
+                if not line:
+                    continue
+
+                span_data = json.loads(line)
+                trace_id = span_data["trace_id"]
+
+                if trace_id not in traces_data:
+                    traces_data[trace_id] = []
+                traces_data[trace_id].append(span_data)
+
+        # Build traces
+        traces = []
+        for trace_id, spans_data in traces_data.items():
+            spans = [Span.from_dict(span_data) for span_data in spans_data]
+            traces.append(Trace(trace_id, spans))
+
+        return traces
+
+    @staticmethod
+    def from_span_list(spans: list[Span]) -> Trace:
+        """Create trace from list of spans.
+
+        Args:
+            spans: List of spans
+
+        Returns:
+            Trace
+
+        Raises:
+            ValueError: If spans don't share same trace_id
+        """
+        if not spans:
+            raise ValueError("Empty span list")
+
+        trace_id = spans[0].trace_id
+        if not all(s.trace_id == trace_id for s in spans):
+            raise ValueError("All spans must have the same trace_id")
+
+        return Trace(trace_id, spans)

prela/replay/result.py

@@ -0,0 +1,196 @@
+"""Result data structures for replay execution."""
+
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+from typing import Any
+
+
+@dataclass
+class ReplayedSpan:
+    """Result of replaying a single span.
+
+    Tracks what was executed, whether it was modified, and the output.
+    """
+
+    original_span_id: str
+    span_type: str
+    name: str
+    input: Any
+    output: Any
+    was_modified: bool = False
+    modification_details: str | None = None
+    duration_ms: float = 0.0
+    tokens_used: int = 0
+    cost_usd: float = 0.0
+    error: str | None = None
+    retry_count: int = 0  # Number of retry attempts for API calls
+
+    def to_dict(self) -> dict[str, Any]:
+        """Convert to dictionary representation."""
+        return {
+            "original_span_id": self.original_span_id,
+            "span_type": self.span_type,
+            "name": self.name,
+            "input": self.input,
+            "output": self.output,
+            "was_modified": self.was_modified,
+            "modification_details": self.modification_details,
+            "duration_ms": self.duration_ms,
+            "tokens_used": self.tokens_used,
+            "cost_usd": self.cost_usd,
+            "error": self.error,
+            "retry_count": self.retry_count,
+        }
+
+
+@dataclass
+class ReplayResult:
+    """Complete result of replaying a trace.
+
+    Contains all replayed spans, aggregated metrics, and final output.
+    """
+
+    trace_id: str
+    spans: list[ReplayedSpan] = field(default_factory=list)
+    total_duration_ms: float = 0.0
+    total_tokens: int = 0
+    total_cost_usd: float = 0.0
+    final_output: Any = None
+    errors: list[str] = field(default_factory=list)
+    modifications_applied: dict[str, Any] = field(default_factory=dict)
+
+    @property
+    def success(self) -> bool:
+        """Check if replay completed without errors."""
+        return len(self.errors) == 0 and all(s.error is None for s in self.spans)
+
+    @property
+    def modified_span_count(self) -> int:
+        """Count how many spans were modified."""
+        return sum(1 for s in self.spans if s.was_modified)
+
+    def to_dict(self) -> dict[str, Any]:
+        """Convert to dictionary representation."""
+        return {
+            "trace_id": self.trace_id,
+            "spans": [s.to_dict() for s in self.spans],
+            "total_duration_ms": self.total_duration_ms,
+            "total_tokens": self.total_tokens,
+            "total_cost_usd": self.total_cost_usd,
+            "final_output": self.final_output,
+            "errors": self.errors,
+            "modifications_applied": self.modifications_applied,
+            "success": self.success,
+            "modified_span_count": self.modified_span_count,
+        }
+
+
+@dataclass
+class SpanDifference:
+    """Difference between two span executions.
+
+    Captures what changed between original and modified replay.
+    """
+
+    span_name: str
+    span_type: str
+    field: str
+    original_value: Any
+    modified_value: Any
+    semantic_similarity: float | None = None
+    exact_match: bool = False
+
+    def to_dict(self) -> dict[str, Any]:
+        """Convert to dictionary representation."""
+        return {
+            "span_name": self.span_name,
+            "span_type": self.span_type,
+            "field": self.field,
+            "original_value": self.original_value,
+            "modified_value": self.modified_value,
+            "semantic_similarity": self.semantic_similarity,
+            "exact_match": self.exact_match,
+        }
+
+
+@dataclass
+class ReplayComparison:
+    """Comparison between two replay results.
+
+    Highlights differences and provides summary statistics.
+    """
+
+    original: ReplayResult
+    modified: ReplayResult
+    differences: list[SpanDifference] = field(default_factory=list)
+    summary: str = ""
+    semantic_similarity_available: bool = False  # Whether sentence-transformers is available
+    semantic_similarity_model: str | None = None  # Model used for similarity (if available)
+
+    @property
+    def identical_spans(self) -> int:
+        """Count spans with identical outputs."""
+        return len(self.original.spans) - len(self.differences)
+
+    @property
+    def changed_spans(self) -> int:
+        """Count spans with different outputs."""
+        return len(set(d.span_name for d in self.differences))
+
+    @property
+    def total_cost_delta(self) -> float:
+        """Calculate cost difference."""
+        return self.modified.total_cost_usd - self.original.total_cost_usd
+
+    @property
+    def total_tokens_delta(self) -> int:
+        """Calculate token usage difference."""
+        return self.modified.total_tokens - self.original.total_tokens
+
+    def generate_summary(self) -> str:
+        """Generate human-readable summary of differences."""
+        total_spans = len(self.original.spans)
+        changed = self.changed_spans
+        identical = self.identical_spans
+
+        # Calculate percentages (handle zero total_spans)
+        identical_pct = (identical / total_spans * 100) if total_spans > 0 else 0.0
+        changed_pct = (changed / total_spans * 100) if total_spans > 0 else 0.0
+
+        lines = [
+            f"Replay Comparison Summary",
+            f"=" * 50,
+            f"Total Spans: {total_spans}",
+            f"Identical: {identical} ({identical_pct:.1f}%)",
+            f"Changed: {changed} ({changed_pct:.1f}%)",
+            f"",
+            f"Cost: ${self.original.total_cost_usd:.4f} → ${self.modified.total_cost_usd:.4f} "
+            f"({'+' if self.total_cost_delta > 0 else ''}{self.total_cost_delta:.4f})",
+            f"Tokens: {self.original.total_tokens} → {self.modified.total_tokens} "
+            f"({'+' if self.total_tokens_delta > 0 else ''}{self.total_tokens_delta})",
+        ]
+
+        if self.differences:
+            lines.append("")
+            lines.append("Key Differences:")
+            for diff in self.differences[:5]:  # Show top 5
+                lines.append(f"  • {diff.span_name} ({diff.field})")
+                if diff.semantic_similarity is not None:
+                    lines.append(f"    Similarity: {diff.semantic_similarity:.2%}")
+
+        self.summary = "\n".join(lines)
+        return self.summary
+
+    def to_dict(self) -> dict[str, Any]:
+        """Convert to dictionary representation."""
+        return {
+            "original": self.original.to_dict(),
+            "modified": self.modified.to_dict(),
+            "differences": [d.to_dict() for d in self.differences],
+            "summary": self.summary or self.generate_summary(),
+            "identical_spans": self.identical_spans,
+            "changed_spans": self.changed_spans,
+            "total_cost_delta": self.total_cost_delta,
+            "total_tokens_delta": self.total_tokens_delta,
+        }