Implement agent using `openai-agents-python` (#9)

* enh(agents): Add `RUNNERS` registry.

Update cli to accept agent_type arg.

* enh(tracing): Decouple into `get_trace_provider` and `setup_tracing`.

Pass `agent_type` to `setup_tracing` in order to use the right instrumentor.

* feat(agents): Add basic agent using openai-agents-python.

* fix(tests): Install openai deps

* Address PR comments.

- Drop Literal in favor of explicit and centralized validation.
- Add docstring and types to `setup_tracing` .

* fix(docs): Add RUNNERS to API Reference.

.github/workflows/tests.yaml

@@ -28,7 +28,7 @@ jobs:
           cache: "pip"
 
       - name: Install
-        run: pip install -e '.[tests]'
+        run: pip install -e '.[openai,tests]'
 
       - name: Run tests
         run: pytest -v tests

docs/api.md

@@ -2,6 +2,10 @@
 
 ::: surf_spot_finder.config.Config
 
+::: surf_spot_finder.agents.RUNNERS
+
+::: surf_spot_finder.agents.openai
+
 ::: surf_spot_finder.agents.smolagents
 
 ::: surf_spot_finder.tracing

pyproject.toml

@@ -18,6 +18,11 @@ dependencies = [
 ]
 
 [project.optional-dependencies]
+openai = [
+  "openai-agents",
+  "openinference-instrumentation-openai"
+]
+
 demo = [
   "gradio",
   "spaces"

src/surf_spot_finder/agents/__init__.py

@@ -0,0 +1,12 @@
+from .openai import run_openai_agent
+from .smolagents import run_smolagent
+
+RUNNERS = {
+    "openai": run_openai_agent,
+    "smolagents": run_smolagent,
+}
+
+
+def validate_agent_type(value) -> str:
+    if value not in RUNNERS:
+        raise ValueError(f"agent_type must be one of {RUNNERS.keys()}")

src/surf_spot_finder/agents/openai.py

@@ -0,0 +1,76 @@
+import os
+from typing import Optional, TYPE_CHECKING
+
+from loguru import logger
+
+if TYPE_CHECKING:
+    from agents import RunResult
+
+
+@logger.catch(reraise=True)
+def run_openai_agent(
+    model_id: str,
+    prompt: str,
+    name: str = "surf-spot-finder",
+    instructions: Optional[str] = None,
+    api_key_var: Optional[str] = None,
+    base_url: Optional[str] = None,
+) -> "RunResult":
+    """Runs an OpenAI agent with the given prompt and configuration.
+
+    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 agent. Defaults to "surf-spot-finder".
+        instructions (Optional[str], optional): Initial instructions to give the agent.
+            Defaults to None.
+        api_key_var (Optional[str], optional): The name of the environment variable
+            containing the OpenAI API key. If provided, along with `base_url`, an
+            external OpenAI client will be used. Defaults to None.
+        base_url (Optional[str], optional): The base URL for the OpenAI API.
+            Required if `api_key_var` is provided to use an external OpenAI client.
+            Defaults to None.
+
+    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.
+    """
+    from agents import (
+        Agent,
+        AsyncOpenAI,
+        OpenAIChatCompletionsModel,
+        Runner,
+        WebSearchTool,
+    )
+
+    if api_key_var and base_url:
+        external_client = AsyncOpenAI(
+            api_key=os.environ[api_key_var],
+            base_url=base_url,
+        )
+        agent = Agent(
+            name=name,
+            instructions=instructions,
+            model=OpenAIChatCompletionsModel(
+                model=model_id,
+                openai_client=external_client,
+            ),
+            tools=[WebSearchTool()],
+        )
+    else:
+        agent = Agent(
+            model=model_id,
+            instructions=instructions,
+            name=name,
+            tools=[WebSearchTool()],
+        )
+    result = Runner.run_sync(agent, prompt)
+    logger.info(result.final_output)
+    return result

src/surf_spot_finder/agents/smolagents.py

@@ -16,7 +16,9 @@ def run_smolagent(
 ) -> "CodeAgent":
     """
     Create and configure a Smolagents CodeAgent with the specified model.
+
     See https://docs.litellm.ai/docs/providers for details on available LiteLLM providers.
+
     Args:
         model_id (str): Model identifier using LiteLLM syntax (e.g., 'openai/o1', 'anthropic/claude-3-sonnet')
         prompt (str): Prompt to provide to the model
@@ -36,13 +38,6 @@ def run_smolagent(
         LiteLLMModel,
         ToolCollection,
     )
-
-    model = LiteLLMModel(
-        model_id=model_id,
-        api_base=api_base if api_base else None,
-        api_key=os.environ[api_key_var] if api_key_var else None,
-    )
-
     from mcp import StdioServerParameters
 
     model = LiteLLMModel(

src/surf_spot_finder/cli.py

@@ -7,8 +7,8 @@ from surf_spot_finder.config import (
     Config,
     DEFAULT_PROMPT,
 )
-from surf_spot_finder.agents.smolagents import run_smolagent
-from surf_spot_finder.tracing import setup_tracing
+from surf_spot_finder.agents import RUNNERS
+from surf_spot_finder.tracing import get_tracer_provider, setup_tracing
 
 
 @logger.catch(reraise=True)
@@ -17,6 +17,7 @@ def find_surf_spot(
     date: str,
     max_driving_hours: int,
     model_id: str,
+    agent_type: str = "smolagents",
     api_key_var: Optional[str] = None,
     prompt: str = DEFAULT_PROMPT,
     json_tracer: bool = True,
@@ -28,6 +29,7 @@ def find_surf_spot(
         date=date,
         max_driving_hours=max_driving_hours,
         model_id=model_id,
+        agent_type=agent_type,
         api_key_var=api_key_var,
         prompt=prompt,
         json_tracer=json_tracer,
@@ -35,13 +37,14 @@ def find_surf_spot(
     )
 
     logger.info("Setting up tracing")
-    setup_tracing(project_name="surf-spot-finder", json_tracer=config.json_tracer)
+    tracer_provider = get_tracer_provider(
+        project_name="surf-spot-finder", json_tracer=config.json_tracer
+    )
+    setup_tracing(tracer_provider, config.agent_type)
 
-    logger.info("Running agent")
-    run_smolagent(
+    logger.info(f"Running {config.agent_type} agent")
+    RUNNERS[config.agent_type](
         model_id=config.model_id,
-        api_key_var=config.api_key_var,
-        api_base=config.api_base,
         prompt=config.prompt.format(
             LOCATION=config.location,
             MAX_DRIVING_HOURS=config.max_driving_hours,

src/surf_spot_finder/config.py

@@ -21,12 +21,20 @@ def validate_prompt(value) -> str:
     return value
 
 
+def validate_agent_type(value) -> str:
+    from surf_spot_finder.agents import validate_agent_type
+
+    validate_agent_type(value)
+    return value
+
+
 class Config(BaseModel):
     prompt: Annotated[str, AfterValidator(validate_prompt)]
     location: str
     max_driving_hours: PositiveInt
     date: FutureDatetime
     model_id: str
+    agent_type: Annotated[str, AfterValidator(validate_agent_type)]
     api_key_var: Optional[str] = None
     json_tracer: bool = True
     api_base: Optional[str] = None

src/surf_spot_finder/tracing.py

@@ -1,8 +1,7 @@
-from datetime import datetime
 import os
+from datetime import datetime
 
 from opentelemetry import trace
-from openinference.instrumentation.smolagents import SmolagentsInstrumentor
 from opentelemetry.sdk.trace import TracerProvider
 from opentelemetry.sdk.trace.export import SimpleSpanProcessor
 from opentelemetry.sdk.trace.export import SpanExporter
@@ -24,9 +23,9 @@ class JsonFileSpanExporter(SpanExporter):
         pass
 
 
-def setup_tracing(project_name: str, json_tracer: bool) -> TracerProvider:
+def get_tracer_provider(project_name: str, json_tracer: bool) -> TracerProvider:
     """
-    Set up tracing configuration based on the selected mode.
+    Create a tracer_provider based on the selected mode.
 
     Args:
         project_name: Name of the project for tracing
@@ -52,6 +51,27 @@ def setup_tracing(project_name: str, json_tracer: bool) -> TracerProvider:
     else:
         tracer_provider = register(project_name=project_name)
 
-    SmolagentsInstrumentor().instrument(tracer_provider=tracer_provider)
-
     return tracer_provider
+
+
+def setup_tracing(tracer_provider: TracerProvider, agent_type: str) -> None:
+    """Setup tracing for `agent_type` by instrumenting `trace_provider`.
+
+    Args:
+        tracer_provider (TracerProvider): The configured tracer provider from
+            [get_tracer_provider][surf_spot_finder.tracing.get_tracer_provider].
+        agent_type (str): The type of agent being used.
+            Must be one of the supported types in [RUNNERS][surf_spot_finder.agents.RUNNERS].
+    """
+    from surf_spot_finder.agents import validate_agent_type
+
+    validate_agent_type(agent_type)
+
+    if agent_type == "openai":
+        from openinference.instrumentation.openai import OpenAIInstrumentor
+
+        OpenAIInstrumentor().instrument(tracer_provider=tracer_provider)
+    elif agent_type == "smolagents":
+        from openinference.instrumentation.smolagents import SmolagentsInstrumentor
+
+        SmolagentsInstrumentor().instrument(tracer_provider=tracer_provider)

tests/unit/agents/test_unit_openai.py

@@ -0,0 +1,59 @@
+import os
+import pytest
+from unittest.mock import patch, MagicMock
+
+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
+
+
+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=[
+            mock_agents_module["WebSearchTool"].return_value,
+        ],
+    )
+
+
+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"}):
+        run_openai_agent(
+            "gpt-4o", "Test prompt", base_url="FOO", api_key_var="TEST_API_KEY"
+        )
+        mock_agents_module["AsyncOpenAI"].assert_called_once_with(
+            api_key="test-key-12345",
+            base_url="FOO",
+        )
+        mock_agents_module["OpenAIChatCompletionsModel"].assert_called_once()
+
+
+def test_run_smolagent_environment_error():
+    """Test that passing a bad api_key_var throws an 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"
+            )