Add audio input/output support to OpenAiChatPromptDriver

Author: collindutter

griptape/common/__init__.py

@@ -4,8 +4,11 @@
 from .prompt_stack.contents.base_message_content import BaseMessageContent
 from .prompt_stack.contents.base_delta_message_content import BaseDeltaMessageContent
 from .prompt_stack.contents.text_delta_message_content import TextDeltaMessageContent
+from .prompt_stack.contents.audio_delta_message_content import AudioDeltaMessageContent
+from .prompt_stack.contents.audio_transcript_delta_message_content import AudioTranscriptDeltaMessageContent
 from .prompt_stack.contents.text_message_content import TextMessageContent
 from .prompt_stack.contents.image_message_content import ImageMessageContent
+from .prompt_stack.contents.audio_message_content import AudioMessageContent
 from .prompt_stack.contents.action_call_delta_message_content import ActionCallDeltaMessageContent
 from .prompt_stack.contents.action_call_message_content import ActionCallMessageContent
 from .prompt_stack.contents.action_result_message_content import ActionResultMessageContent
@@ -30,8 +33,11 @@
     "DeltaMessage",
     "Message",
     "TextDeltaMessageContent",
+    "AudioDeltaMessageContent",
+    "AudioTranscriptDeltaMessageContent",
     "TextMessageContent",
     "ImageMessageContent",
+    "AudioMessageContent",
     "GenericMessageContent",
     "ActionCallDeltaMessageContent",
     "ActionCallMessageContent",

griptape/common/prompt_stack/contents/audio_delta_message_content.py

@@ -0,0 +1,24 @@
+from __future__ import annotations
+
+from typing import Optional
+
+from attrs import define, field
+
+from griptape.common import BaseDeltaMessageContent
+
+
+@define
+class AudioDeltaMessageContent(BaseDeltaMessageContent):
+    """A delta message content for audio data.
+
+    Attributes:
+        id: The ID of the audio data.
+        data: Base64 encoded audio data.
+        transcript: The transcript of the audio data.
+        expires_at: The Unix timestamp (in seconds) for when this audio data will no longer be accessible.
+    """
+
+    id: Optional[str] = field(default=None, kw_only=True, metadata={"serializable": True})
+    data: Optional[str] = field(kw_only=True, metadata={"serializable": True})
+    transcript: Optional[str] = field(default=None, kw_only=True, metadata={"serializable": True})
+    expires_at: Optional[int] = field(default=None, kw_only=True, metadata={"serializable": True})

griptape/common/prompt_stack/contents/audio_message_content.py

@@ -0,0 +1,43 @@
+from __future__ import annotations
+
+import base64
+from typing import TYPE_CHECKING
+
+from attrs import define, field
+
+from griptape.artifacts import AudioArtifact
+from griptape.common import (
+    AudioDeltaMessageContent,
+    BaseDeltaMessageContent,
+    BaseMessageContent,
+)
+
+if TYPE_CHECKING:
+    from collections.abc import Sequence
+
+
+@define
+class AudioMessageContent(BaseMessageContent):
+    artifact: AudioArtifact = field(metadata={"serializable": True})
+
+    @classmethod
+    def from_deltas(cls, deltas: Sequence[BaseDeltaMessageContent]) -> AudioMessageContent:
+        audio_deltas = [delta for delta in deltas if isinstance(delta, AudioDeltaMessageContent)]
+        audio_data = [delta.data for delta in audio_deltas if delta.data is not None]
+        transcript_data = [delta.transcript for delta in audio_deltas if delta.transcript is not None]
+        expires_at = next(delta.expires_at for delta in audio_deltas if delta.expires_at is not None)
+        audio_id = next(delta.id for delta in audio_deltas if delta.id is not None)
+
+        audio_transcript = "".join(data for data in transcript_data)
+
+        artifact = AudioArtifact(
+            value=b"".join(base64.b64decode(data) for data in audio_data),
+            format="wav",
+            meta={
+                "audio_id": audio_id,
+                "expires_at": expires_at,
+                "transcript": audio_transcript,
+            },
+        )
+
+        return cls(artifact=artifact)

griptape/common/prompt_stack/contents/audio_transcript_delta_message_content.py

@@ -0,0 +1,10 @@
+from __future__ import annotations
+
+from attrs import define, field
+
+from griptape.common import BaseDeltaMessageContent
+
+
+@define
+class AudioTranscriptDeltaMessageContent(BaseDeltaMessageContent):
+    text: str = field(metadata={"serializable": True})

griptape/common/prompt_stack/messages/message.py

@@ -37,6 +37,9 @@ def has_any_content_type(self, content_type: type[T]) -> bool:
     def get_content_type(self, content_type: type[T]) -> list[T]:
         return [content for content in self.content if isinstance(content, content_type)]
 
+    def exclude_content_type(self, content_type: type[T] | tuple[type[T]]) -> list[BaseMessageContent]:
+        return [content for content in self.content if not isinstance(content, content_type)]
+
     def is_text(self) -> bool:
         return all(isinstance(content, TextMessageContent) for content in self.content)
 

griptape/common/prompt_stack/prompt_stack.py

@@ -6,6 +6,7 @@
 
 from griptape.artifacts import (
     ActionArtifact,
+    AudioArtifact,
     BaseArtifact,
     GenericArtifact,
     ImageArtifact,
@@ -15,6 +16,7 @@
 from griptape.common import (
     ActionCallMessageContent,
     ActionResultMessageContent,
+    AudioMessageContent,
     BaseMessageContent,
     GenericMessageContent,
     ImageMessageContent,
@@ -77,6 +79,8 @@ def __to_message_content(self, artifact: str | BaseArtifact) -> list[BaseMessage
             return [TextMessageContent(artifact)]
         elif isinstance(artifact, ImageArtifact):
             return [ImageMessageContent(artifact)]
+        elif isinstance(artifact, AudioArtifact):
+            return [AudioMessageContent(artifact)]
         elif isinstance(artifact, GenericArtifact):
             return [GenericMessageContent(artifact)]
         elif isinstance(artifact, ActionArtifact):

griptape/drivers/prompt/base_prompt_driver.py

@@ -9,6 +9,8 @@
 from griptape.common import (
     ActionCallDeltaMessageContent,
     ActionCallMessageContent,
+    AudioDeltaMessageContent,
+    AudioMessageContent,
     BaseDeltaMessageContent,
     DeltaMessage,
     Message,
@@ -19,6 +21,7 @@
 )
 from griptape.events import (
     ActionChunkEvent,
+    AudioChunkEvent,
     EventBus,
     FinishPromptEvent,
     StartPromptEvent,
@@ -177,6 +180,8 @@ def __process_stream(self, prompt_stack: PromptStack) -> Message:
                     delta_contents[content.index] = [content]
                 if isinstance(content, TextDeltaMessageContent):
                     EventBus.publish_event(TextChunkEvent(token=content.text, index=content.index))
+                elif isinstance(content, AudioDeltaMessageContent) and content.data is not None:
+                    EventBus.publish_event(AudioChunkEvent(data=content.data))
                 elif isinstance(content, ActionCallDeltaMessageContent):
                     EventBus.publish_event(
                         ActionChunkEvent(
@@ -197,10 +202,13 @@ def __build_message(
         content = []
         for delta_content in delta_contents:
             text_deltas = [delta for delta in delta_content if isinstance(delta, TextDeltaMessageContent)]
+            audio_deltas = [delta for delta in delta_content if isinstance(delta, AudioDeltaMessageContent)]
             action_deltas = [delta for delta in delta_content if isinstance(delta, ActionCallDeltaMessageContent)]
 
             if text_deltas:
                 content.append(TextMessageContent.from_deltas(text_deltas))
+            if audio_deltas:
+                content.append(AudioMessageContent.from_deltas(audio_deltas))
             if action_deltas:
                 content.append(ActionCallMessageContent.from_deltas(action_deltas))
 

griptape/drivers/prompt/openai_chat_prompt_driver.py

@@ -1,18 +1,22 @@
 from __future__ import annotations
 
+import base64
 import json
 import logging
+import time
 from typing import TYPE_CHECKING, Optional
 
 import openai
 from attrs import Factory, define, field
 from schema import Schema
 
-from griptape.artifacts import ActionArtifact, TextArtifact
+from griptape.artifacts import ActionArtifact, AudioArtifact, TextArtifact
 from griptape.common import (
     ActionCallDeltaMessageContent,
     ActionCallMessageContent,
     ActionResultMessageContent,
+    AudioDeltaMessageContent,
+    AudioMessageContent,
     BaseDeltaMessageContent,
     BaseMessageContent,
     DeltaMessage,
@@ -94,6 +98,10 @@ class OpenAiChatPromptDriver(BasePromptDriver):
         ),
         kw_only=True,
     )
+    modalities: list[str] = field(default=Factory(lambda: ["text"]), kw_only=True, metadata={"serializable": True})
+    audio: dict = field(
+        default=Factory(lambda: {"voice": "alloy", "format": "pcm16"}), kw_only=True, metadata={"serializable": True}
+    )
     _client: openai.OpenAI = field(default=None, kw_only=True, alias="client", metadata={"serializable": False})
 
     @lazy_property()
@@ -144,14 +152,18 @@ def try_stream(self, prompt_stack: PromptStack) -> Iterator[DeltaMessage]:
                 choice = chunk.choices[0]
                 delta = choice.delta
 
-                yield DeltaMessage(content=self.__to_prompt_stack_delta_message_content(delta))
+                content = self.__to_prompt_stack_delta_message_content(delta)
+                if content is not None:
+                    yield DeltaMessage(content=content)
 
     def _base_params(self, prompt_stack: PromptStack) -> dict:
         params = {
             "model": self.model,
             "temperature": self.temperature,
             "user": self.user,
             "seed": self.seed,
+            "modalities": self.modalities,
+            "audio": self.audio,
             **({"stop": self.tokenizer.stop_sequences} if self.tokenizer.stop_sequences else {}),
             **({"max_tokens": self.max_tokens} if self.max_tokens is not None else {}),
             **({"stream_options": {"include_usage": True}} if self.stream else {}),
@@ -196,45 +208,44 @@ def __to_openai_messages(self, messages: list[Message]) -> list[dict]:
         openai_messages = []
 
         for message in messages:
-            # If the message only contains textual content we can send it as a single content.
-            if message.is_text():
-                openai_messages.append({"role": self.__to_openai_role(message), "content": message.to_text()})
             # Action results must be sent as separate messages.
-            elif message.has_any_content_type(ActionResultMessageContent):
+
+            action_result_contents = message.get_content_type(ActionResultMessageContent)
+            # Action results must be sent as separate messages.
+            if action_result_contents:
                 openai_messages.extend(
                     {
-                        "role": self.__to_openai_role(message, action_result),
-                        "content": self.__to_openai_message_content(action_result),
-                        "tool_call_id": action_result.action.tag,
+                        "role": self.__to_openai_role(message, action_result_content),
+                        "content": self.__to_openai_message_content(action_result_content),
+                        "tool_call_id": action_result_content.action.tag,
                     }
-                    for action_result in message.get_content_type(ActionResultMessageContent)
+                    for action_result_content in action_result_contents
                 )
 
                 if message.has_any_content_type(TextMessageContent):
                     openai_messages.append({"role": self.__to_openai_role(message), "content": message.to_text()})
             else:
                 openai_message = {
                     "role": self.__to_openai_role(message),
-                    "content": [
-                        self.__to_openai_message_content(content)
-                        for content in [
-                            content for content in message.content if not isinstance(content, ActionCallMessageContent)
-                        ]
-                    ],
+                    "content": [],
                 }
+
+                for content in message.content:
+                    if isinstance(content, ActionCallMessageContent):
+                        if "tool_calls" not in openai_message:
+                            openai_message["tool_calls"] = []
+                        openai_message["tool_calls"].append(self.__to_openai_message_content(content))
+                    elif isinstance(content, AudioMessageContent) and message.is_assistant():
+                        openai_message["audio"] = {
+                            "id": content.artifact.meta["audio_id"],
+                        }
+                    else:
+                        openai_message["content"].append(self.__to_openai_message_content(content))
+
                 # Some OpenAi-compatible services don't accept an empty array for content
                 if not openai_message["content"]:
                     openai_message["content"] = ""
 
-                # Action calls must be attached to the message, not sent as content.
-                action_call_content = [
-                    content for content in message.content if isinstance(content, ActionCallMessageContent)
-                ]
-                if action_call_content:
-                    openai_message["tool_calls"] = [
-                        self.__to_openai_message_content(action_call) for action_call in action_call_content
-                    ]
-
                 openai_messages.append(openai_message)
 
         return openai_messages
@@ -272,6 +283,23 @@ def __to_openai_message_content(self, content: BaseMessageContent) -> str | dict
                 "type": "image_url",
                 "image_url": {"url": f"data:{content.artifact.mime_type};base64,{content.artifact.base64}"},
             }
+        elif isinstance(content, AudioMessageContent):
+            artifact = content.artifact
+
+            # We can't send the audio if it's expired.
+            if int(time.time()) > artifact.meta.get("expires_at", float("inf")):
+                return {
+                    "type": "text",
+                    "text": artifact.meta.get("transcript"),
+                }
+            else:
+                return {
+                    "type": "input_audio",
+                    "input_audio": {
+                        "data": base64.b64encode(artifact.value).decode("utf-8"),
+                        "format": artifact.format,
+                    },
+                }
         elif isinstance(content, ActionCallMessageContent):
             action = content.artifact.value
 
@@ -290,6 +318,20 @@ def __to_prompt_stack_message_content(self, response: ChatCompletionMessage) ->
 
         if response.content is not None:
             content.append(TextMessageContent(TextArtifact(response.content)))
+        if response.audio is not None:
+            content.append(
+                AudioMessageContent(
+                    AudioArtifact(
+                        value=base64.b64decode(response.audio.data),
+                        format="wav",
+                        meta={
+                            "audio_id": response.audio.id,
+                            "transcript": response.audio.transcript,
+                            "expires_at": response.audio.expires_at,
+                        },
+                    )
+                )
+            )
         if response.tool_calls is not None:
             content.extend(
                 [
@@ -309,7 +351,7 @@ def __to_prompt_stack_message_content(self, response: ChatCompletionMessage) ->
 
         return content
 
-    def __to_prompt_stack_delta_message_content(self, content_delta: ChoiceDelta) -> BaseDeltaMessageContent:
+    def __to_prompt_stack_delta_message_content(self, content_delta: ChoiceDelta) -> Optional[BaseDeltaMessageContent]:
         if content_delta.content is not None:
             return TextDeltaMessageContent(content_delta.content)
         elif content_delta.tool_calls is not None:
@@ -334,5 +376,13 @@ def __to_prompt_stack_delta_message_content(self, content_delta: ChoiceDelta) ->
                     raise ValueError(f"Unsupported tool call delta: {tool_call}")
             else:
                 raise ValueError(f"Unsupported tool call delta length: {len(tool_calls)}")
-        else:
-            return TextDeltaMessageContent("")
+        # OpenAi doesn't have types for audio deltas so we need to use hasattr and getattr.
+        elif hasattr(content_delta, "audio") and getattr(content_delta, "audio") is not None:
+            audio_chunk: dict = getattr(content_delta, "audio")
+            return AudioDeltaMessageContent(
+                id=audio_chunk.get("id"),
+                data=audio_chunk.get("data"),
+                expires_at=audio_chunk.get("expires_at"),
+                transcript=audio_chunk.get("transcript"),
+            )
+        return None

griptape/events/__init__.py

@@ -12,6 +12,7 @@
 from .finish_structure_run_event import FinishStructureRunEvent
 from .base_chunk_event import BaseChunkEvent
 from .text_chunk_event import TextChunkEvent
+from .audio_chunk_event import AudioChunkEvent
 from .action_chunk_event import ActionChunkEvent
 from .event_listener import EventListener
 from .start_image_generation_event import StartImageGenerationEvent
@@ -41,6 +42,7 @@
     "FinishStructureRunEvent",
     "BaseChunkEvent",
     "TextChunkEvent",
+    "AudioChunkEvent",
     "ActionChunkEvent",
     "EventListener",
     "StartImageGenerationEvent",