From 150721772507a234f37f84168612087c5a99dd29 Mon Sep 17 00:00:00 2001 From: Krrish Dholakia Date: Sat, 6 Jan 2024 22:50:35 +0530 Subject: [PATCH] fix(factory.py): more logging around the image loading for gemini --- litellm/llms/prompt_templates/factory.py | 36 ++++++++++--------- litellm/tests/test_google_ai_studio_gemini.py | 1 + 2 files changed, 21 insertions(+), 16 deletions(-) diff --git a/litellm/llms/prompt_templates/factory.py b/litellm/llms/prompt_templates/factory.py index 8e9504ee932..b0a417b75df 100644 --- a/litellm/llms/prompt_templates/factory.py +++ b/litellm/llms/prompt_templates/factory.py @@ -386,27 +386,31 @@ def anthropic_pt( return prompt -def _load_image_from_url(image_url: str): - """ - Loads an image from a URL. - - Args: - image_url (str): The URL of the image. - - Returns: - Image: The loaded image. - """ - from io import BytesIO +def _load_image_from_url(image_url): try: from PIL import Image except: raise Exception("gemini image conversion failed please run `pip install Pillow`") - - # Download the image from the URL - response = requests.get(image_url) - image = Image.open(BytesIO(response.content)) + from io import BytesIO + try: + # Send a GET request to the image URL + response = requests.get(image_url) + response.raise_for_status() # Raise an exception for HTTP errors - return image + # Check the response's content type to ensure it is an image + content_type = response.headers.get('content-type') + if not content_type or 'image' not in content_type: + raise ValueError(f"URL does not point to a valid image (content-type: {content_type})") + + # Load the image from the response content + return Image.open(BytesIO(response.content)) + + except requests.RequestException as e: + print(f"Request failed: {e}") + except UnidentifiedImageError: + print("Cannot identify image file (it may not be a supported image format or might be corrupted).") + except ValueError as e: + print(e) def _gemini_vision_convert_messages(messages: list): diff --git a/litellm/tests/test_google_ai_studio_gemini.py b/litellm/tests/test_google_ai_studio_gemini.py index e9aa00d4a3c..db732d7105d 100644 --- a/litellm/tests/test_google_ai_studio_gemini.py +++ b/litellm/tests/test_google_ai_studio_gemini.py @@ -27,6 +27,7 @@ def generate_text(): ] response = litellm.completion(model="gemini/gemini-pro-vision", messages=messages) print(response) + assert isinstance(response.choices[0].message.content, str) == True except Exception as exception: raise Exception("An error occurred during text generation:", exception)