kimi-cli 0.40__py3-none-any.whl → 0.42__py3-none-any.whl

kimi_cli/soul/__init__.py

@@ -1,7 +1,12 @@
-from typing import TYPE_CHECKING, NamedTuple, Protocol, runtime_checkable
+import asyncio
+import contextlib
+from collections.abc import Callable, Coroutine
+from contextvars import ContextVar
+from typing import Any, NamedTuple, Protocol, runtime_checkable
 
-if TYPE_CHECKING:
-    from kimi_cli.soul.wire import Wire
+from kimi_cli.utils.logging import logger
+from kimi_cli.wire import Wire, WireUISide
+from kimi_cli.wire.message import WireMessage
 
 
 class LLMNotSet(Exception):
@@ -42,13 +47,12 @@ class Soul(Protocol):
         """The current status of the soul. The returned value is immutable."""
         ...
 
-    async def run(self, user_input: str, wire: "Wire"):
+    async def run(self, user_input: str):
         """
-        Run the agent with the given user input.
+        Run the agent with the given user input until the max steps or no more tool calls.
 
         Args:
             user_input (str): The user input to the agent.
-            wire (Wire): The wire to send events and requests to the UI loop.
 
         Raises:
             LLMNotSet: When the LLM is not set.
@@ -57,3 +61,95 @@ class Soul(Protocol):
             asyncio.CancelledError: When the run is cancelled by user.
         """
         ...
+
+
+type UILoopFn = Callable[[WireUISide], 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,
+) -> None:
+    """
+    Run the soul with the given user input, connecting it to the UI loop with a wire.
+
+    `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:
+        LLMNotSet: When the LLM is not set.
+        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()
+    wire_token = _current_wire.set(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.ui_side))
+
+    logger.debug("Starting soul run")
+    soul_task = asyncio.create_task(soul.run(user_input))
+
+    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 UI loop")
+        # shutting down the wire should break the UI loop
+        wire.shutdown()
+        try:
+            await asyncio.wait_for(ui_task, timeout=0.5)
+        except asyncio.QueueShutDown:
+            # expected
+            pass
+        except TimeoutError:
+            logger.warning("UI loop timed out")
+
+        _current_wire.reset(wire_token)
+
+
+_current_wire = ContextVar[Wire | None]("current_wire", default=None)
+
+
+def get_wire_or_none() -> Wire | None:
+    """
+    Get the current wire or None.
+    Expect to be not None when called from anywhere in the agent loop.
+    """
+    return _current_wire.get()
+
+
+def wire_send(msg: WireMessage) -> None:
+    """
+    Send a wire message to the current wire.
+    Take this as `print` and `input` for souls.
+    Souls should always use this function to send wire messages.
+    """
+    wire = get_wire_or_none()
+    assert wire is not None, "Wire is expected to be set when soul is running"
+    wire.soul_side.send(msg)

kimi_cli/soul/agent.py

@@ -0,0 +1,152 @@
+import importlib
+import inspect
+import string
+from pathlib import Path
+from typing import Any, NamedTuple
+
+import fastmcp
+from kosong.tooling import CallableTool, CallableTool2, Toolset
+
+from kimi_cli.agentspec import ResolvedAgentSpec, load_agent_spec
+from kimi_cli.config import Config
+from kimi_cli.session import Session
+from kimi_cli.soul.approval import Approval
+from kimi_cli.soul.denwarenji import DenwaRenji
+from kimi_cli.soul.runtime import BuiltinSystemPromptArgs, Runtime
+from kimi_cli.soul.toolset import CustomToolset
+from kimi_cli.tools.mcp import MCPTool
+from kimi_cli.utils.logging import logger
+
+
+class Agent(NamedTuple):
+    """The loaded agent."""
+
+    name: str
+    system_prompt: str
+    toolset: Toolset
+
+
+async def load_agent(
+    agent_file: Path,
+    runtime: Runtime,
+    *,
+    mcp_configs: list[dict[str, Any]],
+) -> Agent:
+    """
+    Load agent from specification file.
+
+    Raises:
+        FileNotFoundError: If the agent spec file does not exist.
+        AgentSpecError: If the agent spec is not valid.
+    """
+    logger.info("Loading agent: {agent_file}", agent_file=agent_file)
+    agent_spec = load_agent_spec(agent_file)
+
+    system_prompt = _load_system_prompt(
+        agent_spec.system_prompt_path,
+        agent_spec.system_prompt_args,
+        runtime.builtin_args,
+    )
+
+    tool_deps = {
+        ResolvedAgentSpec: agent_spec,
+        Runtime: runtime,
+        Config: runtime.config,
+        BuiltinSystemPromptArgs: runtime.builtin_args,
+        Session: runtime.session,
+        DenwaRenji: runtime.denwa_renji,
+        Approval: runtime.approval,
+    }
+    tools = agent_spec.tools
+    if agent_spec.exclude_tools:
+        logger.debug("Excluding tools: {tools}", tools=agent_spec.exclude_tools)
+        tools = [tool for tool in tools if tool not in agent_spec.exclude_tools]
+    toolset = CustomToolset()
+    bad_tools = _load_tools(toolset, tools, tool_deps)
+    if bad_tools:
+        raise ValueError(f"Invalid tools: {bad_tools}")
+
+    assert isinstance(toolset, CustomToolset)
+    if mcp_configs:
+        await _load_mcp_tools(toolset, mcp_configs)
+
+    return Agent(
+        name=agent_spec.name,
+        system_prompt=system_prompt,
+        toolset=toolset,
+    )
+
+
+def _load_system_prompt(
+    path: Path, args: dict[str, str], builtin_args: BuiltinSystemPromptArgs
+) -> str:
+    logger.info("Loading system prompt: {path}", path=path)
+    system_prompt = path.read_text(encoding="utf-8").strip()
+    logger.debug(
+        "Substituting system prompt with builtin args: {builtin_args}, spec args: {spec_args}",
+        builtin_args=builtin_args,
+        spec_args=args,
+    )
+    return string.Template(system_prompt).substitute(builtin_args._asdict(), **args)
+
+
+type ToolType = CallableTool | CallableTool2[Any]
+# TODO: move this to kosong.tooling.simple
+
+
+def _load_tools(
+    toolset: CustomToolset,
+    tool_paths: list[str],
+    dependencies: dict[type[Any], Any],
+) -> list[str]:
+    bad_tools: list[str] = []
+    for tool_path in tool_paths:
+        tool = _load_tool(tool_path, dependencies)
+        if tool:
+            toolset += tool
+        else:
+            bad_tools.append(tool_path)
+    logger.info("Loaded tools: {tools}", tools=[tool.name for tool in toolset.tools])
+    if bad_tools:
+        logger.error("Bad tools: {bad_tools}", bad_tools=bad_tools)
+    return bad_tools
+
+
+def _load_tool(tool_path: str, dependencies: dict[type[Any], Any]) -> ToolType | None:
+    logger.debug("Loading tool: {tool_path}", tool_path=tool_path)
+    module_name, class_name = tool_path.rsplit(":", 1)
+    try:
+        module = importlib.import_module(module_name)
+    except ImportError:
+        return None
+    cls = getattr(module, class_name, None)
+    if cls is None:
+        return None
+    args: list[type[Any]] = []
+    for param in inspect.signature(cls).parameters.values():
+        if param.kind == inspect.Parameter.KEYWORD_ONLY:
+            # once we encounter a keyword-only parameter, we stop injecting dependencies
+            break
+        # all positional parameters should be dependencies to be injected
+        if param.annotation not in dependencies:
+            raise ValueError(f"Tool dependency not found: {param.annotation}")
+        args.append(dependencies[param.annotation])
+    return cls(*args)
+
+
+async def _load_mcp_tools(
+    toolset: CustomToolset,
+    mcp_configs: list[dict[str, Any]],
+):
+    """
+    Raises:
+        ValueError: If the MCP config is not valid.
+        RuntimeError: If the MCP server cannot be connected.
+    """
+    for mcp_config in mcp_configs:
+        logger.info("Loading MCP tools from: {mcp_config}", mcp_config=mcp_config)
+        client = fastmcp.Client(mcp_config)
+        async with client:
+            for tool in await client.list_tools():
+                toolset += MCPTool(tool, client)
+    return toolset

kimi_cli/soul/approval.py

@@ -1,8 +1,8 @@
 import asyncio
 
 from kimi_cli.soul.toolset import get_current_tool_call_or_none
-from kimi_cli.soul.wire import ApprovalRequest, ApprovalResponse
 from kimi_cli.utils.logging import logger
+from kimi_cli.wire.message import ApprovalRequest, ApprovalResponse
 
 
 class Approval:

kimi_cli/soul/compaction.py

@@ -10,8 +10,6 @@ from kimi_cli.llm import LLM
 from kimi_cli.soul.message import system
 from kimi_cli.utils.logging import logger
 
-MAX_PRESERVED_MESSAGES = 2
-
 
 @runtime_checkable
 class Compaction(Protocol):
@@ -33,6 +31,8 @@ class Compaction(Protocol):
 
 
 class SimpleCompaction(Compaction):
+    MAX_PRESERVED_MESSAGES = 2
+
     async def compact(self, messages: Sequence[Message], llm: LLM) -> Sequence[Message]:
         history = list(messages)
         if not history:
@@ -43,11 +43,11 @@ class SimpleCompaction(Compaction):
         for index in range(len(history) - 1, -1, -1):
             if history[index].role in {"user", "assistant"}:
                 n_preserved += 1
-                if n_preserved == MAX_PRESERVED_MESSAGES:
+                if n_preserved == self.MAX_PRESERVED_MESSAGES:
                     preserve_start_index = index
                     break
 
-        if n_preserved < MAX_PRESERVED_MESSAGES:
+        if n_preserved < self.MAX_PRESERVED_MESSAGES:
             return history
 
         to_compact = history[:preserve_start_index]

kimi_cli/soul/kimisoul.py

@@ -16,58 +16,55 @@ from kosong.chat_provider import (
 from kosong.tooling import ToolResult
 from tenacity import RetryCallState, retry_if_exception, stop_after_attempt, wait_exponential_jitter
 
-from kimi_cli.agent import Agent, AgentGlobals
-from kimi_cli.config import LoopControl
-from kimi_cli.soul import LLMNotSet, MaxStepsReached, Soul, StatusSnapshot
+from kimi_cli.soul import LLMNotSet, MaxStepsReached, Soul, StatusSnapshot, wire_send
+from kimi_cli.soul.agent import Agent
 from kimi_cli.soul.compaction import SimpleCompaction
 from kimi_cli.soul.context import Context
 from kimi_cli.soul.message import system, tool_result_to_messages
-from kimi_cli.soul.wire import (
+from kimi_cli.soul.runtime import Runtime
+from kimi_cli.tools.dmail import NAME as SendDMail_NAME
+from kimi_cli.tools.utils import ToolRejectedError
+from kimi_cli.utils.logging import logger
+from kimi_cli.wire.message import (
     CompactionBegin,
     CompactionEnd,
     StatusUpdate,
     StepBegin,
     StepInterrupted,
-    Wire,
-    current_wire,
 )
-from kimi_cli.tools.dmail import NAME as SendDMail_NAME
-from kimi_cli.tools.utils import ToolRejectedError
-from kimi_cli.utils.logging import logger
 
 RESERVED_TOKENS = 50_000
 
 
-class KimiSoul:
+class KimiSoul(Soul):
     """The soul of Kimi CLI."""
 
     def __init__(
         self,
         agent: Agent,
-        agent_globals: AgentGlobals,
+        runtime: Runtime,
         *,
         context: Context,
-        loop_control: LoopControl,
     ):
         """
         Initialize the soul.
 
         Args:
             agent (Agent): The agent to run.
-            agent_globals (AgentGlobals): Global states and parameters.
+            runtime (Runtime): Runtime parameters and states.
             context (Context): The context of the agent.
             loop_control (LoopControl): The control parameters for the agent loop.
         """
         self._agent = agent
-        self._agent_globals = agent_globals
-        self._denwa_renji = agent_globals.denwa_renji
-        self._approval = agent_globals.approval
+        self._runtime = runtime
+        self._denwa_renji = runtime.denwa_renji
+        self._approval = runtime.approval
         self._context = context
-        self._loop_control = loop_control
+        self._loop_control = runtime.config.loop_control
         self._compaction = SimpleCompaction()  # TODO: maybe configurable and composable
         self._reserved_tokens = RESERVED_TOKENS
-        if self._agent_globals.llm is not None:
-            assert self._reserved_tokens <= self._agent_globals.llm.max_context_size
+        if self._runtime.llm is not None:
+            assert self._reserved_tokens <= self._runtime.llm.max_context_size
 
         for tool in agent.toolset.tools:
             if tool.name == SendDMail_NAME:
@@ -82,7 +79,7 @@ class KimiSoul:
 
     @property
     def model(self) -> str:
-        return self._agent_globals.llm.chat_provider.model_name if self._agent_globals.llm else ""
+        return self._runtime.llm.chat_provider.model_name if self._runtime.llm else ""
 
     @property
     def status(self) -> StatusSnapshot:
@@ -90,38 +87,34 @@ class KimiSoul:
 
     @property
     def _context_usage(self) -> float:
-        if self._agent_globals.llm is not None:
-            return self._context.token_count / self._agent_globals.llm.max_context_size
+        if self._runtime.llm is not None:
+            return self._context.token_count / self._runtime.llm.max_context_size
         return 0.0
 
     async def _checkpoint(self):
         await self._context.checkpoint(self._checkpoint_with_user_message)
 
-    async def run(self, user_input: str, wire: Wire):
-        if self._agent_globals.llm is None:
+    async def run(self, user_input: str):
+        if self._runtime.llm is None:
             raise LLMNotSet()
 
         await self._checkpoint()  # this creates the checkpoint 0 on first run
         await self._context.append_message(Message(role="user", content=user_input))
         logger.debug("Appended user message to context")
-        wire_token = current_wire.set(wire)
-        try:
-            await self._agent_loop(wire)
-        finally:
-            current_wire.reset(wire_token)
+        await self._agent_loop()
 
-    async def _agent_loop(self, wire: Wire):
+    async def _agent_loop(self):
         """The main agent loop for one run."""
-        assert self._agent_globals.llm is not None
+        assert self._runtime.llm is not None
 
         async def _pipe_approval_to_wire():
             while True:
                 request = await self._approval.fetch_request()
-                wire.send(request)
+                wire_send(request)
 
         step_no = 1
         while True:
-            wire.send(StepBegin(step_no))
+            wire_send(StepBegin(step_no))
             approval_task = asyncio.create_task(_pipe_approval_to_wire())
             # FIXME: It's possible that a subagent's approval task steals approval request
             # from the main agent. We must ensure that the Task tool will redirect them
@@ -131,24 +124,24 @@ class KimiSoul:
                 # compact the context if needed
                 if (
                     self._context.token_count + self._reserved_tokens
-                    >= self._agent_globals.llm.max_context_size
+                    >= self._runtime.llm.max_context_size
                 ):
                     logger.info("Context too long, compacting...")
-                    wire.send(CompactionBegin())
+                    wire_send(CompactionBegin())
                     await self.compact_context()
-                    wire.send(CompactionEnd())
+                    wire_send(CompactionEnd())
 
                 logger.debug("Beginning step {step_no}", step_no=step_no)
                 await self._checkpoint()
                 self._denwa_renji.set_n_checkpoints(self._context.n_checkpoints)
-                finished = await self._step(wire)
+                finished = await self._step()
             except BackToTheFuture as e:
                 await self._context.revert_to(e.checkpoint_id)
                 await self._checkpoint()
                 await self._context.append_message(e.messages)
                 continue
             except (ChatProviderError, asyncio.CancelledError):
-                wire.send(StepInterrupted())
+                wire_send(StepInterrupted())
                 # break the agent loop
                 raise
             finally:
@@ -161,11 +154,11 @@ class KimiSoul:
             if step_no > self._loop_control.max_steps_per_run:
                 raise MaxStepsReached(self._loop_control.max_steps_per_run)
 
-    async def _step(self, wire: Wire) -> bool:
+    async def _step(self) -> bool:
         """Run an single step and return whether the run should be stopped."""
         # already checked in `run`
-        assert self._agent_globals.llm is not None
-        chat_provider = self._agent_globals.llm.chat_provider
+        assert self._runtime.llm is not None
+        chat_provider = self._runtime.llm.chat_provider
 
         @tenacity.retry(
             retry=retry_if_exception(self._is_retryable_error),
@@ -181,8 +174,8 @@ class KimiSoul:
                 self._agent.system_prompt,
                 self._agent.toolset,
                 self._context.history,
-                on_message_part=wire.send,
-                on_tool_result=wire.send,
+                on_message_part=wire_send,
+                on_tool_result=wire_send,
             )
 
         result = await _kosong_step_with_retry()
@@ -190,7 +183,7 @@ class KimiSoul:
         if result.usage is not None:
             # mark the token count for the context before the step
             await self._context.update_token_count(result.usage.input)
-            wire.send(StatusUpdate(status=self.status))
+            wire_send(StatusUpdate(status=self.status))
 
         # wait for all tool results (may be interrupted)
         results = await result.tool_results()
@@ -260,9 +253,9 @@ class KimiSoul:
             reraise=True,
         )
         async def _compact_with_retry() -> Sequence[Message]:
-            if self._agent_globals.llm is None:
+            if self._runtime.llm is None:
                 raise LLMNotSet()
-            return await self._compaction.compact(self._context.history, self._agent_globals.llm)
+            return await self._compaction.compact(self._context.history, self._runtime.llm)
 
         compacted_messages = await _compact_with_retry()
         await self._context.revert_to(0)

kimi_cli/soul/runtime.py

@@ -0,0 +1,94 @@
+import subprocess
+import sys
+from datetime import datetime
+from pathlib import Path
+from typing import NamedTuple
+
+from kimi_cli.config import Config
+from kimi_cli.llm import LLM
+from kimi_cli.session import Session
+from kimi_cli.soul.approval import Approval
+from kimi_cli.soul.denwarenji import DenwaRenji
+from kimi_cli.utils.logging import logger
+
+
+class BuiltinSystemPromptArgs(NamedTuple):
+    """Builtin system prompt arguments."""
+
+    KIMI_NOW: str
+    """The current datetime."""
+    KIMI_WORK_DIR: Path
+    """The current working directory."""
+    KIMI_WORK_DIR_LS: str
+    """The directory listing of current working directory."""
+    KIMI_AGENTS_MD: str  # TODO: move to first message from system prompt
+    """The content of AGENTS.md."""
+
+
+def load_agents_md(work_dir: Path) -> str | None:
+    paths = [
+        work_dir / "AGENTS.md",
+        work_dir / "agents.md",
+    ]
+    for path in paths:
+        if path.is_file():
+            logger.info("Loaded agents.md: {path}", path=path)
+            return path.read_text(encoding="utf-8").strip()
+    logger.info("No AGENTS.md found in {work_dir}", work_dir=work_dir)
+    return None
+
+
+def _list_work_dir(work_dir: Path) -> str:
+    if sys.platform == "win32":
+        ls = subprocess.run(
+            ["cmd", "/c", "dir", work_dir],
+            capture_output=True,
+            text=True,
+            encoding="utf-8",
+            errors="replace",
+        )
+    else:
+        ls = subprocess.run(
+            ["ls", "-la", work_dir],
+            capture_output=True,
+            text=True,
+            encoding="utf-8",
+            errors="replace",
+        )
+    return ls.stdout.strip()
+
+
+class Runtime(NamedTuple):
+    """Agent globals."""
+
+    config: Config
+    llm: LLM | None
+    session: Session
+    builtin_args: BuiltinSystemPromptArgs
+    denwa_renji: DenwaRenji
+    approval: Approval
+
+    @staticmethod
+    async def create(
+        config: Config,
+        llm: LLM | None,
+        session: Session,
+        yolo: bool,
+    ) -> "Runtime":
+        # FIXME: do these asynchronously
+        ls_output = _list_work_dir(session.work_dir)
+        agents_md = load_agents_md(session.work_dir) or ""
+
+        return Runtime(
+            config=config,
+            llm=llm,
+            session=session,
+            builtin_args=BuiltinSystemPromptArgs(
+                KIMI_NOW=datetime.now().astimezone().isoformat(),
+                KIMI_WORK_DIR=session.work_dir,
+                KIMI_WORK_DIR_LS=ls_output,
+                KIMI_AGENTS_MD=agents_md,
+            ),
+            denwa_renji=DenwaRenji(),
+            approval=Approval(yolo=yolo),
+        )

kimi_cli/tools/dmail/__init__.py

@@ -10,7 +10,7 @@ NAME = "SendDMail"
 
 class SendDMail(CallableTool2):
     name: str = NAME
-    description: str = (Path(__file__).parent / "dmail.md").read_text()
+    description: str = (Path(__file__).parent / "dmail.md").read_text(encoding="utf-8")
     params: type[DMail] = DMail
 
     def __init__(self, denwa_renji: DenwaRenji, **kwargs):

kimi_cli/tools/file/glob.md

@@ -14,4 +14,4 @@ Find files and directories using glob patterns. This tool supports standard glob
 
 **Bad example patterns:**
 - `**`, `**/*.py` - Any pattern starting with '**' will be rejected. Because it would recursively search all directories and subdirectories, which is very likely to yield large result that exceeds your context size. Always use more specific patterns like `src/**/*.py` instead.
-- `node_modules/**/*.js` - Although this does not start with '**', it would still highly possible to yield large result because `node_modules` is well-known to contain too many directories and files. Avoid recursivelly searching in such directories, other examples include `venv`, `.venv`, `__pycache__`, `target`. If you really need to search in a dependency, use more specific patterns like `node_modules/react/src/*` instead.
+- `node_modules/**/*.js` - Although this does not start with '**', it would still highly possible to yield large result because `node_modules` is well-known to contain too many directories and files. Avoid recursively searching in such directories, other examples include `venv`, `.venv`, `__pycache__`, `target`. If you really need to search in a dependency, use more specific patterns like `node_modules/react/src/*` instead.

kimi_cli/tools/file/glob.py

@@ -8,7 +8,7 @@ import aiofiles.os
 from kosong.tooling import CallableTool2, ToolError, ToolOk, ToolReturnType
 from pydantic import BaseModel, Field
 
-from kimi_cli.agent import BuiltinSystemPromptArgs
+from kimi_cli.soul.runtime import BuiltinSystemPromptArgs
 from kimi_cli.tools.utils import load_desc
 
 MAX_MATCHES = 1000
@@ -128,7 +128,7 @@ class Glob(CallableTool2[Params]):
             message = (
                 f"Found {len(matches)} matches for pattern `{params.pattern}`."
                 if len(matches) > 0
-                else "No matches found for pattern `{params.pattern}`."
+                else f"No matches found for pattern `{params.pattern}`."
             )
             if len(matches) > MAX_MATCHES:
                 matches = matches[:MAX_MATCHES]

kimi_cli/tools/file/grep.py

@@ -220,7 +220,7 @@ async def _ensure_rg_path() -> str:
 
 class Grep(CallableTool2[Params]):
     name: str = "Grep"
-    description: str = (Path(__file__).parent / "grep.md").read_text()
+    description: str = (Path(__file__).parent / "grep.md").read_text(encoding="utf-8")
     params: type[Params] = Params
 
     @override