diff --git a/wren-ai-service/src/pipelines/ask/historical_question.py b/wren-ai-service/src/pipelines/ask/historical_question.py index 4af64261b..e51197119 100644 --- a/wren-ai-service/src/pipelines/ask/historical_question.py +++ b/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) diff --git a/wren-ai-service/src/pipelines/indexing/indexing.py b/wren-ai-service/src/pipelines/indexing/indexing.py index 12820260f..45c8182d5 100644 --- a/wren-ai-service/src/pipelines/indexing/indexing.py +++ b/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, diff --git a/wren-ai-service/tests/pytest/pipelines/test_view_converter.py b/wren-ai-service/tests/pytest/pipelines/test_view_converter.py index 010eb8d99..2a9d591de 100644 --- a/wren-ai-service/tests/pytest/pipelines/test_view_converter.py +++ b/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?"