Sync updates from stainless branch: ashwinb/dev

src/llama_stack_client/resources/safety.py

@@ -23,6 +23,7 @@
 )
 from .._base_client import make_request_options
 from ..types.run_shield_response import RunShieldResponse
+from ..types.shared_params.message import Message
 
 __all__ = ["SafetyResource", "AsyncSafetyResource"]
 
@@ -50,7 +51,7 @@ def with_streaming_response(self) -> SafetyResourceWithStreamingResponse:
     def run_shield(
         self,
         *,
-        messages: Iterable[safety_run_shield_params.Message],
+        messages: Iterable[Message],
         params: Dict[str, Union[bool, float, str, Iterable[object], object, None]],
         shield_id: str,
         x_llama_stack_client_version: str | NotGiven = NOT_GIVEN,
@@ -121,7 +122,7 @@ def with_streaming_response(self) -> AsyncSafetyResourceWithStreamingResponse:
     async def run_shield(
         self,
         *,
-        messages: Iterable[safety_run_shield_params.Message],
+        messages: Iterable[Message],
         params: Dict[str, Union[bool, float, str, Iterable[object], object, None]],
         shield_id: str,
         x_llama_stack_client_version: str | NotGiven = NOT_GIVEN,

src/llama_stack_client/resources/tool_runtime/rag_tool.py

@@ -22,9 +22,9 @@
 )
 from ..._base_client import make_request_options
 from ...types.tool_runtime import rag_tool_query_params, rag_tool_insert_params
-from ...types.tool_runtime.query_result import QueryResult
-from ...types.tool_runtime.document_param import DocumentParam
-from ...types.tool_runtime.query_config_param import QueryConfigParam
+from ...types.shared.query_result import QueryResult
+from ...types.shared_params.document import Document
+from ...types.shared_params.query_config import QueryConfig
 from ...types.shared_params.interleaved_content import InterleavedContent
 
 __all__ = ["RagToolResource", "AsyncRagToolResource"]
@@ -54,7 +54,7 @@ def insert(
         self,
         *,
         chunk_size_in_tokens: int,
-        documents: Iterable[DocumentParam],
+        documents: Iterable[Document],
         vector_db_id: str,
         x_llama_stack_client_version: str | NotGiven = NOT_GIVEN,
         x_llama_stack_provider_data: str | NotGiven = NOT_GIVEN,
@@ -108,7 +108,7 @@ def query(
         *,
         content: InterleavedContent,
         vector_db_ids: List[str],
-        query_config: QueryConfigParam | NotGiven = NOT_GIVEN,
+        query_config: QueryConfig | NotGiven = NOT_GIVEN,
         x_llama_stack_client_version: str | NotGiven = NOT_GIVEN,
         x_llama_stack_provider_data: str | NotGiven = NOT_GIVEN,
         # Use the following arguments if you need to pass additional parameters to the API that aren't available via kwargs.
@@ -180,7 +180,7 @@ async def insert(
         self,
         *,
         chunk_size_in_tokens: int,
-        documents: Iterable[DocumentParam],
+        documents: Iterable[Document],
         vector_db_id: str,
         x_llama_stack_client_version: str | NotGiven = NOT_GIVEN,
         x_llama_stack_provider_data: str | NotGiven = NOT_GIVEN,
@@ -234,7 +234,7 @@ async def query(
         *,
         content: InterleavedContent,
         vector_db_ids: List[str],
-        query_config: QueryConfigParam | NotGiven = NOT_GIVEN,
+        query_config: QueryConfig | NotGiven = NOT_GIVEN,
         x_llama_stack_client_version: str | NotGiven = NOT_GIVEN,
         x_llama_stack_provider_data: str | NotGiven = NOT_GIVEN,
         # Use the following arguments if you need to pass additional parameters to the API that aren't available via kwargs.

src/llama_stack_client/types/__init__.py

@@ -9,10 +9,13 @@
 from .shared import (
     URL as URL,
     Message as Message,
+    Document as Document,
     ToolCall as ToolCall,
     ParamType as ParamType,
     ReturnType as ReturnType,
     AgentConfig as AgentConfig,
+    QueryConfig as QueryConfig,
+    QueryResult as QueryResult,
     UserMessage as UserMessage,
     ContentDelta as ContentDelta,
     ScoringResult as ScoringResult,

src/llama_stack_client/types/agents/step_retrieve_response.py

@@ -1,8 +1,9 @@
 # File generated from our OpenAPI spec by Stainless. See CONTRIBUTING.md for details.
 
 from typing import Union
-from typing_extensions import TypeAlias
+from typing_extensions import Annotated, TypeAlias
 
+from ..._utils import PropertyInfo
 from ..._models import BaseModel
 from ..inference_step import InferenceStep
 from ..shield_call_step import ShieldCallStep
@@ -11,7 +12,10 @@
 
 __all__ = ["StepRetrieveResponse", "Step"]
 
-Step: TypeAlias = Union[InferenceStep, ToolExecutionStep, ShieldCallStep, MemoryRetrievalStep]
+Step: TypeAlias = Annotated[
+    Union[InferenceStep, ToolExecutionStep, ShieldCallStep, MemoryRetrievalStep],
+    PropertyInfo(discriminator="step_type"),
+]
 
 
 class StepRetrieveResponse(BaseModel):

src/llama_stack_client/types/agents/turn.py

@@ -2,8 +2,9 @@
 
 from typing import List, Union, Optional
 from datetime import datetime
-from typing_extensions import Literal, TypeAlias
+from typing_extensions import Literal, Annotated, TypeAlias
 
+from ..._utils import PropertyInfo
 from ..._models import BaseModel
 from ..shared.url import URL
 from ..inference_step import InferenceStep
@@ -62,7 +63,10 @@ class OutputAttachment(BaseModel):
     mime_type: str
 
 
-Step: TypeAlias = Union[InferenceStep, ToolExecutionStep, ShieldCallStep, MemoryRetrievalStep]
+Step: TypeAlias = Annotated[
+    Union[InferenceStep, ToolExecutionStep, ShieldCallStep, MemoryRetrievalStep],
+    PropertyInfo(discriminator="step_type"),
+]
 
 
 class Turn(BaseModel):

src/llama_stack_client/types/agents/turn_create_response.py

@@ -1,9 +1,10 @@
 # File generated from our OpenAPI spec by Stainless. See CONTRIBUTING.md for details.
 
 from typing import Dict, List, Union, Optional
-from typing_extensions import Literal, TypeAlias
+from typing_extensions import Literal, Annotated, TypeAlias
 
 from .turn import Turn
+from ..._utils import PropertyInfo
 from ..._models import BaseModel
 from ..inference_step import InferenceStep
 from ..shield_call_step import ShieldCallStep
@@ -16,16 +17,16 @@
     "AgentTurnResponseStreamChunk",
     "AgentTurnResponseStreamChunkEvent",
     "AgentTurnResponseStreamChunkEventPayload",
-    "AgentTurnResponseStreamChunkEventPayloadAgentTurnResponseStepStartPayload",
-    "AgentTurnResponseStreamChunkEventPayloadAgentTurnResponseStepProgressPayload",
-    "AgentTurnResponseStreamChunkEventPayloadAgentTurnResponseStepCompletePayload",
-    "AgentTurnResponseStreamChunkEventPayloadAgentTurnResponseStepCompletePayloadStepDetails",
-    "AgentTurnResponseStreamChunkEventPayloadAgentTurnResponseTurnStartPayload",
-    "AgentTurnResponseStreamChunkEventPayloadAgentTurnResponseTurnCompletePayload",
+    "AgentTurnResponseStreamChunkEventPayloadStepStart",
+    "AgentTurnResponseStreamChunkEventPayloadStepProgress",
+    "AgentTurnResponseStreamChunkEventPayloadStepComplete",
+    "AgentTurnResponseStreamChunkEventPayloadStepCompleteStepDetails",
+    "AgentTurnResponseStreamChunkEventPayloadTurnStart",
+    "AgentTurnResponseStreamChunkEventPayloadTurnComplete",
 ]
 
 
-class AgentTurnResponseStreamChunkEventPayloadAgentTurnResponseStepStartPayload(BaseModel):
+class AgentTurnResponseStreamChunkEventPayloadStepStart(BaseModel):
     event_type: Literal["step_start"]
 
     step_id: str
@@ -35,7 +36,7 @@ class AgentTurnResponseStreamChunkEventPayloadAgentTurnResponseStepStartPayload(
     metadata: Optional[Dict[str, Union[bool, float, str, List[object], object, None]]] = None
 
 
-class AgentTurnResponseStreamChunkEventPayloadAgentTurnResponseStepProgressPayload(BaseModel):
+class AgentTurnResponseStreamChunkEventPayloadStepProgress(BaseModel):
     delta: ContentDelta
 
     event_type: Literal["step_progress"]
@@ -45,39 +46,43 @@ class AgentTurnResponseStreamChunkEventPayloadAgentTurnResponseStepProgressPaylo
     step_type: Literal["inference", "tool_execution", "shield_call", "memory_retrieval"]
 
 
-AgentTurnResponseStreamChunkEventPayloadAgentTurnResponseStepCompletePayloadStepDetails: TypeAlias = Union[
-    InferenceStep, ToolExecutionStep, ShieldCallStep, MemoryRetrievalStep
+AgentTurnResponseStreamChunkEventPayloadStepCompleteStepDetails: TypeAlias = Annotated[
+    Union[InferenceStep, ToolExecutionStep, ShieldCallStep, MemoryRetrievalStep],
+    PropertyInfo(discriminator="step_type"),
 ]
 
 
-class AgentTurnResponseStreamChunkEventPayloadAgentTurnResponseStepCompletePayload(BaseModel):
+class AgentTurnResponseStreamChunkEventPayloadStepComplete(BaseModel):
     event_type: Literal["step_complete"]
 
-    step_details: AgentTurnResponseStreamChunkEventPayloadAgentTurnResponseStepCompletePayloadStepDetails
+    step_details: AgentTurnResponseStreamChunkEventPayloadStepCompleteStepDetails
 
     step_id: str
 
     step_type: Literal["inference", "tool_execution", "shield_call", "memory_retrieval"]
 
 
-class AgentTurnResponseStreamChunkEventPayloadAgentTurnResponseTurnStartPayload(BaseModel):
+class AgentTurnResponseStreamChunkEventPayloadTurnStart(BaseModel):
     event_type: Literal["turn_start"]
 
     turn_id: str
 
 
-class AgentTurnResponseStreamChunkEventPayloadAgentTurnResponseTurnCompletePayload(BaseModel):
+class AgentTurnResponseStreamChunkEventPayloadTurnComplete(BaseModel):
     event_type: Literal["turn_complete"]
 
     turn: Turn
 
 
-AgentTurnResponseStreamChunkEventPayload: TypeAlias = Union[
-    AgentTurnResponseStreamChunkEventPayloadAgentTurnResponseStepStartPayload,
-    AgentTurnResponseStreamChunkEventPayloadAgentTurnResponseStepProgressPayload,
-    AgentTurnResponseStreamChunkEventPayloadAgentTurnResponseStepCompletePayload,
-    AgentTurnResponseStreamChunkEventPayloadAgentTurnResponseTurnStartPayload,
-    AgentTurnResponseStreamChunkEventPayloadAgentTurnResponseTurnCompletePayload,
+AgentTurnResponseStreamChunkEventPayload: TypeAlias = Annotated[
+    Union[
+        AgentTurnResponseStreamChunkEventPayloadStepStart,
+        AgentTurnResponseStreamChunkEventPayloadStepProgress,
+        AgentTurnResponseStreamChunkEventPayloadStepComplete,
+        AgentTurnResponseStreamChunkEventPayloadTurnStart,
+        AgentTurnResponseStreamChunkEventPayloadTurnComplete,
+    ],
+    PropertyInfo(discriminator="event_type"),
 ]
 
 

src/llama_stack_client/types/eval_evaluate_rows_params.py

@@ -13,18 +13,18 @@
 __all__ = [
     "EvalEvaluateRowsParams",
     "TaskConfig",
-    "TaskConfigBenchmarkEvalTaskConfig",
-    "TaskConfigBenchmarkEvalTaskConfigEvalCandidate",
-    "TaskConfigBenchmarkEvalTaskConfigEvalCandidateModelCandidate",
-    "TaskConfigBenchmarkEvalTaskConfigEvalCandidateAgentCandidate",
-    "TaskConfigAppEvalTaskConfig",
-    "TaskConfigAppEvalTaskConfigEvalCandidate",
-    "TaskConfigAppEvalTaskConfigEvalCandidateModelCandidate",
-    "TaskConfigAppEvalTaskConfigEvalCandidateAgentCandidate",
-    "TaskConfigAppEvalTaskConfigScoringParams",
-    "TaskConfigAppEvalTaskConfigScoringParamsLlmAsJudgeScoringFnParams",
-    "TaskConfigAppEvalTaskConfigScoringParamsRegexParserScoringFnParams",
-    "TaskConfigAppEvalTaskConfigScoringParamsBasicScoringFnParams",
+    "TaskConfigBenchmark",
+    "TaskConfigBenchmarkEvalCandidate",
+    "TaskConfigBenchmarkEvalCandidateModel",
+    "TaskConfigBenchmarkEvalCandidateAgent",
+    "TaskConfigApp",
+    "TaskConfigAppEvalCandidate",
+    "TaskConfigAppEvalCandidateModel",
+    "TaskConfigAppEvalCandidateAgent",
+    "TaskConfigAppScoringParams",
+    "TaskConfigAppScoringParamsLlmAsJudge",
+    "TaskConfigAppScoringParamsRegexParser",
+    "TaskConfigAppScoringParamsBasic",
 ]
 
 
@@ -40,7 +40,7 @@ class EvalEvaluateRowsParams(TypedDict, total=False):
     x_llama_stack_provider_data: Annotated[str, PropertyInfo(alias="X-LlamaStack-Provider-Data")]
 
 
-class TaskConfigBenchmarkEvalTaskConfigEvalCandidateModelCandidate(TypedDict, total=False):
+class TaskConfigBenchmarkEvalCandidateModel(TypedDict, total=False):
     model: Required[str]
 
     sampling_params: Required[SamplingParams]
@@ -50,27 +50,26 @@ class TaskConfigBenchmarkEvalTaskConfigEvalCandidateModelCandidate(TypedDict, to
     system_message: SystemMessage
 
 
-class TaskConfigBenchmarkEvalTaskConfigEvalCandidateAgentCandidate(TypedDict, total=False):
+class TaskConfigBenchmarkEvalCandidateAgent(TypedDict, total=False):
     config: Required[AgentConfig]
 
     type: Required[Literal["agent"]]
 
 
-TaskConfigBenchmarkEvalTaskConfigEvalCandidate: TypeAlias = Union[
-    TaskConfigBenchmarkEvalTaskConfigEvalCandidateModelCandidate,
-    TaskConfigBenchmarkEvalTaskConfigEvalCandidateAgentCandidate,
+TaskConfigBenchmarkEvalCandidate: TypeAlias = Union[
+    TaskConfigBenchmarkEvalCandidateModel, TaskConfigBenchmarkEvalCandidateAgent
 ]
 
 
-class TaskConfigBenchmarkEvalTaskConfig(TypedDict, total=False):
-    eval_candidate: Required[TaskConfigBenchmarkEvalTaskConfigEvalCandidate]
+class TaskConfigBenchmark(TypedDict, total=False):
+    eval_candidate: Required[TaskConfigBenchmarkEvalCandidate]
 
     type: Required[Literal["benchmark"]]
 
     num_examples: int
 
 
-class TaskConfigAppEvalTaskConfigEvalCandidateModelCandidate(TypedDict, total=False):
+class TaskConfigAppEvalCandidateModel(TypedDict, total=False):
     model: Required[str]
 
     sampling_params: Required[SamplingParams]
@@ -80,18 +79,16 @@ class TaskConfigAppEvalTaskConfigEvalCandidateModelCandidate(TypedDict, total=Fa
     system_message: SystemMessage
 
 
-class TaskConfigAppEvalTaskConfigEvalCandidateAgentCandidate(TypedDict, total=False):
+class TaskConfigAppEvalCandidateAgent(TypedDict, total=False):
     config: Required[AgentConfig]
 
     type: Required[Literal["agent"]]
 
 
-TaskConfigAppEvalTaskConfigEvalCandidate: TypeAlias = Union[
-    TaskConfigAppEvalTaskConfigEvalCandidateModelCandidate, TaskConfigAppEvalTaskConfigEvalCandidateAgentCandidate
-]
+TaskConfigAppEvalCandidate: TypeAlias = Union[TaskConfigAppEvalCandidateModel, TaskConfigAppEvalCandidateAgent]
 
 
-class TaskConfigAppEvalTaskConfigScoringParamsLlmAsJudgeScoringFnParams(TypedDict, total=False):
+class TaskConfigAppScoringParamsLlmAsJudge(TypedDict, total=False):
     judge_model: Required[str]
 
     type: Required[Literal["llm_as_judge"]]
@@ -103,35 +100,33 @@ class TaskConfigAppEvalTaskConfigScoringParamsLlmAsJudgeScoringFnParams(TypedDic
     prompt_template: str
 
 
-class TaskConfigAppEvalTaskConfigScoringParamsRegexParserScoringFnParams(TypedDict, total=False):
+class TaskConfigAppScoringParamsRegexParser(TypedDict, total=False):
     type: Required[Literal["regex_parser"]]
 
     aggregation_functions: List[Literal["average", "median", "categorical_count", "accuracy"]]
 
     parsing_regexes: List[str]
 
 
-class TaskConfigAppEvalTaskConfigScoringParamsBasicScoringFnParams(TypedDict, total=False):
+class TaskConfigAppScoringParamsBasic(TypedDict, total=False):
     type: Required[Literal["basic"]]
 
     aggregation_functions: List[Literal["average", "median", "categorical_count", "accuracy"]]
 
 
-TaskConfigAppEvalTaskConfigScoringParams: TypeAlias = Union[
-    TaskConfigAppEvalTaskConfigScoringParamsLlmAsJudgeScoringFnParams,
-    TaskConfigAppEvalTaskConfigScoringParamsRegexParserScoringFnParams,
-    TaskConfigAppEvalTaskConfigScoringParamsBasicScoringFnParams,
+TaskConfigAppScoringParams: TypeAlias = Union[
+    TaskConfigAppScoringParamsLlmAsJudge, TaskConfigAppScoringParamsRegexParser, TaskConfigAppScoringParamsBasic
 ]
 
 
-class TaskConfigAppEvalTaskConfig(TypedDict, total=False):
-    eval_candidate: Required[TaskConfigAppEvalTaskConfigEvalCandidate]
+class TaskConfigApp(TypedDict, total=False):
+    eval_candidate: Required[TaskConfigAppEvalCandidate]
 
-    scoring_params: Required[Dict[str, TaskConfigAppEvalTaskConfigScoringParams]]
+    scoring_params: Required[Dict[str, TaskConfigAppScoringParams]]
 
     type: Required[Literal["app"]]
 
     num_examples: int
 
 
-TaskConfig: TypeAlias = Union[TaskConfigBenchmarkEvalTaskConfig, TaskConfigAppEvalTaskConfig]
+TaskConfig: TypeAlias = Union[TaskConfigBenchmark, TaskConfigApp]