306 lines
11 KiB
Python
306 lines
11 KiB
Python
import json
|
|
import asyncio
|
|
from typing import AsyncGenerator
|
|
import httpx
|
|
from langchain_core.messages import (
|
|
AIMessage,
|
|
AIMessageChunk,
|
|
BaseMessage,
|
|
HumanMessage,
|
|
SystemMessage,
|
|
ToolMessage,
|
|
)
|
|
|
|
from ...config import settings
|
|
|
|
KIA_BASE_URL = "https://api.kie.ai"
|
|
|
|
# Google Search tool shorthand — mutually exclusive with function calling
|
|
KIA_GOOGLE_SEARCH_TOOL = {"type": "function", "function": {"name": "googleSearch"}}
|
|
|
|
|
|
def _to_kia_messages(messages: list[BaseMessage]) -> list[dict]:
|
|
"""
|
|
Convert LangChain messages to KIA message format.
|
|
|
|
Supported content formats per KIA docs:
|
|
- Plain text → str
|
|
- Text-only list → joined into plain string
|
|
- Multimodal (image/video/audio/pdf) → list with unified format:
|
|
[{"type": "image_url", "image_url": {"url": "..."}}]
|
|
All media types (image, video, audio, PDF) use the same image_url structure.
|
|
"""
|
|
result = []
|
|
for msg in messages:
|
|
if isinstance(msg, SystemMessage):
|
|
result.append({"role": "system", "content": _normalize_content(msg.content)})
|
|
elif isinstance(msg, HumanMessage):
|
|
result.append({"role": "user", "content": _normalize_content(msg.content)})
|
|
elif isinstance(msg, AIMessage):
|
|
result.append({"role": "assistant", "content": _normalize_content(msg.content)})
|
|
elif isinstance(msg, ToolMessage):
|
|
result.append({
|
|
"role": "tool",
|
|
"content": msg.content,
|
|
"tool_call_id": msg.tool_call_id,
|
|
})
|
|
else:
|
|
result.append({"role": "user", "content": str(msg.content)})
|
|
return result
|
|
|
|
|
|
def _normalize_content(content) -> str | list:
|
|
"""
|
|
- str → returned as-is
|
|
- list of text-only dicts → joined into plain string
|
|
- list with image_url/media → returned as-is (multimodal)
|
|
"""
|
|
if isinstance(content, str):
|
|
return content
|
|
|
|
if isinstance(content, list):
|
|
is_text_only = all(
|
|
isinstance(item, dict) and item.get("type") == "text"
|
|
for item in content
|
|
)
|
|
if is_text_only:
|
|
return " ".join(item.get("text", "") for item in content)
|
|
return content # multimodal — keep as list
|
|
|
|
return str(content)
|
|
|
|
|
|
class KiaAIService:
|
|
"""
|
|
Calls KIA AI API (https://api.kie.ai) directly via httpx SSE streaming.
|
|
"""
|
|
|
|
# ------------------------------------------------------------------
|
|
# Internal helpers
|
|
# ------------------------------------------------------------------
|
|
|
|
def _build_payload(
|
|
self,
|
|
model: str,
|
|
messages: list[BaseMessage],
|
|
reasoning_effort: str | None,
|
|
include_thoughts: bool,
|
|
google_search: bool,
|
|
tools: list[dict] | None,
|
|
response_format: dict | None,
|
|
) -> tuple[str, dict]:
|
|
"""Build the KIA API URL and request payload."""
|
|
model_name = model.removeprefix("kia/")
|
|
url = f"{KIA_BASE_URL}/{model_name}/v1/chat/completions"
|
|
kia_messages = _to_kia_messages(messages)
|
|
|
|
print(f"[KiaAI] POST {url}")
|
|
print(f"[KiaAI] msgs={len(kia_messages)} | reasoning={reasoning_effort} | thoughts={include_thoughts} | gsearch={google_search}")
|
|
for i, m in enumerate(kia_messages):
|
|
preview = m["content"][:120] if isinstance(m["content"], str) else str(m["content"])[:120]
|
|
print(f"[KiaAI] msg[{i}] role={m['role']} | {preview!r}")
|
|
|
|
payload: dict = {"messages": kia_messages, "stream": True}
|
|
if reasoning_effort:
|
|
payload["reasoning_effort"] = reasoning_effort
|
|
if include_thoughts:
|
|
payload["include_thoughts"] = True
|
|
if google_search:
|
|
payload["tools"] = [KIA_GOOGLE_SEARCH_TOOL]
|
|
elif tools:
|
|
payload["tools"] = tools
|
|
elif response_format:
|
|
payload["response_format"] = response_format
|
|
|
|
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
|
|
|
|
|
|
# Compatibility stub — KIA models bypass LangChain and use KiaAIService directly
|
|
def get_kia_model(model: str):
|
|
return None
|