diff --git a/src/services/kiaai/chat.py b/src/services/kiaai/chat.py index d3ed16a..a27dff6 100644 --- a/src/services/kiaai/chat.py +++ b/src/services/kiaai/chat.py @@ -13,31 +13,26 @@ from ...config import settings KIA_BASE_URL = "https://api.kie.ai" -# Models that support reasoning_effort -REASONING_MODELS = { - "gpt-5-2", "claude-sonnet-4-5", "claude-opus-4-5", - "gemini-3-pro", "gemini-2.5-flash", "gemini-2.5-pro", -} - -def _to_openai_messages(messages: list[BaseMessage]) -> list[dict]: +def _to_kia_messages(messages: list[BaseMessage]) -> list[dict]: """ - Convert LangChain messages to KIA/OpenAI-compatible dict format. + Convert LangChain messages to KIA message format. - Supports multimodal HumanMessages where content is already a list - (e.g. [{"type": "text", "text": "..."}, {"type": "image_url", ...}]). - For plain text-only content lists, normalizes to a plain string. + HumanMessage.content can be: + - str → {"role": "user", "content": "..."} + - list → {"role": "user", "content": [{"type": "text", ...}, {"type": "image_url", ...}]} """ result = [] for msg in messages: if isinstance(msg, SystemMessage): - result.append({"role": "system", "content": _normalize_content(msg.content)}) + result.append({"role": "system", "content": msg.content}) elif isinstance(msg, HumanMessage): - result.append({"role": "user", "content": _normalize_content(msg.content)}) + # Multimodal content (text + image) stays as list; plain string stays as string + result.append({"role": "user", "content": msg.content}) elif isinstance(msg, AIMessage): - result.append({"role": "assistant", "content": _normalize_content(msg.content)}) + result.append({"role": "assistant", "content": msg.content}) elif isinstance(msg, ToolMessage): result.append({ @@ -52,37 +47,21 @@ def _to_openai_messages(messages: list[BaseMessage]) -> list[dict]: return result -def _normalize_content(content) -> str | list: - """ - If content is a list containing only text-type items, return plain string. - If it contains images or other types, return the list as-is (multimodal). - """ - if isinstance(content, str): - return content - - if isinstance(content, list): - # Check if all items are plain text type - 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) - # Multimodal — keep as list - return content - - return str(content) - - class KiaAIService: """ - Service for calling KIA AI models directly via httpx SSE streaming. + Calls KIA AI API directly via httpx streaming (SSE). - KIA always responds with SSE stream regardless of stream=false in payload. - We collect all delta chunks and return a complete AIMessage. + Based on the official Python sample: + url = "https://api.kie.ai/{model}/v1/chat/completions" + payload = {"messages": [...], "stream": True, ...} + response = requests.post(url, headers=headers, data=json.dumps(payload), stream=True) + for line in response.iter_lines(): + if line: + decoded_line = line.decode('utf-8') + if decoded_line.startswith('data: '): + data = json.loads(decoded_line[6:]) - Endpoint pattern: - POST https://api.kie.ai/{modelName}/v1/chat/completions + Retries on KIA 500 errors and returns a fully assembled AIMessage. """ async def invoke( @@ -90,50 +69,70 @@ class KiaAIService: model: str, messages: list[BaseMessage], reasoning_effort: str | None = None, + include_thoughts: bool = False, + tools: list[dict] | None = None, max_retries: int = 3, retry_delay: float = 2.0, ) -> AIMessage: """ - Call KIA API and return an assembled AIMessage. - Retries up to max_retries times on 500 server errors. + Args: + model: e.g. 'kia/gemini-3-flash' or 'gemini-3-flash' + messages: LangChain message list + reasoning_effort: 'low' | 'medium' | 'high' + include_thoughts: stream internal reasoning steps + tools: list of OpenAI-style tool dicts + max_retries: retry count on 500 errors + retry_delay: seconds between retries """ model_name = model.removeprefix("kia/") url = f"{KIA_BASE_URL}/{model_name}/v1/chat/completions" - openai_messages = _to_openai_messages(messages) + kia_messages = _to_kia_messages(messages) print(f"[KiaAI] POST {url}") - print(f"[KiaAI] messages_count={len(openai_messages)} | reasoning_effort={reasoning_effort}") - for i, m in enumerate(openai_messages): + print(f"[KiaAI] messages_count={len(kia_messages)} | reasoning_effort={reasoning_effort} | include_thoughts={include_thoughts}") + 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 exactly like the official sample payload: dict = { - "model": model_name, - "messages": openai_messages, + "messages": kia_messages, "stream": True, } - if reasoning_effort and model_name in REASONING_MODELS: + + if reasoning_effort: payload["reasoning_effort"] = reasoning_effort - last_error: Exception | None = None + if include_thoughts: + payload["include_thoughts"] = True + if tools: + payload["tools"] = tools + + 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 {attempt}/{max_retries} after {retry_delay}s...") + print(f"[KiaAI] Retry {attempt}/{max_retries} in {retry_delay}s ...") await asyncio.sleep(retry_delay) - try: - result = await self._do_request(url, payload) - return result + 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 _do_request(self, url: str, payload: dict) -> AIMessage: - """Execute a single SSE streaming request to KIA and assemble AIMessage.""" + async def _stream_request(self, url: str, payload: dict) -> AIMessage: + """ + Execute one SSE request and assemble the full AIMessage. + + Mirrors the official sample's per-line handling: + if decoded_line.startswith('data: '): + data = json.loads(decoded_line[6:]) + """ full_content = "" usage: dict = {} @@ -141,48 +140,53 @@ class KiaAIService: async with client.stream( "POST", url, - json=payload, + content=json.dumps(payload), # mirrors: data=json.dumps(payload) headers={ - "Authorization": f"Bearer {settings.kia_api_key}", "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.strip(): + if not line: continue - raw_data = line.removeprefix("data:").strip() + # mirrors: if decoded_line.startswith('data: '): + if not line.startswith("data:"): + print(f"[KiaAI] non-data line: {line!r}") + continue - if raw_data == "[DONE]": - print("[KiaAI] Stream [DONE]") + raw_data = line[6:] if line.startswith("data: ") else line[5:] + + if raw_data.strip() == "[DONE]": + print("[KiaAI] [DONE]") break try: - chunk = json.loads(raw_data) + # mirrors: data = json.loads(decoded_line[6:]) + data = json.loads(raw_data) except json.JSONDecodeError: - print(f"[KiaAI] Skipping non-JSON line: {raw_data!r}") + print(f"[KiaAI] JSON decode error: {raw_data!r}") continue - # KIA server error inside stream body - if chunk.get("code") and int(chunk["code"]) >= 500: - error_msg = chunk.get("msg", "Unknown KIA server error") - print(f"[KiaAI] Server error chunk: code={chunk['code']} msg={error_msg!r}") - raise RuntimeError(f"KIA API error {chunk['code']}: {error_msg}") + # KIA server error embedded in stream + 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}") - choices = chunk.get("choices") or [] - for choice in choices: + # Accumulate delta content + for choice in (data.get("choices") or []): delta = choice.get("delta") or {} - content_piece = delta.get("content") or "" - full_content += content_piece + full_content += delta.get("content") or "" - if chunk.get("usage"): - usage = chunk["usage"] + if data.get("usage"): + usage = data["usage"] - print(f"[KiaAI] Final content (first 200): '{full_content[:200]}'") - print(f"[KiaAI] usage={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 = { @@ -190,10 +194,9 @@ class KiaAIService: "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 +# Compatibility stub def get_kia_model(model: str): return None