diff --git a/application/llm/google_ai.py b/application/llm/google_ai.py index 5e33550c..180ed43b 100644 --- a/application/llm/google_ai.py +++ b/application/llm/google_ai.py @@ -1,5 +1,8 @@ from google import genai from google.genai import types +import os +import logging +import mimetypes from application.llm.base import BaseLLM @@ -9,6 +12,141 @@ class GoogleLLM(BaseLLM): super().__init__(*args, **kwargs) self.api_key = api_key self.user_api_key = user_api_key + self.client = genai.Client(api_key=self.api_key) + + def get_supported_attachment_types(self): + """ + Return a list of MIME types supported by Google Gemini for file uploads. + + Returns: + list: List of supported MIME types + """ + return [ + 'application/pdf', + 'image/png', + 'image/jpeg', + 'image/jpg', + 'image/webp', + 'image/gif' + ] + + def prepare_messages_with_attachments(self, messages, attachments=None): + """ + Process attachments using Google AI's file API for more efficient handling. + + Args: + messages (list): List of message dictionaries. + attachments (list): List of attachment dictionaries with content and metadata. + + Returns: + list: Messages formatted with file references for Google AI API. + """ + + if not attachments: + return messages + + prepared_messages = messages.copy() + + # Find the user message to attach files to the last one + user_message_index = None + for i in range(len(prepared_messages) - 1, -1, -1): + if prepared_messages[i].get("role") == "user": + user_message_index = i + break + + + if user_message_index is None: + user_message = {"role": "user", "content": []} + prepared_messages.append(user_message) + user_message_index = len(prepared_messages) - 1 + + if isinstance(prepared_messages[user_message_index].get("content"), str): + text_content = prepared_messages[user_message_index]["content"] + prepared_messages[user_message_index]["content"] = [ + {"type": "text", "text": text_content} + ] + elif not isinstance(prepared_messages[user_message_index].get("content"), list): + prepared_messages[user_message_index]["content"] = [] + + file_uris = [] + for attachment in attachments: + mime_type = attachment.get('mime_type') + if not mime_type: + file_path = attachment.get('path') + if file_path: + mime_type = mimetypes.guess_type(file_path)[0] or 'application/octet-stream' + + if mime_type in self.get_supported_attachment_types(): + try: + file_uri = self._upload_file_to_google(attachment) + logging.info(f"GoogleLLM: Successfully uploaded file, got URI: {file_uri}") + file_uris.append((file_uri, mime_type)) + except Exception as e: + logging.error(f"GoogleLLM: Error uploading file: {e}") + if 'content' in attachment: + prepared_messages[user_message_index]["content"].append({ + "type": "text", + "text": f"[File could not be processed: {attachment.get('path', 'unknown')}]" + }) + + if file_uris: + logging.info(f"GoogleLLM: Adding {len(file_uris)} file URIs to message") + prepared_messages[user_message_index]["content"].append({ + "type": "file_uris", + "file_uris": file_uris + }) + + return prepared_messages + + def _upload_file_to_google(self, attachment): + """ + Upload a file to Google AI and return the file URI. + + Args: + attachment (dict): Attachment dictionary with path and metadata. + + Returns: + str: Google AI file URI for the uploaded file. + """ + if 'google_file_uri' in attachment: + return attachment['google_file_uri'] + + file_path = attachment.get('path') + if not file_path: + raise ValueError("No file path provided in attachment") + + if not os.path.isabs(file_path): + current_dir = os.path.dirname( + os.path.dirname(os.path.dirname(os.path.abspath(__file__))) + ) + file_path = os.path.join(current_dir, "application", file_path) + + if not os.path.exists(file_path): + raise FileNotFoundError(f"File not found: {file_path}") + + mime_type = attachment.get('mime_type') + if not mime_type: + mime_type = mimetypes.guess_type(file_path)[0] or 'application/octet-stream' + + try: + response = self.client.files.upload(file=file_path) + + file_uri = response.uri + + from application.core.mongo_db import MongoDB + mongo = MongoDB.get_client() + db = mongo["docsgpt"] + attachments_collection = db["attachments"] + if '_id' in attachment: + attachments_collection.update_one( + {"_id": attachment['_id']}, + {"$set": {"google_file_uri": file_uri}} + ) + + return file_uri + except Exception as e: + logging.error(f"Error uploading file to Google AI: {e}") + raise def _clean_messages_google(self, messages): cleaned_messages = [] @@ -26,7 +164,7 @@ class GoogleLLM(BaseLLM): elif isinstance(content, list): for item in content: if "text" in item: - parts.append(types.Part.from_text(item["text"])) + parts.append(types.Part.from_text(text=item["text"])) elif "function_call" in item: parts.append( types.Part.from_function_call( @@ -41,6 +179,14 @@ class GoogleLLM(BaseLLM): response=item["function_response"]["response"], ) ) + elif "type" in item and item["type"] == "file_uris": + for file_uri, mime_type in item["file_uris"]: + parts.append( + types.Part.from_uri( + file_uri=file_uri, + mime_type=mime_type + ) + ) else: raise ValueError( f"Unexpected content dictionary format:{item}"