kimi-cli 0.44__py3-none-any.whl → 0.78__py3-none-any.whl

kimi_cli/tools/utils.py

@@ -2,17 +2,34 @@ import re
 import string
 from pathlib import Path
 
-from kosong.tooling import ToolError, ToolOk
+from jinja2 import Environment
+from kosong.tooling import BriefDisplayBlock, DisplayBlock, ToolError, ToolReturnValue
+from kosong.utils.typing import JsonType
 
 
 def load_desc(path: Path, substitutions: dict[str, str] | None = None) -> str:
     """Load a tool description from a file, with optional substitutions."""
     description = path.read_text(encoding="utf-8")
     if substitutions:
-        description = string.Template(description).substitute(substitutions)
+        description = string.Template(description).safe_substitute(substitutions)
     return description
 
 
+def load_desc_jinja(path: Path, context: dict[str, object] | None = None) -> str:
+    """Load a tool description from a file, rendered via Jinja2."""
+    description = path.read_text(encoding="utf-8")
+    env = Environment(
+        autoescape=False,
+        keep_trailing_newline=True,
+        lstrip_blocks=True,
+        trim_blocks=True,
+        variable_start_string="${",
+        variable_end_string="}",
+    )
+    template = env.from_string(description)
+    return template.render(context or {})
+
+
 def truncate_line(line: str, max_length: int, marker: str = "...") -> str:
     """
     Truncate a line if it exceeds `max_length`, preserving the beginning and the line break.
@@ -53,6 +70,23 @@ class ToolResultBuilder:
         self._n_chars = 0
         self._n_lines = 0
         self._truncation_happened = False
+        self._display: list[DisplayBlock] = []
+        self._extras: dict[str, JsonType] | None = None
+
+    @property
+    def is_full(self) -> bool:
+        """Check if output buffer is full due to character limit."""
+        return self._n_chars >= self.max_chars
+
+    @property
+    def n_chars(self) -> int:
+        """Get current character count."""
+        return self._n_chars
+
+    @property
+    def n_lines(self) -> int:
+        """Get current line count."""
+        return self._n_lines
 
     def write(self, text: str) -> int:
         """
@@ -93,8 +127,18 @@ class ToolResultBuilder:
 
         return chars_written
 
-    def ok(self, message: str = "", *, brief: str = "") -> ToolOk:
-        """Create a ToolOk result with the current output."""
+    def display(self, *blocks: DisplayBlock) -> None:
+        """Add display blocks to the tool result."""
+        self._display.extend(blocks)
+
+    def extras(self, **extras: JsonType) -> None:
+        """Add extra data to the tool result."""
+        if self._extras is None:
+            self._extras = {}
+        self._extras.update(extras)
+
+    def ok(self, message: str = "", *, brief: str = "") -> ToolReturnValue:
+        """Create a ToolReturnValue with is_error=False and the current output."""
         output = "".join(self._buffer)
 
         final_message = message
@@ -106,11 +150,16 @@ class ToolResultBuilder:
                 final_message += f" {truncation_msg}"
             else:
                 final_message = truncation_msg
+        return ToolReturnValue(
+            is_error=False,
+            output=output,
+            message=final_message,
+            display=([BriefDisplayBlock(text=brief)] if brief else []) + self._display,
+            extras=self._extras,
+        )
 
-        return ToolOk(output=output, message=final_message, brief=brief)
-
-    def error(self, message: str, *, brief: str) -> ToolError:
-        """Create a ToolError result with the current output."""
+    def error(self, message: str, *, brief: str) -> ToolReturnValue:
+        """Create a ToolReturnValue with is_error=True and the current output."""
         output = "".join(self._buffer)
 
         final_message = message
@@ -121,22 +170,13 @@ class ToolResultBuilder:
             else:
                 final_message = truncation_msg
 
-        return ToolError(output=output, message=final_message, brief=brief)
-
-    @property
-    def is_full(self) -> bool:
-        """Check if output buffer is full due to character limit."""
-        return self._n_chars >= self.max_chars
-
-    @property
-    def n_chars(self) -> int:
-        """Get current character count."""
-        return self._n_chars
-
-    @property
-    def n_lines(self) -> int:
-        """Get current line count."""
-        return self._n_lines
+        return ToolReturnValue(
+            is_error=True,
+            output=output,
+            message=final_message,
+            display=([BriefDisplayBlock(text=brief)] if brief else []) + self._display,
+            extras=self._extras,
+        )
 
 
 class ToolRejectedError(ToolError):

kimi_cli/tools/web/fetch.py

@@ -3,11 +3,15 @@ from typing import override
 
 import aiohttp
 import trafilatura
-from kosong.tooling import CallableTool2, ToolReturnType
+from kosong.tooling import CallableTool2, ToolReturnValue
 from pydantic import BaseModel, Field
 
+from kimi_cli.config import Config
+from kimi_cli.constant import USER_AGENT
+from kimi_cli.soul.toolset import get_current_tool_call_or_none
 from kimi_cli.tools.utils import ToolResultBuilder, load_desc
 from kimi_cli.utils.aiohttp import new_client_session
+from kimi_cli.utils.logging import logger
 
 
 class Params(BaseModel):
@@ -19,10 +23,23 @@ class FetchURL(CallableTool2[Params]):
     description: str = load_desc(Path(__file__).parent / "fetch.md", {})
     params: type[Params] = Params
 
+    def __init__(self, config: Config):
+        super().__init__()
+        self._service_config = config.services.moonshot_fetch
+
     @override
-    async def __call__(self, params: Params) -> ToolReturnType:
+    async def __call__(self, params: Params) -> ToolReturnValue:
+        if self._service_config:
+            ret = await self._fetch_with_service(params)
+            if not ret.is_error:
+                return ret
+            logger.warning("Failed to fetch URL via service: {error}", error=ret.message)
+            # fallback to local fetch if service fetch fails
+        return await self.fetch_with_http_get(params)
+
+    @staticmethod
+    async def fetch_with_http_get(params: Params) -> ToolReturnValue:
         builder = ToolResultBuilder(max_line_length=None)
-
         try:
             async with (
                 new_client_session() as session,
@@ -45,7 +62,12 @@ class FetchURL(CallableTool2[Params]):
                         brief=f"HTTP {response.status} error",
                     )
 
-                html = await response.text()
+                resp_text = await response.text()
+
+                content_type = response.headers.get(aiohttp.hdrs.CONTENT_TYPE, "").lower()
+                if content_type.startswith(("text/plain", "text/markdown")):
+                    builder.write(resp_text)
+                    return builder.ok("The returned content is the full content of the page.")
         except aiohttp.ClientError as e:
             return builder.error(
                 (
@@ -55,14 +77,14 @@ class FetchURL(CallableTool2[Params]):
                 brief="Network error",
             )
 
-        if not html:
+        if not resp_text:
             return builder.ok(
                 "The response body is empty.",
                 brief="Empty response body",
             )
 
         extracted_text = trafilatura.extract(
-            html,
+            resp_text,
             include_comments=True,
             include_tables=True,
             include_formatting=False,
@@ -83,13 +105,46 @@ class FetchURL(CallableTool2[Params]):
         builder.write(extracted_text)
         return builder.ok("The returned content is the main text content extracted from the page.")
 
+    async def _fetch_with_service(self, params: Params) -> ToolReturnValue:
+        assert self._service_config is not None
 
-if __name__ == "__main__":
-    import asyncio
+        tool_call = get_current_tool_call_or_none()
+        assert tool_call is not None, "Tool call is expected to be set"
 
-    async def main():
-        fetch_url_tool = FetchURL()
-        result = await fetch_url_tool(Params(url="https://trafilatura.readthedocs.io/en/latest/"))
-        print(result)
+        builder = ToolResultBuilder(max_line_length=None)
+        headers = {
+            "User-Agent": USER_AGENT,
+            "Authorization": f"Bearer {self._service_config.api_key.get_secret_value()}",
+            "Accept": "text/markdown",
+            "X-Msh-Tool-Call-Id": tool_call.id,
+            **(self._service_config.custom_headers or {}),
+        }
 
-    asyncio.run(main())
+        try:
+            async with (
+                new_client_session() as session,
+                session.post(
+                    self._service_config.base_url,
+                    headers=headers,
+                    json={"url": params.url},
+                ) as response,
+            ):
+                if response.status != 200:
+                    return builder.error(
+                        f"Failed to fetch URL via service. Status: {response.status}.",
+                        brief="Failed to fetch URL via fetch service",
+                    )
+
+                content = await response.text()
+                builder.write(content)
+                return builder.ok(
+                    "The returned content is the main content extracted from the page."
+                )
+        except aiohttp.ClientError as e:
+            return builder.error(
+                (
+                    f"Failed to fetch URL via service due to network error: {str(e)}. "
+                    "This may indicate the service is unreachable."
+                ),
+                brief="Network error when calling fetch service",
+            )

kimi_cli/tools/web/search.py

@@ -1,12 +1,13 @@
 from pathlib import Path
 from typing import override
 
-from kosong.tooling import CallableTool2, ToolReturnType
+from kosong.tooling import CallableTool2, ToolReturnValue
 from pydantic import BaseModel, Field, ValidationError
 
 from kimi_cli.config import Config
 from kimi_cli.constant import USER_AGENT
 from kimi_cli.soul.toolset import get_current_tool_call_or_none
+from kimi_cli.tools import SkipThisTool
 from kimi_cli.tools.utils import ToolResultBuilder, load_desc
 from kimi_cli.utils.aiohttp import new_client_session
 
@@ -39,19 +40,16 @@ class SearchWeb(CallableTool2[Params]):
     description: str = load_desc(Path(__file__).parent / "search.md", {})
     params: type[Params] = Params
 
-    def __init__(self, config: Config, **kwargs):
-        super().__init__(**kwargs)
-        if config.services.moonshot_search is not None:
-            self._base_url = config.services.moonshot_search.base_url
-            self._api_key = config.services.moonshot_search.api_key.get_secret_value()
-            self._custom_headers = config.services.moonshot_search.custom_headers or {}
-        else:
-            self._base_url = ""
-            self._api_key = ""
-            self._custom_headers = {}
+    def __init__(self, config: Config):
+        super().__init__()
+        if config.services.moonshot_search is None:
+            raise SkipThisTool()
+        self._base_url = config.services.moonshot_search.base_url
+        self._api_key = config.services.moonshot_search.api_key.get_secret_value()
+        self._custom_headers = config.services.moonshot_search.custom_headers or {}
 
     @override
-    async def __call__(self, params: Params) -> ToolReturnType:
+    async def __call__(self, params: Params) -> ToolReturnValue:
         builder = ToolResultBuilder(max_line_length=None)
 
         if not self._base_url or not self._api_key: