Handling Streaming Failures

import random

def full_jitter_delay(attempt: int, base: float = 0.5, cap: float = 2.0) -> float:
    """Full-jitter backoff: random() * min(cap, base * 2**attempt)."""
    return random.random() * min(cap, base * (2**attempt))

result = stream_with_recovery(
    "https://platform.ai.gloo.com/ai/v2/chat/completions",
    {"Authorization": f"Bearer {access_token}"},
    {
        "messages": [{"role": "user", "content": "Explain Romans 8 in simple terms."}],
        "auto_routing": True,
        "stream": True,
    },
    on_delta=lambda text: print(text, end="", flush=True),
)
# result.status is "complete", "content_filter", "interrupted", or "failed"
# result.text always contains everything safe to display

# Standard library plus `requests`:

"""Streaming recovery client for Gloo AI Completions V2."""

import json
import random
import time
from dataclasses import dataclass

import requests

LIVE_PRE_TOKEN_RETRIES = 2     # full-retry attempts before any content is visible
BACKGROUND_RETRIES = 3         # full-retry attempts for background/batch jobs
LIVE_BACKOFF_CAP_S = 2.0       # keep live UX delays short
BACKGROUND_BACKOFF_CAP_S = 30.0


def full_jitter_delay(attempt: int, base: float = 0.5, cap: float = LIVE_BACKOFF_CAP_S) -> float:
    """Full-jitter backoff: random() * min(cap, base * 2**attempt)."""
    return random.random() * min(cap, base * (2**attempt))


@dataclass
class StreamOutcome:
    """Everything one stream attempt tells you, including observability fields."""

    partial_text: str = ""
    received_content: bool = False
    tool_call_emitted: bool = False
    finish_reason: str | None = None
    error: dict | None = None
    trace_id: str | None = None
    secondary_trace_id: str | None = None  # set when error-level and top-level IDs differ
    transport_error: str | None = None
    http_status: int | None = None

    @property
    def completed(self) -> bool:
        return self.finish_reason in ("stop", "length", "tool_calls")

    @property
    def content_filtered(self) -> bool:
        return self.finish_reason == "content_filter"

    @property
    def retryable(self) -> bool:
        if self.transport_error is not None:
            return True  # transport drops are plausibly transient
        return bool(self.error and self.error.get("retryable"))


@dataclass
class RecoveryResult:
    # "complete" | "content_filter" | "interrupted" | "failed"
    status: str
    text: str
    last_outcome: StreamOutcome
    used_continuation: bool = False
    full_retries_used: int = 0


def stream_once(url: str, headers: dict, body: dict, on_delta=None) -> StreamOutcome:
    """Run one streaming request and fold every SSE event into a StreamOutcome."""
    outcome = StreamOutcome()
    try:
        with requests.post(url, headers=headers, json=body, stream=True, timeout=120) as response:
            outcome.http_status = response.status_code
            if response.status_code != 200:
                # Pre-stream HTTP error: no SSE events will follow.
                outcome.transport_error = f"HTTP {response.status_code}"
                return outcome
            # Lowest-precedence trace source; SSE events override it below.
            outcome.trace_id = response.headers.get("x-sentry-trace-id")
            for line in response.iter_lines(decode_unicode=True):
                if not line or not line.startswith("data:"):
                    continue
                payload = line[len("data:"):].strip()
                if payload == "[DONE]":
                    break
                _apply_event(outcome, json.loads(payload), on_delta)
                if outcome.finish_reason is not None:
                    break
    except requests.RequestException as exc:
        outcome.transport_error = str(exc)
    return outcome


def _apply_event(outcome: StreamOutcome, event: dict, on_delta=None) -> None:
    # trace_id precedence: error-level, then top-level, then response header.
    error_trace = (event.get("error") or {}).get("trace_id")
    top_trace = event.get("trace_id")
    outcome.trace_id = error_trace or top_trace or outcome.trace_id
    if error_trace and top_trace and error_trace != top_trace:
        outcome.secondary_trace_id = top_trace  # log both, send both to support

    choice = (event.get("choices") or [{}])[0]
    delta = choice.get("delta") or {}
    if delta.get("content"):
        outcome.partial_text += delta["content"]
        outcome.received_content = True
        if on_delta:
            on_delta(delta["content"])
    if delta.get("tool_calls") or delta.get("function_call"):
        outcome.tool_call_emitted = True
    if choice.get("finish_reason"):
        outcome.finish_reason = choice["finish_reason"]
    if event.get("error"):
        outcome.error = event["error"]


def continuation_messages(messages: list[dict], partial_text: str) -> list[dict]:
    """Provider-neutral continuation: partial output goes in a NEW user message."""
    return [
        *messages,
        {
            "role": "user",
            "content": (
                "The previous streamed answer was interrupted after this text:\n\n"
                f"{partial_text}\n\n"
                "Continue from that point without repeating the text above."
            ),
        },
    ]


def stream_with_recovery(
    url: str, headers: dict, body: dict, *, background: bool = False, on_delta=None
) -> RecoveryResult:
    """One streamed completion with bounded recovery.

    Live mode: 2 full retries before the first token; after visible content,
    exactly one automatic Continue, then surface user controls.
    Background mode: up to 3 full retries; partial output is discarded safely
    because nothing was shown to a user.
    """
    max_full_retries = BACKGROUND_RETRIES if background else LIVE_PRE_TOKEN_RETRIES
    backoff_cap = BACKGROUND_BACKOFF_CAP_S if background else LIVE_BACKOFF_CAP_S

    full_retries = 0
    while True:
        outcome = stream_once(url, headers, body, on_delta)

        if outcome.completed:
            return RecoveryResult("complete", outcome.partial_text, outcome,
                                  full_retries_used=full_retries)
        if outcome.content_filtered:
            # Terminal non-error signal: surface it; never retry or continue.
            return RecoveryResult("content_filter", outcome.partial_text, outcome,
                                  full_retries_used=full_retries)
        if not outcome.retryable:
            # retryable=false: do not retry unchanged.
            return RecoveryResult("failed", outcome.partial_text, outcome,
                                  full_retries_used=full_retries)

        if outcome.received_content and not background:
            # Content is on screen: a full retry would erase or duplicate it.
            return _recover_after_partial(url, headers, body, outcome, on_delta)

        if full_retries >= max_full_retries:
            return RecoveryResult("failed", outcome.partial_text, outcome,
                                  full_retries_used=full_retries)
        time.sleep(full_jitter_delay(full_retries, cap=backoff_cap))
        full_retries += 1


def _recover_after_partial(
    url: str, headers: dict, body: dict, first: StreamOutcome, on_delta=None
) -> RecoveryResult:
    """Post-content recovery: keep the partial text, try one automatic Continue."""
    if first.tool_call_emitted:
        # Side-effect guard: an automatic attempt could replay the tool call.
        # Mark the answer interrupted and hand control to the user.
        return RecoveryResult("interrupted", first.partial_text, first)

    time.sleep(full_jitter_delay(0))
    continuation_body = {**body, "messages": continuation_messages(body["messages"], first.partial_text)}
    second = stream_once(url, headers, continuation_body, on_delta)
    if second.completed:
        return RecoveryResult("complete", first.partial_text + second.partial_text,
                              second, used_continuation=True)

    # The one automatic Continue failed. Keep everything visible, mark the
    # answer interrupted, and surface user-initiated Continue / Try again.
    return RecoveryResult("interrupted", first.partial_text + second.partial_text,
                          second, used_continuation=True)