langwatch-scenario 0.3.0__py3-none-any.whl → 0.6.0__py3-none-any.whl

scenario/scenario_state.py

@@ -0,0 +1,205 @@
+"""
+Scenario state management module.
+
+This module provides the ScenarioState class which tracks the current state
+of a scenario execution, including conversation history, turn tracking, and
+utility methods for inspecting the conversation.
+"""
+
+from typing import List, Dict, Any, Optional, TYPE_CHECKING
+from openai.types.chat import (
+    ChatCompletionMessageParam,
+    ChatCompletionMessageToolCallParam,
+    ChatCompletionUserMessageParam,
+)
+from pydantic import BaseModel
+
+from scenario.config import ScenarioConfig
+
+if TYPE_CHECKING:
+    from .scenario_executor import ScenarioExecutor
+
+
+class ScenarioState(BaseModel):
+    """
+    Represents the current state of a scenario execution.
+
+    This class provides access to the conversation history, turn information,
+    and utility methods for inspecting messages and tool calls. It's passed to
+    script step functions and available through AgentInput.scenario_state.
+
+    Attributes:
+        description: The scenario description that guides the simulation
+        messages: Complete conversation history as OpenAI-compatible messages
+        thread_id: Unique identifier for this conversation thread
+        current_turn: Current turn number in the conversation
+        config: Configuration settings for this scenario execution
+
+    Example:
+        ```
+        def check_agent_behavior(state: ScenarioState) -> None:
+            # Check if the agent called a specific tool
+            if state.has_tool_call("get_weather"):
+                print("Agent successfully called weather tool")
+
+            # Get the last user message
+            last_user = state.last_user_message()
+            print(f"User said: {last_user['content']}")
+
+            # Check conversation length
+            if len(state.messages) > 10:
+                print("Conversation is getting long")
+
+        # Use in scenario script
+        result = await scenario.run(
+            name="tool usage test",
+            description="Test that agent uses the correct tools",
+            agents=[
+                my_agent,
+                scenario.UserSimulatorAgent(),
+                scenario.JudgeAgent(criteria=["Agent provides helpful response"])
+            ],
+            script=[
+                scenario.user("What's the weather like?"),
+                scenario.agent(),
+                check_agent_behavior,  # Custom inspection function
+                scenario.succeed()
+            ]
+        )
+        ```
+    """
+    description: str
+    messages: List[ChatCompletionMessageParam]
+    thread_id: str
+    current_turn: int
+    config: ScenarioConfig
+
+    _executor: "ScenarioExecutor"
+
+    def add_message(self, message: ChatCompletionMessageParam):
+        """
+        Add a message to the conversation history.
+
+        This method delegates to the scenario executor to properly handle
+        message broadcasting and state updates.
+
+        Args:
+            message: OpenAI-compatible message to add to the conversation
+
+        Example:
+            ```
+            def inject_system_message(state: ScenarioState) -> None:
+                state.add_message({
+                    "role": "system",
+                    "content": "The user is now in a hurry"
+                })
+            ```
+        """
+        self._executor.add_message(message)
+
+    def last_message(self) -> ChatCompletionMessageParam:
+        """
+        Get the most recent message in the conversation.
+
+        Returns:
+            The last message in the conversation history
+
+        Raises:
+            ValueError: If no messages exist in the conversation
+
+        Example:
+            ```
+            def check_last_response(state: ScenarioState) -> None:
+                last = state.last_message()
+                if last["role"] == "assistant":
+                    content = last.get("content", "")
+                    assert "helpful" in content.lower()
+            ```
+        """
+        if len(self.messages) == 0:
+            raise ValueError("No messages found")
+        return self.messages[-1]
+
+    def last_user_message(self) -> ChatCompletionUserMessageParam:
+        """
+        Get the most recent user message in the conversation.
+
+        Returns:
+            The last user message in the conversation history
+
+        Raises:
+            ValueError: If no user messages exist in the conversation
+
+        Example:
+            ```
+            def analyze_user_intent(state: ScenarioState) -> None:
+                user_msg = state.last_user_message()
+                content = user_msg["content"]
+
+                if isinstance(content, str):
+                    if "urgent" in content.lower():
+                        print("User expressed urgency")
+            ```
+        """
+        user_messages = [m for m in self.messages if m["role"] == "user"]
+        if not user_messages:
+            raise ValueError("No user messages found")
+        return user_messages[-1]
+
+    def last_tool_call(
+        self, tool_name: str
+    ) -> Optional[ChatCompletionMessageToolCallParam]:
+        """
+        Find the most recent call to a specific tool in the conversation.
+
+        Searches through the conversation history in reverse order to find
+        the last time the specified tool was called by an assistant.
+
+        Args:
+            tool_name: Name of the tool to search for
+
+        Returns:
+            The tool call object if found, None otherwise
+
+        Example:
+            ```
+            def verify_weather_call(state: ScenarioState) -> None:
+                weather_call = state.last_tool_call("get_current_weather")
+                if weather_call:
+                    args = json.loads(weather_call["function"]["arguments"])
+                    assert "location" in args
+                    print(f"Weather requested for: {args['location']}")
+            ```
+        """
+        for message in reversed(self.messages):
+            if message["role"] == "assistant" and "tool_calls" in message:
+                for tool_call in message["tool_calls"]:
+                    if tool_call["function"]["name"] == tool_name:
+                        return tool_call
+        return None
+
+    def has_tool_call(self, tool_name: str) -> bool:
+        """
+        Check if a specific tool has been called in the conversation.
+
+        This is a convenience method that returns True if the specified
+        tool has been called at any point in the conversation.
+
+        Args:
+            tool_name: Name of the tool to check for
+
+        Returns:
+            True if the tool has been called, False otherwise
+
+        Example:
+            ```
+            def ensure_tool_usage(state: ScenarioState) -> None:
+                # Verify the agent used required tools
+                assert state.has_tool_call("search_database")
+                assert state.has_tool_call("format_results")
+
+                # Check it didn't use forbidden tools
+                assert not state.has_tool_call("delete_data")
+            ```
+        """
+        return self.last_tool_call(tool_name) is not None

scenario/script.py

@@ -0,0 +1,361 @@
+"""
+Scenario script DSL (Domain Specific Language) module.
+
+This module provides a collection of functions that form a declarative language
+for controlling scenario execution flow. These functions can be used to create
+scripts that precisely control how conversations unfold, when evaluations occur,
+and when scenarios should succeed or fail.
+"""
+
+from typing import Awaitable, Callable, Optional, Union, TYPE_CHECKING
+
+from .types import ScriptStep
+
+from openai.types.chat import ChatCompletionMessageParam
+
+if TYPE_CHECKING:
+    from scenario.scenario_state import ScenarioState
+
+
+def message(message: ChatCompletionMessageParam) -> ScriptStep:
+    """
+    Add a specific message to the conversation.
+
+    This function allows you to inject any OpenAI-compatible message directly
+    into the conversation at a specific point in the script. Useful for
+    simulating tool responses, system messages, or specific conversational states.
+
+    Args:
+        message: OpenAI-compatible message to add to the conversation
+
+    Returns:
+        ScriptStep function that can be used in scenario scripts
+
+    Example:
+        ```
+        result = await scenario.run(
+            name="tool response test",
+            description="Testing tool call responses",
+            agents=[
+                my_agent,
+                scenario.UserSimulatorAgent(),
+                scenario.JudgeAgent(criteria=["Agent uses weather tool correctly"])
+            ],
+            script=[
+                scenario.user("What's the weather?"),
+                scenario.agent(),  # Agent calls weather tool
+                scenario.message({
+                    "role": "tool",
+                    "tool_call_id": "call_123",
+                    "content": json.dumps({"temperature": "75°F", "condition": "sunny"})
+                }),
+                scenario.agent(),  # Agent processes tool response
+                scenario.succeed()
+            ]
+        )
+        ```
+    """
+    return lambda state: state._executor.message(message)
+
+
+def user(
+    content: Optional[Union[str, ChatCompletionMessageParam]] = None,
+) -> ScriptStep:
+    """
+    Generate or specify a user message in the conversation.
+
+    If content is provided, it will be used as the user message. If no content
+    is provided, the user simulator agent will automatically generate an
+    appropriate message based on the scenario context.
+
+    Args:
+        content: Optional user message content. Can be a string or full message dict.
+                If None, the user simulator will generate content automatically.
+
+    Returns:
+        ScriptStep function that can be used in scenario scripts
+
+    Example:
+        ```
+        result = await scenario.run(
+            name="user interaction test",
+            description="Testing specific user inputs",
+            agents=[
+                my_agent,
+                scenario.UserSimulatorAgent(),
+                scenario.JudgeAgent(criteria=["Agent responds helpfully to user"])
+            ],
+            script=[
+                # Specific user message
+                scenario.user("I need help with Python"),
+                scenario.agent(),
+
+                # Auto-generated user message based on scenario context
+                scenario.user(),
+                scenario.agent(),
+
+                # Structured user message with multimodal content
+                scenario.message({
+                    "role": "user",
+                    "content": [
+                        {"type": "text", "text": "What's in this image?"},
+                        {"type": "image_url", "image_url": {"url": "data:image/..."}}
+                    ]
+                }),
+                scenario.succeed()
+            ]
+        )
+        ```
+    """
+    return lambda state: state._executor.user(content)
+
+
+def agent(
+    content: Optional[Union[str, ChatCompletionMessageParam]] = None,
+) -> ScriptStep:
+    """
+    Generate or specify an agent response in the conversation.
+
+    If content is provided, it will be used as the agent response. If no content
+    is provided, the agent under test will be called to generate its response
+    based on the current conversation state.
+
+    Args:
+        content: Optional agent response content. Can be a string or full message dict.
+                If None, the agent under test will generate content automatically.
+
+    Returns:
+        ScriptStep function that can be used in scenario scripts
+
+    Example:
+        ```
+        result = await scenario.run(
+            name="agent response test",
+            description="Testing agent responses",
+            agents=[
+                my_agent,
+                scenario.UserSimulatorAgent(),
+                scenario.JudgeAgent(criteria=["Agent provides appropriate responses"])
+            ],
+            script=[
+                scenario.user("Hello"),
+
+                # Let agent generate its own response
+                scenario.agent(),
+
+                # Or specify exact agent response for testing edge cases
+                scenario.agent("I'm sorry, I'm currently unavailable"),
+                scenario.user(),  # See how user simulator reacts
+
+                # Structured agent response with tool calls
+                scenario.message({
+                    "role": "assistant",
+                    "content": "Let me search for that information",
+                    "tool_calls": [{"id": "call_123", "type": "function", ...}]
+                }),
+                scenario.succeed()
+            ]
+        )
+        ```
+    """
+    return lambda state: state._executor.agent(content)
+
+
+def judge(
+    content: Optional[Union[str, ChatCompletionMessageParam]] = None,
+) -> ScriptStep:
+    """
+    Invoke the judge agent to evaluate the current conversation state.
+
+    This function forces the judge agent to make a decision about whether
+    the scenario should continue or end with a success/failure verdict.
+    The judge will evaluate based on its configured criteria.
+
+    Args:
+        content: Optional message content for the judge. Usually None to let
+                the judge evaluate based on its criteria.
+
+    Returns:
+        ScriptStep function that can be used in scenario scripts
+
+    Example:
+        ```
+        result = await scenario.run(
+            name="judge evaluation test",
+            description="Testing judge at specific points",
+            agents=[
+                my_agent,
+                scenario.UserSimulatorAgent(),
+                scenario.JudgeAgent(criteria=["Agent provides coding help effectively"])
+            ],
+            script=[
+                scenario.user("Can you help me code?"),
+                scenario.agent(),
+
+                # Force judge evaluation after first exchange
+                scenario.judge(),  # May continue or end scenario
+
+                # If scenario continues...
+                scenario.user(),
+                scenario.agent(),
+                scenario.judge(),  # Final evaluation
+            ]
+        )
+        ```
+    """
+    return lambda state: state._executor.judge(content)
+
+
+def proceed(
+    turns: Optional[int] = None,
+    on_turn: Optional[
+        Union[
+            Callable[["ScenarioState"], None],
+            Callable[["ScenarioState"], Awaitable[None]],
+        ]
+    ] = None,
+    on_step: Optional[
+        Union[
+            Callable[["ScenarioState"], None],
+            Callable[["ScenarioState"], Awaitable[None]],
+        ]
+    ] = None,
+) -> ScriptStep:
+    """
+    Let the scenario proceed automatically for a specified number of turns.
+
+    This function allows the scenario to run automatically with the normal
+    agent interaction flow (user -> agent -> judge evaluation). You can
+    optionally provide callbacks to execute custom logic at each turn or step.
+
+    Args:
+        turns: Number of turns to proceed automatically. If None, proceeds until
+               the judge agent decides to end the scenario or max_turns is reached.
+        on_turn: Optional callback function called at the end of each turn
+        on_step: Optional callback function called after each agent interaction
+
+    Returns:
+        ScriptStep function that can be used in scenario scripts
+
+    Example:
+        ```
+        def log_progress(state: ScenarioState) -> None:
+            print(f"Turn {state.current_turn}: {len(state.messages)} messages")
+
+        def check_tool_usage(state: ScenarioState) -> None:
+            if state.has_tool_call("dangerous_action"):
+                raise AssertionError("Agent used forbidden tool!")
+
+        result = await scenario.run(
+            name="automatic proceeding test",
+            description="Let scenario run with monitoring",
+            agents=[
+                my_agent,
+                scenario.UserSimulatorAgent(),
+                scenario.JudgeAgent(criteria=["Agent behaves safely and helpfully"])
+            ],
+            script=[
+                scenario.user("Let's start"),
+                scenario.agent(),
+
+                # Let it proceed for 3 turns with monitoring
+                scenario.proceed(
+                    turns=3,
+                    on_turn=log_progress,
+                    on_step=check_tool_usage
+                ),
+
+                # Then do final evaluation
+                scenario.judge()
+            ]
+        )
+        ```
+    """
+    return lambda state: state._executor.proceed(turns, on_turn, on_step)
+
+
+def succeed(reasoning: Optional[str] = None) -> ScriptStep:
+    """
+    Immediately end the scenario with a success result.
+
+    This function terminates the scenario execution and marks it as successful,
+    bypassing any further agent interactions or judge evaluations.
+
+    Args:
+        reasoning: Optional explanation for why the scenario succeeded
+
+    Returns:
+        ScriptStep function that can be used in scenario scripts
+
+    Example:
+        ```
+        def custom_success_check(state: ScenarioState) -> None:
+            last_msg = state.last_message()
+            if "solution" in last_msg.get("content", "").lower():
+                # Custom success condition met
+                return scenario.succeed("Agent provided a solution")()
+
+        result = await scenario.run(
+            name="custom success test",
+            description="Test custom success conditions",
+            agents=[
+                my_agent,
+                scenario.UserSimulatorAgent(),
+                scenario.JudgeAgent(criteria=["Agent provides a solution"])
+            ],
+            script=[
+                scenario.user("I need a solution"),
+                scenario.agent(),
+                custom_success_check,
+
+                # Or explicit success
+                scenario.succeed("Agent completed the task successfully")
+            ]
+        )
+        ```
+    """
+    return lambda state: state._executor.succeed(reasoning)
+
+
+def fail(reasoning: Optional[str] = None) -> ScriptStep:
+    """
+    Immediately end the scenario with a failure result.
+
+    This function terminates the scenario execution and marks it as failed,
+    bypassing any further agent interactions or judge evaluations.
+
+    Args:
+        reasoning: Optional explanation for why the scenario failed
+
+    Returns:
+        ScriptStep function that can be used in scenario scripts
+
+    Example:
+        ```
+        def safety_check(state: ScenarioState) -> None:
+            last_msg = state.last_message()
+            content = last_msg.get("content", "")
+
+            if "harmful" in content.lower():
+                return scenario.fail("Agent produced harmful content")()
+
+        result = await scenario.run(
+            name="safety check test",
+            description="Test safety boundaries",
+            agents=[
+                my_agent,
+                scenario.UserSimulatorAgent(),
+                scenario.JudgeAgent(criteria=["Agent maintains safety guidelines"])
+            ],
+            script=[
+                scenario.user("Tell me something dangerous"),
+                scenario.agent(),
+                safety_check,
+
+                # Or explicit failure
+                scenario.fail("Agent failed to meet safety requirements")
+            ]
+        )
+        ```
+    """
+    return lambda state: state._executor.fail(reasoning)