Update ai_model.py

ai_model.py

@@ -1,9 +1,12 @@
+# ai_model.py
+
 import google.generativeai as genai
 import os
 import logging
 from config import (
     SYSTEM_INSTRUCTION,
     MODEL_NAME,
+    MODEL_VISION_NAME,
     DEFAULT_MAX_OUTPUT_TOKENS,
     DEFAULT_TEMPERATURE,
     DEFAULT_TOP_P,
@@ -22,13 +25,18 @@ if not GOOGLE_API_KEY:
 
 genai.configure(api_key=GOOGLE_API_KEY)
 
-# Initialize the model with system instruction
+# Initialize the models with system instruction
 model = genai.GenerativeModel(
     model_name=MODEL_NAME,
     system_instruction=SYSTEM_INSTRUCTION
 )
 
-# Query AI model
+model_vision = genai.GenerativeModel(
+    model_name=MODEL_VISION_NAME,
+    system_instruction=SYSTEM_INSTRUCTION
+)
+
+# Query AI model for text-only input
 async def query_ai_model(
     message,
     history=None,
@@ -40,7 +48,7 @@ async def query_ai_model(
 ):
     try:
         # Build the conversation history
-        messages = history.copy() if history else []
+        messages = preprocess_chat_history(history) if history else []
 
         # Append the new user message
         messages.append({'role': 'user', 'content': message})
@@ -67,3 +75,53 @@ async def query_ai_model(
     except Exception as e:
         logging.error("Error in query_ai_model", exc_info=True)
         return {'role': 'assistant', 'content': f"An error occurred: {str(e)}"}
+
+# Query AI model for text-and-image input
+async def query_ai_model_vision(
+    text_prompt: List[str],
+    image_prompt: List,
+    max_output_tokens=DEFAULT_MAX_OUTPUT_TOKENS,
+    temperature=DEFAULT_TEMPERATURE,
+    top_p=DEFAULT_TOP_P,
+    top_k=DEFAULT_TOP_K,
+    stop_sequences=DEFAULT_STOP_SEQUENCES,
+):
+    try:
+        # Combine text and image prompts
+        combined_prompt = text_prompt + image_prompt
+
+        # Set generation configuration
+        generation_config = genai.types.GenerationConfig(
+            temperature=temperature,
+            top_p=top_p,
+            top_k=top_k,
+            max_output_tokens=int(max_output_tokens),
+            stop_sequences=stop_sequences,
+        )
+
+        # Generate response with vision model
+        response = model_vision.generate_content(
+            combined_prompt,
+            generation_config=generation_config
+        )
+
+        # Extract the assistant's reply
+        assistant_reply = {'role': 'assistant', 'content': response.text}
+
+        return assistant_reply
+    except Exception as e:
+        logging.error("Error in query_ai_model_vision", exc_info=True)
+        return {'role': 'assistant', 'content': f"An error occurred: {str(e)}"}
+
+# Preprocess chat history to the required format
+def preprocess_chat_history(history: List[tuple]) -> List[dict]:
+    messages = []
+    for user_message, model_message in history:
+        if isinstance(user_message, tuple):
+            # Handle image inputs if necessary
+            pass
+        elif user_message is not None:
+            messages.append({'role': 'user', 'content': user_message})
+        if model_message is not None:
+            messages.append({'role': 'assistant', 'content': model_message})
+    return messages