ink_bridge.py

Go to the documentation of this file.

# Copyright (c) 2025-present Aquin Labs Private Limited. All Rights Reserved.
"""JSONL bridge for driving the Python chat agent from another process.
 
stdin  (commands):  {"type":"user","text":"..."} | {"type":"cancel"} | {"type":"exit"}
stdout (events):    ready | token | tool.start | tool.result | turn.done | error | bye
"""
 
from __future__ import annotations
 
import json
import os
import sys
import threading
from typing import Any
 
 
def _emit(event: dict[str, Any]) -> None:
    sys.stdout.write(json.dumps(event, ensure_ascii=False) + "\n")
    sys.stdout.flush()
 
 
def run_ink_chat_bridge() -> None:
    """Session loop: load model once, then serve turns until exit/EOF."""
    from aquin.cli import _BASE_URL, _inject_session_mode
    from aquin.compute.model_loader import get_active_model_id, load_model, resolve_model_id
    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
 
    active = (get_active_model_id() or "").strip()
    if not active:
        _emit({
            "type": "error",
            "message": "No model loaded. Run: aquin load model <id>",
            "fatal": True,
        })
        _emit({"type": "bye"})
        return
    try:
        locked_model = resolve_model_id(active)
    except ValueError:
        locked_model = active
 
    state = _load_state()
    api_key = state.get("api_key") or resolve_api_key()
    if not api_key:
        _emit({
            "type": "error",
            "message": "Cloud session not configured for chat.",
            "fatal": True,
        })
        _emit({"type": "bye"})
        return
 
    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):
            _emit({
                "type": "error",
                "message": "No tools available for this model mode. Run: aquin load model <id>",
                "fatal": True,
            })
            _emit({"type": "bye"})
            return
    except ValueError as exc:
        _emit({"type": "error", "message": str(exc), "fatal": True})
        _emit({"type": "bye"})
        return
 
    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:
        _emit({"type": "error", "message": str(exc), "fatal": True})
        _emit({"type": "bye"})
        return
 
    _start_local_server()
 
    _emit({
        "type": "ready",
        "model": locked_model,
        "mode": chat_mode,
        "mode_label": mode_label(chat_mode),
    })
 
    cancel_event = threading.Event()
    turn_lock = threading.Lock()
 
    def _run_turn(user_text: str) -> None:
        cancel_event.clear()
 
        def _on_token(piece: str) -> None:
            if cancel_event.is_set():
                raise KeyboardInterrupt()
            _emit({"type": "token", "text": piece})
 
        def _on_tool_start(name: str, tool_args: dict) -> None:
            if cancel_event.is_set():
                raise KeyboardInterrupt()
            _emit({"type": "tool.start", "name": name, "args": tool_args or {}})
 
        def _on_tool_result(name: str, result: Any, is_error: bool) -> None:
            _ = result
            _emit({"type": "tool.result", "name": name, "error": bool(is_error)})
 
        try:
            run_agent_turn(
                session_id="",
                user_message=user_text,
                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,
            )
            if cancel_event.is_set():
                _emit({"type": "error", "message": "Cancelled.", "fatal": False})
            else:
                _emit({"type": "turn.done"})
        except KeyboardInterrupt:
            _emit({"type": "error", "message": "Cancelled.", "fatal": False})
            _emit({"type": "turn.done"})
        except Exception as exc:
            msg = str(exc)
            if "401" in msg or "403" in msg:
                msg = "Auth error. Set required env credentials."
            elif "connect" in msg.lower() or "connection" in msg.lower() or "network" in msg.lower():
                msg = "Could not reach Aquin. Check your internet connection."
            else:
                msg = "Something went wrong. Try again."
            _emit({"type": "error", "message": msg, "fatal": False})
            _emit({"type": "turn.done"})
 
    for line in sys.stdin:
        line = line.strip()
        if not line:
            continue
        try:
            cmd = json.loads(line)
        except json.JSONDecodeError:
            _emit({"type": "error", "message": "bad command JSON", "fatal": False})
            continue
 
        kind = cmd.get("type")
        if kind == "exit":
            break
        if kind == "cancel":
            cancel_event.set()
            continue
        if kind != "user":
            _emit({"type": "error", "message": f"unknown command: {kind}", "fatal": False})
            continue
 
        text = str(cmd.get("text") or "").strip()
        if not text:
            continue
 
        if not turn_lock.acquire(blocking=False):
            _emit({"type": "error", "message": "Turn already in progress.", "fatal": False})
            continue
 
        try:
            _run_turn(text)
        finally:
            turn_lock.release()
 
    _emit({"type": "bye"})