Update ai_model.py

ai_model.py

@@ -1,80 +1,21 @@
-# ai_model.py
-
-import google.generativeai as genai
 import os
-import logging
-from config import (
-    SYSTEM_INSTRUCTION,
-    MODEL_NAME,
-    DEFAULT_MAX_OUTPUT_TOKENS,
-    DEFAULT_TEMPERATURE,
-    DEFAULT_TOP_P,
-    DEFAULT_TOP_K,
-    DEFAULT_STOP_SEQUENCES,
-)
+import google.generativeai as genai
 import asyncio
-from typing import List, Optional  # <-- Added this import
 
-# Configure logging
-logging.basicConfig(level=logging.ERROR)
-
-# Load Google API key from environment
+# Load API key from environment variable
 GOOGLE_API_KEY = os.environ.get("GOOGLE_API_KEY")
 if not GOOGLE_API_KEY:
-    raise ValueError("GOOGLE_API_KEY is not set. Please provide your API key.")
+    raise ValueError("Google API key not set in the environment variables.")
 
+# Configure Google AI model
 genai.configure(api_key=GOOGLE_API_KEY)
+model = genai.GenerativeModel('gemini-pro')
 
-# Initialize the model with system instruction
-model = genai.GenerativeModel(
-    model_name=MODEL_NAME,
-    system_instruction=SYSTEM_INSTRUCTION
-)
-
-# Query AI model
-async def query_ai_model(
-    message: str,
-    history: Optional[List[dict]] = None,
-    max_output_tokens: int = DEFAULT_MAX_OUTPUT_TOKENS,
-    temperature: float = DEFAULT_TEMPERATURE,
-    top_p: float = DEFAULT_TOP_P,
-    top_k: int = DEFAULT_TOP_K,
-    stop_sequences: Optional[List[str]] = DEFAULT_STOP_SEQUENCES,
-):
+# Function to query AI model
+async def query_ai_model(prompt: str) -> str:
     try:
-        # Build the conversation history
-        messages = history.copy() if history else []
-        messages.append({'role': 'user', 'content': message})
-
-        # Set generation configuration
-        generation_config = genai.types.GenerationConfig(
-            temperature=temperature,
-            top_p=top_p,
-            top_k=top_k,
-            max_output_tokens=max_output_tokens,
-            stop_sequences=stop_sequences,
-        )
-
-        # Generate response
-        response = model.generate_content(
-            messages=messages,
-            generation_config=generation_config
-        )
-
-        # Extract the assistant's reply
-        assistant_reply = {'role': 'assistant', 'content': response.text}
-
-        return assistant_reply
+        response = model.generate_content(prompt)
+        return response.text
     except Exception as e:
-        logging.error("Error in query_ai_model", 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 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
+        print(f"Error querying AI model: {e}")
+        return "Sorry, I couldn't process your request."