Update app.py

app.py

@@ -6,29 +6,44 @@ import json
 import uuid
 import os
 
+default_urls = ["https://api.anthropic.com", "https://api.openai.com/v1"]
+
 # List of available Claude models
 claude_models = ["claude-3-opus-20240229", "claude-3-sonnet-20240229", "claude-3-haiku-20240307"]
 
 # List of available OpenAI models
 openai_models = ["gpt-4", "gpt-4-32k", "gpt-3.5-turbo", "gpt-4-0125-preview", "gpt-4-turbo-preview", "gpt-4-1106-preview", "gpt-4-0613"]
 
-def generate_response(api_key, model, user_prompt):
-    if api_key.startswith("sk-ant-"):
-        # Set the Anthropic API key
-        client = Anthropic(api_key=api_key)
-        system_prompt_path = __file__.replace("app.py", "json.txt")
-    elif api_key.startswith("sk-"):
-        # Set the OpenAI API key
-        client = OpenAI(api_key=api_key)
-        system_prompt_path = __file__.replace("app.py", "json.txt")
+both_models = claude_models + openai_models
+
+def generate_response(endpoint, api_key, model, user_prompt):
+    print(endpoint)
+    if endpoint in default_urls:
+        #check api keys as normal
+        if api_key.startswith("sk-ant-"):
+            client = Anthropic(api_key=api_key, base_url=endpoint)
+            system_prompt_path = __file__.replace("app.py", "json.txt")
+        elif api_key.startswith("sk-"):
+            client = OpenAI(api_key=api_key, base_url=endpoint)
+            system_prompt_path = __file__.replace("app.py", "json.txt")
+        else:
+            print("Invalid API key")
+            return "Invalid API key"
     else:
-        return "Invalid API key prefix. Please use 'sk-ant-' for Anthropic or 'sk-' for OpenAI.", "", None
+        if model in claude_models:
+            # Set the Anthropic API key
+            client = Anthropic(api_key=api_key, base_url=endpoint)
+            system_prompt_path = __file__.replace("app.py", "json.txt")
+        else:
+            # Set the OpenAI API key
+            client = OpenAI(api_key=api_key, base_url=endpoint)
+            system_prompt_path = __file__.replace("app.py", "json.txt")
 
     # Read the system prompt from a text file
     with open(system_prompt_path, "r") as file:
         system_prompt = file.read()
 
-    if api_key.startswith("sk-ant-"):
+    if model in claude_models:
         # Generate a response using the selected Anthropic model
         response = client.messages.create(
             system=system_prompt,
@@ -90,23 +105,33 @@ def extract_json(generated_output):
         print(e)
         return None, None
 
-def generate_second_response(api_key, model, generated_output):
-    if api_key.startswith("sk-ant-"):
-        # Set the Anthropic API key
-        client = Anthropic(api_key=api_key)
-        system_prompt_path = __file__.replace("app.py", "diffusion.txt")
-    elif api_key.startswith("sk-"):
-        # Set the OpenAI API key
-        client = OpenAI(api_key=api_key)
-        system_prompt_path = __file__.replace("app.py", "diffusion.txt")
+def generate_second_response(endpoint, api_key, model, generated_output):
+    if endpoint in default_urls:
+        #check api keys as normal
+        if api_key.startswith("sk-ant-"):
+            client = Anthropic(api_key=api_key, base_url=endpoint)
+            system_prompt_path = __file__.replace("app.py", "diffusion.txt")
+        elif api_key.startswith("sk-"):
+            client = OpenAI(api_key=api_key, base_url=endpoint)
+            system_prompt_path = __file__.replace("app.py", "diffusion.txt")
+        else:
+            print("Invalid API key")
+            return "Invalid API key"
     else:
-        return "Invalid API key prefix. Please use 'sk-ant-' for Anthropic or 'sk-' for OpenAI."
+        if model in claude_models:
+            # Set the Anthropic API key
+            client = Anthropic(api_key=api_key, base_url=endpoint)
+            system_prompt_path = __file__.replace("app.py", "diffusion.txt")
+        else:
+            # Set the OpenAI API key
+            client = OpenAI(api_key=api_key, base_url=endpoint)
+            system_prompt_path = __file__.replace("app.py", "diffusion.txt")
 
     # Read the system prompt from a text file
     with open(system_prompt_path, "r") as file:
         system_prompt = file.read()
 
-    if api_key.startswith("sk-ant-"):
+    if model in claude_models:
         # Generate a second response using the selected Anthropic model and the previously generated output
         response = client.messages.create(
             system=system_prompt,
@@ -138,6 +163,7 @@ with gr.Blocks() as demo:
 
     with gr.Row():
         with gr.Column():
+            endpoint = gr.Textbox(label="Endpoint", value="https://api.anthropic.com")
             api_key = gr.Textbox(label="API Key", type="password", placeholder="sk-ant-api03-... or sk-...")
             model_dropdown = gr.Dropdown(choices=[], label="Select a model")
             user_prompt = gr.Textbox(label="User Prompt", value="Make me a card for a panther made of translucent pastel colored goo. Its color never changes once it exists but each 'copy' has a different color. The creature comes out of a small jar, seemingly defying physics with its size. It is the size of a real panther, and as strong as one too. By default its female but is able to change gender. It can even split into multiple copies of itself if needed with no change in its own size or mass. Its outside is normally lightly squishy but solid, but on command it can become viscous like non-newtonian fluids. Be descriptive when describing this character, and make sure to describe all of its features in char_persona just like you do in description. Make sure to describe commonly used features in detail (visual, smell, taste, touch, etc).")
@@ -157,15 +183,15 @@ with gr.Blocks() as demo:
 
     def update_models(api_key):
         if api_key.startswith("sk-ant-"):
-            return gr.Dropdown(choices=claude_models)
+            return gr.Dropdown(choices=claude_models), gr.Textbox(label="Endpoint", value="https://api.anthropic.com")
         elif api_key.startswith("sk-"):
-            return gr.Dropdown(choices=openai_models)
+            return gr.Dropdown(choices=openai_models), gr.Textbox(label="Endpoint", value="https://api.openai.com/v1")
         else:
-            return gr.Dropdown(choices=[])
+            return gr.Dropdown(choices=both_models), gr.Textbox(label="Endpoint", value="https://api.anthropic.com")
 
-    api_key.change(update_models, inputs=api_key, outputs=model_dropdown)
+    api_key.change(update_models, inputs=api_key, outputs=[model_dropdown, endpoint])
 
-    generate_button.click(generate_response, inputs=[api_key, model_dropdown, user_prompt], outputs=[generated_output, json_output, json_download])
-    generate_button_2.click(generate_second_response, inputs=[api_key, model_dropdown, generated_output], outputs=generated_output_2)
+    generate_button.click(generate_response, inputs=[endpoint, api_key, model_dropdown, user_prompt], outputs=[generated_output, json_output, json_download])
+    generate_button_2.click(generate_second_response, inputs=[endpoint, api_key, model_dropdown, generated_output], outputs=generated_output_2)
 
 demo.launch()