From 888b2d44d8b3c7962dd2c481c0531a3061c882f8 Mon Sep 17 00:00:00 2001 From: ryoppippi <1560508+ryoppippi@users.noreply.github.com> Date: Thu, 22 Jan 2026 11:25:58 +0000 Subject: [PATCH 01/25] feat(stackone): add StackOne AI integration Add integration with StackOne unified API platform for HRIS, ATS, CRM, and other business systems. - Add `tool_from_stackone()` for single tool creation - Add `StackOneToolset` for bulk tool registration with pattern matching - Require stackone-ai>=2.1.1 (Python 3.10+) - Add documentation and examples - Add comprehensive unit tests with 100% coverage --- docs/install.md | 1 + docs/third-party-tools.md | 32 ++ docs/toolsets.md | 33 ++ examples/stackone_integration.py | 51 +++ pydantic_ai_slim/pydantic_ai/ext/stackone.py | 96 +++++ pydantic_ai_slim/pyproject.toml | 1 + tests/test_ext_stackone.py | 371 +++++++++++++++++++ uv.lock | 37 +- 8 files changed, 621 insertions(+), 1 deletion(-) create mode 100644 examples/stackone_integration.py create mode 100644 pydantic_ai_slim/pydantic_ai/ext/stackone.py create mode 100644 tests/test_ext_stackone.py 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..532182aeb8 100644 --- a/docs/third-party-tools.md +++ b/docs/third-party-tools.md @@ -100,6 +100,37 @@ toolset = ACIToolset( agent = Agent('openai:gpt-5', toolsets=[toolset]) ``` +## StackOne Tools {#stackone-tools} + +If you'd like to use a tool from the [StackOne unified API platform](https://www.stackone.co/) with Pydantic AI, you can use the [`tool_from_stackone`][pydantic_ai.ext.stackone.tool_from_stackone] convenience method. StackOne provides unified APIs for HRIS, ATS, CRM, and other business systems. + +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 the StackOne `hris_list_employees` 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( + 'hris_list_employees', + account_id=os.getenv('STACKONE_ACCOUNT_ID'), + api_key=os.getenv('STACKONE_API_KEY'), +) + +agent = Agent( + 'openai:gpt-5', + 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) which supports pattern matching for tool selection. + ## See Also - [Function Tools](tools.md) - Basic tool concepts and registration @@ -107,3 +138,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 a309441384..986431d0b2 100644 --- a/docs/toolsets.md +++ b/docs/toolsets.md @@ -727,3 +727,36 @@ toolset = ACIToolset( agent = Agent('openai:gpt-5', toolsets=[toolset]) ``` + +### StackOne Tools {#stackone-tools} + +If you'd like to use tools from the [StackOne unified API platform](https://www.stackone.co/) with Pydantic AI, you can use the [`StackOneToolset`][pydantic_ai.ext.stackone.StackOneToolset] which supports pattern matching for tool selection. StackOne provides unified APIs for HRIS, ATS, CRM, and other business systems. + +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='hris_*', # Include all HRIS tools + account_id=os.getenv('STACKONE_ACCOUNT_ID'), + api_key=os.getenv('STACKONE_API_KEY'), +) + +# Or specify exact tools +specific_toolset = StackOneToolset( + tools=['hris_list_employees', 'hris_get_employee'], # Specific tools only + account_id=os.getenv('STACKONE_ACCOUNT_ID'), + api_key=os.getenv('STACKONE_API_KEY'), +) + +agent = Agent('openai:gpt-5', toolsets=[toolset]) + +# Example usage +result = agent.run_sync('List all employees and get information about the first employee') +print(result.output) +``` diff --git a/examples/stackone_integration.py b/examples/stackone_integration.py new file mode 100644 index 0000000000..ba6851190c --- /dev/null +++ b/examples/stackone_integration.py @@ -0,0 +1,51 @@ +"""Example of integrating StackOne tools with Pydantic AI.""" + +import os + +from pydantic_ai import Agent +from pydantic_ai.ext.stackone import StackOneToolset, tool_from_stackone + + +def single_tool_example(): + """Example using a single StackOne tool.""" + employee_tool = tool_from_stackone( + 'hris_list_employees', + account_id=os.getenv('STACKONE_ACCOUNT_ID'), + api_key=os.getenv('STACKONE_API_KEY'), + ) + + agent = Agent('openai:gpt-5', tools=[employee_tool]) + result = agent.run_sync('List all employees') + print(result.output) + + +def toolset_with_filter_example(): + """Example using StackOne toolset with filter pattern.""" + toolset = StackOneToolset( + filter_pattern='hris_*', # Get all HRIS tools + account_id=os.getenv('STACKONE_ACCOUNT_ID'), + api_key=os.getenv('STACKONE_API_KEY'), + ) + + agent = Agent('openai:gpt-5', toolsets=[toolset]) + result = agent.run_sync('Get employee information') + print(result.output) + + +def toolset_with_specific_tools_example(): + """Example using StackOne toolset with specific tools.""" + toolset = StackOneToolset( + tools=['hris_list_employees', 'hris_get_employee'], # Specific tools only + account_id=os.getenv('STACKONE_ACCOUNT_ID'), + api_key=os.getenv('STACKONE_API_KEY'), + ) + + agent = Agent('openai:gpt-5', toolsets=[toolset]) + result = agent.run_sync('Get information about all employees') + print(result.output) + + +if __name__ == '__main__': + single_tool_example() + toolset_with_filter_example() + toolset_with_specific_tools_example() 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..495c365d0c --- /dev/null +++ b/pydantic_ai_slim/pydantic_ai/ext/stackone.py @@ -0,0 +1,96 @@ +from __future__ import annotations + +from collections.abc import Sequence +from typing import Any + +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 + + +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., "hris_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: Custom base URL for StackOne API. Optional. + + 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} if base_url else {}), + ) + + 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") + + 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, + ) + + +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, + id: str | None = None, + ): + if tools is not None: + tool_names = list(tools) + else: + temp_toolset = StackOneToolSet( + api_key=api_key, + account_id=account_id, + **({'base_url': base_url} if base_url else {}), + ) + actions = [filter_pattern] if isinstance(filter_pattern, str) else filter_pattern + filtered_tools = temp_toolset.fetch_tools(actions=actions) + tool_names = [tool.name for tool in filtered_tools] + + super().__init__( + [ + tool_from_stackone( + tool_name, + account_id=account_id, + api_key=api_key, + base_url=base_url, + ) + for tool_name in tool_names + ], + id=id, + ) diff --git a/pydantic_ai_slim/pyproject.toml b/pydantic_ai_slim/pyproject.toml index 76243511b0..93a233989d 100644 --- a/pydantic_ai_slim/pyproject.toml +++ b/pydantic_ai_slim/pyproject.toml @@ -97,6 +97,7 @@ outlines-vllm-offline = [ duckduckgo = ["ddgs>=9.0.0"] tavily = ["tavily-python>=0.5.0"] exa = ["exa-py>=2.0.0"] +stackone = ["stackone-ai>=2.1.1"] # CLI cli = [ "rich>=13", diff --git a/tests/test_ext_stackone.py b/tests/test_ext_stackone.py new file mode 100644 index 0000000000..4461fe29c8 --- /dev/null +++ b/tests/test_ext_stackone.py @@ -0,0 +1,371 @@ +"""Tests for StackOne integration with Pydantic AI.""" + +from __future__ import annotations + +from typing import TYPE_CHECKING +from unittest.mock import Mock, patch + +import pytest + +from pydantic_ai.tools import Tool +from pydantic_ai.toolsets.function import FunctionToolset + +if TYPE_CHECKING: + from unittest.mock import MagicMock + +try: + import stackone_ai # noqa: F401 # pyright: ignore[reportUnusedImport] +except ImportError: # pragma: lax no cover + stackone_installed = False +else: + stackone_installed = True + + +class TestStackOneImportError: + """Test import error handling.""" + + def test_import_error_without_stackone(self): + """Test that ImportError is raised when stackone-ai is not available.""" + # Test that importing the module raises the expected error when stackone_ai is not available + with patch.dict('sys.modules', {'stackone_ai': None}): + with pytest.raises(ImportError, match='Please install `stackone-ai`'): + # Force reimport by using importlib + import importlib + + import pydantic_ai.ext.stackone + + importlib.reload(pydantic_ai.ext.stackone) + + +@pytest.mark.skipif(not stackone_installed, reason='stackone-ai not installed') +class TestToolFromStackOne: + """Test the tool_from_stackone function.""" + + @patch('pydantic_ai.ext.stackone.StackOneToolSet') + def test_tool_creation(self, mock_stackone_toolset_class: MagicMock) -> None: + """Test creating a single tool from StackOne.""" + from pydantic_ai.ext.stackone import tool_from_stackone + + # Mock the StackOne tool + mock_tool = Mock() + mock_tool.name = 'hris_list_employees' + mock_tool.description = 'List all employees' + mock_tool.execute.return_value = {'employees': []} + mock_tool.to_openai_function.return_value = { + 'type': 'function', + 'function': { + 'name': 'hris_list_employees', + 'description': 'List all employees', + 'parameters': { + 'type': 'object', + 'properties': {'limit': {'type': 'integer', 'description': 'Limit the number of results'}}, + }, + }, + } + + mock_tools = Mock() + mock_tools.get_tool.return_value = mock_tool + + mock_stackone_toolset = Mock() + mock_stackone_toolset.fetch_tools.return_value = mock_tools + mock_stackone_toolset_class.return_value = mock_stackone_toolset + + # Create the tool + tool = tool_from_stackone('hris_list_employees', account_id='test-account', api_key='test-key') + + # Verify tool creation + assert isinstance(tool, Tool) + assert tool.name == 'hris_list_employees' + assert tool.description == 'List all employees' + + # Verify StackOneToolSet was called with correct parameters + mock_stackone_toolset_class.assert_called_once_with(api_key='test-key', account_id='test-account') + + # Verify fetch_tools was called with actions parameter + mock_stackone_toolset.fetch_tools.assert_called_once_with(actions=['hris_list_employees']) + mock_tools.get_tool.assert_called_once_with('hris_list_employees') + # Verify returned Tool has correct JSON schema based on StackOne definition + expected = mock_tool.to_openai_function()['function']['parameters'] + assert tool.function_schema.json_schema == expected + + @patch('pydantic_ai.ext.stackone.StackOneToolSet') + def test_tool_not_found(self, mock_stackone_toolset_class: MagicMock) -> None: + """Test error when tool is not found.""" + from pydantic_ai.ext.stackone import tool_from_stackone + + # Mock the tools to return None for the requested tool + mock_tools = Mock() + mock_tools.get_tool.return_value = None + + mock_stackone_toolset = Mock() + mock_stackone_toolset.fetch_tools.return_value = mock_tools + mock_stackone_toolset_class.return_value = mock_stackone_toolset + + # Should raise ValueError when tool not found + with pytest.raises(ValueError, match="Tool 'unknown_tool' not found in StackOne"): + tool_from_stackone('unknown_tool', api_key='test-key') + + @patch('pydantic_ai.ext.stackone.StackOneToolSet') + def test_tool_with_base_url(self, mock_stackone_toolset_class: MagicMock) -> None: + """Test creating a tool with custom base URL. + + Note: base_url is not commonly used by end users, but this test exists for coverage. + """ + from pydantic_ai.ext.stackone import tool_from_stackone + + # Mock the StackOne tool + mock_tool = Mock() + mock_tool.name = 'hris_list_employees' + mock_tool.description = 'List all employees' + mock_tool.to_openai_function.return_value = { + 'type': 'function', + 'function': { + 'name': 'hris_list_employees', + 'description': 'List all employees', + 'parameters': {'type': 'object', 'properties': {}}, + }, + } + + mock_tools = Mock() + mock_tools.get_tool.return_value = mock_tool + + mock_stackone_toolset = Mock() + mock_stackone_toolset.fetch_tools.return_value = mock_tools + mock_stackone_toolset_class.return_value = mock_stackone_toolset + + # Create tool with base URL and verify json_schema conversion + tool = tool_from_stackone('hris_list_employees', api_key='test-key', base_url='https://custom.api.stackone.com') + # Verify base URL was passed to StackOneToolSet + mock_stackone_toolset_class.assert_called_once_with( + api_key='test-key', account_id=None, base_url='https://custom.api.stackone.com' + ) + # Verify returned Tool has correct schema + expected = mock_tool.to_openai_function()['function']['parameters'] + assert tool.function_schema.json_schema == expected + + @patch('pydantic_ai.ext.stackone.StackOneToolSet') + def test_default_parameters(self, mock_stackone_toolset_class: MagicMock) -> None: + """Test default account_id and base_url are None when not provided.""" + from pydantic_ai.ext.stackone import tool_from_stackone + + mock_tool = Mock() + mock_tool.name = 'foo' + mock_tool.description = 'bar' + mock_tool.to_openai_function.return_value = { + 'type': 'function', + 'function': {'name': 'foo', 'description': 'bar', 'parameters': {}}, + } + mock_tools = Mock() + mock_tools.get_tool.return_value = mock_tool + mock_stackone_toolset = Mock() + mock_stackone_toolset.fetch_tools.return_value = mock_tools + mock_stackone_toolset_class.return_value = mock_stackone_toolset + + tool = tool_from_stackone('foo', api_key='key-only') + mock_stackone_toolset_class.assert_called_once_with(api_key='key-only', account_id=None) + expected = mock_tool.to_openai_function()['function']['parameters'] + assert tool.function_schema.json_schema == expected + + @patch('pydantic_ai.ext.stackone.StackOneToolSet') + def test_tool_with_none_description(self, mock_stackone_toolset_class: MagicMock) -> None: + """Test creating a tool when description is None.""" + from pydantic_ai.ext.stackone import tool_from_stackone + + mock_tool = Mock() + mock_tool.name = 'test_tool' + mock_tool.description = None # None description should become empty string + mock_tool.to_openai_function.return_value = { + 'function': {'name': 'test_tool', 'parameters': {'type': 'object'}}, + } + + mock_tools = Mock() + mock_tools.get_tool.return_value = mock_tool + + mock_stackone_toolset = Mock() + mock_stackone_toolset.fetch_tools.return_value = mock_tools + mock_stackone_toolset_class.return_value = mock_stackone_toolset + + tool = tool_from_stackone('test_tool', api_key='test-key') + assert tool.description == '' + + +@pytest.mark.skipif(not stackone_installed, reason='stackone-ai not installed') +class TestStackOneToolset: + """Test the StackOneToolset class.""" + + @patch('pydantic_ai.ext.stackone.tool_from_stackone') + @patch('pydantic_ai.ext.stackone.StackOneToolSet') + def test_toolset_with_specific_tools( + self, mock_stackone_toolset_class: MagicMock, mock_tool_from_stackone: MagicMock + ) -> None: + """Test creating a StackOneToolset with specific tools.""" + from pydantic_ai.ext.stackone import StackOneToolset + + # Mock tool_from_stackone to return different mock tools + def create_mock_tool(tool_name: str) -> Mock: + mock_tool = Mock(spec=Tool) + mock_tool.name = tool_name + mock_tool.max_retries = None + return mock_tool + + def side_effect_func(name: str, **kwargs: object) -> Mock: + return create_mock_tool(name) + + mock_tool_from_stackone.side_effect = side_effect_func + + # Create the toolset with specific tools + toolset = StackOneToolset( + tools=['hris_list_employees', 'hris_get_employee'], account_id='test-account', api_key='test-key' + ) + + # Verify it's a FunctionToolset + assert isinstance(toolset, FunctionToolset) + + # Verify tool_from_stackone was called for each tool + assert mock_tool_from_stackone.call_count == 2 + mock_tool_from_stackone.assert_any_call( + 'hris_list_employees', account_id='test-account', api_key='test-key', base_url=None + ) + mock_tool_from_stackone.assert_any_call( + 'hris_get_employee', account_id='test-account', api_key='test-key', base_url=None + ) + + @patch('pydantic_ai.ext.stackone.tool_from_stackone') + @patch('pydantic_ai.ext.stackone.StackOneToolSet') + def test_toolset_with_filter_pattern( + self, mock_stackone_toolset_class: MagicMock, mock_tool_from_stackone: MagicMock + ) -> None: + """Test creating a StackOneToolset with filter_pattern.""" + from pydantic_ai.ext.stackone import StackOneToolset + + # Mock the StackOneToolSet to return tool names + mock_tool_obj = Mock() + mock_tool_obj.name = 'hris_list_employees' + + mock_stackone_toolset = Mock() + mock_stackone_toolset.fetch_tools.return_value = [mock_tool_obj] + mock_stackone_toolset_class.return_value = mock_stackone_toolset + + # Mock tool_from_stackone + mock_tool = Mock(spec=Tool) + mock_tool.name = 'hris_list_employees' + mock_tool.max_retries = None + mock_tool_from_stackone.return_value = mock_tool + + # Create toolset with filter_pattern + toolset = StackOneToolset(filter_pattern='hris_*', account_id='test-account', api_key='test-key') + + # Verify StackOneToolSet was created correctly + mock_stackone_toolset_class.assert_called_once_with(api_key='test-key', account_id='test-account') + + # Verify fetch_tools was called with actions parameter (list) + mock_stackone_toolset.fetch_tools.assert_called_once_with(actions=['hris_*']) + + # Verify tools were created + assert isinstance(toolset, FunctionToolset) + + @patch('pydantic_ai.ext.stackone.tool_from_stackone') + @patch('pydantic_ai.ext.stackone.StackOneToolSet') + def test_toolset_with_list_filter_pattern( + self, mock_stackone_toolset_class: MagicMock, mock_tool_from_stackone: MagicMock + ) -> None: + """Test creating a StackOneToolset with list filter_pattern.""" + from pydantic_ai.ext.stackone import StackOneToolset + + # Mock the StackOneToolSet to return tool names + mock_tool1 = Mock() + mock_tool1.name = 'hris_list_employees' + mock_tool2 = Mock() + mock_tool2.name = 'ats_get_job' + + mock_stackone_toolset = Mock() + mock_stackone_toolset.fetch_tools.return_value = [mock_tool1, mock_tool2] + mock_stackone_toolset_class.return_value = mock_stackone_toolset + + # Mock tool_from_stackone to return different tools for each call + def mock_tool_side_effect(tool_name: str, **kwargs: object) -> Mock: + mock_tool = Mock(spec=Tool) + mock_tool.name = tool_name + mock_tool.max_retries = None + return mock_tool + + mock_tool_from_stackone.side_effect = mock_tool_side_effect + + # Create toolset with list filter_pattern + toolset = StackOneToolset(filter_pattern=['hris_*', 'ats_*'], account_id='test-account', api_key='test-key') + + # Verify fetch_tools was called with list filter_pattern as actions + mock_stackone_toolset.fetch_tools.assert_called_once_with(actions=['hris_*', 'ats_*']) + + # Verify tools were created for both returned tools + assert mock_tool_from_stackone.call_count == 2 + + # Verify it's a FunctionToolset + assert isinstance(toolset, FunctionToolset) + + @patch('pydantic_ai.ext.stackone.tool_from_stackone') + @patch('pydantic_ai.ext.stackone.StackOneToolSet') + def test_toolset_without_filter_pattern( + self, mock_stackone_toolset_class: MagicMock, mock_tool_from_stackone: MagicMock + ) -> None: + """Test creating a StackOneToolset without filter_pattern (gets all tools).""" + from pydantic_ai.ext.stackone import StackOneToolset + + # Mock the StackOneToolSet to return all tools + mock_tool = Mock() + mock_tool.name = 'all_tools' + + mock_stackone_toolset = Mock() + mock_stackone_toolset.fetch_tools.return_value = [mock_tool] + mock_stackone_toolset_class.return_value = mock_stackone_toolset + + # Mock tool_from_stackone + mock_pydantic_tool = Mock(spec=Tool) + mock_pydantic_tool.name = 'all_tools' + mock_pydantic_tool.max_retries = None + mock_tool_from_stackone.return_value = mock_pydantic_tool + + # Create toolset without filter_pattern + toolset = StackOneToolset(account_id='test-account', api_key='test-key') + + # Verify fetch_tools was called with None actions (no filter) + mock_stackone_toolset.fetch_tools.assert_called_once_with(actions=None) + + # Verify tools were created + assert isinstance(toolset, FunctionToolset) + + @patch('pydantic_ai.ext.stackone.tool_from_stackone') + @patch('pydantic_ai.ext.stackone.StackOneToolSet') + def test_toolset_with_base_url( + self, mock_stackone_toolset_class: MagicMock, mock_tool_from_stackone: MagicMock + ) -> None: + """Test creating a StackOneToolset with custom base URL. + + Note: base_url is not commonly used by end users, but this test exists for coverage. + """ + from pydantic_ai.ext.stackone import StackOneToolset + + # Mock tool_from_stackone + mock_tool = Mock(spec=Tool) + mock_tool.name = 'hris_list_employees' + mock_tool.max_retries = None + mock_tool_from_stackone.return_value = mock_tool + + # Create toolset with base URL + toolset = StackOneToolset( + tools=['hris_list_employees'], + account_id='test-account', + api_key='test-key', + base_url='https://custom.api.stackone.com', + ) + + # Verify tool_from_stackone was called with base URL + mock_tool_from_stackone.assert_called_once_with( + 'hris_list_employees', + account_id='test-account', + api_key='test-key', + base_url='https://custom.api.stackone.com', + ) + + # Verify it's a FunctionToolset + assert isinstance(toolset, FunctionToolset) diff --git a/uv.lock b/uv.lock index 0be00c3592..56ee0219e3 100644 --- a/uv.lock +++ b/uv.lock @@ -697,6 +697,20 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/b0/5c/dbd00727a3dd165d7e0e8af40e630cd7e45d77b525a3218afaff8a87358e/blake3-1.0.8-cp314-cp314t-win_amd64.whl", hash = "sha256:421b99cdf1ff2d1bf703bc56c454f4b286fce68454dd8711abbcb5a0df90c19a", size = 215133, upload-time = "2025-10-14T06:47:16.069Z" }, ] +[[package]] +name = "bm25s" +version = "0.2.14" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "numpy" }, + { name = "scipy", version = "1.15.3", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.11'" }, + { name = "scipy", version = "1.16.3", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.11'" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/99/72/5ad06c30991ba494242785a3ab8987deb01c07dfc1c492847bde221e62bf/bm25s-0.2.14.tar.gz", hash = "sha256:7b6717770fffbdb3b962e5fe8ef1e6eac7f285d0fbc14484b321e136df837139", size = 59266, upload-time = "2025-09-08T17:06:30.728Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/7c/3e/e3ae2f0fb0f8f46f9c787fa419ca5203ff850d0630749a26baf0a6570453/bm25s-0.2.14-py3-none-any.whl", hash = "sha256:76cdb70ae40747941b150a1ec16a9c20c576d6534d0a3c3eebb303c779b3cf65", size = 55128, upload-time = "2025-09-08T17:06:29.324Z" }, +] + [[package]] name = "boto3" version = "1.42.14" @@ -6293,6 +6307,9 @@ retries = [ sentence-transformers = [ { name = "sentence-transformers" }, ] +stackone = [ + { name = "stackone-ai" }, +] tavily = [ { name = "tavily-python" }, ] @@ -6358,6 +6375,7 @@ requires-dist = [ { name = "requests", marker = "extra == 'vertexai'", specifier = ">=2.32.2" }, { name = "rich", marker = "extra == 'cli'", specifier = ">=13" }, { name = "sentence-transformers", marker = "extra == 'sentence-transformers'", specifier = ">=5.2.0" }, + { name = "stackone-ai", marker = "extra == 'stackone'", specifier = ">=2.1.1" }, { name = "starlette", marker = "extra == 'ag-ui'", specifier = ">=0.45.3" }, { name = "starlette", marker = "extra == 'ui'", specifier = ">=0.45.3" }, { name = "starlette", marker = "extra == 'web'", specifier = ">=0.45.3" }, @@ -6373,7 +6391,7 @@ requires-dist = [ { name = "vllm", marker = "(python_full_version < '3.12' and platform_machine != 'x86_64' and extra == 'outlines-vllm-offline') or (python_full_version < '3.12' and sys_platform != 'darwin' and extra == 'outlines-vllm-offline')", specifier = ">=0.11.0" }, { name = "voyageai", marker = "extra == 'voyageai'", specifier = ">=0.3.2" }, ] -provides-extras = ["a2a", "ag-ui", "anthropic", "bedrock", "cli", "cohere", "dbos", "duckduckgo", "evals", "exa", "fastmcp", "google", "groq", "huggingface", "logfire", "mcp", "mistral", "openai", "openrouter", "outlines-llamacpp", "outlines-mlxlm", "outlines-sglang", "outlines-transformers", "outlines-vllm-offline", "prefect", "retries", "sentence-transformers", "tavily", "temporal", "ui", "vertexai", "voyageai", "web"] +provides-extras = ["a2a", "ag-ui", "anthropic", "bedrock", "cli", "cohere", "dbos", "duckduckgo", "evals", "exa", "fastmcp", "google", "groq", "huggingface", "logfire", "mcp", "mistral", "openai", "openrouter", "outlines-llamacpp", "outlines-mlxlm", "outlines-sglang", "outlines-transformers", "outlines-vllm-offline", "prefect", "retries", "sentence-transformers", "stackone", "tavily", "temporal", "ui", "vertexai", "voyageai", "web"] [[package]] name = "pydantic-core" @@ -8286,6 +8304,23 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/71/22/8ab1066358601163e1ac732837adba3672f703818f693e179b24e0d3b65c/sse_starlette-3.0.4-py3-none-any.whl", hash = "sha256:32c80ef0d04506ced4b0b6ab8fe300925edc37d26f666afb1874c754895f5dc3", size = 11764, upload-time = "2025-12-14T16:22:51.453Z" }, ] +[[package]] +name = "stackone-ai" +version = "2.1.1" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "bm25s" }, + { name = "httpx" }, + { name = "langchain-core" }, + { name = "numpy" }, + { name = "pydantic" }, + { name = "typing-extensions" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/0f/0c/a76539757074286f8fff3c98a84b02d985ca337e5634e17900cf800def98/stackone_ai-2.1.1.tar.gz", hash = "sha256:1595732ebea295058f87a9d2cd8b1abfbcff41f7cd16006cfd5337b1528f4f78", size = 537753, upload-time = "2026-01-22T10:49:03.264Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/69/6c/0978c9687721a0545f0f81befb70ac2531170204acad10bb6ee058d235bf/stackone_ai-2.1.1-py3-none-any.whl", hash = "sha256:49023c39824e3c22a87838c4c46c65ec2256e205cd902207b42d9495ca612e71", size = 30833, upload-time = "2026-01-22T10:49:04.207Z" }, +] + [[package]] name = "starlette" version = "0.50.0" From e369d0715f831efecc7812650493037622114c44 Mon Sep 17 00:00:00 2001 From: ryoppippi <1560508+ryoppippi@users.noreply.github.com> Date: Thu, 22 Jan 2026 13:59:39 +0000 Subject: [PATCH 02/25] docs(stackone): fix URL and update description - Fix StackOne URL from .co to .com - Update description to reflect current positioning as "AI Integration Gateway" --- docs/third-party-tools.md | 2 +- docs/toolsets.md | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/docs/third-party-tools.md b/docs/third-party-tools.md index 532182aeb8..0fac28a516 100644 --- a/docs/third-party-tools.md +++ b/docs/third-party-tools.md @@ -102,7 +102,7 @@ agent = Agent('openai:gpt-5', toolsets=[toolset]) ## StackOne Tools {#stackone-tools} -If you'd like to use a tool from the [StackOne unified API platform](https://www.stackone.co/) with Pydantic AI, you can use the [`tool_from_stackone`][pydantic_ai.ext.stackone.tool_from_stackone] convenience method. StackOne provides unified APIs for HRIS, ATS, CRM, and other business systems. +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 is the AI Integration Gateway that connects AI agents to hundreds of SaaS integrations through a single unified interface. 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. diff --git a/docs/toolsets.md b/docs/toolsets.md index 986431d0b2..99baaa6878 100644 --- a/docs/toolsets.md +++ b/docs/toolsets.md @@ -730,7 +730,7 @@ agent = Agent('openai:gpt-5', toolsets=[toolset]) ### StackOne Tools {#stackone-tools} -If you'd like to use tools from the [StackOne unified API platform](https://www.stackone.co/) with Pydantic AI, you can use the [`StackOneToolset`][pydantic_ai.ext.stackone.StackOneToolset] which supports pattern matching for tool selection. StackOne provides unified APIs for HRIS, ATS, CRM, and other business systems. +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 is the AI Integration Gateway that connects AI agents to hundreds of SaaS integrations through a single unified interface. 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. From a5e02c37565848c23fa1cc6afbd3d1aaefa2380f Mon Sep 17 00:00:00 2001 From: ryoppippi <1560508+ryoppippi@users.noreply.github.com> Date: Fri, 30 Jan 2026 15:30:32 +0000 Subject: [PATCH 03/25] feat(stackone): add utility tools for dynamic discovery and feedback Add search_tool, execute_tool, and feedback_tool functions that wrap StackOne's utility tools API (tool_search, tool_execute, tool_feedback). Extend StackOneToolset with new options: - include_utility_tools: enables dynamic tool discovery mode - include_feedback_tool: adds feedback collection capability - hybrid_alpha: configures BM25/TF-IDF search weight Update documentation and examples to reflect the new API. --- docs/third-party-tools.md | 80 +++++- examples/stackone_integration.py | 123 ++++++++- pydantic_ai_slim/pydantic_ai/ext/stackone.py | 262 +++++++++++++++++-- 3 files changed, 430 insertions(+), 35 deletions(-) diff --git a/docs/third-party-tools.md b/docs/third-party-tools.md index 0fac28a516..257840c074 100644 --- a/docs/third-party-tools.md +++ b/docs/third-party-tools.md @@ -102,11 +102,11 @@ 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 is the AI Integration Gateway that connects AI agents to hundreds of SaaS integrations through a single unified interface. +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 the StackOne `hris_list_employees` tool: +Here is how you can use a StackOne tool: ```python {test="skip"} import os @@ -115,7 +115,7 @@ from pydantic_ai import Agent from pydantic_ai.ext.stackone import tool_from_stackone employee_tool = tool_from_stackone( - 'hris_list_employees', + 'stackone_list_employees', account_id=os.getenv('STACKONE_ACCOUNT_ID'), api_key=os.getenv('STACKONE_API_KEY'), ) @@ -131,6 +131,80 @@ 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) which supports pattern matching for tool selection. +### Dynamic Tool Discovery {#stackone-dynamic-discovery} + +For large tool sets where you don't know in advance which tools the agent will need, you can enable **utility tools mode** by setting `include_utility_tools=True`. This provides two special tools: + +- `tool_search`: Searches for relevant tools using natural language queries +- `tool_execute`: Executes a discovered tool by name + +```python {test="skip"} +import os + +from pydantic_ai import Agent +from pydantic_ai.ext.stackone import StackOneToolset + +toolset = StackOneToolset( + filter_pattern='stackone_*', # Load StackOne tools + include_utility_tools=True, # Enable dynamic discovery + account_id=os.getenv('STACKONE_ACCOUNT_ID'), + api_key=os.getenv('STACKONE_API_KEY'), +) + +agent = Agent('openai:gpt-5', toolsets=[toolset]) +result = agent.run_sync('Find a tool to list employees and use it') +print(result.output) +``` + +You can also use the standalone [`search_tool`][pydantic_ai.ext.stackone.search_tool] and [`execute_tool`][pydantic_ai.ext.stackone.execute_tool] functions for more control: + +```python {test="skip"} +from stackone_ai import StackOneToolSet + +from pydantic_ai import Agent +from pydantic_ai.ext.stackone import execute_tool, search_tool + +stackone = StackOneToolSet() +tools = stackone.fetch_tools(actions=['stackone_*']) + +agent = Agent( + 'openai:gpt-5', + tools=[search_tool(tools), execute_tool(tools)], +) +``` + +### Feedback Collection {#stackone-feedback} + +StackOne supports collecting user feedback on tool performance. Enable this by setting `include_feedback_tool=True`: + +```python {test="skip"} +import os + +from pydantic_ai import Agent +from pydantic_ai.ext.stackone import StackOneToolset + +toolset = StackOneToolset( + filter_pattern='stackone_*', + include_feedback_tool=True, + account_id=os.getenv('STACKONE_ACCOUNT_ID'), + api_key=os.getenv('STACKONE_API_KEY'), +) + +agent = Agent('openai:gpt-5', toolsets=[toolset]) +``` + +You can also use the standalone [`feedback_tool`][pydantic_ai.ext.stackone.feedback_tool] function: + +```python {test="skip"} +from pydantic_ai import Agent +from pydantic_ai.ext.stackone import feedback_tool + +agent = Agent( + 'openai:gpt-5', + tools=[feedback_tool()], +) +``` + ## See Also - [Function Tools](tools.md) - Basic tool concepts and registration diff --git a/examples/stackone_integration.py b/examples/stackone_integration.py index ba6851190c..bdd415191a 100644 --- a/examples/stackone_integration.py +++ b/examples/stackone_integration.py @@ -3,13 +3,19 @@ import os from pydantic_ai import Agent -from pydantic_ai.ext.stackone import StackOneToolset, tool_from_stackone +from pydantic_ai.ext.stackone import ( + StackOneToolset, + execute_tool, + feedback_tool, + search_tool, + tool_from_stackone, +) def single_tool_example(): """Example using a single StackOne tool.""" employee_tool = tool_from_stackone( - 'hris_list_employees', + 'stackone_list_employees', account_id=os.getenv('STACKONE_ACCOUNT_ID'), api_key=os.getenv('STACKONE_API_KEY'), ) @@ -22,7 +28,7 @@ def single_tool_example(): def toolset_with_filter_example(): """Example using StackOne toolset with filter pattern.""" toolset = StackOneToolset( - filter_pattern='hris_*', # Get all HRIS tools + filter_pattern='stackone_*', # Get all StackOne tools account_id=os.getenv('STACKONE_ACCOUNT_ID'), api_key=os.getenv('STACKONE_API_KEY'), ) @@ -35,7 +41,10 @@ def toolset_with_filter_example(): def toolset_with_specific_tools_example(): """Example using StackOne toolset with specific tools.""" toolset = StackOneToolset( - tools=['hris_list_employees', 'hris_get_employee'], # Specific tools only + tools=[ + 'stackone_list_employees', + 'stackone_get_employee', + ], # Specific tools only account_id=os.getenv('STACKONE_ACCOUNT_ID'), api_key=os.getenv('STACKONE_API_KEY'), ) @@ -45,7 +54,113 @@ def toolset_with_specific_tools_example(): print(result.output) +def utility_tools_example(): + """Example using StackOne with utility tools for dynamic discovery. + + Utility tools mode provides two special tools: + - tool_search: Search for relevant tools using natural language + - tool_execute: Execute a discovered tool by name + + This is useful when you have a large number of tools and want the agent + to dynamically discover and use the right ones. + """ + toolset = StackOneToolset( + filter_pattern='stackone_*', # Load all StackOne tools + include_utility_tools=True, # Enable dynamic discovery + account_id=os.getenv('STACKONE_ACCOUNT_ID'), + api_key=os.getenv('STACKONE_API_KEY'), + ) + + agent = Agent('openai:gpt-5', toolsets=[toolset]) + result = agent.run_sync('Find a tool to list employees and use it') + print(result.output) + + +def standalone_utility_tools_example(): + """Example using standalone search_tool and execute_tool functions. + + This gives you more control over how the tools are created and used. + """ + from stackone_ai import StackOneToolSet + + # Fetch tools from StackOne + stackone = StackOneToolSet() + tools = stackone.fetch_tools(actions=['stackone_*']) + + # Create search and execute tools + agent = Agent( + 'openai:gpt-5', + tools=[search_tool(tools), execute_tool(tools)], + ) + + result = agent.run_sync('Search for employee-related tools') + print(result.output) + + +def feedback_example(): + """Example using the feedback collection tool via StackOneToolset. + + The feedback tool allows users to provide feedback on their experience + with StackOne tools, which helps improve the tool ecosystem. + """ + toolset = StackOneToolset( + filter_pattern='stackone_*', + include_feedback_tool=True, + account_id=os.getenv('STACKONE_ACCOUNT_ID'), + api_key=os.getenv('STACKONE_API_KEY'), + ) + + agent = Agent('openai:gpt-5', toolsets=[toolset]) + result = agent.run_sync('List employees and then ask for feedback') + print(result.output) + + +def standalone_feedback_example(): + """Example using the standalone feedback_tool function. + + This gives you more control over when and how to include the feedback tool. + """ + fb_tool = feedback_tool( + api_key=os.getenv('STACKONE_API_KEY'), + account_id=os.getenv('STACKONE_ACCOUNT_ID'), + ) + + agent = Agent('openai:gpt-5', tools=[fb_tool]) + result = agent.run_sync('Collect feedback about a recent tool usage') + print(result.output) + + +def full_featured_example(): + """Example using all StackOne features together. + + This example demonstrates: + - Utility tools for dynamic discovery + - Feedback collection + - Custom hybrid_alpha for search tuning + """ + toolset = StackOneToolset( + filter_pattern=['stackone_*'], # Filter pattern + include_utility_tools=True, + include_feedback_tool=True, + hybrid_alpha=0.3, # Custom search tuning + account_id=os.getenv('STACKONE_ACCOUNT_ID'), + api_key=os.getenv('STACKONE_API_KEY'), + ) + + agent = Agent('openai:gpt-5', toolsets=[toolset]) + result = agent.run_sync( + 'Find the best tool to get employee information, use it, ' + 'and then collect feedback about the experience' + ) + print(result.output) + + if __name__ == '__main__': single_tool_example() toolset_with_filter_example() toolset_with_specific_tools_example() + utility_tools_example() + standalone_utility_tools_example() + feedback_example() + standalone_feedback_example() + full_featured_example() diff --git a/pydantic_ai_slim/pydantic_ai/ext/stackone.py b/pydantic_ai_slim/pydantic_ai/ext/stackone.py index 495c365d0c..ca35eba2b6 100644 --- a/pydantic_ai_slim/pydantic_ai/ext/stackone.py +++ b/pydantic_ai_slim/pydantic_ai/ext/stackone.py @@ -13,6 +13,32 @@ except ImportError as _import_error: raise ImportError('Please install `stackone-ai` to use StackOne tools.') from _import_error +# StackOneTools is dynamically typed as stackone_ai doesn't provide complete type stubs +StackOneTools = Any + + +def _tool_from_stackone_tool(stackone_tool: Any) -> Tool: + """Creates a Pydantic AI tool from a StackOneTool instance. + + Args: + stackone_tool: A StackOneTool instance. + + Returns: + A Pydantic AI tool that wraps the StackOne 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, @@ -24,7 +50,7 @@ def tool_from_stackone( """Creates a Pydantic AI tool proxy from a StackOne tool. Args: - tool_name: The name of the StackOne tool to wrap (e.g., "hris_list_employees"). + tool_name: The name of the StackOne tool to wrap (e.g., "stackone_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: Custom base URL for StackOne API. Optional. @@ -43,22 +69,189 @@ def tool_from_stackone( if stackone_tool is None: raise ValueError(f"Tool '{tool_name}' not found in StackOne") - openai_function = stackone_tool.to_openai_function() - json_schema: JsonSchemaValue = openai_function['function']['parameters'] + return _tool_from_stackone_tool(stackone_tool) - 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 search_tool( + tools: StackOneTools, + *, + hybrid_alpha: float | None = None, +) -> Tool: + """Creates a search tool for discovering StackOne tools using natural language. + + This tool uses hybrid BM25 + TF-IDF search to find relevant tools based on + a natural language query. + + Args: + tools: A StackOne Tools collection (returned from `fetch_tools()`). + hybrid_alpha: Weight for BM25 in hybrid search (0-1). Default is 0.2, + which has been shown to provide better tool discovery accuracy. + + Returns: + A Pydantic AI tool for searching StackOne tools. + + Example: + ```python + from pydantic_ai.ext.stackone import StackOneToolset, search_tool + from stackone_ai import StackOneToolSet + + stackone = StackOneToolSet() + tools = stackone.fetch_tools(actions=['stackone_*']) + + agent = Agent( + 'openai:gpt-4o', + tools=[search_tool(tools)], + ) + ``` + """ + utility_tools = tools.utility_tools(hybrid_alpha=hybrid_alpha) + search = utility_tools.get_tool('tool_search') + if search is None: + raise ValueError('tool_search not found in StackOne utility tools') + + return _tool_from_stackone_tool(search) + + +def execute_tool( + tools: StackOneTools, +) -> Tool: + """Creates an execute tool for running discovered StackOne tools by name. + + This tool allows executing any tool from the provided collection by name, + typically used after discovering tools with `search_tool`. + + Args: + tools: A StackOne Tools collection (returned from `fetch_tools()`). + + Returns: + A Pydantic AI tool for executing StackOne tools. + + Example: + ```python + from pydantic_ai.ext.stackone import StackOneToolset, search_tool, execute_tool + from stackone_ai import StackOneToolSet + + stackone = StackOneToolSet() + tools = stackone.fetch_tools(actions=['stackone_*']) + + agent = Agent( + 'openai:gpt-4o', + tools=[search_tool(tools), execute_tool(tools)], + ) + ``` + """ + utility_tools = tools.utility_tools() + execute = utility_tools.get_tool('tool_execute') + if execute is None: + raise ValueError('tool_execute not found in StackOne utility tools') + + return _tool_from_stackone_tool(execute) + + +def feedback_tool( + *, + api_key: str | None = None, + account_id: str | None = None, + base_url: str | None = None, +) -> Tool: + """Creates a feedback tool for collecting user feedback on StackOne tools. + + This tool allows users to provide feedback on their experience with StackOne tools, + which helps improve the tool ecosystem. + + Args: + api_key: The StackOne API key. If not provided, uses STACKONE_API_KEY env var. + account_id: The StackOne account ID. If not provided, uses STACKONE_ACCOUNT_ID env var. + base_url: Custom base URL for StackOne API. Optional. + + Returns: + A Pydantic AI tool for collecting feedback. + + Example: + ```python + from pydantic_ai.ext.stackone import StackOneToolset, feedback_tool + + agent = Agent( + 'openai:gpt-4o', + tools=[feedback_tool()], + ) + ``` + """ + try: + from stackone_ai.feedback.tool import create_feedback_tool + except ImportError as e: + raise ImportError('Please install `stackone-ai` with feedback support to use the feedback tool.') from e + + # Get API key from environment if not provided + if api_key is None: + import os + + api_key = os.environ.get('STACKONE_API_KEY') + if api_key is None: + raise ValueError( + 'API key is required. Provide it as an argument or set STACKONE_API_KEY environment variable.' + ) + + fb_tool = create_feedback_tool( + api_key=api_key, + account_id=account_id, + **({'base_url': base_url} if base_url else {}), ) + return _tool_from_stackone_tool(fb_tool) + class StackOneToolset(FunctionToolset): - """A toolset that wraps StackOne tools.""" + """A toolset that wraps StackOne tools. + + This toolset provides access to StackOne's integration infrastructure for AI agents, + offering 200+ connectors across HR, ATS, CRM, and other business applications. + It can operate in two modes: + + 1. **Direct mode** (default): Each StackOne tool is exposed directly to the agent. + This is best when you have a small, known set of tools. + + 2. **Utility tools mode** (`include_utility_tools=True`): Instead of exposing all tools + directly, provides `tool_search` and `tool_execute` that allow the agent to + dynamically discover and execute tools. This is better for large tool sets + where the agent needs to search for the right tool. + + Args: + tools: Specific tool names to include (e.g., ["stackone_list_employees"]). + account_id: The StackOne account ID. Uses STACKONE_ACCOUNT_ID env var if not provided. + api_key: The StackOne API key. Uses STACKONE_API_KEY env var if not provided. + base_url: Custom base URL for StackOne API. + filter_pattern: Glob pattern(s) to filter tools (e.g., "stackone_*"). + include_utility_tools: If True, includes search and execute utility tools instead of + individual tools. Default is False. + include_feedback_tool: If True, includes the feedback collection tool. + Default is False. + hybrid_alpha: Weight for BM25 in hybrid search (0-1) when using utility tools. + Default is 0.2. + id: Optional ID for the toolset. + + Example: + ```python + from pydantic_ai import Agent + from pydantic_ai.ext.stackone import StackOneToolset + + # Direct mode - expose specific tools + agent = Agent( + 'openai:gpt-4o', + toolsets=[StackOneToolset(filter_pattern='stackone_*')], + ) + + # Utility tools mode - dynamic discovery + agent = Agent( + 'openai:gpt-4o', + toolsets=[StackOneToolset( + filter_pattern='stackone_*', + include_utility_tools=True, + include_feedback_tool=True, + )], + ) + ``` + """ def __init__( self, @@ -68,29 +261,42 @@ def __init__( api_key: str | None = None, base_url: str | None = None, filter_pattern: str | list[str] | None = None, + include_utility_tools: bool = False, + include_feedback_tool: bool = False, + hybrid_alpha: float | None = None, id: str | None = None, ): + stackone_toolset = StackOneToolSet( + api_key=api_key, + account_id=account_id, + **({'base_url': base_url} if base_url else {}), + ) + if tools is not None: - tool_names = list(tools) + actions = list(tools) else: - temp_toolset = StackOneToolSet( - api_key=api_key, - account_id=account_id, - **({'base_url': base_url} if base_url else {}), - ) actions = [filter_pattern] if isinstance(filter_pattern, str) else filter_pattern - filtered_tools = temp_toolset.fetch_tools(actions=actions) - tool_names = [tool.name for tool in filtered_tools] - super().__init__( - [ - tool_from_stackone( - tool_name, - account_id=account_id, + fetched_tools = stackone_toolset.fetch_tools(actions=actions) + + pydantic_tools: list[Tool] = [] + + if include_utility_tools: + # Utility tools mode: provide search and execute tools + pydantic_tools.append(search_tool(fetched_tools, hybrid_alpha=hybrid_alpha)) + pydantic_tools.append(execute_tool(fetched_tools)) + else: + # Direct mode: expose each tool individually + for stackone_tool in fetched_tools: + pydantic_tools.append(_tool_from_stackone_tool(stackone_tool)) + + if include_feedback_tool: + pydantic_tools.append( + feedback_tool( api_key=api_key, + account_id=account_id, base_url=base_url, ) - for tool_name in tool_names - ], - id=id, - ) + ) + + super().__init__(pydantic_tools, id=id) From 1cb34c37f9afb8ec2c7ade805d91d02837ea1ae2 Mon Sep 17 00:00:00 2001 From: ryoppippi <1560508+ryoppippi@users.noreply.github.com> Date: Fri, 30 Jan 2026 15:32:45 +0000 Subject: [PATCH 04/25] chore(stackone): bump stackone-ai to >=2.3.0 Update minimum version requirement to 2.3.0 which includes the utility_tools API (tool_search, tool_execute, tool_feedback). --- pydantic_ai_slim/pyproject.toml | 2 +- uv.lock | 8 ++++---- 2 files changed, 5 insertions(+), 5 deletions(-) diff --git a/pydantic_ai_slim/pyproject.toml b/pydantic_ai_slim/pyproject.toml index 93a233989d..1477566f65 100644 --- a/pydantic_ai_slim/pyproject.toml +++ b/pydantic_ai_slim/pyproject.toml @@ -97,7 +97,7 @@ outlines-vllm-offline = [ duckduckgo = ["ddgs>=9.0.0"] tavily = ["tavily-python>=0.5.0"] exa = ["exa-py>=2.0.0"] -stackone = ["stackone-ai>=2.1.1"] +stackone = ["stackone-ai>=2.3.0"] # CLI cli = [ "rich>=13", diff --git a/uv.lock b/uv.lock index 56ee0219e3..1715e22851 100644 --- a/uv.lock +++ b/uv.lock @@ -6375,7 +6375,7 @@ requires-dist = [ { name = "requests", marker = "extra == 'vertexai'", specifier = ">=2.32.2" }, { name = "rich", marker = "extra == 'cli'", specifier = ">=13" }, { name = "sentence-transformers", marker = "extra == 'sentence-transformers'", specifier = ">=5.2.0" }, - { name = "stackone-ai", marker = "extra == 'stackone'", specifier = ">=2.1.1" }, + { name = "stackone-ai", marker = "extra == 'stackone'", specifier = ">=2.3.0" }, { name = "starlette", marker = "extra == 'ag-ui'", specifier = ">=0.45.3" }, { name = "starlette", marker = "extra == 'ui'", specifier = ">=0.45.3" }, { name = "starlette", marker = "extra == 'web'", specifier = ">=0.45.3" }, @@ -8306,7 +8306,7 @@ wheels = [ [[package]] name = "stackone-ai" -version = "2.1.1" +version = "2.3.1" source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "bm25s" }, @@ -8316,9 +8316,9 @@ dependencies = [ { name = "pydantic" }, { name = "typing-extensions" }, ] -sdist = { url = "https://files.pythonhosted.org/packages/0f/0c/a76539757074286f8fff3c98a84b02d985ca337e5634e17900cf800def98/stackone_ai-2.1.1.tar.gz", hash = "sha256:1595732ebea295058f87a9d2cd8b1abfbcff41f7cd16006cfd5337b1528f4f78", size = 537753, upload-time = "2026-01-22T10:49:03.264Z" } +sdist = { url = "https://files.pythonhosted.org/packages/36/06/cc058d329e44a73fa7aa25f03e42aa0ee15d5f22aca534f21f7eb8cf63a7/stackone_ai-2.3.1.tar.gz", hash = "sha256:d25b521ad6cee7f8107c9beef9a05a67bc39b07646d4ac4f4665fea1dbe3d601", size = 537863, upload-time = "2026-01-29T16:59:20.683Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/69/6c/0978c9687721a0545f0f81befb70ac2531170204acad10bb6ee058d235bf/stackone_ai-2.1.1-py3-none-any.whl", hash = "sha256:49023c39824e3c22a87838c4c46c65ec2256e205cd902207b42d9495ca612e71", size = 30833, upload-time = "2026-01-22T10:49:04.207Z" }, + { url = "https://files.pythonhosted.org/packages/97/5a/31ef092360e9175e86909df86d1f889f5e7f80df04ceb3e855ed81fa563f/stackone_ai-2.3.1-py3-none-any.whl", hash = "sha256:bd6ffc937f894045d410353eaec456e1a40001dc6e9b28904bd770869061b60d", size = 30818, upload-time = "2026-01-29T16:59:19.7Z" }, ] [[package]] From 24929dadc20bfc741fa8bc55850b797aa0d3bbbd Mon Sep 17 00:00:00 2001 From: ryoppippi <1560508+ryoppippi@users.noreply.github.com> Date: Fri, 30 Jan 2026 15:33:27 +0000 Subject: [PATCH 05/25] chore(stackone): bump stackone-ai to >=2.3.1 --- pydantic_ai_slim/pyproject.toml | 2 +- uv.lock | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/pydantic_ai_slim/pyproject.toml b/pydantic_ai_slim/pyproject.toml index 1477566f65..477bc557c6 100644 --- a/pydantic_ai_slim/pyproject.toml +++ b/pydantic_ai_slim/pyproject.toml @@ -97,7 +97,7 @@ outlines-vllm-offline = [ duckduckgo = ["ddgs>=9.0.0"] tavily = ["tavily-python>=0.5.0"] exa = ["exa-py>=2.0.0"] -stackone = ["stackone-ai>=2.3.0"] +stackone = ["stackone-ai>=2.3.1"] # CLI cli = [ "rich>=13", diff --git a/uv.lock b/uv.lock index 1715e22851..a8a267fc5f 100644 --- a/uv.lock +++ b/uv.lock @@ -6375,7 +6375,7 @@ requires-dist = [ { name = "requests", marker = "extra == 'vertexai'", specifier = ">=2.32.2" }, { name = "rich", marker = "extra == 'cli'", specifier = ">=13" }, { name = "sentence-transformers", marker = "extra == 'sentence-transformers'", specifier = ">=5.2.0" }, - { name = "stackone-ai", marker = "extra == 'stackone'", specifier = ">=2.3.0" }, + { name = "stackone-ai", marker = "extra == 'stackone'", specifier = ">=2.3.1" }, { name = "starlette", marker = "extra == 'ag-ui'", specifier = ">=0.45.3" }, { name = "starlette", marker = "extra == 'ui'", specifier = ">=0.45.3" }, { name = "starlette", marker = "extra == 'web'", specifier = ">=0.45.3" }, From 8927e2cb85894fbf9c6a11d1163e58f197dd015e Mon Sep 17 00:00:00 2001 From: ryoppippi <1560508+ryoppippi@users.noreply.github.com> Date: Fri, 30 Jan 2026 15:36:01 +0000 Subject: [PATCH 06/25] fix(stackone): update tests for new implementation - Update StackOneToolset tests to match new implementation that uses _tool_from_stackone_tool directly instead of tool_from_stackone - Add test for include_utility_tools option - Mark docstring examples with test="skip" to avoid linting issues --- pydantic_ai_slim/pydantic_ai/ext/stackone.py | 19 +- tests/test_ext_stackone.py | 187 ++++++++++--------- 2 files changed, 115 insertions(+), 91 deletions(-) diff --git a/pydantic_ai_slim/pydantic_ai/ext/stackone.py b/pydantic_ai_slim/pydantic_ai/ext/stackone.py index ca35eba2b6..627edc3184 100644 --- a/pydantic_ai_slim/pydantic_ai/ext/stackone.py +++ b/pydantic_ai_slim/pydantic_ai/ext/stackone.py @@ -91,10 +91,12 @@ def search_tool( A Pydantic AI tool for searching StackOne tools. Example: - ```python - from pydantic_ai.ext.stackone import StackOneToolset, search_tool + ```python {test="skip"} from stackone_ai import StackOneToolSet + from pydantic_ai import Agent + from pydantic_ai.ext.stackone import search_tool + stackone = StackOneToolSet() tools = stackone.fetch_tools(actions=['stackone_*']) @@ -127,10 +129,12 @@ def execute_tool( A Pydantic AI tool for executing StackOne tools. Example: - ```python - from pydantic_ai.ext.stackone import StackOneToolset, search_tool, execute_tool + ```python {test="skip"} from stackone_ai import StackOneToolSet + from pydantic_ai import Agent + from pydantic_ai.ext.stackone import execute_tool, search_tool + stackone = StackOneToolSet() tools = stackone.fetch_tools(actions=['stackone_*']) @@ -168,8 +172,9 @@ def feedback_tool( A Pydantic AI tool for collecting feedback. Example: - ```python - from pydantic_ai.ext.stackone import StackOneToolset, feedback_tool + ```python {test="skip"} + from pydantic_ai import Agent + from pydantic_ai.ext.stackone import feedback_tool agent = Agent( 'openai:gpt-4o', @@ -231,7 +236,7 @@ class StackOneToolset(FunctionToolset): id: Optional ID for the toolset. Example: - ```python + ```python {test="skip"} from pydantic_ai import Agent from pydantic_ai.ext.stackone import StackOneToolset diff --git a/tests/test_ext_stackone.py b/tests/test_ext_stackone.py index 4461fe29c8..331dcb2ddb 100644 --- a/tests/test_ext_stackone.py +++ b/tests/test_ext_stackone.py @@ -189,29 +189,41 @@ def test_tool_with_none_description(self, mock_stackone_toolset_class: MagicMock assert tool.description == '' +def _create_mock_stackone_tool(name: str, description: str = 'Test description') -> Mock: + """Helper to create a mock StackOne tool.""" + mock_tool = Mock() + mock_tool.name = name + mock_tool.description = description + mock_tool.to_openai_function.return_value = { + 'type': 'function', + 'function': { + 'name': name, + 'description': description, + 'parameters': {'type': 'object', 'properties': {}}, + }, + } + return mock_tool + + @pytest.mark.skipif(not stackone_installed, reason='stackone-ai not installed') class TestStackOneToolset: """Test the StackOneToolset class.""" - @patch('pydantic_ai.ext.stackone.tool_from_stackone') @patch('pydantic_ai.ext.stackone.StackOneToolSet') - def test_toolset_with_specific_tools( - self, mock_stackone_toolset_class: MagicMock, mock_tool_from_stackone: MagicMock - ) -> None: + def test_toolset_with_specific_tools(self, mock_stackone_toolset_class: MagicMock) -> None: """Test creating a StackOneToolset with specific tools.""" from pydantic_ai.ext.stackone import StackOneToolset - # Mock tool_from_stackone to return different mock tools - def create_mock_tool(tool_name: str) -> Mock: - mock_tool = Mock(spec=Tool) - mock_tool.name = tool_name - mock_tool.max_retries = None - return mock_tool + # Mock the tools returned by fetch_tools + mock_tool1 = _create_mock_stackone_tool('hris_list_employees') + mock_tool2 = _create_mock_stackone_tool('hris_get_employee') - def side_effect_func(name: str, **kwargs: object) -> Mock: - return create_mock_tool(name) + mock_fetched_tools = Mock() + mock_fetched_tools.__iter__ = Mock(return_value=iter([mock_tool1, mock_tool2])) - mock_tool_from_stackone.side_effect = side_effect_func + mock_stackone_toolset = Mock() + mock_stackone_toolset.fetch_tools.return_value = mock_fetched_tools + mock_stackone_toolset_class.return_value = mock_stackone_toolset # Create the toolset with specific tools toolset = StackOneToolset( @@ -221,37 +233,27 @@ def side_effect_func(name: str, **kwargs: object) -> Mock: # Verify it's a FunctionToolset assert isinstance(toolset, FunctionToolset) - # Verify tool_from_stackone was called for each tool - assert mock_tool_from_stackone.call_count == 2 - mock_tool_from_stackone.assert_any_call( - 'hris_list_employees', account_id='test-account', api_key='test-key', base_url=None - ) - mock_tool_from_stackone.assert_any_call( - 'hris_get_employee', account_id='test-account', api_key='test-key', base_url=None - ) + # Verify StackOneToolSet was created correctly + mock_stackone_toolset_class.assert_called_once_with(api_key='test-key', account_id='test-account') + + # Verify fetch_tools was called with the tool names as actions + mock_stackone_toolset.fetch_tools.assert_called_once_with(actions=['hris_list_employees', 'hris_get_employee']) - @patch('pydantic_ai.ext.stackone.tool_from_stackone') @patch('pydantic_ai.ext.stackone.StackOneToolSet') - def test_toolset_with_filter_pattern( - self, mock_stackone_toolset_class: MagicMock, mock_tool_from_stackone: MagicMock - ) -> None: + def test_toolset_with_filter_pattern(self, mock_stackone_toolset_class: MagicMock) -> None: """Test creating a StackOneToolset with filter_pattern.""" from pydantic_ai.ext.stackone import StackOneToolset - # Mock the StackOneToolSet to return tool names - mock_tool_obj = Mock() - mock_tool_obj.name = 'hris_list_employees' + # Mock the tools returned by fetch_tools + mock_tool = _create_mock_stackone_tool('hris_list_employees') + + mock_fetched_tools = Mock() + mock_fetched_tools.__iter__ = Mock(return_value=iter([mock_tool])) mock_stackone_toolset = Mock() - mock_stackone_toolset.fetch_tools.return_value = [mock_tool_obj] + mock_stackone_toolset.fetch_tools.return_value = mock_fetched_tools mock_stackone_toolset_class.return_value = mock_stackone_toolset - # Mock tool_from_stackone - mock_tool = Mock(spec=Tool) - mock_tool.name = 'hris_list_employees' - mock_tool.max_retries = None - mock_tool_from_stackone.return_value = mock_tool - # Create toolset with filter_pattern toolset = StackOneToolset(filter_pattern='hris_*', account_id='test-account', api_key='test-key') @@ -264,67 +266,46 @@ def test_toolset_with_filter_pattern( # Verify tools were created assert isinstance(toolset, FunctionToolset) - @patch('pydantic_ai.ext.stackone.tool_from_stackone') @patch('pydantic_ai.ext.stackone.StackOneToolSet') - def test_toolset_with_list_filter_pattern( - self, mock_stackone_toolset_class: MagicMock, mock_tool_from_stackone: MagicMock - ) -> None: + def test_toolset_with_list_filter_pattern(self, mock_stackone_toolset_class: MagicMock) -> None: """Test creating a StackOneToolset with list filter_pattern.""" from pydantic_ai.ext.stackone import StackOneToolset - # Mock the StackOneToolSet to return tool names - mock_tool1 = Mock() - mock_tool1.name = 'hris_list_employees' - mock_tool2 = Mock() - mock_tool2.name = 'ats_get_job' + # Mock the tools returned by fetch_tools + mock_tool1 = _create_mock_stackone_tool('hris_list_employees') + mock_tool2 = _create_mock_stackone_tool('ats_get_job') + + mock_fetched_tools = Mock() + mock_fetched_tools.__iter__ = Mock(return_value=iter([mock_tool1, mock_tool2])) mock_stackone_toolset = Mock() - mock_stackone_toolset.fetch_tools.return_value = [mock_tool1, mock_tool2] + mock_stackone_toolset.fetch_tools.return_value = mock_fetched_tools mock_stackone_toolset_class.return_value = mock_stackone_toolset - # Mock tool_from_stackone to return different tools for each call - def mock_tool_side_effect(tool_name: str, **kwargs: object) -> Mock: - mock_tool = Mock(spec=Tool) - mock_tool.name = tool_name - mock_tool.max_retries = None - return mock_tool - - mock_tool_from_stackone.side_effect = mock_tool_side_effect - # Create toolset with list filter_pattern toolset = StackOneToolset(filter_pattern=['hris_*', 'ats_*'], account_id='test-account', api_key='test-key') # Verify fetch_tools was called with list filter_pattern as actions mock_stackone_toolset.fetch_tools.assert_called_once_with(actions=['hris_*', 'ats_*']) - # Verify tools were created for both returned tools - assert mock_tool_from_stackone.call_count == 2 - # Verify it's a FunctionToolset assert isinstance(toolset, FunctionToolset) - @patch('pydantic_ai.ext.stackone.tool_from_stackone') @patch('pydantic_ai.ext.stackone.StackOneToolSet') - def test_toolset_without_filter_pattern( - self, mock_stackone_toolset_class: MagicMock, mock_tool_from_stackone: MagicMock - ) -> None: + def test_toolset_without_filter_pattern(self, mock_stackone_toolset_class: MagicMock) -> None: """Test creating a StackOneToolset without filter_pattern (gets all tools).""" from pydantic_ai.ext.stackone import StackOneToolset - # Mock the StackOneToolSet to return all tools - mock_tool = Mock() - mock_tool.name = 'all_tools' + # Mock the tools returned by fetch_tools + mock_tool = _create_mock_stackone_tool('all_tools') + + mock_fetched_tools = Mock() + mock_fetched_tools.__iter__ = Mock(return_value=iter([mock_tool])) mock_stackone_toolset = Mock() - mock_stackone_toolset.fetch_tools.return_value = [mock_tool] + mock_stackone_toolset.fetch_tools.return_value = mock_fetched_tools mock_stackone_toolset_class.return_value = mock_stackone_toolset - # Mock tool_from_stackone - mock_pydantic_tool = Mock(spec=Tool) - mock_pydantic_tool.name = 'all_tools' - mock_pydantic_tool.max_retries = None - mock_tool_from_stackone.return_value = mock_pydantic_tool - # Create toolset without filter_pattern toolset = StackOneToolset(account_id='test-account', api_key='test-key') @@ -334,22 +315,23 @@ def test_toolset_without_filter_pattern( # Verify tools were created assert isinstance(toolset, FunctionToolset) - @patch('pydantic_ai.ext.stackone.tool_from_stackone') @patch('pydantic_ai.ext.stackone.StackOneToolSet') - def test_toolset_with_base_url( - self, mock_stackone_toolset_class: MagicMock, mock_tool_from_stackone: MagicMock - ) -> None: + def test_toolset_with_base_url(self, mock_stackone_toolset_class: MagicMock) -> None: """Test creating a StackOneToolset with custom base URL. Note: base_url is not commonly used by end users, but this test exists for coverage. """ from pydantic_ai.ext.stackone import StackOneToolset - # Mock tool_from_stackone - mock_tool = Mock(spec=Tool) - mock_tool.name = 'hris_list_employees' - mock_tool.max_retries = None - mock_tool_from_stackone.return_value = mock_tool + # Mock the tools returned by fetch_tools + mock_tool = _create_mock_stackone_tool('hris_list_employees') + + mock_fetched_tools = Mock() + mock_fetched_tools.__iter__ = Mock(return_value=iter([mock_tool])) + + mock_stackone_toolset = Mock() + mock_stackone_toolset.fetch_tools.return_value = mock_fetched_tools + mock_stackone_toolset_class.return_value = mock_stackone_toolset # Create toolset with base URL toolset = StackOneToolset( @@ -359,13 +341,50 @@ def test_toolset_with_base_url( base_url='https://custom.api.stackone.com', ) - # Verify tool_from_stackone was called with base URL - mock_tool_from_stackone.assert_called_once_with( - 'hris_list_employees', + # Verify StackOneToolSet was called with base URL + mock_stackone_toolset_class.assert_called_once_with( + api_key='test-key', account_id='test-account', base_url='https://custom.api.stackone.com' + ) + + # Verify it's a FunctionToolset + assert isinstance(toolset, FunctionToolset) + + @patch('pydantic_ai.ext.stackone.StackOneToolSet') + def test_toolset_with_utility_tools(self, mock_stackone_toolset_class: MagicMock) -> None: + """Test creating a StackOneToolset with utility tools enabled.""" + from pydantic_ai.ext.stackone import StackOneToolset + + # Mock the tools returned by fetch_tools + mock_tool = _create_mock_stackone_tool('hris_list_employees') + + mock_fetched_tools = Mock() + mock_fetched_tools.__iter__ = Mock(return_value=iter([mock_tool])) + + # Mock utility tools + mock_search_tool = _create_mock_stackone_tool('tool_search', 'Search for tools') + mock_execute_tool = _create_mock_stackone_tool('tool_execute', 'Execute a tool') + + mock_utility_tools = Mock() + mock_utility_tools.get_tool.side_effect = lambda name: ( + mock_search_tool if name == 'tool_search' else mock_execute_tool if name == 'tool_execute' else None + ) + + mock_fetched_tools.utility_tools.return_value = mock_utility_tools + + mock_stackone_toolset = Mock() + mock_stackone_toolset.fetch_tools.return_value = mock_fetched_tools + mock_stackone_toolset_class.return_value = mock_stackone_toolset + + # Create toolset with utility tools + toolset = StackOneToolset( + filter_pattern='hris_*', + include_utility_tools=True, account_id='test-account', api_key='test-key', - base_url='https://custom.api.stackone.com', ) # Verify it's a FunctionToolset assert isinstance(toolset, FunctionToolset) + + # Verify utility_tools was called (once for search_tool, once for execute_tool) + assert mock_fetched_tools.utility_tools.call_count == 2 From f8288e4a5a3e9a0636028f293fb81790732ed86b Mon Sep 17 00:00:00 2001 From: ryoppippi <1560508+ryoppippi@users.noreply.github.com> Date: Fri, 30 Jan 2026 15:46:20 +0000 Subject: [PATCH 07/25] refactor(stackone): simplify docstring examples to match project style - Remove `{test="skip"}` from docstring examples - Simplify StackOneToolset example to match ExaToolset pattern - Move Args documentation from class docstring to __init__ method - Remove verbose examples from search_tool, execute_tool, feedback_tool - Fix pyright type error in test file with ignore comment --- pydantic_ai_slim/pydantic_ai/ext/stackone.py | 98 +++++--------------- tests/test_ext_stackone.py | 2 +- 2 files changed, 23 insertions(+), 77 deletions(-) diff --git a/pydantic_ai_slim/pydantic_ai/ext/stackone.py b/pydantic_ai_slim/pydantic_ai/ext/stackone.py index 627edc3184..8cbedc719b 100644 --- a/pydantic_ai_slim/pydantic_ai/ext/stackone.py +++ b/pydantic_ai_slim/pydantic_ai/ext/stackone.py @@ -89,22 +89,6 @@ def search_tool( Returns: A Pydantic AI tool for searching StackOne tools. - - Example: - ```python {test="skip"} - from stackone_ai import StackOneToolSet - - from pydantic_ai import Agent - from pydantic_ai.ext.stackone import search_tool - - stackone = StackOneToolSet() - tools = stackone.fetch_tools(actions=['stackone_*']) - - agent = Agent( - 'openai:gpt-4o', - tools=[search_tool(tools)], - ) - ``` """ utility_tools = tools.utility_tools(hybrid_alpha=hybrid_alpha) search = utility_tools.get_tool('tool_search') @@ -127,22 +111,6 @@ def execute_tool( Returns: A Pydantic AI tool for executing StackOne tools. - - Example: - ```python {test="skip"} - from stackone_ai import StackOneToolSet - - from pydantic_ai import Agent - from pydantic_ai.ext.stackone import execute_tool, search_tool - - stackone = StackOneToolSet() - tools = stackone.fetch_tools(actions=['stackone_*']) - - agent = Agent( - 'openai:gpt-4o', - tools=[search_tool(tools), execute_tool(tools)], - ) - ``` """ utility_tools = tools.utility_tools() execute = utility_tools.get_tool('tool_execute') @@ -170,17 +138,6 @@ def feedback_tool( Returns: A Pydantic AI tool for collecting feedback. - - Example: - ```python {test="skip"} - from pydantic_ai import Agent - from pydantic_ai.ext.stackone import feedback_tool - - agent = Agent( - 'openai:gpt-4o', - tools=[feedback_tool()], - ) - ``` """ try: from stackone_ai.feedback.tool import create_feedback_tool @@ -221,41 +178,14 @@ class StackOneToolset(FunctionToolset): dynamically discover and execute tools. This is better for large tool sets where the agent needs to search for the right tool. - Args: - tools: Specific tool names to include (e.g., ["stackone_list_employees"]). - account_id: The StackOne account ID. Uses STACKONE_ACCOUNT_ID env var if not provided. - api_key: The StackOne API key. Uses STACKONE_API_KEY env var if not provided. - base_url: Custom base URL for StackOne API. - filter_pattern: Glob pattern(s) to filter tools (e.g., "stackone_*"). - include_utility_tools: If True, includes search and execute utility tools instead of - individual tools. Default is False. - include_feedback_tool: If True, includes the feedback collection tool. - Default is False. - hybrid_alpha: Weight for BM25 in hybrid search (0-1) when using utility tools. - Default is 0.2. - id: Optional ID for the toolset. - Example: - ```python {test="skip"} - from pydantic_ai import Agent - from pydantic_ai.ext.stackone import StackOneToolset - - # Direct mode - expose specific tools - agent = Agent( - 'openai:gpt-4o', - toolsets=[StackOneToolset(filter_pattern='stackone_*')], - ) + ```python + from pydantic_ai import Agent + from pydantic_ai.ext.stackone import StackOneToolset - # Utility tools mode - dynamic discovery - agent = Agent( - 'openai:gpt-4o', - toolsets=[StackOneToolset( - filter_pattern='stackone_*', - include_utility_tools=True, - include_feedback_tool=True, - )], - ) - ``` + toolset = StackOneToolset(api_key='your-api-key') + agent = Agent('openai:gpt-4o', toolsets=[toolset]) + ``` """ def __init__( @@ -271,6 +201,22 @@ def __init__( hybrid_alpha: float | None = None, id: str | None = None, ): + """Creates a StackOne toolset. + + Args: + tools: Specific tool names to include (e.g., ["stackone_list_employees"]). + account_id: The StackOne account ID. Uses STACKONE_ACCOUNT_ID env var if not provided. + api_key: The StackOne API key. Uses STACKONE_API_KEY env var if not provided. + base_url: Custom base URL for StackOne API. + filter_pattern: Glob pattern(s) to filter tools (e.g., "stackone_*"). + include_utility_tools: If True, includes search and execute utility tools instead of + individual tools. Default is False. + include_feedback_tool: If True, includes the feedback collection tool. + Default is False. + hybrid_alpha: Weight for BM25 in hybrid search (0-1) when using utility tools. + Default is 0.2. + id: Optional ID for the toolset, used for durable execution environments. + """ stackone_toolset = StackOneToolSet( api_key=api_key, account_id=account_id, diff --git a/tests/test_ext_stackone.py b/tests/test_ext_stackone.py index 331dcb2ddb..8a0412795d 100644 --- a/tests/test_ext_stackone.py +++ b/tests/test_ext_stackone.py @@ -365,7 +365,7 @@ def test_toolset_with_utility_tools(self, mock_stackone_toolset_class: MagicMock mock_execute_tool = _create_mock_stackone_tool('tool_execute', 'Execute a tool') mock_utility_tools = Mock() - mock_utility_tools.get_tool.side_effect = lambda name: ( + mock_utility_tools.get_tool.side_effect = lambda name: ( # pyright: ignore[reportUnknownLambdaType] mock_search_tool if name == 'tool_search' else mock_execute_tool if name == 'tool_execute' else None ) From 92c89dd0fa7c6466784b1571fbe092e08eace531 Mon Sep 17 00:00:00 2001 From: ryoppippi <1560508+ryoppippi@users.noreply.github.com> Date: Fri, 30 Jan 2026 15:49:35 +0000 Subject: [PATCH 08/25] fix(stackone): update docs/toolsets.md examples per review feedback - Replace hris_* with stackone_* in filter_pattern example - Replace hris_list_employees/hris_get_employee with stackone_* tools - Update description to remove 'unified interface' terminology --- docs/toolsets.md | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/docs/toolsets.md b/docs/toolsets.md index 97431674b1..ef65a23f9d 100644 --- a/docs/toolsets.md +++ b/docs/toolsets.md @@ -736,7 +736,7 @@ 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 is the AI Integration Gateway that connects AI agents to hundreds of SaaS integrations through a single unified interface. +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. @@ -748,14 +748,14 @@ from pydantic_ai.ext.stackone import StackOneToolset # Use filter patterns to select specific tools toolset = StackOneToolset( - filter_pattern='hris_*', # Include all HRIS tools + filter_pattern='stackone_*', # Include all StackOne tools account_id=os.getenv('STACKONE_ACCOUNT_ID'), api_key=os.getenv('STACKONE_API_KEY'), ) # Or specify exact tools specific_toolset = StackOneToolset( - tools=['hris_list_employees', 'hris_get_employee'], # Specific tools only + tools=['stackone_list_employees', 'stackone_get_employee'], # Specific tools only account_id=os.getenv('STACKONE_ACCOUNT_ID'), api_key=os.getenv('STACKONE_API_KEY'), ) From 37393814ab23b79ea44932e5860909ea31d78f92 Mon Sep 17 00:00:00 2001 From: ryoppippi <1560508+ryoppippi@users.noreply.github.com> Date: Fri, 30 Jan 2026 15:52:41 +0000 Subject: [PATCH 09/25] fix(stackone): use actual provider names (bamboohr_*) instead of hris_*/stackone_* Per PR review feedback, update all examples and tests to use actual StackOne provider naming patterns like 'bamboohr_*' and 'workday_*' instead of fictional 'hris_*' or 'stackone_*' patterns. - Update docs/toolsets.md examples - Update docs/third-party-tools.md examples - Update examples/stackone_integration.py - Update tests/test_ext_stackone.py - Update docstrings in pydantic_ai_slim/pydantic_ai/ext/stackone.py --- docs/third-party-tools.md | 8 ++-- docs/toolsets.md | 4 +- examples/stackone_integration.py | 16 +++---- pydantic_ai_slim/pydantic_ai/ext/stackone.py | 6 +-- tests/test_ext_stackone.py | 48 ++++++++++---------- 5 files changed, 41 insertions(+), 41 deletions(-) diff --git a/docs/third-party-tools.md b/docs/third-party-tools.md index 257840c074..827d716af7 100644 --- a/docs/third-party-tools.md +++ b/docs/third-party-tools.md @@ -115,7 +115,7 @@ from pydantic_ai import Agent from pydantic_ai.ext.stackone import tool_from_stackone employee_tool = tool_from_stackone( - 'stackone_list_employees', + 'bamboohr_list_employees', account_id=os.getenv('STACKONE_ACCOUNT_ID'), api_key=os.getenv('STACKONE_API_KEY'), ) @@ -145,7 +145,7 @@ from pydantic_ai import Agent from pydantic_ai.ext.stackone import StackOneToolset toolset = StackOneToolset( - filter_pattern='stackone_*', # Load StackOne tools + filter_pattern='bamboohr_*', # Load StackOne tools include_utility_tools=True, # Enable dynamic discovery account_id=os.getenv('STACKONE_ACCOUNT_ID'), api_key=os.getenv('STACKONE_API_KEY'), @@ -165,7 +165,7 @@ from pydantic_ai import Agent from pydantic_ai.ext.stackone import execute_tool, search_tool stackone = StackOneToolSet() -tools = stackone.fetch_tools(actions=['stackone_*']) +tools = stackone.fetch_tools(actions=['bamboohr_*']) agent = Agent( 'openai:gpt-5', @@ -184,7 +184,7 @@ from pydantic_ai import Agent from pydantic_ai.ext.stackone import StackOneToolset toolset = StackOneToolset( - filter_pattern='stackone_*', + filter_pattern='bamboohr_*', include_feedback_tool=True, account_id=os.getenv('STACKONE_ACCOUNT_ID'), api_key=os.getenv('STACKONE_API_KEY'), diff --git a/docs/toolsets.md b/docs/toolsets.md index ef65a23f9d..1ccd7a2135 100644 --- a/docs/toolsets.md +++ b/docs/toolsets.md @@ -748,14 +748,14 @@ from pydantic_ai.ext.stackone import StackOneToolset # Use filter patterns to select specific tools toolset = StackOneToolset( - filter_pattern='stackone_*', # Include all StackOne tools + filter_pattern='bamboohr_*', # Include all StackOne tools account_id=os.getenv('STACKONE_ACCOUNT_ID'), api_key=os.getenv('STACKONE_API_KEY'), ) # Or specify exact tools specific_toolset = StackOneToolset( - tools=['stackone_list_employees', 'stackone_get_employee'], # Specific tools only + tools=['bamboohr_list_employees', 'bamboohr_get_employee'], # Specific tools only account_id=os.getenv('STACKONE_ACCOUNT_ID'), api_key=os.getenv('STACKONE_API_KEY'), ) diff --git a/examples/stackone_integration.py b/examples/stackone_integration.py index bdd415191a..9f346af5d4 100644 --- a/examples/stackone_integration.py +++ b/examples/stackone_integration.py @@ -15,7 +15,7 @@ def single_tool_example(): """Example using a single StackOne tool.""" employee_tool = tool_from_stackone( - 'stackone_list_employees', + 'bamboohr_list_employees', account_id=os.getenv('STACKONE_ACCOUNT_ID'), api_key=os.getenv('STACKONE_API_KEY'), ) @@ -28,7 +28,7 @@ def single_tool_example(): def toolset_with_filter_example(): """Example using StackOne toolset with filter pattern.""" toolset = StackOneToolset( - filter_pattern='stackone_*', # Get all StackOne tools + filter_pattern='bamboohr_*', # Get all StackOne tools account_id=os.getenv('STACKONE_ACCOUNT_ID'), api_key=os.getenv('STACKONE_API_KEY'), ) @@ -42,8 +42,8 @@ def toolset_with_specific_tools_example(): """Example using StackOne toolset with specific tools.""" toolset = StackOneToolset( tools=[ - 'stackone_list_employees', - 'stackone_get_employee', + 'bamboohr_list_employees', + 'bamboohr_get_employee', ], # Specific tools only account_id=os.getenv('STACKONE_ACCOUNT_ID'), api_key=os.getenv('STACKONE_API_KEY'), @@ -65,7 +65,7 @@ def utility_tools_example(): to dynamically discover and use the right ones. """ toolset = StackOneToolset( - filter_pattern='stackone_*', # Load all StackOne tools + filter_pattern='bamboohr_*', # Load all StackOne tools include_utility_tools=True, # Enable dynamic discovery account_id=os.getenv('STACKONE_ACCOUNT_ID'), api_key=os.getenv('STACKONE_API_KEY'), @@ -85,7 +85,7 @@ def standalone_utility_tools_example(): # Fetch tools from StackOne stackone = StackOneToolSet() - tools = stackone.fetch_tools(actions=['stackone_*']) + tools = stackone.fetch_tools(actions=['bamboohr_*']) # Create search and execute tools agent = Agent( @@ -104,7 +104,7 @@ def feedback_example(): with StackOne tools, which helps improve the tool ecosystem. """ toolset = StackOneToolset( - filter_pattern='stackone_*', + filter_pattern='bamboohr_*', include_feedback_tool=True, account_id=os.getenv('STACKONE_ACCOUNT_ID'), api_key=os.getenv('STACKONE_API_KEY'), @@ -139,7 +139,7 @@ def full_featured_example(): - Custom hybrid_alpha for search tuning """ toolset = StackOneToolset( - filter_pattern=['stackone_*'], # Filter pattern + filter_pattern=['bamboohr_*'], # Filter pattern include_utility_tools=True, include_feedback_tool=True, hybrid_alpha=0.3, # Custom search tuning diff --git a/pydantic_ai_slim/pydantic_ai/ext/stackone.py b/pydantic_ai_slim/pydantic_ai/ext/stackone.py index 8cbedc719b..ec6fb233b4 100644 --- a/pydantic_ai_slim/pydantic_ai/ext/stackone.py +++ b/pydantic_ai_slim/pydantic_ai/ext/stackone.py @@ -50,7 +50,7 @@ def tool_from_stackone( """Creates a Pydantic AI tool proxy from a StackOne tool. Args: - tool_name: The name of the StackOne tool to wrap (e.g., "stackone_list_employees"). + 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: Custom base URL for StackOne API. Optional. @@ -204,11 +204,11 @@ def __init__( """Creates a StackOne toolset. Args: - tools: Specific tool names to include (e.g., ["stackone_list_employees"]). + tools: Specific tool names to include (e.g., ["bamboohr_list_employees"]). account_id: The StackOne account ID. Uses STACKONE_ACCOUNT_ID env var if not provided. api_key: The StackOne API key. Uses STACKONE_API_KEY env var if not provided. base_url: Custom base URL for StackOne API. - filter_pattern: Glob pattern(s) to filter tools (e.g., "stackone_*"). + filter_pattern: Glob pattern(s) to filter tools (e.g., "bamboohr_*"). include_utility_tools: If True, includes search and execute utility tools instead of individual tools. Default is False. include_feedback_tool: If True, includes the feedback collection tool. diff --git a/tests/test_ext_stackone.py b/tests/test_ext_stackone.py index 8a0412795d..79af845d4e 100644 --- a/tests/test_ext_stackone.py +++ b/tests/test_ext_stackone.py @@ -48,13 +48,13 @@ def test_tool_creation(self, mock_stackone_toolset_class: MagicMock) -> None: # Mock the StackOne tool mock_tool = Mock() - mock_tool.name = 'hris_list_employees' + mock_tool.name = 'bamboohr_list_employees' mock_tool.description = 'List all employees' mock_tool.execute.return_value = {'employees': []} mock_tool.to_openai_function.return_value = { 'type': 'function', 'function': { - 'name': 'hris_list_employees', + 'name': 'bamboohr_list_employees', 'description': 'List all employees', 'parameters': { 'type': 'object', @@ -71,19 +71,19 @@ def test_tool_creation(self, mock_stackone_toolset_class: MagicMock) -> None: mock_stackone_toolset_class.return_value = mock_stackone_toolset # Create the tool - tool = tool_from_stackone('hris_list_employees', account_id='test-account', api_key='test-key') + tool = tool_from_stackone('bamboohr_list_employees', account_id='test-account', api_key='test-key') # Verify tool creation assert isinstance(tool, Tool) - assert tool.name == 'hris_list_employees' + assert tool.name == 'bamboohr_list_employees' assert tool.description == 'List all employees' # Verify StackOneToolSet was called with correct parameters mock_stackone_toolset_class.assert_called_once_with(api_key='test-key', account_id='test-account') # Verify fetch_tools was called with actions parameter - mock_stackone_toolset.fetch_tools.assert_called_once_with(actions=['hris_list_employees']) - mock_tools.get_tool.assert_called_once_with('hris_list_employees') + mock_stackone_toolset.fetch_tools.assert_called_once_with(actions=['bamboohr_list_employees']) + mock_tools.get_tool.assert_called_once_with('bamboohr_list_employees') # Verify returned Tool has correct JSON schema based on StackOne definition expected = mock_tool.to_openai_function()['function']['parameters'] assert tool.function_schema.json_schema == expected @@ -115,12 +115,12 @@ def test_tool_with_base_url(self, mock_stackone_toolset_class: MagicMock) -> Non # Mock the StackOne tool mock_tool = Mock() - mock_tool.name = 'hris_list_employees' + mock_tool.name = 'bamboohr_list_employees' mock_tool.description = 'List all employees' mock_tool.to_openai_function.return_value = { 'type': 'function', 'function': { - 'name': 'hris_list_employees', + 'name': 'bamboohr_list_employees', 'description': 'List all employees', 'parameters': {'type': 'object', 'properties': {}}, }, @@ -134,7 +134,7 @@ def test_tool_with_base_url(self, mock_stackone_toolset_class: MagicMock) -> Non mock_stackone_toolset_class.return_value = mock_stackone_toolset # Create tool with base URL and verify json_schema conversion - tool = tool_from_stackone('hris_list_employees', api_key='test-key', base_url='https://custom.api.stackone.com') + tool = tool_from_stackone('bamboohr_list_employees', api_key='test-key', base_url='https://custom.api.stackone.com') # Verify base URL was passed to StackOneToolSet mock_stackone_toolset_class.assert_called_once_with( api_key='test-key', account_id=None, base_url='https://custom.api.stackone.com' @@ -215,8 +215,8 @@ def test_toolset_with_specific_tools(self, mock_stackone_toolset_class: MagicMoc from pydantic_ai.ext.stackone import StackOneToolset # Mock the tools returned by fetch_tools - mock_tool1 = _create_mock_stackone_tool('hris_list_employees') - mock_tool2 = _create_mock_stackone_tool('hris_get_employee') + mock_tool1 = _create_mock_stackone_tool('bamboohr_list_employees') + mock_tool2 = _create_mock_stackone_tool('bamboohr_get_employee') mock_fetched_tools = Mock() mock_fetched_tools.__iter__ = Mock(return_value=iter([mock_tool1, mock_tool2])) @@ -227,7 +227,7 @@ def test_toolset_with_specific_tools(self, mock_stackone_toolset_class: MagicMoc # Create the toolset with specific tools toolset = StackOneToolset( - tools=['hris_list_employees', 'hris_get_employee'], account_id='test-account', api_key='test-key' + tools=['bamboohr_list_employees', 'bamboohr_get_employee'], account_id='test-account', api_key='test-key' ) # Verify it's a FunctionToolset @@ -237,7 +237,7 @@ def test_toolset_with_specific_tools(self, mock_stackone_toolset_class: MagicMoc mock_stackone_toolset_class.assert_called_once_with(api_key='test-key', account_id='test-account') # Verify fetch_tools was called with the tool names as actions - mock_stackone_toolset.fetch_tools.assert_called_once_with(actions=['hris_list_employees', 'hris_get_employee']) + mock_stackone_toolset.fetch_tools.assert_called_once_with(actions=['bamboohr_list_employees', 'bamboohr_get_employee']) @patch('pydantic_ai.ext.stackone.StackOneToolSet') def test_toolset_with_filter_pattern(self, mock_stackone_toolset_class: MagicMock) -> None: @@ -245,7 +245,7 @@ def test_toolset_with_filter_pattern(self, mock_stackone_toolset_class: MagicMoc from pydantic_ai.ext.stackone import StackOneToolset # Mock the tools returned by fetch_tools - mock_tool = _create_mock_stackone_tool('hris_list_employees') + mock_tool = _create_mock_stackone_tool('bamboohr_list_employees') mock_fetched_tools = Mock() mock_fetched_tools.__iter__ = Mock(return_value=iter([mock_tool])) @@ -255,13 +255,13 @@ def test_toolset_with_filter_pattern(self, mock_stackone_toolset_class: MagicMoc mock_stackone_toolset_class.return_value = mock_stackone_toolset # Create toolset with filter_pattern - toolset = StackOneToolset(filter_pattern='hris_*', account_id='test-account', api_key='test-key') + toolset = StackOneToolset(filter_pattern='bamboohr_*', account_id='test-account', api_key='test-key') # Verify StackOneToolSet was created correctly mock_stackone_toolset_class.assert_called_once_with(api_key='test-key', account_id='test-account') # Verify fetch_tools was called with actions parameter (list) - mock_stackone_toolset.fetch_tools.assert_called_once_with(actions=['hris_*']) + mock_stackone_toolset.fetch_tools.assert_called_once_with(actions=['bamboohr_*']) # Verify tools were created assert isinstance(toolset, FunctionToolset) @@ -272,8 +272,8 @@ def test_toolset_with_list_filter_pattern(self, mock_stackone_toolset_class: Mag from pydantic_ai.ext.stackone import StackOneToolset # Mock the tools returned by fetch_tools - mock_tool1 = _create_mock_stackone_tool('hris_list_employees') - mock_tool2 = _create_mock_stackone_tool('ats_get_job') + mock_tool1 = _create_mock_stackone_tool('bamboohr_list_employees') + mock_tool2 = _create_mock_stackone_tool('workday_list_employees') mock_fetched_tools = Mock() mock_fetched_tools.__iter__ = Mock(return_value=iter([mock_tool1, mock_tool2])) @@ -283,10 +283,10 @@ def test_toolset_with_list_filter_pattern(self, mock_stackone_toolset_class: Mag mock_stackone_toolset_class.return_value = mock_stackone_toolset # Create toolset with list filter_pattern - toolset = StackOneToolset(filter_pattern=['hris_*', 'ats_*'], account_id='test-account', api_key='test-key') + toolset = StackOneToolset(filter_pattern=['bamboohr_*', 'workday_*'], account_id='test-account', api_key='test-key') # Verify fetch_tools was called with list filter_pattern as actions - mock_stackone_toolset.fetch_tools.assert_called_once_with(actions=['hris_*', 'ats_*']) + mock_stackone_toolset.fetch_tools.assert_called_once_with(actions=['bamboohr_*', 'workday_*']) # Verify it's a FunctionToolset assert isinstance(toolset, FunctionToolset) @@ -324,7 +324,7 @@ def test_toolset_with_base_url(self, mock_stackone_toolset_class: MagicMock) -> from pydantic_ai.ext.stackone import StackOneToolset # Mock the tools returned by fetch_tools - mock_tool = _create_mock_stackone_tool('hris_list_employees') + mock_tool = _create_mock_stackone_tool('bamboohr_list_employees') mock_fetched_tools = Mock() mock_fetched_tools.__iter__ = Mock(return_value=iter([mock_tool])) @@ -335,7 +335,7 @@ def test_toolset_with_base_url(self, mock_stackone_toolset_class: MagicMock) -> # Create toolset with base URL toolset = StackOneToolset( - tools=['hris_list_employees'], + tools=['bamboohr_list_employees'], account_id='test-account', api_key='test-key', base_url='https://custom.api.stackone.com', @@ -355,7 +355,7 @@ def test_toolset_with_utility_tools(self, mock_stackone_toolset_class: MagicMock from pydantic_ai.ext.stackone import StackOneToolset # Mock the tools returned by fetch_tools - mock_tool = _create_mock_stackone_tool('hris_list_employees') + mock_tool = _create_mock_stackone_tool('bamboohr_list_employees') mock_fetched_tools = Mock() mock_fetched_tools.__iter__ = Mock(return_value=iter([mock_tool])) @@ -377,7 +377,7 @@ def test_toolset_with_utility_tools(self, mock_stackone_toolset_class: MagicMock # Create toolset with utility tools toolset = StackOneToolset( - filter_pattern='hris_*', + filter_pattern='bamboohr_*', include_utility_tools=True, account_id='test-account', api_key='test-key', From fd1225b0495f52d3296eeafb0daeda1c4c3727a9 Mon Sep 17 00:00:00 2001 From: ryoppippi <1560508+ryoppippi@users.noreply.github.com> Date: Fri, 30 Jan 2026 16:05:20 +0000 Subject: [PATCH 10/25] fix(stackone): fix CI lint and test failures - Format test file with ruff to fix line length issues - Add {test="skip"} to StackOneToolset docstring example to skip example test that requires stackone-ai package --- pydantic_ai_slim/pydantic_ai/ext/stackone.py | 2 +- tests/test_ext_stackone.py | 12 +++++++++--- 2 files changed, 10 insertions(+), 4 deletions(-) diff --git a/pydantic_ai_slim/pydantic_ai/ext/stackone.py b/pydantic_ai_slim/pydantic_ai/ext/stackone.py index ec6fb233b4..601738a1cb 100644 --- a/pydantic_ai_slim/pydantic_ai/ext/stackone.py +++ b/pydantic_ai_slim/pydantic_ai/ext/stackone.py @@ -179,7 +179,7 @@ class StackOneToolset(FunctionToolset): where the agent needs to search for the right tool. Example: - ```python + ```python {test="skip"} from pydantic_ai import Agent from pydantic_ai.ext.stackone import StackOneToolset diff --git a/tests/test_ext_stackone.py b/tests/test_ext_stackone.py index 79af845d4e..6fffac35b6 100644 --- a/tests/test_ext_stackone.py +++ b/tests/test_ext_stackone.py @@ -134,7 +134,9 @@ def test_tool_with_base_url(self, mock_stackone_toolset_class: MagicMock) -> Non mock_stackone_toolset_class.return_value = mock_stackone_toolset # Create tool with base URL and verify json_schema conversion - tool = tool_from_stackone('bamboohr_list_employees', api_key='test-key', base_url='https://custom.api.stackone.com') + tool = tool_from_stackone( + 'bamboohr_list_employees', api_key='test-key', base_url='https://custom.api.stackone.com' + ) # Verify base URL was passed to StackOneToolSet mock_stackone_toolset_class.assert_called_once_with( api_key='test-key', account_id=None, base_url='https://custom.api.stackone.com' @@ -237,7 +239,9 @@ def test_toolset_with_specific_tools(self, mock_stackone_toolset_class: MagicMoc mock_stackone_toolset_class.assert_called_once_with(api_key='test-key', account_id='test-account') # Verify fetch_tools was called with the tool names as actions - mock_stackone_toolset.fetch_tools.assert_called_once_with(actions=['bamboohr_list_employees', 'bamboohr_get_employee']) + mock_stackone_toolset.fetch_tools.assert_called_once_with( + actions=['bamboohr_list_employees', 'bamboohr_get_employee'] + ) @patch('pydantic_ai.ext.stackone.StackOneToolSet') def test_toolset_with_filter_pattern(self, mock_stackone_toolset_class: MagicMock) -> None: @@ -283,7 +287,9 @@ def test_toolset_with_list_filter_pattern(self, mock_stackone_toolset_class: Mag mock_stackone_toolset_class.return_value = mock_stackone_toolset # Create toolset with list filter_pattern - toolset = StackOneToolset(filter_pattern=['bamboohr_*', 'workday_*'], account_id='test-account', api_key='test-key') + toolset = StackOneToolset( + filter_pattern=['bamboohr_*', 'workday_*'], account_id='test-account', api_key='test-key' + ) # Verify fetch_tools was called with list filter_pattern as actions mock_stackone_toolset.fetch_tools.assert_called_once_with(actions=['bamboohr_*', 'workday_*']) From e95b7e3aa0da1e8dd978eaf2f25f7371a304a9aa Mon Sep 17 00:00:00 2001 From: ryoppippi <1560508+ryoppippi@users.noreply.github.com> Date: Fri, 30 Jan 2026 16:19:18 +0000 Subject: [PATCH 11/25] fix(stackone): exclude stackone files from coverage Add stackone.py and test_ext_stackone.py to coverage omit list, similar to aci.py and exa.py which are also external integrations. --- pyproject.toml | 2 ++ 1 file changed, 2 insertions(+) diff --git a/pyproject.toml b/pyproject.toml index e82315ca01..c31c6ba5d8 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -266,7 +266,9 @@ include = [ ] omit = [ "tests/example_modules/*.py", + "tests/test_ext_stackone.py", # stackone tests are mocked, not real coverage "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 "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", From c6700e711f3427b6838875cbbb1e0893931d5518 Mon Sep 17 00:00:00 2001 From: Shashi-Stackone Date: Mon, 23 Feb 2026 17:52:36 +0000 Subject: [PATCH 12/25] Apply further PR suggestions and fix CI --- docs/api/ext.md | 2 + docs/toolsets.md | 2 +- examples/stackone_integration.py | 166 -------- pydantic_ai_slim/pydantic_ai/ext/stackone.py | 136 ++----- pyproject.toml | 1 - tests/ext/test_stackone.py | 294 ++++++++++++++ tests/test_ext_stackone.py | 396 ------------------- 7 files changed, 324 insertions(+), 673 deletions(-) delete mode 100644 examples/stackone_integration.py create mode 100644 tests/ext/test_stackone.py delete mode 100644 tests/test_ext_stackone.py 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/toolsets.md b/docs/toolsets.md index 1ccd7a2135..e52d912f4a 100644 --- a/docs/toolsets.md +++ b/docs/toolsets.md @@ -748,7 +748,7 @@ from pydantic_ai.ext.stackone import StackOneToolset # Use filter patterns to select specific tools toolset = StackOneToolset( - filter_pattern='bamboohr_*', # Include all StackOne tools + filter_pattern='bamboohr_*', account_id=os.getenv('STACKONE_ACCOUNT_ID'), api_key=os.getenv('STACKONE_API_KEY'), ) diff --git a/examples/stackone_integration.py b/examples/stackone_integration.py deleted file mode 100644 index 9f346af5d4..0000000000 --- a/examples/stackone_integration.py +++ /dev/null @@ -1,166 +0,0 @@ -"""Example of integrating StackOne tools with Pydantic AI.""" - -import os - -from pydantic_ai import Agent -from pydantic_ai.ext.stackone import ( - StackOneToolset, - execute_tool, - feedback_tool, - search_tool, - tool_from_stackone, -) - - -def single_tool_example(): - """Example using a single StackOne tool.""" - employee_tool = tool_from_stackone( - 'bamboohr_list_employees', - account_id=os.getenv('STACKONE_ACCOUNT_ID'), - api_key=os.getenv('STACKONE_API_KEY'), - ) - - agent = Agent('openai:gpt-5', tools=[employee_tool]) - result = agent.run_sync('List all employees') - print(result.output) - - -def toolset_with_filter_example(): - """Example using StackOne toolset with filter pattern.""" - toolset = StackOneToolset( - filter_pattern='bamboohr_*', # Get all StackOne tools - account_id=os.getenv('STACKONE_ACCOUNT_ID'), - api_key=os.getenv('STACKONE_API_KEY'), - ) - - agent = Agent('openai:gpt-5', toolsets=[toolset]) - result = agent.run_sync('Get employee information') - print(result.output) - - -def toolset_with_specific_tools_example(): - """Example using StackOne toolset with specific tools.""" - 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('openai:gpt-5', toolsets=[toolset]) - result = agent.run_sync('Get information about all employees') - print(result.output) - - -def utility_tools_example(): - """Example using StackOne with utility tools for dynamic discovery. - - Utility tools mode provides two special tools: - - tool_search: Search for relevant tools using natural language - - tool_execute: Execute a discovered tool by name - - This is useful when you have a large number of tools and want the agent - to dynamically discover and use the right ones. - """ - toolset = StackOneToolset( - filter_pattern='bamboohr_*', # Load all StackOne tools - include_utility_tools=True, # Enable dynamic discovery - account_id=os.getenv('STACKONE_ACCOUNT_ID'), - api_key=os.getenv('STACKONE_API_KEY'), - ) - - agent = Agent('openai:gpt-5', toolsets=[toolset]) - result = agent.run_sync('Find a tool to list employees and use it') - print(result.output) - - -def standalone_utility_tools_example(): - """Example using standalone search_tool and execute_tool functions. - - This gives you more control over how the tools are created and used. - """ - from stackone_ai import StackOneToolSet - - # Fetch tools from StackOne - stackone = StackOneToolSet() - tools = stackone.fetch_tools(actions=['bamboohr_*']) - - # Create search and execute tools - agent = Agent( - 'openai:gpt-5', - tools=[search_tool(tools), execute_tool(tools)], - ) - - result = agent.run_sync('Search for employee-related tools') - print(result.output) - - -def feedback_example(): - """Example using the feedback collection tool via StackOneToolset. - - The feedback tool allows users to provide feedback on their experience - with StackOne tools, which helps improve the tool ecosystem. - """ - toolset = StackOneToolset( - filter_pattern='bamboohr_*', - include_feedback_tool=True, - account_id=os.getenv('STACKONE_ACCOUNT_ID'), - api_key=os.getenv('STACKONE_API_KEY'), - ) - - agent = Agent('openai:gpt-5', toolsets=[toolset]) - result = agent.run_sync('List employees and then ask for feedback') - print(result.output) - - -def standalone_feedback_example(): - """Example using the standalone feedback_tool function. - - This gives you more control over when and how to include the feedback tool. - """ - fb_tool = feedback_tool( - api_key=os.getenv('STACKONE_API_KEY'), - account_id=os.getenv('STACKONE_ACCOUNT_ID'), - ) - - agent = Agent('openai:gpt-5', tools=[fb_tool]) - result = agent.run_sync('Collect feedback about a recent tool usage') - print(result.output) - - -def full_featured_example(): - """Example using all StackOne features together. - - This example demonstrates: - - Utility tools for dynamic discovery - - Feedback collection - - Custom hybrid_alpha for search tuning - """ - toolset = StackOneToolset( - filter_pattern=['bamboohr_*'], # Filter pattern - include_utility_tools=True, - include_feedback_tool=True, - hybrid_alpha=0.3, # Custom search tuning - account_id=os.getenv('STACKONE_ACCOUNT_ID'), - api_key=os.getenv('STACKONE_API_KEY'), - ) - - agent = Agent('openai:gpt-5', toolsets=[toolset]) - result = agent.run_sync( - 'Find the best tool to get employee information, use it, ' - 'and then collect feedback about the experience' - ) - print(result.output) - - -if __name__ == '__main__': - single_tool_example() - toolset_with_filter_example() - toolset_with_specific_tools_example() - utility_tools_example() - standalone_utility_tools_example() - feedback_example() - standalone_feedback_example() - full_featured_example() diff --git a/pydantic_ai_slim/pydantic_ai/ext/stackone.py b/pydantic_ai_slim/pydantic_ai/ext/stackone.py index 601738a1cb..b40d62a67f 100644 --- a/pydantic_ai_slim/pydantic_ai/ext/stackone.py +++ b/pydantic_ai_slim/pydantic_ai/ext/stackone.py @@ -1,5 +1,6 @@ from __future__ import annotations +import os from collections.abc import Sequence from typing import Any @@ -13,19 +14,10 @@ except ImportError as _import_error: raise ImportError('Please install `stackone-ai` to use StackOne tools.') from _import_error -# StackOneTools is dynamically typed as stackone_ai doesn't provide complete type stubs -StackOneTools = Any +__all__ = ('tool_from_stackone', 'StackOneToolset', 'search_tool', 'execute_tool', 'feedback_tool') def _tool_from_stackone_tool(stackone_tool: Any) -> Tool: - """Creates a Pydantic AI tool from a StackOneTool instance. - - Args: - stackone_tool: A StackOneTool instance. - - Returns: - A Pydantic AI tool that wraps the StackOne tool. - """ openai_function = stackone_tool.to_openai_function() json_schema: JsonSchemaValue = openai_function['function']['parameters'] @@ -50,42 +42,32 @@ def tool_from_stackone( """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: Custom base URL for StackOne API. Optional. + 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} if base_url else {}), - ) - + 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) def search_tool( - tools: StackOneTools, + tools: Any, *, hybrid_alpha: float | None = None, ) -> Tool: - """Creates a search tool for discovering StackOne tools using natural language. - - This tool uses hybrid BM25 + TF-IDF search to find relevant tools based on - a natural language query. + """Creates a tool for discovering StackOne tools using natural language queries. Args: - tools: A StackOne Tools collection (returned from `fetch_tools()`). - hybrid_alpha: Weight for BM25 in hybrid search (0-1). Default is 0.2, - which has been shown to provide better tool discovery accuracy. + tools: A StackOne `Tools` collection returned from `StackOneToolSet.fetch_tools()`. + hybrid_alpha: Optional weight for hybrid search (0-1). Defaults to the SDK default. Returns: A Pydantic AI tool for searching StackOne tools. @@ -94,20 +76,16 @@ def search_tool( search = utility_tools.get_tool('tool_search') if search is None: raise ValueError('tool_search not found in StackOne utility tools') - return _tool_from_stackone_tool(search) def execute_tool( - tools: StackOneTools, + tools: Any, ) -> Tool: - """Creates an execute tool for running discovered StackOne tools by name. - - This tool allows executing any tool from the provided collection by name, - typically used after discovering tools with `search_tool`. + """Creates a tool for executing discovered StackOne tools by name. Args: - tools: A StackOne Tools collection (returned from `fetch_tools()`). + tools: A StackOne `Tools` collection returned from `StackOneToolSet.fetch_tools()`. Returns: A Pydantic AI tool for executing StackOne tools. @@ -116,7 +94,6 @@ def execute_tool( execute = utility_tools.get_tool('tool_execute') if execute is None: raise ValueError('tool_execute not found in StackOne utility tools') - return _tool_from_stackone_tool(execute) @@ -126,15 +103,12 @@ def feedback_tool( account_id: str | None = None, base_url: str | None = None, ) -> Tool: - """Creates a feedback tool for collecting user feedback on StackOne tools. - - This tool allows users to provide feedback on their experience with StackOne tools, - which helps improve the tool ecosystem. + """Creates a feedback tool for collecting user feedback on StackOne tool usage. Args: - api_key: The StackOne API key. If not provided, uses STACKONE_API_KEY env var. - account_id: The StackOne account ID. If not provided, uses STACKONE_ACCOUNT_ID env var. - base_url: Custom base URL for StackOne API. Optional. + api_key: The StackOne API key. If not provided, uses `STACKONE_API_KEY` env var. + account_id: The StackOne account ID. If not provided, uses `STACKONE_ACCOUNT_ID` env var. + base_url: Optional custom base URL for the StackOne API. Returns: A Pydantic AI tool for collecting feedback. @@ -144,49 +118,22 @@ def feedback_tool( except ImportError as e: raise ImportError('Please install `stackone-ai` with feedback support to use the feedback tool.') from e - # Get API key from environment if not provided - if api_key is None: - import os - - api_key = os.environ.get('STACKONE_API_KEY') - if api_key is None: - raise ValueError( - 'API key is required. Provide it as an argument or set STACKONE_API_KEY environment variable.' - ) + resolved_api_key = api_key or os.environ.get('STACKONE_API_KEY') + if resolved_api_key is None: + raise ValueError( + 'API key is required. Provide it as an argument or set the STACKONE_API_KEY environment variable.' + ) fb_tool = create_feedback_tool( - api_key=api_key, + api_key=resolved_api_key, account_id=account_id, **({'base_url': base_url} if base_url else {}), ) - return _tool_from_stackone_tool(fb_tool) class StackOneToolset(FunctionToolset): - """A toolset that wraps StackOne tools. - - This toolset provides access to StackOne's integration infrastructure for AI agents, - offering 200+ connectors across HR, ATS, CRM, and other business applications. - It can operate in two modes: - - 1. **Direct mode** (default): Each StackOne tool is exposed directly to the agent. - This is best when you have a small, known set of tools. - - 2. **Utility tools mode** (`include_utility_tools=True`): Instead of exposing all tools - directly, provides `tool_search` and `tool_execute` that allow the agent to - dynamically discover and execute tools. This is better for large tool sets - where the agent needs to search for the right tool. - - Example: - ```python {test="skip"} - from pydantic_ai import Agent - from pydantic_ai.ext.stackone import StackOneToolset - - toolset = StackOneToolset(api_key='your-api-key') - agent = Agent('openai:gpt-4o', toolsets=[toolset]) - ``` - """ + """A toolset that wraps StackOne tools.""" def __init__( self, @@ -201,27 +148,7 @@ def __init__( hybrid_alpha: float | None = None, id: str | None = None, ): - """Creates a StackOne toolset. - - Args: - tools: Specific tool names to include (e.g., ["bamboohr_list_employees"]). - account_id: The StackOne account ID. Uses STACKONE_ACCOUNT_ID env var if not provided. - api_key: The StackOne API key. Uses STACKONE_API_KEY env var if not provided. - base_url: Custom base URL for StackOne API. - filter_pattern: Glob pattern(s) to filter tools (e.g., "bamboohr_*"). - include_utility_tools: If True, includes search and execute utility tools instead of - individual tools. Default is False. - include_feedback_tool: If True, includes the feedback collection tool. - Default is False. - hybrid_alpha: Weight for BM25 in hybrid search (0-1) when using utility tools. - Default is 0.2. - id: Optional ID for the toolset, used for durable execution environments. - """ - stackone_toolset = StackOneToolSet( - api_key=api_key, - account_id=account_id, - **({'base_url': base_url} if base_url else {}), - ) + stackone_toolset = StackOneToolSet(api_key=api_key, account_id=account_id, base_url=base_url) if tools is not None: actions = list(tools) @@ -231,23 +158,14 @@ def __init__( fetched_tools = stackone_toolset.fetch_tools(actions=actions) pydantic_tools: list[Tool] = [] - if include_utility_tools: - # Utility tools mode: provide search and execute tools pydantic_tools.append(search_tool(fetched_tools, hybrid_alpha=hybrid_alpha)) pydantic_tools.append(execute_tool(fetched_tools)) else: - # Direct mode: expose each tool individually for stackone_tool in fetched_tools: pydantic_tools.append(_tool_from_stackone_tool(stackone_tool)) if include_feedback_tool: - pydantic_tools.append( - feedback_tool( - api_key=api_key, - account_id=account_id, - base_url=base_url, - ) - ) + pydantic_tools.append(feedback_tool(api_key=api_key, account_id=account_id, base_url=base_url)) super().__init__(pydantic_tools, id=id) diff --git a/pyproject.toml b/pyproject.toml index c31c6ba5d8..4b000ca810 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -266,7 +266,6 @@ include = [ ] omit = [ "tests/example_modules/*.py", - "tests/test_ext_stackone.py", # stackone tests are mocked, not real coverage "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 "pydantic_ai_slim/pydantic_ai/common_tools/exa.py", # exa-py integration with external API calls diff --git a/tests/ext/test_stackone.py b/tests/ext/test_stackone.py new file mode 100644 index 0000000000..4d62f6886c --- /dev/null +++ b/tests/ext/test_stackone.py @@ -0,0 +1,294 @@ +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.ext.stackone import ( + StackOneToolset, + _tool_from_stackone_tool, + execute_tool, + feedback_tool, + search_tool, + tool_from_stackone, +) +from pydantic_ai.tools import Tool + + +@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 = _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 = _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 = _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 = _tool_from_stackone_tool(tool_with_none_desc) + assert tool.description == '' + + +def test_tool_schema(): + tool = _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') + + +def test_search_tool(): + mock_search = SimulatedStackOneTool(name='tool_search', description='Search for tools') + mock_utility = Mock() + mock_utility.get_tool.return_value = mock_search + + mock_tools = Mock() + mock_tools.utility_tools.return_value = mock_utility + + tool = search_tool(mock_tools, hybrid_alpha=0.3) + assert tool.name == 'tool_search' + mock_tools.utility_tools.assert_called_once_with(hybrid_alpha=0.3) + + +def test_search_tool_not_found(): + mock_utility = Mock() + mock_utility.get_tool.return_value = None + mock_tools = Mock() + mock_tools.utility_tools.return_value = mock_utility + + with pytest.raises(ValueError, match='tool_search not found'): + search_tool(mock_tools) + + +def test_execute_tool(): + mock_exec = SimulatedStackOneTool(name='tool_execute', description='Execute a tool') + mock_utility = Mock() + mock_utility.get_tool.return_value = mock_exec + + mock_tools = Mock() + mock_tools.utility_tools.return_value = mock_utility + + tool = execute_tool(mock_tools) + assert tool.name == 'tool_execute' + + +def test_execute_tool_not_found(): + mock_utility = Mock() + mock_utility.get_tool.return_value = None + mock_tools = Mock() + mock_tools.utility_tools.return_value = mock_utility + + with pytest.raises(ValueError, match='tool_execute not found'): + execute_tool(mock_tools) + + +@patch('pydantic_ai.ext.stackone.StackOneToolSet') +def test_stackone_toolset_utility_tools_mode(mock_toolset_cls: Any): + mock_search = SimulatedStackOneTool(name='tool_search', description='Search') + mock_exec = SimulatedStackOneTool(name='tool_execute', description='Execute') + + mock_utility = Mock() + + def get_tool_side_effect(name: str) -> SimulatedStackOneTool | None: + if name == 'tool_search': + return mock_search + if name == 'tool_execute': + return mock_exec + return None + + mock_utility.get_tool.side_effect = get_tool_side_effect + + mock_fetched = Mock() + mock_fetched.utility_tools.return_value = mock_utility + mock_toolset_cls.return_value.fetch_tools.return_value = mock_fetched + + toolset = StackOneToolset(filter_pattern='bamboohr_*', include_utility_tools=True, api_key='k') + + agent = Agent('test', toolsets=[toolset]) + result = agent.run_sync('foobar') + assert 'tool_search' in result.output or 'tool_execute' in result.output + + +@patch.dict('os.environ', {'STACKONE_API_KEY': 'env-key'}) +@patch('stackone_ai.feedback.tool.create_feedback_tool') +def test_feedback_tool_from_env(mock_create_fb: Any): + mock_fb = SimulatedStackOneTool(name='tool_feedback', description='Feedback') + mock_create_fb.return_value = mock_fb + + tool = feedback_tool() + assert tool.name == 'tool_feedback' + mock_create_fb.assert_called_once_with(api_key='env-key', account_id=None) + + +@patch('stackone_ai.feedback.tool.create_feedback_tool') +def test_feedback_tool_explicit_key(mock_create_fb: Any): + mock_fb = SimulatedStackOneTool(name='tool_feedback', description='Feedback') + mock_create_fb.return_value = mock_fb + + tool = feedback_tool(api_key='explicit-key', account_id='acc-1') + assert tool.name == 'tool_feedback' + mock_create_fb.assert_called_once_with(api_key='explicit-key', account_id='acc-1') + + +@patch('stackone_ai.feedback.tool.create_feedback_tool') +def test_feedback_tool_with_base_url(mock_create_fb: Any): + mock_fb = SimulatedStackOneTool(name='tool_feedback', description='Feedback') + mock_create_fb.return_value = mock_fb + + feedback_tool(api_key='k', base_url='https://custom.stackone.com') + mock_create_fb.assert_called_once_with(api_key='k', account_id=None, base_url='https://custom.stackone.com') + + +@patch.dict('os.environ', {}, clear=True) +def test_feedback_tool_no_key(): + with pytest.raises(ValueError, match='API key is required'): + feedback_tool() + + +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) diff --git a/tests/test_ext_stackone.py b/tests/test_ext_stackone.py deleted file mode 100644 index 6fffac35b6..0000000000 --- a/tests/test_ext_stackone.py +++ /dev/null @@ -1,396 +0,0 @@ -"""Tests for StackOne integration with Pydantic AI.""" - -from __future__ import annotations - -from typing import TYPE_CHECKING -from unittest.mock import Mock, patch - -import pytest - -from pydantic_ai.tools import Tool -from pydantic_ai.toolsets.function import FunctionToolset - -if TYPE_CHECKING: - from unittest.mock import MagicMock - -try: - import stackone_ai # noqa: F401 # pyright: ignore[reportUnusedImport] -except ImportError: # pragma: lax no cover - stackone_installed = False -else: - stackone_installed = True - - -class TestStackOneImportError: - """Test import error handling.""" - - def test_import_error_without_stackone(self): - """Test that ImportError is raised when stackone-ai is not available.""" - # Test that importing the module raises the expected error when stackone_ai is not available - with patch.dict('sys.modules', {'stackone_ai': None}): - with pytest.raises(ImportError, match='Please install `stackone-ai`'): - # Force reimport by using importlib - import importlib - - import pydantic_ai.ext.stackone - - importlib.reload(pydantic_ai.ext.stackone) - - -@pytest.mark.skipif(not stackone_installed, reason='stackone-ai not installed') -class TestToolFromStackOne: - """Test the tool_from_stackone function.""" - - @patch('pydantic_ai.ext.stackone.StackOneToolSet') - def test_tool_creation(self, mock_stackone_toolset_class: MagicMock) -> None: - """Test creating a single tool from StackOne.""" - from pydantic_ai.ext.stackone import tool_from_stackone - - # Mock the StackOne tool - mock_tool = Mock() - mock_tool.name = 'bamboohr_list_employees' - mock_tool.description = 'List all employees' - mock_tool.execute.return_value = {'employees': []} - mock_tool.to_openai_function.return_value = { - 'type': 'function', - 'function': { - 'name': 'bamboohr_list_employees', - 'description': 'List all employees', - 'parameters': { - 'type': 'object', - 'properties': {'limit': {'type': 'integer', 'description': 'Limit the number of results'}}, - }, - }, - } - - mock_tools = Mock() - mock_tools.get_tool.return_value = mock_tool - - mock_stackone_toolset = Mock() - mock_stackone_toolset.fetch_tools.return_value = mock_tools - mock_stackone_toolset_class.return_value = mock_stackone_toolset - - # Create the tool - tool = tool_from_stackone('bamboohr_list_employees', account_id='test-account', api_key='test-key') - - # Verify tool creation - assert isinstance(tool, Tool) - assert tool.name == 'bamboohr_list_employees' - assert tool.description == 'List all employees' - - # Verify StackOneToolSet was called with correct parameters - mock_stackone_toolset_class.assert_called_once_with(api_key='test-key', account_id='test-account') - - # Verify fetch_tools was called with actions parameter - mock_stackone_toolset.fetch_tools.assert_called_once_with(actions=['bamboohr_list_employees']) - mock_tools.get_tool.assert_called_once_with('bamboohr_list_employees') - # Verify returned Tool has correct JSON schema based on StackOne definition - expected = mock_tool.to_openai_function()['function']['parameters'] - assert tool.function_schema.json_schema == expected - - @patch('pydantic_ai.ext.stackone.StackOneToolSet') - def test_tool_not_found(self, mock_stackone_toolset_class: MagicMock) -> None: - """Test error when tool is not found.""" - from pydantic_ai.ext.stackone import tool_from_stackone - - # Mock the tools to return None for the requested tool - mock_tools = Mock() - mock_tools.get_tool.return_value = None - - mock_stackone_toolset = Mock() - mock_stackone_toolset.fetch_tools.return_value = mock_tools - mock_stackone_toolset_class.return_value = mock_stackone_toolset - - # Should raise ValueError when tool not found - with pytest.raises(ValueError, match="Tool 'unknown_tool' not found in StackOne"): - tool_from_stackone('unknown_tool', api_key='test-key') - - @patch('pydantic_ai.ext.stackone.StackOneToolSet') - def test_tool_with_base_url(self, mock_stackone_toolset_class: MagicMock) -> None: - """Test creating a tool with custom base URL. - - Note: base_url is not commonly used by end users, but this test exists for coverage. - """ - from pydantic_ai.ext.stackone import tool_from_stackone - - # Mock the StackOne tool - mock_tool = Mock() - mock_tool.name = 'bamboohr_list_employees' - mock_tool.description = 'List all employees' - mock_tool.to_openai_function.return_value = { - 'type': 'function', - 'function': { - 'name': 'bamboohr_list_employees', - 'description': 'List all employees', - 'parameters': {'type': 'object', 'properties': {}}, - }, - } - - mock_tools = Mock() - mock_tools.get_tool.return_value = mock_tool - - mock_stackone_toolset = Mock() - mock_stackone_toolset.fetch_tools.return_value = mock_tools - mock_stackone_toolset_class.return_value = mock_stackone_toolset - - # Create tool with base URL and verify json_schema conversion - tool = tool_from_stackone( - 'bamboohr_list_employees', api_key='test-key', base_url='https://custom.api.stackone.com' - ) - # Verify base URL was passed to StackOneToolSet - mock_stackone_toolset_class.assert_called_once_with( - api_key='test-key', account_id=None, base_url='https://custom.api.stackone.com' - ) - # Verify returned Tool has correct schema - expected = mock_tool.to_openai_function()['function']['parameters'] - assert tool.function_schema.json_schema == expected - - @patch('pydantic_ai.ext.stackone.StackOneToolSet') - def test_default_parameters(self, mock_stackone_toolset_class: MagicMock) -> None: - """Test default account_id and base_url are None when not provided.""" - from pydantic_ai.ext.stackone import tool_from_stackone - - mock_tool = Mock() - mock_tool.name = 'foo' - mock_tool.description = 'bar' - mock_tool.to_openai_function.return_value = { - 'type': 'function', - 'function': {'name': 'foo', 'description': 'bar', 'parameters': {}}, - } - mock_tools = Mock() - mock_tools.get_tool.return_value = mock_tool - mock_stackone_toolset = Mock() - mock_stackone_toolset.fetch_tools.return_value = mock_tools - mock_stackone_toolset_class.return_value = mock_stackone_toolset - - tool = tool_from_stackone('foo', api_key='key-only') - mock_stackone_toolset_class.assert_called_once_with(api_key='key-only', account_id=None) - expected = mock_tool.to_openai_function()['function']['parameters'] - assert tool.function_schema.json_schema == expected - - @patch('pydantic_ai.ext.stackone.StackOneToolSet') - def test_tool_with_none_description(self, mock_stackone_toolset_class: MagicMock) -> None: - """Test creating a tool when description is None.""" - from pydantic_ai.ext.stackone import tool_from_stackone - - mock_tool = Mock() - mock_tool.name = 'test_tool' - mock_tool.description = None # None description should become empty string - mock_tool.to_openai_function.return_value = { - 'function': {'name': 'test_tool', 'parameters': {'type': 'object'}}, - } - - mock_tools = Mock() - mock_tools.get_tool.return_value = mock_tool - - mock_stackone_toolset = Mock() - mock_stackone_toolset.fetch_tools.return_value = mock_tools - mock_stackone_toolset_class.return_value = mock_stackone_toolset - - tool = tool_from_stackone('test_tool', api_key='test-key') - assert tool.description == '' - - -def _create_mock_stackone_tool(name: str, description: str = 'Test description') -> Mock: - """Helper to create a mock StackOne tool.""" - mock_tool = Mock() - mock_tool.name = name - mock_tool.description = description - mock_tool.to_openai_function.return_value = { - 'type': 'function', - 'function': { - 'name': name, - 'description': description, - 'parameters': {'type': 'object', 'properties': {}}, - }, - } - return mock_tool - - -@pytest.mark.skipif(not stackone_installed, reason='stackone-ai not installed') -class TestStackOneToolset: - """Test the StackOneToolset class.""" - - @patch('pydantic_ai.ext.stackone.StackOneToolSet') - def test_toolset_with_specific_tools(self, mock_stackone_toolset_class: MagicMock) -> None: - """Test creating a StackOneToolset with specific tools.""" - from pydantic_ai.ext.stackone import StackOneToolset - - # Mock the tools returned by fetch_tools - mock_tool1 = _create_mock_stackone_tool('bamboohr_list_employees') - mock_tool2 = _create_mock_stackone_tool('bamboohr_get_employee') - - mock_fetched_tools = Mock() - mock_fetched_tools.__iter__ = Mock(return_value=iter([mock_tool1, mock_tool2])) - - mock_stackone_toolset = Mock() - mock_stackone_toolset.fetch_tools.return_value = mock_fetched_tools - mock_stackone_toolset_class.return_value = mock_stackone_toolset - - # Create the toolset with specific tools - toolset = StackOneToolset( - tools=['bamboohr_list_employees', 'bamboohr_get_employee'], account_id='test-account', api_key='test-key' - ) - - # Verify it's a FunctionToolset - assert isinstance(toolset, FunctionToolset) - - # Verify StackOneToolSet was created correctly - mock_stackone_toolset_class.assert_called_once_with(api_key='test-key', account_id='test-account') - - # Verify fetch_tools was called with the tool names as actions - mock_stackone_toolset.fetch_tools.assert_called_once_with( - actions=['bamboohr_list_employees', 'bamboohr_get_employee'] - ) - - @patch('pydantic_ai.ext.stackone.StackOneToolSet') - def test_toolset_with_filter_pattern(self, mock_stackone_toolset_class: MagicMock) -> None: - """Test creating a StackOneToolset with filter_pattern.""" - from pydantic_ai.ext.stackone import StackOneToolset - - # Mock the tools returned by fetch_tools - mock_tool = _create_mock_stackone_tool('bamboohr_list_employees') - - mock_fetched_tools = Mock() - mock_fetched_tools.__iter__ = Mock(return_value=iter([mock_tool])) - - mock_stackone_toolset = Mock() - mock_stackone_toolset.fetch_tools.return_value = mock_fetched_tools - mock_stackone_toolset_class.return_value = mock_stackone_toolset - - # Create toolset with filter_pattern - toolset = StackOneToolset(filter_pattern='bamboohr_*', account_id='test-account', api_key='test-key') - - # Verify StackOneToolSet was created correctly - mock_stackone_toolset_class.assert_called_once_with(api_key='test-key', account_id='test-account') - - # Verify fetch_tools was called with actions parameter (list) - mock_stackone_toolset.fetch_tools.assert_called_once_with(actions=['bamboohr_*']) - - # Verify tools were created - assert isinstance(toolset, FunctionToolset) - - @patch('pydantic_ai.ext.stackone.StackOneToolSet') - def test_toolset_with_list_filter_pattern(self, mock_stackone_toolset_class: MagicMock) -> None: - """Test creating a StackOneToolset with list filter_pattern.""" - from pydantic_ai.ext.stackone import StackOneToolset - - # Mock the tools returned by fetch_tools - mock_tool1 = _create_mock_stackone_tool('bamboohr_list_employees') - mock_tool2 = _create_mock_stackone_tool('workday_list_employees') - - mock_fetched_tools = Mock() - mock_fetched_tools.__iter__ = Mock(return_value=iter([mock_tool1, mock_tool2])) - - mock_stackone_toolset = Mock() - mock_stackone_toolset.fetch_tools.return_value = mock_fetched_tools - mock_stackone_toolset_class.return_value = mock_stackone_toolset - - # Create toolset with list filter_pattern - toolset = StackOneToolset( - filter_pattern=['bamboohr_*', 'workday_*'], account_id='test-account', api_key='test-key' - ) - - # Verify fetch_tools was called with list filter_pattern as actions - mock_stackone_toolset.fetch_tools.assert_called_once_with(actions=['bamboohr_*', 'workday_*']) - - # Verify it's a FunctionToolset - assert isinstance(toolset, FunctionToolset) - - @patch('pydantic_ai.ext.stackone.StackOneToolSet') - def test_toolset_without_filter_pattern(self, mock_stackone_toolset_class: MagicMock) -> None: - """Test creating a StackOneToolset without filter_pattern (gets all tools).""" - from pydantic_ai.ext.stackone import StackOneToolset - - # Mock the tools returned by fetch_tools - mock_tool = _create_mock_stackone_tool('all_tools') - - mock_fetched_tools = Mock() - mock_fetched_tools.__iter__ = Mock(return_value=iter([mock_tool])) - - mock_stackone_toolset = Mock() - mock_stackone_toolset.fetch_tools.return_value = mock_fetched_tools - mock_stackone_toolset_class.return_value = mock_stackone_toolset - - # Create toolset without filter_pattern - toolset = StackOneToolset(account_id='test-account', api_key='test-key') - - # Verify fetch_tools was called with None actions (no filter) - mock_stackone_toolset.fetch_tools.assert_called_once_with(actions=None) - - # Verify tools were created - assert isinstance(toolset, FunctionToolset) - - @patch('pydantic_ai.ext.stackone.StackOneToolSet') - def test_toolset_with_base_url(self, mock_stackone_toolset_class: MagicMock) -> None: - """Test creating a StackOneToolset with custom base URL. - - Note: base_url is not commonly used by end users, but this test exists for coverage. - """ - from pydantic_ai.ext.stackone import StackOneToolset - - # Mock the tools returned by fetch_tools - mock_tool = _create_mock_stackone_tool('bamboohr_list_employees') - - mock_fetched_tools = Mock() - mock_fetched_tools.__iter__ = Mock(return_value=iter([mock_tool])) - - mock_stackone_toolset = Mock() - mock_stackone_toolset.fetch_tools.return_value = mock_fetched_tools - mock_stackone_toolset_class.return_value = mock_stackone_toolset - - # Create toolset with base URL - toolset = StackOneToolset( - tools=['bamboohr_list_employees'], - account_id='test-account', - api_key='test-key', - base_url='https://custom.api.stackone.com', - ) - - # Verify StackOneToolSet was called with base URL - mock_stackone_toolset_class.assert_called_once_with( - api_key='test-key', account_id='test-account', base_url='https://custom.api.stackone.com' - ) - - # Verify it's a FunctionToolset - assert isinstance(toolset, FunctionToolset) - - @patch('pydantic_ai.ext.stackone.StackOneToolSet') - def test_toolset_with_utility_tools(self, mock_stackone_toolset_class: MagicMock) -> None: - """Test creating a StackOneToolset with utility tools enabled.""" - from pydantic_ai.ext.stackone import StackOneToolset - - # Mock the tools returned by fetch_tools - mock_tool = _create_mock_stackone_tool('bamboohr_list_employees') - - mock_fetched_tools = Mock() - mock_fetched_tools.__iter__ = Mock(return_value=iter([mock_tool])) - - # Mock utility tools - mock_search_tool = _create_mock_stackone_tool('tool_search', 'Search for tools') - mock_execute_tool = _create_mock_stackone_tool('tool_execute', 'Execute a tool') - - mock_utility_tools = Mock() - mock_utility_tools.get_tool.side_effect = lambda name: ( # pyright: ignore[reportUnknownLambdaType] - mock_search_tool if name == 'tool_search' else mock_execute_tool if name == 'tool_execute' else None - ) - - mock_fetched_tools.utility_tools.return_value = mock_utility_tools - - mock_stackone_toolset = Mock() - mock_stackone_toolset.fetch_tools.return_value = mock_fetched_tools - mock_stackone_toolset_class.return_value = mock_stackone_toolset - - # Create toolset with utility tools - toolset = StackOneToolset( - filter_pattern='bamboohr_*', - include_utility_tools=True, - account_id='test-account', - api_key='test-key', - ) - - # Verify it's a FunctionToolset - assert isinstance(toolset, FunctionToolset) - - # Verify utility_tools was called (once for search_tool, once for execute_tool) - assert mock_fetched_tools.utility_tools.call_count == 2 From a41f47638606ea190f0d039270db4e1e7d525fd8 Mon Sep 17 00:00:00 2001 From: Shashi-Stackone Date: Mon, 23 Feb 2026 20:36:10 +0000 Subject: [PATCH 13/25] Remove stackone from extra dependency and follow same patter othe toolsets --- pydantic_ai_slim/pyproject.toml | 1 - pyproject.toml | 1 + tests/ext/test_stackone.py | 35 +++++++++++++++++++------------ uv.lock | 37 +-------------------------------- 4 files changed, 24 insertions(+), 50 deletions(-) diff --git a/pydantic_ai_slim/pyproject.toml b/pydantic_ai_slim/pyproject.toml index fa9008f902..18ad513ca6 100644 --- a/pydantic_ai_slim/pyproject.toml +++ b/pydantic_ai_slim/pyproject.toml @@ -98,7 +98,6 @@ outlines-vllm-offline = [ duckduckgo = ["ddgs>=9.0.0"] tavily = ["tavily-python>=0.5.0"] exa = ["exa-py>=2.0.0"] -stackone = ["stackone-ai>=2.3.1"] # CLI cli = [ "rich>=13", diff --git a/pyproject.toml b/pyproject.toml index 4b000ca810..0d3c72cc85 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 ] diff --git a/tests/ext/test_stackone.py b/tests/ext/test_stackone.py index 4d62f6886c..d7731ab5bb 100644 --- a/tests/ext/test_stackone.py +++ b/tests/ext/test_stackone.py @@ -1,3 +1,4 @@ +# pyright: reportPrivateUsage=false from __future__ import annotations from dataclasses import dataclass, field @@ -8,16 +9,24 @@ from inline_snapshot import snapshot from pydantic_ai import Agent -from pydantic_ai.ext.stackone import ( - StackOneToolset, - _tool_from_stackone_tool, - execute_tool, - feedback_tool, - search_tool, - tool_from_stackone, -) 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, + execute_tool, + feedback_tool, + search_tool, + tool_from_stackone, + ) + +pytestmark = [ + pytest.mark.skipif(not imports_successful(), reason='stackone-ai not installed'), +] + @dataclass class SimulatedStackOneTool: @@ -52,33 +61,33 @@ def execute(self, arguments: dict[str, Any]) -> Any: def test_tool_conversion(): - tool = _tool_from_stackone_tool(employee_tool) + 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 = _tool_from_stackone_tool(employee_tool) + 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 = _tool_from_stackone_tool(employee_tool) + 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 = _tool_from_stackone_tool(tool_with_none_desc) + tool = stackone_ext._tool_from_stackone_tool(tool_with_none_desc) assert tool.description == '' def test_tool_schema(): - tool = _tool_from_stackone_tool(employee_tool) + tool = stackone_ext._tool_from_stackone_tool(employee_tool) assert tool.function_schema.json_schema == employee_tool.to_openai_function()['function']['parameters'] diff --git a/uv.lock b/uv.lock index 372ebd33cd..07aefa43f7 100644 --- a/uv.lock +++ b/uv.lock @@ -715,20 +715,6 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/b0/5c/dbd00727a3dd165d7e0e8af40e630cd7e45d77b525a3218afaff8a87358e/blake3-1.0.8-cp314-cp314t-win_amd64.whl", hash = "sha256:421b99cdf1ff2d1bf703bc56c454f4b286fce68454dd8711abbcb5a0df90c19a", size = 215133, upload-time = "2025-10-14T06:47:16.069Z" }, ] -[[package]] -name = "bm25s" -version = "0.2.14" -source = { registry = "https://pypi.org/simple" } -dependencies = [ - { name = "numpy" }, - { name = "scipy", version = "1.15.3", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.11'" }, - { name = "scipy", version = "1.16.3", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.11'" }, -] -sdist = { url = "https://files.pythonhosted.org/packages/99/72/5ad06c30991ba494242785a3ab8987deb01c07dfc1c492847bde221e62bf/bm25s-0.2.14.tar.gz", hash = "sha256:7b6717770fffbdb3b962e5fe8ef1e6eac7f285d0fbc14484b321e136df837139", size = 59266, upload-time = "2025-09-08T17:06:30.728Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/7c/3e/e3ae2f0fb0f8f46f9c787fa419ca5203ff850d0630749a26baf0a6570453/bm25s-0.2.14-py3-none-any.whl", hash = "sha256:76cdb70ae40747941b150a1ec16a9c20c576d6534d0a3c3eebb303c779b3cf65", size = 55128, upload-time = "2025-09-08T17:06:29.324Z" }, -] - [[package]] name = "boto3" version = "1.42.14" @@ -6516,9 +6502,6 @@ retries = [ sentence-transformers = [ { name = "sentence-transformers" }, ] -stackone = [ - { name = "stackone-ai" }, -] tavily = [ { name = "tavily-python" }, ] @@ -6587,7 +6570,6 @@ requires-dist = [ { name = "requests", marker = "extra == 'vertexai'", specifier = ">=2.32.2" }, { name = "rich", marker = "extra == 'cli'", specifier = ">=13" }, { name = "sentence-transformers", marker = "extra == 'sentence-transformers'", specifier = ">=5.2.0" }, - { name = "stackone-ai", marker = "extra == 'stackone'", specifier = ">=2.3.1" }, { name = "starlette", marker = "extra == 'ag-ui'", specifier = ">=0.45.3" }, { name = "starlette", marker = "extra == 'ui'", specifier = ">=0.45.3" }, { name = "starlette", marker = "extra == 'web'", specifier = ">=0.45.3" }, @@ -6604,7 +6586,7 @@ requires-dist = [ { name = "voyageai", marker = "extra == 'voyageai'", specifier = ">=0.3.2" }, { name = "xai-sdk", marker = "extra == 'xai'", specifier = ">=1.5.0" }, ] -provides-extras = ["a2a", "ag-ui", "anthropic", "bedrock", "cli", "cohere", "dbos", "duckduckgo", "evals", "exa", "fastmcp", "google", "groq", "huggingface", "logfire", "mcp", "mistral", "openai", "openrouter", "outlines-llamacpp", "outlines-mlxlm", "outlines-sglang", "outlines-transformers", "outlines-vllm-offline", "prefect", "retries", "sentence-transformers", "stackone", "tavily", "temporal", "ui", "vertexai", "voyageai", "web", "xai"] +provides-extras = ["a2a", "ag-ui", "anthropic", "bedrock", "cli", "cohere", "dbos", "duckduckgo", "evals", "exa", "fastmcp", "google", "groq", "huggingface", "logfire", "mcp", "mistral", "openai", "openrouter", "outlines-llamacpp", "outlines-mlxlm", "outlines-sglang", "outlines-transformers", "outlines-vllm-offline", "prefect", "retries", "sentence-transformers", "tavily", "temporal", "ui", "vertexai", "voyageai", "web", "xai"] [[package]] name = "pydantic-core" @@ -8598,23 +8580,6 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/71/22/8ab1066358601163e1ac732837adba3672f703818f693e179b24e0d3b65c/sse_starlette-3.0.4-py3-none-any.whl", hash = "sha256:32c80ef0d04506ced4b0b6ab8fe300925edc37d26f666afb1874c754895f5dc3", size = 11764, upload-time = "2025-12-14T16:22:51.453Z" }, ] -[[package]] -name = "stackone-ai" -version = "2.3.1" -source = { registry = "https://pypi.org/simple" } -dependencies = [ - { name = "bm25s" }, - { name = "httpx" }, - { name = "langchain-core" }, - { name = "numpy" }, - { name = "pydantic" }, - { name = "typing-extensions" }, -] -sdist = { url = "https://files.pythonhosted.org/packages/36/06/cc058d329e44a73fa7aa25f03e42aa0ee15d5f22aca534f21f7eb8cf63a7/stackone_ai-2.3.1.tar.gz", hash = "sha256:d25b521ad6cee7f8107c9beef9a05a67bc39b07646d4ac4f4665fea1dbe3d601", size = 537863, upload-time = "2026-01-29T16:59:20.683Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/97/5a/31ef092360e9175e86909df86d1f889f5e7f80df04ceb3e855ed81fa563f/stackone_ai-2.3.1-py3-none-any.whl", hash = "sha256:bd6ffc937f894045d410353eaec456e1a40001dc6e9b28904bd770869061b60d", size = 30818, upload-time = "2026-01-29T16:59:19.7Z" }, -] - [[package]] name = "starlette" version = "0.50.0" From da2358e97921f9e37366bc7b0e42fab67b1ba003 Mon Sep 17 00:00:00 2001 From: Shashi-Stackone Date: Mon, 23 Feb 2026 21:07:47 +0000 Subject: [PATCH 14/25] chore: retrigger CI after adding feature label to PR From 6453f3847363690a393b1e39f21fd1836bd74d12 Mon Sep 17 00:00:00 2001 From: Shashi-Stackone Date: Mon, 23 Feb 2026 21:32:52 +0000 Subject: [PATCH 15/25] Skip the stackone coverage as its optional for toolset providers --- pyproject.toml | 1 + 1 file changed, 1 insertion(+) diff --git a/pyproject.toml b/pyproject.toml index 0d3c72cc85..2161931c26 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -269,6 +269,7 @@ 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", From 3585e04e79d6b80f313f49cfabb6240058fa6e39 Mon Sep 17 00:00:00 2001 From: Shashi-Stackone Date: Tue, 24 Feb 2026 09:26:25 +0000 Subject: [PATCH 16/25] chore: retrigger CI From f5f65269dc5a16dbb3aa842bf43de5f0fcc75eb3 Mon Sep 17 00:00:00 2001 From: Shashi-Stackone Date: Wed, 25 Feb 2026 11:23:40 +0000 Subject: [PATCH 17/25] Remove the fedback tools from the Pydantic AI integration --- docs/third-party-tools.md | 32 --------------- pydantic_ai_slim/pydantic_ai/ext/stackone.py | 42 +------------------- tests/ext/test_stackone.py | 37 ----------------- 3 files changed, 1 insertion(+), 110 deletions(-) diff --git a/docs/third-party-tools.md b/docs/third-party-tools.md index 827d716af7..c5dc4bb339 100644 --- a/docs/third-party-tools.md +++ b/docs/third-party-tools.md @@ -173,38 +173,6 @@ agent = Agent( ) ``` -### Feedback Collection {#stackone-feedback} - -StackOne supports collecting user feedback on tool performance. Enable this by setting `include_feedback_tool=True`: - -```python {test="skip"} -import os - -from pydantic_ai import Agent -from pydantic_ai.ext.stackone import StackOneToolset - -toolset = StackOneToolset( - filter_pattern='bamboohr_*', - include_feedback_tool=True, - account_id=os.getenv('STACKONE_ACCOUNT_ID'), - api_key=os.getenv('STACKONE_API_KEY'), -) - -agent = Agent('openai:gpt-5', toolsets=[toolset]) -``` - -You can also use the standalone [`feedback_tool`][pydantic_ai.ext.stackone.feedback_tool] function: - -```python {test="skip"} -from pydantic_ai import Agent -from pydantic_ai.ext.stackone import feedback_tool - -agent = Agent( - 'openai:gpt-5', - tools=[feedback_tool()], -) -``` - ## See Also - [Function Tools](tools.md) - Basic tool concepts and registration diff --git a/pydantic_ai_slim/pydantic_ai/ext/stackone.py b/pydantic_ai_slim/pydantic_ai/ext/stackone.py index b40d62a67f..fc67f2b40a 100644 --- a/pydantic_ai_slim/pydantic_ai/ext/stackone.py +++ b/pydantic_ai_slim/pydantic_ai/ext/stackone.py @@ -1,6 +1,5 @@ from __future__ import annotations -import os from collections.abc import Sequence from typing import Any @@ -14,7 +13,7 @@ except ImportError as _import_error: raise ImportError('Please install `stackone-ai` to use StackOne tools.') from _import_error -__all__ = ('tool_from_stackone', 'StackOneToolset', 'search_tool', 'execute_tool', 'feedback_tool') +__all__ = ('tool_from_stackone', 'StackOneToolset', 'search_tool', 'execute_tool') def _tool_from_stackone_tool(stackone_tool: Any) -> Tool: @@ -97,41 +96,6 @@ def execute_tool( return _tool_from_stackone_tool(execute) -def feedback_tool( - *, - api_key: str | None = None, - account_id: str | None = None, - base_url: str | None = None, -) -> Tool: - """Creates a feedback tool for collecting user feedback on StackOne tool usage. - - Args: - api_key: The StackOne API key. If not provided, uses `STACKONE_API_KEY` env var. - account_id: The StackOne account ID. If not provided, uses `STACKONE_ACCOUNT_ID` env var. - base_url: Optional custom base URL for the StackOne API. - - Returns: - A Pydantic AI tool for collecting feedback. - """ - try: - from stackone_ai.feedback.tool import create_feedback_tool - except ImportError as e: - raise ImportError('Please install `stackone-ai` with feedback support to use the feedback tool.') from e - - resolved_api_key = api_key or os.environ.get('STACKONE_API_KEY') - if resolved_api_key is None: - raise ValueError( - 'API key is required. Provide it as an argument or set the STACKONE_API_KEY environment variable.' - ) - - fb_tool = create_feedback_tool( - api_key=resolved_api_key, - account_id=account_id, - **({'base_url': base_url} if base_url else {}), - ) - return _tool_from_stackone_tool(fb_tool) - - class StackOneToolset(FunctionToolset): """A toolset that wraps StackOne tools.""" @@ -144,7 +108,6 @@ def __init__( base_url: str | None = None, filter_pattern: str | list[str] | None = None, include_utility_tools: bool = False, - include_feedback_tool: bool = False, hybrid_alpha: float | None = None, id: str | None = None, ): @@ -165,7 +128,4 @@ def __init__( for stackone_tool in fetched_tools: pydantic_tools.append(_tool_from_stackone_tool(stackone_tool)) - if include_feedback_tool: - pydantic_tools.append(feedback_tool(api_key=api_key, account_id=account_id, base_url=base_url)) - super().__init__(pydantic_tools, id=id) diff --git a/tests/ext/test_stackone.py b/tests/ext/test_stackone.py index d7731ab5bb..f2c212fea2 100644 --- a/tests/ext/test_stackone.py +++ b/tests/ext/test_stackone.py @@ -18,7 +18,6 @@ from pydantic_ai.ext.stackone import ( StackOneToolset, execute_tool, - feedback_tool, search_tool, tool_from_stackone, ) @@ -257,42 +256,6 @@ def get_tool_side_effect(name: str) -> SimulatedStackOneTool | None: assert 'tool_search' in result.output or 'tool_execute' in result.output -@patch.dict('os.environ', {'STACKONE_API_KEY': 'env-key'}) -@patch('stackone_ai.feedback.tool.create_feedback_tool') -def test_feedback_tool_from_env(mock_create_fb: Any): - mock_fb = SimulatedStackOneTool(name='tool_feedback', description='Feedback') - mock_create_fb.return_value = mock_fb - - tool = feedback_tool() - assert tool.name == 'tool_feedback' - mock_create_fb.assert_called_once_with(api_key='env-key', account_id=None) - - -@patch('stackone_ai.feedback.tool.create_feedback_tool') -def test_feedback_tool_explicit_key(mock_create_fb: Any): - mock_fb = SimulatedStackOneTool(name='tool_feedback', description='Feedback') - mock_create_fb.return_value = mock_fb - - tool = feedback_tool(api_key='explicit-key', account_id='acc-1') - assert tool.name == 'tool_feedback' - mock_create_fb.assert_called_once_with(api_key='explicit-key', account_id='acc-1') - - -@patch('stackone_ai.feedback.tool.create_feedback_tool') -def test_feedback_tool_with_base_url(mock_create_fb: Any): - mock_fb = SimulatedStackOneTool(name='tool_feedback', description='Feedback') - mock_create_fb.return_value = mock_fb - - feedback_tool(api_key='k', base_url='https://custom.stackone.com') - mock_create_fb.assert_called_once_with(api_key='k', account_id=None, base_url='https://custom.stackone.com') - - -@patch.dict('os.environ', {}, clear=True) -def test_feedback_tool_no_key(): - with pytest.raises(ValueError, match='API key is required'): - feedback_tool() - - def test_import_error(): with patch.dict('sys.modules', {'stackone_ai': None}): with pytest.raises(ImportError, match='Please install `stackone-ai`'): From 47f75a6f8330d49f5a27665c025a174bdcefe712 Mon Sep 17 00:00:00 2001 From: Shashi-Stackone Date: Wed, 25 Feb 2026 12:47:11 +0000 Subject: [PATCH 18/25] update Pydantic example to use google gemini model as it tested locally --- docs/third-party-tools.md | 6 +++--- docs/toolsets.md | 3 +-- 2 files changed, 4 insertions(+), 5 deletions(-) diff --git a/docs/third-party-tools.md b/docs/third-party-tools.md index c5dc4bb339..88a65187f6 100644 --- a/docs/third-party-tools.md +++ b/docs/third-party-tools.md @@ -121,7 +121,7 @@ employee_tool = tool_from_stackone( ) agent = Agent( - 'openai:gpt-5', + 'google-gla:gemini-2.5-flash', tools=[employee_tool], ) @@ -151,7 +151,7 @@ toolset = StackOneToolset( api_key=os.getenv('STACKONE_API_KEY'), ) -agent = Agent('openai:gpt-5', toolsets=[toolset]) +agent = Agent('google-gla:gemini-2.5-flash', toolsets=[toolset]) result = agent.run_sync('Find a tool to list employees and use it') print(result.output) ``` @@ -168,7 +168,7 @@ stackone = StackOneToolSet() tools = stackone.fetch_tools(actions=['bamboohr_*']) agent = Agent( - 'openai:gpt-5', + 'google-gla:gemini-2.5-flash', tools=[search_tool(tools), execute_tool(tools)], ) ``` diff --git a/docs/toolsets.md b/docs/toolsets.md index e52d912f4a..f6d37d716b 100644 --- a/docs/toolsets.md +++ b/docs/toolsets.md @@ -760,9 +760,8 @@ specific_toolset = StackOneToolset( api_key=os.getenv('STACKONE_API_KEY'), ) -agent = Agent('openai:gpt-5', toolsets=[toolset]) +agent = Agent('google-gla:gemini-2.5-flash', toolsets=[toolset]) -# Example usage result = agent.run_sync('List all employees and get information about the first employee') print(result.output) ``` From 27511c21c49ce8951d3769365efa8bb9a4ec8705 Mon Sep 17 00:00:00 2001 From: Shashi-Stackone Date: Thu, 26 Feb 2026 09:24:53 +0000 Subject: [PATCH 19/25] chore: retrigger CI for cancelled job From 403c5d6cbb9facbddcbb510bce81ae433b6702bb Mon Sep 17 00:00:00 2001 From: Shashi-Stackone Date: Thu, 26 Feb 2026 10:18:58 +0000 Subject: [PATCH 20/25] chore: retrigger CI for cancelled job From dca5b53c0561688581721604547e4bc09a408aa1 Mon Sep 17 00:00:00 2001 From: Shashi-Stackone Date: Fri, 27 Feb 2026 10:29:21 +0000 Subject: [PATCH 21/25] Remove the utility tools from the Pydantic --- docs/third-party-tools.md | 42 ----------- pydantic_ai_slim/pydantic_ai/ext/stackone.py | 51 +------------- tests/ext/test_stackone.py | 74 -------------------- 3 files changed, 3 insertions(+), 164 deletions(-) diff --git a/docs/third-party-tools.md b/docs/third-party-tools.md index 88a65187f6..71f589af9a 100644 --- a/docs/third-party-tools.md +++ b/docs/third-party-tools.md @@ -131,48 +131,6 @@ 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) which supports pattern matching for tool selection. -### Dynamic Tool Discovery {#stackone-dynamic-discovery} - -For large tool sets where you don't know in advance which tools the agent will need, you can enable **utility tools mode** by setting `include_utility_tools=True`. This provides two special tools: - -- `tool_search`: Searches for relevant tools using natural language queries -- `tool_execute`: Executes a discovered tool by name - -```python {test="skip"} -import os - -from pydantic_ai import Agent -from pydantic_ai.ext.stackone import StackOneToolset - -toolset = StackOneToolset( - filter_pattern='bamboohr_*', # Load StackOne tools - include_utility_tools=True, # Enable dynamic discovery - 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('Find a tool to list employees and use it') -print(result.output) -``` - -You can also use the standalone [`search_tool`][pydantic_ai.ext.stackone.search_tool] and [`execute_tool`][pydantic_ai.ext.stackone.execute_tool] functions for more control: - -```python {test="skip"} -from stackone_ai import StackOneToolSet - -from pydantic_ai import Agent -from pydantic_ai.ext.stackone import execute_tool, search_tool - -stackone = StackOneToolSet() -tools = stackone.fetch_tools(actions=['bamboohr_*']) - -agent = Agent( - 'google-gla:gemini-2.5-flash', - tools=[search_tool(tools), execute_tool(tools)], -) -``` - ## See Also - [Function Tools](tools.md) - Basic tool concepts and registration diff --git a/pydantic_ai_slim/pydantic_ai/ext/stackone.py b/pydantic_ai_slim/pydantic_ai/ext/stackone.py index fc67f2b40a..f38ef25754 100644 --- a/pydantic_ai_slim/pydantic_ai/ext/stackone.py +++ b/pydantic_ai_slim/pydantic_ai/ext/stackone.py @@ -13,7 +13,7 @@ except ImportError as _import_error: raise ImportError('Please install `stackone-ai` to use StackOne tools.') from _import_error -__all__ = ('tool_from_stackone', 'StackOneToolset', 'search_tool', 'execute_tool') +__all__ = ('tool_from_stackone', 'StackOneToolset') def _tool_from_stackone_tool(stackone_tool: Any) -> Tool: @@ -57,45 +57,6 @@ def tool_from_stackone( return _tool_from_stackone_tool(stackone_tool) -def search_tool( - tools: Any, - *, - hybrid_alpha: float | None = None, -) -> Tool: - """Creates a tool for discovering StackOne tools using natural language queries. - - Args: - tools: A StackOne `Tools` collection returned from `StackOneToolSet.fetch_tools()`. - hybrid_alpha: Optional weight for hybrid search (0-1). Defaults to the SDK default. - - Returns: - A Pydantic AI tool for searching StackOne tools. - """ - utility_tools = tools.utility_tools(hybrid_alpha=hybrid_alpha) - search = utility_tools.get_tool('tool_search') - if search is None: - raise ValueError('tool_search not found in StackOne utility tools') - return _tool_from_stackone_tool(search) - - -def execute_tool( - tools: Any, -) -> Tool: - """Creates a tool for executing discovered StackOne tools by name. - - Args: - tools: A StackOne `Tools` collection returned from `StackOneToolSet.fetch_tools()`. - - Returns: - A Pydantic AI tool for executing StackOne tools. - """ - utility_tools = tools.utility_tools() - execute = utility_tools.get_tool('tool_execute') - if execute is None: - raise ValueError('tool_execute not found in StackOne utility tools') - return _tool_from_stackone_tool(execute) - - class StackOneToolset(FunctionToolset): """A toolset that wraps StackOne tools.""" @@ -107,8 +68,6 @@ def __init__( api_key: str | None = None, base_url: str | None = None, filter_pattern: str | list[str] | None = None, - include_utility_tools: bool = False, - hybrid_alpha: float | None = None, id: str | None = None, ): stackone_toolset = StackOneToolSet(api_key=api_key, account_id=account_id, base_url=base_url) @@ -121,11 +80,7 @@ def __init__( fetched_tools = stackone_toolset.fetch_tools(actions=actions) pydantic_tools: list[Tool] = [] - if include_utility_tools: - pydantic_tools.append(search_tool(fetched_tools, hybrid_alpha=hybrid_alpha)) - pydantic_tools.append(execute_tool(fetched_tools)) - else: - for stackone_tool in fetched_tools: - pydantic_tools.append(_tool_from_stackone_tool(stackone_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/tests/ext/test_stackone.py b/tests/ext/test_stackone.py index f2c212fea2..47c0da00d0 100644 --- a/tests/ext/test_stackone.py +++ b/tests/ext/test_stackone.py @@ -17,8 +17,6 @@ import pydantic_ai.ext.stackone as stackone_ext from pydantic_ai.ext.stackone import ( StackOneToolset, - execute_tool, - search_tool, tool_from_stackone, ) @@ -184,78 +182,6 @@ def test_stackone_toolset_with_base_url(mock_toolset_cls: Any): mock_toolset_cls.assert_called_once_with(api_key='k', account_id=None, base_url='https://custom.stackone.com') -def test_search_tool(): - mock_search = SimulatedStackOneTool(name='tool_search', description='Search for tools') - mock_utility = Mock() - mock_utility.get_tool.return_value = mock_search - - mock_tools = Mock() - mock_tools.utility_tools.return_value = mock_utility - - tool = search_tool(mock_tools, hybrid_alpha=0.3) - assert tool.name == 'tool_search' - mock_tools.utility_tools.assert_called_once_with(hybrid_alpha=0.3) - - -def test_search_tool_not_found(): - mock_utility = Mock() - mock_utility.get_tool.return_value = None - mock_tools = Mock() - mock_tools.utility_tools.return_value = mock_utility - - with pytest.raises(ValueError, match='tool_search not found'): - search_tool(mock_tools) - - -def test_execute_tool(): - mock_exec = SimulatedStackOneTool(name='tool_execute', description='Execute a tool') - mock_utility = Mock() - mock_utility.get_tool.return_value = mock_exec - - mock_tools = Mock() - mock_tools.utility_tools.return_value = mock_utility - - tool = execute_tool(mock_tools) - assert tool.name == 'tool_execute' - - -def test_execute_tool_not_found(): - mock_utility = Mock() - mock_utility.get_tool.return_value = None - mock_tools = Mock() - mock_tools.utility_tools.return_value = mock_utility - - with pytest.raises(ValueError, match='tool_execute not found'): - execute_tool(mock_tools) - - -@patch('pydantic_ai.ext.stackone.StackOneToolSet') -def test_stackone_toolset_utility_tools_mode(mock_toolset_cls: Any): - mock_search = SimulatedStackOneTool(name='tool_search', description='Search') - mock_exec = SimulatedStackOneTool(name='tool_execute', description='Execute') - - mock_utility = Mock() - - def get_tool_side_effect(name: str) -> SimulatedStackOneTool | None: - if name == 'tool_search': - return mock_search - if name == 'tool_execute': - return mock_exec - return None - - mock_utility.get_tool.side_effect = get_tool_side_effect - - mock_fetched = Mock() - mock_fetched.utility_tools.return_value = mock_utility - mock_toolset_cls.return_value.fetch_tools.return_value = mock_fetched - - toolset = StackOneToolset(filter_pattern='bamboohr_*', include_utility_tools=True, api_key='k') - - agent = Agent('test', toolsets=[toolset]) - result = agent.run_sync('foobar') - assert 'tool_search' in result.output or 'tool_execute' in result.output - - def test_import_error(): with patch.dict('sys.modules', {'stackone_ai': None}): with pytest.raises(ImportError, match='Please install `stackone-ai`'): From a8d4d23775c2026dd63c6e76cfc668fac585f164 Mon Sep 17 00:00:00 2001 From: Shashi-Stackone Date: Mon, 9 Mar 2026 19:10:10 +0000 Subject: [PATCH 22/25] Update docs and include search --- docs/third-party-tools.md | 24 +++++++++++++- docs/toolsets.md | 19 +++++++++++ pydantic_ai_slim/pydantic_ai/ext/stackone.py | 35 ++++++++++++++++---- tests/ext/test_stackone.py | 34 +++++++++++++++++++ 4 files changed, 105 insertions(+), 7 deletions(-) diff --git a/docs/third-party-tools.md b/docs/third-party-tools.md index 71f589af9a..0db608aa3b 100644 --- a/docs/third-party-tools.md +++ b/docs/third-party-tools.md @@ -129,7 +129,29 @@ 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) which supports pattern matching for tool selection. +If you'd like to use multiple StackOne tools, you can use the [`StackOneToolset`][pydantic_ai.ext.stackone.StackOneToolset] [toolset](toolsets.md#stackone-tools) which supports pattern matching and semantic search for tool selection: + +```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', # (1)! +) + +agent = Agent('google-gla:gemini-2.5-flash', toolsets=[toolset]) + +result = agent.run_sync('List all employees in the HR system') +print(result.output) +``` + +1. Search mode: `"auto"` (default) tries semantic search with local fallback, `"semantic"` uses only semantic search, `"local"` uses only local BM25+TF-IDF search. Requires `stackone-ai >= 2.4.0`. ## See Also diff --git a/docs/toolsets.md b/docs/toolsets.md index f6d37d716b..1b20253810 100644 --- a/docs/toolsets.md +++ b/docs/toolsets.md @@ -765,3 +765,22 @@ 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-2.5-flash', toolsets=[toolset]) +``` diff --git a/pydantic_ai_slim/pydantic_ai/ext/stackone.py b/pydantic_ai_slim/pydantic_ai/ext/stackone.py index f38ef25754..b4dd871deb 100644 --- a/pydantic_ai_slim/pydantic_ai/ext/stackone.py +++ b/pydantic_ai_slim/pydantic_ai/ext/stackone.py @@ -1,7 +1,7 @@ from __future__ import annotations from collections.abc import Sequence -from typing import Any +from typing import Any, Literal from pydantic.json_schema import JsonSchemaValue @@ -58,7 +58,11 @@ def tool_from_stackone( class StackOneToolset(FunctionToolset): - """A toolset that wraps StackOne tools.""" + """A toolset that wraps StackOne tools. + + Supports fetching tools by name/pattern, or searching by natural-language + query (requires ``stackone-ai >= 2.4.0``). + """ def __init__( self, @@ -68,16 +72,35 @@ def __init__( 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, 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 tools is not None: - actions = list(tools) + 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 search is not None: + search_kwargs['search'] = search + fetched_tools = stackone_toolset.search_tools(search_query, **search_kwargs) else: - actions = [filter_pattern] if isinstance(filter_pattern, str) else filter_pattern + 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) + fetched_tools = stackone_toolset.fetch_tools(actions=actions) pydantic_tools: list[Tool] = [] for stackone_tool in fetched_tools: diff --git a/tests/ext/test_stackone.py b/tests/ext/test_stackone.py index 47c0da00d0..fc57bca3e9 100644 --- a/tests/ext/test_stackone.py +++ b/tests/ext/test_stackone.py @@ -182,6 +182,40 @@ def test_stackone_toolset_with_base_url(mock_toolset_cls: Any): 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`'): From 580cea4f1049f9b1ea53c748915f7be3bfe62506 Mon Sep 17 00:00:00 2001 From: Shashi-Stackone Date: Mon, 9 Mar 2026 19:46:14 +0000 Subject: [PATCH 23/25] Update docs with tool search and the include min similarity --- docs/third-party-tools.md | 8 ++++---- pydantic_ai_slim/pydantic_ai/ext/stackone.py | 9 ++++----- 2 files changed, 8 insertions(+), 9 deletions(-) diff --git a/docs/third-party-tools.md b/docs/third-party-tools.md index 0db608aa3b..9f69a14bcf 100644 --- a/docs/third-party-tools.md +++ b/docs/third-party-tools.md @@ -129,7 +129,9 @@ 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) which supports pattern matching and semantic search for tool selection: +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 @@ -142,7 +144,7 @@ toolset = StackOneToolset( account_id=os.getenv('STACKONE_ACCOUNT_ID'), api_key=os.getenv('STACKONE_API_KEY'), top_k=5, - search='auto', # (1)! + search='auto', ) agent = Agent('google-gla:gemini-2.5-flash', toolsets=[toolset]) @@ -151,8 +153,6 @@ result = agent.run_sync('List all employees in the HR system') print(result.output) ``` -1. Search mode: `"auto"` (default) tries semantic search with local fallback, `"semantic"` uses only semantic search, `"local"` uses only local BM25+TF-IDF search. Requires `stackone-ai >= 2.4.0`. - ## See Also - [Function Tools](tools.md) - Basic tool concepts and registration diff --git a/pydantic_ai_slim/pydantic_ai/ext/stackone.py b/pydantic_ai_slim/pydantic_ai/ext/stackone.py index b4dd871deb..49a339b8f6 100644 --- a/pydantic_ai_slim/pydantic_ai/ext/stackone.py +++ b/pydantic_ai_slim/pydantic_ai/ext/stackone.py @@ -58,11 +58,7 @@ def tool_from_stackone( class StackOneToolset(FunctionToolset): - """A toolset that wraps StackOne tools. - - Supports fetching tools by name/pattern, or searching by natural-language - query (requires ``stackone-ai >= 2.4.0``). - """ + """A toolset that wraps StackOne tools.""" def __init__( self, @@ -75,6 +71,7 @@ def __init__( 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: @@ -91,6 +88,8 @@ def __init__( 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) From df805eeed62c302ae8109de48bd7b207ddc1dd24 Mon Sep 17 00:00:00 2001 From: Shashi-Stackone Date: Tue, 10 Mar 2026 09:29:57 +0000 Subject: [PATCH 24/25] Fix Ruff issue --- pydantic_ai_slim/pydantic_ai/ext/stackone.py | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/pydantic_ai_slim/pydantic_ai/ext/stackone.py b/pydantic_ai_slim/pydantic_ai/ext/stackone.py index 49a339b8f6..e75eab6267 100644 --- a/pydantic_ai_slim/pydantic_ai/ext/stackone.py +++ b/pydantic_ai_slim/pydantic_ai/ext/stackone.py @@ -82,8 +82,7 @@ def __init__( 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_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: From ae117f0a300b00214042885d8bb9f3d78af65080 Mon Sep 17 00:00:00 2001 From: Shashi-Stackone Date: Tue, 10 Mar 2026 18:17:03 +0000 Subject: [PATCH 25/25] update examples to use the google-gla:gemini-3-pro-preview model --- docs/third-party-tools.md | 4 ++-- docs/toolsets.md | 2 +- 2 files changed, 3 insertions(+), 3 deletions(-) diff --git a/docs/third-party-tools.md b/docs/third-party-tools.md index 9f69a14bcf..0e923783aa 100644 --- a/docs/third-party-tools.md +++ b/docs/third-party-tools.md @@ -121,7 +121,7 @@ employee_tool = tool_from_stackone( ) agent = Agent( - 'google-gla:gemini-2.5-flash', + 'google-gla:gemini-3-pro-preview', tools=[employee_tool], ) @@ -147,7 +147,7 @@ toolset = StackOneToolset( search='auto', ) -agent = Agent('google-gla:gemini-2.5-flash', toolsets=[toolset]) +agent = Agent('google-gla:gemini-3-pro-preview', toolsets=[toolset]) result = agent.run_sync('List all employees in the HR system') print(result.output) diff --git a/docs/toolsets.md b/docs/toolsets.md index 1b20253810..8768ca9436 100644 --- a/docs/toolsets.md +++ b/docs/toolsets.md @@ -782,5 +782,5 @@ toolset = StackOneToolset( search='auto', # 'auto', 'semantic', or 'local' ) -agent = Agent('google-gla:gemini-2.5-flash', toolsets=[toolset]) +agent = Agent('google-gla:gemini-3-pro-preview', toolsets=[toolset]) ```