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

scenario/cache.py

@@ -1,3 +1,12 @@
+"""
+Caching module for deterministic scenario testing.
+
+This module provides caching functionality to make scenario tests deterministic
+and repeatable. It caches LLM calls and other non-deterministic operations based
+on scenario configuration and function arguments, enabling consistent test results
+across multiple runs.
+"""
+
 from contextvars import ContextVar
 import inspect
 import os
@@ -8,29 +17,100 @@ from joblib import Memory
 import json
 
 import wrapt
-from scenario.utils import SerializableWithStringFallback
+from scenario.types import AgentInput
+from scenario._utils.utils import SerializableWithStringFallback
 
 if TYPE_CHECKING:
-    from scenario.scenario import Scenario
+    from scenario.scenario_executor import ScenarioExecutor
 
 
 context_scenario = ContextVar("scenario")
 
+
 def get_cache() -> Memory:
-    """Get a cross-platform cache directory for scenario."""
+    """
+    Get a cross-platform cache directory for scenario execution.
+
+    Creates and returns a joblib Memory instance configured to use a
+    cross-platform cache directory. The cache location can be customized
+    via the SCENARIO_CACHE_DIR environment variable.
+
+    Returns:
+        Memory instance configured with the appropriate cache directory
+
+    Example:
+        ```
+        # Default cache location: ~/.scenario/cache
+        cache = get_cache()
+
+        # Custom cache location via environment variable
+        os.environ["SCENARIO_CACHE_DIR"] = "/tmp/my_scenario_cache"
+        cache = get_cache()
+        ```
+    """
     home_dir = str(Path.home())
     cache_dir = os.path.join(home_dir, ".scenario", "cache")
 
     return Memory(location=os.environ.get("SCENARIO_CACHE_DIR", cache_dir), verbose=0)
 
+
 memory = get_cache()
 
+
 def scenario_cache(ignore=[]):
+    """
+    Decorator for caching function calls during scenario execution.
+
+    This decorator caches function calls based on the scenario's cache_key,
+    scenario configuration, and function arguments. It enables deterministic
+    testing by ensuring the same inputs always produce the same outputs,
+    making tests repeatable and faster on subsequent runs.
+
+    Args:
+        ignore: List of argument names to exclude from the cache key computation.
+                Commonly used to ignore 'self' for instance methods or other
+                non-deterministic arguments.
+
+    Returns:
+        Decorator function that can be applied to any function or method
+
+    Example:
+        ```
+        import scenario
+
+        class MyAgent:
+            @scenario.cache(ignore=["self"])
+            def invoke(self, message: str, context: dict) -> str:
+                # This LLM call will be cached
+                response = llm_client.complete(
+                    model="gpt-4",
+                    messages=[{"role": "user", "content": message}]
+                )
+                return response.choices[0].message.content
+
+        # Usage in tests
+        scenario.configure(cache_key="my-test-suite-v1")
+
+        # First run: makes actual LLM calls and caches results
+        result1 = await scenario.run(...)
+
+        # Second run: uses cached results, much faster
+        result2 = await scenario.run(...)
+        # result1 and result2 will be identical
+        ```
+
+    Note:
+        - Caching only occurs when a cache_key is set in the scenario configuration
+        - The cache key is computed from scenario config, function arguments, and cache_key
+        - AgentInput objects are specially handled to exclude thread_id from caching
+        - Both sync and async functions are supported
+    """
+
     @wrapt.decorator
     def wrapper(wrapped: Callable, instance=None, args=[], kwargs={}):
-        scenario: "Scenario" = context_scenario.get()
+        scenario: "ScenarioExecutor" = context_scenario.get()
 
-        if not scenario.cache_key:
+        if not scenario.config.cache_key:
             return wrapped(*args, **kwargs)
 
         sig = inspect.signature(wrapped)
@@ -43,20 +123,65 @@ def scenario_cache(ignore=[]):
             if arg in all_args:
                 del all_args[arg]
 
+        for key, value in all_args.items():
+            if isinstance(value, AgentInput):
+                scenario_state = value.scenario_state.model_dump(exclude={"thread_id"})
+                all_args[key] = value.model_dump(exclude={"thread_id"})
+                all_args[key]["scenario_state"] = scenario_state
+
         cache_key = json.dumps(
             {
-                "cache_key": scenario.cache_key,
-                "scenario": scenario.model_dump(exclude={"agent"}),
+                "cache_key": scenario.config.cache_key,
+                "scenario": scenario.config.model_dump(exclude={"agents"}),
                 "all_args": all_args,
             },
             cls=SerializableWithStringFallback,
         )
 
-        return _cached_call(wrapped, args, kwargs, cache_key=cache_key)
+        # if is an async function, we need to wrap it in a sync function
+        if inspect.iscoroutinefunction(wrapped):
+            return _async_cached_call(wrapped, args, kwargs, cache_key=cache_key)
+        else:
+            return _cached_call(wrapped, args, kwargs, cache_key=cache_key)
 
     return wrapper
 
 
 @memory.cache(ignore=["func", "args", "kwargs"])
 def _cached_call(func: Callable, args, kwargs, cache_key):
-    return func(*args, **kwargs)
+    """
+    Internal function for caching synchronous function calls.
+
+    This function is used internally by the scenario_cache decorator
+    to cache synchronous function calls using joblib.Memory.
+
+    Args:
+        func: The function to call and cache
+        args: Positional arguments for the function
+        kwargs: Keyword arguments for the function
+        cache_key: Cache key for deterministic caching
+
+    Returns:
+        The result of calling func(*args, **kwargs)
+    """
+    return func(*args, **kwargs)
+
+
+@memory.cache(ignore=["func", "args", "kwargs"])
+async def _async_cached_call(func: Callable, args, kwargs, cache_key):
+    """
+    Internal function for caching asynchronous function calls.
+
+    This function is used internally by the scenario_cache decorator
+    to cache asynchronous function calls using joblib.Memory.
+
+    Args:
+        func: The async function to call and cache
+        args: Positional arguments for the function
+        kwargs: Keyword arguments for the function
+        cache_key: Cache key for deterministic caching
+
+    Returns:
+        The result of calling await func(*args, **kwargs)
+    """
+    return await func(*args, **kwargs)

scenario/config.py

@@ -1,33 +1,166 @@
 """
 Configuration module for Scenario.
+
+This module provides configuration classes for customizing the behavior of the
+Scenario testing framework, including LLM model settings, execution parameters,
+and debugging options.
 """
 
-from typing import TYPE_CHECKING, Any, Optional, Type, Union
+from typing import Optional, Union, ClassVar
 from pydantic import BaseModel
 
-if TYPE_CHECKING:
-    from scenario.scenario_agent_adapter import ScenarioAgentAdapter
 
-    ScenarioAgentType = ScenarioAgentAdapter
-else:
-    ScenarioAgentType = Any
+class ModelConfig(BaseModel):
+    """
+    Configuration for LLM model settings.
+
+    This class encapsulates all the parameters needed to configure an LLM model
+    for use with user simulator and judge agents in the Scenario framework.
+
+    Attributes:
+        model: The model identifier (e.g., "openai/gpt-4.1-mini", "anthropic/claude-3-sonnet")
+        api_key: Optional API key for the model provider
+        temperature: Sampling temperature for response generation (0.0 = deterministic, 1.0 = creative)
+        max_tokens: Maximum number of tokens to generate in responses
+
+    Example:
+        ```
+        model_config = ModelConfig(
+            model="openai/gpt-4.1-mini",
+            api_key="your-api-key",
+            temperature=0.1,
+            max_tokens=1000
+        )
+        ```
+    """
+
+    model: str
+    api_key: Optional[str] = None
+    temperature: float = 0.0
+    max_tokens: Optional[int] = None
 
 
 class ScenarioConfig(BaseModel):
     """
-    Configuration class for the Scenario library.
+    Global configuration class for the Scenario testing framework.
 
-    This allows users to set global configuration settings for the library,
-    such as the LLM provider and model to use for the testing agent.
+    This class allows users to set default behavior and parameters that apply
+    to all scenario executions, including the LLM model to use for simulator
+    and judge agents, execution limits, and debugging options.
+
+    Attributes:
+        default_model: Default LLM model configuration for agents (can be string or ModelConfig)
+        max_turns: Maximum number of conversation turns before scenario times out
+        verbose: Whether to show detailed output during execution (True/False or verbosity level)
+        cache_key: Key for caching scenario results to ensure deterministic behavior
+        debug: Whether to enable debug mode with step-by-step interaction
+
+    Example:
+        ```
+        # Configure globally for all scenarios
+        scenario.configure(
+            default_model="openai/gpt-4.1-mini",
+            max_turns=15,
+            verbose=True,
+            cache_key="my-test-suite-v1",
+            debug=False
+        )
+
+        # Or create a specific config instance
+        config = ScenarioConfig(
+            default_model=ModelConfig(
+                model="openai/gpt-4.1-mini",
+                temperature=0.2
+            ),
+            max_turns=20
+        )
+        ```
     """
 
-    testing_agent: Optional[Type[ScenarioAgentType]] = None
+    default_model: Optional[Union[str, ModelConfig]] = None
     max_turns: Optional[int] = 10
     verbose: Optional[Union[bool, int]] = True
     cache_key: Optional[str] = None
     debug: Optional[bool] = False
 
+    default_config: ClassVar[Optional["ScenarioConfig"]] = None
+
+    @classmethod
+    def configure(
+        cls,
+        default_model: Optional[str] = None,
+        max_turns: Optional[int] = None,
+        verbose: Optional[Union[bool, int]] = None,
+        cache_key: Optional[str] = None,
+        debug: Optional[bool] = None,
+    ) -> None:
+        """
+        Set global configuration settings for all scenario executions.
+
+        This method allows you to configure default behavior that will be applied
+        to all scenarios unless explicitly overridden in individual scenario runs.
+
+        Args:
+            default_model: Default LLM model identifier for user simulator and judge agents
+            max_turns: Maximum number of conversation turns before timeout (default: 10)
+            verbose: Enable verbose output during scenario execution
+            cache_key: Cache key for deterministic scenario behavior across runs
+            debug: Enable debug mode for step-by-step execution with user intervention
+
+        Example:
+            ```
+            import scenario
+
+            # Set up default configuration
+            scenario.configure(
+                default_model="openai/gpt-4.1-mini",
+                max_turns=15,
+                verbose=True,
+                debug=False
+            )
+
+            # All subsequent scenario runs will use these defaults
+            result = await scenario.run(
+                name="my test",
+                description="Test scenario",
+                agents=[my_agent, scenario.UserSimulatorAgent(), scenario.JudgeAgent()]
+            )
+            ```
+        """
+        existing_config = cls.default_config or ScenarioConfig()
+
+        cls.default_config = existing_config.merge(
+            ScenarioConfig(
+                default_model=default_model,
+                max_turns=max_turns,
+                verbose=verbose,
+                cache_key=cache_key,
+                debug=debug,
+            )
+        )
+
     def merge(self, other: "ScenarioConfig") -> "ScenarioConfig":
+        """
+        Merge this configuration with another configuration.
+
+        Values from the other configuration will override values in this
+        configuration where they are not None.
+
+        Args:
+            other: Another ScenarioConfig instance to merge with
+
+        Returns:
+            A new ScenarioConfig instance with merged values
+
+        Example:
+            ```
+            base_config = ScenarioConfig(max_turns=10, verbose=True)
+            override_config = ScenarioConfig(max_turns=20)
+
+            merged = base_config.merge(override_config)
+            # Result: max_turns=20, verbose=True
+            ```
+        """
         return ScenarioConfig(
             **{
                 **self.items(),
@@ -36,4 +169,17 @@ class ScenarioConfig(BaseModel):
         )
 
     def items(self):
+        """
+        Get configuration items as a dictionary.
+
+        Returns:
+            Dictionary of configuration key-value pairs, excluding None values
+
+        Example:
+            ```
+            config = ScenarioConfig(max_turns=15, verbose=True)
+            items = config.items()
+            # Result: {"max_turns": 15, "verbose": True}
+            ```
+        """
         return {k: getattr(self, k) for k in self.model_dump(exclude_none=True).keys()}

scenario/events/__init__.py

@@ -0,0 +1,66 @@
+"""
+Scenario events module for handling event publishing, processing, and reporting.
+
+This module provides event models, an event bus for processing, and utilities
+for converting between different message formats.
+"""
+
+# Core event types and models
+from .events import (
+    ScenarioEvent,
+    ScenarioRunStartedEvent,
+    ScenarioRunStartedEventMetadata,
+    ScenarioRunFinishedEvent,
+    ScenarioRunFinishedEventResults,
+    ScenarioRunFinishedEventVerdict,
+    ScenarioRunFinishedEventStatus,
+    ScenarioMessageSnapshotEvent,
+    MessageType,
+)
+
+# Event processing infrastructure
+from .event_bus import ScenarioEventBus
+from .event_reporter import EventReporter
+
+# Message utilities and types
+from .messages import (
+    Message,
+    UserMessage,
+    AssistantMessage,
+    SystemMessage,
+    ToolMessage,
+    ToolCall,
+    FunctionCall,
+)
+
+# Utility functions
+from .utils import convert_messages_to_ag_ui_messages
+
+__all__ = [
+    # Event types
+    "ScenarioEvent",
+    "ScenarioRunStartedEvent",
+    "ScenarioRunStartedEventMetadata",
+    "ScenarioRunFinishedEvent", 
+    "ScenarioRunFinishedEventResults",
+    "ScenarioRunFinishedEventVerdict",
+    "ScenarioRunFinishedEventStatus",
+    "ScenarioMessageSnapshotEvent",
+    "MessageType",
+    
+    # Event processing
+    "ScenarioEventBus",
+    "EventReporter",
+    
+    # Messages
+    "Message",
+    "UserMessage",
+    "AssistantMessage", 
+    "SystemMessage",
+    "ToolMessage",
+    "ToolCall",
+    "FunctionCall",
+    
+    # Utils
+    "convert_messages_to_ag_ui_messages",
+] 

scenario/events/event_bus.py

@@ -0,0 +1,175 @@
+from rx.subject.subject import Subject
+from rx import operators as ops
+from typing import Optional
+from datetime import datetime, UTC
+from .events import ScenarioEvent, ScenarioRunFinishedEvent
+from .event_reporter import EventReporter
+from typing import Any
+
+import asyncio
+
+
+class ScenarioEventBus:
+    """
+    Manages scenario event publishing, subscription, and processing pipeline using RxPY.
+
+    The EventBus provides a centralized event processing system that handles scenario
+    events asynchronously with retry logic and concurrent processing. It automatically
+    manages the event stream lifecycle and ensures all events are processed before
+    completion.
+
+    Events are processed concurrently to improve performance, and failed event
+    processing is automatically retried with exponential backoff.
+
+    Attributes:
+        _events: RxPY Subject for event stream management
+        _event_reporter: EventReporter instance for HTTP posting of events
+        _processing_complete: Async event to signal when all events are processed
+        _processing_task: Background task for event processing
+        _max_retries: Maximum number of retry attempts for failed event processing
+
+    Example:
+        ```python
+        # Create event bus with custom reporter
+        reporter = EventReporter(endpoint="https://api.langwatch.ai")
+        event_bus = ScenarioEventBus(event_reporter=reporter, max_retries=5)
+
+        # Start listening for events
+        await event_bus.listen()
+
+        # Publish events
+        event_bus.publish(scenario_started_event)
+        event_bus.publish(message_snapshot_event)
+        event_bus.publish(scenario_finished_event)  # This completes the stream
+
+        # Wait for all events to be processed
+        await event_bus.drain()
+        ```
+    """
+
+    def __init__(
+        self, event_reporter: Optional[EventReporter] = None, max_retries: int = 3
+    ):
+        """
+        Initialize the event bus with optional event reporter and retry configuration.
+
+        Args:
+            event_reporter: Optional EventReporter for HTTP posting of events.
+                          If not provided, a default EventReporter will be created.
+            max_retries: Maximum number of retry attempts for failed event processing.
+                       Defaults to 3 attempts with exponential backoff.
+        """
+        self._events = Subject()
+        # Use default EventReporter if none provided
+        self._event_reporter: EventReporter = event_reporter or EventReporter()
+        self._processing_complete = asyncio.Event()
+        self._processing_task: Optional[asyncio.Task[Any]] = None
+        self._max_retries = max_retries
+
+    def publish(self, event: ScenarioEvent) -> None:
+        """
+        Publishes an event into the processing pipeline.
+
+        This method adds an event to the RxPY stream for processing. The event
+        timestamp is automatically set to the current time in milliseconds if
+        not already provided. Publishing a ScenarioRunFinishedEvent automatically
+        completes the event stream.
+
+        Args:
+            event: The scenario event to publish. Must be a valid ScenarioEvent type.
+
+        Note:
+            Events are processed asynchronously in the background. Use `drain()`
+            to wait for all events to be processed after publishing.
+        """
+        # Convert to Unix timestamp in milliseconds
+        event.timestamp = int(datetime.now(UTC).timestamp() * 1000)
+        self._events.on_next(event)
+
+        if isinstance(event, ScenarioRunFinishedEvent):
+            self._events.on_completed()
+
+    async def listen(self) -> None:
+        """
+        Begins listening for and processing events.
+
+        This method sets up the RxPY event processing pipeline with concurrent
+        processing and automatic retry logic. It should be called before publishing
+        any events to ensure proper event handling.
+
+        The processing pipeline:
+        1. Receives events from the publish() method
+        2. Processes each event concurrently using asyncio tasks
+        3. Automatically retries failed events with exponential backoff
+        4. Completes when a ScenarioRunFinishedEvent is published
+
+        Note:
+            This method is idempotent - calling it multiple times has no effect
+            if the processing pipeline is already active.
+        """
+        if self._processing_task is not None:
+            return
+
+        async def process_single_event(event: ScenarioEvent, attempt: int = 1) -> bool:
+            """
+            Process a single event with retry logic.
+
+            Args:
+                event: The event to process
+                attempt: Current attempt number (1-based)
+
+            Returns:
+                True if processing succeeded, False if all retries failed
+            """
+            try:
+                if self._event_reporter:
+                    await self._event_reporter.post_event(event)
+                return True
+            except Exception as e:
+                if attempt >= self._max_retries:
+                    print(f"Failed to process event after {attempt} attempts: {e}")
+                    return False
+                print(
+                    f"Error processing event (attempt {attempt}/{self._max_retries}): {e}"
+                )
+                await asyncio.sleep(0.1 * (2 ** (attempt - 1)))
+                return await process_single_event(event, attempt + 1)
+
+        def process_event(event: ScenarioEvent) -> asyncio.Task[bool]:
+            """Create an asyncio task to process an event concurrently."""
+            loop = asyncio.get_event_loop()
+            return loop.create_task(process_single_event(event))
+
+        # Set up the event processing pipeline with concurrent processing
+        self._events.pipe(ops.flat_map(lambda event: process_event(event))).subscribe(
+            on_next=lambda success: None,
+            on_completed=lambda: self._processing_complete.set(),
+            on_error=lambda e: print(f"Unexpected error in event stream: {e}"),
+        )
+
+    async def drain(self) -> None:
+        """
+        Waits for all events to be processed after the stream is completed.
+
+        This method blocks until all events in the processing pipeline have been
+        handled. It should be called after publishing all events to ensure
+        proper cleanup and that no events are lost.
+
+        Note:
+            This method will wait indefinitely if the event stream has not been
+            completed (i.e., if no ScenarioRunFinishedEvent has been published).
+        """
+        await self._processing_complete.wait()
+
+    def is_completed(self) -> bool:
+        """
+        Returns whether the event bus has completed processing all events.
+
+        This method provides a non-blocking way to check if all events have
+        been processed. It's useful for monitoring the state of the event bus
+        without blocking execution.
+
+        Returns:
+            True if all events have been processed, False otherwise
+        """
+        return self._processing_complete.is_set()