modified the chatui

.gitignore

@@ -0,0 +1,2 @@
+rag/__pycache__/rag_pipeline.cpython-311.pyc
+utils/__pycache__/prompts.cpython-311.pyc

app.py

@@ -1,8 +1,8 @@
 import gradio as gr
-import json
 from rag.rag_pipeline import RAGPipeline
 from utils.prompts import highlight_prompt, evidence_based_prompt, sample_questions
 from config import STUDY_FILES
+import json
 
 # Cache for RAG pipelines
 rag_cache = {}
@@ -20,14 +20,11 @@ def get_rag_pipeline(study_name):
 
 def chat_function(message, history, study_name, prompt_type):
     rag = get_rag_pipeline(study_name)
-
-    if prompt_type == "Highlight":
-        prompt = highlight_prompt
-    elif prompt_type == "Evidence-based":
-        prompt = evidence_based_prompt
-    else:
-        prompt = None
-
+    prompt = (
+        highlight_prompt
+        if prompt_type == "Highlight"
+        else evidence_based_prompt if prompt_type == "Evidence-based" else None
+    )
     response = rag.query(message, prompt_template=prompt)
     return response.response
 
@@ -42,26 +39,33 @@ def get_study_info(study_name):
         return "Invalid study name"
 
 
-with gr.Blocks() as demo:
+with gr.Blocks(css="#chatbot {height: 600px; overflow-y: auto;}") as demo:
     gr.Markdown("# RAG Pipeline Demo")
 
     with gr.Row():
-        study_dropdown = gr.Dropdown(
-            choices=list(STUDY_FILES.keys()),
-            label="Select Study",
-            value=list(STUDY_FILES.keys())[0],
-        )
-        study_info = gr.Markdown()
-
-    prompt_type = gr.Radio(
-        ["Default", "Highlight", "Evidence-based"],
-        label="Prompt Type",
-        value="Default",
-    )
+        with gr.Column(scale=3):
+            chatbot = gr.Chatbot(elem_id="chatbot")
+            with gr.Row():
+                msg = gr.Textbox(
+                    show_label=False, placeholder="Enter your message here...", scale=4
+                )
+                send_btn = gr.Button("Send", scale=1)
+
+        with gr.Column(scale=1):
+            study_dropdown = gr.Dropdown(
+                choices=list(STUDY_FILES.keys()),
+                label="Select Study",
+                value=list(STUDY_FILES.keys())[0],
+            )
+            study_info = gr.Markdown()
+            prompt_type = gr.Radio(
+                ["Default", "Highlight", "Evidence-based"],
+                label="Prompt Type",
+                value="Default",
+            )
+            clear = gr.Button("Clear Chat")
 
-    chatbot = gr.Chatbot()
-    msg = gr.Textbox()
-    clear = gr.Button("Clear")
+    gr.Examples(examples=sample_questions[list(STUDY_FILES.keys())[0]], inputs=msg)
 
     def user(user_message, history):
         return "", history + [[user_message, None]]
@@ -75,15 +79,16 @@ with gr.Blocks() as demo:
     msg.submit(user, [msg, chatbot], [msg, chatbot], queue=False).then(
         bot, [chatbot, study_dropdown, prompt_type], chatbot
     )
+    send_btn.click(user, [msg, chatbot], [msg, chatbot], queue=False).then(
+        bot, [chatbot, study_dropdown, prompt_type], chatbot
+    )
     clear.click(lambda: None, None, chatbot, queue=False)
 
     study_dropdown.change(
         fn=get_study_info,
         inputs=study_dropdown,
         outputs=study_info,
-    ).then(lambda: None, None, chatbot, queue=False)
-
-    gr.Examples(examples=sample_questions[list(STUDY_FILES.keys())[0]], inputs=msg)
+    )
 
 if __name__ == "__main__":
     demo.launch(share=True, debug=True)

rag/rag_pipeline.py

@@ -70,11 +70,12 @@ class RAGPipeline:
                 "---------------------\n"
                 "{context_str}\n"
                 "---------------------\n"
-                "Given this information, please answer the question provided in the context. "
-                "Include all relevant information from the provided context. "
-                "If information comes from multiple sources, please mention all of them. "
-                "If the information is not available in the context, please state that clearly. "
-                "When quoting specific information, please use square brackets to indicate the source, e.g. [1], [2], etc."
+                "Given this information, please answer the question: {query_str}\n"
+                "Provide an answer to the question using evidence from the context above. "
+                "Cite sources using square brackets for EVERY piece of information, e.g. [1], [2], etc. "
+                "Even if there's only one source, still include the citation. "
+                "If you're unsure about a source, use [?]. "
+                "Ensure that EVERY statement from the context is properly cited."
             )
 
         # This is a hack to index all the documents in the store :)

utils/prompts.py

@@ -111,5 +111,8 @@ evidence_based_prompt = PromptTemplate(
     "---------------------\n"
     "Given this information, please answer the question: {query_str}\n"
     "Provide an answer to the question using evidence from the context above. "
-    "Cite sources using square brackets."
+    "Cite sources using square brackets for EVERY piece of information, e.g. [1], [2], etc. "
+    "Even if there's only one source, still include the citation. "
+    "If you're unsure about a source, use [?]. "
+    "Ensure that EVERY statement from the context is properly cited."
 )