registry.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.
 
from __future__ import annotations
 
from typing import Any, Callable
 
TOOL_REGISTRY: dict[str, Callable] = {}
_TOOL_SCHEMAS: list[dict[str, Any]] = []
 
 
def register(name: str, schema: dict[str, Any] | None = None):
    def decorator(fn: Callable):
        TOOL_REGISTRY[name] = fn
        if schema is not None:
            _TOOL_SCHEMAS.append(schema)
        return fn
    return decorator
 
 
def dispatch(name: str, args: dict[str, Any], ctx: dict[str, Any]) -> Any:
    handler = TOOL_REGISTRY.get(name)
    if handler is None:
        _load_stubs()
        handler = TOOL_REGISTRY.get(name)
    if handler is None:
        # Stale/partial registry (e.g. old daemon) — bridge may still have the handler.
        from aquin.compute.bridge import TOOL_ROUTES
 
        if name in TOOL_ROUTES:
            from .executor import execute_tool
 
            return execute_tool(name, args, ctx)
        raise NotImplementedError(f"Tool '{name}' not yet ported to engine")
    return handler(args, ctx)
 
 
def list_tools() -> list[str]:
    return list(TOOL_REGISTRY.keys())
 
 
def get_tool_schemas() -> list[dict[str, Any]]:
    return list(_TOOL_SCHEMAS)
 
 
def _load_stubs() -> None:
    from . import executor as _executor  # noqa: F401  # must load before stubs
    from . import inspect as _inspect  # noqa: F401
    from . import session_memory as _session_memory  # noqa: F401
    from . import stubs as _stubs  # noqa: F401
    _ = _executor, _inspect, _session_memory, _stubs  # silence unused warnings
 
_load_stubs()