loader.py

"""Plugin loader and method validation."""

from __future__ import annotations

import importlib
import warnings
from dataclasses import dataclass
from typing import Any

from openagents.config.schema import (
    AgentDefinition,
    ContextAssemblerRef,
    DiagnosticsRef,
    EventBusRef,
    MemoryRef,
    PatternRef,
    PluginRef,
    RuntimeRef,
    SessionRef,
    SkillsRef,
    ToolExecutorRef,
    ToolRef,
)
from openagents.errors.exceptions import PluginLoadError
from openagents.errors.suggestions import near_match
from openagents.interfaces.skills import SkillsPlugin
from openagents.plugins.registry import get_builtin_plugin_class, list_builtin_plugins


@dataclass
class LoadedAgentPlugins:
    memory: Any
    pattern: Any
    tool_executor: Any | None
    context_assembler: Any | None
    tools: dict[str, Any]


@dataclass
class LoadedRuntimeComponents:
    runtime: Any
    session: Any
    events: Any
    skills: Any
    diagnostics: Any = None


def _import_symbol(path: str) -> Any:
    if "." not in path:
        raise PluginLoadError(
            f"Invalid impl path: '{path}'",
            hint="impl path must be 'module.path:Symbol' or 'module.path.Symbol'",
        )
    module_name, attr_name = path.rsplit(".", 1)
    try:
        module = importlib.import_module(module_name)
    except Exception as exc:  # pragma: no cover - defensive
        raise PluginLoadError(
            f"Failed to import module '{module_name}'",
            hint=f"check the spelling and that '{module_name}' is importable on this Python path",
        ) from exc
    try:
        return getattr(module, attr_name)
    except AttributeError as exc:
        raise PluginLoadError(
            f"Module '{module_name}' has no symbol '{attr_name}'",
            hint=f"check the spelling of '{attr_name}' in module '{module_name}'",
        ) from exc


def _instantiate(factory: Any, config: dict[str, Any]) -> Any:
    if not callable(factory):
        return factory
    try:
        return factory(config=config)
    except TypeError as exc:
        raise PluginLoadError(f"Could not instantiate plugin from {factory!r}: {exc}") from exc


def _validate_class_methods(factory: Any, *, required_methods: tuple[str, ...], where: str) -> None:
    if not isinstance(factory, type):
        return
    for method_name in required_methods:
        if not callable(getattr(factory, method_name, None)):
            raise PluginLoadError(f"{where} '{factory.__name__}' must implement '{method_name}'")


def _load_plugin_impl(kind: str, ref: PluginRef, *, required_methods: tuple[str, ...] = ()) -> Any:
    if ref.impl:
        symbol = _import_symbol(ref.impl)
        _validate_class_methods(symbol, required_methods=required_methods, where=f"{kind} plugin")
        return _instantiate(symbol, ref.config)
    if ref.type:
        if kind == "context_assembler" and ref.type == "summarizing":
            raise PluginLoadError(
                "context_assembler type 'summarizing' was renamed to 'truncating' in 0.3.0 "
                "because the old implementation only truncated without summarizing. "
                "Rename to 'truncating', or set impl= to your own LLM-based summarizer."
            )
        plugin_cls = get_builtin_plugin_class(kind, ref.type)
        if plugin_cls is None:
            available = list_builtin_plugins(kind)
            guess = near_match(ref.type, available)
            if guess:
                hint_text = f"Did you mean '{guess}'? Available {kind} plugins: {available}"
            else:
                hint_text = f"Available {kind} plugins: {available}"
            raise PluginLoadError(
                f"Unknown {kind} plugin type: '{ref.type}'",
                hint=hint_text,
            )
        _validate_class_methods(plugin_cls, required_methods=required_methods, where=f"{kind} plugin")
        return _instantiate(plugin_cls, ref.config)
    raise PluginLoadError(
        f"{kind} plugin must set one of 'type' or 'impl'",
        hint="add 'type: \"<name>\"' for a builtin or 'impl: \"module.path:Class\"' for a custom plugin",
    )


def load_plugin(
    kind: str,
    ref: PluginRef,
    *,
    required_methods: tuple[str, ...] = (),
) -> Any:
    """Load a child plugin from a PluginRef.

    Public entry point used by combinator builtins that compose other
    plugins (memory.chain, tool_executor.retry, events.file_logging) and
    by external custom combinators.
    """
    return _load_plugin_impl(kind, ref, required_methods=required_methods)


def _load_plugin(kind: str, ref: PluginRef, *, required_methods: tuple[str, ...] = ()) -> Any:
    """Deprecated alias for :func:`load_plugin`.

    Kept for one release so external combinator plugins keep working.
    Emits a DeprecationWarning at call time.
    """
    warnings.warn(
        "openagents.plugins.loader._load_plugin is deprecated; use openagents.plugins.loader.load_plugin",
        DeprecationWarning,
        stacklevel=2,
    )
    return _load_plugin_impl(kind, ref, required_methods=required_methods)


def load_memory_plugin(ref: MemoryRef) -> Any:
    return _load_plugin_impl("memory", ref)


def load_pattern_plugin(ref: PatternRef) -> Any:
    return _load_plugin_impl("pattern", ref, required_methods=("execute", "react"))


def load_tool_plugin(ref: ToolRef) -> Any:
    return _load_plugin_impl("tool", ref, required_methods=("invoke",))


def load_tool_executor_plugin(ref: ToolExecutorRef | None) -> Any | None:
    if ref is None:
        return None
    plugin = _load_plugin_impl("tool_executor", ref, required_methods=("execute", "execute_stream"))
    if not callable(getattr(plugin, "execute", None)):
        raise PluginLoadError(f"tool executor '{type(plugin).__name__}' must implement 'execute'")
    if not callable(getattr(plugin, "execute_stream", None)):
        raise PluginLoadError(f"tool executor '{type(plugin).__name__}' must implement 'execute_stream'")
    return plugin


def load_context_assembler_plugin(ref: ContextAssemblerRef | None) -> Any | None:
    if ref is None:
        return None
    plugin = _load_plugin_impl("context_assembler", ref, required_methods=("assemble", "finalize"))
    if not callable(getattr(plugin, "assemble", None)):
        raise PluginLoadError(f"context assembler '{type(plugin).__name__}' must implement 'assemble'")
    if not callable(getattr(plugin, "finalize", None)):
        raise PluginLoadError(f"context assembler '{type(plugin).__name__}' must implement 'finalize'")
    return plugin


def load_agent_plugins(agent: AgentDefinition) -> LoadedAgentPlugins:
    memory = load_memory_plugin(agent.memory)
    pattern = load_pattern_plugin(agent.pattern)
    tool_executor = load_tool_executor_plugin(agent.tool_executor)
    context_assembler = load_context_assembler_plugin(agent.context_assembler)

    tools: dict[str, Any] = {}
    for tool_ref in agent.tools:
        if not tool_ref.enabled:
            continue
        tools[tool_ref.id] = load_tool_plugin(tool_ref)

    return LoadedAgentPlugins(
        memory=memory,
        pattern=pattern,
        tool_executor=tool_executor,
        context_assembler=context_assembler,
        tools=tools,
    )


def load_runtime_plugin(ref: RuntimeRef) -> Any:
    """Load a runtime plugin."""
    return _load_plugin_impl("runtime", ref, required_methods=("run",))


def load_session_plugin(ref: SessionRef) -> Any:
    """Load a session manager plugin."""
    return _load_plugin_impl("session", ref, required_methods=("session",))


def load_events_plugin(ref: EventBusRef) -> Any:
    """Load an event bus plugin."""
    return _load_plugin_impl("events", ref, required_methods=("emit", "subscribe"))


def load_skills_plugin(ref: SkillsRef | None) -> SkillsPlugin:
    from openagents.config.schema import SkillsRef as DefaultSkillsRef

    actual = ref or DefaultSkillsRef(type="local")
    plugin = _load_plugin_impl("skills", actual, required_methods=("prepare_session", "load_references", "run_skill"))
    for method_name in ("prepare_session", "load_references", "run_skill"):
        if not callable(getattr(plugin, method_name, None)):
            raise PluginLoadError(f"skills component '{type(plugin).__name__}' must implement '{method_name}'")
    return plugin


def load_diagnostics_plugin(ref: DiagnosticsRef | None) -> Any:
    """Load a diagnostics plugin. Returns NullDiagnosticsPlugin when ref is None."""
    from openagents.plugins.builtin.diagnostics.null_plugin import NullDiagnosticsPlugin

    if ref is None:
        return NullDiagnosticsPlugin()
    return _load_plugin_impl("diagnostics", ref)


def load_agent_router_plugin(config: Any) -> Any:
    """Create a DefaultAgentRouter from a MultiAgentConfig.

    Returns None when the config is absent or disabled.
    """
    if config is None or not getattr(config, "enabled", False):
        return None
    from openagents.plugins.builtin.agent_router.default import DefaultAgentRouter

    return DefaultAgentRouter(config=config.model_dump())


def load_runtime_components(
    runtime_ref: RuntimeRef,
    session_ref: SessionRef,
    events_ref: EventBusRef,
    skills_ref: SkillsRef | None,
    diagnostics_ref: DiagnosticsRef | None = None,
) -> LoadedRuntimeComponents:
    """Load all runtime components from config references.

    ``AppConfig`` guarantees non-None defaults for runtime/session/events
    via pydantic field defaults, so callers always pass real refs for
    those three. ``skills_ref`` remains optional because
    :func:`load_skills_plugin` owns its own default (``type=local``).
    ``diagnostics_ref`` is optional; None yields a NullDiagnosticsPlugin.
    """
    events = load_events_plugin(events_ref)
    session = load_session_plugin(session_ref)

    skills = load_skills_plugin(skills_ref)
    if hasattr(skills, "_session_manager"):
        skills._session_manager = session

    diagnostics = load_diagnostics_plugin(diagnostics_ref)

    runtime = load_runtime_plugin(runtime_ref)

    if hasattr(runtime, "_event_bus"):
        runtime._event_bus = events
    if hasattr(runtime, "_session_manager"):
        runtime._session_manager = session
    if hasattr(runtime, "_diagnostics"):
        runtime._diagnostics = diagnostics

    return LoadedRuntimeComponents(
        runtime=runtime,
        session=session,
        events=events,
        skills=skills,
        diagnostics=diagnostics,
    )