hunjing/backend/app/ai_provider/pipeline.py

"""配置驱动的 AI 提供商五阶段参数管道。

根据 ProviderConfig.parameter_mappings 在分发前对参数进行
转换、归一化、过滤和注入。第五阶段 HandlerStep 处理跨参数
依赖等复杂场景。
"""

from __future__ import annotations

from dataclasses import dataclass, field
from typing import Any

from app.ai_provider.config.models import ModelConfig, ProviderConfig


@dataclass
class PipelineContext:
    """管道的输入上下文。"""

    params: dict[str, Any] = field(default_factory=dict)
    model_config: ModelConfig | None = None
    provider_config: ProviderConfig | None = None
    messages: list[dict[str, Any]] | None = None
    tools: list[dict[str, Any]] | None = None


@dataclass
class PipelineResult:
    """经过管道处理后的结果。"""

    params: dict[str, Any] = field(default_factory=dict)
    filtered_params: dict[str, Any] = field(default_factory=dict)
    messages: list[dict[str, Any]] | None = None
    tools: list[dict[str, Any]] | None = None


class PipelineStep:
    """管道步骤的基类。"""

    def apply(self, ctx: PipelineContext, result: PipelineResult) -> None:
        raise NotImplementedError


class NormalizeStep(PipelineStep):
    """步骤 1：归一化参数范围（夹紧到 min/max）。"""

    def apply(self, ctx: PipelineContext, result: PipelineResult) -> None:
        mappings = ctx.provider_config.parameter_mappings if ctx.provider_config else {}
        for key, value in ctx.params.items():
            mapping = mappings.get(key)
            if mapping:
                if mapping.min is not None:
                    value = max(mapping.min, value)
                if mapping.max is not None:
                    value = min(mapping.max, value)
                result.params[key] = value
            else:
                result.params[key] = value


class FilterStep(PipelineStep):
    """步骤 2：过滤模型/提供商不支持的参数。"""

    def apply(self, ctx: PipelineContext, result: PipelineResult) -> None:
        allowed = (
            set(ctx.provider_config.parameter_mappings.keys())
            if ctx.provider_config
            else set(result.params.keys())
        )
        filtered = {}
        for key in list(result.params.keys()):
            if key not in allowed:
                filtered[key] = result.params.pop(key)
        result.filtered_params = filtered


class TransformStep(PipelineStep):
    """步骤 3：转换参数名称和值为提供商所需格式。"""

    def apply(self, ctx: PipelineContext, result: PipelineResult) -> None:
        mappings = ctx.provider_config.parameter_mappings if ctx.provider_config else {}
        remapped: dict[str, Any] = {}
        for key, value in result.params.items():
            mapping = mappings.get(key)
            if mapping:
                provider_key = mapping.provider_param
                if mapping.transform == "int":
                    value = int(value)
                elif mapping.transform == "str":
                    value = str(value)
                elif mapping.transform == "multiply_1000":
                    value = int(value * 1000)
                remapped[provider_key] = value
            else:
                remapped[key] = value
        result.params = remapped


class InjectStep(PipelineStep):
    """步骤 4：为缺失参数注入提供商默认值。"""

    def apply(self, ctx: PipelineContext, result: PipelineResult) -> None:
        mappings = ctx.provider_config.parameter_mappings if ctx.provider_config else {}
        for mapping in mappings.values():
            provider_key = mapping.provider_param
            if provider_key not in result.params and mapping.default is not None:
                result.params[provider_key] = mapping.default


class HandlerStep(PipelineStep):
    """步骤 5：执行已注册的参数处理器。

    处理 YAML 映射规则无法表达的跨参数依赖和复杂变换：
    - Anthropic thinking 参数要求 max_tokens > budget_tokens
    - 根据模型能力过滤不支持的参数
    - 结构性参数重组
    """

    def apply(self, ctx: PipelineContext, result: PipelineResult) -> None:
        if not ctx.provider_config:
            return
        from app.ai_provider.pipeline_handlers import ParameterHandlerRegistry

        registry = ParameterHandlerRegistry

        # 阶段 A：逐参数处理器（来自 ParameterMapping.handler）
        for key, mapping in ctx.provider_config.parameter_mappings.items():
            if mapping.handler:
                handler = registry.get(mapping.handler)
                handler(ctx, result, key)  # key is the param key

        # 阶段 B：全局处理器（来自 ProviderConfig.parameter_handlers）
        for handler_name in ctx.provider_config.parameter_handlers:
            handler = registry.get(handler_name)
            handler(ctx, result, None)  # global handler: no param_key


class ParameterPipeline:
    """配置驱动的五阶段参数管道。"""

    def __init__(
        self,
        steps: list[PipelineStep],
        provider_config: ProviderConfig | None = None,
    ) -> None:
        self._steps = steps
        self._provider_config = provider_config

    def run(self, ctx: PipelineContext) -> PipelineResult:
        if ctx.provider_config is None and self._provider_config is not None:
            ctx.provider_config = self._provider_config
        result = PipelineResult(params=dict(ctx.params))
        for step in self._steps:
            step.apply(ctx, result)
        return result


class PipelineFactory:
    """工厂：从提供商配置构建管道。"""

    @staticmethod
    def create(provider_config: ProviderConfig) -> ParameterPipeline:
        return ParameterPipeline(
            provider_config=provider_config,
            steps=[
                NormalizeStep(),
                FilterStep(),
                TransformStep(),
                InjectStep(),
                HandlerStep(),
            ],
        )