Merge branch 'main' into content-extension

llama_stack/apis/agents/openai_responses.py

@@ -623,6 +623,62 @@ class OpenAIResponseObjectStreamResponseMcpCallCompleted(BaseModel):
     type: Literal["response.mcp_call.completed"] = "response.mcp_call.completed"
 
 
+@json_schema_type
+class OpenAIResponseContentPartOutputText(BaseModel):
+    type: Literal["output_text"] = "output_text"
+    text: str
+    # TODO: add annotations, logprobs, etc.
+
+
+@json_schema_type
+class OpenAIResponseContentPartRefusal(BaseModel):
+    type: Literal["refusal"] = "refusal"
+    refusal: str
+
+
+OpenAIResponseContentPart = Annotated[
+    OpenAIResponseContentPartOutputText | OpenAIResponseContentPartRefusal,
+    Field(discriminator="type"),
+]
+register_schema(OpenAIResponseContentPart, name="OpenAIResponseContentPart")
+
+
+@json_schema_type
+class OpenAIResponseObjectStreamResponseContentPartAdded(BaseModel):
+    """Streaming event for when a new content part is added to a response item.
+
+    :param response_id: Unique identifier of the response containing this content
+    :param item_id: Unique identifier of the output item containing this content part
+    :param part: The content part that was added
+    :param sequence_number: Sequential number for ordering streaming events
+    :param type: Event type identifier, always "response.content_part.added"
+    """
+
+    response_id: str
+    item_id: str
+    part: OpenAIResponseContentPart
+    sequence_number: int
+    type: Literal["response.content_part.added"] = "response.content_part.added"
+
+
+@json_schema_type
+class OpenAIResponseObjectStreamResponseContentPartDone(BaseModel):
+    """Streaming event for when a content part is completed.
+
+    :param response_id: Unique identifier of the response containing this content
+    :param item_id: Unique identifier of the output item containing this content part
+    :param part: The completed content part
+    :param sequence_number: Sequential number for ordering streaming events
+    :param type: Event type identifier, always "response.content_part.done"
+    """
+
+    response_id: str
+    item_id: str
+    part: OpenAIResponseContentPart
+    sequence_number: int
+    type: Literal["response.content_part.done"] = "response.content_part.done"
+
+
 OpenAIResponseObjectStream = Annotated[
     OpenAIResponseObjectStreamResponseCreated
     | OpenAIResponseObjectStreamResponseOutputItemAdded
@@ -642,6 +698,8 @@ OpenAIResponseObjectStream = Annotated[
     | OpenAIResponseObjectStreamResponseMcpCallInProgress
     | OpenAIResponseObjectStreamResponseMcpCallFailed
     | OpenAIResponseObjectStreamResponseMcpCallCompleted
+    | OpenAIResponseObjectStreamResponseContentPartAdded
+    | OpenAIResponseObjectStreamResponseContentPartDone
     | OpenAIResponseObjectStreamResponseCompleted,
     Field(discriminator="type"),
 ]

llama_stack/apis/batches/__init__.py

@@ -0,0 +1,9 @@
+# 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 .batches import Batches, BatchObject, ListBatchesResponse
+
+__all__ = ["Batches", "BatchObject", "ListBatchesResponse"]

llama_stack/apis/batches/batches.py

@@ -0,0 +1,95 @@
+# 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 Literal, Protocol, runtime_checkable
+
+from pydantic import BaseModel, Field
+
+from llama_stack.schema_utils import json_schema_type, webmethod
+
+try:
+    from openai.types import Batch as BatchObject
+except ImportError as e:
+    raise ImportError("OpenAI package is required for batches API. Please install it with: pip install openai") from e
+
+
+@json_schema_type
+class ListBatchesResponse(BaseModel):
+    """Response containing a list of batch objects."""
+
+    object: Literal["list"] = "list"
+    data: list[BatchObject] = Field(..., description="List of batch objects")
+    first_id: str | None = Field(default=None, description="ID of the first batch in the list")
+    last_id: str | None = Field(default=None, description="ID of the last batch in the list")
+    has_more: bool = Field(default=False, description="Whether there are more batches available")
+
+
+@runtime_checkable
+class Batches(Protocol):
+    """
+    The Batches API enables efficient processing of multiple requests in a single operation,
+    particularly useful for processing large datasets, batch evaluation workflows, and
+    cost-effective inference at scale.
+
+    The API is designed to allow use of openai client libraries for seamless integration.
+
+    This API provides the following extensions:
+     - idempotent batch creation
+
+    Note: This API is currently under active development and may undergo changes.
+    """
+
+    @webmethod(route="/openai/v1/batches", method="POST")
+    async def create_batch(
+        self,
+        input_file_id: str,
+        endpoint: str,
+        completion_window: Literal["24h"],
+        metadata: dict[str, str] | None = None,
+        idempotency_key: str | None = None,
+    ) -> BatchObject:
+        """Create a new batch for processing multiple API requests.
+
+        :param input_file_id: The ID of an uploaded file containing requests for the batch.
+        :param endpoint: The endpoint to be used for all requests in the batch.
+        :param completion_window: The time window within which the batch should be processed.
+        :param metadata: Optional metadata for the batch.
+        :param idempotency_key: Optional idempotency key. When provided, enables idempotent behavior.
+        :returns: The created batch object.
+        """
+        ...
+
+    @webmethod(route="/openai/v1/batches/{batch_id}", method="GET")
+    async def retrieve_batch(self, batch_id: str) -> BatchObject:
+        """Retrieve information about a specific batch.
+
+        :param batch_id: The ID of the batch to retrieve.
+        :returns: The batch object.
+        """
+        ...
+
+    @webmethod(route="/openai/v1/batches/{batch_id}/cancel", method="POST")
+    async def cancel_batch(self, batch_id: str) -> BatchObject:
+        """Cancel a batch that is in progress.
+
+        :param batch_id: The ID of the batch to cancel.
+        :returns: The updated batch object.
+        """
+        ...
+
+    @webmethod(route="/openai/v1/batches", method="GET")
+    async def list_batches(
+        self,
+        after: str | None = None,
+        limit: int = 20,
+    ) -> ListBatchesResponse:
+        """List all batches for the current user.
+
+        :param after: A cursor for pagination; returns batches after this batch ID.
+        :param limit: Number of batches to return (default 20, max 100).
+        :returns: A list of batch objects.
+        """
+        ...

llama_stack/apis/common/errors.py

@@ -72,3 +72,10 @@ class ModelTypeError(TypeError):
             f"Model '{model_name}' is of type '{model_type}' rather than the expected type '{expected_model_type}'"
         )
         super().__init__(message)
+
+
+class ConflictError(ValueError):
+    """raised when an operation cannot be performed due to a conflict with the current state"""
+
+    def __init__(self, message: str) -> None:
+        super().__init__(message)

llama_stack/apis/datatypes.py

@@ -86,6 +86,7 @@ class Api(Enum, metaclass=DynamicApiMeta):
     :cvar inference: Text generation, chat completions, and embeddings
     :cvar safety: Content moderation and safety shields
     :cvar agents: Agent orchestration and execution
+    :cvar batches: Batch processing for asynchronous API requests
     :cvar vector_io: Vector database operations and queries
     :cvar datasetio: Dataset input/output operations
     :cvar scoring: Model output evaluation and scoring
@@ -108,6 +109,7 @@ class Api(Enum, metaclass=DynamicApiMeta):
     inference = "inference"
     safety = "safety"
     agents = "agents"
+    batches = "batches"
     vector_io = "vector_io"
     datasetio = "datasetio"
     scoring = "scoring"

llama_stack/apis/files/files.py

@@ -22,6 +22,7 @@ class OpenAIFilePurpose(StrEnum):
     """
 
     ASSISTANTS = "assistants"
+    BATCH = "batch"
     # TODO: Add other purposes as needed
 
 

llama_stack/apis/inference/inference.py

@@ -473,6 +473,28 @@ class EmbeddingsResponse(BaseModel):
     embeddings: list[list[float]]
 
 
+@json_schema_type
+class RerankData(BaseModel):
+    """A single rerank result from a reranking response.
+
+    :param index: The original index of the document in the input list
+    :param relevance_score: The relevance score from the model output. Values are inverted when applicable so that higher scores indicate greater relevance.
+    """
+
+    index: int
+    relevance_score: float
+
+
+@json_schema_type
+class RerankResponse(BaseModel):
+    """Response from a reranking request.
+
+    :param data: List of rerank result objects, sorted by relevance score (descending)
+    """
+
+    data: list[RerankData]
+
+
 @json_schema_type
 class OpenAIChatCompletionContentPartTextParam(BaseModel):
     """Text content part for OpenAI-compatible chat completion messages.
@@ -1046,6 +1068,7 @@ class InferenceProvider(Protocol):
         :returns: A BatchCompletionResponse with the full completions.
         """
         raise NotImplementedError("Batch completion is not implemented")
+        return  # this is so mypy's safe-super rule will consider the method concrete
 
     @webmethod(route="/inference/chat-completion", method="POST")
     async def chat_completion(
@@ -1110,6 +1133,7 @@ class InferenceProvider(Protocol):
         :returns: A BatchChatCompletionResponse with the full completions.
         """
         raise NotImplementedError("Batch chat completion is not implemented")
+        return  # this is so mypy's safe-super rule will consider the method concrete
 
     @webmethod(route="/inference/embeddings", method="POST")
     async def embeddings(
@@ -1131,6 +1155,25 @@ class InferenceProvider(Protocol):
         """
         ...
 
+    @webmethod(route="/inference/rerank", method="POST", experimental=True)
+    async def rerank(
+        self,
+        model: str,
+        query: str | OpenAIChatCompletionContentPartTextParam | OpenAIChatCompletionContentPartImageParam,
+        items: list[str | OpenAIChatCompletionContentPartTextParam | OpenAIChatCompletionContentPartImageParam],
+        max_num_results: int | None = None,
+    ) -> RerankResponse:
+        """Rerank a list of documents based on their relevance to a query.
+
+        :param model: The identifier of the reranking model to use.
+        :param query: The search query to rank items against. Can be a string, text content part, or image content part. The input must not exceed the model's max input token length.
+        :param items: List of items to rerank. Each item can be a string, text content part, or image content part. Each input must not exceed the model's max input token length.
+        :param max_num_results: (Optional) Maximum number of results to return. Default: returns all.
+        :returns: RerankResponse with indices sorted by relevance score (descending).
+        """
+        raise NotImplementedError("Reranking is not implemented")
+        return  # this is so mypy's safe-super rule will consider the method concrete
+
     @webmethod(route="/openai/v1/completions", method="POST")
     async def openai_completion(
         self,

llama_stack/apis/telemetry/telemetry.py

@@ -386,6 +386,7 @@ class MetricDataPoint(BaseModel):
 
     timestamp: int
     value: float
+    unit: str
 
 
 @json_schema_type
@@ -518,7 +519,7 @@ class Telemetry(Protocol):
         metric_name: str,
         start_time: int,
         end_time: int | None = None,
-        granularity: str | None = "1d",
+        granularity: str | None = None,
         query_type: MetricQueryType = MetricQueryType.RANGE,
         label_matchers: list[MetricLabelMatcher] | None = None,
     ) -> QueryMetricsResponse: