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

kimi_cli/soul/approval.py

@@ -1,25 +1,26 @@
 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:
     def __init__(self, yolo: bool = False):
         self._request_queue = asyncio.Queue[ApprovalRequest]()
         self._yolo = yolo
-        self._auto_approve_actions = set()  # TODO: persist across sessions
+        self._auto_approve_actions: set[str] = set()  # TODO: persist across sessions
         """Set of action names that should automatically be approved."""
 
     def set_yolo(self, yolo: bool) -> None:
         self._yolo = yolo
 
-    async def request(self, action: str, description: str) -> bool:
+    async def request(self, sender: str, action: str, description: str) -> bool:
         """
         Request approval for the given action. Intended to be called by tools.
 
         Args:
+            sender (str): The name of the sender.
             action (str): The action to request approval for.
                 This is used to identify the action for auto-approval.
             description (str): The description of the action. This is used to display to the user.
@@ -47,7 +48,7 @@ class Approval:
         if action in self._auto_approve_actions:
             return True
 
-        request = ApprovalRequest(tool_call.id, action, description)
+        request = ApprovalRequest(tool_call.id, sender, action, description)
         self._request_queue.put_nowait(request)
         response = await request.wait()
         logger.debug("Received approval response: {response}", response=response)
@@ -59,8 +60,6 @@ class Approval:
                 return True
             case ApprovalResponse.REJECT:
                 return False
-            case _:
-                raise ValueError(f"Unknown approval response: {response}")
 
     async def fetch_request(self) -> ApprovalRequest:
         """

kimi_cli/soul/compaction.py

@@ -0,0 +1,106 @@
+from collections.abc import Sequence
+from string import Template
+from typing import TYPE_CHECKING, Protocol, runtime_checkable
+
+from kosong import generate
+from kosong.message import ContentPart, Message, TextPart
+
+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
+
+
+@runtime_checkable
+class Compaction(Protocol):
+    async def compact(self, messages: Sequence[Message], llm: LLM) -> Sequence[Message]:
+        """
+        Compact a sequence of messages into a new sequence of messages.
+
+        Args:
+            messages (Sequence[Message]): The messages to compact.
+            llm (LLM): The LLM to use for compaction.
+
+        Returns:
+            Sequence[Message]: The compacted messages.
+
+        Raises:
+            ChatProviderError: When the chat provider returns an error.
+        """
+        ...
+
+
+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
+
+        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 history
+
+        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 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
+        # TODO: set max completion tokens
+        logger.debug("Compacting context...")
+        result = await generate(
+            chat_provider=llm.chat_provider,
+            system_prompt="You are a helpful assistant that compacts conversation context.",
+            tools=[],
+            history=[compact_message],
+        )
+        if result.usage:
+            logger.debug(
+                "Compaction used {input} input tokens and {output} output tokens",
+                input=result.usage.input,
+                output=result.usage.output,
+            )
+
+        content: list[ContentPart] = [
+            system("Previous context has been compacted. Here is the compaction output:")
+        ]
+        compacted_msg = result.message
+        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_messages.extend(to_preserve)
+        return compacted_messages
+
+
+if TYPE_CHECKING:
+
+    def type_check(simple: SimpleCompaction):
+        _: Compaction = simple

kimi_cli/soul/context.py

@@ -4,7 +4,7 @@ from pathlib import Path
 
 import aiofiles
 import aiofiles.os
-from kosong.base.message import Message
+from kosong.message import Message
 
 from kimi_cli.soul.message import system
 from kimi_cli.utils.logging import logger

kimi_cli/soul/kimisoul.py

@@ -1,55 +1,86 @@
 import asyncio
+from collections.abc import Sequence
+from functools import partial
+from typing import TYPE_CHECKING
 
 import kosong
 import tenacity
 from kosong import StepResult
-from kosong.base.message import Message
 from kosong.chat_provider import (
     APIConnectionError,
+    APIEmptyResponseError,
     APIStatusError,
     APITimeoutError,
     ChatProviderError,
+    ThinkingEffort,
 )
+from kosong.message import ContentPart, ImageURLPart, Message
 from kosong.tooling import ToolResult
 from tenacity import RetryCallState, retry_if_exception, stop_after_attempt, wait_exponential_jitter
 
-from kimi_cli.agent import Agent, AgentGlobals
-from kimi_cli.config import LoopControl
-from kimi_cli.soul import LLMNotSet, MaxStepsReached, Soul, StatusSnapshot
+from kimi_cli.llm import ModelCapability
+from kimi_cli.soul import (
+    LLMNotSet,
+    LLMNotSupported,
+    MaxStepsReached,
+    Soul,
+    StatusSnapshot,
+    wire_send,
+)
+from kimi_cli.soul.agent import Agent
+from kimi_cli.soul.compaction import SimpleCompaction
 from kimi_cli.soul.context import Context
 from kimi_cli.soul.message import system, tool_result_to_messages
-from kimi_cli.soul.wire import StatusUpdate, StepBegin, StepInterrupted, Wire, current_wire
+from kimi_cli.soul.runtime import Runtime
 from kimi_cli.tools.dmail import NAME as SendDMail_NAME
 from kimi_cli.tools.utils import ToolRejectedError
 from kimi_cli.utils.logging import logger
+from kimi_cli.wire.message import (
+    CompactionBegin,
+    CompactionEnd,
+    StatusUpdate,
+    StepBegin,
+    StepInterrupted,
+)
+
+if TYPE_CHECKING:
+
+    def type_check(soul: "KimiSoul"):
+        _: Soul = soul
+
 
+RESERVED_TOKENS = 50_000
 
-class KimiSoul:
+
+class KimiSoul(Soul):
     """The soul of Kimi CLI."""
 
     def __init__(
         self,
         agent: Agent,
-        agent_globals: AgentGlobals,
+        runtime: Runtime,
         *,
         context: Context,
-        loop_control: LoopControl,
     ):
         """
         Initialize the soul.
 
         Args:
             agent (Agent): The agent to run.
-            agent_globals (AgentGlobals): Global states and parameters.
+            runtime (Runtime): Runtime parameters and states.
             context (Context): The context of the agent.
-            loop_control (LoopControl): The control parameters for the agent loop.
         """
         self._agent = agent
-        self._agent_globals = agent_globals
-        self._denwa_renji = agent_globals.denwa_renji
-        self._approval = agent_globals.approval
+        self._runtime = runtime
+        self._denwa_renji = runtime.denwa_renji
+        self._approval = runtime.approval
         self._context = context
-        self._loop_control = loop_control
+        self._loop_control = runtime.config.loop_control
+        self._compaction = SimpleCompaction()  # TODO: maybe configurable and composable
+        self._reserved_tokens = RESERVED_TOKENS
+        if self._runtime.llm is not None:
+            assert self._reserved_tokens <= self._runtime.llm.max_context_size
+        self._thinking_effort: ThinkingEffort = "off"
 
         for tool in agent.toolset.tools:
             if tool.name == SendDMail_NAME:
@@ -63,62 +94,106 @@ class KimiSoul:
         return self._agent.name
 
     @property
-    def model(self) -> str:
-        return self._agent_globals.llm.chat_provider.model_name if self._agent_globals.llm else ""
+    def model_name(self) -> str:
+        return self._runtime.llm.chat_provider.model_name if self._runtime.llm else ""
+
+    @property
+    def model_capabilities(self) -> set[ModelCapability] | None:
+        if self._runtime.llm is None:
+            return None
+        return self._runtime.llm.capabilities
 
     @property
     def status(self) -> StatusSnapshot:
         return StatusSnapshot(context_usage=self._context_usage)
 
+    @property
+    def context(self) -> Context:
+        return self._context
+
     @property
     def _context_usage(self) -> float:
-        if self._agent_globals.llm is not None:
-            return self._context.token_count / self._agent_globals.llm.max_context_size
+        if self._runtime.llm is not None:
+            return self._context.token_count / self._runtime.llm.max_context_size
         return 0.0
 
+    @property
+    def thinking(self) -> bool:
+        """Whether thinking mode is enabled."""
+        return self._thinking_effort != "off"
+
+    def set_thinking(self, enabled: bool) -> None:
+        """
+        Enable/disable thinking mode for the soul.
+
+        Raises:
+            LLMNotSet: When the LLM is not set.
+            LLMNotSupported: When the LLM does not support thinking mode.
+        """
+        if self._runtime.llm is None:
+            raise LLMNotSet()
+        if enabled and "thinking" not in self._runtime.llm.capabilities:
+            raise LLMNotSupported(self._runtime.llm, ["thinking"])
+        self._thinking_effort = "high" if enabled else "off"
+
     async def _checkpoint(self):
         await self._context.checkpoint(self._checkpoint_with_user_message)
 
-    async def run(self, user_input: str, wire: Wire):
-        if self._agent_globals.llm is None:
+    async def run(self, user_input: str | list[ContentPart]):
+        if self._runtime.llm is None:
             raise LLMNotSet()
 
+        if (
+            isinstance(user_input, list)
+            and any(isinstance(part, ImageURLPart) for part in user_input)
+            and "image_in" not in self._runtime.llm.capabilities
+        ):
+            raise LLMNotSupported(self._runtime.llm, ["image_in"])
+
         await self._checkpoint()  # this creates the checkpoint 0 on first run
         await self._context.append_message(Message(role="user", content=user_input))
         logger.debug("Appended user message to context")
-        wire_token = current_wire.set(wire)
-        try:
-            await self._agent_loop(wire)
-        finally:
-            current_wire.reset(wire_token)
+        await self._agent_loop()
 
-    async def _agent_loop(self, wire: Wire):
+    async def _agent_loop(self):
         """The main agent loop for one run."""
+        assert self._runtime.llm is not None
 
         async def _pipe_approval_to_wire():
             while True:
                 request = await self._approval.fetch_request()
-                wire.send(request)
+                wire_send(request)
 
         step_no = 1
         while True:
-            wire.send(StepBegin(step_no))
+            wire_send(StepBegin(step_no))
             approval_task = asyncio.create_task(_pipe_approval_to_wire())
             # FIXME: It's possible that a subagent's approval task steals approval request
             # from the main agent. We must ensure that the Task tool will redirect them
             # to the main wire. See `_SubWire` for more details. Later we need to figure
             # out a better solution.
             try:
+                # compact the context if needed
+                if (
+                    self._context.token_count + self._reserved_tokens
+                    >= self._runtime.llm.max_context_size
+                ):
+                    logger.info("Context too long, compacting...")
+                    wire_send(CompactionBegin())
+                    await self.compact_context()
+                    wire_send(CompactionEnd())
+
+                logger.debug("Beginning step {step_no}", step_no=step_no)
                 await self._checkpoint()
                 self._denwa_renji.set_n_checkpoints(self._context.n_checkpoints)
-                finished = await self._step(wire)
+                finished = await self._step()
             except BackToTheFuture as e:
                 await self._context.revert_to(e.checkpoint_id)
                 await self._checkpoint()
-                await self._context.append_message(e.message)
+                await self._context.append_message(e.messages)
                 continue
             except (ChatProviderError, asyncio.CancelledError):
-                wire.send(StepInterrupted())
+                wire_send(StepInterrupted())
                 # break the agent loop
                 raise
             finally:
@@ -131,34 +206,15 @@ class KimiSoul:
             if step_no > self._loop_control.max_steps_per_run:
                 raise MaxStepsReached(self._loop_control.max_steps_per_run)
 
-    async def _step(self, wire: Wire) -> bool:
+    async def _step(self) -> bool:
         """Run an single step and return whether the run should be stopped."""
         # already checked in `run`
-        assert self._agent_globals.llm is not None
-        chat_provider = self._agent_globals.llm.chat_provider
-
-        def _is_retryable_error(exception: BaseException) -> bool:
-            if isinstance(exception, (APIConnectionError, APITimeoutError)):
-                return True
-            return isinstance(exception, APIStatusError) and exception.status_code in (
-                429,  # Too Many Requests
-                500,  # Internal Server Error
-                502,  # Bad Gateway
-                503,  # Service Unavailable
-            )
-
-        def _retry_log(retry_state: RetryCallState):
-            logger.info(
-                "Retrying step for the {n} time. Waiting {sleep} seconds.",
-                n=retry_state.attempt_number,
-                sleep=retry_state.next_action.sleep
-                if retry_state.next_action is not None
-                else "unknown",
-            )
+        assert self._runtime.llm is not None
+        chat_provider = self._runtime.llm.chat_provider
 
         @tenacity.retry(
-            retry=retry_if_exception(_is_retryable_error),
-            before_sleep=_retry_log,
+            retry=retry_if_exception(self._is_retryable_error),
+            before_sleep=partial(self._retry_log, "step"),
             wait=wait_exponential_jitter(initial=0.3, max=5, jitter=0.5),
             stop=stop_after_attempt(self._loop_control.max_retries_per_step),
             reraise=True,
@@ -166,12 +222,12 @@ class KimiSoul:
         async def _kosong_step_with_retry() -> StepResult:
             # run an LLM step (may be interrupted)
             return await kosong.step(
-                chat_provider,
+                chat_provider.with_thinking(self._thinking_effort),
                 self._agent.system_prompt,
                 self._agent.toolset,
                 self._context.history,
-                on_message_part=wire.send,
-                on_tool_result=wire.send,
+                on_message_part=wire_send,
+                on_tool_result=wire_send,
             )
 
         result = await _kosong_step_with_retry()
@@ -179,7 +235,7 @@ class KimiSoul:
         if result.usage is not None:
             # mark the token count for the context before the step
             await self._context.update_token_count(result.usage.input)
-            wire.send(StatusUpdate(status=self.status))
+            wire_send(StatusUpdate(status=self.status))
 
         # wait for all tool results (may be interrupted)
         results = await result.tool_results()
@@ -202,19 +258,21 @@ class KimiSoul:
             # raise to let the main loop take us back to the future
             raise BackToTheFuture(
                 dmail.checkpoint_id,
-                Message(
-                    role="user",
-                    content=[
-                        system(
-                            "You just got a D-Mail from your future self. "
-                            "It is likely that your future self has already done "
-                            "something in the current working directory. Please read "
-                            "the D-Mail and decide what to do next. You MUST NEVER "
-                            "mention to the user about this information. "
-                            f"D-Mail content:\n\n{dmail.message.strip()}"
-                        )
-                    ],
-                ),
+                [
+                    Message(
+                        role="user",
+                        content=[
+                            system(
+                                "You just got a D-Mail from your future self. "
+                                "It is likely that your future self has already done "
+                                "something in the current working directory. Please read "
+                                "the D-Mail and decide what to do next. You MUST NEVER "
+                                "mention to the user about this information. "
+                                f"D-Mail content:\n\n{dmail.message.strip()}"
+                            )
+                        ],
+                    )
+                ],
             )
 
         return not result.tool_calls
@@ -230,6 +288,54 @@ class KimiSoul:
             logger.debug("Appending tool result to context: {tool_result}", tool_result=tool_result)
             await self._context.append_message(tool_result_to_messages(tool_result))
 
+    async def compact_context(self) -> None:
+        """
+        Compact the context.
+
+        Raises:
+            LLMNotSet: When the LLM is not set.
+            ChatProviderError: When the chat provider returns an error.
+        """
+
+        @tenacity.retry(
+            retry=retry_if_exception(self._is_retryable_error),
+            before_sleep=partial(self._retry_log, "compaction"),
+            wait=wait_exponential_jitter(initial=0.3, max=5, jitter=0.5),
+            stop=stop_after_attempt(self._loop_control.max_retries_per_step),
+            reraise=True,
+        )
+        async def _compact_with_retry() -> Sequence[Message]:
+            if self._runtime.llm is None:
+                raise LLMNotSet()
+            return await self._compaction.compact(self._context.history, self._runtime.llm)
+
+        compacted_messages = await _compact_with_retry()
+        await self._context.revert_to(0)
+        await self._checkpoint()
+        await self._context.append_message(compacted_messages)
+
+    @staticmethod
+    def _is_retryable_error(exception: BaseException) -> bool:
+        if isinstance(exception, (APIConnectionError, APITimeoutError, APIEmptyResponseError)):
+            return True
+        return isinstance(exception, APIStatusError) and exception.status_code in (
+            429,  # Too Many Requests
+            500,  # Internal Server Error
+            502,  # Bad Gateway
+            503,  # Service Unavailable
+        )
+
+    @staticmethod
+    def _retry_log(name: str, retry_state: RetryCallState):
+        logger.info(
+            "Retrying {name} for the {n} time. Waiting {sleep} seconds.",
+            name=name,
+            n=retry_state.attempt_number,
+            sleep=retry_state.next_action.sleep
+            if retry_state.next_action is not None
+            else "unknown",
+        )
+
 
 class BackToTheFuture(Exception):
     """
@@ -237,12 +343,6 @@ class BackToTheFuture(Exception):
     The main agent loop should catch this exception and handle it.
     """
 
-    def __init__(self, checkpoint_id: int, message: Message):
+    def __init__(self, checkpoint_id: int, messages: Sequence[Message]):
         self.checkpoint_id = checkpoint_id
-        self.message = message
-
-
-def __static_type_check(
-    kimi_soul: KimiSoul,
-):
-    _: Soul = kimi_soul
+        self.messages = messages

kimi_cli/soul/message.py

@@ -1,4 +1,4 @@
-from kosong.base.message import ContentPart, Message, TextPart
+from kosong.message import ContentPart, Message, TextPart
 from kosong.tooling import ToolError, ToolOk, ToolResult
 from kosong.tooling.error import ToolRuntimeError
 
@@ -14,7 +14,7 @@ def tool_result_to_messages(tool_result: ToolResult) -> list[Message]:
         message = tool_result.result.message
         if isinstance(tool_result.result, ToolRuntimeError):
             message += "\nThis is an unexpected error and the tool is probably not working."
-        content = [system(message)]
+        content: list[ContentPart] = [system(f"ERROR: {message}")]
         if tool_result.result.output:
             content.append(TextPart(text=tool_result.result.output))
         return [
@@ -26,8 +26,8 @@ def tool_result_to_messages(tool_result: ToolResult) -> list[Message]:
         ]
 
     content = tool_ok_to_message_content(tool_result.result)
-    text_parts = []
-    non_text_parts = []
+    text_parts: list[ContentPart] = []
+    non_text_parts: list[ContentPart] = []
     for part in content:
         if isinstance(part, TextPart):
             text_parts.append(part)
@@ -60,7 +60,7 @@ def tool_result_to_messages(tool_result: ToolResult) -> list[Message]:
 
 def tool_ok_to_message_content(result: ToolOk) -> list[ContentPart]:
     """Convert a tool return value to a list of message content parts."""
-    content = []
+    content: list[ContentPart] = []
     if result.message:
         content.append(system(result.message))
     match output := result.output:
@@ -70,7 +70,7 @@ def tool_ok_to_message_content(result: ToolOk) -> list[ContentPart]:
         case ContentPart():
             content.append(output)
         case _:
-            content.extend(list(output))
+            content.extend(output)
     if not content:
         content.append(system("Tool output is empty."))
     return content