quotation marks -> code block

biocypher · Aug 21, 2024 · a3d01c0 · a3d01c0
1 parent 6a3c83f
commit a3d01c0
Show file tree

Hide file tree

Showing 2 changed files with 49 additions and 38 deletions.
diff --git a/biochatter/database_agent.py b/biochatter/database_agent.py
@@ -83,42 +83,43 @@ def _generate_query(self, query: str):
             return query, results
 
     def _build_response(
-            self, 
-            results: List[Dict], 
-            cypher_query: str, 
-            results_num: Optional[int]=3
-        ) -> List[Document]:
+        self,
+        results: List[Dict],
+        cypher_query: str,
+        results_num: Optional[int] = 3,
+    ) -> List[Document]:
         if len(results) == 0:
-            return [Document(
+            return [
+                Document(
+                    page_content=(
+                        "I didn't find any result in knowledge graph, "
+                        f"but here is the query I used: {cypher_query}. "
+                        "You can ask user to refine the question. "
+                        "Note: please ensure to include the query in a code "
+                        "block in your response so that the user can refine "
+                        "their question effectively."
+                    ),
+                    metadata={"cypher_query": cypher_query},
+                )
+            ]
+
+        clipped_results = results[:results_num] if results_num > 0 else results
+        results_dump = json.dumps(clipped_results)
+
+        return [
+            Document(
                 page_content=(
-                    "I didn't find any result in knowledge graph, "
-                    f"but here is the query I used: {cypher_query}. "
-                    "You can ask user to refine the question. "
-                    "Note: please ensure to include the query in quotation "
-                    "marks in your response so that the user can refine "
+                    "The results retrieved from knowledge graph are: "
+                    f"{results_dump}. "
+                    f"The query used is: {cypher_query}. "
+                    "Note: please ensure to include the query in a code block "
+                    "in your response so that the user can refine "
                     "their question effectively."
                 ),
-                metadata = {
-                    "cypher_query": cypher_query
-                },
-            )]
-
-        clipped_results = results[:results_num] if results_num > 0 else results
-        results_dump = json.dumps(clipped_results)
+                metadata={"cypher_query": cypher_query},
+            )
+        ]
 
-        return [Document(
-            page_content=(
-                "The results retrieved from knowledge graph are: "
-                f"{results_dump}. "
-                f"The query used is: {cypher_query}. "
-                "Note: please ensure to include the query in quotation "
-                "marks in your response so that the user can refine "
-                "their question effectively."
-            ),
-            metadata = {
-                "cypher_query": cypher_query
-            },
-        )]
     def get_query_results(self, query: str, k: int = 3) -> list[Document]:
         """
         Generate a query using the prompt engine and return the results.
@@ -155,7 +156,7 @@ def get_query_results(self, query: str, k: int = 3) -> list[Document]:
         return self._build_response(
             results=results[0], cypher_query=cypher_query, results_num=k
         )
-        
+
     def get_description(self):
         result = self.driver.query("MATCH (n:Schema_info) RETURN n LIMIT 1")
 

diff --git a/test/test_database_agent.py b/test/test_database_agent.py
@@ -44,9 +44,14 @@ def test_get_query_results_with_reflexion():
             result = db_agent.get_query_results("test_query", 3)
 
     # Check if the result is as expected
-    expected_result = db_agent._build_response([
-        {"key": "value"},{"key": "value"},{"key": "value"},
-    ], "test_query")
+    expected_result = db_agent._build_response(
+        [
+            {"key": "value"},
+            {"key": "value"},
+            {"key": "value"},
+        ],
+        "test_query",
+    )
     assert result == expected_result
 
 
@@ -87,7 +92,12 @@ def test_get_query_results_without_reflexion():
             result = db_agent.get_query_results("test_query", 3)
 
     # Check if the result is as expected
-    expected_result = db_agent._build_response([
-        {"key": "value"},{"key": "value"},{"key": "value"},
-    ], "test_query")
+    expected_result = db_agent._build_response(
+        [
+            {"key": "value"},
+            {"key": "value"},
+            {"key": "value"},
+        ],
+        "test_query",
+    )
     assert result == expected_result