add kia ai chat tools17

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"]
+                    stream_result = item
-                    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"])
                 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,

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"]

src/services/kiaai/workflow.py

@@ -22,39 +22,37 @@ def _build_system_message(now, bot) -> SystemMessage:
     )
 
 
-async def kia_stream_and_save(
+async def kia_token_stream(
-    app,
+    messages_history: list,
-    config: dict,
     human_msg: HumanMessage,
     bot,
     now,
     trimmer,
 ) -> AsyncGenerator:
     """
-    KIA direct-stream workflow. Drives KiaAIService.stream() and manages
+    Pure KIA streaming — zero DB operations inside this generator.
-    LangGraph state manually (bypassing app.astream entirely).
+
+    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": "..."}``
+        str  — JSON token chunk (immediately forwarded to client):
-        dict – final metadata (always the last yielded item), keys:
+               ``{"content": "..."}\\n``
-               ``input_tokens``, ``output_tokens``, ``ai_msg_created_at``,
+        dict — single final item (last yield), used by the caller for
-               ``final_msgs``
+               DB persistence:
+               ``{"full_ai_content", "input_tokens", "output_tokens",
+                  "ai_msg_created_at"}``
 
     Raises:
-        Any exception from KiaAIService.stream() is propagated to the caller.
+        Any exception from KiaAIService.stream() propagates 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.
     """
-    # 1. Push human message into graph state
+    # Trim using tiktoken locally — no network call, no DB round-trip
-    await app.aupdate_state(config, {"messages": [human_msg]}, as_node="model")
+    all_msgs = messages_history + [human_msg]
-    state = await app.aget_state(config)
-    all_msgs = state.values.get("messages", [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
+    # Yield buffered result — caller persists this to DB after all tokens sent
-    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 {
+        "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"],
     }