API Updates (#73)

* API Keys passed from Client instead of distro configuration

* delete distribution registry

* Rename the "package" word away

* Introduce a "Router" layer for providers

Some providers need to be factorized and considered as thin routing
layers on top of other providers. Consider two examples:

- The inference API should be a routing layer over inference providers,
  routed using the "model" key
- The memory banks API is another instance where various memory bank
  types will be provided by independent providers (e.g., a vector store
  is served by Chroma while a keyvalue memory can be served by Redis or
  PGVector)

This commit introduces a generalized routing layer for this purpose.

* update `apis_to_serve`

* llama_toolchain -> llama_stack

* Codemod from llama_toolchain -> llama_stack

- added providers/registry
- cleaned up api/ subdirectories and moved impls away
- restructured api/api.py
- from llama_stack.apis.<api> import foo should work now
- update imports to do llama_stack.apis.<api>
- update many other imports
- added __init__, fixed some registry imports
- updated registry imports
- create_agentic_system -> create_agent
- AgenticSystem -> Agent

* Moved some stuff out of common/; re-generated OpenAPI spec

* llama-toolchain -> llama-stack (hyphens)

* add control plane API

* add redis adapter + sqlite provider

* move core -> distribution

* Some more toolchain -> stack changes

* small naming shenanigans

* Removing custom tool and agent utilities and moving them client side

* Move control plane to distribution server for now

* Remove control plane from API list

* no codeshield dependency randomly plzzzzz

* Add "fire" as a dependency

* add back event loggers

* stack configure fixes

* use brave instead of bing in the example client

* add init file so it gets packaged

* add init files so it gets packaged

* Update MANIFEST

* bug fix

---------

Co-authored-by: Hardik Shah <hjshah@fb.com>
Co-authored-by: Xi Yan <xiyan@meta.com>
Co-authored-by: Ashwin Bharambe <ashwin@meta.com>

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

@@ -0,0 +1,95 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the terms described in the LICENSE file in
+# the root directory of this source tree.
+
+import asyncio
+
+from llama_models.sku_list import resolve_model
+
+from llama_stack.distribution.utils.model_utils import model_local_dir
+from llama_stack.apis.safety import *  # noqa
+
+from .config import SafetyConfig
+from .shields import (
+    CodeScannerShield,
+    InjectionShield,
+    JailbreakShield,
+    LlamaGuardShield,
+    PromptGuardShield,
+    ShieldBase,
+    ThirdPartyShield,
+)
+
+
+def resolve_and_get_path(model_name: str) -> str:
+    model = resolve_model(model_name)
+    assert model is not None, f"Could not resolve model {model_name}"
+    model_dir = model_local_dir(model.descriptor())
+    return model_dir
+
+
+class MetaReferenceSafetyImpl(Safety):
+    def __init__(self, config: SafetyConfig) -> None:
+        self.config = config
+
+    async def initialize(self) -> None:
+        shield_cfg = self.config.llama_guard_shield
+        if shield_cfg is not None:
+            model_dir = resolve_and_get_path(shield_cfg.model)
+            _ = LlamaGuardShield.instance(
+                model_dir=model_dir,
+                excluded_categories=shield_cfg.excluded_categories,
+                disable_input_check=shield_cfg.disable_input_check,
+                disable_output_check=shield_cfg.disable_output_check,
+            )
+
+        shield_cfg = self.config.prompt_guard_shield
+        if shield_cfg is not None:
+            model_dir = resolve_and_get_path(shield_cfg.model)
+            _ = PromptGuardShield.instance(model_dir)
+
+    async def run_shields(
+        self,
+        messages: List[Message],
+        shields: List[ShieldDefinition],
+    ) -> RunShieldResponse:
+        shields = [shield_config_to_shield(c, self.config) for c in shields]
+
+        responses = await asyncio.gather(*[shield.run(messages) for shield in shields])
+
+        return RunShieldResponse(responses=responses)
+
+
+def shield_type_equals(a: ShieldType, b: ShieldType):
+    return a == b or a == b.value
+
+
+def shield_config_to_shield(
+    sc: ShieldDefinition, safety_config: SafetyConfig
+) -> ShieldBase:
+    if shield_type_equals(sc.shield_type, BuiltinShield.llama_guard):
+        assert (
+            safety_config.llama_guard_shield is not None
+        ), "Cannot use LlamaGuardShield since not present in config"
+        model_dir = resolve_and_get_path(safety_config.llama_guard_shield.model)
+        return LlamaGuardShield.instance(model_dir=model_dir)
+    elif shield_type_equals(sc.shield_type, BuiltinShield.jailbreak_shield):
+        assert (
+            safety_config.prompt_guard_shield is not None
+        ), "Cannot use Jailbreak Shield since Prompt Guard not present in config"
+        model_dir = resolve_and_get_path(safety_config.prompt_guard_shield.model)
+        return JailbreakShield.instance(model_dir)
+    elif shield_type_equals(sc.shield_type, BuiltinShield.injection_shield):
+        assert (
+            safety_config.prompt_guard_shield is not None
+        ), "Cannot use PromptGuardShield since not present in config"
+        model_dir = resolve_and_get_path(safety_config.prompt_guard_shield.model)
+        return InjectionShield.instance(model_dir)
+    elif shield_type_equals(sc.shield_type, BuiltinShield.code_scanner_guard):
+        return CodeScannerShield.instance()
+    elif shield_type_equals(sc.shield_type, BuiltinShield.third_party_shield):
+        return ThirdPartyShield.instance()
+    else:
+        raise ValueError(f"Unknown shield type: {sc.shield_type}")