add excluded_llm_metadata_keys

This commit is contained in:
seehi 2024-03-07 15:35:28 +08:00
parent f149007752
commit a3b2cf7f0b
4 changed files with 44 additions and 21 deletions

View file

@ -1,6 +1,6 @@
import pytest
from llama_index.core import VectorStoreIndex
from llama_index.core.schema import TextNode
from llama_index.core.schema import NodeWithScore, TextNode
from metagpt.rag.engines import SimpleEngine
from metagpt.rag.retrievers.base import ModifiableRAGRetriever
@ -97,7 +97,8 @@ class TestSimpleEngine:
mock_super_aretrieve = mocker.patch(
"metagpt.rag.engines.simple.RetrieverQueryEngine.aretrieve", new_callable=mocker.AsyncMock
)
mock_super_aretrieve.return_value = ["node_with_score"]
nodes = [NodeWithScore(node=TextNode())]
mock_super_aretrieve.return_value = nodes
# Setup
engine = SimpleEngine(retriever=mocker.MagicMock())
@ -109,7 +110,7 @@ class TestSimpleEngine:
# Assertions
mock_query_bundle.assert_called_once_with(test_query)
mock_super_aretrieve.assert_called_once_with("query_bundle")
assert result == ["node_with_score"]
assert result == nodes
def test_add_docs(self, mocker):
# Mock
@ -157,4 +158,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" in node.metadata
assert "obj_dict" in node.metadata