tests

buster/busterbot.py

@@ -155,7 +155,7 @@ class Buster:
         # Likely that the answer is meaningful, add the top sources
         return score < unk_threshold
 
-    def process_input(self, user_input: str) -> str:
+    def process_input(self, user_input: str) -> Response:
         """
         Main function to process the input question and generate a formatted output.
         """

buster/retriever/pickle.py

@@ -21,7 +21,7 @@ class PickleRetriever(Retriever):
             # Drop the `current` column
             documents.drop(columns=["current"], inplace=True)
 
-        if source is not None and "source" in documents.columns:
+        if source not in [None, ""] and "source" in documents.columns:
             documents = documents[documents.source == source]
 
         return documents

tests/test_chatbot.py

@@ -157,7 +157,7 @@ def test_chatbot_real_data__chatGPT_OOD():
     buster = Buster(cfg=buster_cfg, retriever=retriever)
     response = buster.process_input("What is a good recipe for brocolli soup?")
     assert isinstance(response.completion.text, str)
-    assert response.is_relevant is False
+    assert response.is_relevant == False
 
 
 def test_chatbot_real_data__GPT():
@@ -191,4 +191,4 @@ def test_chatbot_real_data__GPT():
     buster = Buster(cfg=hf_transformers_cfg, retriever=retriever)
     response = buster.process_input("What is a transformer?")
     assert isinstance(response.completion.text, str)
-    assert response.is_relevant is True
+    assert response.is_relevant == True