engine_registry.py 4.89 KB
"""Static runtime metadata registry for managed engines."""

from __future__ import annotations

from typing import Any, Iterable

EngineEntry = dict[str, Any]
EngineCatalog = dict[str, EngineEntry]


def build_default_engine_catalog() -> EngineCatalog:
    """Build the default engine metadata catalog for the web API runtime."""
    return {
        "insight": {
            "name": "insight",
            "available": True,
            "launch_mode": "streamlit",
            "script_path": "apps/engine_console/insight_engine_streamlit_app.py",
            "local_runner": "insight",
        },
        "media": {
            "name": "media",
            "available": True,
            "launch_mode": "streamlit",
            "script_path": "apps/engine_console/media_engine_streamlit_app.py",
            "local_runner": "media",
        },
        "query": {
            "name": "query",
            "available": True,
            "launch_mode": "streamlit",
            "script_path": "apps/engine_console/query_engine_streamlit_app.py",
            "local_runner": "query",
        },
        "forum": {
            "name": "forum",
            "available": True,
            "launch_mode": "callable",
            "script_path": None,
            "local_runner": None,
        },
    }


def _clone_entry(entry: EngineEntry) -> EngineEntry:
    cloned: EngineEntry = {}
    for key, value in entry.items():
        if isinstance(value, list):
            cloned[key] = list(value)
        elif isinstance(value, dict):
            cloned[key] = dict(value)
        else:
            cloned[key] = value
    return cloned


class EngineRuntimeRegistry:
    """Explicit wrapper around the static engine runtime metadata catalog."""

    def __init__(self, catalog: EngineCatalog | None = None) -> None:
        default_catalog = catalog if catalog is not None else build_default_engine_catalog()
        self._default_catalog = {name: _clone_entry(entry) for name, entry in default_catalog.items()}
        self._catalog = {name: _clone_entry(entry) for name, entry in self._default_catalog.items()}

    @property
    def catalog(self) -> EngineCatalog:
        return self._catalog

    def items(self):
        return self._catalog.items()

    def contains(self, engine_name: str) -> bool:
        return engine_name in self._catalog

    def get_entry(self, engine_name: str) -> EngineEntry:
        return self._catalog[engine_name]

    def get_local_runner(self, engine_name: str) -> str | None:
        runner = self.get_entry(engine_name).get("local_runner")
        return str(runner) if runner else None

    def get_script_path(self, engine_name: str) -> str | None:
        script_path = self.get_entry(engine_name).get("script_path")
        return str(script_path) if script_path else None

    def is_available(self, engine_name: str) -> bool:
        return bool(self.get_entry(engine_name).get("available", False))

    def set_available(self, engine_name: str, available: bool) -> None:
        self.get_entry(engine_name)["available"] = bool(available)

    def reset(self) -> None:
        self._catalog = {
            name: _clone_entry(entry)
            for name, entry in self._default_catalog.items()
        }

    def snapshot(self, engine_name: str) -> EngineEntry:
        return _clone_entry(self.get_entry(engine_name))

    def snapshot_all(self) -> EngineCatalog:
        return {name: self.snapshot(name) for name in self._catalog}

    def names(self) -> list[str]:
        return list(self._catalog.keys())

    def engines_for_launch_mode(self, launch_mode: str) -> list[str]:
        return [
            engine_name
            for engine_name, entry in self._catalog.items()
            if entry.get("launch_mode") == launch_mode
        ]

    def streamlit_script_snapshot(self) -> dict[str, str]:
        return {
            engine_name: str(entry["script_path"])
            for engine_name, entry in self._catalog.items()
            if entry.get("launch_mode") == "streamlit" and entry.get("script_path")
        }

    def streamlit_scripts(self) -> dict[str, str]:
        return self.streamlit_script_snapshot()

    def local_runner_snapshot(self, engine_names: Iterable[str] | None = None) -> dict[str, str]:
        names = self.names() if engine_names is None else [str(name) for name in engine_names]
        runners: dict[str, str] = {}
        for engine_name in names:
            if not self.contains(engine_name):
                continue
            runner = self.get_local_runner(engine_name)
            if runner:
                runners[engine_name] = runner
        return runners

    def local_runners(self, engine_names: Iterable[str] | None = None) -> dict[str, str]:
        return self.local_runner_snapshot(engine_names)


ENGINE_RUNTIME_REGISTRY = EngineRuntimeRegistry()


__all__ = [
    "ENGINE_RUNTIME_REGISTRY",
    "EngineCatalog",
    "EngineEntry",
    "EngineRuntimeRegistry",
    "build_default_engine_catalog",
]