sync_dispatch.py

Go to the documentation of this file.

# Copyright (c) 2025-present Aquin Labs Private Limited. All Rights Reserved.
# This file is part of the Aquin Engine. Unauthorized copying, modification,
# distribution, or use of this file, via any medium, is strictly prohibited.
# Proprietary and confidential. See LICENSE for terms.
 
"""Run tools locally and track results in ~/.aquin/commands/."""
from __future__ import annotations
 
import time
import uuid
from typing import Any
 
 
def _resolve_card(tool_name: str, result: Any, card: dict[str, Any] | None) -> dict[str, Any] | None:
    if card is not None:
        return card
    if not isinstance(result, dict):
        return None
    existing = result.get("card")
    if isinstance(existing, dict):
        return existing
    from aquin.compute import card_mapper
 
    mapped = card_mapper.to_card(tool_name, result)
    if mapped is not None:
        return mapped
    if result.get("type") and "data" in result:
        return result
    return None
 
 
def _result_content(tool_name: str, result: Any, card: dict[str, Any] | None) -> Any:
    raw = result.get("content", result) if isinstance(result, dict) else result
    if isinstance(raw, dict):
        from aquin.compute.sync_slim import slim_tool_result_for_sync
 
        return slim_tool_result_for_sync(tool_name, raw, card)
    return raw
 
 
def _track(
    *,
    tool_name: str,
    command: str | None,
    args: dict[str, Any],
    result: Any,
    card: dict[str, Any] | None,
    duration_ms: int,
    exit_ok: bool,
) -> str:
    from aquin.command_log import append_record
 
    slim = _result_content(tool_name, result, card) if isinstance(result, dict) else result
    return append_record(
        tool_name=tool_name,
        command=command,
        args=args,
        result=slim,
        card=card,
        duration_ms=duration_ms,
        exit_ok=exit_ok,
    )
 
 
def _emit_tracked(name: str, record_id: str) -> None:
    print(f"[tracked] {name} -> {record_id[:8]}", flush=True)
 
 
def track_cli_result(
    tool_name: str,
    args: dict[str, Any],
    result: Any,
    *,
    card: dict[str, Any] | None = None,
    command: str | None = None,
    label: str | None = None,
) -> None:
    """Track a standalone CLI command result locally."""
    _ = label
    card = _resolve_card(tool_name, result, card)
    exit_ok = not (isinstance(result, dict) and result.get("error"))
    rid = _track(
        tool_name=tool_name,
        command=command,
        args=args,
        result=result,
        card=card,
        duration_ms=0,
        exit_ok=exit_ok,
    )
    name = command or tool_name
    _emit_tracked(name, rid)
 
 
def dispatch_with_tracking(
    name: str,
    args: dict[str, Any],
    ctx: dict[str, Any],
    *,
    command: str | None = None,
) -> dict[str, Any]:
    from aquin.engine.tools.registry import dispatch, _load_stubs
 
    _load_stubs()
    _ = ctx
    t0 = time.perf_counter()
    try:
        result = dispatch(name, args, ctx)
        exit_ok = True
    except Exception as exc:
        result = {"error": str(exc)}
        exit_ok = False
 
    card = _resolve_card(name, result, None)
    duration_ms = int((time.perf_counter() - t0) * 1000)
    rid = _track(
        tool_name=name,
        command=command,
        args=args,
        result=result,
        card=card,
        duration_ms=duration_ms,
        exit_ok=exit_ok and not (isinstance(result, dict) and result.get("error")),
    )
    _emit_tracked(command or name, rid)
    return result if isinstance(result, dict) else {"content": result}
 
 
def _finish_tracked_dispatch(
    name: str,
    args: dict[str, Any],
    ctx: dict[str, Any],
    result: Any,
    *,
    command: str | None = None,
) -> dict[str, Any]:
    try:
        card = _resolve_card(name, result, None)
        exit_ok = not (isinstance(result, dict) and result.get("error"))
        rid = _track(
            tool_name=name,
            command=command,
            args=args,
            result=result,
            card=card,
            duration_ms=0,
            exit_ok=exit_ok,
        )
        _emit_tracked(command or name, rid)
    except Exception:
        pass
    return result if isinstance(result, dict) else {"content": result}
 
 
def dispatch_model_free(
    name: str,
    args: dict[str, Any],
    ctx: dict[str, Any],
    *,
    command: str | None = None,
) -> dict[str, Any]:
    """Run catalog/session tools that never touch GPU weights."""
    from aquin.session_mode import tool_requires_model
 
    if tool_requires_model(name):
        raise ValueError(f"dispatch_model_free called for GPU tool: {name}")
 
    from aquin.engine.tools.registry import dispatch, _load_stubs
 
    _load_stubs()
    try:
        result = dispatch(name, args, ctx)
    except Exception as exc:
        from aquin.user_errors import friendly_message
 
        result = {"error": friendly_message(exc)}
 
    return _finish_tracked_dispatch(name, args, ctx, result, command=command)
 
 
def run_dispatch(
    name: str,
    args: dict[str, Any],
    ctx: dict[str, Any],
    *,
    command: str | None = None,
    ensure_model: str | None = None,
    timeout: float = 1800,
) -> dict[str, Any]:
    """
    Run a tool through the persistent model daemon when it is available, else load
    in-process and run locally. Either way the result is tracked locally so the
    user sees the [tracked] line with their own context.
    """
    from aquin.session_mode import tool_requires_model
 
    if not tool_requires_model(name):
        return dispatch_model_free(name, args, ctx, command=command)
 
    result: Any = None
    used_daemon = False
 
    try:
        from aquin.engine import model_daemon
 
        if model_daemon.is_running():
            data = model_daemon.dispatch(name, args, ctx, timeout=timeout)
            if isinstance(data, dict):
                if data.get("ok"):
                    result = data.get("result")
                    used_daemon = True
                elif "error" in data:
                    err = str(data.get("error") or "")
                    low = err.lower()
                    if "not yet ported" in low or "unknown tool" in low:
                        used_daemon = False
                        result = None
                    else:
                        result = {"error": err}
                        used_daemon = True
    except Exception:
        used_daemon = False
 
    if not used_daemon:
        if ensure_model:
            try:
                from aquin.compute.model_loader import load_model, resolve_model_id
                load_model(resolve_model_id(ensure_model))
            except Exception as exc:
                from aquin.user_errors import friendly_message
 
                result = {"error": friendly_message(exc)}
        if result is None:
            from aquin.engine.tools.registry import dispatch, _load_stubs
 
            _load_stubs()
            try:
                result = dispatch(name, args, ctx)
            except Exception as exc:
                from aquin.user_errors import friendly_message
 
                result = {"error": friendly_message(exc)}
 
    return _finish_tracked_dispatch(name, args, ctx, result, command=command)
 
 
# Legacy names — same behavior, no cloud sync
dispatch_with_sync = dispatch_with_tracking
 
 
def sync_cli_result(
    ctx: dict[str, Any],
    tool_name: str,
    args: dict[str, Any],
    result: Any,
    *,
    card: dict[str, Any] | None = None,
    command: str | None = None,
) -> None:
    """Track a catalog/sync CLI result. ``ctx`` is accepted for API compat (session sync is optional)."""
    _ = ctx
    if command is None:
        from aquin.command_log import _TOOL_TO_CLI
 
        command = _TOOL_TO_CLI.get(tool_name, tool_name)
    track_cli_result(tool_name, args, result, card=card, command=command)
 
 
def session_ctx_from_state() -> dict[str, str]:
    """Auth + base URL from engine state (optional session id if present)."""
    import os
 
    from aquin.engine.main import _load_state, resolve_api_key
 
    state = _load_state()
    sid = (state.get("active_session_id") or "").strip()
    api_key = resolve_api_key(allow_missing=True)
    base_url = (
        state.get("base_url") or os.environ.get("AQUIN_BASE_URL", "https://api.aquin.app")
    ).rstrip("/")
    return {"session_id": sid, "api_key": api_key, "base_url": base_url}
 
 
def require_session_ctx(*, label: str) -> dict[str, str]:
    """Only required for aquin chat agent API — not for compute."""
    from aquin.engine.main import resolve_api_key
    import sys
 
    ctx = session_ctx_from_state()
    api_key = (ctx.get("api_key") or resolve_api_key(allow_missing=True)).strip()
    if not api_key:
        print(f"[auth] {label} needs a cloud session token — not available in local framework mode.", file=sys.stderr)
        sys.exit(1)
    sid = (ctx.get("session_id") or "").strip()
    if not sid:
        print(
            f"[chat] Session id required for {label}.\n"
            "       Open a session tab in the web app and run:\n"
            "       aquin config send-session <session-uuid>",
            file=sys.stderr,
        )
        sys.exit(1)
    return ctx
 
 
def require_active_session(ctx: dict[str, Any], *, label: str = "this command") -> None:
    """No-op for compute — kept so old imports don't break."""
    _ = ctx, label
 
 
def push_tool_pair(**kwargs: Any) -> bool:
    """Legacy — track locally instead of pushing to cloud."""
    tool_call_id = kwargs.get("tool_call_id") or str(uuid.uuid4())
    _ = tool_call_id
    track_cli_result(
        kwargs.get("tool_name") or "tool",
        kwargs.get("args") or {},
        kwargs.get("result"),
        card=kwargs.get("card"),
    )
    return True