fix completion change (#96)

* fix completion change * tests
jerpint · May 10, 2023 · b31fac3 · b31fac3
1 parent e5f6c22
commit b31fac3
Show file tree

Hide file tree

Showing 3 changed files with 3 additions and 3 deletions.
diff --git a/buster/busterbot.py b/buster/busterbot.py
@@ -70,7 +70,7 @@ def from_dict(cls, answer_dict: dict):
         if isinstance(answer_dict["matched_documents"], str):
             answer_dict["matched_documents"] = pd.read_json(answer_dict["matched_documents"], orient="index")
         elif isinstance(answer_dict["matched_documents"], dict):
-            answer_dict["matched_documents"] = pd.DataFrame(answer_dict["matched_documents"])
+            answer_dict["matched_documents"] = pd.DataFrame(answer_dict["matched_documents"]).T
         else:
             raise ValueError(f"Unknown type for matched_documents: {type(answer_dict['matched_documents'])}")
         answer_dict["completion"] = Completion.from_dict(answer_dict["completion"])

diff --git a/buster/completers/base.py b/buster/completers/base.py
@@ -93,7 +93,7 @@ def generate_response(self, system_prompt, user_input):
 
         completor = self.complete(prompt=prompt, **self.completion_kwargs)
 
-        self.completion = Completion(completor, self.error)
+        self.completion = Completion(completor=completor, error=self.error)
 
         return self.completion
 

diff --git a/tests/test_chatbot.py b/tests/test_chatbot.py
@@ -28,7 +28,7 @@ def complete(self):
         return
 
     def generate_response(self, user_input, system_prompt) -> Completion:
-        return Completion(self.expected_answer, error=False)
+        return Completion(completor=self.expected_answer, error=False)
 
 
 class MockRetriever(Retriever):