244 lines
9.4 KiB
Python
244 lines
9.4 KiB
Python
import json
|
|
import asyncio
|
|
import httpx
|
|
from langchain_core.messages import (
|
|
AIMessage,
|
|
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.
|
|
|
|
Supported parameters (per KIA docs):
|
|
- messages (required) list of role/content objects
|
|
- stream always True — KIA streams by default
|
|
- tools googleSearch OR function calling (mutually exclusive)
|
|
- include_thoughts include model's internal reasoning in response
|
|
- reasoning_effort 'low' | 'high' (default: 'high')
|
|
- response_format JSON schema (mutually exclusive with function calling)
|
|
|
|
Note: tools + response_format are mutually exclusive.
|
|
"""
|
|
|
|
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:
|
|
"""
|
|
Args:
|
|
model: 'kia/gemini-3-flash', 'kia/gpt-5-2', etc.
|
|
messages: LangChain message list
|
|
reasoning_effort: 'low' | 'high'
|
|
include_thoughts: include model reasoning steps in response
|
|
google_search: enable Google Search tool (mutually exclusive with tools/response_format)
|
|
tools: custom function-calling tools (mutually exclusive with google_search/response_format)
|
|
response_format: JSON schema for structured output (mutually exclusive with tools)
|
|
max_retries: retry count on 5xx errors
|
|
retry_delay: seconds between retries
|
|
"""
|
|
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] messages_count={len(kia_messages)} | reasoning_effort={reasoning_effort} | include_thoughts={include_thoughts} | google_search={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']} | content={preview!r}")
|
|
|
|
# --- Build payload per KIA docs ---
|
|
payload: dict = {
|
|
"messages": kia_messages,
|
|
"stream": True,
|
|
}
|
|
|
|
# reasoning_effort: 'low' | 'high'
|
|
if reasoning_effort:
|
|
payload["reasoning_effort"] = reasoning_effort
|
|
|
|
# include_thoughts: show internal model reasoning
|
|
if include_thoughts:
|
|
payload["include_thoughts"] = True
|
|
|
|
# tools — google_search and function calling are mutually exclusive
|
|
# both are mutually exclusive with response_format
|
|
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']}")
|
|
|
|
last_error: Exception | None = None
|
|
for attempt in range(1, max_retries + 1):
|
|
if attempt > 1:
|
|
print(f"[KiaAI] Retry {attempt}/{max_retries} in {retry_delay}s ...")
|
|
await asyncio.sleep(retry_delay)
|
|
try:
|
|
return await self._stream_request(url, payload)
|
|
except RuntimeError as e:
|
|
last_error = e
|
|
print(f"[KiaAI] Attempt {attempt} failed: {e}")
|
|
|
|
raise last_error
|
|
|
|
async def _stream_request(self, url: str, payload: dict) -> AIMessage:
|
|
"""
|
|
Execute one SSE request and assemble the full AIMessage.
|
|
|
|
Per KIA docs, delta chunks arrive as:
|
|
data: {"choices":[{"delta":{"content":"..."},"index":0}], ...}
|
|
data: {"choices":[],"usage":{...}}
|
|
data: [DONE]
|
|
"""
|
|
full_content = ""
|
|
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}")
|
|
resp.raise_for_status()
|
|
|
|
async for line in resp.aiter_lines():
|
|
if not line:
|
|
continue
|
|
|
|
# Non-data lines may carry KIA error JSON
|
|
if not line.startswith("data:"):
|
|
try:
|
|
err = json.loads(line)
|
|
if err.get("code") and int(err["code"]) >= 400:
|
|
msg = err.get("msg", "Unknown error")
|
|
print(f"[KiaAI] Error line: code={err['code']} msg={msg!r}")
|
|
raise RuntimeError(f"KIA error {err['code']}: {msg}")
|
|
except (json.JSONDecodeError, ValueError):
|
|
pass
|
|
print(f"[KiaAI] non-data line: {line!r}")
|
|
continue
|
|
|
|
# Strip "data: " prefix (6 chars) or "data:" (5 chars)
|
|
raw_data = line[6:] if line.startswith("data: ") else line[5:]
|
|
|
|
if raw_data.strip() == "[DONE]":
|
|
print("[KiaAI] [DONE]")
|
|
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: code={data['code']} msg={msg!r}")
|
|
raise RuntimeError(f"KIA error {data['code']}: {msg}")
|
|
|
|
# Accumulate delta content from choices
|
|
for choice in (data.get("choices") or []):
|
|
delta = choice.get("delta") or {}
|
|
full_content += delta.get("content") or ""
|
|
|
|
# Capture usage from the final summary chunk
|
|
if data.get("usage"):
|
|
usage = data["usage"]
|
|
|
|
print(f"[KiaAI] Done | content_len={len(full_content)} | usage={usage}")
|
|
print(f"[KiaAI] content preview: '{full_content[:200]}'")
|
|
|
|
ai_msg = AIMessage(content=full_content)
|
|
ai_msg.usage_metadata = {
|
|
"input_tokens": usage.get("prompt_tokens", 0),
|
|
"output_tokens": usage.get("completion_tokens", 0),
|
|
"total_tokens": usage.get("total_tokens", 0),
|
|
}
|
|
return ai_msg
|
|
|
|
|
|
# Compatibility stub — KIA models bypass LangChain and use KiaAIService.invoke() directly
|
|
def get_kia_model(model: str):
|
|
return None
|