diff --git a/src/services/kiaai/chat.py b/src/services/kiaai/chat.py index 04040fa..2e17d23 100644 --- a/src/services/kiaai/chat.py +++ b/src/services/kiaai/chat.py @@ -115,198 +115,20 @@ class KiaAIService: print(f"[KiaAI] payload keys: {[k for k in payload if k != 'messages']}") return url, payload - # async def _iter_sse( - # self, url: str, payload: dict - # ) -> AsyncGenerator[AIMessageChunk, None]: - # """ - # 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 = {} - - # async with httpx.AsyncClient(timeout=180.0) as client: - # async with client.stream( - # "POST", - # url, - # content=json.dumps(payload), - # headers={ - # "Content-Type": "application/json", - # "Authorization": f"Bearer {settings.kia_api_key}", - # }, - # ) as resp: - # print(f"[KiaAI] HTTP {resp.status_code} | {url}") - # resp.raise_for_status() - - # line_count = 0 - # buf = "" - # done = False - - # 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 - - # 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"] - - # if done: - # break - - # # Final chunk carries usage metadata - # print(f"[KiaAI] Stream done | total_lines={line_count} | usage={usage}") - # yield AIMessageChunk( - # content="", - # usage_metadata={ - # "input_tokens": usage.get("prompt_tokens", 0), - # "output_tokens": usage.get("completion_tokens", 0), - # "total_tokens": usage.get("total_tokens", 0), - # }, - # ) - - # # ------------------------------------------------------------------ - # # Public API - # # ------------------------------------------------------------------ - - # async def stream( - # self, - # model: str, - # messages: list[BaseMessage], - # reasoning_effort: str | None = None, - # include_thoughts: bool = False, - # google_search: bool = False, - # tools: list[dict] | None = None, - # response_format: dict | None = None, - # max_retries: int = 3, - # retry_delay: float = 2.0, - # ) -> AsyncGenerator[AIMessageChunk, None]: - # """ - # Token-by-token async generator. - # Yields AIMessageChunk(content=token) per SSE delta. - # Last chunk: AIMessageChunk(content="", usage_metadata={...}). - # Retries on server errors only if no tokens were yielded yet. - # """ - # url, payload = self._build_payload( - # model, messages, reasoning_effort, include_thoughts, - # google_search, tools, response_format, - # ) - - # last_error: Exception | None = None - # for attempt in range(1, max_retries + 1): - # if attempt > 1: - # print(f"[KiaAI] Stream retry {attempt}/{max_retries} in {retry_delay}s ...") - # await asyncio.sleep(retry_delay) - - # yielded_any = False - # try: - # async for chunk in self._iter_sse(url, payload): - # yielded_any = True - # yield chunk - # return # success - # except RuntimeError as e: - # if yielded_any: - # raise # can't retry after content was already sent - # last_error = e - # print(f"[KiaAI] Attempt {attempt} failed (no content yet): {e}") - - # raise last_error - - # async def invoke( - # self, - # model: str, - # messages: list[BaseMessage], - # reasoning_effort: str | None = None, - # include_thoughts: bool = False, - # google_search: bool = False, - # tools: list[dict] | None = None, - # response_format: dict | None = None, - # max_retries: int = 3, - # retry_delay: float = 2.0, - # ) -> AIMessage: - # """Collects all streaming chunks into a single AIMessage.""" - # full_content = "" - # usage: dict = {} - # async for chunk in self.stream( - # model=model, - # messages=messages, - # reasoning_effort=reasoning_effort, - # include_thoughts=include_thoughts, - # google_search=google_search, - # tools=tools, - # response_format=response_format, - # max_retries=max_retries, - # retry_delay=retry_delay, - # ): - # full_content += chunk.content or "" - # if chunk.usage_metadata: - # usage = chunk.usage_metadata - - # print(f"[KiaAI] invoke() done | content_len={len(full_content)} | usage={usage}") - # ai_msg = AIMessage(content=full_content) - # ai_msg.usage_metadata = usage - # return ai_msg - async def _iter_sse( self, url: str, payload: dict ) -> AsyncGenerator[AIMessageChunk, None]: """ - Real-time SSE parser — yields tokens immediately as they arrive (no line buffering), - but still collects usage metadata from final chunk. + 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 = {} + async with httpx.AsyncClient(timeout=180.0) as client: async with client.stream( "POST", @@ -317,49 +139,75 @@ class KiaAIService: "Authorization": f"Bearer {settings.kia_api_key}", }, ) as resp: + print(f"[KiaAI] HTTP {resp.status_code} | {url}") resp.raise_for_status() - buffer = "" + line_count = 0 + buf = "" done = False - async for raw in resp.aiter_bytes(chunk_size=32): - text = raw.decode("utf-8", errors="replace") - buffer += text - # پردازش خطوط کامل یا partial JSON - while "\n" in buffer: - line, buffer = buffer.split("\n", 1) + 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") - if not line.startswith("data:"): + 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 + 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 توکن‌ها فوراً - for choice in data.get("choices", []): + # 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) - # ذخیره usage از chunk نهایی + # Capture usage from final summary chunk if data.get("usage"): usage = data["usage"] if done: break - # Yield final chunk with usage metadata + # Final chunk carries usage metadata + print(f"[KiaAI] Stream done | total_lines={line_count} | usage={usage}") yield AIMessageChunk( content="", usage_metadata={ @@ -370,7 +218,7 @@ class KiaAIService: ) # ------------------------------------------------------------------ - # Public streaming API + # Public API # ------------------------------------------------------------------ async def stream( @@ -386,7 +234,10 @@ class KiaAIService: retry_delay: float = 2.0, ) -> AsyncGenerator[AIMessageChunk, None]: """ - Real-time token streaming to client, retries on server errors if no tokens yielded yet. + Token-by-token async generator. + Yields AIMessageChunk(content=token) per SSE delta. + Last chunk: AIMessageChunk(content="", usage_metadata={...}). + Retries on server errors only if no tokens were yielded yet. """ url, payload = self._build_payload( model, messages, reasoning_effort, include_thoughts, @@ -396,22 +247,22 @@ class KiaAIService: last_error: Exception | None = None for attempt in range(1, max_retries + 1): if attempt > 1: + print(f"[KiaAI] Stream retry {attempt}/{max_retries} in {retry_delay}s ...") await asyncio.sleep(retry_delay) + yielded_any = False try: async for chunk in self._iter_sse(url, payload): yielded_any = True yield chunk - return + return # success except RuntimeError as e: if yielded_any: - raise + raise # can't retry after content was already sent last_error = e - raise last_error + print(f"[KiaAI] Attempt {attempt} failed (no content yet): {e}") - # ------------------------------------------------------------------ - # Full message collection - # ------------------------------------------------------------------ + raise last_error async def invoke( self, @@ -425,9 +276,7 @@ class KiaAIService: max_retries: int = 3, retry_delay: float = 2.0, ) -> AIMessage: - """ - Collects all streamed chunks into a complete AIMessage. - """ + """Collects all streaming chunks into a single AIMessage.""" full_content = "" usage: dict = {} async for chunk in self.stream( @@ -445,6 +294,7 @@ class KiaAIService: if chunk.usage_metadata: usage = chunk.usage_metadata + print(f"[KiaAI] invoke() done | content_len={len(full_content)} | usage={usage}") ai_msg = AIMessage(content=full_content) ai_msg.usage_metadata = usage return ai_msg