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

llama_stack/providers/utils/common/data_schema_validator.py

@@ -0,0 +1,103 @@
+# 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
+
+from llama_stack.apis.common.type_system import (
+    ChatCompletionInputType,
+    CompletionInputType,
+    StringType,
+)
+from llama_stack.core.datatypes import Api
+
+
+class ColumnName(Enum):
+    input_query = "input_query"
+    expected_answer = "expected_answer"
+    chat_completion_input = "chat_completion_input"
+    completion_input = "completion_input"
+    generated_answer = "generated_answer"
+    context = "context"
+    dialog = "dialog"
+    function = "function"
+    language = "language"
+    id = "id"
+    ground_truth = "ground_truth"
+
+
+VALID_SCHEMAS_FOR_SCORING = [
+    {
+        ColumnName.input_query.value: StringType(),
+        ColumnName.expected_answer.value: StringType(),
+        ColumnName.generated_answer.value: StringType(),
+    },
+    {
+        ColumnName.input_query.value: StringType(),
+        ColumnName.expected_answer.value: StringType(),
+        ColumnName.generated_answer.value: StringType(),
+        ColumnName.context.value: StringType(),
+    },
+    {
+        ColumnName.input_query.value: StringType(),
+        ColumnName.expected_answer.value: StringType(),
+        ColumnName.generated_answer.value: StringType(),
+        ColumnName.function.value: StringType(),
+        ColumnName.language.value: StringType(),
+        ColumnName.id.value: StringType(),
+        ColumnName.ground_truth.value: StringType(),
+    },
+]
+
+VALID_SCHEMAS_FOR_EVAL = [
+    {
+        ColumnName.input_query.value: StringType(),
+        ColumnName.expected_answer.value: StringType(),
+        ColumnName.chat_completion_input.value: ChatCompletionInputType(),
+    },
+    {
+        ColumnName.input_query.value: StringType(),
+        ColumnName.expected_answer.value: StringType(),
+        ColumnName.completion_input.value: CompletionInputType(),
+    },
+    {
+        ColumnName.input_query.value: StringType(),
+        ColumnName.expected_answer.value: StringType(),
+        ColumnName.generated_answer.value: StringType(),
+        ColumnName.function.value: StringType(),
+        ColumnName.language.value: StringType(),
+        ColumnName.id.value: StringType(),
+        ColumnName.ground_truth.value: StringType(),
+    },
+]
+
+
+def get_valid_schemas(api_str: str):
+    if api_str == Api.scoring.value:
+        return VALID_SCHEMAS_FOR_SCORING
+    elif api_str == Api.eval.value:
+        return VALID_SCHEMAS_FOR_EVAL
+    else:
+        raise ValueError(f"Invalid API string: {api_str}")
+
+
+def validate_dataset_schema(
+    dataset_schema: dict[str, Any],
+    expected_schemas: list[dict[str, Any]],
+):
+    if dataset_schema not in expected_schemas:
+        raise ValueError(f"Dataset {dataset_schema} does not have a correct input schema in {expected_schemas}")
+
+
+def validate_row_schema(
+    input_row: dict[str, Any],
+    expected_schemas: list[dict[str, Any]],
+):
+    for schema in expected_schemas:
+        if all(key in input_row for key in schema):
+            return
+
+    raise ValueError(f"Input row {input_row} does not match any of the expected schemas in {expected_schemas}")

llama_stack/providers/utils/datasetio/__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/datasetio/url_utils.py

@@ -0,0 +1,47 @@
+# 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 base64
+import io
+from urllib.parse import unquote
+
+from llama_stack.providers.utils.memory.vector_store import parse_data_url
+
+
+async def get_dataframe_from_uri(uri: str):
+    import pandas
+
+    df = None
+    if uri.endswith(".csv"):
+        # Moving to its own thread to avoid io from blocking the eventloop
+        # This isn't ideal as it moves more then just the IO to a new thread
+        # but it is as close as we can easly get
+        df = await asyncio.to_thread(pandas.read_csv, uri)
+    elif uri.endswith(".xlsx"):
+        df = await asyncio.to_thread(pandas.read_excel, uri)
+    elif uri.startswith("data:"):
+        parts = parse_data_url(uri)
+        data = parts["data"]
+        if parts["is_base64"]:
+            data = base64.b64decode(data)
+        else:
+            data = unquote(data)
+            encoding = parts["encoding"] or "utf-8"
+            data = data.encode(encoding)
+
+        mime_type = parts["mimetype"]
+        mime_category = mime_type.split("/")[0]
+        data_bytes = io.BytesIO(data)
+
+        if mime_category == "text":
+            df = pandas.read_csv(data_bytes)
+        else:
+            df = pandas.read_excel(data_bytes)
+    else:
+        raise ValueError(f"Unsupported file type: {uri}")
+
+    return df

llama_stack/providers/utils/files/__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/files/form_data.py

@@ -0,0 +1,69 @@
+# 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 json
+
+from fastapi import Request
+from pydantic import BaseModel, ValidationError
+
+from llama_stack.apis.files import ExpiresAfter
+
+
+async def parse_pydantic_from_form[T: BaseModel](request: Request, field_name: str, model_class: type[T]) -> T | None:
+    """
+    Generic parser to extract a Pydantic model from multipart form data.
+    Handles both bracket notation (field[attr1], field[attr2]) and JSON string format.
+
+    Args:
+        request: The FastAPI request object
+        field_name: The name of the field in the form data (e.g., "expires_after")
+        model_class: The Pydantic model class to parse into
+
+    Returns:
+        An instance of model_class if parsing succeeds, None otherwise
+
+    Example:
+        expires_after = await parse_pydantic_from_form(
+            request, "expires_after", ExpiresAfter
+        )
+    """
+    form = await request.form()
+
+    # Check for bracket notation first (e.g., expires_after[anchor], expires_after[seconds])
+    bracket_data = {}
+    prefix = f"{field_name}["
+    for key in form.keys():
+        if key.startswith(prefix) and key.endswith("]"):
+            # Extract the attribute name from field_name[attr]
+            attr = key[len(prefix) : -1]
+            bracket_data[attr] = form[key]
+
+    if bracket_data:
+        try:
+            return model_class(**bracket_data)
+        except (ValidationError, TypeError):
+            pass
+
+    # Check for JSON string format
+    if field_name in form:
+        value = form[field_name]
+        if isinstance(value, str):
+            try:
+                data = json.loads(value)
+                return model_class(**data)
+            except (json.JSONDecodeError, TypeError, ValidationError):
+                pass
+
+    return None
+
+
+async def parse_expires_after(request: Request) -> ExpiresAfter | None:
+    """
+    Dependency to parse expires_after from multipart form data.
+    Handles both bracket notation (expires_after[anchor], expires_after[seconds])
+    and JSON string format.
+    """
+    return await parse_pydantic_from_form(request, "expires_after", ExpiresAfter)

llama_stack/providers/utils/inference/__init__.py

@@ -4,10 +4,8 @@
 # This source code is licensed under the terms described in the LICENSE file in
 # the root directory of this source tree.
 
-from typing import List
-
-from llama_models.datatypes import *  # noqa: F403
-from llama_models.sku_list import all_registered_models
+from llama_stack.models.llama.sku_list import all_registered_models
+from llama_stack.models.llama.sku_types import *  # noqa: F403
 
 
 def is_supported_safety_model(model: Model) -> bool:
@@ -22,12 +20,15 @@ def is_supported_safety_model(model: Model) -> bool:
     ]
 
 
-def supported_inference_models() -> List[str]:
+def supported_inference_models() -> list[Model]:
     return [
-        m.descriptor()
+        m
         for m in all_registered_models()
         if (
-            m.model_family in {ModelFamily.llama3_1, ModelFamily.llama3_2}
+            m.model_family in {ModelFamily.llama3_1, ModelFamily.llama3_2, ModelFamily.llama3_3, ModelFamily.llama4}
             or is_supported_safety_model(m)
         )
     ]
+
+
+ALL_HUGGINGFACE_REPOS_TO_MODEL_DESCRIPTOR = {m.huggingface_repo: m.descriptor() for m in all_registered_models()}

llama_stack/providers/utils/inference/embedding_mixin.py

@@ -0,0 +1,101 @@
+# 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 base64
+import platform
+import struct
+from typing import TYPE_CHECKING
+
+import torch
+
+from llama_stack.log import get_logger
+
+if TYPE_CHECKING:
+    from sentence_transformers import SentenceTransformer
+
+from llama_stack.apis.inference import (
+    ModelStore,
+    OpenAIEmbeddingData,
+    OpenAIEmbeddingsRequestWithExtraBody,
+    OpenAIEmbeddingsResponse,
+    OpenAIEmbeddingUsage,
+)
+
+EMBEDDING_MODELS = {}
+
+DARWIN = "Darwin"
+
+
+log = get_logger(name=__name__, category="providers::utils")
+
+
+class SentenceTransformerEmbeddingMixin:
+    model_store: ModelStore
+
+    async def openai_embeddings(
+        self,
+        params: OpenAIEmbeddingsRequestWithExtraBody,
+    ) -> OpenAIEmbeddingsResponse:
+        # Convert input to list format if it's a single string
+        input_list = [params.input] if isinstance(params.input, str) else params.input
+        if not input_list:
+            raise ValueError("Empty list not supported")
+
+        # Get the model and generate embeddings
+        embedding_model = await self._load_sentence_transformer_model(params.model)
+        embeddings = await asyncio.to_thread(embedding_model.encode, input_list, show_progress_bar=False)
+
+        # Convert embeddings to the requested format
+        data = []
+        for i, embedding in enumerate(embeddings):
+            if params.encoding_format == "base64":
+                # Convert float array to base64 string
+                float_bytes = struct.pack(f"{len(embedding)}f", *embedding)
+                embedding_value = base64.b64encode(float_bytes).decode("ascii")
+            else:
+                # Default to float format
+                embedding_value = embedding.tolist()
+
+            data.append(
+                OpenAIEmbeddingData(
+                    embedding=embedding_value,
+                    index=i,
+                )
+            )
+
+        # Not returning actual token usage
+        usage = OpenAIEmbeddingUsage(prompt_tokens=-1, total_tokens=-1)
+        return OpenAIEmbeddingsResponse(
+            data=data,
+            model=params.model,
+            usage=usage,
+        )
+
+    async def _load_sentence_transformer_model(self, model: str) -> "SentenceTransformer":
+        global EMBEDDING_MODELS
+
+        loaded_model = EMBEDDING_MODELS.get(model)
+        if loaded_model is not None:
+            return loaded_model
+
+        log.info(f"Loading sentence transformer for {model}...")
+
+        def _load_model():
+            from sentence_transformers import SentenceTransformer
+
+            platform_name = platform.system()
+            if platform_name == DARWIN:
+                # PyTorch's OpenMP kernels can segfault on macOS when spawned from background
+                # threads with the default parallel settings, so force a single-threaded CPU run.
+                log.debug(f"Constraining torch threads on {platform_name} to a single worker")
+                torch.set_num_threads(1)
+
+            return SentenceTransformer(model, trust_remote_code=True)
+
+        loaded_model = await asyncio.to_thread(_load_model)
+        EMBEDDING_MODELS[model] = loaded_model
+        return loaded_model

llama_stack/providers/utils/inference/inference_store.py

@@ -0,0 +1,264 @@
+# 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
+from typing import Any
+
+from sqlalchemy.exc import IntegrityError
+
+from llama_stack.apis.inference import (
+    ListOpenAIChatCompletionResponse,
+    OpenAIChatCompletion,
+    OpenAICompletionWithInputMessages,
+    OpenAIMessageParam,
+    Order,
+)
+from llama_stack.core.datatypes import AccessRule
+from llama_stack.core.storage.datatypes import InferenceStoreReference, StorageBackendType
+from llama_stack.log import get_logger
+
+from ..sqlstore.api import ColumnDefinition, ColumnType
+from ..sqlstore.authorized_sqlstore import AuthorizedSqlStore
+from ..sqlstore.sqlstore import _SQLSTORE_BACKENDS, sqlstore_impl
+
+logger = get_logger(name=__name__, category="inference")
+
+
+class InferenceStore:
+    def __init__(
+        self,
+        reference: InferenceStoreReference,
+        policy: list[AccessRule],
+    ):
+        self.reference = reference
+        self.sql_store = None
+        self.policy = policy
+        self.enable_write_queue = True
+
+        # Async write queue and worker control
+        self._queue: asyncio.Queue[tuple[OpenAIChatCompletion, list[OpenAIMessageParam]]] | None = None
+        self._worker_tasks: list[asyncio.Task[Any]] = []
+        self._max_write_queue_size: int = reference.max_write_queue_size
+        self._num_writers: int = max(1, reference.num_writers)
+
+    async def initialize(self):
+        """Create the necessary tables if they don't exist."""
+        base_store = sqlstore_impl(self.reference)
+        self.sql_store = AuthorizedSqlStore(base_store, self.policy)
+
+        # Disable write queue for SQLite since WAL mode handles concurrency
+        # Keep it enabled for other backends (like Postgres) for performance
+        backend_config = _SQLSTORE_BACKENDS.get(self.reference.backend)
+        if backend_config and backend_config.type == StorageBackendType.SQL_SQLITE:
+            self.enable_write_queue = False
+            logger.debug("Write queue disabled for SQLite (WAL mode handles concurrency)")
+
+        await self.sql_store.create_table(
+            "chat_completions",
+            {
+                "id": ColumnDefinition(type=ColumnType.STRING, primary_key=True),
+                "created": ColumnType.INTEGER,
+                "model": ColumnType.STRING,
+                "choices": ColumnType.JSON,
+                "input_messages": ColumnType.JSON,
+            },
+        )
+
+        if self.enable_write_queue:
+            self._queue = asyncio.Queue(maxsize=self._max_write_queue_size)
+            for _ in range(self._num_writers):
+                self._worker_tasks.append(asyncio.create_task(self._worker_loop()))
+            logger.debug(
+                f"Inference store write queue enabled with {self._num_writers} writers, max queue size {self._max_write_queue_size}"
+            )
+
+    async def shutdown(self) -> None:
+        if not self._worker_tasks:
+            return
+        if self._queue is not None:
+            await self._queue.join()
+        for t in self._worker_tasks:
+            if not t.done():
+                t.cancel()
+        for t in self._worker_tasks:
+            try:
+                await t
+            except asyncio.CancelledError:
+                pass
+        self._worker_tasks.clear()
+
+    async def flush(self) -> None:
+        """Wait for all queued writes to complete. Useful for testing."""
+        if self.enable_write_queue and self._queue is not None:
+            await self._queue.join()
+
+    async def _ensure_workers_started(self) -> None:
+        """Ensure the async write queue workers run on the current loop."""
+        if not self.enable_write_queue:
+            return
+
+        if self._queue is None:
+            self._queue = asyncio.Queue(maxsize=self._max_write_queue_size)
+            logger.debug(
+                f"Inference store write queue created with max size {self._max_write_queue_size} "
+                f"and {self._num_writers} writers"
+            )
+
+        if not self._worker_tasks:
+            loop = asyncio.get_running_loop()
+            for _ in range(self._num_writers):
+                task = loop.create_task(self._worker_loop())
+                self._worker_tasks.append(task)
+
+    async def store_chat_completion(
+        self, chat_completion: OpenAIChatCompletion, input_messages: list[OpenAIMessageParam]
+    ) -> None:
+        if self.enable_write_queue:
+            await self._ensure_workers_started()
+            if self._queue is None:
+                raise ValueError("Inference store is not initialized")
+            try:
+                self._queue.put_nowait((chat_completion, input_messages))
+            except asyncio.QueueFull:
+                logger.warning(
+                    f"Write queue full; adding chat completion id={getattr(chat_completion, 'id', '<unknown>')}"
+                )
+                await self._queue.put((chat_completion, input_messages))
+        else:
+            await self._write_chat_completion(chat_completion, input_messages)
+
+    async def _worker_loop(self) -> None:
+        assert self._queue is not None
+        while True:
+            try:
+                item = await self._queue.get()
+            except asyncio.CancelledError:
+                break
+            chat_completion, input_messages = item
+            try:
+                await self._write_chat_completion(chat_completion, input_messages)
+            except Exception as e:  # noqa: BLE001
+                logger.error(f"Error writing chat completion: {e}")
+            finally:
+                self._queue.task_done()
+
+    async def _write_chat_completion(
+        self, chat_completion: OpenAIChatCompletion, input_messages: list[OpenAIMessageParam]
+    ) -> None:
+        if self.sql_store is None:
+            raise ValueError("Inference store is not initialized")
+
+        data = chat_completion.model_dump()
+        record_data = {
+            "id": data["id"],
+            "created": data["created"],
+            "model": data["model"],
+            "choices": data["choices"],
+            "input_messages": [message.model_dump() for message in input_messages],
+        }
+
+        try:
+            await self.sql_store.insert(
+                table="chat_completions",
+                data=record_data,
+            )
+        except IntegrityError as e:
+            # Duplicate chat completion IDs can be generated during tests especially if they are replaying
+            # recorded responses across different tests. No need to warn or error under those circumstances.
+            # In the wild, this is not likely to happen at all (no evidence) so we aren't really hiding any problem.
+
+            # Check if it's a unique constraint violation
+            error_message = str(e.orig) if e.orig else str(e)
+            if self._is_unique_constraint_error(error_message):
+                # Update the existing record instead
+                await self.sql_store.update(table="chat_completions", data=record_data, where={"id": data["id"]})
+            else:
+                # Re-raise if it's not a unique constraint error
+                raise
+
+    def _is_unique_constraint_error(self, error_message: str) -> bool:
+        """Check if the error is specifically a unique constraint violation."""
+        error_lower = error_message.lower()
+        return any(
+            indicator in error_lower
+            for indicator in [
+                "unique constraint failed",  # SQLite
+                "duplicate key",  # PostgreSQL
+                "unique violation",  # PostgreSQL alternative
+                "duplicate entry",  # MySQL
+            ]
+        )
+
+    async def list_chat_completions(
+        self,
+        after: str | None = None,
+        limit: int | None = 50,
+        model: str | None = None,
+        order: Order | None = Order.desc,
+    ) -> ListOpenAIChatCompletionResponse:
+        """
+        List chat completions from the database.
+
+        :param after: The ID of the last chat completion to return.
+        :param limit: The maximum number of chat completions to return.
+        :param model: The model to filter by.
+        :param order: The order to sort the chat completions by.
+        """
+        if not self.sql_store:
+            raise ValueError("Inference store is not initialized")
+
+        if not order:
+            order = Order.desc
+
+        where_conditions = {}
+        if model:
+            where_conditions["model"] = model
+
+        paginated_result = await self.sql_store.fetch_all(
+            table="chat_completions",
+            where=where_conditions if where_conditions else None,
+            order_by=[("created", order.value)],
+            cursor=("id", after) if after else None,
+            limit=limit,
+        )
+
+        data = [
+            OpenAICompletionWithInputMessages(
+                id=row["id"],
+                created=row["created"],
+                model=row["model"],
+                choices=row["choices"],
+                input_messages=row["input_messages"],
+            )
+            for row in paginated_result.data
+        ]
+        return ListOpenAIChatCompletionResponse(
+            data=data,
+            has_more=paginated_result.has_more,
+            first_id=data[0].id if data else "",
+            last_id=data[-1].id if data else "",
+        )
+
+    async def get_chat_completion(self, completion_id: str) -> OpenAICompletionWithInputMessages:
+        if not self.sql_store:
+            raise ValueError("Inference store is not initialized")
+
+        row = await self.sql_store.fetch_one(
+            table="chat_completions",
+            where={"id": completion_id},
+        )
+
+        if not row:
+            # SecureSqlStore will return None if record doesn't exist OR access is denied
+            # This provides security by not revealing whether the record exists
+            raise ValueError(f"Chat completion with id {completion_id} not found") from None
+
+        return OpenAICompletionWithInputMessages(
+            id=row["id"],
+            created=row["created"],
+            model=row["model"],
+            choices=row["choices"],
+            input_messages=row["input_messages"],
+        )