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