wehub-resource-sync
da6a09ff09
Auto Bump and Release / auto-bump-and-release (push) Has been cancelled
style-check / pre-commit (push) Has been cancelled
unit-test / unit testing with python (push) Has been cancelled
unit-test / unit testing with python 3.10 (push) Has been cancelled
unit-test / unit testing with python 3.11 (push) Has been cancelled
369 lines
14 KiB
Python
369 lines
14 KiB
Python
import json
|
|
import os
|
|
|
|
import pytest
|
|
|
|
from kotaemon.base import DocumentWithEmbedding
|
|
from kotaemon.storages import (
|
|
ChromaVectorStore,
|
|
InMemoryVectorStore,
|
|
MilvusVectorStore,
|
|
QdrantVectorStore,
|
|
SimpleFileVectorStore,
|
|
)
|
|
|
|
|
|
class TestChromaVectorStore:
|
|
def test_add(self, tmp_path):
|
|
"""Test that the DB add correctly"""
|
|
db = ChromaVectorStore(path=str(tmp_path))
|
|
|
|
embeddings = [[0.1, 0.2, 0.3], [0.4, 0.5, 0.6]]
|
|
metadatas = [{"a": 1, "b": 2}, {"a": 3, "b": 4}]
|
|
ids = ["1", "2"]
|
|
|
|
assert db._collection.count() == 0, "Expected empty collection"
|
|
output = db.add(embeddings=embeddings, metadatas=metadatas, ids=ids)
|
|
assert output == ids, "Expected output to be the same as ids"
|
|
assert db._collection.count() == 2, "Expected 2 added entries"
|
|
|
|
def test_add_from_docs(self, tmp_path):
|
|
db = ChromaVectorStore(path=str(tmp_path))
|
|
|
|
embeddings = [[0.1, 0.2, 0.3], [0.4, 0.5, 0.6]]
|
|
metadatas = [{"a": 1, "b": 2}, {"a": 3, "b": 4}]
|
|
documents = [
|
|
DocumentWithEmbedding(embedding=embedding, metadata=metadata)
|
|
for embedding, metadata in zip(embeddings, metadatas)
|
|
]
|
|
assert db._collection.count() == 0, "Expected empty collection"
|
|
output = db.add(documents)
|
|
assert len(output) == 2, "Expected outputting 2 ids"
|
|
assert db._collection.count() == 2, "Expected 2 added entries"
|
|
|
|
def test_delete(self, tmp_path):
|
|
db = ChromaVectorStore(path=str(tmp_path))
|
|
|
|
embeddings = [[0.1, 0.2, 0.3], [0.4, 0.5, 0.6], [0.7, 0.8, 0.9]]
|
|
metadatas = [{"a": 1, "b": 2}, {"a": 3, "b": 4}, {"a": 5, "b": 6}]
|
|
ids = ["a", "b", "c"]
|
|
|
|
db.add(embeddings=embeddings, metadatas=metadatas, ids=ids)
|
|
assert db._collection.count() == 3, "Expected 3 added entries"
|
|
db.delete(ids=["a", "b"])
|
|
assert db._collection.count() == 1, "Expected 1 remaining entry"
|
|
db.delete(ids=["c"])
|
|
assert db._collection.count() == 0, "Expected 0 remaining entry"
|
|
|
|
def test_query(self, tmp_path):
|
|
db = ChromaVectorStore(path=str(tmp_path))
|
|
|
|
embeddings = [[0.1, 0.2, 0.3], [0.4, 0.5, 0.6], [0.7, 0.8, 0.9]]
|
|
metadatas = [{"a": 1, "b": 2}, {"a": 3, "b": 4}, {"a": 5, "b": 6}]
|
|
ids = ["a", "b", "c"]
|
|
|
|
db.add(embeddings=embeddings, metadatas=metadatas, ids=ids)
|
|
|
|
_, sim, out_ids = db.query(embedding=[0.1, 0.2, 0.3], top_k=1)
|
|
assert sim[0] - 1.0 < 1e-6
|
|
assert out_ids == ["a"]
|
|
|
|
_, _, out_ids = db.query(embedding=[0.42, 0.52, 0.53], top_k=1)
|
|
assert out_ids == ["b"]
|
|
|
|
def test_save_load_delete(self, tmp_path):
|
|
"""Test that save/load func behave correctly."""
|
|
embeddings = [[0.1, 0.2, 0.3], [0.4, 0.5, 0.6], [0.7, 0.8, 0.9]]
|
|
metadatas = [{"a": 1, "b": 2}, {"a": 3, "b": 4}, {"a": 5, "b": 6}]
|
|
ids = ["1", "2", "3"]
|
|
db = ChromaVectorStore(path=str(tmp_path))
|
|
db.add(embeddings=embeddings, metadatas=metadatas, ids=ids)
|
|
|
|
db2 = ChromaVectorStore(path=str(tmp_path))
|
|
assert (
|
|
db2._collection.count() == 3
|
|
), "load function does not load data completely"
|
|
|
|
# test delete collection function
|
|
db2.drop()
|
|
# reinit the chroma with the same collection name
|
|
db2 = ChromaVectorStore(path=str(tmp_path))
|
|
assert (
|
|
db2._collection.count() == 0
|
|
), "delete collection function does not work correctly"
|
|
|
|
|
|
class TestInMemoryVectorStore:
|
|
def test_add(self):
|
|
"""Test that add func adds correctly."""
|
|
|
|
embeddings = [[0.1, 0.2, 0.3], [0.4, 0.5, 0.6]]
|
|
metadatas = [{"a": 1, "b": 2}, {"a": 3, "b": 4}]
|
|
ids = ["1", "2"]
|
|
db = InMemoryVectorStore()
|
|
|
|
output = db.add(embeddings=embeddings, metadatas=metadatas, ids=ids)
|
|
assert output == ids, "Excepted output to be the same as ids"
|
|
|
|
def test_save_load_delete(self, tmp_path):
|
|
"""Test that delete func deletes correctly."""
|
|
embeddings = [[0.1, 0.2, 0.3], [0.4, 0.5, 0.6], [0.7, 0.8, 0.9]]
|
|
metadatas = [{"a": 1, "b": 2}, {"a": 3, "b": 4}, {"a": 5, "b": 6}]
|
|
ids = ["1", "2", "3"]
|
|
db = InMemoryVectorStore()
|
|
db.add(embeddings=embeddings, metadatas=metadatas, ids=ids)
|
|
db.delete(["3"])
|
|
db.save(save_path=tmp_path / "test_save_load_delete.json")
|
|
with open(tmp_path / "test_save_load_delete.json") as f:
|
|
data = json.load(f)
|
|
assert (
|
|
"1" and "2" in data["text_id_to_ref_doc_id"]
|
|
), "save function does not save data completely"
|
|
assert (
|
|
"3" not in data["text_id_to_ref_doc_id"]
|
|
), "delete function does not delete data completely"
|
|
db2 = InMemoryVectorStore()
|
|
db2.load(load_path=tmp_path / "test_save_load_delete.json")
|
|
assert db2.get("2") == [
|
|
0.4,
|
|
0.5,
|
|
0.6,
|
|
], "load function does not load data completely"
|
|
|
|
|
|
class TestSimpleFileVectorStore:
|
|
def test_add_delete(self, tmp_path):
|
|
"""Test that delete func deletes correctly."""
|
|
embeddings = [[0.1, 0.2, 0.3], [0.4, 0.5, 0.6], [0.7, 0.8, 0.9]]
|
|
metadatas = [{"a": 1, "b": 2}, {"a": 3, "b": 4}, {"a": 5, "b": 6}]
|
|
ids = ["1", "2", "3"]
|
|
collection_name = "test_save_load_delete"
|
|
db = SimpleFileVectorStore(path=tmp_path, collection_name=collection_name)
|
|
db.add(embeddings=embeddings, metadatas=metadatas, ids=ids)
|
|
db.delete(["3"])
|
|
with open(tmp_path / collection_name) as f:
|
|
data = json.load(f)
|
|
assert (
|
|
"1" and "2" in data["text_id_to_ref_doc_id"]
|
|
), "save function does not save data completely"
|
|
assert (
|
|
"3" not in data["text_id_to_ref_doc_id"]
|
|
), "delete function does not delete data completely"
|
|
db2 = SimpleFileVectorStore(path=tmp_path, collection_name=collection_name)
|
|
assert db2.get("2") == [
|
|
0.4,
|
|
0.5,
|
|
0.6,
|
|
], "load function does not load data completely"
|
|
|
|
os.remove(tmp_path / collection_name)
|
|
|
|
|
|
class TestMilvusVectorStore:
|
|
def test_add(self, tmp_path):
|
|
"""Test that the DB add correctly"""
|
|
db = MilvusVectorStore(
|
|
path=str(tmp_path),
|
|
overwrite=True,
|
|
)
|
|
|
|
embeddings = [[0.1, 0.2, 0.3], [0.4, 0.5, 0.6]]
|
|
metadatas = [{"a": 1, "b": 2}, {"a": 3, "b": 4}]
|
|
ids = ["1", "2"]
|
|
|
|
assert db.count() == 0, "Expected empty collection"
|
|
output = db.add(embeddings=embeddings, metadatas=metadatas, ids=ids)
|
|
assert output == ids, "Expected output to be the same as ids"
|
|
assert db.count() == 2, "Expected 2 added entries"
|
|
|
|
def test_add_from_docs(self, tmp_path):
|
|
db = MilvusVectorStore(
|
|
path=str(tmp_path),
|
|
overwrite=True,
|
|
)
|
|
|
|
embeddings = [[0.1, 0.2, 0.3], [0.4, 0.5, 0.6]]
|
|
metadatas = [{"a": 1, "b": 2}, {"a": 3, "b": 4}]
|
|
documents = [
|
|
DocumentWithEmbedding(embedding=embedding, metadata=metadata)
|
|
for embedding, metadata in zip(embeddings, metadatas)
|
|
]
|
|
assert db.count() == 0, "Expected empty collection"
|
|
output = db.add(documents)
|
|
assert len(output) == 2, "Expected outputting 2 ids"
|
|
assert db.count() == 2, "Expected 2 added entries"
|
|
|
|
def test_delete(self, tmp_path):
|
|
db = MilvusVectorStore(
|
|
path=str(tmp_path),
|
|
overwrite=True,
|
|
)
|
|
|
|
embeddings = [[0.1, 0.2, 0.3], [0.4, 0.5, 0.6], [0.7, 0.8, 0.9]]
|
|
metadatas = [{"a": 1, "b": 2}, {"a": 3, "b": 4}, {"a": 5, "b": 6}]
|
|
ids = ["a", "b", "c"]
|
|
|
|
db.add(embeddings=embeddings, metadatas=metadatas, ids=ids)
|
|
assert db.count() == 3, "Expected 3 added entries"
|
|
db.delete(ids=["a", "b"])
|
|
assert db.count() == 1, "Expected 1 remaining entry"
|
|
db.delete(ids=["c"])
|
|
assert db.count() == 0, "Expected 0 remaining entry"
|
|
|
|
def test_query(self, tmp_path):
|
|
db = MilvusVectorStore(path=str(tmp_path), overwrite=True)
|
|
import numpy as np
|
|
|
|
embeddings = np.array([[0.1, 0.2, 0.3], [0.4, 0.5, 0.6], [0.7, 0.8, 0.9]])
|
|
norms = np.linalg.norm(embeddings, axis=1)
|
|
normalized_embeddings = (embeddings / norms[:, np.newaxis]).tolist()
|
|
|
|
metadatas = [{"a": 1, "b": 2}, {"a": 3, "b": 4}, {"a": 5, "b": 6}]
|
|
ids = ["a", "b", "c"]
|
|
|
|
db.add(embeddings=normalized_embeddings, metadatas=metadatas, ids=ids)
|
|
|
|
_, sim, out_ids = db.query(embedding=normalized_embeddings[0], top_k=1)
|
|
assert sim[0] - 1.0 < 1e-6
|
|
assert out_ids == ["a"]
|
|
|
|
query_embedding = [
|
|
normalized_embeddings[1][0] + 0.02,
|
|
normalized_embeddings[1][1] + 0.02,
|
|
normalized_embeddings[1][2] + 0.02,
|
|
]
|
|
_, _, out_ids = db.query(embedding=query_embedding, top_k=1)
|
|
assert out_ids == ["b"]
|
|
|
|
def test_save_load_delete(self, tmp_path):
|
|
"""Test that save/load func behave correctly."""
|
|
embeddings = [[0.1, 0.2, 0.3], [0.4, 0.5, 0.6], [0.7, 0.8, 0.9]]
|
|
metadatas = [{"a": 1, "b": 2}, {"a": 3, "b": 4}, {"a": 5, "b": 6}]
|
|
ids = ["1", "2", "3"]
|
|
db = MilvusVectorStore(path=str(tmp_path), overwrite=True)
|
|
db.add(embeddings=embeddings, metadatas=metadatas, ids=ids)
|
|
|
|
db2 = MilvusVectorStore(path=str(tmp_path), overrides=False)
|
|
assert db2.count() == 3, "load function does not load data completely"
|
|
|
|
# test delete collection function
|
|
db2.drop()
|
|
# reinit the milvus with the same collection name
|
|
db2 = MilvusVectorStore(path=str(tmp_path), overwrite=False)
|
|
assert db2.count() == 0, "delete collection function does not work correctly"
|
|
|
|
|
|
class TestQdrantVectorStore:
|
|
def test_add(self):
|
|
from qdrant_client import QdrantClient
|
|
|
|
db = QdrantVectorStore(collection_name="test", client=QdrantClient(":memory:"))
|
|
|
|
embeddings = [[0.1, 0.2, 0.3], [0.4, 0.5, 0.6]]
|
|
metadatas = [{"a": 1, "b": 2}, {"a": 3, "b": 4}]
|
|
ids = [
|
|
"0f0611b3-2d9c-4818-ab69-1f1c4cf66693",
|
|
"90aba5d3-f4f8-47c6-bad9-5ea457442e07",
|
|
]
|
|
|
|
output = db.add(embeddings=embeddings, metadatas=metadatas, ids=ids)
|
|
assert output == ids, "Expected output to be the same as ids"
|
|
assert db.count() == 2, "Expected 2 added entries"
|
|
|
|
def test_add_from_docs(self, tmp_path):
|
|
from qdrant_client import QdrantClient
|
|
|
|
db = QdrantVectorStore(collection_name="test", client=QdrantClient(":memory:"))
|
|
|
|
embeddings = [[0.1, 0.2, 0.3], [0.4, 0.5, 0.6]]
|
|
metadatas = [{"a": 1, "b": 2}, {"a": 3, "b": 4}]
|
|
documents = [
|
|
DocumentWithEmbedding(embedding=embedding, metadata=metadata)
|
|
for embedding, metadata in zip(embeddings, metadatas)
|
|
]
|
|
|
|
output = db.add(documents)
|
|
assert len(output) == 2, "Expected outputting 2 ids"
|
|
assert db.count() == 2, "Expected 2 added entries"
|
|
|
|
def test_delete(self, tmp_path):
|
|
from qdrant_client import QdrantClient
|
|
|
|
db = QdrantVectorStore(collection_name="test", client=QdrantClient(":memory:"))
|
|
|
|
embeddings = [[0.1, 0.2, 0.3], [0.4, 0.5, 0.6], [0.7, 0.8, 0.9]]
|
|
metadatas = [{"a": 1, "b": 2}, {"a": 3, "b": 4}, {"a": 5, "b": 6}]
|
|
ids = [
|
|
"0f0611b3-2d9c-4818-ab69-1f1c4cf66693",
|
|
"90aba5d3-f4f8-47c6-bad9-5ea457442e07",
|
|
"6bed07c3-d284-47a3-a711-c3f9186755b8",
|
|
]
|
|
|
|
db.add(embeddings=embeddings, metadatas=metadatas, ids=ids)
|
|
assert db.count() == 3, "Expected 3 added entries"
|
|
db.delete(
|
|
ids=[
|
|
"0f0611b3-2d9c-4818-ab69-1f1c4cf66693",
|
|
"90aba5d3-f4f8-47c6-bad9-5ea457442e07",
|
|
]
|
|
)
|
|
assert db.count() == 1, "Expected 1 remaining entry"
|
|
db.delete(ids=["6bed07c3-d284-47a3-a711-c3f9186755b8"])
|
|
assert db.count() == 0, "Expected 0 remaining entry"
|
|
|
|
def test_query(self, tmp_path):
|
|
from qdrant_client import QdrantClient
|
|
|
|
db = QdrantVectorStore(collection_name="test", client=QdrantClient(":memory:"))
|
|
|
|
embeddings = [[0.1, 0.2, 0.3], [0.4, 0.5, 0.6], [0.7, 0.8, 0.9]]
|
|
metadatas = [{"a": 1, "b": 2}, {"a": 3, "b": 4}, {"a": 5, "b": 6}]
|
|
ids = [
|
|
"0f0611b3-2d9c-4818-ab69-1f1c4cf66693",
|
|
"90aba5d3-f4f8-47c6-bad9-5ea457442e07",
|
|
"6bed07c3-d284-47a3-a711-c3f9186755b8",
|
|
]
|
|
|
|
db.add(embeddings=embeddings, metadatas=metadatas, ids=ids)
|
|
|
|
_, sim, out_ids = db.query(embedding=[0.1, 0.2, 0.3], top_k=1)
|
|
assert sim[0] - 1.0 < 1e-6
|
|
assert out_ids == ["0f0611b3-2d9c-4818-ab69-1f1c4cf66693"]
|
|
|
|
_, _, out_ids = db.query(embedding=[0.4, 0.5, 0.6], top_k=1)
|
|
assert out_ids == ["90aba5d3-f4f8-47c6-bad9-5ea457442e07"]
|
|
|
|
def test_save_load_delete(self, tmp_path):
|
|
"""Test that save/load func behave correctly."""
|
|
embeddings = [[0.1, 0.2, 0.3], [0.4, 0.5, 0.6], [0.7, 0.8, 0.9]]
|
|
metadatas = [{"a": 1, "b": 2}, {"a": 3, "b": 4}, {"a": 5, "b": 6}]
|
|
ids = [
|
|
"0f0611b3-2d9c-4818-ab69-1f1c4cf66693",
|
|
"90aba5d3-f4f8-47c6-bad9-5ea457442e07",
|
|
"6bed07c3-d284-47a3-a711-c3f9186755b8",
|
|
]
|
|
from qdrant_client import QdrantClient
|
|
|
|
db = QdrantVectorStore(
|
|
collection_name="test", client=QdrantClient(path=tmp_path)
|
|
)
|
|
db.add(embeddings=embeddings, metadatas=metadatas, ids=ids)
|
|
del db
|
|
|
|
db2 = QdrantVectorStore(
|
|
collection_name="test", client=QdrantClient(path=tmp_path)
|
|
)
|
|
assert db2.count() == 3
|
|
|
|
db2.drop()
|
|
del db2
|
|
|
|
db2 = QdrantVectorStore(
|
|
collection_name="test", client=QdrantClient(path=tmp_path)
|
|
)
|
|
|
|
with pytest.raises(Exception):
|
|
# Since no docs were added, the collection should not exist yet
|
|
# and thus the count function should raise an exception
|
|
db2.count()
|