Merge branch 'main' into aidand-391-guided-decoding-vllm_3

llama_stack/apis/agents/agents.py

@@ -23,6 +23,7 @@ from llama_models.schema_utils import json_schema_type, webmethod
 from pydantic import BaseModel, ConfigDict, Field
 from typing_extensions import Annotated
 
+from llama_stack.providers.utils.telemetry.trace_protocol import trace_protocol
 from llama_models.llama3.api.datatypes import *  # noqa: F403
 from llama_stack.apis.common.deployment_types import *  # noqa: F403
 from llama_stack.apis.inference import *  # noqa: F403
@@ -418,6 +419,7 @@ class AgentStepResponse(BaseModel):
 
 
 @runtime_checkable
+@trace_protocol
 class Agents(Protocol):
     @webmethod(route="/agents/create")
     async def create_agent(

llama_stack/apis/datasetio/datasetio.py

@@ -37,3 +37,8 @@ class DatasetIO(Protocol):
         page_token: Optional[str] = None,
         filter_condition: Optional[str] = None,
     ) -> PaginatedRowsResult: ...
+
+    @webmethod(route="/datasetio/append-rows", method="POST")
+    async def append_rows(
+        self, dataset_id: str, rows: List[Dict[str, Any]]
+    ) -> None: ...

llama_stack/apis/datasets/client.py

@@ -78,6 +78,21 @@ class DatasetsClient(Datasets):
 
             return [DatasetDefWithProvider(**x) for x in response.json()]
 
+    async def unregister_dataset(
+        self,
+        dataset_id: str,
+    ) -> None:
+        async with httpx.AsyncClient() as client:
+            response = await client.delete(
+                f"{self.base_url}/datasets/unregister",
+                params={
+                    "dataset_id": dataset_id,
+                },
+                headers={"Content-Type": "application/json"},
+                timeout=60,
+            )
+            response.raise_for_status()
+
 
 async def run_main(host: str, port: int):
     client = DatasetsClient(f"http://{host}:{port}")

llama_stack/apis/datasets/datasets.py

@@ -64,3 +64,9 @@ class Datasets(Protocol):
 
     @webmethod(route="/datasets/list", method="GET")
     async def list_datasets(self) -> List[Dataset]: ...
+
+    @webmethod(route="/datasets/unregister", method="POST")
+    async def unregister_dataset(
+        self,
+        dataset_id: str,
+    ) -> None: ...

llama_stack/apis/inference/inference.py

@@ -21,6 +21,8 @@ from llama_models.schema_utils import json_schema_type, webmethod
 from pydantic import BaseModel, Field
 from typing_extensions import Annotated
 
+from llama_stack.providers.utils.telemetry.trace_protocol import trace_protocol
+
 from llama_models.llama3.api.datatypes import *  # noqa: F403
 from llama_stack.apis.models import *  # noqa: F403
 
@@ -220,6 +222,7 @@ class ModelStore(Protocol):
 
 
 @runtime_checkable
+@trace_protocol
 class Inference(Protocol):
     model_store: ModelStore
 

llama_stack/apis/memory/memory.py

@@ -16,6 +16,7 @@ from pydantic import BaseModel, Field
 
 from llama_models.llama3.api.datatypes import *  # noqa: F403
 from llama_stack.apis.memory_banks import *  # noqa: F403
+from llama_stack.providers.utils.telemetry.trace_protocol import trace_protocol
 
 
 @json_schema_type
@@ -43,6 +44,7 @@ class MemoryBankStore(Protocol):
 
 
 @runtime_checkable
+@trace_protocol
 class Memory(Protocol):
     memory_bank_store: MemoryBankStore
 

llama_stack/apis/memory_banks/memory_banks.py

@@ -20,6 +20,7 @@ from llama_models.schema_utils import json_schema_type, webmethod
 from pydantic import BaseModel, Field
 
 from llama_stack.apis.resource import Resource, ResourceType
+from llama_stack.providers.utils.telemetry.trace_protocol import trace_protocol
 
 
 @json_schema_type
@@ -129,6 +130,7 @@ class MemoryBankInput(BaseModel):
 
 
 @runtime_checkable
+@trace_protocol
 class MemoryBanks(Protocol):
     @webmethod(route="/memory-banks/list", method="GET")
     async def list_memory_banks(self) -> List[MemoryBank]: ...

llama_stack/apis/models/models.py

@@ -10,6 +10,7 @@ from llama_models.schema_utils import json_schema_type, webmethod
 from pydantic import BaseModel, ConfigDict, Field
 
 from llama_stack.apis.resource import Resource, ResourceType
+from llama_stack.providers.utils.telemetry.trace_protocol import trace_protocol
 
 
 class CommonModelFields(BaseModel):
@@ -43,6 +44,7 @@ class ModelInput(CommonModelFields):
 
 
 @runtime_checkable
+@trace_protocol
 class Models(Protocol):
     @webmethod(route="/models/list", method="GET")
     async def list_models(self) -> List[Model]: ...

llama_stack/apis/safety/client.py

@@ -17,6 +17,8 @@ from llama_models.llama3.api.datatypes import *  # noqa: F403
 from pydantic import BaseModel
 from termcolor import cprint
 
+from llama_stack.apis.version import LLAMA_STACK_API_VERSION
+
 from llama_stack.distribution.datatypes import RemoteProviderConfig
 
 from llama_stack.apis.safety import *  # noqa: F403
@@ -45,7 +47,7 @@ class SafetyClient(Safety):
     ) -> RunShieldResponse:
         async with httpx.AsyncClient() as client:
             response = await client.post(
-                f"{self.base_url}/safety/run_shield",
+                f"{self.base_url}/{LLAMA_STACK_API_VERSION}/safety/run-shield",
                 json=dict(
                     shield_id=shield_id,
                     messages=[encodable_dict(m) for m in messages],
@@ -91,7 +93,7 @@ async def run_main(host: str, port: int, image_path: str = None):
     ]:
         cprint(f"User>{message.content}", "green")
         response = await client.run_shield(
-            shield_id="llama_guard",
+            shield_id="meta-llama/Llama-Guard-3-1B",
             messages=[message],
         )
         print(response)

llama_stack/apis/safety/safety.py

@@ -10,6 +10,8 @@ from typing import Any, Dict, List, Protocol, runtime_checkable
 from llama_models.schema_utils import json_schema_type, webmethod
 from pydantic import BaseModel
 
+from llama_stack.providers.utils.telemetry.trace_protocol import trace_protocol
+
 from llama_models.llama3.api.datatypes import *  # noqa: F403
 from llama_stack.apis.shields import *  # noqa: F403
 
@@ -43,6 +45,7 @@ class ShieldStore(Protocol):
 
 
 @runtime_checkable
+@trace_protocol
 class Safety(Protocol):
     shield_store: ShieldStore
 

llama_stack/apis/shields/shields.py

@@ -10,6 +10,7 @@ from llama_models.schema_utils import json_schema_type, webmethod
 from pydantic import BaseModel
 
 from llama_stack.apis.resource import Resource, ResourceType
+from llama_stack.providers.utils.telemetry.trace_protocol import trace_protocol
 
 
 class CommonShieldFields(BaseModel):
@@ -38,6 +39,7 @@ class ShieldInput(CommonShieldFields):
 
 
 @runtime_checkable
+@trace_protocol
 class Shields(Protocol):
     @webmethod(route="/shields/list", method="GET")
     async def list_shields(self) -> List[Shield]: ...

llama_stack/apis/telemetry/telemetry.py

@@ -6,12 +6,24 @@
 
 from datetime import datetime
 from enum import Enum
-from typing import Any, Dict, Literal, Optional, Protocol, runtime_checkable, Union
+from typing import (
+    Any,
+    Dict,
+    List,
+    Literal,
+    Optional,
+    Protocol,
+    runtime_checkable,
+    Union,
+)
 
 from llama_models.schema_utils import json_schema_type, webmethod
 from pydantic import BaseModel, Field
 from typing_extensions import Annotated
 
+# Add this constant near the top of the file, after the imports
+DEFAULT_TTL_DAYS = 7
+
 
 @json_schema_type
 class SpanStatus(Enum):
@@ -29,6 +41,11 @@ class Span(BaseModel):
     end_time: Optional[datetime] = None
     attributes: Optional[Dict[str, Any]] = Field(default_factory=dict)
 
+    def set_attribute(self, key: str, value: Any):
+        if self.attributes is None:
+            self.attributes = {}
+        self.attributes[key] = value
+
 
 @json_schema_type
 class Trace(BaseModel):
@@ -123,10 +140,66 @@ Event = Annotated[
 ]
 
 
+@json_schema_type
+class EvalTrace(BaseModel):
+    session_id: str
+    step: str
+    input: str
+    output: str
+    expected_output: str
+
+
+@json_schema_type
+class SpanWithChildren(Span):
+    children: List["SpanWithChildren"] = Field(default_factory=list)
+    status: Optional[SpanStatus] = None
+
+
+@json_schema_type
+class QueryCondition(BaseModel):
+    key: str
+    op: Literal["eq", "ne", "gt", "lt"]
+    value: Any
+
+
 @runtime_checkable
 class Telemetry(Protocol):
-    @webmethod(route="/telemetry/log-event")
-    async def log_event(self, event: Event) -> None: ...
 
-    @webmethod(route="/telemetry/get-trace", method="GET")
-    async def get_trace(self, trace_id: str) -> Trace: ...
+    @webmethod(route="/telemetry/log-event")
+    async def log_event(
+        self, event: Event, ttl_seconds: int = DEFAULT_TTL_DAYS * 86400
+    ) -> None: ...
+
+    @webmethod(route="/telemetry/query-traces", method="POST")
+    async def query_traces(
+        self,
+        attribute_filters: Optional[List[QueryCondition]] = None,
+        limit: Optional[int] = 100,
+        offset: Optional[int] = 0,
+        order_by: Optional[List[str]] = None,
+    ) -> List[Trace]: ...
+
+    @webmethod(route="/telemetry/get-span-tree", method="POST")
+    async def get_span_tree(
+        self,
+        span_id: str,
+        attributes_to_return: Optional[List[str]] = None,
+        max_depth: Optional[int] = None,
+    ) -> SpanWithChildren: ...
+
+    @webmethod(route="/telemetry/query-spans", method="POST")
+    async def query_spans(
+        self,
+        attribute_filters: List[QueryCondition],
+        attributes_to_return: List[str],
+        max_depth: Optional[int] = None,
+    ) -> List[Span]: ...
+
+    @webmethod(route="/telemetry/save-spans-to-dataset", method="POST")
+    async def save_spans_to_dataset(
+        self,
+        attribute_filters: List[QueryCondition],
+        attributes_to_save: List[str],
+        dataset_id: str,
+        max_depth: Optional[int] = None,
+    ) -> None: ...

llama_stack/distribution/build.py

@@ -10,6 +10,7 @@ from typing import List
 
 import pkg_resources
 from pydantic import BaseModel
+from termcolor import cprint
 
 from llama_stack.distribution.utils.exec import run_with_pty
 
@@ -45,7 +46,7 @@ class ApiInput(BaseModel):
 
 
 def get_provider_dependencies(
-    config_providers: Dict[str, List[Provider]]
+    config_providers: Dict[str, List[Provider]],
 ) -> tuple[list[str], list[str]]:
     """Get normal and special dependencies from provider configuration."""
     all_providers = get_provider_registry()
@@ -90,11 +91,12 @@ def get_provider_dependencies(
 def print_pip_install_help(providers: Dict[str, List[Provider]]):
     normal_deps, special_deps = get_provider_dependencies(providers)
 
-    print(
-        f"Please install needed dependencies using the following commands:\n\n\tpip install {' '.join(normal_deps)}"
+    cprint(
+        f"Please install needed dependencies using the following commands:\n\npip install {' '.join(normal_deps)}",
+        "yellow",
     )
     for special_dep in special_deps:
-        log.info(f"\tpip install {special_dep}")
+        cprint(f"pip install {special_dep}", "yellow")
     print()
 
 

llama_stack/distribution/library_client.py

@@ -0,0 +1,272 @@
+# 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 asyncio
+import inspect
+import queue
+import threading
+from concurrent.futures import ThreadPoolExecutor
+from pathlib import Path
+from typing import Any, Generator, get_args, get_origin, Optional, TypeVar
+
+import yaml
+from llama_stack_client import AsyncLlamaStackClient, LlamaStackClient, NOT_GIVEN
+from pydantic import TypeAdapter
+from rich.console import Console
+
+from termcolor import cprint
+
+from llama_stack.distribution.build import print_pip_install_help
+from llama_stack.distribution.configure import parse_and_maybe_upgrade_config
+from llama_stack.distribution.resolver import ProviderRegistry
+from llama_stack.distribution.server.endpoints import get_all_api_endpoints
+from llama_stack.distribution.stack import (
+    construct_stack,
+    get_stack_run_config_from_template,
+    replace_env_vars,
+)
+
+T = TypeVar("T")
+
+
+def stream_across_asyncio_run_boundary(
+    async_gen_maker,
+    pool_executor: ThreadPoolExecutor,
+) -> Generator[T, None, None]:
+    result_queue = queue.Queue()
+    stop_event = threading.Event()
+
+    async def consumer():
+        # make sure we make the generator in the event loop context
+        gen = await async_gen_maker()
+        try:
+            async for item in gen:
+                result_queue.put(item)
+        except Exception as e:
+            print(f"Error in generator {e}")
+            result_queue.put(e)
+        except asyncio.CancelledError:
+            return
+        finally:
+            result_queue.put(StopIteration)
+            stop_event.set()
+
+    def run_async():
+        # Run our own loop to avoid double async generator cleanup which is done
+        # by asyncio.run()
+        loop = asyncio.new_event_loop()
+        asyncio.set_event_loop(loop)
+        try:
+            task = loop.create_task(consumer())
+            loop.run_until_complete(task)
+        finally:
+            # Handle pending tasks like a generator's athrow()
+            pending = asyncio.all_tasks(loop)
+            if pending:
+                loop.run_until_complete(
+                    asyncio.gather(*pending, return_exceptions=True)
+                )
+            loop.close()
+
+    future = pool_executor.submit(run_async)
+
+    try:
+        # yield results as they come in
+        while not stop_event.is_set() or not result_queue.empty():
+            try:
+                item = result_queue.get(timeout=0.1)
+                if item is StopIteration:
+                    break
+                if isinstance(item, Exception):
+                    raise item
+                yield item
+            except queue.Empty:
+                continue
+    finally:
+        future.result()
+
+
+class LlamaStackAsLibraryClient(LlamaStackClient):
+    def __init__(
+        self,
+        config_path_or_template_name: str,
+        custom_provider_registry: Optional[ProviderRegistry] = None,
+    ):
+        super().__init__()
+        self.async_client = AsyncLlamaStackAsLibraryClient(
+            config_path_or_template_name, custom_provider_registry
+        )
+        self.pool_executor = ThreadPoolExecutor(max_workers=4)
+
+    def initialize(self):
+        asyncio.run(self.async_client.initialize())
+
+    def get(self, *args, **kwargs):
+        if kwargs.get("stream"):
+            return stream_across_asyncio_run_boundary(
+                lambda: self.async_client.get(*args, **kwargs),
+                self.pool_executor,
+            )
+        else:
+            return asyncio.run(self.async_client.get(*args, **kwargs))
+
+    def post(self, *args, **kwargs):
+        if kwargs.get("stream"):
+            return stream_across_asyncio_run_boundary(
+                lambda: self.async_client.post(*args, **kwargs),
+                self.pool_executor,
+            )
+        else:
+            return asyncio.run(self.async_client.post(*args, **kwargs))
+
+
+class AsyncLlamaStackAsLibraryClient(AsyncLlamaStackClient):
+    def __init__(
+        self,
+        config_path_or_template_name: str,
+        custom_provider_registry: Optional[ProviderRegistry] = None,
+    ):
+        super().__init__()
+
+        if config_path_or_template_name.endswith(".yaml"):
+            config_path = Path(config_path_or_template_name)
+            if not config_path.exists():
+                raise ValueError(f"Config file {config_path} does not exist")
+            config_dict = replace_env_vars(yaml.safe_load(config_path.read_text()))
+            config = parse_and_maybe_upgrade_config(config_dict)
+        else:
+            # template
+            config = get_stack_run_config_from_template(config_path_or_template_name)
+
+        self.config_path_or_template_name = config_path_or_template_name
+        self.config = config
+        self.custom_provider_registry = custom_provider_registry
+
+    async def initialize(self):
+        try:
+            self.impls = await construct_stack(
+                self.config, self.custom_provider_registry
+            )
+        except ModuleNotFoundError as e:
+            cprint(
+                "Using llama-stack as a library requires installing dependencies depending on the template (providers) you choose.\n",
+                "yellow",
+            )
+            print_pip_install_help(self.config.providers)
+            raise e
+
+        console = Console()
+        console.print(f"Using config [blue]{self.config_path_or_template_name}[/blue]:")
+        console.print(yaml.dump(self.config.model_dump(), indent=2))
+
+        endpoints = get_all_api_endpoints()
+        endpoint_impls = {}
+        for api, api_endpoints in endpoints.items():
+            for endpoint in api_endpoints:
+                impl = self.impls[api]
+                func = getattr(impl, endpoint.name)
+                endpoint_impls[endpoint.route] = func
+
+        self.endpoint_impls = endpoint_impls
+
+    async def get(
+        self,
+        path: str,
+        *,
+        stream=False,
+        **kwargs,
+    ):
+        if not self.endpoint_impls:
+            raise ValueError("Client not initialized")
+
+        if stream:
+            return self._call_streaming(path, "GET")
+        else:
+            return await self._call_non_streaming(path, "GET")
+
+    async def post(
+        self,
+        path: str,
+        *,
+        body: dict = None,
+        stream=False,
+        **kwargs,
+    ):
+        if not self.endpoint_impls:
+            raise ValueError("Client not initialized")
+
+        if stream:
+            return self._call_streaming(path, "POST", body)
+        else:
+            return await self._call_non_streaming(path, "POST", body)
+
+    async def _call_non_streaming(self, path: str, method: str, body: dict = None):
+        func = self.endpoint_impls.get(path)
+        if not func:
+            raise ValueError(f"No endpoint found for {path}")
+
+        body = self._convert_body(path, body)
+        return await func(**body)
+
+    async def _call_streaming(self, path: str, method: str, body: dict = None):
+        func = self.endpoint_impls.get(path)
+        if not func:
+            raise ValueError(f"No endpoint found for {path}")
+
+        body = self._convert_body(path, body)
+        async for chunk in await func(**body):
+            yield chunk
+
+    def _convert_body(self, path: str, body: Optional[dict] = None) -> dict:
+        if not body:
+            return {}
+
+        func = self.endpoint_impls[path]
+        sig = inspect.signature(func)
+
+        # Strip NOT_GIVENs to use the defaults in signature
+        body = {k: v for k, v in body.items() if v is not NOT_GIVEN}
+
+        # Convert parameters to Pydantic models where needed
+        converted_body = {}
+        for param_name, param in sig.parameters.items():
+            if param_name in body:
+                value = body.get(param_name)
+                converted_body[param_name] = self._convert_param(
+                    param.annotation, value
+                )
+        return converted_body
+
+    def _convert_param(self, annotation: Any, value: Any) -> Any:
+        if isinstance(annotation, type) and annotation in {str, int, float, bool}:
+            return value
+
+        origin = get_origin(annotation)
+        if origin is list:
+            item_type = get_args(annotation)[0]
+            try:
+                return [self._convert_param(item_type, item) for item in value]
+            except Exception:
+                print(f"Error converting list {value}")
+                return value
+
+        elif origin is dict:
+            key_type, val_type = get_args(annotation)
+            try:
+                return {k: self._convert_param(val_type, v) for k, v in value.items()}
+            except Exception:
+                print(f"Error converting dict {value}")
+                return value
+
+        try:
+            # Handle Pydantic models and discriminated unions
+            return TypeAdapter(annotation).validate_python(value)
+        except Exception as e:
+            cprint(
+                f"Warning: direct client failed to convert parameter {value} into {annotation}: {e}",
+                "yellow",
+            )
+            return value

llama_stack/distribution/request_headers.py

@@ -35,7 +35,7 @@ class NeedsRequestProviderData:
             provider_data = validator(**val)
             return provider_data
         except Exception as e:
-            log.error("Error parsing provider data", e)
+            log.error(f"Error parsing provider data: {e}")
 
 
 def set_request_provider_data(headers: Dict[str, str]):

llama_stack/distribution/routers/routers.py

@@ -222,6 +222,12 @@ class DatasetIORouter(DatasetIO):
             filter_condition=filter_condition,
         )
 
+    async def append_rows(self, dataset_id: str, rows: List[Dict[str, Any]]) -> None:
+        return await self.routing_table.get_provider_impl(dataset_id).append_rows(
+            dataset_id=dataset_id,
+            rows=rows,
+        )
+
 
 class ScoringRouter(Scoring):
     def __init__(

llama_stack/distribution/routers/routing_tables.py

@@ -57,6 +57,8 @@ async def unregister_object_from_provider(obj: RoutableObject, p: Any) -> None:
         return await p.unregister_memory_bank(obj.identifier)
     elif api == Api.inference:
         return await p.unregister_model(obj.identifier)
+    elif api == Api.datasetio:
+        return await p.unregister_dataset(obj.identifier)
     else:
         raise ValueError(f"Unregister not supported for {api}")
 
@@ -354,6 +356,12 @@ class DatasetsRoutingTable(CommonRoutingTableImpl, Datasets):
         )
         await self.register_object(dataset)
 
+    async def unregister_dataset(self, dataset_id: str) -> None:
+        dataset = await self.get_dataset(dataset_id)
+        if dataset is None:
+            raise ValueError(f"Dataset {dataset_id} not found")
+        await self.unregister_object(dataset)
+
 
 class ScoringFunctionsRoutingTable(CommonRoutingTableImpl, ScoringFunctions):
     async def list_scoring_functions(self) -> List[ScoringFn]:

llama_stack/distribution/server/server.py

@@ -43,9 +43,9 @@ from llama_stack.distribution.stack import (
     replace_env_vars,
     validate_env_pair,
 )
-from llama_stack.providers.inline.meta_reference.telemetry.console import (
-    ConsoleConfig,
-    ConsoleTelemetryImpl,
+from llama_stack.providers.inline.telemetry.meta_reference import (
+    TelemetryAdapter,
+    TelemetryConfig,
 )
 
 from .endpoints import get_all_api_endpoints
@@ -217,7 +217,7 @@ class TracingMiddleware:
 
     async def __call__(self, scope, receive, send):
         path = scope["path"]
-        await start_trace(path, {"location": "server"})
+        await start_trace(path, {"__location__": "server"})
         try:
             return await self.app(scope, receive, send)
         finally:
@@ -290,7 +290,7 @@ def main():
     if Api.telemetry in impls:
         setup_logger(impls[Api.telemetry])
     else:
-        setup_logger(ConsoleTelemetryImpl(ConsoleConfig()))
+        setup_logger(TelemetryAdapter(TelemetryConfig()))
 
     all_endpoints = get_all_api_endpoints()
 

llama_stack/distribution/tests/library_client_test.py

@@ -0,0 +1,103 @@
+# 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 argparse
+import os
+
+from llama_stack.distribution.library_client import LlamaStackAsLibraryClient
+from llama_stack_client.lib.agents.agent import Agent
+from llama_stack_client.lib.agents.event_logger import EventLogger as AgentEventLogger
+from llama_stack_client.lib.inference.event_logger import EventLogger
+from llama_stack_client.types import UserMessage
+from llama_stack_client.types.agent_create_params import AgentConfig
+
+
+def main(config_path: str):
+    client = LlamaStackAsLibraryClient(config_path)
+    client.initialize()
+
+    models = client.models.list()
+    print("\nModels:")
+    for model in models:
+        print(model)
+
+    if not models:
+        print("No models found, skipping chat completion test")
+        return
+
+    model_id = models[0].identifier
+    response = client.inference.chat_completion(
+        messages=[UserMessage(content="What is the capital of France?", role="user")],
+        model_id=model_id,
+        stream=False,
+    )
+    print("\nChat completion response (non-stream):")
+    print(response)
+
+    response = client.inference.chat_completion(
+        messages=[UserMessage(content="What is the capital of France?", role="user")],
+        model_id=model_id,
+        stream=True,
+    )
+
+    print("\nChat completion response (stream):")
+    for log in EventLogger().log(response):
+        log.print()
+
+    print("\nAgent test:")
+    agent_config = AgentConfig(
+        model=model_id,
+        instructions="You are a helpful assistant",
+        sampling_params={
+            "strategy": "greedy",
+            "temperature": 1.0,
+            "top_p": 0.9,
+        },
+        tools=(
+            [
+                {
+                    "type": "brave_search",
+                    "engine": "brave",
+                    "api_key": os.getenv("BRAVE_SEARCH_API_KEY"),
+                }
+            ]
+            if os.getenv("BRAVE_SEARCH_API_KEY")
+            else []
+        ),
+        tool_choice="auto",
+        tool_prompt_format="json",
+        input_shields=[],
+        output_shields=[],
+        enable_session_persistence=False,
+    )
+    agent = Agent(client, agent_config)
+    user_prompts = [
+        "Hello",
+        "Which players played in the winning team of the NBA western conference semifinals of 2024, please use tools",
+    ]
+
+    session_id = agent.create_session("test-session")
+
+    for prompt in user_prompts:
+        response = agent.create_turn(
+            messages=[
+                {
+                    "role": "user",
+                    "content": prompt,
+                }
+            ],
+            session_id=session_id,
+        )
+
+        for log in AgentEventLogger().log(response):
+            log.print()
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    parser.add_argument("config_path", help="Path to the config YAML file")
+    args = parser.parse_args()
+    main(args.config_path)

llama_stack/distribution/ui/README.md

@@ -0,0 +1,42 @@
+# (Experimental) LLama Stack UI
+
+## Docker Setup
+
+:warning: This is a work in progress.
+
+## Developer Setup
+
+1. Start up Llama Stack API server. More details [here](https://llama-stack.readthedocs.io/en/latest/getting_started/index.html).
+
+```
+llama stack build --template together --image-type conda
+
+llama stack run together
+```
+
+2. (Optional) Register datasets and eval tasks as resources. If you want to run pre-configured evaluation flows (e.g. Evaluations (Generation + Scoring) Page).
+
+```bash
+$ llama-stack-client datasets register \
+--dataset-id "mmlu" \
+--provider-id "huggingface" \
+--url "https://huggingface.co/datasets/llamastack/evals" \
+--metadata '{"path": "llamastack/evals", "name": "evals__mmlu__details", "split": "train"}' \
+--schema '{"input_query": {"type": "string"}, "expected_answer": {"type": "string", "chat_completion_input": {"type": "string"}}}'
+```
+
+```bash
+$ llama-stack-client eval_tasks register \
+--eval-task-id meta-reference-mmlu \
+--provider-id meta-reference \
+--dataset-id mmlu \
+--scoring-functions basic::regex_parser_multiple_choice_answer
+```
+
+3. Start Streamlit UI
+
+```bash
+cd llama_stack/distribution/ui
+pip install -r requirements.txt
+streamlit run app.py
+```

llama_stack/distribution/ui/app.py

@@ -0,0 +1,57 @@
+# 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 streamlit as st
+
+
+def main():
+    # Evaluation pages
+    application_evaluation_page = st.Page(
+        "page/evaluations/app_eval.py",
+        title="Evaluations (Scoring)",
+        icon="📊",
+        default=False,
+    )
+    native_evaluation_page = st.Page(
+        "page/evaluations/native_eval.py",
+        title="Evaluations (Generation + Scoring)",
+        icon="📊",
+        default=False,
+    )
+
+    # Playground pages
+    chat_page = st.Page(
+        "page/playground/chat.py", title="Chat", icon="💬", default=True
+    )
+    rag_page = st.Page("page/playground/rag.py", title="RAG", icon="💬", default=False)
+
+    # Distribution pages
+    resources_page = st.Page(
+        "page/distribution/resources.py", title="Resources", icon="🔍", default=False
+    )
+    provider_page = st.Page(
+        "page/distribution/providers.py",
+        title="API Providers",
+        icon="🔍",
+        default=False,
+    )
+
+    pg = st.navigation(
+        {
+            "Playground": [
+                chat_page,
+                rag_page,
+                application_evaluation_page,
+                native_evaluation_page,
+            ],
+            "Inspect": [provider_page, resources_page],
+        },
+        expanded=False,
+    )
+    pg.run()
+
+
+if __name__ == "__main__":
+    main()

llama_stack/distribution/ui/modules/__init__.py

@@ -0,0 +1,5 @@
+# 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.

llama_stack/distribution/ui/modules/api.py

@@ -0,0 +1,36 @@
+# 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 typing import Optional
+
+from llama_stack_client import LlamaStackClient
+
+
+class LlamaStackApi:
+    def __init__(self):
+        self.client = LlamaStackClient(
+            base_url=os.environ.get("LLAMA_STACK_ENDPOINT", "http://localhost:5000"),
+            provider_data={
+                "fireworks_api_key": os.environ.get("FIREWORKS_API_KEY", ""),
+                "together_api_key": os.environ.get("TOGETHER_API_KEY", ""),
+                "openai_api_key": os.environ.get("OPENAI_API_KEY", ""),
+            },
+        )
+
+    def run_scoring(
+        self, row, scoring_function_ids: list[str], scoring_params: Optional[dict]
+    ):
+        """Run scoring on a single row"""
+        if not scoring_params:
+            scoring_params = {fn_id: None for fn_id in scoring_function_ids}
+        return self.client.scoring.score(
+            input_rows=[row], scoring_functions=scoring_params
+        )
+
+
+llama_stack_api = LlamaStackApi()

llama_stack/distribution/ui/modules/utils.py

@@ -0,0 +1,42 @@
+# 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 base64
+import os
+
+import pandas as pd
+import streamlit as st
+
+
+def process_dataset(file):
+    if file is None:
+        return "No file uploaded", None
+
+    try:
+        # Determine file type and read accordingly
+        file_ext = os.path.splitext(file.name)[1].lower()
+        if file_ext == ".csv":
+            df = pd.read_csv(file)
+        elif file_ext in [".xlsx", ".xls"]:
+            df = pd.read_excel(file)
+        else:
+            return "Unsupported file format. Please upload a CSV or Excel file.", None
+
+        return df
+
+    except Exception as e:
+        st.error(f"Error processing file: {str(e)}")
+        return None
+
+
+def data_url_from_file(file) -> str:
+    file_content = file.getvalue()
+    base64_content = base64.b64encode(file_content).decode("utf-8")
+    mime_type = file.type
+
+    data_url = f"data:{mime_type};base64,{base64_content}"
+
+    return data_url

llama_stack/distribution/ui/page/__init__.py

@@ -0,0 +1,5 @@
+# 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.

llama_stack/distribution/ui/page/distribution/datasets.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.
+
+import streamlit as st
+from modules.api import llama_stack_api
+
+
+def datasets():
+    st.header("Datasets")
+
+    datasets_info = {
+        d.identifier: d.to_dict() for d in llama_stack_api.client.datasets.list()
+    }
+
+    selected_dataset = st.selectbox("Select a dataset", list(datasets_info.keys()))
+    st.json(datasets_info[selected_dataset], expanded=True)

llama_stack/distribution/ui/page/distribution/eval_tasks.py

@@ -0,0 +1,22 @@
+# 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 streamlit as st
+from modules.api import llama_stack_api
+
+
+def eval_tasks():
+    # Eval Tasks Section
+    st.header("Eval Tasks")
+
+    eval_tasks_info = {
+        d.identifier: d.to_dict() for d in llama_stack_api.client.eval_tasks.list()
+    }
+
+    selected_eval_task = st.selectbox(
+        "Select an eval task", list(eval_tasks_info.keys()), key="eval_task_inspect"
+    )
+    st.json(eval_tasks_info[selected_eval_task], expanded=True)

llama_stack/distribution/ui/page/distribution/memory_banks.py

@@ -0,0 +1,23 @@
+# 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 streamlit as st
+from modules.api import llama_stack_api
+
+
+def memory_banks():
+    st.header("Memory Banks")
+    memory_banks_info = {
+        m.identifier: m.to_dict() for m in llama_stack_api.client.memory_banks.list()
+    }
+
+    if len(memory_banks_info) > 0:
+        selected_memory_bank = st.selectbox(
+            "Select a memory bank", list(memory_banks_info.keys())
+        )
+        st.json(memory_banks_info[selected_memory_bank])
+    else:
+        st.info("No memory banks found")

llama_stack/distribution/ui/page/distribution/models.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.
+
+import streamlit as st
+from modules.api import llama_stack_api
+
+
+def models():
+    # Models Section
+    st.header("Models")
+    models_info = {
+        m.identifier: m.to_dict() for m in llama_stack_api.client.models.list()
+    }
+
+    selected_model = st.selectbox("Select a model", list(models_info.keys()))
+    st.json(models_info[selected_model])

llama_stack/distribution/ui/page/distribution/providers.py

@@ -0,0 +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 streamlit as st
+from modules.api import llama_stack_api
+
+
+def providers():
+    st.header("🔍 API Providers")
+    apis_providers_info = llama_stack_api.client.providers.list()
+    # selected_api = st.selectbox("Select an API", list(apis_providers_info.keys()))
+    for api in apis_providers_info.keys():
+        st.markdown(f"###### {api}")
+        st.dataframe([p.to_dict() for p in apis_providers_info[api]], width=500)
+
+
+providers()

llama_stack/distribution/ui/page/distribution/resources.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 page.distribution.datasets import datasets
+from page.distribution.eval_tasks import eval_tasks
+from page.distribution.memory_banks import memory_banks
+from page.distribution.models import models
+from page.distribution.scoring_functions import scoring_functions
+from page.distribution.shields import shields
+
+from streamlit_option_menu import option_menu
+
+
+def resources_page():
+    options = [
+        "Models",
+        "Memory Banks",
+        "Shields",
+        "Scoring Functions",
+        "Datasets",
+        "Eval Tasks",
+    ]
+    icons = ["magic", "memory", "shield", "file-bar-graph", "database", "list-task"]
+    selected_resource = option_menu(
+        None,
+        options,
+        icons=icons,
+        orientation="horizontal",
+        styles={
+            "nav-link": {
+                "font-size": "12px",
+            },
+        },
+    )
+    if selected_resource == "Eval Tasks":
+        eval_tasks()
+    elif selected_resource == "Memory Banks":
+        memory_banks()
+    elif selected_resource == "Datasets":
+        datasets()
+    elif selected_resource == "Models":
+        models()
+    elif selected_resource == "Scoring Functions":
+        scoring_functions()
+    elif selected_resource == "Shields":
+        shields()
+
+
+resources_page()

llama_stack/distribution/ui/page/distribution/scoring_functions.py

@@ -0,0 +1,22 @@
+# 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 streamlit as st
+from modules.api import llama_stack_api
+
+
+def scoring_functions():
+    st.header("Scoring Functions")
+
+    scoring_functions_info = {
+        s.identifier: s.to_dict()
+        for s in llama_stack_api.client.scoring_functions.list()
+    }
+
+    selected_scoring_function = st.selectbox(
+        "Select a scoring function", list(scoring_functions_info.keys())
+    )
+    st.json(scoring_functions_info[selected_scoring_function], expanded=True)

llama_stack/distribution/ui/page/distribution/shields.py

@@ -0,0 +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 streamlit as st
+from modules.api import llama_stack_api
+
+
+def shields():
+    # Shields Section
+    st.header("Shields")
+
+    shields_info = {
+        s.identifier: s.to_dict() for s in llama_stack_api.client.shields.list()
+    }
+
+    selected_shield = st.selectbox("Select a shield", list(shields_info.keys()))
+    st.json(shields_info[selected_shield])

llama_stack/distribution/ui/page/evaluations/__init__.py

@@ -0,0 +1,5 @@
+# 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.

llama_stack/distribution/ui/page/evaluations/app_eval.py

@@ -0,0 +1,148 @@
+# 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 pandas as pd
+import streamlit as st
+
+from modules.api import llama_stack_api
+from modules.utils import process_dataset
+
+
+def application_evaluation_page():
+
+    st.set_page_config(page_title="Evaluations (Scoring)", page_icon="🦙")
+    st.title("📊 Evaluations (Scoring)")
+
+    # File uploader
+    uploaded_file = st.file_uploader("Upload Dataset", type=["csv", "xlsx", "xls"])
+
+    if uploaded_file is None:
+        st.error("No file uploaded")
+        return
+
+    # Process uploaded file
+    df = process_dataset(uploaded_file)
+    if df is None:
+        st.error("Error processing file")
+        return
+
+    # Display dataset information
+    st.success("Dataset loaded successfully!")
+
+    # Display dataframe preview
+    st.subheader("Dataset Preview")
+    st.dataframe(df)
+
+    # Select Scoring Functions to Run Evaluation On
+    st.subheader("Select Scoring Functions")
+    scoring_functions = llama_stack_api.client.scoring_functions.list()
+    scoring_functions = {sf.identifier: sf for sf in scoring_functions}
+    scoring_functions_names = list(scoring_functions.keys())
+    selected_scoring_functions = st.multiselect(
+        "Choose one or more scoring functions",
+        options=scoring_functions_names,
+        help="Choose one or more scoring functions.",
+    )
+
+    available_models = llama_stack_api.client.models.list()
+    available_models = [m.identifier for m in available_models]
+
+    scoring_params = {}
+    if selected_scoring_functions:
+        st.write("Selected:")
+        for scoring_fn_id in selected_scoring_functions:
+            scoring_fn = scoring_functions[scoring_fn_id]
+            st.write(f"- **{scoring_fn_id}**: {scoring_fn.description}")
+            new_params = None
+            if scoring_fn.params:
+                new_params = {}
+                for param_name, param_value in scoring_fn.params.to_dict().items():
+                    if param_name == "type":
+                        new_params[param_name] = param_value
+                        continue
+
+                    if param_name == "judge_model":
+                        value = st.selectbox(
+                            f"Select **{param_name}** for {scoring_fn_id}",
+                            options=available_models,
+                            index=0,
+                            key=f"{scoring_fn_id}_{param_name}",
+                        )
+                        new_params[param_name] = value
+                    else:
+                        value = st.text_area(
+                            f"Enter value for **{param_name}** in {scoring_fn_id} in valid JSON format",
+                            value=json.dumps(param_value, indent=2),
+                            height=80,
+                        )
+                        try:
+                            new_params[param_name] = json.loads(value)
+                        except json.JSONDecodeError:
+                            st.error(
+                                f"Invalid JSON for **{param_name}** in {scoring_fn_id}"
+                            )
+
+                st.json(new_params)
+            scoring_params[scoring_fn_id] = new_params
+
+        # Add run evaluation button & slider
+        total_rows = len(df)
+        num_rows = st.slider("Number of rows to evaluate", 1, total_rows, total_rows)
+
+        if st.button("Run Evaluation"):
+            progress_text = "Running evaluation..."
+            progress_bar = st.progress(0, text=progress_text)
+            rows = df.to_dict(orient="records")
+            if num_rows < total_rows:
+                rows = rows[:num_rows]
+
+            # Create separate containers for progress text and results
+            progress_text_container = st.empty()
+            results_container = st.empty()
+            output_res = {}
+            for i, r in enumerate(rows):
+                # Update progress
+                progress = i / len(rows)
+                progress_bar.progress(progress, text=progress_text)
+
+                # Run evaluation for current row
+                score_res = llama_stack_api.run_scoring(
+                    r,
+                    scoring_function_ids=selected_scoring_functions,
+                    scoring_params=scoring_params,
+                )
+
+                for k in r.keys():
+                    if k not in output_res:
+                        output_res[k] = []
+                    output_res[k].append(r[k])
+
+                for fn_id in selected_scoring_functions:
+                    if fn_id not in output_res:
+                        output_res[fn_id] = []
+                    output_res[fn_id].append(score_res.results[fn_id].score_rows[0])
+
+                # Display current row results using separate containers
+                progress_text_container.write(
+                    f"Expand to see current processed result ({i+1}/{len(rows)})"
+                )
+                results_container.json(
+                    score_res.to_json(),
+                    expanded=2,
+                )
+
+            progress_bar.progress(1.0, text="Evaluation complete!")
+
+            # Display results in dataframe
+            if output_res:
+                output_df = pd.DataFrame(output_res)
+                st.subheader("Evaluation Results")
+                st.dataframe(output_df)
+
+
+application_evaluation_page()

llama_stack/distribution/ui/page/evaluations/native_eval.py

@@ -0,0 +1,257 @@
+# 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 pandas as pd
+
+import streamlit as st
+
+from modules.api import llama_stack_api
+
+
+def select_eval_task_1():
+    # Select Eval Tasks
+    st.subheader("1. Choose An Eval Task")
+    eval_tasks = llama_stack_api.client.eval_tasks.list()
+    eval_tasks = {et.identifier: et for et in eval_tasks}
+    eval_tasks_names = list(eval_tasks.keys())
+    selected_eval_task = st.selectbox(
+        "Choose an eval task.",
+        options=eval_tasks_names,
+        help="Choose an eval task. Each eval task is parameterized by a dataset, and list of scoring functions.",
+    )
+    with st.expander("View Eval Task"):
+        st.json(eval_tasks[selected_eval_task], expanded=True)
+
+    st.session_state["selected_eval_task"] = selected_eval_task
+    st.session_state["eval_tasks"] = eval_tasks
+    if st.button("Confirm", key="confirm_1"):
+        st.session_state["selected_eval_task_1_next"] = True
+
+
+def define_eval_candidate_2():
+    if not st.session_state.get("selected_eval_task_1_next", None):
+        return
+
+    st.subheader("2. Define Eval Candidate")
+    st.info(
+        """
+        Define the configurations for the evaluation candidate model or agent used for generation.
+        Select "model" if you want to run generation with inference API, or "agent" if you want to run generation with agent API through specifying AgentConfig.
+        """
+    )
+    with st.expander("Define Eval Candidate", expanded=True):
+        # Define Eval Candidate
+        candidate_type = st.radio("Candidate Type", ["model", "agent"])
+
+        available_models = llama_stack_api.client.models.list()
+        available_models = [model.identifier for model in available_models]
+        selected_model = st.selectbox(
+            "Choose a model",
+            available_models,
+            index=0,
+        )
+
+        # Sampling Parameters
+        st.markdown("##### Sampling Parameters")
+        strategy = st.selectbox(
+            "Strategy",
+            ["greedy", "top_p", "top_k"],
+            index=0,
+        )
+        temperature = st.slider(
+            "Temperature",
+            min_value=0.0,
+            max_value=1.0,
+            value=0.0,
+            step=0.1,
+            help="Controls the randomness of the response. Higher values make the output more creative and unexpected, lower values make it more conservative and predictable",
+        )
+        top_p = st.slider(
+            "Top P",
+            min_value=0.0,
+            max_value=1.0,
+            value=0.95,
+            step=0.1,
+        )
+        max_tokens = st.slider(
+            "Max Tokens",
+            min_value=0,
+            max_value=4096,
+            value=512,
+            step=1,
+            help="The maximum number of tokens to generate",
+        )
+        repetition_penalty = st.slider(
+            "Repetition Penalty",
+            min_value=1.0,
+            max_value=2.0,
+            value=1.0,
+            step=0.1,
+            help="Controls the likelihood for generating the same word or phrase multiple times in the same sentence or paragraph. 1 implies no penalty, 2 will strongly discourage model to repeat words or phrases.",
+        )
+        if candidate_type == "model":
+            eval_candidate = {
+                "type": "model",
+                "model": selected_model,
+                "sampling_params": {
+                    "strategy": strategy,
+                    "temperature": temperature,
+                    "top_p": top_p,
+                    "max_tokens": max_tokens,
+                    "repetition_penalty": repetition_penalty,
+                },
+            }
+        elif candidate_type == "agent":
+            system_prompt = st.text_area(
+                "System Prompt",
+                value="You are a helpful AI assistant.",
+                help="Initial instructions given to the AI to set its behavior and context",
+            )
+            tools_json = st.text_area(
+                "Tools Configuration (JSON)",
+                value=json.dumps(
+                    [
+                        {
+                            "type": "brave_search",
+                            "engine": "brave",
+                            "api_key": "ENTER_BRAVE_API_KEY_HERE",
+                        }
+                    ]
+                ),
+                help="Enter tool configurations in JSON format. Each tool should have a name, description, and parameters.",
+                height=200,
+            )
+            try:
+                tools = json.loads(tools_json)
+            except json.JSONDecodeError:
+                st.error("Invalid JSON format for tools configuration")
+                tools = []
+            eval_candidate = {
+                "type": "agent",
+                "config": {
+                    "model": selected_model,
+                    "instructions": system_prompt,
+                    "tools": tools,
+                    "tool_choice": "auto",
+                    "tool_prompt_format": "json",
+                    "input_shields": [],
+                    "output_shields": [],
+                    "enable_session_persistence": False,
+                },
+            }
+        st.session_state["eval_candidate"] = eval_candidate
+
+    if st.button("Confirm", key="confirm_2"):
+        st.session_state["selected_eval_candidate_2_next"] = True
+
+
+def run_evaluation_3():
+    if not st.session_state.get("selected_eval_candidate_2_next", None):
+        return
+
+    st.subheader("3. Run Evaluation")
+    # Add info box to explain configurations being used
+    st.info(
+        """
+        Review the configurations that will be used for this evaluation run, make any necessary changes, and then click the "Run Evaluation" button.
+        """
+    )
+    selected_eval_task = st.session_state["selected_eval_task"]
+    eval_tasks = st.session_state["eval_tasks"]
+    eval_candidate = st.session_state["eval_candidate"]
+
+    dataset_id = eval_tasks[selected_eval_task].dataset_id
+    rows = llama_stack_api.client.datasetio.get_rows_paginated(
+        dataset_id=dataset_id,
+        rows_in_page=-1,
+    )
+    total_rows = len(rows.rows)
+    # Add number of examples control
+    num_rows = st.number_input(
+        "Number of Examples to Evaluate",
+        min_value=1,
+        max_value=total_rows,
+        value=5,
+        help="Number of examples from the dataset to evaluate. ",
+    )
+
+    eval_task_config = {
+        "type": "benchmark",
+        "eval_candidate": eval_candidate,
+        "scoring_params": {},
+    }
+
+    with st.expander("View Evaluation Task", expanded=True):
+        st.json(eval_tasks[selected_eval_task], expanded=True)
+    with st.expander("View Evaluation Task Configuration", expanded=True):
+        st.json(eval_task_config, expanded=True)
+
+    # Add run button and handle evaluation
+    if st.button("Run Evaluation"):
+
+        progress_text = "Running evaluation..."
+        progress_bar = st.progress(0, text=progress_text)
+        rows = rows.rows
+        if num_rows < total_rows:
+            rows = rows[:num_rows]
+
+        # Create separate containers for progress text and results
+        progress_text_container = st.empty()
+        results_container = st.empty()
+        output_res = {}
+        for i, r in enumerate(rows):
+            # Update progress
+            progress = i / len(rows)
+            progress_bar.progress(progress, text=progress_text)
+            # Run evaluation for current row
+            eval_res = llama_stack_api.client.eval.evaluate_rows(
+                task_id=selected_eval_task,
+                input_rows=[r],
+                scoring_functions=eval_tasks[selected_eval_task].scoring_functions,
+                task_config=eval_task_config,
+            )
+
+            for k in r.keys():
+                if k not in output_res:
+                    output_res[k] = []
+                output_res[k].append(r[k])
+
+            for k in eval_res.generations[0].keys():
+                if k not in output_res:
+                    output_res[k] = []
+                output_res[k].append(eval_res.generations[0][k])
+
+            for scoring_fn in eval_tasks[selected_eval_task].scoring_functions:
+                if scoring_fn not in output_res:
+                    output_res[scoring_fn] = []
+                output_res[scoring_fn].append(eval_res.scores[scoring_fn].score_rows[0])
+
+            progress_text_container.write(
+                f"Expand to see current processed result ({i+1}/{len(rows)})"
+            )
+            results_container.json(eval_res, expanded=2)
+
+        progress_bar.progress(1.0, text="Evaluation complete!")
+        # Display results in dataframe
+        if output_res:
+            output_df = pd.DataFrame(output_res)
+            st.subheader("Evaluation Results")
+            st.dataframe(output_df)
+
+
+def native_evaluation_page():
+
+    st.set_page_config(page_title="Evaluations (Generation + Scoring)", page_icon="🦙")
+    st.title("📊 Evaluations (Generation + Scoring)")
+
+    select_eval_task_1()
+    define_eval_candidate_2()
+    run_evaluation_3()
+
+
+native_evaluation_page()

llama_stack/distribution/ui/page/playground/__init__.py

@@ -0,0 +1,5 @@
+# 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.

llama_stack/distribution/ui/page/playground/chat.py

@@ -0,0 +1,123 @@
+# 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 streamlit as st
+from modules.api import llama_stack_api
+
+# Sidebar configurations
+with st.sidebar:
+    st.header("Configuration")
+    available_models = llama_stack_api.client.models.list()
+    available_models = [model.identifier for model in available_models]
+    selected_model = st.selectbox(
+        "Choose a model",
+        available_models,
+        index=0,
+    )
+
+    temperature = st.slider(
+        "Temperature",
+        min_value=0.0,
+        max_value=1.0,
+        value=0.0,
+        step=0.1,
+        help="Controls the randomness of the response. Higher values make the output more creative and unexpected, lower values make it more conservative and predictable",
+    )
+
+    top_p = st.slider(
+        "Top P",
+        min_value=0.0,
+        max_value=1.0,
+        value=0.95,
+        step=0.1,
+    )
+
+    max_tokens = st.slider(
+        "Max Tokens",
+        min_value=0,
+        max_value=4096,
+        value=512,
+        step=1,
+        help="The maximum number of tokens to generate",
+    )
+
+    repetition_penalty = st.slider(
+        "Repetition Penalty",
+        min_value=1.0,
+        max_value=2.0,
+        value=1.0,
+        step=0.1,
+        help="Controls the likelihood for generating the same word or phrase multiple times in the same sentence or paragraph. 1 implies no penalty, 2 will strongly discourage model to repeat words or phrases.",
+    )
+
+    stream = st.checkbox("Stream", value=True)
+    system_prompt = st.text_area(
+        "System Prompt",
+        value="You are a helpful AI assistant.",
+        help="Initial instructions given to the AI to set its behavior and context",
+    )
+
+    # Add clear chat button to sidebar
+    if st.button("Clear Chat", use_container_width=True):
+        st.session_state.messages = []
+        st.rerun()
+
+
+# Main chat interface
+st.title("🦙 Chat")
+
+
+# Initialize chat history
+if "messages" not in st.session_state:
+    st.session_state.messages = []
+
+# Display chat messages
+for message in st.session_state.messages:
+    with st.chat_message(message["role"]):
+        st.markdown(message["content"])
+
+# Chat input
+if prompt := st.chat_input("Example: What is Llama Stack?"):
+    # Add user message to chat history
+    st.session_state.messages.append({"role": "user", "content": prompt})
+
+    # Display user message
+    with st.chat_message("user"):
+        st.markdown(prompt)
+
+    # Display assistant response
+    with st.chat_message("assistant"):
+        message_placeholder = st.empty()
+        full_response = ""
+
+        response = llama_stack_api.client.inference.chat_completion(
+            messages=[
+                {"role": "system", "content": system_prompt},
+                {"role": "user", "content": prompt},
+            ],
+            model_id=selected_model,
+            stream=stream,
+            sampling_params={
+                "temperature": temperature,
+                "top_p": top_p,
+                "max_tokens": max_tokens,
+                "repetition_penalty": repetition_penalty,
+            },
+        )
+
+        if stream:
+            for chunk in response:
+                if chunk.event.event_type == "progress":
+                    full_response += chunk.event.delta
+                message_placeholder.markdown(full_response + "▌")
+            message_placeholder.markdown(full_response)
+        else:
+            full_response = response
+            message_placeholder.markdown(full_response.completion_message.content)
+
+        st.session_state.messages.append(
+            {"role": "assistant", "content": full_response}
+        )

llama_stack/distribution/ui/page/playground/rag.py

@@ -0,0 +1,188 @@
+# 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 streamlit as st
+from llama_stack_client.lib.agents.agent import Agent
+from llama_stack_client.lib.agents.event_logger import EventLogger
+from llama_stack_client.types.agent_create_params import AgentConfig
+from llama_stack_client.types.memory_insert_params import Document
+
+from modules.api import llama_stack_api
+from modules.utils import data_url_from_file
+
+
+def rag_chat_page():
+    st.title("🦙 RAG")
+
+    with st.sidebar:
+        # File/Directory Upload Section
+        st.subheader("Upload Documents")
+        uploaded_files = st.file_uploader(
+            "Upload file(s) or directory",
+            accept_multiple_files=True,
+            type=["txt", "pdf", "doc", "docx"],  # Add more file types as needed
+        )
+        # Process uploaded files
+        if uploaded_files:
+            st.success(f"Successfully uploaded {len(uploaded_files)} files")
+            # Add memory bank name input field
+            memory_bank_name = st.text_input(
+                "Memory Bank Name",
+                value="rag_bank",
+                help="Enter a unique identifier for this memory bank",
+            )
+            if st.button("Create Memory Bank"):
+                documents = [
+                    Document(
+                        document_id=uploaded_file.name,
+                        content=data_url_from_file(uploaded_file),
+                    )
+                    for i, uploaded_file in enumerate(uploaded_files)
+                ]
+
+                providers = llama_stack_api.client.providers.list()
+                llama_stack_api.client.memory_banks.register(
+                    memory_bank_id=memory_bank_name,  # Use the user-provided name
+                    params={
+                        "embedding_model": "all-MiniLM-L6-v2",
+                        "chunk_size_in_tokens": 512,
+                        "overlap_size_in_tokens": 64,
+                    },
+                    provider_id=providers["memory"][0].provider_id,
+                )
+
+                # insert documents using the custom bank name
+                llama_stack_api.client.memory.insert(
+                    bank_id=memory_bank_name,  # Use the user-provided name
+                    documents=documents,
+                )
+                st.success("Memory bank created successfully!")
+
+        st.subheader("Configure Agent")
+        # select memory banks
+        memory_banks = llama_stack_api.client.memory_banks.list()
+        memory_banks = [bank.identifier for bank in memory_banks]
+        selected_memory_banks = st.multiselect(
+            "Select Memory Banks",
+            memory_banks,
+        )
+        memory_bank_configs = [
+            {"bank_id": bank_id, "type": "vector"} for bank_id in selected_memory_banks
+        ]
+
+        available_models = llama_stack_api.client.models.list()
+        available_models = [model.identifier for model in available_models]
+        selected_model = st.selectbox(
+            "Choose a model",
+            available_models,
+            index=0,
+        )
+        system_prompt = st.text_area(
+            "System Prompt",
+            value="You are a helpful assistant. ",
+            help="Initial instructions given to the AI to set its behavior and context",
+        )
+        temperature = st.slider(
+            "Temperature",
+            min_value=0.0,
+            max_value=1.0,
+            value=0.0,
+            step=0.1,
+            help="Controls the randomness of the response. Higher values make the output more creative and unexpected, lower values make it more conservative and predictable",
+        )
+
+        top_p = st.slider(
+            "Top P",
+            min_value=0.0,
+            max_value=1.0,
+            value=0.95,
+            step=0.1,
+        )
+
+        # Add clear chat button to sidebar
+        if st.button("Clear Chat", use_container_width=True):
+            st.session_state.messages = []
+            st.rerun()
+
+    # Chat Interface
+    if "messages" not in st.session_state:
+        st.session_state.messages = []
+
+    # Display chat history
+    for message in st.session_state.messages:
+        with st.chat_message(message["role"]):
+            st.markdown(message["content"])
+
+    selected_model = llama_stack_api.client.models.list()[0].identifier
+
+    agent_config = AgentConfig(
+        model=selected_model,
+        instructions=system_prompt,
+        sampling_params={
+            "strategy": "greedy",
+            "temperature": temperature,
+            "top_p": top_p,
+        },
+        tools=[
+            {
+                "type": "memory",
+                "memory_bank_configs": memory_bank_configs,
+                "query_generator_config": {"type": "default", "sep": " "},
+                "max_tokens_in_context": 4096,
+                "max_chunks": 10,
+            }
+        ],
+        tool_choice="auto",
+        tool_prompt_format="json",
+        input_shields=[],
+        output_shields=[],
+        enable_session_persistence=False,
+    )
+
+    agent = Agent(llama_stack_api.client, agent_config)
+    session_id = agent.create_session("rag-session")
+
+    # Chat input
+    if prompt := st.chat_input("Ask a question about your documents"):
+        # Add user message to chat history
+        st.session_state.messages.append({"role": "user", "content": prompt})
+
+        # Display user message
+        with st.chat_message("user"):
+            st.markdown(prompt)
+
+        response = agent.create_turn(
+            messages=[
+                {
+                    "role": "user",
+                    "content": prompt,
+                }
+            ],
+            session_id=session_id,
+        )
+
+        # Display assistant response
+        with st.chat_message("assistant"):
+            retrieval_message_placeholder = st.empty()
+            message_placeholder = st.empty()
+            full_response = ""
+            retrieval_response = ""
+            for log in EventLogger().log(response):
+                log.print()
+                if log.role == "memory_retrieval":
+                    retrieval_response += log.content.replace("====", "").strip()
+                    retrieval_message_placeholder.info(retrieval_response)
+                else:
+                    full_response += log.content
+                    message_placeholder.markdown(full_response + "▌")
+            message_placeholder.markdown(full_response)
+
+            st.session_state.messages.append(
+                {"role": "assistant", "content": full_response}
+            )
+
+
+rag_chat_page()

llama_stack/distribution/ui/requirements.txt

@@ -0,0 +1,4 @@
+streamlit
+pandas
+llama-stack-client>=0.0.55
+streamlit-option-menu

llama_stack/distribution/utils/model_utils.py

@@ -4,11 +4,10 @@
 # This source code is licensed under the terms described in the LICENSE file in
 # the root directory of this source tree.
 
-import os
+from pathlib import Path
 
 from .config_dirs import DEFAULT_CHECKPOINT_DIR
 
 
 def model_local_dir(descriptor: str) -> str:
-    path = os.path.join(DEFAULT_CHECKPOINT_DIR, descriptor)
-    return path.replace(":", "-")
+    return str(Path(DEFAULT_CHECKPOINT_DIR) / (descriptor.replace(":", "-")))

llama_stack/providers/datatypes.py

@@ -63,6 +63,8 @@ class MemoryBanksProtocolPrivate(Protocol):
 class DatasetsProtocolPrivate(Protocol):
     async def register_dataset(self, dataset: Dataset) -> None: ...
 
+    async def unregister_dataset(self, dataset_id: str) -> None: ...
+
 
 class ScoringFunctionsProtocolPrivate(Protocol):
     async def list_scoring_functions(self) -> List[ScoringFn]: ...

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

@@ -144,87 +144,91 @@ class ChatAgent(ShieldRunnerMixin):
     async def create_session(self, name: str) -> str:
         return await self.storage.create_session(name)
 
-    @tracing.span("create_and_execute_turn")
     async def create_and_execute_turn(
         self, request: AgentTurnCreateRequest
     ) -> AsyncGenerator:
-        assert request.stream is True, "Non-streaming not supported"
+        with tracing.span("create_and_execute_turn") as span:
+            span.set_attribute("session_id", request.session_id)
+            span.set_attribute("agent_id", self.agent_id)
+            span.set_attribute("request", request.model_dump_json())
+            assert request.stream is True, "Non-streaming not supported"
 
-        session_info = await self.storage.get_session_info(request.session_id)
-        if session_info is None:
-            raise ValueError(f"Session {request.session_id} not found")
+            session_info = await self.storage.get_session_info(request.session_id)
+            if session_info is None:
+                raise ValueError(f"Session {request.session_id} not found")
 
-        turns = await self.storage.get_session_turns(request.session_id)
+            turns = await self.storage.get_session_turns(request.session_id)
 
-        messages = []
-        if self.agent_config.instructions != "":
-            messages.append(SystemMessage(content=self.agent_config.instructions))
+            messages = []
+            if self.agent_config.instructions != "":
+                messages.append(SystemMessage(content=self.agent_config.instructions))
 
-        for i, turn in enumerate(turns):
-            messages.extend(self.turn_to_messages(turn))
+            for i, turn in enumerate(turns):
+                messages.extend(self.turn_to_messages(turn))
 
-        messages.extend(request.messages)
+            messages.extend(request.messages)
 
-        turn_id = str(uuid.uuid4())
-        start_time = datetime.now()
-        yield AgentTurnResponseStreamChunk(
-            event=AgentTurnResponseEvent(
-                payload=AgentTurnResponseTurnStartPayload(
-                    turn_id=turn_id,
+            turn_id = str(uuid.uuid4())
+            span.set_attribute("turn_id", turn_id)
+            start_time = datetime.now()
+            yield AgentTurnResponseStreamChunk(
+                event=AgentTurnResponseEvent(
+                    payload=AgentTurnResponseTurnStartPayload(
+                        turn_id=turn_id,
+                    )
                 )
             )
-        )
 
-        steps = []
-        output_message = None
-        async for chunk in self.run(
-            session_id=request.session_id,
-            turn_id=turn_id,
-            input_messages=messages,
-            attachments=request.attachments or [],
-            sampling_params=self.agent_config.sampling_params,
-            stream=request.stream,
-        ):
-            if isinstance(chunk, CompletionMessage):
-                log.info(
-                    f"{chunk.role.capitalize()}: {chunk.content}",
-                )
-                output_message = chunk
-                continue
-
-            assert isinstance(
-                chunk, AgentTurnResponseStreamChunk
-            ), f"Unexpected type {type(chunk)}"
-            event = chunk.event
-            if (
-                event.payload.event_type
-                == AgentTurnResponseEventType.step_complete.value
+            steps = []
+            output_message = None
+            async for chunk in self.run(
+                session_id=request.session_id,
+                turn_id=turn_id,
+                input_messages=messages,
+                attachments=request.attachments or [],
+                sampling_params=self.agent_config.sampling_params,
+                stream=request.stream,
             ):
-                steps.append(event.payload.step_details)
+                if isinstance(chunk, CompletionMessage):
+                    log.info(
+                        f"{chunk.role.capitalize()}: {chunk.content}",
+                    )
+                    output_message = chunk
+                    continue
 
-            yield chunk
+                assert isinstance(
+                    chunk, AgentTurnResponseStreamChunk
+                ), f"Unexpected type {type(chunk)}"
+                event = chunk.event
+                if (
+                    event.payload.event_type
+                    == AgentTurnResponseEventType.step_complete.value
+                ):
+                    steps.append(event.payload.step_details)
 
-        assert output_message is not None
+                yield chunk
 
-        turn = Turn(
-            turn_id=turn_id,
-            session_id=request.session_id,
-            input_messages=request.messages,
-            output_message=output_message,
-            started_at=start_time,
-            completed_at=datetime.now(),
-            steps=steps,
-        )
-        await self.storage.add_turn_to_session(request.session_id, turn)
+            assert output_message is not None
 
-        chunk = AgentTurnResponseStreamChunk(
-            event=AgentTurnResponseEvent(
-                payload=AgentTurnResponseTurnCompletePayload(
-                    turn=turn,
+            turn = Turn(
+                turn_id=turn_id,
+                session_id=request.session_id,
+                input_messages=request.messages,
+                output_message=output_message,
+                started_at=start_time,
+                completed_at=datetime.now(),
+                steps=steps,
+            )
+            await self.storage.add_turn_to_session(request.session_id, turn)
+
+            chunk = AgentTurnResponseStreamChunk(
+                event=AgentTurnResponseEvent(
+                    payload=AgentTurnResponseTurnCompletePayload(
+                        turn=turn,
+                    )
                 )
             )
-        )
-        yield chunk
+            yield chunk
 
     async def run(
         self,
@@ -273,7 +277,6 @@ class ChatAgent(ShieldRunnerMixin):
 
         yield final_response
 
-    @tracing.span("run_shields")
     async def run_multiple_shields_wrapper(
         self,
         turn_id: str,
@@ -281,23 +284,47 @@ class ChatAgent(ShieldRunnerMixin):
         shields: List[str],
         touchpoint: str,
     ) -> AsyncGenerator:
-        if len(shields) == 0:
-            return
+        with tracing.span("run_shields") as span:
+            span.set_attribute("turn_id", turn_id)
+            span.set_attribute("input", [m.model_dump_json() for m in messages])
+            if len(shields) == 0:
+                span.set_attribute("output", "no shields")
+                return
 
-        step_id = str(uuid.uuid4())
-        try:
-            yield AgentTurnResponseStreamChunk(
-                event=AgentTurnResponseEvent(
-                    payload=AgentTurnResponseStepStartPayload(
-                        step_type=StepType.shield_call.value,
-                        step_id=step_id,
-                        metadata=dict(touchpoint=touchpoint),
+            step_id = str(uuid.uuid4())
+            try:
+                yield AgentTurnResponseStreamChunk(
+                    event=AgentTurnResponseEvent(
+                        payload=AgentTurnResponseStepStartPayload(
+                            step_type=StepType.shield_call.value,
+                            step_id=step_id,
+                            metadata=dict(touchpoint=touchpoint),
+                        )
                     )
                 )
-            )
-            await self.run_multiple_shields(messages, shields)
+                await self.run_multiple_shields(messages, shields)
+
+            except SafetyException as e:
+                yield AgentTurnResponseStreamChunk(
+                    event=AgentTurnResponseEvent(
+                        payload=AgentTurnResponseStepCompletePayload(
+                            step_type=StepType.shield_call.value,
+                            step_details=ShieldCallStep(
+                                step_id=step_id,
+                                turn_id=turn_id,
+                                violation=e.violation,
+                            ),
+                        )
+                    )
+                )
+                span.set_attribute("output", e.violation.model_dump_json())
+
+                yield CompletionMessage(
+                    content=str(e),
+                    stop_reason=StopReason.end_of_turn,
+                )
+                yield False
 
-        except SafetyException as e:
             yield AgentTurnResponseStreamChunk(
                 event=AgentTurnResponseEvent(
                     payload=AgentTurnResponseStepCompletePayload(
@@ -305,30 +332,12 @@ class ChatAgent(ShieldRunnerMixin):
                         step_details=ShieldCallStep(
                             step_id=step_id,
                             turn_id=turn_id,
-                            violation=e.violation,
+                            violation=None,
                         ),
                     )
                 )
             )
-
-            yield CompletionMessage(
-                content=str(e),
-                stop_reason=StopReason.end_of_turn,
-            )
-            yield False
-
-        yield AgentTurnResponseStreamChunk(
-            event=AgentTurnResponseEvent(
-                payload=AgentTurnResponseStepCompletePayload(
-                    step_type=StepType.shield_call.value,
-                    step_details=ShieldCallStep(
-                        step_id=step_id,
-                        turn_id=turn_id,
-                        violation=None,
-                    ),
-                )
-            )
-        )
+            span.set_attribute("output", "no violations")
 
     async def _run(
         self,
@@ -356,10 +365,15 @@ class ChatAgent(ShieldRunnerMixin):
 
             # TODO: find older context from the session and either replace it
             # or append with a sliding window. this is really a very simplistic implementation
-            with tracing.span("retrieve_rag_context"):
+            with tracing.span("retrieve_rag_context") as span:
                 rag_context, bank_ids = await self._retrieve_context(
                     session_id, input_messages, attachments
                 )
+                span.set_attribute(
+                    "input", [m.model_dump_json() for m in input_messages]
+                )
+                span.set_attribute("output", rag_context)
+                span.set_attribute("bank_ids", bank_ids)
 
             step_id = str(uuid.uuid4())
             yield AgentTurnResponseStreamChunk(
@@ -416,7 +430,7 @@ class ChatAgent(ShieldRunnerMixin):
             content = ""
             stop_reason = None
 
-            with tracing.span("inference"):
+            with tracing.span("inference") as span:
                 async for chunk in await self.inference_api.chat_completion(
                     self.agent_config.model,
                     input_messages,
@@ -436,7 +450,6 @@ class ChatAgent(ShieldRunnerMixin):
                     if isinstance(delta, ToolCallDelta):
                         if delta.parse_status == ToolCallParseStatus.success:
                             tool_calls.append(delta.content)
-
                         if stream:
                             yield AgentTurnResponseStreamChunk(
                                 event=AgentTurnResponseEvent(
@@ -466,6 +479,13 @@ class ChatAgent(ShieldRunnerMixin):
 
                     if event.stop_reason is not None:
                         stop_reason = event.stop_reason
+                span.set_attribute("stop_reason", stop_reason)
+                span.set_attribute(
+                    "input", [m.model_dump_json() for m in input_messages]
+                )
+                span.set_attribute(
+                    "output", f"content: {content} tool_calls: {tool_calls}"
+                )
 
             stop_reason = stop_reason or StopReason.out_of_tokens
 
@@ -549,7 +569,13 @@ class ChatAgent(ShieldRunnerMixin):
                         )
                     )
 
-                    with tracing.span("tool_execution"):
+                    with tracing.span(
+                        "tool_execution",
+                        {
+                            "tool_name": tool_call.tool_name,
+                            "input": message.model_dump_json(),
+                        },
+                    ) as span:
                         result_messages = await execute_tool_call_maybe(
                             self.tools_dict,
                             [message],
@@ -558,6 +584,7 @@ class ChatAgent(ShieldRunnerMixin):
                             len(result_messages) == 1
                         ), "Currently not supporting multiple messages"
                         result_message = result_messages[0]
+                        span.set_attribute("output", result_message.model_dump_json())
 
                     yield AgentTurnResponseStreamChunk(
                         event=AgentTurnResponseEvent(

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

@@ -6,9 +6,12 @@
 
 import json
 import logging
+import shutil
 import uuid
 from typing import AsyncGenerator
 
+from termcolor import colored
+
 from llama_stack.apis.inference import Inference
 from llama_stack.apis.memory import Memory
 from llama_stack.apis.memory_banks import MemoryBanks
@@ -44,6 +47,15 @@ class MetaReferenceAgentsImpl(Agents):
     async def initialize(self) -> None:
         self.persistence_store = await kvstore_impl(self.config.persistence_store)
 
+        # check if "bwrap" is available
+        if not shutil.which("bwrap"):
+            print(
+                colored(
+                    "Warning: `bwrap` is not available. Code interpreter tool will not work correctly.",
+                    "yellow",
+                )
+            )
+
     async def create_agent(
         self,
         agent_config: AgentConfig,

llama_stack/providers/inline/datasetio/__init__.py

@@ -0,0 +1,5 @@
+# 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.

llama_stack/providers/inline/datasetio/localfs/datasetio.py

@@ -3,14 +3,17 @@
 #
 # This source code is licensed under the terms described in the LICENSE file in
 # the root directory of this source tree.
-from typing import Optional
+from typing import Any, Dict, List, Optional
 
 import pandas
 from llama_models.llama3.api.datatypes import *  # noqa: F403
 
 from llama_stack.apis.datasetio import *  # noqa: F403
+import base64
+import os
 from abc import ABC, abstractmethod
 from dataclasses import dataclass
+from urllib.parse import urlparse
 
 from llama_stack.providers.datatypes import DatasetsProtocolPrivate
 from llama_stack.providers.utils.datasetio.url_utils import get_dataframe_from_url
@@ -97,6 +100,9 @@ class LocalFSDatasetIOImpl(DatasetIO, DatasetsProtocolPrivate):
             dataset_impl=dataset_impl,
         )
 
+    async def unregister_dataset(self, dataset_id: str) -> None:
+        del self.dataset_infos[dataset_id]
+
     async def get_rows_paginated(
         self,
         dataset_id: str,
@@ -128,3 +134,41 @@ class LocalFSDatasetIOImpl(DatasetIO, DatasetsProtocolPrivate):
             total_count=len(rows),
             next_page_token=str(end),
         )
+
+    async def append_rows(self, dataset_id: str, rows: List[Dict[str, Any]]) -> None:
+        dataset_info = self.dataset_infos.get(dataset_id)
+        if dataset_info is None:
+            raise ValueError(f"Dataset with id {dataset_id} not found")
+
+        dataset_impl = dataset_info.dataset_impl
+        dataset_impl.load()
+
+        new_rows_df = pandas.DataFrame(rows)
+        new_rows_df = dataset_impl._validate_dataset_schema(new_rows_df)
+        dataset_impl.df = pandas.concat(
+            [dataset_impl.df, new_rows_df], ignore_index=True
+        )
+
+        url = str(dataset_info.dataset_def.url)
+        parsed_url = urlparse(url)
+
+        if parsed_url.scheme == "file" or not parsed_url.scheme:
+            file_path = parsed_url.path
+            os.makedirs(os.path.dirname(file_path), exist_ok=True)
+            dataset_impl.df.to_csv(file_path, index=False)
+        elif parsed_url.scheme == "data":
+            # For data URLs, we need to update the base64-encoded content
+            if not parsed_url.path.startswith("text/csv;base64,"):
+                raise ValueError("Data URL must be a base64-encoded CSV")
+
+            csv_buffer = dataset_impl.df.to_csv(index=False)
+            base64_content = base64.b64encode(csv_buffer.encode("utf-8")).decode(
+                "utf-8"
+            )
+            dataset_info.dataset_def.url = URL(
+                uri=f"data:text/csv;base64,{base64_content}"
+            )
+        else:
+            raise ValueError(
+                f"Unsupported URL scheme: {parsed_url.scheme}. Only file:// and data: URLs are supported for writing."
+            )

llama_stack/providers/inline/eval/__init__.py

@@ -0,0 +1,5 @@
+# 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.

llama_stack/providers/inline/eval/meta_reference/config.py

@@ -3,12 +3,13 @@
 #
 # This source code is licensed under the terms described in the LICENSE file in
 # the root directory of this source tree.
+from pydantic import BaseModel
+
 from llama_stack.distribution.utils.config_dirs import RUNTIME_BASE_DIR
 from llama_stack.providers.utils.kvstore.config import (
     KVStoreConfig,
     SqliteKVStoreConfig,
 )
-from pydantic import BaseModel
 
 
 class MetaReferenceEvalConfig(BaseModel):

llama_stack/providers/inline/eval/meta_reference/eval.py

@@ -4,7 +4,9 @@
 # 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 Any, Dict, List, Optional
 from llama_models.llama3.api.datatypes import *  # noqa: F403
+from tqdm import tqdm
 
 from .....apis.common.job_types import Job
 from .....apis.eval.eval import Eval, EvalTaskConfig, EvaluateResponse, JobStatus
@@ -17,7 +19,6 @@ from llama_stack.apis.inference import Inference
 from llama_stack.apis.scoring import Scoring
 from llama_stack.providers.datatypes import EvalTasksProtocolPrivate
 from llama_stack.providers.utils.kvstore import kvstore_impl
-from tqdm import tqdm
 
 from .config import MetaReferenceEvalConfig
 

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

@@ -27,7 +27,6 @@ from llama_stack.providers.utils.memory.vector_store import (
     BankWithIndex,
     EmbeddingIndex,
 )
-from llama_stack.providers.utils.telemetry import tracing
 
 from .config import FaissImplConfig
 
@@ -95,7 +94,6 @@ class FaissIndex(EmbeddingIndex):
 
         await self.kvstore.delete(f"faiss_index:v1::{self.bank_id}")
 
-    @tracing.span(name="add_chunks")
     async def add_chunks(self, chunks: List[Chunk], embeddings: NDArray):
         indexlen = len(self.id_by_index)
         for i, chunk in enumerate(chunks):

llama_stack/providers/inline/meta_reference/telemetry/__init__.py

@@ -1,15 +0,0 @@
-# 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 ConsoleConfig
-
-
-async def get_provider_impl(config: ConsoleConfig, _deps):
-    from .console import ConsoleTelemetryImpl
-
-    impl = ConsoleTelemetryImpl(config)
-    await impl.initialize()
-    return impl

llama_stack/providers/inline/meta_reference/telemetry/config.py

@@ -1,21 +0,0 @@
-# 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 llama_models.schema_utils import json_schema_type
-
-from pydantic import BaseModel
-
-
-class LogFormat(Enum):
-    TEXT = "text"
-    JSON = "json"
-
-
-@json_schema_type
-class ConsoleConfig(BaseModel):
-    log_format: LogFormat = LogFormat.TEXT

llama_stack/providers/inline/meta_reference/telemetry/console.py

@@ -5,7 +5,7 @@
 # the root directory of this source tree.
 
 import json
-from typing import Optional
+from typing import List, Optional
 
 from .config import LogFormat
 
@@ -49,8 +49,27 @@ class ConsoleTelemetryImpl(Telemetry):
         if formatted:
             print(formatted)
 
-    async def get_trace(self, trace_id: str) -> Trace:
-        raise NotImplementedError()
+    async def query_traces(
+        self,
+        attribute_conditions: Optional[List[QueryCondition]] = None,
+        attribute_keys_to_return: Optional[List[str]] = None,
+        limit: Optional[int] = 100,
+        offset: Optional[int] = 0,
+        order_by: Optional[List[str]] = None,
+    ) -> List[Trace]:
+        raise NotImplementedError("Console telemetry does not support trace querying")
+
+    async def get_spans(
+        self,
+        span_id: str,
+        attribute_conditions: Optional[List[QueryCondition]] = None,
+        attribute_keys_to_return: Optional[List[str]] = None,
+        max_depth: Optional[int] = None,
+        limit: Optional[int] = 100,
+        offset: Optional[int] = 0,
+        order_by: Optional[List[str]] = None,
+    ) -> SpanWithChildren:
+        raise NotImplementedError("Console telemetry does not support span querying")
 
 
 COLORS = {

llama_stack/providers/inline/scoring/__init__.py

@@ -0,0 +1,5 @@
+# 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.

llama_stack/providers/inline/scoring/braintrust/__init__.py

@@ -5,11 +5,17 @@
 # the root directory of this source tree.
 from typing import Dict
 
+from pydantic import BaseModel
+
 from llama_stack.distribution.datatypes import Api, ProviderSpec
 
 from .config import BraintrustScoringConfig
 
 
+class BraintrustProviderDataValidator(BaseModel):
+    openai_api_key: str
+
+
 async def get_provider_impl(
     config: BraintrustScoringConfig,
     deps: Dict[Api, ProviderSpec],

llama_stack/providers/inline/scoring/braintrust/braintrust.py

@@ -12,9 +12,12 @@ from llama_stack.apis.common.type_system import *  # noqa: F403
 from llama_stack.apis.datasetio import *  # noqa: F403
 from llama_stack.apis.datasets import *  # noqa: F403
 
-# from .scoring_fn.braintrust_scoring_fn import BraintrustScoringFn
+import os
+
 from autoevals.llm import Factuality
 from autoevals.ragas import AnswerCorrectness
+
+from llama_stack.distribution.request_headers import NeedsRequestProviderData
 from llama_stack.providers.datatypes import ScoringFunctionsProtocolPrivate
 
 from llama_stack.providers.utils.scoring.aggregation_utils import aggregate_average
@@ -24,7 +27,9 @@ from .scoring_fn.fn_defs.answer_correctness import answer_correctness_fn_def
 from .scoring_fn.fn_defs.factuality import factuality_fn_def
 
 
-class BraintrustScoringImpl(Scoring, ScoringFunctionsProtocolPrivate):
+class BraintrustScoringImpl(
+    Scoring, ScoringFunctionsProtocolPrivate, NeedsRequestProviderData
+):
     def __init__(
         self,
         config: BraintrustScoringConfig,
@@ -79,12 +84,25 @@ class BraintrustScoringImpl(Scoring, ScoringFunctionsProtocolPrivate):
                     f"Dataset {dataset_id} does not have a '{required_column}' column of type 'string'."
                 )
 
+    async def set_api_key(self) -> None:
+        # api key is in the request headers
+        if self.config.openai_api_key is None:
+            provider_data = self.get_request_provider_data()
+            if provider_data is None or not provider_data.openai_api_key:
+                raise ValueError(
+                    'Pass OpenAI API Key in the header X-LlamaStack-ProviderData as { "openai_api_key": <your api key>}'
+                )
+            self.config.openai_api_key = provider_data.openai_api_key
+
+        os.environ["OPENAI_API_KEY"] = self.config.openai_api_key
+
     async def score_batch(
         self,
         dataset_id: str,
         scoring_functions: List[str],
         save_results_dataset: bool = False,
     ) -> ScoreBatchResponse:
+        await self.set_api_key()
         await self.validate_scoring_input_dataset_schema(dataset_id=dataset_id)
         all_rows = await self.datasetio_api.get_rows_paginated(
             dataset_id=dataset_id,
@@ -105,6 +123,7 @@ class BraintrustScoringImpl(Scoring, ScoringFunctionsProtocolPrivate):
     async def score_row(
         self, input_row: Dict[str, Any], scoring_fn_identifier: Optional[str] = None
     ) -> ScoringResultRow:
+        await self.set_api_key()
         assert scoring_fn_identifier is not None, "scoring_fn_identifier cannot be None"
         expected_answer = input_row["expected_answer"]
         generated_answer = input_row["generated_answer"]
@@ -118,6 +137,7 @@ class BraintrustScoringImpl(Scoring, ScoringFunctionsProtocolPrivate):
     async def score(
         self, input_rows: List[Dict[str, Any]], scoring_functions: List[str]
     ) -> ScoreResponse:
+        await self.set_api_key()
         res = {}
         for scoring_fn_id in scoring_functions:
             if scoring_fn_id not in self.supported_fn_defs_registry:

llama_stack/providers/inline/scoring/braintrust/config.py

@@ -6,4 +6,8 @@
 from llama_stack.apis.scoring import *  # noqa: F401, F403
 
 
-class BraintrustScoringConfig(BaseModel): ...
+class BraintrustScoringConfig(BaseModel):
+    openai_api_key: Optional[str] = Field(
+        default=None,
+        description="The OpenAI API Key",
+    )

llama_stack/providers/inline/scoring/braintrust/scoring_fn/fn_defs/answer_correctness.py

@@ -10,7 +10,7 @@ from llama_stack.apis.scoring_functions import ScoringFn
 
 answer_correctness_fn_def = ScoringFn(
     identifier="braintrust::answer-correctness",
-    description="Test whether an output is factual, compared to an original (`expected`) value. One of Braintrust LLM basd scorer https://github.com/braintrustdata/autoevals/blob/main/py/autoevals/llm.py",
+    description="Scores the correctness of the answer based on the ground truth.. One of Braintrust LLM basd scorer https://github.com/braintrustdata/autoevals/blob/main/py/autoevals/llm.py",
     params=None,
     provider_id="braintrust",
     provider_resource_id="answer-correctness",

llama_stack/providers/inline/scoring/llm_as_judge/scoring_fn/fn_defs/llm_as_judge_base.py

@@ -5,7 +5,7 @@
 # the root directory of this source tree.
 
 from llama_stack.apis.common.type_system import NumberType
-from llama_stack.apis.scoring_functions import ScoringFn
+from llama_stack.apis.scoring_functions import LLMAsJudgeScoringFnParams, ScoringFn
 
 
 llm_as_judge_base = ScoringFn(
@@ -14,4 +14,8 @@ llm_as_judge_base = ScoringFn(
     return_type=NumberType(),
     provider_id="llm-as-judge",
     provider_resource_id="llm-as-judge-base",
+    params=LLMAsJudgeScoringFnParams(
+        judge_model="meta-llama/Llama-3.1-405B-Instruct",
+        prompt_template="Enter custom LLM as Judge Prompt Template",
+    ),
 )

llama_stack/providers/inline/telemetry/__init__.py

@@ -0,0 +1,5 @@
+# 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.

llama_stack/providers/inline/telemetry/meta_reference/__init__.py

@@ -0,0 +1,18 @@
+# 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 typing import Any, Dict
+
+from .config import TelemetryConfig, TelemetrySink
+from .telemetry import TelemetryAdapter
+
+__all__ = ["TelemetryConfig", "TelemetryAdapter", "TelemetrySink"]
+
+
+async def get_provider_impl(config: TelemetryConfig, deps: Dict[str, Any]):
+    impl = TelemetryAdapter(config, deps)
+    await impl.initialize()
+    return impl

llama_stack/providers/inline/telemetry/meta_reference/config.py

@@ -0,0 +1,45 @@
+# 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 Any, Dict, List
+
+from pydantic import BaseModel, Field
+
+from llama_stack.distribution.utils.config_dirs import RUNTIME_BASE_DIR
+
+
+class TelemetrySink(str, Enum):
+    OTEL = "otel"
+    SQLITE = "sqlite"
+    CONSOLE = "console"
+
+
+class TelemetryConfig(BaseModel):
+    otel_endpoint: str = Field(
+        default="http://localhost:4318/v1/traces",
+        description="The OpenTelemetry collector endpoint URL",
+    )
+    service_name: str = Field(
+        default="llama-stack",
+        description="The service name to use for telemetry",
+    )
+    sinks: List[TelemetrySink] = Field(
+        default=[TelemetrySink.CONSOLE, TelemetrySink.SQLITE],
+        description="List of telemetry sinks to enable (possible values: otel, sqlite, console)",
+    )
+    sqlite_db_path: str = Field(
+        default=(RUNTIME_BASE_DIR / "trace_store.db").as_posix(),
+        description="The path to the SQLite database to use for storing traces",
+    )
+
+    @classmethod
+    def sample_run_config(cls, **kwargs) -> Dict[str, Any]:
+        return {
+            "service_name": "${env.OTEL_SERVICE_NAME:llama-stack}",
+            "sinks": "${env.TELEMETRY_SINKS:['console', 'sqlite']}",
+            "sqlite_db_path": "${env.SQLITE_DB_PATH:${runtime.base_dir}/trace_store.db}",
+        }

llama_stack/providers/inline/telemetry/meta_reference/console_span_processor.py

@@ -0,0 +1,117 @@
+# 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
+from datetime import datetime
+
+from opentelemetry.sdk.trace import ReadableSpan
+from opentelemetry.sdk.trace.export import SpanProcessor
+from opentelemetry.trace.status import StatusCode
+
+# Colors for console output
+COLORS = {
+    "reset": "\033[0m",
+    "bold": "\033[1m",
+    "dim": "\033[2m",
+    "red": "\033[31m",
+    "green": "\033[32m",
+    "yellow": "\033[33m",
+    "blue": "\033[34m",
+    "magenta": "\033[35m",
+    "cyan": "\033[36m",
+    "white": "\033[37m",
+}
+
+
+class ConsoleSpanProcessor(SpanProcessor):
+
+    def __init__(self, print_attributes: bool = False):
+        self.print_attributes = print_attributes
+
+    def on_start(self, span: ReadableSpan, parent_context=None) -> None:
+        if span.attributes and span.attributes.get("__autotraced__"):
+            return
+
+        timestamp = datetime.utcfromtimestamp(span.start_time / 1e9).strftime(
+            "%H:%M:%S.%f"
+        )[:-3]
+
+        print(
+            f"{COLORS['dim']}{timestamp}{COLORS['reset']} "
+            f"{COLORS['magenta']}[START]{COLORS['reset']} "
+            f"{COLORS['dim']}{span.name}{COLORS['reset']}"
+        )
+
+    def on_end(self, span: ReadableSpan) -> None:
+        if span.attributes and span.attributes.get("__autotraced__"):
+            return
+
+        timestamp = datetime.utcfromtimestamp(span.end_time / 1e9).strftime(
+            "%H:%M:%S.%f"
+        )[:-3]
+
+        span_context = (
+            f"{COLORS['dim']}{timestamp}{COLORS['reset']} "
+            f"{COLORS['magenta']}[END]{COLORS['reset']} "
+            f"{COLORS['dim']}{span.name}{COLORS['reset']}"
+        )
+
+        if span.status.status_code == StatusCode.ERROR:
+            span_context += f"{COLORS['reset']} {COLORS['red']}[ERROR]{COLORS['reset']}"
+        elif span.status.status_code != StatusCode.UNSET:
+            span_context += f"{COLORS['reset']} [{span.status.status_code}]"
+
+        duration_ms = (span.end_time - span.start_time) / 1e6
+        span_context += f"{COLORS['reset']} ({duration_ms:.2f}ms)"
+
+        print(span_context)
+
+        if self.print_attributes and span.attributes:
+            for key, value in span.attributes.items():
+                if key.startswith("__"):
+                    continue
+                str_value = str(value)
+                if len(str_value) > 1000:
+                    str_value = str_value[:997] + "..."
+                print(f"    {COLORS['dim']}{key}: {str_value}{COLORS['reset']}")
+
+        for event in span.events:
+            event_time = datetime.utcfromtimestamp(event.timestamp / 1e9).strftime(
+                "%H:%M:%S.%f"
+            )[:-3]
+
+            severity = event.attributes.get("severity", "info")
+            message = event.attributes.get("message", event.name)
+            if isinstance(message, (dict, list)):
+                message = json.dumps(message, indent=2)
+
+            severity_colors = {
+                "error": f"{COLORS['bold']}{COLORS['red']}",
+                "warn": f"{COLORS['bold']}{COLORS['yellow']}",
+                "info": COLORS["white"],
+                "debug": COLORS["dim"],
+            }
+            msg_color = severity_colors.get(severity, COLORS["white"])
+
+            print(
+                f" {event_time} "
+                f"{msg_color}[{severity.upper()}] "
+                f"{message}{COLORS['reset']}"
+            )
+
+            if event.attributes:
+                for key, value in event.attributes.items():
+                    if key.startswith("__") or key in ["message", "severity"]:
+                        continue
+                    print(f"   {COLORS['dim']}{key}: {value}{COLORS['reset']}")
+
+    def shutdown(self) -> None:
+        """Shutdown the processor."""
+        pass
+
+    def force_flush(self, timeout_millis: float = None) -> bool:
+        """Force flush any pending spans."""
+        return True

llama_stack/providers/inline/telemetry/meta_reference/sqlite_span_processor.py

@@ -0,0 +1,242 @@
+# 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 os
+import sqlite3
+import threading
+from datetime import datetime, timedelta
+from typing import Dict
+
+from opentelemetry.sdk.trace import SpanProcessor
+from opentelemetry.trace import Span
+
+
+class SQLiteSpanProcessor(SpanProcessor):
+    def __init__(self, conn_string, ttl_days=30):
+        """Initialize the SQLite span processor with a connection string."""
+        self.conn_string = conn_string
+        self.ttl_days = ttl_days
+        self.cleanup_task = None
+        self._thread_local = threading.local()
+        self._connections: Dict[int, sqlite3.Connection] = {}
+        self._lock = threading.Lock()
+        self.setup_database()
+
+    def _get_connection(self) -> sqlite3.Connection:
+        """Get a thread-specific database connection."""
+        thread_id = threading.get_ident()
+        with self._lock:
+            if thread_id not in self._connections:
+                conn = sqlite3.connect(self.conn_string)
+                self._connections[thread_id] = conn
+            return self._connections[thread_id]
+
+    def setup_database(self):
+        """Create the necessary tables if they don't exist."""
+        # Create directory if it doesn't exist
+        os.makedirs(os.path.dirname(self.conn_string), exist_ok=True)
+
+        conn = self._get_connection()
+        cursor = conn.cursor()
+
+        cursor.execute(
+            """
+            CREATE TABLE IF NOT EXISTS traces (
+                trace_id TEXT PRIMARY KEY,
+                service_name TEXT,
+                root_span_id TEXT,
+                start_time TIMESTAMP,
+                end_time TIMESTAMP,
+                created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP
+            )
+        """
+        )
+
+        cursor.execute(
+            """
+            CREATE TABLE IF NOT EXISTS spans (
+                span_id TEXT PRIMARY KEY,
+                trace_id TEXT REFERENCES traces(trace_id),
+                parent_span_id TEXT,
+                name TEXT,
+                start_time TIMESTAMP,
+                end_time TIMESTAMP,
+                attributes TEXT,
+                status TEXT,
+                kind TEXT
+            )
+        """
+        )
+
+        cursor.execute(
+            """
+            CREATE TABLE IF NOT EXISTS span_events (
+                id INTEGER PRIMARY KEY AUTOINCREMENT,
+                span_id TEXT REFERENCES spans(span_id),
+                name TEXT,
+                timestamp TIMESTAMP,
+                attributes TEXT
+            )
+        """
+        )
+
+        cursor.execute(
+            """
+            CREATE INDEX IF NOT EXISTS idx_traces_created_at
+            ON traces(created_at)
+        """
+        )
+
+        conn.commit()
+        cursor.close()
+
+        # Start periodic cleanup in a separate thread
+        self.cleanup_task = threading.Thread(target=self._periodic_cleanup, daemon=True)
+        self.cleanup_task.start()
+
+    def _cleanup_old_data(self):
+        """Delete records older than TTL."""
+        try:
+            conn = self._get_connection()
+            cutoff_date = (datetime.now() - timedelta(days=self.ttl_days)).isoformat()
+            cursor = conn.cursor()
+
+            # Delete old span events
+            cursor.execute(
+                """
+                DELETE FROM span_events
+                WHERE span_id IN (
+                    SELECT span_id FROM spans
+                    WHERE trace_id IN (
+                        SELECT trace_id FROM traces
+                        WHERE created_at < ?
+                    )
+                )
+            """,
+                (cutoff_date,),
+            )
+
+            # Delete old spans
+            cursor.execute(
+                """
+                DELETE FROM spans
+                WHERE trace_id IN (
+                    SELECT trace_id FROM traces
+                    WHERE created_at < ?
+                )
+            """,
+                (cutoff_date,),
+            )
+
+            # Delete old traces
+            cursor.execute("DELETE FROM traces WHERE created_at < ?", (cutoff_date,))
+
+            conn.commit()
+            cursor.close()
+        except Exception as e:
+            print(f"Error during cleanup: {e}")
+
+    def _periodic_cleanup(self):
+        """Run cleanup periodically."""
+        import time
+
+        while True:
+            time.sleep(3600)  # Sleep for 1 hour
+            self._cleanup_old_data()
+
+    def on_start(self, span: Span, parent_context=None):
+        """Called when a span starts."""
+        pass
+
+    def on_end(self, span: Span):
+        """Called when a span ends. Export the span data to SQLite."""
+        try:
+            conn = self._get_connection()
+            cursor = conn.cursor()
+
+            trace_id = format(span.get_span_context().trace_id, "032x")
+            span_id = format(span.get_span_context().span_id, "016x")
+            service_name = span.resource.attributes.get("service.name", "unknown")
+
+            parent_span_id = None
+            parent_context = span.parent
+            if parent_context:
+                parent_span_id = format(parent_context.span_id, "016x")
+
+            # Insert into traces
+            cursor.execute(
+                """
+                INSERT INTO traces (
+                    trace_id, service_name, root_span_id, start_time, end_time
+                ) VALUES (?, ?, ?, ?, ?)
+                ON CONFLICT(trace_id) DO UPDATE SET
+                    root_span_id = COALESCE(root_span_id, excluded.root_span_id),
+                    start_time = MIN(excluded.start_time, start_time),
+                    end_time = MAX(excluded.end_time, end_time)
+            """,
+                (
+                    trace_id,
+                    service_name,
+                    (span_id if not parent_span_id else None),
+                    datetime.fromtimestamp(span.start_time / 1e9).isoformat(),
+                    datetime.fromtimestamp(span.end_time / 1e9).isoformat(),
+                ),
+            )
+
+            # Insert into spans
+            cursor.execute(
+                """
+                INSERT INTO spans (
+                    span_id, trace_id, parent_span_id, name,
+                    start_time, end_time, attributes, status,
+                    kind
+                ) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?)
+            """,
+                (
+                    span_id,
+                    trace_id,
+                    parent_span_id,
+                    span.name,
+                    datetime.fromtimestamp(span.start_time / 1e9).isoformat(),
+                    datetime.fromtimestamp(span.end_time / 1e9).isoformat(),
+                    json.dumps(dict(span.attributes)),
+                    span.status.status_code.name,
+                    span.kind.name,
+                ),
+            )
+
+            for event in span.events:
+                cursor.execute(
+                    """
+                    INSERT INTO span_events (
+                        span_id, name, timestamp, attributes
+                    ) VALUES (?, ?, ?, ?)
+                """,
+                    (
+                        span_id,
+                        event.name,
+                        datetime.fromtimestamp(event.timestamp / 1e9).isoformat(),
+                        json.dumps(dict(event.attributes)),
+                    ),
+                )
+
+            conn.commit()
+            cursor.close()
+        except Exception as e:
+            print(f"Error exporting span to SQLite: {e}")
+
+    def shutdown(self):
+        """Cleanup any resources."""
+        with self._lock:
+            for conn in self._connections.values():
+                if conn:
+                    conn.close()
+            self._connections.clear()
+
+    def force_flush(self, timeout_millis=30000):
+        """Force export of spans."""
+        pass

llama_stack/providers/inline/telemetry/meta_reference/telemetry.py

@@ -5,6 +5,7 @@
 # the root directory of this source tree.
 
 import threading
+from typing import Any, Dict, List, Optional
 
 from opentelemetry import metrics, trace
 from opentelemetry.exporter.otlp.proto.http.metric_exporter import OTLPMetricExporter
@@ -16,10 +17,23 @@ from opentelemetry.sdk.trace import TracerProvider
 from opentelemetry.sdk.trace.export import BatchSpanProcessor
 from opentelemetry.semconv.resource import ResourceAttributes
 
+from llama_stack.providers.inline.telemetry.meta_reference.console_span_processor import (
+    ConsoleSpanProcessor,
+)
+
+from llama_stack.providers.inline.telemetry.meta_reference.sqlite_span_processor import (
+    SQLiteSpanProcessor,
+)
+from llama_stack.providers.utils.telemetry import (
+    SQLiteTraceStore,
+    TelemetryDatasetMixin,
+)
 
 from llama_stack.apis.telemetry import *  # noqa: F403
 
-from .config import OpenTelemetryConfig
+from llama_stack.distribution.datatypes import Api
+
+from .config import TelemetryConfig, TelemetrySink
 
 _GLOBAL_STORAGE = {
     "active_spans": {},
@@ -45,9 +59,10 @@ def is_tracing_enabled(tracer):
         return span.is_recording()
 
 
-class OpenTelemetryAdapter(Telemetry):
-    def __init__(self, config: OpenTelemetryConfig):
+class TelemetryAdapter(TelemetryDatasetMixin, Telemetry):
+    def __init__(self, config: TelemetryConfig, deps: Dict[str, Any]) -> None:
         self.config = config
+        self.datasetio_api = deps[Api.datasetio]
 
         resource = Resource.create(
             {
@@ -57,22 +72,29 @@ class OpenTelemetryAdapter(Telemetry):
 
         provider = TracerProvider(resource=resource)
         trace.set_tracer_provider(provider)
-        otlp_exporter = OTLPSpanExporter(
-            endpoint=self.config.otel_endpoint,
-        )
-        span_processor = BatchSpanProcessor(otlp_exporter)
-        trace.get_tracer_provider().add_span_processor(span_processor)
-        # Set up metrics
-        metric_reader = PeriodicExportingMetricReader(
-            OTLPMetricExporter(
+        if TelemetrySink.OTEL in self.config.sinks:
+            otlp_exporter = OTLPSpanExporter(
                 endpoint=self.config.otel_endpoint,
             )
-        )
-        metric_provider = MeterProvider(
-            resource=resource, metric_readers=[metric_reader]
-        )
-        metrics.set_meter_provider(metric_provider)
-        self.meter = metrics.get_meter(__name__)
+            span_processor = BatchSpanProcessor(otlp_exporter)
+            trace.get_tracer_provider().add_span_processor(span_processor)
+            metric_reader = PeriodicExportingMetricReader(
+                OTLPMetricExporter(
+                    endpoint=self.config.otel_endpoint,
+                )
+            )
+            metric_provider = MeterProvider(
+                resource=resource, metric_readers=[metric_reader]
+            )
+            metrics.set_meter_provider(metric_provider)
+            self.meter = metrics.get_meter(__name__)
+        if TelemetrySink.SQLITE in self.config.sinks:
+            trace.get_tracer_provider().add_span_processor(
+                SQLiteSpanProcessor(self.config.sqlite_db_path)
+            )
+            self.trace_store = SQLiteTraceStore(self.config.sqlite_db_path)
+        if TelemetrySink.CONSOLE in self.config.sinks:
+            trace.get_tracer_provider().add_span_processor(ConsoleSpanProcessor())
         self._lock = _global_lock
 
     async def initialize(self) -> None:
@@ -83,15 +105,17 @@ class OpenTelemetryAdapter(Telemetry):
         trace.get_tracer_provider().shutdown()
         metrics.get_meter_provider().shutdown()
 
-    async def log_event(self, event: Event) -> None:
+    async def log_event(self, event: Event, ttl_seconds: int = 604800) -> None:
         if isinstance(event, UnstructuredLogEvent):
-            self._log_unstructured(event)
+            self._log_unstructured(event, ttl_seconds)
         elif isinstance(event, MetricEvent):
             self._log_metric(event)
         elif isinstance(event, StructuredLogEvent):
-            self._log_structured(event)
+            self._log_structured(event, ttl_seconds)
+        else:
+            raise ValueError(f"Unknown event type: {event}")
 
-    def _log_unstructured(self, event: UnstructuredLogEvent) -> None:
+    def _log_unstructured(self, event: UnstructuredLogEvent, ttl_seconds: int) -> None:
         with self._lock:
             # Use global storage instead of instance storage
             span_id = string_to_span_id(event.span_id)
@@ -104,6 +128,7 @@ class OpenTelemetryAdapter(Telemetry):
                     attributes={
                         "message": event.message,
                         "severity": event.severity.value,
+                        "__ttl__": ttl_seconds,
                         **event.attributes,
                     },
                     timestamp=timestamp_ns,
@@ -154,11 +179,14 @@ class OpenTelemetryAdapter(Telemetry):
             )
         return _GLOBAL_STORAGE["up_down_counters"][name]
 
-    def _log_structured(self, event: StructuredLogEvent) -> None:
+    def _log_structured(self, event: StructuredLogEvent, ttl_seconds: int) -> None:
         with self._lock:
             span_id = string_to_span_id(event.span_id)
             trace_id = string_to_trace_id(event.trace_id)
             tracer = trace.get_tracer(__name__)
+            if event.attributes is None:
+                event.attributes = {}
+            event.attributes["__ttl__"] = ttl_seconds
 
             if isinstance(event.payload, SpanStartPayload):
                 # Check if span already exists to prevent duplicates
@@ -170,7 +198,6 @@ class OpenTelemetryAdapter(Telemetry):
                     parent_span_id = string_to_span_id(event.payload.parent_span_id)
                     parent_span = _GLOBAL_STORAGE["active_spans"].get(parent_span_id)
 
-                # Create a new trace context with the trace_id
                 context = trace.Context(trace_id=trace_id)
                 if parent_span:
                     context = trace.set_span_in_context(parent_span, context)
@@ -179,14 +206,9 @@ class OpenTelemetryAdapter(Telemetry):
                     name=event.payload.name,
                     context=context,
                     attributes=event.attributes or {},
-                    start_time=int(event.timestamp.timestamp() * 1e9),
                 )
                 _GLOBAL_STORAGE["active_spans"][span_id] = span
 
-                # Set as current span using context manager
-                with trace.use_span(span, end_on_exit=False):
-                    pass  # Let the span continue beyond this block
-
             elif isinstance(event.payload, SpanEndPayload):
                 span = _GLOBAL_STORAGE["active_spans"].get(span_id)
                 if span:
@@ -199,10 +221,33 @@ class OpenTelemetryAdapter(Telemetry):
                         else trace.Status(status_code=trace.StatusCode.ERROR)
                     )
                     span.set_status(status)
-                    span.end(end_time=int(event.timestamp.timestamp() * 1e9))
-
-                    # Remove from active spans
+                    span.end()
                     _GLOBAL_STORAGE["active_spans"].pop(span_id, None)
+            else:
+                raise ValueError(f"Unknown structured log event: {event}")
 
-    async def get_trace(self, trace_id: str) -> Trace:
-        raise NotImplementedError("Trace retrieval not implemented yet")
+    async def query_traces(
+        self,
+        attribute_filters: Optional[List[QueryCondition]] = None,
+        limit: Optional[int] = 100,
+        offset: Optional[int] = 0,
+        order_by: Optional[List[str]] = None,
+    ) -> List[Trace]:
+        return await self.trace_store.query_traces(
+            attribute_filters=attribute_filters,
+            limit=limit,
+            offset=offset,
+            order_by=order_by,
+        )
+
+    async def get_span_tree(
+        self,
+        span_id: str,
+        attributes_to_return: Optional[List[str]] = None,
+        max_depth: Optional[int] = None,
+    ) -> SpanWithChildren:
+        return await self.trace_store.get_span_tree(
+            span_id=span_id,
+            attributes_to_return=attributes_to_return,
+            max_depth=max_depth,
+        )

llama_stack/providers/registry/inference.py

@@ -61,6 +61,17 @@ def available_providers() -> List[ProviderSpec]:
                 config_class="llama_stack.providers.remote.inference.sample.SampleConfig",
             ),
         ),
+        remote_provider_spec(
+            api=Api.inference,
+            adapter=AdapterSpec(
+                adapter_type="cerebras",
+                pip_packages=[
+                    "cerebras_cloud_sdk",
+                ],
+                module="llama_stack.providers.remote.inference.cerebras",
+                config_class="llama_stack.providers.remote.inference.cerebras.CerebrasImplConfig",
+            ),
+        ),
         remote_provider_spec(
             api=Api.inference,
             adapter=AdapterSpec(

llama_stack/providers/registry/scoring.py

@@ -44,5 +44,6 @@ def available_providers() -> List[ProviderSpec]:
                 Api.datasetio,
                 Api.datasets,
             ],
+            provider_data_validator="llama_stack.providers.inline.scoring.braintrust.BraintrustProviderDataValidator",
         ),
     ]

llama_stack/providers/registry/telemetry.py

@@ -14,9 +14,13 @@ def available_providers() -> List[ProviderSpec]:
         InlineProviderSpec(
             api=Api.telemetry,
             provider_type="inline::meta-reference",
-            pip_packages=[],
-            module="llama_stack.providers.inline.meta_reference.telemetry",
-            config_class="llama_stack.providers.inline.meta_reference.telemetry.ConsoleConfig",
+            pip_packages=[
+                "opentelemetry-sdk",
+                "opentelemetry-exporter-otlp-proto-http",
+            ],
+            api_dependencies=[Api.datasetio],
+            module="llama_stack.providers.inline.telemetry.meta_reference",
+            config_class="llama_stack.providers.inline.telemetry.meta_reference.config.TelemetryConfig",
         ),
         remote_provider_spec(
             api=Api.telemetry,
@@ -27,18 +31,4 @@ def available_providers() -> List[ProviderSpec]:
                 config_class="llama_stack.providers.remote.telemetry.sample.SampleConfig",
             ),
         ),
-        remote_provider_spec(
-            api=Api.telemetry,
-            adapter=AdapterSpec(
-                adapter_type="opentelemetry-jaeger",
-                pip_packages=[
-                    "opentelemetry-api",
-                    "opentelemetry-sdk",
-                    "opentelemetry-exporter-jaeger",
-                    "opentelemetry-semantic-conventions",
-                ],
-                module="llama_stack.providers.remote.telemetry.opentelemetry",
-                config_class="llama_stack.providers.remote.telemetry.opentelemetry.OpenTelemetryConfig",
-            ),
-        ),
     ]

llama_stack/providers/remote/datasetio/__init__.py

@@ -0,0 +1,5 @@
+# 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.

llama_stack/providers/remote/datasetio/huggingface/config.py

@@ -3,12 +3,13 @@
 #
 # This source code is licensed under the terms described in the LICENSE file in
 # the root directory of this source tree.
+from pydantic import BaseModel
+
 from llama_stack.distribution.utils.config_dirs import RUNTIME_BASE_DIR
 from llama_stack.providers.utils.kvstore.config import (
     KVStoreConfig,
     SqliteKVStoreConfig,
 )
-from pydantic import BaseModel
 
 
 class HuggingfaceDatasetIOConfig(BaseModel):

llama_stack/providers/remote/datasetio/huggingface/huggingface.py

@@ -3,12 +3,13 @@
 #
 # This source code is licensed under the terms described in the LICENSE file in
 # the root directory of this source tree.
-from typing import Optional
+from typing import Any, Dict, List, Optional
 
 from llama_stack.apis.datasetio import *  # noqa: F403
 
 
 import datasets as hf_datasets
+
 from llama_stack.providers.datatypes import DatasetsProtocolPrivate
 from llama_stack.providers.utils.datasetio.url_utils import get_dataframe_from_url
 from llama_stack.providers.utils.kvstore import kvstore_impl
@@ -63,6 +64,11 @@ class HuggingfaceDatasetIOImpl(DatasetIO, DatasetsProtocolPrivate):
         )
         self.dataset_infos[dataset_def.identifier] = dataset_def
 
+    async def unregister_dataset(self, dataset_id: str) -> None:
+        key = f"{DATASETS_PREFIX}{dataset_id}"
+        await self.kvstore.delete(key=key)
+        del self.dataset_infos[dataset_id]
+
     async def get_rows_paginated(
         self,
         dataset_id: str,
@@ -94,3 +100,22 @@ class HuggingfaceDatasetIOImpl(DatasetIO, DatasetsProtocolPrivate):
             total_count=len(rows),
             next_page_token=str(end),
         )
+
+    async def append_rows(self, dataset_id: str, rows: List[Dict[str, Any]]) -> None:
+        dataset_def = self.dataset_infos[dataset_id]
+        loaded_dataset = load_hf_dataset(dataset_def)
+
+        # Convert rows to HF Dataset format
+        new_dataset = hf_datasets.Dataset.from_list(rows)
+
+        # Concatenate the new rows with existing dataset
+        updated_dataset = hf_datasets.concatenate_datasets(
+            [loaded_dataset, new_dataset]
+        )
+
+        if dataset_def.metadata.get("path", None):
+            updated_dataset.push_to_hub(dataset_def.metadata["path"])
+        else:
+            raise NotImplementedError(
+                "Uploading to URL-based datasets is not supported yet"
+            )

llama_stack/providers/remote/inference/cerebras/__init__.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 .config import CerebrasImplConfig
+
+
+async def get_adapter_impl(config: CerebrasImplConfig, _deps):
+    from .cerebras import CerebrasInferenceAdapter
+
+    assert isinstance(
+        config, CerebrasImplConfig
+    ), f"Unexpected config type: {type(config)}"
+
+    impl = CerebrasInferenceAdapter(config)
+
+    await impl.initialize()
+
+    return impl

llama_stack/providers/remote/inference/cerebras/cerebras.py

@@ -0,0 +1,191 @@
+# 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 typing import AsyncGenerator
+
+from cerebras.cloud.sdk import AsyncCerebras
+
+from llama_models.llama3.api.chat_format import ChatFormat
+
+from llama_models.llama3.api.datatypes import Message
+from llama_models.llama3.api.tokenizer import Tokenizer
+
+from llama_stack.apis.inference import *  # noqa: F403
+
+from llama_models.datatypes import CoreModelId
+
+from llama_stack.providers.utils.inference.model_registry import (
+    build_model_alias,
+    ModelRegistryHelper,
+)
+from llama_stack.providers.utils.inference.openai_compat import (
+    get_sampling_options,
+    process_chat_completion_response,
+    process_chat_completion_stream_response,
+    process_completion_response,
+    process_completion_stream_response,
+)
+from llama_stack.providers.utils.inference.prompt_adapter import (
+    chat_completion_request_to_prompt,
+    completion_request_to_prompt,
+)
+
+from .config import CerebrasImplConfig
+
+
+model_aliases = [
+    build_model_alias(
+        "llama3.1-8b",
+        CoreModelId.llama3_1_8b_instruct.value,
+    ),
+    build_model_alias(
+        "llama3.1-70b",
+        CoreModelId.llama3_1_70b_instruct.value,
+    ),
+]
+
+
+class CerebrasInferenceAdapter(ModelRegistryHelper, Inference):
+    def __init__(self, config: CerebrasImplConfig) -> None:
+        ModelRegistryHelper.__init__(
+            self,
+            model_aliases=model_aliases,
+        )
+        self.config = config
+        self.formatter = ChatFormat(Tokenizer.get_instance())
+
+        self.client = AsyncCerebras(
+            base_url=self.config.base_url, api_key=self.config.api_key
+        )
+
+    async def initialize(self) -> None:
+        return
+
+    async def shutdown(self) -> None:
+        pass
+
+    async def completion(
+        self,
+        model_id: str,
+        content: InterleavedTextMedia,
+        sampling_params: Optional[SamplingParams] = SamplingParams(),
+        response_format: Optional[ResponseFormat] = None,
+        stream: Optional[bool] = False,
+        logprobs: Optional[LogProbConfig] = None,
+    ) -> AsyncGenerator:
+        model = await self.model_store.get_model(model_id)
+        request = CompletionRequest(
+            model=model.provider_resource_id,
+            content=content,
+            sampling_params=sampling_params,
+            response_format=response_format,
+            stream=stream,
+            logprobs=logprobs,
+        )
+        if stream:
+            return self._stream_completion(
+                request,
+            )
+        else:
+            return await self._nonstream_completion(request)
+
+    async def _nonstream_completion(
+        self, request: CompletionRequest
+    ) -> CompletionResponse:
+        params = self._get_params(request)
+
+        r = await self.client.completions.create(**params)
+
+        return process_completion_response(r, self.formatter)
+
+    async def _stream_completion(self, request: CompletionRequest) -> AsyncGenerator:
+        params = self._get_params(request)
+
+        stream = await self.client.completions.create(**params)
+
+        async for chunk in process_completion_stream_response(stream, self.formatter):
+            yield chunk
+
+    async def chat_completion(
+        self,
+        model_id: str,
+        messages: List[Message],
+        sampling_params: Optional[SamplingParams] = SamplingParams(),
+        tools: Optional[List[ToolDefinition]] = None,
+        tool_choice: Optional[ToolChoice] = ToolChoice.auto,
+        tool_prompt_format: Optional[ToolPromptFormat] = ToolPromptFormat.json,
+        response_format: Optional[ResponseFormat] = None,
+        stream: Optional[bool] = False,
+        logprobs: Optional[LogProbConfig] = None,
+    ) -> AsyncGenerator:
+        model = await self.model_store.get_model(model_id)
+        request = ChatCompletionRequest(
+            model=model.provider_resource_id,
+            messages=messages,
+            sampling_params=sampling_params,
+            tools=tools or [],
+            tool_choice=tool_choice,
+            tool_prompt_format=tool_prompt_format,
+            response_format=response_format,
+            stream=stream,
+            logprobs=logprobs,
+        )
+
+        if stream:
+            return self._stream_chat_completion(request)
+        else:
+            return await self._nonstream_chat_completion(request)
+
+    async def _nonstream_chat_completion(
+        self, request: CompletionRequest
+    ) -> CompletionResponse:
+        params = self._get_params(request)
+
+        r = await self.client.completions.create(**params)
+
+        return process_chat_completion_response(r, self.formatter)
+
+    async def _stream_chat_completion(
+        self, request: CompletionRequest
+    ) -> AsyncGenerator:
+        params = self._get_params(request)
+
+        stream = await self.client.completions.create(**params)
+
+        async for chunk in process_chat_completion_stream_response(
+            stream, self.formatter
+        ):
+            yield chunk
+
+    def _get_params(
+        self, request: Union[ChatCompletionRequest, CompletionRequest]
+    ) -> dict:
+        if request.sampling_params and request.sampling_params.top_k:
+            raise ValueError("`top_k` not supported by Cerebras")
+
+        prompt = ""
+        if type(request) == ChatCompletionRequest:
+            prompt = chat_completion_request_to_prompt(
+                request, self.get_llama_model(request.model), self.formatter
+            )
+        elif type(request) == CompletionRequest:
+            prompt = completion_request_to_prompt(request, self.formatter)
+        else:
+            raise ValueError(f"Unknown request type {type(request)}")
+
+        return {
+            "model": request.model,
+            "prompt": prompt,
+            "stream": request.stream,
+            **get_sampling_options(request.sampling_params),
+        }
+
+    async def embeddings(
+        self,
+        model_id: str,
+        contents: List[InterleavedTextMedia],
+    ) -> EmbeddingsResponse:
+        raise NotImplementedError()

llama_stack/providers/remote/inference/cerebras/config.py

@@ -0,0 +1,32 @@
+# 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 typing import Any, Dict, Optional
+
+from llama_models.schema_utils import json_schema_type
+from pydantic import BaseModel, Field
+
+DEFAULT_BASE_URL = "https://api.cerebras.ai"
+
+
+@json_schema_type
+class CerebrasImplConfig(BaseModel):
+    base_url: str = Field(
+        default=os.environ.get("CEREBRAS_BASE_URL", DEFAULT_BASE_URL),
+        description="Base URL for the Cerebras API",
+    )
+    api_key: Optional[str] = Field(
+        default=os.environ.get("CEREBRAS_API_KEY"),
+        description="Cerebras API Key",
+    )
+
+    @classmethod
+    def sample_run_config(cls, **kwargs) -> Dict[str, Any]:
+        return {
+            "base_url": DEFAULT_BASE_URL,
+            "api_key": "${env.CEREBRAS_API_KEY}",
+        }

llama_stack/providers/remote/inference/nvidia/config.py

@@ -35,7 +35,9 @@ class NVIDIAConfig(BaseModel):
     """
 
     url: str = Field(
-        default="https://integrate.api.nvidia.com",
+        default_factory=lambda: os.getenv(
+            "NVIDIA_BASE_URL", "https://integrate.api.nvidia.com"
+        ),
         description="A base url for accessing the NVIDIA NIM",
     )
     api_key: Optional[str] = Field(

llama_stack/providers/remote/inference/ollama/ollama.py

@@ -180,7 +180,6 @@ class OllamaInferenceAdapter(Inference, ModelsProtocolPrivate):
     async def _nonstream_completion(self, request: CompletionRequest) -> AsyncGenerator:
         params = await self._get_params(request)
         r = await self.client.generate(**params)
-        assert isinstance(r, dict)
 
         choice = OpenAICompatCompletionChoice(
             finish_reason=r["done_reason"] if r["done"] else None,
@@ -270,7 +269,6 @@ class OllamaInferenceAdapter(Inference, ModelsProtocolPrivate):
             r = await self.client.chat(**params)
         else:
             r = await self.client.generate(**params)
-        assert isinstance(r, dict)
 
         if "message" in r:
             choice = OpenAICompatCompletionChoice(

llama_stack/providers/remote/inference/tgi/tgi.py

@@ -89,8 +89,9 @@ class _HfAdapter(Inference, ModelsProtocolPrivate):
         stream: Optional[bool] = False,
         logprobs: Optional[LogProbConfig] = None,
     ) -> AsyncGenerator:
+        model = await self.model_store.get_model(model_id)
         request = CompletionRequest(
-            model=model_id,
+            model=model.provider_resource_id,
             content=content,
             sampling_params=sampling_params,
             response_format=response_format,
@@ -194,8 +195,9 @@ class _HfAdapter(Inference, ModelsProtocolPrivate):
         stream: Optional[bool] = False,
         logprobs: Optional[LogProbConfig] = None,
     ) -> AsyncGenerator:
+        model = await self.model_store.get_model(model_id)
         request = ChatCompletionRequest(
-            model=model_id,
+            model=model.provider_resource_id,
             messages=messages,
             sampling_params=sampling_params,
             tools=tools or [],
@@ -249,7 +251,7 @@ class _HfAdapter(Inference, ModelsProtocolPrivate):
 
     def _get_params(self, request: ChatCompletionRequest) -> dict:
         prompt, input_tokens = chat_completion_request_to_model_input_info(
-            request, self.formatter
+            request, self.register_helper.get_llama_model(request.model), self.formatter
         )
         return dict(
             prompt=prompt,

llama_stack/providers/remote/telemetry/opentelemetry/__init__.py

@@ -1,15 +0,0 @@
-# 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 OpenTelemetryConfig
-
-
-async def get_adapter_impl(config: OpenTelemetryConfig, _deps):
-    from .opentelemetry import OpenTelemetryAdapter
-
-    impl = OpenTelemetryAdapter(config)
-    await impl.initialize()
-    return impl

llama_stack/providers/remote/telemetry/opentelemetry/config.py

@@ -1,27 +0,0 @@
-# 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 typing import Any, Dict
-
-from pydantic import BaseModel, Field
-
-
-class OpenTelemetryConfig(BaseModel):
-    otel_endpoint: str = Field(
-        default="http://localhost:4318/v1/traces",
-        description="The OpenTelemetry collector endpoint URL",
-    )
-    service_name: str = Field(
-        default="llama-stack",
-        description="The service name to use for telemetry",
-    )
-
-    @classmethod
-    def sample_run_config(cls, **kwargs) -> Dict[str, Any]:
-        return {
-            "otel_endpoint": "${env.OTEL_ENDPOINT:http://localhost:4318/v1/traces}",
-            "service_name": "${env.OTEL_SERVICE_NAME:llama-stack}",
-        }

llama_stack/providers/tests/datasetio/test_datasetio.py

@@ -81,6 +81,18 @@ class TestDatasetIO:
         assert len(response) == 1
         assert response[0].identifier == "test_dataset"
 
+        with pytest.raises(Exception) as exc_info:
+            # unregister a dataset that does not exist
+            await datasets_impl.unregister_dataset("test_dataset2")
+
+        await datasets_impl.unregister_dataset("test_dataset")
+        response = await datasets_impl.list_datasets()
+        assert isinstance(response, list)
+        assert len(response) == 0
+
+        with pytest.raises(Exception) as exc_info:
+            await datasets_impl.unregister_dataset("test_dataset")
+
     @pytest.mark.asyncio
     async def test_get_rows_paginated(self, datasetio_stack):
         datasetio_impl, datasets_impl = datasetio_stack

llama_stack/providers/tests/eval/conftest.py

@@ -6,10 +6,14 @@
 
 import pytest
 
+from ..agents.fixtures import AGENTS_FIXTURES
+
 from ..conftest import get_provider_fixture_overrides
 
 from ..datasetio.fixtures import DATASETIO_FIXTURES
 from ..inference.fixtures import INFERENCE_FIXTURES
+from ..memory.fixtures import MEMORY_FIXTURES
+from ..safety.fixtures import SAFETY_FIXTURES
 from ..scoring.fixtures import SCORING_FIXTURES
 from .fixtures import EVAL_FIXTURES
 
@@ -20,6 +24,9 @@ DEFAULT_PROVIDER_COMBINATIONS = [
             "scoring": "basic",
             "datasetio": "localfs",
             "inference": "fireworks",
+            "agents": "meta_reference",
+            "safety": "llama_guard",
+            "memory": "faiss",
         },
         id="meta_reference_eval_fireworks_inference",
         marks=pytest.mark.meta_reference_eval_fireworks_inference,
@@ -30,6 +37,9 @@ DEFAULT_PROVIDER_COMBINATIONS = [
             "scoring": "basic",
             "datasetio": "localfs",
             "inference": "together",
+            "agents": "meta_reference",
+            "safety": "llama_guard",
+            "memory": "faiss",
         },
         id="meta_reference_eval_together_inference",
         marks=pytest.mark.meta_reference_eval_together_inference,
@@ -40,6 +50,9 @@ DEFAULT_PROVIDER_COMBINATIONS = [
             "scoring": "basic",
             "datasetio": "huggingface",
             "inference": "together",
+            "agents": "meta_reference",
+            "safety": "llama_guard",
+            "memory": "faiss",
         },
         id="meta_reference_eval_together_inference_huggingface_datasetio",
         marks=pytest.mark.meta_reference_eval_together_inference_huggingface_datasetio,
@@ -75,6 +88,9 @@ def pytest_generate_tests(metafunc):
             "scoring": SCORING_FIXTURES,
             "datasetio": DATASETIO_FIXTURES,
             "inference": INFERENCE_FIXTURES,
+            "agents": AGENTS_FIXTURES,
+            "safety": SAFETY_FIXTURES,
+            "memory": MEMORY_FIXTURES,
         }
         combinations = (
             get_provider_fixture_overrides(metafunc.config, available_fixtures)

llama_stack/providers/tests/eval/fixtures.py

@@ -40,14 +40,30 @@ async def eval_stack(request):
 
     providers = {}
     provider_data = {}
-    for key in ["datasetio", "eval", "scoring", "inference"]:
+    for key in [
+        "datasetio",
+        "eval",
+        "scoring",
+        "inference",
+        "agents",
+        "safety",
+        "memory",
+    ]:
         fixture = request.getfixturevalue(f"{key}_{fixture_dict[key]}")
         providers[key] = fixture.providers
         if fixture.provider_data:
             provider_data.update(fixture.provider_data)
 
     test_stack = await construct_stack_for_test(
-        [Api.eval, Api.datasetio, Api.inference, Api.scoring],
+        [
+            Api.eval,
+            Api.datasetio,
+            Api.inference,
+            Api.scoring,
+            Api.agents,
+            Api.safety,
+            Api.memory,
+        ],
         providers,
         provider_data,
     )

llama_stack/providers/tests/inference/fixtures.py

@@ -17,9 +17,11 @@ from llama_stack.providers.inline.inference.meta_reference import (
 )
 from llama_stack.providers.remote.inference.bedrock import BedrockConfig
 
+from llama_stack.providers.remote.inference.cerebras import CerebrasImplConfig
 from llama_stack.providers.remote.inference.fireworks import FireworksImplConfig
 from llama_stack.providers.remote.inference.nvidia import NVIDIAConfig
 from llama_stack.providers.remote.inference.ollama import OllamaImplConfig
+from llama_stack.providers.remote.inference.tgi import TGIImplConfig
 from llama_stack.providers.remote.inference.together import TogetherImplConfig
 from llama_stack.providers.remote.inference.vllm import VLLMInferenceAdapterConfig
 from llama_stack.providers.tests.resolver import construct_stack_for_test
@@ -63,6 +65,21 @@ def inference_meta_reference(inference_model) -> ProviderFixture:
     )
 
 
+@pytest.fixture(scope="session")
+def inference_cerebras() -> ProviderFixture:
+    return ProviderFixture(
+        providers=[
+            Provider(
+                provider_id="cerebras",
+                provider_type="remote::cerebras",
+                config=CerebrasImplConfig(
+                    api_key=get_env_or_fail("CEREBRAS_API_KEY"),
+                ).model_dump(),
+            )
+        ],
+    )
+
+
 @pytest.fixture(scope="session")
 def inference_ollama(inference_model) -> ProviderFixture:
     inference_model = (
@@ -156,6 +173,22 @@ def inference_nvidia() -> ProviderFixture:
     )
 
 
+@pytest.fixture(scope="session")
+def inference_tgi() -> ProviderFixture:
+    return ProviderFixture(
+        providers=[
+            Provider(
+                provider_id="tgi",
+                provider_type="remote::tgi",
+                config=TGIImplConfig(
+                    url=get_env_or_fail("TGI_URL"),
+                    api_token=os.getenv("TGI_API_TOKEN", None),
+                ).model_dump(),
+            )
+        ],
+    )
+
+
 def get_model_short_name(model_name: str) -> str:
     """Convert model name to a short test identifier.
 
@@ -189,7 +222,9 @@ INFERENCE_FIXTURES = [
     "vllm_remote",
     "remote",
     "bedrock",
+    "cerebras",
     "nvidia",
+    "tgi",
 ]
 
 

llama_stack/providers/tests/inference/test_text_inference.py

@@ -94,6 +94,7 @@ class TestInference:
             "remote::tgi",
             "remote::together",
             "remote::fireworks",
+            "remote::cerebras",
         ):
             pytest.skip("Other inference providers don't support completion() yet")
 
@@ -140,6 +141,8 @@ class TestInference:
             "remote::together",
             "remote::fireworks",
             "remote::vllm",
+            "remote::cerebras",
+
         ):
             pytest.skip(
                 "Other inference providers don't support structured output in completions yet"
@@ -213,7 +216,15 @@ class TestInference:
         response = await inference_impl.chat_completion(
             model_id=inference_model,
             messages=[
-                SystemMessage(content="You are a helpful assistant."),
+                # we include context about Michael Jordan in the prompt so that the test is
+                # focused on the funtionality of the model and not on the information embedded
+                # in the model. Llama 3.2 3B Instruct tends to think MJ played for 14 seasons.
+                SystemMessage(
+                    content=(
+                        "You are a helpful assistant.\n\n"
+                        "Michael Jordan was born in 1963. He played basketball for the Chicago Bulls for 15 seasons."
+                    )
+                ),
                 UserMessage(content="Please give me information about Michael Jordan."),
             ],
             stream=False,

llama_stack/providers/tests/scoring/fixtures.py

@@ -10,9 +10,10 @@ import pytest_asyncio
 from llama_stack.apis.models import ModelInput
 
 from llama_stack.distribution.datatypes import Api, Provider
-
+from llama_stack.providers.inline.scoring.braintrust import BraintrustScoringConfig
 from llama_stack.providers.tests.resolver import construct_stack_for_test
 from ..conftest import ProviderFixture, remote_stack_fixture
+from ..env import get_env_or_fail
 
 
 @pytest.fixture(scope="session")
@@ -40,7 +41,9 @@ def scoring_braintrust() -> ProviderFixture:
             Provider(
                 provider_id="braintrust",
                 provider_type="inline::braintrust",
-                config={},
+                config=BraintrustScoringConfig(
+                    openai_api_key=get_env_or_fail("OPENAI_API_KEY"),
+                ).model_dump(),
             )
         ],
     )

llama_stack/providers/utils/scoring/__init__.py

@@ -0,0 +1,5 @@
+# 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.

llama_stack/providers/utils/telemetry/__init__.py

@@ -3,3 +3,6 @@
 #
 # This source code is licensed under the terms described in the LICENSE file in
 # the root directory of this source tree.
+
+from .dataset_mixin import TelemetryDatasetMixin  # noqa: F401
+from .sqlite_trace_store import SQLiteTraceStore, TraceStore  # noqa: F401

llama_stack/providers/utils/telemetry/dataset_mixin.py

@@ -0,0 +1,87 @@
+# 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 typing import List, Optional
+
+from llama_stack.apis.datasetio import DatasetIO
+from llama_stack.apis.telemetry import QueryCondition, Span, SpanWithChildren
+
+
+class TelemetryDatasetMixin:
+    """Mixin class that provides dataset-related functionality for telemetry providers."""
+
+    datasetio_api: DatasetIO
+
+    async def save_spans_to_dataset(
+        self,
+        attribute_filters: List[QueryCondition],
+        attributes_to_save: List[str],
+        dataset_id: str,
+        max_depth: Optional[int] = None,
+    ) -> None:
+        spans = await self.query_spans(
+            attribute_filters=attribute_filters,
+            attributes_to_return=attributes_to_save,
+            max_depth=max_depth,
+        )
+
+        rows = [
+            {
+                "trace_id": span.trace_id,
+                "span_id": span.span_id,
+                "parent_span_id": span.parent_span_id,
+                "name": span.name,
+                "start_time": span.start_time,
+                "end_time": span.end_time,
+                **{attr: span.attributes.get(attr) for attr in attributes_to_save},
+            }
+            for span in spans
+        ]
+
+        await self.datasetio_api.append_rows(dataset_id=dataset_id, rows=rows)
+
+    async def query_spans(
+        self,
+        attribute_filters: List[QueryCondition],
+        attributes_to_return: List[str],
+        max_depth: Optional[int] = None,
+    ) -> List[Span]:
+        traces = await self.query_traces(attribute_filters=attribute_filters)
+        spans = []
+
+        for trace in traces:
+            span_tree = await self.get_span_tree(
+                span_id=trace.root_span_id,
+                attributes_to_return=attributes_to_return,
+                max_depth=max_depth,
+            )
+
+            def extract_spans(span: SpanWithChildren) -> List[Span]:
+                result = []
+                if span.attributes and all(
+                    attr in span.attributes and span.attributes[attr] is not None
+                    for attr in attributes_to_return
+                ):
+                    result.append(
+                        Span(
+                            trace_id=trace.root_span_id,
+                            span_id=span.span_id,
+                            parent_span_id=span.parent_span_id,
+                            name=span.name,
+                            start_time=span.start_time,
+                            end_time=span.end_time,
+                            attributes=span.attributes,
+                        )
+                    )
+
+                for child in span.children:
+                    result.extend(extract_spans(child))
+
+                return result
+
+            spans.extend(extract_spans(span_tree))
+
+        return spans

llama_stack/providers/utils/telemetry/sqlite_trace_store.py

@@ -0,0 +1,180 @@
+# 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
+from datetime import datetime
+from typing import List, Optional, Protocol
+
+import aiosqlite
+
+from llama_stack.apis.telemetry import QueryCondition, SpanWithChildren, Trace
+
+
+class TraceStore(Protocol):
+
+    async def query_traces(
+        self,
+        attribute_filters: Optional[List[QueryCondition]] = None,
+        limit: Optional[int] = 100,
+        offset: Optional[int] = 0,
+        order_by: Optional[List[str]] = None,
+    ) -> List[Trace]: ...
+
+    async def get_span_tree(
+        self,
+        span_id: str,
+        attributes_to_return: Optional[List[str]] = None,
+        max_depth: Optional[int] = None,
+    ) -> SpanWithChildren: ...
+
+
+class SQLiteTraceStore(TraceStore):
+    def __init__(self, conn_string: str):
+        self.conn_string = conn_string
+
+    async def query_traces(
+        self,
+        attribute_filters: Optional[List[QueryCondition]] = None,
+        limit: Optional[int] = 100,
+        offset: Optional[int] = 0,
+        order_by: Optional[List[str]] = None,
+    ) -> List[Trace]:
+
+        def build_where_clause() -> tuple[str, list]:
+            if not attribute_filters:
+                return "", []
+
+            ops_map = {"eq": "=", "ne": "!=", "gt": ">", "lt": "<"}
+
+            conditions = [
+                f"json_extract(s.attributes, '$.{condition.key}') {ops_map[condition.op]} ?"
+                for condition in attribute_filters
+            ]
+            params = [condition.value for condition in attribute_filters]
+            where_clause = " WHERE " + " AND ".join(conditions)
+            return where_clause, params
+
+        def build_order_clause() -> str:
+            if not order_by:
+                return ""
+
+            order_clauses = []
+            for field in order_by:
+                desc = field.startswith("-")
+                clean_field = field[1:] if desc else field
+                order_clauses.append(f"t.{clean_field} {'DESC' if desc else 'ASC'}")
+            return " ORDER BY " + ", ".join(order_clauses)
+
+        # Build the main query
+        base_query = """
+            WITH matching_traces AS (
+                SELECT DISTINCT t.trace_id
+                FROM traces t
+                JOIN spans s ON t.trace_id = s.trace_id
+                {where_clause}
+            ),
+            filtered_traces AS (
+                SELECT t.trace_id, t.root_span_id, t.start_time, t.end_time
+                FROM matching_traces mt
+                JOIN traces t ON mt.trace_id = t.trace_id
+                LEFT JOIN spans s ON t.trace_id = s.trace_id
+                {order_clause}
+            )
+            SELECT DISTINCT trace_id, root_span_id, start_time, end_time
+            FROM filtered_traces
+            LIMIT {limit} OFFSET {offset}
+        """
+
+        where_clause, params = build_where_clause()
+        query = base_query.format(
+            where_clause=where_clause,
+            order_clause=build_order_clause(),
+            limit=limit,
+            offset=offset,
+        )
+
+        # Execute query and return results
+        async with aiosqlite.connect(self.conn_string) as conn:
+            conn.row_factory = aiosqlite.Row
+            async with conn.execute(query, params) as cursor:
+                rows = await cursor.fetchall()
+                return [
+                    Trace(
+                        trace_id=row["trace_id"],
+                        root_span_id=row["root_span_id"],
+                        start_time=datetime.fromisoformat(row["start_time"]),
+                        end_time=datetime.fromisoformat(row["end_time"]),
+                    )
+                    for row in rows
+                ]
+
+    async def get_span_tree(
+        self,
+        span_id: str,
+        attributes_to_return: Optional[List[str]] = None,
+        max_depth: Optional[int] = None,
+    ) -> SpanWithChildren:
+        # Build the attributes selection
+        attributes_select = "s.attributes"
+        if attributes_to_return:
+            json_object = ", ".join(
+                f"'{key}', json_extract(s.attributes, '$.{key}')"
+                for key in attributes_to_return
+            )
+            attributes_select = f"json_object({json_object})"
+
+        # SQLite CTE query with filtered attributes
+        query = f"""
+        WITH RECURSIVE span_tree AS (
+            SELECT s.*, 1 as depth, {attributes_select} as filtered_attributes
+            FROM spans s
+            WHERE s.span_id = ?
+
+            UNION ALL
+
+            SELECT s.*, st.depth + 1, {attributes_select} as filtered_attributes
+            FROM spans s
+            JOIN span_tree st ON s.parent_span_id = st.span_id
+            WHERE (? IS NULL OR st.depth < ?)
+        )
+        SELECT *
+        FROM span_tree
+        ORDER BY depth, start_time
+        """
+
+        async with aiosqlite.connect(self.conn_string) as conn:
+            conn.row_factory = aiosqlite.Row
+            async with conn.execute(query, (span_id, max_depth, max_depth)) as cursor:
+                rows = await cursor.fetchall()
+
+                if not rows:
+                    raise ValueError(f"Span {span_id} not found")
+
+                # Build span tree
+                spans_by_id = {}
+                root_span = None
+
+                for row in rows:
+                    span = SpanWithChildren(
+                        span_id=row["span_id"],
+                        trace_id=row["trace_id"],
+                        parent_span_id=row["parent_span_id"],
+                        name=row["name"],
+                        start_time=datetime.fromisoformat(row["start_time"]),
+                        end_time=datetime.fromisoformat(row["end_time"]),
+                        attributes=json.loads(row["filtered_attributes"]),
+                        status=row["status"].lower(),
+                        children=[],
+                    )
+
+                    spans_by_id[span.span_id] = span
+
+                    if span.span_id == span_id:
+                        root_span = span
+                    elif span.parent_span_id in spans_by_id:
+                        spans_by_id[span.parent_span_id].children.append(span)
+
+                return root_span

llama_stack/providers/utils/telemetry/trace_protocol.py

@@ -0,0 +1,129 @@
+# 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 asyncio
+import inspect
+import json
+from functools import wraps
+from typing import Any, AsyncGenerator, Callable, Type, TypeVar
+
+from pydantic import BaseModel
+
+T = TypeVar("T")
+
+
+def serialize_value(value: Any) -> str:
+    """Helper function to serialize values to string representation."""
+    try:
+        if isinstance(value, BaseModel):
+            return value.model_dump_json()
+        elif isinstance(value, list) and value and isinstance(value[0], BaseModel):
+            return json.dumps([item.model_dump_json() for item in value])
+        elif hasattr(value, "to_dict"):
+            return json.dumps(value.to_dict())
+        elif isinstance(value, (dict, list, int, float, str, bool)):
+            return json.dumps(value)
+        else:
+            return str(value)
+    except Exception:
+        return str(value)
+
+
+def trace_protocol(cls: Type[T]) -> Type[T]:
+    """
+    A class decorator that automatically traces all methods in a protocol/base class
+    and its inheriting classes.
+    """
+
+    def trace_method(method: Callable) -> Callable:
+        from llama_stack.providers.utils.telemetry import tracing
+
+        is_async = asyncio.iscoroutinefunction(method)
+        is_async_gen = inspect.isasyncgenfunction(method)
+
+        def create_span_context(self: Any, *args: Any, **kwargs: Any) -> tuple:
+            class_name = self.__class__.__name__
+            method_name = method.__name__
+
+            span_type = (
+                "async_generator" if is_async_gen else "async" if is_async else "sync"
+            )
+            span_attributes = {
+                "__autotraced__": True,
+                "__class__": class_name,
+                "__method__": method_name,
+                "__type__": span_type,
+                "__args__": serialize_value(args),
+            }
+
+            return class_name, method_name, span_attributes
+
+        @wraps(method)
+        async def async_gen_wrapper(
+            self: Any, *args: Any, **kwargs: Any
+        ) -> AsyncGenerator:
+            class_name, method_name, span_attributes = create_span_context(
+                self, *args, **kwargs
+            )
+
+            with tracing.span(f"{class_name}.{method_name}", span_attributes) as span:
+                try:
+                    count = 0
+                    async for item in method(self, *args, **kwargs):
+                        yield item
+                        count += 1
+                finally:
+                    span.set_attribute("chunk_count", count)
+
+        @wraps(method)
+        async def async_wrapper(self: Any, *args: Any, **kwargs: Any) -> Any:
+            class_name, method_name, span_attributes = create_span_context(
+                self, *args, **kwargs
+            )
+
+            with tracing.span(f"{class_name}.{method_name}", span_attributes) as span:
+                try:
+                    result = await method(self, *args, **kwargs)
+                    span.set_attribute("output", serialize_value(result))
+                    return result
+                except Exception as e:
+                    span.set_attribute("error", str(e))
+                    raise
+
+        @wraps(method)
+        def sync_wrapper(self: Any, *args: Any, **kwargs: Any) -> Any:
+            class_name, method_name, span_attributes = create_span_context(
+                self, *args, **kwargs
+            )
+
+            with tracing.span(f"{class_name}.{method_name}", span_attributes) as span:
+                try:
+                    result = method(self, *args, **kwargs)
+                    span.set_attribute("output", serialize_value(result))
+                    return result
+                except Exception as _e:
+                    raise
+
+        if is_async_gen:
+            return async_gen_wrapper
+        elif is_async:
+            return async_wrapper
+        else:
+            return sync_wrapper
+
+    original_init_subclass = getattr(cls, "__init_subclass__", None)
+
+    def __init_subclass__(cls_child, **kwargs):  # noqa: N807
+        if original_init_subclass:
+            original_init_subclass(**kwargs)
+
+        for name, method in vars(cls_child).items():
+            if inspect.isfunction(method) and not name.startswith("_"):
+                setattr(cls_child, name, trace_method(method))  # noqa: B010
+
+    cls.__init_subclass__ = classmethod(__init_subclass__)
+
+    return cls

llama_stack/providers/utils/telemetry/tracing.py

@@ -69,7 +69,7 @@ class TraceContext:
         self.logger = logger
         self.trace_id = trace_id
 
-    def push_span(self, name: str, attributes: Dict[str, Any] = None):
+    def push_span(self, name: str, attributes: Dict[str, Any] = None) -> Span:
         current_span = self.get_current_span()
         span = Span(
             span_id=generate_short_uuid(),
@@ -94,6 +94,7 @@ class TraceContext:
         )
 
         self.spans.append(span)
+        return span
 
     def pop_span(self, status: SpanStatus = SpanStatus.OK):
         span = self.spans.pop()
@@ -203,12 +204,13 @@ class SpanContextManager:
     def __init__(self, name: str, attributes: Dict[str, Any] = None):
         self.name = name
         self.attributes = attributes
+        self.span = None
 
     def __enter__(self):
         global CURRENT_TRACE_CONTEXT
         context = CURRENT_TRACE_CONTEXT
         if context:
-            context.push_span(self.name, self.attributes)
+            self.span = context.push_span(self.name, self.attributes)
         return self
 
     def __exit__(self, exc_type, exc_value, traceback):
@@ -217,11 +219,24 @@ class SpanContextManager:
         if context:
             context.pop_span()
 
+    def set_attribute(self, key: str, value: Any):
+        if self.span:
+            if self.span.attributes is None:
+                self.span.attributes = {}
+            self.span.attributes[key] = value
+
     async def __aenter__(self):
-        return self.__enter__()
+        global CURRENT_TRACE_CONTEXT
+        context = CURRENT_TRACE_CONTEXT
+        if context:
+            self.span = context.push_span(self.name, self.attributes)
+        return self
 
     async def __aexit__(self, exc_type, exc_value, traceback):
-        self.__exit__(exc_type, exc_value, traceback)
+        global CURRENT_TRACE_CONTEXT
+        context = CURRENT_TRACE_CONTEXT
+        if context:
+            context.pop_span()
 
     def __call__(self, func: Callable):
         @wraps(func)
@@ -246,3 +261,11 @@ class SpanContextManager:
 
 def span(name: str, attributes: Dict[str, Any] = None):
     return SpanContextManager(name, attributes)
+
+
+def get_current_span() -> Optional[Span]:
+    global CURRENT_TRACE_CONTEXT
+    context = CURRENT_TRACE_CONTEXT
+    if context:
+        return context.get_current_span()
+    return None

llama_stack/templates/bedrock/bedrock.py

@@ -6,6 +6,9 @@
 
 from pathlib import Path
 
+from llama_stack.distribution.datatypes import Provider
+
+from llama_stack.providers.inline.memory.faiss.config import FaissImplConfig
 from llama_stack.templates.template import DistributionTemplate, RunConfigSettings
 
 
@@ -16,10 +19,19 @@ def get_distribution_template() -> DistributionTemplate:
         "safety": ["remote::bedrock"],
         "agents": ["inline::meta-reference"],
         "telemetry": ["inline::meta-reference"],
+        "eval": ["inline::meta-reference"],
+        "datasetio": ["remote::huggingface", "inline::localfs"],
+        "scoring": ["inline::basic", "inline::llm-as-judge", "inline::braintrust"],
     }
+    name = "bedrock"
+    memory_provider = Provider(
+        provider_id="faiss",
+        provider_type="inline::faiss",
+        config=FaissImplConfig.sample_run_config(f"distributions/{name}"),
+    )
 
     return DistributionTemplate(
-        name="bedrock",
+        name=name,
         distro_type="self_hosted",
         description="Use AWS Bedrock for running LLM inference and safety",
         docker_image=None,
@@ -27,7 +39,11 @@ def get_distribution_template() -> DistributionTemplate:
         providers=providers,
         default_models=[],
         run_configs={
-            "run.yaml": RunConfigSettings(),
+            "run.yaml": RunConfigSettings(
+                provider_overrides={
+                    "memory": [memory_provider],
+                },
+            ),
         },
         run_config_env_vars={
             "LLAMASTACK_PORT": (

llama_stack/templates/bedrock/build.yaml

@@ -16,4 +16,13 @@ distribution_spec:
     - inline::meta-reference
     telemetry:
     - inline::meta-reference
+    eval:
+    - inline::meta-reference
+    datasetio:
+    - remote::huggingface
+    - inline::localfs
+    scoring:
+    - inline::basic
+    - inline::llm-as-judge
+    - inline::braintrust
 image_type: conda

llama_stack/templates/bedrock/run.yaml

@@ -4,9 +4,12 @@ docker_image: null
 conda_env: bedrock
 apis:
 - agents
+- datasetio
+- eval
 - inference
 - memory
 - safety
+- scoring
 - telemetry
 providers:
   inference:
@@ -37,6 +40,27 @@ providers:
   - provider_id: meta-reference
     provider_type: inline::meta-reference
     config: {}
+  eval:
+  - provider_id: meta-reference
+    provider_type: inline::meta-reference
+    config: {}
+  datasetio:
+  - provider_id: huggingface
+    provider_type: remote::huggingface
+    config: {}
+  - provider_id: localfs
+    provider_type: inline::localfs
+    config: {}
+  scoring:
+  - provider_id: basic
+    provider_type: inline::basic
+    config: {}
+  - provider_id: llm-as-judge
+    provider_type: inline::llm-as-judge
+    config: {}
+  - provider_id: braintrust
+    provider_type: inline::braintrust
+    config: {}
 metadata_store:
   namespace: null
   type: sqlite

llama_stack/templates/cerebras/__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 .cerebras import get_distribution_template  # noqa: F401

llama_stack/templates/cerebras/build.yaml

@@ -0,0 +1,17 @@
+version: '2'
+name: cerebras
+distribution_spec:
+  description: Use Cerebras for running LLM inference
+  docker_image: null
+  providers:
+    inference:
+    - remote::cerebras
+    safety:
+    - inline::llama-guard
+    memory:
+    - inline::meta-reference
+    agents:
+    - inline::meta-reference
+    telemetry:
+    - inline::meta-reference
+image_type: conda