chore(wren-ai-service): fix view candidate indexing and retrieval (#635)

* fix conflict

* fix test

wren-ai-service/src/pipelines/ask/historical_question.py

@@ -1,4 +1,3 @@
-import ast
 import logging
 import sys
 from pathlib import Path
@@ -26,10 +25,12 @@ class ScoreFilter:
     @component.output_types(
         documents=List[Document],
     )
-    def run(self, documents: List[Document], score: float = 0.8):
+    def run(self, documents: List[Document], score: float = 0.9):
         return {
-            "documents": list(
-                filter(lambda document: document.score >= score, documents)
+            "documents": sorted(
+                filter(lambda document: document.score >= score, documents),
+                key=lambda document: document.score,
+                reverse=True,
             )
         }
 
@@ -44,12 +45,11 @@ class OutputFormatter:
         logger.debug(f"historical_question_output_formatter: {documents}")
 
         for doc in documents:
-            content = ast.literal_eval(doc.content)
             formatted = {
-                "question": content.get("question"),
-                "summary": content.get("summary"),
-                "statement": content.get("statement"),
-                "viewId": content.get("viewId"),
+                "question": doc.content,
+                "summary": doc.meta.get("summary"),
+                "statement": doc.meta.get("statement"),
+                "viewId": doc.meta.get("viewId"),
             }
             list.append(formatted)
 

wren-ai-service/src/pipelines/indexing/indexing.py

@@ -98,25 +98,31 @@ class ViewConverter:
 
     @component.output_types(documents=List[Document])
     def run(self, mdl: Dict[str, Any], id: Optional[str] = None) -> None:
-        def _format(view: Dict[str, Any]) -> List[str]:
+        def _get_content(view: Dict[str, Any]) -> str:
             properties = view.get("properties", {})
-            return str(
-                {
-                    "question": properties.get("question", ""),
-                    "summary": properties.get("summary", ""),
-                    "statement": view.get("statement", ""),
-                    "viewId": properties.get("viewId", ""),
-                }
-            )
+            return str(properties.get("question", ""))
 
-        converted_views = [_format(view) for view in mdl["views"]]
+        def _get_meta(view: Dict[str, Any]) -> Dict[str, Any]:
+            properties = view.get("properties", {})
+            return {
+                "summary": properties.get("summary", ""),
+                "statement": view.get("statement", ""),
+                "viewId": properties.get("viewId", ""),
+            }
+
+        converted_views = [
+            {"content": _get_content(view), "meta": _get_meta(view)}
+            for view in mdl["views"]
+        ]
 
         return {
             "documents": [
                 Document(
                     id=str(uuid.uuid4()),
-                    meta={"project_id": id} if id else {},
-                    content=converted_view,
+                    meta={"project_id": id, **converted_view["meta"]}
+                    if id
+                    else {**converted_view["meta"]},
+                    content=converted_view["content"],
                 )
                 for converted_view in tqdm(
                     converted_views,

wren-ai-service/tests/pytest/pipelines/test_view_converter.py

@@ -31,14 +31,12 @@ def test_single_view():
     assert len(actual["documents"]) == 1
 
     document: Document = actual["documents"][0]
-    assert document.content == str(
-        {
-            "question": "How many books are there?",
-            "summary": "Retrieve the number of books",
-            "statement": "SELECT * FROM book",
-            "viewId": "fake-id-1",
-        }
-    )
+    assert document.meta == {
+        "summary": "Retrieve the number of books",
+        "statement": "SELECT * FROM book",
+        "viewId": "fake-id-1",
+    }
+    assert document.content == "How many books are there?"
 
 
 def test_view_missing_properties():
@@ -56,14 +54,12 @@ def test_view_missing_properties():
     assert len(actual["documents"]) == 1
 
     document: Document = actual["documents"][0]
-    assert document.content == str(
-        {
-            "question": "",
-            "summary": "",
-            "statement": "SELECT * FROM book",
-            "viewId": "",
-        }
-    )
+    assert document.meta == {
+        "summary": "",
+        "statement": "SELECT * FROM book",
+        "viewId": "",
+    }
+    assert document.content == ""
 
 
 def test_view_missing_question():
@@ -85,14 +81,12 @@ def test_view_missing_question():
     assert len(actual["documents"]) == 1
 
     document: Document = actual["documents"][0]
-    assert document.content == str(
-        {
-            "question": "",
-            "summary": "Retrieve the number of books",
-            "statement": "SELECT * FROM book",
-            "viewId": "fake-id-1",
-        }
-    )
+    assert document.meta == {
+        "summary": "Retrieve the number of books",
+        "statement": "SELECT * FROM book",
+        "viewId": "fake-id-1",
+    }
+    assert document.content == ""
 
 
 def test_view_missing_summary():
@@ -114,14 +108,12 @@ def test_view_missing_summary():
     assert len(actual["documents"]) == 1
 
     document: Document = actual["documents"][0]
-    assert document.content == str(
-        {
-            "question": "How many books are there?",
-            "summary": "",
-            "statement": "SELECT * FROM book",
-            "viewId": "fake-id-1",
-        }
-    )
+    assert document.meta == {
+        "summary": "",
+        "statement": "SELECT * FROM book",
+        "viewId": "fake-id-1",
+    }
+    assert document.content == "How many books are there?"
 
 
 def test_view_missing_id():
@@ -143,14 +135,12 @@ def test_view_missing_id():
     assert len(actual["documents"]) == 1
 
     document: Document = actual["documents"][0]
-    assert document.content == str(
-        {
-            "question": "How many books are there?",
-            "summary": "Retrieve the number of books",
-            "statement": "SELECT * FROM book",
-            "viewId": "",
-        }
-    )
+    assert document.meta == {
+        "summary": "Retrieve the number of books",
+        "statement": "SELECT * FROM book",
+        "viewId": "",
+    }
+    assert document.content == "How many books are there?"
 
 
 def test_multi_views():
@@ -182,21 +172,17 @@ def test_multi_views():
     assert len(actual["documents"]) == 2
 
     document_1: Document = actual["documents"][0]
-    assert document_1.content == str(
-        {
-            "question": "How many books are there?",
-            "summary": "Retrieve the number of books",
-            "statement": "SELECT * FROM book",
-            "viewId": "fake-id-1",
-        }
-    )
+    assert document_1.meta == {
+        "summary": "Retrieve the number of books",
+        "statement": "SELECT * FROM book",
+        "viewId": "fake-id-1",
+    }
+    assert document_1.content == "How many books are there?"
 
     document_2: Document = actual["documents"][1]
-    assert document_2.content == str(
-        {
-            "question": "How many books are there?",
-            "summary": "Retrieve the number of books",
-            "statement": "SELECT * FROM book",
-            "viewId": "fake-id-2",
-        }
-    )
+    assert document_2.meta == {
+        "summary": "Retrieve the number of books",
+        "statement": "SELECT * FROM book",
+        "viewId": "fake-id-2",
+    }
+    assert document_2.content == "How many books are there?"