feat: add PydanticAI agents support (#755)

Co-authored-by: GlockPL <[email protected]>

Author: akotyla

packages/ragbits-agents/CHANGELOG.md

@@ -2,6 +2,7 @@
 
 ## Unreleased
 - Add tool_choice parameter to agent interface (#738)
+- Add PydanticAI agents support (#755)
 
 ## 1.2.1 (2025-08-04)
 
@@ -16,6 +17,7 @@
 - ragbits-core updated to version v1.2.0
 - Add native openai tools support (#621)
 - add Context to Agents (#715)
+- add PydanticAI agnts support (#755)
 
 ## 1.1.0 (2025-07-09)
 

packages/ragbits-agents/src/ragbits/agents/_main.py

@@ -3,11 +3,13 @@
 from contextlib import suppress
 from copy import deepcopy
 from dataclasses import dataclass
+from datetime import timedelta
 from inspect import iscoroutinefunction
 from types import ModuleType, SimpleNamespace
 from typing import ClassVar, Generic, cast, overload
 
 from pydantic import BaseModel, Field
+from typing_extensions import Self
 
 from ragbits import agents
 from ragbits.agents.exceptions import (
@@ -20,7 +22,7 @@
     AgentToolNotAvailableError,
     AgentToolNotSupportedError,
 )
-from ragbits.agents.mcp.server import MCPServer
+from ragbits.agents.mcp.server import MCPServer, MCPServerStdio, MCPServerStreamableHttp
 from ragbits.agents.mcp.utils import get_tools
 from ragbits.agents.tool import Tool, ToolCallResult, ToolChoice
 from ragbits.core.audit.traces import trace
@@ -34,6 +36,10 @@
 
 with suppress(ImportError):
     from a2a.types import AgentCapabilities, AgentCard, AgentSkill
+    from pydantic_ai import Agent as PydanticAIAgent
+    from pydantic_ai import mcp
+
+    from ragbits.core.llms import LiteLLM
 
 
 @dataclass
@@ -579,3 +585,106 @@ async def _extract_agent_skills(self) -> list["AgentSkill"]:
             )
             for tool in all_tools.values()
         ]
+
+    @requires_dependencies("pydantic_ai")
+    def to_pydantic_ai(self) -> "PydanticAIAgent":
+        """
+        Convert ragbits agent instance into a `pydantic_ai.Agent` representation.
+
+        Returns:
+            PydanticAIAgent: The equivalent Pydantic-based agent configuration.
+
+        Raises:
+            ValueError: If the `prompt` is not a string or a `Prompt` instance.
+        """
+        mcp_servers: list[mcp.MCPServerStdio | mcp.MCPServerHTTP] = []
+
+        if not self.prompt:
+            raise ValueError("Prompt is required but was None.")
+
+        if isinstance(self.prompt, str):
+            system_prompt = self.prompt
+        else:
+            if not self.prompt.system_prompt:
+                raise ValueError("System prompt is required but was None.")
+            system_prompt = self.prompt.system_prompt
+
+        for mcp_server in self.mcp_servers:
+            if isinstance(mcp_server, MCPServerStdio):
+                mcp_servers.append(
+                    mcp.MCPServerStdio(
+                        command=mcp_server.params.command, args=mcp_server.params.args, env=mcp_server.params.env
+                    )
+                )
+            elif isinstance(mcp_server, MCPServerStreamableHttp):
+                timeout = mcp_server.params["timeout"]
+                sse_timeout = mcp_server.params["sse_read_timeout"]
+
+                mcp_servers.append(
+                    mcp.MCPServerHTTP(
+                        url=mcp_server.params["url"],
+                        headers=mcp_server.params["headers"],
+                        timeout=timeout.total_seconds() if isinstance(timeout, timedelta) else timeout,
+                        sse_read_timeout=sse_timeout.total_seconds()
+                        if isinstance(sse_timeout, timedelta)
+                        else sse_timeout,
+                    )
+                )
+        return PydanticAIAgent(
+            model=self.llm.model_name,
+            system_prompt=system_prompt,
+            tools=[tool.to_pydantic_ai() for tool in self.tools],
+            mcp_servers=mcp_servers,
+        )
+
+    @classmethod
+    @requires_dependencies("pydantic_ai")
+    def from_pydantic_ai(cls, pydantic_ai_agent: "PydanticAIAgent") -> Self:
+        """
+        Construct an agent instance from a `pydantic_ai.Agent` representation.
+
+        Args:
+            pydantic_ai_agent: A Pydantic-based agent configuration.
+
+        Returns:
+            An instance of the agent class initialized from the Pydantic representation.
+        """
+        mcp_servers: list[MCPServerStdio | MCPServerStreamableHttp] = []
+        for mcp_server in pydantic_ai_agent._mcp_servers:
+            if isinstance(mcp_server, mcp.MCPServerStdio):
+                mcp_servers.append(
+                    MCPServerStdio(
+                        params={
+                            "command": mcp_server.command,
+                            "args": list(mcp_server.args),
+                            "env": mcp_server.env or {},
+                        }
+                    )
+                )
+            elif isinstance(mcp_server, mcp.MCPServerHTTP):
+                headers = mcp_server.headers or {}
+
+                mcp_servers.append(
+                    MCPServerStreamableHttp(
+                        params={
+                            "url": mcp_server.url,
+                            "headers": {str(k): str(v) for k, v in headers.items()},
+                            "sse_read_timeout": mcp_server.sse_read_timeout,
+                            "timeout": mcp_server.timeout,
+                        }
+                    )
+                )
+
+        if not pydantic_ai_agent.model:
+            raise ValueError("Missing LLM in `pydantic_ai.Agent` instance")
+        elif isinstance(pydantic_ai_agent.model, str):
+            model_name = pydantic_ai_agent.model
+        else:
+            model_name = pydantic_ai_agent.model.model_name
+
+        return cls(
+            llm=LiteLLM(model_name=model_name),  # type: ignore[arg-type]
+            prompt="\n".join(pydantic_ai_agent._system_prompts),
+            tools=[tool.function for _, tool in pydantic_ai_agent._function_tools.items()],
+            mcp_servers=cast(list[MCPServer], mcp_servers),
+        )

packages/ragbits-agents/src/ragbits/agents/tool.py

@@ -1,11 +1,16 @@
 from collections.abc import Callable
+from contextlib import suppress
 from dataclasses import dataclass
 from typing import Any, Literal
 
 from typing_extensions import Self
 
+from ragbits.core.utils.decorators import requires_dependencies
 from ragbits.core.utils.function_schema import convert_function_to_function_schema, get_context_variable_name
 
+with suppress(ImportError):
+    from pydantic_ai import Tool as PydanticAITool
+
 
 @dataclass
 class ToolCallResult:
@@ -77,5 +82,19 @@ def to_function_schema(self) -> dict[str, Any]:
             },
         }
 
+    @requires_dependencies("pydantic_ai")
+    def to_pydantic_ai(self) -> "PydanticAITool":
+        """
+        Convert ragbits tool to a Pydantic AI Tool.
+
+        Returns:
+            A `pydantic_ai.tools.Tool` object.
+        """
+        return PydanticAITool(
+            function=self.on_tool_call,
+            name=self.name,
+            description=self.description,
+        )
+
 
 ToolChoice = Literal["auto", "none", "required"] | Callable