enh(cli): Add `from_config` option. (#17)

* 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

* enh(cli): Add `from_config` option.

- Add `examples` dir with input configs.
- Move prompts to `surf_spot_finder.prompts`.

* fix api.md

* fix(test_unit_smolagents): Don't rely on patching `sys`.

docs/api.md

@@ -1,13 +1,19 @@
 # API Reference
 
+::: surf_spot_finder.cli
+
 ::: surf_spot_finder.config.Config
 
 ::: surf_spot_finder.agents.RUNNERS
 
 ::: surf_spot_finder.agents.openai
 
-::: surf_spot_finder.agents.openai.DEFAULT_MULTIAGENT_INSTRUCTIONS
-
 ::: surf_spot_finder.agents.smolagents
 
+::: surf_spot_finder.prompts.openai.SINGLE_AGENT_SYSTEM_PROMPT
+
+::: surf_spot_finder.prompts.openai.MULTI_AGENT_SYSTEM_PROMPT
+
+::: surf_spot_finder.prompts.shared.INPUT_PROMPT
+
 ::: surf_spot_finder.tracing

examples/openai_multi_agent.yaml

@@ -0,0 +1,6 @@
+location: Pontevedra
+date: 2025-03-22 12:00
+max_driving_hours: 2
+model_id: o3-mini
+agent_type: openai_multi_agent
+# input_prompt_template:

examples/openai_single_agent.yaml

@@ -0,0 +1,6 @@
+location: Pontevedra
+date: 2025-03-22 12:00
+max_driving_hours: 2
+model_id: o3-mini
+agent_type: openai
+# input_prompt_template:

examples/smolagents_single_agent.yaml

@@ -0,0 +1,7 @@
+location: Pontevedra
+date: 2025-03-22 12:00
+max_driving_hours: 2
+model_id: openai/gpt-3.5-turbo
+api_key_var: OPENAI_API_KEY
+agent_type: smolagents
+# input_prompt_template:

src/surf_spot_finder/agents/openai.py

@@ -17,6 +17,12 @@ from smolagents import (
 )
 
 
+from surf_spot_finder.prompts.openai import (
+    SINGLE_AGENT_SYSTEM_PROMPT,
+    MULTI_AGENT_SYSTEM_PROMPT,
+)
+
+
 @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.
@@ -69,7 +75,7 @@ def run_openai_agent(
     model_id: str,
     prompt: str,
     name: str = "surf-spot-finder",
-    instructions: Optional[str] = None,
+    instructions: Optional[str] = SINGLE_AGENT_SYSTEM_PROMPT,
     api_key_var: Optional[str] = None,
     base_url: Optional[str] = None,
 ) -> RunResult:
@@ -87,7 +93,7 @@ def run_openai_agent(
         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.
+            Defaults to [SINGLE_AGENT_SYSTEM_PROMPT][surf_spot_finder.prompts.openai.SINGLE_AGENT_SYSTEM_PROMPT].
         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.
@@ -126,27 +132,12 @@ def run_openai_agent(
     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,
+    instructions: Optional[str] = MULTI_AGENT_SYSTEM_PROMPT,
 ) -> RunResult:
     """Runs multiple OpenAI agents orchestrated by a main agent.
 
@@ -162,7 +153,7 @@ def run_openai_multi_agent(
         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].
+            Defaults to [MULTI_AGENT_SYSTEM_PROMPT][surf_spot_finder.prompts.openai.MULTI_AGENT_SYSTEM_PROMPT].
 
     Returns:
         RunResult: A RunResult object containing the output of the agent run.

src/surf_spot_finder/agents/smolagents.py

@@ -1,10 +1,16 @@
 import os
-from typing import Optional, TYPE_CHECKING
+from typing import Optional
 
 from loguru import logger
 
-if TYPE_CHECKING:
-    from smolagents import CodeAgent
+from smolagents import (
+    CodeAgent,
+    DuckDuckGoSearchTool,
+    LiteLLMModel,
+    ToolCollection,
+)
+from mcp import StdioServerParameters
+from surf_spot_finder.prompts.smolagents import SYSTEM_PROMPT
 
 
 @logger.catch(reraise=True)
@@ -13,7 +19,7 @@ def run_smolagent(
     prompt: str,
     api_key_var: Optional[str] = None,
     api_base: Optional[str] = None,
-) -> "CodeAgent":
+) -> CodeAgent:
     """
     Create and configure a Smolagents CodeAgent with the specified model.
 
@@ -29,17 +35,10 @@ def run_smolagent(
         CodeAgent: Configured agent ready to process requests
 
     Example:
+
         >>> agent = run_smolagent("anthropic/claude-3-haiku", "my prompt here", "ANTHROPIC_API_KEY", None, None)
         >>> agent.run("Find surf spots near San Diego")
     """
-    from smolagents import (  # pylint: disable=import-outside-toplevel
-        CodeAgent,
-        DuckDuckGoSearchTool,
-        LiteLLMModel,
-        ToolCollection,
-    )
-    from mcp import StdioServerParameters
-    from surf_spot_finder.agents.prompts.smolagents import SYSTEM_PROMPT
 
     model = LiteLLMModel(
         model_id=model_id,

src/surf_spot_finder/cli.py

@@ -1,40 +1,72 @@
+from pathlib import Path
 from typing import Optional
 
+import yaml
 from fire import Fire
 from loguru import logger
 
 from surf_spot_finder.config import (
     Config,
-    DEFAULT_PROMPT,
 )
 from surf_spot_finder.agents import RUNNERS
+from surf_spot_finder.prompts.shared import INPUT_PROMPT
 from surf_spot_finder.tracing import get_tracer_provider, setup_tracing
 
 
 @logger.catch(reraise=True)
 def find_surf_spot(
-    location: str,
-    date: str,
-    max_driving_hours: int,
-    model_id: str,
+    location: Optional[str] = None,
+    date: Optional[str] = None,
+    max_driving_hours: Optional[int] = None,
+    model_id: Optional[str] = None,
     agent_type: str = "smolagents",
     api_key_var: Optional[str] = None,
-    prompt: str = DEFAULT_PROMPT,
+    input_prompt_template: str = INPUT_PROMPT,
     json_tracer: bool = True,
     api_base: Optional[str] = None,
+    from_config: Optional[str] = None,
 ):
-    logger.info("Loading config")
-    config = Config(
-        location=location,
-        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,
-        api_base=api_base,
-    )
+    """Find the best surf spot based on the given criteria.
+
+    Args:
+        location (str): The location to search around.
+            Required if `from_config` is not provided.
+        date (str): The date to search for.
+            Required if `from_config` is not provided.
+        max_driving_hours (int): The maximum driving hours from the location.
+            Required if `from_config` is not provided.
+        model_id (str): The ID of the model to use.
+            Required if `from_config` is not provided.
+
+            If using `agent_type=smolagents`, use LiteLLM syntax (e.g., 'openai/o1', 'anthropic/claude-3-sonnet').
+            If using `agent_type={openai,openai_multi_agent}`, use OpenAI syntax (e.g., 'o1').
+        agent_type (str, optional): The type of agent to use.
+            Must be one of the supported types in [RUNNERS][surf_spot_finder.agents.RUNNERS].
+        api_key_var (Optional[str], optional): The name of the environment variable containing the API key.
+        input_prompt_template (str, optional): The template for the imput_prompt.
+
+            Must contain the following placeholders: `{LOCATION}`, `{MAX_DRIVING_HOURS}`, and `{DATE}`.
+        json_tracer (bool, optional): Whether to use the custom JSON file exporter.
+        api_base (Optional[str], optional): The base URL for the API.
+        from_config (Optional[str], optional): Path to a YAML config file.
+
+            If provided, all other arguments will be ignored.
+    """
+    if from_config:
+        logger.info(f"Loading {from_config}")
+        config = Config.model_validate(yaml.safe_load(Path(from_config).read_text()))
+    else:
+        config = Config(
+            location=location,
+            date=date,
+            max_driving_hours=max_driving_hours,
+            model_id=model_id,
+            agent_type=agent_type,
+            api_key_var=api_key_var,
+            prompt=input_prompt_template,
+            json_tracer=json_tracer,
+            api_base=api_base,
+        )
 
     logger.info("Setting up tracing")
     tracer_provider, _ = get_tracer_provider(
@@ -45,7 +77,7 @@ def find_surf_spot(
     logger.info(f"Running {config.agent_type} agent")
     RUNNERS[config.agent_type](
         model_id=config.model_id,
-        prompt=config.prompt.format(
+        prompt=config.input_prompt_template.format(
             LOCATION=config.location,
             MAX_DRIVING_HOURS=config.max_driving_hours,
             DATE=config.date,

src/surf_spot_finder/config.py

@@ -1,14 +1,7 @@
 from typing import Annotated, Optional
 from pydantic import AfterValidator, BaseModel, FutureDatetime, PositiveInt
-from datetime import datetime
 
-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}?"
-)
+from surf_spot_finder.prompts.shared import INPUT_PROMPT
 
 
 def validate_prompt(value) -> str:
@@ -26,7 +19,9 @@ def validate_agent_type(value) -> str:
 
 
 class Config(BaseModel):
-    prompt: Annotated[str, AfterValidator(validate_prompt)]
+    input_prompt_template: Annotated[str, AfterValidator(validate_prompt)] = (
+        INPUT_PROMPT
+    )
     location: str
     max_driving_hours: PositiveInt
     date: FutureDatetime

src/surf_spot_finder/prompts/openai.py

@@ -0,0 +1,20 @@
+SINGLE_AGENT_SYSTEM_PROMPT = """
+You will be asked to perform a task.
+
+Before solving the task, plan a sequence of actions using the available tools.
+Then, execute the sequence of actions using the tools.
+""".strip()
+
+MULTI_AGENT_SYSTEM_PROMPT = """
+You will be asked to perform a task.
+
+Always follow this steps:
+
+First, before solving the task, 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.
+""".strip()

src/surf_spot_finder/prompts/shared.py

@@ -0,0 +1,5 @@
+INPUT_PROMPT = """
+According to the forecast, what will be the best spot to surf around {LOCATION},
+in a {MAX_DRIVING_HOURS} hour driving radius,
+at {DATE}?"
+""".strip()

tests/unit/agents/test_unit_openai.py

@@ -9,7 +9,10 @@ from surf_spot_finder.agents.openai import (
     search_web,
     user_verification,
     visit_webpage,
-    DEFAULT_MULTIAGENT_INSTRUCTIONS,
+)
+from surf_spot_finder.prompts.openai import (
+    SINGLE_AGENT_SYSTEM_PROMPT,
+    MULTI_AGENT_SYSTEM_PROMPT,
 )
 
 
@@ -23,7 +26,7 @@ def test_run_openai_agent_default():
         run_openai_agent("gpt-4o", "Test prompt")
         mock_agent.assert_called_once_with(
             model="gpt-4o",
-            instructions=None,
+            instructions=SINGLE_AGENT_SYSTEM_PROMPT,
             name="surf-spot-finder",
             tools=[search_web, visit_webpage],
         )
@@ -91,7 +94,7 @@ def test_run_openai_multiagent():
 
         mock_agent.assert_any_call(
             model="gpt-4o",
-            instructions=DEFAULT_MULTIAGENT_INSTRUCTIONS,
+            instructions=MULTI_AGENT_SYSTEM_PROMPT,
             name="surf-spot-finder",
             # TODO: add more elaborated checks
             handoffs=ANY,

tests/unit/agents/test_unit_smolagents.py

@@ -1,103 +1,84 @@
 import os
 import pytest
 from unittest.mock import patch, MagicMock
+import contextlib
 
 from surf_spot_finder.agents.smolagents import run_smolagent
 
 
 @pytest.fixture
-def mock_smolagents_imports():
-    """Mock the smolagents imports to avoid actual instantiation."""
-    mock_code_agent = MagicMock()
-    mock_ddg_tool = MagicMock()
-    mock_litellm_model = MagicMock()
+def common_patches():
+    litellm_model_mock = MagicMock()
+    code_agent_mock = MagicMock()
+    patch_context = contextlib.ExitStack()
     mock_tool_collection = MagicMock()
 
-    # Configure the mock tool collection to work as a context manager
     mock_tool_collection.from_mcp.return_value.__enter__.return_value = (
         mock_tool_collection
     )
     mock_tool_collection.from_mcp.return_value.__exit__.return_value = None
     mock_tool_collection.tools = ["mock_tool"]
+    patch_context.enter_context(
+        patch("surf_spot_finder.agents.smolagents.StdioServerParameters", MagicMock())
+    )
+    patch_context.enter_context(
+        patch("surf_spot_finder.agents.smolagents.CodeAgent", code_agent_mock)
+    )
+    patch_context.enter_context(
+        patch("surf_spot_finder.agents.smolagents.LiteLLMModel", litellm_model_mock)
+    )
+    patch_context.enter_context(
+        patch("surf_spot_finder.agents.smolagents.ToolCollection", mock_tool_collection)
+    )
+    yield patch_context, litellm_model_mock, code_agent_mock, mock_tool_collection
+    patch_context.close()
+
 
-    with patch.dict(
-        "sys.modules",
-        {
-            "smolagents": MagicMock(
-                CodeAgent=mock_code_agent,
-                DuckDuckGoSearchTool=mock_ddg_tool,
-                LiteLLMModel=mock_litellm_model,
-                ToolCollection=mock_tool_collection,
-            ),
-            "mcp": MagicMock(
-                StdioServerParameters=MagicMock(),
-            ),
-        },
-    ):
-        yield {
-            "CodeAgent": mock_code_agent,
-            "DuckDuckGoSearchTool": mock_ddg_tool,
-            "LiteLLMModel": mock_litellm_model,
-            "ToolCollection": mock_tool_collection,
-        }
-
-
-@pytest.mark.usefixtures("mock_smolagents_imports")
-def test_run_smolagent_with_api_key_var():
-    """Test smolagent creation with an API key from environment variable."""
-    # The patch.dict(os.environ, {"TEST_API_KEY": "test-key-12345"})
-    # is a testing construct that temporarily modifies the environment variables
-    # for the duration of the test.
-    # some tests use TEST_API_KEY while others don't
-    with patch.dict(os.environ, {"TEST_API_KEY": "test-key-12345"}):
-        from smolagents import CodeAgent, LiteLLMModel
+def test_run_smolagent_with_api_key_var(common_patches):
+    patch_context, litellm_model_mock, code_agent_mock, *_ = common_patches
 
+    with patch_context, patch.dict(os.environ, {"TEST_API_KEY": "test-key-12345"}):
         run_smolagent("openai/gpt-4", "Test prompt", api_key_var="TEST_API_KEY")
 
-        LiteLLMModel.assert_called()
-        model_call_kwargs = LiteLLMModel.call_args[1]
+        litellm_model_mock.assert_called()
+        model_call_kwargs = litellm_model_mock.call_args[1]
         assert model_call_kwargs["model_id"] == "openai/gpt-4"
         assert model_call_kwargs["api_key"] == "test-key-12345"
         assert model_call_kwargs["api_base"] is None
 
-        CodeAgent.assert_called_once()
-        CodeAgent.return_value.run.assert_called_once_with("Test prompt")
+        code_agent_mock.assert_called_once()
+        code_agent_mock.return_value.run.assert_called_once_with("Test prompt")
 
 
-@pytest.mark.usefixtures("mock_smolagents_imports")
-def test_run_smolagent_with_custom_api_base():
-    """Test smolagent creation with a custom API base."""
-    with patch.dict(os.environ, {"TEST_API_KEY": "test-key-12345"}):
-        from smolagents import LiteLLMModel
+def test_run_smolagent_with_custom_api_base(common_patches):
+    patch_context, litellm_model_mock, *_ = common_patches
 
-        # Act
+    with patch_context, patch.dict(os.environ, {"TEST_API_KEY": "test-key-12345"}):
         run_smolagent(
             "anthropic/claude-3-sonnet",
             "Test prompt",
             api_key_var="TEST_API_KEY",
             api_base="https://custom-api.example.com",
         )
-        last_call = LiteLLMModel.call_args_list[-1]
+        last_call = litellm_model_mock.call_args_list[-1]
 
         assert last_call[1]["model_id"] == "anthropic/claude-3-sonnet"
         assert last_call[1]["api_key"] == "test-key-12345"
         assert last_call[1]["api_base"] == "https://custom-api.example.com"
 
 
-@pytest.mark.usefixtures("mock_smolagents_imports")
-def test_run_smolagent_without_api_key():
-    """You should be able to run the smolagent without an API key."""
-    from smolagents import LiteLLMModel
+def test_run_smolagent_without_api_key(common_patches):
+    patch_context, litellm_model_mock, *_ = common_patches
 
-    run_smolagent("ollama_chat/deepseek-r1", "Test prompt")
+    with patch_context:
+        run_smolagent("ollama_chat/deepseek-r1", "Test prompt")
 
-    last_call = LiteLLMModel.call_args_list[-1]
+    last_call = litellm_model_mock.call_args_list[-1]
     assert last_call[1]["model_id"] == "ollama_chat/deepseek-r1"
     assert last_call[1]["api_key"] is None
 
 
 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_smolagent("test-model", "Test prompt", api_key_var="MISSING_KEY")