From 529144649e6967b9b78004581c107dd908f4adc2 Mon Sep 17 00:00:00 2001 From: Nicolas Suzor Date: Wed, 20 Dec 2023 15:58:39 +1000 Subject: [PATCH] community[patch]: add png support for vertexai._parse_chat_history_gemini() (#14788) - **Description:** Modify community chat model vertexai to handle png and other image types encoded in base64 - **Dependencies:** added `import re` but no new dependencies. This addresses a problem where the vertexai method _parse_chat_history_gemini() was only recognizing image uris in jpeg format. I made a simple change to cover other extension types. --- .../langchain_community/chat_models/vertexai.py | 15 +++++++++++++-- 1 file changed, 13 insertions(+), 2 deletions(-) diff --git a/libs/community/langchain_community/chat_models/vertexai.py b/libs/community/langchain_community/chat_models/vertexai.py index bb48aecb477..8cc493a1fc0 100644 --- a/libs/community/langchain_community/chat_models/vertexai.py +++ b/libs/community/langchain_community/chat_models/vertexai.py @@ -3,6 +3,7 @@ from __future__ import annotations import base64 import logging +import re from dataclasses import dataclass, field from typing import TYPE_CHECKING, Any, Dict, Iterator, List, Optional, Union, cast @@ -105,8 +106,18 @@ def _parse_chat_history_gemini( path = part["image_url"]["url"] if path.startswith("gs://"): image = load_image_from_gcs(path=path, project=project) - elif path.startswith("data:image/jpeg;base64,"): - image = Image.from_bytes(base64.b64decode(path[23:])) + elif path.startswith("data:image/"): + # extract base64 component from image uri + try: + encoded = re.search(r"data:image/\w{2,4};base64,(.*)", path).group( + 1 + ) + except AttributeError: + raise ValueError( + "Invalid image uri. It should be in the format " + "data:image/;base64,." + ) + image = Image.from_bytes(base64.b64decode(encoded)) else: image = Image.load_from_file(path) else: