Consider first class support for Pydantic AI

[shea-parkes]

I fully understand why y'all are developing chatlas. It's nice to have a posit-y GenAI middleware in Python that mostly mirrors y'all's ellmer R package.

And I appreciate the tasteful coupling y'all have done between chatlas and shinychat.

But I just felt obligated to open this issue to at least raise the idea of leaning into Pydantic AI instead. Either in addition to chatlas or as a full replacement.

For the simplest use cases, it is pretty easy to swap pydantic_ai for chatlas and still leverage shinychat. But y'all have put in the (tasteful) bindings in here on chatlas.ContentToolResult. So anyone wishing to adopt pydantic_ai instead of chatlas will need to shim their own means of displaying tool calls. It might be nice if y'all centrally did the work of displaying pydantic_ai tool calls with shinychat.

Some potential reasons to consider this:

• Pydantic AI is getting a lot of investment and developer focus. Might as well piggy back that.
  • They're staying ahead of chatlas on a pile of things like A2A.
• Pydantic AI has much tighter bindings to FastMCP. Chatlas does obviously work with FastMCP... but not to the same extent.

[shea-parkes]

And a bit of GenAI slop from Claude Opus to smell out plausibility. Hopefully I'll get a chance to make it real in the new year:

First the bridge module:

"""Bridge module to integrate pydantic_ai tool calls with shinychat UI."""

from __future__ import annotations

import json
from typing import TYPE_CHECKING, Optional, Any
from collections. abc import AsyncIterable

from htmltools import TagList
from pydantic import BaseModel, Field
from pydantic_ai import (
    AgentStreamEvent,
    FunctionToolCallEvent,
    FunctionToolResultEvent,
    RunContext,
)

# Import the shinychat tool display components
from shinychat._chat_normalize_chatlas import (
    ToolRequestComponent,
    ToolResultComponent,
    ToolResultDisplay,
)
from shinychat._chat_normalize import message_content
from shinychat._chat_types import ChatMessage

if TYPE_CHECKING: 
    from shinychat import Chat


class PydanticAIToolRequest(BaseModel):
    """Wrapper class to represent a pydantic_ai tool request for shinychat."""
    
    id: str = Field(description="The tool call ID - must match for request/result pairing")
    name: str = Field(description="The name of the tool being called")
    arguments: dict[str, Any] = Field(description="The arguments passed to the tool")
    tool_title: Optional[str] = Field(default=None, description="Display title for the tool")
    icon: Optional[Any] = Field(default=None, description="Icon to display for the tool")
    
    model_config = {"arbitrary_types_allowed": True}
    
    def to_component(self) -> ToolRequestComponent:
        """Convert to shinychat ToolRequestComponent."""
        return ToolRequestComponent(
            request_id=self.id,  # This ID is crucial for replacement
            tool_name=self.name,
            tool_title=self.tool_title,
            arguments=json.dumps(self.arguments),
            icon=self.icon,
            intent=self.arguments.get("_intent") if isinstance(self.arguments, dict) else None,
            expanded=False,
        )


class PydanticAIToolResult(BaseModel):
    """Wrapper class to represent a pydantic_ai tool result for shinychat."""
    
    id: str = Field(description="The tool call ID - must match the request for replacement")
    request: PydanticAIToolRequest = Field(description="The original tool request")
    result: Any = Field(description="The result returned by the tool")
    error: Optional[Exception] = Field(default=None, description="Any error that occurred during tool execution")
    display: ToolResultDisplay = Field(default_factory=ToolResultDisplay, description="Display configuration for the result")
    
    model_config = {"arbitrary_types_allowed": True}
    
    def to_component(self) -> ToolResultComponent:
        """Convert to shinychat ToolResultComponent."""
        # Format the request call for display
        request_call = json.dumps(
            {
                "id": self.id,
                "name": self.request.name,
                "arguments": self.request.arguments,
            },
            indent=2,
        )
        
        # Determine value and value_type based on display settings
        if self.error is not None:
            value = str(self.error)
            value_type = "code"
        elif self.display. html is not None:
            value = self.display.html
            value_type = "html"
        elif self.display.markdown is not None:
            value = self.display.markdown
            value_type = "markdown"
        elif self. display.text is not None:
            value = self.display.text
            value_type = "text"
        else:
            value = str(self.result)
            value_type = "code"
            
        return ToolResultComponent(
            request_id=self. id,  # SAME ID as the request - this triggers replacement
            request_call=request_call,
            tool_name=self.request.name,
            tool_title=self.display.title or self.request.tool_title,
            status="success" if self.error is None else "error",
            value=value,
            value_type=value_type,
            icon=self.display. icon or self.request.icon,
            intent=self.request.arguments.get("_intent") if isinstance(self.request.arguments, dict) else None,
            show_request=self.display.show_request,
            expanded=self.display.open,
        )


# Register message content handlers for our Pydantic models
@message_content. register
def _(message: PydanticAIToolRequest) -> ChatMessage:
    return ChatMessage(content=message.to_component())


@message_content.register
def _(message: PydanticAIToolResult) -> ChatMessage:
    return ChatMessage(content=message.to_component())


class ShinyEventStreamHandler(BaseModel):
    """Event stream handler that bridges pydantic_ai events to shinychat UI."""
    
    chat: Any = Field(description="The shinychat Chat instance to update")
    display_config: dict[str, ToolResultDisplay] = Field(
        default_factory=dict,
        description="Mapping of tool names to ToolResultDisplay configurations"
    )
    active_tool_calls: dict[str, PydanticAIToolRequest] = Field(
        default_factory=dict,
        description="Currently active tool calls being processed"
    )
    
    model_config = {"arbitrary_types_allowed": True}
    
    async def __call__(
        self,
        ctx: RunContext,
        event_stream: AsyncIterable[AgentStreamEvent],
    ):
        """Process the event stream and update the chat UI."""
        async for event in event_stream: 
            await self.handle_event(event)
            
    async def handle_event(self, event: AgentStreamEvent):
        """Handle a single event from the stream."""
        
        if isinstance(event, FunctionToolCallEvent):
            # Create and display the tool request
            # The event already has fully accumulated arguments
            request = PydanticAIToolRequest(
                id=event. part.tool_call_id,
                name=event.part. tool_name,
                arguments=event.part.args,  # Already complete! 
                # You might want to extract title/icon from tool metadata if available
            )
            self.active_tool_calls[event.part.tool_call_id] = request
            
            # Append the tool request to the chat
            # This will show the "pending" state with the request_id
            await self.chat. append_message(request)
            
        elif isinstance(event, FunctionToolResultEvent):
            # Get the corresponding request
            request = self.active_tool_calls.get(event.tool_call_id)
            if request:
                # Get display config for this tool if available
                display = self.display_config.get(request.name, ToolResultDisplay())
                
                # Create and display the tool result
                # CRITICAL: The id MUST match the request id for replacement to work
                result = PydanticAIToolResult(
                    id=event.tool_call_id,  # Same ID as request - triggers replacement! 
                    request=request,
                    result=event.result. content,
                    error=None,  # TODO: Check if pydantic_ai provides error info
                    display=display,
                )
                
                # Append the tool result to the chat
                # The frontend will see the matching request_id and REPLACE the request
                await self.chat.append_message(result)
                
                # Clean up
                del self.active_tool_calls[event.tool_call_id]
        
        # We can safely ignore these events for tool display purposes: 
        # - PartStartEvent: Just signals a new part is starting
        # - PartDeltaEvent with ToolCallPartDelta: Incremental argument building
        # - TextPartDelta, ThinkingPartDelta:  Handled by stream_text()
        # - FinalResultEvent: Just signals the final result is starting


# Convenience function to create a configured event stream handler
def create_shinychat_stream_handler(
    chat: "Chat",
    tool_displays: Optional[dict[str, dict]] = None
) -> ShinyEventStreamHandler: 
    """
    Create an event stream handler for use with pydantic_ai agents.
    
    The handler will automatically replace pending tool requests with their results
    when they complete, matching the behavior of chatlas integration. 
    
    Args:
        chat: The shinychat Chat instance
        tool_displays: Optional dict mapping tool names to display configs
                      Each config can have:  title, icon, show_request, open, html, markdown, text
    
    Returns:
        ShinyEventStreamHandler instance
        
    Example:
        ```python
        from shinychat import Chat
        from pydantic_ai import Agent
        
        chat = Chat()
        agent = Agent(...)
        
        handler = create_shinychat_stream_handler(
            chat,
            tool_displays={
                "weather_forecast": {
                    "title":  "Weather Forecast",
                    "icon": "☀️",
                    "open": True,
                }
            }
        )
        
        async with agent.run_stream(prompt, event_stream_handler=handler) as run:
            async for text in run.stream_text():
                await chat.append_message_stream(text)
        ```
    """
    display_config = {}
    if tool_displays:
        for tool_name, config in tool_displays.items():
            display_config[tool_name] = ToolResultDisplay(**config)
            
    return ShinyEventStreamHandler(chat=chat, display_config=display_config)

And the corresponding example code:

"""Example of using pydantic_ai with shinychat UI."""

import asyncio
from datetime import date
from shiny import App, ui
from shinychat import Chat
from pydantic_ai import Agent, RunContext

# Import the bridge module
from pydantic_ai_shinychat_bridge import create_shinychat_stream_handler

# Create your pydantic_ai agent
weather_agent = Agent(
    'openai: gpt-4',
    system_prompt='Providing weather forecasts at the locations the user provides.',
)

@weather_agent.tool
async def weather_forecast(
    ctx: RunContext,
    location:  str,
    forecast_date:  date,
) -> str:
    """Get the weather forecast for a location."""
    # Simulate some processing time
    await asyncio.sleep(1)
    return f'The forecast in {location} on {forecast_date} is 24°C and sunny.'

@weather_agent.tool
async def get_temperature(
    ctx: RunContext,
    location: str,
) -> str:
    """Get current temperature for a location."""
    return f'Current temperature in {location} is 22°C.'

# Create Shiny app
app_ui = ui.page_fluid(
    ui.h2("Weather Assistant with Tool Display"),
    Chat(id="chat", messages=[]),
)

def server(input, output, session):
    chat = Chat. get_instance("chat")
    
    @chat.on_user_submit
    async def handle_user_message():
        user_msg = chat.user_input()
        
        # Create the event stream handler with custom tool displays
        handler = create_shinychat_stream_handler(
            chat,
            tool_displays={
                "weather_forecast": {
                    "title":  "📅 Weather Forecast",
                    "icon": "☀️",
                    "open": True,
                },
                "get_temperature": {
                    "title": "🌡️ Current Temperature",
                    "icon":  "🌡️",
                    "markdown": "Checking temperature sensors...",
                    "open": True,
                }
            }
        )
        
        # Run the agent with the stream handler
        try:
            async with weather_agent.run_stream(
                user_msg, 
                event_stream_handler=handler
            ) as run:
                # Stream the text response
                message_id = await chat.append_message_stream()
                async for text_chunk in run.stream_text():
                    await chat.append_message_stream(text_chunk, message_id=message_id)
                    
        except Exception as e:
            await chat.append_message(f"Error: {str(e)}")

app = App(app_ui, server)

[cpsievert]

Hi @shea-parkes, thanks for sharing. Just FYI, we do have a collection of templates on how to use Pydantic AI, LangChain, etc. with shinychat. Those templates are here https://github.com/posit-dev/py-shiny-templates/tree/main/gen-ai/interactive-chat-explorer

You will notice, however, that we don't currently have tool displays in those templates -- I would be very welcoming of a PR if you find an abstraction that you're happy with!

This also reminds me -- we had plans to get the templates that I linked to up on https://shiny.posit.co/py/templates. Any idea if we have any wip for that @karangattu?

[shea-parkes]

Alright, had a bit of quiet time on this Friday. Have a first proof of concept working for tool calls. Still some todos (e.g. supporting errors), but I think the remaining work isn't that crazy.

I'm focusing just on the OpenAI case. I'm not sure how much the pydantic_ai event structures change with different providers.

The much shrunk "bridge" code:

"""
Bridge module to integrate pydantic_ai with shinychat

Focuses on streaming use case, with tool call integration
"""

import json
import typing

import pydantic
import pydantic_ai
from shinychat._chat_normalize import message_content_chunk
from shinychat._chat_normalize_chatlas import ToolRequestComponent
from shinychat._chat_normalize_chatlas import ToolResultComponent
from shinychat._chat_normalize_chatlas import ToolResultDisplay
from shinychat._chat_types import ChatMessage


class PydanticAIToolRequest(pydantic.BaseModel):
    """Wrapper class to represent a pydantic_ai tool request for shinychat."""

    tool_call_id: str = pydantic.Field(
        description="The tool call ID - must match for request/result pairing"
    )
    name: str = pydantic.Field(description="The name of the tool being called")
    arguments: dict[str, typing.Any] = pydantic.Field(
        description="The arguments passed to the tool"
    )
    tool_title: str | None = pydantic.Field(
        default=None, description="Display title for the tool"
    )
    icon: typing.Any | None = pydantic.Field(
        default=None, description="Icon to display for the tool"
    )

    def to_component(self) -> ToolRequestComponent:
        """Convert to shinychat.ToolRequestComponent."""
        return ToolRequestComponent(
            request_id=self.tool_call_id,
            tool_name=self.name,
            tool_title=self.tool_title,
            arguments=json.dumps(self.arguments, indent=2),
            icon=self.icon,
            intent=None,
            expanded=False,
        )


class PydanticAIToolResult(pydantic.BaseModel):
    """Wrapper class to represent a pydantic_ai tool result for shinychat."""

    tool_call_id: str = pydantic.Field(
        description="The tool call ID - must match the request for replacement"
    )
    request: PydanticAIToolRequest = pydantic.Field(
        description="The original tool request"
    )
    result: typing.Any = pydantic.Field(description="The result returned by the tool")
    error: typing.Any | None = pydantic.Field(
        default=None,
        description="Any error that occurred during tool execution",
    )
    display: ToolResultDisplay = pydantic.Field(
        default_factory=ToolResultDisplay,
        description="Display configuration for the result",
    )

    def to_component(self) -> ToolResultComponent:
        """Convert to shinychat ToolResultComponent."""
        # Format the request call for display
        request_call = json.dumps(
            {
                "id": self.request.tool_call_id,
                "name": self.request.name,
                "arguments": self.request.arguments,
            },
            indent=2,
        )

        # Determine value and value_type based on display settings
        if self.error is not None:
            value = str(self.error)
            value_type = "code"
        elif self.display.html is not None:
            value = self.display.html
            value_type = "html"
        elif self.display.markdown is not None:
            value = self.display.markdown
            value_type = "markdown"
        elif self.display.text is not None:
            value = self.display.text
            value_type = "text"
        else:
            value = str(self.result)
            value_type = "code"

        return ToolResultComponent(
            request_id=self.tool_call_id,
            request_call=request_call,
            tool_name=self.request.name,
            tool_title=self.display.title or self.request.tool_title,
            status="success" if self.error is None else "error",
            value=value,
            value_type=value_type,
            icon=self.display.icon or self.request.icon,
            intent=None,
            show_request=self.display.show_request,
            expanded=self.display.open,
        )


# Register new classes with message_content_chunk single dispatcher
@message_content_chunk.register
def _(chunk: PydanticAIToolRequest) -> ChatMessage:
    return ChatMessage(content=chunk.to_component())


@message_content_chunk.register
def _(chunk: PydanticAIToolResult) -> ChatMessage:
    return ChatMessage(content=chunk.to_component())


async def transform_events_for_shinychat(event_stream):
    """Transform pydantic_ai.Agent.run_event_stream() for shinychat"""
    tool_requests: dict[str, PydanticAIToolRequest] = {}
    async for event in event_stream:
        if isinstance(event, pydantic_ai.FunctionToolCallEvent):
            assert isinstance(event.part.args, str)
            request = PydanticAIToolRequest(
                tool_call_id=event.part.tool_call_id,
                name=event.part.tool_name,
                arguments=json.loads(event.part.args),
            )
            tool_requests[request.tool_call_id] = request
            yield request
        elif isinstance(event, pydantic_ai.FunctionToolResultEvent):
            yield PydanticAIToolResult(
                tool_call_id=event.result.tool_call_id,
                request=tool_requests[event.result.tool_call_id],
                result=event.result.content,
                error=None,  # TODO: implement error displaying
            )
        elif isinstance(event, pydantic_ai.PartStartEvent):
            if isinstance(event.part, pydantic_ai.TextPart):
                yield event.part.content
        elif isinstance(event, pydantic_ai.PartDeltaEvent):
            if isinstance(event.delta, pydantic_ai.TextPartDelta):
                yield event.delta.content_delta

Example usage:

def server(input, output, session):
    ui_chat = ui.Chat(id="mvaa_chat")

    @ui_chat.on_user_submit
    async def _(user_input: str):
        """Handle a submission"""
        await ui_chat.append_message_stream(
            pydantic_chat.transform_events_for_shinychat(
                agent.run_stream_events(user_input)
            )
        )

It does seem pretty promising. And pydantic_ai has a lot more functionality than chatlas, so I'll probably continue with hardening it for my own use case. I may get to the point of a pull request here, but I'm not sure how robust y'all would like the support.

[shea-parkes]

Got some time in this weekend to slim this down a lot. Dropped all the shimming classes and just used ToolRequestComponent and ToolResultComponent directly. Added some "meh" formatting for thinking outputs.

I'm not really sure how OpenAI-centric this is, I haven't tested it on other providers. Probably won't for quite some time.

Probably wouldn't include get_tool_titles() stuff in a PR into here (unless some maintainer says otherwise). Similarly might want/need to drop the parts about capturing the remote state ID. I'm using that to leverage OpenAI's responses API caching.

Let me know if this is looking like something y'all would want to accept into here and I can try to get a plausible PR up. Thanks.

"""
Bridge module to integrate pydantic_ai with shinychat

Focuses on streaming use case, with tool call integration
Aims to support tools defined via FastMCP or Pydantic AI
"""

import contextlib
import json
import typing

import pydantic
import pydantic_ai
import shiny
from pydantic_ai.models.test import TestModel
from shinychat._chat_normalize_chatlas import ToolRequestComponent
from shinychat._chat_normalize_chatlas import ToolResultComponent
from shinychat._chat_types import ChatMessageDict


async def get_tool_titles(agent) -> dict[str, str]:
    """Get all tool titles from this agent"""
    titles: dict[str, str] = {}
    _tools = await agent._get_toolset().get_tools(
        pydantic_ai.RunContext(
            deps=None,
            model=TestModel(),
            usage=pydantic_ai.RunUsage(),
        )
    )
    for tool_name, tool in _tools.items():
        with contextlib.suppress(AttributeError, KeyError, TypeError):
            titles[tool_name] = tool.tool_def.metadata["annotations"]["title"]  # type: ignore
    return titles


def transform_message_for_pydantic(
    message: ChatMessageDict,
) -> pydantic_ai.ModelMessage:
    """Transform a shinychat.ChatMessageDict to a pydantic_ai.ModelMessage"""
    role = message.get("role")
    content = message.get("content", "")
    if role == "user":
        return pydantic_ai.ModelRequest(
            parts=[pydantic_ai.UserPromptPart(content=content)],
        )
    if role == "assistant":
        return pydantic_ai.ModelResponse(
            parts=[pydantic_ai.TextPart(content=content)],
        )
    raise ValueError(f"Unsupported role of {role}")


class EventTransformer(pydantic.BaseModel):
    """Helper class that transforms pydantic_ai.Agent.run_event_stream() for shinychat"""

    model_config = pydantic.ConfigDict(arbitrary_types_allowed=True)

    tool_titles: typing.Annotated[
        dict[str, str],
        pydantic.Field(description="A mapping of user-facing Titles for tools"),
    ]
    remote_state_id: typing.Annotated[
        shiny.reactive.Value[str | None],
        pydantic.Field(
            description="A reactive value that contains the remote state_id"
        ),
    ]

    async def transform_events_for_shinychat( 
        self,
        event_stream,
        *,
        update_remote_state_id: bool = True,
        include_tool_calls: bool = True,
        include_thinking: bool = True,
    ):
        """Transform pydantic_ai.Agent.run_event_stream() for shinychat"""
        tool_requests: dict[str, ToolRequestComponent] = {}
        async for event in event_stream:
            if (
                isinstance(event, pydantic_ai.FunctionToolCallEvent)
                and include_tool_calls
            ):
                assert isinstance(event.part.args, str)
                arguments = json.loads(event.part.args)
                request = ToolRequestComponent(
                    request_id=event.part.tool_call_id,
                    tool_name=event.part.tool_name,
                    tool_title=self.tool_titles.get(event.part.tool_name),
                    arguments=json.dumps(arguments, indent=2),
                    icon=None,
                    intent=arguments.get("_intent"),
                    expanded=False,
                )
                tool_requests[request.request_id] = request
                yield request
            elif (
                isinstance(event, pydantic_ai.FunctionToolResultEvent)
                and include_tool_calls
            ):
                failure = isinstance(event.result, pydantic_ai.RetryPromptPart) or (
                    event.result.metadata is not None
                    and event.result.metadata.get("is_error", False)
                )
                request = tool_requests[event.result.tool_call_id]
                arguments = json.loads(request.arguments)
                request_call = json.dumps(
                    {
                        "tool_call_id": request.request_id,
                        "tool_name": request.tool_name,
                        "arguments": arguments,
                    },
                    indent=2,
                )
                yield ToolResultComponent(
                    request_id=event.result.tool_call_id,
                    request_call=request_call,
                    tool_name=request.tool_name,
                    tool_title=request.tool_title,
                    status="success" if not failure else "error",
                    value=json.dumps(event.result.content, indent=2),
                    value_type="code",
                    icon=None,
                    intent=arguments.get("_intent"),
                    show_request=True,
                    expanded=False,
                )
            elif isinstance(event, pydantic_ai.PartStartEvent):
                if isinstance(event.part, pydantic_ai.TextPart):
                    yield "\n\n" + event.part.content
                elif (
                    isinstance(event.part, pydantic_ai.ThinkingPart)
                    and include_thinking
                ):
                    yield "\n\n> ***Thought Process***: " + event.part.content.replace(
                        "\n", "\n> "
                    )
            elif isinstance(event, pydantic_ai.PartDeltaEvent):
                if isinstance(event.delta, pydantic_ai.TextPartDelta):
                    yield event.delta.content_delta
                elif (
                    isinstance(event.delta, pydantic_ai.ThinkingPartDelta)
                    and include_thinking
                    and event.delta.content_delta is not None
                ):
                    yield event.delta.content_delta.replace("\n", "\n> ")
            elif (
                isinstance(event, pydantic_ai.AgentRunResultEvent)
                and isinstance(event.result, pydantic_ai.AgentRunResult)
                and update_remote_state_id
            ):
                self.remote_state_id.set(event.result.response.provider_response_id)

[shea-parkes]

And I see work is already going on a better Thinking support with #76 and #167 . That will be good to adapt to.