Add GoogleGenAI FileAPI support for large files (#19853)

* Add GoogleGenAI FileAPI support for large files

* Bump version

* Update dependencies

* Bump version to 0.4.0

* Bump version

* Bump version to 0.5.0

---------

Co-authored-by: Florian Bacho <[email protected]>

Author: Florian-BACHO

llama-index-integrations/llms/llama-index-llms-google-genai/llama_index/llms/google_genai/base.py

@@ -1,5 +1,6 @@
 """Google's hosted Gemini API."""
 
+import asyncio
 import functools
 import os
 import typing
@@ -52,6 +53,7 @@
     prepare_chat_params,
     handle_streaming_flexible_model,
     create_retry_decorator,
+    delete_uploaded_files,
 )
 
 import google.genai
@@ -135,6 +137,10 @@ class GoogleGenAI(FunctionCallingLLM):
         default=None,
         description="Google GenAI tool to use for the model to augment responses.",
     )
+    use_file_api: bool = Field(
+        default=True,
+        description="Whether or not to use the FileAPI for large files (>20MB).",
+    )
 
     _max_tokens: int = PrivateAttr()
     _client: google.genai.Client = PrivateAttr()
@@ -157,6 +163,7 @@ def __init__(
         is_function_calling_model: bool = True,
         cached_content: Optional[str] = None,
         built_in_tool: Optional[types.Tool] = None,
+        use_file_api: bool = True,
         **kwargs: Any,
     ):
         # API keys are optional. The API can be authorised via OAuth (detected
@@ -205,6 +212,7 @@ def __init__(
             max_retries=max_retries,
             cached_content=cached_content,
             built_in_tool=built_in_tool,
+            use_file_api=use_file_api,
             **kwargs,
         )
 
@@ -297,9 +305,21 @@ def _chat(self, messages: Sequence[ChatMessage], **kwargs: Any):
             **kwargs.pop("generation_config", {}),
         }
         params = {**kwargs, "generation_config": generation_config}
-        next_msg, chat_kwargs = prepare_chat_params(self.model, messages, **params)
+        next_msg, chat_kwargs = asyncio.run(
+            prepare_chat_params(
+                self.model, messages, self.use_file_api, self._client, **params
+            )
+        )
         chat = self._client.chats.create(**chat_kwargs)
-        response = chat.send_message(next_msg.parts)
+        response = chat.send_message(
+            next_msg.parts if isinstance(next_msg, types.Content) else next_msg
+        )
+
+        if self.use_file_api:
+            asyncio.run(
+                delete_uploaded_files([*chat_kwargs["history"], next_msg], self._client)
+            )
+
         return chat_from_gemini_response(response)
 
     @llm_retry_decorator
@@ -309,9 +329,19 @@ async def _achat(self, messages: Sequence[ChatMessage], **kwargs: Any):
             **kwargs.pop("generation_config", {}),
         }
         params = {**kwargs, "generation_config": generation_config}
-        next_msg, chat_kwargs = prepare_chat_params(self.model, messages, **params)
+        next_msg, chat_kwargs = await prepare_chat_params(
+            self.model, messages, self.use_file_api, self._client, **params
+        )
         chat = self._client.aio.chats.create(**chat_kwargs)
-        response = await chat.send_message(next_msg.parts)
+        response = await chat.send_message(
+            next_msg.parts if isinstance(next_msg, types.Content) else next_msg
+        )
+
+        if self.use_file_api:
+            await delete_uploaded_files(
+                [*chat_kwargs["history"], next_msg], self._client
+            )
+
         return chat_from_gemini_response(response)
 
     @llm_chat_callback()
@@ -332,9 +362,15 @@ def _stream_chat(
             **kwargs.pop("generation_config", {}),
         }
         params = {**kwargs, "generation_config": generation_config}
-        next_msg, chat_kwargs = prepare_chat_params(self.model, messages, **params)
+        next_msg, chat_kwargs = asyncio.run(
+            prepare_chat_params(
+                self.model, messages, self.use_file_api, self._client, **params
+            )
+        )
         chat = self._client.chats.create(**chat_kwargs)
-        response = chat.send_message_stream(next_msg.parts)
+        response = chat.send_message_stream(
+            next_msg.parts if isinstance(next_msg, types.Content) else next_msg
+        )
 
         def gen() -> ChatResponseGen:
             content = ""
@@ -361,6 +397,13 @@ def gen() -> ChatResponseGen:
                 llama_resp.message.additional_kwargs["tool_calls"] = existing_tool_calls
                 yield llama_resp
 
+            if self.use_file_api:
+                asyncio.run(
+                    delete_uploaded_files(
+                        [*chat_kwargs["history"], next_msg], self._client
+                    )
+                )
+
         return gen()
 
     @llm_chat_callback()
@@ -377,14 +420,18 @@ async def _astream_chat(
             **kwargs.pop("generation_config", {}),
         }
         params = {**kwargs, "generation_config": generation_config}
-        next_msg, chat_kwargs = prepare_chat_params(self.model, messages, **params)
+        next_msg, chat_kwargs = await prepare_chat_params(
+            self.model, messages, self.use_file_api, self._client, **params
+        )
         chat = self._client.aio.chats.create(**chat_kwargs)
 
         async def gen() -> ChatResponseAsyncGen:
             content = ""
             existing_tool_calls = []
             thoughts = ""
-            async for r in await chat.send_message_stream(next_msg.parts):
+            async for r in await chat.send_message_stream(
+                next_msg.parts if isinstance(next_msg, types.Content) else next_msg
+            ):
                 if candidates := r.candidates:
                     if not candidates:
                         continue
@@ -412,6 +459,11 @@ async def gen() -> ChatResponseAsyncGen:
                             )
                             yield llama_resp
 
+            if self.use_file_api:
+                await delete_uploaded_files(
+                    [*chat_kwargs["history"], next_msg], self._client
+                )
+
         return gen()
 
     @llm_chat_callback()
@@ -529,9 +581,15 @@ def structured_predict_without_function_calling(
         llm_kwargs = llm_kwargs or {}
 
         messages = prompt.format_messages(**prompt_args)
+        contents = [
+            asyncio.run(
+                chat_message_to_gemini(message, self.use_file_api, self._client)
+            )
+            for message in messages
+        ]
         response = self._client.models.generate_content(
             model=self.model,
-            contents=list(map(chat_message_to_gemini, messages)),
+            contents=contents,
             **{
                 **llm_kwargs,
                 **{
@@ -543,6 +601,9 @@ def structured_predict_without_function_calling(
             },
         )
 
+        if self.use_file_api:
+            asyncio.run(delete_uploaded_files(contents, self._client))
+
         if isinstance(response.parsed, BaseModel):
             return response.parsed
         else:
@@ -570,13 +631,21 @@ def structured_predict(
             generation_config["response_schema"] = output_cls
 
             messages = prompt.format_messages(**prompt_args)
-            contents = list(map(chat_message_to_gemini, messages))
+            contents = [
+                asyncio.run(
+                    chat_message_to_gemini(message, self.use_file_api, self._client)
+                )
+                for message in messages
+            ]
             response = self._client.models.generate_content(
                 model=self.model,
                 contents=contents,
                 config=generation_config,
             )
 
+            if self.use_file_api:
+                asyncio.run(delete_uploaded_files(contents, self._client))
+
             if isinstance(response.parsed, BaseModel):
                 return response.parsed
             else:
@@ -609,13 +678,21 @@ async def astructured_predict(
             generation_config["response_schema"] = output_cls
 
             messages = prompt.format_messages(**prompt_args)
-            contents = list(map(chat_message_to_gemini, messages))
+            contents = await asyncio.gather(
+                *[
+                    chat_message_to_gemini(message, self.use_file_api, self._client)
+                    for message in messages
+                ]
+            )
             response = await self._client.aio.models.generate_content(
                 model=self.model,
                 contents=contents,
                 config=generation_config,
             )
 
+            if self.use_file_api:
+                await delete_uploaded_files(contents, self._client)
+
             if isinstance(response.parsed, BaseModel):
                 return response.parsed
             else:
@@ -648,7 +725,12 @@ def stream_structured_predict(
             generation_config["response_schema"] = output_cls
 
             messages = prompt.format_messages(**prompt_args)
-            contents = list(map(chat_message_to_gemini, messages))
+            contents = [
+                asyncio.run(
+                    chat_message_to_gemini(message, self.use_file_api, self._client)
+                )
+                for message in messages
+            ]
 
             def gen() -> Generator[Union[Model, FlexibleModel], None, None]:
                 flexible_model = create_flexible_model(output_cls)
@@ -672,6 +754,9 @@ def gen() -> Generator[Union[Model, FlexibleModel], None, None]:
                         if streaming_model:
                             yield streaming_model
 
+                if self.use_file_api:
+                    asyncio.run(delete_uploaded_files(contents, self._client))
+
             return gen()
         else:
             return super().stream_structured_predict(
@@ -700,7 +785,12 @@ async def astream_structured_predict(
             generation_config["response_schema"] = output_cls
 
             messages = prompt.format_messages(**prompt_args)
-            contents = list(map(chat_message_to_gemini, messages))
+            contents = await asyncio.gather(
+                *[
+                    chat_message_to_gemini(message, self.use_file_api, self._client)
+                    for message in messages
+                ]
+            )
 
             async def gen() -> AsyncGenerator[Union[Model, FlexibleModel], None]:
                 flexible_model = create_flexible_model(output_cls)
@@ -724,6 +814,9 @@ async def gen() -> AsyncGenerator[Union[Model, FlexibleModel], None]:
                         if streaming_model:
                             yield streaming_model
 
+                if self.use_file_api:
+                    await delete_uploaded_files(contents, self._client)
+
             return gen()
         else:
             return await super().astream_structured_predict(