Merge branch 'meta-llama:main' into qdrant

llama_stack/providers/adapters/memory/chroma/chroma.py

@@ -5,16 +5,17 @@
 # the root directory of this source tree.
 
 import json
-import uuid
 from typing import List
 from urllib.parse import urlparse
 
 import chromadb
 from numpy.typing import NDArray
 
-from llama_stack.apis.memory import *  # noqa: F403
-from llama_stack.distribution.datatypes import RoutableProvider
+from pydantic import parse_obj_as
 
+from llama_stack.apis.memory import *  # noqa: F403
+
+from llama_stack.providers.datatypes import MemoryBanksProtocolPrivate
 from llama_stack.providers.utils.memory.vector_store import (
     BankWithIndex,
     EmbeddingIndex,
@@ -65,7 +66,7 @@ class ChromaIndex(EmbeddingIndex):
         return QueryDocumentsResponse(chunks=chunks, scores=scores)
 
 
-class ChromaMemoryAdapter(Memory, RoutableProvider):
+class ChromaMemoryAdapter(Memory, MemoryBanksProtocolPrivate):
     def __init__(self, url: str) -> None:
         print(f"Initializing ChromaMemoryAdapter with url: {url}")
         url = url.rstrip("/")
@@ -93,56 +94,43 @@ class ChromaMemoryAdapter(Memory, RoutableProvider):
     async def shutdown(self) -> None:
         pass
 
-    async def validate_routing_keys(self, routing_keys: List[str]) -> None:
-        print(f"[chroma] Registering memory bank routing keys: {routing_keys}")
-        pass
-
-    async def create_memory_bank(
+    async def register_memory_bank(
         self,
-        name: str,
-        config: MemoryBankConfig,
-        url: Optional[URL] = None,
-    ) -> MemoryBank:
-        bank_id = str(uuid.uuid4())
-        bank = MemoryBank(
-            bank_id=bank_id,
-            name=name,
-            config=config,
-            url=url,
-        )
-        collection = await self.client.create_collection(
-            name=bank_id,
-            metadata={"bank": bank.json()},
+        memory_bank: MemoryBankDef,
+    ) -> None:
+        assert (
+            memory_bank.type == MemoryBankType.vector.value
+        ), f"Only vector banks are supported {memory_bank.type}"
+
+        collection = await self.client.get_or_create_collection(
+            name=memory_bank.identifier,
+            metadata={"bank": memory_bank.json()},
         )
         bank_index = BankWithIndex(
-            bank=bank, index=ChromaIndex(self.client, collection)
+            bank=memory_bank, index=ChromaIndex(self.client, collection)
         )
-        self.cache[bank_id] = bank_index
-        return bank
-
-    async def get_memory_bank(self, bank_id: str) -> Optional[MemoryBank]:
-        bank_index = await self._get_and_cache_bank_index(bank_id)
-        if bank_index is None:
-            return None
-        return bank_index.bank
-
-    async def _get_and_cache_bank_index(self, bank_id: str) -> Optional[BankWithIndex]:
-        if bank_id in self.cache:
-            return self.cache[bank_id]
+        self.cache[memory_bank.identifier] = bank_index
 
+    async def list_memory_banks(self) -> List[MemoryBankDef]:
         collections = await self.client.list_collections()
         for collection in collections:
-            if collection.name == bank_id:
-                print(collection.metadata)
-                bank = MemoryBank(**json.loads(collection.metadata["bank"]))
-                index = BankWithIndex(
-                    bank=bank,
-                    index=ChromaIndex(self.client, collection),
-                )
-                self.cache[bank_id] = index
-                return index
+            try:
+                data = json.loads(collection.metadata["bank"])
+                bank = parse_obj_as(MemoryBankDef, data)
+            except Exception:
+                import traceback
 
-        return None
+                traceback.print_exc()
+                print(f"Failed to parse bank: {collection.metadata}")
+                continue
+
+            index = BankWithIndex(
+                bank=bank,
+                index=ChromaIndex(self.client, collection),
+            )
+            self.cache[bank.identifier] = index
+
+        return [i.bank for i in self.cache.values()]
 
     async def insert_documents(
         self,
@@ -150,7 +138,7 @@ class ChromaMemoryAdapter(Memory, RoutableProvider):
         documents: List[MemoryBankDocument],
         ttl_seconds: Optional[int] = None,
     ) -> None:
-        index = await self._get_and_cache_bank_index(bank_id)
+        index = self.cache.get(bank_id, None)
         if not index:
             raise ValueError(f"Bank {bank_id} not found")
 
@@ -162,7 +150,7 @@ class ChromaMemoryAdapter(Memory, RoutableProvider):
         query: InterleavedTextMedia,
         params: Optional[Dict[str, Any]] = None,
     ) -> QueryDocumentsResponse:
-        index = await self._get_and_cache_bank_index(bank_id)
+        index = self.cache.get(bank_id, None)
         if not index:
             raise ValueError(f"Bank {bank_id} not found")
 

llama_stack/providers/adapters/memory/pgvector/pgvector.py

@@ -4,7 +4,6 @@
 # This source code is licensed under the terms described in the LICENSE file in
 # the root directory of this source tree.
 
-import uuid
 from typing import List, Tuple
 
 import psycopg2
@@ -12,11 +11,11 @@ from numpy.typing import NDArray
 from psycopg2 import sql
 from psycopg2.extras import execute_values, Json
 
-from pydantic import BaseModel
+from pydantic import BaseModel, parse_obj_as
 
 from llama_stack.apis.memory import *  # noqa: F403
-from llama_stack.distribution.datatypes import RoutableProvider
 
+from llama_stack.providers.datatypes import MemoryBanksProtocolPrivate
 from llama_stack.providers.utils.memory.vector_store import (
     ALL_MINILM_L6_V2_DIMENSION,
     BankWithIndex,
@@ -46,23 +45,17 @@ def upsert_models(cur, keys_models: List[Tuple[str, BaseModel]]):
     execute_values(cur, query, values, template="(%s, %s)")
 
 
-def load_models(cur, keys: List[str], cls):
+def load_models(cur, cls):
     query = "SELECT key, data FROM metadata_store"
-    if keys:
-        placeholders = ",".join(["%s"] * len(keys))
-        query += f" WHERE key IN ({placeholders})"
-        cur.execute(query, keys)
-    else:
-        cur.execute(query)
-
+    cur.execute(query)
     rows = cur.fetchall()
-    return [cls(**row["data"]) for row in rows]
+    return [parse_obj_as(cls, row["data"]) for row in rows]
 
 
 class PGVectorIndex(EmbeddingIndex):
-    def __init__(self, bank: MemoryBank, dimension: int, cursor):
+    def __init__(self, bank: MemoryBankDef, dimension: int, cursor):
         self.cursor = cursor
-        self.table_name = f"vector_store_{bank.name}"
+        self.table_name = f"vector_store_{bank.identifier}"
 
         self.cursor.execute(
             f"""
@@ -119,7 +112,7 @@ class PGVectorIndex(EmbeddingIndex):
         return QueryDocumentsResponse(chunks=chunks, scores=scores)
 
 
-class PGVectorMemoryAdapter(Memory, RoutableProvider):
+class PGVectorMemoryAdapter(Memory, MemoryBanksProtocolPrivate):
     def __init__(self, config: PGVectorConfig) -> None:
         print(f"Initializing PGVectorMemoryAdapter -> {config.host}:{config.port}")
         self.config = config
@@ -161,57 +154,37 @@ class PGVectorMemoryAdapter(Memory, RoutableProvider):
     async def shutdown(self) -> None:
         pass
 
-    async def validate_routing_keys(self, routing_keys: List[str]) -> None:
-        print(f"[pgvector] Registering memory bank routing keys: {routing_keys}")
-        pass
-
-    async def create_memory_bank(
+    async def register_memory_bank(
         self,
-        name: str,
-        config: MemoryBankConfig,
-        url: Optional[URL] = None,
-    ) -> MemoryBank:
-        bank_id = str(uuid.uuid4())
-        bank = MemoryBank(
-            bank_id=bank_id,
-            name=name,
-            config=config,
-            url=url,
-        )
+        memory_bank: MemoryBankDef,
+    ) -> None:
+        assert (
+            memory_bank.type == MemoryBankType.vector.value
+        ), f"Only vector banks are supported {memory_bank.type}"
+
         upsert_models(
             self.cursor,
             [
-                (bank.bank_id, bank),
+                (memory_bank.identifier, memory_bank),
             ],
         )
+
         index = BankWithIndex(
-            bank=bank,
-            index=PGVectorIndex(bank, ALL_MINILM_L6_V2_DIMENSION, self.cursor),
+            bank=memory_bank,
+            index=PGVectorIndex(memory_bank, ALL_MINILM_L6_V2_DIMENSION, self.cursor),
         )
-        self.cache[bank_id] = index
-        return bank
+        self.cache[memory_bank.identifier] = index
 
-    async def get_memory_bank(self, bank_id: str) -> Optional[MemoryBank]:
-        bank_index = await self._get_and_cache_bank_index(bank_id)
-        if bank_index is None:
-            return None
-        return bank_index.bank
-
-    async def _get_and_cache_bank_index(self, bank_id: str) -> Optional[BankWithIndex]:
-        if bank_id in self.cache:
-            return self.cache[bank_id]
-
-        banks = load_models(self.cursor, [bank_id], MemoryBank)
-        if not banks:
-            return None
-
-        bank = banks[0]
-        index = BankWithIndex(
-            bank=bank,
-            index=PGVectorIndex(bank, ALL_MINILM_L6_V2_DIMENSION, self.cursor),
-        )
-        self.cache[bank_id] = index
-        return index
+    async def list_memory_banks(self) -> List[MemoryBankDef]:
+        banks = load_models(self.cursor, MemoryBankDef)
+        for bank in banks:
+            if bank.identifier not in self.cache:
+                index = BankWithIndex(
+                    bank=bank,
+                    index=PGVectorIndex(bank, ALL_MINILM_L6_V2_DIMENSION, self.cursor),
+                )
+                self.cache[bank.identifier] = index
+        return banks
 
     async def insert_documents(
         self,
@@ -219,7 +192,7 @@ class PGVectorMemoryAdapter(Memory, RoutableProvider):
         documents: List[MemoryBankDocument],
         ttl_seconds: Optional[int] = None,
     ) -> None:
-        index = await self._get_and_cache_bank_index(bank_id)
+        index = self.cache.get(bank_id, None)
         if not index:
             raise ValueError(f"Bank {bank_id} not found")
 
@@ -231,7 +204,7 @@ class PGVectorMemoryAdapter(Memory, RoutableProvider):
         query: InterleavedTextMedia,
         params: Optional[Dict[str, Any]] = None,
     ) -> QueryDocumentsResponse:
-        index = await self._get_and_cache_bank_index(bank_id)
+        index = self.cache.get(bank_id, None)
         if not index:
             raise ValueError(f"Bank {bank_id} not found")
 

llama_stack/providers/adapters/memory/sample/sample.py

@@ -9,14 +9,12 @@ from .config import SampleConfig
 
 from llama_stack.apis.memory import *  # noqa: F403
 
-from llama_stack.distribution.datatypes import RoutableProvider
 
-
-class SampleMemoryImpl(Memory, RoutableProvider):
+class SampleMemoryImpl(Memory):
     def __init__(self, config: SampleConfig):
         self.config = config
 
-    async def validate_routing_keys(self, routing_keys: list[str]) -> None:
+    async def register_memory_bank(self, memory_bank: MemoryBankDef) -> None:
         # these are the memory banks the Llama Stack will use to route requests to this provider
         # perform validation here if necessary
         pass

llama_stack/providers/adapters/memory/weaviate/__init__.py

@@ -1,8 +1,15 @@
-from .config import WeaviateConfig
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the terms described in the LICENSE file in
+# the root directory of this source tree.
+
+from .config import WeaviateConfig, WeaviateRequestProviderData  # noqa: F401
+
 
 async def get_adapter_impl(config: WeaviateConfig, _deps):
     from .weaviate import WeaviateMemoryAdapter
 
     impl = WeaviateMemoryAdapter(config)
     await impl.initialize()
-    return impl
+    return impl

llama_stack/providers/adapters/memory/weaviate/config.py

@@ -4,15 +4,13 @@
 # This source code is licensed under the terms described in the LICENSE file in
 # the root directory of this source tree.
 
-from llama_models.schema_utils import json_schema_type
-from pydantic import BaseModel, Field
+from pydantic import BaseModel
+
 
 class WeaviateRequestProviderData(BaseModel):
-    # if there _is_ provider data, it must specify the API KEY
-    # if you want it to be optional, use Optional[str]
     weaviate_api_key: str
     weaviate_cluster_url: str
 
-@json_schema_type
+
 class WeaviateConfig(BaseModel):
-    collection: str = Field(default="MemoryBank")
+    pass

llama_stack/providers/adapters/memory/weaviate/weaviate.py

@@ -1,14 +1,20 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the terms described in the LICENSE file in
+# the root directory of this source tree.
 import json
-import uuid
-from typing import List, Optional, Dict, Any
-from numpy.typing import NDArray
+
+from typing import Any, Dict, List, Optional
 
 import weaviate
 import weaviate.classes as wvc
+from numpy.typing import NDArray
 from weaviate.classes.init import Auth
 
-from llama_stack.apis.memory import *
-from llama_stack.distribution.request_headers import get_request_provider_data
+from llama_stack.apis.memory import *  # noqa: F403
+from llama_stack.distribution.request_headers import NeedsRequestProviderData
+from llama_stack.providers.datatypes import MemoryBanksProtocolPrivate
 from llama_stack.providers.utils.memory.vector_store import (
     BankWithIndex,
     EmbeddingIndex,
@@ -16,162 +22,154 @@ from llama_stack.providers.utils.memory.vector_store import (
 
 from .config import WeaviateConfig, WeaviateRequestProviderData
 
+
 class WeaviateIndex(EmbeddingIndex):
-    def __init__(self, client: weaviate.Client, collection: str):
+    def __init__(self, client: weaviate.Client, collection_name: str):
         self.client = client
-        self.collection = collection
+        self.collection_name = collection_name
 
     async def add_chunks(self, chunks: List[Chunk], embeddings: NDArray):
-        assert len(chunks) == len(embeddings), f"Chunk length {len(chunks)} does not match embedding length {len(embeddings)}"
+        assert len(chunks) == len(
+            embeddings
+        ), f"Chunk length {len(chunks)} does not match embedding length {len(embeddings)}"
 
         data_objects = []
         for i, chunk in enumerate(chunks):
-            
-            data_objects.append(wvc.data.DataObject(
-                properties={
-                    "chunk_content": chunk,
-                },
-                vector = embeddings[i].tolist()
-            ))
+            data_objects.append(
+                wvc.data.DataObject(
+                    properties={
+                        "chunk_content": chunk.json(),
+                    },
+                    vector=embeddings[i].tolist(),
+                )
+            )
 
         # Inserting chunks into a prespecified Weaviate collection
-        assert self.collection is not None, "Collection name must be specified"
-        my_collection = self.client.collections.get(self.collection)
-        
-        await my_collection.data.insert_many(data_objects)
+        collection = self.client.collections.get(self.collection_name)
 
+        # TODO: make this async friendly
+        collection.data.insert_many(data_objects)
 
     async def query(self, embedding: NDArray, k: int) -> QueryDocumentsResponse:
-        assert self.collection is not None, "Collection name must be specified"
+        collection = self.client.collections.get(self.collection_name)
 
-        my_collection = self.client.collections.get(self.collection)
-        
-        results = my_collection.query.near_vector(
-            near_vector = embedding.tolist(),
-            limit = k,
-            return_meta_data = wvc.query.MetadataQuery(distance=True)
+        results = collection.query.near_vector(
+            near_vector=embedding.tolist(),
+            limit=k,
+            return_metadata=wvc.query.MetadataQuery(distance=True),
         )
 
         chunks = []
         scores = []
         for doc in results.objects:
+            chunk_json = doc.properties["chunk_content"]
             try:
-                chunk = doc.properties["chunk_content"]
-                chunks.append(chunk)
-                scores.append(1.0 / doc.metadata.distance)
-            
-            except Exception as e:
+                chunk_dict = json.loads(chunk_json)
+                chunk = Chunk(**chunk_dict)
+            except Exception:
                 import traceback
+
                 traceback.print_exc()
-                print(f"Failed to parse document: {e}")
+                print(f"Failed to parse document: {chunk_json}")
+                continue
+
+            chunks.append(chunk)
+            scores.append(1.0 / doc.metadata.distance)
 
         return QueryDocumentsResponse(chunks=chunks, scores=scores)
 
 
-class WeaviateMemoryAdapter(Memory):
+class WeaviateMemoryAdapter(
+    Memory, NeedsRequestProviderData, MemoryBanksProtocolPrivate
+):
     def __init__(self, config: WeaviateConfig) -> None:
         self.config = config
-        self.client = None
+        self.client_cache = {}
         self.cache = {}
 
     def _get_client(self) -> weaviate.Client:
-            request_provider_data = get_request_provider_data()
-            
-            if request_provider_data is not None:
-                assert isinstance(request_provider_data, WeaviateRequestProviderData)
-            
-            # Connect to Weaviate Cloud
-            return weaviate.connect_to_weaviate_cloud(
-                cluster_url = request_provider_data.weaviate_cluster_url,
-                auth_credentials = Auth.api_key(request_provider_data.weaviate_api_key),
-                )
+        provider_data = self.get_request_provider_data()
+        assert provider_data is not None, "Request provider data must be set"
+        assert isinstance(provider_data, WeaviateRequestProviderData)
+
+        key = f"{provider_data.weaviate_cluster_url}::{provider_data.weaviate_api_key}"
+        if key in self.client_cache:
+            return self.client_cache[key]
+
+        client = weaviate.connect_to_weaviate_cloud(
+            cluster_url=provider_data.weaviate_cluster_url,
+            auth_credentials=Auth.api_key(provider_data.weaviate_api_key),
+        )
+        self.client_cache[key] = client
+        return client
 
     async def initialize(self) -> None:
-        try:
-            self.client = self._get_client()
-
-            # Create collection if it doesn't exist
-            if not self.client.collections.exists(self.config.collection):
-               self.client.collections.create(
-                    name = self.config.collection,
-                    vectorizer_config = wvc.config.Configure.Vectorizer.none(),
-                    properties=[
-                        wvc.config.Property(
-                        name="chunk_content",
-                        data_type=wvc.config.DataType.TEXT,
-                        ),
-                    ]
-                )
-
-        except Exception as e:
-            import traceback
-            traceback.print_exc()
-            raise RuntimeError("Could not connect to Weaviate server") from e
+        pass
 
     async def shutdown(self) -> None:
-        self.client = self._get_client()
+        for client in self.client_cache.values():
+            client.close()
 
-        if self.client:
-            self.client.close()
-
-    async def create_memory_bank(
+    async def register_memory_bank(
         self,
-        name: str,
-        config: MemoryBankConfig,
-        url: Optional[URL] = None,
-    ) -> MemoryBank:
-        bank_id = str(uuid.uuid4())
-        bank = MemoryBank(
-            bank_id=bank_id,
-            name=name,
-            config=config,
-            url=url,
-        )
-        self.client = self._get_client()
-        
-        # Store the bank as a new collection in Weaviate
-        self.client.collections.create(
-            name=bank_id
-        )
+        memory_bank: MemoryBankDef,
+    ) -> None:
+        assert (
+            memory_bank.type == MemoryBankType.vector.value
+        ), f"Only vector banks are supported {memory_bank.type}"
+
+        client = self._get_client()
+
+        # Create collection if it doesn't exist
+        if not client.collections.exists(memory_bank.identifier):
+            client.collections.create(
+                name=memory_bank.identifier,
+                vectorizer_config=wvc.config.Configure.Vectorizer.none(),
+                properties=[
+                    wvc.config.Property(
+                        name="chunk_content",
+                        data_type=wvc.config.DataType.TEXT,
+                    ),
+                ],
+            )
 
         index = BankWithIndex(
-            bank=bank,
-            index=WeaviateIndex(cleint = self.client, collection = bank_id),
+            bank=memory_bank,
+            index=WeaviateIndex(client=client, collection_name=memory_bank.identifier),
         )
-        self.cache[bank_id] = index
-        return bank
+        self.cache[memory_bank.identifier] = index
 
-    async def get_memory_bank(self, bank_id: str) -> Optional[MemoryBank]:
-        bank_index = await self._get_and_cache_bank_index(bank_id)
-        if bank_index is None:
-            return None
-        return bank_index.bank
+    async def list_memory_banks(self) -> List[MemoryBankDef]:
+        # 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
+        # yet available. We need to figure out a way to make this work.
+        return [i.bank for i in self.cache.values()]
 
     async def _get_and_cache_bank_index(self, bank_id: str) -> Optional[BankWithIndex]:
-        
-        self.client = self._get_client()
-
         if bank_id in self.cache:
             return self.cache[bank_id]
 
-        collections = await self.client.collections.list_all().keys()
+        bank = await self.memory_bank_store.get_memory_bank(bank_id)
+        if not bank:
+            raise ValueError(f"Bank {bank_id} not found")
 
-        for collection in collections:
-            if collection == bank_id:
-                bank = MemoryBank(**json.loads(collection.metadata["bank"]))
-                index = BankWithIndex(
-                    bank=bank,
-                    index=WeaviateIndex(self.client, collection),
-                )
-                self.cache[bank_id] = index
-                return index
+        client = self._get_client()
+        if not client.collections.exists(bank_id):
+            raise ValueError(f"Collection with name `{bank_id}` not found")
 
-        return None
+        index = BankWithIndex(
+            bank=bank,
+            index=WeaviateIndex(client=client, collection_name=bank_id),
+        )
+        self.cache[bank_id] = index
+        return index
 
     async def insert_documents(
         self,
         bank_id: str,
         documents: List[MemoryBankDocument],
+        ttl_seconds: Optional[int] = None,
     ) -> None:
         index = await self._get_and_cache_bank_index(bank_id)
         if not index:
@@ -189,4 +187,4 @@ class WeaviateMemoryAdapter(Memory):
         if not index:
             raise ValueError(f"Bank {bank_id} not found")
 
-        return await index.query_documents(query, params)
+        return await index.query_documents(query, params)