Update ai_model.py

ai_model.py

@@ -1,64 +1,65 @@
 import google.generativeai as genai
 import os
-from typing import List, Dict, Union
-from PIL import Image
+from config import (
+    SYSTEM_MESSAGE,
+    MODEL_NAME,
+    DEFAULT_MAX_NEW_TOKENS,
+    DEFAULT_TEMPERATURE,
+    DEFAULT_TOP_P,
+    DEFAULT_TOP_K,
+    DEFAULT_STOP_SEQUENCES,
+)
 
 # Load Google API key from environment
 GOOGLE_API_KEY = os.environ.get("GOOGLE_API_KEY")
-genai.configure(api_key=GOOGLE_API_KEY)
-
-IMAGE_CACHE_DIRECTORY = "/tmp"
-IMAGE_WIDTH = 512
-
-def preprocess_image(image: Image.Image) -> Image.Image:
-    image_height = int(image.height * IMAGE_WIDTH / image.width)
-    return image.resize((IMAGE_WIDTH, image_height))
+if not GOOGLE_API_KEY:
+    raise ValueError("GOOGLE_API_KEY is not set. Please provide your API key.")
 
-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
+genai.configure(api_key=GOOGLE_API_KEY)
 
+# Query AI model
 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,
+    message,
+    history=None,
+    system_message=SYSTEM_MESSAGE,
+    max_new_tokens=DEFAULT_MAX_NEW_TOKENS,
+    temperature=DEFAULT_TEMPERATURE,
+    top_p=DEFAULT_TOP_P,
+    top_k=DEFAULT_TOP_K,
+    stop_sequences=DEFAULT_STOP_SEQUENCES,
 ):
     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,
+        # Build the conversation history in the required format
+        messages = []
+
+        if system_message:
+            messages.append({'role': 'system', 'content': system_message})
+
+        if history:
+            for user_msg, bot_reply in history:
+                if user_msg:
+                    messages.append({'role': 'user', 'content': user_msg})
+                if bot_reply:
+                    messages.append({'role': 'assistant', 'content': bot_reply})
+
+        # Append the new user message
+        messages.append({'role': 'user', 'content': message})
+
+        # Set parameters
+        parameters = {
+            'temperature': temperature,
+            'top_p': top_p,
+            'top_k': top_k,
+            'max_output_tokens': int(max_new_tokens),
+            'stop_sequences': stop_sequences,
+        }
+
+        # Generate response
+        response = genai.generate_chat(
+            model=MODEL_NAME,
+            messages=messages,
+            **parameters
         )
-        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
+        return response.candidates[0]['content']
     except Exception as e:
         return f"An error occurred: {str(e)}"