core[minor]: Add support for DocumentIndex in the index api (#25100)

Support document index in the index api.
This commit is contained in:
Eugene Yurtsev
2024-08-06 15:30:49 -04:00
committed by GitHub
parent 264ab96980
commit d283f452cc
3 changed files with 183 additions and 39 deletions

View File

@@ -1386,7 +1386,14 @@ def test_indexing_custom_batch_size(
with patch.object(vector_store, "add_documents") as mock_add_documents:
index(docs, record_manager, vector_store, batch_size=batch_size)
args, kwargs = mock_add_documents.call_args
assert args == (docs,)
docs_with_id = [
Document(
page_content="This is a test document.",
metadata={"source": "1"},
id=ids[0],
)
]
assert args == (docs_with_id,)
assert kwargs == {"ids": ids, "batch_size": batch_size}
@@ -1407,5 +1414,12 @@ async def test_aindexing_custom_batch_size(
with patch.object(vector_store, "aadd_documents") as mock_add_documents:
await aindex(docs, arecord_manager, vector_store, batch_size=batch_size)
args, kwargs = mock_add_documents.call_args
assert args == (docs,)
docs_with_id = [
Document(
page_content="This is a test document.",
metadata={"source": "1"},
id=ids[0],
)
]
assert args == (docs_with_id,)
assert kwargs == {"ids": ids, "batch_size": batch_size}