google: fix VisionImageRawFrame context

This commit is contained in:
Aleix Conchillo Flaqué
2024-12-10 17:08:49 -08:00
parent a618bd3fa6
commit 9e7ecccf1e

View File

@@ -320,6 +320,15 @@ class GoogleContextAggregatorPair:
class GoogleLLMContext(OpenAILLMContext): class GoogleLLMContext(OpenAILLMContext):
def __init__(
self,
messages: list[dict] | None = None,
tools: list[dict] | None = None,
tool_choice: dict | None = None,
):
super().__init__(messages=messages, tools=tools, tool_choice=tool_choice)
self.system_message = None
@staticmethod @staticmethod
def upgrade_to_google(obj: OpenAILLMContext) -> "GoogleLLMContext": def upgrade_to_google(obj: OpenAILLMContext) -> "GoogleLLMContext":
if isinstance(obj, OpenAILLMContext) and not isinstance(obj, GoogleLLMContext): if isinstance(obj, OpenAILLMContext) and not isinstance(obj, GoogleLLMContext):
@@ -371,9 +380,8 @@ class GoogleLLMContext(OpenAILLMContext):
parts = [] parts = []
if text: if text:
parts.append(glm.Part(text=text)) parts.append(glm.Part(text=text))
parts.append( parts.append(glm.Part(inline_data=glm.Blob(mime_type="image/jpeg", data=buffer.getvalue())))
glm.Part(inline_data=glm.Blob(mime_type="image/jpeg", data=buffer.getvalue())),
)
self.add_message(glm.Content(role="user", parts=parts)) self.add_message(glm.Content(role="user", parts=parts))
def add_audio_frames_message(self, *, audio_frames: list[AudioRawFrame], text: str = None): def add_audio_frames_message(self, *, audio_frames: list[AudioRawFrame], text: str = None):
@@ -649,12 +657,14 @@ class GoogleLLMService(LLMService):
context = None context = None
if isinstance(frame, OpenAILLMContextFrame): if isinstance(frame, OpenAILLMContextFrame):
context: GoogleLLMContext = GoogleLLMContext.upgrade_to_google(frame.context) context = GoogleLLMContext.upgrade_to_google(frame.context)
elif isinstance(frame, LLMMessagesFrame): elif isinstance(frame, LLMMessagesFrame):
context = GoogleLLMContext(frame.messages) context = GoogleLLMContext(frame.messages)
elif isinstance(frame, VisionImageRawFrame): elif isinstance(frame, VisionImageRawFrame):
# todo: fix this context = GoogleLLMContext()
context = OpenAILLMContext.from_image_frame(frame) context.add_image_frame_message(
format=frame.format, size=frame.size, image=frame.image, text=frame.text
)
elif isinstance(frame, LLMUpdateSettingsFrame): elif isinstance(frame, LLMUpdateSettingsFrame):
await self._update_settings(frame.settings) await self._update_settings(frame.settings)
else: else: