diff --git a/examples/python/customized-tools/README.md b/examples/python/customized-tools/README.md new file mode 100644 index 0000000..6b603fe --- /dev/null +++ b/examples/python/customized-tools/README.md @@ -0,0 +1,21 @@ +# Example: Customized Tools + +This example shows how to define a custom tool and load it through an agent file +when using the Kimi Agent SDK. + +## Run + +```sh +cd examples/python/customized-tools +uv sync --reinstall + +# configure your API key +export KIMI_API_KEY=your-api-key +export KIMI_BASE_URL=https://api.moonshot.ai/v1 +export KIMI_MODEL_NAME=kimi-k2-thinking-turbo + +uv run main.py +``` + +The agent file `myagent.yaml` registers the custom tool `my_tools.ls:Ls` and +reuses the default tool set. diff --git a/examples/python/customized-tools/main.py b/examples/python/customized-tools/main.py new file mode 100644 index 0000000..2c91010 --- /dev/null +++ b/examples/python/customized-tools/main.py @@ -0,0 +1,19 @@ +import asyncio +from pathlib import Path + +from kimi_agent_sdk import prompt + + +async def main() -> None: + agent_file = Path(__file__).parent / "myagent.yaml" + async for msg in prompt( + "What tools do you have?", + agent_file=agent_file, + yolo=True, + ): + print(msg.extract_text(), end="", flush=True) + print() + + +if __name__ == "__main__": + asyncio.run(main()) diff --git a/examples/python/customized-tools/my_tools/__init__.py b/examples/python/customized-tools/my_tools/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/examples/python/customized-tools/my_tools/ls.py b/examples/python/customized-tools/my_tools/ls.py new file mode 100644 index 0000000..ccec1ce --- /dev/null +++ b/examples/python/customized-tools/my_tools/ls.py @@ -0,0 +1,29 @@ +from kimi_agent_sdk import CallableTool2, ToolError, ToolOk, ToolReturnValue +from pydantic import BaseModel, Field + + +class Params(BaseModel): + directory: str = Field( + default=".", + description="The directory to list files from.", + ) + + +class Ls(CallableTool2): + name: str = "Ls" + description: str = "List files in a directory." + params: type[Params] = Params + + async def __call__(self, params: Params) -> ToolReturnValue: + import os + + try: + files = os.listdir(params.directory) + output = "\n".join(files) + return ToolOk(output=output) + except Exception as exc: + return ToolError( + output="", + message=str(exc), + brief="Failed to list files", + ) diff --git a/examples/python/customized-tools/myagent.yaml b/examples/python/customized-tools/myagent.yaml new file mode 100644 index 0000000..39e6160 --- /dev/null +++ b/examples/python/customized-tools/myagent.yaml @@ -0,0 +1,15 @@ +version: 1 +agent: + extend: default + tools: + - "kimi_cli.tools.multiagent:Task" + - "kimi_cli.tools.todo:SetTodoList" + - "kimi_cli.tools.shell:Shell" + - "kimi_cli.tools.file:ReadFile" + - "kimi_cli.tools.file:Glob" + - "kimi_cli.tools.file:Grep" + - "kimi_cli.tools.file:WriteFile" + - "kimi_cli.tools.file:StrReplaceFile" + - "kimi_cli.tools.web:SearchWeb" + - "kimi_cli.tools.web:FetchURL" + - "my_tools.ls:Ls" # custom tool diff --git a/examples/python/customized-tools/pyproject.toml b/examples/python/customized-tools/pyproject.toml new file mode 100644 index 0000000..d111d08 --- /dev/null +++ b/examples/python/customized-tools/pyproject.toml @@ -0,0 +1,10 @@ +[project] +name = "customized-tools" +version = "0.1.0" +description = "Custom tool example for Kimi Agent SDK" +readme = "README.md" +requires-python = ">=3.12" +dependencies = ["kimi-agent-sdk"] + +[tool.uv.sources] +kimi-agent-sdk = { path = "../../../python" } diff --git a/guides/python/customized-tools.md b/guides/python/customized-tools.md new file mode 100644 index 0000000..b910bbb --- /dev/null +++ b/guides/python/customized-tools.md @@ -0,0 +1,112 @@ +# Customized Tools + +Kimi Agent SDK is a thin wrapper around Kimi Code (Kimi CLI), so custom tools +are defined exactly the same way: write a Python tool class, register it in an +agent file, and pass that agent file to `prompt()` or `Session.create(...)`. + +If you already have a Kimi CLI agent file and tools, you can reuse them as-is. + +## Step 1: Implement a tool + +Create a tool class with a Pydantic parameter model and return `ToolOk` or +`ToolError`: + +```python +from kimi_agent_sdk import CallableTool2, ToolError, ToolOk, ToolReturnValue +from pydantic import BaseModel, Field + + +class Params(BaseModel): + directory: str = Field( + default=".", + description="The directory to list files from.", + ) + + +class Ls(CallableTool2): + name: str = "Ls" + description: str = "List files in a directory." + params: type[Params] = Params + + async def __call__(self, params: Params) -> ToolReturnValue: + import os + + try: + files = os.listdir(params.directory) + return ToolOk(output="\n".join(files)) + except Exception as exc: + return ToolError( + output="", + message=str(exc), + brief="Failed to list files", + ) +``` + +## Step 2: Make the tool importable + +Ensure your module is importable by the Python process running the SDK: + +``` +my_tools/ + __init__.py + ls.py +``` + +Options: + +- Install your project package into the current environment. +- Or add the project root to `PYTHONPATH` when running your script. + +## Step 3: Register the tool in an agent file + +Add your tool path (`module:ClassName`) to `tools`. Note that `tools` replaces +the inherited list, so include every tool you want to keep. + +```yaml +version: 1 +agent: + extend: default + tools: + - "kimi_cli.tools.multiagent:Task" + - "kimi_cli.tools.todo:SetTodoList" + - "kimi_cli.tools.shell:Shell" + - "kimi_cli.tools.file:ReadFile" + - "kimi_cli.tools.file:Glob" + - "kimi_cli.tools.file:Grep" + - "kimi_cli.tools.file:WriteFile" + - "kimi_cli.tools.file:StrReplaceFile" + - "kimi_cli.tools.web:SearchWeb" + - "kimi_cli.tools.web:FetchURL" + - "my_tools.ls:Ls" # custom tool +``` + +For full agent file format, see the +[Kimi Code agent docs](https://moonshotai.github.io/kimi-cli/en/customization/agents.html#custom-agent-files). + +## Step 4: Use the agent file in Python + +Pass the agent file path to `prompt()` or `Session.create(...)`: + +```python +import asyncio +from pathlib import Path + +from kimi_agent_sdk import prompt + + +async def main() -> None: + async for msg in prompt( + "What tools do you have?", + agent_file=Path("myagent.yaml"), + yolo=True, + ): + print(msg.extract_text(), end="", flush=True) + print() + + +asyncio.run(main()) +``` + +If you prefer the low-level API, use `Session.create(agent_file=...)` instead. + +For full code examples, see [here](../../examples/python/customized-tools). diff --git a/guides/python/prompt.md b/guides/python/prompt.md index f635daf..c574a20 100644 --- a/guides/python/prompt.md +++ b/guides/python/prompt.md @@ -60,7 +60,8 @@ async def prompt( - `approval_handler_fn`: Callback that receives an [`ApprovalRequest`](https://moonshotai.github.io/kimi-cli/en/customization/wire-mode.html#approvalrequest) and must call `request.resolve(...)` with `"approve"`, `"approve_for_session"`, or `"reject"`. Required when `yolo=False`. -- `agent_file`: Path to a CLI agent spec file (tools, prompts, subagents). +- `agent_file`: Path to a CLI agent spec file (tools, prompts, subagents), for more details see + [Custom Agent Files](https://moonshotai.github.io/kimi-cli/en/customization/agents.html#custom-agent-files). - `mcp_configs`: MCP configuration objects or raw dictionaries. Matches the Kimi Code [MCP schema](https://moonshotai.github.io/kimi-cli/en/customization/mcp.html) (for example, `mcp.json`). - `skills_dir`: Directory containing agent skills to load. It should be a valid KaosPath object. diff --git a/guides/python/quickstart.md b/guides/python/quickstart.md index 527ea44..a69e34e 100644 --- a/guides/python/quickstart.md +++ b/guides/python/quickstart.md @@ -160,3 +160,9 @@ access: - You can run multiple prompts within the same session. - You can resume previous sessions with `Session.resume(...)`. - It exposes raw [Wire](https://moonshotai.github.io/kimi-cli/en/customization/wire-mode.html#wire-mode) messages and approvals for fine-grained control. + +## What's Next + +- [Prompt API](./prompt.md) - Deep dive into the high-level `prompt()` helper +- [Session API](./session.md) - Manage sessions, approvals, and raw Wire messages +- [Customized Tools](./customized-tools.md) - Add your own tools via agent files diff --git a/guides/python/session.md b/guides/python/session.md index c83d260..9e851c5 100644 --- a/guides/python/session.md +++ b/guides/python/session.md @@ -115,7 +115,8 @@ Parameters: - `thinking`: Enable thinking mode for supported models. - `yolo`: Auto-approve all approval requests. When enabled, you may not receive `ApprovalRequest` messages. -- `agent_file`: Agent specification file path. +- `agent_file`: Agent specification file path, for more details see + [Custom Agent Files](https://moonshotai.github.io/kimi-cli/en/customization/agents.html#custom-agent-files). - `mcp_configs`: MCP server configs (same schema as Kimi Code). - `skills_dir`: Skills directory (KaosPath). - `max_steps_per_turn`: Maximum number of steps allowed per turn. diff --git a/python/src/kimi_agent_sdk/__init__.py b/python/src/kimi_agent_sdk/__init__.py index 8528b1a..0fd3ec9 100644 --- a/python/src/kimi_agent_sdk/__init__.py +++ b/python/src/kimi_agent_sdk/__init__.py @@ -30,7 +30,6 @@ TokenUsage, ToolCallPart, ToolResult, - ToolReturnValue, TurnBegin, WireMessage, is_event, @@ -53,6 +52,7 @@ ToolCall, VideoURLPart, ) +from kosong.tooling import CallableTool2, ToolError, ToolOk, ToolReturnValue from kimi_agent_sdk._approval import ApprovalHandlerFn from kimi_agent_sdk._exception import PromptValidationError, SessionStateError @@ -98,6 +98,9 @@ "TokenUsage", "is_event", "is_request", + "CallableTool2", + "ToolOk", + "ToolError", # Exceptions "KimiAgentException", "ConfigError", diff --git a/python/tests/test_customized_tools.py b/python/tests/test_customized_tools.py new file mode 100644 index 0000000..453366d --- /dev/null +++ b/python/tests/test_customized_tools.py @@ -0,0 +1,118 @@ +from __future__ import annotations + +import asyncio +import sys +from pathlib import Path +from typing import Any, cast + +import pytest +from kaos.path import KaosPath +from kosong.tooling import ( + CallableTool2 as KosongCallableTool2, +) +from kosong.tooling import ( + ToolError as KosongToolError, +) +from kosong.tooling import ( + ToolOk as KosongToolOk, +) +from kosong.tooling import ( + ToolReturnValue as KosongToolReturnValue, +) +from pydantic import BaseModel, Field + +from kimi_agent_sdk import ( + CallableTool2, + Config, + Session, + ToolError, + ToolOk, + ToolReturnValue, +) + + +def test_tooling_exports_match_kosong() -> None: + assert CallableTool2 is KosongCallableTool2 + assert ToolOk is KosongToolOk + assert ToolError is KosongToolError + assert ToolReturnValue is KosongToolReturnValue + + +def test_custom_tool_uses_sdk_exports() -> None: + class Params(BaseModel): + directory: str = Field(default=".") + + class Ls(CallableTool2[Params]): + name: str = "Ls" + description: str = "List files in a directory." + params: type[Params] = Params + + async def __call__(self, params: Params) -> ToolReturnValue: + return ToolOk(output="ok") + + result = asyncio.run(Ls()(Params())) + assert isinstance(result, ToolReturnValue) + assert result.is_error is False + + +@pytest.mark.asyncio +async def test_agent_loads_custom_tool(tmp_path: Path) -> None: + tools_dir = tmp_path / "my_tools" + tools_dir.mkdir() + (tools_dir / "__init__.py").write_text("", encoding="utf-8") + (tools_dir / "ls.py").write_text( + "\n".join( + [ + "from pydantic import BaseModel, Field", + "from kimi_agent_sdk import CallableTool2, ToolOk, ToolReturnValue", + "", + "", + "class Params(BaseModel):", + " directory: str = Field(default='.')", + "", + "", + "class Ls(CallableTool2):", + " name: str = 'Ls'", + " description: str = 'List files in a directory.'", + " params: type[Params] = Params", + "", + " async def __call__(self, params: Params) -> ToolReturnValue:", + " return ToolOk(output='ok')", + "", + ] + ), + encoding="utf-8", + ) + (tmp_path / "system.md").write_text("You are a test agent.", encoding="utf-8") + agent_file = tmp_path / "agent.yaml" + agent_file.write_text( + "\n".join( + [ + "version: 1", + "agent:", + " name: test-agent", + " system_prompt_path: ./system.md", + " tools:", + ' - "my_tools.ls:Ls"', + "", + ] + ), + encoding="utf-8", + ) + + sys.path.insert(0, str(tmp_path)) + try: + async with await Session.create( + work_dir=KaosPath(str(tmp_path)), + config=Config(), + agent_file=agent_file, + yolo=True, + ) as session: + cli = cast(Any, session)._cli + toolset = cli.soul.agent.toolset + tool = toolset.find("Ls") + assert tool is not None + assert tool.__class__.__module__ == "my_tools.ls" + finally: + if sys.path and sys.path[0] == str(tmp_path): + sys.path.pop(0) diff --git a/scripts/check_version_tag.py b/scripts/check_version_tag.py index 945d222..a945469 100644 --- a/scripts/check_version_tag.py +++ b/scripts/check_version_tag.py @@ -22,7 +22,9 @@ def load_project_version(pyproject_path: Path) -> str: def main() -> int: - parser = argparse.ArgumentParser(description="Validate tag version against pyproject.") + parser = argparse.ArgumentParser( + description="Validate tag version against pyproject." + ) parser.add_argument("--pyproject", type=Path, required=True) parser.add_argument("--expected-version", required=True) args = parser.parse_args()