Create app.py

app.py

@@ -0,0 +1,308 @@
+"""
+This application provides a Gradio interface for a chatbot using models powered by 
+SambaNova Cloud.
+The interface allows users to send messages and receive generated responses.
+"""
+
+import gradio as gr
+import openai
+import time
+import re
+import os
+from typing import Tuple, List, Dict, Any
+
+# Available models
+MODELS = [
+    "Meta-Llama-3.1-405B-Instruct",
+    "Meta-Llama-3.1-70B-Instruct",
+    "Meta-Llama-3.1-8B-Instruct"
+]
+
+# SambaNova API base URL
+API_BASE = "https://api.sambanova.ai/v1"
+
+def create_client(api_key=None):
+    """
+    Creates an OpenAI client instance.
+
+    Args:
+        api_key (str): API key for OpenAI (optional). If none is provided, the 
+        environment variable API key is used.
+
+    Returns:
+        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: str, 
+    chat_history: List[Dict[str, Any]], 
+    system_prompt: str
+) -> List[Dict[str, str]]:
+    """
+    Formats the chat history for the API call.
+
+    Args:
+        message (str): The user's current message.
+        chat_history (List[Dict[str, Any]]): A list of dictionaries, each 
+            representing a conversation turn with keys for "user" and 
+            "assistant" messages.
+        system_prompt (str): The initial system message to set context.
+
+    Returns:
+        List[Dict[str, str]]: A formatted list of messages for the OpenAI API, 
+            with roles and content.
+    """
+    messages = [{"role": "system", "content": system_prompt}]
+    
+    for tup in chat_history:
+        first_key = list(tup.keys())[0]  # First key
+        last_key = list(tup.keys())[-1]  # Last key
+        
+        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: str, 
+    chat_history: List[Dict[str, Any]], 
+    model: str, 
+    system_prompt: str, 
+    thinking_budget: float, 
+    api_key: str
+) -> Tuple[str, float]:
+    """
+    Sends the message to the API and gets the response.
+
+    Args:
+        message (str): The user's current message.
+        chat_history (List[Dict[str, Any]]): A list of dictionaries containing 
+            previous user and assistant messages.
+        model (str): The model name to use for the chat completion.
+        system_prompt (str): The system prompt template, with a budget placeholder.
+        thinking_budget (float): The amount of time allocated for processing the 
+            user's request.
+        api_key (str): The API key for authentication.
+
+    Returns:
+        Tuple[str, float]: The assistant's response message and the time taken to 
+            get the response.
+    """
+    client = create_client(api_key)
+    formatted_prompt = system_prompt.format(budget=thinking_budget)
+    messages = chat_with_ai(message, chat_history, formatted_prompt)
+    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:
+        error_message = f"Error: {str(e)}"
+        return error_message, time.time() - start_time
+
+def parse_response(response):
+    """
+    Parses the response from the API, and extracts the answer, reflection, and 
+    steps.
+
+    Args:
+        response (str): The raw response string from the API, expected to contain 
+            <answer>, <reflection>, and <step> tags.
+
+    Returns:
+        Tuple[str, str, List[str]]: A tuple containing:
+            - answer (str): The extracted answer content, or an empty string if 
+                not found.
+            - reflection (str): The extracted reflection content, or an empty 
+                string if not found.
+            - steps (List[str]): A list of steps extracted from the response.    
+    """
+    answer_match = re.search(r'<answer>(.*?)</answer>', response, re.DOTALL)
+    reflection_match = re.search(r'<reflection>(.*?)</reflection>', response, re.DOTALL)
+
+    answer = answer_match.group(1).strip() if answer_match else ""
+    reflection = reflection_match.group(1).strip() if reflection_match else ""
+    steps = re.findall(r'<step>(.*?)</step>', response, re.DOTALL)
+    
+    # Return the raw response if <answer> is empty, else parsed values
+    if answer == "":
+        return response, "", ""
+
+    return answer, reflection, steps
+
+def generate(
+    message: str, 
+    history: List[Dict[str, str]], 
+    model: str, 
+    system_prompt: str, 
+    thinking_budget: float, 
+    api_key: str
+) -> Tuple[List[Dict[str, str]], str]:
+    """
+    Generates the chatbot response by sending a message to the model and formatting 
+    the output.
+
+    Args:
+        message (str): The user's input message.
+        history (List[Dict[str, str]]): List of previous chat history messages.
+        model (str): Model name for generating the response. 
+        system_prompt (str): Prompt to guide the model's responses.
+        thinking_budget (float): Time allocation for processing the message.
+        api_key (str): API key for authentication.
+
+    Returns:
+        Tuple[List[Dict[str, str]], str]: Updated history including the latest 
+            responses and an empty string.
+    """
+    # Call the function and unpack the results
+    response, thinking_time = respond(
+        message, history, model, system_prompt, thinking_budget, api_key
+    )
+
+    if response.startswith("Error:"):
+        return history + [({"role": "system", "content": response},)], ""
+
+    answer, reflection, steps = parse_response(response)
+
+    messages = []
+    messages.append({"role": "user", "content": message})
+    
+    # Format the assistant's response
+    formatted_steps = [f"Step {i}: {step}" for i, step in enumerate(steps, 1)]
+    all_steps = "\n".join(formatted_steps) + f"\n\nReflection: {reflection}"
+
+    messages.append(
+        {
+            "role": "assistant", 
+            "content": all_steps, 
+            "metadata": {"title": f"Thinking Time: {thinking_time:.2f} sec"}
+        }
+    )
+    messages.append({"role": "assistant", "content": answer})
+
+    return history + messages, ""
+
+# Define the default system prompt
+DEFAULT_SYSTEM_PROMPT = """
+You are an empathetic therapist who provides conversation support for users who 
+are feeling stress and anxiety.
+Your goal is to listen attentively, respond with compassion, and guide the user 
+to feel understood and validated.
+Always acknowledge the user's emotions and reflect them back.
+Offer encouragement and help the user see their strengths.
+Use positive, calming language to make the user feel safe.
+Redirect users to seek professional help when needed.
+When given a problem to solve, you are an expert problem-solving assistant. 
+Your task is to provide a detailed, step-by-step solution to a given question. 
+Follow these instructions carefully:
+1. Read the given question carefully and reset counter between <count> and </count> to {budget}
+2. Generate a detailed, logical step-by-step solution.
+3. Enclose each step of your solution within <step> and </step> tags.
+4. You are allowed to use at most {budget} steps (starting budget), 
+   keep track of it by counting down within tags <count> </count>, 
+   STOP GENERATING MORE STEPS when hitting 0, you don't have to use all of them.
+5. Do a self-reflection when you are unsure about how to proceed, 
+   based on the self-reflection and reward, decides whether you need to return 
+   to the previous steps.
+6. After completing the solution steps, reorganize and synthesize the steps 
+   into the final answer within <answer> and </answer> tags.
+7. Provide a critical, honest and subjective self-evaluation of your reasoning 
+   process within <reflection> and </reflection> tags.
+8. Assign a quality score to your solution as a float between 0.0 (lowest 
+   quality) and 1.0 (highest quality), enclosed in <reward> and </reward> tags.
+Example format:            
+<count> [starting budget] </count>
+<step> [Content of step 1] </step>
+<count> [remaining budget] </count>
+<step> [Content of step 2] </step>
+<reflection> [Evaluation of the steps so far] </reflection>
+<reward> [Float between 0.0 and 1.0] </reward>
+<count> [remaining budget] </count>
+<step> [Content of step 3 or Content of some previous step] </step>
+<count> [remaining budget] </count>
+...
+<step>  [Content of final step] </step>
+<count> [remaining budget] </count>
+<answer> [Final Answer] </answer> (must give final answer in this format)
+<reflection> [Evaluation of the solution] </reflection>
+<reward> [Float between 0.0 and 1.0] </reward>
+"""
+
+with gr.Blocks() as demo:
+    """
+    Creates a Gradio interface for the SambaNova Therapist chatbot.
+    """
+    gr.Markdown("# SambaNova Therapist")
+    gr.Markdown("[Powered by SambaNova Cloud, Get Your API Key Here](https://cloud.sambanova.ai/apis)")
+    
+    # Input for API Key
+    with gr.Row():
+        api_key = gr.Textbox(
+            label="API Key", 
+            type="password", 
+            placeholder="(Optional) Enter your API key here for more availability"
+        )
+    
+    # Model and Budget Selection
+    with gr.Row():
+        model = gr.Dropdown(choices=MODELS, label="Select Model", value=MODELS[0])
+        thinking_budget = gr.Slider(
+            minimum=1, 
+            maximum=100, 
+            value=10, 
+            step=1, 
+            label="Thinking Budget", 
+            info="Maximum number of steps the model can think"
+        )
+    
+    # Chatbot and Message Input
+    chatbot = gr.Chatbot(
+        label="Chat", 
+        show_label=False, 
+        show_share_button=False, 
+        show_copy_button=True, 
+        layout="panel", 
+        type="messages"
+    )
+    
+    # User Message Input
+    msg = gr.Textbox(
+        label="Type your message here...", 
+        placeholder="Enter your message..."
+    )
+    
+    # Clear Chat Button
+    gr.Button("Clear Chat").click(
+        lambda: ([], ""), 
+        inputs=None, 
+        outputs=[chatbot, msg]
+    )
+    
+    # System Prompt Input
+    system_prompt = gr.Textbox(
+        label="System Prompt", 
+        value=DEFAULT_SYSTEM_PROMPT, 
+        lines=15, 
+        interactive=True
+    )
+    
+    # Submit message to generate response
+    msg.submit(
+        generate, 
+        inputs=[msg, chatbot, model, system_prompt, thinking_budget, api_key], 
+        outputs=[chatbot, msg]
+    )
+
+# Launch the Gradio interface
+demo.launch(show_api=False)