diff --git a/src/services/kiaai/chat.py b/src/services/kiaai/chat.py index 2e17d23..04040fa 100644 --- a/src/services/kiaai/chat.py +++ b/src/services/kiaai/chat.py @@ -115,20 +115,198 @@ 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]: """ - 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. + Real-time SSE parser — yields tokens immediately as they arrive (no line buffering), + but still collects usage metadata from final chunk. """ usage: dict = {} - async with httpx.AsyncClient(timeout=180.0) as client: async with client.stream( "POST", @@ -139,75 +317,49 @@ class KiaAIService: "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 = "" + buffer = "" done = False - async for raw in resp.aiter_bytes(chunk_size=32): - buf += raw.decode("utf-8", errors="replace") + text = raw.decode("utf-8", errors="replace") + buffer += text - while "\n" in buf: - line, buf = buf.split("\n", 1) + # پردازش خطوط کامل یا partial JSON + while "\n" in buffer: + line, buffer = buffer.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 []): + # yield توکن‌ها فوراً + for choice in data.get("choices", []): delta = choice.get("delta") or {} token = delta.get("content") or "" if token: yield AIMessageChunk(content=token) - # Capture usage from final summary chunk + # ذخیره usage از 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 final chunk with usage metadata yield AIMessageChunk( content="", usage_metadata={ @@ -218,7 +370,7 @@ class KiaAIService: ) # ------------------------------------------------------------------ - # Public API + # Public streaming API # ------------------------------------------------------------------ async def stream( @@ -234,10 +386,7 @@ class KiaAIService: 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. + Real-time token streaming to client, retries on server errors if no tokens yielded yet. """ url, payload = self._build_payload( model, messages, reasoning_effort, include_thoughts, @@ -247,23 +396,23 @@ 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 # success + return except RuntimeError as e: if yielded_any: - raise # can't retry after content was already sent + raise last_error = e - print(f"[KiaAI] Attempt {attempt} failed (no content yet): {e}") - raise last_error + # ------------------------------------------------------------------ + # Full message collection + # ------------------------------------------------------------------ + async def invoke( self, model: str, @@ -276,7 +425,9 @@ class KiaAIService: max_retries: int = 3, retry_delay: float = 2.0, ) -> AIMessage: - """Collects all streaming chunks into a single AIMessage.""" + """ + Collects all streamed chunks into a complete AIMessage. + """ full_content = "" usage: dict = {} async for chunk in self.stream( @@ -294,7 +445,6 @@ 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