Skip to content

Commit

Permalink
fix: remote_url doesn't work for gemini (langgenius#5090)
Browse files Browse the repository at this point in the history
  • Loading branch information
rerorero authored Jun 12, 2024
1 parent cf54eb9 commit badb334
Showing 1 changed file with 21 additions and 6 deletions.
27 changes: 21 additions & 6 deletions api/core/model_runtime/model_providers/google/llm/llm.py
Original file line number Diff line number Diff line change
@@ -1,18 +1,22 @@
import base64
import json
import logging
import mimetypes
from collections.abc import Generator
from typing import Optional, Union
from typing import Optional, Union, cast

import google.ai.generativelanguage as glm
import google.api_core.exceptions as exceptions
import google.generativeai as genai
import google.generativeai.client as client
import requests
from google.generativeai.types import ContentType, GenerateContentResponse, HarmBlockThreshold, HarmCategory
from google.generativeai.types.content_types import to_part

from core.model_runtime.entities.llm_entities import LLMResult, LLMResultChunk, LLMResultChunkDelta
from core.model_runtime.entities.message_entities import (
AssistantPromptMessage,
ImagePromptMessageContent,
PromptMessage,
PromptMessageContentType,
PromptMessageTool,
Expand Down Expand Up @@ -361,11 +365,22 @@ def _format_message_to_glm_content(self, message: PromptMessage) -> ContentType:
for c in message.content:
if c.type == PromptMessageContentType.TEXT:
glm_content['parts'].append(to_part(c.data))
else:
metadata, data = c.data.split(',', 1)
mime_type = metadata.split(';', 1)[0].split(':')[1]
blob = {"inline_data":{"mime_type":mime_type,"data":data}}
elif c.type == PromptMessageContentType.IMAGE:
message_content = cast(ImagePromptMessageContent, c)
if message_content.data.startswith("data:"):
metadata, base64_data = c.data.split(',', 1)
mime_type = metadata.split(';', 1)[0].split(':')[1]
else:
# fetch image data from url
try:
image_content = requests.get(message_content.data).content
mime_type, _ = mimetypes.guess_type(message_content.data)
base64_data = base64.b64encode(image_content).decode('utf-8')
except Exception as ex:
raise ValueError(f"Failed to fetch image data from url {message_content.data}, {ex}")
blob = {"inline_data":{"mime_type":mime_type,"data":base64_data}}
glm_content['parts'].append(blob)

return glm_content
elif isinstance(message, AssistantPromptMessage):
glm_content = {
Expand Down Expand Up @@ -444,4 +459,4 @@ def _invoke_error_mapping(self) -> dict[type[InvokeError], list[type[Exception]]
exceptions.RequestRangeNotSatisfiable,
exceptions.Cancelled,
]
}
}

0 comments on commit badb334

Please sign in to comment.