llama-stack-api 0.4.2__py3-none-any.whl → 0.4.4__py3-none-any.whl

llama_stack_api/schema_utils.py

@@ -0,0 +1,208 @@
+# 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
+    required_scope: 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,
+    required_scope: 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 required_scope: Required scope for this endpoint (e.g., 'monitoring.viewer').
+    :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,
+            required_scope=required_scope,
+            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

llama_stack_api/scoring.py

@@ -0,0 +1,93 @@
+# 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, Protocol, runtime_checkable
+
+from pydantic import BaseModel
+
+from llama_stack_api.schema_utils import json_schema_type, webmethod
+from llama_stack_api.scoring_functions import ScoringFn, ScoringFnParams
+from llama_stack_api.version import LLAMA_STACK_API_V1
+
+# mapping of metric to value
+ScoringResultRow = dict[str, Any]
+
+
+@json_schema_type
+class ScoringResult(BaseModel):
+    """
+    A scoring result for a single row.
+
+    :param score_rows: The scoring result for each row. Each row is a map of column name to value.
+    :param aggregated_results: Map of metric name to aggregated value
+    """
+
+    score_rows: list[ScoringResultRow]
+    # aggregated metrics to value
+    aggregated_results: dict[str, Any]
+
+
+@json_schema_type
+class ScoreBatchResponse(BaseModel):
+    """Response from batch scoring operations on datasets.
+
+    :param dataset_id: (Optional) The identifier of the dataset that was scored
+    :param results: A map of scoring function name to ScoringResult
+    """
+
+    dataset_id: str | None = None
+    results: dict[str, ScoringResult]
+
+
+@json_schema_type
+class ScoreResponse(BaseModel):
+    """
+    The response from scoring.
+
+    :param results: A map of scoring function name to ScoringResult.
+    """
+
+    # each key in the dict is a scoring function name
+    results: dict[str, ScoringResult]
+
+
+class ScoringFunctionStore(Protocol):
+    def get_scoring_function(self, scoring_fn_id: str) -> ScoringFn: ...
+
+
+@runtime_checkable
+class Scoring(Protocol):
+    scoring_function_store: ScoringFunctionStore
+
+    @webmethod(route="/scoring/score-batch", method="POST", level=LLAMA_STACK_API_V1)
+    async def score_batch(
+        self,
+        dataset_id: str,
+        scoring_functions: dict[str, ScoringFnParams | None],
+        save_results_dataset: bool = False,
+    ) -> ScoreBatchResponse:
+        """Score a batch of rows.
+
+        :param dataset_id: The ID of the dataset to score.
+        :param scoring_functions: The scoring functions to use for the scoring.
+        :param save_results_dataset: Whether to save the results to a dataset.
+        :returns: A ScoreBatchResponse.
+        """
+        ...
+
+    @webmethod(route="/scoring/score", method="POST", level=LLAMA_STACK_API_V1)
+    async def score(
+        self,
+        input_rows: list[dict[str, Any]],
+        scoring_functions: dict[str, ScoringFnParams | None],
+    ) -> ScoreResponse:
+        """Score a list of rows.
+
+        :param input_rows: The rows to score.
+        :param scoring_functions: The scoring functions to use for the scoring.
+        :returns: A ScoreResponse object containing rows and aggregated results.
+        """
+        ...

llama_stack_api/scoring_functions.py

@@ -0,0 +1,211 @@
+# 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.
+
+# TODO: use enum.StrEnum when we drop support for python 3.10
+from enum import StrEnum
+from typing import (
+    Annotated,
+    Any,
+    Literal,
+    Protocol,
+    runtime_checkable,
+)
+
+from pydantic import BaseModel, Field
+
+from llama_stack_api.common.type_system import ParamType
+from llama_stack_api.resource import Resource, ResourceType
+from llama_stack_api.schema_utils import json_schema_type, register_schema, webmethod
+from llama_stack_api.version import LLAMA_STACK_API_V1
+
+
+# Perhaps more structure can be imposed on these functions. Maybe they could be associated
+# with standard metrics so they can be rolled up?
+@json_schema_type
+class ScoringFnParamsType(StrEnum):
+    """Types of scoring function parameter configurations.
+    :cvar llm_as_judge: Use an LLM model to evaluate and score responses
+    :cvar regex_parser: Use regex patterns to extract and score specific parts of responses
+    :cvar basic: Basic scoring with simple aggregation functions
+    """
+
+    llm_as_judge = "llm_as_judge"
+    regex_parser = "regex_parser"
+    basic = "basic"
+
+
+@json_schema_type
+class AggregationFunctionType(StrEnum):
+    """Types of aggregation functions for scoring results.
+    :cvar average: Calculate the arithmetic mean of scores
+    :cvar weighted_average: Calculate a weighted average of scores
+    :cvar median: Calculate the median value of scores
+    :cvar categorical_count: Count occurrences of categorical values
+    :cvar accuracy: Calculate accuracy as the proportion of correct answers
+    """
+
+    average = "average"
+    weighted_average = "weighted_average"
+    median = "median"
+    categorical_count = "categorical_count"
+    accuracy = "accuracy"
+
+
+@json_schema_type
+class LLMAsJudgeScoringFnParams(BaseModel):
+    """Parameters for LLM-as-judge scoring function configuration.
+    :param type: The type of scoring function parameters, always llm_as_judge
+    :param judge_model: Identifier of the LLM model to use as a judge for scoring
+    :param prompt_template: (Optional) Custom prompt template for the judge model
+    :param judge_score_regexes: Regexes to extract the answer from generated response
+    :param aggregation_functions: Aggregation functions to apply to the scores of each row
+    """
+
+    type: Literal[ScoringFnParamsType.llm_as_judge] = ScoringFnParamsType.llm_as_judge
+    judge_model: str
+    prompt_template: str | None = None
+    judge_score_regexes: list[str] = Field(
+        description="Regexes to extract the answer from generated response",
+        default_factory=lambda: [],
+    )
+    aggregation_functions: list[AggregationFunctionType] = Field(
+        description="Aggregation functions to apply to the scores of each row",
+        default_factory=lambda: [],
+    )
+
+
+@json_schema_type
+class RegexParserScoringFnParams(BaseModel):
+    """Parameters for regex parser scoring function configuration.
+    :param type: The type of scoring function parameters, always regex_parser
+    :param parsing_regexes: Regex to extract the answer from generated response
+    :param aggregation_functions: Aggregation functions to apply to the scores of each row
+    """
+
+    type: Literal[ScoringFnParamsType.regex_parser] = ScoringFnParamsType.regex_parser
+    parsing_regexes: list[str] = Field(
+        description="Regex to extract the answer from generated response",
+        default_factory=lambda: [],
+    )
+    aggregation_functions: list[AggregationFunctionType] = Field(
+        description="Aggregation functions to apply to the scores of each row",
+        default_factory=lambda: [],
+    )
+
+
+@json_schema_type
+class BasicScoringFnParams(BaseModel):
+    """Parameters for basic scoring function configuration.
+    :param type: The type of scoring function parameters, always basic
+    :param aggregation_functions: Aggregation functions to apply to the scores of each row
+    """
+
+    type: Literal[ScoringFnParamsType.basic] = ScoringFnParamsType.basic
+    aggregation_functions: list[AggregationFunctionType] = Field(
+        description="Aggregation functions to apply to the scores of each row",
+        default_factory=list,
+    )
+
+
+ScoringFnParams = Annotated[
+    LLMAsJudgeScoringFnParams | RegexParserScoringFnParams | BasicScoringFnParams,
+    Field(discriminator="type"),
+]
+register_schema(ScoringFnParams, name="ScoringFnParams")
+
+
+class CommonScoringFnFields(BaseModel):
+    description: str | None = None
+    metadata: dict[str, Any] = Field(
+        default_factory=dict,
+        description="Any additional metadata for this definition",
+    )
+    return_type: ParamType = Field(
+        description="The return type of the deterministic function",
+    )
+    params: ScoringFnParams | None = Field(
+        description="The parameters for the scoring function for benchmark eval, these can be overridden for app eval",
+        default=None,
+    )
+
+
+@json_schema_type
+class ScoringFn(CommonScoringFnFields, Resource):
+    """A scoring function resource for evaluating model outputs.
+    :param type: The resource type, always scoring_function
+    """
+
+    type: Literal[ResourceType.scoring_function] = ResourceType.scoring_function
+
+    @property
+    def scoring_fn_id(self) -> str:
+        return self.identifier
+
+    @property
+    def provider_scoring_fn_id(self) -> str | None:
+        return self.provider_resource_id
+
+
+class ScoringFnInput(CommonScoringFnFields, BaseModel):
+    scoring_fn_id: str
+    provider_id: str | None = None
+    provider_scoring_fn_id: str | None = None
+
+
+@json_schema_type
+class ListScoringFunctionsResponse(BaseModel):
+    data: list[ScoringFn]
+
+
+@runtime_checkable
+class ScoringFunctions(Protocol):
+    @webmethod(route="/scoring-functions", method="GET", level=LLAMA_STACK_API_V1)
+    async def list_scoring_functions(self) -> ListScoringFunctionsResponse:
+        """List all scoring functions.
+
+        :returns: A ListScoringFunctionsResponse.
+        """
+        ...
+
+    @webmethod(route="/scoring-functions/{scoring_fn_id:path}", method="GET", level=LLAMA_STACK_API_V1)
+    async def get_scoring_function(self, scoring_fn_id: str, /) -> ScoringFn:
+        """Get a scoring function by its ID.
+
+        :param scoring_fn_id: The ID of the scoring function to get.
+        :returns: A ScoringFn.
+        """
+        ...
+
+    @webmethod(route="/scoring-functions", method="POST", level=LLAMA_STACK_API_V1, deprecated=True)
+    async def register_scoring_function(
+        self,
+        scoring_fn_id: str,
+        description: str,
+        return_type: ParamType,
+        provider_scoring_fn_id: str | None = None,
+        provider_id: str | None = None,
+        params: ScoringFnParams | None = None,
+    ) -> None:
+        """Register a scoring function.
+
+        :param scoring_fn_id: The ID of the scoring function to register.
+        :param description: The description of the scoring function.
+        :param return_type: The return type of the scoring function.
+        :param provider_scoring_fn_id: The ID of the provider scoring function to use for the scoring function.
+        :param provider_id: The ID of the provider to use for the scoring function.
+        :param params: The parameters for the scoring function for benchmark eval, these can be overridden for app eval.
+        """
+        ...
+
+    @webmethod(
+        route="/scoring-functions/{scoring_fn_id:path}", method="DELETE", level=LLAMA_STACK_API_V1, deprecated=True
+    )
+    async def unregister_scoring_function(self, scoring_fn_id: str) -> None:
+        """Unregister a scoring function.
+
+        :param scoring_fn_id: The ID of the scoring function to unregister.
+        """
+        ...

llama_stack_api/shields.py

@@ -0,0 +1,93 @@
+# 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, Literal, Protocol, runtime_checkable
+
+from pydantic import BaseModel
+
+from llama_stack_api.resource import Resource, ResourceType
+from llama_stack_api.schema_utils import json_schema_type, webmethod
+from llama_stack_api.version import LLAMA_STACK_API_V1
+
+
+class CommonShieldFields(BaseModel):
+    params: dict[str, Any] | None = None
+
+
+@json_schema_type
+class Shield(CommonShieldFields, Resource):
+    """A safety shield resource that can be used to check content.
+
+    :param params: (Optional) Configuration parameters for the shield
+    :param type: The resource type, always shield
+    """
+
+    type: Literal[ResourceType.shield] = ResourceType.shield
+
+    @property
+    def shield_id(self) -> str:
+        return self.identifier
+
+    @property
+    def provider_shield_id(self) -> str | None:
+        return self.provider_resource_id
+
+
+class ShieldInput(CommonShieldFields):
+    shield_id: str
+    provider_id: str | None = None
+    provider_shield_id: str | None = None
+
+
+@json_schema_type
+class ListShieldsResponse(BaseModel):
+    data: list[Shield]
+
+
+@runtime_checkable
+class Shields(Protocol):
+    @webmethod(route="/shields", method="GET", level=LLAMA_STACK_API_V1)
+    async def list_shields(self) -> ListShieldsResponse:
+        """List all shields.
+
+        :returns: A ListShieldsResponse.
+        """
+        ...
+
+    @webmethod(route="/shields/{identifier:path}", method="GET", level=LLAMA_STACK_API_V1)
+    async def get_shield(self, identifier: str) -> Shield:
+        """Get a shield by its identifier.
+
+        :param identifier: The identifier of the shield to get.
+        :returns: A Shield.
+        """
+        ...
+
+    @webmethod(route="/shields", method="POST", level=LLAMA_STACK_API_V1, deprecated=True)
+    async def register_shield(
+        self,
+        shield_id: str,
+        provider_shield_id: str | None = None,
+        provider_id: str | None = None,
+        params: dict[str, Any] | None = None,
+    ) -> Shield:
+        """Register a shield.
+
+        :param shield_id: The identifier of the shield to register.
+        :param provider_shield_id: The identifier of the shield in the provider.
+        :param provider_id: The identifier of the provider.
+        :param params: The parameters of the shield.
+        :returns: A Shield.
+        """
+        ...
+
+    @webmethod(route="/shields/{identifier:path}", method="DELETE", level=LLAMA_STACK_API_V1, deprecated=True)
+    async def unregister_shield(self, identifier: str) -> None:
+        """Unregister a shield.
+
+        :param identifier: The identifier of the shield to unregister.
+        """
+        ...