ai_model.py

import os
import google.generativeai as genai
import asyncio

# 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 not set in the environment variables.")

# Configure Google AI model
genai.configure(api_key=GOOGLE_API_KEY)
model = genai.GenerativeModel('gemini-pro')

# Function to query AI model
async def query_ai_model(prompt: str) -> str:
    try:
        response = model.generate_content(prompt)
        return response.text
    except Exception as e:
        print(f"Error querying AI model: {e}")
        return "Sorry, I couldn't process your request."