github/langchain
1
0
Fork 0
mirror of https://github.com/hwchase17/langchain.git synced 2025-07-12 15:59:56 +00:00
Code Issues Packages Projects Releases Wiki Activity

chroma[patch]: ruff fixes and rules (#31900)

Browse Source 
* bump ruff deps
* add more thorough ruff rules
* fix said rules
This commit is contained in:
Mason Daugherty 2025-07-07 21:45:19 -04:00 committed by GitHub
parent 2a7645300c
commit 38bd1abb8c
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
7 changed files with 195 additions and 114 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

138
libs/partners/chroma/langchain_chroma/vectorstores.py
View File

@ -79,11 +79,14 @@ def cosine_similarity(X: Matrix, Y: Matrix) -> np.ndarray:
X = np.array(X) X = np.array(X)
Y = np.array(Y) Y = np.array(Y)
if X.shape[1] != Y.shape[1]: if X.shape[1] != Y.shape[1]:
raise ValueError( msg = (
"Number of columns in X and Y must be the same. X has shape" "Number of columns in X and Y must be the same. X has shape"
f"{X.shape} " f"{X.shape} "
f"and Y has shape {Y.shape}." f"and Y has shape {Y.shape}."
) )
raise ValueError(
msg,
)
X_norm = np.linalg.norm(X, axis=1) X_norm = np.linalg.norm(X, axis=1)
Y_norm = np.linalg.norm(Y, axis=1) Y_norm = np.linalg.norm(Y, axis=1)
@ -285,7 +288,7 @@ class Chroma(VectorStore):
collection_metadata: Optional[dict] = None, collection_metadata: Optional[dict] = None,
client: Optional[chromadb.ClientAPI] = None, client: Optional[chromadb.ClientAPI] = None,
relevance_score_fn: Optional[Callable[[float], float]] = None, relevance_score_fn: Optional[Callable[[float], float]] = None,
create_collection_if_not_exists: Optional[bool] = True, create_collection_if_not_exists: Optional[bool] = True, # noqa: FBT002
) -> None: ) -> None:
"""Initialize with a Chroma client. """Initialize with a Chroma client.
@ -351,10 +354,13 @@ class Chroma(VectorStore):
def _collection(self) -> chromadb.Collection: def _collection(self) -> chromadb.Collection:
"""Returns the underlying Chroma collection or throws an exception.""" """Returns the underlying Chroma collection or throws an exception."""
if self._chroma_collection is None: if self._chroma_collection is None:
raise ValueError( msg = (
"Chroma collection not initialized. " "Chroma collection not initialized. "
"Use `reset_collection` to re-create and initialize the collection. " "Use `reset_collection` to re-create and initialize the collection. "
) )
raise ValueError(
msg,
)
return self._chroma_collection return self._chroma_collection
@property @property
@ -392,10 +398,10 @@ class Chroma(VectorStore):
""" """
return self._collection.query( return self._collection.query(
query_texts=query_texts, query_texts=query_texts,
query_embeddings=query_embeddings, # type: ignore query_embeddings=query_embeddings, # type: ignore[arg-type]
n_results=n_results, n_results=n_results,
where=where, # type: ignore where=where, # type: ignore[arg-type]
where_document=where_document, # type: ignore where_document=where_document, # type: ignore[arg-type]
**kwargs, **kwargs,
) )
@ -432,11 +438,12 @@ class Chroma(VectorStore):
if ids is None: if ids is None:
ids = [str(uuid.uuid4()) for _ in uris] ids = [str(uuid.uuid4()) for _ in uris]
else: else:
ids = [id if id is not None else str(uuid.uuid4()) for id in ids] ids = [id_ if id_ is not None else str(uuid.uuid4()) for id_ in ids]
embeddings = None embeddings = None
# Set embeddings # Set embeddings
if self._embedding_function is not None and hasattr( if self._embedding_function is not None and hasattr(
self._embedding_function, "embed_image" self._embedding_function,
"embed_image",
): ):
embeddings = self._embedding_function.embed_image(uris=uris) embeddings = self._embedding_function.embed_image(uris=uris)
if metadatas: if metadatas:
@ -461,8 +468,8 @@ class Chroma(VectorStore):
ids_with_metadata = [ids[idx] for idx in non_empty_ids] ids_with_metadata = [ids[idx] for idx in non_empty_ids]
try: try:
self._collection.upsert( self._collection.upsert(
metadatas=metadatas, # type: ignore metadatas=metadatas, # type: ignore[arg-type]
embeddings=embeddings_with_metadatas, # type: ignore embeddings=embeddings_with_metadatas, # type: ignore[arg-type]
documents=images_with_metadatas, documents=images_with_metadatas,
ids=ids_with_metadata, ids=ids_with_metadata,
) )
@ -473,7 +480,6 @@ class Chroma(VectorStore):
"langchain_community.vectorstores.utils.filter_complex_metadata." "langchain_community.vectorstores.utils.filter_complex_metadata."
) )
raise ValueError(e.args[0] + "\n\n" + msg) raise ValueError(e.args[0] + "\n\n" + msg)
else:
raise e raise e
if empty_ids: if empty_ids:
images_without_metadatas = [b64_texts[j] for j in empty_ids] images_without_metadatas = [b64_texts[j] for j in empty_ids]
@ -519,7 +525,7 @@ class Chroma(VectorStore):
if ids is None: if ids is None:
ids = [str(uuid.uuid4()) for _ in texts] ids = [str(uuid.uuid4()) for _ in texts]
else: else:
ids = [id if id is not None else str(uuid.uuid4()) for id in ids] ids = [id_ if id_ is not None else str(uuid.uuid4()) for id_ in ids]
embeddings = None embeddings = None
texts = list(texts) texts = list(texts)
@ -549,8 +555,8 @@ class Chroma(VectorStore):
ids_with_metadata = [ids[idx] for idx in non_empty_ids] ids_with_metadata = [ids[idx] for idx in non_empty_ids]
try: try:
self._collection.upsert( self._collection.upsert(
metadatas=metadatas, # type: ignore metadatas=metadatas, # type: ignore[arg-type]
embeddings=embeddings_with_metadatas, # type: ignore embeddings=embeddings_with_metadatas, # type: ignore[arg-type]
documents=texts_with_metadatas, documents=texts_with_metadatas,
ids=ids_with_metadata, ids=ids_with_metadata,
) )
@ -561,7 +567,6 @@ class Chroma(VectorStore):
"langchain_community.vectorstores.utils.filter_complex_metadata." "langchain_community.vectorstores.utils.filter_complex_metadata."
) )
raise ValueError(e.args[0] + "\n\n" + msg) raise ValueError(e.args[0] + "\n\n" + msg)
else:
raise e raise e
if empty_ids: if empty_ids:
texts_without_metadatas = [texts[j] for j in empty_ids] texts_without_metadatas = [texts[j] for j in empty_ids]
@ -570,13 +575,13 @@ class Chroma(VectorStore):
) )
ids_without_metadatas = [ids[j] for j in empty_ids] ids_without_metadatas = [ids[j] for j in empty_ids]
self._collection.upsert( self._collection.upsert(
embeddings=embeddings_without_metadatas, # type: ignore embeddings=embeddings_without_metadatas, # type: ignore[arg-type]
documents=texts_without_metadatas, documents=texts_without_metadatas,
ids=ids_without_metadatas, ids=ids_without_metadatas,
) )
else: else:
self._collection.upsert( self._collection.upsert(
embeddings=embeddings, # type: ignore embeddings=embeddings, # type: ignore[arg-type]
documents=texts, documents=texts,
ids=ids, ids=ids,
) )
@ -586,7 +591,7 @@ class Chroma(VectorStore):
self, self,
query: str, query: str,
k: int = DEFAULT_K, k: int = DEFAULT_K,
filter: Optional[dict[str, str]] = None, filter: Optional[dict[str, str]] = None, # noqa: A002
**kwargs: Any, **kwargs: Any,
) -> list[Document]: ) -> list[Document]:
"""Run similarity search with Chroma. """Run similarity search with Chroma.
@ -601,7 +606,10 @@ class Chroma(VectorStore):
List of documents most similar to the query text. List of documents most similar to the query text.
""" """
docs_and_scores = self.similarity_search_with_score( docs_and_scores = self.similarity_search_with_score(
query, k, filter=filter, **kwargs query,
k,
filter=filter,
**kwargs,
) )
return [doc for doc, _ in docs_and_scores] return [doc for doc, _ in docs_and_scores]
@ -609,7 +617,7 @@ class Chroma(VectorStore):
self, self,
embedding: list[float], embedding: list[float],
k: int = DEFAULT_K, k: int = DEFAULT_K,
filter: Optional[dict[str, str]] = None, filter: Optional[dict[str, str]] = None, # noqa: A002
where_document: Optional[dict[str, str]] = None, where_document: Optional[dict[str, str]] = None,
**kwargs: Any, **kwargs: Any,
) -> list[Document]: ) -> list[Document]:
@ -639,7 +647,7 @@ class Chroma(VectorStore):
self, self,
embedding: list[float], embedding: list[float],
k: int = DEFAULT_K, k: int = DEFAULT_K,
filter: Optional[dict[str, str]] = None, filter: Optional[dict[str, str]] = None, # noqa: A002
where_document: Optional[dict[str, str]] = None, where_document: Optional[dict[str, str]] = None,
**kwargs: Any, **kwargs: Any,
) -> list[tuple[Document, float]]: ) -> list[tuple[Document, float]]:
@ -670,7 +678,7 @@ class Chroma(VectorStore):
self, self,
query: str, query: str,
k: int = DEFAULT_K, k: int = DEFAULT_K,
filter: Optional[dict[str, str]] = None, filter: Optional[dict[str, str]] = None, # noqa: A002
where_document: Optional[dict[str, str]] = None, where_document: Optional[dict[str, str]] = None,
**kwargs: Any, **kwargs: Any,
) -> list[tuple[Document, float]]: ) -> list[tuple[Document, float]]:
@ -712,7 +720,7 @@ class Chroma(VectorStore):
self, self,
query: str, query: str,
k: int = DEFAULT_K, k: int = DEFAULT_K,
filter: Optional[dict[str, str]] = None, filter: Optional[dict[str, str]] = None, # noqa: A002
where_document: Optional[dict[str, str]] = None, where_document: Optional[dict[str, str]] = None,
**kwargs: Any, **kwargs: Any,
) -> list[tuple[Document, np.ndarray]]: ) -> list[tuple[Document, np.ndarray]]:
@ -780,22 +788,24 @@ class Chroma(VectorStore):
if distance == "cosine": if distance == "cosine":
return self._cosine_relevance_score_fn return self._cosine_relevance_score_fn
elif distance == "l2": if distance == "l2":
return self._euclidean_relevance_score_fn return self._euclidean_relevance_score_fn
elif distance == "ip": if distance == "ip":
return self._max_inner_product_relevance_score_fn return self._max_inner_product_relevance_score_fn
else: msg = (
raise ValueError(
"No supported normalization function" "No supported normalization function"
f" for distance metric of type: {distance}." f" for distance metric of type: {distance}."
"Consider providing relevance_score_fn to Chroma constructor." "Consider providing relevance_score_fn to Chroma constructor."
) )
raise ValueError(
msg,
)
def similarity_search_by_image( def similarity_search_by_image(
self, self,
uri: str, uri: str,
k: int = DEFAULT_K, k: int = DEFAULT_K,
filter: Optional[dict[str, str]] = None, filter: Optional[dict[str, str]] = None, # noqa: A002
**kwargs: Any, **kwargs: Any,
) -> list[Document]: ) -> list[Document]:
"""Search for similar images based on the given image URI. """Search for similar images based on the given image URI.
@ -817,29 +827,29 @@ class Chroma(VectorStore):
ValueError: If the embedding function does not support image embeddings. ValueError: If the embedding function does not support image embeddings.
""" """
if self._embedding_function is None or not hasattr( if self._embedding_function is None or not hasattr(
self._embedding_function, "embed_image" self._embedding_function,
"embed_image",
): ):
raise ValueError("The embedding function must support image embedding.") msg = "The embedding function must support image embedding."
raise ValueError(msg)
# Obtain image embedding # Obtain image embedding
# Assuming embed_image returns a single embedding # Assuming embed_image returns a single embedding
image_embedding = self._embedding_function.embed_image(uris=[uri]) image_embedding = self._embedding_function.embed_image(uris=[uri])
# Perform similarity search based on the obtained embedding # Perform similarity search based on the obtained embedding
results = self.similarity_search_by_vector( return self.similarity_search_by_vector(
embedding=image_embedding, embedding=image_embedding,
k=k, k=k,
filter=filter, filter=filter,
**kwargs, **kwargs,
) )
return results
def similarity_search_by_image_with_relevance_score( def similarity_search_by_image_with_relevance_score(
self, self,
uri: str, uri: str,
k: int = DEFAULT_K, k: int = DEFAULT_K,
filter: Optional[dict[str, str]] = None, filter: Optional[dict[str, str]] = None, # noqa: A002
**kwargs: Any, **kwargs: Any,
) -> list[tuple[Document, float]]: ) -> list[tuple[Document, float]]:
"""Search for similar images based on the given image URI. """Search for similar images based on the given image URI.
@ -861,31 +871,31 @@ class Chroma(VectorStore):
ValueError: If the embedding function does not support image embeddings. ValueError: If the embedding function does not support image embeddings.
""" """
if self._embedding_function is None or not hasattr( if self._embedding_function is None or not hasattr(
self._embedding_function, "embed_image" self._embedding_function,
"embed_image",
): ):
raise ValueError("The embedding function must support image embedding.") msg = "The embedding function must support image embedding."
raise ValueError(msg)
# Obtain image embedding # Obtain image embedding
# Assuming embed_image returns a single embedding # Assuming embed_image returns a single embedding
image_embedding = self._embedding_function.embed_image(uris=[uri]) image_embedding = self._embedding_function.embed_image(uris=[uri])
# Perform similarity search based on the obtained embedding # Perform similarity search based on the obtained embedding
results = self.similarity_search_by_vector_with_relevance_scores( return self.similarity_search_by_vector_with_relevance_scores(
embedding=image_embedding, embedding=image_embedding,
k=k, k=k,
filter=filter, filter=filter,
**kwargs, **kwargs,
) )
return results
def max_marginal_relevance_search_by_vector( def max_marginal_relevance_search_by_vector(
self, self,
embedding: list[float], embedding: list[float],
k: int = DEFAULT_K, k: int = DEFAULT_K,
fetch_k: int = 20, fetch_k: int = 20,
lambda_mult: float = 0.5, lambda_mult: float = 0.5,
filter: Optional[dict[str, str]] = None, filter: Optional[dict[str, str]] = None, # noqa: A002
where_document: Optional[dict[str, str]] = None, where_document: Optional[dict[str, str]] = None,
**kwargs: Any, **kwargs: Any,
) -> list[Document]: ) -> list[Document]:
@ -928,8 +938,7 @@ class Chroma(VectorStore):
candidates = _results_to_docs(results) candidates = _results_to_docs(results)
selected_results = [r for i, r in enumerate(candidates) if i in mmr_selected] return [r for i, r in enumerate(candidates) if i in mmr_selected]
return selected_results
def max_marginal_relevance_search( def max_marginal_relevance_search(
self, self,
@ -937,7 +946,7 @@ class Chroma(VectorStore):
k: int = DEFAULT_K, k: int = DEFAULT_K,
fetch_k: int = 20, fetch_k: int = 20,
lambda_mult: float = 0.5, lambda_mult: float = 0.5,
filter: Optional[dict[str, str]] = None, filter: Optional[dict[str, str]] = None, # noqa: A002
where_document: Optional[dict[str, str]] = None, where_document: Optional[dict[str, str]] = None,
**kwargs: Any, **kwargs: Any,
) -> list[Document]: ) -> list[Document]:
@ -966,8 +975,9 @@ class Chroma(VectorStore):
ValueError: If the embedding function is not provided. ValueError: If the embedding function is not provided.
""" """
if self._embedding_function is None: if self._embedding_function is None:
msg = "For MMR search, you must specify an embedding function on creation."
raise ValueError( raise ValueError(
"For MMR search, you must specify an embedding function on creation." msg,
) )
embedding = self._embedding_function.embed_query(query) embedding = self._embedding_function.embed_query(query)
@ -1032,7 +1042,7 @@ class Chroma(VectorStore):
if include is not None: if include is not None:
kwargs["include"] = include kwargs["include"] = include
return self._collection.get(**kwargs) # type: ignore return self._collection.get(**kwargs) # type: ignore[arg-type, return-value]
def get_by_ids(self, ids: Sequence[str], /) -> list[Document]: def get_by_ids(self, ids: Sequence[str], /) -> list[Document]:
"""Get documents by their IDs. """Get documents by their IDs.
@ -1062,7 +1072,9 @@ class Chroma(VectorStore):
return [ return [
Document(page_content=doc, metadata=meta, id=doc_id) Document(page_content=doc, metadata=meta, id=doc_id)
for doc, meta, doc_id in zip( for doc, meta, doc_id in zip(
results["documents"], results["metadatas"], results["ids"] results["documents"],
results["metadatas"],
results["ids"],
) )
] ]
@ -1075,7 +1087,6 @@ class Chroma(VectorStore):
""" """
return self.update_documents([document_id], [document]) return self.update_documents([document_id], [document])
# type: ignore
def update_documents(self, ids: list[str], documents: list[Document]) -> None: def update_documents(self, ids: list[str], documents: list[Document]) -> None:
"""Update a document in the collection. """Update a document in the collection.
@ -1089,24 +1100,27 @@ class Chroma(VectorStore):
text = [document.page_content for document in documents] text = [document.page_content for document in documents]
metadata = [document.metadata for document in documents] metadata = [document.metadata for document in documents]
if self._embedding_function is None: if self._embedding_function is None:
msg = "For update, you must specify an embedding function on creation."
raise ValueError( raise ValueError(
"For update, you must specify an embedding function on creation." msg,
) )
embeddings = self._embedding_function.embed_documents(text) embeddings = self._embedding_function.embed_documents(text)
if hasattr( if hasattr(
self._collection._client, "get_max_batch_size" self._collection._client, # noqa: SLF001
"get_max_batch_size",
) or hasattr( # for Chroma 0.5.1 and above ) or hasattr( # for Chroma 0.5.1 and above
self._collection._client, "max_batch_size" self._collection._client, # noqa: SLF001
"max_batch_size",
): # for Chroma 0.4.10 and above ): # for Chroma 0.4.10 and above
from chromadb.utils.batch_utils import create_batches from chromadb.utils.batch_utils import create_batches
for batch in create_batches( for batch in create_batches(
api=self._collection._client, api=self._collection._client, # noqa: SLF001
ids=ids, ids=ids,
metadatas=metadata, # type: ignore metadatas=metadata, # type: ignore[arg-type]
documents=text, documents=text,
embeddings=embeddings, # type: ignore embeddings=embeddings, # type: ignore[arg-type]
): ):
self._collection.update( self._collection.update(
ids=batch[0], ids=batch[0],
@ -1117,9 +1131,9 @@ class Chroma(VectorStore):
else: else:
self._collection.update( self._collection.update(
ids=ids, ids=ids,
embeddings=embeddings, # type: ignore embeddings=embeddings, # type: ignore[arg-type]
documents=text, documents=text,
metadatas=metadata, # type: ignore metadatas=metadata, # type: ignore[arg-type]
) )
@classmethod @classmethod
@ -1170,23 +1184,25 @@ class Chroma(VectorStore):
if ids is None: if ids is None:
ids = [str(uuid.uuid4()) for _ in texts] ids = [str(uuid.uuid4()) for _ in texts]
else: else:
ids = [id if id is not None else str(uuid.uuid4()) for id in ids] ids = [id_ if id_ is not None else str(uuid.uuid4()) for id_ in ids]
if hasattr( if hasattr(
chroma_collection._client, "get_max_batch_size" chroma_collection._client, # noqa: SLF001
"get_max_batch_size",
) or hasattr( # for Chroma 0.5.1 and above ) or hasattr( # for Chroma 0.5.1 and above
chroma_collection._client, "max_batch_size" chroma_collection._client, # noqa: SLF001
"max_batch_size",
): # for Chroma 0.4.10 and above ): # for Chroma 0.4.10 and above
from chromadb.utils.batch_utils import create_batches from chromadb.utils.batch_utils import create_batches
for batch in create_batches( for batch in create_batches(
api=chroma_collection._client, api=chroma_collection._client, # noqa: SLF001
ids=ids, ids=ids,
metadatas=metadatas, # type: ignore metadatas=metadatas, # type: ignore[arg-type]
documents=texts, documents=texts,
): ):
chroma_collection.add_texts( chroma_collection.add_texts(
texts=batch[3] if batch[3] else [], texts=batch[3] if batch[3] else [],
metadatas=batch[2] if batch[2] else None, # type: ignore metadatas=batch[2] if batch[2] else None, # type: ignore[arg-type]
ids=batch[0], ids=batch[0],
) )
else: else:

49
libs/partners/chroma/pyproject.toml
View File

@ -62,8 +62,52 @@ disallow_untyped_defs = true
target-version = "py39" target-version = "py39"
[tool.ruff.lint] [tool.ruff.lint]
select = ["E", "F", "I", "T201", "D", "UP", "S"] select = [
ignore = [ "UP007", ] "A", # flake8-builtins
"ASYNC", # flake8-async
"C4", # flake8-comprehensions
"COM", # flake8-commas
"D", # pydocstyle
"DOC", # pydoclint
"E", # pycodestyle error
"EM", # flake8-errmsg
"F", # pyflakes
"FA", # flake8-future-annotations
"FBT", # flake8-boolean-trap
"FLY", # flake8-flynt
"I", # isort
"ICN", # flake8-import-conventions
"INT", # flake8-gettext
"ISC", # isort-comprehensions
"PGH", # pygrep-hooks
"PIE", # flake8-pie
"PERF", # flake8-perf
"PYI", # flake8-pyi
"Q", # flake8-quotes
"RET", # flake8-return
"RSE", # flake8-rst-docstrings
"RUF", # ruff
"S", # flake8-bandit
"SLF", # flake8-self
"SLOT", # flake8-slots
"SIM", # flake8-simplify
"T10", # flake8-debugger
"T20", # flake8-print
"TID", # flake8-tidy-imports
"UP", # pyupgrade
"W", # pycodestyle warning
"YTT", # flake8-2020
]
ignore = [
"D100", # Missing docstring in public module
"D101", # Missing docstring in public class
"D102", # Missing docstring in public method
"D103", # Missing docstring in public function
"D104", # Missing docstring in public package
"D105", # Missing docstring in magic method
"D107", # Missing docstring in __init__
"UP007", # pyupgrade: non-pep604-annotation-union
]
[tool.coverage.run] [tool.coverage.run]
omit = ["tests/*"] omit = ["tests/*"]
@ -86,4 +130,5 @@ convention = "google"
"tests/**/*.py" = [ "tests/**/*.py" = [
"S101", # Tests need assertions "S101", # Tests need assertions
"S311", # Standard pseudo-random generators are not suitable for cryptographic purposes "S311", # Standard pseudo-random generators are not suitable for cryptographic purposes
"SLF001", # Private member access in tests
] ]

2
libs/partners/chroma/scripts/check_imports.py
View File

@ -10,7 +10,7 @@ if __name__ == "__main__":
for file in files: for file in files:
try: try:
SourceFileLoader("x", file).load_module() SourceFileLoader("x", file).load_module()
except Exception: except Exception: # noqa: PERF203
has_failure = True has_failure = True
print(file) # noqa: T201 print(file) # noqa: T201
traceback.print_exc() traceback.print_exc()

2
libs/partners/chroma/tests/integration_tests/fake_embeddings.py
View File

@ -44,7 +44,7 @@ class ConsistentFakeEmbeddings(FakeEmbeddings):
if text not in self.known_texts: if text not in self.known_texts:
self.known_texts.append(text) self.known_texts.append(text)
vector = [1.0] * (self.dimensionality - 1) + [ vector = [1.0] * (self.dimensionality - 1) + [
float(self.known_texts.index(text)) float(self.known_texts.index(text)),
] ]
out_vectors.append(vector) out_vectors.append(vector)
return out_vectors return out_vectors

1
libs/partners/chroma/tests/integration_tests/test_compile.py
View File

@ -4,4 +4,3 @@ import pytest # type: ignore[import-not-found]
@pytest.mark.compile @pytest.mark.compile
def test_placeholder() -> None: def test_placeholder() -> None:
"""Used for compiling integration tests without running any real tests.""" """Used for compiling integration tests without running any real tests."""
pass

98
libs/partners/chroma/tests/integration_tests/test_vectorstores.py
View File

@ -28,8 +28,8 @@ class MyEmbeddingFunction:
def __init__(self, fak: Fak): def __init__(self, fak: Fak):
self.fak = fak self.fak = fak
def __call__(self, input: Embeddable) -> list[list[float]]: def __call__(self, input_: Embeddable) -> list[list[float]]:
texts = cast(list[str], input) texts = cast(list[str], input_)
return self.fak.embed_documents(texts=texts) return self.fak.embed_documents(texts=texts)
@ -44,7 +44,9 @@ def test_chroma() -> None:
"""Test end to end construction and search.""" """Test end to end construction and search."""
texts = ["foo", "bar", "baz"] texts = ["foo", "bar", "baz"]
docsearch = Chroma.from_texts( docsearch = Chroma.from_texts(
collection_name="test_collection", texts=texts, embedding=FakeEmbeddings() collection_name="test_collection",
texts=texts,
embedding=FakeEmbeddings(),
) )
output = docsearch.similarity_search("foo", k=1) output = docsearch.similarity_search("foo", k=1)
@ -92,7 +94,9 @@ async def test_chroma_async() -> None:
"""Test end to end construction and search.""" """Test end to end construction and search."""
texts = ["foo", "bar", "baz"] texts = ["foo", "bar", "baz"]
docsearch = Chroma.from_texts( docsearch = Chroma.from_texts(
collection_name="test_collection", texts=texts, embedding=FakeEmbeddings() collection_name="test_collection",
texts=texts,
embedding=FakeEmbeddings(),
) )
output = await docsearch.asimilarity_search("foo", k=1) output = await docsearch.asimilarity_search("foo", k=1)
@ -173,7 +177,7 @@ def test_chroma_with_metadatas_with_scores_and_ids() -> None:
output = docsearch.similarity_search_with_score("foo", k=1) output = docsearch.similarity_search_with_score("foo", k=1)
docsearch.delete_collection() docsearch.delete_collection()
assert output == [ assert output == [
(Document(page_content="foo", metadata={"page": "0"}, id="id_0"), 0.0) (Document(page_content="foo", metadata={"page": "0"}, id="id_0"), 0.0),
] ]
@ -211,11 +215,12 @@ def test_chroma_with_metadatas_with_scores_using_vector() -> None:
) )
embedded_query = embeddings.embed_query("foo") embedded_query = embeddings.embed_query("foo")
output = docsearch.similarity_search_by_vector_with_relevance_scores( output = docsearch.similarity_search_by_vector_with_relevance_scores(
embedding=embedded_query, k=1 embedding=embedded_query,
k=1,
) )
docsearch.delete_collection() docsearch.delete_collection()
assert output == [ assert output == [
(Document(page_content="foo", metadata={"page": "0"}, id="id_0"), 0.0) (Document(page_content="foo", metadata={"page": "0"}, id="id_0"), 0.0),
] ]
@ -235,10 +240,10 @@ def test_chroma_search_filter() -> None:
output2 = docsearch.similarity_search("far", k=1, filter={"first_letter": "b"}) output2 = docsearch.similarity_search("far", k=1, filter={"first_letter": "b"})
docsearch.delete_collection() docsearch.delete_collection()
assert output1 == [ assert output1 == [
Document(page_content="far", metadata={"first_letter": "f"}, id="id_0") Document(page_content="far", metadata={"first_letter": "f"}, id="id_0"),
] ]
assert output2 == [ assert output2 == [
Document(page_content="bar", metadata={"first_letter": "b"}, id="id_1") Document(page_content="bar", metadata={"first_letter": "b"}, id="id_1"),
] ]
@ -255,17 +260,21 @@ def test_chroma_search_filter_with_scores() -> None:
ids=ids, ids=ids,
) )
output1 = docsearch.similarity_search_with_score( output1 = docsearch.similarity_search_with_score(
"far", k=1, filter={"first_letter": "f"} "far",
k=1,
filter={"first_letter": "f"},
) )
output2 = docsearch.similarity_search_with_score( output2 = docsearch.similarity_search_with_score(
"far", k=1, filter={"first_letter": "b"} "far",
k=1,
filter={"first_letter": "b"},
) )
docsearch.delete_collection() docsearch.delete_collection()
assert output1 == [ assert output1 == [
(Document(page_content="far", metadata={"first_letter": "f"}, id="id_0"), 0.0) (Document(page_content="far", metadata={"first_letter": "f"}, id="id_0"), 0.0),
] ]
assert output2 == [ assert output2 == [
(Document(page_content="bar", metadata={"first_letter": "b"}, id="id_1"), 1.0) (Document(page_content="bar", metadata={"first_letter": "b"}, id="id_1"), 1.0),
] ]
@ -365,7 +374,9 @@ def test_chroma_mmr() -> None:
"""Test end to end construction and search.""" """Test end to end construction and search."""
texts = ["foo", "bar", "baz"] texts = ["foo", "bar", "baz"]
docsearch = Chroma.from_texts( docsearch = Chroma.from_texts(
collection_name="test_collection", texts=texts, embedding=FakeEmbeddings() collection_name="test_collection",
texts=texts,
embedding=FakeEmbeddings(),
) )
output = docsearch.max_marginal_relevance_search("foo", k=1) output = docsearch.max_marginal_relevance_search("foo", k=1)
docsearch.delete_collection() docsearch.delete_collection()
@ -379,7 +390,9 @@ def test_chroma_mmr_by_vector() -> None:
texts = ["foo", "bar", "baz"] texts = ["foo", "bar", "baz"]
embeddings = FakeEmbeddings() embeddings = FakeEmbeddings()
docsearch = Chroma.from_texts( docsearch = Chroma.from_texts(
collection_name="test_collection", texts=texts, embedding=embeddings collection_name="test_collection",
texts=texts,
embedding=embeddings,
) )
embedded_query = embeddings.embed_query("foo") embedded_query = embeddings.embed_query("foo")
output = docsearch.max_marginal_relevance_search_by_vector(embedded_query, k=1) output = docsearch.max_marginal_relevance_search_by_vector(embedded_query, k=1)
@ -393,7 +406,9 @@ def test_chroma_with_include_parameter() -> None:
"""Test end to end construction and include parameter.""" """Test end to end construction and include parameter."""
texts = ["foo", "bar", "baz"] texts = ["foo", "bar", "baz"]
docsearch = Chroma.from_texts( docsearch = Chroma.from_texts(
collection_name="test_collection", texts=texts, embedding=FakeEmbeddings() collection_name="test_collection",
texts=texts,
embedding=FakeEmbeddings(),
) )
output1 = docsearch.get(include=["embeddings"]) output1 = docsearch.get(include=["embeddings"])
output2 = docsearch.get() output2 = docsearch.get()
@ -424,7 +439,7 @@ def test_chroma_update_document() -> None:
embedding=embedding, embedding=embedding,
ids=[document_id], ids=[document_id],
) )
old_embedding = docsearch._collection.peek()["embeddings"][ # type: ignore old_embedding = docsearch._collection.peek()["embeddings"][ # type: ignore[index]
docsearch._collection.peek()["ids"].index(document_id) docsearch._collection.peek()["ids"].index(document_id)
] ]
@ -441,7 +456,7 @@ def test_chroma_update_document() -> None:
output = docsearch.similarity_search(updated_content, k=1) output = docsearch.similarity_search(updated_content, k=1)
# Assert that the new embedding is correct # Assert that the new embedding is correct
new_embedding = docsearch._collection.peek()["embeddings"][ # type: ignore new_embedding = docsearch._collection.peek()["embeddings"][ # type: ignore[index]
docsearch._collection.peek()["ids"].index(document_id) docsearch._collection.peek()["ids"].index(document_id)
] ]
@ -449,7 +464,7 @@ def test_chroma_update_document() -> None:
# Assert that the updated document is returned by the search # Assert that the updated document is returned by the search
assert output == [ assert output == [
Document(page_content=updated_content, metadata={"page": "0"}, id=document_id) Document(page_content=updated_content, metadata={"page": "0"}, id=document_id),
] ]
assert list(new_embedding) == list(embedding.embed_documents([updated_content])[0]) assert list(new_embedding) == list(embedding.embed_documents([updated_content])[0])
@ -470,7 +485,9 @@ def test_chroma_update_document_with_id() -> None:
# Create an instance of Document with initial content and metadata # Create an instance of Document with initial content and metadata
original_doc = Document( original_doc = Document(
page_content=initial_content, metadata={"page": "0"}, id=document_id page_content=initial_content,
metadata={"page": "0"},
id=document_id,
) )
# Initialize a Chroma instance with the original document # Initialize a Chroma instance with the original document
@ -479,7 +496,7 @@ def test_chroma_update_document_with_id() -> None:
documents=[original_doc], documents=[original_doc],
embedding=embedding, embedding=embedding,
) )
old_embedding = docsearch._collection.peek()["embeddings"][ # type: ignore old_embedding = docsearch._collection.peek()["embeddings"][ # type: ignore[index]
docsearch._collection.peek()["ids"].index(document_id) docsearch._collection.peek()["ids"].index(document_id)
] ]
@ -488,7 +505,9 @@ def test_chroma_update_document_with_id() -> None:
# Create a new Document instance with the updated content and the same id # Create a new Document instance with the updated content and the same id
updated_doc = Document( updated_doc = Document(
page_content=updated_content, metadata={"page": "0"}, id=document_id page_content=updated_content,
metadata={"page": "0"},
id=document_id,
) )
# Update the document in the Chroma instance # Update the document in the Chroma instance
@ -498,7 +517,7 @@ def test_chroma_update_document_with_id() -> None:
output = docsearch.similarity_search(updated_content, k=1) output = docsearch.similarity_search(updated_content, k=1)
# Assert that the new embedding is correct # Assert that the new embedding is correct
new_embedding = docsearch._collection.peek()["embeddings"][ # type: ignore new_embedding = docsearch._collection.peek()["embeddings"][ # type: ignore[index]
docsearch._collection.peek()["ids"].index(document_id) docsearch._collection.peek()["ids"].index(document_id)
] ]
@ -506,7 +525,7 @@ def test_chroma_update_document_with_id() -> None:
# Assert that the updated document is returned by the search # Assert that the updated document is returned by the search
assert output == [ assert output == [
Document(page_content=updated_content, metadata={"page": "0"}, id=document_id) Document(page_content=updated_content, metadata={"page": "0"}, id=document_id),
] ]
assert list(new_embedding) == list(embedding.embed_documents([updated_content])[0]) assert list(new_embedding) == list(embedding.embed_documents([updated_content])[0])
@ -568,7 +587,8 @@ def test_chroma_add_documents_mixed_metadata() -> None:
assert actual_ids == ids assert actual_ids == ids
assert sorted(search, key=lambda d: d.page_content) == sorted( assert sorted(search, key=lambda d: d.page_content) == sorted(
docs, key=lambda d: d.page_content docs,
key=lambda d: d.page_content,
) )
@ -582,9 +602,7 @@ def is_api_accessible(url: str) -> bool:
def batch_support_chroma_version() -> bool: def batch_support_chroma_version() -> bool:
major, minor, patch = chromadb.__version__.split(".") major, minor, patch = chromadb.__version__.split(".")
if int(major) == 0 and int(minor) >= 4 and int(patch) >= 10: return bool(int(major) == 0 and int(minor) >= 4 and int(patch) >= 10)
return True
return False
@pytest.mark.requires("chromadb") @pytest.mark.requires("chromadb")
@ -601,9 +619,9 @@ def test_chroma_large_batch() -> None:
embedding_function = MyEmbeddingFunction(fak=Fak(size=255)) embedding_function = MyEmbeddingFunction(fak=Fak(size=255))
col = client.get_or_create_collection( col = client.get_or_create_collection(
"my_collection", "my_collection",
embedding_function=embedding_function, # type: ignore embedding_function=embedding_function, # type: ignore[arg-type]
) )
docs = ["This is a test document"] * (client.get_max_batch_size() + 100) # type: ignore docs = ["This is a test document"] * (client.get_max_batch_size() + 100)
db = Chroma.from_texts( db = Chroma.from_texts(
client=client, client=client,
collection_name=col.name, collection_name=col.name,
@ -629,9 +647,9 @@ def test_chroma_large_batch_update() -> None:
embedding_function = MyEmbeddingFunction(fak=Fak(size=255)) embedding_function = MyEmbeddingFunction(fak=Fak(size=255))
col = client.get_or_create_collection( col = client.get_or_create_collection(
"my_collection", "my_collection",
embedding_function=embedding_function, # type: ignore embedding_function=embedding_function, # type: ignore[arg-type]
) )
docs = ["This is a test document"] * (client.get_max_batch_size() + 100) # type: ignore docs = ["This is a test document"] * (client.get_max_batch_size() + 100)
ids = [str(uuid.uuid4()) for _ in range(len(docs))] ids = [str(uuid.uuid4()) for _ in range(len(docs))]
db = Chroma.from_texts( db = Chroma.from_texts(
client=client, client=client,
@ -642,11 +660,12 @@ def test_chroma_large_batch_update() -> None:
) )
new_docs = [ new_docs = [
Document( Document(
page_content="This is a new test document", metadata={"doc_id": f"{i}"} page_content="This is a new test document",
metadata={"doc_id": f"{i}"},
) )
for i in range(len(docs) - 10) for i in range(len(docs) - 10)
] ]
new_ids = [_id for _id in ids[: len(new_docs)]] new_ids = list(ids[: len(new_docs)])
db.update_documents(ids=new_ids, documents=new_docs) db.update_documents(ids=new_ids, documents=new_docs)
db.delete_collection() db.delete_collection()
@ -658,14 +677,15 @@ def test_chroma_large_batch_update() -> None:
reason="API not accessible", reason="API not accessible",
) )
@pytest.mark.skipif( @pytest.mark.skipif(
batch_support_chroma_version(), reason="ChromaDB version does not support batching" batch_support_chroma_version(),
reason="ChromaDB version does not support batching",
) )
def test_chroma_legacy_batching() -> None: def test_chroma_legacy_batching() -> None:
client = chromadb.HttpClient() client = chromadb.HttpClient()
embedding_function = Fak(size=255) embedding_function = Fak(size=255)
col = client.get_or_create_collection( col = client.get_or_create_collection(
"my_collection", "my_collection",
embedding_function=MyEmbeddingFunction, # type: ignore embedding_function=MyEmbeddingFunction, # type: ignore[arg-type]
) )
docs = ["This is a test document"] * 100 docs = ["This is a test document"] * 100
db = Chroma.from_texts( db = Chroma.from_texts(
@ -683,7 +703,9 @@ def test_create_collection_if_not_exist_default() -> None:
"""Tests existing behaviour without the new create_collection_if_not_exists flag.""" """Tests existing behaviour without the new create_collection_if_not_exists flag."""
texts = ["foo", "bar", "baz"] texts = ["foo", "bar", "baz"]
docsearch = Chroma.from_texts( docsearch = Chroma.from_texts(
collection_name="test_collection", texts=texts, embedding=FakeEmbeddings() collection_name="test_collection",
texts=texts,
embedding=FakeEmbeddings(),
) )
assert docsearch._client.get_collection("test_collection") is not None assert docsearch._client.get_collection("test_collection") is not None
docsearch.delete_collection() docsearch.delete_collection()
@ -798,7 +820,7 @@ def test_delete_where_clause(client: chromadb.ClientAPI) -> None:
[ [
Document(page_content="foo", metadata={"test": "bar"}), Document(page_content="foo", metadata={"test": "bar"}),
Document(page_content="bar", metadata={"test": "foo"}), Document(page_content="bar", metadata={"test": "foo"}),
] ],
) )
assert vectorstore._collection.count() == 2 assert vectorstore._collection.count() == 2
vectorstore.delete(where={"test": "bar"}) vectorstore.delete(where={"test": "bar"})

3
libs/partners/chroma/tests/unit_tests/test_standard.py
View File

@ -9,11 +9,10 @@ from langchain_chroma import Chroma
class TestChromaStandard(VectorStoreIntegrationTests): class TestChromaStandard(VectorStoreIntegrationTests):
@pytest.fixture() @pytest.fixture()
def vectorstore(self) -> Generator[VectorStore, None, None]: # type: ignore def vectorstore(self) -> Generator[VectorStore, None, None]: # type: ignore[override]
"""Get an empty vectorstore for unit tests.""" """Get an empty vectorstore for unit tests."""
store = Chroma(embedding_function=self.get_embeddings()) store = Chroma(embedding_function=self.get_embeddings())
try: try:
yield store yield store
finally: finally:
store.delete_collection() store.delete_collection()
pass