From b909d54e70d7dddfd856489822ad574be5af6e84 Mon Sep 17 00:00:00 2001
From: Kaiwei Zhang <134007383+kwei-zhang@users.noreply.github.com>
Date: Fri, 13 Dec 2024 16:58:34 -0500
Subject: [PATCH] chroma[patch]: Update logic for assigning ids

---
 libs/partners/chroma/langchain_chroma/vectorstores.py  | 10 ++++++----
 .../{integration_tests => unit_tests}/test_standard.py |  0
 2 files changed, 6 insertions(+), 4 deletions(-)
 rename libs/partners/chroma/tests/{integration_tests => unit_tests}/test_standard.py (100%)

diff --git a/libs/partners/chroma/langchain_chroma/vectorstores.py b/libs/partners/chroma/langchain_chroma/vectorstores.py
index afb9191c60a..bc87f9e088f 100644
--- a/libs/partners/chroma/langchain_chroma/vectorstores.py
+++ b/libs/partners/chroma/langchain_chroma/vectorstores.py
@@ -433,6 +433,8 @@ class Chroma(VectorStore):
         # Populate IDs
         if ids is None:
             ids = [str(uuid.uuid4()) for _ in uris]
+        else:
+            ids = [id if id is not None else str(uuid.uuid4()) for id in ids]
         embeddings = None
         # Set embeddings
         if self._embedding_function is not None and hasattr(
@@ -519,10 +521,8 @@ class Chroma(VectorStore):
         if ids is None:
             ids = [str(uuid.uuid4()) for _ in texts]
         else:
-            # Assign strings to any null IDs
-            for idx, _id in enumerate(ids):
-                if _id is None:
-                    ids[idx] = str(uuid.uuid4())
+            ids = [id if id is not None else str(uuid.uuid4()) for id in ids]
+
         embeddings = None
         texts = list(texts)
         if self._embedding_function is not None:
@@ -1169,6 +1169,8 @@ class Chroma(VectorStore):
         )
         if ids is None:
             ids = [str(uuid.uuid4()) for _ in texts]
+        else:
+            ids = [id if id is not None else str(uuid.uuid4()) for id in ids]
         if hasattr(
             chroma_collection._client, "get_max_batch_size"
         ) or hasattr(  # for Chroma 0.5.1 and above
diff --git a/libs/partners/chroma/tests/integration_tests/test_standard.py b/libs/partners/chroma/tests/unit_tests/test_standard.py
similarity index 100%
rename from libs/partners/chroma/tests/integration_tests/test_standard.py
rename to libs/partners/chroma/tests/unit_tests/test_standard.py