diff --git a/astrbot/core/provider/sources/gemini_source.py b/astrbot/core/provider/sources/gemini_source.py index 992c4e6b9b..f4cba79805 100644 --- a/astrbot/core/provider/sources/gemini_source.py +++ b/astrbot/core/provider/sources/gemini_source.py @@ -289,7 +289,7 @@ async def _prepare_query_config( ), ) - def _prepare_conversation(self, payloads: dict) -> list[types.Content]: + async def _prepare_conversation(self, payloads: dict) -> list[types.Content]: """准备 Gemini SDK 的 Content 列表""" def create_text_part(text: str) -> types.Part: @@ -298,11 +298,21 @@ def create_text_part(text: str) -> types.Part: logger.warning("Text content is empty, added a space as placeholder.") return types.Part.from_text(text=content_a) - def process_image_url(image_url_dict: dict) -> types.Part: + async def process_image_url(image_url_dict: dict) -> types.Part: url = image_url_dict["url"] - mime_type = url.split(":")[1].split(";")[0] - image_bytes = base64.b64decode(url.split(",", 1)[1]) - return types.Part.from_bytes(data=image_bytes, mime_type=mime_type) + image_data = await resolve_media_ref_to_base64_data( + url, + media_type="image", + strict=True, + ) + if image_data is None: + raise ValueError( + f"Failed to resolve Gemini history image: {describe_media_ref(url)}" + ) + return types.Part.from_bytes( + data=base64.b64decode(image_data.base64_data), + mime_type=image_data.mime_type, + ) def process_audio_url(audio_url_dict: dict) -> types.Part: url = audio_url_dict["url"] @@ -327,18 +337,14 @@ def append_or_extend( if role == "user": if isinstance(content, list): - parts = [ - ( - types.Part.from_text(text=item["text"] or " ") - if item["type"] == "text" - else ( - process_image_url(item["image_url"]) - if item["type"] == "image_url" - else process_audio_url(item["audio_url"]) - ) - ) - for item in content - ] + parts = [] + for item in content: + if item["type"] == "text": + parts.append(types.Part.from_text(text=item["text"] or " ")) + elif item["type"] == "image_url": + parts.append(await process_image_url(item["image_url"])) + else: + parts.append(process_audio_url(item["audio_url"])) else: parts = [create_text_part(content)] append_or_extend(gemini_contents, parts, types.UserContent) @@ -597,7 +603,7 @@ async def _query( if self.provider_config.get("gm_resp_image_modal", False): modalities.append("IMAGE") - conversation = self._prepare_conversation(payloads) + conversation = await self._prepare_conversation(payloads) temperature = payloads.get("temperature", 0.7) result: types.GenerateContentResponse | None = None @@ -692,7 +698,7 @@ async def _query_stream( None, ) model = payloads.get("model", self.get_model()) - conversation = self._prepare_conversation(payloads) + conversation = await self._prepare_conversation(payloads) result = None while True: diff --git a/tests/test_gemini_source.py b/tests/test_gemini_source.py index 7b47bcb5e1..4c33de6bca 100644 --- a/tests/test_gemini_source.py +++ b/tests/test_gemini_source.py @@ -10,10 +10,11 @@ from astrbot.core.provider.sources.gemini_source import ProviderGoogleGenAI -def test_gemini_prepare_conversation_removes_leading_model_content(): +@pytest.mark.asyncio +async def test_gemini_prepare_conversation_removes_leading_model_content(): provider = ProviderGoogleGenAI.__new__(ProviderGoogleGenAI) - contents = provider._prepare_conversation( + contents = await provider._prepare_conversation( { "messages": [ {"role": "assistant", "content": "stale assistant turn"}, @@ -28,10 +29,11 @@ def test_gemini_prepare_conversation_removes_leading_model_content(): assert contents[0].parts[-1].text == "current user turn" -def test_gemini_prepare_conversation_keeps_normal_user_first_history(): +@pytest.mark.asyncio +async def test_gemini_prepare_conversation_keeps_normal_user_first_history(): provider = ProviderGoogleGenAI.__new__(ProviderGoogleGenAI) - contents = provider._prepare_conversation( + contents = await provider._prepare_conversation( { "messages": [ {"role": "user", "content": "first user turn"}, @@ -50,10 +52,11 @@ def test_gemini_prepare_conversation_keeps_normal_user_first_history(): assert contents[-1].parts[-1].text == "current user turn" -def test_gemini_prepare_conversation_preserves_user_model_history(): +@pytest.mark.asyncio +async def test_gemini_prepare_conversation_preserves_user_model_history(): provider = ProviderGoogleGenAI.__new__(ProviderGoogleGenAI) - contents = provider._prepare_conversation( + contents = await provider._prepare_conversation( { "messages": [ {"role": "user", "content": "user turn"}, @@ -70,6 +73,40 @@ def test_gemini_prepare_conversation_preserves_user_model_history(): assert contents[-1].parts[-1].text == "assistant turn" +@pytest.mark.asyncio +async def test_gemini_prepare_conversation_resolves_local_history_image(tmp_path): + image_path = tmp_path / "history.webp" + image_bytes = ( + b"RIFF\x16\x00\x00\x00WEBPVP8L\x0a\x00\x00\x00" + b"/\x00\x00\x00\x10\x07\x10\x11\x11\x88\x88\xfe\x07" + ) + image_path.write_bytes(image_bytes) + provider = ProviderGoogleGenAI.__new__(ProviderGoogleGenAI) + + contents = await provider._prepare_conversation( + { + "messages": [ + { + "role": "user", + "content": [ + {"type": "text", "text": "historical image"}, + { + "type": "image_url", + "image_url": {"url": str(image_path)}, + }, + ], + } + ] + } + ) + + assert contents[0].parts is not None + image_part = contents[0].parts[1] + assert image_part.inline_data is not None + assert image_part.inline_data.mime_type == "image/webp" + assert image_part.inline_data.data == image_bytes + + def test_gemini_empty_output_raises_empty_model_output_error(): llm_response = LLMResponse(role="assistant")