初始化提交

2026-08-21 16:00:13 +08:00
commit 887c1f449e
@@ -0,0 +1,463 @@
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import Any
+
+from app.ai_provider.config.loader import (
+    load_models_config,
+    load_pricing_config,
+    load_providers_config,
+)
+from app.ai_provider.config.models import ModelConfig, ProviderConfig
+from app.ai_provider.core.billing import (
+    BillingFactory,
+    BillingTracker,
+    BudgetController,
+    PricingInfo,
+    PricingRegistry,
+)
+from app.ai_provider.core.registry import ModelRegistry, ProviderRegistry
+from app.ai_provider.core.resilience import CircuitBreaker
+from app.ai_provider.llm.anthropic_format import (
+    AnthropicFormatAdapter,
+    AnthropicFormatAdapterFactory,
+)
+from app.ai_provider.llm.openai_format import OpenAIFormatAdapter, OpenAIFormatAdapterFactory
+from app.ai_provider.llm.stream import StreamHandler
+from app.ai_provider.routing import (
+    CapabilityPermissionController,
+    ModelSelector,
+    MultiModalRouter,
+)
+from app.ai_provider.types.chat import ChatResult, TokenUsage
+from app.ai_provider.types.errors import AIProviderError
+from app.ai_provider.types.messages import (
+    AIMessage,
+    BaseMessage,
+    HumanMessage,
+    SystemMessage,
+)
+from app.core.config import settings
+
+
+@dataclass
+class _ResolvedEndpoint:
+    """端点到适配器和凭证的解析结果。"""
+
+    adapter: OpenAIFormatAdapter | AnthropicFormatAdapter
+    model: ModelConfig
+    provider: ProviderConfig
+    api_key: str
+    base_url: str | None
+    breaker: CircuitBreaker | None = None
+
+
+class ProviderManager:
+    """所有 AI 提供商操作的统一入口。
+
+    取代旧的 ModelRouter + get_model_router() 模式。
+    组合了注册表、适配器、池、计费、弹性和路由。
+    """
+
+    def __init__(
+        self,
+        model_registry: ModelRegistry,
+        provider_registry: ProviderRegistry,
+        billing_tracker: BillingTracker | None = None,
+        budget_controller: BudgetController | None = None,
+        pricing_registry: PricingRegistry | None = None,
+    ) -> None:
+        self.model_registry = model_registry
+        self.provider_registry = provider_registry
+        self._billing = billing_tracker or BillingFactory.create_tracker()
+        self._budget = budget_controller
+        self._pricing_registry = pricing_registry or PricingRegistry()
+        self._adapter_cache: dict[str, OpenAIFormatAdapter | AnthropicFormatAdapter] = {}
+        self._model_selector = ModelSelector(model_registry)
+        self._multi_modal_router = MultiModalRouter()
+        self._permission_controller = CapabilityPermissionController()
+
+    # ── 模型解析 ──────────────────────────────────────────────
+
+    def resolve_model(self, model_hint: str | None = None) -> str:
+        if model_hint is None:
+            return self.model_registry.get_default_name()
+        return self.model_registry.resolve_name(model_hint)
+
+    def get_model_config(self, name: str) -> ModelConfig:
+        return self.model_registry.get(name)
+
+    def get_provider_config(self, name: str) -> ProviderConfig:
+        return self.provider_registry.get(name)
+
+    # ── 适配器解析 ────────────────────────────────────────────
+
+    def get_adapter(self, model_name: str) -> OpenAIFormatAdapter | AnthropicFormatAdapter:
+        if model_name in self._adapter_cache:
+            return self._adapter_cache[model_name]
+
+        model = self.model_registry.get(model_name)
+        provider = self.provider_registry.get(model.provider)
+
+        if model.format.value == "openai":
+            adapter = OpenAIFormatAdapterFactory.create(provider)
+        else:
+            adapter = AnthropicFormatAdapterFactory.create(provider)
+
+        self._adapter_cache[model_name] = adapter
+        return adapter
+
+    # ── 消息辅助 ───────────────────────────────────────────────
+
+    @staticmethod
+    def to_message_list(
+        system_messages: list[str],
+        history: list[tuple[str, str]],
+    ) -> list[BaseMessage]:
+        messages: list[BaseMessage] = []
+        for msg in system_messages:
+            if msg:
+                messages.append(SystemMessage(content=msg))
+        for role, content in history:
+            if role == "system":
+                messages.append(SystemMessage(content=content))
+            elif role == "user":
+                messages.append(HumanMessage(content=content))
+            elif role == "assistant":
+                messages.append(AIMessage(content=content))
+        return messages
+
+    # ── Embedding ──────────────────────────────────────────────────────
+
+    def create_embedding(self, model_name: str) -> Any:
+        """按模型名创建 Embedding 适配器。
+
+        优先读取 model endpoint 级 base_url/api_key 覆盖；无 endpoint 时
+        回退 provider 级凭证。
+        """
+        from app.ai_provider.embedding.openai_format import OpenAIFormatEmbeddingAdapter
+
+        model = self.model_registry.get(model_name)
+        provider = self.provider_registry.get(model.provider)
+
+        # 尝试从 model endpoint 读取 base_url / api_key 覆盖
+        base_url: str | None = None
+        api_key: str = ""
+        if model.endpoints and len(model.endpoints) > 0:
+            ep = model.endpoints[0]  # 取第一个活跃 endpoint
+            if ep.base_url:
+                base_url = ep.base_url
+            elif ep.base_url_env and hasattr(settings, ep.base_url_env.lower()):
+                base_url = getattr(settings, ep.base_url_env.lower())
+            if ep.api_key:
+                api_key = ep.api_key
+            elif ep.api_key_env and hasattr(settings, ep.api_key_env.lower()):
+                api_key = getattr(settings, ep.api_key_env.lower(), "")
+
+        if not base_url:
+            _, base_url = self._resolve_credentials(provider)
+        if not base_url:
+            base_url = provider.base_url
+        if not api_key:
+            api_key, _ = self._resolve_credentials(provider)
+
+        return OpenAIFormatEmbeddingAdapter(
+            base_url=base_url or "",
+            api_key=api_key or "",
+            model=model.model,
+            request_timeout_seconds=model.timeout_seconds,
+        )
+
+    # ── Reranker ───────────────────────────────────────────────────────
+
+    def create_reranker(self, model_name: str) -> Any:
+        """按模型名创建 Reranker 适配器。"""
+        from app.ai_provider.reranker.cohere_format import CohereFormatRerankerAdapter
+
+        model = self.model_registry.get(model_name)
+        provider = self.provider_registry.get(model.provider)
+        api_key, base_url = self._resolve_credentials(provider)
+        if not base_url:
+            base_url = provider.base_url
+        return CohereFormatRerankerAdapter(
+            base_url=base_url or "",
+            api_key=api_key,
+            model=model.model,
+            request_timeout_seconds=model.timeout_seconds,
+        )
+
+    # ── 对话 ──────────────────────────────────────────────────────────
+
+    def _resolve_endpoint(self, model_name: str) -> _ResolvedEndpoint:
+        """解析模型名→(适配器, 凭证, 熔断器)。
+
+        处理池选择、熔断初始化、凭证回退。
+        """
+        model = self.model_registry.get(model_name)
+        provider = self.provider_registry.get(model.provider)
+        adapter = self.get_adapter(model_name)
+
+        breaker: CircuitBreaker | None = None
+
+        if model.endpoint_pool_enabled and model.endpoints:
+            from app.ai_provider.llm.pool import PoolFactory
+
+            pool = PoolFactory.create_pool(
+                logical_model_name=model_name,
+                endpoints=[ep.model_dump(exclude_none=True) for ep in model.endpoints],
+                strategy_name=model.routing_strategy,
+            )
+            selected = pool.select(set())
+            if selected is not None:
+                api_key = selected.get("api_key") or (
+                    settings.__getattribute__(provider.api_key_env.lower())
+                    if hasattr(settings, provider.api_key_env.lower())
+                    else ""
+                )
+                base_url_env = selected.get("base_url_env") or provider.base_url_env
+                base_url = selected.get("base_url") or (
+                    settings.__getattribute__(base_url_env.lower())
+                    if hasattr(settings, base_url_env.lower())
+                    else None
+                )
+                # 熔断器
+                from app.ai_provider.core.resilience import ResilienceFactory
+
+                breaker = ResilienceFactory.create_circuit_breaker(
+                    name=f"{model_name}@{selected.get('name', 'default')}",
+                    failure_threshold=3,
+                    cooldown_seconds=30,
+                )
+                if not breaker.allow_request():
+                    raise AIProviderError(
+                        message=f"Circuit breaker open for {selected.get('name', 'unknown')}",
+                        code="CIRCUIT_BREAKER_OPEN",
+                        status_code=503,
+                    )
+                return _ResolvedEndpoint(
+                    adapter=adapter,
+                    model=model,
+                    provider=provider,
+                    api_key=api_key,
+                    base_url=base_url,
+                    breaker=breaker,
+                )
+            # 无活跃 endpoint，回退 provider 凭证（无熔断）
+            api_key, base_url = self._resolve_credentials(provider)
+            return _ResolvedEndpoint(
+                adapter=adapter,
+                model=model,
+                provider=provider,
+                api_key=api_key,
+                base_url=base_url,
+            )
+
+        # 无池，直接走 provider 凭证
+        api_key, base_url = self._resolve_credentials(provider)
+        return _ResolvedEndpoint(
+            adapter=adapter,
+            model=model,
+            provider=provider,
+            api_key=api_key,
+            base_url=base_url,
+        )
+
+    async def chat(
+        self,
+        messages: list[BaseMessage],
+        model_hint: str | None = None,
+        params: dict[str, Any] | None = None,
+        tools: list[dict[str, Any]] | None = None,
+        tool_choice: str | dict[str, Any] | None = None,
+    ) -> ChatResult:
+        model_name = self.resolve_model(model_hint)
+        if self._budget and not self._budget.allow_request():
+            raise AIProviderError(
+                message="Budget limit exceeded", code="BUDGET_EXCEEDED", status_code=429
+            )
+
+        ep = self._resolve_endpoint(model_name)
+
+        try:
+            result = await ep.adapter.chat(
+                model=ep.model,
+                messages=messages,
+                api_key=ep.api_key,
+                base_url=ep.base_url,
+                params=params,
+                tools=tools,
+                tool_choice=tool_choice,
+            )
+            if ep.breaker:
+                ep.breaker.record_success()
+        except Exception:
+            if ep.breaker:
+                ep.breaker.record_failure()
+            raise
+
+        self._record_usage(
+            model_name, ep.provider.name, result.usage, self._pricing_registry.get(model_name)
+        )
+        return result
+
+    async def chat_stream(
+        self,
+        messages: list[BaseMessage],
+        model_hint: str | None = None,
+        params: dict[str, Any] | None = None,
+        tools: list[dict[str, Any]] | None = None,
+        tool_choice: str | dict[str, Any] | None = None,
+    ):
+        model_name = self.resolve_model(model_hint)
+        if self._budget and not self._budget.allow_request():
+            raise AIProviderError(
+                message="Budget limit exceeded", code="BUDGET_EXCEEDED", status_code=429
+            )
+
+        ep = self._resolve_endpoint(model_name)
+
+        try:
+            stream = ep.adapter.chat_stream(
+                model=ep.model,
+                messages=messages,
+                api_key=ep.api_key,
+                base_url=ep.base_url,
+                params=params,
+                tools=tools,
+                tool_choice=tool_choice,
+            )
+            handler = StreamHandler()
+            async for chunk in stream:
+                handler.add_chunk(chunk)
+                yield chunk
+
+            if ep.breaker:
+                ep.breaker.record_success()
+        except Exception:
+            if ep.breaker:
+                ep.breaker.record_failure()
+            raise
+
+        result = handler.to_result(model_name=model_name)
+        self._record_usage(
+            model_name, ep.provider.name, result.usage, self._pricing_registry.get(model_name)
+        )
+
+    def _resolve_credentials(self, provider: ProviderConfig) -> tuple[str, str | None]:
+        """从 settings 解析 provider 级别的 API key 和 base URL。"""
+        api_key = (
+            settings.__getattribute__(provider.api_key_env.lower())
+            if hasattr(settings, provider.api_key_env.lower())
+            else ""
+        )
+        base_url = (
+            settings.__getattribute__(provider.base_url_env.lower())
+            if hasattr(settings, provider.base_url_env.lower())
+            else None
+        )
+        return api_key, base_url
+
+    # ── 路由辅助 ───────────────────────────────────────────────
+
+    def select_model_for_context(
+        self,
+        needs_vision: bool = False,
+        needs_thinking: bool = False,
+    ) -> str:
+        capabilities: set[str] = set()
+        if needs_vision:
+            capabilities.add("vision")
+        if needs_thinking:
+            capabilities.add("thinking")
+        model = self._model_selector.select({"capabilities": capabilities})
+        return model.name
+
+    def get_adapter_for_tool_binding(
+        self,
+        model_name: str,
+        tools: list[dict[str, Any]],
+        tool_choice: str | dict[str, Any] | None = None,
+    ):
+        """返回 (adapter, request_body_base) 用于工具绑定。
+
+        为 Agent 运行时保留独立的 bind_tools 步骤。
+        """
+        adapter = self.get_adapter(model_name)
+        body = {"model": self.model_registry.get(model_name).model}
+        adapter.bind_tools(body, tools, tool_choice=tool_choice)
+        return adapter, body
+
+    # ── 用量与计费 ───────────────────────────────────────────────
+
+    def _record_usage(
+        self,
+        model_name: str,
+        provider: str,
+        usage: TokenUsage | None,
+        pricing: PricingInfo | None = None,
+    ) -> None:
+        record = self._billing.record(
+            model_name=model_name,
+            provider=provider,
+            usage=usage,
+            pricing=pricing,
+        )
+        if self._budget and record.cost > 0:
+            self._budget.add_cost_cny(
+                record.cost if record.currency == "CNY" else record.cost * 7.3
+            )
+
+    def get_usage_records(self, **kwargs) -> list:
+        return self._billing.get_records(**kwargs)
+
+    def get_budget_status(self) -> dict[str, bool]:
+        if self._budget:
+            return {"within_daily": self._budget.allow_request()}
+        return {"within_daily": True}
+
+    # ── 生命周期 ─────────────────────────────────────────────────────
+
+    async def close(self) -> None:
+        for adapter in self._adapter_cache.values():
+            await adapter.close()
+
+    async def __aenter__(self) -> ProviderManager:
+        return self
+
+    async def __aexit__(self, *exc_info: object) -> None:
+        await self.close()
+
+
+class ProviderManagerFactory:
+    """工厂：从配置文件创建 ProviderManager。"""
+
+    @staticmethod
+    def create(
+        models_path: str | None = None,
+        providers_path: str | None = None,
+        aliases: dict[str, str] | None = None,
+        daily_budget_cny: float = 0.0,
+    ) -> ProviderManager:
+        from pathlib import Path
+
+        models = load_models_config(Path(models_path) if models_path else None)
+        providers = load_providers_config(Path(providers_path) if providers_path else None)
+        pricing_data = load_pricing_config()
+        from app.ai_provider.core.registry import RegistryFactory
+
+        registries = RegistryFactory.create(
+            models, providers, aliases=aliases, pricing_data=pricing_data
+        )
+        billing = BillingFactory.create_tracker()
+        budget = (
+            BillingFactory.create_budget_controller(daily_budget_cny=daily_budget_cny)
+            if daily_budget_cny > 0
+            else None
+        )
+        return ProviderManager(
+            model_registry=registries.model_registry,
+            provider_registry=registries.provider_registry,
+            billing_tracker=billing,
+            budget_controller=budget,
+            pricing_registry=registries.pricing_registry,
+        )