diff --git a/src/services/kiaai/chat.py b/src/services/kiaai/chat.py index d55c110..2e17d23 100644 --- a/src/services/kiaai/chat.py +++ b/src/services/kiaai/chat.py @@ -122,6 +122,10 @@ class KiaAIService: Core SSE parser — yields one AIMessageChunk per content delta token. Final chunk: AIMessageChunk(content="", usage_metadata={...}). Raises RuntimeError on KIA server errors (triggers retry in stream()). + + Uses aiter_bytes(chunk_size=32) instead of aiter_lines() so every + byte received from the network is processed immediately without + waiting for httpx's internal 64 KB read buffer to fill up. """ usage: dict = {} @@ -139,57 +143,69 @@ class KiaAIService: resp.raise_for_status() line_count = 0 - async for line in resp.aiter_lines(): - line_count += 1 - if line_count <= 3: - print(f"[KiaAI] raw[{line_count}]: {line!r}") + buf = "" + done = False - if not line: - continue + async for raw in resp.aiter_bytes(chunk_size=32): + buf += raw.decode("utf-8", errors="replace") + + while "\n" in buf: + line, buf = buf.split("\n", 1) + line = line.rstrip("\r") + line_count += 1 + if line_count <= 3: + print(f"[KiaAI] raw[{line_count}]: {line!r}") + + if not line: + continue + + # Non-data lines may carry KIA error JSON + if not line.startswith("data:"): + try: + err = json.loads(line) + code = err.get("code") + if code and int(code) >= 400: + msg = err.get("msg", "Unknown error") + print(f"[KiaAI] Error line: code={code} msg={msg!r}") + raise RuntimeError(f"KIA error {code}: {msg}") + except (json.JSONDecodeError, ValueError): + pass + print(f"[KiaAI] non-data line: {line!r}") + continue + + raw_data = line[6:] if line.startswith("data: ") else line[5:] + + if raw_data.strip() == "[DONE]": + print("[KiaAI] [DONE]") + done = True + break - # Non-data lines may carry KIA error JSON - if not line.startswith("data:"): try: - err = json.loads(line) - code = err.get("code") - if code and int(code) >= 400: - msg = err.get("msg", "Unknown error") - print(f"[KiaAI] Error line: code={code} msg={msg!r}") - raise RuntimeError(f"KIA error {code}: {msg}") - except (json.JSONDecodeError, ValueError): - pass - print(f"[KiaAI] non-data line: {line!r}") - continue + data = json.loads(raw_data) + except json.JSONDecodeError: + print(f"[KiaAI] JSON decode error: {raw_data!r}") + continue - raw_data = line[6:] if line.startswith("data: ") else line[5:] + # KIA server error embedded in stream body + if data.get("code") and int(data["code"]) >= 500: + msg = data.get("msg", "Unknown error") + print(f"[KiaAI] Error chunk: {data['code']} {msg!r}") + raise RuntimeError(f"KIA error {data['code']}: {msg}") - if raw_data.strip() == "[DONE]": - print("[KiaAI] [DONE]") + # Yield content tokens + for choice in (data.get("choices") or []): + delta = choice.get("delta") or {} + token = delta.get("content") or "" + if token: + yield AIMessageChunk(content=token) + + # Capture usage from final summary chunk + if data.get("usage"): + usage = data["usage"] + + if done: break - try: - data = json.loads(raw_data) - except json.JSONDecodeError: - print(f"[KiaAI] JSON decode error: {raw_data!r}") - continue - - # KIA server error embedded in stream body - if data.get("code") and int(data["code"]) >= 500: - msg = data.get("msg", "Unknown error") - print(f"[KiaAI] Error chunk: {data['code']} {msg!r}") - raise RuntimeError(f"KIA error {data['code']}: {msg}") - - # Yield content tokens - for choice in (data.get("choices") or []): - delta = choice.get("delta") or {} - token = delta.get("content") or "" - if token: - yield AIMessageChunk(content=token) - - # Capture usage from final summary chunk - if data.get("usage"): - usage = data["usage"] - # Final chunk carries usage metadata print(f"[KiaAI] Stream done | total_lines={line_count} | usage={usage}") yield AIMessageChunk(