waldiez 0.4.7__py3-none-any.whl → 0.4.9__py3-none-any.whl

waldiez/io/_ws.py

@@ -0,0 +1,199 @@
+# SPDX-License-Identifier: Apache-2.0.
+# Copyright (c) 2024 - 2025 Waldiez and contributors.
+# flake8: noqa: E501
+# pylint: disable=line-too-long
+"""WebSocket IOStream implementation for AsyncIO."""
+
+import asyncio
+from typing import Any, Protocol
+
+HAS_WS_LIB = False
+
+try:
+    from starlette.websockets import WebSocket  # type: ignore[unused-ignore, unused-import, import-not-found, import-untyped]  # noqa
+
+    HAS_WS_LIB = True  # pyright: ignore
+except ImportError:
+    pass
+
+try:
+    import websockets  # type: ignore[unused-ignore, unused-import, import-not-found, import-untyped]  # noqa
+
+    HAS_WS_LIB = True  # pyright: ignore
+except ImportError:  # pragma: no cover
+    pass
+
+
+class WebSocketConnection(Protocol):
+    """Protocol for WebSocket connections."""
+
+    async def send_message(self, message: str) -> None:
+        """Send a message over the WebSocket connection.
+
+        Parameters
+        ----------
+        message : str
+            The message to send.
+        """
+
+    async def receive_message(
+        self,
+        timeout: float = 120,
+    ) -> str:  # pyright: ignore
+        """Receive a message from the WebSocket connection.
+
+        Parameters
+        ----------
+        timeout : float, optional
+            The timeout for receiving the message. Defaults to 120 seconds.
+            If the timeout is reached, an empty string is returned.
+            If an error occurs, an empty string is returned.
+
+        Returns
+        -------
+        str
+            The received message.
+        """
+
+
+class WebSocketsAdapter:
+    """Adapter for websockets library connections."""
+
+    def __init__(self, websocket: "websockets.ServerConnection"):
+        """Initialize the adapter.
+
+        Parameters
+        ----------
+        websocket : ServerConnection
+            The websockets library connection.
+        """
+        self.websocket = websocket
+
+    async def send_message(self, message: str) -> None:
+        """Send a message using websockets library.
+
+        Parameters
+        ----------
+        message : str
+            The message to send.
+        """
+        await self.websocket.send(message)
+
+    async def receive_message(self, timeout: float = 120) -> str:
+        """Receive a message using websockets library.
+
+        Parameters
+        ----------
+        timeout : float, optional
+            The timeout for receiving the message. Defaults to 120 seconds.
+
+        Returns
+        -------
+        str
+            The received message, decoded as a string.
+        """
+        # pylint: disable=too-many-try-statements
+        try:
+            response = await asyncio.wait_for(
+                self.websocket.recv(), timeout=timeout
+            )
+            if isinstance(response, bytes):
+                return response.decode("utf-8")
+            return response if isinstance(response, str) else str(response)
+        except asyncio.TimeoutError:
+            return ""
+        except BaseException:  # pylint: disable=broad-exception-caught
+            # Handle other exceptions that may occur during receive
+            return ""
+
+
+class StarletteAdapter:
+    """Adapter for Starlette/FastAPI WebSocket connections."""
+
+    def __init__(self, websocket: "WebSocket"):
+        """Initialize the adapter.
+
+        Parameters
+        ----------
+        websocket : WebSocket
+            The Starlette/FastAPI WebSocket connection.
+        """
+        self.websocket = websocket
+
+    async def send_message(self, message: str) -> None:
+        """Send a message using Starlette WebSocket.
+
+        Parameters
+        ----------
+        message : str
+            The message to send.
+        """
+        await self.websocket.send_text(message)
+
+    async def receive_message(self, timeout: float = 120) -> str:
+        """Receive a message using Starlette WebSocket.
+
+        Parameters
+        ----------
+        timeout : float, optional
+            The timeout for receiving the message. Defaults to 120 seconds.
+
+        Returns
+        -------
+        str
+            The received message, decoded as a string.
+        """
+        try:
+            return await asyncio.wait_for(
+                self.websocket.receive_text(), timeout=timeout
+            )
+        except asyncio.TimeoutError:
+            return ""
+        except BaseException:  # pylint: disable=broad-exception-caught
+            # Handle other exceptions that may occur during receive
+            return ""
+
+
+def create_websocket_adapter(websocket: Any) -> WebSocketConnection:
+    """Create an appropriate adapter for the given WebSocket connection.
+
+    Parameters
+    ----------
+    websocket : Any
+        The WebSocket connection (websockets or Starlette).
+
+    Returns
+    -------
+    WebSocketConnection
+        An adapter that implements the WebSocketConnection protocol.
+
+    Raises
+    ------
+    ValueError
+        If the WebSocket type is not supported.
+    """
+    # Check for Starlette/FastAPI WebSocket
+    if hasattr(websocket, "send_text") and hasattr(websocket, "receive_text"):
+        return StarletteAdapter(websocket)
+
+    # Check for websockets library
+    if hasattr(websocket, "send") and hasattr(websocket, "recv"):
+        return WebSocketsAdapter(websocket)
+
+    raise ValueError(
+        "Unsupported WebSocket type. "
+        "Must be either websockets.ServerConnection "
+        "or starlette.websockets.WebSocket. "
+        f"Received: {type(websocket)}"
+    )
+
+
+def is_websocket_available() -> bool:
+    """Check if any WebSocket library is available.
+
+    Returns
+    -------
+    bool
+        True if websockets or starlette is available, False otherwise.
+    """
+    return HAS_WS_LIB

waldiez/io/models/__init__.py

@@ -0,0 +1,60 @@
+# SPDX-License-Identifier: Apache-2.0.
+# Copyright (c) 2024 - 2025 Waldiez and contributors.
+
+"""Models package for structured I/O."""
+
+# Base models
+from .base import PrintMessage, StructuredBase, UserInputRequest
+
+# Constants and mappings
+from .constants import (
+    CONTENT_MAPPING,
+    ContentMappingEntry,
+    ContentTypeKey,
+    MediaContent,
+)
+
+# Content models
+from .content import (
+    AudioContent,
+    AudioMediaContent,
+    FileContent,
+    FileMediaContent,
+    ImageContent,
+    ImageMediaContent,
+    ImageUrlMediaContent,
+    TextMediaContent,
+    VideoContent,
+    VideoMediaContent,
+)
+
+# User models
+from .user_input import UserInputData
+from .user_response import UserResponse
+
+__all__ = [
+    # Base models
+    "StructuredBase",
+    "UserInputRequest",
+    "PrintMessage",
+    # Content base types
+    "ImageContent",
+    "VideoContent",
+    "AudioContent",
+    "FileContent",
+    "MediaContent",
+    # Content models
+    "TextMediaContent",
+    "ImageMediaContent",
+    "ImageUrlMediaContent",
+    "VideoMediaContent",
+    "AudioMediaContent",
+    "FileMediaContent",
+    # User models
+    "UserInputData",
+    "UserResponse",
+    # Constants
+    "ContentTypeKey",
+    "ContentMappingEntry",
+    "CONTENT_MAPPING",
+]

waldiez/io/models/base.py

@@ -0,0 +1,66 @@
+# SPDX-License-Identifier: Apache-2.0.
+# Copyright (c) 2024 - 2025 Waldiez and contributors.
+
+"""Base models for structured data."""
+
+from typing import Any
+
+from pydantic import BaseModel, ConfigDict, Field
+
+from ..utils import MessageType, gen_id, now
+
+
+class StructuredBase(BaseModel):
+    """Base model for structured data."""
+
+    id: str = Field(default_factory=gen_id)
+    type: MessageType
+    timestamp: str = Field(default_factory=now)
+    model_config = ConfigDict(
+        extra="ignore",
+        arbitrary_types_allowed=True,
+    )
+
+
+class UserInputRequest(StructuredBase):
+    """User input prompt model."""
+
+    type: MessageType = "input_request"
+    request_id: str
+    prompt: str
+    password: bool = False
+
+
+class PrintMessage(StructuredBase):
+    """Model for print messages."""
+
+    type: MessageType = "print"
+    data: str
+
+    @classmethod
+    def create(cls, *args: Any, **kwargs: Any) -> "PrintMessage":
+        """Create a new print message.
+
+        Parameters
+        ----------
+        *args : Any
+            Positional arguments (not used).
+        **kwargs : Any
+            Keyword arguments
+
+        Returns
+        -------
+        PrintMessage
+            A new print message instance with the provided data.
+        """
+        message = " ".join(str(arg) for arg in args)
+        if "file" in kwargs:
+            file = kwargs.pop("file")
+            if hasattr(file, "getvalue"):
+                io_value = file.getvalue()
+                if isinstance(io_value, bytes):
+                    io_value = io_value.decode("utf-8", errors="replace")
+                message += io_value
+        end = kwargs.get("end", "\n")
+        message += end
+        return cls(data=message)

waldiez/io/models/constants.py

@@ -0,0 +1,78 @@
+# SPDX-License-Identifier: Apache-2.0.
+# Copyright (c) 2024 - 2025 Waldiez and contributors.
+
+"""Constants and type mappings for content models."""
+
+from typing import Type, TypedDict, Union
+
+from typing_extensions import Literal
+
+from .content.audio import AudioMediaContent
+from .content.file import FileMediaContent
+from .content.image import ImageMediaContent, ImageUrlMediaContent
+from .content.text import TextMediaContent
+from .content.video import VideoMediaContent
+
+MediaContent = Union[
+    TextMediaContent,
+    ImageMediaContent,
+    ImageUrlMediaContent,
+    VideoMediaContent,
+    AudioMediaContent,
+    FileMediaContent,
+]
+
+ContentTypeKey = Literal["text", "image", "image_url", "video", "audio", "file"]
+"""Possible content types for the mapping."""
+
+
+class ContentMappingEntry(TypedDict):
+    """TypedDict for content mapping entries.
+
+    Attributes
+    ----------
+    fields: List[str]
+        List of possible field names for the content type
+    cls: Type[MediaContent]
+        The class to instantiate for this content type
+    required_field: str
+        The field name required by the class constructor
+    """
+
+    fields: list[str]
+    cls: Type[MediaContent]
+    required_field: str
+
+
+CONTENT_MAPPING: dict[ContentTypeKey, ContentMappingEntry] = {
+    "text": {
+        "fields": ["text"],
+        "cls": TextMediaContent,
+        "required_field": "text",
+    },
+    "image": {
+        "fields": ["image"],
+        "cls": ImageMediaContent,
+        "required_field": "image",
+    },
+    "image_url": {
+        "fields": ["image_url", "url"],
+        "cls": ImageUrlMediaContent,
+        "required_field": "image_url",
+    },
+    "video": {
+        "fields": ["video"],
+        "cls": VideoMediaContent,
+        "required_field": "video",
+    },
+    "audio": {
+        "fields": ["audio"],
+        "cls": AudioMediaContent,
+        "required_field": "audio",
+    },
+    "file": {
+        "fields": ["file"],
+        "cls": FileMediaContent,
+        "required_field": "file",
+    },
+}

waldiez/io/models/content/__init__.py

@@ -0,0 +1,23 @@
+# SPDX-License-Identifier: Apache-2.0.
+# Copyright (c) 2024 - 2025 Waldiez and contributors.
+"""Content models for media types."""
+
+from .audio import AudioMediaContent
+from .base import AudioContent, FileContent, ImageContent, VideoContent
+from .file import FileMediaContent
+from .image import ImageMediaContent, ImageUrlMediaContent
+from .text import TextMediaContent
+from .video import VideoMediaContent
+
+__all__ = [
+    "AudioMediaContent",
+    "AudioContent",
+    "VideoMediaContent",
+    "VideoContent",
+    "ImageMediaContent",
+    "ImageUrlMediaContent",
+    "ImageContent",
+    "FileMediaContent",
+    "FileContent",
+    "TextMediaContent",
+]

waldiez/io/models/content/audio.py

@@ -0,0 +1,43 @@
+# SPDX-License-Identifier: Apache-2.0.
+# Copyright (c) 2024 - 2025 Waldiez and contributors.
+# pylint: disable=unused-argument
+"""Audio media content models."""
+
+from pathlib import Path
+from typing import Literal
+
+from pydantic import BaseModel
+
+from .base import AudioContent
+
+
+class AudioMediaContent(BaseModel):
+    """Audio media content."""
+
+    type: Literal["audio"] = "audio"
+    audio: AudioContent
+
+    def to_string(
+        self,
+        uploads_root: Path | None,
+        base_name: str | None = None,
+    ) -> str:
+        """Convert the content to a string.
+
+        Parameters
+        ----------
+        uploads_root : Path | None
+            The root directory for storing images, optional.
+        base_name : str | None
+            The base name for the image file, optional.
+
+        Returns
+        -------
+        str
+            The string representation of the content.
+        """
+        if self.audio.url:
+            return f"<audio src='{self.audio.url}'></audio>"
+        if self.audio.file:
+            return f"<audio src='{self.audio.file}'></audio>"
+        return str(self.audio.file)

waldiez/io/models/content/base.py

@@ -0,0 +1,45 @@
+# SPDX-License-Identifier: Apache-2.0.
+# Copyright (c) 2024 - 2025 Waldiez and contributors.
+"""Base content models for media types."""
+
+from typing import Any, Optional
+
+from pydantic import BaseModel
+
+
+class ImageContent(BaseModel):
+    """Image content model."""
+
+    url: Optional[str] = None
+    file: Optional[Any] = None  # File type not directly mappable in Python
+    alt: Optional[str] = None
+
+
+class VideoContent(BaseModel):
+    """Video content model."""
+
+    url: Optional[str] = None
+    file: Optional[Any] = None
+    duration: Optional[int] = None
+    thumbnailUrl: Optional[str] = None
+    mimeType: Optional[str] = None
+
+
+class AudioContent(BaseModel):
+    """Audio content model."""
+
+    url: Optional[str] = None
+    file: Optional[Any] = None
+    duration: Optional[int] = None
+    transcript: Optional[str] = None
+
+
+class FileContent(BaseModel):
+    """File content model."""
+
+    url: Optional[str] = None
+    file: Optional[Any] = None
+    name: str
+    size: Optional[int] = None
+    type: Optional[str] = None
+    previewUrl: Optional[str] = None

waldiez/io/models/content/file.py

@@ -0,0 +1,43 @@
+# SPDX-License-Identifier: Apache-2.0.
+# Copyright (c) 2024 - 2025 Waldiez and contributors.
+# pylint: disable=unused-argument
+"""File media content models."""
+
+from pathlib import Path
+from typing import Literal
+
+from pydantic import BaseModel
+
+from .base import FileContent
+
+
+class FileMediaContent(BaseModel):
+    """File media content."""
+
+    type: Literal["file"] = "file"
+    file: FileContent
+
+    def to_string(
+        self,
+        uploads_root: Path | None,
+        base_name: str | None = None,
+    ) -> str:
+        """Convert the content to a string.
+
+        Parameters
+        ----------
+        uploads_root : Path | None
+            The root directory for storing images, optional.
+        base_name : str | None
+            The base name for the image file, optional.
+
+        Returns
+        -------
+        str
+            The string representation of the content.
+        """
+        if self.file.url:
+            return f"<a href='{self.file.url}'>{self.file.name}</a>"
+        if self.file.file:
+            return f"<a href='{self.file.file}'>{self.file.name}</a>"
+        return str(self.file.file)

waldiez/io/models/content/image.py

@@ -0,0 +1,96 @@
+# SPDX-License-Identifier: Apache-2.0.
+# Copyright (c) 2024 - 2025 Waldiez and contributors.
+# pylint: disable=unused-argument
+"""Image media content models."""
+
+from pathlib import Path
+
+from pydantic import BaseModel
+from typing_extensions import Literal
+
+from ...utils import get_image
+from .base import ImageContent
+
+
+class ImageMediaContent(BaseModel):
+    """Image media content."""
+
+    type: Literal["image"] = "image"
+    image: ImageContent
+
+    def to_string(
+        self,
+        uploads_root: Path | None,
+        base_name: str | None = None,
+    ) -> str:
+        """Convert the content to a string.
+
+        Parameters
+        ----------
+        uploads_root : Path | None
+            The root directory for storing images, optional.
+        base_name : str | None
+            The base name for the image file, optional.
+
+        Returns
+        -------
+        str
+            The string representation of the content.
+        """
+        if self.image.url:
+            image = get_image(
+                uploads_root=uploads_root,
+                image_data=self.image.url,
+                base_name=base_name,
+            )
+            return f"<img {image}>"
+        if self.image.file:
+            image = get_image(
+                uploads_root=uploads_root,
+                image_data=self.image.file,
+                base_name=base_name,
+            )
+            return f"<img {image}>"
+        return str(self.image.file)
+
+
+class ImageUrlMediaContent(BaseModel):
+    """Image URL media content."""
+
+    type: Literal["image_url"] = "image_url"
+    image_url: ImageContent
+
+    def to_string(
+        self,
+        uploads_root: Path | None,
+        base_name: str | None = None,
+    ) -> str:
+        """Convert the content to a string.
+
+        Parameters
+        ----------
+        uploads_root : Path | None
+            The root directory for storing images, optional.
+        base_name : str | None
+            The base name for the image file, optional.
+
+        Returns
+        -------
+        str
+            The string representation of the content.
+        """
+        if self.image_url.url:
+            image = get_image(
+                uploads_root=uploads_root,
+                image_data=self.image_url.url,
+                base_name=base_name,
+            )
+            return f"<img {image}>"
+        if self.image_url.file:
+            image = get_image(
+                uploads_root=uploads_root,
+                image_data=self.image_url.file,
+                base_name=base_name,
+            )
+            return f"<img {image}>"
+        return str(self.image_url.file)

waldiez/io/models/content/text.py

@@ -0,0 +1,37 @@
+# SPDX-License-Identifier: Apache-2.0.
+# Copyright (c) 2024 - 2025 Waldiez and contributors.
+# pylint: disable=unused-argument
+"""Text media content models."""
+
+from pathlib import Path
+
+from pydantic import BaseModel
+from typing_extensions import Literal
+
+
+class TextMediaContent(BaseModel):
+    """Text media content."""
+
+    type: Literal["text"] = "text"
+    text: str
+
+    def to_string(
+        self,
+        uploads_root: Path | None = None,
+        base_name: str | None = None,
+    ) -> str:
+        """Convert the content to a string.
+
+        Parameters
+        ----------
+        uploads_root : Path | None
+            The root directory for storing images, optional.
+        base_name : str | None
+            The base name for the image file, optional.
+
+        Returns
+        -------
+        str
+            The string representation of the content.
+        """
+        return self.text