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

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,12 +18,17 @@ class LLMProvider(BaseModel):
     """API base URL"""
     api_key: SecretStr
     """API key"""
+    custom_headers: dict[str, str] | None = None
+    """Custom headers to include in API requests"""
 
     @field_serializer("api_key", when_used="json")
     def dump_secret(self, v: SecretStr):
         return v.get_secret_value()
 
 
+LLMModelCapability = Literal["image_in"]
+
+
 class LLMModel(BaseModel):
     """LLM model configuration."""
 
@@ -32,6 +38,8 @@ class LLMModel(BaseModel):
     """Model name"""
     max_context_size: int
     """Maximum context size (unit: tokens)"""
+    capabilities: set[LLMModelCapability] | None = None
+    """Model capabilities"""
 
 
 class LoopControl(BaseModel):
@@ -50,6 +58,8 @@ class MoonshotSearchConfig(BaseModel):
     """Base URL for Moonshot Search service."""
     api_key: SecretStr
     """API key for Moonshot Search service."""
+    custom_headers: dict[str, str] | None = None
+    """Custom headers to include in API requests."""
 
     @field_serializer("api_key", when_used="json")
     def dump_secret(self, v: SecretStr):
@@ -99,13 +109,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 +137,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/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

@@ -7,26 +7,47 @@ 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.config import LLMModel, LLMModelCapability, LLMProvider
 from kimi_cli.constant import USER_AGENT
 
 
 class LLM(NamedTuple):
     chat_provider: ChatProvider
     max_context_size: int
+    capabilities: set[LLMModelCapability]
+    # TODO: these additional fields should be moved to ChatProvider
 
+    @property
+    def model_name(self) -> str:
+        return self.chat_provider.model_name
+
+    @property
+    def supports_image_in(self) -> bool:
+        return "image_in" in self.capabilities
+
+
+def augment_provider_with_env_vars(provider: LLMProvider, model: LLMModel) -> dict[str, str]:
+    """Override provider/model settings from environment variables.
+
+    Returns:
+        Mapping of environment variables that were applied.
+    """
+    applied: dict[str, str] = {}
 
-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
+                applied["KIMI_BASE_URL"] = base_url
             if api_key := os.getenv("KIMI_API_KEY"):
                 provider.api_key = SecretStr(api_key)
+                applied["KIMI_API_KEY"] = "******"
             if model_name := os.getenv("KIMI_MODEL_NAME"):
                 model.model = model_name
+                applied["KIMI_MODEL_NAME"] = model.model
             if max_context_size := os.getenv("KIMI_MODEL_MAX_CONTEXT_SIZE"):
                 model.max_context_size = int(max_context_size)
+                applied["KIMI_MODEL_MAX_CONTEXT_SIZE"] = str(model.max_context_size)
         case "openai_legacy":
             if base_url := os.getenv("OPENAI_BASE_URL"):
                 provider.base_url = base_url
@@ -35,6 +56,8 @@ def augment_provider_with_env_vars(provider: LLMProvider, model: LLMModel):
         case _:
             pass
 
+    return applied
+
 
 def create_llm(
     provider: LLMProvider,
@@ -52,6 +75,7 @@ def create_llm(
                 stream=stream,
                 default_headers={
                     "User-Agent": USER_AGENT,
+                    **(provider.custom_headers or {}),
                 },
             )
             if session_id:
@@ -74,4 +98,8 @@ def create_llm(
                 ),
             )
 
-    return LLM(chat_provider=chat_provider, max_context_size=model.max_context_size)
+    return LLM(
+        chat_provider=chat_provider,
+        max_context_size=model.max_context_size,
+        capabilities=model.capabilities or set(),
+    )

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/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,
+        )

kimi_cli/soul/__init__.py

@@ -4,6 +4,9 @@ from collections.abc import Callable, Coroutine
 from contextvars import ContextVar
 from typing import Any, NamedTuple, Protocol, runtime_checkable
 
+from kosong.base.message import ContentPart
+
+from kimi_cli.llm import LLM
 from kimi_cli.utils.logging import logger
 from kimi_cli.wire import Wire, WireUISide
 from kimi_cli.wire.message import WireMessage
@@ -15,6 +18,19 @@ class LLMNotSet(Exception):
     pass
 
 
+class LLMNotSupported(Exception):
+    """Raised when the LLM does not have required capabilities."""
+
+    def __init__(self, llm: LLM, capabilities: list[str]):
+        self.llm = llm
+        self.capabilities = capabilities
+        capabilities_str = "capability" if len(capabilities) == 1 else "capabilities"
+        super().__init__(
+            f"The LLM model '{llm.model_name}' does not support required {capabilities_str}: "
+            f"{', '.join(capabilities)}."
+        )
+
+
 class MaxStepsReached(Exception):
     """Raised when the maximum number of steps is reached."""
 
@@ -47,15 +63,16 @@ class Soul(Protocol):
         """The current status of the soul. The returned value is immutable."""
         ...
 
-    async def run(self, user_input: str):
+    async def run(self, user_input: str | list[ContentPart]):
         """
         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.
+            user_input (str | list[ContentPart]): The user input to the agent.
 
         Raises:
             LLMNotSet: When the LLM is not set.
+            LLMNotSupported: When the LLM does not have required capabilities.
             ChatProviderError: When the LLM provider returns an error.
             MaxStepsReached: When the maximum number of steps is reached.
             asyncio.CancelledError: When the run is cancelled by user.
@@ -73,7 +90,7 @@ class RunCancelled(Exception):
 
 async def run_soul(
     soul: "Soul",
-    user_input: str,
+    user_input: str | list[ContentPart],
     ui_loop_fn: UILoopFn,
     cancel_event: asyncio.Event,
 ) -> None:
@@ -85,6 +102,7 @@ async def run_soul(
 
     Raises:
         LLMNotSet: When the LLM is not set.
+        LLMNotSupported: When the LLM does not have required capabilities.
         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.
@@ -125,7 +143,7 @@ async def run_soul(
         try:
             await asyncio.wait_for(ui_task, timeout=0.5)
         except asyncio.QueueShutDown:
-            # expected
+            logger.debug("UI loop shut down")
             pass
         except TimeoutError:
             logger.warning("UI loop timed out")

kimi_cli/soul/agent.py

@@ -9,10 +9,10 @@ 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.session 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.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
@@ -26,27 +26,18 @@ class Agent(NamedTuple):
     toolset: Toolset
 
 
-async def load_agent_with_mcp(
+async def load_agent(
     agent_file: Path,
-    globals_: AgentGlobals,
+    runtime: Runtime,
+    *,
     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.
+        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)
@@ -54,17 +45,17 @@ def load_agent(
     system_prompt = _load_system_prompt(
         agent_spec.system_prompt_path,
         agent_spec.system_prompt_args,
-        globals_.builtin_args,
+        runtime.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,
+        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:
@@ -75,6 +66,10 @@ def load_agent(
     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,

kimi_cli/soul/context.py

@@ -19,11 +19,6 @@ 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: