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

kimi_cli/cli.py

@@ -0,0 +1,249 @@
+import asyncio
+import json
+import sys
+from collections.abc import Callable
+from pathlib import Path
+from typing import Any, Literal, get_args
+
+import click
+
+from kimi_cli import KimiCLI
+from kimi_cli.constant import VERSION
+from kimi_cli.session import 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
+
+
+UIMode = Literal["shell", "print", "acp"]
+
+
+@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=None,
+    help="Custom agent specification file. Default: builtin default agent.",
+)
+@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(get_args(UIMode)),
+    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`. Note: print mode implicitly adds `--yolo`.",
+)
+@click.option(
+    "--acp",
+    "ui",
+    flag_value="acp",
+    help="Start ACP server. Shortcut for `--ui acp`.",
+)
+@click.option(
+    "--input-format",
+    type=click.Choice(get_args(InputFormat)),
+    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(get_args(OutputFormat)),
+    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 | None,
+    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."""
+
+    def _noop_echo(*args: Any, **kwargs: Any):
+        pass
+
+    echo: Callable[..., None] = click.echo if verbose else _noop_echo
+
+    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 = Session.continue_(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 = Session.create(work_dir)
+        echo(f"✓ Created new session: {session.id}")
+    echo(f"✓ Session history file: {session.history_file}")
+
+    if command is not None:
+        command = command.strip()
+        if not command:
+            raise click.BadOptionUsage("--command", "Command cannot be empty")
+
+    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(encoding="utf-8")) 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
+
+    async def _run() -> bool:
+        instance = await KimiCLI.create(
+            session,
+            yolo=yolo or (ui == "print"),  # print mode implies yolo
+            stream=ui != "print",  # use non-streaming mode only for print UI
+            mcp_configs=mcp_configs,
+            model_name=model_name,
+            agent_file=agent_file,
+        )
+        match ui:
+            case "shell":
+                return await instance.run_shell_mode(command)
+            case "print":
+                return await instance.run_print_mode(
+                    input_format or "text",
+                    output_format or "text",
+                    command,
+                )
+            case "acp":
+                if command is not None:
+                    logger.warning("ACP server ignores command argument")
+                return await instance.run_acp_server()
+
+    while True:
+        try:
+            succeeded = asyncio.run(_run())
+            if not succeeded:
+                sys.exit(1)
+            break
+        except Reload:
+            continue
+
+
+def main():
+    kimi()
+
+
+if __name__ == "__main__":
+    main()

kimi_cli/config.py

@@ -4,6 +4,7 @@ from typing import Literal, Self
 
 from pydantic import BaseModel, Field, SecretStr, ValidationError, field_serializer, model_validator
 
+from kimi_cli.exception import ConfigError
 from kimi_cli.share import get_share_dir
 from kimi_cli.utils.logging import logger
 
@@ -17,6 +18,8 @@ class LLMProvider(BaseModel):
     """API base URL"""
     api_key: SecretStr
     """API key"""
+    custom_headers: dict[str, str] = Field(default_factory=dict)
+    """Custom headers to include in API requests"""
 
     @field_serializer("api_key", when_used="json")
     def dump_secret(self, v: SecretStr):
@@ -50,6 +53,8 @@ class MoonshotSearchConfig(BaseModel):
     """Base URL for Moonshot Search service."""
     api_key: SecretStr
     """API key for Moonshot Search service."""
+    custom_headers: dict[str, str] = Field(default_factory=dict)
+    """Custom headers to include in API requests."""
 
     @field_serializer("api_key", when_used="json")
     def dump_secret(self, v: SecretStr):
@@ -99,13 +104,21 @@ def get_default_config() -> Config:
     )
 
 
-def load_config() -> Config:
-    """Load configuration from config file.
+def load_config(config_file: Path | None = None) -> Config:
+    """
+    Load configuration from config file.
+    If the config file does not exist, create it with default configuration.
+
+    Args:
+        config_file (Path | None): Path to the configuration file. If None, use default path.
 
     Returns:
         Validated Config object.
+
+    Raises:
+        ConfigError: If the configuration file is invalid.
     """
-    config_file = get_config_file()
+    config_file = config_file or get_config_file()
     logger.debug("Loading config from file: {file}", file=config_file)
 
     if not config_file.exists():
@@ -119,20 +132,21 @@ def load_config() -> Config:
         with open(config_file, encoding="utf-8") as f:
             data = json.load(f)
         return Config(**data)
-    except (json.JSONDecodeError, ValidationError) as e:
-        raise ConfigError(f"Invalid configuration file: {config_file}") from e
-
+    except json.JSONDecodeError as e:
+        raise ConfigError(f"Invalid JSON in configuration file: {e}") from e
+    except ValidationError as e:
+        raise ConfigError(f"Invalid configuration file: {e}") from e
 
-class ConfigError(Exception):
-    """Configuration error."""
-
-    def __init__(self, message: str):
-        super().__init__(message)
 
+def save_config(config: Config, config_file: Path | None = None):
+    """
+    Save configuration to config file.
 
-def save_config(config: Config):
-    """Save configuration to config file."""
-    config_file = get_config_file()
+    Args:
+        config (Config): Config object to save.
+        config_file (Path | None): Path to the configuration file. If None, use default path.
+    """
+    config_file = config_file or get_config_file()
     logger.debug("Saving config to file: {file}", file=config_file)
     with open(config_file, "w", encoding="utf-8") as f:
         f.write(config.model_dump_json(indent=2, exclude_none=True))

kimi_cli/constant.py

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

kimi_cli/exception.py

@@ -0,0 +1,16 @@
+class KimiCLIException(Exception):
+    """Base exception class for Kimi CLI."""
+
+    pass
+
+
+class ConfigError(KimiCLIException):
+    """Configuration error."""
+
+    pass
+
+
+class AgentSpecError(KimiCLIException):
+    """Agent specification error."""
+
+    pass

kimi_cli/llm.py

@@ -1,8 +1,78 @@
+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,
+                    **provider.custom_headers,
+                },
+            )
+            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/metadata.py

@@ -1,8 +1,6 @@
 import json
-import uuid
 from hashlib import md5
 from pathlib import Path
-from typing import NamedTuple
 
 from pydantic import BaseModel, Field
 
@@ -33,10 +31,12 @@ class WorkDirMeta(BaseModel):
 class Metadata(BaseModel):
     """Kimi metadata structure."""
 
-    work_dirs: list[WorkDirMeta] = Field(default_factory=list, description="Work directory list")
+    work_dirs: list[WorkDirMeta] = Field(
+        default_factory=list[WorkDirMeta], description="Work directory list"
+    )
 
 
-def _load_metadata() -> Metadata:
+def load_metadata() -> Metadata:
     metadata_file = get_metadata_file()
     logger.debug("Loading metadata from file: {file}", file=metadata_file)
     if not metadata_file.exists():
@@ -47,71 +47,8 @@ def _load_metadata() -> Metadata:
         return Metadata(**data)
 
 
-def _save_metadata(metadata: Metadata):
+def save_metadata(metadata: Metadata):
     metadata_file = get_metadata_file()
     logger.debug("Saving metadata to file: {file}", file=metadata_file)
     with open(metadata_file, "w", encoding="utf-8") as f:
         json.dump(metadata.model_dump(), f, indent=2, ensure_ascii=False)
-
-
-class Session(NamedTuple):
-    """A session of a work directory."""
-
-    id: str
-    work_dir: WorkDirMeta
-    history_file: Path
-
-
-def new_session(work_dir: Path, _history_file: Path | None = None) -> Session:
-    """Create a new session for a work directory."""
-    logger.debug("Creating new session for work directory: {work_dir}", work_dir=work_dir)
-
-    metadata = _load_metadata()
-    work_dir_meta = next((wd for wd in metadata.work_dirs if wd.path == str(work_dir)), None)
-    if work_dir_meta is None:
-        work_dir_meta = WorkDirMeta(path=str(work_dir))
-        metadata.work_dirs.append(work_dir_meta)
-
-    session_id = str(uuid.uuid4())
-    if _history_file is None:
-        history_file = work_dir_meta.sessions_dir / f"{session_id}.jsonl"
-        work_dir_meta.last_session_id = session_id
-    else:
-        logger.warning("Using provided history file: {history_file}", history_file=_history_file)
-        _history_file.parent.mkdir(parents=True, exist_ok=True)
-        if _history_file.exists():
-            assert _history_file.is_file()
-        history_file = _history_file
-
-    if history_file.exists():
-        # truncate if exists
-        logger.warning(
-            "History file already exists, truncating: {history_file}", history_file=history_file
-        )
-        history_file.unlink()
-        history_file.touch()
-
-    _save_metadata(metadata)
-    return Session(id=session_id, work_dir=work_dir_meta, history_file=history_file)
-
-
-def continue_session(work_dir: Path) -> Session | None:
-    """Get the last session for a work directory."""
-    logger.debug("Continuing session for work directory: {work_dir}", work_dir=work_dir)
-
-    metadata = _load_metadata()
-    work_dir_meta = next((wd for wd in metadata.work_dirs if wd.path == str(work_dir)), None)
-    if work_dir_meta is None:
-        logger.debug("Work directory never been used")
-        return None
-    if work_dir_meta.last_session_id is None:
-        logger.debug("Work directory never had a session")
-        return None
-
-    logger.debug(
-        "Found last session for work directory: {session_id}",
-        session_id=work_dir_meta.last_session_id,
-    )
-    session_id = work_dir_meta.last_session_id
-    history_file = work_dir_meta.sessions_dir / f"{session_id}.jsonl"
-    return Session(id=session_id, work_dir=work_dir_meta, history_file=history_file)

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/session.py

@@ -0,0 +1,81 @@
+import uuid
+from pathlib import Path
+from typing import NamedTuple
+
+from kimi_cli.metadata import WorkDirMeta, load_metadata, save_metadata
+from kimi_cli.utils.logging import logger
+
+
+class Session(NamedTuple):
+    """A session of a work directory."""
+
+    id: str
+    work_dir: Path
+    history_file: Path
+
+    @staticmethod
+    def create(work_dir: Path, _history_file: Path | None = None) -> "Session":
+        """Create a new session for a work directory."""
+        logger.debug("Creating new session for work directory: {work_dir}", work_dir=work_dir)
+
+        metadata = load_metadata()
+        work_dir_meta = next((wd for wd in metadata.work_dirs if wd.path == str(work_dir)), None)
+        if work_dir_meta is None:
+            work_dir_meta = WorkDirMeta(path=str(work_dir))
+            metadata.work_dirs.append(work_dir_meta)
+
+        session_id = str(uuid.uuid4())
+        if _history_file is None:
+            history_file = work_dir_meta.sessions_dir / f"{session_id}.jsonl"
+            work_dir_meta.last_session_id = session_id
+        else:
+            logger.warning(
+                "Using provided history file: {history_file}", history_file=_history_file
+            )
+            _history_file.parent.mkdir(parents=True, exist_ok=True)
+            if _history_file.exists():
+                assert _history_file.is_file()
+            history_file = _history_file
+
+        if history_file.exists():
+            # truncate if exists
+            logger.warning(
+                "History file already exists, truncating: {history_file}", history_file=history_file
+            )
+            history_file.unlink()
+            history_file.touch()
+
+        save_metadata(metadata)
+
+        return Session(
+            id=session_id,
+            work_dir=work_dir,
+            history_file=history_file,
+        )
+
+    @staticmethod
+    def continue_(work_dir: Path) -> "Session | None":
+        """Get the last session for a work directory."""
+        logger.debug("Continuing session for work directory: {work_dir}", work_dir=work_dir)
+
+        metadata = load_metadata()
+        work_dir_meta = next((wd for wd in metadata.work_dirs if wd.path == str(work_dir)), None)
+        if work_dir_meta is None:
+            logger.debug("Work directory never been used")
+            return None
+        if work_dir_meta.last_session_id is None:
+            logger.debug("Work directory never had a session")
+            return None
+
+        logger.debug(
+            "Found last session for work directory: {session_id}",
+            session_id=work_dir_meta.last_session_id,
+        )
+        session_id = work_dir_meta.last_session_id
+        history_file = work_dir_meta.sessions_dir / f"{session_id}.jsonl"
+
+        return Session(
+            id=session_id,
+            work_dir=work_dir,
+            history_file=history_file,
+        )