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
This commit is contained in:
Ashwin Bharambe 2025-01-28 13:19:47 -08:00 committed by GitHub
parent ec3ebb5bcf
commit f2feb7d15c
No known key found for this signature in database
GPG key ID: B5690EEEBB952194

View file

@ -46,11 +46,12 @@ class ChromaIndex(EmbeddingIndex):
embeddings embeddings
), f"Chunk length {len(chunks)} does not match embedding length {len(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( await maybe_await(
self.collection.add( self.collection.add(
documents=[chunk.model_dump_json() for chunk in chunks], documents=[chunk.model_dump_json() for chunk in chunks],
embeddings=embeddings, 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, self,
vector_db_id: str, vector_db_id: str,
chunks: List[Chunk], chunks: List[Chunk],
embeddings: NDArray, ttl_seconds: Optional[int] = None,
) -> None: ) -> None:
index = await self._get_and_cache_vector_db_index(vector_db_id) 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( async def query_chunks(
self, self,