Update app.py

app.py

@@ -1,6 +1,7 @@
 import gradio as gr
 import openai
 from openai import OpenAI
+import google.generativeai as genai
 import os
 import io
 import base64
@@ -8,8 +9,9 @@ import base64
 # Set API key
 #base_url = os.environ.get("OPENAI_API_BASE")
 api_key = ""
-client = OpenAI(api_key=api_key)
 
+endpoints = ''
+    
 # Define the model to be used
 MODEL = os.environ.get("MODEL")
 
@@ -36,6 +38,17 @@ footer {
 
 LICENSE = 'MODEL: ' + MODEL + ' LOADED'
 
+
+def endpoints(api_key):
+    if api_key not None:
+        if api_key[:3] == "sk-":
+            return 'OPENAI'
+        else:
+            return 'GOOGLE'
+    return
+
+endpoints = endpoints(api_key)
+
 def read(filename):
     with open(filename) as f:
         data = f.read()
@@ -44,7 +57,8 @@ def read(filename):
 SYS_PROMPT = read('system_prompt.txt')
 
 def process_text(text_input, unit):
-    if text_input:
+    if text_input and endpoints == 'OPENAI':
+        client = OpenAI(api_key=api_key)
         completion = client.chat.completions.create(
             model=MODEL,
             messages=[
@@ -53,6 +67,10 @@ def process_text(text_input, unit):
             ]
         )
         return completion.choices[0].message.content
+    elif text_input and endpoints == "GOOGLE":
+        genai.configure(api_key=GOOGLE_API_KEY)
+        model = genai.GenerativeModel(MODEL),
+        return response = model.generate_content(SYS_PROMPT + text_input).text
     return ""
 
 def encode_image_to_base64(image_input):
@@ -62,7 +80,8 @@ def encode_image_to_base64(image_input):
     return img_str
 
 def process_image(image_input, unit):
-    if image_input is not None:
+    if image_input is not None and endpoints == 'OPENAI':
+        client = OpenAI(api_key=api_key)
         #with open(image_input.name, "rb") as f:
         #    base64_image = base64.b64encode(f.read()).decode("utf-8")
         base64_image = encode_image_to_base64(image_input)
@@ -83,6 +102,10 @@ def process_image(image_input, unit):
             max_tokens=1024,
         )
         return response.choices[0].message.content
+    elif image_input is not None and ENDPOINTS == "GOOGLE":
+        genai.configure(api_key=GOOGLE_API_KEY)
+        model = genai.GenerativeModel(MODEL),
+        return response = model.generate_content(image_input).text
 
 
 def main(text_input="", image_input=None, unit=""):
@@ -102,7 +125,7 @@ with gr.Blocks(theme='shivi/calm_seafoam', css=css, title="Medster - Medical Dia
     with gr.Row():
         output_box = gr.Markdown(label="Diagnosis")  # Create an output textbox
     with gr.Row():
-        api_key = gr.Textbox(label="OpenAI API Key") # Input API key
+        api_key = gr.Textbox(label="API Key") # Input API key
     with gr.Row():
         image_input = gr.Image(type="pil", label="Upload Image")  # Create an image upload button
         text_input = gr.Textbox(label="Submit")  # Create a text input box