Function wrap_openai_stream_with_logging

app/utils/request_logging.py:187–264

Wrap OpenAI SSE stream and persist completion metadata.

(
    stream: AsyncGenerator[str, None],
    *,
    provider: str,
    model: str,
    source_info: RequestSourceInfo,
    started_at: float,
)

async def wrap_openai_stream_with_logging(
    stream: AsyncGenerator[str, None],
    *,
    provider: str,
    model: str,
    source_info: RequestSourceInfo,
    started_at: float,
) -> AsyncGenerator[str, None]:
    """Wrap OpenAI SSE stream and persist completion metadata."""
    success = True
    status_code = 200
    error_message: Optional[str] = None
    first_token_time = 0.0
    usage = {
        "input_tokens": 0,
        "output_tokens": 0,
        "cache_creation_tokens": 0,
        "cache_read_tokens": 0,
        "total_tokens": 0,
    }

    try:
        async for chunk in stream:
            if chunk.startswith("data: "):
                payload_text = chunk[6:].strip()
                if payload_text and payload_text != "[DONE]":
                    try:
                        payload = json.loads(payload_text)
                    except json.JSONDecodeError:
                        payload = None

                    if isinstance(payload, dict):
                        if "error" in payload:
                            success = False
                            error = payload.get("error") or {}
                            error_message = (
                                error.get("message")
                                or "Unknown stream error"
                            )
                            status_code = int(error.get("code") or 500)
                        else:
                            if (
                                not first_token_time
                                and _openai_payload_has_output(payload)
                            ):
                                first_token_time = max(
                                    0.0,
                                    time.perf_counter() - started_at,
                                )
                            if payload.get("usage"):
                                usage = _merge_usage(
                                    usage,
                                    extract_openai_usage(payload),
                                    include_cache_in_total=False,
                                )

            yield chunk
    except Exception as exc: