tactus 0.31.2__py3-none-any.whl

tactus/tracing/trace_manager.py

@@ -0,0 +1,417 @@
+"""
+Trace Manager - API for managing execution traces and debugging sessions.
+
+Provides operations for querying, filtering, and analyzing procedure execution traces.
+"""
+
+import logging
+from typing import Optional, List, Dict, Any
+from pathlib import Path
+
+from tactus.protocols.storage import StorageBackend
+from tactus.protocols.models import ExecutionRun, CheckpointEntry, Breakpoint
+
+logger = logging.getLogger(__name__)
+
+
+class TraceManager:
+    """
+    Manages execution traces and debugging sessions.
+
+    Provides API for:
+    - Listing and querying execution runs
+    - Accessing checkpoint data
+    - Managing breakpoints
+    - Comparing runs
+    - Exporting traces
+    """
+
+    def __init__(self, storage: StorageBackend):
+        """
+        Initialize trace manager.
+
+        Args:
+            storage: Storage backend for trace persistence
+        """
+        self.storage = storage
+        logger.info("TraceManager initialized")
+
+    # Run Management
+
+    def list_runs(
+        self, procedure_name: Optional[str] = None, limit: Optional[int] = None
+    ) -> List[ExecutionRun]:
+        """
+        List all execution runs, optionally filtered by procedure name.
+
+        Args:
+            procedure_name: Optional procedure name filter
+            limit: Optional limit on number of runs returned
+
+        Returns:
+            List of execution runs, sorted by start time (newest first)
+        """
+        logger.debug(f"Listing runs (procedure={procedure_name}, limit={limit})")
+
+        runs = self.storage.list_runs(procedure_name=procedure_name)
+
+        if limit:
+            runs = runs[:limit]
+
+        return runs
+
+    def get_run(self, run_id: str) -> ExecutionRun:
+        """
+        Get complete run data.
+
+        Args:
+            run_id: Run identifier
+
+        Returns:
+            Complete execution run with all checkpoints
+
+        Raises:
+            FileNotFoundError: If run not found
+        """
+        logger.debug(f"Getting run {run_id}")
+        return self.storage.load_run(run_id)
+
+    def get_checkpoint(self, run_id: str, position: int) -> CheckpointEntry:
+        """
+        Get specific checkpoint from a run.
+
+        Args:
+            run_id: Run identifier
+            position: Checkpoint position (0-indexed)
+
+        Returns:
+            Checkpoint entry
+
+        Raises:
+            FileNotFoundError: If run not found
+            IndexError: If position out of range
+        """
+        logger.debug(f"Getting checkpoint {position} from run {run_id}")
+
+        run = self.get_run(run_id)
+
+        if position < 0 or position >= len(run.execution_log):
+            raise IndexError(
+                f"Checkpoint position {position} out of range (0-{len(run.execution_log) - 1})"
+            )
+
+        return run.execution_log[position]
+
+    def get_checkpoints(
+        self, run_id: str, start: int = 0, end: Optional[int] = None
+    ) -> List[CheckpointEntry]:
+        """
+        Get range of checkpoints from a run.
+
+        Args:
+            run_id: Run identifier
+            start: Start position (inclusive, 0-indexed)
+            end: End position (exclusive, None = end of log)
+
+        Returns:
+            List of checkpoint entries
+        """
+        logger.debug(f"Getting checkpoints {start}:{end} from run {run_id}")
+
+        run = self.get_run(run_id)
+
+        if end is None:
+            return run.execution_log[start:]
+        else:
+            return run.execution_log[start:end]
+
+    # Breakpoint Management
+
+    def set_breakpoint(self, file: str, line: int, condition: Optional[str] = None) -> Breakpoint:
+        """
+        Set a breakpoint at file:line.
+
+        Args:
+            file: File path
+            line: Line number (1-indexed)
+            condition: Optional Python expression to evaluate
+
+        Returns:
+            Created breakpoint
+        """
+        import uuid
+
+        logger.info(f"Setting breakpoint at {file}:{line}")
+
+        breakpoint = Breakpoint(
+            breakpoint_id=str(uuid.uuid4()),
+            file=file,
+            line=line,
+            condition=condition,
+            enabled=True,
+            hit_count=0,
+        )
+
+        # Load existing breakpoints for this file
+        # Extract procedure name from file path
+        procedure_name = Path(file).stem
+        breakpoints = self.storage.load_breakpoints(procedure_name)
+
+        # Add new breakpoint
+        breakpoints.append(breakpoint)
+
+        # Save updated list
+        self.storage.save_breakpoints(procedure_name, breakpoints)
+
+        return breakpoint
+
+    def remove_breakpoint(self, breakpoint_id: str) -> None:
+        """
+        Remove a breakpoint.
+
+        Args:
+            breakpoint_id: Breakpoint identifier
+        """
+        logger.info(f"Removing breakpoint {breakpoint_id}")
+
+        # We need to search all breakpoint files
+        # For now, this is a simple implementation that loads all
+        # TODO: Optimize with an index if this becomes a bottleneck
+
+        # This is a placeholder - in practice we'd need to know which file
+        # For now, we'll need to search or maintain an index
+        raise NotImplementedError("Remove breakpoint requires file index")
+
+    def list_breakpoints(self, file: Optional[str] = None) -> List[Breakpoint]:
+        """
+        List all breakpoints, optionally filtered by file.
+
+        Args:
+            file: Optional file path filter
+
+        Returns:
+            List of breakpoints
+        """
+        logger.debug(f"Listing breakpoints (file={file})")
+
+        if file:
+            procedure_name = Path(file).stem
+            return self.storage.load_breakpoints(procedure_name)
+        else:
+            # List all breakpoints across all procedures
+            # This requires iterating through breakpoint directory
+            # For now, require a file parameter
+            raise NotImplementedError("Listing all breakpoints requires file parameter")
+
+    def toggle_breakpoint(self, breakpoint_id: str, enabled: bool) -> None:
+        """
+        Enable or disable a breakpoint.
+
+        Args:
+            breakpoint_id: Breakpoint identifier
+            enabled: Whether to enable or disable
+        """
+        logger.info(f"Toggling breakpoint {breakpoint_id} to {enabled}")
+
+        # Similar to remove_breakpoint, we need an index or file reference
+        raise NotImplementedError("Toggle breakpoint requires file index")
+
+    # Query/Analysis
+
+    def find_checkpoint_after_line(
+        self, run_id: str, file: str, line: int
+    ) -> Optional[CheckpointEntry]:
+        """
+        Find nearest checkpoint after specified line.
+
+        This is used for breakpoint mapping: when a user sets a breakpoint at
+        a specific line, we find the next checkpoint that will be hit.
+
+        Args:
+            run_id: Run identifier
+            file: File path
+            line: Line number (1-indexed)
+
+        Returns:
+            Next checkpoint after the line, or None if no checkpoint found
+        """
+        logger.debug(f"Finding checkpoint after {file}:{line} in run {run_id}")
+
+        run = self.get_run(run_id)
+
+        # Search for first checkpoint with source location >= line
+        for checkpoint in run.execution_log:
+            if checkpoint.source_location:
+                if (
+                    checkpoint.source_location.file == file
+                    and checkpoint.source_location.line >= line
+                ):
+                    return checkpoint
+
+        return None
+
+    def find_checkpoints_by_type(self, run_id: str, checkpoint_type: str) -> List[CheckpointEntry]:
+        """
+        Find all checkpoints of a specific type.
+
+        Args:
+            run_id: Run identifier
+            checkpoint_type: Checkpoint type (agent_turn, model_predict, etc.)
+
+        Returns:
+            List of matching checkpoints
+        """
+        logger.debug(f"Finding checkpoints of type '{checkpoint_type}' in run {run_id}")
+
+        run = self.get_run(run_id)
+
+        return [cp for cp in run.execution_log if cp.type == checkpoint_type]
+
+    def compare_runs(self, run_id1: str, run_id2: str) -> Dict[str, Any]:
+        """
+        Compare two runs for debugging non-determinism.
+
+        Args:
+            run_id1: First run identifier
+            run_id2: Second run identifier
+
+        Returns:
+            Comparison results with differences
+        """
+        logger.info(f"Comparing runs {run_id1} vs {run_id2}")
+
+        run1 = self.get_run(run_id1)
+        run2 = self.get_run(run_id2)
+
+        comparison = {
+            "run1": {
+                "id": run_id1,
+                "procedure": run1.procedure_name,
+                "status": run1.status,
+                "checkpoint_count": len(run1.execution_log),
+            },
+            "run2": {
+                "id": run_id2,
+                "procedure": run2.procedure_name,
+                "status": run2.status,
+                "checkpoint_count": len(run2.execution_log),
+            },
+            "differences": [],
+        }
+
+        # Compare checkpoint counts
+        if len(run1.execution_log) != len(run2.execution_log):
+            comparison["differences"].append(
+                {
+                    "type": "checkpoint_count_mismatch",
+                    "run1_count": len(run1.execution_log),
+                    "run2_count": len(run2.execution_log),
+                }
+            )
+
+        # Compare checkpoints position by position
+        for i in range(min(len(run1.execution_log), len(run2.execution_log))):
+            cp1 = run1.execution_log[i]
+            cp2 = run2.execution_log[i]
+
+            # Compare types
+            if cp1.type != cp2.type:
+                comparison["differences"].append(
+                    {
+                        "type": "checkpoint_type_mismatch",
+                        "position": i,
+                        "run1_type": cp1.type,
+                        "run2_type": cp2.type,
+                    }
+                )
+
+            # Compare source locations
+            if cp1.source_location and cp2.source_location:
+                if cp1.source_location.line != cp2.source_location.line:
+                    comparison["differences"].append(
+                        {
+                            "type": "source_location_mismatch",
+                            "position": i,
+                            "run1_line": cp1.source_location.line,
+                            "run2_line": cp2.source_location.line,
+                        }
+                    )
+
+            # Compare results (simple equality check)
+            if cp1.result != cp2.result:
+                comparison["differences"].append(
+                    {
+                        "type": "result_mismatch",
+                        "position": i,
+                        "checkpoint_type": cp1.type,
+                    }
+                )
+
+        return comparison
+
+    def export_trace(self, run_id: str, format: str = "json") -> str:
+        """
+        Export trace for external analysis.
+
+        Args:
+            run_id: Run identifier
+            format: Export format (json, csv, etc.)
+
+        Returns:
+            Exported trace data as string
+
+        Raises:
+            ValueError: If format not supported
+        """
+        logger.info(f"Exporting run {run_id} as {format}")
+
+        run = self.get_run(run_id)
+
+        if format == "json":
+            import json
+
+            return json.dumps(run.model_dump(), indent=2, default=str)
+        else:
+            raise ValueError(f"Unsupported export format: {format}")
+
+    def get_statistics(self, run_id: str) -> Dict[str, Any]:
+        """
+        Get statistics about a run.
+
+        Args:
+            run_id: Run identifier
+
+        Returns:
+            Statistics dictionary with checkpoint counts, timing, etc.
+        """
+        logger.debug(f"Getting statistics for run {run_id}")
+
+        run = self.get_run(run_id)
+
+        # Count checkpoints by type
+        type_counts: Dict[str, int] = {}
+        total_duration_ms = 0.0
+
+        for checkpoint in run.execution_log:
+            type_counts[checkpoint.type] = type_counts.get(checkpoint.type, 0) + 1
+            if checkpoint.duration_ms:
+                total_duration_ms += checkpoint.duration_ms
+
+        # Calculate timing stats
+        if run.start_time and run.end_time:
+            total_time_sec = (run.end_time - run.start_time).total_seconds()
+        else:
+            total_time_sec = None
+
+        return {
+            "run_id": run_id,
+            "procedure": run.procedure_name,
+            "status": run.status,
+            "total_checkpoints": len(run.execution_log),
+            "checkpoints_by_type": type_counts,
+            "total_duration_ms": total_duration_ms,
+            "total_time_sec": total_time_sec,
+            "has_source_locations": sum(
+                1 for cp in run.execution_log if cp.source_location is not None
+            ),
+        }

tactus/utils/__init__.py

@@ -0,0 +1 @@
+"""Utility modules for Tactus."""

tactus/utils/cost_calculator.py

@@ -0,0 +1,72 @@
+"""
+Cost calculator for LLM usage.
+
+Calculates costs based on token usage and model pricing.
+"""
+
+from typing import Dict, Any, Optional
+from .model_pricing import get_model_pricing, normalize_model_name
+
+
+class CostCalculator:
+    """
+    Calculate LLM costs from token usage and model information.
+
+    Aligned with pydantic-ai's usage tracking.
+    """
+
+    def calculate_cost(
+        self,
+        model_name: str,
+        provider: Optional[str],
+        prompt_tokens: int,
+        completion_tokens: int,
+        cache_tokens: Optional[int] = None,
+    ) -> Dict[str, Any]:
+        """
+        Calculate cost for a single LLM call.
+
+        Args:
+            model_name: Model identifier
+            provider: Provider name (openai, anthropic, bedrock, google)
+            prompt_tokens: Number of prompt tokens
+            completion_tokens: Number of completion tokens
+            cache_tokens: Number of cached tokens (if applicable)
+
+        Returns:
+            Dict with:
+                - prompt_cost: Cost for prompt tokens
+                - completion_cost: Cost for completion tokens
+                - cache_cost: Cost savings from cache (if applicable)
+                - total_cost: Total cost
+                - model: Normalized model name
+                - provider: Detected provider
+                - pricing_found: Whether pricing was found (False = using defaults)
+        """
+        # Normalize model name and get provider
+        normalized_model, detected_provider = normalize_model_name(model_name, provider)
+
+        # Get pricing
+        pricing = get_model_pricing(model_name, provider)
+
+        # Calculate costs (pricing is per million tokens)
+        prompt_cost = (prompt_tokens / 1_000_000) * pricing["input"]
+        completion_cost = (completion_tokens / 1_000_000) * pricing["output"]
+
+        # Calculate cache savings if applicable
+        cache_cost = None
+        if cache_tokens and cache_tokens > 0:
+            # Cached tokens typically cost 10% of input tokens
+            cache_cost = (cache_tokens / 1_000_000) * pricing["input"] * 0.9
+
+        total_cost = prompt_cost + completion_cost
+
+        return {
+            "prompt_cost": prompt_cost,
+            "completion_cost": completion_cost,
+            "cache_cost": cache_cost,
+            "total_cost": total_cost,
+            "model": normalized_model,
+            "provider": detected_provider,
+            "pricing_found": True,  # Could track if we used DEFAULT_PRICING
+        }

tactus/utils/model_pricing.py

@@ -0,0 +1,132 @@
+"""
+Model pricing data for cost calculation.
+
+Prices are per million tokens in USD.
+Data sourced from provider documentation and pricing pages.
+"""
+
+from typing import Dict, Optional
+
+# Pricing per million tokens (USD)
+MODEL_PRICING: Dict[str, Dict[str, Dict[str, float]]] = {
+    "openai": {
+        "gpt-4o": {"input": 2.50, "output": 10.00},
+        "gpt-4o-mini": {"input": 0.15, "output": 0.60},
+        "gpt-4o-2024-11-20": {"input": 2.50, "output": 10.00},
+        "gpt-4o-2024-08-06": {"input": 2.50, "output": 10.00},
+        "gpt-4o-2024-05-13": {"input": 5.00, "output": 15.00},
+        "gpt-4-turbo": {"input": 10.00, "output": 30.00},
+        "gpt-4-turbo-preview": {"input": 10.00, "output": 30.00},
+        "gpt-4": {"input": 30.00, "output": 60.00},
+        "gpt-3.5-turbo": {"input": 0.50, "output": 1.50},
+        "o1": {"input": 15.00, "output": 60.00},
+        "o1-mini": {"input": 1.10, "output": 4.40},
+        "o3-mini": {"input": 1.10, "output": 4.40},
+    },
+    "anthropic": {
+        "claude-3-5-sonnet": {"input": 3.00, "output": 15.00},
+        "claude-3-5-sonnet-20241022": {"input": 3.00, "output": 15.00},
+        "claude-3-5-haiku": {"input": 0.80, "output": 4.00},
+        "claude-3-5-haiku-20241022": {"input": 0.80, "output": 4.00},
+        "claude-3-opus": {"input": 15.00, "output": 75.00},
+        "claude-3-sonnet": {"input": 3.00, "output": 15.00},
+        "claude-3-haiku": {"input": 0.25, "output": 1.25},
+        "claude-sonnet-4": {"input": 3.00, "output": 15.00},
+        "claude-sonnet-4.5": {"input": 3.00, "output": 15.00},
+        "claude-opus-4": {"input": 15.00, "output": 75.00},
+    },
+    "bedrock": {
+        # Anthropic models on Bedrock
+        "anthropic.claude-3-5-sonnet-20241022-v2:0": {"input": 3.00, "output": 15.00},
+        "anthropic.claude-3-5-sonnet-20240620-v1:0": {"input": 3.00, "output": 15.00},
+        "anthropic.claude-3-5-haiku-20241022-v1:0": {"input": 0.80, "output": 4.00},
+        "anthropic.claude-3-opus-20240229-v1:0": {"input": 15.00, "output": 75.00},
+        "anthropic.claude-3-sonnet-20240229-v1:0": {"input": 3.00, "output": 15.00},
+        "anthropic.claude-3-haiku-20240307-v1:0": {"input": 0.25, "output": 1.25},
+    },
+    "google": {
+        "gemini-2.5-flash": {"input": 0.30, "output": 2.50},
+        "gemini-2.5-pro": {"input": 1.25, "output": 10.00},
+        "gemini-2.0-flash": {"input": 0.10, "output": 0.40},
+        "gemini-1.5-pro": {"input": 1.25, "output": 5.00},
+        "gemini-1.5-flash": {"input": 0.075, "output": 0.30},
+    },
+}
+
+# Default pricing for unknown models (conservative estimate)
+DEFAULT_PRICING = {"input": 10.00, "output": 30.00}
+
+
+def normalize_model_name(model_name: str, provider: Optional[str] = None) -> tuple[str, str]:
+    """
+    Normalize model name and extract provider.
+
+    Handles formats like:
+    - "gpt-4o" -> ("gpt-4o", "openai")
+    - "openai:gpt-4o" -> ("gpt-4o", "openai")
+    - "anthropic.claude-3-5-sonnet-20241022-v2:0" -> (full name, "bedrock")
+
+    Args:
+        model_name: Model identifier
+        provider: Optional provider hint
+
+    Returns:
+        Tuple of (normalized_model_name, provider)
+    """
+    # Extract provider from model name if present
+    if ":" in model_name:
+        parts = model_name.split(":", 1)
+        detected_provider = parts[0].lower()
+        model_only = parts[1]
+        return (model_only, detected_provider)
+
+    # Check for Bedrock format (anthropic.claude-...)
+    if model_name.startswith("anthropic."):
+        return (model_name, "bedrock")
+
+    # Use provided provider or try to infer
+    if provider:
+        return (model_name, provider.lower())
+
+    # Infer provider from model name patterns
+    if model_name.startswith("gpt-") or model_name.startswith("o1") or model_name.startswith("o3"):
+        return (model_name, "openai")
+    elif model_name.startswith("claude-"):
+        return (model_name, "anthropic")
+    elif model_name.startswith("gemini-"):
+        return (model_name, "google")
+
+    # Default to openai if unknown
+    return (model_name, provider or "openai")
+
+
+def get_model_pricing(model_name: str, provider: Optional[str] = None) -> Dict[str, float]:
+    """
+    Get pricing for a model.
+
+    Args:
+        model_name: Model identifier
+        provider: Optional provider
+
+    Returns:
+        Dict with 'input' and 'output' pricing per million tokens
+    """
+    normalized_model, detected_provider = normalize_model_name(model_name, provider)
+
+    # Look up pricing
+    provider_pricing = MODEL_PRICING.get(detected_provider, {})
+    pricing = provider_pricing.get(normalized_model)
+
+    if pricing:
+        return pricing
+
+    # Try without version suffix (e.g., "gpt-4o-2024-11-20" -> "gpt-4o")
+    base_model = normalized_model.split("-")[0:2]  # Get first two parts
+    if len(base_model) >= 2:
+        base_name = "-".join(base_model)
+        pricing = provider_pricing.get(base_name)
+        if pricing:
+            return pricing
+
+    # Return default pricing with warning
+    return DEFAULT_PRICING