Update app.py

app.py

@@ -1,18 +1,25 @@
 import os
 from typing import List, Tuple
-import openai  # Assuming you're using OpenAI's API (make sure to install the OpenAI package)
 from flask import Flask, request, jsonify
+from google.cloud import vertex_ai  # Ensure to install the Google Cloud SDK (vertex-ai)
 
 # Initialize Flask app
 app = Flask(__name__)
 
-# Set the OpenAI API key
-openai.api_key = os.getenv("OPENAI_API_KEY")
+# Set the Google Cloud project ID and location (Make sure to replace with your own)
+project_id = os.getenv("GOOGLE_CLOUD_PROJECT_ID")  # Make sure to set this in your environment
+location = os.getenv("GOOGLE_CLOUD_LOCATION", "us-central1")  # Default location if not set
 
-# Define a system message
+# Initialize Vertex AI client
+vertex_ai_client = vertex_ai.PredictionServiceClient(client_options={"api_endpoint": f"{location}-aiplatform.googleapis.com"})
+
+# Define the endpoint for your model deployment
+endpoint = "projects/{project_id}/locations/{location}/endpoints/{endpoint_id}"  # Replace with your actual endpoint ID
+
+# Define a system message (if necessary)
 SYSTEM_MESSAGE = "You are a helpful assistant."
 
-# Function to generate AI response
+# Function to generate AI response using Google Gemini (Vertex AI)
 def generate_response(
     user_input: str, 
     history: List[Tuple[str, str]], 
@@ -21,7 +28,7 @@ def generate_response(
     top_p: float = 1.0
 ) -> str:
     """
-    Generates a response from the AI model.
+    Generates a response using the Google Gemini (Vertex AI) API.
     Args:
         user_input: The user's input message.
         history: A list of tuples containing the conversation history 
@@ -33,41 +40,32 @@ def generate_response(
         str: The generated response from the AI model.
     """
     try:
-        # Build the message list with system message and history
-        messages = [{"role": "system", "content": SYSTEM_MESSAGE}]
-        
-        # Iterate through the history list and format accordingly
+        # Prepare the history and current input for the model
+        conversation = [{"role": "system", "content": SYSTEM_MESSAGE}]
         for user_message, assistant_message in history:
-            messages.append({"role": "user", "content": user_message})
-            messages.append({"role": "assistant", "content": assistant_message})
-        
+            conversation.append({"role": "user", "content": user_message})
+            conversation.append({"role": "assistant", "content": assistant_message})
+
         # Add the current user input
-        messages.append({"role": "user", "content": user_input})
+        conversation.append({"role": "user", "content": user_input})
+
+        # Prepare the payload for the request to Vertex AI
+        instances = [{"content": conversation}]
+        parameters = {
+            "temperature": temperature,
+            "max_output_tokens": max_tokens,
+            "top_p": top_p,
+        }
 
-        # Generate response from the model
-        response = ""
-        for msg in openai.ChatCompletion.create(
-            model="gpt-3.5-turbo",  # You can use any model you prefer
-            messages=messages,
-            max_tokens=max_tokens,
-            temperature=temperature,
-            top_p=top_p,
-            stream=True
-        ):
-            # Check if 'choices' is present and non-empty
-            if msg and 'choices' in msg and msg['choices']:
-                # Ensure the 'delta' and 'content' properties exist before using them
-                token = msg['choices'][0].get('delta', {}).get('content', '')
-                if token:
-                    response += token
-            else:
-                # Handle unexpected response format or empty choices
-                print("Warning: Unexpected response format or empty 'choices'.")
-                break
-        return response or "Sorry, I couldn't generate a response. Please try again."
+        # Send the request to the Vertex AI API
+        response = vertex_ai_client.predict(endpoint=endpoint, instances=instances, parameters=parameters)
+
+        # Extract the response from the API output
+        ai_response = response.predictions[0].get('content', 'Sorry, I couldn’t generate a response.')
+        
+        return ai_response
 
     except Exception as e:
-        # Log the error for debugging purposes
         print(f"An error occurred: {e}")
         return "Error: An unexpected error occurred while processing your request."