llama-stack 0.0.42__py3-none-any.whl → 0.3.4__py3-none-any.whl

llama_stack/providers/utils/scoring/base_scoring_fn.py

@@ -0,0 +1,114 @@
+# 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 abc import ABC, abstractmethod
+from typing import Any
+
+from llama_stack.apis.scoring import ScoringFnParams, ScoringResultRow
+from llama_stack.apis.scoring_functions import ScoringFn
+from llama_stack.providers.utils.scoring.aggregation_utils import aggregate_metrics
+
+
+class BaseScoringFn(ABC):
+    """
+    Base interface class for Scoring Functions.
+    Each scoring function needs to implement the following methods:
+    - score_row(self, row)
+    - aggregate(self, scoring_fn_results)
+    """
+
+    def __init__(self, *args, **kwargs) -> None:
+        super().__init__(*args, **kwargs)
+
+    def __str__(self) -> str:
+        return self.__class__.__name__
+
+    @abstractmethod
+    async def score_row(
+        self,
+        input_row: dict[str, Any],
+        scoring_fn_identifier: str | None = None,
+        scoring_params: ScoringFnParams | None = None,
+    ) -> ScoringResultRow:
+        raise NotImplementedError()
+
+    @abstractmethod
+    async def aggregate(
+        self,
+        scoring_results: list[ScoringResultRow],
+        scoring_fn_identifier: str | None = None,
+        scoring_params: ScoringFnParams | None = None,
+    ) -> dict[str, Any]:
+        raise NotImplementedError()
+
+    @abstractmethod
+    async def score(
+        self,
+        input_rows: list[dict[str, Any]],
+        scoring_fn_identifier: str | None = None,
+        scoring_params: ScoringFnParams | None = None,
+    ) -> list[ScoringResultRow]:
+        raise NotImplementedError()
+
+
+class RegisteredBaseScoringFn(BaseScoringFn):
+    """
+    Interface for native scoring functions that are registered in LlamaStack.
+    """
+
+    def __init__(self, *args, **kwargs) -> None:
+        super().__init__(*args, **kwargs)
+        self.supported_fn_defs_registry = {}
+
+    def __str__(self) -> str:
+        return self.__class__.__name__
+
+    def get_supported_scoring_fn_defs(self) -> list[ScoringFn]:
+        return list(self.supported_fn_defs_registry.values())
+
+    def register_scoring_fn_def(self, scoring_fn: ScoringFn) -> None:
+        if scoring_fn.identifier in self.supported_fn_defs_registry:
+            raise ValueError(f"Scoring function def with identifier {scoring_fn.identifier} already exists.")
+        self.supported_fn_defs_registry[scoring_fn.identifier] = scoring_fn
+
+    def unregister_scoring_fn_def(self, scoring_fn_id: str) -> None:
+        if scoring_fn_id not in self.supported_fn_defs_registry:
+            raise ValueError(f"Scoring function def with identifier {scoring_fn_id} does not exist.")
+        del self.supported_fn_defs_registry[scoring_fn_id]
+
+    @abstractmethod
+    async def score_row(
+        self,
+        input_row: dict[str, Any],
+        scoring_fn_identifier: str | None = None,
+        scoring_params: ScoringFnParams | None = None,
+    ) -> ScoringResultRow:
+        raise NotImplementedError()
+
+    async def aggregate(
+        self,
+        scoring_results: list[ScoringResultRow],
+        scoring_fn_identifier: str | None = None,
+        scoring_params: ScoringFnParams | None = None,
+    ) -> dict[str, Any]:
+        params = self.supported_fn_defs_registry[scoring_fn_identifier].params
+        if scoring_params is not None:
+            if params is None:
+                params = scoring_params
+            else:
+                params.aggregation_functions = scoring_params.aggregation_functions
+
+        aggregation_functions = []
+        if params and hasattr(params, "aggregation_functions") and params.aggregation_functions:
+            aggregation_functions.extend(params.aggregation_functions)
+        return aggregate_metrics(scoring_results, aggregation_functions)
+
+    async def score(
+        self,
+        input_rows: list[dict[str, Any]],
+        scoring_fn_identifier: str | None = None,
+        scoring_params: ScoringFnParams | None = None,
+    ) -> list[ScoringResultRow]:
+        return [await self.score_row(input_row, scoring_fn_identifier, scoring_params) for input_row in input_rows]

llama_stack/providers/utils/scoring/basic_scoring_utils.py

@@ -0,0 +1,26 @@
+# 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 contextlib
+import signal
+from collections.abc import Iterator
+from types import FrameType
+
+
+class TimeoutError(Exception):
+    pass
+
+
+@contextlib.contextmanager
+def time_limit(seconds: float) -> Iterator[None]:
+    def signal_handler(signum: int, frame: FrameType | None) -> None:
+        raise TimeoutError("Timed out!")
+
+    signal.setitimer(signal.ITIMER_REAL, seconds)
+    signal.signal(signal.SIGALRM, signal_handler)
+    try:
+        yield
+    finally:
+        signal.setitimer(signal.ITIMER_REAL, 0)

llama_stack/providers/utils/sqlstore/__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/utils/sqlstore/api.py

@@ -0,0 +1,128 @@
+# 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 Mapping, Sequence
+from enum import Enum
+from typing import Any, Literal, Protocol
+
+from pydantic import BaseModel
+
+from llama_stack.apis.common.responses import PaginatedResponse
+
+
+class ColumnType(Enum):
+    INTEGER = "INTEGER"
+    STRING = "STRING"
+    TEXT = "TEXT"
+    FLOAT = "FLOAT"
+    BOOLEAN = "BOOLEAN"
+    JSON = "JSON"
+    DATETIME = "DATETIME"
+
+
+class ColumnDefinition(BaseModel):
+    type: ColumnType
+    primary_key: bool = False
+    nullable: bool = True
+    default: Any = None
+
+
+class SqlStore(Protocol):
+    """
+    A protocol for a SQL store.
+    """
+
+    async def create_table(self, table: str, schema: Mapping[str, ColumnType | ColumnDefinition]) -> None:
+        """
+        Create a table.
+        """
+        pass
+
+    async def insert(self, table: str, data: Mapping[str, Any] | Sequence[Mapping[str, Any]]) -> None:
+        """
+        Insert a row or batch of rows into a table.
+        """
+        pass
+
+    async def fetch_all(
+        self,
+        table: str,
+        where: Mapping[str, Any] | None = None,
+        where_sql: str | None = None,
+        limit: int | None = None,
+        order_by: list[tuple[str, Literal["asc", "desc"]]] | None = None,
+        cursor: tuple[str, str] | None = None,
+    ) -> PaginatedResponse:
+        """
+        Fetch all rows from a table with optional cursor-based pagination.
+
+        :param table: The table name
+        :param where: Simple key-value WHERE conditions
+        :param where_sql: Raw SQL WHERE clause for complex queries
+        :param limit: Maximum number of records to return
+        :param order_by: List of (column, order) tuples for sorting
+        :param cursor: Tuple of (key_column, cursor_id) for pagination (None for first page)
+                      Requires order_by with exactly one column when used
+        :return: PaginatedResult with data and has_more flag
+
+        Note: Cursor pagination only supports single-column ordering for simplicity.
+        Multi-column ordering is allowed without cursor but will raise an error with cursor.
+        """
+        pass
+
+    async def fetch_one(
+        self,
+        table: str,
+        where: Mapping[str, Any] | None = None,
+        where_sql: str | None = None,
+        order_by: list[tuple[str, Literal["asc", "desc"]]] | None = None,
+    ) -> dict[str, Any] | None:
+        """
+        Fetch one row from a table.
+        """
+        pass
+
+    async def update(
+        self,
+        table: str,
+        data: Mapping[str, Any],
+        where: Mapping[str, Any],
+    ) -> None:
+        """
+        Update a row in a table.
+        """
+        pass
+
+    async def delete(
+        self,
+        table: str,
+        where: Mapping[str, Any],
+    ) -> None:
+        """
+        Delete a row from a table.
+        """
+        pass
+
+    async def add_column_if_not_exists(
+        self,
+        table: str,
+        column_name: str,
+        column_type: ColumnType,
+        nullable: bool = True,
+    ) -> None:
+        """
+        Add a column to an existing table if the column doesn't already exist.
+
+        This is useful for table migrations when adding new functionality.
+        If the table doesn't exist, this method should do nothing.
+        If the column already exists, this method should do nothing.
+
+        :param table: Table name
+        :param column_name: Name of the column to add
+        :param column_type: Type of the column to add
+        :param nullable: Whether the column should be nullable (default: True)
+        """
+        pass

llama_stack/providers/utils/sqlstore/authorized_sqlstore.py

@@ -0,0 +1,319 @@
+# 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 Mapping, Sequence
+from typing import Any, Literal
+
+from llama_stack.core.access_control.access_control import default_policy, is_action_allowed
+from llama_stack.core.access_control.conditions import ProtectedResource
+from llama_stack.core.access_control.datatypes import AccessRule, Action, Scope
+from llama_stack.core.datatypes import User
+from llama_stack.core.request_headers import get_authenticated_user
+from llama_stack.core.storage.datatypes import StorageBackendType
+from llama_stack.log import get_logger
+
+from .api import ColumnDefinition, ColumnType, PaginatedResponse, SqlStore
+
+logger = get_logger(name=__name__, category="providers::utils")
+
+# Hardcoded copy of the default policy that our SQL filtering implements
+# WARNING: If default_policy() changes, this constant must be updated accordingly
+# or SQL filtering will fall back to conservative mode (safe but less performant)
+#
+# This policy represents: "Permit all actions when user is in owners list for ALL attribute categories"
+# The corresponding SQL logic is implemented in _build_default_policy_where_clause():
+# - Public records (no access_attributes) are always accessible
+# - Records with access_attributes require user to match ALL categories that exist in the resource
+# - Missing categories in the resource are treated as "no restriction" (allow)
+# - Within each category, user needs ANY matching value (OR logic)
+# - Between categories, user needs ALL categories to match (AND logic)
+SQL_OPTIMIZED_POLICY = [
+    AccessRule(
+        permit=Scope(actions=list(Action)),
+        when=["user in owners roles", "user in owners teams", "user in owners projects", "user in owners namespaces"],
+    ),
+]
+
+
+def _enhance_item_with_access_control(item: Mapping[str, Any], current_user: User | None) -> Mapping[str, Any]:
+    """Add access control attributes to a data item."""
+    enhanced = dict(item)
+    if current_user:
+        enhanced["owner_principal"] = current_user.principal
+        enhanced["access_attributes"] = current_user.attributes
+    else:
+        # IMPORTANT: Use empty string and null value (not None) to match public access filter
+        # The public access filter in _get_public_access_conditions() expects:
+        # - owner_principal = '' (empty string)
+        # - access_attributes = null (JSON null, which serializes to the string 'null')
+        # Setting them to None (SQL NULL) will cause rows to be filtered out on read.
+        enhanced["owner_principal"] = ""
+        enhanced["access_attributes"] = None  # Pydantic/JSON will serialize this as JSON null
+    return enhanced
+
+
+class SqlRecord(ProtectedResource):
+    def __init__(self, record_id: str, table_name: str, owner: User):
+        self.type = f"sql_record::{table_name}"
+        self.identifier = record_id
+        self.owner = owner
+
+
+class AuthorizedSqlStore:
+    """
+    Authorization layer for SqlStore that provides access control functionality.
+
+    This class composes a base SqlStore and adds authorization methods that handle
+    access control policies, user attribute capture, and SQL filtering optimization.
+    """
+
+    def __init__(self, sql_store: SqlStore, policy: list[AccessRule]):
+        """
+        Initialize the authorization layer.
+
+        :param sql_store: Base SqlStore implementation to wrap
+        :param policy: Access control policy to use for authorization
+        """
+        self.sql_store = sql_store
+        self.policy = policy
+        self._detect_database_type()
+        self._validate_sql_optimized_policy()
+
+    def _detect_database_type(self) -> None:
+        """Detect the database type from the underlying SQL store."""
+        if not hasattr(self.sql_store, "config"):
+            raise ValueError("SqlStore must have a config attribute to be used with AuthorizedSqlStore")
+
+        self.database_type = self.sql_store.config.type.value
+        if self.database_type not in [StorageBackendType.SQL_POSTGRES.value, StorageBackendType.SQL_SQLITE.value]:
+            raise ValueError(f"Unsupported database type: {self.database_type}")
+
+    def _validate_sql_optimized_policy(self) -> None:
+        """Validate that SQL_OPTIMIZED_POLICY matches the actual default_policy().
+
+        This ensures that if default_policy() changes, we detect the mismatch and
+        can update our SQL filtering logic accordingly.
+        """
+        actual_default = default_policy()
+
+        if SQL_OPTIMIZED_POLICY != actual_default:
+            logger.warning(
+                f"SQL_OPTIMIZED_POLICY does not match default_policy(). "
+                f"SQL filtering will use conservative mode. "
+                f"Expected: {SQL_OPTIMIZED_POLICY}, Got: {actual_default}",
+            )
+
+    async def create_table(self, table: str, schema: Mapping[str, ColumnType | ColumnDefinition]) -> None:
+        """Create a table with built-in access control support."""
+
+        enhanced_schema = dict(schema)
+        if "access_attributes" not in enhanced_schema:
+            enhanced_schema["access_attributes"] = ColumnType.JSON
+        if "owner_principal" not in enhanced_schema:
+            enhanced_schema["owner_principal"] = ColumnType.STRING
+
+        await self.sql_store.create_table(table, enhanced_schema)
+        await self.sql_store.add_column_if_not_exists(table, "access_attributes", ColumnType.JSON)
+        await self.sql_store.add_column_if_not_exists(table, "owner_principal", ColumnType.STRING)
+
+    async def insert(self, table: str, data: Mapping[str, Any] | Sequence[Mapping[str, Any]]) -> None:
+        """Insert a row or batch of rows with automatic access control attribute capture."""
+        current_user = get_authenticated_user()
+        enhanced_data: Mapping[str, Any] | Sequence[Mapping[str, Any]]
+        if isinstance(data, Mapping):
+            enhanced_data = _enhance_item_with_access_control(data, current_user)
+        else:
+            enhanced_data = [_enhance_item_with_access_control(item, current_user) for item in data]
+        await self.sql_store.insert(table, enhanced_data)
+
+    async def fetch_all(
+        self,
+        table: str,
+        where: Mapping[str, Any] | None = None,
+        limit: int | None = None,
+        order_by: list[tuple[str, Literal["asc", "desc"]]] | None = None,
+        cursor: tuple[str, str] | None = None,
+    ) -> PaginatedResponse:
+        """Fetch all rows with automatic access control filtering."""
+        access_where = self._build_access_control_where_clause(self.policy)
+        rows = await self.sql_store.fetch_all(
+            table=table,
+            where=where,
+            where_sql=access_where,
+            limit=limit,
+            order_by=order_by,
+            cursor=cursor,
+        )
+
+        current_user = get_authenticated_user()
+        filtered_rows = []
+
+        for row in rows.data:
+            stored_access_attrs = row.get("access_attributes")
+            stored_owner_principal = row.get("owner_principal") or ""
+
+            record_id = row.get("id", "unknown")
+            sql_record = SqlRecord(
+                str(record_id), table, User(principal=stored_owner_principal, attributes=stored_access_attrs)
+            )
+
+            if is_action_allowed(self.policy, Action.READ, sql_record, current_user):
+                filtered_rows.append(row)
+
+        return PaginatedResponse(
+            data=filtered_rows,
+            has_more=rows.has_more,
+        )
+
+    async def fetch_one(
+        self,
+        table: str,
+        where: Mapping[str, Any] | None = None,
+        order_by: list[tuple[str, Literal["asc", "desc"]]] | None = None,
+    ) -> dict[str, Any] | None:
+        """Fetch one row with automatic access control checking."""
+        results = await self.fetch_all(
+            table=table,
+            where=where,
+            limit=1,
+            order_by=order_by,
+        )
+
+        return results.data[0] if results.data else None
+
+    async def update(self, table: str, data: Mapping[str, Any], where: Mapping[str, Any]) -> None:
+        """Update rows with automatic access control attribute capture."""
+        enhanced_data = dict(data)
+
+        current_user = get_authenticated_user()
+        if current_user:
+            enhanced_data["owner_principal"] = current_user.principal
+            enhanced_data["access_attributes"] = current_user.attributes
+        else:
+            # IMPORTANT: Use empty string for owner_principal to match public access filter
+            enhanced_data["owner_principal"] = ""
+            enhanced_data["access_attributes"] = None  # Will serialize as JSON null
+
+        await self.sql_store.update(table, enhanced_data, where)
+
+    async def delete(self, table: str, where: Mapping[str, Any]) -> None:
+        """Delete rows with automatic access control filtering."""
+        await self.sql_store.delete(table, where)
+
+    def _build_access_control_where_clause(self, policy: list[AccessRule]) -> str:
+        """Build SQL WHERE clause for access control filtering.
+
+        Only applies SQL filtering for the default policy to ensure correctness.
+        For custom policies, uses conservative filtering to avoid blocking legitimate access.
+        """
+        current_user = get_authenticated_user()
+
+        if not policy or policy == SQL_OPTIMIZED_POLICY:
+            return self._build_default_policy_where_clause(current_user)
+        else:
+            return self._build_conservative_where_clause()
+
+    def _json_extract(self, column: str, path: str) -> str:
+        """Extract JSON value (keeping JSON type).
+
+        Args:
+            column: The JSON column name
+            path: The JSON path (e.g., 'roles', 'teams')
+
+        Returns:
+            SQL expression to extract JSON value
+        """
+        if self.database_type == StorageBackendType.SQL_POSTGRES.value:
+            return f"{column}->'{path}'"
+        elif self.database_type == StorageBackendType.SQL_SQLITE.value:
+            return f"JSON_EXTRACT({column}, '$.{path}')"
+        else:
+            raise ValueError(f"Unsupported database type: {self.database_type}")
+
+    def _json_extract_text(self, column: str, path: str) -> str:
+        """Extract JSON value as text.
+
+        Args:
+            column: The JSON column name
+            path: The JSON path (e.g., 'roles', 'teams')
+
+        Returns:
+            SQL expression to extract JSON value as text
+        """
+        if self.database_type == StorageBackendType.SQL_POSTGRES.value:
+            return f"{column}->>'{path}'"
+        elif self.database_type == StorageBackendType.SQL_SQLITE.value:
+            return f"JSON_EXTRACT({column}, '$.{path}')"
+        else:
+            raise ValueError(f"Unsupported database type: {self.database_type}")
+
+    def _get_public_access_conditions(self) -> list[str]:
+        """Get the SQL conditions for public access.
+
+        Public records are those with:
+        - owner_principal = '' (empty string)
+        - access_attributes is either SQL NULL or JSON null
+
+        Note: Different databases serialize None differently:
+        - SQLite: None → JSON null (text = 'null')
+        - Postgres: None → SQL NULL (IS NULL)
+        """
+        conditions = ["owner_principal = ''"]
+        if self.database_type == StorageBackendType.SQL_POSTGRES.value:
+            # Accept both SQL NULL and JSON null for Postgres compatibility
+            # This handles both old rows (SQL NULL) and new rows (JSON null)
+            conditions.append("(access_attributes IS NULL OR access_attributes::text = 'null')")
+        elif self.database_type == StorageBackendType.SQL_SQLITE.value:
+            # SQLite serializes None as JSON null
+            conditions.append("(access_attributes IS NULL OR access_attributes = 'null')")
+        else:
+            raise ValueError(f"Unsupported database type: {self.database_type}")
+        return conditions
+
+    def _build_default_policy_where_clause(self, current_user: User | None) -> str:
+        """Build SQL WHERE clause for the default policy.
+
+        Default policy: permit all actions when user in owners [roles, teams, projects, namespaces]
+        This means user must match ALL attribute categories that exist in the resource.
+        """
+        base_conditions = self._get_public_access_conditions()
+        user_attr_conditions = []
+
+        if current_user and current_user.attributes:
+            for attr_key, user_values in current_user.attributes.items():
+                if user_values:
+                    value_conditions = []
+                    for value in user_values:
+                        # Check if JSON array contains the value
+                        escaped_value = value.replace("'", "''")
+                        json_text = self._json_extract_text("access_attributes", attr_key)
+                        value_conditions.append(f"({json_text} LIKE '%\"{escaped_value}\"%')")
+
+                    if value_conditions:
+                        # Check if the category is missing (NULL)
+                        category_missing = f"{self._json_extract('access_attributes', attr_key)} IS NULL"
+                        user_matches_category = f"({' OR '.join(value_conditions)})"
+                        user_attr_conditions.append(f"({category_missing} OR {user_matches_category})")
+
+            if user_attr_conditions:
+                all_requirements_met = f"({' AND '.join(user_attr_conditions)})"
+                base_conditions.append(all_requirements_met)
+
+        return f"({' OR '.join(base_conditions)})"
+
+    def _build_conservative_where_clause(self) -> str:
+        """Conservative SQL filtering for custom policies.
+
+        Only filters records we're 100% certain would be denied by any reasonable policy.
+        """
+        current_user = get_authenticated_user()
+
+        if not current_user:
+            # Only allow public records
+            base_conditions = self._get_public_access_conditions()
+            return f"({' OR '.join(base_conditions)})"
+
+        return "1=1"