Update app.py

app.py

@@ -23,8 +23,6 @@ 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"
 ]
 
@@ -36,6 +34,12 @@ MODEL_TOKEN_LIMITS = {
     "meta-llama/Meta-Llama-3.1-70B-Instruct": 8192,
 }
 
+DEFAULT_SYSTEM_PROMPT = """You are a world-class financial AI assistant, capable of complex reasoning and reflection.
+Reason through the query inside <thinking> tags, and then provide your final response inside <output> tags.
+Providing comprehensive and accurate information based on web search results is essential.
+Your goal is to synthesize the given context into a coherent and detailed response that directly addresses the user's query.
+Please ensure that your response is well-structured, factual."""
+
 def get_embeddings():
     return HuggingFaceEmbeddings(model_name="sentence-transformers/stsb-roberta-large")
 
@@ -50,14 +54,14 @@ 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):
+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):
+        for response in respond(message, history, model, temperature, num_calls, use_embeddings, system_prompt):
             history[-1] = (message, response)
             yield history
     except gr.CancelledError:
@@ -67,21 +71,22 @@ def chatbot_interface(message, history, model, temperature, num_calls):
         history[-1] = (message, f"An unexpected error occurred: {str(e)}")
         yield history
 
-def retry_last_response(history, model, temperature, num_calls):
+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)
+    return chatbot_interface(last_user_msg, history, model, temperature, num_calls, use_embeddings, system_prompt)
 
-def respond(message, history, model, temperature, num_calls):
+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}")
 
     try:
-        for main_content, sources in get_response_with_search(message, model, num_calls=num_calls, temperature=temperature):
+        for main_content, sources in get_response_with_search(message, model, num_calls, temperature, use_embeddings, system_prompt):
             response = f"{main_content}\n\n{sources}"
             first_line = response.split('\n')[0] if response else ''
             yield response
@@ -100,21 +105,20 @@ def create_web_search_vectors(search_results):
     
     return FAISS.from_documents(documents, embed)
 
-import json
-
-def get_response_with_search(query, model, num_calls=3, temperature=0.2):
+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)
-    web_search_database = create_web_search_vectors(search_results)
-    
-    if not web_search_database:
-        yield "No web search results available. Please try again.", ""
-        return
-    
-    retriever = web_search_database.as_retriever(search_kwargs={"k": 5})
-    relevant_docs = retriever.get_relevant_documents(query)
-    
-    context = "\n".join([doc.page_content for doc in relevant_docs])
     
+    if use_embeddings:
+        web_search_database = create_web_search_vectors(search_results)
+        if not web_search_database:
+            yield "No web search results available. Please try again.", ""
+            return
+        retriever = web_search_database.as_retriever(search_kwargs={"k": 5})
+        relevant_docs = retriever.get_relevant_documents(query)
+        context = "\n".join([doc.page_content for doc in relevant_docs])
+    else:
+        context = "\n".join([f"{result['title']}\n{result['body']}\nSource: {result['href']}" for result in search_results])
+
     prompt = f"""Using the following context from web search results:
 {context}
 Write a detailed and complete research document that fulfills the following user request: '{query}'
@@ -136,7 +140,10 @@ After writing the document, please provide a list of sources used in your respon
     for i in range(num_calls):
         try:
             response = client.chat_completion(
-                messages=[{"role": "user", "content": prompt}],
+                messages=[
+                    {"role": "system", "content": system_prompt},
+                    {"role": "user", "content": prompt}
+                ],
                 max_tokens=max_new_tokens,
                 temperature=temperature,
                 stream=False,
@@ -181,55 +188,71 @@ 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, and number of API calls for fine-tuned responses.\n"
-                "3. For any queries, feel free to reach out @[email protected] or discord - shreyas094\n\n"
+                "2. You can adjust the model, temperature, number of API calls, and whether to use embeddings for fine-tuned responses.\n"
+                "3. You can also customize the system prompt to guide my behavior.\n"
+                "4. 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=[
-        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"),
-    ],
-    title="AI-powered Web Search Assistant",
-    description="Ask questions and get answers from web search results.",
-    theme=gr.themes.Soft(
-        primary_hue="orange",
-        secondary_hue="amber",
-        neutral_hue="gray",
-        font=[gr.themes.GoogleFont("Exo"), "ui-sans-serif", "system-ui", "sans-serif"]
-    ).set(
-        body_background_fill_dark="#0c0505",
-        block_background_fill_dark="#0c0505",
-        block_border_width="1px",
-        block_title_background_fill_dark="#1b0f0f",
-        input_background_fill_dark="#140b0b",
-        button_secondary_background_fill_dark="#140b0b",
-        border_color_accent_dark="#1b0f0f",
-        border_color_primary_dark="#1b0f0f",
-        background_fill_secondary_dark="#0c0505",
-        color_accent_soft_dark="transparent",
-        code_background_fill_dark="#140b0b"
-    ),
-    css=css,
-    examples=[
-        ["What are the latest developments in artificial intelligence?"],
-        ["Can you explain the basics of quantum computing?"],
-        ["What are the current global economic trends?"]
-    ],
-    cache_examples=False,
-    analytics_enabled=False,
-    textbox=gr.Textbox(placeholder="Ask a question", container=False, scale=7),
-    chatbot = gr.Chatbot(  
+with gr.Blocks(css=css, theme=gr.themes.Soft(
+    primary_hue="orange",
+    secondary_hue="amber",
+    neutral_hue="gray",
+    font=[gr.themes.GoogleFont("Exo"), "ui-sans-serif", "system-ui", "sans-serif"]
+).set(
+    body_background_fill_dark="#0c0505",
+    block_background_fill_dark="#0c0505",
+    block_border_width="1px",
+    block_title_background_fill_dark="#1b0f0f",
+    input_background_fill_dark="#140b0b",
+    button_secondary_background_fill_dark="#140b0b",
+    border_color_accent_dark="#1b0f0f",
+    border_color_primary_dark="#1b0f0f",
+    background_fill_secondary_dark="#0c0505",
+    color_accent_soft_dark="transparent",
+    code_background_fill_dark="#140b0b"
+)) as demo:
+    gr.Markdown("# AI-powered Web Search Assistant")
+    gr.Markdown("Ask questions and get answers from web search results.")
+
+    chatbot = gr.Chatbot(
         show_copy_button=True,
         likeable=True,
         layout="bubble",
         height=400,
         value=initial_conversation()
     )
-)
+
+    with gr.Row():
+        msg = gr.Textbox(placeholder="Ask a question", container=False, scale=7)
+        submit = gr.Button("Submit")
+
+    with gr.Accordion("Advanced Options", open=False):
+        model = gr.Dropdown(choices=MODELS, label="Select Model", value=MODELS[2])
+        temperature = gr.Slider(minimum=0.1, maximum=1.0, value=0.2, step=0.1, label="Temperature")
+        num_calls = gr.Slider(minimum=1, maximum=5, value=1, step=1, label="Number of API Calls")
+        use_embeddings = gr.Checkbox(label="Use Embeddings", value=True)
+        system_prompt = gr.Textbox(label="System Prompt", value=DEFAULT_SYSTEM_PROMPT, lines=5)
+
+    clear = gr.Button("Clear")
+    retry = gr.Button("Retry Last Response")
+
+    # Set up event handlers
+    submit.click(chatbot_interface, [msg, chatbot, model, temperature, num_calls, use_embeddings, system_prompt], chatbot)
+    msg.submit(chatbot_interface, [msg, chatbot, model, temperature, num_calls, use_embeddings, system_prompt], chatbot)
+    clear.click(lambda: None, None, chatbot, queue=False)
+    retry.click(retry_last_response, [chatbot, model, temperature, num_calls, use_embeddings, system_prompt], chatbot)
+
+    chatbot.like(vote, None, None)
+
+    gr.Examples(
+        examples=[
+            ["What are the latest developments in artificial intelligence?"],
+            ["Can you explain the basics of quantum computing?"],
+            ["What are the current global economic trends?"]
+        ],
+        inputs=msg
+    )
 
 if __name__ == "__main__":
     demo.launch(share=True)