From 15dcc4ea5eac658801b051b28c2595159ad3e4d9 Mon Sep 17 00:00:00 2001 From: Xi Yan Date: Thu, 30 Jan 2025 18:03:02 -0800 Subject: [PATCH] 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. --- docs/openapi_generator/pyopenapi/generator.py | 35 +- docs/resources/llama-stack-spec.html | 500 +++++++++--------- docs/resources/llama-stack-spec.yaml | 317 +++++------ 3 files changed, 433 insertions(+), 419 deletions(-) diff --git a/docs/openapi_generator/pyopenapi/generator.py b/docs/openapi_generator/pyopenapi/generator.py index 202d3732b..05c609867 100644 --- a/docs/openapi_generator/pyopenapi/generator.py +++ b/docs/openapi_generator/pyopenapi/generator.py @@ -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 diff --git a/docs/resources/llama-stack-spec.html b/docs/resources/llama-stack-spec.html index 5b3771340..85ae516f5 100644 --- a/docs/resources/llama-stack-spec.html +++ b/docs/resources/llama-stack-spec.html @@ -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": { diff --git a/docs/resources/llama-stack-spec.yaml b/docs/resources/llama-stack-spec.yaml index 01232c001..2a95acf38 100644 --- a/docs/resources/llama-stack-spec.yaml +++ b/docs/resources/llama-stack-spec.yaml @@ -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: