[API Updates] Model / shield / memory-bank routing + agent persistence + support for private headers (#92)

This is yet another of those large PRs (hopefully we will have less and less of them as things mature fast). This one introduces substantial improvements and some simplifications to the stack.

Most important bits:

* Agents reference implementation now has support for session / turn persistence. The default implementation uses sqlite but there's also support for using Redis.

* We have re-architected the structure of the Stack APIs to allow for more flexible routing. The motivating use cases are:
  - routing model A to ollama and model B to a remote provider like Together
  - routing shield A to local impl while shield B to a remote provider like Bedrock
  - routing a vector memory bank to Weaviate while routing a keyvalue memory bank to Redis

* Support for provider specific parameters to be passed from the clients. A client can pass data using `x_llamastack_provider_data` parameter which can be type-checked and provided to the Adapter implementations.

llama_stack/providers/utils/kvstore/__init__.py

@@ -0,0 +1,7 @@
+# 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 .kvstore import *  # noqa: F401, F403

llama_stack/providers/utils/kvstore/api.py

@@ -0,0 +1,21 @@
+# 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 datetime import datetime
+from typing import List, Optional, Protocol
+
+
+class KVStore(Protocol):
+    # TODO: make the value type bytes instead of str
+    async def set(
+        self, key: str, value: str, expiration: Optional[datetime] = None
+    ) -> None: ...
+
+    async def get(self, key: str) -> Optional[str]: ...
+
+    async def delete(self, key: str) -> None: ...
+
+    async def range(self, start_key: str, end_key: str) -> List[str]: ...

llama_stack/providers/utils/kvstore/config.py

@@ -0,0 +1,55 @@
+# 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 enum import Enum
+from typing import Literal, Optional, Union
+
+from pydantic import BaseModel, Field
+from typing_extensions import Annotated
+
+from llama_stack.distribution.utils.config_dirs import RUNTIME_BASE_DIR
+
+
+class KVStoreType(Enum):
+    redis = "redis"
+    sqlite = "sqlite"
+    postgres = "postgres"
+
+
+class CommonConfig(BaseModel):
+    namespace: Optional[str] = Field(
+        default=None,
+        description="All keys will be prefixed with this namespace",
+    )
+
+
+class RedisKVStoreConfig(CommonConfig):
+    type: Literal[KVStoreType.redis.value] = KVStoreType.redis.value
+    host: str = "localhost"
+    port: int = 6379
+
+
+class SqliteKVStoreConfig(CommonConfig):
+    type: Literal[KVStoreType.sqlite.value] = KVStoreType.sqlite.value
+    db_path: str = Field(
+        default=(RUNTIME_BASE_DIR / "kvstore.db").as_posix(),
+        description="File path for the sqlite database",
+    )
+
+
+class PostgresKVStoreConfig(CommonConfig):
+    type: Literal[KVStoreType.postgres.value] = KVStoreType.postgres.value
+    host: str = "localhost"
+    port: int = 5432
+    db: str = "llamastack"
+    user: str
+    password: Optional[str] = None
+
+
+KVStoreConfig = Annotated[
+    Union[RedisKVStoreConfig, SqliteKVStoreConfig, PostgresKVStoreConfig],
+    Field(discriminator="type", default=KVStoreType.sqlite.value),
+]

llama_stack/providers/utils/kvstore/kvstore.py

@@ -0,0 +1,51 @@
+# 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 .api import *  # noqa: F403
+from .config import *  # noqa: F403
+
+
+def kvstore_dependencies():
+    return ["aiosqlite", "psycopg2-binary", "redis"]
+
+
+class InmemoryKVStoreImpl(KVStore):
+    def __init__(self):
+        self._store = {}
+
+    async def initialize(self) -> None:
+        pass
+
+    async def get(self, key: str) -> Optional[str]:
+        return self._store.get(key)
+
+    async def set(self, key: str, value: str) -> None:
+        self._store[key] = value
+
+    async def range(self, start_key: str, end_key: str) -> List[str]:
+        return [
+            self._store[key]
+            for key in self._store.keys()
+            if key >= start_key and key < end_key
+        ]
+
+
+async def kvstore_impl(config: KVStoreConfig) -> KVStore:
+    if config.type == KVStoreType.redis.value:
+        from .redis import RedisKVStoreImpl
+
+        impl = RedisKVStoreImpl(config)
+    elif config.type == KVStoreType.sqlite.value:
+        from .sqlite import SqliteKVStoreImpl
+
+        impl = SqliteKVStoreImpl(config)
+    elif config.type == KVStoreType.postgres.value:
+        raise NotImplementedError()
+    else:
+        raise ValueError(f"Unknown kvstore type {config.type}")
+
+    await impl.initialize()
+    return impl

llama_stack/providers/utils/kvstore/redis/__init__.py

@@ -0,0 +1,7 @@
+# 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 .redis import RedisKVStoreImpl  # noqa: F401

llama_stack/providers/utils/kvstore/redis/redis.py

@@ -0,0 +1,52 @@
+# 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 datetime import datetime
+from typing import List, Optional
+
+from redis.asyncio import Redis
+
+from ..api import *  # noqa: F403
+from ..config import RedisKVStoreConfig
+
+
+class RedisKVStoreImpl(KVStore):
+    def __init__(self, config: RedisKVStoreConfig):
+        self.config = config
+
+    async def initialize(self) -> None:
+        self.redis = Redis.from_url(self.config.url)
+
+    def _namespaced_key(self, key: str) -> str:
+        if not self.config.namespace:
+            return key
+        return f"{self.config.namespace}:{key}"
+
+    async def set(
+        self, key: str, value: str, expiration: Optional[datetime] = None
+    ) -> None:
+        key = self._namespaced_key(key)
+        await self.redis.set(key, value)
+        if expiration:
+            await self.redis.expireat(key, expiration)
+
+    async def get(self, key: str) -> Optional[str]:
+        key = self._namespaced_key(key)
+        value = await self.redis.get(key)
+        if value is None:
+            return None
+        ttl = await self.redis.ttl(key)
+        return value
+
+    async def delete(self, key: str) -> None:
+        key = self._namespaced_key(key)
+        await self.redis.delete(key)
+
+    async def range(self, start_key: str, end_key: str) -> List[str]:
+        start_key = self._namespaced_key(start_key)
+        end_key = self._namespaced_key(end_key)
+
+        return await self.redis.zrangebylex(start_key, end_key)

llama_stack/providers/utils/kvstore/sqlite/__init__.py

@@ -0,0 +1,7 @@
+# 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 .sqlite import SqliteKVStoreImpl  # noqa: F401

llama_stack/providers/utils/kvstore/sqlite/config.py

@@ -0,0 +1,19 @@
+# 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 llama_models.schema_utils import json_schema_type
+from pydantic import BaseModel, Field
+
+
+@json_schema_type
+class SqliteControlPlaneConfig(BaseModel):
+    db_path: str = Field(
+        description="File path for the sqlite database",
+    )
+    table_name: str = Field(
+        default="llamastack_control_plane",
+        description="Table into which all the keys will be placed",
+    )

llama_stack/providers/utils/kvstore/sqlite/sqlite.py

@@ -0,0 +1,73 @@
+# 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 os
+
+from datetime import datetime
+from typing import List, Optional
+
+import aiosqlite
+
+from ..api import *  # noqa: F403
+from ..config import SqliteKVStoreConfig
+
+
+class SqliteKVStoreImpl(KVStore):
+    def __init__(self, config: SqliteKVStoreConfig):
+        self.db_path = config.db_path
+        self.table_name = "kvstore"
+
+    async def initialize(self):
+        os.makedirs(os.path.dirname(self.db_path), exist_ok=True)
+        async with aiosqlite.connect(self.db_path) as db:
+            await db.execute(
+                f"""
+                CREATE TABLE IF NOT EXISTS {self.table_name} (
+                    key TEXT PRIMARY KEY,
+                    value TEXT,
+                    expiration TIMESTAMP
+                )
+            """
+            )
+            await db.commit()
+
+    async def set(
+        self, key: str, value: str, expiration: Optional[datetime] = None
+    ) -> None:
+        async with aiosqlite.connect(self.db_path) as db:
+            await db.execute(
+                f"INSERT OR REPLACE INTO {self.table_name} (key, value, expiration) VALUES (?, ?, ?)",
+                (key, value, expiration),
+            )
+            await db.commit()
+
+    async def get(self, key: str) -> Optional[str]:
+        async with aiosqlite.connect(self.db_path) as db:
+            async with db.execute(
+                f"SELECT value, expiration FROM {self.table_name} WHERE key = ?", (key,)
+            ) as cursor:
+                row = await cursor.fetchone()
+                if row is None:
+                    return None
+                value, expiration = row
+                return value
+
+    async def delete(self, key: str) -> None:
+        async with aiosqlite.connect(self.db_path) as db:
+            await db.execute(f"DELETE FROM {self.table_name} WHERE key = ?", (key,))
+            await db.commit()
+
+    async def range(self, start_key: str, end_key: str) -> List[str]:
+        async with aiosqlite.connect(self.db_path) as db:
+            async with db.execute(
+                f"SELECT key, value, expiration FROM {self.table_name} WHERE key >= ? AND key <= ?",
+                (start_key, end_key),
+            ) as cursor:
+                result = []
+                async for row in cursor:
+                    _, value, _ = row
+                    result.append(value)
+                return result

llama_stack/providers/utils/memory/vector_store.py

@@ -16,6 +16,7 @@ import httpx
 import numpy as np
 from numpy.typing import NDArray
 from pypdf import PdfReader
+from termcolor import cprint
 
 from llama_models.llama3.api.datatypes import *  # noqa: F403
 from llama_models.llama3.api.tokenizer import Tokenizer
@@ -160,6 +161,8 @@ class BankWithIndex:
                 self.bank.config.overlap_size_in_tokens
                 or (self.bank.config.chunk_size_in_tokens // 4),
             )
+            if not chunks:
+                continue
             embeddings = model.encode([x.content for x in chunks]).astype(np.float32)
 
             await self.index.add_chunks(chunks, embeddings)

llama_stack/providers/utils/telemetry/tracing.py

@@ -12,7 +12,7 @@ import threading
 import uuid
 from datetime import datetime
 from functools import wraps
-from typing import Any, Dict, List
+from typing import Any, Callable, Dict, List
 
 
 from llama_stack.apis.telemetry import *  # noqa: F403
@@ -196,33 +196,40 @@ class TelemetryHandler(logging.Handler):
         pass
 
 
-def span(name: str, attributes: Dict[str, Any] = None):
-    def decorator(func):
+class SpanContextManager:
+    def __init__(self, name: str, attributes: Dict[str, Any] = None):
+        self.name = name
+        self.attributes = attributes
+
+    def __enter__(self):
+        global CURRENT_TRACE_CONTEXT
+        context = CURRENT_TRACE_CONTEXT
+        if context:
+            context.push_span(self.name, self.attributes)
+        return self
+
+    def __exit__(self, exc_type, exc_value, traceback):
+        global CURRENT_TRACE_CONTEXT
+        context = CURRENT_TRACE_CONTEXT
+        if context:
+            context.pop_span()
+
+    async def __aenter__(self):
+        return self.__enter__()
+
+    async def __aexit__(self, exc_type, exc_value, traceback):
+        self.__exit__(exc_type, exc_value, traceback)
+
+    def __call__(self, func: Callable):
         @wraps(func)
         def sync_wrapper(*args, **kwargs):
-            try:
-                global CURRENT_TRACE_CONTEXT
-
-                context = CURRENT_TRACE_CONTEXT
-                if context:
-                    context.push_span(name, attributes)
-                result = func(*args, **kwargs)
-            finally:
-                context.pop_span()
-            return result
+            with self:
+                return func(*args, **kwargs)
 
         @wraps(func)
         async def async_wrapper(*args, **kwargs):
-            try:
-                global CURRENT_TRACE_CONTEXT
-
-                context = CURRENT_TRACE_CONTEXT
-                if context:
-                    context.push_span(name, attributes)
-                result = await func(*args, **kwargs)
-            finally:
-                context.pop_span()
-            return result
+            async with self:
+                return await func(*args, **kwargs)
 
         @wraps(func)
         def wrapper(*args, **kwargs):
@@ -233,4 +240,6 @@ def span(name: str, attributes: Dict[str, Any] = None):
 
         return wrapper
 
-    return decorator
+
+def span(name: str, attributes: Dict[str, Any] = None):
+    return SpanContextManager(name, attributes)