kimi-cli 0.35__py3-none-any.whl → 0.52__py3-none-any.whl

kimi_cli/llm.py

@@ -1,8 +1,149 @@
-from typing import NamedTuple
+import os
+from dataclasses import dataclass
+from typing import TYPE_CHECKING, Literal, cast, get_args
 
-from kosong.base.chat_provider import ChatProvider
+from kosong.chat_provider import ChatProvider
+from pydantic import SecretStr
 
+from kimi_cli.constant import USER_AGENT
 
-class LLM(NamedTuple):
+if TYPE_CHECKING:
+    from kimi_cli.config import LLMModel, LLMProvider
+
+type ProviderType = Literal["kimi", "openai_legacy", "openai_responses", "anthropic", "_chaos"]
+
+type ModelCapability = Literal["image_in", "thinking"]
+ALL_MODEL_CAPABILITIES: set[ModelCapability] = set(get_args(ModelCapability))
+
+
+@dataclass(slots=True)
+class LLM:
     chat_provider: ChatProvider
     max_context_size: int
+    capabilities: set[ModelCapability]
+
+    @property
+    def model_name(self) -> str:
+        return self.chat_provider.model_name
+
+
+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] = {}
+
+    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_name
+            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"] = max_context_size
+            if capabilities := os.getenv("KIMI_MODEL_CAPABILITIES"):
+                caps_lower = (cap.strip().lower() for cap in capabilities.split(",") if cap.strip())
+                model.capabilities = set(
+                    cast(ModelCapability, cap)
+                    for cap in caps_lower
+                    if cap in get_args(ModelCapability)
+                )
+                applied["KIMI_MODEL_CAPABILITIES"] = capabilities
+        case "openai_legacy" | "openai_responses":
+            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
+
+    return applied
+
+
+def create_llm(
+    provider: "LLMProvider",
+    model: "LLMModel",
+    *,
+    stream: bool = True,
+    session_id: str | None = None,
+) -> LLM:
+    match provider.type:
+        case "kimi":
+            from kosong.chat_provider.kimi import 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 or {}),
+                },
+            )
+            if session_id:
+                chat_provider = chat_provider.with_generation_kwargs(prompt_cache_key=session_id)
+        case "openai_legacy":
+            from kosong.contrib.chat_provider.openai_legacy import OpenAILegacy
+
+            chat_provider = OpenAILegacy(
+                model=model.model,
+                base_url=provider.base_url,
+                api_key=provider.api_key.get_secret_value(),
+                stream=stream,
+            )
+        case "openai_responses":
+            from kosong.contrib.chat_provider.openai_responses import OpenAIResponses
+
+            chat_provider = OpenAIResponses(
+                model=model.model,
+                base_url=provider.base_url,
+                api_key=provider.api_key.get_secret_value(),
+                stream=stream,
+            )
+        case "anthropic":
+            from kosong.contrib.chat_provider.anthropic import Anthropic
+
+            chat_provider = Anthropic(
+                model=model.model,
+                base_url=provider.base_url,
+                api_key=provider.api_key.get_secret_value(),
+                stream=stream,
+                default_max_tokens=50000,
+            )
+        case "_chaos":
+            from kosong.chat_provider.chaos import ChaosChatProvider, ChaosConfig
+
+            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,
+        capabilities=_derive_capabilities(provider, model),
+    )
+
+
+def _derive_capabilities(provider: "LLMProvider", model: "LLMModel") -> set[ModelCapability]:
+    capabilities = model.capabilities or set()
+    if provider.type != "kimi":
+        return capabilities
+
+    if model.model == "kimi-for-coding" or "thinking" in model.model:
+        capabilities.add("thinking")
+    return capabilities

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
 
@@ -25,7 +23,7 @@ class WorkDirMeta(BaseModel):
 
     @property
     def sessions_dir(self) -> Path:
-        path = get_share_dir() / "sessions" / md5(self.path.encode()).hexdigest()
+        path = get_share_dir() / "sessions" / md5(self.path.encode(encoding="utf-8")).hexdigest()
         path.mkdir(parents=True, exist_ok=True)
         return path
 
@@ -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

@@ -0,0 +1,4 @@
+from pathlib import Path
+
+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,103 @@
+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"
+        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,
+        )
+
+    def mark_as_last(self) -> None:
+        """Mark this session as the last completed session for its work directory."""
+        metadata = load_metadata()
+        work_dir_meta = next(
+            (wd for wd in metadata.work_dirs if wd.path == str(self.work_dir)), None
+        )
+
+        if work_dir_meta is None:
+            logger.warning(
+                "Work directory metadata missing when marking last session, recreating: {work_dir}",
+                work_dir=self.work_dir,
+            )
+            work_dir_meta = WorkDirMeta(path=str(self.work_dir))
+            metadata.work_dirs.append(work_dir_meta)
+
+        work_dir_meta.last_session_id = self.id
+        logger.debug(
+            "Updated last session for work directory: {work_dir} -> {session_id}",
+            work_dir=self.work_dir,
+            session_id=self.id,
+        )
+        save_metadata(metadata)

kimi_cli/soul/__init__.py

@@ -1,7 +1,17 @@
-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 TYPE_CHECKING, Any, NamedTuple, Protocol, runtime_checkable
+
+from kosong.message import ContentPart
+
+from kimi_cli.utils.logging import logger
+from kimi_cli.wire import Wire, WireUISide
+from kimi_cli.wire.message import WireMessage
 
 if TYPE_CHECKING:
-    from kimi_cli.soul.wire import Wire
+    from kimi_cli.llm import LLM, ModelCapability
 
 
 class LLMNotSet(Exception):
@@ -10,6 +20,19 @@ class LLMNotSet(Exception):
     pass
 
 
+class LLMNotSupported(Exception):
+    """Raised when the LLM does not have required capabilities."""
+
+    def __init__(self, llm: "LLM", capabilities: "list[ModelCapability]"):
+        self.llm = llm
+        self.capabilities = capabilities
+        capabilities_str = "capability" if len(capabilities) == 1 else "capabilities"
+        super().__init__(
+            f"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."""
 
@@ -33,8 +56,13 @@ class Soul(Protocol):
         ...
 
     @property
-    def model(self) -> str:
-        """The LLM model used by the soul. Empty string indicates no LLM configured."""
+    def model_name(self) -> str:
+        """The name of the LLM model used by the soul. Empty string indicates no LLM configured."""
+        ...
+
+    @property
+    def model_capabilities(self) -> "set[ModelCapability] | None":
+        """The capabilities of the LLM model used by the soul. None indicates no LLM configured."""
         ...
 
     @property
@@ -42,18 +70,111 @@ 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 | list[ContentPart]):
         """
-        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.
+            user_input (str | list[ContentPart]): The user input to the agent.
 
         Raises:
-            ChatProviderNotSet: When the chat provider is not set.
+            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.
         """
         ...
+
+
+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 | list[ContentPart],
+    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.
+        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.
+    """
+    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:
+            logger.debug("UI loop shut down")
+            pass
+        except TimeoutError:
+            logger.warning("UI loop timed out")
+        finally:
+            _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,159 @@
+import importlib
+import inspect
+import string
+from pathlib import Path
+from typing import Any, NamedTuple
+
+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 import SkipThisTool
+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:
+        try:
+            tool = _load_tool(tool_path, dependencies)
+        except SkipThisTool:
+            logger.info("Skipping tool: {tool_path}", tool_path=tool_path)
+            continue
+        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.
+    """
+    import fastmcp
+
+    from kimi_cli.tools.mcp import MCPTool
+
+    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