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

kimi_cli/cli.py

@@ -0,0 +1,235 @@
+import asyncio
+import json
+import sys
+from pathlib import Path
+
+import click
+
+from kimi_cli import UIMode, kimi_run
+from kimi_cli.agentspec import DEFAULT_AGENT_FILE
+from kimi_cli.config import ConfigError, load_config
+from kimi_cli.constant import VERSION
+from kimi_cli.metadata import continue_session, new_session
+from kimi_cli.share import get_share_dir
+from kimi_cli.ui.print import InputFormat, OutputFormat
+from kimi_cli.utils.logging import logger
+
+
+class Reload(Exception):
+    """Reload configuration."""
+
+    pass
+
+
+@click.command(context_settings=dict(help_option_names=["-h", "--help"]))
+@click.version_option(VERSION)
+@click.option(
+    "--verbose",
+    is_flag=True,
+    default=False,
+    help="Print verbose information. Default: no.",
+)
+@click.option(
+    "--debug",
+    is_flag=True,
+    default=False,
+    help="Log debug information. Default: no.",
+)
+@click.option(
+    "--agent-file",
+    type=click.Path(exists=True, file_okay=True, dir_okay=False, path_type=Path),
+    default=DEFAULT_AGENT_FILE,
+    help="Custom agent specification file. Default: builtin Kimi Koder.",
+)
+@click.option(
+    "--model",
+    "-m",
+    "model_name",
+    type=str,
+    default=None,
+    help="LLM model to use. Default: default model set in config file.",
+)
+@click.option(
+    "--work-dir",
+    "-w",
+    type=click.Path(exists=True, file_okay=False, dir_okay=True, path_type=Path),
+    default=Path.cwd(),
+    help="Working directory for the agent. Default: current directory.",
+)
+@click.option(
+    "--continue",
+    "-C",
+    "continue_",
+    is_flag=True,
+    default=False,
+    help="Continue the previous session for the working directory. Default: no.",
+)
+@click.option(
+    "--command",
+    "-c",
+    "--query",
+    "-q",
+    "command",
+    type=str,
+    default=None,
+    help="User query to the agent. Default: prompt interactively.",
+)
+@click.option(
+    "--ui",
+    "ui",
+    type=click.Choice(["shell", "print", "acp"]),
+    default="shell",
+    help="UI mode to use. Default: shell.",
+)
+@click.option(
+    "--print",
+    "ui",
+    flag_value="print",
+    help="Run in print mode. Shortcut for `--ui print`.",
+)
+@click.option(
+    "--acp",
+    "ui",
+    flag_value="acp",
+    help="Start ACP server. Shortcut for `--ui acp`.",
+)
+@click.option(
+    "--input-format",
+    type=click.Choice(["text", "stream-json"]),
+    default=None,
+    help=(
+        "Input format to use. Must be used with `--print` "
+        "and the input must be piped in via stdin. "
+        "Default: text."
+    ),
+)
+@click.option(
+    "--output-format",
+    type=click.Choice(["text", "stream-json"]),
+    default=None,
+    help="Output format to use. Must be used with `--print`. Default: text.",
+)
+@click.option(
+    "--mcp-config-file",
+    type=click.Path(exists=True, file_okay=True, dir_okay=False, path_type=Path),
+    multiple=True,
+    help=(
+        "MCP config file to load. Add this option multiple times to specify multiple MCP configs. "
+        "Default: none."
+    ),
+)
+@click.option(
+    "--mcp-config",
+    type=str,
+    multiple=True,
+    help=(
+        "MCP config JSON to load. Add this option multiple times to specify multiple MCP configs. "
+        "Default: none."
+    ),
+)
+@click.option(
+    "--yolo",
+    "--yes",
+    "-y",
+    "--auto-approve",
+    "yolo",
+    is_flag=True,
+    default=False,
+    help="Automatically approve all actions. Default: no.",
+)
+def kimi(
+    verbose: bool,
+    debug: bool,
+    agent_file: Path,
+    model_name: str | None,
+    work_dir: Path,
+    continue_: bool,
+    command: str | None,
+    ui: UIMode,
+    input_format: InputFormat | None,
+    output_format: OutputFormat | None,
+    mcp_config_file: list[Path],
+    mcp_config: list[str],
+    yolo: bool,
+):
+    """Kimi, your next CLI agent."""
+    echo = click.echo if verbose else lambda *args, **kwargs: None
+
+    logger.add(
+        get_share_dir() / "logs" / "kimi.log",
+        level="DEBUG" if debug else "INFO",
+        rotation="06:00",
+        retention="10 days",
+    )
+
+    work_dir = work_dir.absolute()
+
+    if continue_:
+        session = continue_session(work_dir)
+        if session is None:
+            raise click.BadOptionUsage(
+                "--continue", "No previous session found for the working directory"
+            )
+        echo(f"✓ Continuing previous session: {session.id}")
+    else:
+        session = new_session(work_dir)
+        echo(f"✓ Created new session: {session.id}")
+    echo(f"✓ Session history file: {session.history_file}")
+
+    if input_format is not None and ui != "print":
+        raise click.BadOptionUsage(
+            "--input-format",
+            "Input format is only supported for print UI",
+        )
+    if output_format is not None and ui != "print":
+        raise click.BadOptionUsage(
+            "--output-format",
+            "Output format is only supported for print UI",
+        )
+
+    try:
+        mcp_configs = [json.loads(conf.read_text()) for conf in mcp_config_file]
+    except json.JSONDecodeError as e:
+        raise click.BadOptionUsage("--mcp-config-file", f"Invalid JSON: {e}") from e
+
+    try:
+        mcp_configs += [json.loads(conf) for conf in mcp_config]
+    except json.JSONDecodeError as e:
+        raise click.BadOptionUsage("--mcp-config", f"Invalid JSON: {e}") from e
+
+    while True:
+        try:
+            try:
+                config = load_config()
+            except ConfigError as e:
+                raise click.ClickException(f"Failed to load config: {e}") from e
+            echo(f"✓ Loaded config: {config}")
+
+            succeeded = asyncio.run(
+                kimi_run(
+                    config=config,
+                    model_name=model_name,
+                    work_dir=work_dir,
+                    session=session,
+                    command=command,
+                    agent_file=agent_file,
+                    ui=ui,
+                    input_format=input_format,
+                    output_format=output_format,
+                    mcp_configs=mcp_configs,
+                    yolo=yolo,
+                )
+            )
+            if not succeeded:
+                sys.exit(1)
+            break
+        except Reload:
+            continue
+
+
+def main():
+    kimi()
+
+
+if __name__ == "__main__":
+    main()

kimi_cli/constant.py

@@ -0,0 +1,4 @@
+import importlib.metadata
+
+VERSION = importlib.metadata.version("kimi-cli")
+USER_AGENT = f"KimiCLI/{VERSION}"

kimi_cli/llm.py

@@ -1,8 +1,77 @@
+import os
 from typing import NamedTuple
 
 from kosong.base.chat_provider import ChatProvider
+from kosong.chat_provider.chaos import ChaosChatProvider, ChaosConfig
+from kosong.chat_provider.kimi import Kimi
+from kosong.chat_provider.openai_legacy import OpenAILegacy
+from pydantic import SecretStr
+
+from kimi_cli.config import LLMModel, LLMProvider
+from kimi_cli.constant import USER_AGENT
 
 
 class LLM(NamedTuple):
     chat_provider: ChatProvider
     max_context_size: int
+
+
+def augment_provider_with_env_vars(provider: LLMProvider, model: LLMModel):
+    match provider.type:
+        case "kimi":
+            if base_url := os.getenv("KIMI_BASE_URL"):
+                provider.base_url = base_url
+            if api_key := os.getenv("KIMI_API_KEY"):
+                provider.api_key = SecretStr(api_key)
+            if model_name := os.getenv("KIMI_MODEL_NAME"):
+                model.model = model_name
+            if max_context_size := os.getenv("KIMI_MODEL_MAX_CONTEXT_SIZE"):
+                model.max_context_size = int(max_context_size)
+        case "openai_legacy":
+            if base_url := os.getenv("OPENAI_BASE_URL"):
+                provider.base_url = base_url
+            if api_key := os.getenv("OPENAI_API_KEY"):
+                provider.api_key = SecretStr(api_key)
+        case _:
+            pass
+
+
+def create_llm(
+    provider: LLMProvider,
+    model: LLMModel,
+    *,
+    stream: bool = True,
+    session_id: str | None = None,
+) -> LLM:
+    match provider.type:
+        case "kimi":
+            chat_provider = Kimi(
+                model=model.model,
+                base_url=provider.base_url,
+                api_key=provider.api_key.get_secret_value(),
+                stream=stream,
+                default_headers={
+                    "User-Agent": USER_AGENT,
+                },
+            )
+            if session_id:
+                chat_provider = chat_provider.with_generation_kwargs(prompt_cache_key=session_id)
+        case "openai_legacy":
+            chat_provider = OpenAILegacy(
+                model=model.model,
+                base_url=provider.base_url,
+                api_key=provider.api_key.get_secret_value(),
+                stream=stream,
+            )
+        case "_chaos":
+            chat_provider = ChaosChatProvider(
+                model=model.model,
+                base_url=provider.base_url,
+                api_key=provider.api_key.get_secret_value(),
+                chaos_config=ChaosConfig(
+                    error_probability=0.8,
+                    error_types=[429, 500, 503],
+                ),
+            )
+
+    return LLM(chat_provider=chat_provider, max_context_size=model.max_context_size)

kimi_cli/prompts/__init__.py

@@ -1,4 +1,4 @@
 from pathlib import Path
 
-INIT = (Path(__file__).parent / "init.md").read_text()
-COMPACT = (Path(__file__).parent / "compact.md").read_text()
+INIT = (Path(__file__).parent / "init.md").read_text(encoding="utf-8")
+COMPACT = (Path(__file__).parent / "compact.md").read_text(encoding="utf-8")

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,157 @@
+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.metadata import Session
+from kimi_cli.soul.approval import Approval
+from kimi_cli.soul.denwarenji import DenwaRenji
+from kimi_cli.soul.globals import AgentGlobals, BuiltinSystemPromptArgs
+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_with_mcp(
+    agent_file: Path,
+    globals_: AgentGlobals,
+    mcp_configs: list[dict[str, Any]],
+) -> Agent:
+    agent = load_agent(agent_file, globals_)
+    assert isinstance(agent.toolset, CustomToolset)
+    if mcp_configs:
+        await _load_mcp_tools(agent.toolset, mcp_configs)
+    return agent
+
+
+def load_agent(
+    agent_file: Path,
+    globals_: AgentGlobals,
+) -> Agent:
+    """
+    Load agent from specification file.
+
+    Raises:
+        ValueError: 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,
+        globals_.builtin_args,
+    )
+
+    tool_deps = {
+        ResolvedAgentSpec: agent_spec,
+        AgentGlobals: globals_,
+        Config: globals_.config,
+        BuiltinSystemPromptArgs: globals_.builtin_args,
+        Session: globals_.session,
+        DenwaRenji: globals_.denwa_renji,
+        Approval: globals_.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}")
+
+    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/context.py

@@ -19,6 +19,11 @@ class Context:
         self._next_checkpoint_id: int = 0
         """The ID of the next checkpoint, starting from 0, incremented after each checkpoint."""
 
+    @property
+    def file_backend(self) -> Path:
+        """The JSONL file backend of the context."""
+        return self._file_backend
+
     async def restore(self) -> bool:
         logger.debug("Restoring context from file: {file_backend}", file_backend=self._file_backend)
         if self._history: