From ac244d18c174b105d9410285099d53d2d5b8c31b Mon Sep 17 00:00:00 2001 From: Dinesh Yeduguru Date: Wed, 20 Nov 2024 00:03:15 -0800 Subject: [PATCH] for pgvector --- llama_stack/providers/remote/memory/pgvector/pgvector.py | 9 ++++++++- 1 file changed, 8 insertions(+), 1 deletion(-) diff --git a/llama_stack/providers/remote/memory/pgvector/pgvector.py b/llama_stack/providers/remote/memory/pgvector/pgvector.py index 44c2a8fe1..46a85bf1e 100644 --- a/llama_stack/providers/remote/memory/pgvector/pgvector.py +++ b/llama_stack/providers/remote/memory/pgvector/pgvector.py @@ -215,6 +215,13 @@ class PGVectorMemoryAdapter(Memory, MemoryBanksProtocolPrivate): ) -> QueryDocumentsResponse: index = self.cache.get(bank_id, None) if not index: - raise ValueError(f"Bank {bank_id} not found") + 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)