updated main app

scripts/app.py

@@ -1,5 +1,13 @@
+import os
+from pathlib import Path
+
 import gradio as gr
+import nest_asyncio
+import yaml
+from trulens.core import TruSession
+from trulens.dashboard import run_dashboard
 
+from src.mythesis_chatbot.evaluation import get_prebuilt_trulens_recorder
 from src.mythesis_chatbot.rag_setup import (
     SupportedRags,
     automerging_retrieval_setup,
@@ -7,68 +15,117 @@ from src.mythesis_chatbot.rag_setup import (
     sentence_window_retrieval_setup,
 )
 
-input_file = "./data/Master_Thesis.pdf"
-save_dir = "./data/indices/"
-
-automerging_engine = automerging_retrieval_setup(
-    input_file=input_file,
-    save_dir=save_dir,
-    llm_openai_model="gpt-4o-mini",
-    embed_model="BAAI/bge-small-en-v1.5",
-    chunk_sizes=[2048, 512, 128],
-    similarity_top_k=6,
-    rerank_model="cross-encoder/ms-marco-MiniLM-L-2-v2",
-    rerank_top_n=2,
-)
+input_file_dir = Path(__file__).parents[1] / "data/"
+save_dir = Path(__file__).parents[1] / "data/indices/"
+config_dir = Path(__file__).parents[1] / "configs/"
+welcome_message_path = Path(__file__).parents[1] / "spaces/welcome_message.md"
 
-sentence_window_engine = sentence_window_retrieval_setup(
-    input_file=input_file,
-    save_dir=save_dir,
-    llm_openai_model="gpt-4o-mini",
-    embed_model="BAAI/bge-small-en-v1.5",
-    sentence_window_size=3,
-    similarity_top_k=6,
-    rerank_model="cross-encoder/ms-marco-MiniLM-L-2-v2",
-    rerank_top_n=2,
-)
+# Enables running async code inside an existing event loop without crashing.
+nest_asyncio.apply()
 
-basic_engine = basic_rag_setup(
-    input_file=input_file,
-    save_dir=save_dir,
-    llm_openai_model="gpt-4o-mini",
-    embed_model="BAAI/bge-small-en-v1.5",
-    similarity_top_k=6,
-    rerank_model="cross-encoder/ms-marco-MiniLM-L-2-v2",
-    rerank_top_n=2,
-)
+tru = TruSession(database_url=os.getenv("SUPABASE_CONNECTION_STRING"))
+run_dashboard(tru)
+
+
+class ChatBot:
+    def __init__(
+        self,
+        input_file_dir,
+        save_dir,
+        config_dir,
+    ):
+        self.recorder = None
+        self.previous_rag_mode = None
+        self.recorder = None
+
+        with open(os.path.join(config_dir, "basic.yaml")) as f:
+            self.basic_config = yaml.safe_load(f)
+        with open(os.path.join(config_dir, "auto_merging.yaml")) as f:
+            self.automerging_config = yaml.safe_load(f)
+        with open(os.path.join(config_dir, "sentence_window.yaml")) as f:
+            self.sentence_window_config = yaml.safe_load(f)
+
+        self.basic_engine = basic_rag_setup(
+            input_file=os.path.join(input_file_dir, self.basic_config["source_doc"]),
+            save_dir=save_dir,
+            **self.basic_config,
+        )
+        self.automerging_engine = automerging_retrieval_setup(
+            input_file=os.path.join(
+                input_file_dir, self.automerging_config["source_doc"]
+            ),
+            save_dir=save_dir,
+            **self.automerging_config,
+        )
+        self.sentence_window_engine = sentence_window_retrieval_setup(
+            input_file=os.path.join(
+                input_file_dir, self.sentence_window_config["source_doc"]
+            ),
+            save_dir=save_dir,
+            **self.sentence_window_config,
+        )
+
+    def __call__(self, query: str, rag_mode: SupportedRags):
 
+        match rag_mode:
+            case "classic retrieval":
 
-def chat_bot(query: str, rag_mode: SupportedRags) -> str:
-    if rag_mode == "basic":
-        return basic_engine.query(query).response
-    if rag_mode == "auto-merging retrieval":
-        return automerging_engine.query(query).response
-    if rag_mode == "sentence window retrieval":
-        return sentence_window_engine.query(query).response
+                if self.previous_rag_mode != rag_mode:
+                    self.previous_rag_mode = rag_mode
+                    self.recorder = get_prebuilt_trulens_recorder(
+                        self.basic_engine, self.basic_config
+                    )
 
+                with self.recorder as recording:  # noqa: F841
+                    response = self.basic_engine.query(query)
 
+            case "auto-merging retrieval":
+                if self.previous_rag_mode != rag_mode:
+                    self.previous_rag_mode = rag_mode
+                    self.recorder = get_prebuilt_trulens_recorder(
+                        self.automerging_engine, self.automerging_config
+                    )
+
+                with self.recorder as recording:  # noqa: F841
+                    response = self.automerging_engine.query(query)
+
+            case "sentence window retrieval":
+                if self.previous_rag_mode != rag_mode:
+                    self.previous_rag_mode = rag_mode
+                    self.recorder = get_prebuilt_trulens_recorder(
+                        self.sentence_window_engine, self.sentence_window_config
+                    )
+
+                with self.recorder as recording:  # noqa: F841
+                    response = self.sentence_window_engine.query(query)
+
+        return response.response
+
+
+chat_bot = ChatBot(input_file_dir, save_dir, config_dir)
 default_message = (
-    "Ask a about a topic that is discussed in my master thesis."
-    " E.g., what is epistemic uncertainty?"
+    "Ask about a topic that is discussed in my master thesis."
+    " E.g., what is this master thesis about? Or what is epistemic uncertainty?"
 )
 
+with open(welcome_message_path, encoding="utf-8") as f:
+    description = f.read()
+
 gradio_app = gr.Interface(
     fn=chat_bot,
     inputs=[
-        gr.Textbox(placeholder=default_message),
+        gr.Textbox(placeholder=default_message, label="Query"),
         gr.Dropdown(
-            choices=["basic", "sentence window retrieval", "auto-merging retrieval"],
+            choices=SupportedRags.__args__,
             label="RAG mode",
-            value="basic",
+            value=SupportedRags.__args__[0],
         ),
     ],
-    outputs=["text"],
+    outputs=[
+        gr.Textbox(label="Answer"),
+    ],
+    title="RAG powered chatbot",
+    description=description,
 )
 
-if __name__ == "__main__":
-    gradio_app.launch()
+gradio_app.launch()