kimi-cli 0.35__py3-none-any.whl

kimi_cli/tools/utils.py

@@ -0,0 +1,150 @@
+import re
+import string
+from pathlib import Path
+
+from kosong.tooling import ToolError, ToolOk
+
+
+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()
+    if substitutions:
+        description = string.Template(description).substitute(substitutions)
+    return description
+
+
+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.
+    The output may be longer than `max_length` if it is too short to fit the marker.
+    """
+    if len(line) <= max_length:
+        return line
+
+    # Find line breaks at the end of the line
+    m = re.search(r"[\r\n]+$", line)
+    linebreak = m.group(0) if m else ""
+    end = marker + linebreak
+    max_length = max(max_length, len(end))
+    return line[: max_length - len(end)] + end
+
+
+# Default output limits
+DEFAULT_MAX_CHARS = 50_000
+DEFAULT_MAX_LINE_LENGTH = 2000
+
+
+class ToolResultBuilder:
+    """
+    Builder for tool results with character and line limits.
+    """
+
+    def __init__(
+        self,
+        max_chars: int = DEFAULT_MAX_CHARS,
+        max_line_length: int | None = DEFAULT_MAX_LINE_LENGTH,
+    ):
+        self.max_chars = max_chars
+        self.max_line_length = max_line_length
+        self._marker = "[...truncated]"
+        if max_line_length is not None:
+            assert max_line_length > len(self._marker)
+        self._buffer: list[str] = []
+        self._n_chars = 0
+        self._n_lines = 0
+        self._truncation_happened = False
+
+    def write(self, text: str) -> int:
+        """
+        Write text to the output buffer.
+
+        Returns:
+            int: Number of characters actually written
+        """
+        if self.is_full:
+            return 0
+
+        lines = text.splitlines(keepends=True)
+        if not lines:
+            return 0
+
+        chars_written = 0
+
+        for line in lines:
+            if self.is_full:
+                break
+
+            original_line = line
+            remaining_chars = self.max_chars - self._n_chars
+            limit = (
+                min(remaining_chars, self.max_line_length)
+                if self.max_line_length is not None
+                else remaining_chars
+            )
+            line = truncate_line(line, limit, self._marker)
+            if line != original_line:
+                self._truncation_happened = True
+
+            self._buffer.append(line)
+            chars_written += len(line)
+            self._n_chars += len(line)
+            if line.endswith("\n"):
+                self._n_lines += 1
+
+        return chars_written
+
+    def ok(self, message: str = "", *, brief: str = "") -> ToolOk:
+        """Create a ToolOk result with the current output."""
+        output = "".join(self._buffer)
+
+        final_message = message
+        if final_message and not final_message.endswith("."):
+            final_message += "."
+        truncation_msg = "Output is truncated to fit in the message."
+        if self._truncation_happened:
+            if final_message:
+                final_message += f" {truncation_msg}"
+            else:
+                final_message = truncation_msg
+
+        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."""
+        output = "".join(self._buffer)
+
+        final_message = message
+        if self._truncation_happened:
+            truncation_msg = "Output is truncated to fit in the message."
+            if final_message:
+                final_message += f" {truncation_msg}"
+            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
+
+
+class ToolRejectedError(ToolError):
+    def __init__(self):
+        super().__init__(
+            message=(
+                "The tool call is rejected by the user. "
+                "Please follow the new instructions from the user."
+            ),
+            brief="Rejected by user",
+        )

kimi_cli/tools/web/__init__.py

@@ -0,0 +1,4 @@
+from .fetch import FetchURL
+from .search import SearchWeb
+
+__all__ = ("SearchWeb", "FetchURL")

kimi_cli/tools/web/fetch.md

@@ -0,0 +1 @@
+Fetch a web page from a URL and extract main text content from it.

kimi_cli/tools/web/fetch.py

@@ -0,0 +1,94 @@
+from pathlib import Path
+from typing import override
+
+import aiohttp
+import trafilatura
+from kosong.tooling import CallableTool2, ToolReturnType
+from pydantic import BaseModel, Field
+
+from kimi_cli.tools.utils import ToolResultBuilder, load_desc
+
+
+class Params(BaseModel):
+    url: str = Field(description="The URL to fetch content from.")
+
+
+class FetchURL(CallableTool2[Params]):
+    name: str = "FetchURL"
+    description: str = load_desc(Path(__file__).parent / "fetch.md", {})
+    params: type[Params] = Params
+
+    @override
+    async def __call__(self, params: Params) -> ToolReturnType:
+        builder = ToolResultBuilder(max_line_length=None)
+
+        try:
+            async with (
+                aiohttp.ClientSession() as session,
+                session.get(
+                    params.url,
+                    headers={
+                        "User-Agent": (
+                            "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 "
+                            "(KHTML, like Gecko) Chrome/91.0.4472.124 Safari/537.36"
+                        ),
+                    },
+                ) as response,
+            ):
+                if response.status >= 400:
+                    return builder.error(
+                        (
+                            f"Failed to fetch URL. Status: {response.status}. "
+                            f"This may indicate the page is not accessible or the server is down."
+                        ),
+                        brief=f"HTTP {response.status} error",
+                    )
+
+                html = await response.text()
+        except aiohttp.ClientError as e:
+            return builder.error(
+                (
+                    f"Failed to fetch URL due to network error: {str(e)}. "
+                    "This may indicate the URL is invalid or the server is unreachable."
+                ),
+                brief="Network error",
+            )
+
+        if not html:
+            return builder.ok(
+                "The response body is empty.",
+                brief="Empty response body",
+            )
+
+        extracted_text = trafilatura.extract(
+            html,
+            include_comments=True,
+            include_tables=True,
+            include_formatting=False,
+            output_format="txt",
+            with_metadata=True,
+        )
+
+        if not extracted_text:
+            return builder.error(
+                (
+                    "Failed to extract meaningful content from the page. "
+                    "This may indicate the page content is not suitable for text extraction, "
+                    "or the page requires JavaScript to render its content."
+                ),
+                brief="No content extracted",
+            )
+
+        builder.write(extracted_text)
+        return builder.ok("The returned content is the main text content extracted from the page.")
+
+
+if __name__ == "__main__":
+    import asyncio
+
+    async def main():
+        fetch_url_tool = FetchURL()
+        result = await fetch_url_tool(Params(url="https://trafilatura.readthedocs.io/en/latest/"))
+        print(result)
+
+    asyncio.run(main())

kimi_cli/tools/web/search.md

@@ -0,0 +1 @@
+WebSearch tool allows you to search on the internet to get latest information, including news, documents, release notes, blog posts, papers, etc.

kimi_cli/tools/web/search.py

@@ -0,0 +1,126 @@
+from pathlib import Path
+from typing import override
+
+import aiohttp
+from kosong.tooling import CallableTool2, ToolReturnType
+from pydantic import BaseModel, Field, ValidationError
+
+import kimi_cli
+from kimi_cli.config import Config
+from kimi_cli.soul.toolset import get_current_tool_call_or_none
+from kimi_cli.tools.utils import ToolResultBuilder, load_desc
+
+
+class Params(BaseModel):
+    query: str = Field(description="The query text to search for.")
+    limit: int = Field(
+        description=(
+            "The number of results to return. "
+            "Typically you do not need to set this value. "
+            "When the results do not contain what you need, "
+            "you probably want to give a more concrete query."
+        ),
+        default=5,
+        ge=1,
+        le=20,
+    )
+    include_content: bool = Field(
+        description=(
+            "Whether to include the content of the web pages in the results. "
+            "It can consume a large amount of tokens when this is set to True. "
+            "You should avoid enabling this when `limit` is set to a large value."
+        ),
+        default=False,
+    )
+
+
+class SearchWeb(CallableTool2[Params]):
+    name: str = "SearchWeb"
+    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()
+        else:
+            self._base_url = ""
+            self._api_key = ""
+
+    @override
+    async def __call__(self, params: Params) -> ToolReturnType:
+        builder = ToolResultBuilder(max_line_length=None)
+
+        if not self._base_url or not self._api_key:
+            return builder.error(
+                "Search service is not configured. You may want to try other methods to search.",
+                brief="Search service not configured",
+            )
+
+        tool_call = get_current_tool_call_or_none()
+        assert tool_call is not None, "Tool call is expected to be set"
+
+        async with (
+            aiohttp.ClientSession() as session,
+            session.post(
+                self._base_url,
+                headers={
+                    "User-Agent": kimi_cli.USER_AGENT,
+                    "Authorization": f"Bearer {self._api_key}",
+                    "X-Msh-Tool-Call-Id": tool_call.id,
+                },
+                json={
+                    "text_query": params.query,
+                    "limit": params.limit,
+                    "enable_page_crawling": params.include_content,
+                    "timeout_seconds": 30,
+                },
+            ) as response,
+        ):
+            if response.status != 200:
+                return builder.error(
+                    (
+                        f"Failed to search. Status: {response.status}. "
+                        "This may indicates that the search service is currently unavailable."
+                    ),
+                    brief="Failed to search",
+                )
+
+            try:
+                results = Response(**await response.json()).search_results
+            except ValidationError as e:
+                return builder.error(
+                    (
+                        f"Failed to parse search results. Error: {e}. "
+                        "This may indicates that the search service is currently unavailable."
+                    ),
+                    brief="Failed to parse search results",
+                )
+
+        for i, result in enumerate(results):
+            if i > 0:
+                builder.write("---\n\n")
+            builder.write(
+                f"Title: {result.title}\nDate: {result.date}\n"
+                f"URL: {result.url}\nSummary: {result.snippet}\n\n"
+            )
+            if result.content:
+                builder.write(f"{result.content}\n\n")
+
+        return builder.ok()
+
+
+class SearchResult(BaseModel):
+    site_name: str
+    title: str
+    url: str
+    snippet: str
+    content: str = ""
+    date: str = ""
+    icon: str = ""
+    mime: str = ""
+
+
+class Response(BaseModel):
+    search_results: list[SearchResult]

kimi_cli/ui/__init__.py

@@ -0,0 +1,68 @@
+import asyncio
+import contextlib
+from collections.abc import Callable, Coroutine
+from typing import Any
+
+from kimi_cli.soul import Soul
+from kimi_cli.soul.wire import Wire
+from kimi_cli.utils.logging import logger
+
+type UILoopFn = Callable[[Wire], Coroutine[Any, Any, None]]
+"""A long-running async function to visualize the agent behavior."""
+
+
+class RunCancelled(Exception):
+    """The run was cancelled by the cancel event."""
+
+
+async def run_soul(
+    soul: Soul,
+    user_input: str,
+    ui_loop_fn: UILoopFn,
+    cancel_event: asyncio.Event,
+):
+    """
+    Run the soul with the given user input.
+
+    `cancel_event` is a outside handle that can be used to cancel the run. When the event is set,
+    the run will be gracefully stopped and a `RunCancelled` will be raised.
+
+    Raises:
+        ChatProviderError: When the LLM provider returns an error.
+        MaxStepsReached: When the maximum number of steps is reached.
+        RunCancelled: When the run is cancelled by the cancel event.
+    """
+    wire = Wire()
+    logger.debug("Starting UI loop with function: {ui_loop_fn}", ui_loop_fn=ui_loop_fn)
+
+    ui_task = asyncio.create_task(ui_loop_fn(wire))
+    soul_task = asyncio.create_task(soul.run(user_input, wire))
+
+    cancel_event_task = asyncio.create_task(cancel_event.wait())
+    await asyncio.wait(
+        [soul_task, cancel_event_task],
+        return_when=asyncio.FIRST_COMPLETED,
+    )
+
+    try:
+        if cancel_event.is_set():
+            logger.debug("Cancelling the run task")
+            soul_task.cancel()
+            try:
+                await soul_task
+            except asyncio.CancelledError:
+                raise RunCancelled from None
+        else:
+            assert soul_task.done()  # either stop event is set or the run task is done
+            cancel_event_task.cancel()
+            with contextlib.suppress(asyncio.CancelledError):
+                await cancel_event_task
+            soul_task.result()  # this will raise if any exception was raised in the run task
+    finally:
+        logger.debug("Shutting down the visualization loop")
+        # shutting down the event queue should break the visualization loop
+        wire.shutdown()
+        try:
+            await asyncio.wait_for(ui_task, timeout=0.5)
+        except TimeoutError:
+            logger.warning("Visualization loop timed out")