tracereplay 0.1.0__py3-none-any.whl

tracereplay/__init__.py

@@ -0,0 +1,240 @@
+"""TraceReplay: Record once. Replay reliably. Test without repeated API calls.
+
+TraceReplay is an open-source testing platform for applications powered by
+Large Language Models (LLMs) and autonomous AI agents.
+"""
+
+from __future__ import annotations
+
+from tracereplay import chaos, clock, guardrails, semantic
+from tracereplay.agent import AgentStep, AgentStepType, AgentTracker, record_tool
+from tracereplay.chaos import ChaosConfig, ChaosEngine, Fault, FaultType
+from tracereplay.clock import (
+    ClockConfig,
+    ClockState,
+    DeterministicClock,
+    freeze_time,
+)
+from tracereplay.config import TraceReplayConfig, configure, get_config, reset_config
+from tracereplay.core import (
+    Recorder,
+    Replayer,
+    TraceReplaySession,
+    get_current_session,
+    session,
+)
+from tracereplay.exceptions import (
+    BudgetExceededError,
+    ChaosContextLengthError,
+    ChaosError,
+    ChaosMalformedOutputError,
+    ChaosRateLimitError,
+    ChaosServerError,
+    ChaosTimeoutError,
+    ConfigurationError,
+    CostBudgetExceededError,
+    FixtureNotFoundError,
+    InvalidFixtureError,
+    LatencyBudgetExceededError,
+    LiveFallbackProhibitedError,
+    RecordingError,
+    RequestMismatchError,
+    SemanticDriftError,
+    SequenceMismatchError,
+    TokenBudgetExceededError,
+    ToolExecutionError,
+    TraceReplayError,
+    UnsupportedOperationError,
+)
+from tracereplay.guardrails import (
+    GuardrailAction,
+    GuardrailConfig,
+    GuardrailEngine,
+    GuardrailReport,
+    GuardrailType,
+    GuardrailViolation,
+    calculate_token_cost,
+    cost_budget,
+    latency_budget,
+    token_budget,
+)
+from tracereplay.matcher import (
+    MatchResult,
+    RequestMatcher,
+    RequestNormalizer,
+    compute_request_hash,
+    generate_diff,
+)
+from tracereplay.models import (
+    Interaction,
+    OperationType,
+    ReplayMode,
+    RequestPayload,
+    ResponsePayload,
+    SessionRecord,
+    TokenUsage,
+    ToolCallRecord,
+)
+from tracereplay.providers import (
+    AnthropicAdapter,
+    AnthropicMessageResponse,
+    BaseProviderAdapter,
+    ChatCompletionResponse,
+    DeepSeekAdapter,
+    GeminiAdapter,
+    GeminiGenerateContentResponse,
+    GenericOpenAIAdapter,
+    GroqAdapter,
+    MistralAdapter,
+    MistralChatCompletionResponse,
+    OllamaAdapter,
+    OpenAIAdapter,
+    patch,
+    unpatch,
+)
+from tracereplay.sanitizer import Sanitizer, sanitize_data
+from tracereplay.semantic import (
+    DriftAction,
+    DriftReport,
+    DriftStatus,
+    SemanticConfig,
+    SemanticDriftRecord,
+    SemanticMatcher,
+    calculate_semantic_similarity,
+    colorize_diff,
+    cosine_similarity_vectors,
+    generate_inline_token_diff,
+    generate_semantic_diff,
+    levenshtein_similarity,
+    token_cosine_similarity,
+    token_jaccard_similarity,
+    tokenize_text,
+)
+from tracereplay.storage import FixtureResolver, JsonStorageEngine, sanitize_filename
+
+__version__ = "0.1.0"
+__author__ = "TraceReplay Contributors"
+__license__ = "MIT"
+
+__all__ = [
+    # Metadata
+    "__version__",
+    "__author__",
+    "__license__",
+    # Configuration
+    "TraceReplayConfig",
+    "get_config",
+    "configure",
+    "reset_config",
+    # Core Engines & Sessions
+    "session",
+    "get_current_session",
+    "TraceReplaySession",
+    "Recorder",
+    "Replayer",
+    # Providers & Patching
+    "patch",
+    "unpatch",
+    "BaseProviderAdapter",
+    "OpenAIAdapter",
+    "AnthropicAdapter",
+    "GeminiAdapter",
+    "GroqAdapter",
+    "DeepSeekAdapter",
+    "OllamaAdapter",
+    "MistralAdapter",
+    "GenericOpenAIAdapter",
+    "ChatCompletionResponse",
+    "MistralChatCompletionResponse",
+    "AnthropicMessageResponse",
+    "GeminiGenerateContentResponse",
+    # Models & Enums
+    "ReplayMode",
+    "OperationType",
+    "TokenUsage",
+    "ToolCallRecord",
+    "RequestPayload",
+    "ResponsePayload",
+    "Interaction",
+    "SessionRecord",
+    # Exceptions
+    "TraceReplayError",
+    "ConfigurationError",
+    "FixtureNotFoundError",
+    "RequestMismatchError",
+    "InvalidFixtureError",
+    "LiveFallbackProhibitedError",
+    "UnsupportedOperationError",
+    "SequenceMismatchError",
+    "RecordingError",
+    "ToolExecutionError",
+    "ChaosError",
+    "ChaosRateLimitError",
+    "ChaosTimeoutError",
+    "ChaosServerError",
+    "ChaosContextLengthError",
+    "ChaosMalformedOutputError",
+    "BudgetExceededError",
+    "TokenBudgetExceededError",
+    "CostBudgetExceededError",
+    "LatencyBudgetExceededError",
+    "SemanticDriftError",
+    # Chaos Engineering
+    "chaos",
+    "ChaosConfig",
+    "ChaosEngine",
+    "Fault",
+    "FaultType",
+    # Deterministic Clock & Freeze Time
+    "clock",
+    "ClockConfig",
+    "ClockState",
+    "DeterministicClock",
+    "freeze_time",
+    # Guardrails
+    "guardrails",
+    "GuardrailConfig",
+    "GuardrailAction",
+    "GuardrailType",
+    "GuardrailViolation",
+    "GuardrailReport",
+    "GuardrailEngine",
+    "calculate_token_cost",
+    "token_budget",
+    "cost_budget",
+    "latency_budget",
+    # Semantic Matching & Drift Detection
+    "semantic",
+    "DriftAction",
+    "DriftReport",
+    "DriftStatus",
+    "SemanticConfig",
+    "SemanticDriftRecord",
+    "SemanticMatcher",
+    "calculate_semantic_similarity",
+    "colorize_diff",
+    "cosine_similarity_vectors",
+    "generate_inline_token_diff",
+    "generate_semantic_diff",
+    "levenshtein_similarity",
+    "token_cosine_similarity",
+    "token_jaccard_similarity",
+    "tokenize_text",
+    # Storage & Sanitizer
+    "Sanitizer",
+    "sanitize_data",
+    "JsonStorageEngine",
+    "FixtureResolver",
+    "sanitize_filename",
+    # Matcher
+    "RequestMatcher",
+    "RequestNormalizer",
+    "MatchResult",
+    "compute_request_hash",
+    "generate_diff",
+    # Agent & Tool Recording
+    "AgentTracker",
+    "AgentStep",
+    "AgentStepType",
+    "record_tool",
+]

tracereplay/__main__.py

@@ -0,0 +1,8 @@
+"""Entry point for python -m tracereplay."""
+
+from __future__ import annotations
+
+from tracereplay.cli import main
+
+if __name__ == "__main__":
+    main()

tracereplay/agent/__init__.py

@@ -0,0 +1,13 @@
+"""Multi-turn agent execution tracker and tool recording package."""
+
+from __future__ import annotations
+
+from tracereplay.agent.tool_recorder import record_tool
+from tracereplay.agent.tracker import AgentStep, AgentStepType, AgentTracker
+
+__all__ = [
+    "AgentStep",
+    "AgentStepType",
+    "AgentTracker",
+    "record_tool",
+]

tracereplay/agent/tool_recorder.py

@@ -0,0 +1,213 @@
+"""Decorator and capture utilities for recording and replaying local Python tools."""
+
+from __future__ import annotations
+
+import asyncio
+import functools
+import inspect
+import time
+from collections.abc import Callable
+from typing import Any, TypeVar, cast, overload
+
+from tracereplay.core.session import get_current_session
+from tracereplay.models import ReplayMode, ToolCallRecord
+
+F = TypeVar("F", bound=Callable[..., Any])
+
+
+@overload
+def record_tool(func: F) -> F: ...
+
+
+@overload
+def record_tool(
+    func: None = None,
+    *,
+    name: str | None = None,
+    ignore_args: list[str] | None = None,
+) -> Callable[[F], F]: ...
+
+
+def record_tool(
+    func: F | None = None,
+    *,
+    name: str | None = None,
+    ignore_args: list[str] | None = None,
+) -> F | Callable[[F], F]:
+    """Decorator to intercept, record, and re-execute local Python tool functions.
+
+    In RECORD mode:
+        Executes the genuine Python function, captures arguments and return value,
+        and persists the execution to the session fixture.
+
+    In REPLAY mode:
+        Bypasses the actual Python function completely, verifying sequential
+        arguments and returning the deterministic recorded output.
+
+    Usage:
+        @record_tool
+        def calculate_tax(amount: float, rate: float = 0.05) -> float:
+            return amount * rate
+
+        @record_tool(name="fetch_user_data", ignore_args=["db_session"])
+        async def fetch_user(user_id: str, db_session: Any = None) -> dict:
+            return await db.query(user_id)
+    """
+
+    def decorator(fn: F) -> F:
+        tool_name = name or fn.__name__
+        ignored = set(ignore_args or [])
+        sig = inspect.signature(fn)
+
+        def _extract_args(args: tuple[Any, ...], kwargs: dict[str, Any]) -> dict[str, Any]:
+            bound = sig.bind(*args, **kwargs)
+            bound.apply_defaults()
+            extracted: dict[str, Any] = {}
+            for k, v in bound.arguments.items():
+                if k not in ignored:
+                    extracted[k] = v
+            return extracted
+
+        if asyncio.iscoroutinefunction(fn):
+
+            @functools.wraps(fn)
+            async def async_wrapper(*args: Any, **kwargs: Any) -> Any:
+                sess = get_current_session()
+                if sess is not None:
+                    sess.check_fork_trigger(tool_name=tool_name)
+                if sess is None or sess.mode == ReplayMode.PASSTHROUGH:
+                    res = await fn(*args, **kwargs)
+                    if sess is not None:
+                        sess.notify_tool_call_completed()
+                    return res
+
+                call_args = _extract_args(args, kwargs)
+
+                # AI Chaos Engineering check
+                if sess is not None and sess.chaos_engine is not None:
+                    await sess.chaos_engine.aevaluate_tool_call(
+                        tool_name=tool_name,
+                        arguments=call_args,
+                        step_index=sess._tool_step_count,
+                    )
+
+                if sess.mode == ReplayMode.REPLAY:
+                    start = time.perf_counter()
+                    res = await sess.areplay_tool_call(tool_name, call_args)
+                    duration_ms = (time.perf_counter() - start) * 1000.0
+                    if sess.guardrail_engine is not None:
+                        sess.guardrail_engine.evaluate_tool_execution(
+                            tool_name=tool_name,
+                            duration_ms=duration_ms,
+                            step_index=sess._total_step_count,
+                        )
+                    sess.notify_tool_call_completed()
+                    return res
+
+                # RECORD mode
+                start = time.perf_counter()
+                try:
+                    res = await fn(*args, **kwargs)
+                    duration_ms = (time.perf_counter() - start) * 1000.0
+                    record = ToolCallRecord(
+                        function_name=tool_name,
+                        arguments=call_args,
+                        result=res,
+                        duration_ms=duration_ms,
+                    )
+                    sess.record_tool_call(record)
+                    if sess.guardrail_engine is not None:
+                        sess.guardrail_engine.evaluate_tool_execution(
+                            tool_name=tool_name,
+                            duration_ms=duration_ms,
+                            step_index=sess._total_step_count,
+                        )
+                    sess.notify_tool_call_completed()
+                    return res
+                except Exception as err:
+                    duration_ms = (time.perf_counter() - start) * 1000.0
+                    record = ToolCallRecord(
+                        function_name=tool_name,
+                        arguments=call_args,
+                        result=None,
+                        duration_ms=duration_ms,
+                        error=str(err),
+                    )
+                    sess.record_tool_call(record)
+                    raise
+
+            return cast(F, async_wrapper)
+
+        else:
+
+            @functools.wraps(fn)
+            def sync_wrapper(*args: Any, **kwargs: Any) -> Any:
+                sess = get_current_session()
+                if sess is not None:
+                    sess.check_fork_trigger(tool_name=tool_name)
+                if sess is None or sess.mode == ReplayMode.PASSTHROUGH:
+                    res = fn(*args, **kwargs)
+                    if sess is not None:
+                        sess.notify_tool_call_completed()
+                    return res
+
+                call_args = _extract_args(args, kwargs)
+
+                # AI Chaos Engineering check
+                if sess is not None and sess.chaos_engine is not None:
+                    sess.chaos_engine.evaluate_tool_call(
+                        tool_name=tool_name,
+                        arguments=call_args,
+                        step_index=sess._tool_step_count,
+                    )
+
+                if sess.mode == ReplayMode.REPLAY:
+                    start = time.perf_counter()
+                    res = sess.replay_tool_call(tool_name, call_args)
+                    duration_ms = (time.perf_counter() - start) * 1000.0
+                    if sess.guardrail_engine is not None:
+                        sess.guardrail_engine.evaluate_tool_execution(
+                            tool_name=tool_name,
+                            duration_ms=duration_ms,
+                            step_index=sess._total_step_count,
+                        )
+                    sess.notify_tool_call_completed()
+                    return res
+
+                # RECORD mode
+                start = time.perf_counter()
+                try:
+                    res = fn(*args, **kwargs)
+                    duration_ms = (time.perf_counter() - start) * 1000.0
+                    record = ToolCallRecord(
+                        function_name=tool_name,
+                        arguments=call_args,
+                        result=res,
+                        duration_ms=duration_ms,
+                    )
+                    sess.record_tool_call(record)
+                    if sess.guardrail_engine is not None:
+                        sess.guardrail_engine.evaluate_tool_execution(
+                            tool_name=tool_name,
+                            duration_ms=duration_ms,
+                            step_index=sess._total_step_count,
+                        )
+                    sess.notify_tool_call_completed()
+                    return res
+                except Exception as err:
+                    duration_ms = (time.perf_counter() - start) * 1000.0
+                    record = ToolCallRecord(
+                        function_name=tool_name,
+                        arguments=call_args,
+                        result=None,
+                        duration_ms=duration_ms,
+                        error=str(err),
+                    )
+                    sess.record_tool_call(record)
+                    raise
+
+            return cast(F, sync_wrapper)
+
+    if func is not None:
+        return decorator(func)
+    return decorator

tracereplay/agent/tracker.py

@@ -0,0 +1,192 @@
+"""Agent execution tracker modeling multi-turn decisions and tool execution flows."""
+
+from __future__ import annotations
+
+from datetime import datetime, timezone
+from enum import Enum
+from typing import Any
+
+from pydantic import BaseModel, ConfigDict, Field
+
+
+def _utc_now() -> datetime:
+    return datetime.now(timezone.utc)
+
+
+class AgentStepType(str, Enum):
+    """Classification of an agent execution step in a multi-turn graph."""
+
+    USER_INPUT = "user_input"
+    LLM_DECISION = "llm_decision"
+    TOOL_EXECUTION = "tool_execution"
+    TOOL_RESULT = "tool_result"
+    FINAL_ANSWER = "final_answer"
+
+
+class AgentStep(BaseModel):
+    """Single transition node within an agent reasoning and execution graph."""
+
+    model_config = ConfigDict(frozen=True)
+
+    step_index: int
+    step_type: AgentStepType
+    name: str
+    data: dict[str, Any] = Field(default_factory=dict)
+    duration_ms: float = 0.0
+    timestamp: datetime = Field(default_factory=_utc_now)
+
+
+class AgentTracker:
+    """Models and validates multi-turn autonomous AI agent execution graphs.
+
+    Graph flow:
+        User Input -> LLM Decision -> Tool Execution -> Tool Result -> Next Decision -> Final Answer
+    """
+
+    def __init__(self, agent_name: str = "default_agent") -> None:
+        self.agent_name = agent_name
+        self.steps: list[AgentStep] = []
+
+    @property
+    def total_steps(self) -> int:
+        return len(self.steps)
+
+    def add_step(
+        self,
+        step_type: AgentStepType,
+        name: str,
+        data: dict[str, Any] | None = None,
+        duration_ms: float = 0.0,
+    ) -> AgentStep:
+        """Append a new execution step to the graph."""
+        step = AgentStep(
+            step_index=len(self.steps),
+            step_type=step_type,
+            name=name,
+            data=data or {},
+            duration_ms=duration_ms,
+        )
+        self.steps.append(step)
+        return step
+
+    def add_user_input(
+        self,
+        prompt: str,
+        metadata: dict[str, Any] | None = None,
+    ) -> AgentStep:
+        """Record the initial or follow-up user prompt initiating a reasoning turn."""
+        data: dict[str, Any] = {"prompt": prompt}
+        if metadata:
+            data.update(metadata)
+        return self.add_step(AgentStepType.USER_INPUT, name="user_input", data=data)
+
+    def add_llm_decision(
+        self,
+        decision: Any,
+        tool_calls: list[Any] | None = None,
+        metadata: dict[str, Any] | None = None,
+    ) -> AgentStep:
+        """Record an LLM reasoning output or tool invocation decision."""
+        data: dict[str, Any] = {"decision": str(decision)}
+        if tool_calls is not None:
+            data["tool_calls"] = tool_calls
+        if metadata:
+            data.update(metadata)
+        return self.add_step(AgentStepType.LLM_DECISION, name="llm_decision", data=data)
+
+    def add_tool_execution(
+        self,
+        tool_name: str,
+        arguments: dict[str, Any],
+        metadata: dict[str, Any] | None = None,
+    ) -> AgentStep:
+        """Record the invocation of an external or local tool with arguments."""
+        data: dict[str, Any] = {"tool_name": tool_name, "arguments": arguments}
+        if metadata:
+            data.update(metadata)
+        return self.add_step(
+            AgentStepType.TOOL_EXECUTION,
+            name=f"execute_{tool_name}",
+            data=data,
+        )
+
+    def add_tool_result(
+        self,
+        tool_name: str,
+        result: Any,
+        duration_ms: float = 0.0,
+        error: str | None = None,
+        metadata: dict[str, Any] | None = None,
+    ) -> AgentStep:
+        """Record the return value or failure of a completed tool execution."""
+        data: dict[str, Any] = {"tool_name": tool_name, "result": result}
+        if error:
+            data["error"] = error
+        if metadata:
+            data.update(metadata)
+        return self.add_step(
+            AgentStepType.TOOL_RESULT,
+            name=f"result_{tool_name}",
+            data=data,
+            duration_ms=duration_ms,
+        )
+
+    def add_final_answer(
+        self,
+        answer: str,
+        metadata: dict[str, Any] | None = None,
+    ) -> AgentStep:
+        """Record the agent's completed final response returned to the user."""
+        data: dict[str, Any] = {"answer": answer}
+        if metadata:
+            data.update(metadata)
+        return self.add_step(AgentStepType.FINAL_ANSWER, name="final_answer", data=data)
+
+    def visualize_flow(self) -> str:
+        """Generate a human-readable diagram of the agent's execution flow."""
+        if not self.steps:
+            return f"Agent Flow [{self.agent_name}]: (empty)"
+
+        lines: list[str] = [f"Agent Flow [{self.agent_name}]:"]
+        for i, step in enumerate(self.steps):
+            detail = ""
+            if step.step_type == AgentStepType.USER_INPUT:
+                prompt_preview = str(step.data.get("prompt", ""))[:40]
+                detail = f'prompt="{prompt_preview}"'
+            elif step.step_type == AgentStepType.LLM_DECISION:
+                calls = step.data.get("tool_calls")
+                detail = f"invokes {len(calls)} tool(s)" if calls else "generates response"
+            elif step.step_type == AgentStepType.TOOL_EXECUTION:
+                tname = step.data.get("tool_name", "")
+                args = step.data.get("arguments", {})
+                detail = f"{tname}({args})"
+            elif step.step_type == AgentStepType.TOOL_RESULT:
+                tname = step.data.get("tool_name", "")
+                res = str(step.data.get("result", ""))[:30]
+                detail = f"{tname} -> {res} ({step.duration_ms:.1f}ms)"
+            elif step.step_type == AgentStepType.FINAL_ANSWER:
+                ans_preview = str(step.data.get("answer", ""))[:40]
+                detail = f'answer="{ans_preview}"'
+
+            step_line = f"  [{step.step_index}] {step.step_type.value.upper()}: {detail}"
+            lines.append(step_line)
+            if i < len(self.steps) - 1:
+                lines.append("       |")
+                lines.append("       v")
+
+        return "\n".join(lines)
+
+    def to_dict(self) -> dict[str, Any]:
+        """Serialize the complete agent execution trace to dictionary format."""
+        return {
+            "agent_name": self.agent_name,
+            "steps": [step.model_dump(mode="json") for step in self.steps],
+        }
+
+    @classmethod
+    def from_dict(cls, data: dict[str, Any]) -> AgentTracker:
+        """Reconstruct an AgentTracker instance from serialized dictionary data."""
+        tracker = cls(agent_name=data.get("agent_name", "default_agent"))
+        for s in data.get("steps", []):
+            tracker.steps.append(AgentStep.model_validate(s))
+        return tracker