llama-stack-api 0.4.3__py3-none-any.whl → 0.5.0rc1__py3-none-any.whl

llama_stack_api/router_utils.py

@@ -0,0 +1,160 @@
+# 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.
+
+"""Utilities for creating FastAPI routers with standard error responses.
+
+This module provides standard error response definitions for FastAPI routers.
+These responses use OpenAPI $ref references to component responses defined
+in the OpenAPI specification.
+"""
+
+import inspect
+from collections.abc import Callable
+from typing import Annotated, Any, TypeVar
+
+from fastapi import Path, Query
+from pydantic import BaseModel
+
+# OpenAPI extension key to mark routes that don't require authentication.
+# Use this in FastAPI route decorators: @router.get("/health", openapi_extra={PUBLIC_ROUTE_KEY: True})
+PUBLIC_ROUTE_KEY = "x-public"
+
+
+standard_responses: dict[int | str, dict[str, Any]] = {
+    400: {"$ref": "#/components/responses/BadRequest400"},
+    429: {"$ref": "#/components/responses/TooManyRequests429"},
+    500: {"$ref": "#/components/responses/InternalServerError500"},
+    "default": {"$ref": "#/components/responses/DefaultError"},
+}
+
+T = TypeVar("T", bound=BaseModel)
+
+
+def create_query_dependency[T: BaseModel](model_class: type[T]) -> Callable[..., T]:
+    """Create a FastAPI dependency function from a Pydantic model for query parameters.
+
+    FastAPI does not natively support using Pydantic models as query parameters
+    without a dependency function. Using a dependency function typically leads to
+    duplication: field types, default values, and descriptions must be repeated in
+    `Query(...)` annotations even though they already exist in the Pydantic model.
+
+    This function automatically generates a dependency function that extracts query parameters
+    from the request and constructs an instance of the Pydantic model. The descriptions and
+    defaults are automatically extracted from the model's Field definitions, making the model
+    the single source of truth.
+
+    Args:
+        model_class: The Pydantic model class to create a dependency for
+
+    Returns:
+        A dependency function that can be used with FastAPI's Depends()
+        ```
+    """
+    # Build function signature dynamically from model fields
+    annotations: dict[str, Any] = {}
+    defaults: dict[str, Any] = {}
+
+    for field_name, field_info in model_class.model_fields.items():
+        # Extract description from Field
+        description = field_info.description
+
+        # Create Query annotation with description from model
+        query_annotation = Query(description=description) if description else Query()
+
+        # Create Annotated type with Query
+        field_type = field_info.annotation
+        annotations[field_name] = Annotated[field_type, query_annotation]
+
+        # Set default value from model
+        if field_info.default is not inspect.Parameter.empty:
+            defaults[field_name] = field_info.default
+
+    # Create the dependency function dynamically
+    def dependency_func(**kwargs: Any) -> T:
+        return model_class(**kwargs)
+
+    # Set function signature
+    sig_params = []
+    for field_name, field_type in annotations.items():
+        default = defaults.get(field_name, inspect.Parameter.empty)
+        param = inspect.Parameter(
+            field_name,
+            inspect.Parameter.POSITIONAL_OR_KEYWORD,
+            default=default,
+            annotation=field_type,
+        )
+        sig_params.append(param)
+
+    # These attributes are set dynamically at runtime. While mypy can't verify them statically,
+    # they are standard Python function attributes that exist on all callable objects at runtime.
+    # Setting them allows FastAPI to properly introspect the function signature for dependency injection.
+    dependency_func.__signature__ = inspect.Signature(sig_params)  # type: ignore[attr-defined]
+    dependency_func.__annotations__ = annotations  # type: ignore[attr-defined]
+    dependency_func.__name__ = f"get_{model_class.__name__.lower()}_request"  # type: ignore[attr-defined]
+
+    return dependency_func
+
+
+def create_path_dependency[T: BaseModel](model_class: type[T]) -> Callable[..., T]:
+    """Create a FastAPI dependency function from a Pydantic model for path parameters.
+
+    FastAPI requires path parameters to be explicitly annotated with `Path()`. When using
+    a Pydantic model that contains path parameters, you typically need a dependency function
+    that extracts the path parameter and constructs the model. This leads to duplication:
+    the parameter name, type, and description must be repeated in `Path(...)` annotations
+    even though they already exist in the Pydantic model.
+
+    This function automatically generates a dependency function that extracts path parameters
+    from the request and constructs an instance of the Pydantic model. The descriptions are
+    automatically extracted from the model's Field definitions, making the model the single
+    source of truth.
+
+    Args:
+        model_class: The Pydantic model class to create a dependency for. The model should
+            have exactly one field that represents the path parameter.
+
+    Returns:
+        A dependency function that can be used with FastAPI's Depends()
+        ```
+    """
+    # Get the single field from the model (path parameter models typically have one field)
+    if len(model_class.model_fields) != 1:
+        raise ValueError(
+            f"Path parameter model {model_class.__name__} must have exactly one field, "
+            f"but has {len(model_class.model_fields)} fields"
+        )
+
+    field_name, field_info = next(iter(model_class.model_fields.items()))
+
+    # Extract description from Field
+    description = field_info.description
+
+    # Create Path annotation with description from model
+    path_annotation = Path(description=description) if description else Path()
+
+    # Create Annotated type with Path
+    field_type = field_info.annotation
+    annotations: dict[str, Any] = {field_name: Annotated[field_type, path_annotation]}
+
+    # Create the dependency function dynamically
+    def dependency_func(**kwargs: Any) -> T:
+        return model_class(**kwargs)
+
+    # Set function signature
+    param = inspect.Parameter(
+        field_name,
+        inspect.Parameter.POSITIONAL_OR_KEYWORD,
+        annotation=annotations[field_name],
+    )
+
+    # These attributes are set dynamically at runtime. While mypy can't verify them statically,
+    # they are standard Python function attributes that exist on all callable objects at runtime.
+    # Setting them allows FastAPI to properly introspect the function signature for dependency injection.
+    dependency_func.__signature__ = inspect.Signature([param])  # type: ignore[attr-defined]
+    dependency_func.__annotations__ = annotations  # type: ignore[attr-defined]
+    dependency_func.__name__ = f"get_{model_class.__name__.lower()}_request"  # type: ignore[attr-defined]
+
+    return dependency_func

llama_stack_api/safety/__init__.py

@@ -0,0 +1,37 @@
+# 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.
+
+"""Safety API protocol and models.
+
+This module contains the Safety protocol definition for content moderation and safety shields.
+Pydantic models are defined in llama_stack_api.safety.models.
+The FastAPI router is defined in llama_stack_api.safety.fastapi_routes.
+"""
+
+from . import fastapi_routes
+from .api import Safety
+from .datatypes import (
+    ModerationObject,
+    ModerationObjectResults,
+    RunShieldResponse,
+    SafetyViolation,
+    ShieldStore,
+    ViolationLevel,
+)
+from .models import RunModerationRequest, RunShieldRequest
+
+__all__ = [
+    "Safety",
+    "ShieldStore",
+    "ModerationObject",
+    "ModerationObjectResults",
+    "ViolationLevel",
+    "SafetyViolation",
+    "RunShieldResponse",
+    "RunShieldRequest",
+    "RunModerationRequest",
+    "fastapi_routes",
+]

llama_stack_api/safety/api.py

@@ -0,0 +1,29 @@
+# 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 Protocol, runtime_checkable
+
+from llama_stack_api.safety.datatypes import ModerationObject, RunShieldResponse, ShieldStore
+
+from .models import RunModerationRequest, RunShieldRequest
+
+
+@runtime_checkable
+class Safety(Protocol):
+    """Safety API for content moderation and safety shields.
+
+    OpenAI-compatible Moderations API with additional shield capabilities.
+    """
+
+    shield_store: ShieldStore
+
+    async def run_shield(self, request: RunShieldRequest) -> RunShieldResponse:
+        """Run a safety shield on messages."""
+        ...
+
+    async def run_moderation(self, request: RunModerationRequest) -> ModerationObject:
+        """Classify if inputs are potentially harmful."""
+        ...

llama_stack_api/safety/datatypes.py

@@ -0,0 +1,83 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the terms described in the LICENSE file in
+# the root directory of this source tree.
+
+from enum import Enum
+from typing import Any, Protocol
+
+from pydantic import BaseModel, Field
+
+from llama_stack_api.schema_utils import json_schema_type
+from llama_stack_api.shields import GetShieldRequest, Shield
+
+
+@json_schema_type
+class ModerationObjectResults(BaseModel):
+    """A moderation result object containing flagged status and category information."""
+
+    flagged: bool = Field(..., description="Whether any of the below categories are flagged")
+    categories: dict[str, bool] | None = Field(
+        None, description="A dictionary of the categories, and whether they are flagged or not"
+    )
+    category_applied_input_types: dict[str, list[str]] | None = Field(
+        None, description="A dictionary of the categories along with the input type(s) that the score applies to"
+    )
+    category_scores: dict[str, float] | None = Field(
+        None, description="A dictionary of the categories along with their scores as predicted by model"
+    )
+    user_message: str | None = Field(None, description="A message to convey to the user about the moderation result")
+    metadata: dict[str, Any] = Field(default_factory=dict, description="Additional metadata about the moderation")
+
+
+@json_schema_type
+class ModerationObject(BaseModel):
+    """A moderation object containing the results of content classification."""
+
+    id: str = Field(..., description="The unique identifier for the moderation request")
+    model: str = Field(..., description="The model used to generate the moderation results")
+    results: list[ModerationObjectResults] = Field(..., description="A list of moderation result objects")
+
+
+@json_schema_type
+class ViolationLevel(Enum):
+    """Severity level of a safety violation."""
+
+    INFO = "info"  # Informational level violation that does not require action
+    WARN = "warn"  # Warning level violation that suggests caution but allows continuation
+    ERROR = "error"  # Error level violation that requires blocking or intervention
+
+
+@json_schema_type
+class SafetyViolation(BaseModel):
+    """Details of a safety violation detected by content moderation."""
+
+    violation_level: ViolationLevel = Field(..., description="Severity level of the violation")
+    user_message: str | None = Field(None, description="Message to convey to the user about the violation")
+    metadata: dict[str, Any] = Field(
+        default_factory=dict, description="Additional metadata including specific violation codes"
+    )
+
+
+@json_schema_type
+class RunShieldResponse(BaseModel):
+    """Response from running a safety shield."""
+
+    violation: SafetyViolation | None = Field(None, description="Safety violation detected by the shield, if any")
+
+
+class ShieldStore(Protocol):
+    """Protocol for accessing shields."""
+
+    async def get_shield(self, request: GetShieldRequest) -> Shield: ...
+
+
+__all__ = [
+    "ModerationObjectResults",
+    "ModerationObject",
+    "ViolationLevel",
+    "SafetyViolation",
+    "RunShieldResponse",
+    "ShieldStore",
+]

llama_stack_api/safety/fastapi_routes.py

@@ -0,0 +1,55 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the terms described in the LICENSE file in
+# the root directory of this source tree.
+
+from typing import Annotated
+
+from fastapi import APIRouter, Body
+
+from llama_stack_api.router_utils import standard_responses
+from llama_stack_api.version import LLAMA_STACK_API_V1
+
+from .api import Safety
+from .datatypes import ModerationObject, RunShieldResponse
+from .models import RunModerationRequest, RunShieldRequest
+
+
+def create_router(impl: Safety) -> APIRouter:
+    """Create a FastAPI router for the Safety API."""
+    router = APIRouter(
+        prefix=f"/{LLAMA_STACK_API_V1}",
+        tags=["Safety"],
+        responses=standard_responses,
+    )
+
+    @router.post(
+        "/safety/run-shield",
+        response_model=RunShieldResponse,
+        summary="Run Shield",
+        description="Run a safety shield on messages to check for policy violations.",
+        responses={
+            200: {"description": "The shield response indicating any violations detected."},
+        },
+    )
+    async def run_shield(
+        request: Annotated[RunShieldRequest, Body(...)],
+    ) -> RunShieldResponse:
+        return await impl.run_shield(request)
+
+    @router.post(
+        "/moderations",
+        response_model=ModerationObject,
+        summary="Create Moderation",
+        description="Classifies if text inputs are potentially harmful. OpenAI-compatible endpoint.",
+        responses={
+            200: {"description": "The moderation results for the input."},
+        },
+    )
+    async def run_moderation(
+        request: Annotated[RunModerationRequest, Body(...)],
+    ) -> ModerationObject:
+        return await impl.run_moderation(request)
+
+    return router

llama_stack_api/safety/models.py

@@ -0,0 +1,38 @@
+# 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, Field
+
+from llama_stack_api.inference import OpenAIMessageParam
+from llama_stack_api.schema_utils import json_schema_type
+
+
+@json_schema_type
+class RunShieldRequest(BaseModel):
+    """Request model for running a safety shield."""
+
+    shield_id: str = Field(..., description="The identifier of the shield to run", min_length=1)
+    messages: list[OpenAIMessageParam] = Field(..., description="The messages to run the shield on")
+
+
+@json_schema_type
+class RunModerationRequest(BaseModel):
+    """Request model for running content moderation."""
+
+    input: str | list[str] = Field(
+        ...,
+        description="Input (or inputs) to classify. Can be a single string or an array of strings.",
+    )
+    model: str | None = Field(
+        None,
+        description="The content moderation model to use. If not specified, the default shield will be used.",
+    )
+
+
+__all__ = [
+    "RunShieldRequest",
+    "RunModerationRequest",
+]

llama_stack_api/schema_utils.py

@@ -0,0 +1,251 @@
+# 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 collections.abc import Callable, Iterable
+from dataclasses import dataclass
+from typing import Any, Literal, TypeVar
+
+
+class ExtraBodyField[T]:
+    """
+    Marker annotation for parameters that arrive via extra_body in the client SDK.
+
+    These parameters:
+    - Will NOT appear in the generated client SDK method signature
+    - WILL be documented in OpenAPI spec under x-llama-stack-extra-body-params
+    - MUST be passed via the extra_body parameter in client SDK calls
+    - WILL be available in server-side method signature with proper typing
+
+    Example:
+        ```python
+        async def create_openai_response(
+            self,
+            input: str,
+            model: str,
+            shields: Annotated[
+                list[str] | None, ExtraBodyField("List of shields to apply")
+            ] = None,
+        ) -> ResponseObject:
+            # shields is available here with proper typing
+            if shields:
+                print(f"Using shields: {shields}")
+        ```
+
+        Client usage:
+        ```python
+        client.responses.create(
+            input="hello", model="llama-3", extra_body={"shields": ["shield-1"]}
+        )
+        ```
+    """
+
+    def __init__(self, description: str | None = None):
+        self.description = description
+
+
+SchemaSource = Literal["json_schema_type", "registered_schema", "dynamic_schema"]
+
+
+@dataclass(frozen=True)
+class SchemaInfo:
+    """Metadata describing a schema entry exposed to OpenAPI generation."""
+
+    name: str
+    type: Any
+    source: SchemaSource
+
+
+_json_schema_types: dict[type, SchemaInfo] = {}
+
+
+def json_schema_type(cls):
+    """
+    Decorator to mark a Pydantic model for top-level component registration.
+
+    Models marked with this decorator will be registered as top-level components
+    in the OpenAPI schema, while unmarked models will be inlined.
+
+    This provides control over schema registration to avoid unnecessary indirection
+    for simple one-off types while keeping complex reusable types as components.
+    """
+    cls._llama_stack_schema_type = True
+    schema_name = getattr(cls, "__name__", f"Anonymous_{id(cls)}")
+    cls._llama_stack_schema_name = schema_name
+    _json_schema_types.setdefault(cls, SchemaInfo(name=schema_name, type=cls, source="json_schema_type"))
+    return cls
+
+
+# Global registries for schemas discoverable by the generator
+_registered_schemas: dict[Any, SchemaInfo] = {}
+_dynamic_schema_types: dict[type, SchemaInfo] = {}
+
+
+def register_schema(schema_type, name: str | None = None):
+    """
+    Register a schema type for top-level component registration.
+
+    This replicates the behavior of strong_typing's register_schema function.
+    It's used for union types and other complex types that should appear as
+    top-level components in the OpenAPI schema.
+
+    Args:
+        schema_type: The type to register (e.g., union types, Annotated types)
+        name: Optional name for the schema in the OpenAPI spec. If not provided,
+              uses the type's __name__ or a generated name.
+    """
+    if name is None:
+        name = getattr(schema_type, "__name__", f"Anonymous_{id(schema_type)}")
+
+    # Store the registration information in a global registry
+    # since union types don't allow setting attributes
+    _registered_schemas[schema_type] = SchemaInfo(name=name, type=schema_type, source="registered_schema")
+
+    return schema_type
+
+
+def get_registered_schema_info(schema_type: Any) -> SchemaInfo | None:
+    """Return the registration metadata for a schema type if present."""
+    return _registered_schemas.get(schema_type)
+
+
+def iter_registered_schema_types() -> Iterable[SchemaInfo]:
+    """Iterate over all explicitly registered schema entries."""
+    return tuple(_registered_schemas.values())
+
+
+def iter_json_schema_types() -> Iterable[type]:
+    """Iterate over all Pydantic models decorated with @json_schema_type."""
+    return tuple(info.type for info in _json_schema_types.values())
+
+
+def iter_dynamic_schema_types() -> Iterable[type]:
+    """Iterate over dynamic models registered at generation time."""
+    return tuple(info.type for info in _dynamic_schema_types.values())
+
+
+def register_dynamic_schema_type(schema_type: type, name: str | None = None) -> type:
+    """Register a dynamic model generated at runtime for schema inclusion."""
+    schema_name = name if name is not None else getattr(schema_type, "__name__", f"Anonymous_{id(schema_type)}")
+    _dynamic_schema_types[schema_type] = SchemaInfo(name=schema_name, type=schema_type, source="dynamic_schema")
+    return schema_type
+
+
+def clear_dynamic_schema_types() -> None:
+    """Clear dynamic schema registrations."""
+    _dynamic_schema_types.clear()
+
+
+@dataclass
+class WebMethod:
+    level: str | None = None
+    route: str | None = None
+    public: bool = False
+    request_examples: list[Any] | None = None
+    response_examples: list[Any] | None = None
+    method: str | None = None
+    raw_bytes_request_body: bool | None = False
+    # A descriptive name of the corresponding span created by tracing
+    descriptive_name: str | None = None
+    deprecated: bool | None = False
+    require_authentication: bool | None = True
+
+
+CallableT = TypeVar("CallableT", bound=Callable[..., Any])
+
+
+def webmethod(
+    route: str | None = None,
+    method: str | None = None,
+    level: str | None = None,
+    public: bool | None = False,
+    request_examples: list[Any] | None = None,
+    response_examples: list[Any] | None = None,
+    raw_bytes_request_body: bool | None = False,
+    descriptive_name: str | None = None,
+    deprecated: bool | None = False,
+    require_authentication: bool | None = True,
+) -> Callable[[CallableT], CallableT]:
+    """
+    Decorator that supplies additional metadata to an endpoint operation function.
+
+    :param route: The URL path pattern associated with this operation which path parameters are substituted into.
+    :param public: True if the operation can be invoked without prior authentication.
+    :param request_examples: Sample requests that the operation might take. Pass a list of objects, not JSON.
+    :param response_examples: Sample responses that the operation might produce. Pass a list of objects, not JSON.
+    :param require_authentication: Whether this endpoint requires authentication (default True).
+    """
+
+    def wrap(func: CallableT) -> CallableT:
+        webmethod_obj = WebMethod(
+            route=route,
+            method=method,
+            level=level,
+            public=public or False,
+            request_examples=request_examples,
+            response_examples=response_examples,
+            raw_bytes_request_body=raw_bytes_request_body,
+            descriptive_name=descriptive_name,
+            deprecated=deprecated,
+            require_authentication=require_authentication if require_authentication is not None else True,
+        )
+
+        # Store all webmethods in a list to support multiple decorators
+        if not hasattr(func, "__webmethods__"):
+            func.__webmethods__ = []  # type: ignore
+        func.__webmethods__.append(webmethod_obj)  # type: ignore
+
+        # Keep the last one as __webmethod__ for backwards compatibility
+        func.__webmethod__ = webmethod_obj  # type: ignore
+        return func
+
+    return wrap
+
+
+def remove_null_from_anyof(schema: dict, *, add_nullable: bool = False) -> None:
+    """Remove null type from anyOf and optionally add nullable flag.
+
+    Converts Pydantic's default OpenAPI 3.1 style:
+        anyOf: [{type: X, enum: [...]}, {type: null}]
+
+    To flattened format:
+        type: X
+        enum: [...]
+        nullable: true  # only if add_nullable=True
+
+    Args:
+        schema: The JSON schema dict to modify in-place
+        add_nullable: If True, adds 'nullable: true' when null was present.
+                      Use True for OpenAPI 3.0 compatibility with OpenAI's spec.
+    """
+    # Handle anyOf format: anyOf: [{type: string, enum: [...]}, {type: null}]
+    if "anyOf" in schema:
+        non_null = [s for s in schema["anyOf"] if s.get("type") != "null"]
+        has_null = len(non_null) < len(schema["anyOf"])
+
+        if len(non_null) == 1:
+            # Flatten to single type
+            only_schema = non_null[0]
+            schema.pop("anyOf")
+            schema.update(only_schema)
+            if has_null and add_nullable:
+                schema["nullable"] = True
+
+    # Handle OpenAPI 3.1 format: type: ['string', 'null']
+    elif isinstance(schema.get("type"), list) and "null" in schema["type"]:
+        has_null = "null" in schema["type"]
+        schema["type"].remove("null")
+        if len(schema["type"]) == 1:
+            schema["type"] = schema["type"][0]
+        if has_null and add_nullable:
+            schema["nullable"] = True
+
+
+def nullable_openai_style(schema: dict) -> None:
+    """Shorthand for remove_null_from_anyof with add_nullable=True.
+
+    Use this for fields that need OpenAPI 3.0 nullable style to match OpenAI's spec.
+    """
+    remove_null_from_anyof(schema, add_nullable=True)