fix up openai vision and gemini implementation

src/pipecat/processors/aggregators/openai_llm_context.py

@@ -6,6 +6,7 @@
 
 from dataclasses import dataclass
 import io
+import json
 
 from typing import List
 
@@ -21,6 +22,17 @@ from openai.types.chat import (
     ChatCompletionMessageParam
 )
 
+# JSON custom encoder to handle bytes arrays so that we can log contexts
+# with images to the console.
+
+
+class CustomEncoder(json.JSONEncoder):
+    def default(self, obj):
+        if isinstance(obj, io.BytesIO):
+            # Convert the first 8 bytes to an ASCII hex string
+            return (f"{obj.getbuffer()[0:8].hex()}...")
+        return super().default(obj)
+
 
 class OpenAILLMContext:
 
@@ -66,7 +78,7 @@ class OpenAILLMContext:
         context.add_message({
             "content": frame.text,
             "role": "user",
-            "data": buffer.getvalue(),
+            "data": buffer,
             "mime_type": "image/jpeg"
         })
         return context
@@ -77,6 +89,10 @@ class OpenAILLMContext:
     def get_messages(self) -> List[ChatCompletionMessageParam]:
         return self.messages
 
+    def get_messages_json(self) -> str:
+        return json.dumps(self.messages, cls=CustomEncoder)
+        # return json.dumps(self.messages)
+
     def set_tool_choice(
         self, tool_choice: ChatCompletionToolChoiceOptionParam | NotGiven
     ):

src/pipecat/services/google.py

@@ -1,7 +1,8 @@
-import google.generativeai as gai
-import google.ai.generativelanguage as glm
+
+import json
 import os
 import asyncio
+import time
 
 from typing import List
 
@@ -10,14 +11,26 @@ from pipecat.frames.frames import (
     TextFrame,
     VisionImageRawFrame,
     LLMMessagesFrame,
+    LLMFullResponseStartFrame,
     LLMResponseStartFrame,
-    LLMResponseEndFrame)
+    LLMResponseEndFrame,
+    LLMFullResponseEndFrame
+)
 from pipecat.processors.frame_processor import FrameDirection
 from pipecat.services.ai_services import LLMService
 from pipecat.processors.aggregators.openai_llm_context import OpenAILLMContext, OpenAILLMContextFrame
 
 from loguru import logger
 
+try:
+    import google.generativeai as gai
+    import google.ai.generativelanguage as glm
+except ModuleNotFoundError as e:
+    logger.error(f"Exception: {e}")
+    logger.error(
+        "In order to use Google AI, you need to `pip install pipecat-ai[google]`. Also, set `GOOGLE_API_KEY` environment variable.")
+    raise Exception(f"Missing module: {e}")
+
 
 class GoogleLLMService(LLMService):
     """This class implements inference with Google's AI models
@@ -54,7 +67,7 @@ class GoogleLLMService(LLMService):
                 parts.append(
                     glm.Part(inline_data=glm.Blob(
                         mime_type=message["mime_type"],
-                        data=message["data"]
+                        data=message["data"].getvalue()
                     )))
             google_messages.append({"role": role, "parts": parts})
 
@@ -66,19 +79,25 @@ class GoogleLLMService(LLMService):
             await asyncio.sleep(0)
 
     async def _process_context(self, context: OpenAILLMContext):
+        await self.push_frame(LLMFullResponseStartFrame())
         try:
+            logger.debug(f"Generating chat: {context.get_messages_json()}")
+
             messages = self._get_messages_from_openai_context(context)
 
-            await self.push_frame(LLMResponseStartFrame())
+            start_time = time.time()
             response = self._client.generate_content(messages, stream=True)
+            logger.debug(f"Google LLM TTFB: {time.time() - start_time}")
 
             async for chunk in self._async_generator_wrapper(response):
-                logger.debug(f"Pushing inference text: {chunk.text}")
+                await self.push_frame(LLMResponseStartFrame())
                 await self.push_frame(TextFrame(chunk.text))
+                await self.push_frame(LLMResponseEndFrame())
 
-            await self.push_frame(LLMResponseEndFrame())
         except Exception as e:
             logger.error(f"Exception: {e}")
+        finally:
+            await self.push_frame(LLMFullResponseEndFrame())
 
     async def process_frame(self, frame: Frame, direction: FrameDirection):
         context = None

src/pipecat/services/openai.py

@@ -68,12 +68,14 @@ class BaseOpenAILLMService(LLMService):
     async def _stream_chat_completions(
         self, context: OpenAILLMContext
     ) -> AsyncStream[ChatCompletionChunk]:
+        logger.debug(f"Generating chat: {context.get_messages_json()}")
+
         messages: List[ChatCompletionMessageParam] = context.get_messages()
 
         # base64 encode any images
         for message in messages:
             if message.get("mime_type") == "image/jpeg":
-                encoded_image = base64.b64encode(message["data"]).decode("utf-8")
+                encoded_image = base64.b64encode(message["data"].getvalue()).decode("utf-8")
                 text = message["content"]
                 message["content"] = [
                     {"type": "text", "text": text},
@@ -82,9 +84,6 @@ class BaseOpenAILLMService(LLMService):
                 del message["data"]
                 del message["mime_type"]
 
-        messages_for_log = json.dumps(messages)
-        logger.debug(f"Generating chat: {messages_for_log}")
-
         start_time = time.time()
         chunks: AsyncStream[ChatCompletionChunk] = (
             await self._client.chat.completions.create(
@@ -101,10 +100,6 @@ class BaseOpenAILLMService(LLMService):
         return chunks
 
     async def _chat_completions(self, messages) -> str | None:
-        messages_for_log = json.dumps(messages)
-
-        logger.debug(f"Generating chat: {messages_for_log}")
-
         response: ChatCompletion = await self._client.chat.completions.create(
             model=self._model, stream=False, messages=messages
         )