Update app.py

app.py

@@ -6,12 +6,12 @@ import io
 import base64
 
 # Set API key and organization ID from environment variables
-api_key = os.environ.get("OPENAI_API_KEY")
+api_key = os.environ.get("API_KEY")
 #base_url = os.environ.get("OPENAI_API_BASE")
-client = OpenAI(api_key=api_key)
 
 # Define the model to be used
 MODEL = os.environ.get("MODEL")
+MODEL_NAME = MODEL.split("/")[-1] if "/" in MODEL else MODEL
 
 def read(filename):
     with open(filename) as f:
@@ -42,18 +42,33 @@ footer {
 """
 
 
-LICENSE = '采用 ' + MODEL + ' 模型'
-    
+LICENSE = '采用 ' + MODEL_NAME + ' 模型'
+
+def endpoints(api_key):
+    if api_key is not None:
+        if api_key.startwith('sk-'):
+            return 'OPENAI'
+        else:
+            return 'GOOGLE'
+
 def process_text(text_input, unit):
-    if text_input:
+    endpoint = endpoints(api_key)
+    if text_input and endpoint == 'OPENAI':
+        client = OpenAI(api_key=api_key)
         completion = client.chat.completions.create(
             model=MODEL,
             messages=[
-                {"role": "system", "content": f" You are a experienced {unit} doctor." + SYS_PROMPT},
+                {"role": "system", "content": f" You are a experienced {unit} doctor AI assistant." + SYS_PROMPT},
                 {"role": "user", "content": f"Hello! Could you solve {text_input}?"}
             ]
         )
         return completion.choices[0].message.content
+    elif text_input and endpoint == 'GOOGLE':
+        genai.configure(api_key=api_key)
+        model = genai.GenerativeModel(model_name=MODEL)
+        prompt = f" You are a experienced {unit} doctor AI assistant." + SYS_PROMPT + f"Could you solve {text_input}?"
+        response = model.generate_content(prompt)
+        return response.text        
     return ""
 
 def encode_image_to_base64(image_input):
@@ -63,14 +78,16 @@ def encode_image_to_base64(image_input):
     return img_str
 
 def process_image(image_input, unit):
-    if image_input is not None:
+    endpoint = endpoints(api_key)
+    if image_input is not None and endpoint == 'OPENAI':
         #with open(image_input.name, "rb") as f:
         #    base64_image = base64.b64encode(f.read()).decode("utf-8")
+        client = OpenAI(api_key=api_key)
         base64_image = encode_image_to_base64(image_input)
         response = client.chat.completions.create(
             model=MODEL,
             messages=[
-                {"role": "system", "content": f" You are a experienced {unit} doctor." + SYS_PROMPT},
+                {"role": "system", "content": f" You are a experienced {unit} doctor AI assistant." + SYS_PROMPT},
                 {"role": "user", "content": [
                     {"type": "text", "text": "Help me understand what is in this picture and analysis."},
                     {"type": "image_url", 
@@ -84,6 +101,12 @@ def process_image(image_input, unit):
             max_tokens=1024,
         )
         return response.choices[0].message.content
+    elif image_input is not None and endpoint == 'GOOGLE':
+        genai.configure(api_key=api_key)
+        model = genai.GenerativeModel(model_name=MODEL)
+        prompt = f" You are a experienced {unit} doctor AI assistant." + SYS_PROMPT + "Help me understand what is in this picture and analysis."
+        response = model.generate_content([prompt, image_input],request_options={"timeout": 60})
+        return response.text        
 
 
 def main(text_input="", image_input=None, unit=""):