(feat:attach) strategy specific to google genai

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}"