File size: 3,183 Bytes
c61646b
 
c2fdfed
c61646b
ee1c4f8
 
 
 
 
 
 
 
 
 
c61646b
ee1c4f8
 
c61646b
ee1c4f8
 
 
 
 
 
 
 
 
 
 
c61646b
 
ee1c4f8
 
 
 
 
 
 
 
 
 
 
 
 
c61646b
 
 
ee1c4f8
c61646b
 
 
ee1c4f8
c61646b
 
ee1c4f8
c61646b
 
 
 
 
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
import os
import pytest
from unittest.mock import patch, MagicMock, ANY

from surf_spot_finder.agents.openai import (
    final_answer,
    run_openai_agent,
    run_openai_multi_agent,
    search_web,
    user_verification,
    visit_webpage,
    DEFAULT_MULTIAGENT_INSTRUCTIONS,
)


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=None,
            name="surf-spot-finder",
            tools=[search_web, visit_webpage],
        )


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", base_url="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", base_url="FOO", api_key_var="MISSING_KEY"
            )


def test_run_openai_multiagent():
    mock_agent = MagicMock()

    with (
        patch("surf_spot_finder.agents.openai.Agent", mock_agent),
        patch("surf_spot_finder.agents.openai.Runner", MagicMock()),
    ):
        run_openai_multi_agent("gpt-4o", "Test prompt")
        mock_agent.assert_any_call(
            model="gpt-4o",
            instructions="Display the current output to the user, then ask for verification.",
            name="user-verification-agent",
            tools=[user_verification],
        )

        mock_agent.assert_any_call(
            model="gpt-4o",
            instructions="Find relevant information about the provided task by combining web searches with visiting webpages.",
            name="search-web-agent",
            tools=[search_web, visit_webpage],
        )

        mock_agent.assert_any_call(
            model="gpt-4o",
            instructions=None,
            name="communication-agent",
            tools=[final_answer],
        )

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