[VectorStore] Improvement: Improve chroma vector store (#28524)

- Complete unit test
- Fix spelling error
This commit is contained in:
ZhangShenao 2024-12-06 00:58:32 +08:00 committed by GitHub
parent 7d44316d92
commit d26555c682
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
2 changed files with 16 additions and 1 deletions

View File

@ -758,7 +758,7 @@ class Chroma(VectorStore):
The most similar documents will have the lowest relevance score. Default The most similar documents will have the lowest relevance score. Default
relevance score function is euclidean distance. Distance metric must be relevance score function is euclidean distance. Distance metric must be
provided in `collection_metadata` during initizalition of Chroma object. provided in `collection_metadata` during initialization of Chroma object.
Example: collection_metadata={"hnsw:space": "cosine"}. Available distance Example: collection_metadata={"hnsw:space": "cosine"}. Available distance
metrics are: 'cosine', 'l2' and 'ip'. metrics are: 'cosine', 'l2' and 'ip'.

View File

@ -13,3 +13,18 @@ def test_initialization() -> None:
texts=texts, texts=texts,
embedding=FakeEmbeddings(size=10), embedding=FakeEmbeddings(size=10),
) )
def test_similarity_search() -> None:
"""Test similarity search by Chroma."""
texts = ["foo", "bar", "baz"]
metadatas = [{"page": str(i)} for i in range(len(texts))]
docsearch = Chroma.from_texts(
collection_name="test_collection",
texts=texts,
embedding=FakeEmbeddings(size=10),
metadatas=metadatas,
)
output = docsearch.similarity_search("foo", k=1)
docsearch.delete_collection()
assert len(output) == 1