understudy 0.1.0__py3-none-any.whl

understudy/__init__.py

@@ -0,0 +1,67 @@
+"""understudy: simulation and trace-based evaluation for agentic systems.
+
+The simulated user is an understudy standing in for a real user.
+You write scenes, run rehearsals, and check the performance —
+not by reading the script, but by inspecting what actually happened.
+"""
+
+from .check import CheckItem, CheckResult, check
+from .judges import Judge, JudgeResult
+from .mocks import MockToolkit, ToolError
+from .models import Expectations, Persona, PersonaPreset, Scene
+from .prompts import (
+    ADVERSARIAL_ROBUSTNESS,
+    FACTUAL_GROUNDING,
+    INSTRUCTION_FOLLOWING,
+    POLICY_COMPLIANCE,
+    TASK_COMPLETION,
+    TONE_EMPATHY,
+    TOOL_USAGE_CORRECTNESS,
+)
+from .runner import AgentApp, AgentResponse, run
+from .storage import RunStorage
+from .suite import SceneResult, Suite, SuiteResults
+from .trace import AgentTransfer, ToolCall, Trace, Turn
+
+__version__ = "0.1.0"
+
+__all__ = [
+    # models
+    "Scene",
+    "Persona",
+    "PersonaPreset",
+    "Expectations",
+    # trace
+    "Trace",
+    "Turn",
+    "ToolCall",
+    "AgentTransfer",
+    # runner
+    "run",
+    "AgentApp",
+    "AgentResponse",
+    # check
+    "check",
+    "CheckResult",
+    "CheckItem",
+    # suite
+    "Suite",
+    "SuiteResults",
+    "SceneResult",
+    # storage
+    "RunStorage",
+    # judges
+    "Judge",
+    "JudgeResult",
+    # mocks
+    "MockToolkit",
+    "ToolError",
+    # rubrics
+    "TOOL_USAGE_CORRECTNESS",
+    "POLICY_COMPLIANCE",
+    "TONE_EMPATHY",
+    "ADVERSARIAL_ROBUSTNESS",
+    "TASK_COMPLETION",
+    "FACTUAL_GROUNDING",
+    "INSTRUCTION_FOLLOWING",
+]

understudy/adk/__init__.py

@@ -0,0 +1,194 @@
+"""ADK adapter: wraps Google ADK agents for use with understudy."""
+
+from datetime import UTC, datetime
+from typing import Any
+
+from ..mocks import MockToolkit
+from ..runner import AgentApp, AgentResponse
+from ..trace import AgentTransfer, ToolCall
+
+
+def _create_mock_callback(mocks: MockToolkit | None):
+    """Create a before_tool_callback that returns mock responses.
+
+    Args:
+        mocks: MockToolkit instance or None.
+
+    Returns:
+        A callback function compatible with google-adk's before_tool_callback.
+        Returns dict to bypass real tool execution (mock response).
+        Returns None to allow normal execution.
+    """
+
+    def callback(tool, args: dict[str, Any], tool_context) -> dict | None:
+        if mocks is None:
+            return None
+        tool_name = getattr(tool, "name", None) or getattr(tool, "__name__", str(tool))
+        if mocks.get_handler(tool_name):
+            try:
+                result = mocks.call(tool_name, **args)
+                return result
+            except Exception as e:
+                return {"error": str(e)}
+        return None
+
+    return callback
+
+
+class ADKApp(AgentApp):
+    """Wraps a Google ADK Agent for use with understudy.
+
+    Usage:
+        from google.adk import Agent
+        from understudy.adk import ADKApp
+
+        agent = Agent(model="gemini-2.5-flash", name="my_agent", ...)
+        app = ADKApp(agent=agent)
+        trace = run(app, scene)
+    """
+
+    def __init__(self, agent: Any, session_id: str | None = None):
+        """
+        Args:
+            agent: A google.adk.Agent instance.
+            session_id: Optional session ID. If None, a random one is generated.
+        """
+        self.agent = agent
+        self.session_id = session_id
+        self._runner = None
+        self._session = None
+        self._mocks: MockToolkit | None = None
+        self._current_agent: str | None = None
+        self._agent_transfers: list[AgentTransfer] = []
+
+    def start(self, mocks: MockToolkit | None = None) -> None:
+        """Initialize the ADK session."""
+        try:
+            from google.adk import Runner
+            from google.adk.sessions import InMemorySessionService
+        except ImportError as e:
+            raise ImportError(
+                "google-adk package required. Install with: pip install understudy[adk]"
+            ) from e
+        import uuid
+
+        self._mocks = mocks
+        self._current_agent = getattr(self.agent, "name", None)
+        self._agent_transfers = []
+        self._session_id = self.session_id or str(uuid.uuid4())
+
+        session_service = InMemorySessionService()
+        if mocks:
+            self.agent.before_tool_callback = _create_mock_callback(mocks)
+
+        self._runner = Runner(
+            agent=self.agent,
+            app_name="understudy_test",
+            session_service=session_service,
+        )
+        self._session = session_service.create_session_sync(
+            app_name="understudy_test",
+            user_id="understudy_user",
+            session_id=self._session_id,
+        )
+
+    def send(self, message: str) -> AgentResponse:
+        """Send a user message to the ADK agent and capture the response."""
+        try:
+            from google.genai import types
+        except ImportError as e:
+            raise ImportError(
+                "google-adk package required. Install with: pip install understudy[adk]"
+            ) from e
+
+        user_content = types.Content(
+            role="user",
+            parts=[types.Part(text=message)],
+        )
+
+        tool_calls: list[ToolCall] = []
+        agent_text_parts: list[str] = []
+        terminal_state: str | None = None
+        current_agent_name = self._current_agent
+
+        for event in self._runner.run(
+            user_id="understudy_user",
+            session_id=self._session.id,
+            new_message=user_content,
+        ):
+            # track agent attribution from event.author
+            if hasattr(event, "author") and event.author:
+                event_agent = event.author
+                if event_agent != current_agent_name and current_agent_name:
+                    self._agent_transfers.append(
+                        AgentTransfer(
+                            from_agent=current_agent_name,
+                            to_agent=event_agent,
+                            timestamp=datetime.now(UTC),
+                        )
+                    )
+                current_agent_name = event_agent
+
+            # detect explicit transfer_to_agent actions
+            if (
+                hasattr(event, "actions")
+                and event.actions
+                and hasattr(event.actions, "transfer_to_agent")
+                and event.actions.transfer_to_agent
+            ):
+                target_agent = event.actions.transfer_to_agent
+                if current_agent_name and target_agent != current_agent_name:
+                    self._agent_transfers.append(
+                        AgentTransfer(
+                            from_agent=current_agent_name,
+                            to_agent=target_agent,
+                            timestamp=datetime.now(UTC),
+                        )
+                    )
+                current_agent_name = target_agent
+
+            # capture tool calls using get_function_calls()
+            for fc in event.get_function_calls():
+                call = ToolCall(
+                    tool_name=fc.name,
+                    arguments=dict(fc.args) if fc.args else {},
+                    agent_name=current_agent_name,
+                )
+                tool_calls.append(call)
+
+            # capture function responses and update tool call results
+            for fr in event.get_function_responses():
+                for call in tool_calls:
+                    if call.tool_name == fr.name and call.result is None:
+                        call.result = fr.response
+                        break
+
+            # capture text responses from content parts
+            if hasattr(event, "content") and event.content and hasattr(event.content, "parts"):
+                for part in event.content.parts:
+                    text = getattr(part, "text", None)
+                    if text:
+                        agent_text_parts.append(text)
+
+                        # check for terminal state markers
+                        # convention: agent emits "TERMINAL_STATE: <state>"
+                        if "TERMINAL_STATE:" in text:
+                            state = text.split("TERMINAL_STATE:")[-1].strip()
+                            terminal_state = state.split()[0].strip()
+
+        self._current_agent = current_agent_name
+
+        response = AgentResponse(
+            content=" ".join(agent_text_parts),
+            tool_calls=tool_calls,
+            terminal_state=terminal_state,
+        )
+        response.agent_name = current_agent_name
+        response.agent_transfers = list(self._agent_transfers)
+        return response
+
+    def stop(self) -> None:
+        """Clean up the ADK session."""
+        self._runner = None
+        self._session = None
+        self._mocks = None

understudy/check.py

@@ -0,0 +1,138 @@
+"""Check: validate a trace against scene expectations."""
+
+from dataclasses import dataclass, field
+
+from .models import Expectations
+from .trace import Trace
+
+
+@dataclass
+class CheckResult:
+    """Result of checking a trace against expectations."""
+
+    checks: list["CheckItem"] = field(default_factory=list)
+
+    @property
+    def passed(self) -> bool:
+        return all(c.passed for c in self.checks)
+
+    @property
+    def failed_checks(self) -> list["CheckItem"]:
+        return [c for c in self.checks if not c.passed]
+
+    def summary(self) -> str:
+        lines = []
+        for c in self.checks:
+            mark = "✓" if c.passed else "✗"
+            lines.append(f"  {mark} {c.label}: {c.detail}")
+        return "\n".join(lines)
+
+    def __repr__(self) -> str:
+        n_pass = sum(1 for c in self.checks if c.passed)
+        return f"CheckResult({n_pass}/{len(self.checks)} passed)"
+
+
+@dataclass
+class CheckItem:
+    """A single check result."""
+
+    label: str
+    passed: bool
+    detail: str
+
+
+def check(trace: Trace, expectations: Expectations) -> CheckResult:
+    """Validate a trace against expectations.
+
+    Args:
+        trace: The execution trace from a rehearsal.
+        expectations: The expectations from a scene.
+
+    Returns:
+        A CheckResult with individual check outcomes.
+    """
+    result = CheckResult()
+    called_tools = set(trace.call_sequence())
+
+    # required tools
+    for tool in expectations.required_tools:
+        result.checks.append(
+            CheckItem(
+                label="required_tool",
+                passed=tool in called_tools,
+                detail=f"{tool} {'called' if tool in called_tools else 'NOT called'}",
+            )
+        )
+
+    # forbidden tools
+    for tool in expectations.forbidden_tools:
+        was_called = tool in called_tools
+        result.checks.append(
+            CheckItem(
+                label="forbidden_tool",
+                passed=not was_called,
+                detail=f"{tool} {'CALLED (violation)' if was_called else 'not called'}",
+            )
+        )
+
+    # terminal state
+    if expectations.allowed_terminal_states:
+        in_allowed = trace.terminal_state in expectations.allowed_terminal_states
+        result.checks.append(
+            CheckItem(
+                label="terminal_state",
+                passed=in_allowed,
+                detail=(
+                    f"{trace.terminal_state} ({'allowed' if in_allowed else 'NOT in allowed'})"
+                ),
+            )
+        )
+
+    if expectations.forbidden_terminal_states:
+        in_forbidden = trace.terminal_state in expectations.forbidden_terminal_states
+        result.checks.append(
+            CheckItem(
+                label="forbidden_terminal_state",
+                passed=not in_forbidden,
+                detail=(
+                    f"{trace.terminal_state} "
+                    f"{'FORBIDDEN (violation)' if in_forbidden else 'not forbidden'}"
+                ),
+            )
+        )
+
+    # required agents
+    invoked_agents = set(trace.agents_invoked())
+    for agent in expectations.required_agents:
+        result.checks.append(
+            CheckItem(
+                label="required_agent",
+                passed=agent in invoked_agents,
+                detail=f"{agent} {'invoked' if agent in invoked_agents else 'NOT invoked'}",
+            )
+        )
+
+    # forbidden agents
+    for agent in expectations.forbidden_agents:
+        was_invoked = agent in invoked_agents
+        result.checks.append(
+            CheckItem(
+                label="forbidden_agent",
+                passed=not was_invoked,
+                detail=f"{agent} {'INVOKED (violation)' if was_invoked else 'not invoked'}",
+            )
+        )
+
+    # required agent tools
+    for agent, tools in expectations.required_agent_tools.items():
+        for tool in tools:
+            called = trace.agent_called(agent, tool)
+            result.checks.append(
+                CheckItem(
+                    label="required_agent_tool",
+                    passed=called,
+                    detail=f"{agent}.{tool} {'called' if called else 'NOT called'}",
+                )
+            )
+
+    return result

understudy/cli.py

@@ -0,0 +1,258 @@
+"""CLI: command-line interface for understudy."""
+
+import sys
+from pathlib import Path
+
+import click
+
+from .reports import ReportGenerator
+from .storage import RunStorage
+
+
+@click.group()
+@click.version_option()
+def main():
+    """understudy - Test your AI agents with simulated users."""
+    pass
+
+
+@main.command()
+@click.option(
+    "--runs",
+    "-r",
+    type=click.Path(exists=True, path_type=Path),
+    default=".understudy/runs",
+    help="Path to runs directory",
+)
+@click.option(
+    "--output",
+    "-o",
+    type=click.Path(path_type=Path),
+    default="report.html",
+    help="Output HTML file path",
+)
+def report(runs: Path, output: Path):
+    """Generate a static HTML report from saved runs."""
+    storage = RunStorage(path=runs)
+
+    run_ids = storage.list_runs()
+    if not run_ids:
+        click.echo(f"No runs found in {runs}")
+        sys.exit(1)
+
+    click.echo(f"Found {len(run_ids)} runs")
+
+    generator = ReportGenerator(storage)
+    generator.generate_static_report(output)
+
+    click.echo(f"Report generated: {output}")
+
+
+@main.command()
+@click.option(
+    "--runs",
+    "-r",
+    type=click.Path(exists=True, path_type=Path),
+    default=".understudy/runs",
+    help="Path to runs directory",
+)
+@click.option(
+    "--port",
+    "-p",
+    type=int,
+    default=8080,
+    help="Port to serve on",
+)
+@click.option(
+    "--host",
+    "-h",
+    type=str,
+    default="127.0.0.1",
+    help="Host to bind to",
+)
+def serve(runs: Path, port: int, host: str):
+    """Start an interactive report browser."""
+    storage = RunStorage(path=runs)
+
+    run_ids = storage.list_runs()
+    if not run_ids:
+        click.echo(f"No runs found in {runs}")
+        sys.exit(1)
+
+    click.echo(f"Found {len(run_ids)} runs")
+
+    generator = ReportGenerator(storage)
+    generator.serve(port=port, host=host)
+
+
+@main.command("list")
+@click.option(
+    "--runs",
+    "-r",
+    type=click.Path(exists=True, path_type=Path),
+    default=".understudy/runs",
+    help="Path to runs directory",
+)
+def list_runs(runs: Path):
+    """List all saved runs."""
+    storage = RunStorage(path=runs)
+
+    run_ids = storage.list_runs()
+    if not run_ids:
+        click.echo(f"No runs found in {runs}")
+        return
+
+    click.echo(f"Found {len(run_ids)} runs:\n")
+
+    for run_id in run_ids:
+        data = storage.load(run_id)
+        meta = data.get("metadata", {})
+        status = "PASS" if meta.get("passed") else "FAIL"
+        state = meta.get("terminal_state", "none")
+        turns = meta.get("turn_count", 0)
+        click.echo(f"  [{status}] {run_id} - {state} ({turns} turns)")
+
+
+@main.command()
+@click.option(
+    "--runs",
+    "-r",
+    type=click.Path(exists=True, path_type=Path),
+    default=".understudy/runs",
+    help="Path to runs directory",
+)
+def summary(runs: Path):
+    """Show aggregate metrics for saved runs."""
+    storage = RunStorage(path=runs)
+
+    run_ids = storage.list_runs()
+    if not run_ids:
+        click.echo(f"No runs found in {runs}")
+        return
+
+    stats = storage.get_summary()
+
+    click.echo("understudy Summary")
+    click.echo("=" * 40)
+    click.echo(f"Total Runs:    {stats['total_runs']}")
+    click.echo(f"Pass Rate:     {stats['pass_rate'] * 100:.1f}%")
+    click.echo(f"Avg Turns:     {stats['avg_turns']:.1f}")
+
+    if stats["tool_usage"]:
+        click.echo("\nTool Usage:")
+        for tool, count in sorted(stats["tool_usage"].items(), key=lambda x: -x[1]):
+            click.echo(f"  {tool}: {count}")
+
+    if stats["terminal_states"]:
+        click.echo("\nTerminal States:")
+        for state, count in sorted(stats["terminal_states"].items(), key=lambda x: -x[1]):
+            click.echo(f"  {state}: {count}")
+
+    if stats["agents"]:
+        click.echo("\nAgents:")
+        for agent, count in sorted(stats["agents"].items(), key=lambda x: -x[1]):
+            click.echo(f"  {agent}: {count}")
+
+
+@main.command()
+@click.argument("run_id")
+@click.option(
+    "--runs",
+    "-r",
+    type=click.Path(exists=True, path_type=Path),
+    default=".understudy/runs",
+    help="Path to runs directory",
+)
+def show(run_id: str, runs: Path):
+    """Show details for a specific run."""
+    storage = RunStorage(path=runs)
+
+    try:
+        data = storage.load(run_id)
+    except FileNotFoundError:
+        click.echo(f"Run not found: {run_id}")
+        sys.exit(1)
+
+    meta = data.get("metadata", {})
+    trace = data.get("trace")
+    check = data.get("check")
+
+    click.echo(f"Run: {run_id}")
+    click.echo("=" * 40)
+    click.echo(f"Scene:          {meta.get('scene_id', 'unknown')}")
+    click.echo(f"Status:         {'PASS' if meta.get('passed') else 'FAIL'}")
+    click.echo(f"Terminal State: {meta.get('terminal_state', 'none')}")
+    click.echo(f"Turns:          {meta.get('turn_count', 0)}")
+    click.echo(f"Tools Called:   {', '.join(meta.get('tools_called', []))}")
+    click.echo(f"Agents:         {', '.join(meta.get('agents_invoked', []))}")
+
+    if check and check.get("checks"):
+        click.echo("\nExpectation Checks:")
+        for c in check["checks"]:
+            icon = "+" if c["passed"] else "-"
+            click.echo(f"  {icon} {c['label']}: {c['detail']}")
+
+    if trace:
+        click.echo("\nConversation:")
+        click.echo("-" * 40)
+        for turn in trace.turns:
+            role = turn.agent_name or turn.role.upper()
+            click.echo(f"[{role}]: {turn.content[:100]}{'...' if len(turn.content) > 100 else ''}")
+            for call in turn.tool_calls:
+                click.echo(f"  -> {call.tool_name}({call.arguments})")
+
+
+@main.command()
+@click.argument("run_id")
+@click.option(
+    "--runs",
+    "-r",
+    type=click.Path(exists=True, path_type=Path),
+    default=".understudy/runs",
+    help="Path to runs directory",
+)
+@click.option("--yes", "-y", is_flag=True, help="Skip confirmation")
+def delete(run_id: str, runs: Path, yes: bool):
+    """Delete a specific run."""
+    storage = RunStorage(path=runs)
+
+    try:
+        storage.load(run_id)
+    except FileNotFoundError:
+        click.echo(f"Run not found: {run_id}")
+        sys.exit(1)
+
+    if not yes and not click.confirm(f"Delete run {run_id}?"):
+        return
+
+    storage.delete(run_id)
+    click.echo(f"Deleted: {run_id}")
+
+
+@main.command()
+@click.option(
+    "--runs",
+    "-r",
+    type=click.Path(exists=True, path_type=Path),
+    default=".understudy/runs",
+    help="Path to runs directory",
+)
+@click.option("--yes", "-y", is_flag=True, help="Skip confirmation")
+def clear(runs: Path, yes: bool):
+    """Delete all saved runs."""
+    storage = RunStorage(path=runs)
+
+    run_ids = storage.list_runs()
+    if not run_ids:
+        click.echo(f"No runs found in {runs}")
+        return
+
+    if not yes and not click.confirm(f"Delete all {len(run_ids)} runs?"):
+        return
+
+    storage.clear()
+    click.echo(f"Cleared {len(run_ids)} runs")
+
+
+if __name__ == "__main__":
+    main()