fix rag ut failed cases

This commit is contained in:
betterwang 2024-03-07 14:37:26 +08:00
parent 91c58d04cf
commit f45631e8d4
3 changed files with 5 additions and 5 deletions

1
.gitignore vendored
View file

@ -28,6 +28,7 @@ share/python-wheels/
*.egg
MANIFEST
metagpt/tools/schemas/
examples/data/search_kb/*.json
# PyInstaller
# Usually these files are written by a python scripts from a template

View file

@ -171,6 +171,6 @@ class SimpleEngine(RetrieverQueryEngine):
def _try_reconstruct_obj(nodes: list[NodeWithScore]):
"""If node is object, then dynamically reconstruct object, and save object to node.metadata["obj"]."""
for node in nodes:
if node.metadata.get("is_obj"):
if node.metadata.get("is_obj", False):
obj_cls = import_class(node.metadata["obj_cls_name"], node.metadata["obj_mod_name"])
node.metadata["obj"] = obj_cls(**node.metadata["obj_dict"])

View file

@ -97,8 +97,7 @@ class TestSimpleEngine:
mock_super_aretrieve = mocker.patch(
"metagpt.rag.engines.simple.RetrieverQueryEngine.aretrieve", new_callable=mocker.AsyncMock
)
nodes = [NodeWithScore(node=TextNode())]
mock_super_aretrieve.return_value = nodes
mock_super_aretrieve.return_value = [TextNode(text="node_with_score", metadata={"is_obj": False})]
# Setup
engine = SimpleEngine(retriever=mocker.MagicMock())
@ -110,7 +109,7 @@ class TestSimpleEngine:
# Assertions
mock_query_bundle.assert_called_once_with(test_query)
mock_super_aretrieve.assert_called_once_with("query_bundle")
assert result == nodes
assert result[0].text == "node_with_score"
def test_add_docs(self, mocker):
# Mock
@ -158,4 +157,4 @@ class TestSimpleEngine:
assert mock_retriever.add_nodes.call_count == 1
for node in mock_retriever.add_nodes.call_args[0][0]:
assert isinstance(node, TextNode)
assert "obj_dict" in node.metadata
assert "is_obj" in node.metadata