fixing test

NREL · Dec 20, 2024 · 4ca2970 · 4ca2970
1 parent 4018c59
commit 4ca2970
Show file tree

Hide file tree

Showing 2 changed files with 5 additions and 2 deletions.
diff --git a/elm/wizard.py b/elm/wizard.py
@@ -91,6 +91,9 @@ def engineer_query(self, query,
             returned here
         vector_query_time : float
             measures vector database query time
+        used_index : list
+            Shows the indices of the documents used in making a query to the 
+            vector database
         """
 
         self.messages.append({"role": "user", "content": query})
@@ -228,7 +231,7 @@ def chat(self, query,
         else:
             response_message = response.choices[0].message.content
         finish_completion_time = perf_counter()
-        chat_completion_time = start_completion_time - finish_completion_time
+        chat_completion_time = finish_completion_time - start_completion_time
 
         self.messages.append({'role': 'assistant',
                               'content': response_message})

diff --git a/tests/test_wizard.py b/tests/test_wizard.py
@@ -88,7 +88,7 @@ def test_chunk_and_embed(mocker):
     assert query.endswith(question)
     assert 'source0' in references
     assert isinstance(performance, dict)
-
+    assert performance["chat_completion_time"] > 0
 
 def test_convo_query(mocker):
     """Query with multiple messages