migrate memory banks to Resource and new registration

llama_stack/providers/datatypes.py

@@ -13,7 +13,6 @@ from pydantic import BaseModel, Field
 
 from llama_stack.apis.datasets import DatasetDef
 from llama_stack.apis.eval_tasks import EvalTaskDef
-from llama_stack.apis.memory_banks import MemoryBankDef
 from llama_stack.apis.models import Model
 from llama_stack.apis.scoring_functions import ScoringFnDef
 from llama_stack.apis.shields import Shield
@@ -51,9 +50,9 @@ class ShieldsProtocolPrivate(Protocol):
 
 
 class MemoryBanksProtocolPrivate(Protocol):
-    async def list_memory_banks(self) -> List[MemoryBankDef]: ...
+    async def list_memory_banks(self) -> List[MemoryBank]: ...
 
-    async def register_memory_bank(self, memory_bank: MemoryBankDef) -> None: ...
+    async def register_memory_bank(self, memory_bank: MemoryBank) -> None: ...
 
 
 class DatasetsProtocolPrivate(Protocol):

llama_stack/providers/inline/agents/meta_reference/agent_instance.py

@@ -641,7 +641,7 @@ class ChatAgent(ShieldRunnerMixin):
 
         if session_info.memory_bank_id is None:
             bank_id = f"memory_bank_{session_id}"
-            memory_bank = VectorMemoryBankDef(
+            memory_bank = VectorMemoryBank(
                 identifier=bank_id,
                 embedding_model="all-MiniLM-L6-v2",
                 chunk_size_in_tokens=512,

llama_stack/providers/inline/memory/faiss/faiss.py

@@ -83,7 +83,7 @@ class FaissMemoryImpl(Memory, MemoryBanksProtocolPrivate):
         stored_banks = await self.kvstore.range(start_key, end_key)
 
         for bank_data in stored_banks:
-            bank = VectorMemoryBankDef.model_validate_json(bank_data)
+            bank = VectorMemoryBank.model_validate_json(bank_data)
             index = BankWithIndex(
                 bank=bank, index=FaissIndex(ALL_MINILM_L6_V2_DIMENSION)
             )
@@ -95,10 +95,10 @@ class FaissMemoryImpl(Memory, MemoryBanksProtocolPrivate):
 
     async def register_memory_bank(
         self,
-        memory_bank: MemoryBankDef,
+        memory_bank: MemoryBank,
     ) -> None:
         assert (
-            memory_bank.type == MemoryBankType.vector.value
+            memory_bank.memory_bank_type == MemoryBankType.vector.value
         ), f"Only vector banks are supported {memory_bank.type}"
 
         # Store in kvstore
@@ -114,7 +114,7 @@ class FaissMemoryImpl(Memory, MemoryBanksProtocolPrivate):
         )
         self.cache[memory_bank.identifier] = index
 
-    async def list_memory_banks(self) -> List[MemoryBankDef]:
+    async def list_memory_banks(self) -> List[MemoryBank]:
         return [i.bank for i in self.cache.values()]
 
     async def insert_documents(

llama_stack/providers/remote/memory/qdrant/qdrant.py

@@ -12,6 +12,7 @@ from numpy.typing import NDArray
 from qdrant_client import AsyncQdrantClient, models
 from qdrant_client.models import PointStruct
 
+from llama_stack.apis.memory_banks import *  # noqa: F403
 from llama_stack.providers.datatypes import MemoryBanksProtocolPrivate
 
 from llama_stack.apis.memory import *  # noqa: F403
@@ -112,11 +113,11 @@ class QdrantVectorMemoryAdapter(Memory, MemoryBanksProtocolPrivate):
 
     async def register_memory_bank(
         self,
-        memory_bank: MemoryBankDef,
+        memory_bank: MemoryBank,
     ) -> None:
         assert (
-            memory_bank.type == MemoryBankType.vector.value
-        ), f"Only vector banks are supported {memory_bank.type}"
+            memory_bank.memory_bank_type == MemoryBankType.vector
+        ), f"Only vector banks are supported {memory_bank.memory_bank_type}"
 
         index = BankWithIndex(
             bank=memory_bank,
@@ -125,7 +126,7 @@ class QdrantVectorMemoryAdapter(Memory, MemoryBanksProtocolPrivate):
 
         self.cache[memory_bank.identifier] = index
 
-    async def list_memory_banks(self) -> List[MemoryBankDef]:
+    async def list_memory_banks(self) -> List[MemoryBank]:
         # Qdrant doesn't have collection level metadata to store the bank properties
         # So we only return from the cache value
         return [i.bank for i in self.cache.values()]

llama_stack/providers/remote/memory/weaviate/weaviate.py

@@ -114,11 +114,11 @@ class WeaviateMemoryAdapter(
 
     async def register_memory_bank(
         self,
-        memory_bank: MemoryBankDef,
+        memory_bank: MemoryBank,
     ) -> None:
         assert (
-            memory_bank.type == MemoryBankType.vector.value
-        ), f"Only vector banks are supported {memory_bank.type}"
+            memory_bank.memory_bank_type == MemoryBankType.vector
+        ), f"Only vector banks are supported {memory_bank.memory_bank_type}"
 
         client = self._get_client()
 
@@ -141,7 +141,7 @@ class WeaviateMemoryAdapter(
         )
         self.cache[memory_bank.identifier] = index
 
-    async def list_memory_banks(self) -> List[MemoryBankDef]:
+    async def list_memory_banks(self) -> List[MemoryBank]:
         # TODO: right now the Llama Stack is the source of truth for these banks. That is
         # not ideal. It should be Weaviate which is the source of truth. Unfortunately,
         # list() happens at Stack startup when the Weaviate client (credentials) is not
@@ -157,8 +157,8 @@ class WeaviateMemoryAdapter(
             raise ValueError(f"Bank {bank_id} not found")
 
         client = self._get_client()
-        if not client.collections.exists(bank_id):
-            raise ValueError(f"Collection with name `{bank_id}` not found")
+        if not client.collections.exists(bank.identifier):
+            raise ValueError(f"Collection with name `{bank.identifier}` not found")
 
         index = BankWithIndex(
             bank=bank,

llama_stack/providers/tests/memory/test_memory.py

@@ -43,14 +43,15 @@ def sample_documents():
 
 
 async def register_memory_bank(banks_impl: MemoryBanks):
-    bank = VectorMemoryBankDef(
-        identifier="test_bank",
-        embedding_model="all-MiniLM-L6-v2",
-        chunk_size_in_tokens=512,
-        overlap_size_in_tokens=64,
-    )
 
-    await banks_impl.register_memory_bank(bank)
+    await banks_impl.register_memory_bank(
+        VectorRegistration(
+            memory_bank_id="test_bank",
+            embedding_model="all-MiniLM-L6-v2",
+            chunk_size_in_tokens=512,
+            overlap_size_in_tokens=64,
+        )
+    )
 
 
 class TestMemory:
@@ -68,8 +69,8 @@ class TestMemory:
         # NOTE: this needs you to ensure that you are starting from a clean state
         # but so far we don't have an unregister API unfortunately, so be careful
         _, banks_impl = memory_stack
-        bank = VectorMemoryBankDef(
-            identifier="test_bank_no_provider",
+        bank = VectorRegistration(
+            memory_bank_id="test_bank_no_provider",
             embedding_model="all-MiniLM-L6-v2",
             chunk_size_in_tokens=512,
             overlap_size_in_tokens=64,

llama_stack/providers/utils/memory/vector_store.py

@@ -148,7 +148,7 @@ class EmbeddingIndex(ABC):
 
 @dataclass
 class BankWithIndex:
-    bank: MemoryBankDef
+    bank: VectorMemoryBank
     index: EmbeddingIndex
 
     async def insert_documents(