add _get_and_cache_bank_index

llama_stack/providers/remote/memory/chroma/chroma.py

@@ -147,9 +147,7 @@ class ChromaMemoryAdapter(Memory, MemoryBanksProtocolPrivate):
         documents: List[MemoryBankDocument],
         ttl_seconds: Optional[int] = None,
     ) -> None:
-        index = self.cache.get(bank_id, None)
-        if not index:
-            raise ValueError(f"Bank {bank_id} not found")
+        index = await self._get_and_cache_bank_index(bank_id)
 
         await index.insert_documents(documents)
 
@@ -159,9 +157,14 @@ class ChromaMemoryAdapter(Memory, MemoryBanksProtocolPrivate):
         query: InterleavedTextMedia,
         params: Optional[Dict[str, Any]] = None,
     ) -> QueryDocumentsResponse:
-        index = self.cache.get(bank_id, None)
-        if not index:
-            # if not in cache, try to get from chroma directly
+        index = await self._get_and_cache_bank_index(bank_id)
+
+        return await index.query_documents(query, params)
+
+    async def _get_and_cache_bank_index(self, bank_id: str) -> BankWithIndex:
+        if bank_id in self.cache:
+            return self.cache[bank_id]
+
         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")
@@ -170,5 +173,4 @@ class ChromaMemoryAdapter(Memory, MemoryBanksProtocolPrivate):
             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 index

llama_stack/providers/remote/memory/pgvector/pgvector.py

@@ -201,10 +201,7 @@ class PGVectorMemoryAdapter(Memory, MemoryBanksProtocolPrivate):
         documents: List[MemoryBankDocument],
         ttl_seconds: Optional[int] = None,
     ) -> None:
-        index = self.cache.get(bank_id, None)
-        if not index:
-            raise ValueError(f"Bank {bank_id} not found")
-
+        index = await self._get_and_cache_bank_index(bank_id)
         await index.insert_documents(documents)
 
     async def query_documents(
@@ -213,15 +210,17 @@ class PGVectorMemoryAdapter(Memory, MemoryBanksProtocolPrivate):
         query: InterleavedTextMedia,
         params: Optional[Dict[str, Any]] = None,
     ) -> QueryDocumentsResponse:
-        index = self.cache.get(bank_id, None)
-        if not index:
+        index = await self._get_and_cache_bank_index(bank_id)
+        return await index.query_documents(query, params)
+
+    async def _get_and_cache_bank_index(self, bank_id: str) -> BankWithIndex:
+        if bank_id in self.cache:
+            return self.cache[bank_id]
+
         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")
         index = BankWithIndex(
             bank=bank,
             index=PGVectorIndex(bank, ALL_MINILM_L6_V2_DIMENSION, self.cursor),
         )
         self.cache[bank_id] = index
-
-        return await index.query_documents(query, params)
+        return index