From 4d857ab524425b8da81c907c2d97561842010da4 Mon Sep 17 00:00:00 2001 From: vahidrezvani Date: Thu, 26 Feb 2026 10:20:35 +0330 Subject: [PATCH] add kia ai chat tools17 --- src/services/chatbot.py | 32 +++++++++++++++++---- src/services/kiaai/__init__.py | 4 +-- src/services/kiaai/workflow.py | 51 +++++++++++++--------------------- 3 files changed, 47 insertions(+), 40 deletions(-) diff --git a/src/services/chatbot.py b/src/services/chatbot.py index b77dab2..8146df8 100644 --- a/src/services/chatbot.py +++ b/src/services/chatbot.py @@ -31,7 +31,7 @@ from langchain_core.messages import ( ) from .database.service import DatabaseService -from .kiaai import get_kia_model, kia_stream_and_save +from .kiaai import get_kia_model, kia_token_stream from ..config import settings from ..models import Chatbot, Message, BotStatus, BotType from ..messages import MESSAGES @@ -423,14 +423,17 @@ async def send_message( input_messages = [human_msg] # ── KIA path: token-by-token streaming via kiaai.workflow ────────────── + # Pre-fetch history here (one DB read) so kia_token_stream has zero + # DB operations before the first token — streaming starts immediately. if bot.model.startswith("kia/"): + kia_state = await app.aget_state(config) + messages_history = kia_state.values.get("messages", []) + + stream_result: dict = {} try: - async for item in kia_stream_and_save(app, config, human_msg, bot, now, get_trimmer()): + async for item in kia_token_stream(messages_history, human_msg, bot, now, get_trimmer()): if isinstance(item, dict): - input_tokens = item["input_tokens"] - output_tokens = item["output_tokens"] - ai_msg_created_at = item["ai_msg_created_at"] - human_msg_id, ai_msg_id = get_message_id(item["final_msgs"]) + stream_result = item else: yield item except Exception as e: @@ -438,6 +441,23 @@ async def send_message( yield json.dumps({"error": True, "detail": str(e)}, ensure_ascii=False) return + # ── Post-stream: persist both messages to LangGraph then DB ───────── + ai_msg = AIMessage(content=stream_result["full_ai_content"]) + ai_msg.usage_metadata = { + "input_tokens": stream_result["input_tokens"], + "output_tokens": stream_result["output_tokens"], + "total_tokens": stream_result["input_tokens"] + stream_result["output_tokens"], + } + ai_msg.additional_kwargs["created_at"] = stream_result["ai_msg_created_at"] + + await app.aupdate_state(config, {"messages": [human_msg, ai_msg]}, as_node="model") + final_state = await app.aget_state(config) + human_msg_id, ai_msg_id = get_message_id(final_state.values["messages"]) + + input_tokens = stream_result["input_tokens"] + output_tokens = stream_result["output_tokens"] + ai_msg_created_at = stream_result["ai_msg_created_at"] + await user.save() await Message.create( input_tokens=input_tokens, diff --git a/src/services/kiaai/__init__.py b/src/services/kiaai/__init__.py index f5bf000..91f72f4 100644 --- a/src/services/kiaai/__init__.py +++ b/src/services/kiaai/__init__.py @@ -1,4 +1,4 @@ from .chat import KiaAIService, get_kia_model -from .workflow import kia_stream_and_save +from .workflow import kia_token_stream -__all__ = ["KiaAIService", "get_kia_model", "kia_stream_and_save"] +__all__ = ["KiaAIService", "get_kia_model", "kia_token_stream"] diff --git a/src/services/kiaai/workflow.py b/src/services/kiaai/workflow.py index 1e25f5d..e6d180f 100644 --- a/src/services/kiaai/workflow.py +++ b/src/services/kiaai/workflow.py @@ -22,39 +22,37 @@ def _build_system_message(now, bot) -> SystemMessage: ) -async def kia_stream_and_save( - app, - config: dict, +async def kia_token_stream( + messages_history: list, human_msg: HumanMessage, bot, now, trimmer, ) -> AsyncGenerator: """ - KIA direct-stream workflow. Drives KiaAIService.stream() and manages - LangGraph state manually (bypassing app.astream entirely). + Pure KIA streaming — zero DB operations inside this generator. + + Accepts pre-fetched ``messages_history`` so the KIA request starts + immediately without any database round-trips blocking the first token. Yields: - str – JSON-encoded token chunks: ``{"content": "..."}`` - dict – final metadata (always the last yielded item), keys: - ``input_tokens``, ``output_tokens``, ``ai_msg_created_at``, - ``final_msgs`` + str — JSON token chunk (immediately forwarded to client): + ``{"content": "..."}\\n`` + dict — single final item (last yield), used by the caller for + DB persistence: + ``{"full_ai_content", "input_tokens", "output_tokens", + "ai_msg_created_at"}`` Raises: - Any exception from KiaAIService.stream() is propagated to the caller. - If an exception is raised after tokens have already been yielded, the - graph state will contain the human message but no AI reply. + Any exception from KiaAIService.stream() propagates to the caller. """ - # 1. Push human message into graph state - await app.aupdate_state(config, {"messages": [human_msg]}, as_node="model") - state = await app.aget_state(config) - all_msgs = state.values.get("messages", [human_msg]) + # Trim using tiktoken locally — no network call, no DB round-trip + all_msgs = messages_history + [human_msg] trimmed = await trimmer.ainvoke(all_msgs) - kia_input = [_build_system_message(now, bot)] + trimmed print( - f"[KiaAI] Direct stream | model='{bot.model}' " + f"[KiaAI] Stream start | model='{bot.model}' " f"| web_search={bot.web_search} | history_len={len(trimmed)}" ) @@ -63,7 +61,7 @@ async def kia_stream_and_save( input_tokens = 0 output_tokens = 0 - # 2. Stream tokens — exceptions propagate to the caller + # Stream — each token forwarded to client immediately, buffered in memory async for chunk in KiaAIService().stream( model=bot.model, messages=kia_input, @@ -77,21 +75,10 @@ async def kia_stream_and_save( input_tokens = chunk.usage_metadata.get("input_tokens", 0) output_tokens = chunk.usage_metadata.get("output_tokens", 0) - # 3. Save completed AI message back into 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]}, as_node="model") - - # 4. Yield final metadata for the caller to persist and respond with - final_state = await app.aget_state(config) + # Yield buffered result — caller persists this to DB after all tokens sent yield { + "full_ai_content": full_ai_content, "input_tokens": input_tokens, "output_tokens": output_tokens, "ai_msg_created_at": ai_msg_created_at, - "final_msgs": final_state.values["messages"], }