pythinker-code 2.0.0__py3-none-any.whl

pythinker_code/ui/shell/replay.py

@@ -0,0 +1,215 @@
+from __future__ import annotations
+
+import asyncio
+import contextlib
+from collections import deque
+from collections.abc import Sequence
+from dataclasses import dataclass
+from typing import cast
+
+from pythinker_core.message import ContentPart, Message
+from pythinker_core.tooling import ToolError, ToolOk
+
+from pythinker_code.notifications.llm import is_notification_message
+from pythinker_code.soul.message import is_system_reminder_message
+from pythinker_code.ui.shell.console import console
+from pythinker_code.ui.shell.echo import render_user_echo
+from pythinker_code.ui.shell.visualize import visualize
+from pythinker_code.utils.aioqueue import QueueShutDown
+from pythinker_code.utils.logging import logger
+from pythinker_code.utils.message import message_stringify
+from pythinker_code.utils.slashcmd import parse_slash_command_call
+from pythinker_code.wire import Wire
+from pythinker_code.wire.file import WireFile
+from pythinker_code.wire.types import (
+    Event,
+    StatusUpdate,
+    SteerInput,
+    StepBegin,
+    TextPart,
+    ToolResult,
+    TurnBegin,
+    is_event,
+)
+
+MAX_REPLAY_TURNS = 5
+
+
+@dataclass(slots=True)
+class _ReplayTurn:
+    user_message: Message
+    events: list[Event]
+    n_steps: int = 0
+
+
+async def replay_recent_history(
+    history: Sequence[Message],
+    *,
+    wire_file: WireFile | None = None,
+    show_thinking_stream: bool = False,
+) -> None:
+    """
+    Replay the most recent user-initiated turns from the provided message history or wire file.
+    """
+    if not history:
+        # if the context history is empty,either this is a new session
+        # or the context has been cleared
+        return
+
+    start_idx = _find_replay_start(history)
+    history_turns = (
+        [] if start_idx is None else _build_replay_turns_from_history(history[start_idx:])
+    )
+    turns = await _build_replay_turns_from_wire(wire_file)
+    if not turns or (history_turns and not _same_user_turns(turns, history_turns)):
+        turns = history_turns
+    if not turns:
+        return
+
+    for turn in turns:
+        wire = Wire()
+        console.print(render_user_echo(turn.user_message))
+        ui_task = asyncio.create_task(
+            visualize(
+                wire.ui_side(merge=False),
+                initial_status=StatusUpdate(),
+                show_thinking_stream=show_thinking_stream,
+            )
+        )
+        for event in turn.events:
+            wire.soul_side.send(event)
+            await asyncio.sleep(0)  # yield to UI loop
+        wire.shutdown()
+        with contextlib.suppress(QueueShutDown):
+            await ui_task
+
+
+async def _build_replay_turns_from_wire(wire_file: WireFile | None) -> list[_ReplayTurn]:
+    if wire_file is None or not wire_file.path.exists():
+        return []
+
+    size = wire_file.path.stat().st_size
+    if size > 20 * 1024 * 1024:
+        logger.info(
+            "Wire file too large for replay, skipping: {file} ({size} bytes)",
+            file=wire_file.path,
+            size=size,
+        )
+        return []
+
+    turns: deque[_ReplayTurn] = deque(maxlen=MAX_REPLAY_TURNS)
+    try:
+        async for record in wire_file.iter_records():
+            wire_msg = record.to_wire_message()
+
+            if isinstance(wire_msg, TurnBegin):
+                if _is_clear_command_input(wire_msg.user_input):
+                    turns.clear()
+                    continue
+                turns.append(
+                    _ReplayTurn(
+                        user_message=_message_from_user_input(wire_msg.user_input),
+                        events=[],
+                    )
+                )
+                continue
+
+            if isinstance(wire_msg, SteerInput):
+                turns.append(
+                    _ReplayTurn(
+                        user_message=_message_from_user_input(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.path)
+        return []
+    return list(turns)
+
+
+def _message_from_user_input(user_input: str | list[ContentPart]) -> Message:
+    content = cast(
+        list[ContentPart],
+        list(user_input) if isinstance(user_input, list) else [TextPart(text=user_input)],
+    )
+    return Message(role="user", content=content)
+
+
+def _same_user_turns(lhs: Sequence[_ReplayTurn], rhs: Sequence[_ReplayTurn]) -> bool:
+    return [message_stringify(turn.user_message) for turn in lhs] == [
+        message_stringify(turn.user_message) for turn in rhs
+    ]
+
+
+def _is_clear_command_input(user_input: str | list[ContentPart]) -> bool:
+    if isinstance(user_input, list):
+        text = Message(role="user", content=user_input).extract_text(" ").strip()
+    else:
+        text = str(user_input).strip()
+    call = parse_slash_command_call(text)
+    if call is None:
+        return False
+    return call.name in {"clear", "reset"}
+
+
+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
+    if message.extract_text().startswith("<system>CHECKPOINT"):
+        return False
+    if is_notification_message(message):
+        return False
+    return not is_system_reminder_message(message)
+
+
+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_TURNS messages
+    return indices[max(0, len(indices) - MAX_REPLAY_TURNS)]
+
+
+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 turn
+            if current_turn is not None:
+                turns.append(current_turn)
+            current_turn = _ReplayTurn(user_message=message, events=[])
+        elif message.role == "assistant":
+            if current_turn is None:
+                continue
+            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_turn is None:
+                continue
+            assert message.tool_call_id is not None
+            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_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

pythinker_code/ui/shell/session_picker.py

@@ -0,0 +1,227 @@
+"""Interactive session picker with directory scope toggle.
+
+Provides a full-screen prompt_toolkit Application that lets the user browse
+sessions for the current working directory or across all known directories,
+toggled via ``Ctrl+A``.
+"""
+
+from __future__ import annotations
+
+from pathlib import Path
+from typing import Literal
+
+from prompt_toolkit.application import Application
+from prompt_toolkit.formatted_text import StyleAndTextTuples
+from prompt_toolkit.key_binding import KeyBindings, KeyPressEvent
+from prompt_toolkit.layout import HSplit, Layout, Window
+from prompt_toolkit.layout.controls import FormattedTextControl
+from prompt_toolkit.styles import Style
+from prompt_toolkit.widgets import Box, Frame, RadioList
+from pythinker_host.path import HostPath
+
+from pythinker_code.session import Session
+from pythinker_code.utils.datetime import format_relative_time
+
+SessionScope = Literal["current", "all"]
+
+_EMPTY_SESSION_ID = "__empty__"
+
+
+def _shorten_work_dir(work_dir: str, max_len: int = 30) -> str:
+    """Abbreviate a work directory path for display."""
+    home = str(Path.home())
+    if work_dir.startswith(home):
+        work_dir = "~" + work_dir[len(home) :]
+    if len(work_dir) <= max_len:
+        return work_dir
+    return "..." + work_dir[-(max_len - 3) :]
+
+
+class SessionPickerApp:
+    """Full-screen session picker with Ctrl+A directory scope toggle."""
+
+    def __init__(
+        self,
+        *,
+        work_dir: HostPath,
+        current_session: Session,
+    ) -> None:
+        self._work_dir = work_dir
+        self._current_session = current_session
+        self._scope: SessionScope = "current"
+        self._sessions: list[Session] = []
+        self._result: str | None = None
+        self._reload_version: int = 0
+
+        self._radio_list = RadioList[str](
+            values=[(_EMPTY_SESSION_ID, "Loading...")],
+            default=_EMPTY_SESSION_ID,
+            show_numbers=False,
+            select_on_focus=True,
+            open_character="",
+            select_character="\u276f",
+            close_character="",
+            show_cursor=False,
+            show_scrollbar=False,
+            container_style="class:task-list",
+            checked_style="class:task-list.checked",
+        )
+        self._app = self._build_app()
+
+    async def run(self) -> tuple[str, HostPath] | None:
+        """Run the picker and return ``(session_id, work_dir)``, or *None*."""
+        await self._load_sessions()
+        self._sync_radio_list()
+        result = await self._app.run_async()
+        if result is None:
+            return None
+        if result == _EMPTY_SESSION_ID:
+            return None
+        # Look up the work_dir for the selected session.
+        for s in self._sessions:
+            if s.id == result:
+                return (result, s.work_dir)
+        return None
+
+    # ------------------------------------------------------------------
+    # Data loading
+    # ------------------------------------------------------------------
+
+    async def _load_sessions(self) -> None:
+        current = self._current_session
+
+        if self._scope == "current":
+            sessions = [s for s in await Session.list(self._work_dir) if s.id != current.id]
+        else:
+            sessions = [s for s in await Session.list_all() if s.id != current.id]
+
+        await current.refresh()
+        if not current.is_empty():
+            sessions.insert(0, current)
+        self._sessions = sessions
+
+    def _build_values(self) -> list[tuple[str, str]]:
+        if not self._sessions:
+            return [(_EMPTY_SESSION_ID, "No sessions found.")]
+
+        current_id = self._current_session.id
+        values: list[tuple[str, str]] = []
+        for session in self._sessions:
+            time_str = format_relative_time(session.updated_at)
+            short_id = session.id[:8]
+            marker = " (current)" if session.id == current_id else ""
+
+            title_line = f"{session.title}{marker}"
+            meta_parts = [time_str, short_id]
+            if self._scope == "all":
+                meta_parts.append(_shorten_work_dir(str(session.work_dir)))
+            meta_line = "  " + " \u00b7 ".join(meta_parts)
+            label = f"{title_line}\n{meta_line}"
+
+            values.append((session.id, label))
+        return values
+
+    def _sync_radio_list(self) -> None:
+        values = self._build_values()
+        self._radio_list.values = values
+        default = values[0][0]
+        self._radio_list.current_value = default
+        self._radio_list.current_values = [default]
+        for idx, (val, _) in enumerate(values):
+            if val == default:
+                self._radio_list._selected_index = idx  # pyright: ignore[reportPrivateUsage]
+                break
+
+    # ------------------------------------------------------------------
+    # UI construction
+    # ------------------------------------------------------------------
+
+    def _header_fragments(self) -> StyleAndTextTuples:
+        scope_label = "current directory" if self._scope == "current" else "all directories"
+        total = len(self._sessions)
+        selected = self._radio_list._selected_index + 1  # pyright: ignore[reportPrivateUsage]
+        return [
+            ("class:header.title", f" SESSIONS ({selected} of {total}) "),
+            ("class:header.meta", f" [{scope_label}] "),
+        ]
+
+    def _footer_fragments(self) -> StyleAndTextTuples:
+        scope_action = "show all projects" if self._scope == "current" else "show current project"
+        return [
+            ("class:footer.text", f" Ctrl+A to {scope_action}"),
+            ("class:footer.text", " \u00b7 "),
+            ("class:footer.text", "Enter to select"),
+            ("class:footer.text", " \u00b7 "),
+            ("class:footer.text", "Esc to cancel "),
+        ]
+
+    def _build_app(self) -> Application[str | None]:
+        kb = KeyBindings()
+
+        @kb.add("escape")
+        @kb.add("c-c")
+        def _cancel(event: KeyPressEvent) -> None:
+            event.app.exit(result=None)
+
+        @kb.add("enter", eager=True)
+        def _select(event: KeyPressEvent) -> None:
+            value = self._radio_list.current_value
+            event.app.exit(result=value)
+
+        @kb.add("c-a")
+        def _toggle_scope(event: KeyPressEvent) -> None:
+            self._scope = "all" if self._scope == "current" else "current"
+            self._reload_version += 1
+            event.app.create_background_task(self._reload_and_refresh(event.app))
+
+        # Mark handlers as used
+        _ = (_cancel, _select, _toggle_scope)
+
+        header = Window(
+            FormattedTextControl(self._header_fragments),
+            height=1,
+            style="class:header",
+        )
+        body = Frame(
+            Box(self._radio_list, padding=1),
+            title=lambda: " Sessions ",
+        )
+        footer = Window(
+            FormattedTextControl(self._footer_fragments),
+            height=1,
+            style="class:footer",
+        )
+
+        return Application(
+            layout=Layout(
+                HSplit([header, body, footer]),
+                focused_element=self._radio_list,
+            ),
+            key_bindings=kb,
+            full_screen=True,
+            erase_when_done=True,
+            style=_session_picker_style(),
+        )
+
+    async def _reload_and_refresh(self, app: Application[str | None]) -> None:
+        version = self._reload_version
+
+        # Show loading state
+        self._radio_list.values = [(_EMPTY_SESSION_ID, "Loading...")]
+        self._radio_list.current_value = _EMPTY_SESSION_ID
+        self._radio_list._selected_index = 0  # pyright: ignore[reportPrivateUsage]
+        app.invalidate()
+
+        await self._load_sessions()
+
+        if version != self._reload_version:
+            return  # stale reload; a newer toggle already started
+
+        self._sync_radio_list()
+        app.invalidate()
+
+
+def _session_picker_style() -> Style:
+    from pythinker_code.ui.theme import get_task_browser_style
+
+    return get_task_browser_style()

pythinker_code/ui/shell/setup.py

@@ -0,0 +1,212 @@
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, NamedTuple
+
+import aiohttp
+from prompt_toolkit import PromptSession
+from prompt_toolkit.shortcuts.choice_input import ChoiceInput
+from pydantic import SecretStr
+
+from pythinker_code.auth import PYTHINKER_CODE_PLATFORM_ID
+from pythinker_code.auth.platforms import (
+    PLATFORMS,
+    ModelInfo,
+    Platform,
+    get_platform_by_name,
+    list_models,
+    managed_model_key,
+    managed_provider_key,
+)
+from pythinker_code.config import (
+    LLMModel,
+    LLMProvider,
+    PythinkerAIFetchConfig,
+    PythinkerAISearchConfig,
+    load_config,
+    save_config,
+)
+from pythinker_code.ui.shell.console import console
+from pythinker_code.ui.shell.slash import registry
+from pythinker_code.utils.logging import logger
+
+if TYPE_CHECKING:
+    from pythinker_code.ui.shell import Shell
+
+
+async def select_platform() -> Platform | None:
+    platform_name = await _prompt_choice(
+        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 = get_platform_by_name(platform_name)
+    if platform is None:
+        console.print("[red]Unknown platform[/red]")
+        return None
+    return platform
+
+
+async def setup_platform(platform: Platform) -> bool:
+    result = await _setup_platform(platform)
+    if not result:
+        # error message already printed
+        return False
+
+    _apply_setup_result(result)
+    thinking_label = "on" if result.thinking else "off"
+    console.print("[green]✓ Setup complete![/green]")
+    console.print(f"  Platform: [bold]{result.platform.name}[/bold]")
+    console.print(f"  Model:    [bold]{result.selected_model.id}[/bold]")
+    console.print(f"  Thinking: [bold]{thinking_label}[/bold]")
+    console.print("  Reloading...")
+    return True
+
+
+class _SetupResult(NamedTuple):
+    platform: Platform
+    api_key: SecretStr
+    selected_model: ModelInfo
+    models: list[ModelInfo]
+    thinking: bool
+
+
+async def _setup_platform(platform: Platform) -> _SetupResult | None:
+    # enter the API key
+    api_key = await _prompt_text("Enter your API key", is_password=True)
+    if not api_key:
+        return None
+
+    # list models
+    try:
+        with console.status("[cyan]Verifying API key...[/cyan]"):
+            models = await list_models(platform, api_key)
+    except aiohttp.ClientResponseError as e:
+        logger.error("Failed to get models: {error}", error=e)
+        console.print(f"[red]Failed to get models: {e.message}[/red]")
+        if e.status == 401 and platform.id != PYTHINKER_CODE_PLATFORM_ID:
+            console.print(
+                "[yellow]Hint: If your API key was obtained from Pythinker, "
+                'please select "Pythinker" instead.[/yellow]'
+            )
+        return None
+    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
+
+    # select the model
+    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 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
+
+    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=["on", "off"],
+        )
+        if not thinking_selection:
+            return None
+        thinking = thinking_selection == "on"
+    else:
+        thinking = False
+
+    return _SetupResult(
+        platform=platform,
+        api_key=SecretStr(api_key),
+        selected_model=selected_model,
+        models=models,
+        thinking=thinking,
+    )
+
+
+def _apply_setup_result(result: _SetupResult) -> None:
+    config = load_config()
+    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="pythinker",
+        base_url=result.platform.base_url,
+        api_key=result.api_key,
+    )
+    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.pythinker_ai_search = PythinkerAISearchConfig(
+            base_url=result.platform.search_url,
+            api_key=result.api_key,
+        )
+
+    if result.platform.fetch_url:
+        config.services.pythinker_ai_fetch = PythinkerAIFetchConfig(
+            base_url=result.platform.fetch_url,
+            api_key=result.api_key,
+        )
+
+    save_config(config)
+
+
+async def _prompt_choice(*, header: str, choices: list[str]) -> str | None:
+    if not choices:
+        return None
+
+    try:
+        return await ChoiceInput(
+            message=header,
+            options=[(choice, choice) for choice in choices],
+            default=choices[0],
+        ).prompt_async()
+    except (EOFError, KeyboardInterrupt):
+        return None
+
+
+async def _prompt_text(prompt: str, *, is_password: bool = False) -> str | None:
+    session = PromptSession[str]()
+    try:
+        return str(
+            await session.prompt_async(
+                f" {prompt}: ",
+                is_password=is_password,
+            )
+        ).strip()
+    except (EOFError, KeyboardInterrupt):
+        return None
+
+
+@registry.command
+def reload(app: Shell, args: str):
+    """Reload configuration"""
+    from pythinker_code.cli import Reload
+
+    raise Reload