diff --git a/docs/api/ext.md b/docs/api/ext.md index 7f01b44d45..7ec3a4740f 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 43bceed1b7..bf355d33f8 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"} * `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"} * `fastmcp` - installs [FastMCP](mcp/fastmcp-client.md) dependency `fastmcp` [PyPI ↗](https://pypi.org/project/fastmcp){:target="_blank"} diff --git a/docs/third-party-tools.md b/docs/third-party-tools.md index 6f21655252..0e923783aa 100644 --- a/docs/third-party-tools.md +++ b/docs/third-party-tools.md @@ -100,6 +100,59 @@ toolset = ACIToolset( agent = Agent('openai:gpt-5', toolsets=[toolset]) ``` +## StackOne Tools {#stackone-tools} + +If you'd like to use a tool from [StackOne](https://www.stackone.com/) with Pydantic AI, you can use the [`tool_from_stackone`][pydantic_ai.ext.stackone.tool_from_stackone] convenience method. StackOne provides integration infrastructure for AI agents, enabling them to execute actions across 200+ enterprise applications. + +You will need to install the `stackone-ai` package (requires Python 3.10+), set your StackOne API key in the `STACKONE_API_KEY` environment variable, and provide your StackOne account ID via the `STACKONE_ACCOUNT_ID` environment variable or pass it directly to the function. + +Here is how you can use a StackOne tool: + +```python {test="skip"} +import os + +from pydantic_ai import Agent +from pydantic_ai.ext.stackone import tool_from_stackone + +employee_tool = tool_from_stackone( + 'bamboohr_list_employees', + account_id=os.getenv('STACKONE_ACCOUNT_ID'), + api_key=os.getenv('STACKONE_API_KEY'), +) + +agent = Agent( + 'google-gla:gemini-3-pro-preview', + tools=[employee_tool], +) + +result = agent.run_sync('List all employees in the HR system') +print(result.output) +``` + +If you'd like to use multiple StackOne tools, you can use the [`StackOneToolset`][pydantic_ai.ext.stackone.StackOneToolset] [toolset](toolsets.md#stackone-tools). It supports pattern matching via `filter_pattern`, or natural-language tool search via `search_query` to discover relevant tools across your linked accounts. + +Tool search uses the StackOne SDK's semantic search to match tools by natural-language description instead of exact names. This is useful when you don't know the exact tool names available for your connected accounts. You can control the search behavior with three modes: `auto` (default) tries semantic search first and falls back to local search, `semantic` uses only the semantic search API, and `local` uses only local BM25+TF-IDF ranking. Use `top_k` to limit the number of results and `min_similarity` (0-1) to set a minimum relevance threshold. Requires `stackone-ai >= 2.4.0`. + +```python {test="skip"} +import os + +from pydantic_ai import Agent +from pydantic_ai.ext.stackone import StackOneToolset + +toolset = StackOneToolset( + search_query='list and manage employees', + account_id=os.getenv('STACKONE_ACCOUNT_ID'), + api_key=os.getenv('STACKONE_API_KEY'), + top_k=5, + search='auto', +) + +agent = Agent('google-gla:gemini-3-pro-preview', toolsets=[toolset]) + +result = agent.run_sync('List all employees in the HR system') +print(result.output) +``` + ## See Also - [Function Tools](tools.md) - Basic tool concepts and registration @@ -107,3 +160,4 @@ agent = Agent('openai:gpt-5', 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 1605adb73c..8768ca9436 100644 --- a/docs/toolsets.md +++ b/docs/toolsets.md @@ -733,3 +733,54 @@ toolset = ACIToolset( agent = Agent('openai:gpt-5', toolsets=[toolset]) ``` + +### StackOne Tools {#stackone-tools} + +If you'd like to use tools from [StackOne](https://www.stackone.com/) with Pydantic AI, you can use the [`StackOneToolset`][pydantic_ai.ext.stackone.StackOneToolset] which supports pattern matching for tool selection. StackOne provides integration infrastructure for AI agents, enabling them to execute actions across 200+ enterprise applications. + +You will need to install the `stackone-ai` package (requires Python 3.10+), set your StackOne API key in the `STACKONE_API_KEY` environment variable, and provide your StackOne account ID via the `STACKONE_ACCOUNT_ID` environment variable or pass it directly to the toolset. + +```python {test="skip"} +import os + +from pydantic_ai import Agent +from pydantic_ai.ext.stackone import StackOneToolset + +# Use filter patterns to select specific tools +toolset = StackOneToolset( + filter_pattern='bamboohr_*', + account_id=os.getenv('STACKONE_ACCOUNT_ID'), + api_key=os.getenv('STACKONE_API_KEY'), +) + +# Or specify exact tools +specific_toolset = StackOneToolset( + tools=['bamboohr_list_employees', 'bamboohr_get_employee'], # Specific tools only + account_id=os.getenv('STACKONE_ACCOUNT_ID'), + api_key=os.getenv('STACKONE_API_KEY'), +) + +agent = Agent('google-gla:gemini-2.5-flash', toolsets=[toolset]) + +result = agent.run_sync('List all employees and get information about the first employee') +print(result.output) +``` + +You can also use semantic search to find tools by natural-language query (requires `stackone-ai >= 2.4.0`): + +```python {test="skip"} +import os + +from pydantic_ai import Agent +from pydantic_ai.ext.stackone import StackOneToolset + +toolset = StackOneToolset( + search_query='list and manage employees', + account_id=os.getenv('STACKONE_ACCOUNT_ID'), + api_key=os.getenv('STACKONE_API_KEY'), + top_k=5, + search='auto', # 'auto', 'semantic', or 'local' +) + +agent = Agent('google-gla:gemini-3-pro-preview', toolsets=[toolset]) +``` 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 0000000000..e75eab6267 --- /dev/null +++ b/pydantic_ai_slim/pydantic_ai/ext/stackone.py @@ -0,0 +1,107 @@ +from __future__ import annotations + +from collections.abc import Sequence +from typing import 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 + +__all__ = ('tool_from_stackone', 'StackOneToolset') + + +def _tool_from_stackone_tool(stackone_tool: Any) -> 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_id: 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., `"bamboohr_list_employees"`). + account_id: The StackOne account ID. If not provided, uses `STACKONE_ACCOUNT_ID` env var. + api_key: The StackOne API key. If not provided, uses `STACKONE_API_KEY` env var. + base_url: Optional custom base URL for the StackOne API. + + Returns: + A Pydantic AI tool that corresponds to the StackOne tool. + """ + stackone_toolset = StackOneToolSet(api_key=api_key, account_id=account_id, base_url=base_url) + 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_id: str | None = None, + api_key: str | None = None, + base_url: str | None = None, + filter_pattern: str | list[str] | None = None, + search_query: str | None = None, + search: Literal['auto', 'semantic', 'local'] | None = None, + top_k: int | None = None, + min_similarity: float | None = None, + id: str | None = None, + ): + if search_query is not None and tools is not None: + raise ValueError("Cannot specify both 'tools' and 'search_query'") + + stackone_toolset = StackOneToolSet(api_key=api_key, account_id=account_id, base_url=base_url) + + if search_query is not None: + if not hasattr(stackone_toolset, 'search_tools'): + raise ImportError( + 'search_query requires stackone-ai >= 2.4.0. Install with `pip install stackone-ai>=2.4.0`' + ) + search_kwargs: dict[str, Any] = {} + if top_k is not None: + search_kwargs['top_k'] = top_k + if min_similarity is not None: + search_kwargs['min_similarity'] = min_similarity + if search is not None: + search_kwargs['search'] = search + fetched_tools = stackone_toolset.search_tools(search_query, **search_kwargs) + else: + if tools is not None: + actions = list(tools) + else: + actions = [filter_pattern] if isinstance(filter_pattern, str) else 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 e82315ca01..2161931c26 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -224,6 +224,7 @@ 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 "pydantic_ai_slim/pydantic_ai/embeddings/voyageai.py", # voyageai package has no type stubs ] @@ -267,6 +268,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 0000000000..fc57bca3e9 --- /dev/null +++ b/tests/ext/test_stackone.py @@ -0,0 +1,226 @@ +# 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'] + + +@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_id='test-account') + + assert tool.name == 'bamboohr_list_employees' + mock_toolset_cls.assert_called_once_with(api_key='test-key', account_id='test-account', base_url=None) + 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_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') + + +@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', base_url='https://custom.api.stackone.com') + mock_toolset_cls.assert_called_once_with(api_key='k', account_id=None, base_url='https://custom.api.stackone.com') + + +@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_id='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_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') + 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') + 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') + 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', base_url='https://custom.stackone.com') + mock_toolset_cls.assert_called_once_with(api_key='k', account_id=None, base_url='https://custom.stackone.com') + + +@patch('pydantic_ai.ext.stackone.StackOneToolSet') +def test_stackone_toolset_with_search_query(mock_toolset_cls: Any): + mock_found = Mock() + mock_found.__iter__ = Mock(return_value=iter([employee_tool])) + mock_toolset_cls.return_value.search_tools.return_value = mock_found + + toolset = StackOneToolset( + search_query='list employees', + api_key='test-key', + account_id='test-account', + top_k=5, + search='local', + ) + + mock_toolset_cls.return_value.search_tools.assert_called_once_with( + 'list employees', + top_k=5, + search='local', + ) + agent = Agent('test', toolsets=[toolset]) + result = agent.run_sync('foobar') + assert 'bamboohr_list_employees' in result.output + + +@patch('pydantic_ai.ext.stackone.StackOneToolSet') +def test_stackone_toolset_search_query_mutual_exclusion(mock_toolset_cls: Any): + with pytest.raises(ValueError, match="Cannot specify both 'tools' and 'search_query'"): + StackOneToolset( + tools=['bamboohr_list_employees'], + search_query='list employees', + api_key='test-key', + ) + + +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)