Sync updates from stainless branch: ashwinb/dev

Author: ashwinb

src/llama_stack_client/resources/batch_inference.py

@@ -23,6 +23,7 @@
     async_to_streamed_response_wrapper,
 )
 from .._base_client import make_request_options
+from ..types.shared_params.message import Message
 from ..types.shared.batch_completion import BatchCompletion
 from ..types.shared_params.sampling_params import SamplingParams
 from ..types.shared_params.interleaved_content import InterleavedContent
@@ -54,7 +55,7 @@ def with_streaming_response(self) -> BatchInferenceResourceWithStreamingResponse
     def chat_completion(
         self,
         *,
-        messages_batch: Iterable[Iterable[batch_inference_chat_completion_params.MessagesBatch]],
+        messages_batch: Iterable[Iterable[Message]],
         model: str,
         logprobs: batch_inference_chat_completion_params.Logprobs | NotGiven = NOT_GIVEN,
         sampling_params: SamplingParams | NotGiven = NOT_GIVEN,
@@ -195,7 +196,7 @@ def with_streaming_response(self) -> AsyncBatchInferenceResourceWithStreamingRes
     async def chat_completion(
         self,
         *,
-        messages_batch: Iterable[Iterable[batch_inference_chat_completion_params.MessagesBatch]],
+        messages_batch: Iterable[Iterable[Message]],
         model: str,
         logprobs: batch_inference_chat_completion_params.Logprobs | NotGiven = NOT_GIVEN,
         sampling_params: SamplingParams | NotGiven = NOT_GIVEN,

src/llama_stack_client/resources/inference.py

@@ -30,6 +30,7 @@
 from .._streaming import Stream, AsyncStream
 from .._base_client import make_request_options
 from ..types.embeddings_response import EmbeddingsResponse
+from ..types.shared_params.message import Message
 from ..types.inference_completion_response import InferenceCompletionResponse
 from ..types.shared_params.sampling_params import SamplingParams
 from ..types.shared_params.interleaved_content import InterleavedContent
@@ -62,7 +63,7 @@ def with_streaming_response(self) -> InferenceResourceWithStreamingResponse:
     def chat_completion(
         self,
         *,
-        messages: Iterable[inference_chat_completion_params.Message],
+        messages: Iterable[Message],
         model_id: str,
         logprobs: inference_chat_completion_params.Logprobs | NotGiven = NOT_GIVEN,
         response_format: inference_chat_completion_params.ResponseFormat | NotGiven = NOT_GIVEN,
@@ -106,7 +107,7 @@ def chat_completion(
     def chat_completion(
         self,
         *,
-        messages: Iterable[inference_chat_completion_params.Message],
+        messages: Iterable[Message],
         model_id: str,
         stream: Literal[True],
         logprobs: inference_chat_completion_params.Logprobs | NotGiven = NOT_GIVEN,
@@ -150,7 +151,7 @@ def chat_completion(
     def chat_completion(
         self,
         *,
-        messages: Iterable[inference_chat_completion_params.Message],
+        messages: Iterable[Message],
         model_id: str,
         stream: bool,
         logprobs: inference_chat_completion_params.Logprobs | NotGiven = NOT_GIVEN,
@@ -194,7 +195,7 @@ def chat_completion(
     def chat_completion(
         self,
         *,
-        messages: Iterable[inference_chat_completion_params.Message],
+        messages: Iterable[Message],
         model_id: str,
         logprobs: inference_chat_completion_params.Logprobs | NotGiven = NOT_GIVEN,
         response_format: inference_chat_completion_params.ResponseFormat | NotGiven = NOT_GIVEN,
@@ -472,7 +473,7 @@ def with_streaming_response(self) -> AsyncInferenceResourceWithStreamingResponse
     async def chat_completion(
         self,
         *,
-        messages: Iterable[inference_chat_completion_params.Message],
+        messages: Iterable[Message],
         model_id: str,
         logprobs: inference_chat_completion_params.Logprobs | NotGiven = NOT_GIVEN,
         response_format: inference_chat_completion_params.ResponseFormat | NotGiven = NOT_GIVEN,
@@ -516,7 +517,7 @@ async def chat_completion(
     async def chat_completion(
         self,
         *,
-        messages: Iterable[inference_chat_completion_params.Message],
+        messages: Iterable[Message],
         model_id: str,
         stream: Literal[True],
         logprobs: inference_chat_completion_params.Logprobs | NotGiven = NOT_GIVEN,
@@ -560,7 +561,7 @@ async def chat_completion(
     async def chat_completion(
         self,
         *,
-        messages: Iterable[inference_chat_completion_params.Message],
+        messages: Iterable[Message],
         model_id: str,
         stream: bool,
         logprobs: inference_chat_completion_params.Logprobs | NotGiven = NOT_GIVEN,
@@ -604,7 +605,7 @@ async def chat_completion(
     async def chat_completion(
         self,
         *,
-        messages: Iterable[inference_chat_completion_params.Message],
+        messages: Iterable[Message],
         model_id: str,
         logprobs: inference_chat_completion_params.Logprobs | NotGiven = NOT_GIVEN,
         response_format: inference_chat_completion_params.ResponseFormat | NotGiven = NOT_GIVEN,

src/llama_stack_client/resources/synthetic_data_generation.py

@@ -23,6 +23,7 @@
     async_to_streamed_response_wrapper,
 )
 from .._base_client import make_request_options
+from ..types.shared_params.message import Message
 from ..types.synthetic_data_generation_response import SyntheticDataGenerationResponse
 
 __all__ = ["SyntheticDataGenerationResource", "AsyncSyntheticDataGenerationResource"]
@@ -51,7 +52,7 @@ def with_streaming_response(self) -> SyntheticDataGenerationResourceWithStreamin
     def generate(
         self,
         *,
-        dialogs: Iterable[synthetic_data_generation_generate_params.Dialog],
+        dialogs: Iterable[Message],
         filtering_function: Literal["none", "random", "top_k", "top_p", "top_k_top_p", "sigmoid"],
         model: str | NotGiven = NOT_GIVEN,
         x_llama_stack_client_version: str | NotGiven = NOT_GIVEN,
@@ -122,7 +123,7 @@ def with_streaming_response(self) -> AsyncSyntheticDataGenerationResourceWithStr
     async def generate(
         self,
         *,
-        dialogs: Iterable[synthetic_data_generation_generate_params.Dialog],
+        dialogs: Iterable[Message],
         filtering_function: Literal["none", "random", "top_k", "top_p", "top_k_top_p", "sigmoid"],
         model: str | NotGiven = NOT_GIVEN,
         x_llama_stack_client_version: str | NotGiven = NOT_GIVEN,

src/llama_stack_client/types/__init__.py

@@ -8,6 +8,7 @@
 from .trace import Trace as Trace
 from .shared import (
     URL as URL,
+    Message as Message,
     ToolCall as ToolCall,
     ParamType as ParamType,
     ReturnType as ReturnType,
@@ -19,6 +20,7 @@
     SamplingParams as SamplingParams,
     BatchCompletion as BatchCompletion,
     SafetyViolation as SafetyViolation,
+    CompletionMessage as CompletionMessage,
     InterleavedContent as InterleavedContent,
     ToolParamDefinition as ToolParamDefinition,
     ToolResponseMessage as ToolResponseMessage,

src/llama_stack_client/types/agents/turn.py

@@ -7,12 +7,11 @@
 from ..._models import BaseModel
 from ..shared.url import URL
 from ..inference_step import InferenceStep
-from ..shared.tool_call import ToolCall
 from ..shield_call_step import ShieldCallStep
 from ..shared.user_message import UserMessage
 from ..tool_execution_step import ToolExecutionStep
 from ..memory_retrieval_step import MemoryRetrievalStep
-from ..shared.interleaved_content import InterleavedContent
+from ..shared.completion_message import CompletionMessage
 from ..shared.tool_response_message import ToolResponseMessage
 from ..shared.interleaved_content_item import InterleavedContentItem
 
@@ -24,7 +23,6 @@
     "OutputAttachmentContentImageContentItem",
     "OutputAttachmentContentImageContentItemImage",
     "OutputAttachmentContentTextContentItem",
-    "OutputMessage",
     "Step",
 ]
 
@@ -64,16 +62,6 @@ class OutputAttachment(BaseModel):
     mime_type: str
 
 
-class OutputMessage(BaseModel):
-    content: InterleavedContent
-
-    role: Literal["assistant"]
-
-    stop_reason: Literal["end_of_turn", "end_of_message", "out_of_tokens"]
-
-    tool_calls: List[ToolCall]
-
-
 Step: TypeAlias = Union[InferenceStep, ToolExecutionStep, ShieldCallStep, MemoryRetrievalStep]
 
 
@@ -82,7 +70,7 @@ class Turn(BaseModel):
 
     output_attachments: List[OutputAttachment]
 
-    output_message: OutputMessage
+    output_message: CompletionMessage
 
     session_id: str
 

src/llama_stack_client/types/batch_inference_chat_completion_params.py

@@ -3,22 +3,18 @@
 from __future__ import annotations
 
 from typing import Dict, Union, Iterable
-from typing_extensions import Literal, Required, Annotated, TypeAlias, TypedDict
+from typing_extensions import Literal, Required, Annotated, TypedDict
 
 from .._utils import PropertyInfo
-from .shared_params.tool_call import ToolCall
-from .shared_params.user_message import UserMessage
-from .shared_params.system_message import SystemMessage
+from .shared_params.message import Message
 from .shared_params.sampling_params import SamplingParams
-from .shared_params.interleaved_content import InterleavedContent
 from .shared_params.tool_param_definition import ToolParamDefinition
-from .shared_params.tool_response_message import ToolResponseMessage
 
-__all__ = ["BatchInferenceChatCompletionParams", "MessagesBatch", "MessagesBatchCompletionMessage", "Logprobs", "Tool"]
+__all__ = ["BatchInferenceChatCompletionParams", "Logprobs", "Tool"]
 
 
 class BatchInferenceChatCompletionParams(TypedDict, total=False):
-    messages_batch: Required[Iterable[Iterable[MessagesBatch]]]
+    messages_batch: Required[Iterable[Iterable[Message]]]
 
     model: Required[str]
 
@@ -48,19 +44,6 @@ class BatchInferenceChatCompletionParams(TypedDict, total=False):
     x_llama_stack_provider_data: Annotated[str, PropertyInfo(alias="X-LlamaStack-Provider-Data")]
 
 
-class MessagesBatchCompletionMessage(TypedDict, total=False):
-    content: Required[InterleavedContent]
-
-    role: Required[Literal["assistant"]]
-
-    stop_reason: Required[Literal["end_of_turn", "end_of_message", "out_of_tokens"]]
-
-    tool_calls: Required[Iterable[ToolCall]]
-
-
-MessagesBatch: TypeAlias = Union[UserMessage, SystemMessage, ToolResponseMessage, MessagesBatchCompletionMessage]
-
-
 class Logprobs(TypedDict, total=False):
     top_k: int
 

src/llama_stack_client/types/batch_inference_chat_completion_response.py

@@ -1,24 +1,12 @@
 # File generated from our OpenAPI spec by Stainless. See CONTRIBUTING.md for details.
 
 from typing import List
-from typing_extensions import Literal
 
 from .._models import BaseModel
-from .shared.tool_call import ToolCall
-from .shared.interleaved_content import InterleavedContent
+from .shared.completion_message import CompletionMessage
 
-__all__ = ["BatchInferenceChatCompletionResponse", "CompletionMessageBatch"]
-
-
-class CompletionMessageBatch(BaseModel):
-    content: InterleavedContent
-
-    role: Literal["assistant"]
-
-    stop_reason: Literal["end_of_turn", "end_of_message", "out_of_tokens"]
-
-    tool_calls: List[ToolCall]
+__all__ = ["BatchInferenceChatCompletionResponse"]
 
 
 class BatchInferenceChatCompletionResponse(BaseModel):
-    completion_message_batch: List[CompletionMessageBatch]
+    completion_message_batch: List[CompletionMessage]

src/llama_stack_client/types/inference_chat_completion_params.py

@@ -6,18 +6,12 @@
 from typing_extensions import Literal, Required, Annotated, TypeAlias, TypedDict
 
 from .._utils import PropertyInfo
-from .shared_params.tool_call import ToolCall
-from .shared_params.user_message import UserMessage
-from .shared_params.system_message import SystemMessage
+from .shared_params.message import Message
 from .shared_params.sampling_params import SamplingParams
-from .shared_params.interleaved_content import InterleavedContent
 from .shared_params.tool_param_definition import ToolParamDefinition
-from .shared_params.tool_response_message import ToolResponseMessage
 
 __all__ = [
     "InferenceChatCompletionParamsBase",
-    "Message",
-    "MessageCompletionMessage",
     "Logprobs",
     "ResponseFormat",
     "ResponseFormatUnionMember0",
@@ -61,19 +55,6 @@ class InferenceChatCompletionParamsBase(TypedDict, total=False):
     x_llama_stack_provider_data: Annotated[str, PropertyInfo(alias="X-LlamaStack-Provider-Data")]
 
 
-class MessageCompletionMessage(TypedDict, total=False):
-    content: Required[InterleavedContent]
-
-    role: Required[Literal["assistant"]]
-
-    stop_reason: Required[Literal["end_of_turn", "end_of_message", "out_of_tokens"]]
-
-    tool_calls: Required[Iterable[ToolCall]]
-
-
-Message: TypeAlias = Union[UserMessage, SystemMessage, ToolResponseMessage, MessageCompletionMessage]
-
-
 class Logprobs(TypedDict, total=False):
     top_k: int
 

src/llama_stack_client/types/inference_chat_completion_response.py

@@ -5,31 +5,19 @@
 
 from .._models import BaseModel
 from .token_log_probs import TokenLogProbs
-from .shared.tool_call import ToolCall
 from .shared.content_delta import ContentDelta
-from .shared.interleaved_content import InterleavedContent
+from .shared.completion_message import CompletionMessage
 
 __all__ = [
     "InferenceChatCompletionResponse",
     "ChatCompletionResponse",
-    "ChatCompletionResponseCompletionMessage",
     "ChatCompletionResponseStreamChunk",
     "ChatCompletionResponseStreamChunkEvent",
 ]
 
 
-class ChatCompletionResponseCompletionMessage(BaseModel):
-    content: InterleavedContent
-
-    role: Literal["assistant"]
-
-    stop_reason: Literal["end_of_turn", "end_of_message", "out_of_tokens"]
-
-    tool_calls: List[ToolCall]
-
-
 class ChatCompletionResponse(BaseModel):
-    completion_message: ChatCompletionResponseCompletionMessage
+    completion_message: CompletionMessage
 
     logprobs: Optional[List[TokenLogProbs]] = None
 

src/llama_stack_client/types/inference_step.py

@@ -1,30 +1,19 @@
 # File generated from our OpenAPI spec by Stainless. See CONTRIBUTING.md for details.
 
-from typing import List, Optional
+from typing import Optional
 from datetime import datetime
 from typing_extensions import Literal
 
 from pydantic import Field as FieldInfo
 
 from .._models import BaseModel
-from .shared.tool_call import ToolCall
-from .shared.interleaved_content import InterleavedContent
+from .shared.completion_message import CompletionMessage
 
-__all__ = ["InferenceStep", "ModelResponse"]
-
-
-class ModelResponse(BaseModel):
-    content: InterleavedContent
-
-    role: Literal["assistant"]
-
-    stop_reason: Literal["end_of_turn", "end_of_message", "out_of_tokens"]
-
-    tool_calls: List[ToolCall]
+__all__ = ["InferenceStep"]
 
 
 class InferenceStep(BaseModel):
-    inference_model_response: ModelResponse = FieldInfo(alias="model_response")
+    inference_model_response: CompletionMessage = FieldInfo(alias="model_response")
 
     step_id: str