kimi-cli 0.44__py3-none-any.whl → 0.78__py3-none-any.whl

kimi_cli/soul/agent.py

@@ -1,41 +1,211 @@
-import importlib
-import inspect
+from __future__ import annotations
+
+import asyncio
 import string
+from collections.abc import Mapping
+from dataclasses import asdict, dataclass
+from datetime import datetime
 from pathlib import Path
-from typing import Any, NamedTuple
+from typing import TYPE_CHECKING, Any
 
-from kosong.tooling import CallableTool, CallableTool2, Toolset
+import pydantic
+from kaos.path import KaosPath
+from kosong.tooling import Toolset
 
-from kimi_cli.agentspec import ResolvedAgentSpec, load_agent_spec
+from kimi_cli.agentspec import load_agent_spec
 from kimi_cli.config import Config
+from kimi_cli.exception import MCPConfigError
+from kimi_cli.llm import LLM
 from kimi_cli.session import Session
+from kimi_cli.skill import (
+    Skill,
+    discover_skills_from_roots,
+    get_builtin_skills_dir,
+    get_claude_skills_dir,
+    get_skills_dir,
+    index_skills,
+)
 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.soul.toolset import KimiToolset
+from kimi_cli.utils.environment import Environment
 from kimi_cli.utils.logging import logger
+from kimi_cli.utils.path import list_directory
+
+if TYPE_CHECKING:
+    from fastmcp.mcp_config import MCPConfig
+
+
+@dataclass(frozen=True, slots=True, kw_only=True)
+class BuiltinSystemPromptArgs:
+    """Builtin system prompt arguments."""
+
+    KIMI_NOW: str
+    """The current datetime."""
+    KIMI_WORK_DIR: KaosPath
+    """The absolute path of 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."""
+    KIMI_SKILLS: str
+    """Formatted information about available skills."""
+
+
+async def load_agents_md(work_dir: KaosPath) -> str | None:
+    paths = [
+        work_dir / "AGENTS.md",
+        work_dir / "agents.md",
+    ]
+    for path in paths:
+        if await path.is_file():
+            logger.info("Loaded agents.md: {path}", path=path)
+            return (await path.read_text()).strip()
+    logger.info("No AGENTS.md found in {work_dir}", work_dir=work_dir)
+    return None
+
+
+@dataclass(slots=True, kw_only=True)
+class Runtime:
+    """Agent runtime."""
+
+    config: Config
+    llm: LLM | None  # we do not freeze the `Runtime` dataclass because LLM can be changed
+    session: Session
+    builtin_args: BuiltinSystemPromptArgs
+    denwa_renji: DenwaRenji
+    approval: Approval
+    labor_market: LaborMarket
+    environment: Environment
+    skills: dict[str, Skill]
+
+    @staticmethod
+    async def create(
+        config: Config,
+        llm: LLM | None,
+        session: Session,
+        yolo: bool,
+        skills_dir: Path | None = None,
+    ) -> Runtime:
+        ls_output, agents_md, environment = await asyncio.gather(
+            list_directory(session.work_dir),
+            load_agents_md(session.work_dir),
+            Environment.detect(),
+        )
+
+        # Discover and format skills
+        builtin_skills_dir = get_builtin_skills_dir()
+        if skills_dir is None:
+            skills_dir = get_skills_dir()
+            if not skills_dir.is_dir() and (claude_skills_dir := get_claude_skills_dir()).is_dir():
+                skills_dir = claude_skills_dir
+        skills_roots = [builtin_skills_dir, skills_dir]
+        skills = discover_skills_from_roots(skills_roots)
+        skills_by_name = index_skills(skills)
+        logger.info("Discovered {count} skill(s)", count=len(skills))
+        skills_formatted = "\n".join(
+            (
+                f"- {skill.name}\n"
+                f"  - Path: {skill.skill_md_file}\n"
+                f"  - Description: {skill.description}"
+            )
+            for skill in skills
+        )
+
+        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 or "",
+                KIMI_SKILLS=skills_formatted or "No skills found.",
+            ),
+            denwa_renji=DenwaRenji(),
+            approval=Approval(yolo=yolo),
+            labor_market=LaborMarket(),
+            environment=environment,
+            skills=skills_by_name,
+        )
+
+    def copy_for_fixed_subagent(self) -> Runtime:
+        """Clone runtime for fixed subagent."""
+        return Runtime(
+            config=self.config,
+            llm=self.llm,
+            session=self.session,
+            builtin_args=self.builtin_args,
+            denwa_renji=DenwaRenji(),  # subagent must have its own DenwaRenji
+            approval=self.approval,
+            labor_market=LaborMarket(),  # fixed subagent has its own LaborMarket
+            environment=self.environment,
+            skills=self.skills,
+        )
 
+    def copy_for_dynamic_subagent(self) -> Runtime:
+        """Clone runtime for dynamic subagent."""
+        return Runtime(
+            config=self.config,
+            llm=self.llm,
+            session=self.session,
+            builtin_args=self.builtin_args,
+            denwa_renji=DenwaRenji(),  # subagent must have its own DenwaRenji
+            approval=self.approval,
+            labor_market=self.labor_market,  # dynamic subagent shares LaborMarket with main agent
+            environment=self.environment,
+            skills=self.skills,
+        )
 
-class Agent(NamedTuple):
+
+@dataclass(frozen=True, slots=True, kw_only=True)
+class Agent:
     """The loaded agent."""
 
     name: str
     system_prompt: str
     toolset: Toolset
+    runtime: Runtime
+    """Each agent has its own runtime, which should be derived from its main agent."""
+
+
+class LaborMarket:
+    def __init__(self):
+        self.fixed_subagents: dict[str, Agent] = {}
+        self.fixed_subagent_descs: dict[str, str] = {}
+        self.dynamic_subagents: dict[str, Agent] = {}
+
+    @property
+    def subagents(self) -> Mapping[str, Agent]:
+        """Get all subagents in the labor market."""
+        return {**self.fixed_subagents, **self.dynamic_subagents}
+
+    def add_fixed_subagent(self, name: str, agent: Agent, description: str):
+        """Add a fixed subagent."""
+        self.fixed_subagents[name] = agent
+        self.fixed_subagent_descs[name] = description
+
+    def add_dynamic_subagent(self, name: str, agent: Agent):
+        """Add a dynamic subagent."""
+        self.dynamic_subagents[name] = agent
 
 
 async def load_agent(
     agent_file: Path,
     runtime: Runtime,
     *,
-    mcp_configs: list[dict[str, Any]],
+    mcp_configs: list[MCPConfig] | 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.
+        FileNotFoundError: When the agent file is not found.
+        AgentSpecError(KimiCLIException, ValueError): When the agent specification is invalid.
+        InvalidToolError(KimiCLIException, ValueError): When any tool cannot be loaded.
+        MCPConfigError(KimiCLIException, ValueError): When any MCP configuration is invalid.
+        MCPRuntimeError(KimiCLIException, RuntimeError): When any MCP server cannot be connected.
     """
     logger.info("Loading agent: {agent_file}", agent_file=agent_file)
     agent_spec = load_agent_spec(agent_file)
@@ -46,32 +216,56 @@ async def load_agent(
         runtime.builtin_args,
     )
 
+    # load subagents before loading tools because Task tool depends on LaborMarket on initialization
+    for subagent_name, subagent_spec in agent_spec.subagents.items():
+        logger.debug("Loading subagent: {subagent_name}", subagent_name=subagent_name)
+        subagent = await load_agent(
+            subagent_spec.path,
+            runtime.copy_for_fixed_subagent(),
+            mcp_configs=mcp_configs,
+        )
+        runtime.labor_market.add_fixed_subagent(subagent_name, subagent, subagent_spec.description)
+
+    toolset = KimiToolset()
     tool_deps = {
-        ResolvedAgentSpec: agent_spec,
+        KimiToolset: toolset,
         Runtime: runtime,
+        # TODO: remove all the following dependencies and use Runtime instead
         Config: runtime.config,
         BuiltinSystemPromptArgs: runtime.builtin_args,
         Session: runtime.session,
         DenwaRenji: runtime.denwa_renji,
         Approval: runtime.approval,
+        LaborMarket: runtime.labor_market,
+        Environment: runtime.environment,
     }
     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}")
+    toolset.load_tools(tools, tool_deps)
 
-    assert isinstance(toolset, CustomToolset)
     if mcp_configs:
-        await _load_mcp_tools(toolset, mcp_configs)
+        validated_mcp_configs: list[MCPConfig] = []
+        if mcp_configs:
+            from fastmcp.mcp_config import MCPConfig
+
+            for mcp_config in mcp_configs:
+                try:
+                    validated_mcp_configs.append(
+                        mcp_config
+                        if isinstance(mcp_config, MCPConfig)
+                        else MCPConfig.model_validate(mcp_config)
+                    )
+                except pydantic.ValidationError as e:
+                    raise MCPConfigError(f"Invalid MCP config: {e}") from e
+        await toolset.load_mcp_tools(validated_mcp_configs, runtime)
 
     return Agent(
         name=agent_spec.name,
         system_prompt=system_prompt,
         toolset=toolset,
+        runtime=runtime,
     )
 
 
@@ -85,70 +279,4 @@ def _load_system_prompt(
         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.
-    """
-    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
+    return string.Template(system_prompt).substitute(asdict(builtin_args), **args)

kimi_cli/soul/approval.py

@@ -1,13 +1,33 @@
+from __future__ import annotations
+
 import asyncio
+import uuid
+from dataclasses import dataclass
+from typing import Literal
 
 from kimi_cli.soul.toolset import get_current_tool_call_or_none
+from kimi_cli.utils.aioqueue import Queue
 from kimi_cli.utils.logging import logger
-from kimi_cli.wire.message import ApprovalRequest, ApprovalResponse
+from kimi_cli.wire.types import DisplayBlock
+
+
+@dataclass(frozen=True, slots=True, kw_only=True)
+class Request:
+    id: str
+    tool_call_id: str
+    sender: str
+    action: str
+    description: str
+    display: list[DisplayBlock]
+
+
+type Response = Literal["approve", "approve_for_session", "reject"]
 
 
 class Approval:
     def __init__(self, yolo: bool = False):
-        self._request_queue = asyncio.Queue[ApprovalRequest]()
+        self._request_queue = Queue[Request]()
+        self._requests: dict[str, tuple[Request, asyncio.Future[bool]]] = {}
         self._yolo = yolo
         self._auto_approve_actions: set[str] = set()  # TODO: persist across sessions
         """Set of action names that should automatically be approved."""
@@ -15,7 +35,16 @@ class Approval:
     def set_yolo(self, yolo: bool) -> None:
         self._yolo = yolo
 
-    async def request(self, sender: str, action: str, description: str) -> bool:
+    def is_yolo(self) -> bool:
+        return self._yolo
+
+    async def request(
+        self,
+        sender: str,
+        action: str,
+        description: str,
+        display: list[DisplayBlock] | None = None,
+    ) -> bool:
         """
         Request approval for the given action. Intended to be called by tools.
 
@@ -48,21 +77,61 @@ class Approval:
         if action in self._auto_approve_actions:
             return True
 
-        request = ApprovalRequest(tool_call.id, sender, action, description)
+        request = Request(
+            id=str(uuid.uuid4()),
+            tool_call_id=tool_call.id,
+            sender=sender,
+            action=action,
+            description=description,
+            display=display or [],
+        )
+        approved_future = asyncio.Future[bool]()
         self._request_queue.put_nowait(request)
-        response = await request.wait()
-        logger.debug("Received approval response: {response}", response=response)
-        match response:
-            case ApprovalResponse.APPROVE:
-                return True
-            case ApprovalResponse.APPROVE_FOR_SESSION:
-                self._auto_approve_actions.add(action)
-                return True
-            case ApprovalResponse.REJECT:
-                return False
-
-    async def fetch_request(self) -> ApprovalRequest:
+        self._requests[request.id] = (request, approved_future)
+        return await approved_future
+
+    async def fetch_request(self) -> Request:
         """
         Fetch an approval request from the queue. Intended to be called by the soul.
         """
-        return await self._request_queue.get()
+        while True:
+            request = await self._request_queue.get()
+            if request.action in self._auto_approve_actions:
+                # the action is not auto-approved when the request was created, but now it should be
+                logger.debug(
+                    "Auto-approving previously requested action: {action}", action=request.action
+                )
+                self.resolve_request(request.id, "approve")
+                continue
+
+            return request
+
+    def resolve_request(self, request_id: str, response: Response) -> None:
+        """
+        Resolve an approval request with the given response. Intended to be called by the soul.
+
+        Args:
+            request_id (str): The ID of the request to resolve.
+            response (Response): The response to the request.
+
+        Raises:
+            KeyError: If there is no pending request with the given ID.
+        """
+        request_tuple = self._requests.pop(request_id, None)
+        if request_tuple is None:
+            raise KeyError(f"No pending request with ID {request_id}")
+        request, future = request_tuple
+
+        logger.debug(
+            "Received approval response for request {request_id}: {response}",
+            request_id=request_id,
+            response=response,
+        )
+        match response:
+            case "approve":
+                future.set_result(True)
+            case "approve_for_session":
+                self._auto_approve_actions.add(request.action)
+                future.set_result(True)
+            case "reject":
+                future.set_result(False)

kimi_cli/soul/compaction.py

@@ -1,14 +1,17 @@
+from __future__ import annotations
+
 from collections.abc import Sequence
-from string import Template
-from typing import TYPE_CHECKING, Protocol, runtime_checkable
+from typing import TYPE_CHECKING, NamedTuple, Protocol, runtime_checkable
 
-from kosong.base import generate
-from kosong.base.message import ContentPart, Message, TextPart
+import kosong
+from kosong.message import Message
+from kosong.tooling.empty import EmptyToolset
 
 import kimi_cli.prompts as prompts
 from kimi_cli.llm import LLM
 from kimi_cli.soul.message import system
 from kimi_cli.utils.logging import logger
+from kimi_cli.wire.types import ContentPart, TextPart, ThinkPart
 
 
 @runtime_checkable
@@ -30,76 +33,84 @@ 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:
-            return history
+if TYPE_CHECKING:
 
-        preserve_start_index = len(history)
-        n_preserved = 0
-        for index in range(len(history) - 1, -1, -1):
-            if history[index].role in {"user", "assistant"}:
-                n_preserved += 1
-                if n_preserved == self.MAX_PRESERVED_MESSAGES:
-                    preserve_start_index = index
-                    break
+    def type_check(simple: SimpleCompaction):
+        _: Compaction = simple
 
-        if n_preserved < self.MAX_PRESERVED_MESSAGES:
-            return history
 
-        to_compact = history[:preserve_start_index]
-        to_preserve = history[preserve_start_index:]
+class SimpleCompaction:
+    def __init__(self, max_preserved_messages: int = 2) -> None:
+        self.max_preserved_messages = max_preserved_messages
 
-        if not to_compact:
-            # Let's hope this won't exceed the context size limit
+    async def compact(self, messages: Sequence[Message], llm: LLM) -> Sequence[Message]:
+        compact_message, to_preserve = self.prepare(messages)
+        if compact_message is None:
             return to_preserve
 
-        # Convert history to string for the compact prompt
-        history_text = "\n\n".join(
-            f"## Message {i + 1}\nRole: {msg.role}\nContent: {msg.content}"
-            for i, msg in enumerate(to_compact)
-        )
-
-        # Build the compact prompt using string template
-        compact_template = Template(prompts.COMPACT)
-        compact_prompt = compact_template.substitute(CONTEXT=history_text)
-
-        # Create input message for compaction
-        compact_message = Message(role="user", content=compact_prompt)
-
-        # Call generate to get the compacted context
+        # Call kosong.step to get the compacted context
         # TODO: set max completion tokens
         logger.debug("Compacting context...")
-        compacted_msg, usage = await generate(
+        result = await kosong.step(
             chat_provider=llm.chat_provider,
             system_prompt="You are a helpful assistant that compacts conversation context.",
-            tools=[],
+            toolset=EmptyToolset(),
             history=[compact_message],
         )
-        if usage:
+        if result.usage:
             logger.debug(
                 "Compaction used {input} input tokens and {output} output tokens",
-                input=usage.input,
-                output=usage.output,
+                input=result.usage.input,
+                output=result.usage.output,
             )
 
         content: list[ContentPart] = [
             system("Previous context has been compacted. Here is the compaction output:")
         ]
-        content.extend(
-            [TextPart(text=compacted_msg.content)]
-            if isinstance(compacted_msg.content, str)
-            else compacted_msg.content
-        )
-        compacted_messages: list[Message] = [Message(role="assistant", content=content)]
+        compacted_msg = result.message
+
+        # drop thinking parts if any
+        content.extend(part for part in compacted_msg.content if not isinstance(part, ThinkPart))
+        compacted_messages: list[Message] = [Message(role="user", content=content)]
         compacted_messages.extend(to_preserve)
         return compacted_messages
 
+    class PrepareResult(NamedTuple):
+        compact_message: Message | None
+        to_preserve: Sequence[Message]
 
-if TYPE_CHECKING:
+    def prepare(self, messages: Sequence[Message]) -> PrepareResult:
+        if not messages or self.max_preserved_messages <= 0:
+            return self.PrepareResult(compact_message=None, to_preserve=messages)
 
-    def type_check(simple: SimpleCompaction):
-        _: Compaction = simple
+        history = list(messages)
+        preserve_start_index = len(history)
+        n_preserved = 0
+        for index in range(len(history) - 1, -1, -1):
+            if history[index].role in {"user", "assistant"}:
+                n_preserved += 1
+                if n_preserved == self.max_preserved_messages:
+                    preserve_start_index = index
+                    break
+
+        if n_preserved < self.max_preserved_messages:
+            return self.PrepareResult(compact_message=None, to_preserve=messages)
+
+        to_compact = history[:preserve_start_index]
+        to_preserve = history[preserve_start_index:]
+
+        if not to_compact:
+            # Let's hope this won't exceed the context size limit
+            return self.PrepareResult(compact_message=None, to_preserve=to_preserve)
+
+        # Create input message for compaction
+        compact_message = Message(role="user", content=[])
+        for i, msg in enumerate(to_compact):
+            compact_message.content.append(
+                TextPart(text=f"## Message {i + 1}\nRole: {msg.role}\nContent:\n")
+            )
+            compact_message.content.extend(
+                part for part in msg.content if not isinstance(part, ThinkPart)
+            )
+        compact_message.content.append(TextPart(text="\n" + prompts.COMPACT))
+        return self.PrepareResult(compact_message=compact_message, to_preserve=to_preserve)