From 3e0ca0daf995750bd882c1d7e42f918968c654a6 Mon Sep 17 00:00:00 2001 From: 9qeklajc Date: Tue, 2 Jun 2026 13:58:43 +0200 Subject: [PATCH] revert sse refactoring --- routstr/upstream/base.py | 162 +++++++++++++-------------------------- 1 file changed, 53 insertions(+), 109 deletions(-) diff --git a/routstr/upstream/base.py b/routstr/upstream/base.py index b2e763f7..4215426d 100644 --- a/routstr/upstream/base.py +++ b/routstr/upstream/base.py @@ -2,6 +2,7 @@ from __future__ import annotations import asyncio import json +import re import traceback import uuid from collections.abc import AsyncGenerator, AsyncIterator @@ -716,101 +717,55 @@ class BaseUpstreamProvider: pass try: - sse_buffer = b"" async for chunk in response.aiter_bytes(): - sse_buffer += chunk - logger.debug( - "[chat] SSE chunk from upstream", - extra={"chunk_size": len(chunk), "buffer_size": len(sse_buffer)}, - ) - # SSE events are separated by \n\n (double newline). - # Process complete events and keep incomplete ones in the buffer. - while b"\n\n" in sse_buffer: - event_raw, sse_buffer = sse_buffer.split(b"\n\n", 1) - event_raw = event_raw.strip() - if not event_raw: + # Split chunk into SSE events + parts = re.split(b"data: ", chunk) + for i, part in enumerate(parts): + if not part: continue - logger.debug( - "[chat] SSE event from upstream", - extra={"event": event_raw[:500].decode("utf-8", errors="replace")}, - ) + stripped_part = part.strip() + if not stripped_part: + continue - if event_raw == b"data: [DONE]": - logger.debug("[chat] SSE [DONE] seen") + if stripped_part == b"[DONE]": done_seen = True continue - # Extract the JSON payload from "data: {...}" - data_prefix = b"data: " - if event_raw.startswith(data_prefix): - payload_bytes = event_raw[len(data_prefix):] - else: - payload_bytes = event_raw - try: - obj = json.loads(payload_bytes) - if isinstance(obj, dict): - self._apply_provider_field(obj) - if obj.get("model"): - last_model_seen = str(obj.get("model")) - if requested_model: - obj["model"] = requested_model - if ( - "id" not in obj - or not isinstance(obj["id"], str) - or obj["id"] == "existing-id" - ): - if not hasattr(self, "_current_stream_id"): - self._current_stream_id = ( - f"chatcmpl-{uuid.uuid4()}" - ) - obj["id"] = self._current_stream_id - if isinstance(obj.get("usage"), dict): - # Check if this chunk has actual content (vs. billing-only chunks) - choices = obj.get("choices") or [] - has_content = any( - c.get("delta", {}).get("content") - or c.get("delta", {}).get("reasoning") - for c in choices - ) - if not has_content: - # Billing-only chunk — hold back for cost metadata injection - logger.debug( - "[chat] Holding back usage-only chunk", - extra={"model": obj.get("model"), "completion_tokens": obj.get("usage", {}).get("completion_tokens", 0)}, - ) + # Only parse if it looks like a JSON object to avoid SSE control messages or partials + if part.strip().startswith(b"{") and part.strip().endswith( + b"}" + ): + obj = json.loads(part) + if isinstance(obj, dict): + self._apply_provider_field(obj) + if obj.get("model"): + last_model_seen = str(obj.get("model")) + if requested_model: + obj["model"] = requested_model + if ( + "id" not in obj + or not isinstance(obj["id"], str) + or obj["id"] == "existing-id" + ): + if not hasattr(self, "_current_stream_id"): + self._current_stream_id = ( + f"chatcmpl-{uuid.uuid4()}" + ) + obj["id"] = self._current_stream_id + if isinstance(obj.get("usage"), dict): usage_chunk_data = obj continue - # Content-bearing chunk — yield it, but also save usage for later injection - logger.debug( - "[chat] Content chunk with usage — yielding", - extra={"model": obj.get("model"), "completion_tokens": obj.get("usage", {}).get("completion_tokens", 0)}, - ) - if not usage_chunk_data: - usage_chunk_data = obj - logger.debug( - "[chat] Yielding chunk to client", - extra={"model": obj.get("model")}, - ) - yield b"data: " + json.dumps(obj).encode() + b"\n\n" - continue - except Exception as e: - logger.debug( - "[chat] JSON parse failed for event", - extra={"error": str(e), "event_preview": event_raw[:200].decode("utf-8", errors="replace")}, - ) + yield b"data: " + json.dumps(obj).encode() + b"\n\n" + continue + except Exception: + pass - # If JSON parsing failed but it looks like valid SSE, pass through - if event_raw.startswith(b"data: "): - yield event_raw + b"\n\n" - else: - yield b"data: " + event_raw + b"\n\n" - - logger.debug( - "[chat] Upstream stream ended", - extra={"usage_chunk_data": bool(usage_chunk_data), "last_model": last_model_seen}, - ) + prefix = ( + b"data: " if (i > 0 or chunk.startswith(b"data: ")) else b"" + ) + yield prefix + part async with create_session() as session: fresh_key = await session.get(key.__class__, key.hashed_key) @@ -1113,30 +1068,23 @@ class BaseUpstreamProvider: pass try: - sse_buffer = b"" async for chunk in response.aiter_bytes(): - sse_buffer += chunk - # SSE events are separated by \n\n (double newline). - # Process complete events and keep incomplete ones in the buffer. - while b"\n\n" in sse_buffer: - event_raw, sse_buffer = sse_buffer.split(b"\n\n", 1) - event_raw = event_raw.strip() - if not event_raw: + # Split chunk into SSE events + parts = re.split(b"data: ", chunk) + for i, part in enumerate(parts): + if not part: continue - if event_raw == b"data: [DONE]": + stripped_part = part.strip() + if not stripped_part: + continue + + if stripped_part == b"[DONE]": done_seen = True continue - # Extract the JSON payload from "data: {...}" - data_prefix = b"data: " - if event_raw.startswith(data_prefix): - payload_bytes = event_raw[len(data_prefix):] - else: - payload_bytes = event_raw - try: - obj = json.loads(payload_bytes) + obj = json.loads(part) if isinstance(obj, dict): self._apply_provider_field(obj) if obj.get("model"): @@ -1162,17 +1110,13 @@ class BaseUpstreamProvider: ): usage_chunk_data = obj continue - - yield b"data: " + json.dumps(obj).encode() + b"\n\n" - continue except json.JSONDecodeError: pass - # Pass through non-JSON SSE events - if event_raw.startswith(b"data: "): - yield event_raw + b"\n\n" - else: - yield b"data: " + event_raw + b"\n\n" + prefix = ( + b"data: " if (i > 0 or chunk.startswith(b"data: ")) else b"" + ) + yield prefix + part # Always emit a cost-bearing data chunk async with create_session() as session: