feat(integrations): add native Pydantic AI support (#184)

shashi-stackone · web-flow · commit 8216ab6c0776 · 2026-04-30T12:00:01.000+01:00
* Add Native Pydantic AI Integration to StackOne SDK

* Keep the implementation minimum for now

* Type chcking enhancements for pydantic tools
diff --git a/README.md b/README.md
@@ -26,6 +26,7 @@ StackOne AI provides a unified interface for accessing various SaaS tools throug
   - LangChain Tools
   - CrewAI Tools
   - LangGraph Tool Node
+  - Pydantic AI Toolset
 
 ## Requirements
 
@@ -195,6 +196,43 @@ for tool_call in response.tool_calls:
 
 </details>
 
+<details>
+<summary>Pydantic AI Integration</summary>
+
+StackOne tools convert to Pydantic AI `Tool` instances via `.to_pydantic_ai()`, parallel to `.to_openai()` and `.to_langchain()`:
+
+Prerequisites:
+
+```bash
+pip install 'stackone-ai[pydantic-ai]'
+```
+
+```python
+import os
+from pydantic_ai import Agent
+from stackone_ai import StackOneToolSet
+
+toolset = StackOneToolSet()
+tools = toolset.fetch_tools(
+    actions=["workday_list_workers", "workday_get_worker"],
+    account_ids=[os.environ["STACKONE_ACCOUNT_ID"]],
+).to_pydantic_ai()
+
+agent = Agent("openai:gpt-5.4", tools=tools)
+result = agent.run_sync("List the first 5 employees")
+print(result.output)
+```
+
+For the full catalog (or the meta search/execute tools), use the `.pydantic_ai()` method on `StackOneToolSet` — parallel to `.openai()` / `.langchain()`:
+
+```python
+toolset = StackOneToolSet()
+tools = toolset.pydantic_ai(account_ids=[os.environ["STACKONE_ACCOUNT_ID"]])
+# or pydantic_ai(mode="search_and_execute") for agent-driven discovery
+```
+
+</details>
+
 <details>
 <summary>LangGraph Integration</summary>
 
diff --git a/examples/pydantic_ai_integration.py b/examples/pydantic_ai_integration.py
@@ -3,11 +3,13 @@
 ```bash
 uv run examples/pydantic_ai_integration.py
 ```
+
+Install with `pip install 'stackone-ai[pydantic-ai]'` (or
+`pip install 'stackone-ai[examples]'` to run this file).
 """
 
 from __future__ import annotations
 
-import json
 import os
 
 try:
@@ -18,31 +20,15 @@
     pass
 
 try:
-    from pydantic_ai import Agent, Tool
+    from pydantic_ai import Agent
 except ImportError:
-    print("Install pydantic-ai to run this example: pip install pydantic-ai")
+    print("Install pydantic-ai to run this example: pip install 'stackone-ai[pydantic-ai]'")
     raise SystemExit(1) from None
 
 from stackone_ai import StackOneToolSet
 
 
-def _to_pydantic_ai_tool(stackone_tool) -> Tool:
-    """Convert a StackOneTool to a Pydantic AI Tool with the proper JSON schema."""
-    params_schema = stackone_tool.to_openai_function()["function"]["parameters"]
-
-    def execute(**kwargs: object) -> str:
-        return json.dumps(stackone_tool.execute(kwargs))
-
-    return Tool.from_schema(
-        execute,
-        name=stackone_tool.name,
-        description=stackone_tool.description,
-        json_schema=params_schema,
-    )
-
-
 def pydantic_ai_integration() -> None:
-    account_id = os.getenv("STACKONE_ACCOUNT_ID")
     for var in ["STACKONE_API_KEY", "STACKONE_ACCOUNT_ID", "OPENAI_API_KEY"]:
         if not os.getenv(var):
             print(f"Set {var} to run this example.")
@@ -51,12 +37,14 @@ def pydantic_ai_integration() -> None:
     toolset = StackOneToolSet()
     tools = toolset.fetch_tools(
         actions=["workday_list_workers", "workday_get_worker", "workday_get_current_user"],
-        account_ids=[account_id],
-    )
-    pydantic_tools = [_to_pydantic_ai_tool(t) for t in tools]
-    print(f"Loaded {len(pydantic_tools)} tools for Pydantic AI.")
+        account_ids=[os.environ["STACKONE_ACCOUNT_ID"]],
+    ).to_pydantic_ai()
 
-    agent = Agent("openai:gpt-5.4", system_prompt="You are a helpful HR assistant.", tools=pydantic_tools)
+    agent = Agent(
+        "openai:gpt-5.4",
+        system_prompt="You are a helpful HR assistant.",
+        tools=tools,
+    )
     result = agent.run_sync("List the first 5 employees")
     print(f"Result:\n{result.output}")
 
diff --git a/pyproject.toml b/pyproject.toml
@@ -37,6 +37,7 @@ packages = ["stackone_ai"]
 
 [project.optional-dependencies]
 mcp = ["mcp>=1.3.0"]
+pydantic-ai = ["pydantic-ai-slim>=1.83.0,<2.0.0"]
 examples = [
   "crewai>=0.102.0",
   "langchain-openai>=1.1.14",
diff --git a/stackone_ai/models.py b/stackone_ai/models.py
@@ -6,13 +6,16 @@
 from collections.abc import Sequence
 from datetime import datetime, timezone
 from enum import Enum
-from typing import Annotated, Any, ClassVar, TypeAlias, cast
+from typing import TYPE_CHECKING, Annotated, Any, ClassVar, TypeAlias, cast
 from urllib.parse import quote
 
 import httpx
 from langchain_core.tools import BaseTool
 from pydantic import BaseModel, BeforeValidator, Field, PrivateAttr
 
+if TYPE_CHECKING:
+    from pydantic_ai.tools import Tool as PydanticAITool
+
 # Type aliases for common types
 JsonDict: TypeAlias = dict[str, Any]
 Headers: TypeAlias = dict[str, str]
@@ -467,6 +470,36 @@ def _run(self, **kwargs: Any) -> Any:
 
         return StackOneLangChainTool()
 
+    def to_pydantic_ai_tool(self) -> PydanticAITool:
+        """Convert this tool to a Pydantic AI ``Tool``.
+
+        Requires ``stackone-ai[pydantic-ai]`` (installs ``pydantic-ai-slim``).
+
+        Returns:
+            A ``pydantic_ai.tools.Tool`` ready to pass to ``Agent(tools=[...])``.
+        """
+        try:
+            from pydantic_ai.tools import Tool
+        except ImportError as e:
+            raise ImportError(
+                "Install `pydantic-ai-slim` (or `stackone-ai[pydantic-ai]`) "
+                "to use the Pydantic AI integration."
+            ) from e
+
+        openai_function = self.to_openai_function()
+        json_schema = openai_function["function"]["parameters"]
+        parent_tool = self
+
+        def implementation(**kwargs: Any) -> Any:
+            return parent_tool.execute(kwargs)
+
+        return Tool.from_schema(
+            function=implementation,
+            name=self.name,
+            description=self.description,
+            json_schema=json_schema,
+        )
+
     def set_account_id(self, account_id: str | None) -> None:
         """Set the account ID for this tool
 
@@ -577,3 +610,13 @@ def to_langchain(self) -> Sequence[BaseTool]:
             Sequence of tools in LangChain format
         """
         return [tool.to_langchain() for tool in self.tools]
+
+    def to_pydantic_ai(self) -> list[PydanticAITool]:
+        """Convert all tools to Pydantic AI ``Tool`` instances.
+
+        Requires ``stackone-ai[pydantic-ai]`` (installs ``pydantic-ai-slim``).
+
+        Returns:
+            List of ``pydantic_ai.tools.Tool`` ready to pass to ``Agent(tools=[...])``.
+        """
+        return [tool.to_pydantic_ai_tool() for tool in self.tools]
diff --git a/stackone_ai/toolset.py b/stackone_ai/toolset.py
@@ -11,7 +11,7 @@
 from collections.abc import Coroutine, Sequence
 from dataclasses import dataclass
 from importlib import metadata
-from typing import Any, Literal, TypedDict, TypeVar
+from typing import TYPE_CHECKING, Any, Literal, TypedDict, TypeVar
 
 from pydantic import BaseModel, Field, PrivateAttr, ValidationError, field_validator
 
@@ -32,6 +32,9 @@
 )
 from stackone_ai.utils.normalize import _normalize_action_name
 
+if TYPE_CHECKING:
+    from pydantic_ai.tools import Tool as PydanticAITool
+
 logger = logging.getLogger("stackone.tools")
 
 SearchMode = Literal["auto", "semantic", "local"]
@@ -752,6 +755,46 @@ def langchain(
 
         return self.fetch_tools(account_ids=effective_account_ids).to_langchain()
 
+    def pydantic_ai(
+        self,
+        *,
+        mode: Literal["search_and_execute"] | None = None,
+        account_ids: list[str] | None = None,
+    ) -> list[PydanticAITool]:
+        """Get tools as Pydantic AI ``Tool`` instances.
+
+        Args:
+            mode: Tool mode.
+                ``None`` (default): fetch all tools and convert to Pydantic AI tools.
+                ``"search_and_execute"``: return two meta tools (tool_search + tool_execute)
+                that let the LLM discover and execute tools on-demand.
+            account_ids: Account IDs to scope tools. Overrides the ``execute``
+                config from the constructor.
+
+        Returns:
+            List of Pydantic AI ``Tool`` objects ready to pass to ``Agent(tools=...)``.
+
+        Requires ``stackone-ai[pydantic-ai]`` (installs ``pydantic-ai-slim``).
+
+        Examples::
+
+            # All tools
+            toolset = StackOneToolSet()
+            tools = toolset.pydantic_ai()
+            agent = Agent("openai:gpt-5.4", tools=tools)
+
+            # Meta tools for agent-driven discovery
+            tools = toolset.pydantic_ai(mode="search_and_execute")
+        """
+        effective_account_ids = account_ids or (
+            self._execute_config.get("account_ids") if self._execute_config else None
+        )
+
+        if mode == "search_and_execute":
+            return self._build_tools(account_ids=effective_account_ids).to_pydantic_ai()
+
+        return self.fetch_tools(account_ids=effective_account_ids).to_pydantic_ai()
+
     def execute(
         self,
         tool_name: str,
diff --git a/tests/test_integrations_pydantic_ai.py b/tests/test_integrations_pydantic_ai.py
diff --git a/uv.lock b/uv.lock
