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

llama_stack/providers/utils/sqlstore/sqlalchemy_sqlstore.py

@@ -0,0 +1,343 @@
+# 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 sqlalchemy import (
+    JSON,
+    Boolean,
+    Column,
+    DateTime,
+    Float,
+    Integer,
+    MetaData,
+    String,
+    Table,
+    Text,
+    event,
+    inspect,
+    select,
+    text,
+)
+from sqlalchemy.ext.asyncio import async_sessionmaker, create_async_engine
+from sqlalchemy.ext.asyncio.engine import AsyncEngine
+from sqlalchemy.sql.elements import ColumnElement
+
+from llama_stack.apis.common.responses import PaginatedResponse
+from llama_stack.core.storage.datatypes import SqlAlchemySqlStoreConfig
+from llama_stack.log import get_logger
+
+from .api import ColumnDefinition, ColumnType, SqlStore
+
+logger = get_logger(name=__name__, category="providers::utils")
+
+TYPE_MAPPING: dict[ColumnType, Any] = {
+    ColumnType.INTEGER: Integer,
+    ColumnType.STRING: String,
+    ColumnType.FLOAT: Float,
+    ColumnType.BOOLEAN: Boolean,
+    ColumnType.DATETIME: DateTime,
+    ColumnType.TEXT: Text,
+    ColumnType.JSON: JSON,
+}
+
+
+def _build_where_expr(column: ColumnElement, value: Any) -> ColumnElement:
+    """Return a SQLAlchemy expression for a where condition.
+
+    `value` may be a simple scalar (equality) or a mapping like {">": 123}.
+    The returned expression is a SQLAlchemy ColumnElement usable in query.where(...).
+    """
+    if isinstance(value, Mapping):
+        if len(value) != 1:
+            raise ValueError(f"Operator mapping must have a single operator, got: {value}")
+        op, operand = next(iter(value.items()))
+        if op == "==" or op == "=":
+            return column == operand
+        if op == ">":
+            return column > operand
+        if op == "<":
+            return column < operand
+        if op == ">=":
+            return column >= operand
+        if op == "<=":
+            return column <= operand
+        raise ValueError(f"Unsupported operator '{op}' in where mapping")
+    return column == value
+
+
+class SqlAlchemySqlStoreImpl(SqlStore):
+    def __init__(self, config: SqlAlchemySqlStoreConfig):
+        self.config = config
+        self.async_session = async_sessionmaker(self.create_engine())
+        self.metadata = MetaData()
+
+    def create_engine(self) -> AsyncEngine:
+        # Configure connection args for better concurrency support
+        connect_args = {}
+        if "sqlite" in self.config.engine_str:
+            # SQLite-specific optimizations for concurrent access
+            # With WAL mode, most locks resolve in milliseconds, but allow up to 5s for edge cases
+            connect_args["timeout"] = 5.0
+            connect_args["check_same_thread"] = False  # Allow usage across asyncio tasks
+
+        engine = create_async_engine(
+            self.config.engine_str,
+            pool_pre_ping=True,
+            connect_args=connect_args,
+        )
+
+        # Enable WAL mode for SQLite to support concurrent readers and writers
+        if "sqlite" in self.config.engine_str:
+
+            @event.listens_for(engine.sync_engine, "connect")
+            def set_sqlite_pragma(dbapi_conn, connection_record):
+                cursor = dbapi_conn.cursor()
+                # Enable Write-Ahead Logging for better concurrency
+                cursor.execute("PRAGMA journal_mode=WAL")
+                # Set busy timeout to 5 seconds (retry instead of immediate failure)
+                # With WAL mode, locks should be brief; if we hit 5s there's a bigger issue
+                cursor.execute("PRAGMA busy_timeout=5000")
+                # Use NORMAL synchronous mode for better performance (still safe with WAL)
+                cursor.execute("PRAGMA synchronous=NORMAL")
+                cursor.close()
+
+        return engine
+
+    async def create_table(
+        self,
+        table: str,
+        schema: Mapping[str, ColumnType | ColumnDefinition],
+    ) -> None:
+        if not schema:
+            raise ValueError(f"No columns defined for table '{table}'.")
+
+        sqlalchemy_columns: list[Column] = []
+
+        for col_name, col_props in schema.items():
+            col_type = None
+            is_primary_key = False
+            is_nullable = True
+
+            if isinstance(col_props, ColumnType):
+                col_type = col_props
+            elif isinstance(col_props, ColumnDefinition):
+                col_type = col_props.type
+                is_primary_key = col_props.primary_key
+                is_nullable = col_props.nullable
+
+            sqlalchemy_type = TYPE_MAPPING.get(col_type)
+            if not sqlalchemy_type:
+                raise ValueError(f"Unsupported column type '{col_type}' for column '{col_name}'.")
+
+            sqlalchemy_columns.append(
+                Column(col_name, sqlalchemy_type, primary_key=is_primary_key, nullable=is_nullable)
+            )
+
+        if table not in self.metadata.tables:
+            sqlalchemy_table = Table(table, self.metadata, *sqlalchemy_columns)
+        else:
+            sqlalchemy_table = self.metadata.tables[table]
+
+        engine = self.create_engine()
+        async with engine.begin() as conn:
+            await conn.run_sync(self.metadata.create_all, tables=[sqlalchemy_table], checkfirst=True)
+
+    async def insert(self, table: str, data: Mapping[str, Any] | Sequence[Mapping[str, Any]]) -> None:
+        async with self.async_session() as session:
+            await session.execute(self.metadata.tables[table].insert(), data)
+            await session.commit()
+
+    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:
+        async with self.async_session() as session:
+            table_obj = self.metadata.tables[table]
+            query = select(table_obj)
+
+            if where:
+                for key, value in where.items():
+                    query = query.where(_build_where_expr(table_obj.c[key], value))
+
+            if where_sql:
+                query = query.where(text(where_sql))
+
+            # Handle cursor-based pagination
+            if cursor:
+                # Validate cursor tuple format
+                if not isinstance(cursor, tuple) or len(cursor) != 2:
+                    raise ValueError(f"Cursor must be a tuple of (key_column, cursor_id), got: {cursor}")
+
+                # Require order_by for cursor pagination
+                if not order_by:
+                    raise ValueError("order_by is required when using cursor pagination")
+
+                # Only support single-column ordering for cursor pagination
+                if len(order_by) != 1:
+                    raise ValueError(
+                        f"Cursor pagination only supports single-column ordering, got {len(order_by)} columns"
+                    )
+
+                cursor_key_column, cursor_id = cursor
+                order_column, order_direction = order_by[0]
+
+                # Verify cursor_key_column exists
+                if cursor_key_column not in table_obj.c:
+                    raise ValueError(f"Cursor key column '{cursor_key_column}' not found in table '{table}'")
+
+                # Get cursor value for the order column
+                cursor_query = select(table_obj.c[order_column]).where(table_obj.c[cursor_key_column] == cursor_id)
+                cursor_result = await session.execute(cursor_query)
+                cursor_row = cursor_result.fetchone()
+
+                if not cursor_row:
+                    raise ValueError(f"Record with {cursor_key_column}='{cursor_id}' not found in table '{table}'")
+
+                cursor_value = cursor_row[0]
+
+                # Apply cursor condition based on sort direction
+                if order_direction == "desc":
+                    query = query.where(table_obj.c[order_column] < cursor_value)
+                else:
+                    query = query.where(table_obj.c[order_column] > cursor_value)
+
+            # Apply ordering
+            if order_by:
+                if not isinstance(order_by, list):
+                    raise ValueError(
+                        f"order_by must be a list of tuples (column, order={['asc', 'desc']}), got {order_by}"
+                    )
+                for order in order_by:
+                    if not isinstance(order, tuple):
+                        raise ValueError(
+                            f"order_by must be a list of tuples (column, order={['asc', 'desc']}), got {order_by}"
+                        )
+                    name, order_type = order
+                    if name not in table_obj.c:
+                        raise ValueError(f"Column '{name}' not found in table '{table}'")
+                    if order_type == "asc":
+                        query = query.order_by(table_obj.c[name].asc())
+                    elif order_type == "desc":
+                        query = query.order_by(table_obj.c[name].desc())
+                    else:
+                        raise ValueError(f"Invalid order '{order_type}' for column '{name}'")
+
+            # Fetch limit + 1 to determine has_more
+            fetch_limit = limit
+            if limit:
+                fetch_limit = limit + 1
+
+            if fetch_limit:
+                query = query.limit(fetch_limit)
+
+            result = await session.execute(query)
+            if result.rowcount == 0:
+                rows = []
+            else:
+                rows = [dict(row._mapping) for row in result]
+
+            # Always return pagination result
+            has_more = False
+            if limit and len(rows) > limit:
+                has_more = True
+                rows = rows[:limit]
+
+            return PaginatedResponse(data=rows, has_more=has_more)
+
+    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:
+        result = await self.fetch_all(table, where, where_sql, limit=1, order_by=order_by)
+        if not result.data:
+            return None
+        return result.data[0]
+
+    async def update(
+        self,
+        table: str,
+        data: Mapping[str, Any],
+        where: Mapping[str, Any],
+    ) -> None:
+        if not where:
+            raise ValueError("where is required for update")
+
+        async with self.async_session() as session:
+            stmt = self.metadata.tables[table].update()
+            for key, value in where.items():
+                stmt = stmt.where(_build_where_expr(self.metadata.tables[table].c[key], value))
+            await session.execute(stmt, data)
+            await session.commit()
+
+    async def delete(self, table: str, where: Mapping[str, Any]) -> None:
+        if not where:
+            raise ValueError("where is required for delete")
+
+        async with self.async_session() as session:
+            stmt = self.metadata.tables[table].delete()
+            for key, value in where.items():
+                stmt = stmt.where(_build_where_expr(self.metadata.tables[table].c[key], value))
+            await session.execute(stmt)
+            await session.commit()
+
+    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."""
+        engine = self.create_engine()
+
+        try:
+            async with engine.begin() as conn:
+
+                def check_column_exists(sync_conn):
+                    inspector = inspect(sync_conn)
+
+                    table_names = inspector.get_table_names()
+                    if table not in table_names:
+                        return False, False  # table doesn't exist, column doesn't exist
+
+                    existing_columns = inspector.get_columns(table)
+                    column_names = [col["name"] for col in existing_columns]
+
+                    return True, column_name in column_names  # table exists, column exists or not
+
+                table_exists, column_exists = await conn.run_sync(check_column_exists)
+                if not table_exists or column_exists:
+                    return
+
+                sqlalchemy_type = TYPE_MAPPING.get(column_type)
+                if not sqlalchemy_type:
+                    raise ValueError(f"Unsupported column type '{column_type}' for column '{column_name}'.")
+
+                # Create the ALTER TABLE statement
+                # Note: We need to get the dialect-specific type name
+                dialect = engine.dialect
+                type_impl = sqlalchemy_type()
+                compiled_type = type_impl.compile(dialect=dialect)
+
+                nullable_clause = "" if nullable else " NOT NULL"
+                add_column_sql = text(f"ALTER TABLE {table} ADD COLUMN {column_name} {compiled_type}{nullable_clause}")
+
+                await conn.execute(add_column_sql)
+
+        except Exception as e:
+            # If any error occurs during migration, log it but don't fail
+            # The table creation will handle adding the column
+            logger.error(f"Error adding column {column_name} to table {table}: {e}")
+            pass

llama_stack/providers/utils/sqlstore/sqlstore.py

@@ -0,0 +1,70 @@
+# 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, cast
+
+from pydantic import Field
+
+from llama_stack.core.storage.datatypes import (
+    PostgresSqlStoreConfig,
+    SqliteSqlStoreConfig,
+    SqlStoreReference,
+    StorageBackendConfig,
+    StorageBackendType,
+)
+
+from .api import SqlStore
+
+sql_store_pip_packages = ["sqlalchemy[asyncio]", "aiosqlite", "asyncpg"]
+
+_SQLSTORE_BACKENDS: dict[str, StorageBackendConfig] = {}
+
+
+SqlStoreConfig = Annotated[
+    SqliteSqlStoreConfig | PostgresSqlStoreConfig,
+    Field(discriminator="type"),
+]
+
+
+def get_pip_packages(store_config: dict | SqlStoreConfig) -> list[str]:
+    """Get pip packages for SQL store config, handling both dict and object cases."""
+    if isinstance(store_config, dict):
+        store_type = store_config.get("type")
+        if store_type == StorageBackendType.SQL_SQLITE.value:
+            return SqliteSqlStoreConfig.pip_packages()
+        elif store_type == StorageBackendType.SQL_POSTGRES.value:
+            return PostgresSqlStoreConfig.pip_packages()
+        else:
+            raise ValueError(f"Unknown SQL store type: {store_type}")
+    else:
+        return store_config.pip_packages()
+
+
+def sqlstore_impl(reference: SqlStoreReference) -> SqlStore:
+    backend_name = reference.backend
+
+    backend_config = _SQLSTORE_BACKENDS.get(backend_name)
+    if backend_config is None:
+        raise ValueError(
+            f"Unknown SQL store backend '{backend_name}'. Registered backends: {sorted(_SQLSTORE_BACKENDS)}"
+        )
+
+    if isinstance(backend_config, SqliteSqlStoreConfig | PostgresSqlStoreConfig):
+        from .sqlalchemy_sqlstore import SqlAlchemySqlStoreImpl
+
+        config = cast(SqliteSqlStoreConfig | PostgresSqlStoreConfig, backend_config).model_copy()
+        return SqlAlchemySqlStoreImpl(config)
+    else:
+        raise ValueError(f"Unknown sqlstore type {backend_config.type}")
+
+
+def register_sqlstore_backends(backends: dict[str, StorageBackendConfig]) -> None:
+    """Register the set of available SQL store backends for reference resolution."""
+    global _SQLSTORE_BACKENDS
+
+    _SQLSTORE_BACKENDS.clear()
+    for name, cfg in backends.items():
+        _SQLSTORE_BACKENDS[name] = cfg

llama_stack/providers/utils/telemetry/trace_protocol.py

@@ -0,0 +1,142 @@
+# 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
+import inspect
+import json
+from collections.abc import AsyncGenerator, Callable
+from functools import wraps
+from typing import Any
+
+from pydantic import BaseModel
+
+from llama_stack.models.llama.datatypes import Primitive
+
+
+def serialize_value(value: Any) -> Primitive:
+    return str(_prepare_for_json(value))
+
+
+def _prepare_for_json(value: Any) -> str:
+    """Serialize a single value into JSON-compatible format."""
+    if value is None:
+        return ""
+    elif isinstance(value, str | int | float | bool):
+        return value
+    elif hasattr(value, "_name_"):
+        return value._name_
+    elif isinstance(value, BaseModel):
+        return json.loads(value.model_dump_json())
+    elif isinstance(value, list | tuple | set):
+        return [_prepare_for_json(item) for item in value]
+    elif isinstance(value, dict):
+        return {str(k): _prepare_for_json(v) for k, v in value.items()}
+    else:
+        try:
+            json.dumps(value)
+            return value
+        except Exception:
+            return str(value)
+
+
+def trace_protocol[T](cls: type[T]) -> type[T]:
+    """
+    A class decorator that automatically traces all methods in a protocol/base class
+    and its inheriting classes.
+    """
+
+    def trace_method(method: Callable) -> Callable:
+        is_async = asyncio.iscoroutinefunction(method)
+        is_async_gen = inspect.isasyncgenfunction(method)
+
+        def create_span_context(self: Any, *args: Any, **kwargs: Any) -> tuple:
+            class_name = self.__class__.__name__
+            method_name = method.__name__
+            span_type = "async_generator" if is_async_gen else "async" if is_async else "sync"
+            sig = inspect.signature(method)
+            param_names = list(sig.parameters.keys())[1:]  # Skip 'self'
+            combined_args = {}
+            for i, arg in enumerate(args):
+                param_name = param_names[i] if i < len(param_names) else f"position_{i + 1}"
+                combined_args[param_name] = serialize_value(arg)
+            for k, v in kwargs.items():
+                combined_args[str(k)] = serialize_value(v)
+
+            span_attributes = {
+                "__autotraced__": True,
+                "__class__": class_name,
+                "__method__": method_name,
+                "__type__": span_type,
+                "__args__": json.dumps(combined_args),
+            }
+
+            return class_name, method_name, span_attributes
+
+        @wraps(method)
+        async def async_gen_wrapper(self: Any, *args: Any, **kwargs: Any) -> AsyncGenerator:
+            from llama_stack.providers.utils.telemetry import tracing
+
+            class_name, method_name, span_attributes = create_span_context(self, *args, **kwargs)
+
+            with tracing.span(f"{class_name}.{method_name}", span_attributes) as span:
+                count = 0
+                try:
+                    async for item in method(self, *args, **kwargs):
+                        yield item
+                        count += 1
+                finally:
+                    span.set_attribute("chunk_count", count)
+
+        @wraps(method)
+        async def async_wrapper(self: Any, *args: Any, **kwargs: Any) -> Any:
+            from llama_stack.providers.utils.telemetry import tracing
+
+            class_name, method_name, span_attributes = create_span_context(self, *args, **kwargs)
+
+            with tracing.span(f"{class_name}.{method_name}", span_attributes) as span:
+                try:
+                    result = await method(self, *args, **kwargs)
+                    span.set_attribute("output", serialize_value(result))
+                    return result
+                except Exception as e:
+                    span.set_attribute("error", str(e))
+                    raise
+
+        @wraps(method)
+        def sync_wrapper(self: Any, *args: Any, **kwargs: Any) -> Any:
+            from llama_stack.providers.utils.telemetry import tracing
+
+            class_name, method_name, span_attributes = create_span_context(self, *args, **kwargs)
+
+            with tracing.span(f"{class_name}.{method_name}", span_attributes) as span:
+                try:
+                    result = method(self, *args, **kwargs)
+                    span.set_attribute("output", serialize_value(result))
+                    return result
+                except Exception as e:
+                    span.set_attribute("error", str(e))
+                    raise
+
+        if is_async_gen:
+            return async_gen_wrapper
+        elif is_async:
+            return async_wrapper
+        else:
+            return sync_wrapper
+
+    original_init_subclass = getattr(cls, "__init_subclass__", None)
+
+    def __init_subclass__(cls_child, **kwargs):  # noqa: N807
+        if original_init_subclass:
+            original_init_subclass(**kwargs)
+
+        for name, method in vars(cls_child).items():
+            if inspect.isfunction(method) and not name.startswith("_"):
+                setattr(cls_child, name, trace_method(method))  # noqa: B010
+
+    cls.__init_subclass__ = classmethod(__init_subclass__)
+
+    return cls