From 49a440acfb961bdc4822d4b580ce82e9645851ad Mon Sep 17 00:00:00 2001 From: Dinesh Yeduguru Date: Tue, 19 Nov 2024 23:59:15 -0800 Subject: [PATCH] fall to back to read from chroma when not in cache --- llama_stack/providers/remote/memory/chroma/chroma.py | 10 +++++++++- 1 file changed, 9 insertions(+), 1 deletion(-) diff --git a/llama_stack/providers/remote/memory/chroma/chroma.py b/llama_stack/providers/remote/memory/chroma/chroma.py index ac00fc749..8d602e7e8 100644 --- a/llama_stack/providers/remote/memory/chroma/chroma.py +++ b/llama_stack/providers/remote/memory/chroma/chroma.py @@ -161,6 +161,14 @@ class ChromaMemoryAdapter(Memory, MemoryBanksProtocolPrivate): ) -> QueryDocumentsResponse: index = self.cache.get(bank_id, None) 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)