Create ai_model.py

This file configures and interacts with the AI model (Gemini Pro) through the Google Generative AI API.

ai_model.py

@@ -0,0 +1,18 @@
+import google.generativeai as genai
+import os
+
+# Load Google API key from environment
+GOOGLE_API_KEY = os.getenv("GOOGLE_API_KEY")
+genai.configure(api_key=GOOGLE_API_KEY)
+model = genai.GenerativeModel('gemini-pro')
+
+# Query AI model
+async def query_ai_model(prompt):
+    try:
+        response = model.generate_content([{'role': 'user', 'parts': [prompt]}], stream=True)
+        result = ''
+        for chunk in response:
+            result += chunk.text
+        return result
+    except Exception as e:
+        return f"An error occurred: {str(e)}"