add kia ai chat tools12

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):

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