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

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()

scenario/events/event_reporter.py

@@ -0,0 +1,83 @@
+import logging
+import os
+import httpx
+from typing import Optional
+from .events import ScenarioEvent
+
+
+class EventReporter:
+    """
+    Handles HTTP posting of scenario events to external endpoints.
+
+    Single responsibility: Send events via HTTP to configured endpoints
+    with proper authentication and error handling.
+
+    Args:
+        endpoint (str, optional): The base URL to post events to. Defaults to LANGWATCH_ENDPOINT env var.
+        api_key (str, optional): The API key for authentication. Defaults to LANGWATCH_API_KEY env var.
+
+    Example:
+        event = {
+            "type": "SCENARIO_RUN_STARTED",
+            "batch_run_id": "batch-1",
+            "scenario_id": "scenario-1",
+            "scenario_run_id": "run-1",
+            "metadata": {
+                "name": "test",
+                "description": "test scenario"
+            }
+        }
+
+        reporter = EventReporter(endpoint="https://api.langwatch.ai", api_key="test-api-key")
+        await reporter.post_event(event)
+    """
+
+    def __init__(self, endpoint: Optional[str] = None, api_key: Optional[str] = None):
+        self.endpoint = endpoint or os.getenv("LANGWATCH_ENDPOINT")
+        self.api_key = api_key or os.getenv("LANGWATCH_API_KEY", "")
+        self.logger = logging.getLogger("EventReporter")
+
+    async def post_event(self, event: ScenarioEvent):
+        """
+        Posts an event to the configured endpoint.
+
+        Args:
+            event: A dictionary containing the event data
+
+        Returns:
+            None - logs success/failure internally
+        """
+        event_type = event.type_
+        self.logger.info(f"[{event_type}] Publishing event ({event.scenario_run_id})")
+
+        if not self.endpoint:
+            self.logger.warning(
+                "No LANGWATCH_ENDPOINT configured, skipping event posting"
+            )
+            return
+
+        try:
+            async with httpx.AsyncClient() as client:
+                response = await client.post(
+                    f"{self.endpoint}/api/scenario-events",
+                    json=event.to_dict(),
+                    headers={
+                        "Content-Type": "application/json",
+                        "X-Auth-Token": self.api_key,
+                    },
+                )
+                self.logger.info(f"[{event_type}] POST response status: {response.status_code} ({event.scenario_run_id})")
+                
+                if response.is_success:
+                    data = response.json()
+                    self.logger.info(f"[{event_type}] POST response: {data} ({event.scenario_run_id})")
+                else:
+                    error_text = response.text
+                    self.logger.error(
+                        f"[{event_type}] Event POST failed: status={response.status_code}, "
+                        f"reason={response.reason_phrase}, error={error_text}, "
+                        f"event={event}"
+                    )
+        except Exception as error:
+            self.logger.error(
+                f"[{event_type}] Event POST error: {error}, event={event}, endpoint={self.endpoint}") 

scenario/events/events.py

@@ -0,0 +1,169 @@
+"""
+Exports scenario event models from the generated LangWatch API client,
+renaming the auto-generated types to clean, meaningful names.
+
+This ensures all event types are always in sync with the OpenAPI spec and
+the backend, and provides a single import location for event models.
+
+If you need to add custom logic or helpers, you can extend or wrap these models here.
+"""
+
+from typing import Union, Any, Optional
+from scenario.generated.langwatch_api_client.lang_watch_api_client.models import (
+    PostApiScenarioEventsBodyType0,
+    PostApiScenarioEventsBodyType0Metadata as ScenarioRunStartedEventMetadata,
+    PostApiScenarioEventsBodyType1,
+    PostApiScenarioEventsBodyType1ResultsType0 as ScenarioRunFinishedEventResults,
+    PostApiScenarioEventsBodyType1ResultsType0Verdict as ScenarioRunFinishedEventVerdict,
+    PostApiScenarioEventsBodyType1Status as ScenarioRunFinishedEventStatus,
+    PostApiScenarioEventsBodyType2,
+    # Message types for the snapshot event
+    PostApiScenarioEventsBodyType2MessagesItemType0,
+    PostApiScenarioEventsBodyType2MessagesItemType1,
+    PostApiScenarioEventsBodyType2MessagesItemType2,
+    PostApiScenarioEventsBodyType2MessagesItemType3,
+    PostApiScenarioEventsBodyType2MessagesItemType4,
+)
+
+# Type alias for message types
+MessageType = Union[
+    PostApiScenarioEventsBodyType2MessagesItemType0,
+    PostApiScenarioEventsBodyType2MessagesItemType1, 
+    PostApiScenarioEventsBodyType2MessagesItemType2,
+    PostApiScenarioEventsBodyType2MessagesItemType3,
+    PostApiScenarioEventsBodyType2MessagesItemType4,
+]
+
+class ScenarioRunStartedEvent(PostApiScenarioEventsBodyType0):
+    """
+    Event published when a scenario run begins execution.
+    
+    Automatically sets type_ to "SCENARIO_RUN_STARTED" and includes metadata
+    about the scenario (name, description, etc.).
+    
+    Args:
+        batch_run_id (str): Unique identifier for the batch of scenario runs
+        scenario_id (str): Unique identifier for the scenario definition
+        scenario_run_id (str): Unique identifier for this specific run
+        metadata (ScenarioRunStartedEventMetadata): Scenario details like name and description
+        timestamp (Optional[int], optional): Unix timestamp in milliseconds, auto-generated if not provided
+        raw_event (Optional[Any], optional): Raw event data
+        scenario_set_id (Optional[str], optional): Set identifier, defaults to "default"
+    """
+    def __init__(
+        self,
+        batch_run_id: str,
+        scenario_id: str,
+        scenario_run_id: str,
+        metadata: ScenarioRunStartedEventMetadata,
+        timestamp: int,
+        raw_event: Optional[Any] = None,
+        scenario_set_id: Optional[str] = "default"
+    ):
+        super().__init__(
+            type_="SCENARIO_RUN_STARTED",
+            batch_run_id=batch_run_id,
+            scenario_id=scenario_id,
+            scenario_run_id=scenario_run_id,
+            metadata=metadata,
+            timestamp=timestamp,
+            raw_event=raw_event,
+            scenario_set_id=scenario_set_id or "default"
+        )
+
+class ScenarioRunFinishedEvent(PostApiScenarioEventsBodyType1):
+    """
+    Event published when a scenario run completes execution.
+    
+    Automatically sets type_ to "SCENARIO_RUN_FINISHED" and includes results
+    with verdict (PASS/FAIL/SUCCESS) and reasoning.
+    
+    Args:
+        batch_run_id (str): Unique identifier for the batch of scenario runs
+        scenario_id (str): Unique identifier for the scenario definition
+        scenario_run_id (str): Unique identifier for this specific run
+        status (ScenarioRunFinishedEventStatus): Overall execution status
+        timestamp (Optional[int], optional): Unix timestamp in milliseconds, auto-generated if not provided
+        raw_event (Optional[Any], optional): Raw event data
+        scenario_set_id (Optional[str], optional): Set identifier, defaults to "default"
+        results (Optional[ScenarioRunFinishedEventResults], optional): Verdict and reasoning for the outcome
+    """
+    def __init__(
+        self,
+        batch_run_id: str,
+        scenario_id: str,
+        scenario_run_id: str,
+        status: ScenarioRunFinishedEventStatus,
+        timestamp: int,
+        results: Optional[ScenarioRunFinishedEventResults] = None,
+        raw_event: Optional[Any] = None,
+        scenario_set_id: Optional[str] = "default",
+    ):
+        super().__init__(
+            type_="SCENARIO_RUN_FINISHED",
+            batch_run_id=batch_run_id,
+            scenario_id=scenario_id,
+            scenario_run_id=scenario_run_id,
+            status=status,
+            timestamp=timestamp,
+            raw_event=raw_event,
+            scenario_set_id=scenario_set_id or "default",
+            results=results
+        )
+
+class ScenarioMessageSnapshotEvent(PostApiScenarioEventsBodyType2):
+    """
+    Event published to capture intermediate state during scenario execution.
+    
+    Automatically sets type_ to "SCENARIO_MESSAGE_SNAPSHOT" and allows tracking
+    of messages, context, or other runtime data during scenario processing.
+    
+    Args:
+        batch_run_id (str): Unique identifier for the batch of scenario runs
+        scenario_id (str): Unique identifier for the scenario definition
+        scenario_run_id (str): Unique identifier for this specific run
+        messages (list[MessageType]): List of message objects in the conversation
+        timestamp (Optional[int], optional): Unix timestamp in milliseconds, auto-generated if not provided
+        raw_event (Optional[Any], optional): Raw event data
+        scenario_set_id (Optional[str], optional): Set identifier, defaults to "default"
+    """
+    def __init__(
+        self,
+        batch_run_id: str,
+        scenario_id: str,
+        scenario_run_id: str,
+        messages: list[MessageType],
+        timestamp: int,
+        raw_event: Optional[Any] = None,
+        scenario_set_id: Optional[str] = "default"
+    ):
+        super().__init__(
+            type_="SCENARIO_MESSAGE_SNAPSHOT",
+            batch_run_id=batch_run_id,
+            scenario_id=scenario_id,
+            scenario_run_id=scenario_run_id,
+            messages=messages,
+            timestamp=timestamp,
+            raw_event=raw_event,
+            scenario_set_id=scenario_set_id or "default"
+        )
+
+# Union type for all supported event types
+ScenarioEvent = Union[
+    ScenarioRunStartedEvent,
+    ScenarioRunFinishedEvent, 
+    ScenarioMessageSnapshotEvent
+]
+
+
+__all__ = [
+    "ScenarioEvent",
+    "ScenarioRunStartedEvent",
+    "ScenarioRunStartedEventMetadata",
+    "ScenarioRunFinishedEvent",
+    "ScenarioRunFinishedEventResults",
+    "ScenarioRunFinishedEventVerdict",
+    "ScenarioRunFinishedEventStatus",
+    "ScenarioMessageSnapshotEvent",
+    "MessageType",
+]

scenario/events/messages.py

@@ -0,0 +1,84 @@
+from typing import Union, Optional, List
+from ag_ui.core import (
+    UserMessage as AgUiUserMessage,
+    AssistantMessage as AgUiAssistantMessage,
+    SystemMessage as AgUiSystemMessage,
+    ToolMessage as AgUiToolMessage,
+    ToolCall as AgUiToolCall,
+    FunctionCall as AgUiFunctionCall,
+)
+
+class UserMessage(AgUiUserMessage):
+    """
+    An AG-UI user message extended with the to_dict method.
+    Enforces role='user' and requires content.
+    """
+    def __init__(self, id: str, content: str, name: Optional[str] = None):
+        super().__init__(id=id, role="user", content=content, name=name)
+        
+    def to_dict(self):
+        """Convert the UserMessage to a dictionary representation."""
+        return self.model_dump(exclude_none=True)
+
+class AssistantMessage(AgUiAssistantMessage):
+    """
+    An AG-UI assistant message extended with the to_dict method.
+    Enforces role='assistant' and allows optional content and tool_calls.
+    """
+    def __init__(self, id: str, content: Optional[str] = None, tool_calls: Optional[List['ToolCall']] = None, name: Optional[str] = None):
+        super().__init__(id=id, role="assistant", content=content, tool_calls=tool_calls, name=name)
+        
+    def to_dict(self):
+        """Convert the AssistantMessage to a dictionary representation."""
+        return self.model_dump(exclude_none=True)
+
+class SystemMessage(AgUiSystemMessage):
+    """
+    An AG-UI system message extended with the to_dict method.
+    Enforces role='system' and requires content.
+    """
+    def __init__(self, id: str, content: str, name: Optional[str] = None):
+        super().__init__(id=id, role="system", content=content, name=name)
+        
+    def to_dict(self):
+        """Convert the SystemMessage to a dictionary representation."""
+        return self.model_dump(exclude_none=True)
+
+class ToolMessage(AgUiToolMessage):
+    """
+    An AG-UI tool message extended with the to_dict method.
+    Enforces role='tool' and requires content and tool_call_id.
+    """
+    def __init__(self, id: str, content: str, tool_call_id: str):
+        super().__init__(id=id, role="tool", content=content, tool_call_id=tool_call_id)
+        
+    def to_dict(self):
+        """Convert the ToolMessage to a dictionary representation."""
+        return self.model_dump(exclude_none=True)
+
+class ToolCall(AgUiToolCall):
+    """
+    An AG-UI tool call extended with the to_dict method.
+    Enforces type='function' and requires id and function.
+    """
+    def __init__(self, id: str, function: 'FunctionCall'):
+        super().__init__(id=id, type="function", function=function)
+        
+    def to_dict(self):
+        """Convert the ToolCall to a dictionary representation."""
+        return self.model_dump(exclude_none=True)
+
+class FunctionCall(AgUiFunctionCall):
+    """
+    An AG-UI function call extended with the to_dict method.
+    Requires name and arguments.
+    """
+    def __init__(self, name: str, arguments: str):
+        super().__init__(name=name, arguments=arguments)
+        
+    def to_dict(self):
+        """Convert the FunctionCall to a dictionary representation."""
+        return self.model_dump(exclude_none=True)
+
+# Union type alias for all message types
+Message = Union[UserMessage, AssistantMessage, SystemMessage, ToolMessage, ToolCall, FunctionCall]

scenario/events/utils.py

@@ -0,0 +1,86 @@
+from openai.types.chat.chat_completion_message_param import ChatCompletionMessageParam
+from .messages import UserMessage, AssistantMessage, SystemMessage, ToolMessage, ToolCall, FunctionCall
+from typing import List, Union
+
+import uuid
+
+# Define the correct Message type for the return value
+Message = Union[UserMessage, AssistantMessage, SystemMessage, ToolMessage]
+
+def convert_messages_to_ag_ui_messages(messages: list[ChatCompletionMessageParam]) -> list[Message]:
+    """
+    Converts OpenAI ChatCompletionMessageParam messages to ag_ui Message format.
+    
+    This function transforms messages from OpenAI's format to the ag_ui protocol
+    format for consistent message handling across the scenario framework.
+    
+    Args:
+        messages: List of OpenAI ChatCompletionMessageParam messages
+        
+    Returns:
+        List of ag_ui Message objects
+        
+    Raises:
+        ValueError: If message role is not supported or message format is invalid
+    """
+
+    converted_messages: list[Message] = []
+    
+    for i, message in enumerate(messages):
+        # Generate unique ID for each message
+        message_id = message.get("id") or str(uuid.uuid4())
+        
+        role = message.get("role")
+        content = message.get("content")
+        
+        if role == "user":
+            if not content:
+                raise ValueError(f"User message at index {i} missing required content")
+            converted_messages.append(UserMessage(
+                id=message_id,
+                content=str(content)
+            ))
+        elif role == "assistant":
+            # Handle tool calls if present
+            tool_calls = message.get("tool_calls")
+            ag_ui_tool_calls: List[ToolCall] | None = None
+            
+            if tool_calls:
+                ag_ui_tool_calls = []
+                for tool_call in tool_calls:
+                    ag_ui_tool_calls.append(ToolCall(
+                        id=tool_call.get("id", str(uuid.uuid4())),
+                        function=FunctionCall(
+                            name=tool_call["function"]["name"],
+                            arguments=tool_call["function"]["arguments"]
+                        )
+                    ))
+            
+            converted_messages.append(AssistantMessage(
+                id=message_id,
+                content=str(content) if content else None,
+                tool_calls=ag_ui_tool_calls
+            ))
+        elif role == "system":
+            if not content:
+                raise ValueError(f"System message at index {i} missing required content")
+            converted_messages.append(SystemMessage(
+                id=message_id,
+                content=str(content)
+            ))
+        elif role == "tool":
+            tool_call_id = message.get("tool_call_id")
+            if not tool_call_id:
+                raise ValueError(f"Tool message at index {i} missing required tool_call_id")
+            if not content:
+                raise ValueError(f"Tool message at index {i} missing required content")
+                
+            converted_messages.append(ToolMessage(
+                id=message_id,
+                content=str(content),
+                tool_call_id=tool_call_id
+            ))
+        else:
+            raise ValueError(f"Unsupported message role '{role}' at index {i}")
+    
+    return converted_messages

scenario/judge_agent.py

@@ -19,7 +19,7 @@ from scenario.cache import scenario_cache
 from scenario.agent_adapter import AgentAdapter
 from scenario.config import ModelConfig, ScenarioConfig
 
-from .error_messages import agent_not_configured_error_message
+from ._error_messages import agent_not_configured_error_message
 from .types import AgentInput, AgentReturnTypes, AgentRole, ScenarioResult
 
 
@@ -48,7 +48,7 @@ class JudgeAgent(AgentAdapter):
         system_prompt: Custom system prompt to override default judge behavior
 
     Example:
-        ```python
+        ```
         import scenario
 
         # Basic judge agent with criteria
@@ -133,14 +133,12 @@ class JudgeAgent(AgentAdapter):
             Exception: If no model is configured either in parameters or global config
 
         Example:
-            ```python
+            ```
             # Customer service judge
             cs_judge = JudgeAgent(
                 criteria=[
-                    "Agent is polite and professional",
-                    "Agent addresses the customer's specific concern",
-                    "Agent offers appropriate solutions or next steps",
-                    "Agent does not make promises the company cannot keep"
+                    "Agent replies with the refund policy",
+                    "Agent offers next steps for the customer",
                 ],
                 temperature=0.1
             )
@@ -148,9 +146,8 @@ class JudgeAgent(AgentAdapter):
             # Technical accuracy judge
             tech_judge = JudgeAgent(
                 criteria=[
-                    "Code examples compile without errors",
-                    "Security vulnerabilities are not introduced",
-                    "Best practices are recommended"
+                    "Agent adds a code review pointing out the code compilation errors",
+                    "Agent adds a code review about the missing security headers"
                 ],
                 system_prompt="You are a senior software engineer reviewing code for production use."
             )
@@ -210,24 +207,6 @@ class JudgeAgent(AgentAdapter):
             Exception: If the judge cannot make a valid decision or if there's an
                       error in the evaluation process
 
-        Example:
-            The judge evaluates conversations like this:
-
-            ```
-            Conversation so far:
-            User: "I need help with authentication"
-            Agent: "I can help! What authentication method are you using?"
-            User: "JWT tokens"
-            Agent: "Here's how to implement JWT securely: [detailed code example]"
-
-            Judge evaluation:
-            - ✓ Agent provides helpful responses
-            - ✓ Agent asks relevant follow-up questions
-            - ✓ Security best practices are mentioned
-
-            Decision: CONTINUE (all criteria being met so far)
-            ```
-
         Note:
             - Returns empty list [] to continue the scenario
             - Returns ScenarioResult to end with success/failure