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

llama_stack_api/providers/models.py

@@ -0,0 +1,24 @@
+# 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.
+
+"""Pydantic models for Providers API requests and responses.
+
+This module re-exports models from llama_stack_api.admin.models to ensure
+a single source of truth and avoid type conflicts.
+"""
+
+# Import and re-export shared models from admin
+from llama_stack_api.admin.models import (
+    InspectProviderRequest,
+    ListProvidersResponse,
+    ProviderInfo,
+)
+
+__all__ = [
+    "ProviderInfo",
+    "ListProvidersResponse",
+    "InspectProviderRequest",
+]

llama_stack_api/rag_tool.py

@@ -0,0 +1,168 @@
+# 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, StrEnum
+from typing import Annotated, Any, Literal
+
+from pydantic import BaseModel, Field, field_validator
+
+from llama_stack_api.common.content_types import URL, InterleavedContent
+
+
+class RRFRanker(BaseModel):
+    """
+    Reciprocal Rank Fusion (RRF) ranker configuration.
+
+    :param type: The type of ranker, always "rrf"
+    :param impact_factor: The impact factor for RRF scoring. Higher values give more weight to higher-ranked results.
+                         Must be greater than 0
+    """
+
+    type: Literal["rrf"] = "rrf"
+    impact_factor: float = Field(default=60.0, gt=0.0)  # default of 60 for optimal performance
+
+
+class WeightedRanker(BaseModel):
+    """
+    Weighted ranker configuration that combines vector and keyword scores.
+
+    :param type: The type of ranker, always "weighted"
+    :param alpha: Weight factor between 0 and 1.
+                 0 means only use keyword scores,
+                 1 means only use vector scores,
+                 values in between blend both scores.
+    """
+
+    type: Literal["weighted"] = "weighted"
+    alpha: float = Field(
+        default=0.5,
+        ge=0.0,
+        le=1.0,
+        description="Weight factor between 0 and 1. 0 means only keyword scores, 1 means only vector scores.",
+    )
+
+
+Ranker = Annotated[
+    RRFRanker | WeightedRanker,
+    Field(discriminator="type"),
+]
+
+
+class RAGDocument(BaseModel):
+    """
+    A document to be used for document ingestion in the RAG Tool.
+
+    :param document_id: The unique identifier for the document.
+    :param content: The content of the document.
+    :param mime_type: The MIME type of the document.
+    :param metadata: Additional metadata for the document.
+    """
+
+    document_id: str
+    content: InterleavedContent | URL
+    mime_type: str | None = None
+    metadata: dict[str, Any] = Field(default_factory=dict)
+
+
+class RAGQueryResult(BaseModel):
+    """Result of a RAG query containing retrieved content and metadata.
+
+    :param content: (Optional) The retrieved content from the query
+    :param metadata: Additional metadata about the query result
+    """
+
+    content: InterleavedContent | None = None
+    metadata: dict[str, Any] = Field(default_factory=dict)
+
+
+class RAGQueryGenerator(Enum):
+    """Types of query generators for RAG systems.
+
+    :cvar default: Default query generator using simple text processing
+    :cvar llm: LLM-based query generator for enhanced query understanding
+    :cvar custom: Custom query generator implementation
+    """
+
+    default = "default"
+    llm = "llm"
+    custom = "custom"
+
+
+class RAGSearchMode(StrEnum):
+    """
+    Search modes for RAG query retrieval:
+    - VECTOR: Uses vector similarity search for semantic matching
+    - KEYWORD: Uses keyword-based search for exact matching
+    - HYBRID: Combines both vector and keyword search for better results
+    """
+
+    VECTOR = "vector"
+    KEYWORD = "keyword"
+    HYBRID = "hybrid"
+
+
+class DefaultRAGQueryGeneratorConfig(BaseModel):
+    """Configuration for the default RAG query generator.
+
+    :param type: Type of query generator, always 'default'
+    :param separator: String separator used to join query terms
+    """
+
+    type: Literal["default"] = "default"
+    separator: str = " "
+
+
+class LLMRAGQueryGeneratorConfig(BaseModel):
+    """Configuration for the LLM-based RAG query generator.
+
+    :param type: Type of query generator, always 'llm'
+    :param model: Name of the language model to use for query generation
+    :param template: Template string for formatting the query generation prompt
+    """
+
+    type: Literal["llm"] = "llm"
+    model: str
+    template: str
+
+
+RAGQueryGeneratorConfig = Annotated[
+    DefaultRAGQueryGeneratorConfig | LLMRAGQueryGeneratorConfig,
+    Field(discriminator="type"),
+]
+
+
+class RAGQueryConfig(BaseModel):
+    """
+    Configuration for the RAG query generation.
+
+    :param query_generator_config: Configuration for the query generator.
+    :param max_tokens_in_context: Maximum number of tokens in the context.
+    :param max_chunks: Maximum number of chunks to retrieve.
+    :param chunk_template: Template for formatting each retrieved chunk in the context.
+        Available placeholders: {index} (1-based chunk ordinal), {chunk.content} (chunk content string), {metadata} (chunk metadata dict).
+        Default: "Result {index}\\nContent: {chunk.content}\\nMetadata: {metadata}\\n"
+    :param mode: Search mode for retrieval—either "vector", "keyword", or "hybrid". Default "vector".
+    :param ranker: Configuration for the ranker to use in hybrid search. Defaults to RRF ranker.
+    """
+
+    # This config defines how a query is generated using the messages
+    # for memory bank retrieval.
+    query_generator_config: RAGQueryGeneratorConfig = Field(default=DefaultRAGQueryGeneratorConfig())
+    max_tokens_in_context: int = 4096
+    max_chunks: int = 5
+    chunk_template: str = "Result {index}\nContent: {chunk.content}\nMetadata: {metadata}\n"
+    mode: RAGSearchMode | None = RAGSearchMode.VECTOR
+    ranker: Ranker | None = Field(default=None)  # Only used for hybrid mode
+
+    @field_validator("chunk_template")
+    def validate_chunk_template(cls, v: str) -> str:
+        if "{chunk.content}" not in v:
+            raise ValueError("chunk_template must contain {chunk.content}")
+        if "{index}" not in v:
+            raise ValueError("chunk_template must contain {index}")
+        if len(v) == 0:
+            raise ValueError("chunk_template must not be empty")
+        return v

llama_stack_api/resource.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.
+
+from enum import StrEnum
+
+from pydantic import BaseModel, Field
+
+
+class ResourceType(StrEnum):
+    model = "model"
+    shield = "shield"
+    vector_store = "vector_store"
+    dataset = "dataset"
+    scoring_function = "scoring_function"
+    benchmark = "benchmark"
+    tool = "tool"
+    tool_group = "tool_group"
+    prompt = "prompt"
+    connector = "connector"
+
+
+class Resource(BaseModel):
+    """Base class for all Llama Stack resources"""
+
+    identifier: str = Field(description="Unique identifier for this resource in llama stack")
+
+    provider_resource_id: str | None = Field(
+        default=None,
+        description="Unique identifier for this resource in the provider",
+    )
+
+    provider_id: str = Field(description="ID of the provider that owns this resource")
+
+    type: ResourceType = Field(description="Type of resource (e.g. 'model', 'shield', 'vector_store', etc.)")

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.py

@@ -0,0 +1,132 @@
+# 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, runtime_checkable
+
+from pydantic import BaseModel, Field
+
+from llama_stack_api.inference import OpenAIMessageParam
+from llama_stack_api.schema_utils import json_schema_type, webmethod
+from llama_stack_api.shields import Shield
+from llama_stack_api.version import LLAMA_STACK_API_V1
+
+
+@json_schema_type
+class ModerationObjectResults(BaseModel):
+    """A moderation object.
+    :param flagged: Whether any of the below categories are flagged.
+    :param categories: A list of the categories, and whether they are flagged or not.
+    :param category_applied_input_types: A list of the categories along with the input type(s) that the score applies to.
+    :param category_scores: A list of the categories along with their scores as predicted by model.
+    """
+
+    flagged: bool
+    categories: dict[str, bool] | None = None
+    category_applied_input_types: dict[str, list[str]] | None = None
+    category_scores: dict[str, float] | None = None
+    user_message: str | None = None
+    metadata: dict[str, Any] = Field(default_factory=dict)
+
+
+@json_schema_type
+class ModerationObject(BaseModel):
+    """A moderation object.
+    :param id: The unique identifier for the moderation request.
+    :param model: The model used to generate the moderation results.
+    :param results: A list of moderation objects
+    """
+
+    id: str
+    model: str
+    results: list[ModerationObjectResults]
+
+
+@json_schema_type
+class ViolationLevel(Enum):
+    """Severity level of a safety violation.
+
+    :cvar INFO: Informational level violation that does not require action
+    :cvar WARN: Warning level violation that suggests caution but allows continuation
+    :cvar ERROR: Error level violation that requires blocking or intervention
+    """
+
+    INFO = "info"
+    WARN = "warn"
+    ERROR = "error"
+
+
+@json_schema_type
+class SafetyViolation(BaseModel):
+    """Details of a safety violation detected by content moderation.
+
+    :param violation_level: Severity level of the violation
+    :param user_message: (Optional) Message to convey to the user about the violation
+    :param metadata: Additional metadata including specific violation codes for debugging and telemetry
+    """
+
+    violation_level: ViolationLevel
+
+    # what message should you convey to the user
+    user_message: str | None = None
+
+    # additional metadata (including specific violation codes) more for
+    # debugging, telemetry
+    metadata: dict[str, Any] = Field(default_factory=dict)
+
+
+@json_schema_type
+class RunShieldResponse(BaseModel):
+    """Response from running a safety shield.
+
+    :param violation: (Optional) Safety violation detected by the shield, if any
+    """
+
+    violation: SafetyViolation | None = None
+
+
+class ShieldStore(Protocol):
+    async def get_shield(self, identifier: str) -> Shield: ...
+
+
+@runtime_checkable
+class Safety(Protocol):
+    """Safety
+
+    OpenAI-compatible Moderations API.
+    """
+
+    shield_store: ShieldStore
+
+    @webmethod(route="/safety/run-shield", method="POST", level=LLAMA_STACK_API_V1)
+    async def run_shield(
+        self,
+        shield_id: str,
+        messages: list[OpenAIMessageParam],
+        params: dict[str, Any],
+    ) -> RunShieldResponse:
+        """Run shield.
+
+        Run a shield.
+
+        :param shield_id: The identifier of the shield to run.
+        :param messages: The messages to run the shield on.
+        :param params: The parameters of the shield.
+        :returns: A RunShieldResponse.
+        """
+        ...
+
+    @webmethod(route="/moderations", method="POST", level=LLAMA_STACK_API_V1)
+    async def run_moderation(self, input: str | list[str], model: str | None = None) -> ModerationObject:
+        """Create moderation.
+
+        Classifies if text and/or image inputs are potentially harmful.
+        :param input: Input (or inputs) to classify.
+        Can be a single string, an array of strings, or an array of multi-modal input objects similar to other models.
+        :param model: (Optional) The content moderation model you would like to use.
+        :returns: A moderation object.
+        """
+        ...