zhenxun_bot/zhenxun/services/ai/run/di.py

from collections import defaultdict
from collections.abc import Awaitable, Callable
import inspect
from typing import Annotated, Any, ClassVar, cast

from nonebot.adapters import Bot, Event
from nonebot.matcher import Matcher
from nonebot.utils import is_coroutine_callable
from nonebot_plugin_session import EventSession, extract_session

from zhenxun.services.ai.context.memory.facades import AgentSessionFacade
from zhenxun.services.ai.run.blackboard import BlackboardManager
from zhenxun.utils.utils import infer_plugin_namespace

from .context import ProviderFunc, RunContext, _is_run_context_type
from .hitl import HITLController
from .ui import UIController


class Hidden:
    """
    标记一个参数为隐藏参数。
    在生成大模型工具 Schema 时，带有该标记的 Annotated 参数将被剔除。
    """

    pass


class _InjectMarker:
    """内部魔术标记，用于识别 DI 类型糖"""

    def __init__(self, key: str):
        self.key = key


class _UpstreamResultMarker:
    """内部魔术标记，用于精准捕获上游特定节点的产出"""

    def __init__(self, step_name: str):
        self.step_name = step_name


CurrentUserId = Annotated[str | None, Hidden(), _InjectMarker("user_id")]
CurrentGroupId = Annotated[str | None, Hidden(), _InjectMarker("group_id")]
CurrentPlatform = Annotated[str, Hidden(), _InjectMarker("platform")]
CurrentBot = Annotated[Bot | None, Hidden(), _InjectMarker("bot")]
CurrentEvent = Annotated[Event | None, Hidden(), _InjectMarker("event")]
CurrentMatcher = Annotated[Matcher | None, Hidden(), _InjectMarker("matcher")]
CurrentSession = Annotated[EventSession | None, Hidden(), _InjectMarker("session")]
CurrentUI = Annotated[UIController, Hidden(), _InjectMarker("ui")]
CurrentHITL = Annotated[HITLController, Hidden(), _InjectMarker("hitl")]
CurrentModelName = Annotated[str | None, Hidden(), _InjectMarker("model_name")]
CurrentToolRetries = Annotated[int, Hidden(), _InjectMarker("tool_retries")]
CurrentState = Annotated[dict[str, Any], Hidden(), _InjectMarker("state")]
CurrentSharedState = Annotated[dict[str, Any], Hidden(), _InjectMarker("shared_state")]
CurrentOriginalInput = Annotated[str, Hidden(), _InjectMarker("original_input")]
UpstreamResults = Annotated[dict[str, Any], Hidden(), _InjectMarker("upstream_results")]
CurrentEventPayload = Annotated[Any, Hidden(), _InjectMarker("stream_event")]
CurrentBlackboard = Annotated[
    BlackboardManager | None, Hidden(), _InjectMarker("blackboard")
]
CurrentMemory = Annotated[AgentSessionFacade, Hidden(), _InjectMarker("memory")]
CurrentSandbox = Annotated[Any, Hidden(), _InjectMarker("sandbox")]


class Inject:
    """
    [命名空间] 大模型工具依赖注入类型糖。
    用于在工具参数中快捷获取群聊/用户的上下文，或注册自定义的依赖提供者。
    """

    _providers: ClassVar[dict[str, dict[str, ProviderFunc]]] = defaultdict(dict)

    @classmethod
    def register_provider(
        cls, key: str, provider: ProviderFunc, scope: str | None = None
    ) -> None:
        """
        (底层) 注册一个自定义的依赖提供者。
        """
        ns = scope if scope is not None else infer_plugin_namespace()
        cls._providers[key][ns] = provider

    @classmethod
    def provider(cls, key: str, scope: str | None = None):
        def decorator(func: ProviderFunc) -> ProviderFunc:
            cls.register_provider(key, func, scope)
            return func

        return decorator

    @classmethod
    def bind(cls, key: str, type_hint: Any = Any) -> Any:
        """
        动态生成一个用于类型注解的 Inject 标记。
        示例：DbSession = Inject.bind("db_session", AsyncSession)
        """
        return Annotated[type_hint, Hidden(), _InjectMarker(key)]

    @classmethod
    def UpstreamResult(cls, step_name: str) -> Any:
        """
        自动注入：精准获取工作流中指定前置节点的产出。可作为参数的默认值使用以通过静态检查。
        示例: def my_step(data: str = Inject.UpstreamResult("AgentA")):
        """
        return _UpstreamResultMarker(step_name)

    UserId = CurrentUserId
    """自动注入：触发当前任务的用户 ID"""

    GroupId = CurrentGroupId
    """自动注入：触发当前任务的群聊 ID（私聊时为 None）"""

    Platform = CurrentPlatform
    """自动注入：当前对接的平台名称（如 qq）"""

    Bot = CurrentBot
    """自动注入：当前的 NoneBot Bot 实例"""

    Event = CurrentEvent
    """自动注入：当前的 NoneBot Event 实例"""

    Matcher = CurrentMatcher
    """自动注入：当前的 NoneBot Matcher 实例"""

    Session = CurrentSession
    """自动注入：当前的 nonebot_plugin_session 会话实例"""

    UI = CurrentUI
    """自动注入：前端 UI 控制器实例"""

    HITL = CurrentHITL
    """自动注入：人机协同控制器实例 (提供底层 waiter 交互封装)"""

    ModelName = CurrentModelName
    """自动注入：当前大模型正在执行的底层模型名称"""

    ToolRetries = CurrentToolRetries
    """自动注入：当前工具正在经历的重试次数 (Int)"""

    State = CurrentState
    """自动注入：当前 Agent 轮次隔离的业务状态字典"""

    OriginalInput = CurrentOriginalInput
    """自动注入：当前工作流最初始的用户输入文本"""

    SharedState = CurrentSharedState
    """自动注入：当前 Team 或全局穿透的共享黑板状态字典"""

    UpstreamResults = UpstreamResults
    """自动注入：当前工作流中所有上游节点的产出字典 (Key: Agent Name)"""

    EventPayload = CurrentEventPayload
    """自动注入：当前生命周期钩子流转中的 EventBus 载荷实体"""

    Blackboard = CurrentBlackboard
    """自动注入：当前工作流/团队挂载的强类型黑板 (BlackboardManager) 实例"""

    Memory = CurrentMemory
    """自动注入：当前会话的持久化记忆存取门面 (AgentSessionFacade) 实例"""

    Sandbox = CurrentSandbox
    """自动注入：当前沙箱环境管理器实例"""


class BaseParamResolver:
    """可插拔参数解析器基类协议"""

    def match(
        self,
        name: str,
        param: inspect.Parameter,
        context: RunContext,
    ) -> bool:
        return False

    def static_match(self, param: inspect.Parameter) -> bool:
        return False

    async def resolve(
        self,
        name: str,
        param: inspect.Parameter,
        context: RunContext,
    ) -> Any:
        raise NotImplementedError


class RunContextResolver(BaseParamResolver):
    def match(self, name, param, context) -> bool:
        return _is_run_context_type(param.annotation)

    def static_match(self, param: inspect.Parameter) -> bool:
        return _is_run_context_type(param.annotation)

    async def resolve(self, name, param, context) -> Any:
        return context


class TypeSugarResolver(BaseParamResolver):
    def _get_marker(self, param: inspect.Parameter) -> Any | None:
        anno = param.annotation
        if hasattr(anno, "__metadata__"):
            for arg in anno.__metadata__:
                if isinstance(arg, _InjectMarker | _UpstreamResultMarker):
                    return arg
        default_val = param.default
        if isinstance(default_val, _InjectMarker | _UpstreamResultMarker):
            return default_val
        if hasattr(default_val, "__metadata__"):
            for arg in default_val.__metadata__:
                if isinstance(arg, _InjectMarker | _UpstreamResultMarker):
                    return arg
        return None

    def match(self, name, param, context) -> bool:
        return self._get_marker(param) is not None

    def static_match(self, param: inspect.Parameter) -> bool:
        return self._get_marker(param) is not None

    async def resolve(self, name, param, context) -> Any:
        marker = self._get_marker(param)
        if not marker:
            raise ValueError(f"参数 {name} 缺失 Inject 标记")
        if isinstance(marker, _UpstreamResultMarker):
            return context.upstream_results.get(marker.step_name)

        marker_key = marker.key

        if marker_key in context.di_cache:
            return context.di_cache[marker_key]

        provider_dict = Inject._providers.get(marker_key, {})
        ns = getattr(context.session, "namespace", "global")
        provider = provider_dict.get(ns) or provider_dict.get("global")

        if not provider:
            raise ValueError(f"未知的类型糖标记或未注册的 Provider: {marker_key}")

        if is_coroutine_callable(provider):
            result = await provider(context)
        else:
            result = provider(context)

        context.di_cache[marker_key] = result
        return result


class DependencyInjector:
    """可插拔依赖注入管线 (Resolver Pipeline)"""

    _resolvers: ClassVar[list[BaseParamResolver]] = []

    @classmethod
    def register(cls, resolver: BaseParamResolver) -> None:
        cls._resolvers.append(resolver)

    @classmethod
    def can_resolve_statically(cls, param: inspect.Parameter) -> bool:
        for resolver in cls._resolvers:
            if resolver.static_match(param):
                return True
        return False

    @classmethod
    async def resolve_all(
        cls,
        sig: inspect.Signature,
        call_kwargs: dict[str, Any],
        context: RunContext,
    ) -> dict[str, Any]:
        resolved_kwargs = dict(call_kwargs)
        for k, v in call_kwargs.items():
            context.di_cache[k] = v
        for name, param in sig.parameters.items():
            if name in ("self", "cls") or name in resolved_kwargs:
                continue

            if param.kind in (
                inspect.Parameter.VAR_POSITIONAL,
                inspect.Parameter.VAR_KEYWORD,
            ):
                continue

            resolved = False
            for resolver in cls._resolvers:
                if resolver.match(name, param, context):
                    val = await resolver.resolve(name, param, context)
                    resolved_kwargs[name] = val
                    resolved = True
                    break

            if not resolved:
                if param.default is not inspect.Parameter.empty:
                    continue

                raise ValueError(
                    f"参数 '{name}' 未被大模型提供且缺少显式的依赖注入标记"
                    "(如 Inject.* 或 RunContext)。"
                )

        return resolved_kwargs

    @classmethod
    async def invoke(
        cls,
        func: Callable[..., Any],
        call_kwargs: dict[str, Any],
        context: RunContext,
    ) -> Any:
        """统一执行带有依赖注入的函数 (支持同步/异步)"""
        sig = inspect.signature(func)
        resolved_kwargs = await cls.resolve_all(sig, call_kwargs, context)
        filtered_kwargs = {
            k: v for k, v in resolved_kwargs.items() if k in sig.parameters
        }

        if is_coroutine_callable(func):
            return await cast(Callable[..., Awaitable[Any]], func)(**filtered_kwargs)
        return cast(Callable[..., Any], func)(**filtered_kwargs)


DependencyInjector.register(RunContextResolver())
DependencyInjector.register(TypeSugarResolver())

Inject.register_provider("user_id", lambda ctx: ctx.get_user_id(), scope="global")
Inject.register_provider("group_id", lambda ctx: ctx.get_group_id(), scope="global")
Inject.register_provider("platform", lambda ctx: ctx.get_platform(), scope="global")
Inject.register_provider("bot", lambda ctx: ctx.get_bot(), scope="global")
Inject.register_provider("event", lambda ctx: ctx.get_event(), scope="global")
Inject.register_provider("matcher", lambda ctx: ctx.get_matcher(), scope="global")

Inject.register_provider("hitl", lambda ctx: HITLController(ctx), scope="global")
Inject.register_provider(
    "model_name", lambda ctx: ctx.run.current_model, scope="global"
)

Inject.register_provider(
    "tool_retries", lambda ctx: ctx.call.retry_count, scope="global"
)

Inject.register_provider(
    "original_input", lambda ctx: ctx.run.user_input, scope="global"
)
Inject.register_provider("state", lambda ctx: ctx.state, scope="global")
Inject.register_provider(
    "upstream_results", lambda ctx: ctx.upstream_results, scope="global"
)
Inject.register_provider(
    "shared_state", lambda ctx: ctx.session.shared_state, scope="global"
)
Inject.register_provider("memory", lambda ctx: ctx.session.memory, scope="global")


def _resolve_blackboard(ctx: RunContext):
    bb = ctx.session.blackboard
    if bb is None:
        raise ValueError(
            "Inject.Blackboard 注入失败："
            "当前会话上下文中未挂载 BlackboardManager 实例。"
        )
    return bb


Inject.register_provider("blackboard", _resolve_blackboard, scope="global")


def _resolve_session(ctx):
    bot = ctx.get_bot()
    event = ctx.get_event()
    return extract_session(bot, event) if bot and event else None


Inject.register_provider("session", _resolve_session, scope="global")
Inject.register_provider("ui", lambda ctx: UIController(ctx), scope="global")

from zhenxun.services.ai.sandbox.manager import sandbox_manager

Inject.register_provider("sandbox", lambda ctx: sandbox_manager, scope="global")


__all__ = [
    "DependencyInjector",
    "Hidden",
    "Inject",
]