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

kimi_cli/ui/shell/replay.py

@@ -1,104 +1,169 @@
+from __future__ import annotations
+
 import asyncio
 import contextlib
 import getpass
+from collections import deque
 from collections.abc import Sequence
 from dataclasses import dataclass
+from pathlib import Path
 
-from kosong.base.message import Message, TextPart
+import aiofiles
+from kosong.message import Message
 from kosong.tooling import ToolError, ToolOk
 
-from kimi_cli.soul import StatusSnapshot
 from kimi_cli.ui.shell.console import console
 from kimi_cli.ui.shell.prompt import PROMPT_SYMBOL
 from kimi_cli.ui.shell.visualize import visualize
-from kimi_cli.utils.message import message_extract_text, message_stringify
+from kimi_cli.utils.aioqueue import QueueShutDown
+from kimi_cli.utils.logging import logger
+from kimi_cli.utils.message import message_stringify
 from kimi_cli.wire import Wire
-from kimi_cli.wire.message import ContentPart, StepBegin, ToolCall, ToolResult
-
-MAX_REPLAY_RUNS = 5
+from kimi_cli.wire.serde import WireMessageRecord
+from kimi_cli.wire.types import (
+    Event,
+    StatusUpdate,
+    StepBegin,
+    TextPart,
+    ToolResult,
+    TurnBegin,
+    is_event,
+)
 
-type _ReplayEvent = StepBegin | ToolCall | ContentPart | ToolResult
+MAX_REPLAY_TURNS = 5
 
 
 @dataclass(slots=True)
-class _ReplayRun:
+class _ReplayTurn:
     user_message: Message
-    events: list[_ReplayEvent]
+    events: list[Event]
     n_steps: int = 0
 
 
-async def replay_recent_history(history: Sequence[Message]) -> None:
+async def replay_recent_history(
+    history: Sequence[Message],
+    *,
+    wire_file: Path | None = None,
+) -> None:
     """
-    Replay the most recent user-initiated runs from the provided message history.
+    Replay the most recent user-initiated turns from the provided message history or wire file.
     """
-    start_idx = _find_replay_start(history)
-    if start_idx is None:
+    turns = await _build_replay_turns_from_wire(wire_file)
+    if not turns:
+        start_idx = _find_replay_start(history)
+        if start_idx is None:
+            return
+        turns = _build_replay_turns_from_history(history[start_idx:])
+    if not turns:
         return
 
-    runs = _build_replay_runs(history[start_idx:])
-    if not runs:
-        return
-
-    for run in runs:
+    for turn in turns:
         wire = Wire()
-        console.print(f"{getpass.getuser()}{PROMPT_SYMBOL} {message_stringify(run.user_message)}")
+        console.print(f"{getpass.getuser()}{PROMPT_SYMBOL} {message_stringify(turn.user_message)}")
         ui_task = asyncio.create_task(
-            visualize(wire.ui_side, initial_status=StatusSnapshot(context_usage=0.0))
+            visualize(wire.ui_side(merge=False), initial_status=StatusUpdate())
         )
-        for event in run.events:
+        for event in turn.events:
             wire.soul_side.send(event)
             await asyncio.sleep(0)  # yield to UI loop
         wire.shutdown()
-        with contextlib.suppress(asyncio.QueueShutDown):
+        with contextlib.suppress(QueueShutDown):
             await ui_task
 
 
+async def _build_replay_turns_from_wire(wire_file: Path | None) -> list[_ReplayTurn]:
+    if wire_file is None or not wire_file.exists():
+        return []
+
+    size = wire_file.stat().st_size
+    if size > 20 * 1024 * 1024:
+        logger.info(
+            "Wire file too large for replay, skipping: {file} ({size} bytes)",
+            file=wire_file,
+            size=size,
+        )
+        return []
+
+    turns: deque[_ReplayTurn] = deque(maxlen=MAX_REPLAY_TURNS)
+    try:
+        async with aiofiles.open(wire_file, encoding="utf-8") as f:
+            async for line in f:
+                line = line.strip()
+                if not line:
+                    continue
+                try:
+                    record = WireMessageRecord.model_validate_json(line)
+                    wire_msg = record.to_wire_message()
+                except ValueError:
+                    continue
+
+                if isinstance(wire_msg, TurnBegin):
+                    turns.append(
+                        _ReplayTurn(
+                            user_message=Message(role="user", content=wire_msg.user_input),
+                            events=[],
+                        )
+                    )
+                    continue
+
+                if not is_event(wire_msg) or not turns:
+                    continue
+
+                current_turn = turns[-1]
+                if isinstance(wire_msg, StepBegin):
+                    current_turn.n_steps = wire_msg.n
+                current_turn.events.append(wire_msg)
+    except Exception:
+        logger.exception("Failed to build replay turns from wire file {file}:", file=wire_file)
+        return []
+    return list(turns)
+
+
 def _is_user_message(message: Message) -> bool:
     # FIXME: should consider non-text tool call results which are sent as user messages
     if message.role != "user":
         return False
-    return not message_extract_text(message).startswith("<system>CHECKPOINT")
+    return not message.extract_text().startswith("<system>CHECKPOINT")
 
 
 def _find_replay_start(history: Sequence[Message]) -> int | None:
     indices = [idx for idx, message in enumerate(history) if _is_user_message(message)]
     if not indices:
         return None
-    # only replay last MAX_REPLAY_RUNS messages
-    return indices[max(0, len(indices) - MAX_REPLAY_RUNS)]
+    # only replay last MAX_REPLAY_TURNS messages
+    return indices[max(0, len(indices) - MAX_REPLAY_TURNS)]
 
 
-def _build_replay_runs(history: Sequence[Message]) -> list[_ReplayRun]:
-    runs: list[_ReplayRun] = []
-    current_run: _ReplayRun | None = None
+def _build_replay_turns_from_history(history: Sequence[Message]) -> list[_ReplayTurn]:
+    turns: list[_ReplayTurn] = []
+    current_turn: _ReplayTurn | None = None
     for message in history:
         if _is_user_message(message):
-            # start a new run
-            if current_run is not None:
-                runs.append(current_run)
-            current_run = _ReplayRun(user_message=message, events=[])
+            # start a new turn
+            if current_turn is not None:
+                turns.append(current_turn)
+            current_turn = _ReplayTurn(user_message=message, events=[])
         elif message.role == "assistant":
-            if current_run is None:
+            if current_turn is None:
                 continue
-            current_run.n_steps += 1
-            current_run.events.append(StepBegin(n=current_run.n_steps))
-            if isinstance(message.content, str):
-                current_run.events.append(TextPart(text=message.content))
-            else:
-                current_run.events.extend(message.content)
-            current_run.events.extend(message.tool_calls or [])
+            current_turn.n_steps += 1
+            current_turn.events.append(StepBegin(n=current_turn.n_steps))
+            current_turn.events.extend(message.content)
+            current_turn.events.extend(message.tool_calls or [])
         elif message.role == "tool":
-            if current_run is None:
+            if current_turn is None:
                 continue
             assert message.tool_call_id is not None
-            if isinstance(message.content, list) and any(
+            if any(
                 isinstance(part, TextPart) and part.text.startswith("<system>ERROR")
                 for part in message.content
             ):
                 result = ToolError(message="", output="", brief="")
             else:
                 result = ToolOk(output=message.content)
-            current_run.events.append(ToolResult(tool_call_id=message.tool_call_id, result=result))
-    if current_run is not None:
-        runs.append(current_run)
-    return runs
+            current_turn.events.append(
+                ToolResult(tool_call_id=message.tool_call_id, return_value=result)
+            )
+    if current_turn is not None:
+        turns.append(current_turn)
+    return turns

kimi_cli/ui/shell/setup.py

@@ -1,52 +1,39 @@
+from __future__ import annotations
+
 import asyncio
 from typing import TYPE_CHECKING, NamedTuple
 
-import aiohttp
+from loguru import logger
 from prompt_toolkit import PromptSession
 from prompt_toolkit.shortcuts.choice_input import ChoiceInput
 from pydantic import SecretStr
 
-from kimi_cli.config import LLMModel, LLMProvider, MoonshotSearchConfig, load_config, save_config
+from kimi_cli.config import (
+    LLMModel,
+    LLMProvider,
+    MoonshotFetchConfig,
+    MoonshotSearchConfig,
+    load_config,
+    save_config,
+)
+from kimi_cli.platforms import (
+    PLATFORMS,
+    ModelInfo,
+    Platform,
+    get_platform_by_name,
+    list_models,
+    managed_model_key,
+    managed_provider_key,
+)
 from kimi_cli.ui.shell.console import console
-from kimi_cli.ui.shell.metacmd import meta_command
-from kimi_cli.utils.aiohttp import new_client_session
+from kimi_cli.ui.shell.slash import registry
 
 if TYPE_CHECKING:
-    from kimi_cli.ui.shell import ShellApp
-
-
-class _Platform(NamedTuple):
-    id: str
-    name: str
-    base_url: str
-    search_url: str | None = None
-    allowed_models: list[str] | None = None
-
-
-_PLATFORMS = [
-    _Platform(
-        id="kimi-for-coding",
-        name="Kimi For Coding",
-        base_url="https://api.kimi.com/coding/v1",
-        search_url="https://api.kimi.com/coding/v1/search",
-    ),
-    _Platform(
-        id="moonshot-cn",
-        name="Moonshot AI 开放平台 (moonshot.cn)",
-        base_url="https://api.moonshot.cn/v1",
-        allowed_models=["kimi-k2-turbo-preview", "kimi-k2-0905-preview", "kimi-k2-0711-preview"],
-    ),
-    _Platform(
-        id="moonshot-ai",
-        name="Moonshot AI Open Platform (moonshot.ai)",
-        base_url="https://api.moonshot.ai/v1",
-        allowed_models=["kimi-k2-turbo-preview", "kimi-k2-0905-preview", "kimi-k2-0711-preview"],
-    ),
-]
-
-
-@meta_command
-async def setup(app: "ShellApp", args: list[str]):
+    from kimi_cli.ui.shell import Shell
+
+
+@registry.command
+async def setup(app: Shell, args: str):
     """Setup Kimi CLI"""
     result = await _setup()
     if not result:
@@ -54,17 +41,26 @@ async def setup(app: "ShellApp", args: list[str]):
         return
 
     config = load_config()
-    config.providers[result.platform.id] = LLMProvider(
+    provider_key = managed_provider_key(result.platform.id)
+    model_key = managed_model_key(result.platform.id, result.selected_model.id)
+    config.providers[provider_key] = LLMProvider(
         type="kimi",
         base_url=result.platform.base_url,
         api_key=result.api_key,
     )
-    config.models[result.model_id] = LLMModel(
-        provider=result.platform.id,
-        model=result.model_id,
-        max_context_size=result.max_context_size,
-    )
-    config.default_model = result.model_id
+    for key, model in list(config.models.items()):
+        if model.provider == provider_key:
+            del config.models[key]
+    for model_info in result.models:
+        capabilities = model_info.capabilities or None
+        config.models[managed_model_key(result.platform.id, model_info.id)] = LLMModel(
+            provider=provider_key,
+            model=model_info.id,
+            max_context_size=model_info.context_length,
+            capabilities=capabilities,
+        )
+    config.default_model = model_key
+    config.default_thinking = result.thinking
 
     if result.platform.search_url:
         config.services.moonshot_search = MoonshotSearchConfig(
@@ -72,6 +68,12 @@ async def setup(app: "ShellApp", args: list[str]):
             api_key=result.api_key,
         )
 
+    if result.platform.fetch_url:
+        config.services.moonshot_fetch = MoonshotFetchConfig(
+            base_url=result.platform.fetch_url,
+            api_key=result.api_key,
+        )
+
     save_config(config)
     console.print("[green]✓[/green] Kimi CLI has been setup! Reloading...")
     await asyncio.sleep(1)
@@ -83,23 +85,27 @@ async def setup(app: "ShellApp", args: list[str]):
 
 
 class _SetupResult(NamedTuple):
-    platform: _Platform
+    platform: Platform
     api_key: SecretStr
-    model_id: str
-    max_context_size: int
+    selected_model: ModelInfo
+    models: list[ModelInfo]
+    thinking: bool
 
 
 async def _setup() -> _SetupResult | None:
     # select the API platform
     platform_name = await _prompt_choice(
-        header="Select the API platform",
-        choices=[platform.name for platform in _PLATFORMS],
+        header="Select a platform (↑↓ navigate, Enter select, Ctrl+C cancel):",
+        choices=[platform.name for platform in PLATFORMS],
     )
     if not platform_name:
         console.print("[red]No platform selected[/red]")
         return None
 
-    platform = next(platform for platform in _PLATFORMS if platform.name == platform_name)
+    platform = get_platform_by_name(platform_name)
+    if platform is None:
+        console.print("[red]Unknown platform[/red]")
+        return None
 
     # enter the API key
     api_key = await _prompt_text("Enter your API key", is_password=True)
@@ -107,51 +113,52 @@ async def _setup() -> _SetupResult | None:
         return None
 
     # list models
-    models_url = f"{platform.base_url}/models"
     try:
-        async with (
-            new_client_session() as session,
-            session.get(
-                models_url,
-                headers={
-                    "Authorization": f"Bearer {api_key}",
-                },
-                raise_for_status=True,
-            ) as response,
-        ):
-            resp_json = await response.json()
-    except aiohttp.ClientError as e:
+        models = await list_models(platform, api_key)
+    except Exception as e:
+        logger.error("Failed to get models: {error}", error=e)
         console.print(f"[red]Failed to get models: {e}[/red]")
         return None
 
-    model_dict = {model["id"]: model for model in resp_json["data"]}
-
     # select the model
-    if platform.allowed_models is None:
-        model_ids = [model["id"] for model in resp_json["data"]]
-    else:
-        id_set = set(model["id"] for model in resp_json["data"])
-        model_ids = [model_id for model_id in platform.allowed_models if model_id in id_set]
-
-    if not model_ids:
+    if not models:
         console.print("[red]No models available for the selected platform[/red]")
         return None
 
+    model_map = {model.id: model for model in models}
     model_id = await _prompt_choice(
-        header="Select the model",
-        choices=model_ids,
+        header="Select a model (↑↓ navigate, Enter select, Ctrl+C cancel):",
+        choices=list(model_map),
     )
     if not model_id:
         console.print("[red]No model selected[/red]")
         return None
 
-    model = model_dict[model_id]
+    selected_model = model_map[model_id]
+
+    # Determine thinking mode based on model capabilities
+    capabilities = selected_model.capabilities
+    thinking: bool
+
+    if "always_thinking" in capabilities:
+        thinking = True
+    elif "thinking" in capabilities:
+        thinking_selection = await _prompt_choice(
+            header="Enable thinking mode? (↑↓ navigate, Enter select, Ctrl+C cancel):",
+            choices=["off", "on"],
+        )
+        if not thinking_selection:
+            return None
+        thinking = thinking_selection == "on"
+    else:
+        thinking = False
 
     return _SetupResult(
         platform=platform,
         api_key=SecretStr(api_key),
-        model_id=model_id,
-        max_context_size=model["context_length"],
+        selected_model=selected_model,
+        models=models,
+        thinking=thinking,
     )
 
 
@@ -170,7 +177,7 @@ async def _prompt_choice(*, header: str, choices: list[str]) -> str | None:
 
 
 async def _prompt_text(prompt: str, *, is_password: bool = False) -> str | None:
-    session = PromptSession()
+    session = PromptSession[str]()
     try:
         return str(
             await session.prompt_async(
@@ -182,8 +189,8 @@ async def _prompt_text(prompt: str, *, is_password: bool = False) -> str | None:
         return None
 
 
-@meta_command
-def reload(app: "ShellApp", args: list[str]):
+@registry.command
+def reload(app: Shell, args: str):
     """Reload configuration"""
     from kimi_cli.cli import Reload