diff --git a/src/services/chatbot.py b/src/services/chatbot.py index 3bcabca..4b66d0f 100644 --- a/src/services/chatbot.py +++ b/src/services/chatbot.py @@ -451,6 +451,90 @@ async def send_message( human_msg.additional_kwargs["created_at"] = human_msg_created_at input_messages = [human_msg] + # ── KIA path: stream tokens directly, bypass LangGraph astream ────────── + # LangGraph delivers an AIMessage as ONE chunk (not token-by-token) because + # call_model returns a complete AIMessage. For real streaming we drive + # KiaAIService.stream() here and manage the graph state manually. + if bot.model.startswith("kia/"): + from .kiaai import KiaAIService + + # Push the human message into the graph state first + await app.aupdate_state(config, {"messages": [human_msg]}) + state = await app.aget_state(config) + all_msgs = state.values.get("messages", [human_msg]) + trimmed = await get_trimmer().ainvoke(all_msgs) + + kia_input = [ + SystemMessage( + f"Today: {now.isoformat()}." + "You are an AI assistant designed to answer user queries conversationally, using tools when necessary." + "Follow these instructions carefully: " + "1. **Mathematical Formatting**: When responding to questions involving mathematics, format all mathematical expressions and formulas using Markdown with LaTeX notation for clarity and readability. Use inline LaTeX (e.g., `$x^2 + 2x + 1$`) for expressions within a sentence, and display LaTeX (e.g., `$$x^2 + 2x + 1 = 0$$`) for standalone equations or complex formulas. Ensure the formatting is compatible with Markdown renderers that support LaTeX." + f"{bot.prompt if bot.web_search else ''}" + "Follow these guidelines to ensure a seamless and informative conversation with the user." + ) + ] + trimmed + + print(f"[KiaAI] Direct stream | model='{bot.model}' | web_search={bot.web_search} | history_len={len(trimmed)}") + + full_ai_content = "" + ai_msg_created_at = datetime.now(UTC).isoformat() + + try: + async for chunk in KiaAIService().stream( + model=bot.model, + messages=kia_input, + google_search=bool(bot.web_search), + ): + token = chunk.content or "" + if token: + full_ai_content += token + yield json.dumps({"content": token}, ensure_ascii=False) + if chunk.usage_metadata: + input_tokens = chunk.usage_metadata.get("input_tokens", 0) + output_tokens = chunk.usage_metadata.get("output_tokens", 0) + except Exception as e: + print(f"[KiaAI] Streaming error: {e}") + yield json.dumps({"error": True, "detail": str(e)}, ensure_ascii=False) + return + + # Save AI response back into the graph state + ai_msg = AIMessage(content=full_ai_content) + ai_msg.usage_metadata = { + "input_tokens": input_tokens, + "output_tokens": output_tokens, + "total_tokens": input_tokens + output_tokens, + } + ai_msg.additional_kwargs["created_at"] = ai_msg_created_at + await app.aupdate_state(config, {"messages": [ai_msg]}) + + await user.save() + await Message.create( + input_tokens=input_tokens, + output_tokens=output_tokens, + chatbot_id=chatbot.id, + cost=bot.cost, + cost_from_gift=usage_report["gift"], + cost_from_credit=usage_report["credit"], + cost_from_free=usage_report["free"], + ) + + final_state = await app.aget_state(config) + final_msgs = final_state.values["messages"] + human_msg_id, ai_msg_id = get_message_id(final_msgs) + + yield json.dumps({ + "credit": user.credit, + "free": user.free_credit, + "gift": user.gift_credit if user.gift_credit is not None else 0, + "ai_message_id": ai_msg_id, + "human_message_id": human_msg_id, + "ai_message_created_at": ai_msg_created_at, + "human_message_created_at": human_msg_created_at, + }) + return + + # ── Non-KIA path (LangGraph astream) ──────────────────────────────────── async for chunk, metadata in app.astream({"messages": input_messages}, config, stream_mode="messages"): try: if isinstance(chunk, AIMessage): diff --git a/src/services/kiaai/chat.py b/src/services/kiaai/chat.py index 5b9f44c..d55c110 100644 --- a/src/services/kiaai/chat.py +++ b/src/services/kiaai/chat.py @@ -1,8 +1,10 @@ import json import asyncio +from typing import AsyncGenerator import httpx from langchain_core.messages import ( AIMessage, + AIMessageChunk, BaseMessage, HumanMessage, SystemMessage, @@ -71,69 +73,38 @@ def _normalize_content(content) -> str | list: 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( + # ------------------------------------------------------------------ + # Internal helpers + # ------------------------------------------------------------------ + + def _build_payload( 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 - """ + 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] messages_count={len(kia_messages)} | reasoning_effort={reasoning_effort} | include_thoughts={include_thoughts} | google_search={google_search}") + 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']} | content={preview!r}") + print(f"[KiaAI] msg[{i}] role={m['role']} | {preview!r}") - # --- Build payload per KIA docs --- - payload: dict = { - "messages": kia_messages, - "stream": True, - } - - # reasoning_effort: 'low' | 'high' + payload: dict = {"messages": kia_messages, "stream": True} 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: @@ -142,30 +113,16 @@ class KiaAIService: payload["response_format"] = response_format print(f"[KiaAI] payload keys: {[k for k in payload if k != 'messages']}") + return url, payload - 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: + async def _iter_sse( + self, url: str, payload: dict + ) -> AsyncGenerator[AIMessageChunk, None]: """ - 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] + 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()). """ - full_content = "" usage: dict = {} async with httpx.AsyncClient(timeout=180.0) as client: @@ -178,10 +135,15 @@ class KiaAIService: "Authorization": f"Bearer {settings.kia_api_key}", }, ) as resp: - print(f"[KiaAI] HTTP {resp.status_code}") + print(f"[KiaAI] HTTP {resp.status_code} | {url}") resp.raise_for_status() + line_count = 0 async for line in resp.aiter_lines(): + line_count += 1 + if line_count <= 3: + print(f"[KiaAI] raw[{line_count}]: {line!r}") + if not line: continue @@ -189,16 +151,16 @@ class KiaAIService: if not line.startswith("data:"): try: err = json.loads(line) - if err.get("code") and int(err["code"]) >= 400: + code = err.get("code") + if code and int(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}") + 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 - # 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]": @@ -214,30 +176,114 @@ class KiaAIService: # 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}") + print(f"[KiaAI] Error chunk: {data['code']} {msg!r}") raise RuntimeError(f"KIA error {data['code']}: {msg}") - # Accumulate delta content from choices + # Yield content tokens for choice in (data.get("choices") or []): delta = choice.get("delta") or {} - full_content += delta.get("content") or "" + token = delta.get("content") or "" + if token: + yield AIMessageChunk(content=token) - # Capture usage from the final summary chunk + # Capture usage from 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]}'") + # 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 = { - "input_tokens": usage.get("prompt_tokens", 0), - "output_tokens": usage.get("completion_tokens", 0), - "total_tokens": usage.get("total_tokens", 0), - } + ai_msg.usage_metadata = usage return ai_msg -# Compatibility stub — KIA models bypass LangChain and use KiaAIService.invoke() directly +# Compatibility stub — KIA models bypass LangChain and use KiaAIService directly def get_kia_model(model: str): return None