Update app.py

app.py

@@ -23,6 +23,8 @@ MODELS = [
     "mistralai/Mixtral-8x7B-Instruct-v0.1",
     "mistralai/Mistral-Nemo-Instruct-2407",
     "meta-llama/Meta-Llama-3.1-8B-Instruct",
+    "meta-llama/Meta-Llama-3.1-70B-Instruct",
+    "meta-llama/Meta-Llama-3.1-8B-Instruct",
     "meta-llama/Meta-Llama-3.1-70B-Instruct"
 ]
 
@@ -54,39 +56,40 @@ class CitingSources(BaseModel):
         description="List of sources to cite. Should be an URL of the source."
     )
 
-def chatbot_interface(message, history, model, temperature, num_calls, use_embeddings):
+def chatbot_interface(message, history, model, temperature, num_calls, use_embeddings, system_prompt):
     if not message.strip():
         return "", history
 
     history = history + [(message, "")]
 
     try:
-        for response in respond(message, history, model, temperature, num_calls, use_embeddings):
+        for response in respond(message, history, model, temperature, num_calls, use_embeddings, system_prompt):
             history[-1] = (message, response)
             yield history
-    except gr.CanceledError:
+    except gr.CancelledError:
         yield history
     except Exception as e:
         logging.error(f"Unexpected error in chatbot_interface: {str(e)}")
         history[-1] = (message, f"An unexpected error occurred: {str(e)}")
         yield history
 
-def retry_last_response(history, model, temperature, num_calls, use_embeddings):
+def retry_last_response(history, model, temperature, num_calls, use_embeddings, system_prompt):
     if not history:
         return history
     
     last_user_msg = history[-1][0]
     history = history[:-1]  # Remove the last response
     
-    return chatbot_interface(last_user_msg, history, model, temperature, num_calls, use_embeddings)
+    return chatbot_interface(last_user_msg, history, model, temperature, num_calls, use_embeddings, system_prompt)
 
-def respond(message, history, model, temperature, num_calls, use_embeddings):
+def respond(message, history, model, temperature, num_calls, use_embeddings, system_prompt):
     logging.info(f"User Query: {message}")
     logging.info(f"Model Used: {model}")
     logging.info(f"Use Embeddings: {use_embeddings}")
+    logging.info(f"System Prompt: {system_prompt}")
 
     try:
-        for main_content, sources in get_response_with_search(message, model, num_calls=num_calls, temperature=temperature, use_embeddings=use_embeddings):
+        for main_content, sources in get_response_with_search(message, model, num_calls=num_calls, temperature=temperature, use_embeddings=use_embeddings, system_prompt=system_prompt):
             response = f"{main_content}\n\n{sources}"
             first_line = response.split('\n')[0] if response else ''
             yield response
@@ -105,7 +108,7 @@ def create_web_search_vectors(search_results):
     
     return FAISS.from_documents(documents, embed)
 
-def get_response_with_search(query, model, num_calls=3, temperature=0.2, use_embeddings=True):
+def get_response_with_search(query, model, num_calls=3, temperature=0.2, use_embeddings=True, system_prompt=DEFAULT_SYSTEM_PROMPT):
     search_results = duckduckgo_search(query)
     
     if use_embeddings:
@@ -144,7 +147,7 @@ After writing the document, please provide a list of sources used in your respon
         try:
             response = client.chat_completion(
                 messages=[
-                    {"role": "system", "content": DEFAULT_SYSTEM_PROMPT},
+                    {"role": "system", "content": system_prompt},
                     {"role": "user", "content": prompt}
                 ],
                 max_tokens=max_new_tokens,
@@ -191,19 +194,19 @@ def initial_conversation():
     return [
         (None, "Welcome! I'm your AI assistant for web search. Here's how you can use me:\n\n"
                 "1. Ask me any question, and I'll search the web for information.\n"
-                "2. You can adjust the model, temperature, number of API calls, and whether to use embeddings for fine-tuned responses.\n"
+                "2. You can adjust the model, temperature, number of API calls, whether to use embeddings, and the system prompt for fine-tuned responses.\n"
                 "3. For any queries, feel free to reach out @[email protected] or discord - shreyas094\n\n"
                 "To get started, ask me a question!")
     ]
 
 demo = gr.ChatInterface(
     respond,
-    additional_inputs_accordion=gr.Accordion(label="⚙️ Parameters", open=True, render=False),
     additional_inputs=[
         gr.Dropdown(choices=MODELS, label="Select Model", value=MODELS[2]),
         gr.Slider(minimum=0.1, maximum=1.0, value=0.2, step=0.1, label="Temperature"),
         gr.Slider(minimum=1, maximum=5, value=1, step=1, label="Number of API Calls"),
         gr.Checkbox(label="Use Embeddings", value=True),
+        gr.Textbox(label="System Prompt", lines=5, value=DEFAULT_SYSTEM_PROMPT),
     ],
     title="AI-powered Web Search Assistant",
     description="Ask questions and get answers from web search results.",