464 lines
17 KiB
Python
464 lines
17 KiB
Python
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,
|
|
)
|