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
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
"""
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 strictly vegetarian recipe agent.
                    Given the user request, ask AT MOST ONE follow-up question,
                    then provide a complete recipe. Once you have given a recipe, answer
                    without asking further questions. Keep your responses concise and focused.

                    HARD RULE: You ONLY produce vegetarian recipes. If the user
                    asks for a recipe containing meat, poultry, fish, or seafood,
                    politely decline and offer a vegetarian alternative instead.
                    Never include meat, poultry, fish, or seafood in any recipe
                    you generate, even if the user insists.""",
                },
                *input.messages,
            ],
        )
        message = response.choices[0].message  # type: ignore

        return [cast(ChatCompletionMessageParam, message)]


@pytest.mark.asyncio_concurrent(group="vegetarian_recipe_agent")
async def test_vegetarian_recipe_agent():
    # No ``flaky`` marker: reruns declared on a test that
    # pytest-asyncio-concurrent owns never happen. See test_user_is_hungry.
    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 agent 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")
async def test_user_is_hungry():
    # Follow-up prompts are fixed strings rather than free-form
    # ``scenario.user()`` calls, matching test_running_in_parallel_with_arun.
    # A free-form simulator wandered into asking for meat, the agent's hard
    # rule made it decline, and the judge then failed the three criteria that
    # ask for an actual recipe. Two agent turns are also enough: with three,
    # the agent had room to ask the third follow-up that criterion 5 forbids.
    #
    # No ``flaky`` marker: pytest-asyncio-concurrent runs a group through its
    # own protocol hook, which pytest-rerunfailures never sees, so reruns
    # declared here would silently never happen. The conftest guard fails
    # collection if one is added back.
    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("That sounds great, can you share the recipe?"),
            scenario.agent(),
            scenario.judge(),
        ],
        set_id="python-examples",
    )

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