Repository: run-llama/llama_index
Pull request: Add support for Google Gemini models via VertexAI
# Description

Calling gemini models via VertexAI simplify moving an existing vertexAI application to gemini: all you need to do is to change the model to "gemini-pro" or "gemini-pro-vision".

Fixes #9607

Now it's possible to use VertexAI directly with Gemini models

```
llm = Vertex(temperature=0, model="gemini-pro")
output = llm.complete("Say foo:")


history = [ChatMessage(role="user", content="Say foo:"), ChatMessage(role="assistant", content="Foo with love !"),
               ChatMessage(role="user", content="Please repeat")]
output = llm.chat(history)

```

## Type of Change

Please delete options that are not relevant.

- [ ] New feature (non-breaking change which adds functionality)


# How Has This Been Tested?

Please describe the tests that you ran to verify your changes. Provide instructions so we can reproduce. Please also list any relevant details for your test configuration

- [x] Added new unit/integration tests


# Suggested Checklist:

- [x] I have performed a self-review of my own code
- [x] I have commented my code, particularly in hard-to-understand areas
- [x] I have made corresponding changes to the documentation
- [] I have added Google Colab support for the newly added notebooks.
- [x] My changes generate no new warnings
- [x] I have added tests that prove my fix is effective or that my feature works
- [x] New and existing unit tests pass locally with my changes
- [x] I ran `make format; make lint` to appease the lint gods


Problem statement:
[Feature Request]: Google Gemini model support via Vertex AI
### Feature Description

Currently Gemini is available also via Vertex AI API. I'd like to be able to switch from 
```
from llama_index.llms import Vertex
Vertex(model="chat-bison")
```
to 
```
from llama_index.llms import Vertex
Vertex(model="gemini-pro")
```

It's also noteworthy that this is already doable in LangChain 0.0.350
```
from langchain_community.chat_models import ChatVertexAI
ChatVertexAI(model_name="gemini-pro")
```

### Reason

Mid term reason is to being able to switch existing VertexAI application to Gemini models
Short term reason is that Gemini API are blocked in many countries (EU-UK for example) but the usage with Vertex is allowed

### Value of Feature

_No response_

Hints from issue discussion:
_🤖_
<!-- Greeting -->
Hello @nicoloboschi! 🖐️

I'm [Dosu](https://dosu.dev), an AI here to assist you with your issues and queries about LlamaIndex until a human maintainer can step in. I can help you troubleshoot bugs, answer questions, and guide you on how to contribute to the project. Don't worry, I'm here to make your experience smoother!

Looking forward to our interaction!

<!-- Answer -->
Thank you for your feature request. I understand that you would like to add support for the Google Gemini model via Vertex AI to the LlamaIndex repository, similar to how it's implemented in LangChain 0.0.350. This would allow users to switch from existing VertexAI applications to Gemini models, especially in regions where the Gemini API is blocked.

From the context provided, it appears that the LlamaIndex repository already supports the Google Gemini model via Vertex AI. The `GeminiMultiModal` class in the LlamaIndex repository is designed to interface with Google's Gemini multi-modal models. 

Here is the relevant code snippet that checks the model support:

```py
# Check whether the Gemini Model is supported or not
if model_name not in GEMINI_MM_MODELS:
    raise ValueError(
        f"Invalid model {model_name}. "
        f"Available models are: {GEMINI_MM_MODELS}"
    )
```

And here is the code snippet that initializes the Gemini model:

```py
self._model = genai.GenerativeModel(
    model_name=model_name,
    generation_config=final_gen_config,
    safety_settings=safety_settings,
)
```

So, you should be able to switch from existing VertexAI applications to Gemini models by using this class and its methods. However, if the Gemini API is blocked in certain regions, you might need to use a VPN or proxy to bypass the blockage, as this is a network issue rather than a software one.

If you're still having trouble implementing this feature, could you please provide more details about the issues you're facing? This will help me better understand the problem and provide a more accurate solution.

## Sources
- [llama_index/multi_modal_llms/gemini.py](https://github.com/run-llama/llama_index/blob/main/llama_index/multi_modal_llms/gemini.py)

---

This response is meant to be useful and save you time. It is **not meant to be a precise solution**, but rather a starting point for your own research.

Help me be more useful! Please leave a 👍 if this is helpful and 👎 if it is irrelevant.

If you want to continue the conversation, start your reply with @dosu-bot.

@nicoloboschi that is a good point. feel free to make a pr to support.
I think you can first try to directly using current Vertex to load gemini pro model to see whether it works
Hi @hatianzhang I opened a pr https://github.com/run-llama/llama_index/pull/9624
Please review

Commit message: fix import

Patch to review:
diff --git a/llama_index/llms/vertex.py b/llama_index/llms/vertex.py
index b6f2764a33..7e527e0d6d 100644
--- a/llama_index/llms/vertex.py
+++ b/llama_index/llms/vertex.py
@@ -28,6 +28,7 @@
     acompletion_with_retry,
     completion_with_retry,
     init_vertexai,
+    is_gemini_model,
 )
 from llama_index.types import BaseOutputParser, PydanticProgramMode
 
@@ -47,6 +48,7 @@ class Vertex(LLM):
     iscode: bool = Field(
         default=False, description="Flag to determine if current model is a Code Model"
     )
+    _is_gemini: bool = PrivateAttr()
     _client: Any = PrivateAttr()
     _chat_client: Any = PrivateAttr()
 
@@ -74,6 +76,7 @@ def __init__(
         additional_kwargs = additional_kwargs or {}
         callback_manager = callback_manager or CallbackManager([])
 
+        self._is_gemini = False
         if model in CHAT_MODELS:
             from vertexai.language_models import ChatModel
 
@@ -92,8 +95,14 @@ def __init__(
             from vertexai.language_models import TextGenerationModel
 
             self._client = TextGenerationModel.from_pretrained(model)
+        elif is_gemini_model(model):
+            from vertexai.preview.generative_models import GenerativeModel
+
+            self._client = GenerativeModel(model_name=model)
+            self._chat_client = self._client
+            self._is_gemini = True
         else:
-            raise (ValueError("Model Not Found Please Check the model name"))
+            raise (ValueError(f"Model {model} not found, please verify the model name"))
 
         super().__init__(
             temperature=temperature,
@@ -142,7 +151,7 @@ def _get_all_kwargs(self, **kwargs: Any) -> Dict[str, Any]:
     @llm_chat_callback()
     def chat(self, messages: Sequence[ChatMessage], **kwargs: Any) -> ChatResponse:
         question = messages[-1].content
-        chat_history = _parse_chat_history(messages[:-1])
+        chat_history = _parse_chat_history(messages[:-1], self._is_gemini)
         chat_params = {**chat_history}
 
         kwargs = kwargs if kwargs else {}
@@ -165,9 +174,10 @@ def chat(self, messages: Sequence[ChatMessage], **kwargs: Any) -> ChatResponse:
             prompt=question,
             chat=True,
             stream=False,
+            is_gemini=self._is_gemini,
             params=chat_params,
             max_retries=self.max_retries,
-            **params
+            **params,
         )
 
         return ChatResponse(
@@ -183,7 +193,11 @@ def complete(self, prompt: str, **kwargs: Any) -> CompletionResponse:
             raise (ValueError("candidate_count is not supported by the codey model's"))
 
         completion = completion_with_retry(
-            self._client, prompt, max_retries=self.max_retries, **params
+            self._client,
+            prompt,
+            max_retries=self.max_retries,
+            is_gemini=self._is_gemini,
+            **params,
         )
         return CompletionResponse(text=completion.text, raw=completion.__dict__)
 
@@ -192,7 +206,7 @@ def stream_chat(
         self, messages: Sequence[ChatMessage], **kwargs: Any
     ) -> ChatResponseGen:
         question = messages[-1].content
-        chat_history = _parse_chat_history(messages[:-1])
+        chat_history = _parse_chat_history(messages[:-1], self._is_gemini)
         chat_params = {**chat_history}
         kwargs = kwargs if kwargs else {}
         params = {**self._model_kwargs, **kwargs}
@@ -212,19 +226,17 @@ def stream_chat(
             prompt=question,
             chat=True,
             stream=True,
+            is_gemini=self._is_gemini,
             params=chat_params,
             max_retries=self.max_retries,
-            **params
+            **params,
         )
 
         def gen() -> ChatResponseGen:
             content = ""
             role = MessageRole.ASSISTANT
             for r in response:
-                if "text" in r.__dict__:
-                    content_delta = r.text
-                else:
-                    content_delta = ""
+                content_delta = r.text
                 content += content_delta
                 yield ChatResponse(
                     message=ChatMessage(role=role, content=content),
@@ -245,8 +257,9 @@ def stream_complete(self, prompt: str, **kwargs: Any) -> CompletionResponseGen:
             client=self._client,
             prompt=prompt,
             stream=True,
+            is_gemini=self._is_gemini,
             max_retries=self.max_retries,
-            **params
+            **params,
         )
 
         def gen() -> CompletionResponseGen:
@@ -265,7 +278,7 @@ async def achat(
         self, messages: Sequence[ChatMessage], **kwargs: Any
     ) -> ChatResponse:
         question = messages[-1].content
-        chat_history = _parse_chat_history(messages[:-1])
+        chat_history = _parse_chat_history(messages[:-1], self._is_gemini)
         chat_params = {**chat_history}
         kwargs = kwargs if kwargs else {}
         params = {**self._model_kwargs, **kwargs}
@@ -283,9 +296,10 @@ async def achat(
             client=self._chat_client,
             prompt=question,
             chat=True,
+            is_gemini=self._is_gemini,
             params=chat_params,
             max_retries=self.max_retries,
-            **params
+            **params,
         )
         ##this is due to a bug in vertex AI we have to await twice
         if self.iscode:
@@ -302,7 +316,11 @@ async def acomplete(self, prompt: str, **kwargs: Any) -> CompletionResponse:
         if self.iscode and "candidate_count" in params:
             raise (ValueError("candidate_count is not supported by the codey model's"))
         completion = await acompletion_with_retry(
-            client=self._client, prompt=prompt, max_retries=self.max_retries, **params
+            client=self._client,
+            prompt=prompt,
+            max_retries=self.max_retries,
+            is_gemini=self._is_gemini,
+            **params,
         )
         return CompletionResponse(text=completion.text)
 
diff --git a/llama_index/llms/vertex_utils.py b/llama_index/llms/vertex_utils.py
index e87ff3e4ba..536bf65364 100644
--- a/llama_index/llms/vertex_utils.py
+++ b/llama_index/llms/vertex_utils.py
@@ -1,8 +1,9 @@
 # utils script
+import base64
 
 # generation with retry
 import logging
-from typing import Any, Callable, Optional
+from typing import Any, Callable, Dict, Optional, Union
 
 from tenacity import (
     before_sleep_log,
@@ -20,6 +21,10 @@
 CODE_CHAT_MODELS = ["codechat-bison", "codechat-bison-32k", "codechat-bison@001"]
 
 
+def is_gemini_model(model: str) -> bool:
+    return model.startswith("gemini")
+
+
 logger = logging.getLogger(__name__)
 
 
@@ -49,6 +54,7 @@ def completion_with_retry(
     max_retries: int = 5,
     chat: bool = False,
     stream: bool = False,
+    is_gemini: bool = False,
     params: Any = {},
     **kwargs: Any,
 ) -> Any:
@@ -57,7 +63,15 @@ def completion_with_retry(
 
     @retry_decorator
     def _completion_with_retry(**kwargs: Any) -> Any:
-        if chat:
+        if is_gemini:
+            history = params["message_history"] if "message_history" in params else []
+
+            generation = client.start_chat(history=history)
+            generation_config = dict(kwargs)
+            return generation.send_message(
+                prompt, stream=stream, generation_config=generation_config
+            )
+        elif chat:
             generation = client.start_chat(**params)
             if stream:
                 return generation.send_message_streaming(prompt, **kwargs)
@@ -77,6 +91,7 @@ async def acompletion_with_retry(
     prompt: Optional[str],
     max_retries: int = 5,
     chat: bool = False,
+    is_gemini: bool = False,
     params: Any = {},
     **kwargs: Any,
 ) -> Any:
@@ -85,7 +100,15 @@ async def acompletion_with_retry(
 
     @retry_decorator
     async def _completion_with_retry(**kwargs: Any) -> Any:
-        if chat:
+        if is_gemini:
+            history = params["message_history"] if "message_history" in params else []
+
+            generation = client.start_chat(history=history)
+            generation_config = dict(kwargs)
+            return await generation.send_message_async(
+                prompt, generation_config=generation_config
+            )
+        elif chat:
             generation = client.start_chat(**params)
             return await generation.send_message_async(prompt, **kwargs)
         else:
@@ -123,7 +146,7 @@ def init_vertexai(
     )
 
 
-def _parse_chat_history(history: Any) -> Any:
+def _parse_chat_history(history: Any, is_gemini: bool) -> Any:
     """Parse a sequence of messages into history.
 
     Args:
@@ -141,13 +164,51 @@ def _parse_chat_history(history: Any) -> Any:
     vertex_messages, context = [], None
     for i, message in enumerate(history):
         if i == 0 and message.role == MessageRole.SYSTEM:
-            context = message.content
-        elif message.role == MessageRole.ASSISTANT:
-            vertex_message = ChatMessage(content=message.content, author="bot")
-            vertex_messages.append(vertex_message)
-        elif message.role == MessageRole.USER:
-            vertex_message = ChatMessage(content=message.content, author="user")
-            vertex_messages.append(vertex_message)
+            if is_gemini:
+                raise ValueError("Gemini model don't support system messages")
+        elif message.role == MessageRole.ASSISTANT or message.role == MessageRole.USER:
+            if is_gemini:
+                from vertexai.preview.generative_models import Content, Part
+
+                def _convert_gemini_part_to_prompt(part: Union[str, Dict]) -> Part:
+                    from vertexai.preview.generative_models import Image, Part
+
+                    if isinstance(part, str):
+                        return Part.from_text(part)
+
+                    if not isinstance(part, Dict):
+                        raise ValueError(
+                            f"Message's content is expected to be a dict, got {type(part)}!"
+                        )
+                    if part["type"] == "text":
+                        return Part.from_text(part["text"])
+                    elif part["type"] == "image_url":
+                        path = part["image_url"]["url"]
+                        if path.startswith("gs://"):
+                            raise ValueError("Only local image path is supported!")
+                        elif path.startswith("data:image/jpeg;base64,"):
+                            image = Image.from_bytes(base64.b64decode(path[23:]))
+                        else:
+                            image = Image.load_from_file(path)
+                    else:
+                        raise ValueError("Only text and image_url types are supported!")
+                    return Part.from_image(image)
+
+                raw_content = message.content
+                if isinstance(raw_content, str):
+                    raw_content = [raw_content]
+                parts = [_convert_gemini_part_to_prompt(part) for part in raw_content]
+                vertex_message = Content(
+                    role="user" if message.role == MessageRole.USER else "model",
+                    parts=parts,
+                )
+                vertex_messages.append(vertex_message)
+            else:
+                vertex_message = ChatMessage(
+                    content=message.content,
+                    author="bot" if message.role == MessageRole.ASSISTANT else "user",
+                )
+                vertex_messages.append(vertex_message)
         else:
             raise ValueError(
                 f"Unexpected message with type {type(message)} at the position {i}."
diff --git a/tests/llms/test_vertex.py b/tests/llms/test_vertex.py
index 4d6f006863..17d84fd3cd 100644
--- a/tests/llms/test_vertex.py
+++ b/tests/llms/test_vertex.py
@@ -1,5 +1,5 @@
 import pytest
-from llama_index.llms.types import CompletionResponse
+from llama_index.llms.types import ChatMessage, CompletionResponse
 from llama_index.llms.vertex import Vertex
 from llama_index.llms.vertex_utils import init_vertexai
 
@@ -54,6 +54,7 @@ def test_vertex_stream() -> None:
 
 
 @pytest.mark.skipif(vertex_init is False, reason="vertex not installed")
+@pytest.mark.asyncio()
 async def test_vertex_consistency() -> None:
     llm = Vertex(temperature=0)
     output = llm.complete("Please say foo:")
@@ -61,3 +62,30 @@ async def test_vertex_consistency() -> None:
     async_output = await llm.acomplete("Please say foo:")
     assert output.text == streaming_output[-1].text
     assert output.text == async_output.text
+
+
+@pytest.mark.skipif(vertex_init is False, reason="vertex not installed")
+@pytest.mark.asyncio()
+async def test_vertex_gemini_call() -> None:
+    # complete
+    llm = Vertex(temperature=0, model="gemini-pro")
+    output = llm.complete("Say foo:")
+    assert "foo" in output.text.lower()
+    streaming_output = list(llm.stream_complete("Please say foo:"))
+    assert "foo" in streaming_output[-1].text
+
+    async_output = await llm.acomplete("Please say foo:")
+    assert "foo" in async_output.text
+
+    # chat
+    history = [
+        ChatMessage(role="user", content="Say foo:"),
+        ChatMessage(role="assistant", content="Foo with love !"),
+        ChatMessage(role="user", content="Please repeat"),
+    ]
+    output = llm.chat(history)
+    assert "Foo with love !" in output.message.content
+    streaming_output = list(llm.stream_chat(history))
+    assert "Foo with love !" in streaming_output[-1].message.content
+    async_output = await llm.achat(history)
+    assert "Foo with love !" in async_output.message.content


Code review comments:
diff_hunk: @@ -141,13 +164,51 @@ def _parse_chat_history(history: Any) -> Any:
     vertex_messages, context = [], None
     for i, message in enumerate(history):
         if i == 0 and message.role == MessageRole.SYSTEM:
-            context = message.content
-        elif message.role == MessageRole.ASSISTANT:
-            vertex_message = ChatMessage(content=message.content, author="bot")
-            vertex_messages.append(vertex_message)
-        elif message.role == MessageRole.USER:
-            vertex_message = ChatMessage(content=message.content, author="user")
-            vertex_messages.append(vertex_message)
+            if is_gemini:
+                raise ValueError("Gemini model don't support system messages")
+        elif message.role == MessageRole.ASSISTANT or message.role == MessageRole.USER:
+            if is_gemini:
+                from vertexai.preview.generative_models import Content, Part
+
+                def _convert_gemini_part_to_prompt(part: Union[str, Dict]) -> Part:
path: llama_index/llms/vertex_utils.py
text: @user1:
could you make a gemini_util file for this piece