Update ai_model.py

ai_model.py

@@ -1,42 +1,64 @@
 import google.generativeai as genai
 import os
+from typing import List, Dict, Union
+from PIL import Image
 
 # Load Google API key from environment
 GOOGLE_API_KEY = os.environ.get("GOOGLE_API_KEY")
 genai.configure(api_key=GOOGLE_API_KEY)
 
-# Query AI model
-async def query_ai_model(message, history=None, system_message=None, max_new_tokens=512, temperature=0.7, top_p=0.95):
-    try:
-        # Build the conversation history in the required format
-        messages = []
-
-        if system_message:
-            messages.append({'author': 'system', 'content': system_message})
+IMAGE_CACHE_DIRECTORY = "/tmp"
+IMAGE_WIDTH = 512
 
-        if history:
-            for user_msg, bot_reply in history:
-                if user_msg:
-                    messages.append({'author': 'user', 'content': user_msg})
-                if bot_reply:
-                    messages.append({'author': 'assistant', 'content': bot_reply})
+def preprocess_image(image: Image.Image) -> Image.Image:
+    image_height = int(image.height * IMAGE_WIDTH / image.width)
+    return image.resize((IMAGE_WIDTH, image_height))
 
-        # Append the new user message
-        messages.append({'author': 'user', 'content': message})
+def cache_pil_image(image: Image.Image) -> str:
+    import uuid
+    image_filename = f"{uuid.uuid4()}.jpeg"
+    os.makedirs(IMAGE_CACHE_DIRECTORY, exist_ok=True)
+    image_path = os.path.join(IMAGE_CACHE_DIRECTORY, image_filename)
+    image.save(image_path, "JPEG")
+    return image_path
 
-        # Set parameters
-        parameters = {
-            'temperature': temperature,
-            'top_p': top_p,
-            'max_output_tokens': int(max_new_tokens)
-        }
-
-        # Generate response
-        response = genai.generate_chat(
-            model='chat-bison-001',  # Update model name as necessary
-            messages=messages,
-            **parameters
+async def query_ai_model(
+    messages: List[Dict[str, Union[str, List[str]]]],
+    temperature: float,
+    max_output_tokens: int,
+    stop_sequences: List[str],
+    top_k: int,
+    top_p: float,
+    use_vision: bool = False,
+    image_files: List[str] = None,
+):
+    try:
+        generation_config = genai.types.GenerationConfig(
+            temperature=temperature,
+            max_output_tokens=max_output_tokens,
+            stop_sequences=stop_sequences if stop_sequences else None,
+            top_k=top_k,
+            top_p=top_p,
         )
-        return response.candidates[0]['content']
+        if use_vision and image_files:
+            # For vision model
+            image_prompt = [Image.open(file).convert('RGB') for file in image_files]
+            text_prompt = [msg['parts'][0] for msg in messages if msg['role'] == 'user']
+            model = genai.GenerativeModel('gemini-pro-vision')
+            response = model.generate_content(
+                text_prompt + image_prompt,
+                stream=False,
+                generation_config=generation_config
+            )
+        else:
+            # For text model
+            model = genai.GenerativeModel('gemini-pro')
+            response = model.generate_content(
+                messages,
+                stream=False,
+                generation_config=generation_config
+            )
+        # Since we are not streaming, get the full response text
+        return response.result
     except Exception as e:
         return f"An error occurred: {str(e)}"