From 5a35ed0b63e6eaccdb31ee8055ee90a3f8db6b13 Mon Sep 17 00:00:00 2001 From: vahidrezvani Date: Wed, 25 Feb 2026 14:06:29 +0330 Subject: [PATCH] add kia ai chat tools11 --- src/services/chatbot.py | 90 ++++++++++++++++ src/services/kiaai/chat.py | 205 +++++++++++++++++++++++-------------- 2 files changed, 217 insertions(+), 78 deletions(-) diff --git a/src/services/chatbot.py b/src/services/chatbot.py index 3bcabca..914e9b3 100644 --- a/src/services/chatbot.py +++ b/src/services/chatbot.py @@ -451,6 +451,96 @@ 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 5b9f44c..4da78f5 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, @@ -83,6 +85,86 @@ 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, @@ -96,78 +178,42 @@ class KiaAIService: 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] + Collect all chunks and return a single AIMessage. + Used for non-streaming contexts. """ 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", @@ -178,10 +224,14 @@ 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={url}") 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 @@ -198,7 +248,6 @@ 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]": @@ -217,25 +266,25 @@ class KiaAIService: print(f"[KiaAI] Error chunk: code={data['code']} msg={msg!r}") raise RuntimeError(f"KIA error {data['code']}: {msg}") - # Accumulate delta content from choices + # Extract delta content from choices 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 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 = { + # Emit one final empty chunk carrying usage metadata + print(f"[KiaAI] Stream complete | usage={usage}") + 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 + yield AIMessageChunk(content="", usage_metadata=usage_metadata) # Compatibility stub — KIA models bypass LangChain and use KiaAIService.invoke() directly