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

openbb_pydantic_ai/_serializers.py

@@ -0,0 +1,110 @@
+"""Content serialization utilities for OpenBB Pydantic AI adapter."""
+
+from __future__ import annotations
+
+import json
+from typing import Any, cast
+
+from openbb_ai.models import LlmClientFunctionCallResultMessage
+
+from ._types import SerializedContent
+
+
+class ContentSerializer:
+    """Handles serialization and parsing of content across the adapter."""
+
+    @staticmethod
+    def serialize_result(
+        message: LlmClientFunctionCallResultMessage,
+    ) -> SerializedContent:
+        """Serialize a function call result message into a content dictionary.
+
+        Parameters
+        ----------
+        message : LlmClientFunctionCallResultMessage
+            The function call result message to serialize
+
+        Returns
+        -------
+        SerializedContent
+            A typed dictionary containing input_arguments, data, and
+            optionally extra_state
+        """
+        data: list[Any] = []
+        for item in message.data:
+            if hasattr(item, "model_dump"):
+                data.append(item.model_dump(mode="json", exclude_none=True))
+            else:
+                data.append(item)
+
+        content: SerializedContent = cast(
+            SerializedContent,
+            {
+                "input_arguments": message.input_arguments,
+                "data": data,
+            },
+        )
+        if message.extra_state:
+            content["extra_state"] = message.extra_state
+        return content
+
+    @staticmethod
+    def parse_json(raw_content: str) -> Any:
+        """Parse JSON content, returning the original string if parsing fails.
+
+        Parameters
+        ----------
+        raw_content : str
+            The raw JSON string to parse
+
+        Returns
+        -------
+        Any
+            Parsed JSON object or original string if parsing fails
+        """
+        try:
+            return json.loads(raw_content)
+        except (json.JSONDecodeError, ValueError):
+            return raw_content
+
+    @staticmethod
+    def to_string(content: Any) -> str | None:
+        """Convert content to string with JSON fallback.
+
+        Parameters
+        ----------
+        content : Any
+            Content to stringify
+
+        Returns
+        -------
+        str | None
+            String representation or None if content is None
+        """
+        if content is None:
+            return None
+        if isinstance(content, str):
+            return content
+        try:
+            return json.dumps(content, default=str)
+        except (TypeError, ValueError):
+            return str(content)
+
+    @staticmethod
+    def to_json(value: Any) -> str:
+        """Convert value to JSON string with fallback to str().
+
+        Parameters
+        ----------
+        value : Any
+            Value to convert to JSON
+
+        Returns
+        -------
+        str
+            JSON string representation
+        """
+        try:
+            return json.dumps(value, default=str)
+        except (TypeError, ValueError):
+            return str(value)

openbb_pydantic_ai/_toolsets.py

@@ -0,0 +1,264 @@
+"""Toolset implementations for OpenBB widgets and visualization."""
+
+from __future__ import annotations
+
+import re
+from collections.abc import Mapping, Sequence
+from typing import Any, Literal
+
+from openbb_ai.helpers import chart, table
+from openbb_ai.models import Undefined, Widget, WidgetCollection, WidgetParam
+from pydantic_ai import CallDeferred, Tool
+from pydantic_ai.tools import RunContext
+from pydantic_ai.toolsets import FunctionToolset
+
+from ._dependencies import OpenBBDeps
+
+
+def _base_param_schema(param: WidgetParam) -> dict[str, Any]:
+    """Build the base JSON schema for a widget parameter."""
+    type_mapping: dict[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"},
+    }
+
+    schema = type_mapping.get(param.type, {"type": "string"})
+    schema = dict(schema)  # copy
+    schema["description"] = param.description
+
+    if param.options:
+        schema["enum"] = list(param.options)
+
+    if param.get_options:
+        schema.setdefault(
+            "description",
+            param.description + " (options retrieved dynamically)",
+        )
+
+    if param.default_value is not Undefined.UNDEFINED:
+        schema["default"] = param.default_value
+
+    if param.current_value is not None and param.multi_select is False:
+        schema.setdefault("examples", []).append(param.current_value)
+
+    return schema
+
+
+def _param_schema(param: WidgetParam) -> tuple[dict[str, Any], bool]:
+    """Return the schema for a parameter and whether it's required."""
+    schema = _base_param_schema(param)
+
+    if param.multi_select:
+        schema = {
+            "type": "array",
+            "items": schema,
+            "description": schema.get("description"),
+        }
+
+    is_required = param.default_value is Undefined.UNDEFINED
+    return schema, is_required
+
+
+def _widget_schema(widget: Widget) -> dict[str, Any]:
+    properties: dict[str, Any] = {}
+    required: list[str] = []
+
+    for param in widget.params:
+        schema, is_required = _param_schema(param)
+        properties[param.name] = schema
+        if is_required:
+            required.append(param.name)
+
+    widget_schema: dict[str, Any] = {
+        "type": "object",
+        "title": widget.name,
+        "properties": properties,
+        "additionalProperties": False,
+    }
+    if required:
+        widget_schema["required"] = required
+
+    return widget_schema
+
+
+def _slugify(value: str) -> str:
+    slug = re.sub(r"[^0-9A-Za-z]+", "_", value).strip("_")
+    return slug.lower() or "value"
+
+
+def build_widget_tool_name(widget: Widget) -> str:
+    """Generate a deterministic tool name for a widget.
+
+    The tool name is constructed as: openbb_widget_{origin}_{widget_id}
+    where both origin and widget_id are slugified.
+
+    Parameters
+    ----------
+    widget : Widget
+        The widget to generate a tool name for
+
+    Returns
+    -------
+    str
+        A unique, deterministic tool name string
+    """
+    origin_slug = _slugify(widget.origin)
+    widget_slug = _slugify(widget.widget_id)
+    return f"openbb_widget_{origin_slug}_{widget_slug}"
+
+
+def build_widget_tool(widget: Widget) -> Tool:
+    """Create a deferred tool for a widget.
+
+    This creates a Pydantic AI tool that will be called by the LLM but
+    executed by the OpenBB Workspace frontend (deferred execution).
+
+    Parameters
+    ----------
+    widget : Widget
+        The widget to create a tool for
+
+    Returns
+    -------
+    Tool
+        A Tool configured for deferred execution
+    """
+    tool_name = build_widget_tool_name(widget)
+    schema = _widget_schema(widget)
+    description = widget.description or widget.name
+
+    async def _call_widget(ctx: RunContext[OpenBBDeps], **input_arguments: Any) -> None:
+        # Ensure we have a tool call id for deferred execution
+        if ctx.tool_call_id is None:
+            raise RuntimeError("Deferred widget tools require a tool call id.")
+        raise CallDeferred
+
+    _call_widget.__name__ = f"call_widget_{widget.uuid}"
+
+    return Tool.from_schema(
+        function=_call_widget,
+        name=tool_name,
+        description=description,
+        json_schema=schema,
+        takes_ctx=True,
+    )
+
+
+class WidgetToolset(FunctionToolset[OpenBBDeps]):
+    """Toolset that exposes widgets as deferred tools."""
+
+    def __init__(self, widgets: Sequence[Widget]):
+        super().__init__()
+        self._widgets_by_tool: dict[str, Widget] = {}
+
+        for widget in widgets:
+            tool = build_widget_tool(widget)
+            self.add_tool(tool)
+            self._widgets_by_tool[tool.name] = widget
+
+    @property
+    def widgets_by_tool(self) -> Mapping[str, Widget]:
+        return self._widgets_by_tool
+
+
+class VisualizationToolset(FunctionToolset[OpenBBDeps]):
+    """Toolset exposing helper utilities for charts and tables."""
+
+    def __init__(self) -> None:
+        super().__init__()
+
+        def _create_table(
+            data: list[dict[str, Any]],
+            name: str | None = None,
+            description: str | None = None,
+        ):
+            """Create a table artifact to display in OpenBB Workspace."""
+
+            return table(data=data, name=name, description=description)
+
+        def _create_chart(
+            type: Literal["line", "bar", "scatter", "pie", "donut"],
+            data: list[dict[str, Any]],
+            x_key: str | None = None,
+            y_keys: list[str] | None = None,
+            angle_key: str | None = None,
+            callout_label_key: str | None = None,
+            name: str | None = None,
+            description: str | None = None,
+        ):
+            """Create a chart artifact (line, bar, scatter, pie, donut).
+
+            Raises
+            ------
+            ValueError
+                If required parameters for the given chart ``type`` are missing.
+            """
+
+            if type in {"line", "bar", "scatter"}:
+                if not x_key:
+                    raise ValueError(
+                        "x_key is required for line, bar, and scatter charts"
+                    )
+                if not y_keys:
+                    raise ValueError(
+                        "y_keys is required for line, bar, and scatter charts"
+                    )
+            elif type in {"pie", "donut"}:
+                if not angle_key:
+                    raise ValueError("angle_key is required for pie and donut charts")
+                if not callout_label_key:
+                    raise ValueError(
+                        "callout_label_key is required for pie and donut charts"
+                    )
+
+            return chart(
+                type=type,
+                data=data,
+                x_key=x_key,
+                y_keys=y_keys,
+                angle_key=angle_key,
+                callout_label_key=callout_label_key,
+                name=name,
+                description=description,
+            )
+
+        self.add_function(_create_table, name="openbb_create_table")
+        self.add_function(_create_chart, name="openbb_create_chart")
+
+
+def build_widget_toolsets(
+    collection: WidgetCollection | None,
+) -> tuple[FunctionToolset[OpenBBDeps], ...]:
+    """Create toolsets for each widget priority group plus visualization tools.
+
+    Widgets are organized into separate toolsets by priority (primary, secondary, extra)
+    to allow control over tool selection. The visualization toolset is always
+    included for creating charts and tables.
+
+    Parameters
+    ----------
+    collection : WidgetCollection | None
+        Widget collection with priority groups, or None
+
+    Returns
+    -------
+    tuple[FunctionToolset[OpenBBDeps], ...]
+        Toolsets including widget toolsets and visualization toolset
+    """
+    if collection is None:
+        return (VisualizationToolset(),)
+
+    toolsets: list[FunctionToolset[OpenBBDeps]] = []
+    for widgets in (collection.primary, collection.secondary, collection.extra):
+        if widgets:
+            toolsets.append(WidgetToolset(widgets))
+
+    toolsets.append(VisualizationToolset())
+
+    return tuple(toolsets)

openbb_pydantic_ai/_types.py

@@ -0,0 +1,39 @@
+"""Type definitions and protocols for OpenBB Pydantic AI adapter."""
+
+from __future__ import annotations
+
+import sys
+from typing import Any, Protocol, TypedDict
+
+if sys.version_info >= (3, 11):
+    from typing import NotRequired
+else:
+    from typing_extensions import NotRequired  # type: ignore[assignment]
+
+
+class SerializedContent(TypedDict):
+    """Structure for serialized tool result content."""
+
+    input_arguments: dict[str, Any]
+    data: list[Any]
+    extra_state: NotRequired[dict[str, Any]]
+
+
+class ToolCallMetadata(TypedDict):
+    """Metadata for tracking tool calls in flight."""
+
+    tool_call_id: str
+    widget_uuid: str
+    widget_id: str
+
+
+class TextStreamCallback(Protocol):
+    """Protocol for callbacks that mark text as having been streamed."""
+
+    def __call__(self) -> None:
+        """Mark that text has been streamed."""
+        ...
+
+
+# Type alias for structured detail entries
+DetailEntry = dict[str, Any] | str

openbb_pydantic_ai/_utils.py

@@ -0,0 +1,132 @@
+"""Utility functions for OpenBB Pydantic AI UI adapter."""
+
+from __future__ import annotations
+
+import hashlib
+import json
+from collections.abc import Mapping, Sequence
+from typing import Any
+
+from ._config import (
+    MAX_ARG_DISPLAY_CHARS,
+    MAX_ARG_PREVIEW_ITEMS,
+)
+from ._serializers import ContentSerializer
+
+
+def hash_tool_call(function: str, input_arguments: dict[str, Any]) -> str:
+    """Generate a deterministic hash-based ID for a tool call.
+
+    This creates a unique identifier by hashing the function name and arguments,
+    ensuring consistent tool call IDs across message history and deferred results.
+
+    Parameters
+    ----------
+    function : str
+        The name of the function/tool being called
+    input_arguments : dict[str, Any]
+        The arguments passed to the tool
+
+    Returns
+    -------
+    str
+        A string combining the function name with a 16-character hash digest
+    """
+    payload = json.dumps(
+        {"function": function, "input_arguments": input_arguments},
+        sort_keys=True,
+        default=str,
+    )
+    digest = hashlib.sha256(payload.encode("utf-8")).hexdigest()
+    return f"{function}_{digest[:16]}"
+
+
+def normalize_args(args: Any) -> dict[str, Any]:
+    """Normalize tool call arguments to a dictionary."""
+    if isinstance(args, dict):
+        return args
+    if isinstance(args, str):
+        try:
+            parsed = json.loads(args)
+            if isinstance(parsed, dict):
+                return parsed
+        except ValueError:
+            pass
+    return {}
+
+
+def get_str(mapping: Mapping[str, Any], *keys: str) -> str | None:
+    """Return the first string value found for the given keys."""
+    for key in keys:
+        value = mapping.get(key)
+        if isinstance(value, str):
+            return value
+    return None
+
+
+def get_str_list(mapping: Mapping[str, Any], *keys: str) -> list[str] | None:
+    """Return the first list of strings (or single string) found for the keys."""
+    for key in keys:
+        value = mapping.get(key)
+        if isinstance(value, str):
+            return [value]
+        if isinstance(value, list):
+            items = [item for item in value if isinstance(item, str)]
+            if items:
+                return items
+    return None
+
+
+def _truncate(value: str, max_chars: int = 160) -> str:
+    if len(value) <= max_chars:
+        return value
+    return value[: max_chars - 3] + "..."
+
+
+def _json_dump(value: Any) -> str:
+    return ContentSerializer.to_json(value)
+
+
+def format_arg_value(
+    value: Any,
+    *,
+    max_chars: int = MAX_ARG_DISPLAY_CHARS,
+    max_items: int = MAX_ARG_PREVIEW_ITEMS,
+) -> str:
+    """Summarize nested structures so reasoning details stay readable."""
+
+    if isinstance(value, str):
+        return _truncate(value, max_chars)
+
+    if isinstance(value, (int, float, bool)) or value is None:
+        return _json_dump(value)
+
+    if isinstance(value, Mapping):
+        keys = list(value.keys())
+        preview_keys = keys[:max_items]
+        preview = {k: value[k] for k in preview_keys}
+        suffix = "..." if len(keys) > max_items else ""
+        return _truncate(
+            f"dict(keys={preview_keys}{suffix}, sample={_json_dump(preview)})",
+            max_chars,
+        )
+
+    if isinstance(value, Sequence) and not isinstance(value, (bytes, bytearray)):
+        seq = list(value)
+        preview = seq[:max_items]
+        suffix = "..." if len(seq) > max_items else ""
+        return _truncate(
+            f"list(len={len(seq)}{suffix}, sample={_json_dump(preview)})",
+            max_chars,
+        )
+
+    return _truncate(_json_dump(value), max_chars)
+
+
+def format_args(args: Mapping[str, Any]) -> dict[str, str]:
+    """Format a mapping of arguments into readable key/value strings."""
+
+    formatted: dict[str, str] = {}
+    for key, value in args.items():
+        formatted[key] = format_arg_value(value)
+    return formatted

openbb_pydantic_ai/_widget_registry.py

@@ -0,0 +1,145 @@
+"""Widget registry for centralized widget discovery and lookup."""
+
+from __future__ import annotations
+
+from collections.abc import Iterator, Mapping, Sequence
+from typing import TYPE_CHECKING
+
+from openbb_ai.models import (
+    LlmClientFunctionCallResultMessage,
+    Widget,
+    WidgetCollection,
+)
+
+from ._config import GET_WIDGET_DATA_TOOL_NAME
+
+if TYPE_CHECKING:
+    from pydantic_ai.toolsets import FunctionToolset
+
+    from ._dependencies import OpenBBDeps
+
+
+class WidgetRegistry:
+    """Centralized registry for widget discovery and lookup."""
+
+    def __init__(
+        self,
+        collection: WidgetCollection | None = None,
+        toolsets: Sequence[FunctionToolset[OpenBBDeps]] | None = None,
+    ):
+        """Initialize widget registry from collection and toolsets.
+
+        Parameters
+        ----------
+        collection : WidgetCollection | None
+            Widget collection with priority groups
+        toolsets : Sequence[FunctionToolset[OpenBBDeps]] | None
+            Widget toolsets
+        """
+        self._by_tool_name: dict[str, Widget] = {}
+        self._by_uuid: dict[str, Widget] = {}
+
+        # Build lookup from toolsets
+        if toolsets:
+            for toolset in toolsets:
+                widgets = getattr(toolset, "widgets_by_tool", None)
+                if widgets:
+                    for tool_name, widget in widgets.items():
+                        self._by_tool_name[tool_name] = widget
+                        self._by_uuid[str(widget.uuid)] = widget
+
+        # Also index from collection if provided
+        if collection:
+            for widget in self._iter_collection(collection):
+                self._by_uuid[str(widget.uuid)] = widget
+
+    @staticmethod
+    def _iter_collection(collection: WidgetCollection) -> Iterator[Widget]:
+        """Iterate all widgets in a collection."""
+        for group in (collection.primary, collection.secondary, collection.extra):
+            yield from group
+
+    def find_by_tool_name(self, name: str) -> Widget | None:
+        """Find a widget by its tool name.
+
+        Parameters
+        ----------
+        name : str
+            The tool name to search for
+
+        Returns
+        -------
+        Widget | None
+            The widget if found, None otherwise
+        """
+        return self._by_tool_name.get(name)
+
+    def find_by_uuid(self, uuid: str) -> Widget | None:
+        """Find a widget by its UUID string.
+
+        Parameters
+        ----------
+        uuid : str
+            The UUID to search for
+
+        Returns
+        -------
+        Widget | None
+            The widget if found, None otherwise
+        """
+        return self._by_uuid.get(uuid)
+
+    def find_for_result(
+        self, result: LlmClientFunctionCallResultMessage
+    ) -> Widget | None:
+        """Find the widget that produced a result message.
+
+        Parameters
+        ----------
+        result : LlmClientFunctionCallResultMessage
+            The result message to find a widget for
+
+        Returns
+        -------
+        Widget | None
+            The widget if found, None otherwise
+        """
+        # Check direct tool name match
+        widget = self.find_by_tool_name(result.function)
+        if widget is not None:
+            return widget
+
+        # Check if it's a get_widget_data call
+        if result.function == GET_WIDGET_DATA_TOOL_NAME:
+            data_sources = result.input_arguments.get("data_sources", [])
+            if data_sources:
+                widget_uuid = data_sources[0].get("widget_uuid")
+                if widget_uuid:
+                    return self.find_by_uuid(widget_uuid)
+
+        return None
+
+    def iter_all(self) -> Iterator[Widget]:
+        """Iterate all registered widgets.
+
+        Returns
+        -------
+        Iterator[Widget]
+            Iterator over all widgets
+        """
+        # Use dict to deduplicate by UUID
+        seen = set()
+        for widget in self._by_uuid.values():
+            if str(widget.uuid) not in seen:
+                seen.add(str(widget.uuid))
+                yield widget
+
+    def as_mapping(self) -> Mapping[str, Widget]:
+        """Get widget lookup as a read-only mapping by tool name.
+
+        Returns
+        -------
+        Mapping[str, Widget]
+            Read-only mapping from tool names to widgets
+        """
+        return self._by_tool_name