gemini function calling and partial implementation of standard context stuff

src/pipecat/services/google.py

@@ -5,10 +5,14 @@
 #
 
 import asyncio
+from dataclasses import dataclass
 import json
+import io
 from typing import AsyncGenerator, List, Literal, Optional
 
+
 from loguru import logger
+from PIL import Image
 from pydantic import BaseModel
 
 from pipecat.frames.frames import (
@@ -28,6 +32,10 @@ from pipecat.processors.aggregators.openai_llm_context import (
     OpenAILLMContext,
     OpenAILLMContextFrame,
 )
+from pipecat.services.openai import (
+    OpenAIAssistantContextAggregator,
+    OpenAIUserContextAggregator,
+)
 from pipecat.processors.frame_processor import FrameDirection
 from pipecat.services.ai_services import LLMService, TTSService
 from pipecat.transcriptions.language import Language
@@ -45,6 +53,148 @@ except ModuleNotFoundError as e:
     raise Exception(f"Missing module: {e}")
 
 
+class GoogleUserContextAggregator(OpenAIUserContextAggregator):
+    async def _push_aggregation(self):
+        if len(self._aggregation) > 0:
+            self._context.add_message({"role": "user", "parts": [glm.Part(text=self._aggregation)]})
+
+            # Reset the aggregation. Reset it before pushing it down, otherwise
+            # if the tasks gets cancelled we won't be able to clear things up.
+            self._aggregation = ""
+
+            frame = OpenAILLMContextFrame(self._context)
+            await self.push_frame(frame)
+
+            # Reset our accumulator state.
+            self._reset()
+
+
+class GoogleAssistantContextAggregator(OpenAIAssistantContextAggregator):
+    async def _push_aggregation(self):
+        if not (
+            self._aggregation or self._function_call_result or self._pending_image_frame_message
+        ):
+            return
+
+        run_llm = False
+
+        aggregation = self._aggregation
+        self._reset()
+
+        try:
+            if self._function_call_result:
+                frame = self._function_call_result
+                self._function_call_result = None
+                if frame.result:
+                    logger.debug(f"FunctionCallResultFrame result: {frame.arguments}")
+                    self._context.add_message(
+                        {
+                            "role": "model",
+                            "parts": [
+                                glm.Part(
+                                    function_call=glm.FunctionCall(
+                                        name=frame.function_name, args=frame.arguments
+                                    )
+                                )
+                            ],
+                        }
+                    )
+                    response = frame.result
+                    if isinstance(response, str):
+                        response = {"response": response}
+                    self._context.add_message(
+                        {
+                            "role": "user",
+                            "parts": [
+                                glm.Part(
+                                    function_response=glm.FunctionResponse(
+                                        name=frame.function_name, response=response
+                                    )
+                                )
+                            ],
+                        }
+                    )
+                    run_llm = not bool(self._function_calls_in_progress)
+            else:
+                self._context.add_message({"role": "model", "parts": [glm.Part(text=aggregation)]})
+
+            if self._pending_image_frame_message:
+                frame = self._pending_image_frame_message
+                self._pending_image_frame_message = None
+                self._context.add_image_frame_message(
+                    format=frame.user_image_raw_frame.format,
+                    size=frame.user_image_raw_frame.size,
+                    image=frame.user_image_raw_frame.image,
+                    text=frame.text,
+                )
+                run_llm = True
+
+            if run_llm:
+                await self._user_context_aggregator.push_context_frame()
+
+            frame = OpenAILLMContextFrame(self._context)
+            await self.push_frame(frame)
+
+        except Exception as e:
+            logger.exception(f"Error processing frame: {e}")
+
+
+@dataclass
+class GoogleContextAggregatorPair:
+    _user: "GoogleUserContextAggregator"
+    _assistant: "GoogleAssistantContextAggregator"
+
+    def user(self) -> "GoogleUserContextAggregator":
+        return self._user
+
+    def assistant(self) -> "GoogleAssistantContextAggregator":
+        return self._assistant
+
+
+class GoogleLLMContext(OpenAILLMContext):
+    @staticmethod
+    def upgrade_to_google(obj: OpenAILLMContext) -> "GoogleLLMContext":
+        if isinstance(obj, OpenAILLMContext) and not isinstance(obj, GoogleLLMContext):
+            logger.debug(f"Upgrading to Google: {obj}")
+            obj.__class__ = GoogleLLMContext
+            obj._restructure_from_openai_messages()
+        return obj
+
+    def from_standard_message(self, message):
+        role = message["role"]
+        content = message["content"]
+        if role == "system":
+            role = "user"
+        elif role == "assistant":
+            role = "model"
+
+        parts = []
+        if isinstance(content, str):
+            parts.append(glm.Part(text=content))
+        elif isinstance(content, list):
+            logger.debug("!!!NEED TO IMPL CONTENT LIST")
+
+        message = {"role": role, "parts": parts}
+        return message
+
+    def add_image_frame_message(
+        self, *, format: str, size: tuple[int, int], image: bytes, text: str = None
+    ):
+        buffer = io.BytesIO()
+        Image.frombytes(format, size, image).save(buffer, format="JPEG")
+
+        parts = []
+        if text:
+            parts.append(glm.Part(text=text))
+        parts.append(
+            glm.Part(inline_data=glm.Blob(mime_type="image/jpeg", data=buffer.getvalue())),
+        )
+        self.add_message({"role": "user", "parts": parts})
+
+    def _restructure_from_openai_messages(self):
+        self._messages[:] = [self.from_standard_message(m) for m in self._messages]
+
+
 class GoogleLLMService(LLMService):
     """This class implements inference with Google's AI models
 
@@ -98,20 +248,34 @@ class GoogleLLMService(LLMService):
     async def _process_context(self, context: OpenAILLMContext):
         await self.push_frame(LLMFullResponseStartFrame())
         try:
-            logger.debug(f"Generating chat: {context.get_messages_json()}")
+            logger.debug(f"Generating chat: {context.messages}")
 
-            messages = self._get_messages_from_openai_context(context)
+            # todo: move this into the new context code structure, convert from openai context one time
+            # todo: add system instructions
+            # messages = self._get_messages_from_openai_context(context)
+            messages = context.messages
 
             await self.start_ttfb_metrics()
 
-            response = self._client.generate_content(messages, stream=True)
+            tools = context.tools if context.tools else []
+            response = self._client.generate_content(contents=messages, tools=tools, stream=True)
 
             await self.stop_ttfb_metrics()
 
             async for chunk in self._async_generator_wrapper(response):
+                # todo: usage
                 try:
-                    text = chunk.text
-                    await self.push_frame(TextFrame(text))
+                    for c in chunk.parts:
+                        if c.text:
+                            await self.push_frame(TextFrame(c.text))
+                        elif c.function_call:
+                            args = type(c.function_call).to_dict(c.function_call).get("args", {})
+                            await self.call_function(
+                                context=context,
+                                tool_call_id="what_should_this_be",
+                                function_name=c.function_call.name,
+                                arguments=args,
+                            )
                 except Exception as e:
                     # Google LLMs seem to flag safety issues a lot!
                     if chunk.candidates[0].finish_reason == 3:
@@ -132,10 +296,11 @@ class GoogleLLMService(LLMService):
         context = None
 
         if isinstance(frame, OpenAILLMContextFrame):
-            context: OpenAILLMContext = frame.context
+            context: GoogleLLMContext = GoogleLLMContext.upgrade_to_google(frame.context)
         elif isinstance(frame, LLMMessagesFrame):
-            context = OpenAILLMContext.from_messages(frame.messages)
+            context = GoogleLLMContext(frame.messages)
         elif isinstance(frame, VisionImageRawFrame):
+            # todo: fix this
             context = OpenAILLMContext.from_image_frame(frame)
         elif isinstance(frame, LLMUpdateSettingsFrame):
             await self._update_settings(frame.settings)
@@ -145,6 +310,16 @@ class GoogleLLMService(LLMService):
         if context:
             await self._process_context(context)
 
+    @staticmethod
+    def create_context_aggregator(
+        context: OpenAILLMContext, *, assistant_expect_stripped_words: bool = True
+    ) -> GoogleContextAggregatorPair:
+        user = GoogleUserContextAggregator(context)
+        assistant = GoogleAssistantContextAggregator(
+            user, expect_stripped_words=assistant_expect_stripped_words
+        )
+        return GoogleContextAggregatorPair(_user=user, _assistant=assistant)
+
 
 class GoogleTTSService(TTSService):
     class InputParams(BaseModel):