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

scenario/testing_agent.py

@@ -1,279 +0,0 @@
-"""
-TestingAgent module: defines the testing agent that interacts with the agent under test.
-"""
-
-import json
-import logging
-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.scenario_agent_adapter import ScenarioAgentAdapter
-from scenario.utils import reverse_roles
-
-from .error_messages import testing_agent_not_configured_error_message
-from .types import AgentInput, AgentReturnTypes, ScenarioAgentRole, ScenarioResult
-
-
-logger = logging.getLogger("scenario")
-
-
-class TestingAgent(ScenarioAgentAdapter):
-    """
-    The Testing Agent that interacts with the agent under test.
-
-    This agent is responsible for:
-    1. Generating messages to send to the agent based on the scenario
-    2. Evaluating the responses from the agent against the success/failure criteria
-    3. Determining when to end the test and return a result
-    """
-
-    roles = {ScenarioAgentRole.USER, ScenarioAgentRole.JUDGE}
-
-    model: str = ""
-    api_key: Optional[str] = None
-    temperature: float = 0.0
-    max_tokens: Optional[int] = None
-
-    # 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"])
-    async def call(
-        self,
-        input: AgentInput,
-    ) -> AgentReturnTypes:
-        """
-        Generate the next message in the conversation based on history OR
-        return a ScenarioResult if the test should conclude.
-
-        Returns either:
-          - A string message to send to the agent (if conversation should continue)
-          - A ScenarioResult (if the test should conclude)
-        """
-
-        scenario = input.scenario_state.scenario
-
-        messages = [
-            {
-                "role": "system",
-                "content": f"""
-<role>
-You are pretending to be a user, you are testing an AI Agent (shown as the user role) based on a scenario.
-Approach this naturally, as a human user would, with very short inputs, few words, all lowercase, imperative, not periods, like when they google or talk to chatgpt.
-</role>
-
-<goal>
-Your goal (assistant) is to interact with the Agent Under Test (user) as if you were a human user to see if it can complete the scenario successfully.
-</goal>
-
-<scenario>
-{scenario.description}
-</scenario>
-
-<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 all the criteria have been met
-</execution_flow>
-
-<rules>
-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?"},
-            *input.messages,
-        ]
-
-        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",
-                    "content": """
-System:
-
-<finish_test>
-This is the last message, conversation has reached the maximum number of turns, give your final verdict,
-if you don't have enough information to make a verdict, say inconclusive with max turns reached.
-</finish_test>
-""",
-                }
-            )
-
-        # 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.
-        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",
-                "function": {
-                    "name": "finish_test",
-                    "description": "Complete the test with a final verdict",
-                    "strict": True,
-                    "parameters": {
-                        "type": "object",
-                        "properties": {
-                            "criteria": {
-                                "type": "object",
-                                "properties": {
-                                    criteria_names[idx]: {
-                                        "enum": [True, False, "inconclusive"],
-                                        "description": criterion,
-                                    }
-                                    for idx, criterion in enumerate(scenario.criteria)
-                                },
-                                "required": criteria_names,
-                                "additionalProperties": False,
-                                "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": ["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(
-                model=self.model,
-                messages=messages,
-                temperature=self.temperature,
-                max_tokens=self.max_tokens,
-                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
-                ),
-            ),
-        )
-
-        # Extract the content from the response
-        if hasattr(response, "choices") and len(response.choices) > 0:
-            message = cast(Choices, response.choices[0]).message
-
-            # Check if the LLM chose to use the tool
-            if message.tool_calls:
-                tool_call = message.tool_calls[0]
-                if tool_call.function.name == "finish_test":
-                    # Parse the tool call arguments
-                    try:
-                        args = json.loads(tool_call.function.arguments)
-                        verdict = args.get("verdict", "inconclusive")
-                        reasoning = args.get("reasoning", "No reasoning provided")
-                        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
-                        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")
-
-            # If no tool call use the message content as next message
-            message_content = message.content
-            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"No response from LLM: {response.__repr__()}")
-
-            return {"role": "user", "content": message_content}
-        else:
-            raise Exception(
-                f"Unexpected response format from LLM: {response.__repr__()}"
-            )

scenario/utils.py

@@ -1,264 +0,0 @@
-from contextlib import contextmanager
-import sys
-from typing import (
-    Any,
-    Iterator,
-    List,
-    Literal,
-    Optional,
-    Union,
-    TypeVar,
-    Awaitable,
-    cast,
-)
-from pydantic import BaseModel
-
-import json
-
-import termcolor
-from textwrap import indent
-from openai.types.chat import ChatCompletionMessageParam
-from rich.live import Live
-from rich.spinner import Spinner
-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)
-
-
-class SerializableWithStringFallback(SerializableAndPydanticEncoder):
-    def default(self, o):
-        try:
-            return super().default(o)
-        except:
-            return str(o)
-
-
-def safe_list_at(list, index, default=None):
-    try:
-        return list[index]
-    except:
-        return default
-
-
-def safe_attr_or_key(obj, attr_or_key, default=None):
-    return getattr(obj, attr_or_key, obj.get(attr_or_key))
-
-
-def title_case(string):
-    return " ".join(word.capitalize() for word in string.split("_"))
-
-
-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")
-        if role == "assistant":
-            tool_calls = safe_attr_or_key(msg, "tool_calls")
-            if content:
-                print(scenario_name + termcolor.colored("Agent:", "blue"), content)
-            if tool_calls:
-                for tool_call in tool_calls:
-                    function = safe_attr_or_key(tool_call, "function")
-                    name = safe_attr_or_key(function, "name")
-                    args = safe_attr_or_key(function, "arguments", "{}")
-                    args = _take_maybe_json_first_lines(args)
-                    print(
-                        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(
-                scenario_name + termcolor.colored(f"ToolResult:", "magenta"),
-                f"\n\n{indent(content, ' ' * 4)}\n",
-            )
-        else:
-            print(
-                scenario_name + termcolor.colored(f"{title_case(role)}:", "magenta"),
-                msg.__repr__(),
-            )
-
-
-def _take_maybe_json_first_lines(string, max_lines=5):
-    content = str(string)
-    try:
-        content = json.dumps(json.loads(content), indent=2)
-    except:
-        pass
-    content = content.split("\n")
-    if len(content) > max_lines:
-        content = content[:max_lines] + ["..."]
-    return "\n".join(content)
-
-
-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
-        self.text_before = text
-        self.color = color
-
-    def render(self, time):
-        # Get the original spinner frame
-        spinner_frame = super().render(time)
-        # Create a composite with text first, then spinner
-        return Text(f"{self.text_before} ", style=self.color) + spinner_frame
-
-
-@contextmanager
-def show_spinner(
-    text: str, color: str = "white", enabled: Optional[Union[bool, int]] = None
-):
-    if not enabled:
-        yield
-    else:
-        spinner = TextFirstSpinner("dots", text, color=color)
-        try:
-            with Live(spinner, console=console, refresh_per_second=20):
-                yield
-        # It happens when we are multi-threading, it's fine, just ignore it, you probably don't want multiple spinners at once anyway
-        except LiveError:
-            yield
-
-        # 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