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

scenario/testing_agent.py

@@ -4,25 +4,24 @@ TestingAgent module: defines the testing agent that interacts with the agent und
 
 import json
 import logging
-from typing import TYPE_CHECKING, Dict, List, Any, Optional, Union, cast
-from pydantic import BaseModel
+import re
+from typing import Optional, Type, cast
 
 from litellm import Choices, completion
 from litellm.files.main import ModelResponse
 
 from scenario.cache import scenario_cache
-from scenario.utils import safe_attr_or_key
+from scenario.scenario_agent_adapter import ScenarioAgentAdapter
+from scenario.utils import reverse_roles
 
-from .result import ScenarioResult
-
-if TYPE_CHECKING:
-    from scenario.scenario import Scenario
+from .error_messages import testing_agent_not_configured_error_message
+from .types import AgentInput, AgentReturnTypes, ScenarioAgentRole, ScenarioResult
 
 
 logger = logging.getLogger("scenario")
 
 
-class TestingAgent(BaseModel):
+class TestingAgent(ScenarioAgentAdapter):
     """
     The Testing Agent that interacts with the agent under test.
 
@@ -32,7 +31,9 @@ class TestingAgent(BaseModel):
     3. Determining when to end the test and return a result
     """
 
-    model: str
+    roles = {ScenarioAgentRole.USER, ScenarioAgentRole.JUDGE}
+
+    model: str = ""
     api_key: Optional[str] = None
     temperature: float = 0.0
     max_tokens: Optional[int] = None
@@ -40,14 +41,36 @@ class TestingAgent(BaseModel):
     # To prevent pytest from thinking this is actually a test class
     __test__ = False
 
+    def __init__(self, input: AgentInput):
+        super().__init__(input)
+
+        if not self.model:
+            raise Exception(testing_agent_not_configured_error_message)
+
+    @classmethod
+    def with_config(
+        cls,
+        model: str,
+        api_key: Optional[str] = None,
+        temperature: float = 0.0,
+        max_tokens: Optional[int] = None,
+    ) -> Type["TestingAgent"]:
+        class TestingAgentWithConfig(cls):
+            def __init__(self, input: AgentInput):
+                self.model = model
+                self.api_key = api_key
+                self.temperature = temperature
+                self.max_tokens = max_tokens
+
+                super().__init__(input)
+
+        return TestingAgentWithConfig
+
     @scenario_cache(ignore=["scenario"])
-    def generate_next_message(
+    async def call(
         self,
-        scenario: "Scenario",
-        conversation: List[Dict[str, Any]],
-        first_message: bool = False,
-        last_message: bool = False,
-    ) -> Union[str, ScenarioResult]:
+        input: AgentInput,
+    ) -> AgentReturnTypes:
         """
         Generate the next message in the conversation based on history OR
         return a ScenarioResult if the test should conclude.
@@ -57,6 +80,8 @@ class TestingAgent(BaseModel):
           - A ScenarioResult (if the test should conclude)
         """
 
+        scenario = input.scenario_state.scenario
+
         messages = [
             {
                 "role": "system",
@@ -74,37 +99,34 @@ Your goal (assistant) is to interact with the Agent Under Test (user) as if you
 {scenario.description}
 </scenario>
 
-<strategy>
-{scenario.strategy or "Start with a first message and guide the conversation to play out the scenario."}
-</strategy>
-
-<success_criteria>
-{json.dumps(scenario.success_criteria, indent=2)}
-</success_criteria>
-
-<failure_criteria>
-{json.dumps(scenario.failure_criteria, indent=2)}
-</failure_criteria>
+<criteria>
+{"\n".join([f"{idx + 1}. {criterion}" for idx, criterion in enumerate(scenario.criteria)])}
+</criteria>
 
 <execution_flow>
 1. Generate the first message to start the scenario
 2. After the Agent Under Test (user) responds, generate the next message to send to the Agent Under Test, keep repeating step 2 until criterias match
-3. If the test should end, use the finish_test tool to determine if success or failure criteria have been met
+3. If the test should end, use the finish_test tool to determine if all the criteria have been met
 </execution_flow>
 
 <rules>
-1. Test should end immediately if a failure criteria is triggered
-2. Test should continue until all success criteria have been met
+1. Test should end immediately if a criteria mentioning something the agent should NOT do is met
+2. Test should continue until all scenario goals have been met to try going through all the criteria
 3. DO NOT make any judgment calls that are not explicitly listed in the success or failure criteria, withhold judgement if necessary
 4. DO NOT carry over any requests yourself, YOU ARE NOT the assistant today, wait for the user to do it
 </rules>
 """,
             },
             {"role": "assistant", "content": "Hello, how can I help you today?"},
-            *conversation,
+            *input.messages,
         ]
 
-        if last_message:
+        is_first_message = len(input.messages) == 0
+        is_last_message = (
+            input.scenario_state.current_turn == input.scenario_state.scenario.max_turns
+        )
+
+        if is_last_message:
             messages.append(
                 {
                     "role": "user",
@@ -122,25 +144,17 @@ if you don't have enough information to make a verdict, say inconclusive with ma
         # User to assistant role reversal
         # LLM models are biased to always be the assistant not the user, so we need to do this reversal otherwise models like GPT 4.5 is
         # super confused, and Claude 3.7 even starts throwing exceptions.
-        for message in messages:
-            # Can't reverse tool calls
-            if not safe_attr_or_key(message, "content") or safe_attr_or_key(
-                message, "tool_calls"
-            ):
-                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"
-                elif getattr(message, "role", None) == "assistant":
-                    message.role = "user"
+        messages = reverse_roles(messages)
 
         # Define the tool
+        criteria_names = [
+            re.sub(
+                r"[^a-zA-Z0-9]",
+                "_",
+                criterion.replace(" ", "_").replace("'", "").lower(),
+            )[:70]
+            for criterion in scenario.criteria
+        ]
         tools = [
             {
                 "type": "function",
@@ -151,46 +165,46 @@ if you don't have enough information to make a verdict, say inconclusive with ma
                     "parameters": {
                         "type": "object",
                         "properties": {
-                            "verdict": {
-                                "type": "string",
-                                "enum": ["success", "failure", "inconclusive"],
-                                "description": "The final verdict of the test",
-                            },
-                            "reasoning": {
-                                "type": "string",
-                                "description": "Explanation of why this verdict was reached",
-                            },
-                            "details": {
+                            "criteria": {
                                 "type": "object",
                                 "properties": {
-                                    "met_criteria": {
-                                        "type": "array",
-                                        "items": {"type": "string"},
-                                        "description": "List of success criteria that have been met",
-                                    },
-                                    "unmet_criteria": {
-                                        "type": "array",
-                                        "items": {"type": "string"},
-                                        "description": "List of success criteria that have not been met",
-                                    },
-                                    "triggered_failures": {
-                                        "type": "array",
-                                        "items": {"type": "string"},
-                                        "description": "List of failure criteria that have been triggered",
-                                    },
+                                    criteria_names[idx]: {
+                                        "enum": [True, False, "inconclusive"],
+                                        "description": criterion,
+                                    }
+                                    for idx, criterion in enumerate(scenario.criteria)
                                 },
-                                "required": ["met_criteria", "unmet_criteria", "triggered_failures"],
+                                "required": criteria_names,
                                 "additionalProperties": False,
-                                "description": "Detailed information about criteria evaluation",
+                                "description": "Strict verdict for each criterion",
+                            },
+                            "reasoning": {
+                                "type": "string",
+                                "description": "Explanation of what the final verdict should be",
+                            },
+                            "verdict": {
+                                "type": "string",
+                                "enum": ["success", "failure", "inconclusive"],
+                                "description": "The final verdict of the test",
                             },
                         },
-                        "required": ["verdict", "reasoning", "details"],
+                        "required": ["criteria", "reasoning", "verdict"],
                         "additionalProperties": False,
                     },
                 },
             }
         ]
 
+        enforce_judgment = input.requested_role == ScenarioAgentRole.JUDGE
+        has_criteria = len(scenario.criteria) > 0
+
+        if enforce_judgment and not has_criteria:
+            return ScenarioResult(
+                success=False,
+                messages=[],
+                reasoning="TestingAgent was called as a judge, but it has no criteria to judge against",
+            )
+
         response = cast(
             ModelResponse,
             completion(
@@ -198,8 +212,16 @@ if you don't have enough information to make a verdict, say inconclusive with ma
                 messages=messages,
                 temperature=self.temperature,
                 max_tokens=self.max_tokens,
-                tools=tools if not first_message else None,
-                tool_choice="required" if last_message else None,
+                tools=(
+                    tools
+                    if (not is_first_message or enforce_judgment) and has_criteria
+                    else None
+                ),
+                tool_choice=(
+                    "required"
+                    if (is_last_message or enforce_judgment) and has_criteria
+                    else None
+                ),
             ),
         )
 
@@ -216,36 +238,27 @@ if you don't have enough information to make a verdict, say inconclusive with ma
                         args = json.loads(tool_call.function.arguments)
                         verdict = args.get("verdict", "inconclusive")
                         reasoning = args.get("reasoning", "No reasoning provided")
-                        details = args.get("details", {})
-
-                        met_criteria = details.get("met_criteria", [])
-                        unmet_criteria = details.get("unmet_criteria", [])
-                        triggered_failures = details.get("triggered_failures", [])
+                        criteria = args.get("criteria", {})
+
+                        passed_criteria = [
+                            scenario.criteria[idx]
+                            for idx, criterion in enumerate(criteria.values())
+                            if criterion == True
+                        ]
+                        failed_criteria = [
+                            scenario.criteria[idx]
+                            for idx, criterion in enumerate(criteria.values())
+                            if criterion == False
+                        ]
 
                         # Return the appropriate ScenarioResult based on the verdict
-                        if verdict == "success":
-                            return ScenarioResult.success_result(
-                                conversation=conversation,
-                                reasoning=reasoning,
-                                met_criteria=met_criteria,
-                            )
-                        elif verdict == "failure":
-                            return ScenarioResult.failure_result(
-                                conversation=conversation,
-                                reasoning=reasoning,
-                                met_criteria=met_criteria,
-                                unmet_criteria=unmet_criteria,
-                                triggered_failures=triggered_failures,
-                            )
-                        else:  # inconclusive
-                            return ScenarioResult(
-                                success=False,
-                                conversation=conversation,
-                                reasoning=reasoning,
-                                met_criteria=met_criteria,
-                                unmet_criteria=unmet_criteria,
-                                triggered_failures=triggered_failures,
-                            )
+                        return ScenarioResult(
+                            success=verdict == "success",
+                            messages=messages,
+                            reasoning=reasoning,
+                            passed_criteria=passed_criteria,
+                            failed_criteria=failed_criteria,
+                        )
                     except json.JSONDecodeError:
                         logger.error("Failed to parse tool call arguments")
 
@@ -254,12 +267,13 @@ if you don't have enough information to make a verdict, say inconclusive with ma
             if message_content is None:
                 # If invalid tool call, raise an error
                 if message.tool_calls:
-                    raise Exception(f"Invalid tool call from testing agent: {message.tool_calls.__repr__()}")
+                    raise Exception(
+                        f"Invalid tool call from testing agent: {message.tool_calls.__repr__()}"
+                    )
                 raise Exception(f"No response from LLM: {response.__repr__()}")
 
-            return message_content
+            return {"role": "user", "content": message_content}
         else:
             raise Exception(
                 f"Unexpected response format from LLM: {response.__repr__()}"
             )
-

scenario/types.py

@@ -0,0 +1,96 @@
+from enum import Enum
+from pydantic import BaseModel, Field, SkipValidation
+from typing import (
+    TYPE_CHECKING,
+    Annotated,
+    Any,
+    Awaitable,
+    Callable,
+    Coroutine,
+    Dict,
+    List,
+    Optional,
+    Union,
+)
+
+from openai.types.chat import ChatCompletionMessageParam, ChatCompletionUserMessageParam
+
+# Prevent circular imports + Pydantic breaking
+if TYPE_CHECKING:
+    from scenario.scenario_executor import ScenarioExecutor
+
+    ScenarioExecutorType = ScenarioExecutor
+else:
+    ScenarioExecutorType = Any
+
+
+class ScenarioAgentRole(Enum):
+    USER = "User"
+    AGENT = "Agent"
+    JUDGE = "Judge"
+
+
+class AgentInput(BaseModel):
+    thread_id: str
+    # Prevent pydantic from validating/parsing the messages and causing issues: https://github.com/pydantic/pydantic/issues/9541
+    messages: Annotated[List[ChatCompletionMessageParam], SkipValidation]
+    new_messages: Annotated[List[ChatCompletionMessageParam], SkipValidation]
+    context: Dict[str, Any]
+    requested_role: ScenarioAgentRole
+    scenario_state: ScenarioExecutorType = Field(exclude=True)
+
+    def last_new_user_message(self) -> ChatCompletionUserMessageParam:
+        user_messages = [m for m in self.new_messages if m["role"] == "user"]
+        if not user_messages:
+            raise ValueError(
+                "No new user messages found, did you mean to call the assistant twice? Perhaps change your adapter to use the full messages list instead."
+            )
+        return user_messages[-1]
+
+    def last_new_user_message_str(self) -> str:
+        content = self.last_new_user_message()["content"]
+        if type(content) != str:
+            raise ValueError(
+                f"Last user message is not a string: {content.__repr__()}. Please use the full messages list instead."
+            )
+        return content
+
+
+class ScenarioResult(BaseModel):
+    """
+    Represents the results of a scenario test run.
+
+    Attributes:
+        success: Whether the scenario passed
+        conversation: The conversation history
+        reasoning: Reasoning for the result
+        passed_criteria: List of criteria that were met
+        failed_criteria: List of criteria that were not met
+    """
+
+    success: bool
+    messages: List[ChatCompletionMessageParam]
+    reasoning: Optional[str] = None
+    passed_criteria: List[str] = []
+    failed_criteria: List[str] = []
+    total_time: Optional[float] = None
+    agent_time: Optional[float] = None
+
+    def __repr__(self) -> str:
+        """Provide a concise representation for debugging."""
+        status = "PASSED" if self.success else "FAILED"
+        return f"ScenarioResult(success={self.success}, status={status}, reasoning='{self.reasoning or 'None'}')"
+
+
+AgentReturnTypes = Union[
+    str, ChatCompletionMessageParam, List[ChatCompletionMessageParam], ScenarioResult
+]
+
+# TODO: remove the optional ScenarioResult return type from here, use events instead
+ScriptStep = Union[
+    Callable[["ScenarioExecutor"], None],
+    Callable[["ScenarioExecutor"], Optional[ScenarioResult]],
+    # Async as well
+    Callable[["ScenarioExecutor"], Awaitable[None]],
+    Callable[["ScenarioExecutor"], Awaitable[Optional[ScenarioResult]]],
+]

scenario/utils.py

@@ -1,6 +1,16 @@
 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 json
@@ -14,12 +24,18 @@ 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):
     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)
 
 
@@ -46,7 +62,9 @@ def title_case(string):
     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]
+):
     for msg in messages:
         role = safe_attr_or_key(msg, "role")
         content = safe_attr_or_key(msg, "content")
@@ -61,9 +79,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(
@@ -91,9 +112,12 @@ def _take_maybe_json_first_lines(string, max_lines=5):
 
 console = Console()
 
+
 class TextFirstSpinner(Spinner):
     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 +129,9 @@ 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
+):
     if not enabled:
         yield
     else:
@@ -119,3 +145,120 @@ 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:
+    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]:
+    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,
+                ),
+            )
+        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.
+    """
+
+    for message in messages.copy():
+        # Can't reverse tool calls
+        if not safe_attr_or_key(message, "content") or safe_attr_or_key(
+            message, "tool_calls"
+        ):
+            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
+
+    return messages
+
+
+async def await_if_awaitable(value: T) -> T:
+    if isinstance(value, Awaitable):
+        return await value
+    else:
+        return value

langwatch_scenario-0.1.3.dist-info/RECORD

@@ -1,15 +0,0 @@
-scenario/__init__.py,sha256=LfCjOpbn55jYBBZHyMSZtRAWeCDFn4z4OhAyFnu8aMg,602
-scenario/cache.py,sha256=sYu16SAf-BnVYkWSlEDzpyynJGIQyNYsgMXPgCqEnmk,1719
-scenario/config.py,sha256=5UVBmuQDtni0Yu00bMh5p0xMGsrymYVRftXBGTsi2fI,802
-scenario/error_messages.py,sha256=ZMcAOKJmKaLIinMZ0yBIOgDhPfeJH0uZxIEmolRArtc,2344
-scenario/pytest_plugin.py,sha256=BuBbyKLa-t9AFVn9EETl7OvGSt__dFO7KnbZynfS1UM,5789
-scenario/result.py,sha256=SGF8uYNtkP7cJy4KsshUozZRevmdiyX2TFzr6VreTv8,2717
-scenario/scenario.py,sha256=tYn3Y1sK6_7pg7hFb_5w0TW6nun-za_4F8kqcnrXXU4,4077
-scenario/scenario_executor.py,sha256=c8xV6GoJgO2JoZBWpYPQN5YwwQ3G9iJUtXV9UGSf1q8,7919
-scenario/testing_agent.py,sha256=eS-c_io5cHgzJ88wwRvU_vve-pmB2HsGWN6qwlq0sPg,10865
-scenario/utils.py,sha256=tMESosrxesA1B5zZB3IJ-sNSXDmnpNNib-DHobveVLA,3918
-langwatch_scenario-0.1.3.dist-info/METADATA,sha256=7OIolGcZ3fkCXFmE6JHkckVCeJb1r3yYSYveJ6iE9zw,8801
-langwatch_scenario-0.1.3.dist-info/WHEEL,sha256=pxyMxgL8-pra_rKaQ4drOZAegBVuX-G_4nRHjjgWbmo,91
-langwatch_scenario-0.1.3.dist-info/entry_points.txt,sha256=WlEnJ_gku0i18bIa3DSuGqXRX-QDQLe_s0YmRzK45TI,45
-langwatch_scenario-0.1.3.dist-info/top_level.txt,sha256=45Mn28aedJsetnBMB5xSmrJ-yo701QLH89Zlz4r1clE,9
-langwatch_scenario-0.1.3.dist-info/RECORD,,