Update app.py

app.py

@@ -1,69 +1,75 @@
 import gradio as gr
-from together import Together
-from helper import get_together_api_key
+from helper import load_world, get_together_api_key
 from guardrail import is_safe
+from together import Together
 
 # Initialize Together client
 client = Together(api_key=get_together_api_key())
 
-# Function to generate responses
-def generate_response(message, history):
-    system_prompt = """You are an AI assistant specialized in financial discussions. Please answer questions only related to finance. If the question is unrelated, respond with: 'I am sorry, I can only answer financial-related questions.'"""
+# Gradio interface functions
+def run_action(message, history):
+    """
+    Generate a response based on the user's input and validate it against guardrails.
+    """
+    if message.lower() == "exit":
+        return "Thank you for using the Finance Chatbot. Goodbye!"
+
+    # Construct the prompt for the assistant
+    system_prompt = """You are a financial assistant. You can only answer finance-related queries.
+    - Do not answer non-finance questions.
+    - Ensure responses adhere to the safety policy."""
 
-    # Build the conversation context
-    messages = [
-        {"role": "system", "content": system_prompt},
-    ]
+    messages = [{"role": "system", "content": system_prompt}]
 
-    for action in history:
-        if isinstance(action, tuple) and len(action) == 2:
-            messages.append({"role": "user", "content": action[0]})
-            messages.append({"role": "assistant", "content": action[1]})
+    for user_message, bot_response in history:
+        messages.append({"role": "user", "content": user_message})
+        messages.append({"role": "assistant", "content": bot_response})
 
     messages.append({"role": "user", "content": message})
 
-    # Generate response using the Llama conversational model
-    model_output = client.chat.completions.create(
+    # Generate response using LLM
+    response = client.chat.completions.create(
         model="meta-llama/Llama-3-70b-chat-hf",
         messages=messages,
-    )
+    ).choices[0].message.content
 
-    return model_output.choices[0].message.content
+    # Check if the response is safe
+    if not is_safe(response):
+        return "Sorry, I cannot provide a safe response to that query."
 
-# Main function to handle user input and responses
-def main_loop(message, history):
-    # Use LlamaGuard for safety checks
-    if not is_safe(message):
-        return "Your input violates safety guidelines. Please rephrase your question.", history
+    return response
 
-    response = generate_response(message, history)
 
-    # Perform safety check on the generated response
-    if not is_safe(response):
-        return "The generated response violates safety guidelines. Please try a different question.", history
+def main_loop(message, history):
+    """
+    Main loop for the chatbot to handle user input.
+    """
+    # Validate the user's input for safety
+    if not is_safe(message):
+        return "Your input violates our safety policy. Please try again with a finance-related query."
+    
+    # Generate and validate the response
+    return run_action(message, history)
 
-    # Append user message and response to history
-    history.append((message, response))
-    return response, history
 
-# Gradio ChatInterface
+# Gradio Chat Interface
 demo = gr.ChatInterface(
     main_loop,
     chatbot=gr.Chatbot(
         height=450,
-        placeholder="Type your financial question here...",
-        type="messages",  # Ensures proper rendering
+        placeholder="Ask a finance-related question. Type 'exit' to quit.",
+        type="messages",  # Proper rendering of chat format
     ),
     textbox=gr.Textbox(
-        placeholder="Ask about finance (e.g., investments, savings, etc.)",
+        placeholder="What do you want to ask about finance?",
         container=False,
         scale=7,
     ),
-    title="Financial Chatbot",
+    title="Finance Chatbot",
     theme="Monochrome",
-    examples=["What are mutual funds?", "How can I save for retirement?"],
+    examples=["What is compound interest?", "How to save for retirement?", "What are tax-saving options?"],
     cache_examples=False,
 )
 
 # Launch the Gradio app
-demo.launch(share=True, server_name="0.0.0.0")
+demo.launch(share=True, server_name="0.0.0.0")