add kia ai chat tools9

src/services/kiaai/chat.py

@@ -13,31 +13,26 @@ from ...config import settings
 
 KIA_BASE_URL = "https://api.kie.ai"
 
-# Models that support reasoning_effort
-REASONING_MODELS = {
-    "gpt-5-2", "claude-sonnet-4-5", "claude-opus-4-5",
-    "gemini-3-pro", "gemini-2.5-flash", "gemini-2.5-pro",
-}
 
-
-def _to_openai_messages(messages: list[BaseMessage]) -> list[dict]:
+def _to_kia_messages(messages: list[BaseMessage]) -> list[dict]:
     """
-    Convert LangChain messages to KIA/OpenAI-compatible dict format.
+    Convert LangChain messages to KIA message format.
 
-    Supports multimodal HumanMessages where content is already a list
-    (e.g. [{"type": "text", "text": "..."}, {"type": "image_url", ...}]).
-    For plain text-only content lists, normalizes to a plain string.
+    HumanMessage.content can be:
+    - str  → {"role": "user", "content": "..."}
+    - list → {"role": "user", "content": [{"type": "text", ...}, {"type": "image_url", ...}]}
     """
     result = []
     for msg in messages:
         if isinstance(msg, SystemMessage):
-            result.append({"role": "system", "content": _normalize_content(msg.content)})
+            result.append({"role": "system", "content": msg.content})
 
         elif isinstance(msg, HumanMessage):
-            result.append({"role": "user", "content": _normalize_content(msg.content)})
+            # Multimodal content (text + image) stays as list; plain string stays as string
+            result.append({"role": "user", "content": msg.content})
 
         elif isinstance(msg, AIMessage):
-            result.append({"role": "assistant", "content": _normalize_content(msg.content)})
+            result.append({"role": "assistant", "content": msg.content})
 
         elif isinstance(msg, ToolMessage):
             result.append({
@@ -52,37 +47,21 @@ def _to_openai_messages(messages: list[BaseMessage]) -> list[dict]:
     return result
 
 
-def _normalize_content(content) -> str | list:
-    """
-    If content is a list containing only text-type items, return plain string.
-    If it contains images or other types, return the list as-is (multimodal).
-    """
-    if isinstance(content, str):
-        return content
-
-    if isinstance(content, list):
-        # Check if all items are plain text type
-        is_text_only = all(
-            isinstance(item, dict) and item.get("type") == "text"
-            for item in content
-        )
-        if is_text_only:
-            return " ".join(item.get("text", "") for item in content)
-        # Multimodal — keep as list
-        return content
-
-    return str(content)
-
-
 class KiaAIService:
     """
-    Service for calling KIA AI models directly via httpx SSE streaming.
+    Calls KIA AI API directly via httpx streaming (SSE).
 
-    KIA always responds with SSE stream regardless of stream=false in payload.
-    We collect all delta chunks and return a complete AIMessage.
+    Based on the official Python sample:
+        url = "https://api.kie.ai/{model}/v1/chat/completions"
+        payload = {"messages": [...], "stream": True, ...}
+        response = requests.post(url, headers=headers, data=json.dumps(payload), stream=True)
+        for line in response.iter_lines():
+            if line:
+                decoded_line = line.decode('utf-8')
+                if decoded_line.startswith('data: '):
+                    data = json.loads(decoded_line[6:])
 
-    Endpoint pattern:
-        POST https://api.kie.ai/{modelName}/v1/chat/completions
+    Retries on KIA 500 errors and returns a fully assembled AIMessage.
     """
 
     async def invoke(
@@ -90,50 +69,70 @@ class KiaAIService:
         model: str,
         messages: list[BaseMessage],
         reasoning_effort: str | None = None,
+        include_thoughts: bool = False,
+        tools: list[dict] | None = None,
         max_retries: int = 3,
         retry_delay: float = 2.0,
     ) -> AIMessage:
         """
-        Call KIA API and return an assembled AIMessage.
-        Retries up to max_retries times on 500 server errors.
+        Args:
+            model:            e.g. 'kia/gemini-3-flash' or 'gemini-3-flash'
+            messages:         LangChain message list
+            reasoning_effort: 'low' | 'medium' | 'high'
+            include_thoughts: stream internal reasoning steps
+            tools:            list of OpenAI-style tool dicts
+            max_retries:      retry count on 500 errors
+            retry_delay:      seconds between retries
         """
         model_name = model.removeprefix("kia/")
         url = f"{KIA_BASE_URL}/{model_name}/v1/chat/completions"
 
-        openai_messages = _to_openai_messages(messages)
+        kia_messages = _to_kia_messages(messages)
 
         print(f"[KiaAI] POST {url}")
-        print(f"[KiaAI] messages_count={len(openai_messages)} | reasoning_effort={reasoning_effort}")
-        for i, m in enumerate(openai_messages):
+        print(f"[KiaAI] messages_count={len(kia_messages)} | reasoning_effort={reasoning_effort} | include_thoughts={include_thoughts}")
+        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 exactly like the official sample
         payload: dict = {
-            "model": model_name,
-            "messages": openai_messages,
+            "messages": kia_messages,
             "stream": True,
         }
-        if reasoning_effort and model_name in REASONING_MODELS:
+
+        if reasoning_effort:
             payload["reasoning_effort"] = reasoning_effort
 
-        last_error: Exception | None = None
+        if include_thoughts:
+            payload["include_thoughts"] = True
 
+        if tools:
+            payload["tools"] = tools
+
+        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 {attempt}/{max_retries} after {retry_delay}s...")
+                print(f"[KiaAI] Retry {attempt}/{max_retries} in {retry_delay}s ...")
                 await asyncio.sleep(retry_delay)
-
             try:
-                result = await self._do_request(url, payload)
-                return result
+                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 _do_request(self, url: str, payload: dict) -> AIMessage:
-        """Execute a single SSE streaming request to KIA and assemble AIMessage."""
+    async def _stream_request(self, url: str, payload: dict) -> AIMessage:
+        """
+        Execute one SSE request and assemble the full AIMessage.
+
+        Mirrors the official sample's per-line handling:
+            if decoded_line.startswith('data: '):
+                data = json.loads(decoded_line[6:])
+        """
         full_content = ""
         usage: dict = {}
 
@@ -141,48 +140,53 @@ class KiaAIService:
             async with client.stream(
                 "POST",
                 url,
-                json=payload,
+                content=json.dumps(payload),           # mirrors: data=json.dumps(payload)
                 headers={
-                    "Authorization": f"Bearer {settings.kia_api_key}",
                     "Content-Type": "application/json",
+                    "Authorization": f"Bearer {settings.kia_api_key}",
                 },
             ) as resp:
                 print(f"[KiaAI] HTTP {resp.status_code}")
                 resp.raise_for_status()
 
                 async for line in resp.aiter_lines():
-                    if not line.strip():
+                    if not line:
                         continue
 
-                    raw_data = line.removeprefix("data:").strip()
+                    # mirrors: if decoded_line.startswith('data: '):
+                    if not line.startswith("data:"):
+                        print(f"[KiaAI] non-data line: {line!r}")
+                        continue
 
-                    if raw_data == "[DONE]":
-                        print("[KiaAI] Stream [DONE]")
+                    raw_data = line[6:] if line.startswith("data: ") else line[5:]
+
+                    if raw_data.strip() == "[DONE]":
+                        print("[KiaAI] [DONE]")
                         break
 
                     try:
-                        chunk = json.loads(raw_data)
+                        # mirrors: data = json.loads(decoded_line[6:])
+                        data = json.loads(raw_data)
                     except json.JSONDecodeError:
-                        print(f"[KiaAI] Skipping non-JSON line: {raw_data!r}")
+                        print(f"[KiaAI] JSON decode error: {raw_data!r}")
                         continue
 
-                    # KIA server error inside stream body
-                    if chunk.get("code") and int(chunk["code"]) >= 500:
-                        error_msg = chunk.get("msg", "Unknown KIA server error")
-                        print(f"[KiaAI] Server error chunk: code={chunk['code']} msg={error_msg!r}")
-                        raise RuntimeError(f"KIA API error {chunk['code']}: {error_msg}")
+                    # KIA server error embedded in stream
+                    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}")
+                        raise RuntimeError(f"KIA error {data['code']}: {msg}")
 
-                    choices = chunk.get("choices") or []
-                    for choice in choices:
+                    # Accumulate delta content
+                    for choice in (data.get("choices") or []):
                         delta = choice.get("delta") or {}
-                        content_piece = delta.get("content") or ""
-                        full_content += content_piece
+                        full_content += delta.get("content") or ""
 
-                    if chunk.get("usage"):
-                        usage = chunk["usage"]
+                    if data.get("usage"):
+                        usage = data["usage"]
 
-        print(f"[KiaAI] Final content (first 200): '{full_content[:200]}'")
-        print(f"[KiaAI] usage={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 = {
@@ -190,10 +194,9 @@ class KiaAIService:
             "output_tokens": usage.get("completion_tokens", 0),
             "total_tokens": usage.get("total_tokens", 0),
         }
-
         return ai_msg
 
 
-# Compatibility stub — KIA models bypass LangChain and use KiaAIService.invoke() directly
+# Compatibility stub
 def get_kia_model(model: str):
     return None