langwatch-scenario 0.2.0__py3-none-any.whl → 0.4.0__py3-none-any.whl

scenario/utils.py

@@ -1,7 +1,26 @@
+"""
+Utility functions for scenario execution and message handling.
+
+This module provides various utility functions used throughout the Scenario framework,
+including message formatting, validation, role reversal, and UI components like spinners
+for better user experience during scenario execution.
+"""
+
 from contextlib import contextmanager
 import sys
-from typing import Optional, Union
+from typing import (
+    Any,
+    Iterator,
+    List,
+    Literal,
+    Optional,
+    Union,
+    TypeVar,
+    Awaitable,
+    cast,
+)
 from pydantic import BaseModel
+import copy
 
 import json
 
@@ -14,16 +33,53 @@ from rich.console import Console
 from rich.text import Text
 from rich.errors import LiveError
 
+from scenario.error_messages import message_return_error_message
+from scenario.types import AgentReturnTypes, ScenarioResult
+
+T = TypeVar("T")
 
 
 class SerializableAndPydanticEncoder(json.JSONEncoder):
+    """
+    JSON encoder that handles Pydantic models and iterators.
+
+    This encoder extends the standard JSON encoder to handle Pydantic BaseModel
+    instances and iterator objects, converting them to serializable formats.
+    Used for caching and logging scenarios that contain complex objects.
+
+    Example:
+        ```python
+        data = {
+            "model": SomeBaseModel(field="value"),
+            "iterator": iter([1, 2, 3])
+        }
+        json.dumps(data, cls=SerializableAndPydanticEncoder)
+        ```
+    """
     def default(self, o):
         if isinstance(o, BaseModel):
             return o.model_dump(exclude_unset=True)
+        if isinstance(o, Iterator):
+            return list(o)
         return super().default(o)
 
 
 class SerializableWithStringFallback(SerializableAndPydanticEncoder):
+    """
+    JSON encoder with string fallback for non-serializable objects.
+
+    This encoder extends SerializableAndPydanticEncoder by providing a string
+    fallback for any object that cannot be serialized normally. This ensures
+    that logging and caching operations never fail due to serialization issues.
+
+    Example:
+        ```python
+        # This will work even with complex non-serializable objects
+        data = {"function": lambda x: x, "complex_object": SomeComplexClass()}
+        json.dumps(data, cls=SerializableWithStringFallback)
+        # Result: {"function": "<function <lambda> at 0x...>", "complex_object": "..."}
+        ```
+    """
     def default(self, o):
         try:
             return super().default(o)
@@ -32,6 +88,25 @@ class SerializableWithStringFallback(SerializableAndPydanticEncoder):
 
 
 def safe_list_at(list, index, default=None):
+    """
+    Safely get an item from a list by index with a default fallback.
+
+    Args:
+        list: The list to access
+        index: The index to retrieve
+        default: Value to return if index is out of bounds
+
+    Returns:
+        The item at the index, or the default value if index is invalid
+
+    Example:
+        ```python
+        items = ["a", "b", "c"]
+        print(safe_list_at(items, 1))    # "b"
+        print(safe_list_at(items, 10))   # None
+        print(safe_list_at(items, 10, "default"))  # "default"
+        ```
+    """
     try:
         return list[index]
     except:
@@ -39,14 +114,85 @@ def safe_list_at(list, index, default=None):
 
 
 def safe_attr_or_key(obj, attr_or_key, default=None):
+    """
+    Safely get an attribute or dictionary key from an object.
+
+    Tries to get the value as an attribute first, then as a dictionary key,
+    returning the default if neither exists.
+
+    Args:
+        obj: Object to access (can have attributes or be dict-like)
+        attr_or_key: Name of attribute or key to retrieve
+        default: Value to return if attribute/key doesn't exist
+
+    Returns:
+        The attribute/key value, or the default if not found
+
+    Example:
+        ```python
+        class MyClass:
+            attr = "value"
+
+        obj = MyClass()
+        dict_obj = {"key": "value"}
+
+        print(safe_attr_or_key(obj, "attr"))     # "value"
+        print(safe_attr_or_key(dict_obj, "key")) # "value"
+        print(safe_attr_or_key(obj, "missing"))  # None
+        ```
+    """
     return getattr(obj, attr_or_key, obj.get(attr_or_key))
 
 
 def title_case(string):
+    """
+    Convert snake_case string to Title Case.
+
+    Args:
+        string: Snake_case string to convert
+
+    Returns:
+        String converted to Title Case
+
+    Example:
+        ```python
+        print(title_case("user_simulator_agent"))  # "User Simulator Agent"
+        print(title_case("api_key"))               # "Api Key"
+        ```
+    """
     return " ".join(word.capitalize() for word in string.split("_"))
 
 
-def print_openai_messages(scenario_name: str, messages: list[ChatCompletionMessageParam]):
+def print_openai_messages(
+    scenario_name: str, messages: list[ChatCompletionMessageParam]
+):
+    """
+    Print OpenAI-format messages with colored formatting for readability.
+
+    This function formats and prints conversation messages with appropriate
+    colors and formatting for different message types (user, assistant, tool calls, etc.).
+    Used for verbose output during scenario execution.
+
+    Args:
+        scenario_name: Name of the scenario (used as prefix)
+        messages: List of OpenAI-compatible messages to print
+
+    Example:
+        ```python
+        messages = [
+            {"role": "user", "content": "Hello"},
+            {"role": "assistant", "content": "Hi there!"},
+            {"role": "assistant", "tool_calls": [{"function": {"name": "search"}}]}
+        ]
+        print_openai_messages("Test Scenario", messages)
+        ```
+
+    Note:
+        - User messages are printed in green
+        - Assistant messages are printed in blue
+        - Tool calls are printed in magenta with formatted JSON
+        - Long JSON content is truncated for readability
+    """
     for msg in messages:
         role = safe_attr_or_key(msg, "role")
         content = safe_attr_or_key(msg, "content")
@@ -61,9 +207,12 @@ def print_openai_messages(scenario_name: str, messages: list[ChatCompletionMessa
                     args = safe_attr_or_key(function, "arguments", "{}")
                     args = _take_maybe_json_first_lines(args)
                     print(
-                        scenario_name + termcolor.colored(f"ToolCall({name}):", "magenta"),
+                        scenario_name
+                        + termcolor.colored(f"ToolCall({name}):", "magenta"),
                         f"\n\n{indent(args, ' ' * 4)}\n",
                     )
+        elif role == "user":
+            print(scenario_name + termcolor.colored("User:", "green"), content)
         elif role == "tool":
             content = _take_maybe_json_first_lines(content or msg.__repr__())
             print(
@@ -78,6 +227,19 @@ def print_openai_messages(scenario_name: str, messages: list[ChatCompletionMessa
 
 
 def _take_maybe_json_first_lines(string, max_lines=5):
+    """
+    Truncate string content and format JSON if possible.
+
+    Internal utility function that attempts to format content as JSON
+    and truncates it to a reasonable number of lines for display.
+
+    Args:
+        string: Content to format and truncate
+        max_lines: Maximum number of lines to show
+
+    Returns:
+        Formatted and potentially truncated string
+    """
     content = str(string)
     try:
         content = json.dumps(json.loads(content), indent=2)
@@ -91,9 +253,25 @@ def _take_maybe_json_first_lines(string, max_lines=5):
 
 console = Console()
 
+
 class TextFirstSpinner(Spinner):
+    """
+    Custom spinner that displays text before the spinning animation.
+
+    This class extends Rich's Spinner to show descriptive text followed
+    by the spinning animation, improving the user experience during
+    scenario execution by clearly indicating what operation is happening.
+
+    Args:
+        name: Name of the spinner animation style
+        text: Descriptive text to show before the spinner
+        color: Color for the descriptive text
+        **kwargs: Additional arguments passed to the base Spinner class
+    """
     def __init__(self, name, text: str, color: str, **kwargs):
-        super().__init__(name, "", style="bold white", **kwargs)  # Initialize with empty text
+        super().__init__(
+            name, "", style="bold white", **kwargs
+        )  # Initialize with empty text
         self.text_before = text
         self.color = color
 
@@ -105,7 +283,35 @@ class TextFirstSpinner(Spinner):
 
 
 @contextmanager
-def show_spinner(text: str, color: str = "white", enabled: Optional[Union[bool, int]] = None):
+def show_spinner(
+    text: str, color: str = "white", enabled: Optional[Union[bool, int]] = None
+):
+    """
+    Context manager for displaying a spinner during long-running operations.
+
+    Shows a spinning indicator with descriptive text while code executes
+    within the context. Automatically cleans up the spinner display when
+    the operation completes.
+
+    Args:
+        text: Descriptive text to show next to the spinner
+        color: Color for the descriptive text
+        enabled: Whether to show the spinner (respects verbose settings)
+
+    Example:
+        ```python
+        with show_spinner("Calling agent...", color="blue", enabled=True):
+            response = await agent.call(input_data)
+
+        # Spinner automatically disappears when block completes
+        print("Agent call completed")
+        ```
+
+    Note:
+        - Spinner is automatically cleaned up when context exits
+        - Gracefully handles multi-threading scenarios where multiple spinners might conflict
+        - Cursor positioning ensures clean terminal output
+    """
     if not enabled:
         yield
     else:
@@ -119,3 +325,190 @@ def show_spinner(text: str, color: str = "white", enabled: Optional[Union[bool,
 
         # Cursor up one line
         sys.stdout.write("\033[F")
+        # Erase the line
+        sys.stdout.write("\033[2K")
+
+
+def check_valid_return_type(return_value: Any, class_name: str) -> None:
+    """
+    Validate that an agent's return value is in the expected format.
+
+    This function ensures that agent adapters return values in one of the
+    supported formats (string, OpenAI message, list of messages, or ScenarioResult).
+    It also verifies that the returned data is JSON-serializable for caching.
+
+    Args:
+        return_value: The value returned by an agent's call method
+        class_name: Name of the agent class (for error messages)
+
+    Raises:
+        ValueError: If the return value is not in a supported format
+
+    Example:
+        ```python
+        # Valid return values
+        check_valid_return_type("Hello world", "MyAgent")  # OK
+        check_valid_return_type({"role": "assistant", "content": "Hi"}, "MyAgent")  # OK
+        check_valid_return_type([{"role": "assistant", "content": "Hi"}], "MyAgent")  # OK
+
+        # Invalid return value
+        check_valid_return_type(42, "MyAgent")  # Raises ValueError
+        ```
+    """
+    def _is_valid_openai_message(message: Any) -> bool:
+        return (isinstance(message, dict) and "role" in message) or (
+            isinstance(message, BaseModel) and hasattr(message, "role")
+        )
+
+    if (
+        isinstance(return_value, str)
+        or _is_valid_openai_message(return_value)
+        or (
+            isinstance(return_value, list)
+            and all(_is_valid_openai_message(message) for message in return_value)
+        )
+        or isinstance(return_value, ScenarioResult)
+    ):
+        try:
+            json.dumps(return_value, cls=SerializableAndPydanticEncoder)
+        except:
+            raise ValueError(
+                message_return_error_message(got=return_value, class_name=class_name)
+            )
+
+        return
+
+    raise ValueError(
+        message_return_error_message(got=return_value, class_name=class_name)
+    )
+
+
+def convert_agent_return_types_to_openai_messages(
+    agent_response: AgentReturnTypes, role: Literal["user", "assistant"]
+) -> List[ChatCompletionMessageParam]:
+    """
+    Convert various agent return types to standardized OpenAI message format.
+
+    This function normalizes different return types from agent adapters into
+    a consistent list of OpenAI-compatible messages that can be used throughout
+    the scenario execution pipeline.
+
+    Args:
+        agent_response: Response from an agent adapter call
+        role: The role to assign to string responses ("user" or "assistant")
+
+    Returns:
+        List of OpenAI-compatible messages
+
+    Raises:
+        ValueError: If agent_response is a ScenarioResult (which should be handled separately)
+
+    Example:
+        ```python
+        # String response
+        messages = convert_agent_return_types_to_openai_messages("Hello", "assistant")
+        # Result: [{"role": "assistant", "content": "Hello"}]
+
+        # Dict response
+        response = {"role": "assistant", "content": "Hi", "tool_calls": [...]}
+        messages = convert_agent_return_types_to_openai_messages(response, "assistant")
+        # Result: [{"role": "assistant", "content": "Hi", "tool_calls": [...]}]
+
+        # List response
+        responses = [
+            {"role": "assistant", "content": "Thinking..."},
+            {"role": "assistant", "content": "Here's the answer"}
+        ]
+        messages = convert_agent_return_types_to_openai_messages(responses, "assistant")
+        # Result: Same list, validated and normalized
+        ```
+    """
+    if isinstance(agent_response, ScenarioResult):
+        raise ValueError(
+            "Unexpectedly tried to convert a ScenarioResult to openai messages",
+            agent_response.__repr__(),
+        )
+
+    def convert_maybe_object_to_openai_message(
+        obj: Any,
+    ) -> ChatCompletionMessageParam:
+        if isinstance(obj, dict):
+            return cast(ChatCompletionMessageParam, obj)
+        elif isinstance(obj, BaseModel):
+            return cast(
+                ChatCompletionMessageParam,
+                obj.model_dump(
+                    exclude_unset=True,
+                    exclude_none=True,
+                    exclude_defaults=True,
+                    warnings=False,
+                ),
+            )
+        else:
+            raise ValueError(f"Unexpected agent response type: {type(obj).__name__}")
+
+    def ensure_dict(
+        obj: T,
+    ) -> T:
+        return json.loads(json.dumps(obj, cls=SerializableAndPydanticEncoder))
+
+    if isinstance(agent_response, str):
+        return [
+            (
+                {"role": "user", "content": agent_response}
+                if role == "user"
+                else {"role": "assistant", "content": agent_response}
+            )
+        ]
+    elif isinstance(agent_response, list):
+        return [
+            ensure_dict(convert_maybe_object_to_openai_message(message))
+            for message in agent_response
+        ]
+    else:
+        return [ensure_dict(convert_maybe_object_to_openai_message(agent_response))]
+
+
+def reverse_roles(
+    messages: list[ChatCompletionMessageParam],
+) -> list[ChatCompletionMessageParam]:
+    """
+    Reverses the roles of the messages in the list.
+
+    Args:
+        messages: The list of messages to reverse the roles of.
+    """
+
+    reversed_messages = []
+    for message in messages:
+        message = copy.deepcopy(message)
+        # Can't reverse tool calls
+        if not safe_attr_or_key(message, "content") or safe_attr_or_key(
+            message, "tool_calls"
+        ):
+            # If no content nor tool calls, we should skip it entirely, as anthropic may generate some invalid ones e.g. pure {"role": "assistant"}
+            if safe_attr_or_key(message, "tool_calls"):
+                reversed_messages.append(message)
+            continue
+
+        if type(message) == dict:
+            if message["role"] == "user":
+                message["role"] = "assistant"
+            elif message["role"] == "assistant":
+                message["role"] = "user"
+        else:
+            if getattr(message, "role", None) == "user":
+                message.role = "assistant"  # type: ignore
+            elif getattr(message, "role", None) == "assistant":
+                message.role = "user"  # type: ignore
+
+        reversed_messages.append(message)
+
+    return reversed_messages
+
+
+async def await_if_awaitable(value: T) -> T:
+    if isinstance(value, Awaitable):
+        return await value
+    else:
+        return value