openapi gen return type fix for streaming/non-streaming (#910)

# What does this PR do?

We need to change

```yaml
/v1/inference/chat-completion:
    post:
      responses:
        '200':
          description: >-
            If stream=False, returns a ChatCompletionResponse with the full completion.
            If stream=True, returns an SSE event stream of ChatCompletionResponseStreamChunk
          content:
            text/event-stream:
              schema:
                oneOf:
                  - $ref: '#/components/schemas/ChatCompletionResponse'
                  - $ref: '#/components/schemas/ChatCompletionResponseStreamChunk'
```

into

```yaml
/v1/inference/chat-completion:
    post:
      responses:
        '200':
          description: >-
            If stream=False, returns a ChatCompletionResponse with the full completion.
            If stream=True, returns an SSE event stream of ChatCompletionResponseStreamChunk
          content:
            text/event-stream:
              schema:
                $ref: '#/components/schemas/ChatCompletionResponseStreamChunk'
            application/json:
              schema:
                $ref: '#/components/schemas/ChatCompletionResponse'
```

## Test Plan

**Python**
- tested in SDK sync:
https://github.com/meta-llama/llama-stack-client-python/pull/108

**Node**
- tested w/
https://gist.github.com/yanxi0830/b782f4b91e21dcccdfef8898ce55157e (SDK
udpate follow up)


## Sources

Please link relevant resources if necessary.


## Before submitting

- [ ] This PR fixes a typo or improves the docs (you can dismiss the
other checks if that's the case).
- [ ] Ran pre-commit to handle lint / formatting issues.
- [ ] Read the [contributor
guideline](https://github.com/meta-llama/llama-stack/blob/main/CONTRIBUTING.md),
      Pull Request section?
- [ ] Updated relevant documentation.
- [ ] Wrote necessary unit or integration tests.
This commit is contained in:
Xi Yan 2025-01-30 18:03:02 -08:00 committed by GitHub
parent 2f11c7c203
commit 15dcc4ea5e
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
3 changed files with 433 additions and 419 deletions

View file

@ -177,20 +177,37 @@ class ContentBuilder:
) -> Dict[str, MediaType]:
"Creates the content subtree for a request or response."
def has_iterator_type(t):
if typing.get_origin(t) is typing.Union:
return any(has_iterator_type(a) for a in typing.get_args(t))
def is_iterator_type(t):
return "StreamChunk" in str(t)
def get_media_type(t):
if is_generic_list(t):
return "application/jsonl"
elif is_iterator_type(t):
return "text/event-stream"
else:
# TODO: needs a proper fix where we let all types correctly flow upwards
# and then test against AsyncIterator
return "StreamChunk" in str(t)
return "application/json"
if typing.get_origin(payload_type) is typing.Union:
media_types = []
item_types = []
for x in typing.get_args(payload_type):
media_types.append(get_media_type(x))
item_types.append(x)
if len(set(media_types)) == 1:
# all types have the same media type
return {media_types[0]: self.build_media_type(payload_type, examples)}
else:
# different types have different media types
return {
media_type: self.build_media_type(item_type, examples)
for media_type, item_type in zip(media_types, item_types)
}
if is_generic_list(payload_type):
media_type = "application/jsonl"
item_type = unwrap_generic_list(payload_type)
elif has_iterator_type(payload_type):
item_type = payload_type
media_type = "text/event-stream"
else:
media_type = "application/json"
item_type = payload_type

View file

@ -192,16 +192,14 @@
"200": {
"description": "If stream=False, returns a ChatCompletionResponse with the full completion. If stream=True, returns an SSE event stream of ChatCompletionResponseStreamChunk",
"content": {
"application/json": {
"schema": {
"$ref": "#/components/schemas/ChatCompletionResponse"
}
},
"text/event-stream": {
"schema": {
"oneOf": [
{
"$ref": "#/components/schemas/ChatCompletionResponse"
},
{
"$ref": "#/components/schemas/ChatCompletionResponseStreamChunk"
}
]
"$ref": "#/components/schemas/ChatCompletionResponseStreamChunk"
}
}
}
@ -230,16 +228,14 @@
"200": {
"description": "If stream=False, returns a CompletionResponse with the full completion. If stream=True, returns an SSE event stream of CompletionResponseStreamChunk",
"content": {
"application/json": {
"schema": {
"$ref": "#/components/schemas/CompletionResponse"
}
},
"text/event-stream": {
"schema": {
"oneOf": [
{
"$ref": "#/components/schemas/CompletionResponse"
},
{
"$ref": "#/components/schemas/CompletionResponseStreamChunk"
}
]
"$ref": "#/components/schemas/CompletionResponseStreamChunk"
}
}
}
@ -337,16 +333,14 @@
"200": {
"description": "A single turn in an interaction with an Agentic System. **OR** streamed agent turn completion response.",
"content": {
"application/json": {
"schema": {
"$ref": "#/components/schemas/Turn"
}
},
"text/event-stream": {
"schema": {
"oneOf": [
{
"$ref": "#/components/schemas/Turn"
},
{
"$ref": "#/components/schemas/AgentTurnResponseStreamChunk"
}
]
"$ref": "#/components/schemas/AgentTurnResponseStreamChunk"
}
}
}
@ -3780,235 +3774,6 @@
"messages"
]
},
"AgentTurnResponseEvent": {
"type": "object",
"properties": {
"payload": {
"$ref": "#/components/schemas/AgentTurnResponseEventPayload"
}
},
"additionalProperties": false,
"required": [
"payload"
]
},
"AgentTurnResponseEventPayload": {
"oneOf": [
{
"$ref": "#/components/schemas/AgentTurnResponseStepStartPayload"
},
{
"$ref": "#/components/schemas/AgentTurnResponseStepProgressPayload"
},
{
"$ref": "#/components/schemas/AgentTurnResponseStepCompletePayload"
},
{
"$ref": "#/components/schemas/AgentTurnResponseTurnStartPayload"
},
{
"$ref": "#/components/schemas/AgentTurnResponseTurnCompletePayload"
}
],
"discriminator": {
"propertyName": "event_type",
"mapping": {
"step_start": "#/components/schemas/AgentTurnResponseStepStartPayload",
"step_progress": "#/components/schemas/AgentTurnResponseStepProgressPayload",
"step_complete": "#/components/schemas/AgentTurnResponseStepCompletePayload",
"turn_start": "#/components/schemas/AgentTurnResponseTurnStartPayload",
"turn_complete": "#/components/schemas/AgentTurnResponseTurnCompletePayload"
}
}
},
"AgentTurnResponseStepCompletePayload": {
"type": "object",
"properties": {
"event_type": {
"type": "string",
"const": "step_complete",
"default": "step_complete"
},
"step_type": {
"type": "string",
"enum": [
"inference",
"tool_execution",
"shield_call",
"memory_retrieval"
]
},
"step_id": {
"type": "string"
},
"step_details": {
"oneOf": [
{
"$ref": "#/components/schemas/InferenceStep"
},
{
"$ref": "#/components/schemas/ToolExecutionStep"
},
{
"$ref": "#/components/schemas/ShieldCallStep"
},
{
"$ref": "#/components/schemas/MemoryRetrievalStep"
}
],
"discriminator": {
"propertyName": "step_type",
"mapping": {
"inference": "#/components/schemas/InferenceStep",
"tool_execution": "#/components/schemas/ToolExecutionStep",
"shield_call": "#/components/schemas/ShieldCallStep",
"memory_retrieval": "#/components/schemas/MemoryRetrievalStep"
}
}
}
},
"additionalProperties": false,
"required": [
"event_type",
"step_type",
"step_id",
"step_details"
]
},
"AgentTurnResponseStepProgressPayload": {
"type": "object",
"properties": {
"event_type": {
"type": "string",
"const": "step_progress",
"default": "step_progress"
},
"step_type": {
"type": "string",
"enum": [
"inference",
"tool_execution",
"shield_call",
"memory_retrieval"
]
},
"step_id": {
"type": "string"
},
"delta": {
"$ref": "#/components/schemas/ContentDelta"
}
},
"additionalProperties": false,
"required": [
"event_type",
"step_type",
"step_id",
"delta"
]
},
"AgentTurnResponseStepStartPayload": {
"type": "object",
"properties": {
"event_type": {
"type": "string",
"const": "step_start",
"default": "step_start"
},
"step_type": {
"type": "string",
"enum": [
"inference",
"tool_execution",
"shield_call",
"memory_retrieval"
]
},
"step_id": {
"type": "string"
},
"metadata": {
"type": "object",
"additionalProperties": {
"oneOf": [
{
"type": "null"
},
{
"type": "boolean"
},
{
"type": "number"
},
{
"type": "string"
},
{
"type": "array"
},
{
"type": "object"
}
]
}
}
},
"additionalProperties": false,
"required": [
"event_type",
"step_type",
"step_id"
]
},
"AgentTurnResponseStreamChunk": {
"type": "object",
"properties": {
"event": {
"$ref": "#/components/schemas/AgentTurnResponseEvent"
}
},
"additionalProperties": false,
"required": [
"event"
],
"title": "streamed agent turn completion response."
},
"AgentTurnResponseTurnCompletePayload": {
"type": "object",
"properties": {
"event_type": {
"type": "string",
"const": "turn_complete",
"default": "turn_complete"
},
"turn": {
"$ref": "#/components/schemas/Turn"
}
},
"additionalProperties": false,
"required": [
"event_type",
"turn"
]
},
"AgentTurnResponseTurnStartPayload": {
"type": "object",
"properties": {
"event_type": {
"type": "string",
"const": "turn_start",
"default": "turn_start"
},
"turn_id": {
"type": "string"
}
},
"additionalProperties": false,
"required": [
"event_type",
"turn_id"
]
},
"InferenceStep": {
"type": "object",
"properties": {
@ -4349,6 +4114,235 @@
"error"
]
},
"AgentTurnResponseEvent": {
"type": "object",
"properties": {
"payload": {
"$ref": "#/components/schemas/AgentTurnResponseEventPayload"
}
},
"additionalProperties": false,
"required": [
"payload"
]
},
"AgentTurnResponseEventPayload": {
"oneOf": [
{
"$ref": "#/components/schemas/AgentTurnResponseStepStartPayload"
},
{
"$ref": "#/components/schemas/AgentTurnResponseStepProgressPayload"
},
{
"$ref": "#/components/schemas/AgentTurnResponseStepCompletePayload"
},
{
"$ref": "#/components/schemas/AgentTurnResponseTurnStartPayload"
},
{
"$ref": "#/components/schemas/AgentTurnResponseTurnCompletePayload"
}
],
"discriminator": {
"propertyName": "event_type",
"mapping": {
"step_start": "#/components/schemas/AgentTurnResponseStepStartPayload",
"step_progress": "#/components/schemas/AgentTurnResponseStepProgressPayload",
"step_complete": "#/components/schemas/AgentTurnResponseStepCompletePayload",
"turn_start": "#/components/schemas/AgentTurnResponseTurnStartPayload",
"turn_complete": "#/components/schemas/AgentTurnResponseTurnCompletePayload"
}
}
},
"AgentTurnResponseStepCompletePayload": {
"type": "object",
"properties": {
"event_type": {
"type": "string",
"const": "step_complete",
"default": "step_complete"
},
"step_type": {
"type": "string",
"enum": [
"inference",
"tool_execution",
"shield_call",
"memory_retrieval"
]
},
"step_id": {
"type": "string"
},
"step_details": {
"oneOf": [
{
"$ref": "#/components/schemas/InferenceStep"
},
{
"$ref": "#/components/schemas/ToolExecutionStep"
},
{
"$ref": "#/components/schemas/ShieldCallStep"
},
{
"$ref": "#/components/schemas/MemoryRetrievalStep"
}
],
"discriminator": {
"propertyName": "step_type",
"mapping": {
"inference": "#/components/schemas/InferenceStep",
"tool_execution": "#/components/schemas/ToolExecutionStep",
"shield_call": "#/components/schemas/ShieldCallStep",
"memory_retrieval": "#/components/schemas/MemoryRetrievalStep"
}
}
}
},
"additionalProperties": false,
"required": [
"event_type",
"step_type",
"step_id",
"step_details"
]
},
"AgentTurnResponseStepProgressPayload": {
"type": "object",
"properties": {
"event_type": {
"type": "string",
"const": "step_progress",
"default": "step_progress"
},
"step_type": {
"type": "string",
"enum": [
"inference",
"tool_execution",
"shield_call",
"memory_retrieval"
]
},
"step_id": {
"type": "string"
},
"delta": {
"$ref": "#/components/schemas/ContentDelta"
}
},
"additionalProperties": false,
"required": [
"event_type",
"step_type",
"step_id",
"delta"
]
},
"AgentTurnResponseStepStartPayload": {
"type": "object",
"properties": {
"event_type": {
"type": "string",
"const": "step_start",
"default": "step_start"
},
"step_type": {
"type": "string",
"enum": [
"inference",
"tool_execution",
"shield_call",
"memory_retrieval"
]
},
"step_id": {
"type": "string"
},
"metadata": {
"type": "object",
"additionalProperties": {
"oneOf": [
{
"type": "null"
},
{
"type": "boolean"
},
{
"type": "number"
},
{
"type": "string"
},
{
"type": "array"
},
{
"type": "object"
}
]
}
}
},
"additionalProperties": false,
"required": [
"event_type",
"step_type",
"step_id"
]
},
"AgentTurnResponseStreamChunk": {
"type": "object",
"properties": {
"event": {
"$ref": "#/components/schemas/AgentTurnResponseEvent"
}
},
"additionalProperties": false,
"required": [
"event"
],
"title": "streamed agent turn completion response."
},
"AgentTurnResponseTurnCompletePayload": {
"type": "object",
"properties": {
"event_type": {
"type": "string",
"const": "turn_complete",
"default": "turn_complete"
},
"turn": {
"$ref": "#/components/schemas/Turn"
}
},
"additionalProperties": false,
"required": [
"event_type",
"turn"
]
},
"AgentTurnResponseTurnStartPayload": {
"type": "object",
"properties": {
"event_type": {
"type": "string",
"const": "turn_start",
"default": "turn_start"
},
"turn_id": {
"type": "string"
}
},
"additionalProperties": false,
"required": [
"event_type",
"turn_id"
]
},
"EmbeddingsRequest": {
"type": "object",
"properties": {

View file

@ -113,11 +113,12 @@ paths:
If stream=False, returns a ChatCompletionResponse with the full completion.
If stream=True, returns an SSE event stream of ChatCompletionResponseStreamChunk
content:
application/json:
schema:
$ref: '#/components/schemas/ChatCompletionResponse'
text/event-stream:
schema:
oneOf:
- $ref: '#/components/schemas/ChatCompletionResponse'
- $ref: '#/components/schemas/ChatCompletionResponseStreamChunk'
$ref: '#/components/schemas/ChatCompletionResponseStreamChunk'
tags:
- Inference
summary: >-
@ -137,11 +138,12 @@ paths:
If stream=False, returns a CompletionResponse with the full completion.
If stream=True, returns an SSE event stream of CompletionResponseStreamChunk
content:
application/json:
schema:
$ref: '#/components/schemas/CompletionResponse'
text/event-stream:
schema:
oneOf:
- $ref: '#/components/schemas/CompletionResponse'
- $ref: '#/components/schemas/CompletionResponseStreamChunk'
$ref: '#/components/schemas/CompletionResponseStreamChunk'
tags:
- Inference
summary: >-
@ -202,11 +204,12 @@ paths:
A single turn in an interaction with an Agentic System. **OR** streamed
agent turn completion response.
content:
application/json:
schema:
$ref: '#/components/schemas/Turn'
text/event-stream:
schema:
oneOf:
- $ref: '#/components/schemas/Turn'
- $ref: '#/components/schemas/AgentTurnResponseStreamChunk'
$ref: '#/components/schemas/AgentTurnResponseStreamChunk'
tags:
- Agents
parameters:
@ -2394,154 +2397,6 @@ components:
additionalProperties: false
required:
- messages
AgentTurnResponseEvent:
type: object
properties:
payload:
$ref: '#/components/schemas/AgentTurnResponseEventPayload'
additionalProperties: false
required:
- payload
AgentTurnResponseEventPayload:
oneOf:
- $ref: '#/components/schemas/AgentTurnResponseStepStartPayload'
- $ref: '#/components/schemas/AgentTurnResponseStepProgressPayload'
- $ref: '#/components/schemas/AgentTurnResponseStepCompletePayload'
- $ref: '#/components/schemas/AgentTurnResponseTurnStartPayload'
- $ref: '#/components/schemas/AgentTurnResponseTurnCompletePayload'
discriminator:
propertyName: event_type
mapping:
step_start: '#/components/schemas/AgentTurnResponseStepStartPayload'
step_progress: '#/components/schemas/AgentTurnResponseStepProgressPayload'
step_complete: '#/components/schemas/AgentTurnResponseStepCompletePayload'
turn_start: '#/components/schemas/AgentTurnResponseTurnStartPayload'
turn_complete: '#/components/schemas/AgentTurnResponseTurnCompletePayload'
AgentTurnResponseStepCompletePayload:
type: object
properties:
event_type:
type: string
const: step_complete
default: step_complete
step_type:
type: string
enum:
- inference
- tool_execution
- shield_call
- memory_retrieval
step_id:
type: string
step_details:
oneOf:
- $ref: '#/components/schemas/InferenceStep'
- $ref: '#/components/schemas/ToolExecutionStep'
- $ref: '#/components/schemas/ShieldCallStep'
- $ref: '#/components/schemas/MemoryRetrievalStep'
discriminator:
propertyName: step_type
mapping:
inference: '#/components/schemas/InferenceStep'
tool_execution: '#/components/schemas/ToolExecutionStep'
shield_call: '#/components/schemas/ShieldCallStep'
memory_retrieval: '#/components/schemas/MemoryRetrievalStep'
additionalProperties: false
required:
- event_type
- step_type
- step_id
- step_details
AgentTurnResponseStepProgressPayload:
type: object
properties:
event_type:
type: string
const: step_progress
default: step_progress
step_type:
type: string
enum:
- inference
- tool_execution
- shield_call
- memory_retrieval
step_id:
type: string
delta:
$ref: '#/components/schemas/ContentDelta'
additionalProperties: false
required:
- event_type
- step_type
- step_id
- delta
AgentTurnResponseStepStartPayload:
type: object
properties:
event_type:
type: string
const: step_start
default: step_start
step_type:
type: string
enum:
- inference
- tool_execution
- shield_call
- memory_retrieval
step_id:
type: string
metadata:
type: object
additionalProperties:
oneOf:
- type: 'null'
- type: boolean
- type: number
- type: string
- type: array
- type: object
additionalProperties: false
required:
- event_type
- step_type
- step_id
AgentTurnResponseStreamChunk:
type: object
properties:
event:
$ref: '#/components/schemas/AgentTurnResponseEvent'
additionalProperties: false
required:
- event
title: streamed agent turn completion response.
AgentTurnResponseTurnCompletePayload:
type: object
properties:
event_type:
type: string
const: turn_complete
default: turn_complete
turn:
$ref: '#/components/schemas/Turn'
additionalProperties: false
required:
- event_type
- turn
AgentTurnResponseTurnStartPayload:
type: object
properties:
event_type:
type: string
const: turn_start
default: turn_start
turn_id:
type: string
additionalProperties: false
required:
- event_type
- turn_id
InferenceStep:
type: object
properties:
@ -2765,6 +2620,154 @@ components:
- info
- warn
- error
AgentTurnResponseEvent:
type: object
properties:
payload:
$ref: '#/components/schemas/AgentTurnResponseEventPayload'
additionalProperties: false
required:
- payload
AgentTurnResponseEventPayload:
oneOf:
- $ref: '#/components/schemas/AgentTurnResponseStepStartPayload'
- $ref: '#/components/schemas/AgentTurnResponseStepProgressPayload'
- $ref: '#/components/schemas/AgentTurnResponseStepCompletePayload'
- $ref: '#/components/schemas/AgentTurnResponseTurnStartPayload'
- $ref: '#/components/schemas/AgentTurnResponseTurnCompletePayload'
discriminator:
propertyName: event_type
mapping:
step_start: '#/components/schemas/AgentTurnResponseStepStartPayload'
step_progress: '#/components/schemas/AgentTurnResponseStepProgressPayload'
step_complete: '#/components/schemas/AgentTurnResponseStepCompletePayload'
turn_start: '#/components/schemas/AgentTurnResponseTurnStartPayload'
turn_complete: '#/components/schemas/AgentTurnResponseTurnCompletePayload'
AgentTurnResponseStepCompletePayload:
type: object
properties:
event_type:
type: string
const: step_complete
default: step_complete
step_type:
type: string
enum:
- inference
- tool_execution
- shield_call
- memory_retrieval
step_id:
type: string
step_details:
oneOf:
- $ref: '#/components/schemas/InferenceStep'
- $ref: '#/components/schemas/ToolExecutionStep'
- $ref: '#/components/schemas/ShieldCallStep'
- $ref: '#/components/schemas/MemoryRetrievalStep'
discriminator:
propertyName: step_type
mapping:
inference: '#/components/schemas/InferenceStep'
tool_execution: '#/components/schemas/ToolExecutionStep'
shield_call: '#/components/schemas/ShieldCallStep'
memory_retrieval: '#/components/schemas/MemoryRetrievalStep'
additionalProperties: false
required:
- event_type
- step_type
- step_id
- step_details
AgentTurnResponseStepProgressPayload:
type: object
properties:
event_type:
type: string
const: step_progress
default: step_progress
step_type:
type: string
enum:
- inference
- tool_execution
- shield_call
- memory_retrieval
step_id:
type: string
delta:
$ref: '#/components/schemas/ContentDelta'
additionalProperties: false
required:
- event_type
- step_type
- step_id
- delta
AgentTurnResponseStepStartPayload:
type: object
properties:
event_type:
type: string
const: step_start
default: step_start
step_type:
type: string
enum:
- inference
- tool_execution
- shield_call
- memory_retrieval
step_id:
type: string
metadata:
type: object
additionalProperties:
oneOf:
- type: 'null'
- type: boolean
- type: number
- type: string
- type: array
- type: object
additionalProperties: false
required:
- event_type
- step_type
- step_id
AgentTurnResponseStreamChunk:
type: object
properties:
event:
$ref: '#/components/schemas/AgentTurnResponseEvent'
additionalProperties: false
required:
- event
title: streamed agent turn completion response.
AgentTurnResponseTurnCompletePayload:
type: object
properties:
event_type:
type: string
const: turn_complete
default: turn_complete
turn:
$ref: '#/components/schemas/Turn'
additionalProperties: false
required:
- event_type
- turn
AgentTurnResponseTurnStartPayload:
type: object
properties:
event_type:
type: string
const: turn_start
default: turn_start
turn_id:
type: string
additionalProperties: false
required:
- event_type
- turn_id
EmbeddingsRequest:
type: object
properties: