Update ai_model.py

ai_model.py

@@ -3,6 +3,7 @@
 import google.generativeai as genai
 import os
 import logging
+from typing import List, Dict, Optional
 from config import (
     SYSTEM_INSTRUCTION,
     MODEL_NAME,
@@ -12,7 +13,7 @@ from config import (
     DEFAULT_TOP_K,
     DEFAULT_STOP_SEQUENCES,
 )
-from typing import List, Optional, Dict
+import asyncio
 
 # Configure logging
 logging.basicConfig(level=logging.ERROR)
@@ -33,14 +34,11 @@ model = genai.GenerativeModel(
 # Preprocess chat history to the required format
 def preprocess_chat_history(history: List[tuple]) -> List[Dict[str, str]]:
     messages = []
-    for user_message, model_message in history:
-        if isinstance(user_message, tuple):
-            # If user_message is a tuple (e.g., contains images), skip for text-only model
-            continue
-        elif user_message is not None:
+    for user_message, assistant_message in history:
+        if 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})
+        if assistant_message is not None:
+            messages.append({'role': 'assistant', 'content': assistant_message})
     return messages
 
 # Query AI model
@@ -65,7 +63,7 @@ async def query_ai_model(
             temperature=temperature,
             top_p=top_p,
             top_k=top_k,
-            max_output_tokens=int(max_output_tokens),
+            max_output_tokens=max_output_tokens,
             stop_sequences=stop_sequences,
         )
 
@@ -78,7 +76,4 @@ async def query_ai_model(
         # 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", exc_info=True)
-        return {'role': 'assistant', 'content': f"An error occurred: {str(e)}"}
+        return