95a935fc48 · tduncan/Nvidia-Optimization-Stack-VLLM

@@ -2,6 +2,7 @@

2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project

3

4 import asyncio

5 +import json

6 import time

7 from collections.abc import AsyncGenerator, AsyncIterator

8 from contextlib import AsyncExitStack

@@ -10,10 +11,22 @@ from http import HTTPStatus

11 from typing import Any, Callable, Final, Optional, Union

12

13 import jinja2

14 +import openai.types.responses as openai_responses_types

15 from fastapi import Request

14 -from openai.types.responses import (ResponseFunctionToolCall,

15 - ResponseOutputItem, ResponseOutputMessage,

16 - ResponseOutputText, ResponseReasoningItem)

16 +from openai import BaseModel

17 +# yapf conflicts with isort for this block

18 +# yapf: disable

19 +from openai.types.responses import (ResponseContentPartDoneEvent,

20 + ResponseCreatedEvent,

21 + ResponseFunctionToolCall,

22 + ResponseInProgressEvent,

23 + ResponseOutputItem,

24 + ResponseOutputItemDoneEvent,

25 + ResponseOutputMessage, ResponseOutputText,

26 + ResponseReasoningItem,

27 + ResponseReasoningTextDeltaEvent,

28 + ResponseReasoningTextDoneEvent)

29 +# yapf: enable

30 from openai.types.responses.response_reasoning_item import (

31 Content as ResponseReasoningTextContent)

32 from openai_harmony import Message as OpenAIHarmonyMessage

@@ -330,8 +343,15 @@ class OpenAIServingResponses(OpenAIServing):

343 return response

344

345 if request.stream:

333 - raise NotImplementedError(

334 - "Streaming responses are not supported")

346 + return self.responses_stream_generator(

347 + request,

348 + sampling_params,

349 + result_generator,

350 + context,

351 + model_name,

352 + tokenizer,

353 + request_metadata,

354 + )

355

356 try:

357 return await self.responses_full_generator(

@@ -744,3 +764,423 @@ class OpenAIServingResponses(OpenAIServing):

764 "starting the vLLM server."),

765 status_code=HTTPStatus.BAD_REQUEST,

766 )

767 +

768 + async def responses_stream_generator(

769 + self,

770 + request: ResponsesRequest,

771 + sampling_params: SamplingParams,

772 + result_generator: AsyncIterator[Optional[ConversationContext]],

773 + context: ConversationContext,

774 + model_name: str,

775 + tokenizer: AnyTokenizer,

776 + request_metadata: RequestResponseMetadata,

777 + created_time: Optional[int] = None,

778 + ) -> AsyncGenerator[str, None]:

779 + # TODO:

780 + # 1. Handle disconnect

781 +

782 + if not isinstance(context, StreamingHarmonyContext):

783 + raise NotImplementedError(

784 + "Streaming is not supported for responses API without Harmony."

785 + )

786 +

787 + created_time = created_time or int(time.time())

788 +

789 + sequence_number = 0

790 +

791 + def _send_event(event: BaseModel):

792 + nonlocal sequence_number

793 + # Set sequence_number if the event has this attribute

794 + if hasattr(event, 'sequence_number'):

795 + event.sequence_number = sequence_number

796 + sequence_number += 1

797 + # Get event type from the event's type field if it exists

798 + event_type = getattr(event, 'type', 'unknown')

799 + return (f"event: {event_type}\n"

800 + f"data: {event.model_dump_json(indent=None)}\n\n")

801 +

802 + current_content_index = 0 # FIXME: this number is never changed

803 + current_output_index = 0

804 + current_item_id = "" # FIXME: this number is never changed

805 + sent_output_item_added = False

806 +

807 + initial_response = ResponsesResponse.from_request(

808 + request,

809 + sampling_params,

810 + model_name=model_name,

811 + created_time=created_time,

812 + output=[],

813 + status="in_progress",

814 + usage=None,

815 + ).model_dump()

816 + yield _send_event(

817 + ResponseCreatedEvent(

818 + type="response.created",

819 + sequence_number=-1,

820 + response=initial_response,

821 + ))

822 + yield _send_event(

823 + ResponseInProgressEvent(

824 + type="response.in_progress",

825 + sequence_number=-1,

826 + response=initial_response,

827 + ))

828 +

829 + async for ctx in result_generator:

830 +

831 + assert isinstance(ctx, StreamingHarmonyContext)

832 +

833 + if ctx.is_expecting_start():

834 + current_output_index += 1

835 + sent_output_item_added = False

836 +

837 + if len(ctx.parser.messages) > 0:

838 + previous_item = ctx.parser.messages[-1]

839 + if previous_item.recipient is not None:

840 + # Deal with tool call here

841 + pass

842 + elif previous_item.channel == "analysis":

843 + reasoning_item = ResponseReasoningItem(

844 + type="reasoning",

845 + content=[

846 + ResponseReasoningTextContent(

847 + text=previous_item.content[0].text),

848 + ],

849 + status="completed",

850 + )

851 + yield _send_event(

852 + ResponseReasoningTextDoneEvent(

853 + type="response.reasoning_text.done",

854 + item_id=current_item_id,

855 + sequence_number=-1,

856 + output_index=current_output_index,

857 + content_index=current_content_index,

858 + text=previous_item.content[0].text,

859 + ))

860 + yield _send_event(

861 + ResponseContentPartDoneEvent(

862 + type="response.content_part.done",

863 + item_id=current_item_id,

864 + sequence_number=-1,

865 + output_index=current_output_index,

866 + content_index=current_content_index,

867 + part=reasoning_item,

868 + ))

869 + yield _send_event(

870 + ResponseOutputItemDoneEvent(

871 + type="response.output_item.done",

872 + sequence_number=-1,

873 + output_index=current_output_index,

874 + item=reasoning_item,

875 + ))

876 + elif previous_item.channel == "final":

877 + text_content = ResponseOutputText(

878 + type="output_text",

879 + text=previous_item.content[0].text,

880 + annotations=[],

881 + )

882 + yield _send_event(

883 + openai_responses_types.ResponseTextDoneEvent(

884 + type="response.output_text.done",

885 + sequence_number=-1,

886 + output_index=current_output_index,

887 + content_index=current_content_index,

888 + text=previous_item.content[0].text,

889 + logprobs=[],

890 + item_id=current_item_id,

891 + ))

892 + yield _send_event(

893 + openai_responses_types.

894 + ResponseContentPartDoneEvent(

895 + type="response.content_part.done",

896 + sequence_number=-1,

897 + item_id=current_item_id,

898 + output_index=current_output_index,

899 + content_index=current_content_index,

900 + part=text_content,

901 + ))

902 + yield _send_event(

903 + openai_responses_types.ResponseOutputItemDoneEvent(

904 + type="response.output_item.done",

905 + sequence_number=-1,

906 + output_index=current_output_index,

907 + item=ResponseOutputMessage(

908 + id=current_item_id,

909 + type="message",

910 + role="assistant",

911 + content=[text_content],

912 + status="completed",

913 + ),

914 + ))

915 +

916 + if ctx.parser.last_content_delta:

917 + if (ctx.parser.current_channel == "final"

918 + and ctx.parser.current_recipient is None):

919 + if not sent_output_item_added:

920 + sent_output_item_added = True

921 + yield _send_event(

922 + openai_responses_types.

923 + ResponseOutputItemAddedEvent(

924 + type="response.output_item.added",

925 + sequence_number=-1,

926 + output_index=current_output_index,

927 + item=openai_responses_types.

928 + ResponseOutputMessage(

929 + id=current_item_id,

930 + type="message",

931 + role="assistant",

932 + content=[],

933 + status="in_progress",

934 + ),

935 + ))

936 + yield _send_event(

937 + openai_responses_types.

938 + ResponseContentPartAddedEvent(

939 + type="response.content_part.added",

940 + sequence_number=-1,

941 + output_index=current_output_index,

942 + item_id=current_item_id,

943 + content_index=current_content_index,

944 + part=openai_responses_types.ResponseOutputText(

945 + type="output_text",

946 + text="",

947 + annotations=[],

948 + logprobs=[],

949 + ),

950 + ))

951 + yield _send_event(

952 + openai_responses_types.ResponseTextDeltaEvent(

953 + type="response.output_text.delta",

954 + sequence_number=-1,

955 + content_index=current_content_index,

956 + output_index=current_output_index,

957 + item_id=current_item_id,

958 + delta=ctx.parser.last_content_delta,

959 + # TODO, use logprobs from ctx.last_request_output

960 + logprobs=[],

961 + ))

962 + elif (ctx.parser.current_channel == "analysis"

963 + and ctx.parser.current_recipient is None):

964 + if not sent_output_item_added:

965 + sent_output_item_added = True

966 + yield _send_event(

967 + openai_responses_types.

968 + ResponseOutputItemAddedEvent(

969 + type="response.output_item.added",

970 + sequence_number=-1,

971 + output_index=current_output_index,

972 + item=openai_responses_types.

973 + ResponseReasoningItem(

974 + type="reasoning",

975 + id=current_item_id,

976 + summary=[],

977 + status="in_progress",

978 + ),

979 + ))

980 + yield _send_event(

981 + openai_responses_types.

982 + ResponseContentPartAddedEvent(

983 + type="response.content_part.added",

984 + sequence_number=-1,

985 + output_index=current_output_index,

986 + item_id=current_item_id,

987 + content_index=current_content_index,

988 + part=openai_responses_types.ResponseOutputText(

989 + type="output_text",

990 + text="",

991 + annotations=[],

992 + logprobs=[],

993 + ),

994 + ))

995 + yield _send_event(

996 + ResponseReasoningTextDeltaEvent(

997 + type="response.reasoning_text.delta",

998 + item_id=current_item_id,

999 + output_index=current_output_index,

1000 + content_index=current_content_index,

1001 + delta=ctx.parser.last_content_delta,

1002 + sequence_number=-1,

1003 + ))

1004 +

1005 + if ctx.is_assistant_action_turn() and len(ctx.parser.messages) > 0:

1006 + previous_item = ctx.parser.messages[-1]

1007 + if (self.tool_server is not None

1008 + and self.tool_server.has_tool("browser")

1009 + and previous_item.recipient is not None

1010 + and previous_item.recipient.startswith("browser.")):

1011 + function_name = previous_item.recipient[len("browser."):]

1012 + action = None

1013 + parsed_args = json.loads(previous_item.content[0].text)

1014 + if function_name == "search":

1015 + action = (openai_responses_types.

1016 + response_function_web_search.ActionSearch(

1017 + type="search",

1018 + query=parsed_args["query"],

1019 + ))

1020 + elif function_name == "open":

1021 + action = (

1022 + openai_responses_types.

1023 + response_function_web_search.ActionOpenPage(

1024 + type="open_page",

1025 + # TODO: translate to url

1026 + url=f"cursor:{parsed_args.get('cursor', '')}",

1027 + ))

1028 + elif function_name == "find":

1029 + action = (

1030 + openai_responses_types.

1031 + response_function_web_search.ActionFind(

1032 + type="find",

1033 + pattern=parsed_args["pattern"],

1034 + # TODO: translate to url

1035 + url=f"cursor:{parsed_args.get('cursor', '')}",

1036 + ))

1037 + else:

1038 + raise ValueError(

1039 + f"Unknown function name: {function_name}")

1040 +

1041 + yield _send_event(

1042 + openai_responses_types.ResponseOutputItemAddedEvent(

1043 + type="response.output_item.added",

1044 + sequence_number=-1,

1045 + output_index=current_output_index,

1046 + item=openai_responses_types.

1047 + response_function_web_search.

1048 + ResponseFunctionWebSearch(

1049 + # TODO: generate a unique id for web search call

1050 + type="web_search_call",

1051 + id=current_item_id,

1052 + action=action,

1053 + status="in_progress",

1054 + ),

1055 + ))

1056 + yield _send_event(

1057 + openai_responses_types.

1058 + ResponseWebSearchCallInProgressEvent(

1059 + type="response.web_search_call.in_progress",

1060 + sequence_number=-1,

1061 + output_index=current_output_index,

1062 + item_id=current_item_id,

1063 + ))

1064 + yield _send_event(

1065 + openai_responses_types.

1066 + ResponseWebSearchCallSearchingEvent(

1067 + type="response.web_search_call.searching",

1068 + sequence_number=-1,

1069 + output_index=current_output_index,

1070 + item_id=current_item_id,

1071 + ))

1072 +

1073 + # enqueue

1074 + yield _send_event(

1075 + openai_responses_types.

1076 + ResponseWebSearchCallCompletedEvent(

1077 + type="response.web_search_call.completed",

1078 + sequence_number=-1,

1079 + output_index=current_output_index,

1080 + item_id=current_item_id,

1081 + ))

1082 + yield _send_event(

1083 + openai_responses_types.ResponseOutputItemDoneEvent(

1084 + type="response.output_item.done",

1085 + sequence_number=-1,

1086 + output_index=current_output_index,

1087 + item=openai_responses_types.

1088 + ResponseFunctionWebSearch(

1089 + type="web_search_call",

1090 + id=current_item_id,

1091 + action=action,

1092 + status="completed",

1093 + ),

1094 + ))

1095 +

1096 + if (self.tool_server is not None

1097 + and self.tool_server.has_tool("python")

1098 + and previous_item.recipient is not None

1099 + and previous_item.recipient.startswith("python")):

1100 + yield _send_event(

1101 + openai_responses_types.ResponseOutputItemAddedEvent(

1102 + type="response.output_item.added",

1103 + sequence_number=-1,

1104 + output_index=current_output_index,

1105 + item=openai_responses_types.

1106 + ResponseCodeInterpreterToolCallParam(

1107 + type="code_interpreter_call",

1108 + id=current_item_id,

1109 + code="",

1110 + container_id="auto",

1111 + outputs=[],

1112 + status="in_progress",

1113 + ),

1114 + ))

1115 + yield _send_event(

1116 + openai_responses_types.

1117 + ResponseCodeInterpreterCallInProgressEvent(

1118 + type="response.code_interpreter_call.in_progress",

1119 + sequence_number=-1,

1120 + output_index=current_output_index,

1121 + item_id=current_item_id,

1122 + ))

1123 + # TODO: do we need to add delta event here?

1124 + yield _send_event(

1125 + openai_responses_types.

1126 + ResponseCodeInterpreterCallCodeDoneEvent(

1127 + type="response.code_interpreter_call_code.done",

1128 + sequence_number=-1,

1129 + output_index=current_output_index,

1130 + item_id=current_item_id,

1131 + code=previous_item.content[0].text))

1132 + yield _send_event(

1133 + openai_responses_types.

1134 + ResponseCodeInterpreterCallInterpretingEvent(

1135 + type="response.code_interpreter_call.interpreting",

1136 + sequence_number=-1,

1137 + output_index=current_output_index,

1138 + item_id=current_item_id,

1139 + ))

1140 + yield _send_event(

1141 + openai_responses_types.

1142 + ResponseCodeInterpreterCallCompletedEvent(

1143 + type="response.code_interpreter_call.completed",

1144 + sequence_number=-1,

1145 + output_index=current_output_index,

1146 + item_id=current_item_id,

1147 + ))

1148 + yield _send_event(

1149 + openai_responses_types.ResponseOutputItemDoneEvent(

1150 + type="response.output_item.done",

1151 + sequence_number=-1,

1152 + output_index=current_output_index,

1153 + item=openai_responses_types.

1154 + ResponseCodeInterpreterToolCallParam(

1155 + type="code_interpreter_call",

1156 + id=current_item_id,

1157 + code=previous_item.content[0].text,

1158 + container_id="auto",

1159 + # TODO: add outputs here

1160 + outputs=[],

1161 + status="completed",

1162 + ),

1163 + ))

1164 +

1165 + async def empty_async_generator():

1166 + # A hack to trick Python to think this is a generator but in fact

1167 + # it immediately returns.

1168 + if False:

1169 + yield

1170 +

1171 + final_response = await self.responses_full_generator(

1172 + request,

1173 + sampling_params,

1174 + empty_async_generator(),

1175 + context,

1176 + model_name,

1177 + tokenizer,

1178 + request_metadata,

1179 + created_time=created_time,

1180 + )

1181 + yield _send_event(

1182 + openai_responses_types.ResponseCompletedEvent(

1183 + type="response.completed",

1184 + sequence_number=-1,

1185 + response=final_response.model_dump(),

1186 + ))

[gpt-oss] Support streaming in response API (#22431)