import json import logging import time import uuid from typing import Any, AsyncIterator import httpx from fastapi import HTTPException from config import ( UPSTREAM_URL, UPSTREAM_HEADERS, DEFAULT_MODEL, UPSTREAM_TIMEOUT_SECS, UPSTREAM_CONNECT_TIMEOUT_SECS, UPSTREAM_MAX_CONNECTIONS, UPSTREAM_MAX_KEEPALIVE_CONNECTIONS, UPSTREAM_KEEPALIVE_EXPIRY_SECS, UPSTREAM_PROXY_URL, ) from response_cache import CompletionArtifact logger = logging.getLogger(__name__) _UPSTREAM_CLIENT: httpx.AsyncClient | None = None def create_upstream_client() -> httpx.AsyncClient: transport_kwargs: dict[str, Any] = {} if UPSTREAM_PROXY_URL: transport_kwargs["proxy"] = UPSTREAM_PROXY_URL transport = httpx.AsyncHTTPTransport( retries=0, limits=httpx.Limits( max_connections=UPSTREAM_MAX_CONNECTIONS, max_keepalive_connections=UPSTREAM_MAX_KEEPALIVE_CONNECTIONS, keepalive_expiry=UPSTREAM_KEEPALIVE_EXPIRY_SECS, ), **transport_kwargs, ) return httpx.AsyncClient( timeout=httpx.Timeout(UPSTREAM_TIMEOUT_SECS, connect=UPSTREAM_CONNECT_TIMEOUT_SECS), transport=transport, trust_env=True, http2=False, ) def get_upstream_client() -> httpx.AsyncClient: global _UPSTREAM_CLIENT if _UPSTREAM_CLIENT is None: _UPSTREAM_CLIENT = create_upstream_client() return _UPSTREAM_CLIENT async def close_upstream_client() -> None: global _UPSTREAM_CLIENT if _UPSTREAM_CLIENT is not None: await _UPSTREAM_CLIENT.aclose() _UPSTREAM_CLIENT = None async def iter_upstream_events(payload: dict) -> AsyncIterator[tuple[str, dict]]: client = get_upstream_client() request_started = time.time() event_counts: dict[str, int] = {} raw_line_samples: list[str] = [] parsed_samples: list[dict] = [] skipped_lines = 0 json_errors = 0 byte_count = 0 try: async with client.stream("POST", UPSTREAM_URL, headers=UPSTREAM_HEADERS, json=payload) as r: logger.info( "upstream response status=%s content-type=%s model=%s messages=%s", r.status_code, r.headers.get("content-type"), payload.get("model"), len(payload.get("messages") or []), ) if r.status_code != 200: body = await r.aread() err_body = body.decode("utf-8", errors="replace")[:2000] logger.warning( "upstream non-200 status=%s body=%r headers=%s", r.status_code, err_body[:1000], dict(r.headers), ) yield ("error", {"status_code": r.status_code, "body": err_body or f"upstream returned HTTP {r.status_code}"}) return buffer = b"" async for chunk in r.aiter_bytes(): if not chunk: continue byte_count += len(chunk) buffer += chunk while b"\n" in buffer: line, buffer = buffer.split(b"\n", 1) s = line.decode("utf-8", errors="replace").rstrip("\r") if s and len(raw_line_samples) < 8: raw_line_samples.append(s[:500]) if not s or s.startswith("event:") or s.startswith("ddata:"): skipped_lines += 1 continue if s.startswith("ata:"): s = "d" + s if not s.startswith("data:"): skipped_lines += 1 continue data = s[5:].lstrip() if data == "[DONE]": event_counts["done"] = event_counts.get("done", 0) + 1 yield ("done", {}) continue try: obj = json.loads(data) except Exception as exc: json_errors += 1 if json_errors <= 3: logger.warning("upstream JSON parse failed: %s raw=%r", exc, data[:500]) continue if not isinstance(obj, dict): skipped_lines += 1 continue if len(parsed_samples) < 5: parsed_samples.append(obj) t = obj.get("type") if not t or t == "job_id": skipped_lines += 1 continue if t == "delta" and "text" in obj and "index" not in obj: obj = { "type": "content_block_delta", "index": 0, "delta": {"type": "text_delta", "text": obj.get("text", "")}, } t = "content_block_delta" if t == "error": yield ("error", {"status_code": 502, "body": obj.get("message") or json.dumps(obj, ensure_ascii=False)}) return event_counts[t] = event_counts.get(t, 0) + 1 yield (t, obj) if buffer.strip(): tail = buffer.decode("utf-8", errors="replace").strip() if len(raw_line_samples) < 8: raw_line_samples.append(tail[:500]) logger.info( "upstream stream finished status=%s bytes=%s duration=%.2fs events=%s skipped=%s json_errors=%s samples=%s parsed_samples=%s", r.status_code, byte_count, time.time() - request_started, event_counts, skipped_lines, json_errors, raw_line_samples, parsed_samples, ) except httpx.TimeoutException as exc: logger.exception("upstream timeout after %.2fs model=%s error=%s", time.time() - request_started, payload.get("model"), exc) yield ("error", {"status_code": 504, "body": f"upstream timeout: {type(exc).__name__}: {exc}"}) except httpx.HTTPError as exc: logger.exception("upstream HTTP error after %.2fs model=%s error=%s", time.time() - request_started, payload.get("model"), exc) yield ("error", {"status_code": 502, "body": f"upstream HTTP error: {type(exc).__name__}: {exc}"}) except Exception as exc: logger.exception("upstream unexpected error after %.2fs model=%s error=%s", time.time() - request_started, payload.get("model"), exc) yield ("error", {"status_code": 502, "body": f"upstream unexpected error: {type(exc).__name__}: {exc}"}) async def collect_upstream_text(payload: dict) -> tuple[str, str, dict, str, dict]: """Run through the whole upstream stream and collect: (full_text, model_id, usage, stop_reason, message_start_obj) """ text_parts: list[str] = [] model_id = payload.get("model", DEFAULT_MODEL) stop_reason = "end_turn" usage = {"input_tokens": 0, "output_tokens": 0} msg_start = {} event_counts: dict[str, int] = {} async for event_name, obj in iter_upstream_events(payload): event_counts[event_name] = event_counts.get(event_name, 0) + 1 if event_name == "error": logger.warning("upstream collection failed events=%s error=%s", event_counts, obj) raise HTTPException(status_code=obj.get("status_code", 502), detail=obj.get("body", "upstream error")) if event_name == "message_start": msg_start = obj.get("message", {}) model_id = msg_start.get("model", model_id) if "usage" in msg_start: usage["input_tokens"] = msg_start["usage"].get("input_tokens", 0) elif event_name == "content_block_delta": delta = obj.get("delta", {}) if delta.get("type") == "text_delta": text_parts.append(delta.get("text", "")) elif event_name == "message_delta": d = obj.get("delta", {}) if d.get("stop_reason"): stop_reason = d["stop_reason"] u = obj.get("usage") if u and "output_tokens" in u: usage["output_tokens"] = u["output_tokens"] full_text = "".join(text_parts) if not full_text: logger.warning( "upstream returned no text model=%s events=%s msg_start=%s usage=%s stop_reason=%s", model_id, event_counts, msg_start, usage, stop_reason, ) return (full_text, model_id, usage, stop_reason, msg_start) def build_completion_artifact(raw_text: str, model_id: str, usage: dict, stop_reason: str) -> CompletionArtifact: return CompletionArtifact( schema_version=1, raw_text=raw_text, model_id=model_id, usage_input_tokens=int(usage.get("input_tokens", 0)), usage_output_tokens=int(usage.get("output_tokens", 0)), stop_reason=stop_reason or "end_turn", stored_at=time.time(), ) class LiveArtifactCapture: def __init__(self, fallback_model: str): self.fallback_model = fallback_model self.model_id = fallback_model self.raw_text_parts: list[str] = [] self.usage = {"input_tokens": 0, "output_tokens": 0} self.stop_reason = "end_turn" self.failed = False self.completed = False def observe(self, event_name: str, obj: dict) -> None: if event_name == "error": self.failed = True return if event_name == "done": self.completed = True if not self.raw_text_parts: logger.warning( "live stream completed without text model=%s usage=%s stop_reason=%s", self.model_id, self.usage, self.stop_reason, ) return if event_name == "message_start": message = obj.get("message", {}) self.model_id = message.get("model", self.model_id) if "usage" in message: self.usage["input_tokens"] = message["usage"].get("input_tokens", 0) return if event_name == "content_block_delta": delta = obj.get("delta", {}) if delta.get("type") == "text_delta": self.raw_text_parts.append(delta.get("text", "")) return if event_name == "message_delta": delta = obj.get("delta", {}) if delta.get("stop_reason"): self.stop_reason = delta["stop_reason"] usage = obj.get("usage") if usage and "output_tokens" in usage: self.usage["output_tokens"] = usage["output_tokens"] def is_cacheable(self) -> bool: return self.completed and not self.failed def build(self) -> CompletionArtifact: return build_completion_artifact("".join(self.raw_text_parts), self.model_id, self.usage, self.stop_reason) async def fetch_completion_artifact(payload: dict) -> CompletionArtifact: text, model_id, usage, stop_reason, _ = await collect_upstream_text(payload) return build_completion_artifact(text, model_id, usage, stop_reason)