shinychat 0.1.0__py3-none-any.whl → 0.2.1__py3-none-any.whl

shinychat/_chat_normalize.py

@@ -1,350 +1,315 @@
+from __future__ import annotations
+
 import sys
-from abc import ABC, abstractmethod
-from typing import TYPE_CHECKING, Any, Optional, cast
+from functools import singledispatch
 
-from htmltools import HTML, Tagifiable
+from htmltools import HTML, Tagifiable, TagList
 
+from ._chat_normalize_chatlas import tool_request_contents, tool_result_contents
 from ._chat_types import ChatMessage
 
-if TYPE_CHECKING:
-    from anthropic.types import Message as AnthropicMessage
-    from anthropic.types import MessageStreamEvent
+__all__ = ["message_content", "message_content_chunk"]
+
+
+@singledispatch
+def message_content(message):
+    """
+    Extract content from various message types into a ChatMessage.
+
+    This function uses `singledispatch` to allow for easy extension to support
+    new message types. To add support for a new type, register a new function
+    using the `@message_content.register` decorator.
 
-    if sys.version_info >= (3, 9):
-        from google.generativeai.types.generation_types import (  # pyright: ignore[reportMissingTypeStubs]
-            GenerateContentResponse,
+    Parameters
+    ----------
+    message
+        The message object to extract content from (e.g., ChatCompletion,
+        BaseMessage, etc.).
+
+    Note
+    ----
+    This function is implicitly called by `Chat.append_message()` to support
+    handling of various message types. It is not intended to be called directly
+    by users, but may be useful for debugging or advanced use cases.
+
+    Returns
+    -------
+    ChatMessage
+        A ChatMessage object containing the extracted content and role.
+
+    Raises
+    ------
+    ValueError
+        If the message type is unsupported.
+    """
+    if isinstance(message, (str, HTML)) or message is None:
+        return ChatMessage(content=message)
+    if isinstance(message, dict):
+        if "content" not in message:
+            raise ValueError("Message dictionary must have a 'content' key")
+        return ChatMessage(
+            content=message["content"],
+            role=message.get("role", "assistant"),
         )
-    else:
+    raise ValueError(
+        f"Don't know how to extract content for message type {type(message)}: {message}. "
+        "Consider registering a function to handle this type via `@message_content.register`"
+    )
 
-        class GenerateContentResponse:
-            text: str
 
-    from langchain_core.messages import BaseMessage, BaseMessageChunk
-    from openai.types.chat import ChatCompletion, ChatCompletionChunk
+@singledispatch
+def message_content_chunk(chunk):
+    """
+    Extract content from various message chunk types into a ChatMessage.
+
+    This function uses `singledispatch` to allow for easy extension to support
+    new chunk types. To add support for a new type, register a new function
+    using the `@message_content_chunk.register` decorator.
+
+    Parameters
+    ----------
+    chunk
+        The message chunk object to extract content from (e.g., ChatCompletionChunk,
+        BaseMessageChunk, etc.).
+
+    Note
+    ----
+    This function is implicitly called by `Chat.append_message_stream()` (on every
+    chunk of a message stream). It is not intended to be called directly by
+    users, but may be useful for debugging or advanced use cases.
+
+    Returns
+    -------
+    ChatMessage
+        A ChatMessage object containing the extracted content and role.
+
+    Raises
+    ------
+    ValueError
+        If the chunk type is unsupported.
+    """
+    if isinstance(chunk, (str, HTML)) or chunk is None:
+        return ChatMessage(content=chunk)
+    if isinstance(chunk, dict):
+        if "content" not in chunk:
+            raise ValueError("Chunk dictionary must have a 'content' key")
+        return ChatMessage(
+            content=chunk["content"],
+            role=chunk.get("role", "assistant"),
+        )
+    raise ValueError(
+        f"Don't know how to extract content for message chunk type {type(chunk)}: {chunk}. "
+        "Consider registering a function to handle this type via `@message_content_chunk.register`"
+    )
 
 
-class BaseMessageNormalizer(ABC):
-    @abstractmethod
-    def normalize(self, message: Any) -> ChatMessage:
-        pass
+# ------------------------------------------------------------------
+# Shiny tagifiable content extractor
+# ------------------------------------------------------------------
 
-    @abstractmethod
-    def normalize_chunk(self, chunk: Any) -> ChatMessage:
-        pass
 
-    @abstractmethod
-    def can_normalize(self, message: Any) -> bool:
-        pass
+@message_content.register
+def _(message: Tagifiable):
+    return ChatMessage(content=message)
 
-    @abstractmethod
-    def can_normalize_chunk(self, chunk: Any) -> bool:
-        pass
 
+@message_content_chunk.register
+def _(chunk: Tagifiable):
+    return ChatMessage(content=chunk)
 
-class StringNormalizer(BaseMessageNormalizer):
-    def normalize(self, message: Any) -> ChatMessage:
-        x = cast(Optional[str], message)
-        return ChatMessage(content=x or "", role="assistant")
 
-    def normalize_chunk(self, chunk: Any) -> ChatMessage:
-        x = cast(Optional[str], chunk)
-        return ChatMessage(content=x or "", role="assistant")
+# -----------------------------------------------------------------
+# chatlas tool call display
+# -----------------------------------------------------------------
+try:
+    from chatlas import ContentToolRequest, ContentToolResult, Turn
+    from chatlas.types import Content, ContentText
 
-    def can_normalize(self, message: Any) -> bool:
-        return isinstance(message, (str, HTML)) or message is None
+    @message_content.register
+    def _(message: Content):
+        return ChatMessage(content=str(message))
 
-    def can_normalize_chunk(self, chunk: Any) -> bool:
-        return isinstance(chunk, (str, HTML)) or chunk is None
+    @message_content_chunk.register
+    def _(chunk: Content):
+        return message_content(chunk)
 
+    @message_content.register
+    def _(message: ContentText):
+        return ChatMessage(content=message.text)
 
-class DictNormalizer(BaseMessageNormalizer):
-    def normalize(self, message: Any) -> ChatMessage:
-        x = cast("dict[str, Any]", message)
-        if "content" not in x:
-            raise ValueError("Message must have 'content' key")
-        return ChatMessage(content=x["content"], role=x.get("role", "assistant"))
+    @message_content_chunk.register
+    def _(chunk: ContentText):
+        return message_content(chunk)
 
-    def normalize_chunk(self, chunk: Any) -> ChatMessage:
-        x = cast("dict[str, Any]", chunk)
-        if "content" not in x:
-            raise ValueError("Message must have 'content' key")
-        return ChatMessage(content=x["content"], role=x.get("role", "assistant"))
+    @message_content.register
+    def _(chunk: ContentToolRequest):
+        return ChatMessage(content=tool_request_contents(chunk))
 
-    def can_normalize(self, message: Any) -> bool:
-        return isinstance(message, dict)
+    @message_content_chunk.register
+    def _(chunk: ContentToolRequest):
+        return message_content(chunk)
 
-    def can_normalize_chunk(self, chunk: Any) -> bool:
-        return isinstance(chunk, dict)
+    @message_content.register
+    def _(chunk: ContentToolResult):
+        return ChatMessage(content=tool_result_contents(chunk))
 
+    @message_content_chunk.register
+    def _(chunk: ContentToolResult):
+        return message_content(chunk)
 
-class TagifiableNormalizer(DictNormalizer):
-    def normalize(self, message: Any) -> ChatMessage:
-        x = cast("Tagifiable", message)
-        return super().normalize({"content": x})
+    @message_content.register
+    def _(message: Turn):
+        contents = TagList()
+        for x in message.contents:
+            contents.append(message_content(x).content)
+        return ChatMessage(content=contents)
 
-    def normalize_chunk(self, chunk: Any) -> ChatMessage:
-        x = cast("Tagifiable", chunk)
-        return super().normalize_chunk({"content": x})
+    @message_content_chunk.register
+    def _(chunk: Turn):
+        return message_content(chunk)
 
-    def can_normalize(self, message: Any) -> bool:
-        return isinstance(message, Tagifiable)
+    # N.B., unlike R, Python Chat stores UI state and so can replay
+    # it with additional workarounds. That's why R currently has a
+    # shinychat_contents() method for Chat, but Python doesn't.
+except ImportError:
+    pass
 
-    def can_normalize_chunk(self, chunk: Any) -> bool:
-        return isinstance(chunk, Tagifiable)
+# ------------------------------------------------------------------
+# LangChain content extractor
+# ------------------------------------------------------------------
 
+try:
+    from langchain_core.messages import BaseMessage, BaseMessageChunk
 
-class LangChainNormalizer(BaseMessageNormalizer):
-    def normalize(self, message: Any) -> ChatMessage:
-        x = cast("BaseMessage", message)
-        if isinstance(x.content, list):  # type: ignore
+    @message_content.register
+    def _(message: BaseMessage):
+        if isinstance(message.content, list):
             raise ValueError(
                 "The `message.content` provided seems to represent numerous messages. "
                 "Consider iterating over `message.content` and calling .append_message() on each iteration."
             )
-        return ChatMessage(content=x.content, role="assistant")
+        return ChatMessage(
+            content=message.content,
+            role="assistant",
+        )
 
-    def normalize_chunk(self, chunk: Any) -> ChatMessage:
-        x = cast("BaseMessageChunk", chunk)
-        if isinstance(x.content, list):  # type: ignore
+    @message_content_chunk.register
+    def _(chunk: BaseMessageChunk):
+        if isinstance(chunk.content, list):
             raise ValueError(
-                "The `message.content` provided seems to represent numerous messages. "
-                "Consider iterating over `message.content` and calling .append_message() on each iteration."
+                "The `chunk.content` provided seems to represent numerous message chunks. "
+                "Consider iterating over `chunk.content` and calling .append_message() on each iteration."
             )
-        return ChatMessage(content=x.content, role="assistant")
-
-    def can_normalize(self, message: Any) -> bool:
-        try:
-            from langchain_core.messages import BaseMessage
-
-            return isinstance(message, BaseMessage)
-        except Exception:
-            return False
-
-    def can_normalize_chunk(self, chunk: Any) -> bool:
-        try:
-            from langchain_core.messages import BaseMessageChunk
-
-            return isinstance(chunk, BaseMessageChunk)
-        except Exception:
-            return False
+        return ChatMessage(
+            content=chunk.content,
+            role="assistant",
+        )
+except ImportError:
+    pass
 
 
-class OpenAINormalizer(StringNormalizer):
-    def normalize(self, message: Any) -> ChatMessage:
-        x = cast("ChatCompletion", message)
-        return super().normalize(x.choices[0].message.content)
+# ------------------------------------------------------------------
+# OpenAI content extractor
+# ------------------------------------------------------------------
 
-    def normalize_chunk(self, chunk: Any) -> ChatMessage:
-        x = cast("ChatCompletionChunk", chunk)
-        return super().normalize_chunk(x.choices[0].delta.content)
+try:
+    from openai.types.chat import ChatCompletion, ChatCompletionChunk
 
-    def can_normalize(self, message: Any) -> bool:
-        try:
-            from openai.types.chat import ChatCompletion
+    @message_content.register
+    def _(message: ChatCompletion):
+        return ChatMessage(
+            content=message.choices[0].message.content,
+            role="assistant",
+        )
 
-            return isinstance(message, ChatCompletion)
-        except Exception:
-            return False
+    @message_content_chunk.register
+    def _(chunk: ChatCompletionChunk):
+        return ChatMessage(
+            content=chunk.choices[0].delta.content,
+            role="assistant",
+        )
+except ImportError:
+    pass
 
-    def can_normalize_chunk(self, chunk: Any) -> bool:
-        try:
-            from openai.types.chat import ChatCompletionChunk
 
-            return isinstance(chunk, ChatCompletionChunk)
-        except Exception:
-            return False
+# ------------------------------------------------------------------
+# Anthropic content extractor
+# ------------------------------------------------------------------
 
+try:
+    from anthropic.types import (  # pyright: ignore[reportMissingImports]
+        Message as AnthropicMessage,
+    )
 
-class AnthropicNormalizer(BaseMessageNormalizer):
-    def normalize(self, message: Any) -> ChatMessage:
-        x = cast("AnthropicMessage", message)
-        content = x.content[0]
+    @message_content.register
+    def _(message: AnthropicMessage):
+        content = message.content[0]
         if content.type != "text":
             raise ValueError(
                 f"Anthropic message type {content.type} not supported. "
                 "Only 'text' type is currently supported"
             )
-        return ChatMessage(content=content.text, role="assistant")
-
-    def normalize_chunk(self, chunk: Any) -> ChatMessage:
-        x = cast("MessageStreamEvent", chunk)
-        content = ""
-        if x.type == "content_block_delta":
-            if x.delta.type != "text_delta":
-                raise ValueError(
-                    f"Anthropic message delta type {x.delta.type} not supported. "
-                    "Only 'text_delta' type is supported"
-                )
-            content = x.delta.text
-
-        return ChatMessage(content=content, role="assistant")
-
-    def can_normalize(self, message: Any) -> bool:
-        try:
-            from anthropic.types import Message as AnthropicMessage
-
-            return isinstance(message, AnthropicMessage)
-        except Exception:
-            return False
-
-    def can_normalize_chunk(self, chunk: Any) -> bool:
-        try:
-            from anthropic.types import (
-                RawContentBlockDeltaEvent,
-                RawContentBlockStartEvent,
-                RawContentBlockStopEvent,
-                RawMessageDeltaEvent,
-                RawMessageStartEvent,
-                RawMessageStopEvent,
-            )
+        return ChatMessage(content=content.text)
 
-            # The actual MessageStreamEvent is a generic, so isinstance() can't
-            # be used to check the type. Instead, we manually construct the relevant
-            # union of relevant classes...
-            return (
-                isinstance(chunk, RawContentBlockDeltaEvent)
-                or isinstance(chunk, RawContentBlockStartEvent)
-                or isinstance(chunk, RawContentBlockStopEvent)
-                or isinstance(chunk, RawMessageDeltaEvent)
-                or isinstance(chunk, RawMessageStartEvent)
-                or isinstance(chunk, RawMessageStopEvent)
-            )
-        except Exception:
-            return False
+    # Old versions of singledispatch doesn't seem to support union types
+    if sys.version_info >= (3, 11):
+        from anthropic.types import (  # pyright: ignore[reportMissingImports]
+            RawMessageStreamEvent,
+        )
 
+        @message_content_chunk.register
+        def _(chunk: RawMessageStreamEvent):
+            content = ""
+            if chunk.type == "content_block_delta":
+                if chunk.delta.type != "text_delta":
+                    raise ValueError(
+                        f"Anthropic message delta type {chunk.delta.type} not supported. "
+                        "Only 'text_delta' type is supported"
+                    )
+                content = chunk.delta.text
+
+            return ChatMessage(content=content)
+except ImportError:
+    pass
+
+
+# ------------------------------------------------------------------
+# Google content extractor
+# ------------------------------------------------------------------
+
+try:
+    from google.generativeai.types.generation_types import (
+        GenerateContentResponse,
+    )
 
-class GoogleNormalizer(BaseMessageNormalizer):
-    def normalize(self, message: Any) -> ChatMessage:
-        x = cast("GenerateContentResponse", message)
-        return ChatMessage(content=x.text, role="assistant")
+    @message_content.register
+    def _(message: GenerateContentResponse):
+        return ChatMessage(content=message.text)
 
-    def normalize_chunk(self, chunk: Any) -> ChatMessage:
-        x = cast("GenerateContentResponse", chunk)
-        return ChatMessage(content=x.text, role="assistant")
+    @message_content_chunk.register
+    def _(chunk: GenerateContentResponse):
+        return ChatMessage(content=chunk.text)
 
-    def can_normalize(self, message: Any) -> bool:
-        try:
-            import google.generativeai.types.generation_types as gtypes  # pyright: ignore[reportMissingTypeStubs, reportMissingImports]
+except ImportError:
+    pass
 
-            return isinstance(
-                message,
-                gtypes.GenerateContentResponse,  # pyright: ignore[reportUnknownMemberType]
-            )
-        except Exception:
-            return False
-
-    def can_normalize_chunk(self, chunk: Any) -> bool:
-        return self.can_normalize(chunk)
-
-
-class OllamaNormalizer(DictNormalizer):
-    def normalize(self, message: Any) -> ChatMessage:
-        x = cast("dict[str, Any]", message["message"])
-        return super().normalize(x)
-
-    def normalize_chunk(self, chunk: "dict[str, Any]") -> ChatMessage:
-        msg = cast("dict[str, Any]", chunk["message"])
-        return super().normalize_chunk(msg)
-
-    def can_normalize(self, message: Any) -> bool:
-        try:
-            from ollama import ChatResponse
-
-            # Ollama<0.4 used TypedDict (now it uses pydantic)
-            # https://github.com/ollama/ollama-python/pull/276
-            if isinstance(ChatResponse, dict):
-                return "message" in message and super().can_normalize(
-                    message["message"]
-                )
-            else:
-                return isinstance(message, ChatResponse)
-        except Exception:
-            return False
-
-    def can_normalize_chunk(self, chunk: Any) -> bool:
-        return self.can_normalize(chunk)
-
-
-class NormalizerRegistry:
-    def __init__(self) -> None:
-        # Order of strategies matters (the 1st one that can normalize the message is used)
-        # So make sure to put the most specific strategies first
-        self._strategies: dict[str, BaseMessageNormalizer] = {
-            "openai": OpenAINormalizer(),
-            "anthropic": AnthropicNormalizer(),
-            "google": GoogleNormalizer(),
-            "langchain": LangChainNormalizer(),
-            "ollama": OllamaNormalizer(),
-            "tagify": TagifiableNormalizer(),
-            "dict": DictNormalizer(),
-            "string": StringNormalizer(),
-        }
-
-    def register(
-        self, provider: str, strategy: BaseMessageNormalizer, force: bool = False
-    ) -> None:
-        if provider in self._strategies:
-            if force:
-                del self._strategies[provider]
-            else:
-                raise ValueError(f"Provider {provider} already exists in registry")
-        # Update the strategies dict such that the new strategy is the first to be considered
-        self._strategies = {provider: strategy, **self._strategies}
-
-
-message_normalizer_registry = NormalizerRegistry()
-
-
-def register_custom_normalizer(
-    provider: str, normalizer: BaseMessageNormalizer, force: bool = False
-) -> None:
-    """
-    Register a custom normalizer for handling specific message types.
 
-    Parameters
-    ----------
-    provider : str
-        A unique identifier for this normalizer in the registry
-    normalizer : BaseMessageNormalizer
-        A normalizer instance that can handle your specific message type
-    force : bool, optional
-        Whether to override an existing normalizer with the same provider name,
-        by default False
-
-    Examples
-    --------
-    >>> class MyCustomMessage:
-    ...     def __init__(self, content):
-    ...         self.content = content
-    ...
-    >>> class MyCustomNormalizer(StringNormalizer):
-    ...     def normalize(self, message):
-    ...         return ChatMessage(content=message.content, role="assistant")
-    ...     def can_normalize(self, message):
-    ...         return isinstance(message, MyCustomMessage)
-    ...
-    >>> register_custom_normalizer("my_provider", MyCustomNormalizer())
-    """
-    message_normalizer_registry.register(provider, normalizer, force)
+# ------------------------------------------------------------------
+# Ollama content extractor
+# ------------------------------------------------------------------
 
+try:
+    from ollama import ChatResponse
 
-def normalize_message(message: Any) -> ChatMessage:
-    strategies = message_normalizer_registry._strategies
-    for strategy in strategies.values():
-        if strategy.can_normalize(message):
-            return strategy.normalize(message)
-    raise ValueError(
-        f"Could not find a normalizer for message of type {type(message)}: {message}. "
-        "Consider registering a custom normalizer via shiny.ui._chat_types.registry.register()"
-    )
+    @message_content.register
+    def _(message: ChatResponse):
+        msg = message.message
+        return ChatMessage(msg.content)
 
+    @message_content_chunk.register
+    def _(chunk: ChatResponse):
+        msg = chunk.message
+        return ChatMessage(msg.content)
 
-def normalize_message_chunk(chunk: Any) -> ChatMessage:
-    strategies = message_normalizer_registry._strategies
-    for strategy in strategies.values():
-        if strategy.can_normalize_chunk(chunk):
-            return strategy.normalize_chunk(chunk)
-    raise ValueError(
-        f"Could not find a normalizer for message chunk of type {type(chunk)}: {chunk}. "
-        "Consider registering a custom normalizer via shiny.ui._chat_normalize.register_custom_normalizer()"
-    )
+except ImportError:
+    pass