diff --git a/docs/api/ext.md b/docs/api/ext.md index 7f01b44d..7ec3a474 100644 --- a/docs/api/ext.md +++ b/docs/api/ext.md @@ -3,3 +3,5 @@ ::: pydantic_ai.ext.langchain ::: pydantic_ai.ext.aci + +::: pydantic_ai.ext.stackone diff --git a/docs/install.md b/docs/install.md index f5f3ec36..9427e339 100644 --- a/docs/install.md +++ b/docs/install.md @@ -59,6 +59,7 @@ pip/uv-add "pydantic-ai-slim[openai]" * `duckduckgo` - installs [DuckDuckGo Search Tool](common-tools.md#duckduckgo-search-tool) dependency `ddgs` [PyPI ↗](https://pypi.org/project/ddgs){:target="_blank"} * `tavily` - installs [Tavily Search Tool](common-tools.md#tavily-search-tool) dependency `tavily-python` [PyPI ↗](https://pypi.org/project/tavily-python){:target="_blank"} * `exa` - installs [Exa Search Tool](common-tools.md#exa-search-tool) dependency `exa-py` [PyPI ↗](https://pypi.org/project/exa-py){:target="_blank"} +* `stackone` - installs [StackOne Tools](third-party-tools.md#stackone-tools) dependency `stackone-ai` [PyPI ↗](https://pypi.org/project/stackone-ai){:target="_blank"} * `web-fetch` - installs [Web Fetch Tool](common-tools.md#web-fetch-tool) dependency `markdownify` [PyPI ↗](https://pypi.org/project/markdownify){:target="_blank"} * `cli` - installs [CLI](cli.md) dependencies `rich` [PyPI ↗](https://pypi.org/project/rich){:target="_blank"}, `prompt-toolkit` [PyPI ↗](https://pypi.org/project/prompt-toolkit){:target="_blank"}, and `argcomplete` [PyPI ↗](https://pypi.org/project/argcomplete){:target="_blank"} * `mcp` - installs [MCP](mcp/client.md) dependency `mcp` [PyPI ↗](https://pypi.org/project/mcp){:target="_blank"} diff --git a/docs/third-party-tools.md b/docs/third-party-tools.md index 9595f6d7..599ddafc 100644 --- a/docs/third-party-tools.md +++ b/docs/third-party-tools.md @@ -100,6 +100,84 @@ toolset = ACIToolset( agent = Agent('openai:gpt-5.2', toolsets=[toolset]) ``` +## StackOne Tools {#stackone-tools} + +[StackOne](https://www.stackone.com) is Integration Infrastructure for AI Agents, providing tools for 200+ enterprise applications. Use the [`StackOneToolset`][pydantic_ai.ext.stackone.StackOneToolset] [toolset](toolsets.md#stackone-tools) to give your agent access to every tool for your connected account(s). Requires the `stackone-ai` package (Python 3.10+), a StackOne API key (`STACKONE_API_KEY`), and at least one connected account ID. Set `STACKONE_ACCOUNT_ID` in your environment, or pass `account_ids=` (a string or list of strings) directly. + +```python {test="skip"} +from pydantic_ai import Agent +from pydantic_ai.ext.stackone import StackOneToolset + +toolset = StackOneToolset() + +agent = Agent('openai:gpt-5.4', toolsets=[toolset]) + +result = agent.run_sync('List the first 5 workers') +print(result.output) +``` + +!!! note + Bare `StackOneToolset()` loads every tool for your connected account(s). Large catalogs may exceed the model's per-request tool-count limit. Use `mode='search_and_execute'` for on-demand discovery (see the Search & Execute tab below), or narrow the toolset with `tools=` or `filter_pattern=` (see the Tools tab below). + +For more control, discover tools on demand, narrow the toolset explicitly, or connect multiple accounts: + +=== "Search & Execute" + + Instead of exposing every tool to the agent, `mode='search_and_execute'` exposes only two tools, `tool_search` and `tool_execute`. The model decides the search query at runtime, picks a result, and calls `tool_execute` with the tool name and parameters. Prompt size stays constant regardless of catalog size. Requires `stackone-ai >= 2.5.0`. + + `search_config['method']` picks the search strategy: + + - `auto` (default): tries `semantic` first, falls back to `local` if the StackOne search API is unavailable + - `semantic`: calls the StackOne semantic search API, which uses server-side model embeddings to rank tools by relevance + - `local`: runs a BM25 + TF-IDF hybrid keyword search on-device, with no call to the StackOne API + + ```python {test="skip"} + # Default: the SDK picks the search strategy + toolset = StackOneToolset(mode='search_and_execute') + + # Or force a strategy explicitly + toolset = StackOneToolset( + mode='search_and_execute', + search_config={'method': 'semantic'}, # or 'auto', 'local' + ) + + agent = Agent('openai:gpt-5.4', toolsets=[toolset]) + ``` + +=== "Tools" + + Narrow the toolset with an explicit list of tool names or a glob pattern. You can also wrap a single tool directly with [`tool_from_stackone`][pydantic_ai.ext.stackone.tool_from_stackone], and a toolset can be combined with one or more standalone tools on the same agent. Examples use Workday tool names; any connected vendor works identically. + + ```python {test="skip"} + from pydantic_ai.ext.stackone import tool_from_stackone + + # An explicit list of tools + toolset = StackOneToolset(tools=['workday_list_workers', 'workday_get_worker']) + + # Or a glob pattern + toolset = StackOneToolset(filter_pattern='workday_list_worker*') + + # Or a single tool without a toolset + worker_tool = tool_from_stackone('workday_list_workers') + + # A toolset and standalone tools can be combined on the same agent + agent = Agent('openai:gpt-5.4', toolsets=[toolset], tools=[worker_tool]) + ``` + + For large catalogs, prefer `tools=` or `filter_pattern=` over loading every tool. Unfiltered loading can exceed model tool-count limits. + +=== "Multiple accounts" + + Pass `account_ids` as a list to fetch tools from several connected accounts. The SDK fans out in parallel and merges the results. + + ```python {test="skip"} + toolset = StackOneToolset( + filter_pattern='workday_list_worker*', + account_ids=['workday-acct-1', 'bamboohr-acct-2'], + ) + agent = Agent('openai:gpt-5.4', toolsets=[toolset]) + ``` + ## See Also - [Function Tools](tools.md) - Basic tool concepts and registration @@ -107,3 +185,4 @@ agent = Agent('openai:gpt-5.2', toolsets=[toolset]) - [MCP Client](mcp/client.md) - Using MCP servers with Pydantic AI - [LangChain Toolsets](toolsets.md#langchain-tools) - Using LangChain toolsets - [ACI.dev Toolsets](toolsets.md#aci-tools) - Using ACI.dev toolsets +- [StackOne Toolsets](toolsets.md#stackone-tools) - Using StackOne toolsets diff --git a/docs/toolsets.md b/docs/toolsets.md index 44d8c3a7..c066a74c 100644 --- a/docs/toolsets.md +++ b/docs/toolsets.md @@ -955,3 +955,20 @@ toolset = ACIToolset( agent = Agent('openai:gpt-5.2', toolsets=[toolset]) ``` + +### StackOne Tools {#stackone-tools} + +If you'd like to use tools from [StackOne](https://www.stackone.com)'s Integration Infrastructure (covering 200+ enterprise applications) with Pydantic AI, you can use the [`StackOneToolset`][pydantic_ai.ext.stackone.StackOneToolset]. + +You will need to install the `stackone-ai` package, set `STACKONE_API_KEY` in your environment, and provide at least one connected account ID via `STACKONE_ACCOUNT_ID` or `account_ids=`. + +```python {test="skip"} +from pydantic_ai import Agent +from pydantic_ai.ext.stackone import StackOneToolset + +toolset = StackOneToolset() + +agent = Agent('openai:gpt-5.4', toolsets=[toolset]) +``` + +See [StackOne Tools](third-party-tools.md#stackone-tools) for on-demand search, tool narrowing, and multi-account usage. diff --git a/pydantic_ai_slim/pydantic_ai/ext/stackone.py b/pydantic_ai_slim/pydantic_ai/ext/stackone.py new file mode 100644 index 00000000..b84faf8b --- /dev/null +++ b/pydantic_ai_slim/pydantic_ai/ext/stackone.py @@ -0,0 +1,154 @@ +from __future__ import annotations + +import os +from collections.abc import Sequence +from typing import TYPE_CHECKING, Any, Literal + +from pydantic.json_schema import JsonSchemaValue + +from pydantic_ai.tools import Tool +from pydantic_ai.toolsets.function import FunctionToolset + +try: + from stackone_ai import StackOneToolSet +except ImportError as _import_error: + raise ImportError('Please install `stackone-ai` to use StackOne tools.') from _import_error + +if TYPE_CHECKING: + from stackone_ai import ExecuteToolsConfig, SearchConfig + from stackone_ai.models import StackOneTool + +__all__ = ('tool_from_stackone', 'StackOneToolset') + + +def _resolve_account_ids( + account_ids: str | list[str] | None, + *, + required: bool = True, +) -> list[str]: + """Return account IDs from the explicit arg or `STACKONE_ACCOUNT_ID` env var. + + Accepts a single string, a list of strings, or `None`. Raises `ValueError` + when nothing is provided and `required=True`. + """ + if isinstance(account_ids, str): + return [account_ids] + if account_ids: + return list(account_ids) + env = os.getenv('STACKONE_ACCOUNT_ID') + if env: + return [env] + if required: + raise ValueError( + 'StackOne account ID(s) required. ' + "Pass `account_ids='acct-1'` or `account_ids=['acct-1', 'acct-2']`, " + 'or set the `STACKONE_ACCOUNT_ID` environment variable.' + ) + return [] + + +def _tool_from_stackone_tool(stackone_tool: StackOneTool) -> Tool: + openai_function = stackone_tool.to_openai_function() + json_schema: JsonSchemaValue = openai_function['function']['parameters'] + + def implementation(**kwargs: Any) -> Any: + return stackone_tool.execute(kwargs) + + return Tool.from_schema( + function=implementation, + name=stackone_tool.name, + description=stackone_tool.description or '', + json_schema=json_schema, + ) + + +def tool_from_stackone( + tool_name: str, + *, + account_ids: str | list[str] | None = None, + api_key: str | None = None, + base_url: str | None = None, +) -> Tool: + """Creates a Pydantic AI tool proxy from a StackOne tool. + + Args: + tool_name: The name of the StackOne tool to wrap (e.g., `"workday_list_workers"`). + account_ids: One or more StackOne account IDs. Pass a string for a single account + or a list for multiple. If not provided, uses the `STACKONE_ACCOUNT_ID` + environment variable. + api_key: The StackOne API key. If not provided, uses the `STACKONE_API_KEY` + environment variable. + base_url: Optional custom base URL for the StackOne API. + + Returns: + A Pydantic AI tool that corresponds to the StackOne tool. + """ + resolved = _resolve_account_ids(account_ids) + stackone_toolset = StackOneToolSet(api_key=api_key, base_url=base_url) + stackone_toolset.set_accounts(resolved) + tools = stackone_toolset.fetch_tools(actions=[tool_name]) + stackone_tool = tools.get_tool(tool_name) + if stackone_tool is None: + raise ValueError(f"Tool '{tool_name}' not found in StackOne") + return _tool_from_stackone_tool(stackone_tool) + + +class StackOneToolset(FunctionToolset): + """A toolset that wraps StackOne tools.""" + + def __init__( + self, + tools: Sequence[str] | None = None, + *, + account_ids: str | list[str] | None = None, + api_key: str | None = None, + base_url: str | None = None, + filter_pattern: str | list[str] | None = None, + mode: Literal['search_and_execute'] | None = None, + search_config: SearchConfig | None = None, + execute_config: ExecuteToolsConfig | None = None, + id: str | None = None, + ): + if mode == 'search_and_execute': + if tools is not None or filter_pattern is not None: + raise ValueError("Cannot combine mode='search_and_execute' with 'tools' or 'filter_pattern'") + has_execute_accounts = execute_config is not None and 'account_ids' in execute_config + if has_execute_accounts and account_ids is not None: + raise ValueError("Cannot specify both 'account_ids' and 'execute_config[\"account_ids\"]'") + resolved = [] if has_execute_accounts else _resolve_account_ids(account_ids) + stackone_toolset = StackOneToolSet( + api_key=api_key, + base_url=base_url, + search=search_config or {}, + execute=execute_config, + ) + if resolved: + stackone_toolset.set_accounts(resolved) + if not hasattr(stackone_toolset, '_build_tools'): + raise ImportError( + "mode='search_and_execute' requires stackone-ai >= 2.5.0. " + 'Install with `pip install stackone-ai>=2.5.0`' + ) + meta_tools = stackone_toolset._build_tools() + pydantic_tools = [_tool_from_stackone_tool(t) for t in meta_tools] + super().__init__(pydantic_tools, id=id) + return + + resolved = _resolve_account_ids(account_ids) + stackone_toolset = StackOneToolSet(api_key=api_key, base_url=base_url) + stackone_toolset.set_accounts(resolved) + + if tools is not None: + actions = list(tools) + elif isinstance(filter_pattern, str): + actions = [filter_pattern] + else: + actions = filter_pattern + + fetched_tools = stackone_toolset.fetch_tools(actions=actions) + + pydantic_tools: list[Tool] = [] + for stackone_tool in fetched_tools: + pydantic_tools.append(_tool_from_stackone_tool(stackone_tool)) + + super().__init__(pydantic_tools, id=id) diff --git a/pyproject.toml b/pyproject.toml index 852b3d02..48aa4940 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -227,6 +227,8 @@ executionEnvironments = [ exclude = [ "examples/pydantic_ai_examples/weather_agent_gradio.py", "pydantic_ai_slim/pydantic_ai/ext/aci.py", # aci-sdk is too niche to be added as an (optional) dependency + "pydantic_ai_slim/pydantic_ai/ext/stackone.py", # stackone-ai is too niche to be added as an (optional) dependency + "tests/ext/test_stackone.py", # stackone-ai is not installed in CI "pydantic_ai_slim/pydantic_ai/embeddings/voyageai.py", # voyageai package has no type stubs ] @@ -294,6 +296,8 @@ include = [ omit = [ "tests/example_modules/*.py", "pydantic_ai_slim/pydantic_ai/ext/aci.py", # aci-sdk is too niche to be added as an (optional) dependency + "pydantic_ai_slim/pydantic_ai/ext/stackone.py", # stackone-ai integration with external API calls + "tests/ext/test_stackone.py", # stackone-ai is not installed in CI "pydantic_ai_slim/pydantic_ai/common_tools/exa.py", # exa-py integration with external API calls # TODO(Marcelo): Enable prefect coverage again. "pydantic_ai_slim/pydantic_ai/durable_exec/prefect/*.py", diff --git a/tests/ext/test_stackone.py b/tests/ext/test_stackone.py new file mode 100644 index 00000000..c3d38c42 --- /dev/null +++ b/tests/ext/test_stackone.py @@ -0,0 +1,308 @@ +# pyright: reportPrivateUsage=false +from __future__ import annotations + +from dataclasses import dataclass, field +from typing import Any +from unittest.mock import Mock, patch + +import pytest +from inline_snapshot import snapshot + +from pydantic_ai import Agent +from pydantic_ai.tools import Tool + +from ..conftest import try_import + +with try_import() as imports_successful: + import pydantic_ai.ext.stackone as stackone_ext + from pydantic_ai.ext.stackone import ( + StackOneToolset, + tool_from_stackone, + ) + +pytestmark = [ + pytest.mark.skipif(not imports_successful(), reason='stackone-ai not installed'), +] + + +@dataclass +class SimulatedStackOneTool: + name: str + description: str + parameters: dict[str, Any] = field(default_factory=lambda: {'type': 'object', 'properties': {}}) + + def to_openai_function(self) -> dict[str, Any]: + return { + 'type': 'function', + 'function': { + 'name': self.name, + 'description': self.description, + 'parameters': self.parameters, + }, + } + + def execute(self, arguments: dict[str, Any]) -> Any: + return f'executed {self.name} with {arguments}' + + +employee_tool = SimulatedStackOneTool( + name='bamboohr_list_employees', + description='List all employees from BambooHR', + parameters={ + 'type': 'object', + 'properties': { + 'limit': {'type': 'integer', 'description': 'Max results to return'}, + }, + }, +) + + +def test_tool_conversion(): + tool = stackone_ext._tool_from_stackone_tool(employee_tool) + assert isinstance(tool, Tool) + assert tool.name == 'bamboohr_list_employees' + assert tool.description == 'List all employees from BambooHR' + + +def test_tool_conversion_with_agent(): + tool = stackone_ext._tool_from_stackone_tool(employee_tool) + agent = Agent('test', tools=[tool]) + result = agent.run_sync('foobar') + assert result.output == snapshot('{"bamboohr_list_employees":"executed bamboohr_list_employees with {}"}') + + +def test_tool_execution(): + tool = stackone_ext._tool_from_stackone_tool(employee_tool) + result = tool.function(limit=10) # type: ignore + assert result == snapshot("executed bamboohr_list_employees with {'limit': 10}") + + +def test_tool_none_description(): + tool_with_none_desc = SimulatedStackOneTool(name='test_tool', description=None) # type: ignore + tool = stackone_ext._tool_from_stackone_tool(tool_with_none_desc) + assert tool.description == '' + + +def test_tool_schema(): + tool = stackone_ext._tool_from_stackone_tool(employee_tool) + assert tool.function_schema.json_schema == employee_tool.to_openai_function()['function']['parameters'] + + +def test_tool_from_stackone_missing_account_id_raises(monkeypatch: pytest.MonkeyPatch): + monkeypatch.delenv('STACKONE_ACCOUNT_ID', raising=False) + with pytest.raises(ValueError, match='StackOne account ID'): + tool_from_stackone('workday_list_workers', api_key='test-key') + + +def test_stackone_toolset_missing_account_id_raises(monkeypatch: pytest.MonkeyPatch): + monkeypatch.delenv('STACKONE_ACCOUNT_ID', raising=False) + with pytest.raises(ValueError, match='StackOne account ID'): + StackOneToolset(api_key='test-key') + + +@patch('pydantic_ai.ext.stackone.StackOneToolSet') +def test_tool_from_stackone(mock_toolset_cls: Any): + mock_tools = Mock() + mock_tools.get_tool.return_value = employee_tool + mock_toolset_cls.return_value.fetch_tools.return_value = mock_tools + + tool = tool_from_stackone('bamboohr_list_employees', api_key='test-key', account_ids='test-account') + + assert tool.name == 'bamboohr_list_employees' + mock_toolset_cls.assert_called_once_with(api_key='test-key', base_url=None) + mock_toolset_cls.return_value.set_accounts.assert_called_once_with(['test-account']) + mock_toolset_cls.return_value.fetch_tools.assert_called_once_with(actions=['bamboohr_list_employees']) + + +@patch('pydantic_ai.ext.stackone.StackOneToolSet') +def test_tool_from_stackone_with_multiple_account_ids(mock_toolset_cls: Any): + mock_tools = Mock() + mock_tools.get_tool.return_value = employee_tool + mock_toolset_cls.return_value.fetch_tools.return_value = mock_tools + + tool_from_stackone('bamboohr_list_employees', api_key='test-key', account_ids=['acct-1', 'acct-2']) + mock_toolset_cls.return_value.set_accounts.assert_called_once_with(['acct-1', 'acct-2']) + + +@patch('pydantic_ai.ext.stackone.StackOneToolSet') +def test_tool_from_stackone_not_found(mock_toolset_cls: Any): + mock_tools = Mock() + mock_tools.get_tool.return_value = None + mock_toolset_cls.return_value.fetch_tools.return_value = mock_tools + + with pytest.raises(ValueError, match="Tool 'nonexistent' not found"): + tool_from_stackone('nonexistent', api_key='test-key', account_ids='test-account') + + +@patch('pydantic_ai.ext.stackone.StackOneToolSet') +def test_tool_from_stackone_with_base_url(mock_toolset_cls: Any): + mock_tools = Mock() + mock_tools.get_tool.return_value = employee_tool + mock_toolset_cls.return_value.fetch_tools.return_value = mock_tools + + tool_from_stackone( + 'bamboohr_list_employees', api_key='k', account_ids='acct', base_url='https://custom.api.stackone.com' + ) + mock_toolset_cls.assert_called_once_with(api_key='k', base_url='https://custom.api.stackone.com') + mock_toolset_cls.return_value.set_accounts.assert_called_once_with(['acct']) + + +@patch('pydantic_ai.ext.stackone.StackOneToolSet') +def test_stackone_toolset_with_tools(mock_toolset_cls: Any): + tool2 = SimulatedStackOneTool(name='bamboohr_get_employee', description='Get an employee') + mock_fetched = Mock() + mock_fetched.__iter__ = Mock(return_value=iter([employee_tool, tool2])) + mock_toolset_cls.return_value.fetch_tools.return_value = mock_fetched + + toolset = StackOneToolset( + tools=['bamboohr_list_employees', 'bamboohr_get_employee'], + api_key='test-key', + account_ids='test-account', + ) + + mock_toolset_cls.return_value.set_accounts.assert_called_once_with(['test-account']) + mock_toolset_cls.return_value.fetch_tools.assert_called_once_with( + actions=['bamboohr_list_employees', 'bamboohr_get_employee'] + ) + agent = Agent('test', toolsets=[toolset]) + result = agent.run_sync('foobar') + assert 'bamboohr_list_employees' in result.output or 'bamboohr_get_employee' in result.output + + +@patch('pydantic_ai.ext.stackone.StackOneToolSet') +def test_stackone_toolset_with_multiple_account_ids(mock_toolset_cls: Any): + mock_fetched = Mock() + mock_fetched.__iter__ = Mock(return_value=iter([employee_tool])) + mock_toolset_cls.return_value.fetch_tools.return_value = mock_fetched + + StackOneToolset(tools=['bamboohr_list_employees'], api_key='test-key', account_ids=['acct-1', 'acct-2']) + mock_toolset_cls.return_value.set_accounts.assert_called_once_with(['acct-1', 'acct-2']) + + +@patch('pydantic_ai.ext.stackone.StackOneToolSet') +def test_stackone_toolset_with_filter_pattern(mock_toolset_cls: Any): + mock_fetched = Mock() + mock_fetched.__iter__ = Mock(return_value=iter([employee_tool])) + mock_toolset_cls.return_value.fetch_tools.return_value = mock_fetched + + StackOneToolset(filter_pattern='bamboohr_*', api_key='test-key', account_ids='test-account') + mock_toolset_cls.return_value.fetch_tools.assert_called_once_with(actions=['bamboohr_*']) + + +@patch('pydantic_ai.ext.stackone.StackOneToolSet') +def test_stackone_toolset_with_list_filter_pattern(mock_toolset_cls: Any): + mock_fetched = Mock() + mock_fetched.__iter__ = Mock(return_value=iter([employee_tool])) + mock_toolset_cls.return_value.fetch_tools.return_value = mock_fetched + + StackOneToolset(filter_pattern=['bamboohr_*', 'workday_*'], api_key='test-key', account_ids='test-account') + mock_toolset_cls.return_value.fetch_tools.assert_called_once_with(actions=['bamboohr_*', 'workday_*']) + + +@patch('pydantic_ai.ext.stackone.StackOneToolSet') +def test_stackone_toolset_no_filter(mock_toolset_cls: Any): + mock_fetched = Mock() + mock_fetched.__iter__ = Mock(return_value=iter([employee_tool])) + mock_toolset_cls.return_value.fetch_tools.return_value = mock_fetched + + StackOneToolset(api_key='test-key', account_ids='test-account') + mock_toolset_cls.return_value.fetch_tools.assert_called_once_with(actions=None) + + +@patch('pydantic_ai.ext.stackone.StackOneToolSet') +def test_stackone_toolset_with_base_url(mock_toolset_cls: Any): + mock_fetched = Mock() + mock_fetched.__iter__ = Mock(return_value=iter([employee_tool])) + mock_toolset_cls.return_value.fetch_tools.return_value = mock_fetched + + StackOneToolset( + tools=['bamboohr_list_employees'], api_key='k', account_ids='acct', base_url='https://custom.stackone.com' + ) + mock_toolset_cls.assert_called_once_with(api_key='k', base_url='https://custom.stackone.com') + + +search_meta_tool = SimulatedStackOneTool( + name='tool_search', + description='Search for available tools', + parameters={'type': 'object', 'properties': {'query': {'type': 'string'}}}, +) + +execute_meta_tool = SimulatedStackOneTool( + name='tool_execute', + description='Execute a tool by name', + parameters={ + 'type': 'object', + 'properties': {'tool_name': {'type': 'string'}, 'parameters': {'type': 'object'}}, + }, +) + + +@patch('pydantic_ai.ext.stackone.StackOneToolSet') +def test_stackone_toolset_search_and_execute_mode(mock_toolset_cls: Any): + mock_meta = Mock() + mock_meta.__iter__ = Mock(return_value=iter([search_meta_tool, execute_meta_tool])) + mock_toolset_cls.return_value._build_tools.return_value = mock_meta + + toolset = StackOneToolset(mode='search_and_execute', api_key='test-key', account_ids='test-account') + + mock_toolset_cls.assert_called_once_with(api_key='test-key', base_url=None, search={}, execute=None) + mock_toolset_cls.return_value.set_accounts.assert_called_once_with(['test-account']) + mock_toolset_cls.return_value._build_tools.assert_called_once() + + agent = Agent('test', toolsets=[toolset]) + result = agent.run_sync('foobar') + assert 'tool_search' in result.output or 'tool_execute' in result.output + + +@patch('pydantic_ai.ext.stackone.StackOneToolSet') +def test_stackone_toolset_search_and_execute_with_config(mock_toolset_cls: Any): + mock_meta = Mock() + mock_meta.__iter__ = Mock(return_value=iter([search_meta_tool, execute_meta_tool])) + mock_toolset_cls.return_value._build_tools.return_value = mock_meta + + StackOneToolset( + mode='search_and_execute', + api_key='test-key', + search_config={'method': 'semantic', 'top_k': 10}, + execute_config={'account_ids': ['acct-1']}, + ) + + mock_toolset_cls.assert_called_once_with( + api_key='test-key', + base_url=None, + search={'method': 'semantic', 'top_k': 10}, + execute={'account_ids': ['acct-1']}, + ) + mock_toolset_cls.return_value.set_accounts.assert_not_called() + + +@patch('pydantic_ai.ext.stackone.StackOneToolSet') +def test_stackone_toolset_search_and_execute_mutual_exclusion(mock_toolset_cls: Any): + with pytest.raises(ValueError, match="Cannot combine mode='search_and_execute'"): + StackOneToolset(mode='search_and_execute', tools=['some_tool'], api_key='test-key', account_ids='test-account') + + with pytest.raises(ValueError, match="Cannot combine mode='search_and_execute'"): + StackOneToolset( + mode='search_and_execute', filter_pattern='bamboohr_*', api_key='test-key', account_ids='test-account' + ) + + +@patch('pydantic_ai.ext.stackone.StackOneToolSet') +def test_stackone_toolset_account_ids_and_execute_config_mutual_exclusion(mock_toolset_cls: Any): + with pytest.raises(ValueError, match="Cannot specify both 'account_ids' and 'execute_config"): + StackOneToolset( + mode='search_and_execute', + api_key='test-key', + account_ids='test-account', + execute_config={'account_ids': ['acct-1']}, + ) + + +def test_import_error(): + with patch.dict('sys.modules', {'stackone_ai': None}): + with pytest.raises(ImportError, match='Please install `stackone-ai`'): + import importlib + + import pydantic_ai.ext.stackone + + importlib.reload(pydantic_ai.ext.stackone)