kimi-cli 0.35__py3-none-any.whl

kimi_cli/ui/shell/__init__.py

@@ -0,0 +1,326 @@
+import asyncio
+import signal
+from collections.abc import Awaitable, Coroutine
+from typing import Any
+
+from kosong.base.message import ContentPart, TextPart, ToolCall, ToolCallPart
+from kosong.chat_provider import APIStatusError, ChatProviderError
+from kosong.tooling import ToolResult
+from rich.console import Group, RenderableType
+from rich.panel import Panel
+from rich.table import Table
+from rich.text import Text
+
+from kimi_cli.soul import LLMNotSet, MaxStepsReached, Soul
+from kimi_cli.soul.kimisoul import KimiSoul
+from kimi_cli.soul.wire import (
+    ApprovalRequest,
+    ApprovalResponse,
+    StatusUpdate,
+    StepBegin,
+    StepInterrupted,
+    Wire,
+)
+from kimi_cli.ui import RunCancelled, run_soul
+from kimi_cli.ui.shell.console import console
+from kimi_cli.ui.shell.liveview import StepLiveView
+from kimi_cli.ui.shell.metacmd import get_meta_command
+from kimi_cli.ui.shell.prompt import CustomPromptSession, PromptMode, toast
+from kimi_cli.ui.shell.update import UpdateResult, do_update
+from kimi_cli.utils.logging import logger
+
+
+class Reload(Exception):
+    """Reload configuration."""
+
+    pass
+
+
+class ShellApp:
+    def __init__(self, soul: Soul, welcome_info: dict[str, str] | None = None):
+        self.soul = soul
+        self.welcome_info = welcome_info or {}
+        self._background_tasks: set[asyncio.Task[Any]] = set()
+
+    async def run(self, command: str | None = None) -> bool:
+        if command is not None:
+            # run single command and exit
+            logger.info("Running agent with command: {command}", command=command)
+            return await self._run(command)
+
+        self._start_auto_update_task()
+
+        _print_welcome_info(self.soul.name or "Kimi CLI", self.soul.model, self.welcome_info)
+
+        with CustomPromptSession(lambda: self.soul.status) as prompt_session:
+            while True:
+                try:
+                    user_input = await prompt_session.prompt()
+                except KeyboardInterrupt:
+                    logger.debug("Exiting by KeyboardInterrupt")
+                    console.print("[grey50]Tip: press Ctrl-D or send 'exit' to quit[/grey50]")
+                    continue
+                except EOFError:
+                    logger.debug("Exiting by EOF")
+                    console.print("Bye!")
+                    break
+
+                if not user_input:
+                    logger.debug("Got empty input, skipping")
+                    continue
+                logger.debug("Got user input: {user_input}", user_input=user_input)
+
+                if user_input.command in ["exit", "quit", "/exit", "/quit"]:
+                    logger.debug("Exiting by meta command")
+                    console.print("Bye!")
+                    break
+
+                if user_input.mode == PromptMode.SHELL:
+                    await self._run_shell_command(user_input.command)
+                    continue
+
+                command = user_input.command
+                if command.startswith("/"):
+                    logger.debug("Running meta command: {command}", command=command)
+                    await self._run_meta_command(command[1:])
+                    continue
+
+                logger.info("Running agent command: {command}", command=command)
+                await self._run(command)
+
+        return True
+
+    async def _run_shell_command(self, command: str) -> None:
+        """Run a shell command in foreground."""
+        if not command.strip():
+            return
+
+        logger.info("Running shell command: {cmd}", cmd=command)
+        loop = asyncio.get_running_loop()
+        try:
+            # TODO: For the sake of simplicity, we now use `create_subprocess_shell`.
+            # Later we should consider making this behave like a real shell.
+            proc = await asyncio.create_subprocess_shell(command)
+
+            def _handler():
+                logger.debug("SIGINT received.")
+                proc.terminate()
+
+            loop.add_signal_handler(signal.SIGINT, _handler)
+
+            await proc.wait()
+        except Exception as e:
+            logger.exception("Failed to run shell command:")
+            console.print(f"[red]Failed to run shell command: {e}[/red]")
+        finally:
+            loop.remove_signal_handler(signal.SIGINT)
+
+    async def _run(self, command: str) -> bool:
+        """
+        Run the soul and handle any known exceptions.
+
+        Returns:
+            bool: Whether the run is successful.
+        """
+        cancel_event = asyncio.Event()
+
+        def _handler():
+            logger.debug("SIGINT received.")
+            cancel_event.set()
+
+        loop = asyncio.get_running_loop()
+        loop.add_signal_handler(signal.SIGINT, _handler)
+
+        try:
+            await run_soul(self.soul, command, self._visualize, cancel_event)
+            return True
+        except LLMNotSet:
+            logger.error("LLM not set")
+            console.print("[red]LLM not set, send /setup to configure[/red]")
+        except ChatProviderError as e:
+            logger.exception("LLM provider error:")
+            if isinstance(e, APIStatusError) and e.status_code == 401:
+                console.print("[red]Authorization failed, please check your API key[/red]")
+            elif isinstance(e, APIStatusError) and e.status_code == 402:
+                console.print("[red]Membership expired, please renew your plan[/red]")
+            elif isinstance(e, APIStatusError) and e.status_code == 403:
+                console.print("[red]Quota exceeded, please upgrade your plan or retry later[/red]")
+            else:
+                console.print(f"[red]LLM provider error: {e}[/red]")
+        except MaxStepsReached as e:
+            logger.warning("Max steps reached: {n_steps}", n_steps=e.n_steps)
+            console.print(f"[yellow]Max steps reached: {e.n_steps}[/yellow]")
+        except RunCancelled:
+            logger.info("Cancelled by user")
+            console.print("[red]Interrupted by user[/red]")
+        except Reload:
+            # just propagate
+            raise
+        except BaseException as e:
+            logger.exception("Unknown error:")
+            console.print(f"[red]Unknown error: {e}[/red]")
+            raise  # re-raise unknown error
+        finally:
+            loop.remove_signal_handler(signal.SIGINT)
+        return False
+
+    def _start_auto_update_task(self) -> None:
+        self._add_background_task(self._auto_update_background())
+
+    async def _auto_update_background(self) -> None:
+        toast("checking for updates...", duration=2.0)
+        result = await do_update(print=False, check_only=True)
+        if result == UpdateResult.UPDATE_AVAILABLE:
+            while True:
+                toast("new version found, run `uv tool upgrade ikimi` to upgrade", duration=30.0)
+                await asyncio.sleep(60.0)
+        elif result == UpdateResult.UPDATED:
+            toast("auto updated, restart to use the new version", duration=5.0)
+
+    def _add_background_task(self, coro: Coroutine[Any, Any, Any]) -> asyncio.Task[Any]:
+        task = asyncio.create_task(coro)
+        self._background_tasks.add(task)
+
+        def _cleanup(t: asyncio.Task[Any]) -> None:
+            self._background_tasks.discard(t)
+            try:
+                t.result()
+            except asyncio.CancelledError:
+                pass
+            except Exception:
+                logger.exception("Background task failed:")
+
+        task.add_done_callback(_cleanup)
+        return task
+
+    async def _visualize(self, wire: Wire):
+        """
+        A loop to consume agent events and visualize the agent behavior.
+        This loop never raise any exception except asyncio.CancelledError.
+        """
+        try:
+            # expect a StepBegin
+            assert isinstance(await wire.receive(), StepBegin)
+
+            while True:
+                # spin the moon at the beginning of each step
+                with console.status("", spinner="moon"):
+                    msg = await wire.receive()
+
+                with StepLiveView(self.soul.status) as step:
+                    # visualization loop for one step
+                    while True:
+                        match msg:
+                            case TextPart(text=text):
+                                step.append_text(text)
+                            case ContentPart():
+                                # TODO: support more content parts
+                                step.append_text(f"[{msg.__class__.__name__}]")
+                            case ToolCall():
+                                step.append_tool_call(msg)
+                            case ToolCallPart():
+                                step.append_tool_call_part(msg)
+                            case ToolResult():
+                                step.append_tool_result(msg)
+                            case ApprovalRequest():
+                                msg.resolve(ApprovalResponse.APPROVE)
+                                # TODO(approval): handle approval request
+                            case StatusUpdate(status=status):
+                                step.update_status(status)
+                            case _:
+                                break  # break the step loop
+                        msg = await wire.receive()
+
+                    # cleanup the step live view
+                    if isinstance(msg, StepInterrupted):
+                        step.interrupt()
+                    else:
+                        step.finish()
+
+                if isinstance(msg, StepInterrupted):
+                    # for StepInterrupted, the visualization loop should end immediately
+                    break
+
+                assert isinstance(msg, StepBegin), "expect a StepBegin"
+                # start a new step
+        except asyncio.QueueShutDown:
+            logger.debug("Visualization loop shutting down")
+
+    async def _run_meta_command(self, command_str: str):
+        parts = command_str.split(" ")
+        command_name = parts[0]
+        command_args = parts[1:]
+        command = get_meta_command(command_name)
+        if command is None:
+            console.print(f"Meta command /{command_name} not found")
+            return
+        if command.kimi_soul_only and not isinstance(self.soul, KimiSoul):
+            console.print(f"Meta command /{command_name} not supported")
+            return
+        logger.debug(
+            "Running meta command: {command_name} with args: {command_args}",
+            command_name=command_name,
+            command_args=command_args,
+        )
+        try:
+            ret = command.func(self, command_args)
+            if isinstance(ret, Awaitable):
+                await ret
+        except LLMNotSet:
+            logger.error("LLM not set")
+            console.print("[red]LLM not set, send /setup to configure[/red]")
+        except ChatProviderError as e:
+            logger.exception("LLM provider error:")
+            console.print(f"[red]LLM provider error: {e}[/red]")
+        except Reload:
+            # just propagate
+            raise
+        except BaseException as e:
+            logger.exception("Unknown error:")
+            console.print(f"[red]Unknown error: {e}[/red]")
+            raise  # re-raise unknown error
+
+
+_KIMI_BLUE = "dodger_blue1"
+_LOGO = f"""\
+[{_KIMI_BLUE}]\
+▐█▛█▛█▌
+▐█████▌\
+[{_KIMI_BLUE}]\
+"""
+
+
+def _print_welcome_info(name: str, model: str, info_items: dict[str, str]) -> None:
+    head = Text.from_markup(f"[bold]Welcome to {name}![/bold]")
+    help_text = Text.from_markup("[grey50]Send /help for help information.[/grey50]")
+
+    # Use Table for precise width control
+    logo = Text.from_markup(_LOGO)
+    table = Table(show_header=False, show_edge=False, box=None, padding=(0, 1), expand=False)
+    table.add_column(justify="left")
+    table.add_column(justify="left")
+    table.add_row(logo, Group(head, help_text))
+
+    rows: list[RenderableType] = [table]
+
+    rows.append(Text(""))  # Empty line
+    rows.extend(
+        Text.from_markup(f"[grey50]{key}: {value}[/grey50]") for key, value in info_items.items()
+    )
+    if model:
+        rows.append(Text.from_markup(f"[grey50]Model: {model}[/grey50]"))
+    else:
+        rows.append(
+            Text.from_markup(
+                "[grey50]Model:[/grey50] [yellow]not set, send /setup to configure[/yellow]"
+            )
+        )
+
+    console.print(
+        Panel(
+            Group(*rows),
+            border_style=_KIMI_BLUE,
+            expand=False,
+            padding=(1, 2),
+        )
+    )

kimi_cli/ui/shell/console.py

@@ -0,0 +1,3 @@
+from rich.console import Console
+
+console = Console(highlight=False)

kimi_cli/ui/shell/liveview.py

@@ -0,0 +1,158 @@
+import streamingjson
+from kosong.base.message import ToolCall, ToolCallPart
+from kosong.tooling import ToolError, ToolOk, ToolResult, ToolReturnType
+from rich.console import Group, RenderableType
+from rich.live import Live
+from rich.markup import escape
+from rich.spinner import Spinner
+from rich.text import Text
+
+from kimi_cli.soul import StatusSnapshot
+from kimi_cli.tools import extract_subtitle
+from kimi_cli.ui.shell.console import console
+
+
+class _ToolCallDisplay:
+    def __init__(self, tool_call: ToolCall):
+        self._tool_name = tool_call.function.name
+        self._lexer = streamingjson.Lexer()
+        if tool_call.function.arguments is not None:
+            self._lexer.append_string(tool_call.function.arguments)
+
+        self._title_markup = f"Using [blue]{self._tool_name}[/blue]"
+        self._subtitle = extract_subtitle(self._lexer, self._tool_name)
+        self._finished = False
+        self._spinner = Spinner("dots", text=self._spinner_markup)
+        self.renderable: RenderableType = Group(self._spinner)
+
+    @property
+    def finished(self) -> bool:
+        return self._finished
+
+    @property
+    def _spinner_markup(self) -> str:
+        return self._title_markup + self._subtitle_markup
+
+    @property
+    def _subtitle_markup(self) -> str:
+        subtitle = self._subtitle
+        return f"[grey50]: {escape(subtitle)}[/grey50]" if subtitle else ""
+
+    def append_args_part(self, args_part: str):
+        if self.finished:
+            return
+        self._lexer.append_string(args_part)
+        # TODO: don't extract detail if it's already stable
+        new_subtitle = extract_subtitle(self._lexer, self._tool_name)
+        if new_subtitle and new_subtitle != self._subtitle:
+            self._subtitle = new_subtitle
+            self._spinner.update(text=self._spinner_markup)
+
+    def finish(self, result: ToolReturnType):
+        """
+        Finish the live display of a tool call.
+        After calling this, the `renderable` property should be re-rendered.
+        """
+        self._finished = True
+        sign = "[red]✗[/red]" if isinstance(result, ToolError) else "[green]✓[/green]"
+        lines = [
+            Text.from_markup(f"{sign} Used [blue]{self._tool_name}[/blue]" + self._subtitle_markup)
+        ]
+        if result.brief:
+            lines.append(
+                Text.from_markup(
+                    f"  {result.brief}", style="grey50" if isinstance(result, ToolOk) else "red"
+                )
+            )
+        self.renderable = Group(*lines)
+
+
+class StepLiveView:
+    def __init__(self, status: StatusSnapshot):
+        self._line_buffer = Text("")
+        self._tool_calls: dict[str, _ToolCallDisplay] = {}
+        self._last_tool_call: _ToolCallDisplay | None = None
+        self._status_text: Text | None = Text(
+            self._format_status(status), style="grey50", justify="right"
+        )
+
+    def __enter__(self):
+        self._live = Live(
+            self._compose(),
+            console=console,
+            refresh_per_second=4,
+            transient=False,  # leave the last frame on the screen
+            vertical_overflow="visible",
+        )
+        self._live.__enter__()
+        return self
+
+    def __exit__(self, exc_type, exc_value, traceback):
+        self._live.__exit__(exc_type, exc_value, traceback)
+
+    def _compose(self) -> RenderableType:
+        sections = []
+        if self._line_buffer:
+            sections.append(self._line_buffer)
+        for view in self._tool_calls.values():
+            sections.append(view.renderable)
+        if self._status_text:
+            sections.append(self._status_text)
+        return Group(*sections)
+
+    def _push_out(self, text: Text | str):
+        """
+        Push the text out of the live view to the console.
+        After this, the printed line will not be changed further.
+        """
+        console.print(text)
+
+    def append_text(self, text: str):
+        lines = text.split("\n")
+        prev_is_empty = not self._line_buffer
+        for line in lines[:-1]:
+            self._push_out(self._line_buffer + line)
+            self._line_buffer.plain = ""
+        self._line_buffer.append(lines[-1])
+        if (prev_is_empty and self._line_buffer) or (not prev_is_empty and not self._line_buffer):
+            self._live.update(self._compose())
+
+    def append_tool_call(self, tool_call: ToolCall):
+        self._tool_calls[tool_call.id] = _ToolCallDisplay(tool_call)
+        self._last_tool_call = self._tool_calls[tool_call.id]
+        self._live.update(self._compose())
+
+    def append_tool_call_part(self, tool_call_part: ToolCallPart):
+        if not tool_call_part.arguments_part:
+            return
+        if self._last_tool_call is None:
+            return
+        self._last_tool_call.append_args_part(tool_call_part.arguments_part)
+
+    def append_tool_result(self, tool_result: ToolResult):
+        if view := self._tool_calls.get(tool_result.tool_call_id):
+            view.finish(tool_result.result)
+            self._live.update(self._compose())
+
+    def update_status(self, status: StatusSnapshot):
+        if self._status_text is None:
+            return
+        self._status_text.plain = self._format_status(status)
+
+    def finish(self):
+        for view in self._tool_calls.values():
+            if not view.finished:
+                # this should not happen, but just in case
+                view.finish(ToolOk(output=""))
+        self._live.update(self._compose())
+
+    def interrupt(self):
+        for view in self._tool_calls.values():
+            if not view.finished:
+                view.finish(ToolError(message="", brief="Interrupted"))
+        self._live.update(self._compose())
+
+    @staticmethod
+    def _format_status(status: StatusSnapshot) -> str:
+        bounded = max(0.0, min(status.context_usage, 1.0))
+        return f"context: {bounded:.1%}"