fall to back to read from chroma when not in cache

This commit is contained in:
Dinesh Yeduguru 2024-11-19 23:59:15 -08:00
parent 89f5093dfc
commit 49a440acfb

View file

@ -161,6 +161,14 @@ class ChromaMemoryAdapter(Memory, MemoryBanksProtocolPrivate):
) -> QueryDocumentsResponse: ) -> QueryDocumentsResponse:
index = self.cache.get(bank_id, None) index = self.cache.get(bank_id, None)
if not index: if not index:
raise ValueError(f"Bank {bank_id} not found") # if not in cache, try to get from chroma directly
bank = await self.memory_bank_store.get_memory_bank(bank_id)
if not bank:
raise ValueError(f"Bank {bank_id} not found in Llama Stack")
collection = await self.client.get_collection(bank_id)
if not collection:
raise ValueError(f"Bank {bank_id} not found in Chroma")
index = BankWithIndex(bank=bank, index=ChromaIndex(self.client, collection))
self.cache[bank_id] = index
return await index.query_documents(query, params) return await index.query_documents(query, params)