feat: document search through LLM mock mode, UI polish

app.py

@@ -1,10 +1,14 @@
 import time
+import uuid
 import pandas as pd
 import gradio as gr
-from llama_index import GPTSimpleVectorIndex
+from llama_index import GPTSimpleVectorIndex, MockLLMPredictor, ServiceContext
 
 title = "Confidential forensics tool with ChatGPT"
-examples = ["Who is Phillip Allen?", "What is the project in Austin is about?", "Give me more details about the real estate project"]
+examples = ["Who is Phillip Allen?", "What the project in Austin is about?", "Give me more details about the real estate project"]
+
+llm_predictor = MockLLMPredictor()
+service_context_mock = ServiceContext.from_defaults(llm_predictor=llm_predictor)
 
 index = GPTSimpleVectorIndex.load_from_disk('email.json')
 docs_arr = []
@@ -43,9 +47,27 @@ def respond(message, chat_history):
     chat_history.append((message, bot_message))
     return "", chat_history
 
+def find_doc(opt, msg2):
+    message = ""
+    if len(msg2.strip()) < 1:
+        message = "Oops, it looks like your query was not valid. Please make sure you typed something in your text box and then try again."
+    else:
+        try:
+            resp = index.query(msg2, service_context=service_context_mock)
+            for key, item in resp.extra_info.items():
+                message += f"Document: {key}\nExtra details:\n"
+                for sub_key, sub_item in item.items():
+                    message += f"---- {sub_key}: {sub_item}"
+
+        except Exception as e:
+            message = "An error occured when handling your query, please try again."
+            print(e)
+    return message, ""
+
 with gr.Blocks(title=title) as demo:
     gr.Markdown(
     """
+
     # """ + title + """
     ...
     """)
@@ -72,15 +94,34 @@ with gr.Blocks(title=title) as demo:
         """
         Example of queries
         """)
-        i = 0
         for ex in examples:
-            btn = gr.Button(examples[i])
+            btn = gr.Button(ex)
             btn.click(respond2, [btn, chatbot, msg], [btn, chatbot, msg])
-            i += 1
 
     msg.submit(respond, [msg, chatbot], [msg, chatbot])
     btn_send.click(respond, [msg, chatbot], [msg, chatbot])
     btn_upload.upload(respond_upload, [btn_upload, msg, chatbot], [btn_upload, msg, chatbot])
+    gr.Markdown(
+    """
+
+    ## Search the matching document
+    """)
+    opt = gr.Textbox(
+        show_label=False,
+        placeholder="The document matching with your query will be shown here.",
+        interactive=False,
+        lines=8
+    )
+    with gr.Row():
+        with gr.Column(scale=0.85):
+            msg2 = gr.Textbox(
+                show_label=False,
+                placeholder="Enter text and press enter, or click on Send.",
+            ).style(container=False)
+        with gr.Column(scale=0.15, min_width=0):
+            btn_send2 = gr.Button("Send your query")
 
+    btn_send2.click(find_doc, [opt, msg2], [opt, msg2])
+    
 if __name__ == "__main__":
     demo.launch()