File size: 1,381 Bytes
28a52dc | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 | """Call the assistant-ui upstream LLM endpoint."""
from __future__ import annotations
from collections.abc import AsyncIterator
import httpx
from config import UPSTREAM_HEADERS, UPSTREAM_URL
class UpstreamError(Exception):
def __init__(self, status_code: int, message: str):
self.status_code = status_code
super().__init__(message)
async def call_upstream(payload: dict) -> AsyncIterator[str]:
"""POST to assistant-ui and yield raw SSE lines as they arrive."""
async with httpx.AsyncClient(timeout=httpx.Timeout(120, connect=10)) as client:
async with client.stream(
"POST",
UPSTREAM_URL,
json=payload,
headers=UPSTREAM_HEADERS,
) as resp:
if resp.status_code == 429:
raise UpstreamError(429, "Rate limit exceeded (upstream: 5 req / 30s per IP)")
if resp.status_code >= 400:
body = await resp.aread()
raise UpstreamError(resp.status_code, body.decode(errors="replace")[:200])
async for line in resp.aiter_lines():
yield line
async def call_upstream_full(payload: dict) -> list[str]:
"""POST to assistant-ui and collect all SSE lines (for non-stream mode)."""
lines: list[str] = []
async for line in call_upstream(payload):
lines.append(line)
return lines
|