[TRTLLM-7906][feat] Support multiple post process for Responses API (#9908)

Signed-off-by: Junyi Xu <[email protected]>

Author: JunyiXu-nv

tensorrt_llm/serve/openai_server.py

@@ -51,20 +51,22 @@
                                                 MemoryUpdateRequest, ModelCard,
                                                 ModelList, PromptTokensDetails,
                                                 ResponsesRequest,
+                                                ResponsesResponse,
                                                 UpdateWeightsRequest, UsageInfo,
                                                 to_llm_disaggregated_params)
 from tensorrt_llm.serve.postprocess_handlers import (
     ChatCompletionPostprocArgs, ChatPostprocArgs, CompletionPostprocArgs,
-    chat_harmony_post_processor, chat_harmony_streaming_post_processor,
-    chat_response_post_processor, chat_stream_post_processor,
-    completion_response_post_processor, completion_stream_post_processor)
+    ResponsesAPIPostprocArgs, chat_harmony_post_processor,
+    chat_harmony_streaming_post_processor, chat_response_post_processor,
+    chat_stream_post_processor, completion_response_post_processor,
+    completion_stream_post_processor, responses_api_post_processor,
+    responses_api_streaming_post_processor)
 from tensorrt_llm.serve.responses_utils import (ConversationHistoryStore,
+                                                ResponsesStreamingProcessor,
                                                 ServerArrivalTimeMiddleware)
 from tensorrt_llm.serve.responses_utils import \
     create_response as responses_api_create_response
 from tensorrt_llm.serve.responses_utils import get_steady_clock_now_in_seconds
-from tensorrt_llm.serve.responses_utils import \
-    process_streaming_events as responses_api_process_streaming_events
 from tensorrt_llm.serve.responses_utils import \
     request_preprocess as responses_api_request_preprocess
 from tensorrt_llm.version import __version__ as VERSION
@@ -119,9 +121,8 @@ def __init__(self,
             self.model_config = None
 
         # Enable response storage for Responses API
-        self.enable_store = True
-        if len(os.getenv("TRTLLM_RESPONSES_API_DISABLE_STORE", "")) > 0:
-            self.enable_store = False
+        self.enable_store = (len(os.getenv("TRTLLM_RESPONSES_API_DISABLE_STORE", "")) < 1) and not self.postproc_worker_enabled
+
         self.conversation_store = ConversationHistoryStore()
 
         model_dir = Path(model)
@@ -942,19 +943,39 @@ async def create_streaming_generator(promise: RequestOutput, postproc_params: Po
             return self.create_error_response(message=str(e), err_type="internal_error")
 
     async def openai_responses(self, request: ResponsesRequest, raw_request: Request) -> Response:
-        async def create_stream_response(generator, request: ResponsesRequest, sampling_params) -> AsyncGenerator[str, None]:
-            async for event_data in responses_api_process_streaming_events(
-                request=request,
-                sampling_params=sampling_params,
-                generator=generator,
-                model_name=self.model,
-                conversation_store=self.conversation_store,
-                use_harmony=self.use_harmony,
-                reasoning_parser=self.llm.args.reasoning_parser,
-                tool_parser=self.tool_parser,
-                enable_store=self.enable_store
-            ):
-                yield event_data
+        async def create_response(
+                promise: RequestOutput, postproc_params: PostprocParams) -> ResponsesResponse:
+            await promise.aresult()
+            if self.postproc_worker_enabled:
+                response = promise.outputs[0]._postprocess_result
+            else:
+                args = postproc_params.postproc_args
+                response = await responses_api_create_response(
+                    generator=promise,
+                    request=request,
+                    sampling_params=args.sampling_params,
+                    model_name=self.model,
+                    conversation_store=self.conversation_store,
+                    generation_result=None,
+                    enable_store=self.enable_store and request.store,
+                    use_harmony=self.use_harmony,
+                    reasoning_parser=args.reasoning_parser,
+                    tool_parser=args.tool_parser,
+                )
+
+            return response
+
+        async def create_streaming_generator(promise: RequestOutput, postproc_params: PostprocParams):
+            post_processor, args = postproc_params.post_processor, postproc_params.postproc_args
+            streaming_processor = args.streaming_processor
+            initial_responses = streaming_processor.get_initial_responses()
+            for initial_response in initial_responses:
+                yield initial_response
+
+            async for res in promise:
+                pp_results = res.outputs[0]._postprocess_result if self.postproc_worker_enabled else post_processor(res, args)
+                for pp_res in pp_results:
+                    yield pp_res
 
         try:
             if request.background:
@@ -977,38 +998,61 @@ async def create_stream_response(generator, request: ResponsesRequest, sampling_
                 request=request,
                 prev_response=prev_response,
                 conversation_store=self.conversation_store,
-                enable_store=self.enable_store,
+                enable_store=self.enable_store and request.store,
                 use_harmony=self.use_harmony,
                 tokenizer=self.tokenizer if not self.use_harmony else None,
                 model_config=self.model_config if not self.use_harmony else None,
                 processor=self.processor if not self.use_harmony else None,
             )
 
+            streaming_processor = None
+            if request.stream:
+                # Per-request streaming processor
+                streaming_processor = ResponsesStreamingProcessor(
+                    request=request,
+                    sampling_params=sampling_params,
+                    model_name=self.model,
+                    conversation_store=self.conversation_store,
+                    enable_store=self.enable_store and request.store,
+                    use_harmony=self.use_harmony,
+                    reasoning_parser=self.llm.args.reasoning_parser,
+                    tool_parser=self.tool_parser,
+                )
+
+            postproc_args = ResponsesAPIPostprocArgs(
+                model=self.model,
+                request=request,
+                sampling_params=sampling_params,
+                use_harmony=self.use_harmony,
+                reasoning_parser=self.llm.args.reasoning_parser,
+                tool_parser=self.tool_parser,
+                streaming_processor=streaming_processor,
+            )
+            postproc_params = PostprocParams(
+                post_processor=responses_api_streaming_post_processor
+                if request.stream else responses_api_post_processor,
+                postproc_args=postproc_args,
+            )
             promise = self.llm.generate_async(
                 inputs=input_tokens,
                 sampling_params=sampling_params,
                 streaming=request.stream,
+                _postproc_params=postproc_params if self.postproc_worker_enabled else None,
             )
 
+            if self.postproc_worker_enabled and request.store:
+                logger.warning("Postproc workers are enabled, request will not be stored!")
+
             asyncio.create_task(self.await_disconnected(raw_request, promise))
 
             if request.stream:
                 return StreamingResponse(
-                    create_stream_response(promise, request, sampling_params),
+                    content=create_streaming_generator(promise, postproc_params),
                     media_type="text/event-stream"
                 )
             else:
-                return await responses_api_create_response(
-                    generator=promise,
-                    request=request,
-                    sampling_params=sampling_params,
-                    model_name=self.model,
-                    conversation_store=self.conversation_store,
-                    generation_result=None,
-                    enable_store=self.enable_store,
-                    use_harmony=self.use_harmony,
-                    reasoning_parser=self.llm.args.reasoning_parser,
-                    tool_parser=self.tool_parser)
+                response = await create_response(promise, postproc_params)
+                return JSONResponse(content=response.model_dump())
         except CppExecutorError:
             logger.error(traceback.format_exc())
             # If internal executor error is raised, shutdown the server

tensorrt_llm/serve/postprocess_handlers.py

@@ -1,11 +1,16 @@
 from dataclasses import dataclass, field
 from typing import Any, List, Literal, Optional, Tuple, Union
 
+from tensorrt_llm.serve.responses_utils import ResponsesStreamingProcessor
+from tensorrt_llm.serve.responses_utils import \
+    create_response_non_store as responses_api_create_response_non_store
+
 from .._utils import nvtx_range_debug
 from ..executor import (DetokenizedGenerationResultBase, GenerationResult,
                         GenerationResultBase)
 from ..executor.postproc_worker import PostprocArgs
 from ..executor.result import Logprob, TokenLogprobs
+from ..llmapi import SamplingParams
 from ..llmapi.reasoning_parser import (BaseReasoningParser,
                                        ReasoningParserFactory)
 from ..llmapi.tokenizer import TransformersTokenizer
@@ -26,7 +31,8 @@
                               CompletionResponseStreamChoice,
                               CompletionStreamResponse, DeltaFunctionCall,
                               DeltaMessage, DeltaToolCall, FunctionCall,
-                              PromptTokensDetails, StreamOptions, ToolCall,
+                              PromptTokensDetails, ResponsesRequest,
+                              ResponsesResponse, StreamOptions, ToolCall,
                               UsageInfo, to_disaggregated_params)
 from .tool_parser.base_tool_parser import BaseToolParser
 from .tool_parser.core_types import ToolCallItem
@@ -543,3 +549,42 @@ def chat_harmony_streaming_post_processor(
         num_prompt_tokens=args.num_prompt_tokens,
     )
     return response
+
+
+@dataclass(kw_only=True)
+class ResponsesAPIPostprocArgs(PostprocArgs):
+    model: str
+    request: ResponsesRequest
+    sampling_params: SamplingParams
+    use_harmony: bool
+    reasoning_parser: Optional[str] = None
+    tool_parser: Optional[str] = None
+    streaming_processor: Optional[ResponsesStreamingProcessor] = None
+
+
+@nvtx_range_debug("responses_api_post_processor")
+def responses_api_post_processor(
+        rsp: GenerationResult,
+        args: ResponsesAPIPostprocArgs) -> ResponsesResponse:
+    return responses_api_create_response_non_store(
+        generation_result=rsp,
+        request=args.request,
+        sampling_params=args.sampling_params,
+        model_name=args.model,
+        use_harmony=args.use_harmony,
+        reasoning_parser=args.reasoning_parser,
+        tool_parser=args.tool_parser,
+    )
+
+
+@nvtx_range_debug("responses_api_streaming_post_processor")
+def responses_api_streaming_post_processor(
+        rsp: GenerationResult, args: ResponsesAPIPostprocArgs) -> List[str]:
+    if args.streaming_processor is None:
+        raise ValueError(
+            "streaming_processor is required for streaming post-processing")
+    outputs = args.streaming_processor.process_single_output(rsp)
+    if rsp._done:
+        outputs.append(
+            args.streaming_processor.get_final_response_non_store(rsp))
+    return outputs