vahidrezvani
parent
3b14f898e3
commit
a844bc1b4e
|
|
@ -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
|
||||
|
|
|
|||
Loading…
Reference in New Issue