From 96c8939761b3d4008573a66012dd70e1a37e4c59 Mon Sep 17 00:00:00 2001 From: safina57 Date: Sun, 19 Oct 2025 19:06:56 +0100 Subject: [PATCH] Add OpenAI ChatKit integration with request and response types --- .../pydantic_ai/ui/openai_chatkit/__init__.py | 70 +++++++++ .../pydantic_ai/ui/openai_chatkit/_adapter.py | 141 ++++++++++++++++++ .../ui/openai_chatkit/_request_types.py | 112 ++++++++++++++ .../ui/openai_chatkit/_response_types.py | 139 +++++++++++++++++ 4 files changed, 462 insertions(+) create mode 100644 pydantic_ai_slim/pydantic_ai/ui/openai_chatkit/__init__.py create mode 100644 pydantic_ai_slim/pydantic_ai/ui/openai_chatkit/_adapter.py create mode 100644 pydantic_ai_slim/pydantic_ai/ui/openai_chatkit/_request_types.py create mode 100644 pydantic_ai_slim/pydantic_ai/ui/openai_chatkit/_response_types.py diff --git a/pydantic_ai_slim/pydantic_ai/ui/openai_chatkit/__init__.py b/pydantic_ai_slim/pydantic_ai/ui/openai_chatkit/__init__.py new file mode 100644 index 0000000000..cdd28ade10 --- /dev/null +++ b/pydantic_ai_slim/pydantic_ai/ui/openai_chatkit/__init__.py @@ -0,0 +1,70 @@ +"""OpenAI ChatKit integration for Pydantic AI. + +This module provides integration between Pydantic AI agents and OpenAI's ChatKit UI framework. +ChatKit is a framework for building conversational AI applications with rich UI components. + +Key components: +- ChatKitAdapter: Main adapter for handling ChatKit requests +- ChatKitEventStream: Event stream for converting Pydantic AI events to ChatKit format +- Request/Response types: Official ChatKit types imported from the `chatkit` library + +For more details, see the ChatKit documentation: +https://platform.openai.com/docs/guides/chatkit +""" + +from ._adapter import ChatKitAdapter +from ._event_stream import ChatKitEventStream +from ._request_types import ( + Attachment, + ChatKitReq, + NonStreamingReq, + StreamingReq, + ThreadsAddUserMessageReq, + ThreadsCreateReq, + UserMessageInput, + is_streaming_req, +) +from ._response_types import ( + AssistantMessageContent, + AssistantMessageItem, + ErrorEvent, + ProgressUpdateEvent, + Thread, + ThreadCreatedEvent, + ThreadItem, + ThreadItemAddedEvent, + ThreadItemDoneEvent, + ThreadMetadata, + ThreadStreamEvent, + ThreadUpdatedEvent, + UserMessageItem, +) + +__all__ = [ + # Main classes + 'ChatKitAdapter', + 'ChatKitEventStream', + # Request types + 'ChatKitReq', + 'StreamingReq', + 'NonStreamingReq', + 'ThreadsCreateReq', + 'ThreadsAddUserMessageReq', + 'UserMessageInput', + 'Attachment', + 'is_streaming_req', + # Response types + 'ThreadStreamEvent', + 'ThreadCreatedEvent', + 'ThreadUpdatedEvent', + 'ThreadItemAddedEvent', + 'ThreadItemDoneEvent', + 'ProgressUpdateEvent', + 'ErrorEvent', + 'Thread', + 'ThreadMetadata', + 'ThreadItem', + 'AssistantMessageItem', + 'UserMessageItem', + 'AssistantMessageContent', +] diff --git a/pydantic_ai_slim/pydantic_ai/ui/openai_chatkit/_adapter.py b/pydantic_ai_slim/pydantic_ai/ui/openai_chatkit/_adapter.py new file mode 100644 index 0000000000..12b745612a --- /dev/null +++ b/pydantic_ai_slim/pydantic_ai/ui/openai_chatkit/_adapter.py @@ -0,0 +1,141 @@ +"""OpenAI ChatKit adapter for handling requests. + +This adapter integrates Pydantic AI agents with OpenAI's ChatKit UI framework. + +1. Thread-based conversations: ChatKit works with persistent threads that contain message history. + +2. Single endpoint: All communication happens through one POST endpoint that returns either JSON directly or streams SSE JSON events. + +3. Rich UI components: Supports widgets, progress updates, and client-side tools beyond just text and tool calls. + +4. Multiple formats: Supports both original ChatKit format and newer threads.create format. +""" + +from __future__ import annotations + +from collections.abc import Mapping, Sequence +from dataclasses import dataclass +from functools import cached_property +from typing import TYPE_CHECKING, Any + +from ...agent import AgentDepsT +from ...messages import ( + AudioUrl, + DocumentUrl, + ImageUrl, + ModelMessage, + ModelRequest, + ModelRequestPart, + UserPromptPart, + VideoUrl, +) +from ...output import OutputDataT +from ..adapter import BaseAdapter +from ..event_stream import BaseEventStream +from ._event_stream import ChatKitEventStream + +# Import ChatKit types from the proper modules +from ._request_types import ( + ChatKitReq, + FileAttachment, + ImageAttachment, + ThreadItem, + UserMessageTagContent, + UserMessageTextContent, + request_data_ta, +) +from ._response_types import ( + ThreadStreamEvent, + UserMessageItem, +) + +if TYPE_CHECKING: + try: + from starlette.requests import Request + except ImportError: + pass + +__all__ = ['ChatKitAdapter'] + + +@dataclass +class ChatKitAdapter(BaseAdapter[ChatKitReq, UserMessageItem, ThreadStreamEvent, AgentDepsT, OutputDataT]): + """ChatKit adapter for integrating Pydantic AI agents with ChatKit UI. + + This adapter handles the translation between ChatKit's thread-based protocol + and Pydantic AI's message-based system. + """ + + @classmethod + async def validate_request(cls, request: Request) -> ChatKitReq: + """Validate a ChatKit request, supporting multiple formats.""" + return request_data_ta.validate_json(await request.body()) + + @classmethod + def load_messages(cls, messages: Sequence[ThreadItem]) -> list[ModelMessage]: + """Convert ChatKit UserMessageItem objects to Pydantic AI ModelMessage format.""" + result: list[ModelMessage] = [] + request_parts: list[ModelRequestPart] | None = None + + for item in messages: + # User Messages + if hasattr(item, 'type') and item.type == 'user_message': + if request_parts is None: + request_parts = [] + result.append(ModelRequest(parts=request_parts)) + + # Process content parts + for content in item.content: + if isinstance(content, UserMessageTextContent): + request_parts.append(UserPromptPart(content=content.text)) + elif isinstance(content, UserMessageTagContent): + # For tag content, we'll treat it as text with the tag text + request_parts.append(UserPromptPart(content=content.text)) + + # Process attachments + for attachment in item.attachments: + if isinstance(attachment, FileAttachment): + if attachment.upload_url: + # Determine content type based on mime_type + media_type_prefix = attachment.mime_type.split('/', 1)[0] + match media_type_prefix: + case 'image': + file = ImageUrl(url=str(attachment.upload_url), media_type=attachment.mime_type) + case 'video': + file = VideoUrl(url=str(attachment.upload_url), media_type=attachment.mime_type) + case 'audio': + file = AudioUrl(url=str(attachment.upload_url), media_type=attachment.mime_type) + case _: + file = DocumentUrl(url=str(attachment.upload_url), media_type=attachment.mime_type) + request_parts.append(UserPromptPart(content=[file])) + elif isinstance(attachment, ImageAttachment): + if attachment.upload_url: + # Use the upload URL for the image + file = ImageUrl(url=str(attachment.upload_url), media_type=attachment.mime_type) + request_parts.append(UserPromptPart(content=[file])) + + return result + + def dump_messages(self, messages: Sequence[ModelMessage]) -> list[UserMessageItem]: + """Convert Pydantic AI ModelMessage objects to ChatKit UserMessageItem format.""" + raise NotImplementedError + + @property + def event_stream(self) -> BaseEventStream[ChatKitReq, ThreadStreamEvent, AgentDepsT, OutputDataT]: + """Create the event stream handler for this adapter.""" + return ChatKitEventStream(self.request) + + @cached_property + def messages(self) -> list[ModelMessage]: + """Convert the current request's user message to Pydantic AI format.""" + raise NotImplementedError + + @cached_property + def state(self) -> dict[str, Any] | None: + """Extract state from the ChatKit thread metadata.""" + pass + + @property + def response_headers(self) -> Mapping[str, str] | None: + """Get HTTP response headers for ChatKit compatibility.""" + pass diff --git a/pydantic_ai_slim/pydantic_ai/ui/openai_chatkit/_request_types.py b/pydantic_ai_slim/pydantic_ai/ui/openai_chatkit/_request_types.py new file mode 100644 index 0000000000..1ca3061a32 --- /dev/null +++ b/pydantic_ai_slim/pydantic_ai/ui/openai_chatkit/_request_types.py @@ -0,0 +1,112 @@ +"""OpenAI ChatKit request types. + +This module provides ChatKit-compatible request types. It attempts to import from the +official `chatkit` library. + +For the complete type definitions, see: +https://github.com/openai/chatkit-python/blob/main/chatkit/types.py +""" + +# Try to import from official library first +from chatkit.types import ( + Attachment, + AttachmentBase, + AttachmentCreateParams, + AttachmentDeleteParams, + AttachmentsCreateReq, + AttachmentsDeleteReq, + BaseReq, + ChatKitReq, + FeedbackKind, + FileAttachment, + ImageAttachment, + InferenceOptions, + ItemFeedbackParams, + ItemsFeedbackReq, + ItemsListParams, + ItemsListReq, + NonStreamingReq, + StreamingReq, + ThreadAddClientToolOutputParams, + ThreadAddUserMessageParams, + ThreadCreateParams, + ThreadCustomActionParams, + ThreadDeleteParams, + ThreadGetByIdParams, + ThreadItem, + ThreadListParams, + ThreadRetryAfterItemParams, + ThreadsAddClientToolOutputReq, + ThreadsAddUserMessageReq, + ThreadsCreateReq, + ThreadsCustomActionReq, + ThreadsDeleteReq, + ThreadsGetByIdReq, + ThreadsListReq, + ThreadsRetryAfterItemReq, + ThreadsUpdateReq, + ThreadUpdateParams, + ToolChoice, + UserMessageContent, + UserMessageInput, + UserMessageTagContent, + UserMessageTextContent, + is_streaming_req, +) +from pydantic import TypeAdapter + +request_data_ta: TypeAdapter[ChatKitReq] = TypeAdapter(ChatKitReq) + +__all__ = [ + # Base request types + 'BaseReq', + 'ChatKitReq', + 'StreamingReq', + 'NonStreamingReq', + 'is_streaming_req', + # Specific request types + 'ThreadItem', + 'ThreadsCreateReq', + 'ThreadsGetByIdReq', + 'ThreadsListReq', + 'ThreadsAddUserMessageReq', + 'ThreadsAddClientToolOutputReq', + 'ThreadsCustomActionReq', + 'ThreadsRetryAfterItemReq', + 'ItemsFeedbackReq', + 'AttachmentsCreateReq', + 'AttachmentsDeleteReq', + 'ItemsListReq', + 'ThreadsUpdateReq', + 'ThreadsDeleteReq', + # Parameter types + 'ThreadCreateParams', + 'ThreadGetByIdParams', + 'ThreadListParams', + 'ThreadAddUserMessageParams', + 'ThreadAddClientToolOutputParams', + 'ThreadCustomActionParams', + 'ThreadRetryAfterItemParams', + 'ItemFeedbackParams', + 'AttachmentCreateParams', + 'AttachmentDeleteParams', + 'ItemsListParams', + 'ThreadUpdateParams', + 'ThreadDeleteParams', + # User message types + 'UserMessageInput', + 'UserMessageContent', + 'UserMessageTextContent', + 'UserMessageTagContent', + # Tool and inference types + 'InferenceOptions', + 'ToolChoice', + # Attachment types + 'Attachment', + 'FileAttachment', + 'ImageAttachment', + 'AttachmentBase', + # Misc types + 'FeedbackKind', + 'request_data_ta', +] diff --git a/pydantic_ai_slim/pydantic_ai/ui/openai_chatkit/_response_types.py b/pydantic_ai_slim/pydantic_ai/ui/openai_chatkit/_response_types.py new file mode 100644 index 0000000000..d228570bc9 --- /dev/null +++ b/pydantic_ai_slim/pydantic_ai/ui/openai_chatkit/_response_types.py @@ -0,0 +1,139 @@ +"""OpenAI ChatKit response types (SSE events). + +This module provides ChatKit-compatible response types. It attempts to import from the official `chatkit` library + +For the complete type definitions, see: +https://github.com/openai/chatkit-python/blob/main/chatkit/types.py +""" + +from chatkit.actions import Action +from chatkit.errors import ErrorCode + +# Try to import from official library first +from chatkit.types import ( + ActiveStatus, + Annotation, + AssistantMessageContent, + AssistantMessageContentPartAdded, + AssistantMessageContentPartAnnotationAdded, + AssistantMessageContentPartDone, + AssistantMessageContentPartTextDelta, + AssistantMessageItem, + BaseTask, + ClientToolCallItem, + ClosedStatus, + CustomSummary, + CustomTask, + DurationSummary, + EndOfTurnItem, + EntitySource, + ErrorEvent, + FileSource, + FileTask, + HiddenContextItem, + IconName, + ImageTask, + LockedStatus, + NoticeEvent, + Page, + ProgressUpdateEvent, + SearchTask, + Source, + SourceBase, + Task, + TaskItem, + ThoughtTask, + Thread, + ThreadCreatedEvent, + ThreadItem, + ThreadItemAddedEvent, + ThreadItemDoneEvent, + ThreadItemRemovedEvent, + ThreadItemReplacedEvent, + ThreadItemUpdate, + ThreadItemUpdated, + ThreadMetadata, + ThreadStatus, + ThreadStreamEvent, + ThreadUpdatedEvent, + URLSource, + UserMessageItem, + WidgetComponentUpdated, + WidgetItem, + WidgetRootUpdated, + WidgetStreamingTextValueDelta, + Workflow, + WorkflowItem, + WorkflowSummary, + WorkflowTaskAdded, + WorkflowTaskUpdated, +) + +__all__ = [ + # Stream event types + 'ThreadStreamEvent', + 'ThreadCreatedEvent', + 'ThreadUpdatedEvent', + 'ThreadItemDoneEvent', + 'ThreadItemAddedEvent', + 'ThreadItemUpdated', + 'ThreadItemRemovedEvent', + 'ThreadItemReplacedEvent', + 'ProgressUpdateEvent', + 'ErrorEvent', + 'NoticeEvent', + # Thread item update types + 'AssistantMessageContentPartAdded', + 'AssistantMessageContentPartTextDelta', + 'AssistantMessageContentPartAnnotationAdded', + 'AssistantMessageContentPartDone', + 'WidgetStreamingTextValueDelta', + 'WidgetRootUpdated', + 'WidgetComponentUpdated', + 'WorkflowTaskAdded', + 'WorkflowTaskUpdated', + 'ThreadItemUpdate', + # Thread and thread item types + 'Thread', + 'ThreadMetadata', + 'ThreadItem', + 'ThreadStatus', + 'ActiveStatus', + 'LockedStatus', + 'ClosedStatus', + 'UserMessageItem', + 'AssistantMessageItem', + 'ClientToolCallItem', + 'WidgetItem', + 'TaskItem', + 'WorkflowItem', + 'EndOfTurnItem', + 'HiddenContextItem', + # Assistant message content + 'AssistantMessageContent', + 'Annotation', + # Workflow and task types + 'Workflow', + 'WorkflowSummary', + 'CustomSummary', + 'DurationSummary', + 'Task', + 'BaseTask', + 'CustomTask', + 'SearchTask', + 'ThoughtTask', + 'FileTask', + 'ImageTask', + # Source types + 'Source', + 'SourceBase', + 'FileSource', + 'URLSource', + 'EntitySource', + # Actions + 'Action', + # Misc types + 'IconName', + 'ErrorCode', + 'Page', +]