Send five source documents to the websocket but all of them to the me…

…trics endpoint.
nulib · Aug 26, 2024 · 9db1c1d · 9db1c1d
1 parent d23561b
commit 9db1c1d
Show file tree

Hide file tree

Showing 2 changed files with 21 additions and 44 deletions.
diff --git a/chat/src/helpers/response.py b/chat/src/helpers/response.py
@@ -21,18 +21,20 @@ def debug_response_passthrough(self):
     def original_question_passthrough(self):
         def get_and_send_original_question(docs):
             source_documents = []
-            for doc in docs["context"][:5]:
+            for doc in docs["context"]:
                 doc.metadata = {key: extract_prompt_value(doc.metadata.get(key)) for key in self.config.attributes if key in doc.metadata}
                 source_document = doc.metadata.copy()
                 source_document["content"] = doc.page_content
                 source_documents.append(source_document)
-                
+
             original_question = {
                 "question": self.config.question,
-                "source_documents": source_documents,
+                "source_documents": source_documents[:5]
             }
             self.config.socket.send(original_question)
             self.original_question = original_question
+
+            docs["source_documents"] = source_documents
             return docs
 
         return RunnablePassthrough(get_and_send_original_question)

diff --git a/chat/test/helpers/test_metrics.py b/chat/test/helpers/test_metrics.py
@@ -23,38 +23,7 @@ def setUp(self):
         self.question = "What is your name?"
         self.original_question = {
             "question": self.question,
-            "source_documents": [
-                {
-                    "accession_number": "SourceDoc:1",
-                    "api_link": "https://api.dc.library.northwestern.edu/api/v2/works/881e8cae-67be-4e04-9970-7eafb52b2c5c",
-                    "canonical_link": "https://dc.library.northwestern.edu/items/881e8cae-67be-4e04-9970-7eafb52b2c5c",
-                    "title": "Source Document One!"
-                },
-                {
-                    "accession_number": "SourceDoc:2",
-                    "api_link": "https://api.dc.library.northwestern.edu/api/v2/works/ac0b2a0d-8f80-420a-b1a1-63b6ac2299f1",
-                    "canonical_link": "https://dc.library.northwestern.edu/items/ac0b2a0d-8f80-420a-b1a1-63b6ac2299f1",
-                    "title": "Source Document Two!"
-                },
-                {
-                    "accession_number": "SourceDoc:3",
-                    "api_link": "https://api.dc.library.northwestern.edu/api/v2/works/11569bb5-1b89-4fa9-bdfb-2caf2ded5aa5",
-                    "canonical_link": "https://dc.library.northwestern.edu/items/11569bb5-1b89-4fa9-bdfb-2caf2ded5aa5",
-                    "title": "Source Document Three!"
-                },
-                {
-                    "accession_number": "SourceDoc:4",
-                    "api_link": "https://api.dc.library.northwestern.edu/api/v2/works/211eeeca-d56e-4c6e-9123-1612d72258f9",
-                    "canonical_link": "https://dc.library.northwestern.edu/items/211eeeca-d56e-4c6e-9123-1612d72258f9",
-                    "title": "Source Document Four!"
-                },
-                {
-                    "accession_number": "SourceDoc:5",
-                    "api_link": "https://api.dc.library.northwestern.edu/api/v2/works/10e45e7a-8011-4ac5-97df-efa6a5439d0e",
-                    "canonical_link": "https://dc.library.northwestern.edu/items/10e45e7a-8011-4ac5-97df-efa6a5439d0e",
-                    "title": "Source Document Five!"
-                }
-            ],
+            "source_documents": self.generate_source_documents(20),
         }
         self.event = {
             "body": json.dumps({
@@ -75,6 +44,17 @@ def setUp(self):
         self.response = {
             "output_text": "This is a test response.",
         }
+
+    def generate_source_documents(self, count):
+        return [
+            {
+                "accession_number": f"SourceDoc:{i+1}",
+                "api_link": f"https://api.dc.library.northwestern.edu/api/v2/works/{i+1:0>32}",
+                "canonical_link": f"https://dc.library.northwestern.edu/items/{i+1:0>32}",
+                "title": f"Source Document {i+1}!"
+            }
+            for i in range(count)
+        ]
 
     def test_debug_response(self):
         result = debug_response(self.config, self.response, self.original_question)
@@ -83,15 +63,10 @@ def test_debug_response(self):
         self.assertEqual(result["question"], self.question)
         self.assertEqual(result["ref"], "test")
         self.assertEqual(result["size"], 20)
+        self.assertEqual(len(result["source_documents"]), 20)
         self.assertEqual(
             result["source_documents"],
-            [
-                "https://api.dc.library.northwestern.edu/api/v2/works/881e8cae-67be-4e04-9970-7eafb52b2c5c",
-                "https://api.dc.library.northwestern.edu/api/v2/works/ac0b2a0d-8f80-420a-b1a1-63b6ac2299f1",
-                "https://api.dc.library.northwestern.edu/api/v2/works/11569bb5-1b89-4fa9-bdfb-2caf2ded5aa5",
-                "https://api.dc.library.northwestern.edu/api/v2/works/211eeeca-d56e-4c6e-9123-1612d72258f9",
-                "https://api.dc.library.northwestern.edu/api/v2/works/10e45e7a-8011-4ac5-97df-efa6a5439d0e"
-            ]
+            [doc["api_link"] for doc in self.original_question["source_documents"]]
         )
 
     def test_token_usage(self):
@@ -101,8 +76,8 @@ def test_token_usage(self):
             "answer": 12,
             "prompt": 329,
             "question": 5,
-            "source_documents": 527,
-            "total": 873
+            "source_documents": 1602,
+            "total": 1948
         }
 
         self.assertEqual(result, expected_result)