scenario/python/examples/test_running_in_parallel.py at main · langwatch/scenario · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
"""
Example test for a vegetarian recipe agent.

This example demonstrates testing an AI agent that generates vegetarian recipes.
"""

from typing import cast
import pytest
from dotenv import load_dotenv
import litellm

from openai.types.chat import ChatCompletionMessageParam
from scenario.agent_adapter import AgentAdapter
from scenario.types import AgentInput, AgentReturnTypes

load_dotenv()

import scenario

scenario.configure(default_model="openai/gpt-4.1-mini")


class VegetarianRecipeAgentAdapter(AgentAdapter):
    @scenario.cache()
    async def call(self, input: AgentInput) -> AgentReturnTypes:
        response = litellm.completion(
            model="openai/gpt-4.1-mini",
            messages=[
                {
                    "role": "system",
                    "content": """You are a vegetarian recipe agent.
                    Given the user request, ask AT MOST ONE follow-up question,
                    then provide a complete recipe. Keep your responses concise and focused.""",
                },
                *input.messages,
            ],
        )
        message = response.choices[0].message  # type: ignore

        return [cast(ChatCompletionMessageParam, message)]


@pytest.mark.flaky(reruns=2)
@pytest.mark.asyncio_concurrent(group="vegetarian_recipe_agent")
async def test_vegetarian_recipe_agent():
    # Define the scenario
    result = await scenario.run(
        name="dinner idea",
        description="User is looking for a dinner idea",
        agents=[
            VegetarianRecipeAgentAdapter(),
            scenario.UserSimulatorAgent(),
            scenario.JudgeAgent(
                criteria=[
                    "Recipe agent generates a vegetarian recipe",
                    "Recipe includes a list of ingredients",
                    "Recipe includes step-by-step cooking instructions",
                    "The recipe is vegetarian and does not include meat",
                    "The should NOT ask more than two follow-up questions",
                ]
            ),
        ],
        max_turns=12,
        set_id="python-examples",
    )

    # Assert for pytest to know whether the test passed
    assert result.success


@pytest.mark.agent_test
@pytest.mark.asyncio_concurrent(group="vegetarian_recipe_agent")
@pytest.mark.flaky(reruns=2)
async def test_user_is_hungry():
    # Define the scenario
    result = await scenario.run(
        name="hungry user",
        description="User is very very hungry and wants a big, filling meal",
        agents=[
            VegetarianRecipeAgentAdapter(),
            scenario.UserSimulatorAgent(),
            scenario.JudgeAgent(
                criteria=[
                    "Recipe agent generates a vegetarian recipe",
                    "Recipe includes a list of ingredients",
                    "Recipe includes step-by-step cooking instructions",
                    "The recipe is vegetarian and does not include meat",
                    "The agent should NOT ask more than two follow-up questions",
                ]
            ),
        ],
        script=[
            scenario.user("I'm starving! I need something really filling for dinner tonight"),
            scenario.agent(),
            scenario.user(),
            scenario.agent(),
            scenario.judge(),
        ],
        set_id="python-examples",
    )

    # Assert for pytest to know whether the test passed
    assert result.success