| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | import pydantic |
| 4 | from fastapi import Request |
| 5 | from starlette.concurrency import iterate_in_threadpool |
| 6 | |
| 7 | from vllm.logger import init_logger |
| 8 | |
| 9 | logger = init_logger(__name__) |
| 10 | |
| 11 | |
| 12 | class SSEDecoder: |
| 13 | """Robust Server-Sent Events decoder for streaming responses.""" |
| 14 | |
| 15 | def __init__(self): |
| 16 | self.buffer = "" |
| 17 | self.content_buffer = [] |
| 18 | |
| 19 | def decode_chunk(self, chunk: bytes) -> list[dict]: |
| 20 | """Decode a chunk of SSE data and return parsed events.""" |
| 21 | import json |
| 22 | |
| 23 | try: |
| 24 | chunk_str = chunk.decode("utf-8") |
| 25 | except UnicodeDecodeError: |
| 26 | # Skip malformed chunks |
| 27 | return [] |
| 28 | |
| 29 | self.buffer += chunk_str |
| 30 | events = [] |
| 31 | |
| 32 | # Process complete lines |
| 33 | while "\n" in self.buffer: |
| 34 | line, self.buffer = self.buffer.split("\n", 1) |
| 35 | line = line.rstrip("\r") # Handle CRLF |
| 36 | |
| 37 | if line.startswith("data: "): |
| 38 | data_str = line[6:].strip() |
| 39 | if data_str == "[DONE]": |
| 40 | events.append({"type": "done"}) |
| 41 | elif data_str: |
| 42 | try: |
| 43 | event_data = json.loads(data_str) |
| 44 | events.append({"type": "data", "data": event_data}) |
| 45 | except json.JSONDecodeError: |
| 46 | # Skip malformed JSON |
| 47 | continue |
| 48 | |
| 49 | return events |
| 50 | |
| 51 | def extract_content(self, event_data: dict) -> str: |
| 52 | """Extract content from event data.""" |
| 53 | return _extract_content_from_chunk(event_data) |
| 54 | |
| 55 | def add_content(self, content: str) -> None: |
| 56 | """Add content to the buffer.""" |
| 57 | if content: |
| 58 | self.content_buffer.append(content) |
| 59 | |
| 60 | def get_complete_content(self) -> str: |
| 61 | """Get the complete buffered content.""" |
| 62 | return "".join(self.content_buffer) |
| 63 | |
| 64 | |
| 65 | def _extract_content_from_chunk(chunk_data: dict) -> str: |
| 66 | """Extract content from a streaming response chunk.""" |
| 67 | try: |
| 68 | from vllm.entrypoints.openai.chat_completion.protocol import ( |
| 69 | ChatCompletionStreamResponse, |
| 70 | ) |
| 71 | from vllm.entrypoints.openai.completion.protocol import ( |
| 72 | CompletionStreamResponse, |
| 73 | ) |
| 74 | |
| 75 | # Try using Completion types for type-safe parsing |
| 76 | if chunk_data.get("object") == "chat.completion.chunk": |
| 77 | chat_response = ChatCompletionStreamResponse.model_validate(chunk_data) |
| 78 | if chat_response.choices and chat_response.choices[0].delta.content: |
| 79 | return chat_response.choices[0].delta.content |
| 80 | elif chunk_data.get("object") == "text_completion": |
| 81 | completion_response = CompletionStreamResponse.model_validate(chunk_data) |
| 82 | if completion_response.choices and completion_response.choices[0].text: |
| 83 | return completion_response.choices[0].text |
| 84 | except pydantic.ValidationError: |
| 85 | # Fallback to manual parsing |
| 86 | if "choices" in chunk_data and chunk_data["choices"]: |
| 87 | choice = chunk_data["choices"][0] |
| 88 | if "delta" in choice and choice["delta"].get("content"): |
| 89 | return choice["delta"]["content"] |
| 90 | elif choice.get("text"): |
| 91 | return choice["text"] |
| 92 | return "" |
| 93 | |
| 94 | |
| 95 | def _log_streaming_response(response, response_body: list) -> None: |
| 96 | """Log streaming response with robust SSE parsing.""" |
| 97 | sse_decoder = SSEDecoder() |
| 98 | chunk_count = 0 |
| 99 | |
| 100 | def buffered_iterator(): |
| 101 | nonlocal chunk_count |
| 102 | |
| 103 | for chunk in response_body: |
| 104 | chunk_count += 1 |
| 105 | yield chunk |
| 106 | |
| 107 | # Parse SSE events from chunk |
| 108 | events = sse_decoder.decode_chunk(chunk) |
| 109 | |
| 110 | for event in events: |
| 111 | if event["type"] == "data": |
| 112 | content = sse_decoder.extract_content(event["data"]) |
| 113 | sse_decoder.add_content(content) |
| 114 | elif event["type"] == "done": |
| 115 | # Log complete content when done |
| 116 | full_content = sse_decoder.get_complete_content() |
| 117 | if full_content: |
| 118 | # Truncate if too long |
| 119 | if len(full_content) > 2048: |
| 120 | full_content = full_content[:2048] + "...[truncated]" |
| 121 | logger.info( |
| 122 | "response_body={streaming_complete: content=%r, chunks=%d}", |
| 123 | full_content, |
| 124 | chunk_count, |
| 125 | ) |
| 126 | else: |
| 127 | logger.info( |
| 128 | "response_body={streaming_complete: no_content, chunks=%d}", |
| 129 | chunk_count, |
| 130 | ) |
| 131 | return |
| 132 | |
| 133 | response.body_iterator = iterate_in_threadpool(buffered_iterator()) |
| 134 | logger.info("response_body={streaming_started: chunks=%d}", len(response_body)) |
| 135 | |
| 136 | |
| 137 | def _log_non_streaming_response(response_body: list) -> None: |
| 138 | """Log non-streaming response.""" |
| 139 | try: |
| 140 | decoded_body = response_body[0].decode() |
| 141 | logger.info("response_body={%s}", decoded_body) |
| 142 | except UnicodeDecodeError: |
| 143 | logger.info("response_body={<binary_data>}") |
| 144 | |
| 145 | |
| 146 | async def log_response(request: Request, call_next): |
| 147 | response = await call_next(request) |
| 148 | response_body = [section async for section in response.body_iterator] |
| 149 | response.body_iterator = iterate_in_threadpool(iter(response_body)) |
| 150 | # Check if this is a streaming response by looking at content-type |
| 151 | content_type = response.headers.get("content-type", "") |
| 152 | is_streaming = content_type == "text/event-stream; charset=utf-8" |
| 153 | |
| 154 | # Log response body based on type |
| 155 | if not response_body: |
| 156 | logger.info("response_body={<empty>}") |
| 157 | elif is_streaming: |
| 158 | _log_streaming_response(response, response_body) |
| 159 | else: |
| 160 | _log_non_streaming_response(response_body) |
| 161 | return response |