Merge pull request #652 from pipecat-ai/khk/more-gemini
Gemini new context manager and rewrite to use google data structures internally
This commit is contained in:
Mark Backman
@@ -70,6 +70,8 @@ class OpenAILLMContext:
|
||||
context.add_message(message)
|
||||
return context
|
||||
|
||||
# todo: deprecate from_image_frame. It's only used to create a single-use
|
||||
# context, which isn't useful for most real-world applications.
|
||||
@staticmethod
|
||||
def from_image_frame(frame: VisionImageRawFrame) -> "OpenAILLMContext":
|
||||
"""
|
||||
@@ -77,6 +79,10 @@ class OpenAILLMContext:
|
||||
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.
|
||||
|
||||
NOTE: the above only applies to the deprecated use of this method. The
|
||||
add_image_frame_message() below does the base64 encoding as expected
|
||||
in the OpenAI format.
|
||||
"""
|
||||
context = OpenAILLMContext()
|
||||
buffer = io.BytesIO()
|
||||
|
||||
@@ -5,10 +5,15 @@
|
||||
#
|
||||
|
||||
import asyncio
|
||||
import base64
|
||||
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 +33,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 +54,249 @@ 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(
|
||||
glm.Content(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(
|
||||
glm.Content(
|
||||
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(
|
||||
glm.Content(
|
||||
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(
|
||||
glm.Content(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 set_messages(self, messages: List):
|
||||
self._messages[:] = messages
|
||||
self._restructure_from_openai_messages()
|
||||
|
||||
def get_messages_for_logging(self):
|
||||
msgs = []
|
||||
for message in self.messages:
|
||||
obj = glm.Content.to_dict(message)
|
||||
try:
|
||||
if "parts" in obj:
|
||||
for part in obj["parts"]:
|
||||
if "inline_data" in part:
|
||||
part["inline_data"]["data"] = "..."
|
||||
except Exception as e:
|
||||
logger.debug(f"Error: {e}")
|
||||
msgs.append(obj)
|
||||
return msgs
|
||||
|
||||
def from_standard_message(self, message):
|
||||
role = message["role"]
|
||||
content = message.get("content", [])
|
||||
if role == "system":
|
||||
role = "user"
|
||||
elif role == "assistant":
|
||||
role = "model"
|
||||
|
||||
parts = []
|
||||
if message.get("tool_calls"):
|
||||
for tc in message["tool_calls"]:
|
||||
parts.append(
|
||||
glm.Part(
|
||||
function_call=glm.FunctionCall(
|
||||
name=tc["function"]["name"],
|
||||
args=json.loads(tc["function"]["arguments"]),
|
||||
)
|
||||
)
|
||||
)
|
||||
elif role == "tool":
|
||||
role = "model"
|
||||
parts.append(
|
||||
glm.Part(
|
||||
function_response=glm.FunctionResponse(
|
||||
name="tool_call_result", # seems to work to hard-code the same name every time
|
||||
response=json.loads(message["content"]),
|
||||
)
|
||||
)
|
||||
)
|
||||
elif isinstance(content, str):
|
||||
parts.append(glm.Part(text=content))
|
||||
elif isinstance(content, list):
|
||||
for c in content:
|
||||
if c["type"] == "text":
|
||||
parts.append(glm.Part(text=c["text"]))
|
||||
elif c["type"] == "image_url":
|
||||
parts.append(
|
||||
glm.Part(
|
||||
inline_data=glm.Blob(
|
||||
mime_type="image/jpeg",
|
||||
data=base64.b64decode(c["image_url"]["url"].split(",")[1]),
|
||||
)
|
||||
)
|
||||
)
|
||||
|
||||
message = glm.Content(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(glm.Content(role="user", parts=parts))
|
||||
|
||||
def to_standard_messages(self, obj) -> list:
|
||||
msg = {"role": obj.role, "content": []}
|
||||
if msg["role"] == "model":
|
||||
msg["role"] = "assistant"
|
||||
|
||||
for part in obj.parts:
|
||||
if part.text:
|
||||
msg["content"].append({"type": "text", "text": part.text})
|
||||
elif part.inline_data:
|
||||
encoded = base64.b64encode(part.inline_data.data).decode("utf-8")
|
||||
msg["content"].append(
|
||||
{
|
||||
"type": "image_url",
|
||||
"image_url": {"url": f"data:{part.inline_data.mime_type};base64,{encoded}"},
|
||||
}
|
||||
)
|
||||
elif part.function_call:
|
||||
args = type(part.function_call).to_dict(part.function_call).get("args", {})
|
||||
msg["tool_calls"] = [
|
||||
{
|
||||
"id": part.function_call.name,
|
||||
"type": "function",
|
||||
"function": {
|
||||
"name": part.function_call.name,
|
||||
"arguments": json.dumps(args),
|
||||
},
|
||||
}
|
||||
]
|
||||
|
||||
elif part.function_response:
|
||||
msg["role"] = "tool"
|
||||
resp = (
|
||||
type(part.function_response).to_dict(part.function_response).get("response", {})
|
||||
)
|
||||
msg["tool_call_id"] = part.function_response.name
|
||||
msg["content"] = json.dumps(resp)
|
||||
|
||||
# there might be no content parts for tool_calls messages
|
||||
if not msg["content"]:
|
||||
del msg["content"]
|
||||
return [msg]
|
||||
|
||||
def _restructure_from_openai_messages(self):
|
||||
# first, map across self._messages calling self.from_standard_message(m) to modify messages in place
|
||||
try:
|
||||
self._messages[:] = [self.from_standard_message(m) for m in self._messages]
|
||||
except Exception as e:
|
||||
logger.error(f"Error mapping messages: {e}")
|
||||
# iterate over messages and remove any messages that have an empty content list
|
||||
self._messages = [m for m in self._messages if m.parts]
|
||||
|
||||
|
||||
class GoogleLLMService(LLMService):
|
||||
"""This class implements inference with Google's AI models
|
||||
|
||||
@@ -98,20 +350,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.get_messages_for_logging()}")
|
||||
|
||||
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 +398,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 +412,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):
|
||||
|
||||
Reference in New Issue
Block a user