arcade-mcp/libs/arcade-evals/arcade_evals/capture.py

"""
Capture mode for EvalSuite.

Capture mode runs evaluation cases and records tool calls from the model
without scoring or evaluating them. This is useful for:
- Generating expected tool calls for new test cases
- Debugging model behavior
- Creating baseline recordings
"""

from __future__ import annotations

import json
from dataclasses import dataclass, field
from typing import TYPE_CHECKING, Any

from openai import AsyncOpenAI

if TYPE_CHECKING:
    from arcade_evals.eval import EvalSuite


@dataclass
class CapturedToolCall:
    """
    A captured tool call from the model during capture mode.

    Attributes:
        name: The name of the tool that was called.
        args: The arguments passed to the tool.
    """

    name: str
    args: dict[str, Any] = field(default_factory=dict)

    def to_dict(self) -> dict[str, Any]:
        """Convert to dictionary for JSON serialization."""
        return {"name": self.name, "args": self.args}


@dataclass
class CapturedRun:
    """
    A single capture run for a case, containing tool calls.

    Attributes:
        tool_calls: List of tool calls made by the model in this run.
    """

    tool_calls: list[CapturedToolCall] = field(default_factory=list)

    def to_dict(self) -> dict[str, Any]:
        """Convert to dictionary for JSON serialization."""
        return {"tool_calls": [tc.to_dict() for tc in self.tool_calls]}


@dataclass
class CapturedCase:
    """
    Result of running a single case in capture mode.

    Attributes:
        case_name: The name of the evaluation case.
        user_message: The user message that triggered the tool calls.
        tool_calls: List of tool calls made by the model.
        system_message: The system message (included if include_context is True).
        additional_messages: Additional messages (included if include_context is True).
        track_name: The track name for comparative captures (None for regular cases).
        runs: Optional list of runs (populated when num_runs > 1).
    """

    case_name: str
    user_message: str
    tool_calls: list[CapturedToolCall] = field(default_factory=list)
    system_message: str | None = None
    additional_messages: list[dict[str, Any]] | None = None
    track_name: str | None = None
    runs: list[CapturedRun] = field(default_factory=list)

    @staticmethod
    def _try_parse_json(value: str) -> Any:
        """Try to parse a JSON string, returning the original string if parsing fails."""
        try:
            return json.loads(value)
        except json.JSONDecodeError:
            return value

    @staticmethod
    def _normalize_messages(messages: list[dict[str, Any]]) -> list[dict[str, Any]]:
        """
        Normalize additional_messages by parsing JSON strings into proper objects.

        OpenAI returns:
        - Tool call arguments as JSON strings in assistant messages
        - Tool response content as JSON strings in tool messages

        For cleaner output, we parse these into proper objects.
        """
        normalized = []
        for msg in messages:
            msg_copy = dict(msg)

            # Parse tool call arguments in assistant messages
            if "tool_calls" in msg_copy and isinstance(msg_copy["tool_calls"], list):
                normalized_tool_calls = []
                for tc in msg_copy["tool_calls"]:
                    tc_copy = dict(tc)
                    if "function" in tc_copy and isinstance(tc_copy["function"], dict):
                        func = dict(tc_copy["function"])
                        if "arguments" in func and isinstance(func["arguments"], str):
                            func["arguments"] = CapturedCase._try_parse_json(func["arguments"])
                        tc_copy["function"] = func
                    normalized_tool_calls.append(tc_copy)
                msg_copy["tool_calls"] = normalized_tool_calls

            # Parse content in tool response messages
            if msg_copy.get("role") == "tool" and isinstance(msg_copy.get("content"), str):
                msg_copy["content"] = CapturedCase._try_parse_json(msg_copy["content"])

            normalized.append(msg_copy)
        return normalized

    def to_dict(self, include_context: bool = False) -> dict[str, Any]:
        """Convert to dictionary for JSON serialization."""
        result: dict[str, Any] = {
            "case_name": self.case_name,
            "user_message": self.user_message,
            "tool_calls": [tc.to_dict() for tc in self.tool_calls],
        }
        if self.runs:
            result["runs"] = [run.to_dict() for run in self.runs]
        if self.track_name:
            result["track_name"] = self.track_name
        if include_context:
            result["system_message"] = self.system_message
            # Normalize additional_messages to parse JSON string arguments
            raw_messages = self.additional_messages or []
            result["additional_messages"] = self._normalize_messages(raw_messages)
        return result


@dataclass
class CaptureResult:
    """
    Result of running an EvalSuite in capture mode.

    Attributes:
        suite_name: The name of the evaluation suite.
        model: The model used for capture.
        provider: The provider used (openai, anthropic).
        captured_cases: List of captured cases with tool calls.
    """

    suite_name: str
    model: str
    provider: str
    captured_cases: list[CapturedCase] = field(default_factory=list)

    def to_dict(self, include_context: bool = False) -> dict[str, Any]:
        """Convert to dictionary for JSON serialization."""
        return {
            "suite_name": self.suite_name,
            "model": self.model,
            "provider": self.provider,
            "captured_cases": [c.to_dict(include_context) for c in self.captured_cases],
        }

    def to_json(self, include_context: bool = False, indent: int = 2) -> str:
        """Convert to JSON string."""
        return json.dumps(self.to_dict(include_context), indent=indent)

    def write_to_file(self, file_path: str, include_context: bool = False, indent: int = 2) -> None:
        """Write capture results to a JSON file."""
        with open(file_path, "w", encoding="utf-8") as f:
            f.write(self.to_json(include_context, indent))


# --- Helper functions for running capture mode ---


async def _capture_with_openai(
    suite: EvalSuite,
    api_key: str,
    model: str,
    include_context: bool = False,
    num_runs: int = 1,
    seed: str | int | None = "constant",
) -> CaptureResult:
    """Run capture mode with OpenAI client."""
    async with AsyncOpenAI(api_key=api_key) as client:
        return await suite.capture(
            client,
            model,
            provider="openai",
            include_context=include_context,
            num_runs=num_runs,
            seed=seed,
        )


async def _capture_with_anthropic(
    suite: EvalSuite,
    api_key: str,
    model: str,
    include_context: bool = False,
    num_runs: int = 1,
    seed: str | int | None = "constant",
) -> CaptureResult:
    """Run capture mode with Anthropic client."""
    try:
        from anthropic import AsyncAnthropic
    except ImportError as e:
        raise ImportError(
            "The 'anthropic' package is required for Anthropic provider. "
            "Install it with: pip install anthropic"
        ) from e

    async with AsyncAnthropic(api_key=api_key) as client:
        return await suite.capture(
            client,
            model,
            provider="anthropic",
            include_context=include_context,
            num_runs=num_runs,
            seed=seed,
        )