_runtime.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.
 
"""Shared runtime for the Aquin SDK — same engine path the CLI uses."""
 
from __future__ import annotations
 
import os
from typing import Any
 
 
class AquinError(RuntimeError):
    """Raised when an SDK call returns ``{"error": ...}`` or fails to run."""
 
    def __init__(self, message: str, *, result: dict[str, Any] | None = None):
        super().__init__(message)
        self.result = result or {}
 
 
def build_ctx(*, model_id: str | None = None, extra_state: dict[str, Any] | None = None) -> dict[str, Any]:
    """Build the tool context dict expected by bridge / registry / daemon."""
    from aquin.compute.model_loader import get_active_model_id, resolve_model_id
    from aquin.engine.main import resolve_api_key
    from aquin.engine.session_memory_store import load_local_memory
 
    active = model_id or get_active_model_id() or ""
    if active:
        try:
            active = resolve_model_id(active)
        except ValueError:
            pass
 
    try:
        api_key = resolve_api_key(allow_missing=True) or ""
    except Exception:
        api_key = ""
 
    base_url = os.environ.get("AQUIN_BASE_URL", "https://api.aquin.app").rstrip("/")
    try:
        from aquin.engine.main import _load_state
 
        persisted = _load_state() or {}
        api_key = persisted.get("api_key") or api_key
        base_url = (persisted.get("base_url") or base_url).rstrip("/")
    except Exception:
        pass
 
    local_mem = load_local_memory("local")
    state: dict[str, Any] = {
        "activeModelId": active,
        "memory": local_mem if isinstance(local_mem, dict) else {},
    }
    try:
        from aquin.session_mode import mode_for_model_id
 
        if active:
            state["session_mode"] = mode_for_model_id(active)
    except Exception:
        pass
    if extra_state:
        state.update(extra_state)
 
    return {
        "session_id": "",
        "api_key": api_key,
        "base_url": base_url,
        "cwd": os.getcwd(),
        "state": state,
    }
 
 
def prepare_compute(*, model_id: str | None = None) -> str:
    """Ensure loader shim + local server; return resolved active model id."""
    from aquin.compute.loader_shim import apply as shim_apply
    from aquin.compute.model_loader import get_active_model_id, resolve_model_id
    from aquin.engine.local_server import start as start_local_server
 
    shim_apply()
    start_local_server()
    active = model_id or get_active_model_id() or ""
    if active:
        try:
            active = resolve_model_id(active)
        except ValueError:
            pass
    return active
 
 
def invoke(
    tool_name: str,
    args: dict[str, Any] | None = None,
    *,
    command: str | None = None,
    needs_model: bool = True,
    model_id: str | None = None,
    raise_on_error: bool = True,
) -> dict[str, Any]:
    """
    Run a named compute tool through the same dispatch path as the CLI.
 
    Parameters
    ----------
    tool_name:
        Bridge / registry tool id (e.g. ``run_consistency_eval``).
    args:
        Tool arguments (snake or camelCase accepted by most handlers).
    command:
        Optional CLI-style label for local command tracking.
    needs_model:
        If True, ensure a model is loaded / daemon-ready.
    model_id:
        Override active model for this call.
    raise_on_error:
        If True, raise :class:`AquinError` when the result has ``error``.
    """
    from aquin.engine.sync_dispatch import dispatch_model_free, run_dispatch
 
    tool_args = dict(args or {})
    active = prepare_compute(model_id=model_id) if needs_model else (model_id or "")
    if not needs_model and model_id:
        try:
            from aquin.compute.model_loader import resolve_model_id
 
            active = resolve_model_id(model_id)
        except Exception:
            active = model_id
    if needs_model and active:
        tool_args.setdefault("model_id", active)
 
    ctx = build_ctx(model_id=active or model_id)
 
    if needs_model:
        result = run_dispatch(
            tool_name,
            tool_args,
            ctx,
            command=command or tool_name,
            ensure_model=active or None,
        )
    else:
        result = dispatch_model_free(
            tool_name,
            tool_args,
            ctx,
            command=command or tool_name,
        )
 
    if not isinstance(result, dict):
        result = {"content": result}
 
    if raise_on_error and result.get("error"):
        raise AquinError(str(result["error"]), result=result)
    return result
 
 
def invoke_direct(
    fn: Any,
    *args: Any,
    command: str | None = None,
    tool_name: str | None = None,
    track: bool = True,
    **kwargs: Any,
) -> Any:
    """Call a compute function directly (for paths that skip the tool registry)."""
    result = fn(*args, **kwargs)
    if track:
        try:
            from aquin.engine.sync_dispatch import track_cli_result
 
            track_cli_result(
                tool_name or (command or getattr(fn, "__name__", "direct")),
                kwargs if kwargs else {"args": list(args)},
                result if isinstance(result, dict) else {"result": result},
                command=command,
            )
        except Exception:
            pass
    if isinstance(result, dict) and result.get("error"):
        raise AquinError(str(result["error"]), result=result)
    return result