Compare commits

..

No commits in common. "1e46711ab0696759bbacdf37a9f9ee259245a30f" and "3b14f898e3ab3ad6418bc5fb1315c8b2a6df4a36" have entirely different histories.

3 changed files with 210 additions and 67 deletions

View File

@ -104,11 +104,7 @@ async def send_message(
],
doc=[(settings.base_url + doc_url) if doc_url else None, doc_path],
),
media_type="application/x-ndjson",
headers={
"X-Accel-Buffering": "no",
"Cache-Control": "no-cache",
},
media_type="application/json",
)

View File

@ -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,76 +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):
text = raw.decode("utf-8", errors="replace")
buffer += text
async for raw in resp.aiter_bytes(chunk_size=1):
buf += raw.decode("utf-8", errors="replace")
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)
await asyncio.sleep(0) # flush immediately
# 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={
@ -219,7 +370,7 @@ class KiaAIService:
)
# ------------------------------------------------------------------
# Public API
# Public streaming API
# ------------------------------------------------------------------
async def stream(
@ -235,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,
@ -248,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,
@ -277,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(
@ -295,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

View File

@ -1,5 +1,4 @@
import json
import asyncio
from datetime import datetime, UTC
from typing import AsyncGenerator
@ -72,7 +71,6 @@ async def kia_token_stream(
if token:
full_ai_content += token
yield json.dumps({"content": token}, ensure_ascii=False) + "\n"
await asyncio.sleep(0) # flush token to client immediately
if chunk.usage_metadata:
input_tokens = chunk.usage_metadata.get("input_tokens", 0)
output_tokens = chunk.usage_metadata.get("output_tokens", 0)