kimi-cli 0.35__py3-none-any.whl

kimi_cli/tools/file/replace.py

@@ -0,0 +1,132 @@
+from pathlib import Path
+from typing import override
+
+import aiofiles
+from kosong.tooling import CallableTool2, ToolError, ToolOk, ToolReturnType
+from pydantic import BaseModel, Field
+
+from kimi_cli.agent import BuiltinSystemPromptArgs
+
+
+class Edit(BaseModel):
+    old: str = Field(description="The old string to replace. Can be multi-line.")
+    new: str = Field(description="The new string to replace with. Can be multi-line.")
+    replace_all: bool = Field(description="Whether to replace all occurrences.", default=False)
+
+
+class Params(BaseModel):
+    path: str = Field(description="The absolute path to the file to edit.")
+    edit: Edit | list[Edit] = Field(
+        description=(
+            "The edit(s) to apply to the file. "
+            "You can provide a single edit or a list of edits here."
+        )
+    )
+
+
+class StrReplaceFile(CallableTool2[Params]):
+    name: str = "StrReplaceFile"
+    description: str = (Path(__file__).parent / "replace.md").read_text()
+    params: type[Params] = Params
+
+    def __init__(self, builtin_args: BuiltinSystemPromptArgs, **kwargs):
+        super().__init__(**kwargs)
+        self._work_dir = builtin_args.KIMI_WORK_DIR
+
+    def _validate_path(self, path: Path) -> ToolError | None:
+        """Validate that the path is safe to edit."""
+        # Check for path traversal attempts
+        resolved_path = path.resolve()
+        resolved_work_dir = self._work_dir.resolve()
+
+        # Ensure the path is within work directory
+        if not str(resolved_path).startswith(str(resolved_work_dir)):
+            return ToolError(
+                message=(
+                    f"`{path}` is outside the working directory. "
+                    "You can only edit files within the working directory."
+                ),
+                brief="Path outside working directory",
+            )
+        return None
+
+    def _apply_edit(self, content: str, edit: Edit) -> str:
+        """Apply a single edit to the content."""
+        if edit.replace_all:
+            return content.replace(edit.old, edit.new)
+        else:
+            return content.replace(edit.old, edit.new, 1)
+
+    @override
+    async def __call__(self, params: Params) -> ToolReturnType:
+        try:
+            p = Path(params.path)
+
+            if not p.is_absolute():
+                return ToolError(
+                    message=(
+                        f"`{params.path}` is not an absolute path. "
+                        "You must provide an absolute path to edit a file."
+                    ),
+                    brief="Invalid path",
+                )
+
+            # Validate path safety
+            path_error = self._validate_path(p)
+            if path_error:
+                return path_error
+
+            if not p.exists():
+                return ToolError(
+                    message=f"`{params.path}` does not exist.",
+                    brief="File not found",
+                )
+            if not p.is_file():
+                return ToolError(
+                    message=f"`{params.path}` is not a file.",
+                    brief="Invalid path",
+                )
+
+            # Read the file content
+            async with aiofiles.open(p, encoding="utf-8", errors="replace") as f:
+                content = await f.read()
+
+            original_content = content
+            edits = [params.edit] if isinstance(params.edit, Edit) else params.edit
+
+            # Apply all edits
+            for edit in edits:
+                content = self._apply_edit(content, edit)
+
+            # Check if any changes were made
+            if content == original_content:
+                return ToolError(
+                    message="No replacements were made. The old string was not found in the file.",
+                    brief="No replacements made",
+                )
+
+            # Write the modified content back to the file
+            async with aiofiles.open(p, mode="w", encoding="utf-8") as f:
+                await f.write(content)
+
+            # Count changes for success message
+            total_replacements = 0
+            for edit in edits:
+                if edit.replace_all:
+                    total_replacements += original_content.count(edit.old)
+                else:
+                    total_replacements += 1 if edit.old in original_content else 0
+
+            return ToolOk(
+                output="",
+                message=(
+                    f"File successfully edited. "
+                    f"Applied {len(edits)} edit(s) with {total_replacements} total replacement(s)."
+                ),
+            )
+
+        except Exception as e:
+            return ToolError(
+                message=f"Failed to edit. Error: {e}",
+                brief="Failed to edit file",
+            )

kimi_cli/tools/file/write.md

@@ -0,0 +1,5 @@
+Write content to a file.
+
+**Tips:**
+- When `mode` is not specified, it defaults to `overwrite`. Always write with caution.
+- When the content to write is too long (e.g. > 100 lines), use this tool multiple times instead of a single call. Use `overwrite` mode at the first time, then use `append` mode after the first write.

kimi_cli/tools/file/write.py

@@ -0,0 +1,107 @@
+from pathlib import Path
+from typing import Literal, override
+
+import aiofiles
+from kosong.tooling import CallableTool2, ToolError, ToolOk, ToolReturnType
+from pydantic import BaseModel, Field
+
+from kimi_cli.agent import BuiltinSystemPromptArgs
+
+
+class Params(BaseModel):
+    path: str = Field(description="The absolute path to the file to write")
+    content: str = Field(description="The content to write to the file")
+    mode: Literal["overwrite", "append"] = Field(
+        description=(
+            "The mode to use to write to the file. "
+            "Two modes are supported: `overwrite` for overwriting the whole file and "
+            "`append` for appending to the end of an existing file."
+        ),
+        default="overwrite",
+    )
+
+
+class WriteFile(CallableTool2[Params]):
+    name: str = "WriteFile"
+    description: str = (Path(__file__).parent / "write.md").read_text()
+    params: type[Params] = Params
+
+    def __init__(self, builtin_args: BuiltinSystemPromptArgs, **kwargs):
+        super().__init__(**kwargs)
+        self._work_dir = builtin_args.KIMI_WORK_DIR
+
+    def _validate_path(self, path: Path) -> ToolError | None:
+        """Validate that the path is safe to write."""
+        # Check for path traversal attempts
+        resolved_path = path.resolve()
+        resolved_work_dir = self._work_dir.resolve()
+
+        # Ensure the path is within work directory
+        if not str(resolved_path).startswith(str(resolved_work_dir)):
+            return ToolError(
+                message=(
+                    f"`{path}` is outside the working directory. "
+                    "You can only write files within the working directory."
+                ),
+                brief="Path outside working directory",
+            )
+        return None
+
+    @override
+    async def __call__(self, params: Params) -> ToolReturnType:
+        # TODO: checks:
+        # - check if the path may contain secrets
+        # - check if the file format is writable
+        try:
+            p = Path(params.path)
+
+            if not p.is_absolute():
+                return ToolError(
+                    message=(
+                        f"`{params.path}` is not an absolute path. "
+                        "You must provide an absolute path to write a file."
+                    ),
+                    brief="Invalid path",
+                )
+
+            # Validate path safety
+            path_error = self._validate_path(p)
+            if path_error:
+                return path_error
+
+            if not p.parent.exists():
+                return ToolError(
+                    message=f"`{params.path}` parent directory does not exist.",
+                    brief="Parent directory not found",
+                )
+
+            # Validate mode parameter
+            if params.mode not in ["overwrite", "append"]:
+                return ToolError(
+                    message=(
+                        f"Invalid write mode: `{params.mode}`. "
+                        "Mode must be either `overwrite` or `append`."
+                    ),
+                    brief="Invalid write mode",
+                )
+
+            # Determine file mode for aiofiles
+            file_mode = "w" if params.mode == "overwrite" else "a"
+
+            # Write content to file
+            async with aiofiles.open(p, mode=file_mode, encoding="utf-8") as f:
+                await f.write(params.content)
+
+            # Get file info for success message
+            file_size = p.stat().st_size
+            action = "overwritten" if params.mode == "overwrite" else "appended to"
+            return ToolOk(
+                output="",
+                message=(f"File successfully {action}. Current size: {file_size} bytes."),
+            )
+
+        except Exception as e:
+            return ToolError(
+                message=f"Failed to write to {params.path}. Error: {e}",
+                brief="Failed to write file",
+            )

kimi_cli/tools/mcp.py

@@ -0,0 +1,85 @@
+import fastmcp
+import mcp
+from fastmcp.client.client import CallToolResult
+from kosong.base.message import AudioURLPart, ContentPart, ImageURLPart, TextPart
+from kosong.tooling import CallableTool, ToolOk, ToolReturnType
+
+
+class MCPTool(CallableTool):
+    def __init__(self, mcp_tool: mcp.Tool, client: fastmcp.Client, **kwargs):
+        super().__init__(
+            name=mcp_tool.name,
+            description=mcp_tool.description or "",
+            parameters=mcp_tool.inputSchema,
+            **kwargs,
+        )
+        self._mcp_tool = mcp_tool
+        self._client = client
+
+    async def __call__(self, *args, **kwargs) -> ToolReturnType:
+        async with self._client as client:
+            result = await client.call_tool(self._mcp_tool.name, kwargs, timeout=20)
+            return convert_tool_result(result)
+
+
+def convert_tool_result(result: CallToolResult) -> ToolReturnType:
+    content: list[ContentPart] = []
+    for part in result.content:
+        match part:
+            case mcp.types.TextContent(text=text):
+                content.append(TextPart(text=text))
+            case mcp.types.ImageContent(data=data, mimeType=mimeType):
+                content.append(
+                    ImageURLPart(
+                        image_url=ImageURLPart.ImageURL(url=f"data:{mimeType};base64,{data}")
+                    )
+                )
+            case mcp.types.AudioContent(data=data, mimeType=mimeType):
+                content.append(
+                    AudioURLPart(
+                        audio_url=AudioURLPart.AudioURL(url=f"data:{mimeType};base64,{data}")
+                    )
+                )
+            case mcp.types.EmbeddedResource(
+                resource=mcp.types.BlobResourceContents(uri=_uri, mimeType=mimeType, blob=blob)
+            ):
+                mimeType = mimeType or "application/octet-stream"
+                if mimeType.startswith("image/"):
+                    content.append(
+                        ImageURLPart(
+                            type="image_url",
+                            image_url=ImageURLPart.ImageURL(
+                                url=f"data:{mimeType};base64,{blob}",
+                            ),
+                        )
+                    )
+                elif mimeType.startswith("audio/"):
+                    content.append(
+                        AudioURLPart(
+                            type="audio_url",
+                            audio_url=AudioURLPart.AudioURL(url=f"data:{mimeType};base64,{blob}"),
+                        )
+                    )
+                else:
+                    raise ValueError(f"Unsupported mime type: {mimeType}")
+            case mcp.types.ResourceLink(uri=uri, mimeType=mimeType, description=_description):
+                mimeType = mimeType or "application/octet-stream"
+                if mimeType.startswith("image/"):
+                    content.append(
+                        ImageURLPart(
+                            type="image_url",
+                            image_url=ImageURLPart.ImageURL(url=str(uri)),
+                        )
+                    )
+                elif mimeType.startswith("audio/"):
+                    content.append(
+                        AudioURLPart(
+                            type="audio_url",
+                            audio_url=AudioURLPart.AudioURL(url=str(uri)),
+                        )
+                    )
+                else:
+                    raise ValueError(f"Unsupported mime type: {mimeType}")
+            case _:
+                raise ValueError(f"Unsupported MCP tool result part: {part}")
+    return ToolOk(output=content)

kimi_cli/tools/task/__init__.py

@@ -0,0 +1,156 @@
+from pathlib import Path
+from typing import override
+
+from kosong.tooling import CallableTool2, ToolError, ToolOk, ToolReturnType
+from pydantic import BaseModel, Field
+
+from kimi_cli.agent import Agent, AgentGlobals, AgentSpec, load_agent
+from kimi_cli.soul import MaxStepsReached
+from kimi_cli.soul.context import Context
+from kimi_cli.soul.kimisoul import KimiSoul
+from kimi_cli.soul.wire import ApprovalRequest, Wire, WireMessage, get_wire_or_none
+from kimi_cli.tools.utils import load_desc
+from kimi_cli.utils.message import message_extract_text
+from kimi_cli.utils.path import next_available_rotation
+
+# Maximum continuation attempts for task summary
+MAX_CONTINUE_ATTEMPTS = 1
+
+
+CONTINUE_PROMPT = """
+Your previous response was too brief. Please provide a more comprehensive summary that includes:
+
+1. Specific technical details and implementations
+2. Complete code examples if relevant
+3. Detailed findings and analysis
+4. All important information that should be aware of by the caller
+""".strip()
+
+
+class Params(BaseModel):
+    description: str = Field(description="A short (3-5 word) description of the task")
+    subagent_name: str = Field(
+        description="The name of the specialized subagent to use for this task"
+    )
+    prompt: str = Field(
+        description=(
+            "The task for the subagent to perform. "
+            "You must provide a detailed prompt with all necessary background information "
+            "because the subagent cannot see anything in your context."
+        )
+    )
+
+
+class Task(CallableTool2[Params]):
+    name: str = "Task"
+    params: type[Params] = Params
+
+    def __init__(self, agent_spec: AgentSpec, agent_globals: AgentGlobals, **kwargs):
+        subagents: dict[str, Agent] = {}
+        descs = []
+
+        # load all subagents
+        assert agent_spec.subagents is not None, "Task tool expects subagents"
+        for name, spec in agent_spec.subagents.items():
+            subagents[name] = load_agent(spec.path, agent_globals)
+            descs.append(f"- `{name}`: {spec.description}")
+
+        super().__init__(
+            description=load_desc(
+                Path(__file__).parent / "task.md",
+                {
+                    "SUBAGENTS_MD": "\n".join(descs),
+                },
+            ),
+            **kwargs,
+        )
+
+        self._agent_globals = agent_globals
+        self._session = agent_globals.session
+        self._subagents = subagents
+
+    async def _get_subagent_history_file(self) -> Path:
+        """Generate a unique history file path for subagent."""
+        main_history_file = self._session.history_file
+        subagent_base_name = f"{main_history_file.stem}_sub"
+        main_history_file.parent.mkdir(parents=True, exist_ok=True)  # just in case
+        sub_history_file = await next_available_rotation(
+            main_history_file.parent / f"{subagent_base_name}{main_history_file.suffix}"
+        )
+        assert sub_history_file is not None
+        return sub_history_file
+
+    @override
+    async def __call__(self, params: Params) -> ToolReturnType:
+        if params.subagent_name not in self._subagents:
+            return ToolError(
+                message=f"Subagent not found: {params.subagent_name}",
+                brief="Subagent not found",
+            )
+        agent = self._subagents[params.subagent_name]
+        try:
+            result = await self._run_subagent(agent, params.prompt)
+            return result
+        except Exception as e:
+            return ToolError(
+                message=f"Failed to run subagent: {e}",
+                brief="Failed to run subagent",
+            )
+
+    async def _run_subagent(self, agent: Agent, prompt: str) -> ToolReturnType:
+        """Run subagent with optional continuation for task summary."""
+        subagent_history_file = await self._get_subagent_history_file()
+        context = Context(file_backend=subagent_history_file)
+        soul = KimiSoul(
+            agent,
+            agent_globals=self._agent_globals,
+            context=context,
+            loop_control=self._agent_globals.config.loop_control,
+        )
+        wire = get_wire_or_none()
+        assert wire is not None, "Wire is expected to be set"
+        sub_wire = _SubWire(wire)
+
+        try:
+            await soul.run(prompt, sub_wire)
+        except MaxStepsReached as e:
+            return ToolError(
+                message=(
+                    f"Max steps {e.n_steps} reached when running subagent. "
+                    "Please try splitting the task into smaller subtasks."
+                ),
+                brief="Max steps reached",
+            )
+
+        _error_msg = (
+            "The subagent seemed not to run properly. Maybe you have to do the task yourself."
+        )
+
+        # Check if the subagent context is valid
+        if len(context.history) == 0 or context.history[-1].role != "assistant":
+            return ToolError(message=_error_msg, brief="Failed to run subagent")
+
+        final_response = message_extract_text(context.history[-1])
+
+        # Check if response is too brief, if so, run again with continuation prompt
+        n_attempts_remaining = MAX_CONTINUE_ATTEMPTS
+        if len(final_response) < 200 and n_attempts_remaining > 0:
+            await soul.run(CONTINUE_PROMPT, sub_wire)
+
+            if len(context.history) == 0 or context.history[-1].role != "assistant":
+                return ToolError(message=_error_msg, brief="Failed to run subagent")
+            final_response = message_extract_text(context.history[-1])
+
+        return ToolOk(output=final_response)
+
+
+class _SubWire(Wire):
+    def __init__(self, super_wire: Wire):
+        super().__init__()
+        self._super_wire = super_wire
+
+    @override
+    def send(self, msg: WireMessage):
+        if isinstance(msg, ApprovalRequest):
+            self._super_wire.send(msg)
+        # TODO: visualize subagent behavior by sending other messages in some way

kimi_cli/tools/task/task.md

@@ -0,0 +1,26 @@
+Spawn a subagent to perform a specific task. Subagent will be spawned with a fresh context without any history of yours.
+
+**Context Isolation**
+
+Context isolation is one of the key benefits of using subagents. By delegating tasks to subagents, you can keep your main context clean and focused on the main goal requested by the user.
+
+Here are some scenerios you may want this tool for context isolation:
+
+- You wrote some code and it did not work as expected. In this case you can spawn a subagent to fix the code, asking the subagent to return how it is fixed. This can potentially benefit because the detailed process of fixing the code may not be relevant to your main goal, and may clutter your context.
+- When you need some latest knowledge of a specific library, framework or technology to proceed with your task, you can spawn a subagent to search on the internet for the needed information and return to you the gathered relevant information, for example code examples, API references, etc. This can avoid ton of irrelevant search results in your own context.
+
+DO NOT directly forward the user prompt to Task tool. DO NOT simply spawn Task tool for each todo item. This will cause the user confused because the user cannot see what the subagent do. Only you can see the response from the subagent. So, only spawn subagents for very specific and narrow tasks like fixing a compilation error, or searching for a specific solution.
+
+**Parallel Multi-Tasking**
+
+Parallel multi-tasking is another key benefit of this tool. When the user request involves multiple subtasks that are independent of each other, you can use Task tool multiple times in a single response to let subagents work in parallel for you.
+
+Examples:
+
+- User requests to code, refactor or fix multiple modules/files in a project, and they can be tested independently. In this case you can spawn multiple subagents each working on a different module/file.
+- When you need to analyze a huge codebase (> hundreds of thousands of lines), you can spawn multiple subagents each exploring on a different part of the codebase and gather the summarized results.
+- When you need to search the web for multiple queries, you can spawn multiple subagents for better efficiency.
+
+**Available Subagents:**
+
+${SUBAGENTS_MD}

kimi_cli/tools/test.py

@@ -0,0 +1,55 @@
+import asyncio
+from typing import override
+
+from kosong.tooling import CallableTool2, ToolOk, ToolReturnType
+from pydantic import BaseModel
+
+
+class PlusParams(BaseModel):
+    a: float
+    b: float
+
+
+class Plus(CallableTool2[PlusParams]):
+    name: str = "plus"
+    description: str = "Add two numbers"
+    params: type[PlusParams] = PlusParams
+
+    @override
+    async def __call__(self, params: PlusParams) -> ToolReturnType:
+        return ToolOk(output=str(params.a + params.b))
+
+
+class CompareParams(BaseModel):
+    a: float
+    b: float
+
+
+class Compare(CallableTool2[CompareParams]):
+    name: str = "compare"
+    description: str = "Compare two numbers"
+    params: type[CompareParams] = CompareParams
+
+    @override
+    async def __call__(self, params: CompareParams) -> ToolReturnType:
+        if params.a > params.b:
+            return ToolOk(output="greater")
+        elif params.a < params.b:
+            return ToolOk(output="less")
+        else:
+            return ToolOk(output="equal")
+
+
+class PanicParams(BaseModel):
+    message: str
+
+
+class Panic(CallableTool2[PanicParams]):
+    name: str = "panic"
+    description: str = "Raise an exception to cause the tool call to fail."
+    params: type[PanicParams] = PanicParams
+
+    @override
+    async def __call__(self, params: PanicParams) -> ToolReturnType:
+        await asyncio.sleep(2)
+        raise Exception(f"panicked with a message with {len(params.message)} characters")

kimi_cli/tools/think/__init__.py

@@ -0,0 +1,21 @@
+from pathlib import Path
+from typing import override
+
+from kosong.tooling import CallableTool2, ToolOk, ToolReturnType
+from pydantic import BaseModel, Field
+
+from kimi_cli.tools.utils import load_desc
+
+
+class Params(BaseModel):
+    thought: str = Field(description=("A thought to think about."))
+
+
+class Think(CallableTool2[Params]):
+    name: str = "Think"
+    description: str = load_desc(Path(__file__).parent / "think.md", {})
+    params: type[Params] = Params
+
+    @override
+    async def __call__(self, params: Params) -> ToolReturnType:
+        return ToolOk(output="", message="Thought logged")

kimi_cli/tools/think/think.md

@@ -0,0 +1 @@
+Use the tool to think about something. It will not obtain new information or change the database, but just append the thought to the log. Use it when complex reasoning or some cache memory is needed.

kimi_cli/tools/todo/__init__.py

@@ -0,0 +1,27 @@
+from pathlib import Path
+from typing import Literal, override
+
+from kosong.tooling import CallableTool2, ToolOk, ToolReturnType
+from pydantic import BaseModel, Field
+
+
+class Todo(BaseModel):
+    title: str = Field(description="The title of the todo", min_length=1)
+    status: Literal["Pending", "In Progress", "Done"] = Field(description="The status of the todo")
+
+
+class Params(BaseModel):
+    todos: list[Todo] = Field(description="The updated todo list")
+
+
+class SetTodoList(CallableTool2[Params]):
+    name: str = "SetTodoList"
+    description: str = (Path(__file__).parent / "set_todo_list.md").read_text()
+    params: type[Params] = Params
+
+    @override
+    async def __call__(self, params: Params) -> ToolReturnType:
+        rendered = ""
+        for todo in params.todos:
+            rendered += f"- {todo.title} [{todo.status}]\n"
+        return ToolOk(output=rendered)

kimi_cli/tools/todo/set_todo_list.md

@@ -0,0 +1,15 @@
+Update the whole todo list.
+
+Todo list is a simple yet powerful tool to help you get things done. You typically want to use this tool when the given task involves multiple subtasks/milestones, or, multiple tasks are given in a single request. This tool can help you to break down the task and track the progress.
+
+This is the only todo list tool available to you. That said, each time you want to operate on the todo list, you need to update the whole. Make sure to maintain the todo items and their statuses properly.
+
+Once you finished a subtask/milestone, remember to update the todo list to reflect the progress. Also, you can give yourself a self-encouragement to keep you motivated.
+
+Abusing this tool to track too small steps will just waste your time and make your context messy. For example, here are some cases you should not use this tool:
+
+- When the user just simply ask you a question. E.g. "What language and framework is used in the project?", "What is the best practice for x?"
+- When it only takes a few steps/tool calls to complete the task. E.g. "Fix the unit test function 'test_xxx'", "Refactor the function 'xxx' to make it more solid."
+- When the user prompt is very specific and the only thing you need to do is brainlessly following the instructions. E.g. "Replace xxx to yyy in the file zzz", "Create a file xxx with content yyy."
+
+However, do not get stuck in a rut. Be flexible. Sometimes, you may try to use todo list at first, then realize the task is too simple and you can simply stop using it; or, sometimes, you may realize the task is complex after a few steps and then you can start using todo list to break it down.