fallback.py

Go to the documentation of this file.

# Copyright (c) 2025-present Aquin Labs Private Limited. All Rights Reserved.
"""Rich chat REPL used when the Textual TUI cannot (or should not) run."""
 
from __future__ import annotations
 
import os
import sys
from typing import Any, Literal
 
from aquin.tui.detect import TuiCapability
 
 
ChatUi = Literal["tui", "fallback"]
 
 
def choose_chat_ui(cap: TuiCapability) -> ChatUi:
    """Pick Textual vs Rich from a detect result."""
    return "tui" if cap.tui_ok else "fallback"
 
 
def run_chat_fallback(
    args: list[str] | None = None,
    *,
    capability: TuiCapability | None = None,
) -> None:
    """Run today's Rich ``aquin chat`` REPL (same agent / tools / sync)."""
    _ = args
    _ = capability
 
    # Lazy imports avoid circular load with aquin.cli helpers.
    from aquin.cli import _BASE_URL, _inject_session_mode, _require_loaded_model_id
    from aquin.compute.model_loader import load_model
    from aquin.engine.agent import run_agent_turn
    from aquin.engine.local_server import start as _start_local_server
    from aquin.engine.main import _load_state, resolve_api_key
    from aquin.engine.session_memory_store import load_local_memory
    from aquin.session_mode import get_filtered_tool_schemas, mode_for_model_id, mode_label
    from rich.console import Console
    from rich.prompt import Prompt
    from rich.rule import Rule
    from rich.status import Status
    from rich.text import Text
    from rich.theme import Theme
 
    _theme = Theme({
        "user.label": "bold bright_white",
        "user.text": "white",
        "aquin.label": "bold #facc15",
        "aquin.text": "#e5e7eb",
        "tool.name": "bold #34d399",
        "tool.error": "bold #f87171",
        "tool.args": "dim #9ca3af",
        "meta": "dim #6b7280",
    })
    console = Console(theme=_theme, highlight=False)
 
    locked_model = _require_loaded_model_id()
    state = _load_state()
    api_key = state.get("api_key") or resolve_api_key()
    base_url = (state.get("base_url") or os.environ.get("AQUIN_BASE_URL", _BASE_URL)).rstrip("/")
 
    try:
        chat_mode = mode_for_model_id(locked_model)
        if not get_filtered_tool_schemas(chat_mode):
            console.print(
                "[tool.error]No tools available for this model mode.[/]\n"
                "Run: aquin load model <model-id>",
            )
            sys.exit(1)
    except ValueError as exc:
        console.print(f"[tool.error]{exc}[/]")
        console.print("[meta]Run: aquin load model <model-id>[/]")
        sys.exit(1)
 
    session_state: dict = {
        "memory": load_local_memory("local"),
        "activeModelId": locked_model,
        "session_mode": chat_mode,
    }
    session_state = _inject_session_mode(session_state)
 
    from aquin.compute.loader_shim import apply as _shim_apply
 
    _shim_apply()
    try:
        load_model(locked_model)
    except Exception as exc:
        console.print(f"[tool.error]{exc}[/]")
        sys.exit(1)
 
    _start_local_server()
 
    console.print()
    console.print(f"[meta]{mode_label(chat_mode)}[/]  [meta]{locked_model}[/]")
    console.print(Rule("[aquin.label]aquin chat[/]", style="dim #374151"))
    console.print("[meta]Ctrl+C or empty line to exit[/]\n")
 
    try:
        while True:
            try:
                user_input = Prompt.ask("[user.label]you[/]", console=console).strip()
            except (EOFError, KeyboardInterrupt):
                console.print()
                break
 
            if not user_input:
                break
 
            _first_token = True
            _spinner_active = [True]
            _status = Status("  thinking...", console=console, spinner="dots")
            _status.start()
 
            def _stop_spinner() -> None:
                if _spinner_active[0]:
                    _spinner_active[0] = False
                    _status.stop()
 
            def _on_token(piece: str) -> None:
                nonlocal _first_token
                _stop_spinner()
                if _first_token:
                    console.print(Text("aquin  ", style="aquin.label"), end="")
                    _first_token = False
                sys.stdout.write(piece)
                sys.stdout.flush()
 
            def _on_tool_start(name: str, tool_args: dict) -> None:
                nonlocal _first_token
                _ = tool_args
                _stop_spinner()
                if not _first_token:
                    console.print()
                    _first_token = True
                console.print(Text.assemble(
                    Text("  [tool] ", style="meta"),
                    Text(name, style="tool.name"),
                    Text("  running...", style="meta"),
                ))
 
            def _on_tool_result(name: str, result: Any, is_error: bool) -> None:
                _ = result
                style = "tool.error" if is_error else "tool.name"
                suffix = "  error" if is_error else "  done"
                console.print(Text.assemble(
                    Text("  [tool] ", style="meta"),
                    Text(name, style=style),
                    Text(suffix, style="meta"),
                ))
 
            try:
                run_agent_turn(
                    session_id="",
                    user_message=user_input,
                    state=session_state,
                    api_key=api_key,
                    base_url=base_url,
                    on_token=_on_token,
                    on_tool_start=_on_tool_start,
                    on_tool_result=_on_tool_result,
                )
            except KeyboardInterrupt:
                _stop_spinner()
                console.print()
                raise
            except Exception as exc:
                _stop_spinner()
                msg = str(exc)
                if "401" in msg or "403" in msg:
                    console.print(Text("\n  Auth error. Set HF_TOKEN / required env credentials.", style="tool.error"))
                elif "connect" in msg.lower() or "connection" in msg.lower() or "network" in msg.lower():
                    console.print(Text("\n  Could not reach Aquin. Check your internet connection.", style="tool.error"))
                else:
                    console.print(Text("\n  Something went wrong. Try again.", style="tool.error"))
            finally:
                _stop_spinner()
 
            console.print()
 
    except KeyboardInterrupt:
        console.print()