15 implement simple multi agent workflow (#16)

* Add run_openai_multi_agent.

* config: Simplify DEFAULT_PROMPT.

* fix(test_unit_openai): Import after patch

* Update insrtuctions. Add communication agent

* Iterate on instructions

* enh(tracing): Handle openai_multi_agent. Add `test_setup_tracing`.

* More instructions tuning

* Add test_run_openai_multiagent

* More tweaks

docs/api.md

@@ -6,6 +6,8 @@
 
 ::: surf_spot_finder.agents.openai
 
+::: surf_spot_finder.agents.openai.DEFAULT_MULTIAGENT_INSTRUCTIONS
+
 ::: surf_spot_finder.agents.smolagents
 
 ::: surf_spot_finder.tracing

src/surf_spot_finder/agents/__init__.py

@@ -1,9 +1,10 @@
-from .openai import run_openai_agent
+from .openai import run_openai_agent, run_openai_multi_agent
 from .smolagents import run_smolagent
 
 RUNNERS = {
     "openai": run_openai_agent,
     "smolagents": run_smolagent,
+    "openai_multi_agent": run_openai_multi_agent,
 }
 
 

src/surf_spot_finder/agents/openai.py

@@ -1,10 +1,67 @@
 import os
-from typing import Optional, TYPE_CHECKING
+from typing import Optional
 
+from agents import (
+    Agent,
+    AsyncOpenAI,
+    OpenAIChatCompletionsModel,
+    Runner,
+    RunResult,
+    function_tool,
+)
 from loguru import logger
+from smolagents import (
+    DuckDuckGoSearchTool,
+    VisitWebpageTool,
+    FinalAnswerTool,
+)
 
-if TYPE_CHECKING:
-    from agents import RunResult
+
+@function_tool
+def search_web(query: str) -> str:
+    """Performs a duckduckgo web search based on your query (think a Google search) then returns the top search results.
+
+    Args:
+        query: The search query to perform.
+    """
+    logger.debug(f"Calling search_web: {query}")
+    search_tool = DuckDuckGoSearchTool()
+    return search_tool.forward(query)
+
+
+@function_tool
+def visit_webpage(url: str) -> str:
+    """Visits a webpage at the given url and reads its content as a markdown string. Use this to browse webpages.
+
+    Args:
+        url: The url of the webpage to visit.
+    """
+    logger.debug(f"Calling visit_webpage: {url}")
+    visit_tool = VisitWebpageTool()
+    return visit_tool.forward(url)
+
+
+@function_tool
+def final_answer(answer: str) -> str:
+    """Provides a final answer to the given problem.
+
+    Args:
+        answer: The answer to the problem.
+    """
+    logger.debug("Calling final_answer")
+    final_answer_tool = FinalAnswerTool()
+    return final_answer_tool.forward(answer)
+
+
+@function_tool
+def user_verification(query: str) -> str:
+    """Asks user to verify the given `query`.
+
+    Args:
+        query: The question that requires verification.
+    """
+    logger.debug("Calling user_verification")
+    return input(f"{query} => Type your answer here:")
 
 
 @logger.catch(reraise=True)
@@ -15,7 +72,7 @@ def run_openai_agent(
     instructions: Optional[str] = None,
     api_key_var: Optional[str] = None,
     base_url: Optional[str] = None,
-) -> "RunResult":
+) -> RunResult:
     """Runs an OpenAI agent with the given prompt and configuration.
 
     It leverages the 'agents' library to create and manage the agent
@@ -42,34 +99,6 @@ def run_openai_agent(
         RunResult: A RunResult object containing the output of the agent run.
             See https://openai.github.io/openai-agents-python/ref/result/#agents.result.RunResult.
     """
-    from agents import (
-        Agent,
-        AsyncOpenAI,
-        OpenAIChatCompletionsModel,
-        Runner,
-        function_tool,
-    )
-    from smolagents import DuckDuckGoSearchTool, VisitWebpageTool
-
-    @function_tool
-    def search_web(query: str) -> str:
-        """Performs a duckduckgo web search based on your query (think a Google search) then returns the top search results.
-
-        Args:
-            query: The search query to perform.
-        """
-        search_tool = DuckDuckGoSearchTool()
-        return search_tool.forward(query)
-
-    @function_tool
-    def visit_webpage(url: str) -> str:
-        """Visits a webpage at the given url and reads its content as a markdown string. Use this to browse webpages.
-
-        Args:
-            url: The url of the webpage to visit.
-        """
-        visit_tool = VisitWebpageTool()
-        return visit_tool.forward(url)
 
     if api_key_var and base_url:
         external_client = AsyncOpenAI(
@@ -95,3 +124,88 @@ def run_openai_agent(
     result = Runner.run_sync(agent, prompt)
     logger.info(result.final_output)
     return result
+
+
+DEFAULT_MULTIAGENT_INSTRUCTIONS = """
+You will be asked to perform a task.
+
+Always follow this steps:
+
+First, before solving the task, look at the available agent/tools and plan a sequence of actions using the available tools.
+Second, show the plan of actions and ask for user verification. If the user does not verify the plan, come up with a better plan.
+Third, execute the plan using the available tools, until you get a final answer.
+
+Once you get a final answer, show it and ask for user verification.  If the user does not verify the answer, come up with a better answer.
+
+Finally, use the available handoff tool (`transfer_to_<agent_name>`) to communicate it to the user.
+"""
+
+
+@logger.catch(reraise=True)
+def run_openai_multi_agent(
+    model_id: str,
+    prompt: str,
+    name: str = "surf-spot-finder",
+    instructions: Optional[str] = DEFAULT_MULTIAGENT_INSTRUCTIONS,
+) -> RunResult:
+    """Runs multiple OpenAI agents orchestrated by a main agent.
+
+    It leverages the 'agents' library to create and manage the agent
+    execution.
+
+    See https://openai.github.io/openai-agents-python/ref/agent/ for more details.
+
+
+    Args:
+        model_id (str): The ID of the OpenAI model to use (e.g., "gpt4o").
+            See https://platform.openai.com/docs/api-reference/models.
+        prompt (str): The prompt to be given to the agent.
+        name (str, optional): The name of the main agent. Defaults to "surf-spot-finder".
+        instructions (Optional[str], optional): Initial instructions to give the agent.
+            Defaults to [DEFAULT_MULTIAGENT_INSTRUCTIONS][surf_spot_finder.agents.openai.DEFAULT_MULTIAGENT_INSTRUCTIONS].
+
+    Returns:
+        RunResult: A RunResult object containing the output of the agent run.
+            See https://openai.github.io/openai-agents-python/ref/result/#agents.result.RunResult.
+    """
+    user_verification_agent = Agent(
+        model=model_id,
+        instructions="Display the current output to the user, then ask for verification.",
+        name="user-verification-agent",
+        tools=[user_verification],
+    )
+
+    search_web_agent = Agent(
+        model=model_id,
+        instructions="Find relevant information about the provided task by combining web searches with visiting webpages.",
+        name="search-web-agent",
+        tools=[search_web, visit_webpage],
+    )
+
+    communication_agent = Agent(
+        model=model_id,
+        instructions=None,
+        name="communication-agent",
+        tools=[final_answer],
+    )
+
+    main_agent = Agent(
+        model=model_id,
+        instructions=instructions,
+        name=name,
+        handoffs=[communication_agent],
+        tools=[
+            search_web_agent.as_tool(
+                tool_name="search_web_with_agent",
+                tool_description=search_web_agent.instructions,
+            ),
+            user_verification_agent.as_tool(
+                tool_name="ask_user_verification_with_agent",
+                tool_description=user_verification_agent.instructions,
+            ),
+        ],
+    )
+
+    result = Runner.run_sync(main_agent, prompt)
+    logger.info(result.final_output)
+    return result

src/surf_spot_finder/config.py

@@ -7,12 +7,7 @@ CURRENT_DATE = datetime.now().strftime("%Y-%m-%d")
 DEFAULT_PROMPT = (
     "What will be the best surf spot around {LOCATION}"
     ", in a {MAX_DRIVING_HOURS} hour driving radius"
-    ", at {DATE}? it is currently "
-    + CURRENT_DATE
-    + ". find me the best surf spot and also report back"
-    " on the expected water temperature and wave height."
-    " Please remember that doing a google/duckduckgo search may be useful for finding which sites are relevant,"
-    " but the final answer should be based on information retrieved from https://www.surf-forecast.com."
+    ", at {DATE}?"
 )
 
 

src/surf_spot_finder/tracing.py

@@ -95,7 +95,7 @@ def setup_tracing(tracer_provider: TracerProvider, agent_type: str) -> None:
 
     validate_agent_type(agent_type)
 
-    if agent_type == "openai":
+    if "openai" in agent_type:
         from openinference.instrumentation.openai import OpenAIInstrumentor
 
         OpenAIInstrumentor().instrument(tracer_provider=tracer_provider)

tests/unit/agents/test_unit_openai.py

@@ -2,56 +2,98 @@ import os
 import pytest
 from unittest.mock import patch, MagicMock, ANY
 
-from surf_spot_finder.agents.openai import run_openai_agent
-
-
-@pytest.fixture
-def mock_agents_module():
-    agents_mocks = {
-        name: MagicMock()
-        for name in (
-            "Agent",
-            "AsyncOpenAI",
-            "OpenAIChatCompletionsModel",
-            "Runner",
-            "WebSearchTool",
-        )
-    }
-    with patch.dict(
-        "sys.modules",
-        {
-            "agents": MagicMock(**agents_mocks),
-        },
-    ):
-        yield agents_mocks
+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()
 
-def test_run_openai_agent_default(mock_agents_module):
-    run_openai_agent("gpt-4o", "Test prompt")
-    mock_agents_module["Agent"].assert_called_once_with(
-        model="gpt-4o",
-        instructions=None,
-        name="surf-spot-finder",
-        tools=ANY,
-    )
+    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(mock_agents_module):
-    with patch.dict(os.environ, {"TEST_API_KEY": "test-key-12345"}):
+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"
         )
-        mock_agents_module["AsyncOpenAI"].assert_called_once_with(
+        async_openai_mock.assert_called_once_with(
             api_key="test-key-12345",
             base_url="FOO",
         )
-        mock_agents_module["OpenAIChatCompletionsModel"].assert_called_once()
+        openai_chat_completions_model.assert_called_once()
 
 
-def test_run_smolagent_environment_error():
-    """Test that passing a bad api_key_var throws an error"""
+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,
+        )

tests/unit/test_unit_tracing.py

@@ -32,6 +32,20 @@ def test_get_tracer_provider(tmp_path, json_tracer):
             )
 
 
+@pytest.mark.parametrize(
+    "agent_type,instrumentor",
+    [
+        ("openai", "openai.OpenAIInstrumentor"),
+        ("openai_multi_agent", "openai.OpenAIInstrumentor"),
+        ("smolagents", "smolagents.SmolagentsInstrumentor"),
+    ],
+)
+def test_setup_tracing(agent_type, instrumentor):
+    with patch(f"openinference.instrumentation.{instrumentor}") as mock_instrumentor:
+        setup_tracing(MagicMock(), agent_type)
+        mock_instrumentor.assert_called_once()
+
+
 def test_invalid_agent_type():
     with pytest.raises(ValueError, match="agent_type must be one of"):
         setup_tracing(MagicMock(), "invalid_agent_type")