[memory refactor][5/n] Migrate all vector_io providers (#835)

See https://github.com/meta-llama/llama-stack/issues/827 for the broader
design.

This PR finishes off all the stragglers and migrates everything to the
new naming.

llama_stack/providers/remote/vector_io/chroma/__init__.py

@@ -14,8 +14,8 @@ from .config import ChromaRemoteImplConfig
 async def get_adapter_impl(
     config: ChromaRemoteImplConfig, deps: Dict[Api, ProviderSpec]
 ):
-    from .chroma import ChromaMemoryAdapter
+    from .chroma import ChromaVectorIOAdapter
 
-    impl = ChromaMemoryAdapter(config, deps[Api.inference])
+    impl = ChromaVectorIOAdapter(config, deps[Api.inference])
     await impl.initialize()
     return impl

llama_stack/providers/remote/vector_io/chroma/chroma.py

@@ -6,25 +6,20 @@
 import asyncio
 import json
 import logging
-from typing import List, Optional, Union
+from typing import Any, Dict, List, Optional, Union
 from urllib.parse import urlparse
 
 import chromadb
 from numpy.typing import NDArray
 
 from llama_stack.apis.inference import InterleavedContent
-from llama_stack.apis.memory import (
-    Chunk,
-    Memory,
-    MemoryBankDocument,
-    QueryDocumentsResponse,
-)
-from llama_stack.apis.memory_banks import MemoryBank, MemoryBankType
-from llama_stack.providers.datatypes import Api, MemoryBanksProtocolPrivate
-from llama_stack.providers.inline.memory.chroma import ChromaInlineImplConfig
+from llama_stack.apis.vector_dbs import VectorDB
+from llama_stack.apis.vector_io import Chunk, QueryChunksResponse, VectorIO
+from llama_stack.providers.datatypes import Api, VectorDBsProtocolPrivate
+from llama_stack.providers.inline.vector_io.chroma import ChromaInlineImplConfig
 from llama_stack.providers.utils.memory.vector_store import (
-    BankWithIndex,
     EmbeddingIndex,
+    VectorDBWithIndex,
 )
 from .config import ChromaRemoteImplConfig
 
@@ -61,7 +56,7 @@ class ChromaIndex(EmbeddingIndex):
 
     async def query(
         self, embedding: NDArray, k: int, score_threshold: float
-    ) -> QueryDocumentsResponse:
+    ) -> QueryChunksResponse:
         results = await maybe_await(
             self.collection.query(
                 query_embeddings=[embedding.tolist()],
@@ -85,19 +80,19 @@ class ChromaIndex(EmbeddingIndex):
             chunks.append(chunk)
             scores.append(1.0 / float(dist))
 
-        return QueryDocumentsResponse(chunks=chunks, scores=scores)
+        return QueryChunksResponse(chunks=chunks, scores=scores)
 
     async def delete(self):
         await maybe_await(self.client.delete_collection(self.collection.name))
 
 
-class ChromaMemoryAdapter(Memory, MemoryBanksProtocolPrivate):
+class ChromaVectorIOAdapter(VectorIO, VectorDBsProtocolPrivate):
     def __init__(
         self,
         config: Union[ChromaRemoteImplConfig, ChromaInlineImplConfig],
         inference_api: Api.inference,
     ) -> None:
-        log.info(f"Initializing ChromaMemoryAdapter with url: {config}")
+        log.info(f"Initializing ChromaVectorIOAdapter with url: {config}")
         self.config = config
         self.inference_api = inference_api
 
@@ -123,60 +118,58 @@ class ChromaMemoryAdapter(Memory, MemoryBanksProtocolPrivate):
     async def shutdown(self) -> None:
         pass
 
-    async def register_memory_bank(
+    async def register_vector_db(
         self,
-        memory_bank: MemoryBank,
+        vector_db: VectorDB,
     ) -> None:
-        assert (
-            memory_bank.memory_bank_type == MemoryBankType.vector.value
-        ), f"Only vector banks are supported {memory_bank.memory_bank_type}"
-
         collection = await maybe_await(
             self.client.get_or_create_collection(
-                name=memory_bank.identifier,
-                metadata={"bank": memory_bank.model_dump_json()},
+                name=vector_db.identifier,
+                metadata={"vector_db": vector_db.model_dump_json()},
             )
         )
-        self.cache[memory_bank.identifier] = BankWithIndex(
-            memory_bank, ChromaIndex(self.client, collection), self.inference_api
+        self.cache[vector_db.identifier] = VectorDBWithIndex(
+            vector_db, ChromaIndex(self.client, collection), self.inference_api
         )
 
-    async def unregister_memory_bank(self, memory_bank_id: str) -> None:
-        await self.cache[memory_bank_id].index.delete()
-        del self.cache[memory_bank_id]
+    async def unregister_vector_db(self, vector_db_id: str) -> None:
+        await self.cache[vector_db_id].index.delete()
+        del self.cache[vector_db_id]
 
-    async def insert_documents(
+    async def insert_chunks(
         self,
-        bank_id: str,
-        documents: List[MemoryBankDocument],
-        ttl_seconds: Optional[int] = None,
+        vector_db_id: str,
+        chunks: List[Chunk],
+        embeddings: NDArray,
     ) -> None:
-        index = await self._get_and_cache_bank_index(bank_id)
+        index = await self._get_and_cache_vector_db_index(vector_db_id)
 
-        await index.insert_documents(documents)
+        await index.insert_chunks(chunks, embeddings)
 
-    async def query_documents(
+    async def query_chunks(
         self,
-        bank_id: str,
+        vector_db_id: str,
         query: InterleavedContent,
         params: Optional[Dict[str, Any]] = None,
-    ) -> QueryDocumentsResponse:
-        index = await self._get_and_cache_bank_index(bank_id)
+    ) -> QueryChunksResponse:
+        index = await self._get_and_cache_vector_db_index(vector_db_id)
 
-        return await index.query_documents(query, params)
+        return await index.query_chunks(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]
+    async def _get_and_cache_vector_db_index(
+        self, vector_db_id: str
+    ) -> VectorDBWithIndex:
+        if vector_db_id in self.cache:
+            return self.cache[vector_db_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")
-        collection = await maybe_await(self.client.get_collection(bank_id))
+        vector_db = await self.vector_db_store.get_vector_db(vector_db_id)
+        if not vector_db:
+            raise ValueError(f"Vector DB {vector_db_id} not found in Llama Stack")
+        collection = await maybe_await(self.client.get_collection(vector_db_id))
         if not collection:
-            raise ValueError(f"Bank {bank_id} not found in Chroma")
-        index = BankWithIndex(
-            bank, ChromaIndex(self.client, collection), self.inference_api
+            raise ValueError(f"Vector DB {vector_db_id} not found in Chroma")
+        index = VectorDBWithIndex(
+            vector_db, ChromaIndex(self.client, collection), self.inference_api
         )
-        self.cache[bank_id] = index
+        self.cache[vector_db_id] = index
         return index