Merge branch 'meta-llama:main' into qdrant

.flake8

@@ -21,11 +21,11 @@ ignore =
 optional-ascii-coding = True
 exclude =
     ./.git,
-    ./docs
-    ./build
+    ./docs/*,
+    ./build,
     ./scripts,
     ./venv,
-    *.pyi
-    .pre-commit-config.yaml
-    *.md
+    *.pyi,
+    .pre-commit-config.yaml,
+    *.md,
     .flake8

.github/workflows/pre-commit.yml

@@ -0,0 +1,45 @@
+name: Pre-commit
+
+on: [pull_request]
+
+jobs:
+  pre-commit:
+    runs-on: ubuntu-latest
+
+    steps:
+      - name: Checkout code
+        uses: actions/checkout@d632683dd7b4114ad314bca15554477dd762a938 # v4.2.0
+
+      - name: Set up Python
+        uses: actions/setup-python@f677139bbe7f9c59b41e40162b753c062f5d49a3 # v5.2.0
+        with:
+          python-version: '3.11'
+          cache: pip
+          cache-dependency-path: |
+            **/requirements*.txt
+            .pre-commit-config.yaml
+
+      - name: Install pre-commit
+        run: |
+          python -m pip install --upgrade pip
+          pip install pre-commit
+
+      - name: Fetch base branch
+        run: git fetch origin ${{ github.event.pull_request.base.ref }}:refs/remotes/origin/${{ github.event.pull_request.base.ref }}
+
+      - name: Fetch head commit from PR
+        run: git fetch origin pull/${{ github.event.pull_request.number }}/head:pr-${{ github.event.pull_request.number }}
+
+      - name: Get changed files
+        id: changed-files
+        run: |
+          git diff --name-only origin/${{ github.event.pull_request.base.ref }} pr-${{ github.event.pull_request.number }} > changed_files.txt
+          cat changed_files.txt
+
+      - name: Run pre-commit
+        run: |
+          if [ -s changed_files.txt ]; then
+            pre-commit run --files $(cat changed_files.txt | tr '\n' ' ')
+          else
+            echo "No changed files to run pre-commit on."
+          fi

README.md

@@ -1,3 +1,5 @@
+<img src="https://github.com/user-attachments/assets/2fedfe0f-6df7-4441-98b2-87a1fd95ee1c" width="300" title="Llama Stack Logo" alt="Llama Stack Logo"/>
+
 # Llama Stack
 
 [![PyPI version](https://img.shields.io/pypi/v/llama_stack.svg)](https://pypi.org/project/llama_stack/)
@@ -99,7 +101,7 @@ The `llama` CLI makes it easy to work with the Llama Stack set of tools. Please
 |  **Language** |  **Client SDK** | **Package** |
 | :----: | :----: | :----: |
 | Python |  [llama-stack-client-python](https://github.com/meta-llama/llama-stack-client-python) | [![PyPI version](https://img.shields.io/pypi/v/llama_stack_client.svg)](https://pypi.org/project/llama_stack_client/)
-| Swift  | [llama-stack-client-swift](https://github.com/meta-llama/llama-stack-client-swift) |
+| Swift  | [llama-stack-client-swift](https://github.com/meta-llama/llama-stack-client-swift) | [![Swift Package Index](https://img.shields.io/endpoint?url=https%3A%2F%2Fswiftpackageindex.com%2Fapi%2Fpackages%2Fmeta-llama%2Fllama-stack-client-swift%2Fbadge%3Ftype%3Dswift-versions)](https://swiftpackageindex.com/meta-llama/llama-stack-client-swift)
 | Node   | [llama-stack-client-node](https://github.com/meta-llama/llama-stack-client-node) | [![NPM version](https://img.shields.io/npm/v/llama-stack-client.svg)](https://npmjs.org/package/llama-stack-client)
 | Kotlin | [llama-stack-client-kotlin](https://github.com/meta-llama/llama-stack-client-kotlin) |
 

docs/cli_reference.md

@@ -215,9 +215,8 @@ You can even run `llama model prompt-format` see all of the templates and their
 ```
 llama model prompt-format -m Llama3.2-3B-Instruct
 ```
-<p align="center">
-<img width="719" alt="image" src="https://github.com/user-attachments/assets/c5332026-8c0b-4edc-b438-ec60cd7ca554">
-</p>
+![alt text](resources/prompt-format.png)
+
 
 
 You will be shown a Markdown formatted description of the model interface and how prompts / messages are formatted for various scenarios.

docs/openapi_generator/generate.py

@@ -71,7 +71,7 @@ class LlamaStack(
 
 # TODO: this should be fixed in the generator itself so it reads appropriate annotations
 STREAMING_ENDPOINTS = [
-    "/agentic_system/turn/create",
+    "/agents/turn/create",
     "/inference/chat_completion",
 ]
 

llama_stack/apis/agents/agents.py

@@ -6,7 +6,16 @@
 
 from datetime import datetime
 from enum import Enum
-from typing import Any, Dict, List, Literal, Optional, Protocol, Union
+from typing import (
+    Any,
+    Dict,
+    List,
+    Literal,
+    Optional,
+    Protocol,
+    runtime_checkable,
+    Union,
+)
 
 from llama_models.schema_utils import json_schema_type, webmethod
 
@@ -261,7 +270,7 @@ class Session(BaseModel):
     turns: List[Turn]
     started_at: datetime
 
-    memory_bank: Optional[MemoryBank] = None
+    memory_bank: Optional[MemoryBankDef] = None
 
 
 class AgentConfigCommon(BaseModel):
@@ -404,6 +413,7 @@ class AgentStepResponse(BaseModel):
     step: Step
 
 
+@runtime_checkable
 class Agents(Protocol):
     @webmethod(route="/agents/create")
     async def create_agent(
@@ -411,8 +421,10 @@ class Agents(Protocol):
         agent_config: AgentConfig,
     ) -> AgentCreateResponse: ...
 
+    # This method is not `async def` because it can result in either an
+    # `AsyncGenerator` or a `AgentTurnCreateResponse` depending on the value of `stream`.
     @webmethod(route="/agents/turn/create")
-    async def create_agent_turn(
+    def create_agent_turn(
         self,
         agent_id: str,
         session_id: str,

llama_stack/apis/agents/client.py

@@ -7,7 +7,7 @@
 import asyncio
 import json
 import os
-from typing import AsyncGenerator
+from typing import AsyncGenerator, Optional
 
 import fire
 import httpx
@@ -67,9 +67,17 @@ class AgentsClient(Agents):
             response.raise_for_status()
             return AgentSessionCreateResponse(**response.json())
 
-    async def create_agent_turn(
+    def create_agent_turn(
         self,
         request: AgentTurnCreateRequest,
+    ) -> AsyncGenerator:
+        if request.stream:
+            return self._stream_agent_turn(request)
+        else:
+            return self._nonstream_agent_turn(request)
+
+    async def _stream_agent_turn(
+        self, request: AgentTurnCreateRequest
     ) -> AsyncGenerator:
         async with httpx.AsyncClient() as client:
             async with client.stream(
@@ -93,6 +101,9 @@ class AgentsClient(Agents):
                             print(data)
                             print(f"Error with parsing or validation: {e}")
 
+    async def _nonstream_agent_turn(self, request: AgentTurnCreateRequest):
+        raise NotImplementedError("Non-streaming not implemented yet")
+
 
 async def _run_agent(
     api, model, tool_definitions, tool_prompt_format, user_prompts, attachments=None
@@ -132,8 +143,7 @@ async def _run_agent(
                 log.print()
 
 
-async def run_llama_3_1(host: str, port: int):
-    model = "Llama3.1-8B-Instruct"
+async def run_llama_3_1(host: str, port: int, model: str = "Llama3.1-8B-Instruct"):
     api = AgentsClient(f"http://{host}:{port}")
 
     tool_definitions = [
@@ -173,8 +183,7 @@ async def run_llama_3_1(host: str, port: int):
     await _run_agent(api, model, tool_definitions, ToolPromptFormat.json, user_prompts)
 
 
-async def run_llama_3_2_rag(host: str, port: int):
-    model = "Llama3.2-3B-Instruct"
+async def run_llama_3_2_rag(host: str, port: int, model: str = "Llama3.2-3B-Instruct"):
     api = AgentsClient(f"http://{host}:{port}")
 
     urls = [
@@ -215,8 +224,7 @@ async def run_llama_3_2_rag(host: str, port: int):
     )
 
 
-async def run_llama_3_2(host: str, port: int):
-    model = "Llama3.2-3B-Instruct"
+async def run_llama_3_2(host: str, port: int, model: str = "Llama3.2-3B-Instruct"):
     api = AgentsClient(f"http://{host}:{port}")
 
     # zero shot tools for llama3.2 text models
@@ -262,7 +270,7 @@ async def run_llama_3_2(host: str, port: int):
     )
 
 
-def main(host: str, port: int, run_type: str):
+def main(host: str, port: int, run_type: str, model: Optional[str] = None):
     assert run_type in [
         "tools_llama_3_1",
         "tools_llama_3_2",
@@ -274,7 +282,10 @@ def main(host: str, port: int, run_type: str):
         "tools_llama_3_2": run_llama_3_2,
         "rag_llama_3_2": run_llama_3_2_rag,
     }
-    asyncio.run(fn[run_type](host, port))
+    args = [host, port]
+    if model is not None:
+        args.append(model)
+    asyncio.run(fn[run_type](*args))
 
 
 if __name__ == "__main__":

llama_stack/apis/batch_inference/batch_inference.py

@@ -4,7 +4,7 @@
 # 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, Protocol
+from typing import List, Optional, Protocol, runtime_checkable
 
 from llama_models.schema_utils import json_schema_type, webmethod
 
@@ -47,6 +47,7 @@ class BatchChatCompletionResponse(BaseModel):
     completion_message_batch: List[CompletionMessage]
 
 
+@runtime_checkable
 class BatchInference(Protocol):
     @webmethod(route="/batch_inference/completion")
     async def batch_completion(

llama_stack/apis/inference/client.py

@@ -42,10 +42,10 @@ class InferenceClient(Inference):
     async def shutdown(self) -> None:
         pass
 
-    async def completion(self, request: CompletionRequest) -> AsyncGenerator:
+    def completion(self, request: CompletionRequest) -> AsyncGenerator:
         raise NotImplementedError()
 
-    async def chat_completion(
+    def chat_completion(
         self,
         model: str,
         messages: List[Message],
@@ -66,6 +66,29 @@ class InferenceClient(Inference):
             stream=stream,
             logprobs=logprobs,
         )
+        if stream:
+            return self._stream_chat_completion(request)
+        else:
+            return self._nonstream_chat_completion(request)
+
+    async def _nonstream_chat_completion(
+        self, request: ChatCompletionRequest
+    ) -> ChatCompletionResponse:
+        async with httpx.AsyncClient() as client:
+            response = await client.post(
+                f"{self.base_url}/inference/chat_completion",
+                json=encodable_dict(request),
+                headers={"Content-Type": "application/json"},
+                timeout=20,
+            )
+
+            response.raise_for_status()
+            j = response.json()
+            return ChatCompletionResponse(**j)
+
+    async def _stream_chat_completion(
+        self, request: ChatCompletionRequest
+    ) -> AsyncGenerator:
         async with httpx.AsyncClient() as client:
             async with client.stream(
                 "POST",
@@ -77,7 +100,8 @@ class InferenceClient(Inference):
                 if response.status_code != 200:
                     content = await response.aread()
                     cprint(
-                        f"Error: HTTP {response.status_code} {content.decode()}", "red"
+                        f"Error: HTTP {response.status_code} {content.decode()}",
+                        "red",
                     )
                     return
 
@@ -85,16 +109,11 @@ class InferenceClient(Inference):
                     if line.startswith("data:"):
                         data = line[len("data: ") :]
                         try:
-                            if request.stream:
                             if "error" in data:
                                 cprint(data, "red")
                                 continue
 
-                                yield ChatCompletionResponseStreamChunk(
-                                    **json.loads(data)
-                                )
-                            else:
-                                yield ChatCompletionResponse(**json.loads(data))
+                            yield ChatCompletionResponseStreamChunk(**json.loads(data))
                         except Exception as e:
                             print(data)
                             print(f"Error with parsing or validation: {e}")

llama_stack/apis/inference/inference.py

@@ -6,7 +6,7 @@
 
 from enum import Enum
 
-from typing import List, Literal, Optional, Protocol, Union
+from typing import List, Literal, Optional, Protocol, runtime_checkable, Union
 
 from llama_models.schema_utils import json_schema_type, webmethod
 
@@ -14,6 +14,7 @@ from pydantic import BaseModel, Field
 from typing_extensions import Annotated
 
 from llama_models.llama3.api.datatypes import *  # noqa: F403
+from llama_stack.apis.models import *  # noqa: F403
 
 
 class LogProbConfig(BaseModel):
@@ -172,9 +173,18 @@ class EmbeddingsResponse(BaseModel):
     embeddings: List[List[float]]
 
 
+class ModelStore(Protocol):
+    def get_model(self, identifier: str) -> ModelDef: ...
+
+
+@runtime_checkable
 class Inference(Protocol):
+    model_store: ModelStore
+
+    # This method is not `async def` because it can result in either an
+    # `AsyncGenerator` or a `CompletionResponse` depending on the value of `stream`.
     @webmethod(route="/inference/completion")
-    async def completion(
+    def completion(
         self,
         model: str,
         content: InterleavedTextMedia,
@@ -183,8 +193,10 @@ class Inference(Protocol):
         logprobs: Optional[LogProbConfig] = None,
     ) -> Union[CompletionResponse, CompletionResponseStreamChunk]: ...
 
+    # This method is not `async def` because it can result in either an
+    # `AsyncGenerator` or a `ChatCompletionResponse` depending on the value of `stream`.
     @webmethod(route="/inference/chat_completion")
-    async def chat_completion(
+    def chat_completion(
         self,
         model: str,
         messages: List[Message],

llama_stack/apis/inspect/inspect.py

@@ -4,7 +4,7 @@
 # This source code is licensed under the terms described in the LICENSE file in
 # the root directory of this source tree.
 
-from typing import Dict, List, Protocol
+from typing import Dict, List, Protocol, runtime_checkable
 
 from llama_models.schema_utils import json_schema_type, webmethod
 from pydantic import BaseModel
@@ -12,15 +12,15 @@ from pydantic import BaseModel
 
 @json_schema_type
 class ProviderInfo(BaseModel):
+    provider_id: str
     provider_type: str
-    description: str
 
 
 @json_schema_type
 class RouteInfo(BaseModel):
     route: str
     method: str
-    providers: List[str]
+    provider_types: List[str]
 
 
 @json_schema_type
@@ -29,6 +29,7 @@ class HealthInfo(BaseModel):
     # TODO: add a provider level status
 
 
+@runtime_checkable
 class Inspect(Protocol):
     @webmethod(route="/providers/list", method="GET")
     async def list_providers(self) -> Dict[str, ProviderInfo]: ...

llama_stack/apis/memory/client.py

@@ -5,7 +5,6 @@
 # the root directory of this source tree.
 
 import asyncio
-import json
 import os
 from pathlib import Path
 
@@ -13,11 +12,11 @@ from typing import Any, Dict, List, Optional
 
 import fire
 import httpx
-from termcolor import cprint
 
 from llama_stack.distribution.datatypes import RemoteProviderConfig
 
 from llama_stack.apis.memory import *  # noqa: F403
+from llama_stack.apis.memory_banks.client import MemoryBanksClient
 from llama_stack.providers.utils.memory.file_utils import data_url_from_file
 
 
@@ -35,45 +34,6 @@ class MemoryClient(Memory):
     async def shutdown(self) -> None:
         pass
 
-    async def get_memory_bank(self, bank_id: str) -> Optional[MemoryBank]:
-        async with httpx.AsyncClient() as client:
-            r = await client.get(
-                f"{self.base_url}/memory/get",
-                params={
-                    "bank_id": bank_id,
-                },
-                headers={"Content-Type": "application/json"},
-                timeout=20,
-            )
-            r.raise_for_status()
-            d = r.json()
-            if not d:
-                return None
-            return MemoryBank(**d)
-
-    async def create_memory_bank(
-        self,
-        name: str,
-        config: MemoryBankConfig,
-        url: Optional[URL] = None,
-    ) -> MemoryBank:
-        async with httpx.AsyncClient() as client:
-            r = await client.post(
-                f"{self.base_url}/memory/create",
-                json={
-                    "name": name,
-                    "config": config.dict(),
-                    "url": url,
-                },
-                headers={"Content-Type": "application/json"},
-                timeout=20,
-            )
-            r.raise_for_status()
-            d = r.json()
-            if not d:
-                return None
-            return MemoryBank(**d)
-
     async def insert_documents(
         self,
         bank_id: str,
@@ -113,23 +73,20 @@ class MemoryClient(Memory):
 
 
 async def run_main(host: str, port: int, stream: bool):
-    client = MemoryClient(f"http://{host}:{port}")
+    banks_client = MemoryBanksClient(f"http://{host}:{port}")
 
-    # create a memory bank
-    bank = await client.create_memory_bank(
-        name="test_bank",
-        config=VectorMemoryBankConfig(
-            bank_id="test_bank",
+    bank = VectorMemoryBankDef(
+        identifier="test_bank",
+        provider_id="",
         embedding_model="all-MiniLM-L6-v2",
         chunk_size_in_tokens=512,
         overlap_size_in_tokens=64,
-        ),
     )
-    cprint(json.dumps(bank.dict(), indent=4), "green")
+    await banks_client.register_memory_bank(bank)
 
-    retrieved_bank = await client.get_memory_bank(bank.bank_id)
+    retrieved_bank = await banks_client.get_memory_bank(bank.identifier)
     assert retrieved_bank is not None
-    assert retrieved_bank.config.embedding_model == "all-MiniLM-L6-v2"
+    assert retrieved_bank.embedding_model == "all-MiniLM-L6-v2"
 
     urls = [
         "memory_optimizations.rst",
@@ -160,15 +117,17 @@ async def run_main(host: str, port: int, stream: bool):
         for i, path in enumerate(files)
     ]
 
+    client = MemoryClient(f"http://{host}:{port}")
+
     # insert some documents
     await client.insert_documents(
-        bank_id=bank.bank_id,
+        bank_id=bank.identifier,
         documents=documents,
     )
 
     # query the documents
     response = await client.query_documents(
-        bank_id=bank.bank_id,
+        bank_id=bank.identifier,
         query=[
             "How do I use Lora?",
         ],
@@ -178,7 +137,7 @@ async def run_main(host: str, port: int, stream: bool):
         print(f"Chunk:\n========\n{chunk}\n========\n")
 
     response = await client.query_documents(
-        bank_id=bank.bank_id,
+        bank_id=bank.identifier,
         query=[
             "Tell me more about llama3 and torchtune",
         ],

llama_stack/apis/memory/memory.py

@@ -8,14 +8,14 @@
 #
 # 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, Protocol
+from typing import List, Optional, Protocol, runtime_checkable
 
 from llama_models.schema_utils import json_schema_type, webmethod
 
 from pydantic import BaseModel, Field
-from typing_extensions import Annotated
 
 from llama_models.llama3.api.datatypes import *  # noqa: F403
+from llama_stack.apis.memory_banks import *  # noqa: F403
 
 
 @json_schema_type
@@ -26,44 +26,6 @@ class MemoryBankDocument(BaseModel):
     metadata: Dict[str, Any] = Field(default_factory=dict)
 
 
-@json_schema_type
-class MemoryBankType(Enum):
-    vector = "vector"
-    keyvalue = "keyvalue"
-    keyword = "keyword"
-    graph = "graph"
-
-
-class VectorMemoryBankConfig(BaseModel):
-    type: Literal[MemoryBankType.vector.value] = MemoryBankType.vector.value
-    embedding_model: str
-    chunk_size_in_tokens: int
-    overlap_size_in_tokens: Optional[int] = None
-
-
-class KeyValueMemoryBankConfig(BaseModel):
-    type: Literal[MemoryBankType.keyvalue.value] = MemoryBankType.keyvalue.value
-
-
-class KeywordMemoryBankConfig(BaseModel):
-    type: Literal[MemoryBankType.keyword.value] = MemoryBankType.keyword.value
-
-
-class GraphMemoryBankConfig(BaseModel):
-    type: Literal[MemoryBankType.graph.value] = MemoryBankType.graph.value
-
-
-MemoryBankConfig = Annotated[
-    Union[
-        VectorMemoryBankConfig,
-        KeyValueMemoryBankConfig,
-        KeywordMemoryBankConfig,
-        GraphMemoryBankConfig,
-    ],
-    Field(discriminator="type"),
-]
-
-
 class Chunk(BaseModel):
     content: InterleavedTextMedia
     token_count: int
@@ -76,45 +38,13 @@ class QueryDocumentsResponse(BaseModel):
     scores: List[float]
 
 
-@json_schema_type
-class QueryAPI(Protocol):
-    @webmethod(route="/query_documents")
-    def query_documents(
-        self,
-        query: InterleavedTextMedia,
-        params: Optional[Dict[str, Any]] = None,
-    ) -> QueryDocumentsResponse: ...
-
-
-@json_schema_type
-class MemoryBank(BaseModel):
-    bank_id: str
-    name: str
-    config: MemoryBankConfig
-    # if there's a pre-existing (reachable-from-distribution) store which supports QueryAPI
-    url: Optional[URL] = None
+class MemoryBankStore(Protocol):
+    def get_memory_bank(self, bank_id: str) -> Optional[MemoryBankDef]: ...
 
 
+@runtime_checkable
 class Memory(Protocol):
-    @webmethod(route="/memory/create")
-    async def create_memory_bank(
-        self,
-        name: str,
-        config: MemoryBankConfig,
-        url: Optional[URL] = None,
-    ) -> MemoryBank: ...
-
-    @webmethod(route="/memory/list", method="GET")
-    async def list_memory_banks(self) -> List[MemoryBank]: ...
-
-    @webmethod(route="/memory/get", method="GET")
-    async def get_memory_bank(self, bank_id: str) -> Optional[MemoryBank]: ...
-
-    @webmethod(route="/memory/drop", method="DELETE")
-    async def drop_memory_bank(
-        self,
-        bank_id: str,
-    ) -> str: ...
+    memory_bank_store: MemoryBankStore
 
     # this will just block now until documents are inserted, but it should
     # probably return a Job instance which can be polled for completion
@@ -126,13 +56,6 @@ class Memory(Protocol):
         ttl_seconds: Optional[int] = None,
     ) -> None: ...
 
-    @webmethod(route="/memory/update")
-    async def update_documents(
-        self,
-        bank_id: str,
-        documents: List[MemoryBankDocument],
-    ) -> None: ...
-
     @webmethod(route="/memory/query")
     async def query_documents(
         self,
@@ -140,17 +63,3 @@ class Memory(Protocol):
         query: InterleavedTextMedia,
         params: Optional[Dict[str, Any]] = None,
     ) -> QueryDocumentsResponse: ...
-
-    @webmethod(route="/memory/documents/get", method="GET")
-    async def get_documents(
-        self,
-        bank_id: str,
-        document_ids: List[str],
-    ) -> List[MemoryBankDocument]: ...
-
-    @webmethod(route="/memory/documents/delete", method="DELETE")
-    async def delete_documents(
-        self,
-        bank_id: str,
-        document_ids: List[str],
-    ) -> None: ...

llama_stack/apis/memory_banks/client.py

@@ -5,8 +5,9 @@
 # the root directory of this source tree.
 
 import asyncio
+import json
 
-from typing import List, Optional
+from typing import Any, Dict, List, Optional
 
 import fire
 import httpx
@@ -15,6 +16,27 @@ from termcolor import cprint
 from .memory_banks import *  # noqa: F403
 
 
+def deserialize_memory_bank_def(
+    j: Optional[Dict[str, Any]]
+) -> MemoryBankDefWithProvider:
+    if j is None:
+        return None
+
+    if "type" not in j:
+        raise ValueError("Memory bank type not specified")
+    type = j["type"]
+    if type == MemoryBankType.vector.value:
+        return VectorMemoryBankDef(**j)
+    elif type == MemoryBankType.keyvalue.value:
+        return KeyValueMemoryBankDef(**j)
+    elif type == MemoryBankType.keyword.value:
+        return KeywordMemoryBankDef(**j)
+    elif type == MemoryBankType.graph.value:
+        return GraphMemoryBankDef(**j)
+    else:
+        raise ValueError(f"Unknown memory bank type: {type}")
+
+
 class MemoryBanksClient(MemoryBanks):
     def __init__(self, base_url: str):
         self.base_url = base_url
@@ -25,37 +47,49 @@ class MemoryBanksClient(MemoryBanks):
     async def shutdown(self) -> None:
         pass
 
-    async def list_available_memory_banks(self) -> List[MemoryBankSpec]:
+    async def list_memory_banks(self) -> List[MemoryBankDefWithProvider]:
         async with httpx.AsyncClient() as client:
             response = await client.get(
                 f"{self.base_url}/memory_banks/list",
                 headers={"Content-Type": "application/json"},
             )
             response.raise_for_status()
-            return [MemoryBankSpec(**x) for x in response.json()]
+            return [deserialize_memory_bank_def(x) for x in response.json()]
 
-    async def get_serving_memory_bank(
-        self, bank_type: MemoryBankType
-    ) -> Optional[MemoryBankSpec]:
+    async def register_memory_bank(
+        self, memory_bank: MemoryBankDefWithProvider
+    ) -> None:
+        async with httpx.AsyncClient() as client:
+            response = await client.post(
+                f"{self.base_url}/memory_banks/register",
+                json={
+                    "memory_bank": json.loads(memory_bank.json()),
+                },
+                headers={"Content-Type": "application/json"},
+            )
+            response.raise_for_status()
+
+    async def get_memory_bank(
+        self,
+        identifier: str,
+    ) -> Optional[MemoryBankDefWithProvider]:
         async with httpx.AsyncClient() as client:
             response = await client.get(
                 f"{self.base_url}/memory_banks/get",
                 params={
-                    "bank_type": bank_type.value,
+                    "identifier": identifier,
                 },
                 headers={"Content-Type": "application/json"},
             )
             response.raise_for_status()
             j = response.json()
-            if j is None:
-                return None
-            return MemoryBankSpec(**j)
+            return deserialize_memory_bank_def(j)
 
 
 async def run_main(host: str, port: int, stream: bool):
     client = MemoryBanksClient(f"http://{host}:{port}")
 
-    response = await client.list_available_memory_banks()
+    response = await client.list_memory_banks()
     cprint(f"list_memory_banks response={response}", "green")
 
 

llama_stack/apis/memory_banks/memory_banks.py

@@ -4,29 +4,75 @@
 # 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, Protocol
+from enum import Enum
+from typing import List, Literal, Optional, Protocol, runtime_checkable, Union
 
 from llama_models.schema_utils import json_schema_type, webmethod
 from pydantic import BaseModel, Field
-
-from llama_stack.apis.memory import MemoryBankType
-
-from llama_stack.distribution.datatypes import GenericProviderConfig
+from typing_extensions import Annotated
 
 
 @json_schema_type
-class MemoryBankSpec(BaseModel):
-    bank_type: MemoryBankType
-    provider_config: GenericProviderConfig = Field(
-        description="Provider config for the model, including provider_type, and corresponding config. ",
-    )
+class MemoryBankType(Enum):
+    vector = "vector"
+    keyvalue = "keyvalue"
+    keyword = "keyword"
+    graph = "graph"
 
 
+class CommonDef(BaseModel):
+    identifier: str
+    # Hack: move this out later
+    provider_id: str = ""
+
+
+@json_schema_type
+class VectorMemoryBankDef(CommonDef):
+    type: Literal[MemoryBankType.vector.value] = MemoryBankType.vector.value
+    embedding_model: str
+    chunk_size_in_tokens: int
+    overlap_size_in_tokens: Optional[int] = None
+
+
+@json_schema_type
+class KeyValueMemoryBankDef(CommonDef):
+    type: Literal[MemoryBankType.keyvalue.value] = MemoryBankType.keyvalue.value
+
+
+@json_schema_type
+class KeywordMemoryBankDef(CommonDef):
+    type: Literal[MemoryBankType.keyword.value] = MemoryBankType.keyword.value
+
+
+@json_schema_type
+class GraphMemoryBankDef(CommonDef):
+    type: Literal[MemoryBankType.graph.value] = MemoryBankType.graph.value
+
+
+MemoryBankDef = Annotated[
+    Union[
+        VectorMemoryBankDef,
+        KeyValueMemoryBankDef,
+        KeywordMemoryBankDef,
+        GraphMemoryBankDef,
+    ],
+    Field(discriminator="type"),
+]
+
+MemoryBankDefWithProvider = MemoryBankDef
+
+
+@runtime_checkable
 class MemoryBanks(Protocol):
     @webmethod(route="/memory_banks/list", method="GET")
-    async def list_available_memory_banks(self) -> List[MemoryBankSpec]: ...
+    async def list_memory_banks(self) -> List[MemoryBankDefWithProvider]: ...
 
     @webmethod(route="/memory_banks/get", method="GET")
-    async def get_serving_memory_bank(
-        self, bank_type: MemoryBankType
-    ) -> Optional[MemoryBankSpec]: ...
+    async def get_memory_bank(
+        self, identifier: str
+    ) -> Optional[MemoryBankDefWithProvider]: ...
+
+    @webmethod(route="/memory_banks/register", method="POST")
+    async def register_memory_bank(
+        self, memory_bank: MemoryBankDefWithProvider
+    ) -> None: ...

llama_stack/apis/models/client.py

@@ -5,6 +5,7 @@
 # the root directory of this source tree.
 
 import asyncio
+import json
 
 from typing import List, Optional
 
@@ -25,21 +26,32 @@ class ModelsClient(Models):
     async def shutdown(self) -> None:
         pass
 
-    async def list_models(self) -> List[ModelServingSpec]:
+    async def list_models(self) -> List[ModelDefWithProvider]:
         async with httpx.AsyncClient() as client:
             response = await client.get(
                 f"{self.base_url}/models/list",
                 headers={"Content-Type": "application/json"},
             )
             response.raise_for_status()
-            return [ModelServingSpec(**x) for x in response.json()]
+            return [ModelDefWithProvider(**x) for x in response.json()]
 
-    async def get_model(self, core_model_id: str) -> Optional[ModelServingSpec]:
+    async def register_model(self, model: ModelDefWithProvider) -> None:
+        async with httpx.AsyncClient() as client:
+            response = await client.post(
+                f"{self.base_url}/models/register",
+                json={
+                    "model": json.loads(model.json()),
+                },
+                headers={"Content-Type": "application/json"},
+            )
+            response.raise_for_status()
+
+    async def get_model(self, identifier: str) -> Optional[ModelDefWithProvider]:
         async with httpx.AsyncClient() as client:
             response = await client.get(
                 f"{self.base_url}/models/get",
                 params={
-                    "core_model_id": core_model_id,
+                    "identifier": identifier,
                 },
                 headers={"Content-Type": "application/json"},
             )
@@ -47,7 +59,7 @@ class ModelsClient(Models):
             j = response.json()
             if j is None:
                 return None
-            return ModelServingSpec(**j)
+            return ModelDefWithProvider(**j)
 
 
 async def run_main(host: str, port: int, stream: bool):

llama_stack/apis/models/models.py

@@ -4,29 +4,39 @@
 # 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, Protocol
-
-from llama_models.llama3.api.datatypes import Model
+from typing import Any, Dict, List, Optional, Protocol, runtime_checkable
 
 from llama_models.schema_utils import json_schema_type, webmethod
 from pydantic import BaseModel, Field
 
-from llama_stack.distribution.datatypes import GenericProviderConfig
+
+class ModelDef(BaseModel):
+    identifier: str = Field(
+        description="A unique name for the model type",
+    )
+    llama_model: str = Field(
+        description="Pointer to the underlying core Llama family model. Each model served by Llama Stack must have a core Llama model.",
+    )
+    metadata: Dict[str, Any] = Field(
+        default_factory=dict,
+        description="Any additional metadata for this model",
+    )
 
 
 @json_schema_type
-class ModelServingSpec(BaseModel):
-    llama_model: Model = Field(
-        description="All metadatas associated with llama model (defined in llama_models.models.sku_list).",
-    )
-    provider_config: GenericProviderConfig = Field(
-        description="Provider config for the model, including provider_type, and corresponding config. ",
+class ModelDefWithProvider(ModelDef):
+    provider_id: str = Field(
+        description="The provider ID for this model",
     )
 
 
+@runtime_checkable
 class Models(Protocol):
     @webmethod(route="/models/list", method="GET")
-    async def list_models(self) -> List[ModelServingSpec]: ...
+    async def list_models(self) -> List[ModelDefWithProvider]: ...
 
     @webmethod(route="/models/get", method="GET")
-    async def get_model(self, core_model_id: str) -> Optional[ModelServingSpec]: ...
+    async def get_model(self, identifier: str) -> Optional[ModelDefWithProvider]: ...
+
+    @webmethod(route="/models/register", method="POST")
+    async def register_model(self, model: ModelDefWithProvider) -> None: ...

llama_stack/apis/safety/client.py

@@ -96,12 +96,6 @@ async def run_main(host: str, port: int, image_path: str = None):
         )
         print(response)
 
-        response = await client.run_shield(
-            shield_type="injection_shield",
-            messages=[message],
-        )
-        print(response)
-
 
 def main(host: str, port: int, image: str = None):
     asyncio.run(run_main(host, port, image))

llama_stack/apis/safety/safety.py

@@ -5,12 +5,13 @@
 # the root directory of this source tree.
 
 from enum import Enum
-from typing import Any, Dict, List, Protocol
+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_models.llama3.api.datatypes import *  # noqa: F403
+from llama_stack.apis.shields import *  # noqa: F403
 
 
 @json_schema_type
@@ -37,7 +38,14 @@ class RunShieldResponse(BaseModel):
     violation: Optional[SafetyViolation] = None
 
 
+class ShieldStore(Protocol):
+    def get_shield(self, identifier: str) -> ShieldDef: ...
+
+
+@runtime_checkable
 class Safety(Protocol):
+    shield_store: ShieldStore
+
     @webmethod(route="/safety/run_shield")
     async def run_shield(
         self, shield_type: str, messages: List[Message], params: Dict[str, Any] = None

llama_stack/apis/shields/client.py

@@ -5,6 +5,7 @@
 # the root directory of this source tree.
 
 import asyncio
+import json
 
 from typing import List, Optional
 
@@ -25,16 +26,27 @@ class ShieldsClient(Shields):
     async def shutdown(self) -> None:
         pass
 
-    async def list_shields(self) -> List[ShieldSpec]:
+    async def list_shields(self) -> List[ShieldDefWithProvider]:
         async with httpx.AsyncClient() as client:
             response = await client.get(
                 f"{self.base_url}/shields/list",
                 headers={"Content-Type": "application/json"},
             )
             response.raise_for_status()
-            return [ShieldSpec(**x) for x in response.json()]
+            return [ShieldDefWithProvider(**x) for x in response.json()]
 
-    async def get_shield(self, shield_type: str) -> Optional[ShieldSpec]:
+    async def register_shield(self, shield: ShieldDefWithProvider) -> None:
+        async with httpx.AsyncClient() as client:
+            response = await client.post(
+                f"{self.base_url}/shields/register",
+                json={
+                    "shield": json.loads(shield.json()),
+                },
+                headers={"Content-Type": "application/json"},
+            )
+            response.raise_for_status()
+
+    async def get_shield(self, shield_type: str) -> Optional[ShieldDefWithProvider]:
         async with httpx.AsyncClient() as client:
             response = await client.get(
                 f"{self.base_url}/shields/get",
@@ -49,7 +61,7 @@ class ShieldsClient(Shields):
             if j is None:
                 return None
 
-            return ShieldSpec(**j)
+            return ShieldDefWithProvider(**j)
 
 
 async def run_main(host: str, port: int, stream: bool):

llama_stack/apis/shields/shields.py

@@ -4,25 +4,48 @@
 # 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, Protocol
+from enum import Enum
+from typing import Any, Dict, List, Optional, Protocol, runtime_checkable
 
 from llama_models.schema_utils import json_schema_type, webmethod
 from pydantic import BaseModel, Field
 
-from llama_stack.distribution.datatypes import GenericProviderConfig
-
 
 @json_schema_type
-class ShieldSpec(BaseModel):
-    shield_type: str
-    provider_config: GenericProviderConfig = Field(
-        description="Provider config for the model, including provider_type, and corresponding config. ",
+class ShieldType(Enum):
+    generic_content_shield = "generic_content_shield"
+    llama_guard = "llama_guard"
+    code_scanner = "code_scanner"
+    prompt_guard = "prompt_guard"
+
+
+class ShieldDef(BaseModel):
+    identifier: str = Field(
+        description="A unique identifier for the shield type",
+    )
+    type: str = Field(
+        description="The type of shield this is; the value is one of the ShieldType enum"
+    )
+    params: Dict[str, Any] = Field(
+        default_factory=dict,
+        description="Any additional parameters needed for this shield",
     )
 
 
+@json_schema_type
+class ShieldDefWithProvider(ShieldDef):
+    provider_id: str = Field(
+        description="The provider ID for this shield type",
+    )
+
+
+@runtime_checkable
 class Shields(Protocol):
     @webmethod(route="/shields/list", method="GET")
-    async def list_shields(self) -> List[ShieldSpec]: ...
+    async def list_shields(self) -> List[ShieldDefWithProvider]: ...
 
     @webmethod(route="/shields/get", method="GET")
-    async def get_shield(self, shield_type: str) -> Optional[ShieldSpec]: ...
+    async def get_shield(self, shield_type: str) -> Optional[ShieldDefWithProvider]: ...
+
+    @webmethod(route="/shields/register", method="POST")
+    async def register_shield(self, shield: ShieldDefWithProvider) -> None: ...

llama_stack/apis/telemetry/telemetry.py

@@ -6,7 +6,7 @@
 
 from datetime import datetime
 from enum import Enum
-from typing import Any, Dict, Literal, Optional, Protocol, Union
+from typing import Any, Dict, Literal, Optional, Protocol, runtime_checkable, Union
 
 from llama_models.schema_utils import json_schema_type, webmethod
 from pydantic import BaseModel, Field
@@ -123,6 +123,7 @@ Event = Annotated[
 ]
 
 
+@runtime_checkable
 class Telemetry(Protocol):
     @webmethod(route="/telemetry/log_event")
     async def log_event(self, event: Event) -> None: ...

llama_stack/cli/stack/build.py

@@ -22,7 +22,7 @@ def available_templates_specs() -> List[BuildConfig]:
     import yaml
 
     template_specs = []
-    for p in TEMPLATES_PATH.rglob("*.yaml"):
+    for p in TEMPLATES_PATH.rglob("*build.yaml"):
         with open(p, "r") as f:
             build_config = BuildConfig(**yaml.safe_load(f))
             template_specs.append(build_config)
@@ -105,8 +105,7 @@ class StackBuild(Subcommand):
 
         import yaml
 
-        from llama_stack.distribution.build import ApiInput, build_image, ImageType
-
+        from llama_stack.distribution.build import build_image, ImageType
         from llama_stack.distribution.utils.config_dirs import DISTRIBS_BASE_DIR
         from llama_stack.distribution.utils.serialize import EnumEncoder
         from termcolor import cprint
@@ -151,8 +150,6 @@ class StackBuild(Subcommand):
     def _run_template_list_cmd(self, args: argparse.Namespace) -> None:
         import json
 
-        import yaml
-
         from llama_stack.cli.table import print_table
 
         # eventually, this should query a registry at llama.meta.com/llamastack/distributions
@@ -178,9 +175,12 @@ class StackBuild(Subcommand):
         )
 
     def _run_stack_build_command(self, args: argparse.Namespace) -> None:
+        import textwrap
+
         import yaml
         from llama_stack.distribution.distribution import get_provider_registry
         from prompt_toolkit import prompt
+        from prompt_toolkit.completion import WordCompleter
         from prompt_toolkit.validation import Validator
         from termcolor import cprint
 
@@ -244,26 +244,31 @@ class StackBuild(Subcommand):
             )
 
             cprint(
-                "\n Llama Stack is composed of several APIs working together. Let's configure the providers (implementations) you want to use for these APIs.",
+                textwrap.dedent(
+                    """
+                Llama Stack is composed of several APIs working together. Let's select
+                the provider types (implementations) you want to use for these APIs.
+                """,
+                ),
                 color="green",
             )
 
+            print("Tip: use <TAB> to see options for the providers.\n")
+
             providers = dict()
             for api, providers_for_api in get_provider_registry().items():
+                available_providers = [
+                    x
+                    for x in providers_for_api.keys()
+                    if x not in ("remote", "remote::sample")
+                ]
                 api_provider = prompt(
-                    "> Enter provider for the {} API: (default=meta-reference): ".format(
-                        api.value
-                    ),
+                    "> Enter provider for API {}: ".format(api.value),
+                    completer=WordCompleter(available_providers),
+                    complete_while_typing=True,
                     validator=Validator.from_callable(
-                        lambda x: x in providers_for_api,
-                        error_message="Invalid provider, please enter one of the following: {}".format(
-                            list(providers_for_api.keys())
-                        ),
-                    ),
-                    default=(
-                        "meta-reference"
-                        if "meta-reference" in providers_for_api
-                        else list(providers_for_api.keys())[0]
+                        lambda x: x in available_providers,
+                        error_message="Invalid provider, use <TAB> to see options",
                     ),
                 )
 

llama_stack/cli/stack/configure.py

@@ -71,9 +71,7 @@ class StackConfigure(Subcommand):
         conda_dir = (
             Path(os.path.expanduser("~/.conda/envs")) / f"llamastack-{args.config}"
         )
-        output = subprocess.check_output(
-            ["bash", "-c", "conda info --json -a"]
-        )
+        output = subprocess.check_output(["bash", "-c", "conda info --json"])
         conda_envs = json.loads(output.decode("utf-8"))["envs"]
 
         for x in conda_envs:
@@ -129,7 +127,10 @@ class StackConfigure(Subcommand):
         import yaml
         from termcolor import cprint
 
-        from llama_stack.distribution.configure import configure_api_providers
+        from llama_stack.distribution.configure import (
+            configure_api_providers,
+            parse_and_maybe_upgrade_config,
+        )
         from llama_stack.distribution.utils.serialize import EnumEncoder
 
         builds_dir = BUILDS_BASE_DIR / build_config.image_type
@@ -145,13 +146,14 @@ class StackConfigure(Subcommand):
                 "yellow",
                 attrs=["bold"],
             )
-            config = StackRunConfig(**yaml.safe_load(run_config_file.read_text()))
+            config_dict = yaml.safe_load(run_config_file.read_text())
+            config = parse_and_maybe_upgrade_config(config_dict)
         else:
             config = StackRunConfig(
                 built_at=datetime.now(),
                 image_name=image_name,
-                apis_to_serve=[],
-                api_providers={},
+                apis=list(build_config.distribution_spec.providers.keys()),
+                providers={},
             )
 
         config = configure_api_providers(config, build_config.distribution_spec)

llama_stack/cli/stack/run.py

@@ -7,7 +7,6 @@
 import argparse
 
 from llama_stack.cli.subcommand import Subcommand
-from llama_stack.distribution.datatypes import *  # noqa: F403
 
 
 class StackRun(Subcommand):
@@ -46,10 +45,11 @@ class StackRun(Subcommand):
 
         import pkg_resources
         import yaml
+        from termcolor import cprint
 
         from llama_stack.distribution.build import ImageType
+        from llama_stack.distribution.configure import parse_and_maybe_upgrade_config
         from llama_stack.distribution.utils.config_dirs import BUILDS_BASE_DIR
-
         from llama_stack.distribution.utils.exec import run_with_pty
 
         if not args.config:
@@ -75,8 +75,10 @@ class StackRun(Subcommand):
             )
             return
 
+        cprint(f"Using config `{config_file}`", "green")
         with open(config_file, "r") as f:
-            config = StackRunConfig(**yaml.safe_load(f))
+            config_dict = yaml.safe_load(config_file.read_text())
+            config = parse_and_maybe_upgrade_config(config_dict)
 
         if config.docker_image:
             script = pkg_resources.resource_filename(

llama_stack/cli/tests/test_stack_build.py

@@ -1,105 +0,0 @@
-from argparse import Namespace
-from unittest.mock import MagicMock, patch
-
-import pytest
-from llama_stack.distribution.datatypes import BuildConfig
-from llama_stack.cli.stack.build import StackBuild
-
-
-# temporary while we make the tests work
-pytest.skip(allow_module_level=True)
-
-
-@pytest.fixture
-def stack_build():
-    parser = MagicMock()
-    subparsers = MagicMock()
-    return StackBuild(subparsers)
-
-
-def test_stack_build_initialization(stack_build):
-    assert stack_build.parser is not None
-    assert stack_build.parser.set_defaults.called_once_with(
-        func=stack_build._run_stack_build_command
-    )
-
-
-@patch("llama_stack.distribution.build.build_image")
-def test_run_stack_build_command_with_config(
-    mock_build_image, mock_build_config, stack_build
-):
-    args = Namespace(
-        config="test_config.yaml",
-        template=None,
-        list_templates=False,
-        name=None,
-        image_type="conda",
-    )
-
-    with patch("builtins.open", MagicMock()):
-        with patch("yaml.safe_load") as mock_yaml_load:
-            mock_yaml_load.return_value = {"name": "test_build", "image_type": "conda"}
-            mock_build_config.return_value = MagicMock()
-
-            stack_build._run_stack_build_command(args)
-
-            mock_build_config.assert_called_once()
-            mock_build_image.assert_called_once()
-
-
-@patch("llama_stack.cli.table.print_table")
-def test_run_stack_build_command_list_templates(mock_print_table, stack_build):
-    args = Namespace(list_templates=True)
-
-    stack_build._run_stack_build_command(args)
-
-    mock_print_table.assert_called_once()
-
-
-@patch("prompt_toolkit.prompt")
-@patch("llama_stack.distribution.datatypes.BuildConfig")
-@patch("llama_stack.distribution.build.build_image")
-def test_run_stack_build_command_interactive(
-    mock_build_image, mock_build_config, mock_prompt, stack_build
-):
-    args = Namespace(
-        config=None, template=None, list_templates=False, name=None, image_type=None
-    )
-
-    mock_prompt.side_effect = [
-        "test_name",
-        "conda",
-        "meta-reference",
-        "test description",
-    ]
-    mock_build_config.return_value = MagicMock()
-
-    stack_build._run_stack_build_command(args)
-
-    assert mock_prompt.call_count == 4
-    mock_build_config.assert_called_once()
-    mock_build_image.assert_called_once()
-
-
-@patch("llama_stack.distribution.datatypes.BuildConfig")
-@patch("llama_stack.distribution.build.build_image")
-def test_run_stack_build_command_with_template(
-    mock_build_image, mock_build_config, stack_build
-):
-    args = Namespace(
-        config=None,
-        template="test_template",
-        list_templates=False,
-        name="test_name",
-        image_type="docker",
-    )
-
-    with patch("builtins.open", MagicMock()):
-        with patch("yaml.safe_load") as mock_yaml_load:
-            mock_yaml_load.return_value = {"name": "test_build", "image_type": "conda"}
-            mock_build_config.return_value = MagicMock()
-
-            stack_build._run_stack_build_command(args)
-
-            mock_build_config.assert_called_once()
-            mock_build_image.assert_called_once()

llama_stack/cli/tests/test_stack_config.py

@@ -0,0 +1,133 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the terms described in the LICENSE file in
+# the root directory of this source tree.
+
+from datetime import datetime
+
+import pytest
+import yaml
+from llama_stack.distribution.configure import (
+    LLAMA_STACK_RUN_CONFIG_VERSION,
+    parse_and_maybe_upgrade_config,
+)
+
+
+@pytest.fixture
+def up_to_date_config():
+    return yaml.safe_load(
+        """
+        version: {version}
+        image_name: foo
+        apis_to_serve: []
+        built_at: {built_at}
+        providers:
+          inference:
+            - provider_id: provider1
+              provider_type: meta-reference
+              config: {{}}
+          safety:
+            - provider_id: provider1
+              provider_type: meta-reference
+              config:
+                llama_guard_shield:
+                  model: Llama-Guard-3-1B
+                  excluded_categories: []
+                  disable_input_check: false
+                  disable_output_check: false
+                enable_prompt_guard: false
+          memory:
+            - provider_id: provider1
+              provider_type: meta-reference
+              config: {{}}
+    """.format(
+            version=LLAMA_STACK_RUN_CONFIG_VERSION, built_at=datetime.now().isoformat()
+        )
+    )
+
+
+@pytest.fixture
+def old_config():
+    return yaml.safe_load(
+        """
+        image_name: foo
+        built_at: {built_at}
+        apis_to_serve: []
+        routing_table:
+          inference:
+            - provider_type: remote::ollama
+              config:
+                host: localhost
+                port: 11434
+              routing_key: Llama3.2-1B-Instruct
+            - provider_type: meta-reference
+              config:
+                model: Llama3.1-8B-Instruct
+              routing_key: Llama3.1-8B-Instruct
+          safety:
+            - routing_key: ["shield1", "shield2"]
+              provider_type: meta-reference
+              config:
+                llama_guard_shield:
+                  model: Llama-Guard-3-1B
+                  excluded_categories: []
+                  disable_input_check: false
+                  disable_output_check: false
+                enable_prompt_guard: false
+          memory:
+            - routing_key: vector
+              provider_type: meta-reference
+              config: {{}}
+        api_providers:
+          telemetry:
+            provider_type: noop
+            config: {{}}
+    """.format(
+            built_at=datetime.now().isoformat()
+        )
+    )
+
+
+@pytest.fixture
+def invalid_config():
+    return yaml.safe_load(
+        """
+        routing_table: {}
+        api_providers: {}
+    """
+    )
+
+
+def test_parse_and_maybe_upgrade_config_up_to_date(up_to_date_config):
+    result = parse_and_maybe_upgrade_config(up_to_date_config)
+    assert result.version == LLAMA_STACK_RUN_CONFIG_VERSION
+    assert "inference" in result.providers
+
+
+def test_parse_and_maybe_upgrade_config_old_format(old_config):
+    result = parse_and_maybe_upgrade_config(old_config)
+    assert result.version == LLAMA_STACK_RUN_CONFIG_VERSION
+    assert all(
+        api in result.providers
+        for api in ["inference", "safety", "memory", "telemetry"]
+    )
+    safety_provider = result.providers["safety"][0]
+    assert safety_provider.provider_type == "meta-reference"
+    assert "llama_guard_shield" in safety_provider.config
+
+    inference_providers = result.providers["inference"]
+    assert len(inference_providers) == 2
+    assert set(x.provider_id for x in inference_providers) == {
+        "remote::ollama-00",
+        "meta-reference-01",
+    }
+
+    ollama = inference_providers[0]
+    assert ollama.provider_type == "remote::ollama"
+    assert ollama.config["port"] == 11434
+
+
+def test_parse_and_maybe_upgrade_config_invalid(invalid_config):
+    with pytest.raises(ValueError):
+        parse_and_maybe_upgrade_config(invalid_config)

llama_stack/distribution/configure.py

@@ -3,189 +3,182 @@
 #
 # This source code is licensed under the terms described in the LICENSE file in
 # the root directory of this source tree.
+import textwrap
 
 from typing import Any
 
-from llama_models.sku_list import (
-    llama3_1_family,
-    llama3_2_family,
-    llama3_family,
-    resolve_model,
-    safety_models,
-)
-
-from pydantic import BaseModel
 from llama_stack.distribution.datatypes import *  # noqa: F403
-from prompt_toolkit import prompt
-from prompt_toolkit.validation import Validator
 from termcolor import cprint
 
-from llama_stack.apis.memory.memory import MemoryBankType
 from llama_stack.distribution.distribution import (
     builtin_automatically_routed_apis,
     get_provider_registry,
-    stack_apis,
 )
 from llama_stack.distribution.utils.dynamic import instantiate_class_type
-
 from llama_stack.distribution.utils.prompt_for_config import prompt_for_config
-from llama_stack.providers.impls.meta_reference.safety.config import (
-    MetaReferenceShieldType,
-)
 
 
-ALLOWED_MODELS = (
-    llama3_family() + llama3_1_family() + llama3_2_family() + safety_models()
-)
+from llama_stack.apis.models import *  # noqa: F403
+from llama_stack.apis.shields import *  # noqa: F403
+from llama_stack.apis.memory_banks import *  # noqa: F403
 
 
-def make_routing_entry_type(config_class: Any):
-    class BaseModelWithConfig(BaseModel):
-        routing_key: str
-        config: config_class
-
-    return BaseModelWithConfig
-
-
-def get_builtin_apis(provider_backed_apis: List[str]) -> List[str]:
-    """Get corresponding builtin APIs given provider backed APIs"""
-    res = []
-    for inf in builtin_automatically_routed_apis():
-        if inf.router_api.value in provider_backed_apis:
-            res.append(inf.routing_table_api.value)
-
-    return res
-
-
-# TODO: make sure we can deal with existing configuration values correctly
-# instead of just overwriting them
-def configure_api_providers(
-    config: StackRunConfig, spec: DistributionSpec
-) -> StackRunConfig:
-    apis = config.apis_to_serve or list(spec.providers.keys())
-    # append the bulitin routing APIs
-    apis += get_builtin_apis(apis)
-
-    router_api2builtin_api = {
-        inf.router_api.value: inf.routing_table_api.value
-        for inf in builtin_automatically_routed_apis()
-    }
-
-    config.apis_to_serve = list(set([a for a in apis if a != "telemetry"]))
-
-    apis = [v.value for v in stack_apis()]
-    all_providers = get_provider_registry()
-
-    # configure simple case for with non-routing providers to api_providers
-    for api_str in spec.providers.keys():
-        if api_str not in apis:
-            raise ValueError(f"Unknown API `{api_str}`")
-
-        cprint(f"Configuring API `{api_str}`...", "green", attrs=["bold"])
-        api = Api(api_str)
-
-        p = spec.providers[api_str]
-        cprint(f"=== Configuring provider `{p}` for API {api_str}...", "green")
-
-        if isinstance(p, list):
-            cprint(
-                f"[WARN] Interactive configuration of multiple providers {p} is not supported, configuring {p[0]} only, please manually configure {p[1:]} in routing_table of run.yaml",
-                "yellow",
-            )
-            p = p[0]
-
-        provider_spec = all_providers[api][p]
+def configure_single_provider(
+    registry: Dict[str, ProviderSpec], provider: Provider
+) -> Provider:
+    provider_spec = registry[provider.provider_type]
     config_type = instantiate_class_type(provider_spec.config_class)
     try:
-            provider_config = config.api_providers.get(api_str)
-            if provider_config:
-                existing = config_type(**provider_config.config)
+        if provider.config:
+            existing = config_type(**provider.config)
         else:
             existing = None
     except Exception:
         existing = None
+
     cfg = prompt_for_config(config_type, existing)
-
-        if api_str in router_api2builtin_api:
-            # a routing api, we need to infer and assign it a routing_key and put it in the routing_table
-            routing_key = "<PLEASE_FILL_ROUTING_KEY>"
-            routing_entries = []
-            if api_str == "inference":
-                if hasattr(cfg, "model"):
-                    routing_key = cfg.model
-                else:
-                    routing_key = prompt(
-                        "> Please enter the supported model your provider has for inference: ",
-                        default="Llama3.1-8B-Instruct",
-                        validator=Validator.from_callable(
-                            lambda x: resolve_model(x) is not None,
-                            error_message="Model must be: {}".format(
-                                [x.descriptor() for x in ALLOWED_MODELS]
-                            ),
-                        ),
-                    )
-                routing_entries.append(
-                    RoutableProviderConfig(
-                        routing_key=routing_key,
-                        provider_type=p,
+    return Provider(
+        provider_id=provider.provider_id,
+        provider_type=provider.provider_type,
         config=cfg.dict(),
     )
+
+
+def configure_api_providers(
+    config: StackRunConfig, build_spec: DistributionSpec
+) -> StackRunConfig:
+    is_nux = len(config.providers) == 0
+
+    if is_nux:
+        print(
+            textwrap.dedent(
+                """
+        Llama Stack is composed of several APIs working together. For each API served by the Stack,
+        we need to configure the providers (implementations) you want to use for these APIs.
+"""
+            )
         )
 
-            if api_str == "safety":
-                # TODO: add support for other safety providers, and simplify safety provider config
-                if p == "meta-reference":
-                    routing_entries.append(
-                        RoutableProviderConfig(
-                            routing_key=[s.value for s in MetaReferenceShieldType],
-                            provider_type=p,
-                            config=cfg.dict(),
-                        )
-                    )
+    provider_registry = get_provider_registry()
+    builtin_apis = [a.routing_table_api for a in builtin_automatically_routed_apis()]
+
+    if config.apis:
+        apis_to_serve = config.apis
     else:
+        apis_to_serve = [a.value for a in Api if a not in (Api.telemetry, Api.inspect)]
+
+    for api_str in apis_to_serve:
+        api = Api(api_str)
+        if api in builtin_apis:
+            continue
+        if api not in provider_registry:
+            raise ValueError(f"Unknown API `{api_str}`")
+
+        existing_providers = config.providers.get(api_str, [])
+        if existing_providers:
             cprint(
-                        f"[WARN] Interactive configuration of safety provider {p} is not supported. Please look for `{routing_key}` in run.yaml and replace it appropriately.",
-                        "yellow",
+                f"Re-configuring existing providers for API `{api_str}`...",
+                "green",
                 attrs=["bold"],
             )
-                    routing_entries.append(
-                        RoutableProviderConfig(
-                            routing_key=routing_key,
-                            provider_type=p,
-                            config=cfg.dict(),
-                        )
-                    )
-
-            if api_str == "memory":
-                bank_types = list([x.value for x in MemoryBankType])
-                routing_key = prompt(
-                    "> Please enter the supported memory bank type your provider has for memory: ",
-                    default="vector",
-                    validator=Validator.from_callable(
-                        lambda x: x in bank_types,
-                        error_message="Invalid provider, please enter one of the following: {}".format(
-                            bank_types
-                        ),
-                    ),
-                )
-                routing_entries.append(
-                    RoutableProviderConfig(
-                        routing_key=routing_key,
-                        provider_type=p,
-                        config=cfg.dict(),
-                    )
-                )
-
-            config.routing_table[api_str] = routing_entries
-            config.api_providers[api_str] = PlaceholderProviderConfig(
-                providers=p if isinstance(p, list) else [p]
+            updated_providers = []
+            for p in existing_providers:
+                print(f"> Configuring provider `({p.provider_type})`")
+                updated_providers.append(
+                    configure_single_provider(provider_registry[api], p)
                 )
+                print("")
         else:
-            config.api_providers[api_str] = GenericProviderConfig(
-                provider_type=p,
-                config=cfg.dict(),
-            )
+            # we are newly configuring this API
+            plist = build_spec.providers.get(api_str, [])
+            plist = plist if isinstance(plist, list) else [plist]
 
+            if not plist:
+                raise ValueError(f"No provider configured for API {api_str}?")
+
+            cprint(f"Configuring API `{api_str}`...", "green", attrs=["bold"])
+            updated_providers = []
+            for i, provider_type in enumerate(plist):
+                print(f"> Configuring provider `({provider_type})`")
+                updated_providers.append(
+                    configure_single_provider(
+                        provider_registry[api],
+                        Provider(
+                            provider_id=(
+                                f"{provider_type}-{i:02d}"
+                                if len(plist) > 1
+                                else provider_type
+                            ),
+                            provider_type=provider_type,
+                            config={},
+                        ),
+                    )
+                )
                 print("")
 
+        config.providers[api_str] = updated_providers
+
     return config
+
+
+def upgrade_from_routing_table(
+    config_dict: Dict[str, Any],
+) -> Dict[str, Any]:
+    def get_providers(entries):
+        return [
+            Provider(
+                provider_id=(
+                    f"{entry['provider_type']}-{i:02d}"
+                    if len(entries) > 1
+                    else entry["provider_type"]
+                ),
+                provider_type=entry["provider_type"],
+                config=entry["config"],
+            )
+            for i, entry in enumerate(entries)
+        ]
+
+    providers_by_api = {}
+
+    routing_table = config_dict.get("routing_table", {})
+    for api_str, entries in routing_table.items():
+        providers = get_providers(entries)
+        providers_by_api[api_str] = providers
+
+    provider_map = config_dict.get("api_providers", config_dict.get("provider_map", {}))
+    if provider_map:
+        for api_str, provider in provider_map.items():
+            if isinstance(provider, dict) and "provider_type" in provider:
+                providers_by_api[api_str] = [
+                    Provider(
+                        provider_id=f"{provider['provider_type']}",
+                        provider_type=provider["provider_type"],
+                        config=provider["config"],
+                    )
+                ]
+
+    config_dict["providers"] = providers_by_api
+
+    config_dict.pop("routing_table", None)
+    config_dict.pop("api_providers", None)
+    config_dict.pop("provider_map", None)
+
+    config_dict["apis"] = config_dict["apis_to_serve"]
+    config_dict.pop("apis_to_serve", None)
+
+    return config_dict
+
+
+def parse_and_maybe_upgrade_config(config_dict: Dict[str, Any]) -> StackRunConfig:
+    version = config_dict.get("version", None)
+    if version == LLAMA_STACK_RUN_CONFIG_VERSION:
+        return StackRunConfig(**config_dict)
+
+    if "routing_table" in config_dict:
+        print("Upgrading config...")
+        config_dict = upgrade_from_routing_table(config_dict)
+
+    config_dict["version"] = LLAMA_STACK_RUN_CONFIG_VERSION
+    config_dict["built_at"] = datetime.now().isoformat()
+
+    return StackRunConfig(**config_dict)

llama_stack/distribution/datatypes.py

@@ -11,28 +11,38 @@ from typing import Dict, List, Optional, Union
 from pydantic import BaseModel, Field
 
 from llama_stack.providers.datatypes import *  # noqa: F403
+from llama_stack.apis.models import *  # noqa: F403
+from llama_stack.apis.shields import *  # noqa: F403
+from llama_stack.apis.memory_banks import *  # noqa: F403
+from llama_stack.apis.inference import Inference
+from llama_stack.apis.memory import Memory
+from llama_stack.apis.safety import Safety
 
 
-LLAMA_STACK_BUILD_CONFIG_VERSION = "v1"
-LLAMA_STACK_RUN_CONFIG_VERSION = "v1"
+LLAMA_STACK_BUILD_CONFIG_VERSION = "2"
+LLAMA_STACK_RUN_CONFIG_VERSION = "2"
 
 
 RoutingKey = Union[str, List[str]]
 
 
-class GenericProviderConfig(BaseModel):
-    provider_type: str
-    config: Dict[str, Any]
+RoutableObject = Union[
+    ModelDef,
+    ShieldDef,
+    MemoryBankDef,
+]
 
+RoutableObjectWithProvider = Union[
+    ModelDefWithProvider,
+    ShieldDefWithProvider,
+    MemoryBankDefWithProvider,
+]
 
-class RoutableProviderConfig(GenericProviderConfig):
-    routing_key: RoutingKey
-
-
-class PlaceholderProviderConfig(BaseModel):
-    """Placeholder provider config for API whose provider are defined in routing_table"""
-
-    providers: List[str]
+RoutedProtocol = Union[
+    Inference,
+    Safety,
+    Memory,
+]
 
 
 # Example: /inference, /safety
@@ -53,18 +63,16 @@ class AutoRoutedProviderSpec(ProviderSpec):
 
 
 # Example: /models, /shields
-@json_schema_type
 class RoutingTableProviderSpec(ProviderSpec):
     provider_type: str = "routing_table"
     config_class: str = ""
     docker_image: Optional[str] = None
 
-    inner_specs: List[ProviderSpec]
+    router_api: Api
     module: str
     pip_packages: List[str] = Field(default_factory=list)
 
 
-@json_schema_type
 class DistributionSpec(BaseModel):
     description: Optional[str] = Field(
         default="",
@@ -80,7 +88,12 @@ in the runtime configuration to help route to the correct provider.""",
     )
 
 
-@json_schema_type
+class Provider(BaseModel):
+    provider_id: str
+    provider_type: str
+    config: Dict[str, Any]
+
+
 class StackRunConfig(BaseModel):
     version: str = LLAMA_STACK_RUN_CONFIG_VERSION
     built_at: datetime
@@ -100,36 +113,20 @@ this could be just a hash
         default=None,
         description="Reference to the conda environment if this package refers to a conda environment",
     )
-    apis_to_serve: List[str] = Field(
+    apis: List[str] = Field(
+        default_factory=list,
         description="""
 The list of APIs to serve. If not specified, all APIs specified in the provider_map will be served""",
     )
 
-    api_providers: Dict[
-        str, Union[GenericProviderConfig, PlaceholderProviderConfig]
-    ] = Field(
+    providers: Dict[str, List[Provider]] = Field(
         description="""
-Provider configurations for each of the APIs provided by this package.
+One or more providers to use for each API. The same provider_type (e.g., meta-reference)
+can be instantiated multiple times (with different configs) if necessary.
 """,
     )
-    routing_table: Dict[str, List[RoutableProviderConfig]] = Field(
-        default_factory=dict,
-        description="""
-
-        E.g. The following is a ProviderRoutingEntry for models:
-        - routing_key: Llama3.1-8B-Instruct
-          provider_type: meta-reference
-          config:
-              model: Llama3.1-8B-Instruct
-              quantization: null
-              torch_seed: null
-              max_seq_len: 4096
-              max_batch_size: 1
-        """,
-    )
 
 
-@json_schema_type
 class BuildConfig(BaseModel):
     version: str = LLAMA_STACK_BUILD_CONFIG_VERSION
     name: str

llama_stack/distribution/inspect.py

@@ -6,45 +6,58 @@
 
 from typing import Dict, List
 from llama_stack.apis.inspect import *  # noqa: F403
+from pydantic import BaseModel
 
-
-from llama_stack.distribution.distribution import get_provider_registry
 from llama_stack.distribution.server.endpoints import get_all_api_endpoints
 from llama_stack.providers.datatypes import *  # noqa: F403
+from llama_stack.distribution.datatypes import *  # noqa: F403
 
 
-def is_passthrough(spec: ProviderSpec) -> bool:
-    return isinstance(spec, RemoteProviderSpec) and spec.adapter is None
+class DistributionInspectConfig(BaseModel):
+    run_config: StackRunConfig
+
+
+async def get_provider_impl(config, deps):
+    impl = DistributionInspectImpl(config, deps)
+    await impl.initialize()
+    return impl
 
 
 class DistributionInspectImpl(Inspect):
-    def __init__(self):
+    def __init__(self, config, deps):
+        self.config = config
+        self.deps = deps
+
+    async def initialize(self) -> None:
         pass
 
     async def list_providers(self) -> Dict[str, List[ProviderInfo]]:
+        run_config = self.config.run_config
+
         ret = {}
-        all_providers = get_provider_registry()
-        for api, providers in all_providers.items():
-            ret[api.value] = [
+        for api, providers in run_config.providers.items():
+            ret[api] = [
                 ProviderInfo(
+                    provider_id=p.provider_id,
                     provider_type=p.provider_type,
-                    description="Passthrough" if is_passthrough(p) else "",
                 )
-                for p in providers.values()
+                for p in providers
             ]
 
         return ret
 
     async def list_routes(self) -> Dict[str, List[RouteInfo]]:
+        run_config = self.config.run_config
+
         ret = {}
         all_endpoints = get_all_api_endpoints()
-
         for api, endpoints in all_endpoints.items():
+            providers = run_config.providers.get(api.value, [])
             ret[api.value] = [
                 RouteInfo(
                     route=e.route,
                     method=e.method,
-                    providers=[],
+                    provider_types=[p.provider_type for p in providers],
                 )
                 for e in endpoints
             ]

llama_stack/distribution/resolver.py

@@ -4,146 +4,237 @@
 # This source code is licensed under the terms described in the LICENSE file in
 # the root directory of this source tree.
 import importlib
+import inspect
 
 from typing import Any, Dict, List, Set
 
+from llama_stack.providers.datatypes import *  # noqa: F403
 from llama_stack.distribution.datatypes import *  # noqa: F403
+
+from llama_stack.apis.agents import Agents
+from llama_stack.apis.inference import Inference
+from llama_stack.apis.inspect import Inspect
+from llama_stack.apis.memory import Memory
+from llama_stack.apis.memory_banks import MemoryBanks
+from llama_stack.apis.models import Models
+from llama_stack.apis.safety import Safety
+from llama_stack.apis.shields import Shields
+from llama_stack.apis.telemetry import Telemetry
 from llama_stack.distribution.distribution import (
     builtin_automatically_routed_apis,
     get_provider_registry,
 )
-from llama_stack.distribution.inspect import DistributionInspectImpl
 from llama_stack.distribution.utils.dynamic import instantiate_class_type
 
 
+def api_protocol_map() -> Dict[Api, Any]:
+    return {
+        Api.agents: Agents,
+        Api.inference: Inference,
+        Api.inspect: Inspect,
+        Api.memory: Memory,
+        Api.memory_banks: MemoryBanks,
+        Api.models: Models,
+        Api.safety: Safety,
+        Api.shields: Shields,
+        Api.telemetry: Telemetry,
+    }
+
+
+def additional_protocols_map() -> Dict[Api, Any]:
+    return {
+        Api.inference: ModelsProtocolPrivate,
+        Api.memory: MemoryBanksProtocolPrivate,
+        Api.safety: ShieldsProtocolPrivate,
+    }
+
+
+# TODO: make all this naming far less atrocious. Provider. ProviderSpec. ProviderWithSpec. WTF!
+class ProviderWithSpec(Provider):
+    spec: ProviderSpec
+
+
+# TODO: this code is not very straightforward to follow and needs one more round of refactoring
 async def resolve_impls_with_routing(run_config: StackRunConfig) -> Dict[Api, Any]:
     """
     Does two things:
     - flatmaps, sorts and resolves the providers in dependency order
     - for each API, produces either a (local, passthrough or router) implementation
     """
-    all_providers = get_provider_registry()
-    specs = {}
-    configs = {}
+    all_api_providers = get_provider_registry()
 
-    for api_str, config in run_config.api_providers.items():
+    routing_table_apis = set(
+        x.routing_table_api for x in builtin_automatically_routed_apis()
+    )
+    router_apis = set(x.router_api for x in builtin_automatically_routed_apis())
+
+    providers_with_specs = {}
+
+    for api_str, providers in run_config.providers.items():
         api = Api(api_str)
-
-        # TODO: check that these APIs are not in the routing table part of the config
-        providers = all_providers[api]
-
-        # skip checks for API whose provider config is specified in routing_table
-        if isinstance(config, PlaceholderProviderConfig):
-            continue
-
-        if config.provider_type not in providers:
+        if api in routing_table_apis:
             raise ValueError(
-                f"Provider `{config.provider_type}` is not available for API `{api}`"
+                f"Provider for `{api_str}` is automatically provided and cannot be overridden"
             )
-        specs[api] = providers[config.provider_type]
-        configs[api] = config
 
-    apis_to_serve = run_config.apis_to_serve or set(
-        list(specs.keys()) + list(run_config.routing_table.keys())
+        specs = {}
+        for provider in providers:
+            if provider.provider_type not in all_api_providers[api]:
+                raise ValueError(
+                    f"Provider `{provider.provider_type}` is not available for API `{api}`"
                 )
+
+            p = all_api_providers[api][provider.provider_type]
+            p.deps__ = [a.value for a in p.api_dependencies]
+            spec = ProviderWithSpec(
+                spec=p,
+                **(provider.dict()),
+            )
+            specs[provider.provider_id] = spec
+
+        key = api_str if api not in router_apis else f"inner-{api_str}"
+        providers_with_specs[key] = specs
+
+    apis_to_serve = run_config.apis or set(
+        list(providers_with_specs.keys())
+        + [x.value for x in routing_table_apis]
+        + [x.value for x in router_apis]
+    )
+
     for info in builtin_automatically_routed_apis():
-        source_api = info.routing_table_api
-
-        assert (
-            source_api not in specs
-        ), f"Routing table API {source_api} specified in wrong place?"
-        assert (
-            info.router_api not in specs
-        ), f"Auto-routed API {info.router_api} specified in wrong place?"
-
         if info.router_api.value not in apis_to_serve:
             continue
 
-        if info.router_api.value not in run_config.routing_table:
-            raise ValueError(f"Routing table for `{source_api.value}` is not provided?")
+        available_providers = providers_with_specs[f"inner-{info.router_api.value}"]
 
-        routing_table = run_config.routing_table[info.router_api.value]
-
-        providers = all_providers[info.router_api]
-
-        inner_specs = []
-        inner_deps = []
-        for rt_entry in routing_table:
-            if rt_entry.provider_type not in providers:
-                raise ValueError(
-                    f"Provider `{rt_entry.provider_type}` is not available for API `{api}`"
-                )
-            inner_specs.append(providers[rt_entry.provider_type])
-            inner_deps.extend(providers[rt_entry.provider_type].api_dependencies)
-
-        specs[source_api] = RoutingTableProviderSpec(
-            api=source_api,
+        providers_with_specs[info.routing_table_api.value] = {
+            "__builtin__": ProviderWithSpec(
+                provider_id="__routing_table__",
+                provider_type="__routing_table__",
+                config={},
+                spec=RoutingTableProviderSpec(
+                    api=info.routing_table_api,
+                    router_api=info.router_api,
                     module="llama_stack.distribution.routers",
-            api_dependencies=inner_deps,
-            inner_specs=inner_specs,
+                    api_dependencies=[],
+                    deps__=([f"inner-{info.router_api.value}"]),
+                ),
             )
-        configs[source_api] = routing_table
+        }
 
-        specs[info.router_api] = AutoRoutedProviderSpec(
+        providers_with_specs[info.router_api.value] = {
+            "__builtin__": ProviderWithSpec(
+                provider_id="__autorouted__",
+                provider_type="__autorouted__",
+                config={},
+                spec=AutoRoutedProviderSpec(
                     api=info.router_api,
                     module="llama_stack.distribution.routers",
-            routing_table_api=source_api,
-            api_dependencies=[source_api],
+                    routing_table_api=info.routing_table_api,
+                    api_dependencies=[info.routing_table_api],
+                    deps__=([info.routing_table_api.value]),
+                ),
             )
-        configs[info.router_api] = {}
+        }
 
-    sorted_specs = topological_sort(specs.values())
-    print(f"Resolved {len(sorted_specs)} providers in topological order")
-    for spec in sorted_specs:
-        print(f"  {spec.api}: {spec.provider_type}")
+    sorted_providers = topological_sort(
+        {k: v.values() for k, v in providers_with_specs.items()}
+    )
+    apis = [x[1].spec.api for x in sorted_providers]
+    sorted_providers.append(
+        (
+            "inspect",
+            ProviderWithSpec(
+                provider_id="__builtin__",
+                provider_type="__builtin__",
+                config={
+                    "run_config": run_config.dict(),
+                },
+                spec=InlineProviderSpec(
+                    api=Api.inspect,
+                    provider_type="__builtin__",
+                    config_class="llama_stack.distribution.inspect.DistributionInspectConfig",
+                    module="llama_stack.distribution.inspect",
+                    api_dependencies=apis,
+                    deps__=([x.value for x in apis]),
+                ),
+            ),
+        )
+    )
+
+    print(f"Resolved {len(sorted_providers)} providers")
+    for api_str, provider in sorted_providers:
+        print(f" {api_str} => {provider.provider_id}")
     print("")
-    impls = {}
-    for spec in sorted_specs:
-        api = spec.api
-        deps = {api: impls[api] for api in spec.api_dependencies}
-        impl = await instantiate_provider(spec, deps, configs[api])
 
+    impls = {}
+    inner_impls_by_provider_id = {f"inner-{x.value}": {} for x in router_apis}
+    for api_str, provider in sorted_providers:
+        deps = {a: impls[a] for a in provider.spec.api_dependencies}
+
+        inner_impls = {}
+        if isinstance(provider.spec, RoutingTableProviderSpec):
+            inner_impls = inner_impls_by_provider_id[
+                f"inner-{provider.spec.router_api.value}"
+            ]
+
+        impl = await instantiate_provider(
+            provider,
+            deps,
+            inner_impls,
+        )
+        # TODO: ugh slightly redesign this shady looking code
+        if "inner-" in api_str:
+            inner_impls_by_provider_id[api_str][provider.provider_id] = impl
+        else:
+            api = Api(api_str)
             impls[api] = impl
 
-    impls[Api.inspect] = DistributionInspectImpl()
-    specs[Api.inspect] = InlineProviderSpec(
-        api=Api.inspect,
-        provider_type="__distribution_builtin__",
-        config_class="",
-        module="",
-    )
-
-    return impls, specs
+    return impls
 
 
-def topological_sort(providers: List[ProviderSpec]) -> List[ProviderSpec]:
-    by_id = {x.api: x for x in providers}
+def topological_sort(
+    providers_with_specs: Dict[str, List[ProviderWithSpec]],
+) -> List[ProviderWithSpec]:
+    def dfs(kv, visited: Set[str], stack: List[str]):
+        api_str, providers = kv
+        visited.add(api_str)
 
-    def dfs(a: ProviderSpec, visited: Set[Api], stack: List[Api]):
-        visited.add(a.api)
+        deps = []
+        for provider in providers:
+            for dep in provider.spec.deps__:
+                deps.append(dep)
 
-        for api in a.api_dependencies:
-            if api not in visited:
-                dfs(by_id[api], visited, stack)
+        for dep in deps:
+            if dep not in visited:
+                dfs((dep, providers_with_specs[dep]), visited, stack)
 
-        stack.append(a.api)
+        stack.append(api_str)
 
     visited = set()
     stack = []
 
-    for a in providers:
-        if a.api not in visited:
-            dfs(a, visited, stack)
+    for api_str, providers in providers_with_specs.items():
+        if api_str not in visited:
+            dfs((api_str, providers), visited, stack)
 
-    return [by_id[x] for x in stack]
+    flattened = []
+    for api_str in stack:
+        for provider in providers_with_specs[api_str]:
+            flattened.append((api_str, provider))
+    return flattened
 
 
 # returns a class implementing the protocol corresponding to the Api
 async def instantiate_provider(
-    provider_spec: ProviderSpec,
+    provider: ProviderWithSpec,
     deps: Dict[str, Any],
-    provider_config: Union[GenericProviderConfig, RoutingTable],
+    inner_impls: Dict[str, Any],
 ):
+    protocols = api_protocol_map()
+    additional_protocols = additional_protocols_map()
+
+    provider_spec = provider.spec
     module = importlib.import_module(provider_spec.module)
 
     args = []
@@ -153,9 +244,8 @@ async def instantiate_provider(
         else:
             method = "get_client_impl"
 
-        assert isinstance(provider_config, GenericProviderConfig)
         config_type = instantiate_class_type(provider_spec.config_class)
-        config = config_type(**provider_config.config)
+        config = config_type(**provider.config)
         args = [config, deps]
     elif isinstance(provider_spec, AutoRoutedProviderSpec):
         method = "get_auto_router_impl"
@@ -165,31 +255,69 @@ async def instantiate_provider(
     elif isinstance(provider_spec, RoutingTableProviderSpec):
         method = "get_routing_table_impl"
 
-        assert isinstance(provider_config, List)
-        routing_table = provider_config
-
-        inner_specs = {x.provider_type: x for x in provider_spec.inner_specs}
-        inner_impls = []
-        for routing_entry in routing_table:
-            impl = await instantiate_provider(
-                inner_specs[routing_entry.provider_type],
-                deps,
-                routing_entry,
-            )
-            inner_impls.append((routing_entry.routing_key, impl))
-
         config = None
-        args = [provider_spec.api, inner_impls, routing_table, deps]
+        args = [provider_spec.api, inner_impls, deps]
     else:
         method = "get_provider_impl"
 
-        assert isinstance(provider_config, GenericProviderConfig)
         config_type = instantiate_class_type(provider_spec.config_class)
-        config = config_type(**provider_config.config)
+        config = config_type(**provider.config)
         args = [config, deps]
 
     fn = getattr(module, method)
     impl = await fn(*args)
+    impl.__provider_id__ = provider.provider_id
     impl.__provider_spec__ = provider_spec
     impl.__provider_config__ = config
+
+    check_protocol_compliance(impl, protocols[provider_spec.api])
+    if (
+        not isinstance(provider_spec, AutoRoutedProviderSpec)
+        and provider_spec.api in additional_protocols
+    ):
+        additional_api = additional_protocols[provider_spec.api]
+        check_protocol_compliance(impl, additional_api)
+
     return impl
+
+
+def check_protocol_compliance(obj: Any, protocol: Any) -> None:
+    missing_methods = []
+
+    mro = type(obj).__mro__
+    for name, value in inspect.getmembers(protocol):
+        if inspect.isfunction(value) and hasattr(value, "__webmethod__"):
+            if not hasattr(obj, name):
+                missing_methods.append((name, "missing"))
+            elif not callable(getattr(obj, name)):
+                missing_methods.append((name, "not_callable"))
+            else:
+                # Check if the method signatures are compatible
+                obj_method = getattr(obj, name)
+                proto_sig = inspect.signature(value)
+                obj_sig = inspect.signature(obj_method)
+
+                proto_params = set(proto_sig.parameters)
+                proto_params.discard("self")
+                obj_params = set(obj_sig.parameters)
+                obj_params.discard("self")
+                if not (proto_params <= obj_params):
+                    print(
+                        f"Method {name} incompatible proto: {proto_params} vs. obj: {obj_params}"
+                    )
+                    missing_methods.append((name, "signature_mismatch"))
+                else:
+                    # Check if the method is actually implemented in the class
+                    method_owner = next(
+                        (cls for cls in mro if name in cls.__dict__), None
+                    )
+                    if (
+                        method_owner is None
+                        or method_owner.__name__ == protocol.__name__
+                    ):
+                        missing_methods.append((name, "not_actually_implemented"))
+
+    if missing_methods:
+        raise ValueError(
+            f"Provider `{obj.__provider_id__} ({obj.__provider_spec__.api})` does not implement the following methods:\n{missing_methods}"
+        )

llama_stack/distribution/routers/__init__.py

@@ -4,23 +4,21 @@
 # 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, List, Tuple
+from typing import Any
 
 from llama_stack.distribution.datatypes import *  # noqa: F403
+from .routing_tables import (
+    MemoryBanksRoutingTable,
+    ModelsRoutingTable,
+    ShieldsRoutingTable,
+)
 
 
 async def get_routing_table_impl(
     api: Api,
-    inner_impls: List[Tuple[str, Any]],
-    routing_table_config: Dict[str, List[RoutableProviderConfig]],
+    impls_by_provider_id: Dict[str, RoutedProtocol],
     _deps,
 ) -> Any:
-    from .routing_tables import (
-        MemoryBanksRoutingTable,
-        ModelsRoutingTable,
-        ShieldsRoutingTable,
-    )
-
     api_to_tables = {
         "memory_banks": MemoryBanksRoutingTable,
         "models": ModelsRoutingTable,
@@ -29,7 +27,7 @@ async def get_routing_table_impl(
     if api.value not in api_to_tables:
         raise ValueError(f"API {api.value} not found in router map")
 
-    impl = api_to_tables[api.value](inner_impls, routing_table_config)
+    impl = api_to_tables[api.value](impls_by_provider_id)
     await impl.initialize()
     return impl
 

llama_stack/distribution/routers/routers.py

@@ -14,14 +14,13 @@ from llama_stack.apis.safety import *  # noqa: F403
 
 
 class MemoryRouter(Memory):
-    """Routes to an provider based on the memory bank type"""
+    """Routes to an provider based on the memory bank identifier"""
 
     def __init__(
         self,
         routing_table: RoutingTable,
     ) -> None:
         self.routing_table = routing_table
-        self.bank_id_to_type = {}
 
     async def initialize(self) -> None:
         pass
@@ -29,32 +28,8 @@ class MemoryRouter(Memory):
     async def shutdown(self) -> None:
         pass
 
-    def get_provider_from_bank_id(self, bank_id: str) -> Any:
-        bank_type = self.bank_id_to_type.get(bank_id)
-        if not bank_type:
-            raise ValueError(f"Could not find bank type for {bank_id}")
-
-        provider = self.routing_table.get_provider_impl(bank_type)
-        if not provider:
-            raise ValueError(f"Could not find provider for {bank_type}")
-        return provider
-
-    async def create_memory_bank(
-        self,
-        name: str,
-        config: MemoryBankConfig,
-        url: Optional[URL] = None,
-    ) -> MemoryBank:
-        bank_type = config.type
-        bank = await self.routing_table.get_provider_impl(bank_type).create_memory_bank(
-            name, config, url
-        )
-        self.bank_id_to_type[bank.bank_id] = bank_type
-        return bank
-
-    async def get_memory_bank(self, bank_id: str) -> Optional[MemoryBank]:
-        provider = self.get_provider_from_bank_id(bank_id)
-        return await provider.get_memory_bank(bank_id)
+    async def register_memory_bank(self, memory_bank: MemoryBankDef) -> None:
+        await self.routing_table.register_memory_bank(memory_bank)
 
     async def insert_documents(
         self,
@@ -62,7 +37,7 @@ class MemoryRouter(Memory):
         documents: List[MemoryBankDocument],
         ttl_seconds: Optional[int] = None,
     ) -> None:
-        return await self.get_provider_from_bank_id(bank_id).insert_documents(
+        return await self.routing_table.get_provider_impl(bank_id).insert_documents(
             bank_id, documents, ttl_seconds
         )
 
@@ -72,7 +47,7 @@ class MemoryRouter(Memory):
         query: InterleavedTextMedia,
         params: Optional[Dict[str, Any]] = None,
     ) -> QueryDocumentsResponse:
-        return await self.get_provider_from_bank_id(bank_id).query_documents(
+        return await self.routing_table.get_provider_impl(bank_id).query_documents(
             bank_id, query, params
         )
 
@@ -92,7 +67,10 @@ class InferenceRouter(Inference):
     async def shutdown(self) -> None:
         pass
 
-    async def chat_completion(
+    async def register_model(self, model: ModelDef) -> None:
+        await self.routing_table.register_model(model)
+
+    def chat_completion(
         self,
         model: str,
         messages: List[Message],
@@ -113,27 +91,32 @@ class InferenceRouter(Inference):
             stream=stream,
             logprobs=logprobs,
         )
-        # TODO: we need to fix streaming response to align provider implementations with Protocol.
-        async for chunk in self.routing_table.get_provider_impl(model).chat_completion(
-            **params
-        ):
-            yield chunk
+        provider = self.routing_table.get_provider_impl(model)
+        if stream:
+            return (chunk async for chunk in provider.chat_completion(**params))
+        else:
+            return provider.chat_completion(**params)
 
-    async def completion(
+    def completion(
         self,
         model: str,
         content: InterleavedTextMedia,
         sampling_params: Optional[SamplingParams] = SamplingParams(),
         stream: Optional[bool] = False,
         logprobs: Optional[LogProbConfig] = None,
-    ) -> Union[CompletionResponse, CompletionResponseStreamChunk]:
-        return await self.routing_table.get_provider_impl(model).completion(
+    ) -> AsyncGenerator:
+        provider = self.routing_table.get_provider_impl(model)
+        params = dict(
             model=model,
             content=content,
             sampling_params=sampling_params,
             stream=stream,
             logprobs=logprobs,
         )
+        if stream:
+            return (chunk async for chunk in provider.completion(**params))
+        else:
+            return provider.completion(**params)
 
     async def embeddings(
         self,
@@ -159,6 +142,9 @@ class SafetyRouter(Safety):
     async def shutdown(self) -> None:
         pass
 
+    async def register_shield(self, shield: ShieldDef) -> None:
+        await self.routing_table.register_shield(shield)
+
     async def run_shield(
         self,
         shield_type: str,

llama_stack/distribution/routers/routing_tables.py

@@ -4,9 +4,8 @@
 # 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, List, Optional, Tuple
+from typing import Any, Dict, List, Optional
 
-from llama_models.sku_list import resolve_model
 from llama_models.llama3.api.datatypes import *  # noqa: F403
 
 from llama_stack.apis.models import *  # noqa: F403
@@ -16,129 +15,159 @@ from llama_stack.apis.memory_banks import *  # noqa: F403
 from llama_stack.distribution.datatypes import *  # noqa: F403
 
 
+def get_impl_api(p: Any) -> Api:
+    return p.__provider_spec__.api
+
+
+async def register_object_with_provider(obj: RoutableObject, p: Any) -> None:
+    api = get_impl_api(p)
+    if api == Api.inference:
+        await p.register_model(obj)
+    elif api == Api.safety:
+        await p.register_shield(obj)
+    elif api == Api.memory:
+        await p.register_memory_bank(obj)
+
+
+Registry = Dict[str, List[RoutableObjectWithProvider]]
+
+
+# TODO: this routing table maintains state in memory purely. We need to
+# add persistence to it when we add dynamic registration of objects.
 class CommonRoutingTableImpl(RoutingTable):
     def __init__(
         self,
-        inner_impls: List[Tuple[RoutingKey, Any]],
-        routing_table_config: Dict[str, List[RoutableProviderConfig]],
+        impls_by_provider_id: Dict[str, RoutedProtocol],
     ) -> None:
-        self.unique_providers = []
-        self.providers = {}
-        self.routing_keys = []
-
-        for key, impl in inner_impls:
-            keys = key if isinstance(key, list) else [key]
-            self.unique_providers.append((keys, impl))
-
-            for k in keys:
-                if k in self.providers:
-                    raise ValueError(f"Duplicate routing key {k}")
-                self.providers[k] = impl
-                self.routing_keys.append(k)
-
-        self.routing_table_config = routing_table_config
+        self.impls_by_provider_id = impls_by_provider_id
 
     async def initialize(self) -> None:
-        for keys, p in self.unique_providers:
-            spec = p.__provider_spec__
-            if isinstance(spec, RemoteProviderSpec) and spec.adapter is None:
-                continue
+        self.registry: Registry = {}
 
-            await p.validate_routing_keys(keys)
+        def add_objects(objs: List[RoutableObjectWithProvider]) -> None:
+            for obj in objs:
+                if obj.identifier not in self.registry:
+                    self.registry[obj.identifier] = []
+
+                self.registry[obj.identifier].append(obj)
+
+        for pid, p in self.impls_by_provider_id.items():
+            api = get_impl_api(p)
+            if api == Api.inference:
+                p.model_store = self
+                models = await p.list_models()
+                add_objects(
+                    [ModelDefWithProvider(**m.dict(), provider_id=pid) for m in models]
+                )
+
+            elif api == Api.safety:
+                p.shield_store = self
+                shields = await p.list_shields()
+                add_objects(
+                    [
+                        ShieldDefWithProvider(**s.dict(), provider_id=pid)
+                        for s in shields
+                    ]
+                )
+
+            elif api == Api.memory:
+                p.memory_bank_store = self
+                memory_banks = await p.list_memory_banks()
+
+                # do in-memory updates due to pesky Annotated unions
+                for m in memory_banks:
+                    m.provider_id = pid
+
+                add_objects(memory_banks)
 
     async def shutdown(self) -> None:
-        for _, p in self.unique_providers:
+        for p in self.impls_by_provider_id.values():
             await p.shutdown()
 
-    def get_provider_impl(self, routing_key: str) -> Any:
-        if routing_key not in self.providers:
-            raise ValueError(f"Could not find provider for {routing_key}")
-        return self.providers[routing_key]
+    def get_provider_impl(
+        self, routing_key: str, provider_id: Optional[str] = None
+    ) -> Any:
+        if routing_key not in self.registry:
+            raise ValueError(f"`{routing_key}` not registered")
 
-    def get_routing_keys(self) -> List[str]:
-        return self.routing_keys
+        objs = self.registry[routing_key]
+        for obj in objs:
+            if not provider_id or provider_id == obj.provider_id:
+                return self.impls_by_provider_id[obj.provider_id]
 
-    def get_provider_config(self, routing_key: str) -> Optional[GenericProviderConfig]:
-        for entry in self.routing_table_config:
-            if entry.routing_key == routing_key:
-                return entry
+        raise ValueError(f"Provider not found for `{routing_key}`")
+
+    def get_object_by_identifier(
+        self, identifier: str
+    ) -> Optional[RoutableObjectWithProvider]:
+        objs = self.registry.get(identifier, [])
+        if not objs:
             return None
 
+        # kind of ill-defined behavior here, but we'll just return the first one
+        return objs[0]
+
+    async def register_object(self, obj: RoutableObjectWithProvider):
+        entries = self.registry.get(obj.identifier, [])
+        for entry in entries:
+            if entry.provider_id == obj.provider_id:
+                print(f"`{obj.identifier}` already registered with `{obj.provider_id}`")
+                return
+
+        if obj.provider_id not in self.impls_by_provider_id:
+            raise ValueError(f"Provider `{obj.provider_id}` not found")
+
+        p = self.impls_by_provider_id[obj.provider_id]
+        await register_object_with_provider(obj, p)
+
+        if obj.identifier not in self.registry:
+            self.registry[obj.identifier] = []
+        self.registry[obj.identifier].append(obj)
+
+        # TODO: persist this to a store
+
 
 class ModelsRoutingTable(CommonRoutingTableImpl, Models):
+    async def list_models(self) -> List[ModelDefWithProvider]:
+        objects = []
+        for objs in self.registry.values():
+            objects.extend(objs)
+        return objects
 
-    async def list_models(self) -> List[ModelServingSpec]:
-        specs = []
-        for entry in self.routing_table_config:
-            model_id = entry.routing_key
-            specs.append(
-                ModelServingSpec(
-                    llama_model=resolve_model(model_id),
-                    provider_config=entry,
-                )
-            )
-        return specs
+    async def get_model(self, identifier: str) -> Optional[ModelDefWithProvider]:
+        return self.get_object_by_identifier(identifier)
 
-    async def get_model(self, core_model_id: str) -> Optional[ModelServingSpec]:
-        for entry in self.routing_table_config:
-            if entry.routing_key == core_model_id:
-                return ModelServingSpec(
-                    llama_model=resolve_model(core_model_id),
-                    provider_config=entry,
-                )
-        return None
+    async def register_model(self, model: ModelDefWithProvider) -> None:
+        await self.register_object(model)
 
 
 class ShieldsRoutingTable(CommonRoutingTableImpl, Shields):
+    async def list_shields(self) -> List[ShieldDef]:
+        objects = []
+        for objs in self.registry.values():
+            objects.extend(objs)
+        return objects
 
-    async def list_shields(self) -> List[ShieldSpec]:
-        specs = []
-        for entry in self.routing_table_config:
-            if isinstance(entry.routing_key, list):
-                for k in entry.routing_key:
-                    specs.append(
-                        ShieldSpec(
-                            shield_type=k,
-                            provider_config=entry,
-                        )
-                    )
-            else:
-                specs.append(
-                    ShieldSpec(
-                        shield_type=entry.routing_key,
-                        provider_config=entry,
-                    )
-                )
-        return specs
+    async def get_shield(self, shield_type: str) -> Optional[ShieldDefWithProvider]:
+        return self.get_object_by_identifier(shield_type)
 
-    async def get_shield(self, shield_type: str) -> Optional[ShieldSpec]:
-        for entry in self.routing_table_config:
-            if entry.routing_key == shield_type:
-                return ShieldSpec(
-                    shield_type=entry.routing_key,
-                    provider_config=entry,
-                )
-        return None
+    async def register_shield(self, shield: ShieldDefWithProvider) -> None:
+        await self.register_object(shield)
 
 
 class MemoryBanksRoutingTable(CommonRoutingTableImpl, MemoryBanks):
+    async def list_memory_banks(self) -> List[MemoryBankDefWithProvider]:
+        objects = []
+        for objs in self.registry.values():
+            objects.extend(objs)
+        return objects
 
-    async def list_available_memory_banks(self) -> List[MemoryBankSpec]:
-        specs = []
-        for entry in self.routing_table_config:
-            specs.append(
-                MemoryBankSpec(
-                    bank_type=entry.routing_key,
-                    provider_config=entry,
-                )
-            )
-        return specs
+    async def get_memory_bank(
+        self, identifier: str
+    ) -> Optional[MemoryBankDefWithProvider]:
+        return self.get_object_by_identifier(identifier)
 
-    async def get_serving_memory_bank(self, bank_type: str) -> Optional[MemoryBankSpec]:
-        for entry in self.routing_table_config:
-            if entry.routing_key == bank_type:
-                return MemoryBankSpec(
-                    bank_type=entry.routing_key,
-                    provider_config=entry,
-                )
-        return None
+    async def register_memory_bank(
+        self, memory_bank: MemoryBankDefWithProvider
+    ) -> None:
+        await self.register_object(memory_bank)

llama_stack/distribution/server/endpoints.py

@@ -9,15 +9,7 @@ from typing import Dict, List
 
 from pydantic import BaseModel
 
-from llama_stack.apis.agents import Agents
-from llama_stack.apis.inference import Inference
-from llama_stack.apis.inspect import Inspect
-from llama_stack.apis.memory import Memory
-from llama_stack.apis.memory_banks import MemoryBanks
-from llama_stack.apis.models import Models
-from llama_stack.apis.safety import Safety
-from llama_stack.apis.shields import Shields
-from llama_stack.apis.telemetry import Telemetry
+from llama_stack.distribution.resolver import api_protocol_map
 
 from llama_stack.providers.datatypes import Api
 
@@ -31,18 +23,7 @@ class ApiEndpoint(BaseModel):
 def get_all_api_endpoints() -> Dict[Api, List[ApiEndpoint]]:
     apis = {}
 
-    protocols = {
-        Api.inference: Inference,
-        Api.safety: Safety,
-        Api.agents: Agents,
-        Api.memory: Memory,
-        Api.telemetry: Telemetry,
-        Api.models: Models,
-        Api.shields: Shields,
-        Api.memory_banks: MemoryBanks,
-        Api.inspect: Inspect,
-    }
-
+    protocols = api_protocol_map()
     for api, protocol in protocols.items():
         endpoints = []
         protocol_methods = inspect.getmembers(protocol, predicate=inspect.isfunction)

llama_stack/distribution/server/server.py

@@ -5,18 +5,15 @@
 # the root directory of this source tree.
 
 import asyncio
+import functools
 import inspect
 import json
 import signal
 import traceback
 
-from collections.abc import (
-    AsyncGenerator as AsyncGeneratorABC,
-    AsyncIterator as AsyncIteratorABC,
-)
 from contextlib import asynccontextmanager
 from ssl import SSLError
-from typing import Any, AsyncGenerator, AsyncIterator, Dict, get_type_hints, Optional
+from typing import Any, Dict, Optional
 
 import fire
 import httpx
@@ -29,6 +26,8 @@ from pydantic import BaseModel, ValidationError
 from termcolor import cprint
 from typing_extensions import Annotated
 
+from llama_stack.distribution.distribution import builtin_automatically_routed_apis
+
 from llama_stack.providers.utils.telemetry.tracing import (
     end_trace,
     setup_logger,
@@ -43,20 +42,6 @@ from llama_stack.distribution.resolver import resolve_impls_with_routing
 from .endpoints import get_all_api_endpoints
 
 
-def is_async_iterator_type(typ):
-    if hasattr(typ, "__origin__"):
-        origin = typ.__origin__
-        if isinstance(origin, type):
-            return issubclass(
-                origin,
-                (AsyncIterator, AsyncGenerator, AsyncIteratorABC, AsyncGeneratorABC),
-            )
-        return False
-    return isinstance(
-        typ, (AsyncIterator, AsyncGenerator, AsyncIteratorABC, AsyncGeneratorABC)
-    )
-
-
 def create_sse_event(data: Any) -> str:
     if isinstance(data, BaseModel):
         data = data.json()
@@ -169,11 +154,20 @@ async def passthrough(
         await end_trace(SpanStatus.OK if not erred else SpanStatus.ERROR)
 
 
-def handle_sigint(*args, **kwargs):
+def handle_sigint(app, *args, **kwargs):
     print("SIGINT or CTRL-C detected. Exiting gracefully...")
+
+    async def run_shutdown():
+        for impl in app.__llama_stack_impls__.values():
+            print(f"Shutting down {impl}")
+            await impl.shutdown()
+
+    asyncio.run(run_shutdown())
+
     loop = asyncio.get_event_loop()
     for task in asyncio.all_tasks(loop):
         task.cancel()
+
     loop.stop()
 
 
@@ -181,7 +175,10 @@ def handle_sigint(*args, **kwargs):
 async def lifespan(app: FastAPI):
     print("Starting up")
     yield
+
     print("Shutting down")
+    for impl in app.__llama_stack_impls__.values():
+        await impl.shutdown()
 
 
 def create_dynamic_passthrough(
@@ -193,24 +190,18 @@ def create_dynamic_passthrough(
     return endpoint
 
 
-def create_dynamic_typed_route(func: Any, method: str):
-    hints = get_type_hints(func)
-    response_model = hints.get("return")
+def is_streaming_request(func_name: str, request: Request, **kwargs):
+    # TODO: pass the api method and punt it to the Protocol definition directly
+    return kwargs.get("stream", False)
 
-    # NOTE: I think it is better to just add a method within each Api
-    # "Protocol" / adapter-impl to tell what sort of a response this request
-    # is going to produce. /chat_completion can produce a streaming or
-    # non-streaming response depending on if request.stream is True / False.
-    is_streaming = is_async_iterator_type(response_model)
 
-    if is_streaming:
+async def maybe_await(value):
+    if inspect.iscoroutine(value):
+        return await value
+    return value
 
-        async def endpoint(request: Request, **kwargs):
-            await start_trace(func.__name__)
 
-            set_request_provider_data(request.headers)
-
-            async def sse_generator(event_gen):
+async def sse_generator(event_gen):
     try:
         async for item in event_gen:
             yield create_sse_event(item)
@@ -230,23 +221,23 @@ def create_dynamic_typed_route(func: Any, method: str):
     finally:
         await end_trace()
 
-            return StreamingResponse(
-                sse_generator(func(**kwargs)), media_type="text/event-stream"
-            )
 
-    else:
+def create_dynamic_typed_route(func: Any, method: str):
 
     async def endpoint(request: Request, **kwargs):
         await start_trace(func.__name__)
 
         set_request_provider_data(request.headers)
 
+        is_streaming = is_streaming_request(func.__name__, request, **kwargs)
         try:
-                return (
-                    await func(**kwargs)
-                    if asyncio.iscoroutinefunction(func)
-                    else func(**kwargs)
+            if is_streaming:
+                return StreamingResponse(
+                    sse_generator(func(**kwargs)), media_type="text/event-stream"
                 )
+            else:
+                value = func(**kwargs)
+                return await maybe_await(value)
         except Exception as e:
             traceback.print_exception(e)
             raise translate_exception(e) from e
@@ -285,29 +276,28 @@ def main(
 
     app = FastAPI()
 
-    impls, specs = asyncio.run(resolve_impls_with_routing(config))
+    impls = asyncio.run(resolve_impls_with_routing(config))
     if Api.telemetry in impls:
         setup_logger(impls[Api.telemetry])
 
     all_endpoints = get_all_api_endpoints()
 
-    if config.apis_to_serve:
-        apis_to_serve = set(config.apis_to_serve)
+    if config.apis:
+        apis_to_serve = set(config.apis)
     else:
         apis_to_serve = set(impls.keys())
 
-    apis_to_serve.add(Api.inspect)
+    for inf in builtin_automatically_routed_apis():
+        apis_to_serve.add(inf.routing_table_api.value)
+
+    apis_to_serve.add("inspect")
     for api_str in apis_to_serve:
         api = Api(api_str)
 
         endpoints = all_endpoints[api]
         impl = impls[api]
 
-        provider_spec = specs[api]
-        if (
-            isinstance(provider_spec, RemoteProviderSpec)
-            and provider_spec.adapter is None
-        ):
+        if is_passthrough(impl.__provider_spec__):
             for endpoint in endpoints:
                 url = impl.__provider_config__.url.rstrip("/") + endpoint.route
                 getattr(app, endpoint.method)(endpoint.route)(
@@ -337,7 +327,9 @@ def main(
     print("")
     app.exception_handler(RequestValidationError)(global_exception_handler)
     app.exception_handler(Exception)(global_exception_handler)
-    signal.signal(signal.SIGINT, handle_sigint)
+    signal.signal(signal.SIGINT, functools.partial(handle_sigint, app))
+
+    app.__llama_stack_impls__ = impls
 
     import uvicorn
 

llama_stack/distribution/templates/docker/llamastack-local-cpu/run.yaml

@@ -1,8 +1,9 @@
-built_at: '2024-09-30T09:04:30.533391'
+version: '2'
+built_at: '2024-10-08T17:42:07.505267'
 image_name: local-cpu
 docker_image: local-cpu
 conda_env: null
-apis_to_serve:
+apis:
 - agents
 - inference
 - models
@@ -10,40 +11,32 @@ apis_to_serve:
 - safety
 - shields
 - memory_banks
-api_providers:
+providers:
   inference:
-    providers:
-    - remote::ollama
+  - provider_id: remote::ollama
+    provider_type: remote::ollama
+    config:
+      host: localhost
+      port: 6000
   safety:
-    providers:
-    - meta-reference
+  - provider_id: meta-reference
+    provider_type: meta-reference
+    config:
+      llama_guard_shield: null
+      prompt_guard_shield: null
+  memory:
+  - provider_id: meta-reference
+    provider_type: meta-reference
+    config: {}
   agents:
+  - provider_id: meta-reference
     provider_type: meta-reference
     config:
       persistence_store:
         namespace: null
         type: sqlite
         db_path: ~/.llama/runtime/kvstore.db
-  memory:
-    providers:
-    - meta-reference
   telemetry:
+  - provider_id: meta-reference
     provider_type: meta-reference
     config: {}
-routing_table:
-  inference:
-  - provider_type: remote::ollama
-    config:
-      host: localhost
-      port: 6000
-    routing_key: Llama3.1-8B-Instruct
-  safety:
-  - provider_type: meta-reference
-    config:
-      llama_guard_shield: null
-      prompt_guard_shield: null
-    routing_key: ["llama_guard", "code_scanner_guard", "injection_shield", "jailbreak_shield"]
-  memory:
-  - provider_type: meta-reference
-    config: {}
-    routing_key: vector

llama_stack/distribution/templates/docker/llamastack-local-gpu/run.yaml

@@ -1,8 +1,9 @@
-built_at: '2024-09-30T09:00:56.693751'
+version: '2'
+built_at: '2024-10-08T17:42:33.690666'
 image_name: local-gpu
 docker_image: local-gpu
 conda_env: null
-apis_to_serve:
+apis:
 - memory
 - inference
 - agents
@@ -10,43 +11,35 @@ apis_to_serve:
 - safety
 - models
 - memory_banks
-api_providers:
+providers:
   inference:
-    providers:
-    - meta-reference
-  safety:
-    providers:
-    - meta-reference
-  agents:
+  - provider_id: meta-reference
     provider_type: meta-reference
-    config:
-      persistence_store:
-        namespace: null
-        type: sqlite
-        db_path: ~/.llama/runtime/kvstore.db
-  memory:
-    providers:
-    - meta-reference
-  telemetry:
-    provider_type: meta-reference
-    config: {}
-routing_table:
-  inference:
-  - provider_type: meta-reference
     config:
       model: Llama3.1-8B-Instruct
       quantization: null
       torch_seed: null
       max_seq_len: 4096
       max_batch_size: 1
-    routing_key: Llama3.1-8B-Instruct
   safety:
-  - provider_type: meta-reference
+  - provider_id: meta-reference
+    provider_type: meta-reference
     config:
       llama_guard_shield: null
       prompt_guard_shield: null
-    routing_key: ["llama_guard", "code_scanner_guard", "injection_shield", "jailbreak_shield"]
   memory:
-  - provider_type: meta-reference
+  - provider_id: meta-reference
+    provider_type: meta-reference
+    config: {}
+  agents:
+  - provider_id: meta-reference
+    provider_type: meta-reference
+    config:
+      persistence_store:
+        namespace: null
+        type: sqlite
+        db_path: ~/.llama/runtime/kvstore.db
+  telemetry:
+  - provider_id: meta-reference
+    provider_type: meta-reference
     config: {}
-    routing_key: vector

llama_stack/providers/adapters/inference/bedrock/bedrock.py

@@ -13,7 +13,7 @@ from botocore.config import Config
 from llama_models.llama3.api.chat_format import ChatFormat
 from llama_models.llama3.api.tokenizer import Tokenizer
 
-from llama_stack.providers.utils.inference.routable import RoutableProviderForModels
+from llama_stack.providers.utils.inference.model_registry import ModelRegistryHelper
 
 from llama_stack.apis.inference import *  # noqa: F403
 from llama_stack.providers.adapters.inference.bedrock.config import BedrockConfig
@@ -26,57 +26,16 @@ BEDROCK_SUPPORTED_MODELS = {
 }
 
 
-class BedrockInferenceAdapter(Inference, RoutableProviderForModels):
-
-    @staticmethod
-    def _create_bedrock_client(config: BedrockConfig) -> BaseClient:
-        retries_config = {
-            k: v
-            for k, v in dict(
-                total_max_attempts=config.total_max_attempts,
-                mode=config.retry_mode,
-            ).items()
-            if v is not None
-        }
-
-        config_args = {
-            k: v
-            for k, v in dict(
-                region_name=config.region_name,
-                retries=retries_config if retries_config else None,
-                connect_timeout=config.connect_timeout,
-                read_timeout=config.read_timeout,
-            ).items()
-            if v is not None
-        }
-
-        boto3_config = Config(**config_args)
-
-        session_args = {
-            k: v
-            for k, v in dict(
-                aws_access_key_id=config.aws_access_key_id,
-                aws_secret_access_key=config.aws_secret_access_key,
-                aws_session_token=config.aws_session_token,
-                region_name=config.region_name,
-                profile_name=config.profile_name,
-            ).items()
-            if v is not None
-        }
-
-        boto3_session = boto3.session.Session(**session_args)
-
-        return boto3_session.client("bedrock-runtime", config=boto3_config)
-
+# NOTE: this is not quite tested after the recent refactors
+class BedrockInferenceAdapter(ModelRegistryHelper, Inference):
     def __init__(self, config: BedrockConfig) -> None:
-        RoutableProviderForModels.__init__(
+        ModelRegistryHelper.__init__(
             self, stack_to_provider_models_map=BEDROCK_SUPPORTED_MODELS
         )
         self._config = config
 
-        self._client = BedrockInferenceAdapter._create_bedrock_client(config)
-        tokenizer = Tokenizer.get_instance()
-        self.formatter = ChatFormat(tokenizer)
+        self._client = _create_bedrock_client(config)
+        self.formatter = ChatFormat(Tokenizer.get_instance())
 
     @property
     def client(self) -> BaseClient:
@@ -88,7 +47,7 @@ class BedrockInferenceAdapter(Inference, RoutableProviderForModels):
     async def shutdown(self) -> None:
         self.client.close()
 
-    async def completion(
+    def completion(
         self,
         model: str,
         content: InterleavedTextMedia,
@@ -258,7 +217,7 @@ class BedrockInferenceAdapter(Inference, RoutableProviderForModels):
 
     @staticmethod
     def _tool_parameters_to_input_schema(
-        tool_parameters: Optional[Dict[str, ToolParamDefinition]]
+        tool_parameters: Optional[Dict[str, ToolParamDefinition]],
     ) -> Dict:
         input_schema = {"type": "object"}
         if not tool_parameters:
@@ -324,7 +283,7 @@ class BedrockInferenceAdapter(Inference, RoutableProviderForModels):
             )
         return tool_config
 
-    async def chat_completion(
+    def chat_completion(
         self,
         model: str,
         messages: List[Message],
@@ -443,3 +402,50 @@ class BedrockInferenceAdapter(Inference, RoutableProviderForModels):
                 else:
                     # Ignored
                     pass
+
+    async def embeddings(
+        self,
+        model: str,
+        contents: List[InterleavedTextMedia],
+    ) -> EmbeddingsResponse:
+        raise NotImplementedError()
+
+
+def _create_bedrock_client(config: BedrockConfig) -> BaseClient:
+    retries_config = {
+        k: v
+        for k, v in dict(
+            total_max_attempts=config.total_max_attempts,
+            mode=config.retry_mode,
+        ).items()
+        if v is not None
+    }
+
+    config_args = {
+        k: v
+        for k, v in dict(
+            region_name=config.region_name,
+            retries=retries_config if retries_config else None,
+            connect_timeout=config.connect_timeout,
+            read_timeout=config.read_timeout,
+        ).items()
+        if v is not None
+    }
+
+    boto3_config = Config(**config_args)
+
+    session_args = {
+        k: v
+        for k, v in dict(
+            aws_access_key_id=config.aws_access_key_id,
+            aws_secret_access_key=config.aws_secret_access_key,
+            aws_session_token=config.aws_session_token,
+            region_name=config.region_name,
+            profile_name=config.profile_name,
+        ).items()
+        if v is not None
+    }
+
+    boto3_session = boto3.session.Session(**session_args)
+
+    return boto3_session.client("bedrock-runtime", config=boto3_config)

llama_stack/providers/adapters/inference/databricks/databricks.py

@@ -6,39 +6,41 @@
 
 from typing import AsyncGenerator
 
-from openai import OpenAI
-
 from llama_models.llama3.api.chat_format import ChatFormat
 
-from llama_models.llama3.api.datatypes import Message, StopReason
+from llama_models.llama3.api.datatypes import Message
 from llama_models.llama3.api.tokenizer import Tokenizer
-from llama_models.sku_list import resolve_model
+
+from openai import OpenAI
 
 from llama_stack.apis.inference import *  # noqa: F403
-from llama_stack.providers.utils.inference.augment_messages import (
-    augment_messages_for_tools,
+
+from llama_stack.providers.utils.inference.model_registry import ModelRegistryHelper
+from llama_stack.providers.utils.inference.openai_compat import (
+    get_sampling_options,
+    process_chat_completion_response,
+    process_chat_completion_stream_response,
+)
+from llama_stack.providers.utils.inference.prompt_adapter import (
+    chat_completion_request_to_prompt,
 )
 
 from .config import DatabricksImplConfig
 
+
 DATABRICKS_SUPPORTED_MODELS = {
     "Llama3.1-70B-Instruct": "databricks-meta-llama-3-1-70b-instruct",
     "Llama3.1-405B-Instruct": "databricks-meta-llama-3-1-405b-instruct",
 }
 
 
-class DatabricksInferenceAdapter(Inference):
+class DatabricksInferenceAdapter(ModelRegistryHelper, Inference):
     def __init__(self, config: DatabricksImplConfig) -> None:
-        self.config = config
-        tokenizer = Tokenizer.get_instance()
-        self.formatter = ChatFormat(tokenizer)
-
-    @property
-    def client(self) -> OpenAI:
-        return OpenAI(
-            base_url=self.config.url, 
-            api_key=self.config.api_token
+        ModelRegistryHelper.__init__(
+            self, stack_to_provider_models_map=DATABRICKS_SUPPORTED_MODELS
         )
+        self.config = config
+        self.formatter = ChatFormat(Tokenizer.get_instance())
 
     async def initialize(self) -> None:
         return
@@ -46,47 +48,10 @@ class DatabricksInferenceAdapter(Inference):
     async def shutdown(self) -> None:
         pass
 
-    async def validate_routing_keys(self, routing_keys: list[str]) -> None:
-        # these are the model names the Llama Stack will use to route requests to this provider
-        # perform validation here if necessary
-        pass
-
-    async def completion(self, request: CompletionRequest) -> AsyncGenerator:
+    def completion(self, request: CompletionRequest) -> AsyncGenerator:
         raise NotImplementedError()
 
-    def _messages_to_databricks_messages(self, messages: list[Message]) -> list:
-        databricks_messages = []
-        for message in messages:
-            if message.role == "ipython":
-                role = "tool"
-            else:
-                role = message.role
-            databricks_messages.append({"role": role, "content": message.content})
-
-        return databricks_messages
-
-    def resolve_databricks_model(self, model_name: str) -> str:
-        model = resolve_model(model_name)
-        assert (
-            model is not None
-            and model.descriptor(shorten_default_variant=True)
-            in DATABRICKS_SUPPORTED_MODELS
-        ), f"Unsupported model: {model_name}, use one of the supported models: {','.join(DATABRICKS_SUPPORTED_MODELS.keys())}"
-
-        return DATABRICKS_SUPPORTED_MODELS.get(
-            model.descriptor(shorten_default_variant=True)
-        )
-
-    def get_databricks_chat_options(self, request: ChatCompletionRequest) -> dict:
-        options = {}
-        if request.sampling_params is not None:
-            for attr in {"temperature", "top_p", "top_k", "max_tokens"}:
-                if getattr(request.sampling_params, attr):
-                    options[attr] = getattr(request.sampling_params, attr)
-
-        return options
-
-    async def chat_completion(
+    def chat_completion(
         self,
         model: str,
         messages: List[Message],
@@ -108,150 +73,46 @@ class DatabricksInferenceAdapter(Inference):
             logprobs=logprobs,
         )
 
-        messages = augment_messages_for_tools(request)
-        options = self.get_databricks_chat_options(request)
-        databricks_model = self.resolve_databricks_model(request.model)
-
-        if not request.stream:
-
-            r = self.client.chat.completions.create(
-                model=databricks_model,
-                messages=self._messages_to_databricks_messages(messages),
-                stream=False,
-                **options,
-            )
-
-            stop_reason = None
-            if r.choices[0].finish_reason:
-                if r.choices[0].finish_reason == "stop":
-                    stop_reason = StopReason.end_of_turn
-                elif r.choices[0].finish_reason == "length":
-                    stop_reason = StopReason.out_of_tokens
-
-            completion_message = self.formatter.decode_assistant_message_from_content(
-                r.choices[0].message.content, stop_reason
-            )
-            yield ChatCompletionResponse(
-                completion_message=completion_message,
-                logprobs=None,
-            )
+        client = OpenAI(base_url=self.config.url, api_key=self.config.api_token)
+        if stream:
+            return self._stream_chat_completion(request, client)
         else:
-            yield ChatCompletionResponseStreamChunk(
-                event=ChatCompletionResponseEvent(
-                    event_type=ChatCompletionResponseEventType.start,
-                    delta="",
-                )
-            )
+            return self._nonstream_chat_completion(request, client)
 
-            buffer = ""
-            ipython = False
-            stop_reason = None
+    async def _nonstream_chat_completion(
+        self, request: ChatCompletionRequest, client: OpenAI
+    ) -> ChatCompletionResponse:
+        params = self._get_params(request)
+        r = client.completions.create(**params)
+        return process_chat_completion_response(request, r, self.formatter)
 
-            for chunk in self.client.chat.completions.create(
-                model=databricks_model,
-                messages=self._messages_to_databricks_messages(messages),
-                stream=True,
-                **options,
+    async def _stream_chat_completion(
+        self, request: ChatCompletionRequest, client: OpenAI
+    ) -> AsyncGenerator:
+        params = self._get_params(request)
+
+        async def _to_async_generator():
+            s = client.completions.create(**params)
+            for chunk in s:
+                yield chunk
+
+        stream = _to_async_generator()
+        async for chunk in process_chat_completion_stream_response(
+            request, stream, self.formatter
         ):
-                if chunk.choices[0].finish_reason:
-                    if (
-                        stop_reason is None
-                        and chunk.choices[0].finish_reason == "stop"
-                    ):
-                        stop_reason = StopReason.end_of_turn
-                    elif (
-                        stop_reason is None
-                        and chunk.choices[0].finish_reason == "length"
-                    ):
-                        stop_reason = StopReason.out_of_tokens
-                    break
+            yield chunk
 
-                text = chunk.choices[0].delta.content
+    def _get_params(self, request: ChatCompletionRequest) -> dict:
+        return {
+            "model": self.map_to_provider_model(request.model),
+            "prompt": chat_completion_request_to_prompt(request, self.formatter),
+            "stream": request.stream,
+            **get_sampling_options(request),
+        }
 
-                if text is None:
-                    continue
-
-                # check if its a tool call ( aka starts with <|python_tag|> )
-                if not ipython and text.startswith("<|python_tag|>"):
-                    ipython = True
-                    yield ChatCompletionResponseStreamChunk(
-                        event=ChatCompletionResponseEvent(
-                            event_type=ChatCompletionResponseEventType.progress,
-                            delta=ToolCallDelta(
-                                content="",
-                                parse_status=ToolCallParseStatus.started,
-                            ),
-                        )
-                    )
-                    buffer += text
-                    continue
-
-                if ipython:
-                    if text == "<|eot_id|>":
-                        stop_reason = StopReason.end_of_turn
-                        text = ""
-                        continue
-                    elif text == "<|eom_id|>":
-                        stop_reason = StopReason.end_of_message
-                        text = ""
-                        continue
-
-                    buffer += text
-                    delta = ToolCallDelta(
-                        content=text,
-                        parse_status=ToolCallParseStatus.in_progress,
-                    )
-
-                    yield ChatCompletionResponseStreamChunk(
-                        event=ChatCompletionResponseEvent(
-                            event_type=ChatCompletionResponseEventType.progress,
-                            delta=delta,
-                            stop_reason=stop_reason,
-                        )
-                    )
-                else:
-                    buffer += text
-                    yield ChatCompletionResponseStreamChunk(
-                        event=ChatCompletionResponseEvent(
-                            event_type=ChatCompletionResponseEventType.progress,
-                            delta=text,
-                            stop_reason=stop_reason,
-                        )
-                    )
-
-            # parse tool calls and report errors
-            message = self.formatter.decode_assistant_message_from_content(
-                buffer, stop_reason
-            )
-            parsed_tool_calls = len(message.tool_calls) > 0
-            if ipython and not parsed_tool_calls:
-                yield ChatCompletionResponseStreamChunk(
-                    event=ChatCompletionResponseEvent(
-                        event_type=ChatCompletionResponseEventType.progress,
-                        delta=ToolCallDelta(
-                            content="",
-                            parse_status=ToolCallParseStatus.failure,
-                        ),
-                        stop_reason=stop_reason,
-                    )
-                )
-
-            for tool_call in message.tool_calls:
-                yield ChatCompletionResponseStreamChunk(
-                    event=ChatCompletionResponseEvent(
-                        event_type=ChatCompletionResponseEventType.progress,
-                        delta=ToolCallDelta(
-                            content=tool_call,
-                            parse_status=ToolCallParseStatus.success,
-                        ),
-                        stop_reason=stop_reason,
-                    )
-                )
-
-            yield ChatCompletionResponseStreamChunk(
-                event=ChatCompletionResponseEvent(
-                    event_type=ChatCompletionResponseEventType.complete,
-                    delta="",
-                    stop_reason=stop_reason,
-                )
-            )
+    async def embeddings(
+        self,
+        model: str,
+        contents: List[InterleavedTextMedia],
+    ) -> EmbeddingsResponse:
+        raise NotImplementedError()

llama_stack/providers/adapters/inference/fireworks/fireworks.py

@@ -10,14 +10,19 @@ from fireworks.client import Fireworks
 
 from llama_models.llama3.api.chat_format import ChatFormat
 
-from llama_models.llama3.api.datatypes import Message, StopReason
+from llama_models.llama3.api.datatypes import Message
 from llama_models.llama3.api.tokenizer import Tokenizer
 
-from llama_stack.providers.utils.inference.routable import RoutableProviderForModels
-
 from llama_stack.apis.inference import *  # noqa: F403
-from llama_stack.providers.utils.inference.augment_messages import (
-    augment_messages_for_tools,
+
+from llama_stack.providers.utils.inference.model_registry import ModelRegistryHelper
+from llama_stack.providers.utils.inference.openai_compat import (
+    get_sampling_options,
+    process_chat_completion_response,
+    process_chat_completion_stream_response,
+)
+from llama_stack.providers.utils.inference.prompt_adapter import (
+    chat_completion_request_to_prompt,
 )
 
 from .config import FireworksImplConfig
@@ -27,21 +32,18 @@ FIREWORKS_SUPPORTED_MODELS = {
     "Llama3.1-8B-Instruct": "fireworks/llama-v3p1-8b-instruct",
     "Llama3.1-70B-Instruct": "fireworks/llama-v3p1-70b-instruct",
     "Llama3.1-405B-Instruct": "fireworks/llama-v3p1-405b-instruct",
+    "Llama3.2-1B-Instruct": "fireworks/llama-v3p2-1b-instruct",
+    "Llama3.2-3B-Instruct": "fireworks/llama-v3p2-3b-instruct",
 }
 
 
-class FireworksInferenceAdapter(Inference, RoutableProviderForModels):
+class FireworksInferenceAdapter(ModelRegistryHelper, Inference):
     def __init__(self, config: FireworksImplConfig) -> None:
-        RoutableProviderForModels.__init__(
+        ModelRegistryHelper.__init__(
             self, stack_to_provider_models_map=FIREWORKS_SUPPORTED_MODELS
         )
         self.config = config
-        tokenizer = Tokenizer.get_instance()
-        self.formatter = ChatFormat(tokenizer)
-
-    @property
-    def client(self) -> Fireworks:
-        return Fireworks(api_key=self.config.api_key)
+        self.formatter = ChatFormat(Tokenizer.get_instance())
 
     async def initialize(self) -> None:
         return
@@ -49,7 +51,7 @@ class FireworksInferenceAdapter(Inference, RoutableProviderForModels):
     async def shutdown(self) -> None:
         pass
 
-    async def completion(
+    def completion(
         self,
         model: str,
         content: InterleavedTextMedia,
@@ -59,27 +61,7 @@ class FireworksInferenceAdapter(Inference, RoutableProviderForModels):
     ) -> AsyncGenerator:
         raise NotImplementedError()
 
-    def _messages_to_fireworks_messages(self, messages: list[Message]) -> list:
-        fireworks_messages = []
-        for message in messages:
-            if message.role == "ipython":
-                role = "tool"
-            else:
-                role = message.role
-            fireworks_messages.append({"role": role, "content": message.content})
-
-        return fireworks_messages
-
-    def get_fireworks_chat_options(self, request: ChatCompletionRequest) -> dict:
-        options = {}
-        if request.sampling_params is not None:
-            for attr in {"temperature", "top_p", "top_k", "max_tokens"}:
-                if getattr(request.sampling_params, attr):
-                    options[attr] = getattr(request.sampling_params, attr)
-
-        return options
-
-    async def chat_completion(
+    def chat_completion(
         self,
         model: str,
         messages: List[Message],
@@ -101,147 +83,48 @@ class FireworksInferenceAdapter(Inference, RoutableProviderForModels):
             logprobs=logprobs,
         )
 
-        messages = augment_messages_for_tools(request)
-
-        # accumulate sampling params and other options to pass to fireworks
-        options = self.get_fireworks_chat_options(request)
-        fireworks_model = self.map_to_provider_model(request.model)
-
-        if not request.stream:
-            r = await self.client.chat.completions.acreate(
-                model=fireworks_model,
-                messages=self._messages_to_fireworks_messages(messages),
-                stream=False,
-                **options,
-            )
-            stop_reason = None
-            if r.choices[0].finish_reason:
-                if r.choices[0].finish_reason == "stop":
-                    stop_reason = StopReason.end_of_turn
-                elif r.choices[0].finish_reason == "length":
-                    stop_reason = StopReason.out_of_tokens
-
-            completion_message = self.formatter.decode_assistant_message_from_content(
-                r.choices[0].message.content, stop_reason
-            )
-
-            yield ChatCompletionResponse(
-                completion_message=completion_message,
-                logprobs=None,
-            )
+        client = Fireworks(api_key=self.config.api_key)
+        if stream:
+            return self._stream_chat_completion(request, client)
         else:
-            yield ChatCompletionResponseStreamChunk(
-                event=ChatCompletionResponseEvent(
-                    event_type=ChatCompletionResponseEventType.start,
-                    delta="",
-                )
-            )
+            return self._nonstream_chat_completion(request, client)
 
-            buffer = ""
-            ipython = False
-            stop_reason = None
+    async def _nonstream_chat_completion(
+        self, request: ChatCompletionRequest, client: Fireworks
+    ) -> ChatCompletionResponse:
+        params = self._get_params(request)
+        r = await client.completion.acreate(**params)
+        return process_chat_completion_response(request, r, self.formatter)
 
-            async for chunk in self.client.chat.completions.acreate(
-                model=fireworks_model,
-                messages=self._messages_to_fireworks_messages(messages),
-                stream=True,
+    async def _stream_chat_completion(
+        self, request: ChatCompletionRequest, client: Fireworks
+    ) -> AsyncGenerator:
+        params = self._get_params(request)
+
+        stream = client.completion.acreate(**params)
+        async for chunk in process_chat_completion_stream_response(
+            request, stream, self.formatter
+        ):
+            yield chunk
+
+    def _get_params(self, request: ChatCompletionRequest) -> dict:
+        prompt = chat_completion_request_to_prompt(request, self.formatter)
+        # Fireworks always prepends with BOS
+        if prompt.startswith("<|begin_of_text|>"):
+            prompt = prompt[len("<|begin_of_text|>") :]
+
+        options = get_sampling_options(request)
+        options.setdefault("max_tokens", 512)
+        return {
+            "model": self.map_to_provider_model(request.model),
+            "prompt": prompt,
+            "stream": request.stream,
             **options,
-            ):
-                if chunk.choices[0].finish_reason:
-                    if stop_reason is None and chunk.choices[0].finish_reason == "stop":
-                        stop_reason = StopReason.end_of_turn
-                    elif (
-                        stop_reason is None
-                        and chunk.choices[0].finish_reason == "length"
-                    ):
-                        stop_reason = StopReason.out_of_tokens
-                    break
+        }
 
-                text = chunk.choices[0].delta.content
-                if text is None:
-                    continue
-
-                # check if its a tool call ( aka starts with <|python_tag|> )
-                if not ipython and text.startswith("<|python_tag|>"):
-                    ipython = True
-                    yield ChatCompletionResponseStreamChunk(
-                        event=ChatCompletionResponseEvent(
-                            event_type=ChatCompletionResponseEventType.progress,
-                            delta=ToolCallDelta(
-                                content="",
-                                parse_status=ToolCallParseStatus.started,
-                            ),
-                        )
-                    )
-                    buffer += text
-                    continue
-
-                if ipython:
-                    if text == "<|eot_id|>":
-                        stop_reason = StopReason.end_of_turn
-                        text = ""
-                        continue
-                    elif text == "<|eom_id|>":
-                        stop_reason = StopReason.end_of_message
-                        text = ""
-                        continue
-
-                    buffer += text
-                    delta = ToolCallDelta(
-                        content=text,
-                        parse_status=ToolCallParseStatus.in_progress,
-                    )
-
-                    yield ChatCompletionResponseStreamChunk(
-                        event=ChatCompletionResponseEvent(
-                            event_type=ChatCompletionResponseEventType.progress,
-                            delta=delta,
-                            stop_reason=stop_reason,
-                        )
-                    )
-                else:
-                    buffer += text
-                    yield ChatCompletionResponseStreamChunk(
-                        event=ChatCompletionResponseEvent(
-                            event_type=ChatCompletionResponseEventType.progress,
-                            delta=text,
-                            stop_reason=stop_reason,
-                        )
-                    )
-
-            # parse tool calls and report errors
-            message = self.formatter.decode_assistant_message_from_content(
-                buffer, stop_reason
-            )
-            parsed_tool_calls = len(message.tool_calls) > 0
-            if ipython and not parsed_tool_calls:
-                yield ChatCompletionResponseStreamChunk(
-                    event=ChatCompletionResponseEvent(
-                        event_type=ChatCompletionResponseEventType.progress,
-                        delta=ToolCallDelta(
-                            content="",
-                            parse_status=ToolCallParseStatus.failure,
-                        ),
-                        stop_reason=stop_reason,
-                    )
-                )
-
-            for tool_call in message.tool_calls:
-                yield ChatCompletionResponseStreamChunk(
-                    event=ChatCompletionResponseEvent(
-                        event_type=ChatCompletionResponseEventType.progress,
-                        delta=ToolCallDelta(
-                            content=tool_call,
-                            parse_status=ToolCallParseStatus.success,
-                        ),
-                        stop_reason=stop_reason,
-                    )
-                )
-
-            yield ChatCompletionResponseStreamChunk(
-                event=ChatCompletionResponseEvent(
-                    event_type=ChatCompletionResponseEventType.complete,
-                    delta="",
-                    stop_reason=stop_reason,
-                )
-            )
+    async def embeddings(
+        self,
+        model: str,
+        contents: List[InterleavedTextMedia],
+    ) -> EmbeddingsResponse:
+        raise NotImplementedError()

llama_stack/providers/adapters/inference/ollama/ollama.py

@@ -9,35 +9,38 @@ from typing import AsyncGenerator
 import httpx
 
 from llama_models.llama3.api.chat_format import ChatFormat
-from llama_models.llama3.api.datatypes import Message, StopReason
+from llama_models.llama3.api.datatypes import Message
 from llama_models.llama3.api.tokenizer import Tokenizer
 
 from ollama import AsyncClient
 
 from llama_stack.apis.inference import *  # noqa: F403
-from llama_stack.providers.utils.inference.augment_messages import (
-    augment_messages_for_tools,
-)
-from llama_stack.providers.utils.inference.routable import RoutableProviderForModels
+from llama_stack.providers.datatypes import ModelsProtocolPrivate
 
-# TODO: Eventually this will move to the llama cli model list command
-# mapping of Model SKUs to ollama models
-OLLAMA_SUPPORTED_SKUS = {
+from llama_stack.providers.utils.inference.openai_compat import (
+    get_sampling_options,
+    OpenAICompatCompletionChoice,
+    OpenAICompatCompletionResponse,
+    process_chat_completion_response,
+    process_chat_completion_stream_response,
+)
+from llama_stack.providers.utils.inference.prompt_adapter import (
+    chat_completion_request_to_prompt,
+)
+
+OLLAMA_SUPPORTED_MODELS = {
     "Llama3.1-8B-Instruct": "llama3.1:8b-instruct-fp16",
     "Llama3.1-70B-Instruct": "llama3.1:70b-instruct-fp16",
     "Llama3.2-1B-Instruct": "llama3.2:1b-instruct-fp16",
     "Llama3.2-3B-Instruct": "llama3.2:3b-instruct-fp16",
+    "Llama-Guard-3-8B": "xe/llamaguard3:latest",
 }
 
 
-class OllamaInferenceAdapter(Inference, RoutableProviderForModels):
+class OllamaInferenceAdapter(Inference, ModelsProtocolPrivate):
     def __init__(self, url: str) -> None:
-        RoutableProviderForModels.__init__(
-            self, stack_to_provider_models_map=OLLAMA_SUPPORTED_SKUS
-        )
         self.url = url
-        tokenizer = Tokenizer.get_instance()
-        self.formatter = ChatFormat(tokenizer)
+        self.formatter = ChatFormat(Tokenizer.get_instance())
 
     @property
     def client(self) -> AsyncClient:
@@ -55,7 +58,33 @@ class OllamaInferenceAdapter(Inference, RoutableProviderForModels):
     async def shutdown(self) -> None:
         pass
 
-    async def completion(
+    async def register_model(self, model: ModelDef) -> None:
+        raise ValueError("Dynamic model registration is not supported")
+
+    async def list_models(self) -> List[ModelDef]:
+        ollama_to_llama = {v: k for k, v in OLLAMA_SUPPORTED_MODELS.items()}
+
+        ret = []
+        res = await self.client.ps()
+        for r in res["models"]:
+            if r["model"] not in ollama_to_llama:
+                print(f"Ollama is running a model unknown to Llama Stack: {r['model']}")
+                continue
+
+            llama_model = ollama_to_llama[r["model"]]
+            ret.append(
+                ModelDef(
+                    identifier=llama_model,
+                    llama_model=llama_model,
+                    metadata={
+                        "ollama_model": r["model"],
+                    },
+                )
+            )
+
+        return ret
+
+    def completion(
         self,
         model: str,
         content: InterleavedTextMedia,
@@ -65,32 +94,7 @@ class OllamaInferenceAdapter(Inference, RoutableProviderForModels):
     ) -> AsyncGenerator:
         raise NotImplementedError()
 
-    def _messages_to_ollama_messages(self, messages: list[Message]) -> list:
-        ollama_messages = []
-        for message in messages:
-            if message.role == "ipython":
-                role = "tool"
-            else:
-                role = message.role
-            ollama_messages.append({"role": role, "content": message.content})
-
-        return ollama_messages
-
-    def get_ollama_chat_options(self, request: ChatCompletionRequest) -> dict:
-        options = {}
-        if request.sampling_params is not None:
-            for attr in {"temperature", "top_p", "top_k", "max_tokens"}:
-                if getattr(request.sampling_params, attr):
-                    options[attr] = getattr(request.sampling_params, attr)
-            if (
-                request.sampling_params.repetition_penalty is not None
-                and request.sampling_params.repetition_penalty != 1.0
-            ):
-                options["repeat_penalty"] = request.sampling_params.repetition_penalty
-
-        return options
-
-    async def chat_completion(
+    def chat_completion(
         self,
         model: str,
         messages: List[Message],
@@ -111,156 +115,61 @@ class OllamaInferenceAdapter(Inference, RoutableProviderForModels):
             stream=stream,
             logprobs=logprobs,
         )
-
-        messages = augment_messages_for_tools(request)
-        # accumulate sampling params and other options to pass to ollama
-        options = self.get_ollama_chat_options(request)
-        ollama_model = self.map_to_provider_model(request.model)
-
-        res = await self.client.ps()
-        need_model_pull = True
-        for r in res["models"]:
-            if ollama_model == r["model"]:
-                need_model_pull = False
-                break
-
-        if need_model_pull:
-            print(f"Pulling model: {ollama_model}")
-            status = await self.client.pull(ollama_model)
-            assert (
-                status["status"] == "success"
-            ), f"Failed to pull model {self.model} in ollama"
-
-        if not request.stream:
-            r = await self.client.chat(
-                model=ollama_model,
-                messages=self._messages_to_ollama_messages(messages),
-                stream=False,
-                options=options,
-            )
-            stop_reason = None
-            if r["done"]:
-                if r["done_reason"] == "stop":
-                    stop_reason = StopReason.end_of_turn
-                elif r["done_reason"] == "length":
-                    stop_reason = StopReason.out_of_tokens
-
-            completion_message = self.formatter.decode_assistant_message_from_content(
-                r["message"]["content"], stop_reason
-            )
-            yield ChatCompletionResponse(
-                completion_message=completion_message,
-                logprobs=None,
-            )
+        if stream:
+            return self._stream_chat_completion(request)
         else:
-            yield ChatCompletionResponseStreamChunk(
-                event=ChatCompletionResponseEvent(
-                    event_type=ChatCompletionResponseEventType.start,
-                    delta="",
+            return self._nonstream_chat_completion(request)
+
+    def _get_params(self, request: ChatCompletionRequest) -> dict:
+        return {
+            "model": OLLAMA_SUPPORTED_MODELS[request.model],
+            "prompt": chat_completion_request_to_prompt(request, self.formatter),
+            "options": get_sampling_options(request),
+            "raw": True,
+            "stream": request.stream,
+        }
+
+    async def _nonstream_chat_completion(
+        self, request: ChatCompletionRequest
+    ) -> ChatCompletionResponse:
+        params = 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,
+            text=r["response"],
         )
+        response = OpenAICompatCompletionResponse(
+            choices=[choice],
         )
-            stream = await self.client.chat(
-                model=ollama_model,
-                messages=self._messages_to_ollama_messages(messages),
-                stream=True,
-                options=options,
+        return process_chat_completion_response(request, response, self.formatter)
+
+    async def _stream_chat_completion(
+        self, request: ChatCompletionRequest
+    ) -> AsyncGenerator:
+        params = self._get_params(request)
+
+        async def _generate_and_convert_to_openai_compat():
+            s = await self.client.generate(**params)
+            async for chunk in s:
+                choice = OpenAICompatCompletionChoice(
+                    finish_reason=chunk["done_reason"] if chunk["done"] else None,
+                    text=chunk["response"],
+                )
+                yield OpenAICompatCompletionResponse(
+                    choices=[choice],
                 )
 
-            buffer = ""
-            ipython = False
-            stop_reason = None
+        stream = _generate_and_convert_to_openai_compat()
+        async for chunk in process_chat_completion_stream_response(
+            request, stream, self.formatter
+        ):
+            yield chunk
 
-            async for chunk in stream:
-                if chunk["done"]:
-                    if stop_reason is None and chunk["done_reason"] == "stop":
-                        stop_reason = StopReason.end_of_turn
-                    elif stop_reason is None and chunk["done_reason"] == "length":
-                        stop_reason = StopReason.out_of_tokens
-                    break
-
-                text = chunk["message"]["content"]
-
-                # check if its a tool call ( aka starts with <|python_tag|> )
-                if not ipython and text.startswith("<|python_tag|>"):
-                    ipython = True
-                    yield ChatCompletionResponseStreamChunk(
-                        event=ChatCompletionResponseEvent(
-                            event_type=ChatCompletionResponseEventType.progress,
-                            delta=ToolCallDelta(
-                                content="",
-                                parse_status=ToolCallParseStatus.started,
-                            ),
-                        )
-                    )
-                    buffer += text
-                    continue
-
-                if ipython:
-                    if text == "<|eot_id|>":
-                        stop_reason = StopReason.end_of_turn
-                        text = ""
-                        continue
-                    elif text == "<|eom_id|>":
-                        stop_reason = StopReason.end_of_message
-                        text = ""
-                        continue
-
-                    buffer += text
-                    delta = ToolCallDelta(
-                        content=text,
-                        parse_status=ToolCallParseStatus.in_progress,
-                    )
-
-                    yield ChatCompletionResponseStreamChunk(
-                        event=ChatCompletionResponseEvent(
-                            event_type=ChatCompletionResponseEventType.progress,
-                            delta=delta,
-                            stop_reason=stop_reason,
-                        )
-                    )
-                else:
-                    buffer += text
-                    yield ChatCompletionResponseStreamChunk(
-                        event=ChatCompletionResponseEvent(
-                            event_type=ChatCompletionResponseEventType.progress,
-                            delta=text,
-                            stop_reason=stop_reason,
-                        )
-                    )
-
-            # parse tool calls and report errors
-            message = self.formatter.decode_assistant_message_from_content(
-                buffer, stop_reason
-            )
-            parsed_tool_calls = len(message.tool_calls) > 0
-            if ipython and not parsed_tool_calls:
-                yield ChatCompletionResponseStreamChunk(
-                    event=ChatCompletionResponseEvent(
-                        event_type=ChatCompletionResponseEventType.progress,
-                        delta=ToolCallDelta(
-                            content="",
-                            parse_status=ToolCallParseStatus.failure,
-                        ),
-                        stop_reason=stop_reason,
-                    )
-                )
-
-            for tool_call in message.tool_calls:
-                yield ChatCompletionResponseStreamChunk(
-                    event=ChatCompletionResponseEvent(
-                        event_type=ChatCompletionResponseEventType.progress,
-                        delta=ToolCallDelta(
-                            content=tool_call,
-                            parse_status=ToolCallParseStatus.success,
-                        ),
-                        stop_reason=stop_reason,
-                    )
-                )
-
-            yield ChatCompletionResponseStreamChunk(
-                event=ChatCompletionResponseEvent(
-                    event_type=ChatCompletionResponseEventType.complete,
-                    delta="",
-                    stop_reason=stop_reason,
-                )
-            )
+    async def embeddings(
+        self,
+        model: str,
+        contents: List[InterleavedTextMedia],
+    ) -> EmbeddingsResponse:
+        raise NotImplementedError()

llama_stack/providers/adapters/inference/sample/sample.py

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

llama_stack/providers/adapters/inference/tgi/config.py

@@ -34,7 +34,7 @@ class InferenceEndpointImplConfig(BaseModel):
 
 @json_schema_type
 class InferenceAPIImplConfig(BaseModel):
-    model_id: str = Field(
+    huggingface_repo: str = Field(
         description="The model ID of the model on the Hugging Face Hub (e.g. 'meta-llama/Meta-Llama-3.1-70B-Instruct')",
     )
     api_token: Optional[str] = Field(

llama_stack/providers/adapters/inference/tgi/tgi.py

@@ -6,18 +6,27 @@
 
 
 import logging
-from typing import AsyncGenerator
+from typing import AsyncGenerator, List, Optional
 
 from huggingface_hub import AsyncInferenceClient, HfApi
 from llama_models.llama3.api.chat_format import ChatFormat
-from llama_models.llama3.api.datatypes import StopReason
 from llama_models.llama3.api.tokenizer import Tokenizer
-
-from llama_stack.distribution.datatypes import RoutableProvider
+from llama_models.sku_list import all_registered_models
 
 from llama_stack.apis.inference import *  # noqa: F403
-from llama_stack.providers.utils.inference.augment_messages import (
-    augment_messages_for_tools,
+from llama_stack.apis.models import *  # noqa: F403
+
+from llama_stack.providers.datatypes import ModelDef, ModelsProtocolPrivate
+
+from llama_stack.providers.utils.inference.openai_compat import (
+    get_sampling_options,
+    OpenAICompatCompletionChoice,
+    OpenAICompatCompletionResponse,
+    process_chat_completion_response,
+    process_chat_completion_stream_response,
+)
+from llama_stack.providers.utils.inference.prompt_adapter import (
+    chat_completion_request_to_model_input_info,
 )
 
 from .config import InferenceAPIImplConfig, InferenceEndpointImplConfig, TGIImplConfig
@@ -25,24 +34,39 @@ from .config import InferenceAPIImplConfig, InferenceEndpointImplConfig, TGIImpl
 logger = logging.getLogger(__name__)
 
 
-class _HfAdapter(Inference, RoutableProvider):
+class _HfAdapter(Inference, ModelsProtocolPrivate):
     client: AsyncInferenceClient
     max_tokens: int
     model_id: str
 
     def __init__(self) -> None:
-        self.tokenizer = Tokenizer.get_instance()
-        self.formatter = ChatFormat(self.tokenizer)
+        self.formatter = ChatFormat(Tokenizer.get_instance())
+        self.huggingface_repo_to_llama_model_id = {
+            model.huggingface_repo: model.descriptor()
+            for model in all_registered_models()
+            if model.huggingface_repo
+        }
 
-    async def validate_routing_keys(self, routing_keys: list[str]) -> None:
-        # these are the model names the Llama Stack will use to route requests to this provider
-        # perform validation here if necessary
-        pass
+    async def register_model(self, model: ModelDef) -> None:
+        raise ValueError("Model registration is not supported for HuggingFace models")
+
+    async def list_models(self) -> List[ModelDef]:
+        repo = self.model_id
+        identifier = self.huggingface_repo_to_llama_model_id[repo]
+        return [
+            ModelDef(
+                identifier=identifier,
+                llama_model=identifier,
+                metadata={
+                    "huggingface_repo": repo,
+                },
+            )
+        ]
 
     async def shutdown(self) -> None:
         pass
 
-    async def completion(
+    def completion(
         self,
         model: str,
         content: InterleavedTextMedia,
@@ -52,16 +76,7 @@ class _HfAdapter(Inference, RoutableProvider):
     ) -> AsyncGenerator:
         raise NotImplementedError()
 
-    def get_chat_options(self, request: ChatCompletionRequest) -> dict:
-        options = {}
-        if request.sampling_params is not None:
-            for attr in {"temperature", "top_p", "top_k", "max_tokens"}:
-                if getattr(request.sampling_params, attr):
-                    options[attr] = getattr(request.sampling_params, attr)
-
-        return options
-
-    async def chat_completion(
+    def chat_completion(
         self,
         model: str,
         messages: List[Message],
@@ -83,146 +98,71 @@ class _HfAdapter(Inference, RoutableProvider):
             logprobs=logprobs,
         )
 
-        messages = augment_messages_for_tools(request)
-        model_input = self.formatter.encode_dialog_prompt(messages)
-        prompt = self.tokenizer.decode(model_input.tokens)
+        if stream:
+            return self._stream_chat_completion(request)
+        else:
+            return self._nonstream_chat_completion(request)
 
-        input_tokens = len(model_input.tokens)
+    async def _nonstream_chat_completion(
+        self, request: ChatCompletionRequest
+    ) -> ChatCompletionResponse:
+        params = self._get_params(request)
+        r = await self.client.text_generation(**params)
+
+        choice = OpenAICompatCompletionChoice(
+            finish_reason=r.details.finish_reason,
+            text="".join(t.text for t in r.details.tokens),
+        )
+        response = OpenAICompatCompletionResponse(
+            choices=[choice],
+        )
+        return process_chat_completion_response(request, response, self.formatter)
+
+    async def _stream_chat_completion(
+        self, request: ChatCompletionRequest
+    ) -> AsyncGenerator:
+        params = self._get_params(request)
+
+        async def _generate_and_convert_to_openai_compat():
+            s = await self.client.text_generation(**params)
+            async for chunk in s:
+                token_result = chunk.token
+
+                choice = OpenAICompatCompletionChoice(text=token_result.text)
+                yield OpenAICompatCompletionResponse(
+                    choices=[choice],
+                )
+
+        stream = _generate_and_convert_to_openai_compat()
+        async for chunk in process_chat_completion_stream_response(
+            request, stream, self.formatter
+        ):
+            yield chunk
+
+    def _get_params(self, request: ChatCompletionRequest) -> dict:
+        prompt, input_tokens = chat_completion_request_to_model_input_info(
+            request, self.formatter
+        )
         max_new_tokens = min(
             request.sampling_params.max_tokens or (self.max_tokens - input_tokens),
             self.max_tokens - input_tokens - 1,
         )
-
-        print(f"Calculated max_new_tokens: {max_new_tokens}")
-
-        options = self.get_chat_options(request)
-        if not request.stream:
-            response = await self.client.text_generation(
+        options = get_sampling_options(request)
+        return dict(
             prompt=prompt,
-                stream=False,
+            stream=request.stream,
             details=True,
             max_new_tokens=max_new_tokens,
             stop_sequences=["<|eom_id|>", "<|eot_id|>"],
             **options,
         )
-            stop_reason = None
-            if response.details.finish_reason:
-                if response.details.finish_reason in ["stop", "eos_token"]:
-                    stop_reason = StopReason.end_of_turn
-                elif response.details.finish_reason == "length":
-                    stop_reason = StopReason.out_of_tokens
 
-            completion_message = self.formatter.decode_assistant_message_from_content(
-                response.generated_text,
-                stop_reason,
-            )
-            yield ChatCompletionResponse(
-                completion_message=completion_message,
-                logprobs=None,
-            )
-
-        else:
-            yield ChatCompletionResponseStreamChunk(
-                event=ChatCompletionResponseEvent(
-                    event_type=ChatCompletionResponseEventType.start,
-                    delta="",
-                )
-            )
-            buffer = ""
-            ipython = False
-            stop_reason = None
-            tokens = []
-
-            async for response in await self.client.text_generation(
-                prompt=prompt,
-                stream=True,
-                details=True,
-                max_new_tokens=max_new_tokens,
-                stop_sequences=["<|eom_id|>", "<|eot_id|>"],
-                **options,
-            ):
-                token_result = response.token
-
-                buffer += token_result.text
-                tokens.append(token_result.id)
-
-                if not ipython and buffer.startswith("<|python_tag|>"):
-                    ipython = True
-                    yield ChatCompletionResponseStreamChunk(
-                        event=ChatCompletionResponseEvent(
-                            event_type=ChatCompletionResponseEventType.progress,
-                            delta=ToolCallDelta(
-                                content="",
-                                parse_status=ToolCallParseStatus.started,
-                            ),
-                        )
-                    )
-                    buffer = buffer[len("<|python_tag|>") :]
-                    continue
-
-                if token_result.text == "<|eot_id|>":
-                    stop_reason = StopReason.end_of_turn
-                    text = ""
-                elif token_result.text == "<|eom_id|>":
-                    stop_reason = StopReason.end_of_message
-                    text = ""
-                else:
-                    text = token_result.text
-
-                if ipython:
-                    delta = ToolCallDelta(
-                        content=text,
-                        parse_status=ToolCallParseStatus.in_progress,
-                    )
-                else:
-                    delta = text
-
-                if stop_reason is None:
-                    yield ChatCompletionResponseStreamChunk(
-                        event=ChatCompletionResponseEvent(
-                            event_type=ChatCompletionResponseEventType.progress,
-                            delta=delta,
-                            stop_reason=stop_reason,
-                        )
-                    )
-
-            if stop_reason is None:
-                stop_reason = StopReason.out_of_tokens
-
-            # parse tool calls and report errors
-            message = self.formatter.decode_assistant_message(tokens, stop_reason)
-            parsed_tool_calls = len(message.tool_calls) > 0
-            if ipython and not parsed_tool_calls:
-                yield ChatCompletionResponseStreamChunk(
-                    event=ChatCompletionResponseEvent(
-                        event_type=ChatCompletionResponseEventType.progress,
-                        delta=ToolCallDelta(
-                            content="",
-                            parse_status=ToolCallParseStatus.failure,
-                        ),
-                        stop_reason=stop_reason,
-                    )
-                )
-
-            for tool_call in message.tool_calls:
-                yield ChatCompletionResponseStreamChunk(
-                    event=ChatCompletionResponseEvent(
-                        event_type=ChatCompletionResponseEventType.progress,
-                        delta=ToolCallDelta(
-                            content=tool_call,
-                            parse_status=ToolCallParseStatus.success,
-                        ),
-                        stop_reason=stop_reason,
-                    )
-                )
-
-            yield ChatCompletionResponseStreamChunk(
-                event=ChatCompletionResponseEvent(
-                    event_type=ChatCompletionResponseEventType.complete,
-                    delta="",
-                    stop_reason=stop_reason,
-                )
-            )
+    async def embeddings(
+        self,
+        model: str,
+        contents: List[InterleavedTextMedia],
+    ) -> EmbeddingsResponse:
+        raise NotImplementedError()
 
 
 class TGIAdapter(_HfAdapter):
@@ -236,7 +176,7 @@ class TGIAdapter(_HfAdapter):
 class InferenceAPIAdapter(_HfAdapter):
     async def initialize(self, config: InferenceAPIImplConfig) -> None:
         self.client = AsyncInferenceClient(
-            model=config.model_id, token=config.api_token
+            model=config.huggingface_repo, token=config.api_token
         )
         endpoint_info = await self.client.get_endpoint_info()
         self.max_tokens = endpoint_info["max_total_tokens"]

llama_stack/providers/adapters/inference/together/together.py

@@ -8,17 +8,22 @@ from typing import AsyncGenerator
 
 from llama_models.llama3.api.chat_format import ChatFormat
 
-from llama_models.llama3.api.datatypes import Message, StopReason
+from llama_models.llama3.api.datatypes import Message
 from llama_models.llama3.api.tokenizer import Tokenizer
 
 from together import Together
 
 from llama_stack.apis.inference import *  # noqa: F403
 from llama_stack.distribution.request_headers import NeedsRequestProviderData
-from llama_stack.providers.utils.inference.augment_messages import (
-    augment_messages_for_tools,
+from llama_stack.providers.utils.inference.model_registry import ModelRegistryHelper
+from llama_stack.providers.utils.inference.openai_compat import (
+    get_sampling_options,
+    process_chat_completion_response,
+    process_chat_completion_stream_response,
+)
+from llama_stack.providers.utils.inference.prompt_adapter import (
+    chat_completion_request_to_prompt,
 )
-from llama_stack.providers.utils.inference.routable import RoutableProviderForModels
 
 from .config import TogetherImplConfig
 
@@ -34,19 +39,14 @@ TOGETHER_SUPPORTED_MODELS = {
 
 
 class TogetherInferenceAdapter(
-    Inference, NeedsRequestProviderData, RoutableProviderForModels
+    ModelRegistryHelper, Inference, NeedsRequestProviderData
 ):
     def __init__(self, config: TogetherImplConfig) -> None:
-        RoutableProviderForModels.__init__(
+        ModelRegistryHelper.__init__(
             self, stack_to_provider_models_map=TOGETHER_SUPPORTED_MODELS
         )
         self.config = config
-        tokenizer = Tokenizer.get_instance()
-        self.formatter = ChatFormat(tokenizer)
-
-    @property
-    def client(self) -> Together:
-        return Together(api_key=self.config.api_key)
+        self.formatter = ChatFormat(Tokenizer.get_instance())
 
     async def initialize(self) -> None:
         return
@@ -64,27 +64,7 @@ class TogetherInferenceAdapter(
     ) -> AsyncGenerator:
         raise NotImplementedError()
 
-    def _messages_to_together_messages(self, messages: list[Message]) -> list:
-        together_messages = []
-        for message in messages:
-            if message.role == "ipython":
-                role = "tool"
-            else:
-                role = message.role
-            together_messages.append({"role": role, "content": message.content})
-
-        return together_messages
-
-    def get_together_chat_options(self, request: ChatCompletionRequest) -> dict:
-        options = {}
-        if request.sampling_params is not None:
-            for attr in {"temperature", "top_p", "top_k", "max_tokens"}:
-                if getattr(request.sampling_params, attr):
-                    options[attr] = getattr(request.sampling_params, attr)
-
-        return options
-
-    async def chat_completion(
+    def chat_completion(
         self,
         model: str,
         messages: List[Message],
@@ -95,7 +75,6 @@ class TogetherInferenceAdapter(
         stream: Optional[bool] = False,
         logprobs: Optional[LogProbConfig] = None,
     ) -> AsyncGenerator:
-
         together_api_key = None
         if self.config.api_key is not None:
             together_api_key = self.config.api_key
@@ -108,7 +87,6 @@ class TogetherInferenceAdapter(
             together_api_key = provider_data.together_api_key
 
         client = Together(api_key=together_api_key)
-        # wrapper request to make it easier to pass around (internal only, not exposed to API)
         request = ChatCompletionRequest(
             model=model,
             messages=messages,
@@ -120,146 +98,46 @@ class TogetherInferenceAdapter(
             logprobs=logprobs,
         )
 
-        # accumulate sampling params and other options to pass to together
-        options = self.get_together_chat_options(request)
-        together_model = self.map_to_provider_model(request.model)
-        messages = augment_messages_for_tools(request)
-
-        if not request.stream:
-            # TODO: might need to add back an async here
-            r = client.chat.completions.create(
-                model=together_model,
-                messages=self._messages_to_together_messages(messages),
-                stream=False,
-                **options,
-            )
-            stop_reason = None
-            if r.choices[0].finish_reason:
-                if (
-                    r.choices[0].finish_reason == "stop"
-                    or r.choices[0].finish_reason == "eos"
-                ):
-                    stop_reason = StopReason.end_of_turn
-                elif r.choices[0].finish_reason == "length":
-                    stop_reason = StopReason.out_of_tokens
-
-            completion_message = self.formatter.decode_assistant_message_from_content(
-                r.choices[0].message.content, stop_reason
-            )
-            yield ChatCompletionResponse(
-                completion_message=completion_message,
-                logprobs=None,
-            )
+        if stream:
+            return self._stream_chat_completion(request, client)
         else:
-            yield ChatCompletionResponseStreamChunk(
-                event=ChatCompletionResponseEvent(
-                    event_type=ChatCompletionResponseEventType.start,
-                    delta="",
-                )
-            )
+            return self._nonstream_chat_completion(request, client)
 
-            buffer = ""
-            ipython = False
-            stop_reason = None
+    async def _nonstream_chat_completion(
+        self, request: ChatCompletionRequest, client: Together
+    ) -> ChatCompletionResponse:
+        params = self._get_params(request)
+        r = client.completions.create(**params)
+        return process_chat_completion_response(request, r, self.formatter)
 
-            for chunk in client.chat.completions.create(
-                model=together_model,
-                messages=self._messages_to_together_messages(messages),
-                stream=True,
-                **options,
+    async def _stream_chat_completion(
+        self, request: ChatCompletionRequest, client: Together
+    ) -> AsyncGenerator:
+        params = self._get_params(request)
+
+        # if we shift to TogetherAsyncClient, we won't need this wrapper
+        async def _to_async_generator():
+            s = client.completions.create(**params)
+            for chunk in s:
+                yield chunk
+
+        stream = _to_async_generator()
+        async for chunk in process_chat_completion_stream_response(
+            request, stream, self.formatter
         ):
-                if finish_reason := chunk.choices[0].finish_reason:
-                    if stop_reason is None and finish_reason in ["stop", "eos"]:
-                        stop_reason = StopReason.end_of_turn
-                    elif stop_reason is None and finish_reason == "length":
-                        stop_reason = StopReason.out_of_tokens
-                    break
+            yield chunk
 
-                text = chunk.choices[0].delta.content
-                if text is None:
-                    continue
+    def _get_params(self, request: ChatCompletionRequest) -> dict:
+        return {
+            "model": self.map_to_provider_model(request.model),
+            "prompt": chat_completion_request_to_prompt(request, self.formatter),
+            "stream": request.stream,
+            **get_sampling_options(request),
+        }
 
-                # check if its a tool call ( aka starts with <|python_tag|> )
-                if not ipython and text.startswith("<|python_tag|>"):
-                    ipython = True
-                    yield ChatCompletionResponseStreamChunk(
-                        event=ChatCompletionResponseEvent(
-                            event_type=ChatCompletionResponseEventType.progress,
-                            delta=ToolCallDelta(
-                                content="",
-                                parse_status=ToolCallParseStatus.started,
-                            ),
-                        )
-                    )
-                    buffer += text
-                    continue
-
-                if ipython:
-                    if text == "<|eot_id|>":
-                        stop_reason = StopReason.end_of_turn
-                        text = ""
-                        continue
-                    elif text == "<|eom_id|>":
-                        stop_reason = StopReason.end_of_message
-                        text = ""
-                        continue
-
-                    buffer += text
-                    delta = ToolCallDelta(
-                        content=text,
-                        parse_status=ToolCallParseStatus.in_progress,
-                    )
-
-                    yield ChatCompletionResponseStreamChunk(
-                        event=ChatCompletionResponseEvent(
-                            event_type=ChatCompletionResponseEventType.progress,
-                            delta=delta,
-                            stop_reason=stop_reason,
-                        )
-                    )
-                else:
-                    buffer += text
-                    yield ChatCompletionResponseStreamChunk(
-                        event=ChatCompletionResponseEvent(
-                            event_type=ChatCompletionResponseEventType.progress,
-                            delta=text,
-                            stop_reason=stop_reason,
-                        )
-                    )
-
-            # parse tool calls and report errors
-            message = self.formatter.decode_assistant_message_from_content(
-                buffer, stop_reason
-            )
-            parsed_tool_calls = len(message.tool_calls) > 0
-            if ipython and not parsed_tool_calls:
-                yield ChatCompletionResponseStreamChunk(
-                    event=ChatCompletionResponseEvent(
-                        event_type=ChatCompletionResponseEventType.progress,
-                        delta=ToolCallDelta(
-                            content="",
-                            parse_status=ToolCallParseStatus.failure,
-                        ),
-                        stop_reason=stop_reason,
-                    )
-                )
-
-            for tool_call in message.tool_calls:
-                yield ChatCompletionResponseStreamChunk(
-                    event=ChatCompletionResponseEvent(
-                        event_type=ChatCompletionResponseEventType.progress,
-                        delta=ToolCallDelta(
-                            content=tool_call,
-                            parse_status=ToolCallParseStatus.success,
-                        ),
-                        stop_reason=stop_reason,
-                    )
-                )
-
-            yield ChatCompletionResponseStreamChunk(
-                event=ChatCompletionResponseEvent(
-                    event_type=ChatCompletionResponseEventType.complete,
-                    delta="",
-                    stop_reason=stop_reason,
-                )
-            )
+    async def embeddings(
+        self,
+        model: str,
+        contents: List[InterleavedTextMedia],
+    ) -> EmbeddingsResponse:
+        raise NotImplementedError()

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

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

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

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

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

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

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

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

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

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

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

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

llama_stack/providers/adapters/safety/bedrock/bedrock.py

@@ -7,14 +7,13 @@
 import json
 import logging
 
-import traceback
 from typing import Any, Dict, List
 
 import boto3
 
 from llama_stack.apis.safety import *  # noqa
 from llama_models.llama3.api.datatypes import *  # noqa: F403
-from llama_stack.distribution.datatypes import RoutableProvider
+from llama_stack.providers.datatypes import ShieldsProtocolPrivate
 
 from .config import BedrockSafetyConfig
 
@@ -22,16 +21,17 @@ from .config import BedrockSafetyConfig
 logger = logging.getLogger(__name__)
 
 
-SUPPORTED_SHIELD_TYPES = [
-    "bedrock_guardrail",
+BEDROCK_SUPPORTED_SHIELDS = [
+    ShieldType.generic_content_shield.value,
 ]
 
 
-class BedrockSafetyAdapter(Safety, RoutableProvider):
+class BedrockSafetyAdapter(Safety, ShieldsProtocolPrivate):
     def __init__(self, config: BedrockSafetyConfig) -> None:
         if not config.aws_profile:
             raise ValueError(f"Missing boto_client aws_profile in model info::{config}")
         self.config = config
+        self.registered_shields = []
 
     async def initialize(self) -> None:
         try:
@@ -45,16 +45,23 @@ class BedrockSafetyAdapter(Safety, RoutableProvider):
     async def shutdown(self) -> None:
         pass
 
-    async def validate_routing_keys(self, routing_keys: List[str]) -> None:
-        for key in routing_keys:
-            if key not in SUPPORTED_SHIELD_TYPES:
-                raise ValueError(f"Unknown safety shield type: {key}")
+    async def register_shield(self, shield: ShieldDef) -> None:
+        raise ValueError("Registering dynamic shields is not supported")
+
+    async def list_shields(self) -> List[ShieldDef]:
+        raise NotImplementedError(
+            """
+            `list_shields` not implemented; this should read all guardrails from
+            bedrock and populate guardrailId and guardrailVersion in the ShieldDef.
+        """
+        )
 
     async def run_shield(
         self, shield_type: str, messages: List[Message], params: Dict[str, Any] = None
     ) -> RunShieldResponse:
-        if shield_type not in SUPPORTED_SHIELD_TYPES:
-            raise ValueError(f"Unknown safety shield type: {shield_type}")
+        shield_def = await self.shield_store.get_shield(shield_type)
+        if not shield_def:
+            raise ValueError(f"Unknown shield {shield_type}")
 
         """This is the implementation for the bedrock guardrails. The input to the guardrails is to be of this format
         ```content = [
@@ -69,17 +76,9 @@ class BedrockSafetyAdapter(Safety, RoutableProvider):
 
         They contain content, role . For now we will extract the content and default the "qualifiers": ["query"]
         """
-        try:
-            logger.debug(f"run_shield::{params}::messages={messages}")
-            if "guardrailIdentifier" not in params:
-                raise RuntimeError(
-                    "Error running request for BedrockGaurdrails:Missing GuardrailID in request"
-                )
 
-            if "guardrailVersion" not in params:
-                raise RuntimeError(
-                    "Error running request for BedrockGaurdrails:Missing guardrailVersion in request"
-                )
+        shield_params = shield_def.params
+        logger.debug(f"run_shield::{shield_params}::messages={messages}")
 
         # - convert the messages into format Bedrock expects
         content_messages = []
@@ -90,12 +89,11 @@ class BedrockSafetyAdapter(Safety, RoutableProvider):
         )
 
         response = self.boto_client.apply_guardrail(
-                guardrailIdentifier=params.get("guardrailIdentifier"),
-                guardrailVersion=params.get("guardrailVersion"),
+            guardrailIdentifier=shield_params["guardrailIdentifier"],
+            guardrailVersion=shield_params["guardrailVersion"],
             source="OUTPUT",  # or 'INPUT' depending on your use case
             content=content_messages,
         )
-            logger.debug(f"run_shield:: response: {response}::")
         if response["action"] == "GUARDRAIL_INTERVENED":
             user_message = ""
             metadata = {}
@@ -105,16 +103,11 @@ class BedrockSafetyAdapter(Safety, RoutableProvider):
             for assessment in response["assessments"]:
                 # guardrails returns a list - however for this implementation we will leverage the last values
                 metadata = dict(assessment)
+
             return SafetyViolation(
                 user_message=user_message,
                 violation_level=ViolationLevel.ERROR,
                 metadata=metadata,
             )
 
-        except Exception:
-            error_str = traceback.format_exc()
-            logger.error(
-                f"Error in apply_guardrails:{error_str}:: RETURNING None !!!!!"
-            )
-
         return None

llama_stack/providers/adapters/safety/sample/sample.py

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

llama_stack/providers/adapters/safety/together/together.py

@@ -6,26 +6,21 @@
 from together import Together
 
 from llama_models.llama3.api.datatypes import *  # noqa: F403
-from llama_stack.apis.safety import (
-    RunShieldResponse,
-    Safety,
-    SafetyViolation,
-    ViolationLevel,
-)
-from llama_stack.distribution.datatypes import RoutableProvider
+from llama_stack.apis.safety import *  # noqa: F403
 from llama_stack.distribution.request_headers import NeedsRequestProviderData
+from llama_stack.providers.datatypes import ShieldsProtocolPrivate
 
 from .config import TogetherSafetyConfig
 
 
-SAFETY_SHIELD_TYPES = {
+TOGETHER_SHIELD_MODEL_MAP = {
     "llama_guard": "meta-llama/Meta-Llama-Guard-3-8B",
     "Llama-Guard-3-8B": "meta-llama/Meta-Llama-Guard-3-8B",
     "Llama-Guard-3-11B-Vision": "meta-llama/Llama-Guard-3-11B-Vision-Turbo",
 }
 
 
-class TogetherSafetyImpl(Safety, NeedsRequestProviderData, RoutableProvider):
+class TogetherSafetyImpl(Safety, NeedsRequestProviderData, ShieldsProtocolPrivate):
     def __init__(self, config: TogetherSafetyConfig) -> None:
         self.config = config
 
@@ -35,16 +30,28 @@ class TogetherSafetyImpl(Safety, NeedsRequestProviderData, RoutableProvider):
     async def shutdown(self) -> None:
         pass
 
-    async def validate_routing_keys(self, routing_keys: List[str]) -> None:
-        for key in routing_keys:
-            if key not in SAFETY_SHIELD_TYPES:
-                raise ValueError(f"Unknown safety shield type: {key}")
+    async def register_shield(self, shield: ShieldDef) -> None:
+        raise ValueError("Registering dynamic shields is not supported")
+
+    async def list_shields(self) -> List[ShieldDef]:
+        return [
+            ShieldDef(
+                identifier=ShieldType.llama_guard.value,
+                type=ShieldType.llama_guard.value,
+                params={},
+            )
+        ]
 
     async def run_shield(
         self, shield_type: str, messages: List[Message], params: Dict[str, Any] = None
     ) -> RunShieldResponse:
-        if shield_type not in SAFETY_SHIELD_TYPES:
-            raise ValueError(f"Unknown safety shield type: {shield_type}")
+        shield_def = await self.shield_store.get_shield(shield_type)
+        if not shield_def:
+            raise ValueError(f"Unknown shield {shield_type}")
+
+        model = shield_def.params.get("model", "llama_guard")
+        if model not in TOGETHER_SHIELD_MODEL_MAP:
+            raise ValueError(f"Unsupported safety model: {model}")
 
         together_api_key = None
         if self.config.api_key is not None:
@@ -57,8 +64,6 @@ class TogetherSafetyImpl(Safety, NeedsRequestProviderData, RoutableProvider):
                 )
             together_api_key = provider_data.together_api_key
 
-        model_name = SAFETY_SHIELD_TYPES[shield_type]
-
         # messages can have role assistant or user
         api_messages = []
         for message in messages:
@@ -66,7 +71,7 @@ class TogetherSafetyImpl(Safety, NeedsRequestProviderData, RoutableProvider):
                 api_messages.append({"role": message.role, "content": message.content})
 
         violation = await get_safety_response(
-            together_api_key, model_name, api_messages
+            together_api_key, TOGETHER_SHIELD_MODEL_MAP[model], api_messages
         )
         return RunShieldResponse(violation=violation)
 
@@ -90,7 +95,6 @@ async def get_safety_response(
     if parts[0] == "unsafe":
         return SafetyViolation(
             violation_level=ViolationLevel.ERROR,
-            user_message="unsafe",
             metadata={"violation_type": parts[1]},
         )
 

llama_stack/providers/datatypes.py

@@ -10,6 +10,11 @@ from typing import Any, List, Optional, Protocol
 from llama_models.schema_utils import json_schema_type
 from pydantic import BaseModel, Field
 
+from llama_stack.apis.memory_banks import MemoryBankDef
+
+from llama_stack.apis.models import ModelDef
+from llama_stack.apis.shields import ShieldDef
+
 
 @json_schema_type
 class Api(Enum):
@@ -28,6 +33,24 @@ class Api(Enum):
     inspect = "inspect"
 
 
+class ModelsProtocolPrivate(Protocol):
+    async def list_models(self) -> List[ModelDef]: ...
+
+    async def register_model(self, model: ModelDef) -> None: ...
+
+
+class ShieldsProtocolPrivate(Protocol):
+    async def list_shields(self) -> List[ShieldDef]: ...
+
+    async def register_shield(self, shield: ShieldDef) -> None: ...
+
+
+class MemoryBanksProtocolPrivate(Protocol):
+    async def list_memory_banks(self) -> List[MemoryBankDef]: ...
+
+    async def register_memory_bank(self, memory_bank: MemoryBankDef) -> None: ...
+
+
 @json_schema_type
 class ProviderSpec(BaseModel):
     api: Api
@@ -41,23 +64,14 @@ class ProviderSpec(BaseModel):
         description="Higher-level API surfaces may depend on other providers to provide their functionality",
     )
 
+    # used internally by the resolver; this is a hack for now
+    deps__: List[str] = Field(default_factory=list)
+
 
 class RoutingTable(Protocol):
-    def get_routing_keys(self) -> List[str]: ...
-
     def get_provider_impl(self, routing_key: str) -> Any: ...
 
 
-class RoutableProvider(Protocol):
-    """
-    A provider which sits behind the RoutingTable and can get routed to.
-
-    All Inference / Safety / Memory providers fall into this bucket.
-    """
-
-    async def validate_routing_keys(self, keys: List[str]) -> None: ...
-
-
 @json_schema_type
 class AdapterSpec(BaseModel):
     adapter_type: str = Field(
@@ -154,6 +168,10 @@ as being "Llama Stack compatible"
         return None
 
 
+def is_passthrough(spec: ProviderSpec) -> bool:
+    return isinstance(spec, RemoteProviderSpec) and spec.adapter is None
+
+
 # Can avoid this by using Pydantic computed_field
 def remote_provider_spec(
     api: Api, adapter: Optional[AdapterSpec] = None

llama_stack/providers/impls/ios/inference/LocalInferenceImpl/LocalInference.swift

@@ -34,6 +34,10 @@ public class LocalInference: Inference {
     }
   }
 
+  public func stop() {
+    runnerHolder.runner?.stop()
+  }
+
   public func chatCompletion(request: Components.Schemas.ChatCompletionRequest) -> AsyncStream<Components.Schemas.ChatCompletionResponseStreamChunk> {
     return AsyncStream { continuation in
       runnerQueue.async {

llama_stack/providers/impls/meta_reference/agents/__init__.py

@@ -21,6 +21,7 @@ async def get_provider_impl(
         deps[Api.inference],
         deps[Api.memory],
         deps[Api.safety],
+        deps[Api.memory_banks],
     )
     await impl.initialize()
     return impl

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

@@ -24,6 +24,7 @@ from termcolor import cprint
 from llama_stack.apis.agents import *  # noqa: F403
 from llama_stack.apis.inference import *  # noqa: F403
 from llama_stack.apis.memory import *  # noqa: F403
+from llama_stack.apis.memory_banks import *  # noqa: F403
 from llama_stack.apis.safety import *  # noqa: F403
 
 from llama_stack.providers.utils.kvstore import KVStore
@@ -56,6 +57,7 @@ class ChatAgent(ShieldRunnerMixin):
         agent_config: AgentConfig,
         inference_api: Inference,
         memory_api: Memory,
+        memory_banks_api: MemoryBanks,
         safety_api: Safety,
         persistence_store: KVStore,
     ):
@@ -63,6 +65,7 @@ class ChatAgent(ShieldRunnerMixin):
         self.agent_config = agent_config
         self.inference_api = inference_api
         self.memory_api = memory_api
+        self.memory_banks_api = memory_banks_api
         self.safety_api = safety_api
         self.storage = AgentPersistence(agent_id, persistence_store)
 
@@ -144,6 +147,8 @@ class ChatAgent(ShieldRunnerMixin):
     async def create_and_execute_turn(
         self, request: AgentTurnCreateRequest
     ) -> AsyncGenerator:
+        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")
@@ -635,14 +640,13 @@ class ChatAgent(ShieldRunnerMixin):
             raise ValueError(f"Session {session_id} not found")
 
         if session_info.memory_bank_id is None:
-            memory_bank = await self.memory_api.create_memory_bank(
-                name=f"memory_bank_{session_id}",
-                config=VectorMemoryBankConfig(
+            bank_id = f"memory_bank_{session_id}"
+            memory_bank = VectorMemoryBankDef(
+                identifier=bank_id,
                 embedding_model="all-MiniLM-L6-v2",
                 chunk_size_in_tokens=512,
-                ),
             )
-            bank_id = memory_bank.bank_id
+            await self.memory_banks_api.register_memory_bank(memory_bank)
             await self.storage.add_memory_bank_to_session(session_id, bank_id)
         else:
             bank_id = session_info.memory_bank_id

llama_stack/providers/impls/meta_reference/agents/agents.py

@@ -11,6 +11,7 @@ from typing import AsyncGenerator
 
 from llama_stack.apis.inference import Inference
 from llama_stack.apis.memory import Memory
+from llama_stack.apis.memory_banks import MemoryBanks
 from llama_stack.apis.safety import Safety
 from llama_stack.apis.agents import *  # noqa: F403
 
@@ -30,11 +31,14 @@ class MetaReferenceAgentsImpl(Agents):
         inference_api: Inference,
         memory_api: Memory,
         safety_api: Safety,
+        memory_banks_api: MemoryBanks,
     ):
         self.config = config
         self.inference_api = inference_api
         self.memory_api = memory_api
         self.safety_api = safety_api
+        self.memory_banks_api = memory_banks_api
+
         self.in_memory_store = InmemoryKVStoreImpl()
 
     async def initialize(self) -> None:
@@ -81,6 +85,7 @@ class MetaReferenceAgentsImpl(Agents):
             inference_api=self.inference_api,
             safety_api=self.safety_api,
             memory_api=self.memory_api,
+            memory_banks_api=self.memory_banks_api,
             persistence_store=(
                 self.persistence_store
                 if agent_config.enable_session_persistence
@@ -100,7 +105,7 @@ class MetaReferenceAgentsImpl(Agents):
             session_id=session_id,
         )
 
-    async def create_agent_turn(
+    def create_agent_turn(
         self,
         agent_id: str,
         session_id: str,
@@ -113,16 +118,44 @@ class MetaReferenceAgentsImpl(Agents):
         attachments: Optional[List[Attachment]] = None,
         stream: Optional[bool] = False,
     ) -> AsyncGenerator:
-        agent = await self.get_agent(agent_id)
-
-        # wrapper request to make it easier to pass around (internal only, not exposed to API)
         request = AgentTurnCreateRequest(
             agent_id=agent_id,
             session_id=session_id,
             messages=messages,
             attachments=attachments,
-            stream=stream,
+            stream=True,
         )
+        if stream:
+            return self._create_agent_turn_streaming(request)
+        else:
+            raise NotImplementedError("Non-streaming agent turns not yet implemented")
 
+    async def _create_agent_turn_streaming(
+        self,
+        request: AgentTurnCreateRequest,
+    ) -> AsyncGenerator:
+        agent = await self.get_agent(request.agent_id)
         async for event in agent.create_and_execute_turn(request):
             yield event
+
+    async def get_agents_turn(self, agent_id: str, turn_id: str) -> Turn:
+        raise NotImplementedError()
+
+    async def get_agents_step(
+        self, agent_id: str, turn_id: str, step_id: str
+    ) -> AgentStepResponse:
+        raise NotImplementedError()
+
+    async def get_agents_session(
+        self,
+        agent_id: str,
+        session_id: str,
+        turn_ids: Optional[List[str]] = None,
+    ) -> Session:
+        raise NotImplementedError()
+
+    async def delete_agents_session(self, agent_id: str, session_id: str) -> None:
+        raise NotImplementedError()
+
+    async def delete_agents(self, agent_id: str) -> None:
+        raise NotImplementedError()

llama_stack/providers/impls/meta_reference/agents/rag/context_retriever.py

@@ -63,13 +63,12 @@ async def llm_rag_query_generator(
 
     model = config.model
     message = UserMessage(content=content)
-    response = inference_api.chat_completion(
+    response = await inference_api.chat_completion(
         model=model,
         messages=[message],
         stream=False,
     )
 
-    async for chunk in response:
-        query = chunk.completion_message.content
+    query = response.completion_message.content
 
     return query

llama_stack/providers/impls/meta_reference/agents/tests/test_chat_agent.py

@@ -16,7 +16,7 @@ from llama_stack.apis.agents import *  # noqa: F403
 from ..agents import (
     AGENT_INSTANCES_BY_ID,
     MetaReferenceAgentsImpl,
-    MetaReferenceImplConfig,
+    MetaReferenceInferenceConfig,
 )
 
 
@@ -166,7 +166,7 @@ def mock_memory_api():
 @pytest.fixture
 async def chat_agent(mock_inference_api, mock_safety_api, mock_memory_api):
     impl = MetaReferenceAgentsImpl(
-        config=MetaReferenceImplConfig(),
+        config=MetaReferenceInferenceConfig(),
         inference_api=mock_inference_api,
         safety_api=mock_safety_api,
         memory_api=mock_memory_api,

llama_stack/providers/impls/meta_reference/codeshield/__init__.py

@@ -0,0 +1,15 @@
+# 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 CodeShieldConfig
+
+
+async def get_provider_impl(config: CodeShieldConfig, deps):
+    from .code_scanner import MetaReferenceCodeScannerSafetyImpl
+
+    impl = MetaReferenceCodeScannerSafetyImpl(config, deps)
+    await impl.initialize()
+    return impl

llama_stack/providers/impls/meta_reference/codeshield/code_scanner.py

@@ -0,0 +1,58 @@
+# 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, List
+
+from llama_models.llama3.api.datatypes import interleaved_text_media_as_str, Message
+from termcolor import cprint
+
+from .config import CodeScannerConfig
+
+from llama_stack.apis.safety import *  # noqa: F403
+
+
+class MetaReferenceCodeScannerSafetyImpl(Safety):
+    def __init__(self, config: CodeScannerConfig, deps) -> None:
+        self.config = config
+
+    async def initialize(self) -> None:
+        pass
+
+    async def shutdown(self) -> None:
+        pass
+
+    async def register_shield(self, shield: ShieldDef) -> None:
+        if shield.type != ShieldType.code_scanner.value:
+            raise ValueError(f"Unsupported safety shield type: {shield.type}")
+
+    async def run_shield(
+        self,
+        shield_type: str,
+        messages: List[Message],
+        params: Dict[str, Any] = None,
+    ) -> RunShieldResponse:
+        shield_def = await self.shield_store.get_shield(shield_type)
+        if not shield_def:
+            raise ValueError(f"Unknown shield {shield_type}")
+
+        from codeshield.cs import CodeShield
+
+        text = "\n".join([interleaved_text_media_as_str(m.content) for m in messages])
+        cprint(f"Running CodeScannerShield on {text[50:]}", color="magenta")
+        result = await CodeShield.scan_code(text)
+
+        violation = None
+        if result.is_insecure:
+            violation = SafetyViolation(
+                violation_level=(ViolationLevel.ERROR),
+                user_message="Sorry, I found security concerns in the code.",
+                metadata={
+                    "violation_type": ",".join(
+                        [issue.pattern_id for issue in result.issues_found]
+                    )
+                },
+            )
+        return RunShieldResponse(violation=violation)

llama_stack/providers/impls/meta_reference/codeshield/config.py

@@ -0,0 +1,11 @@
+# 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 pydantic import BaseModel
+
+
+class CodeShieldConfig(BaseModel):
+    pass

llama_stack/providers/impls/meta_reference/inference/__init__.py

@@ -4,16 +4,17 @@
 # This source code is licensed under the terms described in the LICENSE file in
 # the root directory of this source tree.
 
-from .config import MetaReferenceImplConfig  # noqa
+from typing import Union
+
+from .config import MetaReferenceInferenceConfig, MetaReferenceQuantizedInferenceConfig
 
 
-async def get_provider_impl(config: MetaReferenceImplConfig, _deps):
+async def get_provider_impl(
+    config: Union[MetaReferenceInferenceConfig, MetaReferenceQuantizedInferenceConfig],
+    _deps,
+):
     from .inference import MetaReferenceInferenceImpl
 
-    assert isinstance(
-        config, MetaReferenceImplConfig
-    ), f"Unexpected config type: {type(config)}"
-
     impl = MetaReferenceInferenceImpl(config)
     await impl.initialize()
     return impl

llama_stack/providers/impls/meta_reference/inference/config.py

@@ -15,12 +15,11 @@ from pydantic import BaseModel, Field, field_validator
 from llama_stack.providers.utils.inference import supported_inference_models
 
 
-class MetaReferenceImplConfig(BaseModel):
+class MetaReferenceInferenceConfig(BaseModel):
     model: str = Field(
         default="Llama3.1-8B-Instruct",
         description="Model descriptor from `llama model list`",
     )
-    quantization: Optional[QuantizationConfig] = None
     torch_seed: Optional[int] = None
     max_seq_len: int = 4096
     max_batch_size: int = 1
@@ -38,9 +37,9 @@ class MetaReferenceImplConfig(BaseModel):
 
     @property
     def model_parallel_size(self) -> int:
-        # HACK ALERT: this will be fixed when we move inference configuration
-        # to ModelsRegistry and we can explicitly ask for `model_parallel_size`
-        # as configuration there
         resolved = resolve_model(self.model)
-        assert resolved is not None
         return resolved.pth_file_count
+
+
+class MetaReferenceQuantizedInferenceConfig(MetaReferenceInferenceConfig):
+    quantization: QuantizationConfig

llama_stack/providers/impls/meta_reference/inference/generation.py

@@ -11,9 +11,8 @@ import json
 import os
 import sys
 import time
-from dataclasses import dataclass
 from pathlib import Path
-from typing import Generator, List, Optional
+from typing import Generator, List, Optional, Union
 
 import torch
 import torch.nn.functional as F
@@ -35,13 +34,13 @@ from llama_models.llama3.reference_impl.multimodal.model import (
     CrossAttentionTransformer,
 )
 from llama_models.sku_list import resolve_model
-from termcolor import cprint
 
-from llama_stack.apis.inference import QuantizationType
+from pydantic import BaseModel
+from termcolor import cprint
 
 from llama_stack.distribution.utils.model_utils import model_local_dir
 
-from .config import MetaReferenceImplConfig
+from .config import MetaReferenceInferenceConfig, MetaReferenceQuantizedInferenceConfig
 
 
 def model_checkpoint_dir(model) -> str:
@@ -58,8 +57,7 @@ def model_checkpoint_dir(model) -> str:
     return str(checkpoint_dir)
 
 
-@dataclass
-class TokenResult:
+class TokenResult(BaseModel):
     token: int
     text: str
     logprobs: Optional[List[float]] = None
@@ -67,7 +65,11 @@ class TokenResult:
 
 class Llama:
     @staticmethod
-    def build(config: MetaReferenceImplConfig):
+    def build(
+        config: Union[
+            MetaReferenceInferenceConfig, MetaReferenceQuantizedInferenceConfig
+        ]
+    ):
         """
         Build a Llama instance by initializing and loading a model checkpoint.
 
@@ -77,15 +79,6 @@ class Llama:
         """
         model = resolve_model(config.model)
 
-        if (
-            config.quantization
-            and config.quantization.type == QuantizationType.fp8.value
-        ):
-            from .quantization.loader import is_fbgemm_available
-
-            if not is_fbgemm_available():
-                raise ImportError("fbgemm-gpu is required for FP8 quantization")
-
         if not torch.distributed.is_initialized():
             torch.distributed.init_process_group("nccl")
 
@@ -133,12 +126,7 @@ class Llama:
             model_args.vocab_size == tokenizer.n_words
         ), f"model_args vocab = {model_args.vocab_size} but tokenizer vocab = {tokenizer.n_words}"
 
-        fp8 = (
-            config.quantization
-            and config.quantization.type == QuantizationType.fp8.value
-        )
-
-        if fp8:
+        if isinstance(config, MetaReferenceQuantizedInferenceConfig):
             from .quantization.loader import convert_to_quantized_model
 
             # load on CPU in bf16 so that fp8 conversion does not find an

llama_stack/providers/impls/meta_reference/inference/inference.py

@@ -6,18 +6,18 @@
 
 import asyncio
 
-from typing import AsyncIterator, List, Union
+from typing import AsyncGenerator, List
 
 from llama_models.sku_list import resolve_model
 
 from llama_models.llama3.api.datatypes import *  # noqa: F403
 from llama_stack.apis.inference import *  # noqa: F403
-from llama_stack.distribution.datatypes import RoutableProvider
-from llama_stack.providers.utils.inference.augment_messages import (
-    augment_messages_for_tools,
+from llama_stack.providers.datatypes import ModelDef, ModelsProtocolPrivate
+from llama_stack.providers.utils.inference.prompt_adapter import (
+    chat_completion_request_to_messages,
 )
 
-from .config import MetaReferenceImplConfig
+from .config import MetaReferenceInferenceConfig
 from .model_parallel import LlamaModelParallelGenerator
 
 # there's a single model parallel process running serving the model. for now,
@@ -25,8 +25,8 @@ from .model_parallel import LlamaModelParallelGenerator
 SEMAPHORE = asyncio.Semaphore(1)
 
 
-class MetaReferenceInferenceImpl(Inference, RoutableProvider):
-    def __init__(self, config: MetaReferenceImplConfig) -> None:
+class MetaReferenceInferenceImpl(Inference, ModelsProtocolPrivate):
+    def __init__(self, config: MetaReferenceInferenceConfig) -> None:
         self.config = config
         model = resolve_model(config.model)
         if model is None:
@@ -35,21 +35,35 @@ class MetaReferenceInferenceImpl(Inference, RoutableProvider):
         # verify that the checkpoint actually is for this model lol
 
     async def initialize(self) -> None:
+        print(f"Loading model `{self.model.descriptor()}`")
         self.generator = LlamaModelParallelGenerator(self.config)
         self.generator.start()
 
-    async def validate_routing_keys(self, routing_keys: List[str]) -> None:
-        assert (
-            len(routing_keys) == 1
-        ), f"Only one routing key is supported {routing_keys}"
-        assert routing_keys[0] == self.config.model
+    async def register_model(self, model: ModelDef) -> None:
+        raise ValueError("Dynamic model registration is not supported")
+
+    async def list_models(self) -> List[ModelDef]:
+        return [
+            ModelDef(
+                identifier=self.model.descriptor(),
+                llama_model=self.model.descriptor(),
+            )
+        ]
 
     async def shutdown(self) -> None:
         self.generator.stop()
 
-    # hm, when stream=False, we should not be doing SSE :/ which is what the
-    # top-level server is going to do. make the typing more specific here
-    async def chat_completion(
+    def completion(
+        self,
+        model: str,
+        content: InterleavedTextMedia,
+        sampling_params: Optional[SamplingParams] = SamplingParams(),
+        stream: Optional[bool] = False,
+        logprobs: Optional[LogProbConfig] = None,
+    ) -> Union[CompletionResponse, CompletionResponseStreamChunk]:
+        raise NotImplementedError()
+
+    def chat_completion(
         self,
         model: str,
         messages: List[Message],
@@ -59,9 +73,10 @@ class MetaReferenceInferenceImpl(Inference, RoutableProvider):
         tool_prompt_format: Optional[ToolPromptFormat] = ToolPromptFormat.json,
         stream: Optional[bool] = False,
         logprobs: Optional[LogProbConfig] = None,
-    ) -> AsyncIterator[
-        Union[ChatCompletionResponseStreamChunk, ChatCompletionResponse]
-    ]:
+    ) -> AsyncGenerator:
+        if logprobs:
+            assert logprobs.top_k == 1, f"Unexpected top_k={logprobs.top_k}"
+
         # wrapper request to make it easier to pass around (internal only, not exposed to API)
         request = ChatCompletionRequest(
             model=model,
@@ -74,7 +89,6 @@ class MetaReferenceInferenceImpl(Inference, RoutableProvider):
             logprobs=logprobs,
         )
 
-        messages = augment_messages_for_tools(request)
         model = resolve_model(request.model)
         if model is None:
             raise RuntimeError(
@@ -88,8 +102,64 @@ class MetaReferenceInferenceImpl(Inference, RoutableProvider):
         if SEMAPHORE.locked():
             raise RuntimeError("Only one concurrent request is supported")
 
-        async with SEMAPHORE:
         if request.stream:
+            return self._stream_chat_completion(request)
+        else:
+            return self._nonstream_chat_completion(request)
+
+    async def _nonstream_chat_completion(
+        self, request: ChatCompletionRequest
+    ) -> ChatCompletionResponse:
+        async with SEMAPHORE:
+            messages = chat_completion_request_to_messages(request)
+
+            tokens = []
+            logprobs = []
+            stop_reason = None
+
+            for token_result in self.generator.chat_completion(
+                messages=messages,
+                temperature=request.sampling_params.temperature,
+                top_p=request.sampling_params.top_p,
+                max_gen_len=request.sampling_params.max_tokens,
+                logprobs=request.logprobs,
+                tool_prompt_format=request.tool_prompt_format,
+            ):
+                tokens.append(token_result.token)
+
+                if token_result.text == "<|eot_id|>":
+                    stop_reason = StopReason.end_of_turn
+                elif token_result.text == "<|eom_id|>":
+                    stop_reason = StopReason.end_of_message
+
+                if request.logprobs:
+                    assert len(token_result.logprobs) == 1
+
+                    logprobs.append(
+                        TokenLogProbs(
+                            logprobs_by_token={
+                                token_result.text: token_result.logprobs[0]
+                            }
+                        )
+                    )
+
+            if stop_reason is None:
+                stop_reason = StopReason.out_of_tokens
+
+            message = self.generator.formatter.decode_assistant_message(
+                tokens, stop_reason
+            )
+            return ChatCompletionResponse(
+                completion_message=message,
+                logprobs=logprobs if request.logprobs else None,
+            )
+
+    async def _stream_chat_completion(
+        self, request: ChatCompletionRequest
+    ) -> AsyncGenerator:
+        async with SEMAPHORE:
+            messages = chat_completion_request_to_messages(request)
+
             yield ChatCompletionResponseStreamChunk(
                 event=ChatCompletionResponseEvent(
                     event_type=ChatCompletionResponseEventType.start,
@@ -99,10 +169,7 @@ class MetaReferenceInferenceImpl(Inference, RoutableProvider):
 
             tokens = []
             logprobs = []
-
             stop_reason = None
-
-            buffer = ""
             ipython = False
 
             for token_result in self.generator.chat_completion(
@@ -113,10 +180,9 @@ class MetaReferenceInferenceImpl(Inference, RoutableProvider):
                 logprobs=request.logprobs,
                 tool_prompt_format=request.tool_prompt_format,
             ):
-                buffer += token_result.text
                 tokens.append(token_result.token)
 
-                if not ipython and buffer.startswith("<|python_tag|>"):
+                if not ipython and token_result.text.startswith("<|python_tag|>"):
                     ipython = True
                     yield ChatCompletionResponseStreamChunk(
                         event=ChatCompletionResponseEvent(
@@ -127,26 +193,6 @@ class MetaReferenceInferenceImpl(Inference, RoutableProvider):
                             ),
                         )
                     )
-                    buffer = buffer[len("<|python_tag|>") :]
-                    continue
-
-                if not request.stream:
-                    if request.logprobs:
-                        assert (
-                            len(token_result.logprobs) == 1
-                        ), "Expected logprob to contain 1 result for the current token"
-                        assert (
-                            request.logprobs.top_k == 1
-                        ), "Only top_k=1 is supported for LogProbConfig"
-
-                        logprobs.append(
-                            TokenLogProbs(
-                                logprobs_by_token={
-                                    token_result.text: token_result.logprobs[0]
-                                }
-                            )
-                        )
-
                     continue
 
                 if token_result.text == "<|eot_id|>":
@@ -167,23 +213,32 @@ class MetaReferenceInferenceImpl(Inference, RoutableProvider):
                     delta = text
 
                 if stop_reason is None:
+                    if request.logprobs:
+                        assert len(token_result.logprobs) == 1
+
+                        logprobs.append(
+                            TokenLogProbs(
+                                logprobs_by_token={
+                                    token_result.text: token_result.logprobs[0]
+                                }
+                            )
+                        )
                     yield ChatCompletionResponseStreamChunk(
                         event=ChatCompletionResponseEvent(
                             event_type=ChatCompletionResponseEventType.progress,
                             delta=delta,
                             stop_reason=stop_reason,
+                            logprobs=logprobs if request.logprobs else None,
                         )
                     )
 
             if stop_reason is None:
                 stop_reason = StopReason.out_of_tokens
 
-            # TODO(ashwin): parse tool calls separately here and report errors?
-            # if someone breaks the iteration before coming here we are toast
             message = self.generator.formatter.decode_assistant_message(
                 tokens, stop_reason
             )
-            if request.stream:
+
             parsed_tool_calls = len(message.tool_calls) > 0
             if ipython and not parsed_tool_calls:
                 yield ChatCompletionResponseStreamChunk(
@@ -217,9 +272,9 @@ class MetaReferenceInferenceImpl(Inference, RoutableProvider):
                 )
             )
 
-                # TODO(ashwin): what else do we need to send out here when everything finishes?
-            else:
-                yield ChatCompletionResponse(
-                    completion_message=message,
-                    logprobs=logprobs if request.logprobs else None,
-                )
+    async def embeddings(
+        self,
+        model: str,
+        contents: List[InterleavedTextMedia],
+    ) -> EmbeddingsResponse:
+        raise NotImplementedError()

llama_stack/providers/impls/meta_reference/inference/model_parallel.py

@@ -6,7 +6,6 @@
 
 import os
 from copy import deepcopy
-from dataclasses import dataclass
 from functools import partial
 from typing import Generator, List, Optional
 
@@ -15,19 +14,9 @@ from llama_models.llama3.api.datatypes import Message, ToolPromptFormat
 from llama_models.llama3.api.tokenizer import Tokenizer
 from llama_models.sku_list import resolve_model
 
-from .config import MetaReferenceImplConfig
+from .config import MetaReferenceInferenceConfig
 from .generation import Llama, model_checkpoint_dir
-from .parallel_utils import ModelParallelProcessGroup
-
-
-@dataclass
-class InferenceArgs:
-    messages: List[Message]
-    temperature: float
-    top_p: float
-    max_gen_len: int
-    logprobs: bool
-    tool_prompt_format: ToolPromptFormat
+from .parallel_utils import InferenceArgs, ModelParallelProcessGroup
 
 
 class ModelRunner:
@@ -46,7 +35,7 @@ class ModelRunner:
         )
 
 
-def init_model_cb(config: MetaReferenceImplConfig):
+def init_model_cb(config: MetaReferenceInferenceConfig):
     llama = Llama.build(config)
     return ModelRunner(llama)
 
@@ -62,7 +51,7 @@ class LlamaModelParallelGenerator:
     clear at the callsite why we need to use a context manager.
     """
 
-    def __init__(self, config: MetaReferenceImplConfig):
+    def __init__(self, config: MetaReferenceInferenceConfig):
         self.config = config
         self.model = resolve_model(self.config.model)
         # this is a hack because Agent's loop uses this to tokenize and check if input is too long
@@ -102,7 +91,7 @@ class LlamaModelParallelGenerator:
             temperature=temperature,
             top_p=top_p,
             max_gen_len=max_gen_len,
-            logprobs=logprobs,
+            logprobs=logprobs or False,
             tool_prompt_format=tool_prompt_format,
         )
 

llama_stack/providers/impls/meta_reference/inference/parallel_utils.py

@@ -4,14 +4,14 @@
 # This source code is licensed under the terms described in the LICENSE file in
 # the root directory of this source tree.
 
+import json
 import multiprocessing
 import os
-import pickle
 import tempfile
 import time
 import uuid
-
-from typing import Callable, Generator
+from enum import Enum
+from typing import Callable, Generator, List, Literal, Optional, Union
 
 import torch
 
@@ -23,17 +23,106 @@ from fairscale.nn.model_parallel.initialize import (
     get_model_parallel_src_rank,
 )
 
+from llama_models.llama3.api.datatypes import Message, ToolPromptFormat
+
+from pydantic import BaseModel, Field
+
 from torch.distributed.launcher.api import elastic_launch, LaunchConfig
+from typing_extensions import Annotated
+
+from .generation import TokenResult
 
 
-_END_SENTINEL = "__end_sentinel__"
-_CANCEL_SENTINEL = "__cancel_sentinel__"
+class InferenceArgs(BaseModel):
+    messages: List[Message]
+    temperature: float
+    top_p: float
+    max_gen_len: int
+    logprobs: bool
+    tool_prompt_format: ToolPromptFormat
+
+
+class ProcessingMessageName(str, Enum):
+    ready_request = "ready_request"
+    ready_response = "ready_response"
+    end_sentinel = "end_sentinel"
+    cancel_sentinel = "cancel_sentinel"
+    task_request = "task_request"
+    task_response = "task_response"
+    exception_response = "exception_response"
+
+
+class ReadyRequest(BaseModel):
+    type: Literal[ProcessingMessageName.ready_request] = (
+        ProcessingMessageName.ready_request
+    )
+
+
+class ReadyResponse(BaseModel):
+    type: Literal[ProcessingMessageName.ready_response] = (
+        ProcessingMessageName.ready_response
+    )
+
+
+class EndSentinel(BaseModel):
+    type: Literal[ProcessingMessageName.end_sentinel] = (
+        ProcessingMessageName.end_sentinel
+    )
+
+
+class CancelSentinel(BaseModel):
+    type: Literal[ProcessingMessageName.cancel_sentinel] = (
+        ProcessingMessageName.cancel_sentinel
+    )
+
+
+class TaskRequest(BaseModel):
+    type: Literal[ProcessingMessageName.task_request] = (
+        ProcessingMessageName.task_request
+    )
+    task: InferenceArgs
+
+
+class TaskResponse(BaseModel):
+    type: Literal[ProcessingMessageName.task_response] = (
+        ProcessingMessageName.task_response
+    )
+    result: TokenResult
+
+
+class ExceptionResponse(BaseModel):
+    type: Literal[ProcessingMessageName.exception_response] = (
+        ProcessingMessageName.exception_response
+    )
+    error: str
+
+
+ProcessingMessage = Union[
+    ReadyRequest,
+    ReadyResponse,
+    EndSentinel,
+    CancelSentinel,
+    TaskRequest,
+    TaskResponse,
+    ExceptionResponse,
+]
+
+
+class ProcessingMessageWrapper(BaseModel):
+    payload: Annotated[
+        ProcessingMessage,
+        Field(discriminator="type"),
+    ]
 
 
 def mp_rank_0() -> bool:
     return get_model_parallel_rank() == 0
 
 
+def encode_msg(msg: ProcessingMessage) -> bytes:
+    return ProcessingMessageWrapper(payload=msg).model_dump_json().encode("utf-8")
+
+
 def retrieve_requests(reply_socket_url: str):
     if mp_rank_0():
         context = zmq.Context()
@@ -46,20 +135,23 @@ def retrieve_requests(reply_socket_url: str):
                 time.sleep(0.01)
                 continue
 
-            reply_socket.send_multipart([client_id, pickle.dumps("YES READY")])
+            ready_response = ReadyResponse()
+            reply_socket.send_multipart([client_id, encode_msg(ready_response)])
             break
 
-    def send_obj(obj):
-        reply_socket.send_multipart([client_id, pickle.dumps(obj)])
+    def send_obj(obj: ProcessingMessage):
+        reply_socket.send_multipart([client_id, encode_msg(obj)])
 
     while True:
         tasks = [None]
         if mp_rank_0():
-            client_id, task = maybe_get_work(reply_socket)
+            client_id, maybe_task_json = maybe_get_work(reply_socket)
+            if maybe_task_json is not None:
+                task = maybe_parse_message(maybe_task_json)
                 # there is still an unknown unclean GeneratorExit happening resulting in a
                 # cancel sentinel getting queued _after_ we have finished sending everything :/
                 # kind of a hack this is :/
-            if task != _CANCEL_SENTINEL:
+                if task is not None and not isinstance(task, CancelSentinel):
                     tasks = [task]
 
         torch.distributed.broadcast_object_list(
@@ -80,35 +172,36 @@ def retrieve_requests(reply_socket_url: str):
                 for obj in out:
                     updates = [None]
                     if mp_rank_0():
-                        _, update = maybe_get_work(reply_socket)
-                        if update == _CANCEL_SENTINEL:
+                        _, update_json = maybe_get_work(reply_socket)
+                        update = maybe_parse_message(update_json)
+                        if isinstance(update, CancelSentinel):
                             updates = [update]
                         else:
                             # only send the update if it's not cancelled otherwise the object sits in the socket
                             # and gets pulled in the next request lol
-                            send_obj(obj)
+                            send_obj(TaskResponse(result=obj))
 
                     torch.distributed.broadcast_object_list(
                         updates,
                         src=get_model_parallel_src_rank(),
                         group=get_model_parallel_group(),
                     )
-                    if updates[0] == _CANCEL_SENTINEL:
+                    if isinstance(updates[0], CancelSentinel):
                         print("quitting generation loop because request was cancelled")
                         break
 
                 if mp_rank_0():
-                    send_obj(_END_SENTINEL)
+                    send_obj(EndSentinel())
             except Exception as e:
                 print(f"[debug] got exception {e}")
                 import traceback
 
                 traceback.print_exc()
                 if mp_rank_0():
-                    send_obj(e)
+                    send_obj(ExceptionResponse(error=str(e)))
 
     if mp_rank_0():
-        send_obj("DONE")
+        send_obj(EndSentinel())
 
 
 def maybe_get_work(sock: zmq.Socket):
@@ -116,7 +209,7 @@ def maybe_get_work(sock: zmq.Socket):
     client_id = None
     try:
         client_id, obj = sock.recv_multipart(zmq.NOBLOCK)
-        message = pickle.loads(obj)
+        message = obj.decode("utf-8")
     except zmq.ZMQError as e:
         if e.errno != zmq.EAGAIN:
             raise e
@@ -124,6 +217,22 @@ def maybe_get_work(sock: zmq.Socket):
     return client_id, message
 
 
+def maybe_parse_message(maybe_json: Optional[str]) -> Optional[ProcessingMessage]:
+    if maybe_json is None:
+        return None
+    try:
+        return parse_message(maybe_json)
+    except json.JSONDecodeError:
+        return None
+    except ValueError as e:
+        return None
+
+
+def parse_message(json_str: str) -> ProcessingMessage:
+    data = json.loads(json_str)
+    return ProcessingMessageWrapper(**data).payload
+
+
 def worker_process_entrypoint(
     reply_socket_url: str,
     init_model_cb: Callable,
@@ -142,7 +251,8 @@ def worker_process_entrypoint(
             if isinstance(task, str) and task == _END_SENTINEL:
                 break
 
-            result = model(task)
+            assert isinstance(task, TaskRequest)
+            result = model(task.task)
         except StopIteration:
             break
 
@@ -205,9 +315,9 @@ def start_model_parallel_process(
 
     # wait until the model is loaded; rank 0 will send a message to indicate it's ready
 
-    request_socket.send_pyobj("READY?")
-    response = request_socket.recv_pyobj()
-    print(f"Finished model load {response}")
+    request_socket.send(encode_msg(ReadyRequest()))
+    response = request_socket.recv()
+    print("Loaded model...")
 
     return request_socket, process
 
@@ -235,31 +345,36 @@ class ModelParallelProcessGroup:
     def stop(self):
         assert self.started, "process group not started"
         if self.process.is_alive():
-            self.request_socket.send_pyobj(_END_SENTINEL, zmq.NOBLOCK)
+            self.request_socket.send(encode_msg(EndSentinel()), zmq.NOBLOCK)
             self.process.join()
         self.started = False
 
-    def run_inference(self, request) -> Generator:
+    def run_inference(self, inference_args: InferenceArgs) -> Generator:
         assert not self.running, "inference already running"
 
         self.running = True
-        self.request_socket.send_pyobj(request)
+        self.request_socket.send(encode_msg(TaskRequest(task=inference_args)))
         try:
             while True:
-                obj = self.request_socket.recv_pyobj()
-                if obj == _END_SENTINEL:
+                obj_json = self.request_socket.recv()
+                obj = parse_message(obj_json)
+
+                if isinstance(obj, EndSentinel):
                     break
 
-                if isinstance(obj, Exception):
-                    print(f"[debug] got exception {obj}")
-                    raise obj
+                if isinstance(obj, ExceptionResponse):
+                    print(f"[debug] got exception {obj.error}")
+                    raise Exception(obj.error)
+
+                if isinstance(obj, TaskResponse):
+                    yield obj.result
 
-                yield obj
         except GeneratorExit as e:
-            self.request_socket.send_pyobj(_CANCEL_SENTINEL)
+            self.request_socket.send(encode_msg(CancelSentinel()))
             while True:
-                obj = self.request_socket.recv_pyobj()
-                if obj == _END_SENTINEL:
+                obj_json = self.request_socket.send()
+                obj = parse_message(obj_json)
+                if isinstance(obj, EndSentinel):
                     break
         finally:
             self.running = False

llama_stack/providers/impls/meta_reference/inference/quantization/loader.py

@@ -22,19 +22,10 @@ from torch import Tensor
 from llama_stack.apis.inference import QuantizationType
 
 from llama_stack.providers.impls.meta_reference.inference.config import (
-    MetaReferenceImplConfig,
+    MetaReferenceQuantizedInferenceConfig,
 )
 
 
-def is_fbgemm_available() -> bool:
-    try:
-        import fbgemm_gpu.experimental.gen_ai  # noqa: F401
-
-        return True
-    except ImportError:
-        return False
-
-
 def swiglu_wrapper(
     self,
     x: Tensor,
@@ -47,7 +38,7 @@ def swiglu_wrapper(
 
 def convert_to_quantized_model(
     model: Transformer,
-    config: MetaReferenceImplConfig,
+    config: MetaReferenceQuantizedInferenceConfig,
     fp8_activation_scale_ub: Optional[float] = 1200.0,
 ) -> Transformer:
     if config.quantization.type == QuantizationType.bf16.value:

llama_stack/providers/impls/meta_reference/memory/faiss.py

@@ -5,7 +5,6 @@
 # the root directory of this source tree.
 
 import logging
-import uuid
 
 from typing import Any, Dict, List, Optional
 
@@ -14,9 +13,10 @@ import numpy as np
 from numpy.typing import NDArray
 
 from llama_models.llama3.api.datatypes import *  # noqa: F403
-from llama_stack.distribution.datatypes import RoutableProvider
 
 from llama_stack.apis.memory import *  # noqa: F403
+from llama_stack.providers.datatypes import MemoryBanksProtocolPrivate
+
 from llama_stack.providers.utils.memory.vector_store import (
     ALL_MINILM_L6_V2_DIMENSION,
     BankWithIndex,
@@ -63,7 +63,7 @@ class FaissIndex(EmbeddingIndex):
         return QueryDocumentsResponse(chunks=chunks, scores=scores)
 
 
-class FaissMemoryImpl(Memory, RoutableProvider):
+class FaissMemoryImpl(Memory, MemoryBanksProtocolPrivate):
     def __init__(self, config: FaissImplConfig) -> None:
         self.config = config
         self.cache = {}
@@ -72,37 +72,21 @@ class FaissMemoryImpl(Memory, RoutableProvider):
 
     async def shutdown(self) -> None: ...
 
-    async def validate_routing_keys(self, routing_keys: List[str]) -> None:
-        print(f"[faiss] Registering memory bank routing keys: {routing_keys}")
-        pass
-
-    async def create_memory_bank(
+    async def register_memory_bank(
         self,
-        name: str,
-        config: MemoryBankConfig,
-        url: Optional[URL] = None,
-    ) -> MemoryBank:
-        assert url is None, "URL is not supported for this implementation"
+        memory_bank: MemoryBankDef,
+    ) -> None:
         assert (
-            config.type == MemoryBankType.vector.value
-        ), f"Only vector banks are supported {config.type}"
+            memory_bank.type == MemoryBankType.vector.value
+        ), f"Only vector banks are supported {memory_bank.type}"
 
-        bank_id = str(uuid.uuid4())
-        bank = MemoryBank(
-            bank_id=bank_id,
-            name=name,
-            config=config,
-            url=url,
+        index = BankWithIndex(
+            bank=memory_bank, index=FaissIndex(ALL_MINILM_L6_V2_DIMENSION)
         )
-        index = BankWithIndex(bank=bank, index=FaissIndex(ALL_MINILM_L6_V2_DIMENSION))
-        self.cache[bank_id] = index
-        return bank
+        self.cache[memory_bank.identifier] = index
 
-    async def get_memory_bank(self, bank_id: str) -> Optional[MemoryBank]:
-        index = self.cache.get(bank_id)
-        if index is None:
-            return None
-        return index.bank
+    async def list_memory_banks(self) -> List[MemoryBankDef]:
+        return [i.bank for i in self.cache.values()]
 
     async def insert_documents(
         self,

llama_stack/providers/impls/meta_reference/safety/base.py

@@ -44,7 +44,6 @@ def message_content_as_str(message: Message) -> str:
     return interleaved_text_media_as_str(message.content)
 
 
-# For shields that operate on simple strings
 class TextShield(ShieldBase):
     def convert_messages_to_text(self, messages: List[Message]) -> str:
         return "\n".join([message_content_as_str(m) for m in messages])
@@ -56,9 +55,3 @@ class TextShield(ShieldBase):
     @abstractmethod
     async def run_impl(self, text: str) -> ShieldResponse:
         raise NotImplementedError()
-
-
-class DummyShield(TextShield):
-    async def run_impl(self, text: str) -> ShieldResponse:
-        # Dummy return LOW to test e2e
-        return ShieldResponse(is_violation=False)

llama_stack/providers/impls/meta_reference/safety/config.py

@@ -9,23 +9,19 @@ from typing import List, Optional
 
 from llama_models.sku_list import CoreModelId, safety_models
 
-from pydantic import BaseModel, validator
+from pydantic import BaseModel, field_validator
 
 
-class MetaReferenceShieldType(Enum):
-    llama_guard = "llama_guard"
-    code_scanner_guard = "code_scanner_guard"
-    injection_shield = "injection_shield"
-    jailbreak_shield = "jailbreak_shield"
+class PromptGuardType(Enum):
+    injection = "injection"
+    jailbreak = "jailbreak"
 
 
 class LlamaGuardShieldConfig(BaseModel):
     model: str = "Llama-Guard-3-1B"
     excluded_categories: List[str] = []
-    disable_input_check: bool = False
-    disable_output_check: bool = False
 
-    @validator("model")
+    @field_validator("model")
     @classmethod
     def validate_model(cls, model: str) -> str:
         permitted_models = [

llama_stack/providers/impls/meta_reference/safety/llama_guard.py

@@ -113,8 +113,6 @@ class LlamaGuardShield(ShieldBase):
         model: str,
         inference_api: Inference,
         excluded_categories: List[str] = None,
-        disable_input_check: bool = False,
-        disable_output_check: bool = False,
         on_violation_action: OnViolationAction = OnViolationAction.RAISE,
     ):
         super().__init__(on_violation_action)
@@ -132,8 +130,6 @@ class LlamaGuardShield(ShieldBase):
         self.model = model
         self.inference_api = inference_api
         self.excluded_categories = excluded_categories
-        self.disable_input_check = disable_input_check
-        self.disable_output_check = disable_output_check
 
     def check_unsafe_response(self, response: str) -> Optional[str]:
         match = re.match(r"^unsafe\n(.*)$", response)
@@ -180,12 +176,6 @@ class LlamaGuardShield(ShieldBase):
 
     async def run(self, messages: List[Message]) -> ShieldResponse:
         messages = self.validate_messages(messages)
-        if self.disable_input_check and messages[-1].role == Role.user.value:
-            return ShieldResponse(is_violation=False)
-        elif self.disable_output_check and messages[-1].role == Role.assistant.value:
-            return ShieldResponse(
-                is_violation=False,
-            )
 
         if self.model == CoreModelId.llama_guard_3_11b_vision.value:
             shield_input_message = self.build_vision_shield_input(messages)

llama_stack/providers/impls/meta_reference/safety/safety.py

@@ -10,39 +10,50 @@ from llama_stack.distribution.utils.model_utils import model_local_dir
 from llama_stack.apis.inference import *  # noqa: F403
 from llama_stack.apis.safety import *  # noqa: F403
 from llama_models.llama3.api.datatypes import *  # noqa: F403
-from llama_stack.distribution.datatypes import Api, RoutableProvider
+from llama_stack.distribution.datatypes import Api
 
-from llama_stack.providers.impls.meta_reference.safety.shields.base import (
-    OnViolationAction,
-)
+from llama_stack.providers.datatypes import ShieldsProtocolPrivate
 
-from .config import MetaReferenceShieldType, SafetyConfig
+from .base import OnViolationAction, ShieldBase
+from .config import SafetyConfig
+from .llama_guard import LlamaGuardShield
+from .prompt_guard import InjectionShield, JailbreakShield, PromptGuardShield
 
-from .shields import CodeScannerShield, LlamaGuardShield, ShieldBase
 
 PROMPT_GUARD_MODEL = "Prompt-Guard-86M"
 
 
-class MetaReferenceSafetyImpl(Safety, RoutableProvider):
+class MetaReferenceSafetyImpl(Safety, ShieldsProtocolPrivate):
     def __init__(self, config: SafetyConfig, deps) -> None:
         self.config = config
         self.inference_api = deps[Api.inference]
 
+        self.available_shields = []
+        if config.llama_guard_shield:
+            self.available_shields.append(ShieldType.llama_guard.value)
+        if config.enable_prompt_guard:
+            self.available_shields.append(ShieldType.prompt_guard.value)
+
     async def initialize(self) -> None:
         if self.config.enable_prompt_guard:
-            from .shields import PromptGuardShield
-
             model_dir = model_local_dir(PROMPT_GUARD_MODEL)
             _ = PromptGuardShield.instance(model_dir)
 
     async def shutdown(self) -> None:
         pass
 
-    async def validate_routing_keys(self, routing_keys: List[str]) -> None:
-        available_shields = [v.value for v in MetaReferenceShieldType]
-        for key in routing_keys:
-            if key not in available_shields:
-                raise ValueError(f"Unknown safety shield type: {key}")
+    async def register_shield(self, shield: ShieldDef) -> None:
+        raise ValueError("Registering dynamic shields is not supported")
+
+    async def list_shields(self) -> List[ShieldDef]:
+        return [
+            ShieldDef(
+                identifier=shield_type,
+                type=shield_type,
+                params={},
+            )
+            for shield_type in self.available_shields
+        ]
 
     async def run_shield(
         self,
@@ -50,10 +61,11 @@ class MetaReferenceSafetyImpl(Safety, RoutableProvider):
         messages: List[Message],
         params: Dict[str, Any] = None,
     ) -> RunShieldResponse:
-        available_shields = [v.value for v in MetaReferenceShieldType]
-        assert shield_type in available_shields, f"Unknown shield {shield_type}"
+        shield_def = await self.shield_store.get_shield(shield_type)
+        if not shield_def:
+            raise ValueError(f"Unknown shield {shield_type}")
 
-        shield = self.get_shield_impl(MetaReferenceShieldType(shield_type))
+        shield = self.get_shield_impl(shield_def)
 
         messages = messages.copy()
         # some shields like llama-guard require the first message to be a user message
@@ -79,32 +91,22 @@ class MetaReferenceSafetyImpl(Safety, RoutableProvider):
 
         return RunShieldResponse(violation=violation)
 
-    def get_shield_impl(self, typ: MetaReferenceShieldType) -> ShieldBase:
-        cfg = self.config
-        if typ == MetaReferenceShieldType.llama_guard:
-            cfg = cfg.llama_guard_shield
-            assert (
-                cfg is not None
-            ), "Cannot use LlamaGuardShield since not present in config"
-
+    def get_shield_impl(self, shield: ShieldDef) -> ShieldBase:
+        if shield.type == ShieldType.llama_guard.value:
+            cfg = self.config.llama_guard_shield
             return LlamaGuardShield(
                 model=cfg.model,
                 inference_api=self.inference_api,
                 excluded_categories=cfg.excluded_categories,
-                disable_input_check=cfg.disable_input_check,
-                disable_output_check=cfg.disable_output_check,
             )
-        elif typ == MetaReferenceShieldType.jailbreak_shield:
-            from .shields import JailbreakShield
-
-            model_dir = model_local_dir(PROMPT_GUARD_MODEL)
-            return JailbreakShield.instance(model_dir)
-        elif typ == MetaReferenceShieldType.injection_shield:
-            from .shields import InjectionShield
-
+        elif shield.type == ShieldType.prompt_guard.value:
             model_dir = model_local_dir(PROMPT_GUARD_MODEL)
+            subtype = shield.params.get("prompt_guard_type", "injection")
+            if subtype == "injection":
                 return InjectionShield.instance(model_dir)
-        elif typ == MetaReferenceShieldType.code_scanner_guard:
-            return CodeScannerShield.instance()
+            elif subtype == "jailbreak":
+                return JailbreakShield.instance(model_dir)
             else:
-            raise ValueError(f"Unknown shield type: {typ}")
+                raise ValueError(f"Unknown prompt guard type: {subtype}")
+        else:
+            raise ValueError(f"Unknown shield type: {shield.type}")

llama_stack/providers/impls/meta_reference/safety/shields/__init__.py

@@ -1,33 +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.
-
-# supress warnings and spew of logs from hugging face
-import transformers
-
-from .base import (  # noqa: F401
-    DummyShield,
-    OnViolationAction,
-    ShieldBase,
-    ShieldResponse,
-    TextShield,
-)
-from .code_scanner import CodeScannerShield  # noqa: F401
-from .llama_guard import LlamaGuardShield  # noqa: F401
-from .prompt_guard import (  # noqa: F401
-    InjectionShield,
-    JailbreakShield,
-    PromptGuardShield,
-)
-
-transformers.logging.set_verbosity_error()
-
-import os
-
-os.environ["TOKENIZERS_PARALLELISM"] = "false"
-
-import warnings
-
-warnings.filterwarnings("ignore")

llama_stack/providers/impls/meta_reference/safety/shields/code_scanner.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 termcolor import cprint
-
-from .base import ShieldResponse, TextShield
-
-
-class CodeScannerShield(TextShield):
-    async def run_impl(self, text: str) -> ShieldResponse:
-        from codeshield.cs import CodeShield
-
-        cprint(f"Running CodeScannerShield on {text[50:]}", color="magenta")
-        result = await CodeShield.scan_code(text)
-        if result.is_insecure:
-            return ShieldResponse(
-                is_violation=True,
-                violation_type=",".join(
-                    [issue.pattern_id for issue in result.issues_found]
-                ),
-                violation_return_message="Sorry, I found security concerns in the code.",
-            )
-        else:
-            return ShieldResponse(is_violation=False)

llama_stack/providers/impls/vllm/vllm.py

@@ -10,39 +10,25 @@ import uuid
 from typing import Any
 
 from llama_models.llama3.api.chat_format import ChatFormat
-from llama_models.llama3.api.datatypes import (
-    CompletionMessage,
-    InterleavedTextMedia,
-    Message,
-    StopReason,
-    ToolChoice,
-    ToolDefinition,
-    ToolPromptFormat,
-)
+from llama_models.llama3.api.datatypes import *  # noqa: F403
 from llama_models.llama3.api.tokenizer import Tokenizer
 
 from vllm.engine.arg_utils import AsyncEngineArgs
 from vllm.engine.async_llm_engine import AsyncLLMEngine
 from vllm.sampling_params import SamplingParams
 
-from llama_stack.apis.inference import ChatCompletionRequest, Inference
+from llama_stack.apis.inference import *  # noqa: F403
 
-from llama_stack.apis.inference.inference import (
-    ChatCompletionResponse,
-    ChatCompletionResponseEvent,
-    ChatCompletionResponseEventType,
-    ChatCompletionResponseStreamChunk,
-    CompletionResponse,
-    CompletionResponseStreamChunk,
-    EmbeddingsResponse,
-    LogProbConfig,
-    ToolCallDelta,
-    ToolCallParseStatus,
+from llama_stack.providers.utils.inference.model_registry import ModelRegistryHelper
+from llama_stack.providers.utils.inference.openai_compat import (
+    OpenAICompatCompletionChoice,
+    OpenAICompatCompletionResponse,
+    process_chat_completion_response,
+    process_chat_completion_stream_response,
 )
-from llama_stack.providers.utils.inference.augment_messages import (
-    augment_messages_for_tools,
+from llama_stack.providers.utils.inference.prompt_adapter import (
+    chat_completion_request_to_prompt,
 )
-from llama_stack.providers.utils.inference.routable import RoutableProviderForModels
 
 from .config import VLLMConfig
 
@@ -72,10 +58,10 @@ def _vllm_sampling_params(sampling_params: Any) -> SamplingParams:
     if sampling_params.repetition_penalty > 0:
         kwargs["repetition_penalty"] = sampling_params.repetition_penalty
 
-    return SamplingParams().from_optional(**kwargs)
+    return SamplingParams(**kwargs)
 
 
-class VLLMInferenceImpl(Inference, RoutableProviderForModels):
+class VLLMInferenceImpl(ModelRegistryHelper, Inference):
     """Inference implementation for vLLM."""
 
     HF_MODEL_MAPPINGS = {
@@ -109,7 +95,7 @@ class VLLMInferenceImpl(Inference, RoutableProviderForModels):
 
     def __init__(self, config: VLLMConfig):
         Inference.__init__(self)
-        RoutableProviderForModels.__init__(
+        ModelRegistryHelper.__init__(
             self,
             stack_to_provider_models_map=self.HF_MODEL_MAPPINGS,
         )
@@ -148,7 +134,7 @@ class VLLMInferenceImpl(Inference, RoutableProviderForModels):
         if self.engine:
             self.engine.shutdown_background_loop()
 
-    async def completion(
+    def completion(
         self,
         model: str,
         content: InterleavedTextMedia,
@@ -157,17 +143,16 @@ class VLLMInferenceImpl(Inference, RoutableProviderForModels):
         logprobs: LogProbConfig | None = None,
     ) -> CompletionResponse | CompletionResponseStreamChunk:
         log.info("vLLM completion")
-        messages = [Message(role="user", content=content)]
-        async for result in self.chat_completion(
+        messages = [UserMessage(content=content)]
+        return self.chat_completion(
             model=model,
             messages=messages,
             sampling_params=sampling_params,
             stream=stream,
             logprobs=logprobs,
-        ):
-            yield result
+        )
 
-    async def chat_completion(
+    def chat_completion(
         self,
         model: str,
         messages: list[Message],
@@ -194,159 +179,59 @@ class VLLMInferenceImpl(Inference, RoutableProviderForModels):
         )
 
         log.info("Sampling params: %s", sampling_params)
-        vllm_sampling_params = _vllm_sampling_params(sampling_params)
-
-        messages = augment_messages_for_tools(request)
-        log.info("Augmented messages: %s", messages)
-        prompt = "".join([str(message.content) for message in messages])
-
         request_id = _random_uuid()
+
+        prompt = chat_completion_request_to_prompt(request, self.formatter)
+        vllm_sampling_params = _vllm_sampling_params(request.sampling_params)
         results_generator = self.engine.generate(
             prompt, vllm_sampling_params, request_id
         )
-
-        if not stream:
-            # Non-streaming case
-            final_output = None
-            stop_reason = None
-            async for request_output in results_generator:
-                final_output = request_output
-                if stop_reason is None and request_output.outputs:
-                    reason = request_output.outputs[-1].stop_reason
-                    if reason == "stop":
-                        stop_reason = StopReason.end_of_turn
-                    elif reason == "length":
-                        stop_reason = StopReason.out_of_tokens
-
-            if not stop_reason:
-                stop_reason = StopReason.end_of_message
-
-            if final_output:
-                response = "".join([output.text for output in final_output.outputs])
-                yield ChatCompletionResponse(
-                    completion_message=CompletionMessage(
-                        content=response,
-                        stop_reason=stop_reason,
-                    ),
-                    logprobs=None,
-                )
+        if stream:
+            return self._stream_chat_completion(request, results_generator)
         else:
-            # Streaming case
-            yield ChatCompletionResponseStreamChunk(
-                event=ChatCompletionResponseEvent(
-                    event_type=ChatCompletionResponseEventType.start,
-                    delta="",
-                )
-            )
+            return self._nonstream_chat_completion(request, results_generator)
 
-            buffer = ""
-            last_chunk = ""
-            ipython = False
-            stop_reason = None
+    async def _nonstream_chat_completion(
+        self, request: ChatCompletionRequest, results_generator: AsyncGenerator
+    ) -> ChatCompletionResponse:
+        outputs = [o async for o in results_generator]
+        final_output = outputs[-1]
 
+        assert final_output is not None
+        outputs = final_output.outputs
+        finish_reason = outputs[-1].stop_reason
+        choice = OpenAICompatCompletionChoice(
+            finish_reason=finish_reason,
+            text="".join([output.text for output in outputs]),
+        )
+        response = OpenAICompatCompletionResponse(
+            choices=[choice],
+        )
+        return process_chat_completion_response(request, response, self.formatter)
+
+    async def _stream_chat_completion(
+        self, request: ChatCompletionRequest, results_generator: AsyncGenerator
+    ) -> AsyncGenerator:
+        async def _generate_and_convert_to_openai_compat():
             async for chunk in results_generator:
                 if not chunk.outputs:
                     log.warning("Empty chunk received")
                     continue
 
-                if chunk.outputs[-1].stop_reason:
-                    reason = chunk.outputs[-1].stop_reason
-                    if stop_reason is None and reason == "stop":
-                        stop_reason = StopReason.end_of_turn
-                    elif stop_reason is None and reason == "length":
-                        stop_reason = StopReason.out_of_tokens
-                    break
-
                 text = "".join([output.text for output in chunk.outputs])
-
-                # check if its a tool call ( aka starts with <|python_tag|> )
-                if not ipython and text.startswith("<|python_tag|>"):
-                    ipython = True
-                    yield ChatCompletionResponseStreamChunk(
-                        event=ChatCompletionResponseEvent(
-                            event_type=ChatCompletionResponseEventType.progress,
-                            delta=ToolCallDelta(
-                                content="",
-                                parse_status=ToolCallParseStatus.started,
-                            ),
+                choice = OpenAICompatCompletionChoice(
+                    finish_reason=chunk.outputs[-1].stop_reason,
+                    text=text,
                 )
-                    )
-                    buffer += text
-                    continue
-
-                if ipython:
-                    if text == "<|eot_id|>":
-                        stop_reason = StopReason.end_of_turn
-                        text = ""
-                        continue
-                    elif text == "<|eom_id|>":
-                        stop_reason = StopReason.end_of_message
-                        text = ""
-                        continue
-
-                    buffer += text
-                    delta = ToolCallDelta(
-                        content=text,
-                        parse_status=ToolCallParseStatus.in_progress,
+                yield OpenAICompatCompletionResponse(
+                    choices=[choice],
                 )
 
-                    yield ChatCompletionResponseStreamChunk(
-                        event=ChatCompletionResponseEvent(
-                            event_type=ChatCompletionResponseEventType.progress,
-                            delta=delta,
-                            stop_reason=stop_reason,
-                        )
-                    )
-                else:
-                    last_chunk_len = len(last_chunk)
-                    last_chunk = text
-                    yield ChatCompletionResponseStreamChunk(
-                        event=ChatCompletionResponseEvent(
-                            event_type=ChatCompletionResponseEventType.progress,
-                            delta=text[last_chunk_len:],
-                            stop_reason=stop_reason,
-                        )
-                    )
-
-            if not stop_reason:
-                stop_reason = StopReason.end_of_message
-
-            # parse tool calls and report errors
-            message = self.formatter.decode_assistant_message_from_content(
-                buffer, stop_reason
-            )
-            parsed_tool_calls = len(message.tool_calls) > 0
-            if ipython and not parsed_tool_calls:
-                yield ChatCompletionResponseStreamChunk(
-                    event=ChatCompletionResponseEvent(
-                        event_type=ChatCompletionResponseEventType.progress,
-                        delta=ToolCallDelta(
-                            content="",
-                            parse_status=ToolCallParseStatus.failure,
-                        ),
-                        stop_reason=stop_reason,
-                    )
-                )
-
-            for tool_call in message.tool_calls:
-                yield ChatCompletionResponseStreamChunk(
-                    event=ChatCompletionResponseEvent(
-                        event_type=ChatCompletionResponseEventType.progress,
-                        delta=ToolCallDelta(
-                            content=tool_call,
-                            parse_status=ToolCallParseStatus.success,
-                        ),
-                        stop_reason=stop_reason,
-                    )
-                )
-
-            yield ChatCompletionResponseStreamChunk(
-                event=ChatCompletionResponseEvent(
-                    event_type=ChatCompletionResponseEventType.complete,
-                    delta="",
-                    stop_reason=stop_reason,
-                )
-            )
+        stream = _generate_and_convert_to_openai_compat()
+        async for chunk in process_chat_completion_stream_response(
+            request, stream, self.formatter
+        ):
+            yield chunk
 
     async def embeddings(
         self, model: str, contents: list[InterleavedTextMedia]

llama_stack/providers/registry/agents.py

@@ -28,6 +28,7 @@ def available_providers() -> List[ProviderSpec]:
                 Api.inference,
                 Api.safety,
                 Api.memory,
+                Api.memory_banks,
             ],
         ),
         remote_provider_spec(

llama_stack/providers/registry/inference.py

@@ -14,6 +14,21 @@ def available_providers() -> List[ProviderSpec]:
         InlineProviderSpec(
             api=Api.inference,
             provider_type="meta-reference",
+            pip_packages=[
+                "accelerate",
+                "blobfile",
+                "fairscale",
+                "torch",
+                "torchvision",
+                "transformers",
+                "zmq",
+            ],
+            module="llama_stack.providers.impls.meta_reference.inference",
+            config_class="llama_stack.providers.impls.meta_reference.inference.MetaReferenceInferenceConfig",
+        ),
+        InlineProviderSpec(
+            api=Api.inference,
+            provider_type="meta-reference-quantized",
             pip_packages=[
                 "accelerate",
                 "blobfile",
@@ -25,7 +40,7 @@ def available_providers() -> List[ProviderSpec]:
                 "zmq",
             ],
             module="llama_stack.providers.impls.meta_reference.inference",
-            config_class="llama_stack.providers.impls.meta_reference.inference.MetaReferenceImplConfig",
+            config_class="llama_stack.providers.impls.meta_reference.inference.MetaReferenceQuantizedInferenceConfig",
         ),
         remote_provider_spec(
             api=Api.inference,

llama_stack/providers/registry/memory.py

@@ -62,6 +62,7 @@ def available_providers() -> List[ProviderSpec]:
                 adapter_type="weaviate",
                 pip_packages=EMBEDDING_DEPS + ["weaviate-client"],
                 module="llama_stack.providers.adapters.memory.weaviate",
+                config_class="llama_stack.providers.adapters.memory.weaviate.WeaviateConfig",
                 provider_data_validator="llama_stack.providers.adapters.memory.weaviate.WeaviateRequestProviderData",
             ),
         ),

llama_stack/providers/registry/safety.py

@@ -21,7 +21,6 @@ def available_providers() -> List[ProviderSpec]:
             api=Api.safety,
             provider_type="meta-reference",
             pip_packages=[
-                "codeshield",
                 "transformers",
                 "torch --index-url https://download.pytorch.org/whl/cpu",
             ],
@@ -61,4 +60,14 @@ def available_providers() -> List[ProviderSpec]:
                 provider_data_validator="llama_stack.providers.adapters.safety.together.TogetherProviderDataValidator",
             ),
         ),
+        InlineProviderSpec(
+            api=Api.safety,
+            provider_type="meta-reference/codeshield",
+            pip_packages=[
+                "codeshield",
+            ],
+            module="llama_stack.providers.impls.meta_reference.codeshield",
+            config_class="llama_stack.providers.impls.meta_reference.codeshield.CodeShieldConfig",
+            api_dependencies=[],
+        ),
     ]

llama_stack/providers/tests/__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/tests/agents/__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/tests/agents/provider_config_example.yaml

@@ -0,0 +1,34 @@
+providers:
+  inference:
+  - provider_id: together
+    provider_type: remote::together
+    config: {}
+  - provider_id: tgi
+    provider_type: remote::tgi
+    config:
+      url: http://127.0.0.1:7001
+#   - provider_id: meta-reference
+#     provider_type: meta-reference
+#     config:
+#       model: Llama-Guard-3-1B
+#   - provider_id: remote
+#     provider_type: remote
+#     config:
+#       host: localhost
+#       port: 7010
+  safety:
+    - provider_id: together
+      provider_type: remote::together
+      config: {}
+  memory:
+    - provider_id: faiss
+      provider_type: meta-reference
+      config: {}
+  agents:
+    - provider_id: meta-reference
+      provider_type: meta-reference
+      config:
+        persistence_store:
+          namespace: null
+          type: sqlite
+          db_path: /Users/ashwin/.llama/runtime/kvstore.db

llama_stack/providers/tests/agents/test_agents.py

@@ -0,0 +1,210 @@
+# 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
+
+import pytest
+import pytest_asyncio
+
+from llama_stack.apis.agents import *  # noqa: F403
+from llama_stack.providers.tests.resolver import resolve_impls_for_test
+from llama_stack.providers.datatypes import *  # noqa: F403
+
+from dotenv import load_dotenv
+
+# How to run this test:
+#
+# 1. Ensure you have a conda environment with the right dependencies installed.
+#    This includes `pytest` and `pytest-asyncio`.
+#
+# 2. Copy and modify the provider_config_example.yaml depending on the provider you are testing.
+#
+# 3. Run:
+#
+# ```bash
+# PROVIDER_ID=<your_provider> \
+#   PROVIDER_CONFIG=provider_config.yaml \
+#   pytest -s llama_stack/providers/tests/agents/test_agents.py \
+#   --tb=short --disable-warnings
+# ```
+
+load_dotenv()
+
+
+@pytest_asyncio.fixture(scope="session")
+async def agents_settings():
+    impls = await resolve_impls_for_test(
+        Api.agents, deps=[Api.inference, Api.memory, Api.safety]
+    )
+
+    return {
+        "impl": impls[Api.agents],
+        "memory_impl": impls[Api.memory],
+        "common_params": {
+            "model": "Llama3.1-8B-Instruct",
+            "instructions": "You are a helpful assistant.",
+        },
+    }
+
+
+@pytest.fixture
+def sample_messages():
+    return [
+        UserMessage(content="What's the weather like today?"),
+    ]
+
+
+@pytest.fixture
+def search_query_messages():
+    return [
+        UserMessage(content="What are the latest developments in quantum computing?"),
+    ]
+
+
+@pytest.mark.asyncio
+async def test_create_agent_turn(agents_settings, sample_messages):
+    agents_impl = agents_settings["impl"]
+
+    # First, create an agent
+    agent_config = AgentConfig(
+        model=agents_settings["common_params"]["model"],
+        instructions=agents_settings["common_params"]["instructions"],
+        enable_session_persistence=True,
+        sampling_params=SamplingParams(temperature=0.7, top_p=0.95),
+        input_shields=[],
+        output_shields=[],
+        tools=[],
+        max_infer_iters=5,
+    )
+
+    create_response = await agents_impl.create_agent(agent_config)
+    agent_id = create_response.agent_id
+
+    # Create a session
+    session_create_response = await agents_impl.create_agent_session(
+        agent_id, "Test Session"
+    )
+    session_id = session_create_response.session_id
+
+    # Create and execute a turn
+    turn_request = dict(
+        agent_id=agent_id,
+        session_id=session_id,
+        messages=sample_messages,
+        stream=True,
+    )
+
+    turn_response = [
+        chunk async for chunk in agents_impl.create_agent_turn(**turn_request)
+    ]
+
+    assert len(turn_response) > 0
+    assert all(
+        isinstance(chunk, AgentTurnResponseStreamChunk) for chunk in turn_response
+    )
+
+    # Check for expected event types
+    event_types = [chunk.event.payload.event_type for chunk in turn_response]
+    assert AgentTurnResponseEventType.turn_start.value in event_types
+    assert AgentTurnResponseEventType.step_start.value in event_types
+    assert AgentTurnResponseEventType.step_complete.value in event_types
+    assert AgentTurnResponseEventType.turn_complete.value in event_types
+
+    # Check the final turn complete event
+    final_event = turn_response[-1].event.payload
+    assert isinstance(final_event, AgentTurnResponseTurnCompletePayload)
+    assert isinstance(final_event.turn, Turn)
+    assert final_event.turn.session_id == session_id
+    assert final_event.turn.input_messages == sample_messages
+    assert isinstance(final_event.turn.output_message, CompletionMessage)
+    assert len(final_event.turn.output_message.content) > 0
+
+
+@pytest.mark.asyncio
+async def test_create_agent_turn_with_brave_search(
+    agents_settings, search_query_messages
+):
+    agents_impl = agents_settings["impl"]
+
+    if "BRAVE_SEARCH_API_KEY" not in os.environ:
+        pytest.skip("BRAVE_SEARCH_API_KEY not set, skipping test")
+
+    # Create an agent with Brave search tool
+    agent_config = AgentConfig(
+        model=agents_settings["common_params"]["model"],
+        instructions=agents_settings["common_params"]["instructions"],
+        enable_session_persistence=True,
+        sampling_params=SamplingParams(temperature=0.7, top_p=0.95),
+        input_shields=[],
+        output_shields=[],
+        tools=[
+            SearchToolDefinition(
+                type=AgentTool.brave_search.value,
+                api_key=os.environ["BRAVE_SEARCH_API_KEY"],
+                engine=SearchEngineType.brave,
+            )
+        ],
+        tool_choice=ToolChoice.auto,
+        max_infer_iters=5,
+    )
+
+    create_response = await agents_impl.create_agent(agent_config)
+    agent_id = create_response.agent_id
+
+    # Create a session
+    session_create_response = await agents_impl.create_agent_session(
+        agent_id, "Test Session with Brave Search"
+    )
+    session_id = session_create_response.session_id
+
+    # Create and execute a turn
+    turn_request = dict(
+        agent_id=agent_id,
+        session_id=session_id,
+        messages=search_query_messages,
+        stream=True,
+    )
+
+    turn_response = [
+        chunk async for chunk in agents_impl.create_agent_turn(**turn_request)
+    ]
+
+    assert len(turn_response) > 0
+    assert all(
+        isinstance(chunk, AgentTurnResponseStreamChunk) for chunk in turn_response
+    )
+
+    # Check for expected event types
+    event_types = [chunk.event.payload.event_type for chunk in turn_response]
+    assert AgentTurnResponseEventType.turn_start.value in event_types
+    assert AgentTurnResponseEventType.step_start.value in event_types
+    assert AgentTurnResponseEventType.step_complete.value in event_types
+    assert AgentTurnResponseEventType.turn_complete.value in event_types
+
+    # Check for tool execution events
+    tool_execution_events = [
+        chunk
+        for chunk in turn_response
+        if isinstance(chunk.event.payload, AgentTurnResponseStepCompletePayload)
+        and chunk.event.payload.step_details.step_type == StepType.tool_execution.value
+    ]
+    assert len(tool_execution_events) > 0, "No tool execution events found"
+
+    # Check the tool execution details
+    tool_execution = tool_execution_events[0].event.payload.step_details
+    assert isinstance(tool_execution, ToolExecutionStep)
+    assert len(tool_execution.tool_calls) > 0
+    assert tool_execution.tool_calls[0].tool_name == BuiltinTool.brave_search
+    assert len(tool_execution.tool_responses) > 0
+
+    # Check the final turn complete event
+    final_event = turn_response[-1].event.payload
+    assert isinstance(final_event, AgentTurnResponseTurnCompletePayload)
+    assert isinstance(final_event.turn, Turn)
+    assert final_event.turn.session_id == session_id
+    assert final_event.turn.input_messages == search_query_messages
+    assert isinstance(final_event.turn.output_message, CompletionMessage)
+    assert len(final_event.turn.output_message.content) > 0

llama_stack/providers/tests/inference/__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/tests/inference/provider_config_example.yaml

@@ -0,0 +1,24 @@
+providers:
+  - provider_id: test-ollama
+    provider_type: remote::ollama
+    config:
+      host: localhost
+      port: 11434
+  - provider_id: test-tgi
+    provider_type: remote::tgi
+    config:
+      url: http://localhost:7001
+  - provider_id: test-remote
+    provider_type: remote
+    config:
+      host: localhost
+      port: 7002
+  - provider_id: test-together
+    provider_type: remote::together
+    config: {}
+# if a provider needs private keys from the client, they use the
+# "get_request_provider_data" function (see distribution/request_headers.py)
+# this is a place to provide such data.
+provider_data:
+  "test-together":
+    together_api_key: 0xdeadbeefputrealapikeyhere

llama_stack/providers/tests/inference/test_inference.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 itertools
+
+import pytest
+import pytest_asyncio
+
+from llama_models.llama3.api.datatypes import *  # noqa: F403
+from llama_stack.apis.inference import *  # noqa: F403
+
+from llama_stack.distribution.datatypes import *  # noqa: F403
+from llama_stack.providers.tests.resolver import resolve_impls_for_test
+
+# How to run this test:
+#
+# 1. Ensure you have a conda with the right dependencies installed. This is a bit tricky
+#    since it depends on the provider you are testing. On top of that you need
+#    `pytest` and `pytest-asyncio` installed.
+#
+# 2. Copy and modify the provider_config_example.yaml depending on the provider you are testing.
+#
+# 3. Run:
+#
+# ```bash
+# PROVIDER_ID=<your_provider> \
+#   PROVIDER_CONFIG=provider_config.yaml \
+#   pytest -s llama_stack/providers/tests/inference/test_inference.py \
+#   --tb=short --disable-warnings
+# ```
+
+
+def group_chunks(response):
+    return {
+        event_type: list(group)
+        for event_type, group in itertools.groupby(
+            response, key=lambda chunk: chunk.event.event_type
+        )
+    }
+
+
+Llama_8B = "Llama3.1-8B-Instruct"
+Llama_3B = "Llama3.2-3B-Instruct"
+
+
+def get_expected_stop_reason(model: str):
+    return StopReason.end_of_message if "Llama3.1" in model else StopReason.end_of_turn
+
+
+# This is going to create multiple Stack impls without tearing down the previous one
+# Fix that!
+@pytest_asyncio.fixture(
+    scope="session",
+    params=[
+        {"model": Llama_8B},
+        {"model": Llama_3B},
+    ],
+    ids=lambda d: d["model"],
+)
+async def inference_settings(request):
+    model = request.param["model"]
+    impls = await resolve_impls_for_test(
+        Api.inference,
+    )
+
+    return {
+        "impl": impls[Api.inference],
+        "models_impl": impls[Api.models],
+        "common_params": {
+            "model": model,
+            "tool_choice": ToolChoice.auto,
+            "tool_prompt_format": (
+                ToolPromptFormat.json
+                if "Llama3.1" in model
+                else ToolPromptFormat.python_list
+            ),
+        },
+    }
+
+
+@pytest.fixture
+def sample_messages():
+    return [
+        SystemMessage(content="You are a helpful assistant."),
+        UserMessage(content="What's the weather like today?"),
+    ]
+
+
+@pytest.fixture
+def sample_tool_definition():
+    return ToolDefinition(
+        tool_name="get_weather",
+        description="Get the current weather",
+        parameters={
+            "location": ToolParamDefinition(
+                param_type="string",
+                description="The city and state, e.g. San Francisco, CA",
+            ),
+        },
+    )
+
+
+@pytest.mark.asyncio
+async def test_model_list(inference_settings):
+    params = inference_settings["common_params"]
+    models_impl = inference_settings["models_impl"]
+    response = await models_impl.list_models()
+    assert isinstance(response, list)
+    assert len(response) >= 1
+    assert all(isinstance(model, ModelDefWithProvider) for model in response)
+
+    model_def = None
+    for model in response:
+        if model.identifier == params["model"]:
+            model_def = model
+            break
+
+    assert model_def is not None
+    assert model_def.identifier == params["model"]
+
+
+@pytest.mark.asyncio
+async def test_chat_completion_non_streaming(inference_settings, sample_messages):
+    inference_impl = inference_settings["impl"]
+    response = await inference_impl.chat_completion(
+        messages=sample_messages,
+        stream=False,
+        **inference_settings["common_params"],
+    )
+
+    assert isinstance(response, ChatCompletionResponse)
+    assert response.completion_message.role == "assistant"
+    assert isinstance(response.completion_message.content, str)
+    assert len(response.completion_message.content) > 0
+
+
+@pytest.mark.asyncio
+async def test_chat_completion_streaming(inference_settings, sample_messages):
+    inference_impl = inference_settings["impl"]
+    response = [
+        r
+        async for r in inference_impl.chat_completion(
+            messages=sample_messages,
+            stream=True,
+            **inference_settings["common_params"],
+        )
+    ]
+
+    assert len(response) > 0
+    assert all(
+        isinstance(chunk, ChatCompletionResponseStreamChunk) for chunk in response
+    )
+    grouped = group_chunks(response)
+    assert len(grouped[ChatCompletionResponseEventType.start]) == 1
+    assert len(grouped[ChatCompletionResponseEventType.progress]) > 0
+    assert len(grouped[ChatCompletionResponseEventType.complete]) == 1
+
+    end = grouped[ChatCompletionResponseEventType.complete][0]
+    assert end.event.stop_reason == StopReason.end_of_turn
+
+
+@pytest.mark.asyncio
+async def test_chat_completion_with_tool_calling(
+    inference_settings,
+    sample_messages,
+    sample_tool_definition,
+):
+    inference_impl = inference_settings["impl"]
+    messages = sample_messages + [
+        UserMessage(
+            content="What's the weather like in San Francisco?",
+        )
+    ]
+
+    response = await inference_impl.chat_completion(
+        messages=messages,
+        tools=[sample_tool_definition],
+        stream=False,
+        **inference_settings["common_params"],
+    )
+
+    assert isinstance(response, ChatCompletionResponse)
+
+    message = response.completion_message
+
+    # This is not supported in most providers :/ they don't return eom_id / eot_id
+    # stop_reason = get_expected_stop_reason(inference_settings["common_params"]["model"])
+    # assert message.stop_reason == stop_reason
+    assert message.tool_calls is not None
+    assert len(message.tool_calls) > 0
+
+    call = message.tool_calls[0]
+    assert call.tool_name == "get_weather"
+    assert "location" in call.arguments
+    assert "San Francisco" in call.arguments["location"]
+
+
+@pytest.mark.asyncio
+async def test_chat_completion_with_tool_calling_streaming(
+    inference_settings,
+    sample_messages,
+    sample_tool_definition,
+):
+    inference_impl = inference_settings["impl"]
+    messages = sample_messages + [
+        UserMessage(
+            content="What's the weather like in San Francisco?",
+        )
+    ]
+
+    response = [
+        r
+        async for r in inference_impl.chat_completion(
+            messages=messages,
+            tools=[sample_tool_definition],
+            stream=True,
+            **inference_settings["common_params"],
+        )
+    ]
+
+    assert len(response) > 0
+    assert all(
+        isinstance(chunk, ChatCompletionResponseStreamChunk) for chunk in response
+    )
+    grouped = group_chunks(response)
+    assert len(grouped[ChatCompletionResponseEventType.start]) == 1
+    assert len(grouped[ChatCompletionResponseEventType.progress]) > 0
+    assert len(grouped[ChatCompletionResponseEventType.complete]) == 1
+
+    # This is not supported in most providers :/ they don't return eom_id / eot_id
+    # expected_stop_reason = get_expected_stop_reason(
+    #     inference_settings["common_params"]["model"]
+    # )
+    # end = grouped[ChatCompletionResponseEventType.complete][0]
+    # assert end.event.stop_reason == expected_stop_reason
+
+    model = inference_settings["common_params"]["model"]
+    if "Llama3.1" in model:
+        assert all(
+            isinstance(chunk.event.delta, ToolCallDelta)
+            for chunk in grouped[ChatCompletionResponseEventType.progress]
+        )
+        first = grouped[ChatCompletionResponseEventType.progress][0]
+        assert first.event.delta.parse_status == ToolCallParseStatus.started
+
+    last = grouped[ChatCompletionResponseEventType.progress][-1]
+    # assert last.event.stop_reason == expected_stop_reason
+    assert last.event.delta.parse_status == ToolCallParseStatus.success
+    assert isinstance(last.event.delta.content, ToolCall)
+
+    call = last.event.delta.content
+    assert call.tool_name == "get_weather"
+    assert "location" in call.arguments
+    assert "San Francisco" in call.arguments["location"]

llama_stack/providers/tests/inference/test_prompt_adapter.py

@@ -8,7 +8,7 @@ import unittest
 
 from llama_models.llama3.api import *  # noqa: F403
 from llama_stack.inference.api import *  # noqa: F403
-from llama_stack.inference.augment_messages import augment_messages_for_tools
+from llama_stack.inference.prompt_adapter import chat_completion_request_to_messages
 
 MODEL = "Llama3.1-8B-Instruct"
 
@@ -22,7 +22,7 @@ class PrepareMessagesTests(unittest.IsolatedAsyncioTestCase):
                 UserMessage(content=content),
             ],
         )
-        messages = augment_messages_for_tools(request)
+        messages = chat_completion_request_to_messages(request)
         self.assertEqual(len(messages), 2)
         self.assertEqual(messages[-1].content, content)
         self.assertTrue("Cutting Knowledge Date: December 2023" in messages[0].content)
@@ -39,7 +39,7 @@ class PrepareMessagesTests(unittest.IsolatedAsyncioTestCase):
                 ToolDefinition(tool_name=BuiltinTool.brave_search),
             ],
         )
-        messages = augment_messages_for_tools(request)
+        messages = chat_completion_request_to_messages(request)
         self.assertEqual(len(messages), 2)
         self.assertEqual(messages[-1].content, content)
         self.assertTrue("Cutting Knowledge Date: December 2023" in messages[0].content)
@@ -67,7 +67,7 @@ class PrepareMessagesTests(unittest.IsolatedAsyncioTestCase):
             ],
             tool_prompt_format=ToolPromptFormat.json,
         )
-        messages = augment_messages_for_tools(request)
+        messages = chat_completion_request_to_messages(request)
         self.assertEqual(len(messages), 3)
         self.assertTrue("Environment: ipython" in messages[0].content)
 
@@ -97,7 +97,7 @@ class PrepareMessagesTests(unittest.IsolatedAsyncioTestCase):
                 ),
             ],
         )
-        messages = augment_messages_for_tools(request)
+        messages = chat_completion_request_to_messages(request)
         self.assertEqual(len(messages), 3)
 
         self.assertTrue("Environment: ipython" in messages[0].content)
@@ -119,7 +119,7 @@ class PrepareMessagesTests(unittest.IsolatedAsyncioTestCase):
                 ToolDefinition(tool_name=BuiltinTool.code_interpreter),
             ],
         )
-        messages = augment_messages_for_tools(request)
+        messages = chat_completion_request_to_messages(request)
         self.assertEqual(len(messages), 2, messages)
         self.assertTrue(messages[0].content.endswith(system_prompt))
 

llama_stack/providers/tests/memory/__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/tests/memory/provider_config_example.yaml

@@ -0,0 +1,24 @@
+providers:
+  - provider_id: test-faiss
+    provider_type: meta-reference
+    config: {}
+  - provider_id: test-chroma
+    provider_type: remote::chroma
+    config:
+      host: localhost
+      port: 6001
+  - provider_id: test-remote
+    provider_type: remote
+    config:
+      host: localhost
+      port: 7002
+  - provider_id: test-weaviate
+    provider_type: remote::weaviate
+    config: {}
+# if a provider needs private keys from the client, they use the
+# "get_request_provider_data" function (see distribution/request_headers.py)
+# this is a place to provide such data.
+provider_data:
+  "test-weaviate":
+    weaviate_api_key: 0xdeadbeefputrealapikeyhere
+    weaviate_cluster_url: http://foobarbaz

llama_stack/providers/tests/memory/test_memory.py

@@ -0,0 +1,136 @@
+# 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
+
+import pytest
+import pytest_asyncio
+
+from llama_stack.apis.memory import *  # noqa: F403
+from llama_stack.distribution.datatypes import *  # noqa: F403
+from llama_stack.providers.tests.resolver import resolve_impls_for_test
+
+# How to run this test:
+#
+# 1. Ensure you have a conda with the right dependencies installed. This is a bit tricky
+#    since it depends on the provider you are testing. On top of that you need
+#    `pytest` and `pytest-asyncio` installed.
+#
+# 2. Copy and modify the provider_config_example.yaml depending on the provider you are testing.
+#
+# 3. Run:
+#
+# ```bash
+# PROVIDER_ID=<your_provider> \
+#   PROVIDER_CONFIG=provider_config.yaml \
+#   pytest -s llama_stack/providers/tests/memory/test_memory.py \
+#   --tb=short --disable-warnings
+# ```
+
+
+@pytest_asyncio.fixture(scope="session")
+async def memory_settings():
+    impls = await resolve_impls_for_test(
+        Api.memory,
+    )
+    return {
+        "memory_impl": impls[Api.memory],
+        "memory_banks_impl": impls[Api.memory_banks],
+    }
+
+
+@pytest.fixture
+def sample_documents():
+    return [
+        MemoryBankDocument(
+            document_id="doc1",
+            content="Python is a high-level programming language.",
+            metadata={"category": "programming", "difficulty": "beginner"},
+        ),
+        MemoryBankDocument(
+            document_id="doc2",
+            content="Machine learning is a subset of artificial intelligence.",
+            metadata={"category": "AI", "difficulty": "advanced"},
+        ),
+        MemoryBankDocument(
+            document_id="doc3",
+            content="Data structures are fundamental to computer science.",
+            metadata={"category": "computer science", "difficulty": "intermediate"},
+        ),
+        MemoryBankDocument(
+            document_id="doc4",
+            content="Neural networks are inspired by biological neural networks.",
+            metadata={"category": "AI", "difficulty": "advanced"},
+        ),
+    ]
+
+
+async def register_memory_bank(banks_impl: MemoryBanks):
+    bank = VectorMemoryBankDef(
+        identifier="test_bank",
+        embedding_model="all-MiniLM-L6-v2",
+        chunk_size_in_tokens=512,
+        overlap_size_in_tokens=64,
+        provider_id=os.environ["PROVIDER_ID"],
+    )
+
+    await banks_impl.register_memory_bank(bank)
+
+
+@pytest.mark.asyncio
+async def test_banks_list(memory_settings):
+    # NOTE: this needs you to ensure that you are starting from a clean state
+    # but so far we don't have an unregister API unfortunately, so be careful
+    banks_impl = memory_settings["memory_banks_impl"]
+    response = await banks_impl.list_memory_banks()
+    assert isinstance(response, list)
+    assert len(response) == 0
+
+
+@pytest.mark.asyncio
+async def test_query_documents(memory_settings, sample_documents):
+    memory_impl = memory_settings["memory_impl"]
+    banks_impl = memory_settings["memory_banks_impl"]
+
+    with pytest.raises(ValueError):
+        await memory_impl.insert_documents("test_bank", sample_documents)
+
+    await register_memory_bank(banks_impl)
+    await memory_impl.insert_documents("test_bank", sample_documents)
+
+    query1 = "programming language"
+    response1 = await memory_impl.query_documents("test_bank", query1)
+    assert_valid_response(response1)
+    assert any("Python" in chunk.content for chunk in response1.chunks)
+
+    # Test case 3: Query with semantic similarity
+    query3 = "AI and brain-inspired computing"
+    response3 = await memory_impl.query_documents("test_bank", query3)
+    assert_valid_response(response3)
+    assert any("neural networks" in chunk.content.lower() for chunk in response3.chunks)
+
+    # Test case 4: Query with limit on number of results
+    query4 = "computer"
+    params4 = {"max_chunks": 2}
+    response4 = await memory_impl.query_documents("test_bank", query4, params4)
+    assert_valid_response(response4)
+    assert len(response4.chunks) <= 2
+
+    # Test case 5: Query with threshold on similarity score
+    query5 = "quantum computing"  # Not directly related to any document
+    params5 = {"score_threshold": 0.5}
+    response5 = await memory_impl.query_documents("test_bank", query5, params5)
+    assert_valid_response(response5)
+    assert all(score >= 0.5 for score in response5.scores)
+
+
+def assert_valid_response(response: QueryDocumentsResponse):
+    assert isinstance(response, QueryDocumentsResponse)
+    assert len(response.chunks) > 0
+    assert len(response.scores) > 0
+    assert len(response.chunks) == len(response.scores)
+    for chunk in response.chunks:
+        assert isinstance(chunk.content, str)
+        assert chunk.document_id is not None

llama_stack/providers/tests/resolver.py

@@ -0,0 +1,100 @@
+# 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
+from datetime import datetime
+from typing import Any, Dict, List
+
+import yaml
+
+from llama_stack.distribution.datatypes import *  # noqa: F403
+from llama_stack.distribution.configure import parse_and_maybe_upgrade_config
+from llama_stack.distribution.request_headers import set_request_provider_data
+from llama_stack.distribution.resolver import resolve_impls_with_routing
+
+
+async def resolve_impls_for_test(api: Api, deps: List[Api] = None):
+    if "PROVIDER_CONFIG" not in os.environ:
+        raise ValueError(
+            "You must set PROVIDER_CONFIG to a YAML file containing provider config"
+        )
+
+    with open(os.environ["PROVIDER_CONFIG"], "r") as f:
+        config_dict = yaml.safe_load(f)
+
+    providers = read_providers(api, config_dict)
+
+    chosen = choose_providers(providers, api, deps)
+    run_config = dict(
+        built_at=datetime.now(),
+        image_name="test-fixture",
+        apis=[api] + (deps or []),
+        providers=chosen,
+    )
+    run_config = parse_and_maybe_upgrade_config(run_config)
+    impls = await resolve_impls_with_routing(run_config)
+
+    if "provider_data" in config_dict:
+        provider_id = chosen[api.value][0].provider_id
+        provider_data = config_dict["provider_data"].get(provider_id, {})
+        if provider_data:
+            set_request_provider_data(
+                {"X-LlamaStack-ProviderData": json.dumps(provider_data)}
+            )
+
+    return impls
+
+
+def read_providers(api: Api, config_dict: Dict[str, Any]) -> Dict[str, Any]:
+    if "providers" not in config_dict:
+        raise ValueError("Config file should contain a `providers` key")
+
+    providers = config_dict["providers"]
+    if isinstance(providers, dict):
+        return providers
+    elif isinstance(providers, list):
+        return {
+            api.value: providers,
+        }
+    else:
+        raise ValueError(
+            "Config file should contain a list of providers or dict(api to providers)"
+        )
+
+
+def choose_providers(
+    providers: Dict[str, Any], api: Api, deps: List[Api] = None
+) -> Dict[str, Provider]:
+    chosen = {}
+    if api.value not in providers:
+        raise ValueError(f"No providers found for `{api}`?")
+    chosen[api.value] = [pick_provider(api, providers[api.value], "PROVIDER_ID")]
+
+    for dep in deps or []:
+        if dep.value not in providers:
+            raise ValueError(f"No providers specified for `{dep}` in config?")
+        chosen[dep.value] = [Provider(**x) for x in providers[dep.value]]
+
+    return chosen
+
+
+def pick_provider(api: Api, providers: List[Any], key: str) -> Provider:
+    providers_by_id = {x["provider_id"]: x for x in providers}
+    if len(providers_by_id) == 0:
+        raise ValueError(f"No providers found for `{api}` in config file")
+
+    if key in os.environ:
+        provider_id = os.environ[key]
+        if provider_id not in providers_by_id:
+            raise ValueError(f"Provider ID {provider_id} not found in config file")
+        provider = providers_by_id[provider_id]
+    else:
+        provider = list(providers_by_id.values())[0]
+        provider_id = provider["provider_id"]
+        print(f"No provider ID specified, picking first `{provider_id}`")
+
+    return Provider(**provider)