Initial commit of Google Gemini LLM service.
Gemini text input works. We translate from OpenAILLMContext format on the fly in the GoogleLLMService implementation. This commit also implements image input (vision) in both the GoogleLLMService and in the OpenAILLMService. Image input is a hack and needs to be revisited. OpenAI expects images to be uploaded as base64-encoded JPEGs. Google does not require the base64 encoding. Other than for images, we use the OpenAI format as our standard, but base64-encoding the images and then unencoding them in the GoogleLLMService feels wasteful.
This commit is contained in:
@@ -5,10 +5,13 @@
|
||||
#
|
||||
|
||||
from dataclasses import dataclass
|
||||
import io
|
||||
|
||||
from typing import List
|
||||
|
||||
from pipecat.frames.frames import Frame
|
||||
from PIL import Image
|
||||
|
||||
from pipecat.frames.frames import Frame, VisionImageRawFrame
|
||||
|
||||
from openai._types import NOT_GIVEN, NotGiven
|
||||
|
||||
@@ -43,6 +46,31 @@ class OpenAILLMContext:
|
||||
})
|
||||
return context
|
||||
|
||||
@staticmethod
|
||||
def from_image_frame(frame: VisionImageRawFrame) -> "OpenAILLMContext":
|
||||
"""
|
||||
For images, we are deviating from the OpenAI messages shape. OpenAI
|
||||
expects images to be base64 encoded, but other vision models may not.
|
||||
So we'll store the image as bytes and do the base64 encoding as needed
|
||||
in the LLM service.
|
||||
"""
|
||||
context = OpenAILLMContext()
|
||||
buffer = io.BytesIO()
|
||||
Image.frombytes(
|
||||
frame.format,
|
||||
frame.size,
|
||||
frame.image
|
||||
).save(
|
||||
buffer,
|
||||
format="JPEG")
|
||||
context.add_message({
|
||||
"content": frame.text,
|
||||
"role": "user",
|
||||
"data": buffer.getvalue(),
|
||||
"mime_type": "image/jpeg"
|
||||
})
|
||||
return context
|
||||
|
||||
def add_message(self, message: ChatCompletionMessageParam):
|
||||
self.messages.append(message)
|
||||
|
||||
|
||||
96
src/pipecat/services/google.py
Normal file
96
src/pipecat/services/google.py
Normal file
@@ -0,0 +1,96 @@
|
||||
import google.generativeai as gai
|
||||
import google.ai.generativelanguage as glm
|
||||
import os
|
||||
import asyncio
|
||||
|
||||
from typing import List
|
||||
|
||||
from pipecat.frames.frames import (
|
||||
Frame,
|
||||
TextFrame,
|
||||
VisionImageRawFrame,
|
||||
LLMMessagesFrame,
|
||||
LLMResponseStartFrame,
|
||||
LLMResponseEndFrame)
|
||||
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
|
||||
|
||||
|
||||
class GoogleLLMService(LLMService):
|
||||
"""This class implements inference with Google's AI models
|
||||
|
||||
This service translates internally from OpenAILLMContext to the messages format
|
||||
expected by the Google AI model. We are using the OpenAILLMContext as a lingua
|
||||
franca for all LLM services, so that it is easy to switch between different LLMs.
|
||||
"""
|
||||
|
||||
def __init__(self, model="gemini-1.5-flash-latest", api_key=None, **kwargs):
|
||||
super().__init__(**kwargs)
|
||||
self.model = model
|
||||
gai.configure(api_key=api_key or os.environ["GOOGLE_API_KEY"])
|
||||
self.create_client()
|
||||
|
||||
def create_client(self):
|
||||
self._client = gai.GenerativeModel(self.model)
|
||||
|
||||
def _get_messages_from_openai_context(
|
||||
self, context: OpenAILLMContext) -> List[glm.Content]:
|
||||
openai_messages = context.get_messages()
|
||||
google_messages = []
|
||||
|
||||
for message in openai_messages:
|
||||
role = message["role"]
|
||||
content = message["content"]
|
||||
if role == "system":
|
||||
role = "user"
|
||||
elif role == "assistant":
|
||||
role = "model"
|
||||
|
||||
parts = [glm.Part(text=content)]
|
||||
if "mime_type" in message:
|
||||
parts.append(
|
||||
glm.Part(inline_data=glm.Blob(
|
||||
mime_type=message["mime_type"],
|
||||
data=message["data"]
|
||||
)))
|
||||
google_messages.append({"role": role, "parts": parts})
|
||||
|
||||
return google_messages
|
||||
|
||||
async def _async_generator_wrapper(self, sync_generator):
|
||||
for item in sync_generator:
|
||||
yield item
|
||||
await asyncio.sleep(0)
|
||||
|
||||
async def _process_context(self, context: OpenAILLMContext):
|
||||
try:
|
||||
messages = self._get_messages_from_openai_context(context)
|
||||
|
||||
await self.push_frame(LLMResponseStartFrame())
|
||||
response = self._client.generate_content(messages, stream=True)
|
||||
|
||||
async for chunk in self._async_generator_wrapper(response):
|
||||
logger.debug(f"Pushing inference text: {chunk.text}")
|
||||
await self.push_frame(TextFrame(chunk.text))
|
||||
|
||||
await self.push_frame(LLMResponseEndFrame())
|
||||
except Exception as e:
|
||||
logger.error(f"Exception: {e}")
|
||||
|
||||
async def process_frame(self, frame: Frame, direction: FrameDirection):
|
||||
context = None
|
||||
|
||||
if isinstance(frame, OpenAILLMContextFrame):
|
||||
context: OpenAILLMContext = frame.context
|
||||
elif isinstance(frame, LLMMessagesFrame):
|
||||
context = OpenAILLMContext.from_messages(frame.messages)
|
||||
elif isinstance(frame, VisionImageRawFrame):
|
||||
context = OpenAILLMContext.from_image_frame(frame)
|
||||
else:
|
||||
await self.push_frame(frame, direction)
|
||||
|
||||
if context:
|
||||
await self._process_context(context)
|
||||
@@ -8,6 +8,7 @@ import io
|
||||
import json
|
||||
import time
|
||||
import aiohttp
|
||||
import base64
|
||||
|
||||
from PIL import Image
|
||||
|
||||
@@ -22,7 +23,8 @@ from pipecat.frames.frames import (
|
||||
LLMResponseEndFrame,
|
||||
LLMResponseStartFrame,
|
||||
TextFrame,
|
||||
URLImageRawFrame
|
||||
URLImageRawFrame,
|
||||
VisionImageRawFrame
|
||||
)
|
||||
from pipecat.processors.aggregators.openai_llm_context import OpenAILLMContext, OpenAILLMContextFrame
|
||||
from pipecat.processors.frame_processor import FrameDirection
|
||||
@@ -67,8 +69,21 @@ class BaseOpenAILLMService(LLMService):
|
||||
self, context: OpenAILLMContext
|
||||
) -> AsyncStream[ChatCompletionChunk]:
|
||||
messages: List[ChatCompletionMessageParam] = context.get_messages()
|
||||
messages_for_log = json.dumps(messages)
|
||||
logger.debug(f"Generating chat: {messages_for_log}")
|
||||
|
||||
# base64 encode any images
|
||||
for message in messages:
|
||||
if message.get("mime_type") == "image/jpeg":
|
||||
encoded_image = base64.b64encode(message["data"]).decode("utf-8")
|
||||
text = message["content"]
|
||||
message["content"] = [
|
||||
{"type": "text", "text": text},
|
||||
{"type": "image_url", "image_url": {"url": f"data:image/jpeg;base64,{encoded_image}"}}
|
||||
]
|
||||
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] = (
|
||||
@@ -151,6 +166,8 @@ class BaseOpenAILLMService(LLMService):
|
||||
context: OpenAILLMContext = frame.context
|
||||
elif isinstance(frame, LLMMessagesFrame):
|
||||
context = OpenAILLMContext.from_messages(frame.messages)
|
||||
elif isinstance(frame, VisionImageRawFrame):
|
||||
context = OpenAILLMContext.from_image_frame(frame)
|
||||
else:
|
||||
await self.push_frame(frame, direction)
|
||||
|
||||
|
||||
Reference in New Issue
Block a user