openbb-pydantic-ai 0.1.1__py3-none-any.whl

openbb_pydantic_ai/__init__.py

@@ -0,0 +1,52 @@
+"""Pydantic AI UI adapter for OpenBB Workspace."""
+
+from importlib.metadata import PackageNotFoundError, version
+
+from ._adapter import OpenBBAIAdapter
+from ._config import GET_WIDGET_DATA_TOOL_NAME
+from ._dependencies import OpenBBDeps, build_deps_from_request
+from ._event_builder import EventBuilder
+from ._event_stream import OpenBBAIEventStream
+from ._exceptions import (
+    InvalidToolCallError,
+    OpenBBPydanticAIError,
+    SerializationError,
+    WidgetNotFoundError,
+)
+from ._message_transformer import MessageTransformer
+from ._serializers import ContentSerializer
+from ._toolsets import (
+    WidgetToolset,
+    build_widget_tool,
+    build_widget_tool_name,
+    build_widget_toolsets,
+)
+from ._types import SerializedContent, TextStreamCallback
+from ._widget_registry import WidgetRegistry
+
+try:
+    __version__ = version("openbb-pydantic-ai")
+except PackageNotFoundError:
+    __version__ = "0.0.0"
+
+__all__ = [
+    "OpenBBAIAdapter",
+    "OpenBBAIEventStream",
+    "OpenBBDeps",
+    "build_deps_from_request",
+    "WidgetToolset",
+    "build_widget_tool",
+    "build_widget_tool_name",
+    "build_widget_toolsets",
+    "GET_WIDGET_DATA_TOOL_NAME",
+    "EventBuilder",
+    "ContentSerializer",
+    "MessageTransformer",
+    "WidgetRegistry",
+    "OpenBBPydanticAIError",
+    "WidgetNotFoundError",
+    "InvalidToolCallError",
+    "SerializationError",
+    "SerializedContent",
+    "TextStreamCallback",
+]

openbb_pydantic_ai/_adapter.py

@@ -0,0 +1,307 @@
+from __future__ import annotations
+
+from collections.abc import Sequence
+from dataclasses import KW_ONLY, dataclass, field
+from functools import cached_property
+from typing import Any, cast
+
+from openbb_ai.models import (
+    SSE,
+    LlmClientFunctionCallResultMessage,
+    LlmMessage,
+    QueryRequest,
+)
+from pydantic_ai import DeferredToolResults
+from pydantic_ai.messages import (
+    ModelMessage,
+    SystemPromptPart,
+)
+from pydantic_ai.toolsets import AbstractToolset, CombinedToolset, FunctionToolset
+from pydantic_ai.ui import UIAdapter
+
+from ._dependencies import OpenBBDeps, build_deps_from_request
+from ._event_stream import OpenBBAIEventStream
+from ._message_transformer import MessageTransformer
+from ._serializers import ContentSerializer
+from ._toolsets import build_widget_toolsets
+from ._utils import hash_tool_call
+from ._widget_registry import WidgetRegistry
+
+
+@dataclass(slots=True)
+class OpenBBAIAdapter(UIAdapter[QueryRequest, LlmMessage, SSE, OpenBBDeps, Any]):
+    """UI adapter that bridges OpenBB Workspace requests with Pydantic AI."""
+
+    _: KW_ONLY
+    accept: str | None = None
+
+    # Initialized in __post_init__
+    _transformer: MessageTransformer = field(init=False)
+    _registry: WidgetRegistry = field(init=False)
+    _base_messages: list[LlmMessage] = field(init=False, default_factory=list)
+    _pending_results: list[LlmClientFunctionCallResultMessage] = field(
+        init=False, default_factory=list
+    )
+
+    def __post_init__(self) -> None:
+        base, pending = self._split_messages(self.run_input.messages)
+        self._base_messages = base
+        self._pending_results = pending
+
+        # Build tool call ID overrides for consistent IDs
+        tool_call_id_overrides: dict[str, str] = {}
+        for message in self._base_messages:
+            if isinstance(message, LlmClientFunctionCallResultMessage):
+                key = hash_tool_call(message.function, message.input_arguments)
+                tool_call_id = self._tool_call_id_from_result(message)
+                tool_call_id_overrides[key] = tool_call_id
+
+        for message in self._pending_results:
+            key = hash_tool_call(message.function, message.input_arguments)
+            tool_call_id_overrides.setdefault(
+                key,
+                self._tool_call_id_from_result(message),
+            )
+
+        # Initialize transformer and registry
+        self._transformer = MessageTransformer(tool_call_id_overrides)
+        self._registry = WidgetRegistry(
+            collection=self.run_input.widgets,
+            toolsets=self._widget_toolsets,
+        )
+
+    @classmethod
+    def build_run_input(cls, body: bytes) -> QueryRequest:
+        return QueryRequest.model_validate_json(body)
+
+    @classmethod
+    def load_messages(cls, messages: Sequence[LlmMessage]) -> list[ModelMessage]:
+        """Convert OpenBB messages to Pydantic AI messages.
+
+        Note: This creates a transformer without overrides for standalone use.
+        """
+        transformer = MessageTransformer()
+        return transformer.transform_batch(messages)
+
+    @staticmethod
+    def _split_messages(
+        messages: Sequence[LlmMessage],
+    ) -> tuple[list[LlmMessage], list[LlmClientFunctionCallResultMessage]]:
+        """Split messages into base history and pending deferred results.
+
+        Only results after the last AI message are considered pending. Results
+        followed by AI messages were already processed in previous turns.
+
+        Parameters
+        ----------
+        messages : Sequence[LlmMessage]
+            Full message sequence
+
+        Returns
+        -------
+        tuple[list[LlmMessage], list[LlmClientFunctionCallResultMessage]]
+            (base messages, pending results that need processing)
+        """
+        base = list(messages)
+        pending: list[LlmClientFunctionCallResultMessage] = []
+
+        # Treat only the trailing tool results (those after the final assistant
+        # message) as pending. Leave them in the base history so the next model
+        # call still sees the complete tool call/result exchange.
+        idx = len(base) - 1
+        while idx >= 0:
+            message = base[idx]
+            if not isinstance(message, LlmClientFunctionCallResultMessage):
+                break
+            pending.insert(0, cast(LlmClientFunctionCallResultMessage, message))
+            idx -= 1
+
+        return base, pending
+
+    def _tool_call_id_from_result(
+        self, message: LlmClientFunctionCallResultMessage
+    ) -> str:
+        """Extract or generate a tool call ID from a result message."""
+        extra_id = (
+            message.extra_state.get("tool_call_id") if message.extra_state else None
+        )
+        if isinstance(extra_id, str):
+            return extra_id
+        return hash_tool_call(message.function, message.input_arguments)
+
+    @cached_property
+    def deps(self) -> OpenBBDeps:
+        return build_deps_from_request(self.run_input)
+
+    @cached_property
+    def deferred_tool_results(self) -> DeferredToolResults | None:
+        """Build deferred tool results from pending result messages."""
+        if not self._pending_results:
+            return None
+
+        # When those trailing results already sit in the base history, skip
+        # emitting DeferredToolResults; resending them would show up as a
+        # conflicting duplicate tool response upstream.
+        if self._pending_results_are_in_history():
+            return None
+
+        results = DeferredToolResults()
+        for message in self._pending_results:
+            actual_id = self._tool_call_id_from_result(message)
+            serialized = ContentSerializer.serialize_result(message)
+            results.calls[actual_id] = serialized
+        return results
+
+    def _pending_results_are_in_history(self) -> bool:
+        if not self._pending_results:
+            return False
+        pending_len = len(self._pending_results)
+        if pending_len > len(self._base_messages):
+            return False
+        tail = self._base_messages[-pending_len:]
+        return all(
+            orig is pending
+            for orig, pending in zip(tail, self._pending_results, strict=True)
+        )
+
+    @cached_property
+    def _widget_toolsets(self) -> tuple[FunctionToolset[OpenBBDeps], ...]:
+        return build_widget_toolsets(self.run_input.widgets)
+
+    def build_event_stream(self) -> OpenBBAIEventStream:
+        return OpenBBAIEventStream(
+            run_input=self.run_input,
+            widget_registry=self._registry,
+            pending_results=self._pending_results,
+        )
+
+    @cached_property
+    def messages(self) -> list[ModelMessage]:
+        """Build message history with context prompts."""
+        from pydantic_ai.ui import MessagesBuilder
+
+        builder = MessagesBuilder()
+        self._add_context_prompts(builder)
+
+        # Use transformer to convert messages with ID overrides
+        transformed = self._transformer.transform_batch(self._base_messages)
+        for msg in transformed:
+            for part in msg.parts:
+                builder.add(part)
+
+        return builder.messages
+
+    def _add_context_prompts(self, builder) -> None:
+        """Add system prompts with workspace context, URLs, and dashboard info."""
+        lines: list[str] = []
+
+        if self.deps.context:
+            lines.append("Workspace context:")
+            for ctx in self.deps.context:
+                row_count = len(ctx.data.items) if ctx.data and ctx.data.items else 0
+                summary = f"- {ctx.name} ({row_count} rows): {ctx.description}"
+                lines.append(summary)
+
+        if self.deps.urls:
+            joined = ", ".join(self.deps.urls)
+            lines.append(f"Relevant URLs: {joined}")
+
+        workspace_state = self.deps.workspace_state
+        if workspace_state and workspace_state.current_dashboard_info:
+            dashboard = workspace_state.current_dashboard_info
+            lines.append(
+                f"Active dashboard: {dashboard.name} (tab {dashboard.current_tab_id})"
+            )
+
+        if lines:
+            builder.add(SystemPromptPart(content="\n".join(lines)))
+
+    @cached_property
+    def toolset(self) -> AbstractToolset[OpenBBDeps] | None:
+        """Build combined toolset from widget toolsets."""
+        if not self._widget_toolsets:
+            return None
+        if len(self._widget_toolsets) == 1:
+            return self._widget_toolsets[0]
+        combined = CombinedToolset(self._widget_toolsets)
+        return cast(AbstractToolset[OpenBBDeps], combined)
+
+    @cached_property
+    def state(self) -> dict[str, Any] | None:
+        """Extract workspace state as a dictionary."""
+        if self.run_input.workspace_state is None:
+            return None
+        return self.run_input.workspace_state.model_dump(exclude_none=True)
+
+    def run_stream_native(
+        self,
+        *,
+        output_type=None,
+        message_history=None,
+        deferred_tool_results=None,
+        model=None,
+        deps=None,
+        model_settings=None,
+        usage_limits=None,
+        usage=None,
+        infer_name=True,
+        toolsets=None,
+        builtin_tools=None,
+    ):
+        """
+        Run the agent with OpenBB-specific defaults for
+        deps, messages, and deferred results.
+        """
+        deps = deps or self.deps  # type: ignore[assignment]
+        deferred_tool_results = deferred_tool_results or self.deferred_tool_results
+        message_history = message_history or self.messages
+
+        return super().run_stream_native(
+            output_type=output_type,
+            message_history=message_history,
+            deferred_tool_results=deferred_tool_results,
+            model=model,
+            deps=deps,
+            model_settings=model_settings,
+            usage_limits=usage_limits,
+            usage=usage,
+            infer_name=infer_name,
+            toolsets=toolsets,
+            builtin_tools=builtin_tools,
+        )
+
+    def run_stream(
+        self,
+        *,
+        output_type=None,
+        message_history=None,
+        deferred_tool_results=None,
+        model=None,
+        deps=None,
+        model_settings=None,
+        usage_limits=None,
+        usage=None,
+        infer_name=True,
+        toolsets=None,
+        builtin_tools=None,
+        on_complete=None,
+    ):
+        """Run the agent and stream protocol-specific events with OpenBB defaults."""
+        deps = deps or self.deps  # type: ignore[assignment]
+        deferred_tool_results = deferred_tool_results or self.deferred_tool_results
+        message_history = message_history or self.messages
+
+        return super().run_stream(
+            output_type=output_type,
+            message_history=message_history,
+            deferred_tool_results=deferred_tool_results,
+            model=model,
+            deps=deps,
+            model_settings=model_settings,
+            usage_limits=usage_limits,
+            usage=usage,
+            infer_name=infer_name,
+            toolsets=toolsets,
+            builtin_tools=builtin_tools,
+            on_complete=on_complete,
+        )

openbb_pydantic_ai/_config.py

@@ -0,0 +1,58 @@
+"""Centralized configuration for OpenBB Pydantic AI adapter."""
+
+from __future__ import annotations
+
+from typing import Any, Mapping
+
+# Tool name constants
+GET_WIDGET_DATA_TOOL_NAME = "get_widget_data"
+
+# Field exclusion lists for citation and status update details
+CITATION_EXCLUDED_FIELDS = frozenset(
+    [
+        "lastupdated",
+        "source",
+        "id",
+        "uuid",
+        "storedfileuuid",
+        "datakey",
+        "originalfilename",
+        "extension",
+        "category",
+        "subcategory",
+        "transcript_url",
+    ]
+)
+
+STATUS_UPDATE_EXCLUDED_FIELDS = frozenset(
+    [
+        "lastupdated",
+        "source",
+        "id",
+        "uuid",
+        "storedfileuuid",
+        "url",
+        "datakey",
+        "originalfilename",
+        "extension",
+        "category",
+        "subcategory",
+        "transcript_url",
+    ]
+)
+
+# Widget parameter type to JSON schema mapping
+PARAM_TYPE_SCHEMA_MAP: Mapping[str, dict[str, Any]] = {
+    "string": {"type": "string"},
+    "text": {"type": "string"},
+    "number": {"type": "number"},
+    "integer": {"type": "integer"},
+    "boolean": {"type": "boolean"},
+    "date": {"type": "string", "format": "date"},
+    "ticker": {"type": "string"},
+    "endpoint": {"type": "string"},
+}
+
+# Content formatting limits
+MAX_ARG_DISPLAY_CHARS = 160
+MAX_ARG_PREVIEW_ITEMS = 2

openbb_pydantic_ai/_dependencies.py

@@ -0,0 +1,73 @@
+"""Dependency injection container for OpenBB workspace context."""
+
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+from typing import Any, Iterable, Sequence
+
+from openbb_ai.models import (
+    QueryRequest,
+    RawContext,
+    Widget,
+    WidgetCollection,
+    WorkspaceState,
+)
+
+
+@dataclass(slots=True)
+class OpenBBDeps:
+    """Dependency container passed to Pydantic AI runs.
+
+    The dependency bundle exposes OpenBB Workspace specific context so that
+    system prompts, tools, and output validators can access widget metadata or
+    other request scoped information via ``RunContext[OpenBBDeps]``.
+
+    Attributes:
+        widgets: Collection of available widgets organized by priority
+        context: Workspace context data (datasets, documents, etc.)
+        urls: Relevant URLs for the current request
+        workspace_state: Current workspace state including dashboard info
+        timezone: User's timezone (defaults to UTC)
+        state: Serialized workspace state as dictionary
+    """
+
+    widgets: WidgetCollection | None = None
+    context: list[RawContext] | None = None
+    urls: list[str] | None = None
+    workspace_state: WorkspaceState | None = None
+    timezone: str = "UTC"
+    state: dict[str, Any] = field(default_factory=dict)
+
+    def iter_widgets(self) -> Iterable[Widget]:
+        """Yield all widgets across priority groups (primary, secondary, extra)."""
+        if not self.widgets:
+            return
+
+        for group in (self.widgets.primary, self.widgets.secondary, self.widgets.extra):
+            yield from group
+
+    def get_widget_by_uuid(self, widget_uuid: str) -> Widget | None:
+        """Find a widget by its UUID string."""
+        for widget in self.iter_widgets():
+            if str(widget.uuid) == widget_uuid:
+                return widget
+        return None
+
+
+def build_deps_from_request(request: QueryRequest) -> OpenBBDeps:
+    """Create an OpenBBDeps instance from an incoming QueryRequest."""
+    context: Sequence[RawContext] | None = request.context
+    urls: Sequence[str] | None = request.urls
+
+    workspace_state = request.workspace_state
+
+    return OpenBBDeps(
+        widgets=request.widgets,
+        context=list(context) if context is not None else None,
+        urls=list(urls) if urls is not None else None,
+        workspace_state=workspace_state,
+        timezone=request.timezone,
+        state=workspace_state.model_dump(exclude_none=True)
+        if workspace_state is not None
+        else {},
+    )

openbb_pydantic_ai/_event_builder.py

@@ -0,0 +1,183 @@
+"""Event builder utilities for consistent SSE event creation."""
+
+from __future__ import annotations
+
+from typing import Any, Literal
+
+from openbb_ai.models import (
+    Citation,
+    CitationCollection,
+    CitationCollectionSSE,
+    ClientArtifact,
+    MessageArtifactSSE,
+    MessageChunkSSE,
+    MessageChunkSSEData,
+    StatusUpdateSSE,
+    StatusUpdateSSEData,
+)
+
+
+class EventBuilder:
+    """Provides consistent interface for creating SSE events."""
+
+    @staticmethod
+    def reasoning(
+        message: str,
+        *,
+        details: dict[str, Any] | list[dict[str, Any] | str] | None = None,
+        event_type: Literal["INFO", "WARNING", "ERROR"] = "INFO",
+        artifacts: list[ClientArtifact] | None = None,
+        hidden: bool = False,
+    ) -> StatusUpdateSSE:
+        """Create a reasoning/status update event.
+
+        Parameters
+        ----------
+        message : str
+            The reasoning message
+        details : dict | list | None
+            Optional details to include
+        event_type : Literal["INFO", "WARNING", "ERROR"]
+            The type of event
+        artifacts : list[ClientArtifact] | None
+            Optional artifacts to include
+        hidden : bool
+            Whether to hide this event
+
+        Returns
+        -------
+        StatusUpdateSSE
+            A status update SSE event
+        """
+        # Normalize dict details to list format expected by StatusUpdateSSEData
+        normalized_details: list[dict[str, Any] | str] | None = None
+        if details is not None:
+            if isinstance(details, dict):
+                normalized_details = [details]
+            else:
+                # details is already a list[dict[str, Any] | str]
+                normalized_details = details  # type: ignore[assignment]
+
+        return StatusUpdateSSE(
+            data=StatusUpdateSSEData(
+                eventType=event_type,
+                message=message,
+                group="reasoning",
+                details=normalized_details,
+                artifacts=artifacts,
+                hidden=hidden,
+            )
+        )
+
+    @staticmethod
+    def reasoning_with_artifacts(
+        message: str,
+        artifacts: list[ClientArtifact],
+    ) -> StatusUpdateSSE:
+        """Create a reasoning event with inline artifacts.
+
+        Parameters
+        ----------
+        message : str
+            The reasoning message
+        artifacts : list[ClientArtifact]
+            Artifacts to include inline
+
+        Returns
+        -------
+        StatusUpdateSSE
+            A status update SSE event with artifacts
+        """
+        return EventBuilder.reasoning(
+            message,
+            event_type="INFO",
+            artifacts=artifacts,
+        )
+
+    @staticmethod
+    def message(content: str) -> MessageChunkSSE:
+        """Create a message chunk event.
+
+        Parameters
+        ----------
+        content : str
+            The message content
+
+        Returns
+        -------
+        MessageChunkSSE
+            A message chunk SSE event
+        """
+        return MessageChunkSSE(data=MessageChunkSSEData(delta=content))
+
+    @staticmethod
+    def artifact(artifact: ClientArtifact) -> MessageArtifactSSE:
+        """Create an artifact event.
+
+        Parameters
+        ----------
+        artifact : ClientArtifact
+            The artifact to send
+
+        Returns
+        -------
+        MessageArtifactSSE
+            An artifact SSE event
+        """
+        return MessageArtifactSSE(data=artifact)
+
+    @staticmethod
+    def citations(citation_list: list[Citation]) -> CitationCollectionSSE:
+        """Create a citation collection event.
+
+        Parameters
+        ----------
+        citation_list : list[Citation]
+            List of citations
+
+        Returns
+        -------
+        CitationCollectionSSE
+            A citation collection SSE event
+        """
+        return CitationCollectionSSE(data=CitationCollection(citations=citation_list))
+
+    @staticmethod
+    def error(
+        message: str, *, details: dict[str, Any] | None = None
+    ) -> StatusUpdateSSE:
+        """Create an error event.
+
+        Parameters
+        ----------
+        message : str
+            The error message
+        details : dict | None
+            Optional error details
+
+        Returns
+        -------
+        StatusUpdateSSE
+            An error status update SSE event
+        """
+        return EventBuilder.reasoning(message, event_type="ERROR", details=details)
+
+    @staticmethod
+    def warning(
+        message: str, *, details: dict[str, Any] | None = None
+    ) -> StatusUpdateSSE:
+        """Create a warning event.
+
+        Parameters
+        ----------
+        message : str
+            The warning message
+        details : dict | None
+            Optional warning details
+
+        Returns
+        -------
+        StatusUpdateSSE
+            A warning status update SSE event
+        """
+        return EventBuilder.reasoning(message, event_type="WARNING", details=details)