From f2feb7d15c748bfd94d82a04fc6c9c6ca1d08b71 Mon Sep 17 00:00:00 2001 From: Ashwin Bharambe Date: Tue, 28 Jan 2025 13:19:47 -0800 Subject: [PATCH] Fix Chroma adapter (#893) Chroma method had the wrong signature. ## Test Plan Start Chroma: `chroma run --path /tmp/foo/chroma2 --host localhost --port 6001` Modify run.yaml to include Chroma server pointing to localhost:6001 and run `llama stack run` Then: ```bash LLAMA_STACK_BASE_URL=http://localhost:8321 pytest -s -v agents/test_agents.py -k rag ``` passes --- llama_stack/providers/remote/vector_io/chroma/chroma.py | 7 ++++--- 1 file changed, 4 insertions(+), 3 deletions(-) diff --git a/llama_stack/providers/remote/vector_io/chroma/chroma.py b/llama_stack/providers/remote/vector_io/chroma/chroma.py index 724dc3f51..a6c17e979 100644 --- a/llama_stack/providers/remote/vector_io/chroma/chroma.py +++ b/llama_stack/providers/remote/vector_io/chroma/chroma.py @@ -46,11 +46,12 @@ class ChromaIndex(EmbeddingIndex): embeddings ), f"Chunk length {len(chunks)} does not match embedding length {len(embeddings)}" + ids = [f"{c.metadata['document_id']}:chunk-{i}" for i, c in enumerate(chunks)] await maybe_await( self.collection.add( documents=[chunk.model_dump_json() for chunk in chunks], embeddings=embeddings, - ids=[f"{c.document_id}:chunk-{i}" for i, c in enumerate(chunks)], + ids=ids, ) ) @@ -140,11 +141,11 @@ class ChromaVectorIOAdapter(VectorIO, VectorDBsProtocolPrivate): self, vector_db_id: str, chunks: List[Chunk], - embeddings: NDArray, + ttl_seconds: Optional[int] = None, ) -> None: index = await self._get_and_cache_vector_db_index(vector_db_id) - await index.insert_chunks(chunks, embeddings) + await index.insert_chunks(chunks) async def query_chunks( self,