Create app.py

app.py

@@ -0,0 +1,119 @@
+import gradio as gr
+import openai
+import os
+import time
+
+# Available models in increasing order of size/capability.
+MODELS = [
+    "Meta-Llama-3.1-8B-Instruct",
+    "Meta-Llama-3.1-70B-Instruct",
+    "Meta-Llama-3.1-405B-Instruct"
+]
+
+# Sambanova API base URL
+API_BASE = "https://api.sambanova.ai/v1"
+
+def create_client(api_key=None):
+    """Creates an OpenAI client instance."""
+    if api_key:
+        openai.api_key = api_key
+    else:
+        openai.api_key = os.getenv("API_KEY")
+
+    return openai.OpenAI(api_key=openai.api_key, base_url=API_BASE)
+
+def chat_with_ai(message, chat_history):
+    """Formats the chat history for the API call."""
+    messages = [{"role": "system", "content": "You are a helpful assistant."}]
+    for tup in chat_history:
+        first_key = list(tup.keys())[0]
+        last_key = list(tup.keys())[-1]
+        messages.append({"role": "user", "content": tup[first_key]})
+        messages.append({"role": "assistant", "content": tup[last_key]})
+    messages.append({"role": "user", "content": message})
+    return messages
+
+def respond(message, chat_history, model, api_key):
+    """Sends the message to the API and gets the response."""
+    client = create_client(api_key)
+    messages = chat_with_ai(message, chat_history)
+    start_time = time.time()
+
+    try:
+        completion = client.chat.completions.create(model=model, messages=messages)
+        response = completion.choices[0].message.content
+        thinking_time = time.time() - start_time
+        return response, thinking_time
+    except Exception as e:
+        return f"Error: {str(e)}", time.time() - start_time
+
+def generate(message, history, model_index, api_key):
+    """Generates the chatbot response using the current model."""
+    if type(model_index) is not int:
+        model_index = model_index[0]
+    model = MODELS[model_index]  # Select model by index
+    answer, thinking_time = respond(message, history, model, api_key)
+
+    if answer.startswith("Error:"):
+        return history + [({"role": "system", "content": answer},)], ""
+
+    messages = [
+        {"role": "user", "content": message},
+        {"role": "assistant", "content": answer + f"\n\n<sub>press regenerate to make it try harder. Time: {thinking_time:.2f} sec</sub>"},
+    ]
+
+    return history + messages, message
+
+def first_try(message, history, model_index, api_key):
+    model_index = [0, 0]
+    history, message = generate(message, history, model_index, api_key)
+    return  history, message, (0, 0)
+
+def try_harder(history, retry_state, message, api_key):
+    """Tries to generate a response, possibly escalating to a larger model."""
+    model_index, retries = retry_state  # Unpack current state.
+
+    if retries < 1:
+        retries += 1  # Allow one more attempt with the current model.
+    else:
+        # Move to the next larger model, if available.
+        model_index = min(model_index + 1, len(MODELS) - 1)
+        retries = 0  # Reset retries for the new model.
+
+    # Generate a new response with the selected model.
+    history = history[:-2] # remove the last turn of conversation
+    new_history, _ = generate(message, history, model_index, api_key)
+
+    return new_history, (model_index, retries)  # Update retry state.
+
+with gr.Blocks() as demo:
+    gr.Markdown("# 🍎 Please Try Harder")
+    gr.Markdown("[Powered by SambaNova Cloud, Get Your API Key Here](https://cloud.sambanova.ai/apis)")
+
+    with gr.Row():
+        api_key = gr.Textbox(label="API Key", type="password", placeholder="(Optional) Enter your API key here for more availability")
+
+    chatbot = gr.Chatbot(label="Chat", show_label=False, show_share_button=False, layout="panel", type="messages")
+    msg = gr.Textbox(label="Type your message here...", placeholder="Enter your message...")
+
+    # State to track (1) current model index and (2) retries on the same model.
+    retry_state = gr.State(value=(0, 0))  # (model_index, retries)
+
+    # Button to clear chat, resets history and retry state.
+    clear_btn = gr.Button("Clear Chat")
+    clear_btn.click(lambda: ([], (0, 0)), inputs=None, outputs=[chatbot, retry_state])
+
+    # Generate response when a message is submitted.
+    msg.submit(
+        first_try,
+        inputs=[msg, chatbot, retry_state, api_key],
+        outputs=[chatbot, msg, retry_state]
+    )
+
+    chatbot.retry(
+        try_harder,
+        inputs=[chatbot, retry_state, msg, api_key],
+        outputs=[chatbot, retry_state]
+    )
+
+demo.launch(share=True, show_api=False)