diff --git a/src/services/chatbot.py b/src/services/chatbot.py index 914e9b3..3bcabca 100644 --- a/src/services/chatbot.py +++ b/src/services/chatbot.py @@ -451,96 +451,6 @@ async def send_message( human_msg.additional_kwargs["created_at"] = human_msg_created_at input_messages = [human_msg] - is_kia = bot.model.startswith("kia/") - - if is_kia: - # --- KIA path: direct streaming without going through LangGraph astream --- - # LangGraph's stream_mode="messages" delivers chunks only AFTER call_model returns, - # because call_model is a regular async function (not a generator). - # For real token-by-token streaming we drive KiaAIService.stream() directly here. - from .kiaai import KiaAIService - - # Build the full system + history input (same as call_model does) - trimmer = get_trimmer() - # We need the existing graph state to get trimmed history - # app and config are already defined above - - # First: push the human message into the graph state so history is available - await app.aupdate_state(config, {"messages": [human_msg]}) - state = await app.aget_state(config) - all_messages = state.values.get("messages", [human_msg]) - trimmed = await trimmer.ainvoke(all_messages) - - now_state = now - kia_input = [ - SystemMessage( - f"Today: {now_state.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] send_message streaming | model='{bot.model}' | web_search={bot.web_search}") - - full_ai_content = "" - usage_meta: dict = {} - - 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: - usage_meta = chunk.usage_metadata - input_tokens = usage_meta.get("input_tokens", 0) - output_tokens = usage_meta.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 the AI response back into the graph state - ai_msg_created_at = datetime.now(UTC).isoformat() - ai_msg = AIMessage(content=full_ai_content) - ai_msg.usage_metadata = usage_meta - 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_messages = final_state.values["messages"] - human_msg_id, ai_msg_id = get_message_id(final_messages) - - 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 4da78f5..5b9f44c 100644 --- a/src/services/kiaai/chat.py +++ b/src/services/kiaai/chat.py @@ -1,10 +1,8 @@ import json import asyncio -from typing import AsyncGenerator import httpx from langchain_core.messages import ( AIMessage, - AIMessageChunk, BaseMessage, HumanMessage, SystemMessage, @@ -85,86 +83,6 @@ class KiaAIService: Note: tools + response_format are mutually exclusive. """ - async 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 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}") - 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}") - - 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 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]: - """ - Async generator — yields AIMessageChunk per token for real-time streaming. - Last chunk includes usage_metadata. - - Usage: - async for chunk in KiaAIService().stream(model, messages): - yield json.dumps({"content": chunk.content}) - """ - url, payload = await 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) - try: - async for chunk in self._stream_chunks(url, payload): - yield chunk - return - except RuntimeError as e: - last_error = e - print(f"[KiaAI] Stream attempt {attempt} failed: {e}") - - raise last_error - async def invoke( self, model: str, @@ -178,42 +96,78 @@ class KiaAIService: retry_delay: float = 2.0, ) -> AIMessage: """ - Collect all chunks and return a single AIMessage. - Used for non-streaming contexts. + 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 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 - - async def _stream_chunks( - self, url: str, payload: dict - ) -> AsyncGenerator[AIMessageChunk, None]: - """ - Core SSE parser — yields AIMessageChunk per token. - The last chunk (after [DONE]) carries usage_metadata. - Raises RuntimeError on KIA server errors (triggers retry). - """ - usage: dict = {} - async with httpx.AsyncClient(timeout=180.0) as client: async with client.stream( "POST", @@ -224,14 +178,10 @@ class KiaAIService: "Authorization": f"Bearer {settings.kia_api_key}", }, ) as resp: - print(f"[KiaAI] HTTP {resp.status_code} | url={url}") + print(f"[KiaAI] HTTP {resp.status_code}") resp.raise_for_status() - line_count = 0 async for line in resp.aiter_lines(): - line_count += 1 - if line_count <= 5: - print(f"[KiaAI] raw line[{line_count}]: {line!r}") if not line: continue @@ -248,6 +198,7 @@ class KiaAIService: 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]": @@ -266,25 +217,25 @@ class KiaAIService: print(f"[KiaAI] Error chunk: code={data['code']} msg={msg!r}") raise RuntimeError(f"KIA error {data['code']}: {msg}") - # Extract delta content from choices + # Accumulate delta content from choices for choice in (data.get("choices") or []): delta = choice.get("delta") or {} - token = delta.get("content") or "" - if token: - yield AIMessageChunk(content=token) + full_content += delta.get("content") or "" # Capture usage from the final summary chunk if data.get("usage"): usage = data["usage"] - # Emit one final empty chunk carrying usage metadata - print(f"[KiaAI] Stream complete | usage={usage}") - usage_metadata = { + 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), } - yield AIMessageChunk(content="", usage_metadata=usage_metadata) + return ai_msg # Compatibility stub — KIA models bypass LangChain and use KiaAIService.invoke() directly