File size: 3,617 Bytes
c61646b
 
c2fdfed
c61646b
ee1c4f8
 
 
f53eca9
 
 
 
 
ee1c4f8
 
502c6d3
 
 
 
ee1c4f8
 
c61646b
ee1c4f8
 
c61646b
ee1c4f8
 
 
 
 
 
 
502c6d3
ee1c4f8
8aa7d2f
ee1c4f8
c61646b
 
ee1c4f8
 
 
 
 
 
 
 
 
 
 
 
 
c61646b
f53eca9
c61646b
ee1c4f8
c61646b
 
 
ee1c4f8
c61646b
 
ee1c4f8
c61646b
 
 
f53eca9
c61646b
ee1c4f8
 
 
 
f53eca9
ee1c4f8
 
 
 
f53eca9
ee1c4f8
 
 
 
f53eca9
ee1c4f8
f53eca9
 
 
 
ee1c4f8
 
 
 
f53eca9
ee1c4f8
f53eca9
ee1c4f8
 
 
 
f53eca9
ee1c4f8
f53eca9
ee1c4f8
 
 
 
502c6d3
ee1c4f8
 
 
 
 
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
import os
import pytest
from unittest.mock import patch, MagicMock, ANY

from surf_spot_finder.agents.openai import (
    run_openai_agent,
    run_openai_multi_agent,
)
from surf_spot_finder.tools import (
    show_final_answer,
    show_plan,
    ask_user_verification,
    search_web,
    visit_webpage,
)
from surf_spot_finder.prompts.openai import (
    SINGLE_AGENT_SYSTEM_PROMPT,
    MULTI_AGENT_SYSTEM_PROMPT,
)


def test_run_openai_agent_default():
    mock_agent = MagicMock()

    with (
        patch("surf_spot_finder.agents.openai.Agent", mock_agent),
        patch("surf_spot_finder.agents.openai.Runner", MagicMock()),
    ):
        run_openai_agent("gpt-4o", "Test prompt")
        mock_agent.assert_called_once_with(
            model="gpt-4o",
            instructions=SINGLE_AGENT_SYSTEM_PROMPT,
            name="surf-spot-finder",
            tools=ANY,
        )


def test_run_openai_agent_base_url_and_api_key_var():
    async_openai_mock = MagicMock()
    openai_chat_completions_model = MagicMock()
    with (
        patch("surf_spot_finder.agents.openai.Agent", MagicMock()),
        patch("surf_spot_finder.agents.openai.Runner", MagicMock()),
        patch("surf_spot_finder.agents.openai.AsyncOpenAI", async_openai_mock),
        patch(
            "surf_spot_finder.agents.openai.OpenAIChatCompletionsModel",
            openai_chat_completions_model,
        ),
        patch.dict(os.environ, {"TEST_API_KEY": "test-key-12345"}),
    ):
        run_openai_agent(
            "gpt-4o", "Test prompt", api_base="FOO", api_key_var="TEST_API_KEY"
        )
        async_openai_mock.assert_called_once_with(
            api_key="test-key-12345",
            base_url="FOO",
        )
        openai_chat_completions_model.assert_called_once()


def test_run_openai_environment_error():
    with patch.dict(os.environ, {}, clear=True):
        with pytest.raises(KeyError, match="MISSING_KEY"):
            run_openai_agent(
                "test-model", "Test prompt", api_base="FOO", api_key_var="MISSING_KEY"
            )


def test_run_openai_multiagent():
    mock_agent = MagicMock()
    mock_function_tool = MagicMock()

    with (
        patch("surf_spot_finder.agents.openai.Agent", mock_agent),
        patch("surf_spot_finder.agents.openai.Runner", MagicMock()),
        patch("surf_spot_finder.agents.openai.function_tool", mock_function_tool),
    ):
        run_openai_multi_agent("gpt-4o", "Test prompt")
        mock_agent.assert_any_call(
            model="gpt-4o",
            instructions="Interact with the user by showing information and asking for verification.",
            name="user-verification-agent",
            tools=[
                mock_function_tool(show_plan),
                mock_function_tool(ask_user_verification),
            ],
        )

        mock_agent.assert_any_call(
            model="gpt-4o",
            instructions="Find relevant information about the provided task by using your tools.",
            name="search-web-agent",
            tools=[mock_function_tool(search_web), mock_function_tool(visit_webpage)],
        )

        mock_agent.assert_any_call(
            model="gpt-4o",
            instructions="Communicate the final answer to the user.",
            name="communication-agent",
            tools=[mock_function_tool(show_final_answer)],
        )

        mock_agent.assert_any_call(
            model="gpt-4o",
            instructions=MULTI_AGENT_SYSTEM_PROMPT,
            name="surf-spot-finder",
            # TODO: add more elaborated checks
            handoffs=ANY,
            tools=ANY,
        )