llama-stack-api 0.4.4__py3-none-any.whl → 0.5.0rc1__py3-none-any.whl

llama_stack_api/shields/models.py

@@ -0,0 +1,74 @@
+# 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 Shields API requests and responses.
+
+This module defines the request and response models for the Shields API
+using Pydantic with Field descriptions for OpenAPI schema generation.
+"""
+
+from typing import Any, Literal
+
+from pydantic import BaseModel, Field
+
+from llama_stack_api.resource import Resource, ResourceType
+from llama_stack_api.schema_utils import json_schema_type
+
+
+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."""
+
+    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):
+    """Response containing a list of all shields."""
+
+    data: list[Shield] = Field(..., description="List of shield objects")
+
+
+@json_schema_type
+class GetShieldRequest(BaseModel):
+    """Request model for getting a shield by identifier."""
+
+    identifier: str = Field(..., description="The identifier of the shield to get.")
+
+
+@json_schema_type
+class RegisterShieldRequest(BaseModel):
+    """Request model for registering a shield."""
+
+    shield_id: str = Field(..., description="The identifier of the shield to register.")
+    provider_shield_id: str | None = Field(None, description="The identifier of the shield in the provider.")
+    provider_id: str | None = Field(None, description="The identifier of the provider.")
+    params: dict[str, Any] | None = Field(None, description="The parameters of the shield.")
+
+
+@json_schema_type
+class UnregisterShieldRequest(BaseModel):
+    """Request model for unregistering a shield."""
+
+    identifier: str = Field(..., description="The identifier of the shield to unregister.")

llama_stack_api/validators.py

@@ -0,0 +1,46 @@
+# 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.
+
+"""Validators for API request parameters.
+
+This module contains validation functions used by providers to validate
+request parameters that cannot be easily validated using Pydantic alone.
+"""
+
+from llama_stack_api.inference import OpenAIEmbeddingsRequestWithExtraBody
+
+
+def validate_embeddings_input_is_text(
+    params: OpenAIEmbeddingsRequestWithExtraBody,
+) -> None:
+    """
+    Validate that embeddings input contains only text strings, not token arrays.
+
+    Token arrays (list[int] and list[list[int]]) are a newer OpenAI feature
+    that is not universally supported across all embedding providers. This
+    validator should be called by providers that only support text input (str or list[str]).
+
+    :param params: The OpenAI embeddings request parameters
+    :raises ValueError: If input contains token arrays
+    """
+    # Valid: string input
+    if isinstance(params.input, str):
+        return
+
+    # Valid: list of strings
+    if isinstance(params.input, list) and isinstance(params.input[0], str):
+        return
+
+    # If we get here, input is a token array (list[int] or list[list[int]])
+    raise ValueError(
+        f"Model '{params.model}' does not support token arrays. "
+        f"Please provide text input as a string or list of strings instead."
+    )
+
+
+__all__ = [
+    "validate_embeddings_input_is_text",
+]

llama_stack_api/vector_io/__init__.py

@@ -0,0 +1,88 @@
+# 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.
+
+"""VectorIO API module.
+
+This module provides the VectorIO API for vector database operations.
+Protocol definitions are in api.py, Pydantic models are in models.py,
+and FastAPI routes are in fastapi_routes.py.
+"""
+
+# Re-export Protocol classes from api.py
+from . import fastapi_routes
+from .api import VectorIO, VectorStoreTable
+
+# Re-export all Pydantic models from models.py
+from .models import (
+    Chunk,
+    ChunkMetadata,
+    EmbeddedChunk,
+    OpenAICreateVectorStoreFileBatchRequestWithExtraBody,
+    OpenAICreateVectorStoreRequestWithExtraBody,
+    QueryChunksResponse,
+    SearchRankingOptions,
+    VectorStoreChunkingStrategy,
+    VectorStoreChunkingStrategyAuto,
+    VectorStoreChunkingStrategyStatic,
+    VectorStoreChunkingStrategyStaticConfig,
+    VectorStoreContent,
+    VectorStoreCreateRequest,
+    VectorStoreDeleteResponse,
+    VectorStoreFileAttributes,
+    VectorStoreFileBatchObject,
+    VectorStoreFileContentResponse,
+    VectorStoreFileCounts,
+    VectorStoreFileDeleteResponse,
+    VectorStoreFileLastError,
+    VectorStoreFileObject,
+    VectorStoreFilesListInBatchResponse,
+    VectorStoreFileStatus,
+    VectorStoreListFilesResponse,
+    VectorStoreListResponse,
+    VectorStoreModifyRequest,
+    VectorStoreObject,
+    VectorStoreSearchRequest,
+    VectorStoreSearchResponse,
+    VectorStoreSearchResponsePage,
+)
+
+__all__ = [
+    # Protocol classes
+    "VectorIO",
+    "VectorStoreTable",
+    # Pydantic models
+    "Chunk",
+    "ChunkMetadata",
+    "EmbeddedChunk",
+    "OpenAICreateVectorStoreFileBatchRequestWithExtraBody",
+    "OpenAICreateVectorStoreRequestWithExtraBody",
+    "QueryChunksResponse",
+    "SearchRankingOptions",
+    "VectorStoreChunkingStrategy",
+    "VectorStoreChunkingStrategyAuto",
+    "VectorStoreChunkingStrategyStatic",
+    "VectorStoreChunkingStrategyStaticConfig",
+    "VectorStoreContent",
+    "VectorStoreCreateRequest",
+    "VectorStoreDeleteResponse",
+    "VectorStoreFileAttributes",
+    "VectorStoreFileBatchObject",
+    "VectorStoreFileContentResponse",
+    "VectorStoreFileCounts",
+    "VectorStoreFileDeleteResponse",
+    "VectorStoreFileLastError",
+    "VectorStoreFileObject",
+    "VectorStoreFileStatus",
+    "VectorStoreFilesListInBatchResponse",
+    "VectorStoreListFilesResponse",
+    "VectorStoreListResponse",
+    "VectorStoreModifyRequest",
+    "VectorStoreObject",
+    "VectorStoreSearchRequest",
+    "VectorStoreSearchResponse",
+    "VectorStoreSearchResponsePage",
+    "fastapi_routes",
+]

llama_stack_api/vector_io/api.py

@@ -0,0 +1,234 @@
+# 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.
+
+"""VectorIO API protocol definition.
+
+This module contains the VectorIO protocol definition.
+Pydantic models are defined in llama_stack_api.vector_io.models.
+The FastAPI router is defined in llama_stack_api.vector_io.fastapi_routes.
+"""
+
+from typing import Annotated, Any, Protocol, runtime_checkable
+
+from fastapi import Body, Query
+
+from llama_stack_api.inference import InterleavedContent
+from llama_stack_api.vector_stores import VectorStore
+
+from .models import (
+    EmbeddedChunk,
+    OpenAICreateVectorStoreFileBatchRequestWithExtraBody,
+    OpenAICreateVectorStoreRequestWithExtraBody,
+    QueryChunksResponse,
+    SearchRankingOptions,
+    VectorStoreChunkingStrategy,
+    VectorStoreDeleteResponse,
+    VectorStoreFileBatchObject,
+    VectorStoreFileContentResponse,
+    VectorStoreFileDeleteResponse,
+    VectorStoreFileObject,
+    VectorStoreFilesListInBatchResponse,
+    VectorStoreFileStatus,
+    VectorStoreListFilesResponse,
+    VectorStoreListResponse,
+    VectorStoreObject,
+    VectorStoreSearchResponsePage,
+)
+
+
+class VectorStoreTable(Protocol):
+    def get_vector_store(self, vector_store_id: str) -> VectorStore | None: ...
+
+
+@runtime_checkable
+class VectorIO(Protocol):
+    vector_store_table: VectorStoreTable | None = None
+
+    # this will just block now until chunks are inserted, but it should
+    # probably return a Job instance which can be polled for completion
+    async def insert_chunks(
+        self,
+        vector_store_id: str,
+        chunks: list[EmbeddedChunk],
+        ttl_seconds: int | None = None,
+    ) -> None:
+        """Insert embedded chunks into a vector database."""
+        ...
+
+    async def query_chunks(
+        self,
+        vector_store_id: str,
+        query: InterleavedContent,
+        params: dict[str, Any] | None = None,
+    ) -> QueryChunksResponse:
+        """Query chunks from a vector database."""
+        ...
+
+    # OpenAI Vector Stores API endpoints
+    async def openai_create_vector_store(
+        self,
+        params: Annotated[OpenAICreateVectorStoreRequestWithExtraBody, Body(...)],
+    ) -> VectorStoreObject:
+        """Creates a vector store.
+
+        Generate an OpenAI-compatible vector store with the given parameters.
+        """
+        ...
+
+    async def openai_list_vector_stores(
+        self,
+        limit: int | None = 20,
+        order: str | None = "desc",
+        after: str | None = None,
+        before: str | None = None,
+    ) -> VectorStoreListResponse:
+        """Returns a list of vector stores."""
+        ...
+
+    async def openai_retrieve_vector_store(
+        self,
+        vector_store_id: str,
+    ) -> VectorStoreObject:
+        """Retrieves a vector store."""
+        ...
+
+    async def openai_update_vector_store(
+        self,
+        vector_store_id: str,
+        name: str | None = None,
+        expires_after: dict[str, Any] | None = None,
+        metadata: dict[str, Any] | None = None,
+    ) -> VectorStoreObject:
+        """Updates a vector store."""
+        ...
+
+    async def openai_delete_vector_store(
+        self,
+        vector_store_id: str,
+    ) -> VectorStoreDeleteResponse:
+        """Delete a vector store."""
+        ...
+
+    async def openai_search_vector_store(
+        self,
+        vector_store_id: str,
+        query: str | list[str],
+        filters: dict[str, Any] | None = None,
+        max_num_results: int | None = 10,
+        ranking_options: SearchRankingOptions | None = None,
+        rewrite_query: bool | None = False,
+        search_mode: (
+            str | None
+        ) = "vector",  # Using str instead of Literal due to OpenAPI schema generator limitations
+    ) -> VectorStoreSearchResponsePage:
+        """Search for chunks in a vector store.
+
+        Searches a vector store for relevant chunks based on a query and optional file attribute filters.
+        """
+        ...
+
+    async def openai_attach_file_to_vector_store(
+        self,
+        vector_store_id: str,
+        file_id: str,
+        attributes: dict[str, Any] | None = None,
+        chunking_strategy: VectorStoreChunkingStrategy | None = None,
+    ) -> VectorStoreFileObject:
+        """Attach a file to a vector store."""
+        ...
+
+    async def openai_list_files_in_vector_store(
+        self,
+        vector_store_id: str,
+        limit: int | None = 20,
+        order: str | None = "desc",
+        after: str | None = None,
+        before: str | None = None,
+        filter: VectorStoreFileStatus | None = None,
+    ) -> VectorStoreListFilesResponse:
+        """List files in a vector store."""
+        ...
+
+    async def openai_retrieve_vector_store_file(
+        self,
+        vector_store_id: str,
+        file_id: str,
+    ) -> VectorStoreFileObject:
+        """Retrieves a vector store file."""
+        ...
+
+    async def openai_retrieve_vector_store_file_contents(
+        self,
+        vector_store_id: str,
+        file_id: str,
+        include_embeddings: Annotated[bool | None, Query()] = False,
+        include_metadata: Annotated[bool | None, Query()] = False,
+    ) -> VectorStoreFileContentResponse:
+        """Retrieves the contents of a vector store file."""
+        ...
+
+    async def openai_update_vector_store_file(
+        self,
+        vector_store_id: str,
+        file_id: str,
+        attributes: dict[str, Any],
+    ) -> VectorStoreFileObject:
+        """Updates a vector store file."""
+        ...
+
+    async def openai_delete_vector_store_file(
+        self,
+        vector_store_id: str,
+        file_id: str,
+    ) -> VectorStoreFileDeleteResponse:
+        """Delete a vector store file."""
+        ...
+
+    async def openai_create_vector_store_file_batch(
+        self,
+        vector_store_id: str,
+        params: Annotated[OpenAICreateVectorStoreFileBatchRequestWithExtraBody, Body(...)],
+    ) -> VectorStoreFileBatchObject:
+        """Create a vector store file batch.
+
+        Generate an OpenAI-compatible vector store file batch for the given vector store.
+        """
+        ...
+
+    async def openai_retrieve_vector_store_file_batch(
+        self,
+        batch_id: str,
+        vector_store_id: str,
+    ) -> VectorStoreFileBatchObject:
+        """Retrieve a vector store file batch."""
+        ...
+
+    async def openai_list_files_in_vector_store_file_batch(
+        self,
+        batch_id: str,
+        vector_store_id: str,
+        after: str | None = None,
+        before: str | None = None,
+        filter: str | None = None,
+        limit: int | None = 20,
+        order: str | None = "desc",
+    ) -> VectorStoreFilesListInBatchResponse:
+        """Returns a list of vector store files in a batch."""
+        ...
+
+    async def openai_cancel_vector_store_file_batch(
+        self,
+        batch_id: str,
+        vector_store_id: str,
+    ) -> VectorStoreFileBatchObject:
+        """Cancels a vector store file batch."""
+        ...
+
+
+__all__ = [
+    "VectorIO",
+    "VectorStoreTable",
+]