diff --git a/mem0/vector_stores/chroma.py b/mem0/vector_stores/chroma.py index 6bdc9a4bf..b70bd5fc5 100644 --- a/mem0/vector_stores/chroma.py +++ b/mem0/vector_stores/chroma.py @@ -187,7 +187,7 @@ class ChromaDB(VectorStoreBase): """ self.collection.update(ids=vector_id, embeddings=vector, metadatas=payload) - def get(self, vector_id: str) -> OutputData: + def get(self, vector_id: str) -> Optional[OutputData]: """ Retrieve a vector by ID. @@ -195,10 +195,11 @@ class ChromaDB(VectorStoreBase): vector_id (str): ID of the vector to retrieve. Returns: - OutputData: Retrieved vector. + Optional[OutputData]: Retrieved vector, or None if the ID is not found. """ result = self.collection.get(ids=[vector_id]) - return self._parse_output(result)[0] + parsed = self._parse_output(result) + return parsed[0] if parsed else None def list_cols(self) -> List[chromadb.Collection]: """ diff --git a/tests/vector_stores/test_chroma.py b/tests/vector_stores/test_chroma.py index 9c43c24c8..ffa28dbde 100644 --- a/tests/vector_stores/test_chroma.py +++ b/tests/vector_stores/test_chroma.py @@ -154,6 +154,16 @@ def test_get_vector(chromadb_instance): assert result.payload == {"name": "vector1"} +def test_get_missing_vector_returns_none(chromadb_instance): + # Chroma returns empty lists for an unknown id; get() must return None + # rather than raising IndexError (parity with qdrant/pgvector/faiss). + chromadb_instance.collection.get.return_value = {"ids": [], "metadatas": []} + + result = chromadb_instance.get(vector_id="does-not-exist") + + assert result is None + + def test_list_vectors(chromadb_instance): mock_result = { "ids": [["id1", "id2"]],