| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | """Helpers shared by the streaming chat derender tests.""" |
| 4 | |
| 5 | from collections.abc import Callable |
| 6 | |
| 7 | import httpx |
| 8 | |
| 9 | |
| 10 | def assemble_stream(choices: list[dict]) -> dict: |
| 11 | """Fold streamed chat choices into one message. |
| 12 | |
| 13 | A tool call index must keep the first ID it was given. On the derender |
| 14 | side that is what pinning the IDs in `stream_state` guarantees. |
| 15 | """ |
| 16 | content = "" |
| 17 | reasoning = "" |
| 18 | tool_calls: list[dict] = [] |
| 19 | finish_reason = None |
| 20 | for choice in choices: |
| 21 | delta = choice["delta"] |
| 22 | content += delta.get("content") or "" |
| 23 | reasoning += delta.get("reasoning") or "" |
| 24 | for tc in delta.get("tool_calls") or []: |
| 25 | idx = tc["index"] |
| 26 | while len(tool_calls) <= idx: |
| 27 | tool_calls.append({"id": None, "name": None, "arguments": ""}) |
| 28 | if tc.get("id"): |
| 29 | assert tool_calls[idx]["id"] in (None, tc["id"]), ( |
| 30 | f"tool call {idx} changed ID mid stream" |
| 31 | ) |
| 32 | tool_calls[idx]["id"] = tc["id"] |
| 33 | fn = tc.get("function") or {} |
| 34 | if fn.get("name"): |
| 35 | tool_calls[idx]["name"] = fn["name"] |
| 36 | tool_calls[idx]["arguments"] += fn.get("arguments") or "" |
| 37 | finish_reason = choice.get("finish_reason") or finish_reason |
| 38 | return { |
| 39 | "content": content or None, |
| 40 | "reasoning": reasoning or None, |
| 41 | "tool_calls": tool_calls, |
| 42 | "finish_reason": finish_reason, |
| 43 | } |
| 44 | |
| 45 | |
| 46 | async def stream_chat_derender( |
| 47 | client: httpx.AsyncClient, |
| 48 | output_ids: list[int], |
| 49 | chunk_sizes: list[int], |
| 50 | chat_request: dict, |
| 51 | prompt_tokens: int, |
| 52 | prompt_token_ids: list[int], |
| 53 | finish_reason: str = "stop", |
| 54 | on_chunk: Callable[[list[dict]], None] | None = None, |
| 55 | ) -> dict: |
| 56 | """Feed `output_ids` through the streaming chat derender endpoint in |
| 57 | the given `chunk_sizes`, threading `stream_state` across calls and |
| 58 | return the assembled message. |
| 59 | |
| 60 | `finish_reason` goes on the last chunk, so a trailing size of 0 sends it |
| 61 | on a chunk with no tokens. |
| 62 | |
| 63 | If `on_chunk` is given, it is called after every chunk with the |
| 64 | `tool_calls` assembled so far, letting callers assert properties of the |
| 65 | intermediate deltas (e.g. monotonic argument growth) rather than only |
| 66 | the final assembled result. |
| 67 | """ |
| 68 | state = None |
| 69 | choices: list[dict] = [] |
| 70 | pos = 0 |
| 71 | for i, size in enumerate(chunk_sizes): |
| 72 | tids = output_ids[pos : pos + size] |
| 73 | pos += size |
| 74 | is_last = i == len(chunk_sizes) - 1 |
| 75 | resp = await client.post( |
| 76 | "/v1/chat/completions/derender", |
| 77 | json={ |
| 78 | "stream": True, |
| 79 | "model": chat_request["model"], |
| 80 | "generate_chunk": { |
| 81 | "request_id": "stream-test", |
| 82 | "choices": [ |
| 83 | { |
| 84 | "index": 0, |
| 85 | "token_ids": tids, |
| 86 | "finish_reason": finish_reason if is_last else None, |
| 87 | } |
| 88 | ], |
| 89 | }, |
| 90 | "stream_state": state, |
| 91 | "prompt_tokens": prompt_tokens, |
| 92 | "prompt_token_ids": prompt_token_ids, |
| 93 | "chat_request": chat_request, |
| 94 | }, |
| 95 | ) |
| 96 | assert resp.status_code == 200, resp.text |
| 97 | data = resp.json() |
| 98 | state = data["stream_state"] |
| 99 | choices.extend(data["chunk"]["choices"]) |
| 100 | if on_chunk is not None: |
| 101 | on_chunk(assemble_stream(choices)["tool_calls"]) |
| 102 | |
| 103 | return assemble_stream(choices) |