初始化提交
This commit is contained in:
chenxin committed
commit
887c1f449e
963 files changed
+166869
No files matched your search
@@ -0,0 +1,63 @@
|
||||
from app.ai_provider.config.loader import load_models_config, load_providers_config
|
||||
from app.ai_provider.config.models import ModelConfig, ProviderConfig
|
||||
|
||||
# 嵌入 / 重排序
|
||||
from app.ai_provider.embedding import (
|
||||
EmbeddingProtocol,
|
||||
EmbeddingRequest,
|
||||
EmbeddingResponse,
|
||||
EmbeddingResult,
|
||||
EmbeddingService,
|
||||
EmbeddingServiceError,
|
||||
)
|
||||
from app.ai_provider.manager import ProviderManager, ProviderManagerFactory
|
||||
from app.ai_provider.reranker import (
|
||||
RerankerProtocol,
|
||||
RerankerService,
|
||||
RerankerServiceError,
|
||||
RerankRequest,
|
||||
RerankResponse,
|
||||
RerankResult,
|
||||
)
|
||||
from app.ai_provider.types.chat import ChatChunk, ChatResult, TokenUsage
|
||||
from app.ai_provider.types.errors import AIProviderError
|
||||
from app.ai_provider.types.messages import (
|
||||
AIMessage,
|
||||
BaseMessage,
|
||||
HumanMessage,
|
||||
SystemMessage,
|
||||
ToolMessage,
|
||||
)
|
||||
|
||||
__all__ = [
|
||||
# LLM
|
||||
"ProviderManager",
|
||||
"ProviderManagerFactory",
|
||||
"BaseMessage",
|
||||
"SystemMessage",
|
||||
"HumanMessage",
|
||||
"AIMessage",
|
||||
"ToolMessage",
|
||||
"ChatChunk",
|
||||
"ChatResult",
|
||||
"TokenUsage",
|
||||
"AIProviderError",
|
||||
"ModelConfig",
|
||||
"ProviderConfig",
|
||||
"load_models_config",
|
||||
"load_providers_config",
|
||||
# Embedding
|
||||
"EmbeddingProtocol",
|
||||
"EmbeddingRequest",
|
||||
"EmbeddingResponse",
|
||||
"EmbeddingResult",
|
||||
"EmbeddingService",
|
||||
"EmbeddingServiceError",
|
||||
# Reranker
|
||||
"RerankerProtocol",
|
||||
"RerankRequest",
|
||||
"RerankResponse",
|
||||
"RerankResult",
|
||||
"RerankerService",
|
||||
"RerankerServiceError",
|
||||
]
|
||||
@@ -0,0 +1,197 @@
|
||||
"""音频服务 — ASR / TTS 适配器外观。
|
||||
|
||||
用法::
|
||||
|
||||
from app.ai_provider.audio import AudioService
|
||||
|
||||
service = AudioService(factory)
|
||||
result = await service.transcribe(audio_bytes, format="wav")
|
||||
result = await service.synthesize("你好，魂镜", voice="default")
|
||||
|
||||
``AudioService`` 通过注册表中带有 ``asr`` / ``tts`` 能力标签的模型来解析
|
||||
使用哪个模型，然后将创建委托给 ``ModelFactory.create_audio()``。
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from app.ai_provider.audio.base import (
|
||||
ASRProtocol,
|
||||
ASRRequest,
|
||||
ASRResponse,
|
||||
AudioServiceError,
|
||||
TTSProtocol,
|
||||
TTSRequest,
|
||||
TTSResponse,
|
||||
)
|
||||
|
||||
|
||||
class AudioService:
|
||||
"""外观类，负责解析 ASR / TTS 模型并委托给 ProviderManager。
|
||||
|
||||
接受 ProviderManager 或实现了以下接口的适配器：
|
||||
- ``.registry`` → 具有 ``list_capability_models(capability: str)`` 方法
|
||||
- ``.create_audio(name: str)`` → 返回 ASRProtocol / TTSProtocol 实例
|
||||
"""
|
||||
|
||||
def __init__(self, manager) -> None:
|
||||
self._manager = manager
|
||||
self._asr_adapter: ASRProtocol | None = None
|
||||
self._tts_adapter: TTSProtocol | None = None
|
||||
|
||||
# ── Registry 兼容适配 ─────────────────────────────────────────
|
||||
|
||||
@property
|
||||
def _registry(self):
|
||||
"""按需获取 registry，兼容 adapter 和裸 ProviderManager/ModelRegistry。"""
|
||||
mgr = self._manager
|
||||
if hasattr(mgr, "registry"):
|
||||
return mgr.registry
|
||||
if hasattr(mgr, "model_registry"):
|
||||
return mgr.model_registry
|
||||
return mgr
|
||||
|
||||
def _list_capability_models(self, tag: str):
|
||||
reg = self._registry
|
||||
# 适配器接口（list_capability_models）vs ModelRegistry 接口（list_by_capability）
|
||||
if hasattr(reg, "list_capability_models"):
|
||||
return reg.list_capability_models(tag)
|
||||
if hasattr(reg, "list_by_capability"):
|
||||
return reg.list_by_capability(tag)
|
||||
return []
|
||||
|
||||
@property
|
||||
def _factory(self):
|
||||
mgr = self._manager
|
||||
if hasattr(mgr, "create_audio"):
|
||||
return mgr
|
||||
return None
|
||||
|
||||
# ── ASR ──────────────────────────────────────────────────────────────
|
||||
|
||||
@property
|
||||
def asr(self) -> ASRProtocol:
|
||||
"""解析并缓存 ASR 适配器（取注册表中第一个活跃 ASR 模型）。"""
|
||||
if self._asr_adapter is None:
|
||||
self._asr_adapter = self._resolve_asr()
|
||||
return self._asr_adapter
|
||||
|
||||
def asr_for(self, model_name: str) -> ASRProtocol:
|
||||
"""解析并缓存指定模型名的 ASR 适配器（覆盖默认第一个）。"""
|
||||
self._asr_adapter = self._factory.create_audio(model_name)
|
||||
return self._asr_adapter
|
||||
|
||||
async def transcribe(
|
||||
self,
|
||||
audio_data: bytes,
|
||||
format: str = "wav",
|
||||
language: str | None = None,
|
||||
model_name: str | None = None,
|
||||
) -> ASRResponse:
|
||||
"""快捷方法：构造请求并调用 ASR 适配器。
|
||||
|
||||
Args:
|
||||
model_name: 可选，指定使用哪个 ASR 模型（按注册表逻辑名）。
|
||||
省略时取注册表中第一个活跃的 ASR 模型。
|
||||
"""
|
||||
if model_name is not None:
|
||||
adapter = self.asr_for(model_name)
|
||||
else:
|
||||
adapter = self.asr
|
||||
request = ASRRequest(
|
||||
audio_data=audio_data,
|
||||
format=format,
|
||||
language=language,
|
||||
)
|
||||
return await adapter.transcribe(request)
|
||||
|
||||
# ── TTS ──────────────────────────────────────────────────────────────
|
||||
|
||||
@property
|
||||
def tts(self) -> TTSProtocol:
|
||||
"""解析并缓存 TTS 适配器。"""
|
||||
if self._tts_adapter is None:
|
||||
self._tts_adapter = self._resolve_tts()
|
||||
return self._tts_adapter
|
||||
|
||||
async def synthesize(
|
||||
self,
|
||||
text: str,
|
||||
voice: str = "default",
|
||||
response_format: str = "wav",
|
||||
extra: dict | None = None,
|
||||
) -> TTSResponse:
|
||||
"""快捷方法：构造请求并调用 TTS 适配器。"""
|
||||
request = TTSRequest(
|
||||
text=text,
|
||||
voice=voice,
|
||||
response_format=response_format,
|
||||
extra=extra,
|
||||
)
|
||||
return await self.tts.synthesize(request)
|
||||
|
||||
# ── 解析 ───────────────────────────────────────────────────────
|
||||
|
||||
def _resolve_asr(self) -> ASRProtocol:
|
||||
"""查找第一个活跃的 ASR 模型并通过工厂创建。"""
|
||||
models = self._list_capability_models("asr")
|
||||
if not models:
|
||||
raise AudioServiceError(
|
||||
status_code=503,
|
||||
code="audio_asr_no_model",
|
||||
message="No ASR model configured in the model registry.",
|
||||
)
|
||||
factory = self._factory
|
||||
if factory is None:
|
||||
raise AudioServiceError(
|
||||
status_code=501,
|
||||
code="audio_asr_no_factory",
|
||||
message="Audio adapter factory not available (provider manager missing create_audio).",
|
||||
)
|
||||
return factory.create_audio(models[0].name)
|
||||
|
||||
def _resolve_tts(self) -> TTSProtocol:
|
||||
"""查找第一个活跃的 TTS 模型并通过工厂创建。"""
|
||||
models = self._list_capability_models("tts")
|
||||
if not models:
|
||||
raise AudioServiceError(
|
||||
status_code=503,
|
||||
code="audio_tts_no_model",
|
||||
message="No TTS model configured in the model registry.",
|
||||
)
|
||||
factory = self._factory
|
||||
if factory is None:
|
||||
raise AudioServiceError(
|
||||
status_code=501,
|
||||
code="audio_tts_no_factory",
|
||||
message="Audio adapter factory not available (provider manager missing create_audio).",
|
||||
)
|
||||
return factory.create_audio(models[0].name)
|
||||
|
||||
# ── 生命周期 ────────────────────────────────────────────────────────
|
||||
|
||||
async def close(self) -> None:
|
||||
"""释放适配器持有的 HTTP 连接。"""
|
||||
for adapter in (self._asr_adapter, self._tts_adapter):
|
||||
if adapter is None:
|
||||
continue
|
||||
close = getattr(adapter, "close", None)
|
||||
if close is not None:
|
||||
await close()
|
||||
|
||||
async def __aenter__(self) -> AudioService:
|
||||
return self
|
||||
|
||||
async def __aexit__(self, *exc_info: object) -> None:
|
||||
await self.close()
|
||||
|
||||
|
||||
__all__ = [
|
||||
"ASRProtocol",
|
||||
"ASRRequest",
|
||||
"ASRResponse",
|
||||
"AudioService",
|
||||
"AudioServiceError",
|
||||
"TTSProtocol",
|
||||
"TTSRequest",
|
||||
"TTSResponse",
|
||||
]
|
||||
@@ -0,0 +1,263 @@
|
||||
"""MiMo ASR adapter — ``POST /v1/chat/completions`` with ``input_audio``.
|
||||
|
||||
MiMo 的 ASR 不走独立音频端点，而是通过聊天补全接口扩展实现：
|
||||
|
||||
POST {base_url}/chat/completions
|
||||
Content-Type: application/json
|
||||
Authorization: Bearer {api_key}
|
||||
|
||||
{
|
||||
"model": "mimo-v2.5-asr",
|
||||
"messages": [
|
||||
{
|
||||
"role": "user",
|
||||
"content": [
|
||||
{
|
||||
"type": "input_audio",
|
||||
"input_audio": {
|
||||
"data": "data:audio/wav;base64,<base64_encoded_audio>"
|
||||
}
|
||||
}
|
||||
]
|
||||
}
|
||||
],
|
||||
"asr_options": {
|
||||
"language": "auto"
|
||||
}
|
||||
}
|
||||
|
||||
响应：:
|
||||
|
||||
{
|
||||
"choices": [{
|
||||
"message": {
|
||||
"content": "转写后的文本"
|
||||
}
|
||||
}]
|
||||
}
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import base64
|
||||
import hashlib
|
||||
import time
|
||||
|
||||
import httpx
|
||||
|
||||
from app.ai_provider.audio.base import (
|
||||
ASRProtocol,
|
||||
ASRRequest,
|
||||
ASRResponse,
|
||||
AudioAdapterBase,
|
||||
AudioServiceError,
|
||||
)
|
||||
from app.core.logging import logger
|
||||
|
||||
_MIME_MAP = {
|
||||
"wav": "audio/wav",
|
||||
"mp3": "audio/mpeg",
|
||||
"m4a": "audio/mp4",
|
||||
"flac": "audio/flac",
|
||||
}
|
||||
|
||||
# ASR 缓存：默认最多 128 条，TTL 1 小时
|
||||
_ASR_CACHE_MAX_SIZE = 128
|
||||
_ASR_CACHE_TTL_SECONDS = 3600
|
||||
|
||||
|
||||
class _ASRCacheEntry:
|
||||
__slots__ = ("text", "stored_at")
|
||||
|
||||
def __init__(self, text: str) -> None:
|
||||
self.text = text
|
||||
self.stored_at = time.monotonic()
|
||||
|
||||
|
||||
_asr_cache: dict[str, _ASRCacheEntry] = {}
|
||||
|
||||
|
||||
def _asr_cache_key(model: str, audio_data: bytes, language: str, fmt: str) -> str:
|
||||
"""生成 ASR 缓存键：SHA256(model + audio_bytes + language + format)。"""
|
||||
raw = f"{model}\0{language}\0{fmt}\0"
|
||||
raw_bytes = raw.encode() + audio_data
|
||||
return hashlib.sha256(raw_bytes).hexdigest()
|
||||
|
||||
|
||||
def _asr_cache_get(key: str) -> str | None:
|
||||
entry = _asr_cache.get(key)
|
||||
if entry is None:
|
||||
return None
|
||||
if time.monotonic() - entry.stored_at > _ASR_CACHE_TTL_SECONDS:
|
||||
del _asr_cache[key]
|
||||
return None
|
||||
return entry.text
|
||||
|
||||
|
||||
def _asr_cache_set(key: str, text: str) -> None:
|
||||
if len(_asr_cache) >= _ASR_CACHE_MAX_SIZE:
|
||||
sorted_keys = sorted(_asr_cache.keys(), key=lambda k: _asr_cache[k].stored_at)
|
||||
for old_key in sorted_keys[: _ASR_CACHE_MAX_SIZE // 2]:
|
||||
del _asr_cache[old_key]
|
||||
_asr_cache[key] = _ASRCacheEntry(text)
|
||||
|
||||
|
||||
class HunjingASROpenAI(AudioAdapterBase, ASRProtocol):
|
||||
"""MiMo ASR 适配器，通过 ``POST /chat/completions`` 携带 ``input_audio`` 实现。
|
||||
|
||||
将音频字节包装为 chat 消息中的 ``input_audio`` 内容类型的 data URL，
|
||||
然后从响应中提取转写文本。
|
||||
"""
|
||||
|
||||
_adapter_type = "asr"
|
||||
_default_timeout = 300
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
*,
|
||||
base_url: str,
|
||||
api_key: str,
|
||||
model: str = "mimo-v2.5-asr",
|
||||
request_timeout_seconds: int = 300,
|
||||
) -> None:
|
||||
self._base_url = base_url.rstrip("/")
|
||||
self._api_key = api_key
|
||||
self._model = model
|
||||
self._timeout = request_timeout_seconds
|
||||
self._client: httpx.AsyncClient | None = None
|
||||
|
||||
# ── 客户端 ───────────────────────────────────────────────────────────
|
||||
|
||||
@property
|
||||
def _http(self) -> httpx.AsyncClient:
|
||||
if self._client is None or self._client.is_closed:
|
||||
self._client = httpx.AsyncClient(timeout=httpx.Timeout(self._timeout))
|
||||
return self._client
|
||||
|
||||
async def close(self) -> None:
|
||||
if self._client is not None and not self._client.is_closed:
|
||||
await self._client.aclose()
|
||||
|
||||
# ── ASRProtocol ──────────────────────────────────────────────────────
|
||||
|
||||
async def transcribe(self, request: ASRRequest) -> ASRResponse:
|
||||
"""通过 ``POST /chat/completions`` 转写音频。
|
||||
|
||||
将音频编码为 data URL 并以 ``input_audio`` 内容块的形式发送到 chat 消息中。
|
||||
"""
|
||||
# ── 缓存查找 ──────────────────────────────────────────────────
|
||||
cache_key = _asr_cache_key(
|
||||
self._model, request.audio_data, request.language or "auto", request.format
|
||||
)
|
||||
cached = _asr_cache_get(cache_key)
|
||||
if cached is not None:
|
||||
logger.info("audio_asr.cache_hit", model=self._model, text_length=len(cached))
|
||||
return ASRResponse(text=cached)
|
||||
|
||||
# ── 构造请求 ──────────────────────────────────────────────────
|
||||
mime = _MIME_MAP.get(request.format, "audio/wav")
|
||||
b64_audio = base64.b64encode(request.audio_data).decode()
|
||||
data_url = f"data:{mime};base64,{b64_audio}"
|
||||
|
||||
payload: dict = {
|
||||
"model": self._model,
|
||||
"messages": [
|
||||
{
|
||||
"role": "user",
|
||||
"content": [
|
||||
{
|
||||
"type": "input_audio",
|
||||
"input_audio": {
|
||||
"data": data_url,
|
||||
},
|
||||
},
|
||||
],
|
||||
},
|
||||
],
|
||||
"asr_options": {
|
||||
"language": request.language or "auto",
|
||||
},
|
||||
}
|
||||
|
||||
headers = {
|
||||
"Authorization": f"Bearer {self._api_key}",
|
||||
"Content-Type": "application/json",
|
||||
}
|
||||
|
||||
url = f"{self._base_url}/chat/completions"
|
||||
|
||||
logger.info(
|
||||
"audio_asr.request",
|
||||
model=self._model,
|
||||
format=request.format,
|
||||
language=request.language,
|
||||
url=url,
|
||||
size_bytes=len(request.audio_data),
|
||||
)
|
||||
|
||||
# ── API 调用 ────────────────────────────────────────────────
|
||||
try:
|
||||
response = await self._http.post(url, json=payload, headers=headers)
|
||||
except httpx.TimeoutException as exc:
|
||||
raise AudioServiceError(
|
||||
status_code=504,
|
||||
code="audio_asr_timeout",
|
||||
message="ASR request timed out calling chat completions.",
|
||||
) from exc
|
||||
except httpx.RequestError as exc:
|
||||
raise AudioServiceError(
|
||||
status_code=503,
|
||||
code="audio_asr_connection_error",
|
||||
message=f"ASR connection error: {exc}",
|
||||
) from exc
|
||||
|
||||
if response.status_code != 200:
|
||||
body = response.text[:500]
|
||||
raise AudioServiceError(
|
||||
status_code=response.status_code,
|
||||
code="audio_asr_api_error",
|
||||
message=f"ASR API returned {response.status_code}: {body}",
|
||||
)
|
||||
|
||||
# 解析响应，提取转写文本
|
||||
text = self._extract_text(response.json())
|
||||
if not text:
|
||||
raise AudioServiceError(
|
||||
status_code=500,
|
||||
code="audio_asr_empty_response",
|
||||
message="ASR API returned empty transcription.",
|
||||
)
|
||||
|
||||
# ── 写入缓存 ────────────────────────────────────────────────
|
||||
_asr_cache_set(cache_key, text)
|
||||
|
||||
logger.info(
|
||||
"audio_asr.success",
|
||||
model=self._model,
|
||||
text_length=len(text),
|
||||
cache_size=len(_asr_cache),
|
||||
)
|
||||
return ASRResponse(text=text)
|
||||
|
||||
# ── 辅助方法 ──────────────────────────────────────────────────────────
|
||||
|
||||
@staticmethod
|
||||
def _extract_text(body: dict) -> str:
|
||||
"""从 chat completion 响应中提取转写文本。"""
|
||||
try:
|
||||
choices = body.get("choices", [])
|
||||
if not choices:
|
||||
raise AudioServiceError(
|
||||
status_code=500,
|
||||
code="audio_asr_no_choices",
|
||||
message="ASR response has no choices.",
|
||||
)
|
||||
return (choices[0].get("message") or {}).get("content") or ""
|
||||
except AudioServiceError:
|
||||
raise
|
||||
except Exception as exc:
|
||||
raise AudioServiceError(
|
||||
status_code=500,
|
||||
code="audio_asr_response_parse_error",
|
||||
message=f"Failed to parse ASR response: {exc}",
|
||||
) from exc
|
||||
@@ -0,0 +1,245 @@
|
||||
"""音频服务类型 — 请求/响应模型和协议定义。
|
||||
|
||||
ASR（自动语音识别）和 TTS（文本转语音）适配器**不是** LangChain 的
|
||||
``BaseChatModel`` 子类 —— 它们是通过 HTTP 调用的服务，经由 ``AudioService``
|
||||
或直接从 ``ModelFactory`` 消费。
|
||||
|
||||
用法::
|
||||
|
||||
service = AudioService(factory)
|
||||
result = await service.transcribe(audio_data)
|
||||
result = await service.synthesize("你好")
|
||||
|
||||
# 或通过工厂直接调用:
|
||||
tts = factory.create_audio("mimo-v2.5-tts")
|
||||
result = await tts.synthesize(TTSRequest(text="你好"))
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from dataclasses import dataclass
|
||||
from typing import TYPE_CHECKING, Any, Protocol, runtime_checkable
|
||||
|
||||
from app.ai_provider.types.errors import AIProviderError
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from app.ai_provider.config.models import ModelConfig, ModelEndpointConfig
|
||||
|
||||
|
||||
class AudioServiceError(AIProviderError):
|
||||
"""ASR / TTS 服务失败的基础错误。"""
|
||||
|
||||
|
||||
# ── ASR ───────────────────────────────────────────────────────────────────
|
||||
|
||||
|
||||
class ASRRequest:
|
||||
"""将音频转写为文本的请求。
|
||||
|
||||
属性:
|
||||
audio_data: 原始音频字节。
|
||||
format: 音频容器格式（如 ``"wav"``、``"mp3"``、``"m4a"``）。
|
||||
language: 可选的 ISO-639-1 语言提示（如 ``"zh"``、``"en"``）。
|
||||
"""
|
||||
|
||||
__slots__ = ("audio_data", "format", "language")
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
audio_data: bytes,
|
||||
format: str = "wav",
|
||||
language: str | None = None,
|
||||
) -> None:
|
||||
self.audio_data = audio_data
|
||||
self.format = format
|
||||
self.language = language
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class ASRResponse:
|
||||
"""成功转写的结果。"""
|
||||
|
||||
text: str
|
||||
|
||||
|
||||
@runtime_checkable
|
||||
class ASRProtocol(Protocol):
|
||||
"""ASR 适配器的协议。"""
|
||||
|
||||
async def transcribe(self, request: ASRRequest) -> ASRResponse:
|
||||
"""将音频字节转写为文本。
|
||||
|
||||
异常:
|
||||
AudioServiceError: API 调用失败或响应异常时抛出。
|
||||
"""
|
||||
...
|
||||
|
||||
|
||||
# ── TTS ───────────────────────────────────────────────────────────────────
|
||||
|
||||
|
||||
class TTSRequest:
|
||||
"""将文本合成为语音的请求。
|
||||
|
||||
属性:
|
||||
text: 待朗读的文本。
|
||||
voice: 音色标识符（模型决定可用音色）。
|
||||
response_format: 期望的音频格式（``"wav"``、``"mp3"``、``"opus"``）。
|
||||
extra: 厂商特定的额外参数（如音色克隆/设计参数）。
|
||||
"""
|
||||
|
||||
__slots__ = ("text", "voice", "response_format", "extra")
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
text: str,
|
||||
voice: str = "default",
|
||||
response_format: str = "wav",
|
||||
extra: dict | None = None,
|
||||
) -> None:
|
||||
self.text = text
|
||||
self.voice = voice
|
||||
self.response_format = response_format
|
||||
self.extra = extra or {}
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class TTSResponse:
|
||||
"""成功合成的结果。"""
|
||||
|
||||
audio_data: bytes
|
||||
format: str
|
||||
|
||||
|
||||
@runtime_checkable
|
||||
class TTSProtocol(Protocol):
|
||||
"""TTS 适配器的协议。"""
|
||||
|
||||
async def synthesize(self, request: TTSRequest) -> TTSResponse:
|
||||
"""将文本合成为音频字节。
|
||||
|
||||
异常:
|
||||
AudioServiceError: API 调用失败或响应异常时抛出。
|
||||
"""
|
||||
...
|
||||
|
||||
|
||||
# ── 共享工厂基类 ──────────────────────────────────────────────────────────
|
||||
|
||||
|
||||
class AudioAdapterBase:
|
||||
"""ASR/TTS 适配器的基类，提供共享的工厂方法。
|
||||
|
||||
子类必须设置 ``_adapter_type``（如 ``"asr"``、``"tts"``），并在
|
||||
``__init__`` 中接受 ``base_url``、``api_key``、``model``、
|
||||
``request_timeout_seconds`` 参数。
|
||||
"""
|
||||
|
||||
_adapter_type: str = "audio"
|
||||
_default_timeout: int = 120
|
||||
|
||||
@classmethod
|
||||
def from_endpoint(
|
||||
cls,
|
||||
endpoint: ModelEndpointConfig,
|
||||
*,
|
||||
provider: ProviderConfig | None = None,
|
||||
timeout_seconds: int | None = None,
|
||||
) -> Any:
|
||||
"""从已解析的端点配置构建适配器。
|
||||
|
||||
凭证解析优先级：
|
||||
1. endpoint 直填的 ``base_url`` / ``api_key``
|
||||
2. endpoint 的 ``base_url_env`` / ``api_key_env`` 指向的 settings 字段
|
||||
3. provider 级凭证（调用方按需传入 ProviderConfig）
|
||||
"""
|
||||
base_url = _resolve_base_url(endpoint, provider)
|
||||
api_key = _resolve_api_key(endpoint, provider)
|
||||
if not base_url:
|
||||
raise AudioServiceError(
|
||||
status_code=500,
|
||||
code=f"audio_{cls._adapter_type}_endpoint_invalid",
|
||||
message=(
|
||||
f"{cls._adapter_type.upper()} endpoint '{endpoint.name}' "
|
||||
"has no resolved base_url."
|
||||
),
|
||||
)
|
||||
return cls(
|
||||
base_url=base_url,
|
||||
api_key=api_key,
|
||||
model=endpoint.model,
|
||||
request_timeout_seconds=timeout_seconds
|
||||
or endpoint.request_timeout_seconds
|
||||
or cls._default_timeout,
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def from_model_config(
|
||||
cls,
|
||||
config: ModelConfig,
|
||||
*,
|
||||
provider: ProviderConfig | None = None,
|
||||
timeout_seconds: int | None = None,
|
||||
) -> Any:
|
||||
"""从逻辑模型配置构建适配器（使用第一个活跃端点）。"""
|
||||
for ep in config.endpoints:
|
||||
if ep.status == "active":
|
||||
return cls.from_endpoint(
|
||||
ep, provider=provider, timeout_seconds=timeout_seconds
|
||||
)
|
||||
raise AudioServiceError(
|
||||
status_code=500,
|
||||
code=f"audio_{cls._adapter_type}_no_endpoint",
|
||||
message=f"{cls._adapter_type.upper()} model '{config.name}' has no active endpoint.",
|
||||
)
|
||||
|
||||
|
||||
def _resolve_base_url(
|
||||
endpoint: ModelEndpointConfig,
|
||||
provider: ProviderConfig | None,
|
||||
) -> str:
|
||||
"""端点 base_url → endpoint.base_url_env 对应 settings → provider.base_url(_env)。"""
|
||||
if endpoint.base_url:
|
||||
return endpoint.base_url
|
||||
if endpoint.base_url_env:
|
||||
val = _settings_get(endpoint.base_url_env)
|
||||
if val:
|
||||
return val
|
||||
if provider is not None:
|
||||
if provider.base_url:
|
||||
return provider.base_url
|
||||
if provider.base_url_env:
|
||||
val = _settings_get(provider.base_url_env)
|
||||
if val:
|
||||
return val
|
||||
return ""
|
||||
|
||||
|
||||
def _resolve_api_key(
|
||||
endpoint: ModelEndpointConfig,
|
||||
provider: ProviderConfig | None,
|
||||
) -> str:
|
||||
"""端点 api_key → endpoint.api_key_env 对应 settings → provider 级凭证。"""
|
||||
if endpoint.api_key:
|
||||
return endpoint.api_key
|
||||
if endpoint.api_key_env:
|
||||
val = _settings_get(endpoint.api_key_env)
|
||||
if val:
|
||||
return val
|
||||
if provider is not None:
|
||||
# ProviderConfig 的 api_key_env / base_url_env 也走 settings 解析
|
||||
val = _settings_get(provider.api_key_env)
|
||||
if val:
|
||||
return val
|
||||
return ""
|
||||
|
||||
|
||||
def _settings_get(env_name: str) -> str | None:
|
||||
"""按大写 env 名查 settings 上的小写字段（与现有 _resolve_credentials 约定一致）。"""
|
||||
from app.core.config import settings
|
||||
|
||||
key = env_name.lower()
|
||||
if hasattr(settings, key):
|
||||
value = getattr(settings, key)
|
||||
return value if isinstance(value, str) else None
|
||||
return None
|
||||
@@ -0,0 +1,325 @@
|
||||
"""MiMo TTS adapter — ``POST /v1/chat/completions`` with ``audio`` parameter.
|
||||
|
||||
MiMo 的 TTS 不走独立音频端点，而是通过聊天补全接口扩展实现：
|
||||
|
||||
POST {base_url}/chat/completions
|
||||
Content-Type: application/json
|
||||
Authorization: Bearer {api_key}
|
||||
|
||||
{
|
||||
"model": "mimo-v2.5-tts",
|
||||
"messages": [
|
||||
{"role": "user", "content": "（可选）音色风格描述，voicedesign 必填"},
|
||||
{"role": "assistant", "content": "要朗读的文本"}
|
||||
],
|
||||
"audio": {
|
||||
"format": "wav",
|
||||
"voice": "mimo_default"
|
||||
}
|
||||
}
|
||||
|
||||
响应：:
|
||||
|
||||
{
|
||||
"choices": [{
|
||||
"message": {
|
||||
"audio": {
|
||||
"data": "<base64 编码的音频数据>",
|
||||
"id": "...",
|
||||
"expires_at": null,
|
||||
"transcript": null
|
||||
}
|
||||
}
|
||||
}]
|
||||
}
|
||||
|
||||
模型差异：
|
||||
|
||||
- ``mimo-v2.5-tts`` — ``voice`` 可选，默认 ``mimo_default``，仅支持预置音色
|
||||
- ``mimo-v2.5-tts-voiceclone`` — ``voice`` 必填，传入音频样本的 base64（mp3/wav）
|
||||
- ``mimo-v2.5-tts-voicedesign`` — 不支持 ``voice``，\n ``user`` 消息必填（风格描述），支持 ``optimize_text_preview``
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import base64
|
||||
import hashlib
|
||||
import time
|
||||
|
||||
import httpx
|
||||
|
||||
from app.ai_provider.audio.base import (
|
||||
AudioAdapterBase,
|
||||
AudioServiceError,
|
||||
TTSProtocol,
|
||||
TTSRequest,
|
||||
TTSResponse,
|
||||
)
|
||||
from app.core.logging import logger
|
||||
|
||||
# voicedesign 模型的默认风格描述（用户未提供时使用）
|
||||
_DEFAULT_VOICE_DESIGN_PROMPT = "自然、清晰的朗读风格，语气平和，语速适中，适合日常对话场景。"
|
||||
|
||||
# TTS 音频缓存：默认最多 256 条，TTL 1 小时
|
||||
_TTS_CACHE_MAX_SIZE = 256
|
||||
_TTS_CACHE_TTL_SECONDS = 3600
|
||||
|
||||
|
||||
class _TTSCacheEntry:
|
||||
__slots__ = ("data", "format", "stored_at")
|
||||
|
||||
def __init__(self, data: bytes, format: str) -> None:
|
||||
self.data = data
|
||||
self.format = format
|
||||
self.stored_at = time.monotonic()
|
||||
|
||||
|
||||
_tts_cache: dict[str, _TTSCacheEntry] = {}
|
||||
|
||||
|
||||
def _tts_cache_key(model: str, voice: str, fmt: str, text: str, extra_json: str) -> str:
|
||||
"""生成 TTS 缓存键：SHA256(model + voice + format + text + extra)。"""
|
||||
raw = f"{model}\0{voice}\0{fmt}\0{text}\0{extra_json}"
|
||||
return hashlib.sha256(raw.encode()).hexdigest()
|
||||
|
||||
|
||||
def _tts_cache_get(key: str) -> _TTSCacheEntry | None:
|
||||
entry = _tts_cache.get(key)
|
||||
if entry is None:
|
||||
return None
|
||||
if time.monotonic() - entry.stored_at > _TTS_CACHE_TTL_SECONDS:
|
||||
del _tts_cache[key]
|
||||
return None
|
||||
return entry
|
||||
|
||||
|
||||
def _tts_cache_set(key: str, entry: _TTSCacheEntry) -> None:
|
||||
if len(_tts_cache) >= _TTS_CACHE_MAX_SIZE:
|
||||
# 淘汰最旧的一半
|
||||
sorted_keys = sorted(_tts_cache.keys(), key=lambda k: _tts_cache[k].stored_at)
|
||||
for old_key in sorted_keys[: _TTS_CACHE_MAX_SIZE // 2]:
|
||||
del _tts_cache[old_key]
|
||||
_tts_cache[key] = entry
|
||||
|
||||
|
||||
class HunjingTTSOpenAI(AudioAdapterBase, TTSProtocol):
|
||||
"""MiMo TTS 适配器，通过 ``POST /chat/completions`` 携带 ``audio`` 参数实现。
|
||||
|
||||
调用兼容 OpenAI 格式的 chat completions 端点，在请求体中添加 ``audio`` 字段，
|
||||
然后从响应中提取并解码 base64 音频数据。
|
||||
"""
|
||||
|
||||
_adapter_type = "tts"
|
||||
_default_timeout = 120
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
*,
|
||||
base_url: str,
|
||||
api_key: str,
|
||||
model: str = "mimo-v2.5-tts",
|
||||
request_timeout_seconds: int = 120,
|
||||
) -> None:
|
||||
self._base_url = base_url.rstrip("/")
|
||||
self._api_key = api_key
|
||||
self._model = model
|
||||
self._timeout = request_timeout_seconds
|
||||
self._client: httpx.AsyncClient | None = None
|
||||
|
||||
# ── 客户端 ───────────────────────────────────────────────────────────
|
||||
|
||||
@property
|
||||
def _http(self) -> httpx.AsyncClient:
|
||||
if self._client is None or self._client.is_closed:
|
||||
self._client = httpx.AsyncClient(timeout=httpx.Timeout(self._timeout))
|
||||
return self._client
|
||||
|
||||
async def close(self) -> None:
|
||||
if self._client is not None and not self._client.is_closed:
|
||||
await self._client.aclose()
|
||||
|
||||
# ── TTSProtocol ──────────────────────────────────────────────────────
|
||||
|
||||
async def synthesize(self, request: TTSRequest) -> TTSResponse:
|
||||
"""通过 ``POST /chat/completions`` 将文本合成为音频。
|
||||
|
||||
构造 chat 消息和 ``audio`` 参数，匹配 MiMo 兼容 OpenAI 格式的 TTS，
|
||||
然后从响应中解码 base64 音频。
|
||||
|
||||
结果按 ``(model, voice, format, text)`` 在进程内缓存，避免重复 API 调用。
|
||||
"""
|
||||
# ── 缓存查找 ──────────────────────────────────────────────────────
|
||||
cache_key = _tts_cache_key(
|
||||
model=self._model,
|
||||
voice=request.voice,
|
||||
fmt=request.response_format,
|
||||
text=request.text,
|
||||
extra_json=str(sorted((request.extra or {}).items())),
|
||||
)
|
||||
cached = _tts_cache_get(cache_key)
|
||||
if cached is not None:
|
||||
logger.info("audio_tts.cache_hit", model=self._model, text_length=len(request.text))
|
||||
return TTSResponse(audio_data=cached.data, format=cached.format)
|
||||
|
||||
# ── API 调用 ──────────────────────────────────────────────────────
|
||||
url = f"{self._base_url}/chat/completions"
|
||||
|
||||
# 构造 messages
|
||||
messages = self._build_messages(request)
|
||||
|
||||
# 构造 audio 参数
|
||||
audio_params = self._build_audio_params(request)
|
||||
|
||||
payload: dict = {
|
||||
"model": self._model,
|
||||
"messages": messages,
|
||||
"audio": audio_params,
|
||||
"max_completion_tokens": 8192,
|
||||
}
|
||||
# 合并 extra（如 frequency_penalty, temperature 等）
|
||||
if request.extra:
|
||||
payload.update(request.extra)
|
||||
|
||||
headers = {
|
||||
"Authorization": f"Bearer {self._api_key}",
|
||||
"Content-Type": "application/json",
|
||||
}
|
||||
|
||||
logger.info(
|
||||
"audio_tts.request",
|
||||
model=self._model,
|
||||
voice=request.voice,
|
||||
audio_format=request.response_format,
|
||||
text_length=len(request.text),
|
||||
url=url,
|
||||
)
|
||||
|
||||
try:
|
||||
response = await self._http.post(url, json=payload, headers=headers)
|
||||
except httpx.TimeoutException as exc:
|
||||
raise AudioServiceError(
|
||||
status_code=504,
|
||||
code="audio_tts_timeout",
|
||||
message="TTS request timed out calling chat completions.",
|
||||
) from exc
|
||||
except httpx.RequestError as exc:
|
||||
raise AudioServiceError(
|
||||
status_code=503,
|
||||
code="audio_tts_connection_error",
|
||||
message=f"TTS connection error: {exc}",
|
||||
) from exc
|
||||
|
||||
if response.status_code != 200:
|
||||
body = response.text[:500]
|
||||
raise AudioServiceError(
|
||||
status_code=response.status_code,
|
||||
code="audio_tts_api_error",
|
||||
message=f"TTS API returned {response.status_code}: {body}",
|
||||
)
|
||||
|
||||
# 解析响应，提取 base64 音频
|
||||
audio_data = self._extract_audio(response.json())
|
||||
|
||||
# ── 写入缓存 ──────────────────────────────────────────────────────
|
||||
_tts_cache_set(cache_key, _TTSCacheEntry(audio_data, request.response_format))
|
||||
|
||||
logger.info(
|
||||
"audio_tts.success",
|
||||
model=self._model,
|
||||
audio_bytes=len(audio_data),
|
||||
cache_size=len(_tts_cache),
|
||||
)
|
||||
return TTSResponse(audio_data=audio_data, format=request.response_format)
|
||||
|
||||
# ── 辅助方法 ──────────────────────────────────────────────────────────
|
||||
|
||||
def _build_messages(self, request: TTSRequest) -> list[dict]:
|
||||
"""构造 chat messages 数组。
|
||||
|
||||
- voicedesign 模型：user 消息必填（风格描述）
|
||||
- 其他模型：user 消息可选，assistant 消息存放待朗读文本
|
||||
"""
|
||||
is_voicedesign = "voicedesign" in self._model
|
||||
|
||||
if is_voicedesign:
|
||||
# voicedesign：user 消息为风格提示，assistant 消息为朗读文本
|
||||
style_prompt = (request.extra or {}).pop("style_prompt", _DEFAULT_VOICE_DESIGN_PROMPT)
|
||||
return [
|
||||
{"role": "user", "content": style_prompt},
|
||||
{"role": "assistant", "content": request.text},
|
||||
]
|
||||
# 基础 TTS / voiceclone：只传 assistant 消息
|
||||
return [
|
||||
{"role": "assistant", "content": request.text},
|
||||
]
|
||||
|
||||
def _build_audio_params(self, request: TTSRequest) -> dict:
|
||||
"""构造 ``audio`` 参数字段，按模型处理 voice/format/optimize_text_preview 等。"""
|
||||
params: dict = {
|
||||
"format": _resolve_api_format(request.response_format),
|
||||
}
|
||||
|
||||
is_voicedesign = "voicedesign" in self._model
|
||||
|
||||
if is_voicedesign:
|
||||
# voicedesign 不支持 voice 字段
|
||||
optimize = (request.extra or {}).pop("optimize_text_preview", None)
|
||||
if optimize is not None:
|
||||
params["optimize_text_preview"] = optimize
|
||||
else:
|
||||
# 基础 TTS / voiceclone：按需传 voice
|
||||
if request.voice and request.voice != "default":
|
||||
params["voice"] = request.voice
|
||||
else:
|
||||
params["voice"] = "mimo_default"
|
||||
|
||||
return params
|
||||
|
||||
@staticmethod
|
||||
def _extract_audio(body: dict) -> bytes:
|
||||
"""从 chat completion 响应中提取并解码 base64 音频数据。"""
|
||||
try:
|
||||
choices = body.get("choices", [])
|
||||
if not choices:
|
||||
raise AudioServiceError(
|
||||
status_code=500,
|
||||
code="audio_tts_no_choices",
|
||||
message="TTS response has no choices.",
|
||||
)
|
||||
message = choices[0].get("message", {})
|
||||
audio_info = message.get("audio")
|
||||
if not audio_info:
|
||||
raise AudioServiceError(
|
||||
status_code=500,
|
||||
code="audio_tts_no_audio_field",
|
||||
message="TTS response missing audio field in message.",
|
||||
)
|
||||
raw = audio_info.get("data")
|
||||
if not raw:
|
||||
raise AudioServiceError(
|
||||
status_code=500,
|
||||
code="audio_tts_empty_audio_data",
|
||||
message="TTS response has empty audio.data.",
|
||||
)
|
||||
return base64.b64decode(raw)
|
||||
except AudioServiceError:
|
||||
raise
|
||||
except Exception as exc:
|
||||
raise AudioServiceError(
|
||||
status_code=500,
|
||||
code="audio_tts_response_parse_error",
|
||||
message=f"Failed to parse TTS response: {exc}",
|
||||
) from exc
|
||||
|
||||
|
||||
def _resolve_api_format(user_format: str) -> str:
|
||||
"""将用户层面的格式名映射为 API 期望的格式。
|
||||
|
||||
MiMo TTS API 接受：``wav``、``mp3``、``pcm``、``pcm16``。
|
||||
"""
|
||||
fmt = user_format.lower()
|
||||
if fmt in ("pcm", "pcm16"):
|
||||
return "pcm16"
|
||||
if fmt in ("wav", "mp3"):
|
||||
return fmt
|
||||
return "wav"
|
||||
@@ -0,0 +1,231 @@
|
||||
"""yuanqi 本地 ASR 适配器 — ``POST /v1/audio/transcriptions`` (multipart)。
|
||||
|
||||
与 MiMo 走 ``/chat/completions`` + ``input_audio`` 不同，元炁本地推理引擎
|
||||
（yuanqi）对外暴露的是 OpenAI Whisper 兼容的 multipart 端点::
|
||||
|
||||
POST {base_url}/v1/audio/transcriptions
|
||||
Content-Type: multipart/form-data
|
||||
Authorization: Bearer {api_key}
|
||||
|
||||
file: <audio bytes>
|
||||
language: zh (可选)
|
||||
|
||||
响应::
|
||||
|
||||
{"text": "转写后的文本"}
|
||||
|
||||
模型实际为 FunASR SenseVoiceSmall，转写文本已在 yuanqi 侧剥离了
|
||||
``<|zh|><|NEUTRAL|><|Speech|><|withitn|>`` 等内部标签，此处无需二次清洗。
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import hashlib
|
||||
import time
|
||||
|
||||
import httpx
|
||||
|
||||
from app.ai_provider.audio.base import (
|
||||
ASRProtocol,
|
||||
ASRRequest,
|
||||
ASRResponse,
|
||||
AudioAdapterBase,
|
||||
AudioServiceError,
|
||||
)
|
||||
from app.core.logging import logger
|
||||
|
||||
# multipart 使用的 MIME 映射（与 yuanqi _SUPPORTED_FORMATS 白名单对齐）
|
||||
_MIME_MAP = {
|
||||
"wav": "audio/wav",
|
||||
"mp3": "audio/mpeg",
|
||||
"m4a": "audio/mp4",
|
||||
"flac": "audio/flac",
|
||||
"ogg": "audio/ogg",
|
||||
"opus": "audio/opus",
|
||||
"aac": "audio/aac",
|
||||
"webm": "audio/webm",
|
||||
}
|
||||
|
||||
# yuanqi ASR 路径前缀（与 src/endpoints/asr.py 的 router 前缀一致）
|
||||
_ASR_PATH = "/v1/audio/transcriptions"
|
||||
|
||||
# ASR 缓存：默认最多 128 条，TTL 1 小时（与 MiMo ASR 适配器一致）
|
||||
_ASR_CACHE_MAX_SIZE = 128
|
||||
_ASR_CACHE_TTL_SECONDS = 3600
|
||||
|
||||
|
||||
class _ASRCacheEntry:
|
||||
__slots__ = ("text", "stored_at")
|
||||
|
||||
def __init__(self, text: str) -> None:
|
||||
self.text = text
|
||||
self.stored_at = time.monotonic()
|
||||
|
||||
|
||||
_asr_cache: dict[str, _ASRCacheEntry] = {}
|
||||
|
||||
|
||||
def _asr_cache_key(model: str, audio_data: bytes, language: str, fmt: str) -> str:
|
||||
"""生成 ASR 缓存键：SHA256(model + audio_bytes + language + format)。"""
|
||||
raw = f"{model}\0{language}\0{fmt}\0"
|
||||
raw_bytes = raw.encode() + audio_data
|
||||
return hashlib.sha256(raw_bytes).hexdigest()
|
||||
|
||||
|
||||
def _asr_cache_get(key: str) -> str | None:
|
||||
entry = _asr_cache.get(key)
|
||||
if entry is None:
|
||||
return None
|
||||
if time.monotonic() - entry.stored_at > _ASR_CACHE_TTL_SECONDS:
|
||||
del _asr_cache[key]
|
||||
return None
|
||||
return entry.text
|
||||
|
||||
|
||||
def _asr_cache_set(key: str, text: str) -> None:
|
||||
if len(_asr_cache) >= _ASR_CACHE_MAX_SIZE:
|
||||
sorted_keys = sorted(_asr_cache.keys(), key=lambda k: _asr_cache[k].stored_at)
|
||||
for old_key in sorted_keys[: _ASR_CACHE_MAX_SIZE // 2]:
|
||||
del _asr_cache[old_key]
|
||||
_asr_cache[key] = _ASRCacheEntry(text)
|
||||
|
||||
|
||||
class YuanqiASRWhisper(AudioAdapterBase, ASRProtocol):
|
||||
"""yuanqi 本地 ASR 适配器，走 OpenAI Whisper 兼容的 multipart 端点。
|
||||
|
||||
将音频字节以 ``multipart/form-data`` 的 ``file`` 字段上传到 yuanqi 的
|
||||
``/v1/audio/transcriptions``，从 JSON 响应中提取转写文本。
|
||||
"""
|
||||
|
||||
_adapter_type = "asr"
|
||||
_default_timeout = 300
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
*,
|
||||
base_url: str,
|
||||
api_key: str,
|
||||
model: str = "SenseVoiceSmall",
|
||||
request_timeout_seconds: int = 300,
|
||||
) -> None:
|
||||
self._base_url = base_url.rstrip("/")
|
||||
self._api_key = api_key
|
||||
self._model = model
|
||||
self._timeout = request_timeout_seconds
|
||||
self._client: httpx.AsyncClient | None = None
|
||||
|
||||
# ── 客户端 ───────────────────────────────────────────────────────────
|
||||
|
||||
@property
|
||||
def _http(self) -> httpx.AsyncClient:
|
||||
if self._client is None or self._client.is_closed:
|
||||
self._client = httpx.AsyncClient(timeout=httpx.Timeout(self._timeout))
|
||||
return self._client
|
||||
|
||||
async def close(self) -> None:
|
||||
if self._client is not None and not self._client.is_closed:
|
||||
await self._client.aclose()
|
||||
|
||||
# ── ASRProtocol ──────────────────────────────────────────────────────
|
||||
|
||||
async def transcribe(self, request: ASRRequest) -> ASRResponse:
|
||||
"""通过 ``POST /v1/audio/transcriptions`` 转写音频。"""
|
||||
# ── 缓存查找 ──────────────────────────────────────────────────
|
||||
cache_key = _asr_cache_key(
|
||||
self._model, request.audio_data, request.language or "auto", request.format
|
||||
)
|
||||
cached = _asr_cache_get(cache_key)
|
||||
if cached is not None:
|
||||
logger.info("audio_asr.cache_hit", model=self._model, text_length=len(cached))
|
||||
return ASRResponse(text=cached)
|
||||
|
||||
# ── 构造 multipart 请求 ───────────────────────────────────────
|
||||
mime = _MIME_MAP.get(request.format, "audio/wav")
|
||||
filename = f"audio.{request.format}"
|
||||
files = {"file": (filename, request.audio_data, mime)}
|
||||
data: dict[str, str] = {}
|
||||
if request.language:
|
||||
data["language"] = request.language
|
||||
|
||||
# 本地 yuanqi 引擎通常无需鉴权；仅当配置了 api_key 才带 Authorization 头
|
||||
headers: dict[str, str] = {}
|
||||
if self._api_key:
|
||||
headers["Authorization"] = f"Bearer {self._api_key}"
|
||||
|
||||
url = f"{self._base_url}{_ASR_PATH}"
|
||||
|
||||
logger.info(
|
||||
"audio_asr.request",
|
||||
model=self._model,
|
||||
format=request.format,
|
||||
language=request.language,
|
||||
url=url,
|
||||
size_bytes=len(request.audio_data),
|
||||
)
|
||||
|
||||
# ── API 调用 ────────────────────────────────────────────────
|
||||
try:
|
||||
response = await self._http.post(url, files=files, data=data, headers=headers)
|
||||
except httpx.TimeoutException as exc:
|
||||
raise AudioServiceError(
|
||||
status_code=504,
|
||||
code="audio_asr_timeout",
|
||||
message="ASR request timed out calling yuanqi transcriptions.",
|
||||
) from exc
|
||||
except httpx.RequestError as exc:
|
||||
raise AudioServiceError(
|
||||
status_code=503,
|
||||
code="audio_asr_connection_error",
|
||||
message=f"ASR connection error: {exc}",
|
||||
) from exc
|
||||
|
||||
if response.status_code != 200:
|
||||
body = response.text[:500]
|
||||
raise AudioServiceError(
|
||||
status_code=response.status_code,
|
||||
code="audio_asr_api_error",
|
||||
message=f"ASR API returned {response.status_code}: {body}",
|
||||
)
|
||||
|
||||
# 解析响应，提取转写文本
|
||||
text = self._extract_text(response.json())
|
||||
if not text:
|
||||
raise AudioServiceError(
|
||||
status_code=500,
|
||||
code="audio_asr_empty_response",
|
||||
message="ASR API returned empty transcription.",
|
||||
)
|
||||
|
||||
# ── 写入缓存 ────────────────────────────────────────────────
|
||||
_asr_cache_set(cache_key, text)
|
||||
|
||||
logger.info(
|
||||
"audio_asr.success",
|
||||
model=self._model,
|
||||
text_length=len(text),
|
||||
cache_size=len(_asr_cache),
|
||||
)
|
||||
return ASRResponse(text=text)
|
||||
|
||||
# ── 辅助方法 ──────────────────────────────────────────────────────────
|
||||
|
||||
@staticmethod
|
||||
def _extract_text(body: dict) -> str:
|
||||
"""从 Whisper 兼容响应中提取转写文本。"""
|
||||
try:
|
||||
text = body.get("text")
|
||||
if text is None:
|
||||
raise AudioServiceError(
|
||||
status_code=500,
|
||||
code="audio_asr_no_text",
|
||||
message="ASR response has no 'text' field.",
|
||||
)
|
||||
return text or ""
|
||||
except AudioServiceError:
|
||||
raise
|
||||
except Exception as exc:
|
||||
raise AudioServiceError(
|
||||
status_code=500,
|
||||
code="audio_asr_response_parse_error",
|
||||
message=f"Failed to parse ASR response: {exc}",
|
||||
) from exc
|
||||
@@ -0,0 +1,31 @@
|
||||
from app.ai_provider.config.loader import (
|
||||
get_models_config_path,
|
||||
get_pricing_config_path,
|
||||
get_providers_config_path,
|
||||
load_models_config,
|
||||
load_pricing_config,
|
||||
load_providers_config,
|
||||
)
|
||||
from app.ai_provider.config.models import (
|
||||
ModelConfig,
|
||||
ModelEndpointConfig,
|
||||
ParameterMapping,
|
||||
ProviderConfig,
|
||||
derive_capabilities,
|
||||
)
|
||||
from app.ai_provider.core.billing import PricingInfo
|
||||
|
||||
__all__ = [
|
||||
"ModelConfig",
|
||||
"ModelEndpointConfig",
|
||||
"PricingInfo",
|
||||
"ProviderConfig",
|
||||
"ParameterMapping",
|
||||
"derive_capabilities",
|
||||
"load_models_config",
|
||||
"load_providers_config",
|
||||
"load_pricing_config",
|
||||
"get_models_config_path",
|
||||
"get_providers_config_path",
|
||||
"get_pricing_config_path",
|
||||
]
|
||||
@@ -0,0 +1,80 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
|
||||
import yaml
|
||||
from pydantic import BaseModel
|
||||
|
||||
from app.ai_provider.config.models import ModelConfig, ProviderConfig
|
||||
from app.core.config import settings
|
||||
|
||||
|
||||
class _ModelsFile(BaseModel):
|
||||
version: int = 1
|
||||
models: list[ModelConfig]
|
||||
|
||||
|
||||
class _ProvidersFile(BaseModel):
|
||||
version: int = 1
|
||||
providers: dict[str, ProviderConfig]
|
||||
|
||||
|
||||
class _PricingFile(BaseModel):
|
||||
"""pricing.yaml 顶层结构。"""
|
||||
|
||||
version: int = 1
|
||||
pricing: dict[str, dict] = {}
|
||||
|
||||
|
||||
def get_models_config_path() -> Path:
|
||||
"""返回已解析的 models.yaml 路径（project_root/config/models.yaml）。"""
|
||||
return settings.project_root / "config" / "models.yaml"
|
||||
|
||||
|
||||
def get_providers_config_path() -> Path:
|
||||
"""返回已解析的 providers.yaml 路径（project_root/config/providers.yaml）。"""
|
||||
return settings.project_root / "config" / "providers.yaml"
|
||||
|
||||
|
||||
def load_models_config(path: Path | None = None) -> list[ModelConfig]:
|
||||
"""从 YAML 文件加载模型配置。如果未指定路径，使用配置中的默认路径。"""
|
||||
if path is None:
|
||||
path = get_models_config_path()
|
||||
if not path.exists():
|
||||
raise FileNotFoundError(f"Models config not found: {path}")
|
||||
raw: dict[str, Any] = yaml.safe_load(path.read_text(encoding="utf-8")) or {}
|
||||
parsed = _ModelsFile.model_validate(raw)
|
||||
return parsed.models
|
||||
|
||||
|
||||
def load_providers_config(path: Path | None = None) -> dict[str, ProviderConfig]:
|
||||
"""从 YAML 文件加载提供商配置。如果未指定路径，使用配置中的默认路径。"""
|
||||
if path is None:
|
||||
path = get_providers_config_path()
|
||||
if not path.exists():
|
||||
raise FileNotFoundError(f"Providers config not found: {path}")
|
||||
raw: dict[str, Any] = yaml.safe_load(path.read_text(encoding="utf-8")) or {}
|
||||
# 验证前将字典键注入为 provider 的 name 字段
|
||||
providers_raw = raw.get("providers", {})
|
||||
for name, cfg in providers_raw.items():
|
||||
if isinstance(cfg, dict) and "name" not in cfg:
|
||||
cfg["name"] = name
|
||||
parsed = _ProvidersFile.model_validate(raw)
|
||||
return parsed.providers
|
||||
|
||||
|
||||
def get_pricing_config_path() -> Path:
|
||||
"""返回 pricing.yaml 路径（project_root/config/pricing.yaml）。"""
|
||||
return settings.project_root / "config" / "pricing.yaml"
|
||||
|
||||
|
||||
def load_pricing_config(path: Path | None = None) -> dict[str, dict]:
|
||||
"""从 YAML 文件加载定价配置。"""
|
||||
if path is None:
|
||||
path = get_pricing_config_path()
|
||||
if not path.exists():
|
||||
return {}
|
||||
raw: dict[str, Any] = yaml.safe_load(path.read_text(encoding="utf-8")) or {}
|
||||
parsed = _PricingFile.model_validate(raw)
|
||||
return parsed.pricing
|
||||
@@ -0,0 +1,158 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from enum import StrEnum
|
||||
from typing import Any
|
||||
|
||||
from pydantic import ( # noqa: TC002 — Pydantic 模型
|
||||
BaseModel,
|
||||
Field,
|
||||
field_validator,
|
||||
model_validator,
|
||||
)
|
||||
|
||||
|
||||
class SupportedFormats(StrEnum):
|
||||
OPENAI = "openai"
|
||||
ANTHROPIC = "anthropic"
|
||||
|
||||
|
||||
_BASE_CAPABILITIES = {
|
||||
("openai", "gpt"): {"chat", "tool_use", "structured_output"},
|
||||
("openai", "gpt-4"): {"chat", "tool_use", "structured_output", "vision", "reasoning"},
|
||||
("anthropic", "claude"): {
|
||||
"chat",
|
||||
"tool_use",
|
||||
"structured_output",
|
||||
"vision",
|
||||
"reasoning",
|
||||
"long_context",
|
||||
},
|
||||
("deepseek", "deepseek"): {"chat", "tool_use", "structured_output", "reasoning"},
|
||||
("deepseek", "deepseek-v4"): {
|
||||
"chat",
|
||||
"tool_use",
|
||||
"structured_output",
|
||||
"reasoning",
|
||||
"long_context",
|
||||
},
|
||||
}
|
||||
|
||||
|
||||
def derive_capabilities(provider: str, model: str) -> set[str]:
|
||||
"""根据 provider + model 前缀自动推导能力。
|
||||
|
||||
调用者可以在调用此函数后，在 ModelConfig 上覆盖个别标记。
|
||||
"""
|
||||
for (prov_prefix, model_prefix), caps in _BASE_CAPABILITIES.items():
|
||||
if provider.startswith(prov_prefix) and model.startswith(model_prefix):
|
||||
return set(caps)
|
||||
return {"chat"}
|
||||
|
||||
|
||||
class ModelEndpointConfig(BaseModel):
|
||||
name: str
|
||||
model: str
|
||||
base_url: str | None = None
|
||||
base_url_env: str | None = None
|
||||
api_key: str | None = None
|
||||
api_key_env: str | None = None
|
||||
weight: int = 100
|
||||
priority: int = 100
|
||||
rpm_limit: int | None = None
|
||||
tpm_limit: int | None = None
|
||||
max_concurrency: int | None = None
|
||||
request_timeout_seconds: int | None = None
|
||||
status: str = "active"
|
||||
extra_params: dict[str, Any] = Field(default_factory=dict)
|
||||
|
||||
|
||||
class ModelConfig(BaseModel):
|
||||
name: str
|
||||
display_name: str | None = None
|
||||
provider: str
|
||||
model: str
|
||||
format: SupportedFormats # "openai" | "anthropic"
|
||||
chat_format: str | None = None
|
||||
tags: list[str] = Field(default_factory=list)
|
||||
supports_tools: bool | None = None
|
||||
supports_vision: bool | None = None
|
||||
supports_reasoning_effort: bool | None = None
|
||||
supports_thinking: bool | None = None
|
||||
supports_streaming: bool | None = True
|
||||
supports_structured_output: bool | None = None
|
||||
context_window: int = 32000
|
||||
response_reserve: int = 4096
|
||||
max_output_tokens: int | None = None
|
||||
max_thinking_tokens: int | None = None
|
||||
reasoning_effort: str | None = None
|
||||
timeout_seconds: int | None = 60
|
||||
dimensions: int | None = None # 向量维度（embedding 模型专用）
|
||||
routing_strategy: str = "weighted_round_robin"
|
||||
endpoint_pool_enabled: bool = False
|
||||
weight: int = 100
|
||||
endpoints: list[ModelEndpointConfig] = Field(default_factory=list)
|
||||
|
||||
@model_validator(mode="after")
|
||||
def _validate_embedding_dimensions(self) -> ModelConfig:
|
||||
"""embedding 模型必须声明 dimensions。"""
|
||||
if "embedding" in self.tags and self.dimensions is None:
|
||||
raise ValueError(
|
||||
f"Embedding 模型 '{self.name}' 必须设置 dimensions（如 dimensions: 1024）"
|
||||
)
|
||||
return self
|
||||
|
||||
def resolved_capabilities(self) -> dict[str, bool]:
|
||||
base = derive_capabilities(self.provider, self.model)
|
||||
return {
|
||||
"chat": True,
|
||||
"tool_use": self.supports_tools
|
||||
if self.supports_tools is not None
|
||||
else ("tool_use" in base),
|
||||
"vision": self.supports_vision
|
||||
if self.supports_vision is not None
|
||||
else ("vision" in base),
|
||||
"reasoning": self.supports_reasoning_effort
|
||||
if self.supports_reasoning_effort is not None
|
||||
else ("reasoning" in base),
|
||||
"thinking": self.supports_thinking
|
||||
if self.supports_thinking is not None
|
||||
else ("reasoning" in base),
|
||||
"streaming": self.supports_streaming if self.supports_streaming is not None else True,
|
||||
"structured_output": self.supports_structured_output
|
||||
if self.supports_structured_output is not None
|
||||
else ("structured_output" in base),
|
||||
}
|
||||
|
||||
|
||||
class ParameterMapping(BaseModel):
|
||||
provider_param: str
|
||||
min: float | None = None
|
||||
max: float | None = None
|
||||
default: Any = None
|
||||
transform: str | None = None # "multiply_1000", "int", "str"
|
||||
alias_for: str | None = None # map from another param name
|
||||
handler: str | None = None # named handler for complex param transform
|
||||
|
||||
|
||||
class ProviderConfig(BaseModel):
|
||||
name: str
|
||||
format: SupportedFormats
|
||||
base_url_env: str
|
||||
api_key_env: str
|
||||
base_url: str | None = None
|
||||
timeout_seconds: int = 60
|
||||
max_retries: int = 1
|
||||
parameter_mappings: dict[str, ParameterMapping] = Field(default_factory=dict)
|
||||
parameter_handlers: list[str] = Field(
|
||||
default_factory=list
|
||||
) # ordered global handlers per provider
|
||||
extra_headers: dict[str, str] = Field(default_factory=dict)
|
||||
supports_streaming: bool = True
|
||||
|
||||
@field_validator("format")
|
||||
@classmethod
|
||||
def _validate_format(cls, v: str) -> str:
|
||||
allowed = {e.value for e in SupportedFormats}
|
||||
if v not in allowed:
|
||||
raise ValueError(f"format must be one of {allowed}, got '{v}'")
|
||||
return v
|
||||
@@ -0,0 +1,17 @@
|
||||
from app.ai_provider.core.billing import PricingFactory, PricingInfo, PricingRegistry
|
||||
from app.ai_provider.core.registry import (
|
||||
ModelRegistry,
|
||||
ProviderRegistry,
|
||||
RegistryFactory,
|
||||
RegistryHolder,
|
||||
)
|
||||
|
||||
__all__ = [
|
||||
"ModelRegistry",
|
||||
"ProviderRegistry",
|
||||
"RegistryFactory",
|
||||
"RegistryHolder",
|
||||
"PricingInfo",
|
||||
"PricingRegistry",
|
||||
"PricingFactory",
|
||||
]
|
||||
@@ -0,0 +1,253 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from dataclasses import dataclass, field
|
||||
from datetime import date, datetime
|
||||
|
||||
from pydantic import BaseModel
|
||||
|
||||
from app.ai_provider.types.chat import TokenUsage
|
||||
|
||||
|
||||
@dataclass
|
||||
class CostBreakdown:
|
||||
input_cost: float = 0.0
|
||||
output_cost: float = 0.0
|
||||
total: float = 0.0
|
||||
currency: str = "CNY"
|
||||
|
||||
|
||||
class CostCalculator:
|
||||
"""根据价格信息和 token 数量计算成本。
|
||||
|
||||
保持厂商原始单位 —— 不做跨币种换算。
|
||||
|
||||
- 普通 input token 按 ``input_per_1m`` 计费。
|
||||
- 缓存命中的 prompt token 优先按 ``input_cache_read_per_1m`` 计费
|
||||
（若有配置），剩余按普通 input 费率。
|
||||
- 若未配置 cache 费率，所有 prompt token 统一按 input_per_1m 计费。
|
||||
"""
|
||||
|
||||
@staticmethod
|
||||
def calculate(
|
||||
pricing: PricingInfo | None,
|
||||
prompt_tokens: int,
|
||||
completion_tokens: int,
|
||||
prompt_tokens_cached: int = 0,
|
||||
) -> CostBreakdown:
|
||||
if pricing is None:
|
||||
return CostBreakdown(total=0.0)
|
||||
|
||||
# 缓存 token：优先走缓存费率
|
||||
cached = min(prompt_tokens_cached, prompt_tokens)
|
||||
uncached = prompt_tokens - cached
|
||||
|
||||
input_cost = 0.0
|
||||
if uncached > 0 and pricing.input_per_1m is not None:
|
||||
input_cost += (uncached / 1_000_000) * pricing.input_per_1m
|
||||
if cached > 0 and pricing.input_cache_read_per_1m is not None:
|
||||
input_cost += (cached / 1_000_000) * pricing.input_cache_read_per_1m
|
||||
elif cached > 0 and pricing.input_per_1m is not None:
|
||||
# 无缓存费率时统一按普通 input 计
|
||||
input_cost += (cached / 1_000_000) * pricing.input_per_1m
|
||||
|
||||
output_cost = 0.0
|
||||
if pricing.output_per_1m is not None and completion_tokens > 0:
|
||||
output_cost = (completion_tokens / 1_000_000) * pricing.output_per_1m
|
||||
|
||||
return CostBreakdown(
|
||||
input_cost=round(input_cost, 8),
|
||||
output_cost=round(output_cost, 8),
|
||||
total=round(input_cost + output_cost, 8),
|
||||
currency=pricing.currency,
|
||||
)
|
||||
|
||||
|
||||
@dataclass
|
||||
class UsageRecord:
|
||||
model_name: str
|
||||
provider: str
|
||||
prompt_tokens: int = 0
|
||||
completion_tokens: int = 0
|
||||
total_tokens: int = 0
|
||||
currency: str = "CNY"
|
||||
cost: float = 0.0
|
||||
timestamp: datetime = field(default_factory=datetime.now)
|
||||
endpoint_name: str | None = None
|
||||
status: str = "success"
|
||||
error_message: str | None = None
|
||||
|
||||
|
||||
class BillingTracker:
|
||||
"""记录用量并按模型计算成本。"""
|
||||
|
||||
def __init__(self) -> None:
|
||||
self._records: list[UsageRecord] = []
|
||||
|
||||
def record(
|
||||
self,
|
||||
model_name: str,
|
||||
provider: str,
|
||||
usage: TokenUsage | None,
|
||||
pricing: PricingInfo | None = None,
|
||||
endpoint_name: str | None = None,
|
||||
status: str = "success",
|
||||
error_message: str | None = None,
|
||||
) -> UsageRecord:
|
||||
prompt_tokens = usage.prompt_tokens if usage else 0
|
||||
completion_tokens = usage.completion_tokens if usage else 0
|
||||
total = usage.total_tokens if usage else 0
|
||||
cached = usage.prompt_tokens_cached if usage else 0
|
||||
|
||||
cost_info = CostCalculator.calculate(pricing, prompt_tokens, completion_tokens, cached)
|
||||
|
||||
record = UsageRecord(
|
||||
model_name=model_name,
|
||||
provider=provider,
|
||||
prompt_tokens=prompt_tokens,
|
||||
completion_tokens=completion_tokens,
|
||||
total_tokens=total,
|
||||
currency=cost_info.currency,
|
||||
cost=cost_info.total,
|
||||
endpoint_name=endpoint_name,
|
||||
status=status,
|
||||
error_message=error_message,
|
||||
)
|
||||
self._records.append(record)
|
||||
return record
|
||||
|
||||
def get_records(
|
||||
self,
|
||||
since: datetime | None = None,
|
||||
model_name: str | None = None,
|
||||
) -> list[UsageRecord]:
|
||||
records = self._records
|
||||
if since:
|
||||
records = [r for r in records if r.timestamp >= since]
|
||||
if model_name:
|
||||
records = [r for r in records if r.model_name == model_name]
|
||||
return records
|
||||
|
||||
def total_cost(self, currency: str | None = None) -> float:
|
||||
if currency:
|
||||
return sum(r.cost for r in self._records if r.currency == currency)
|
||||
return sum(r.cost for r in self._records)
|
||||
|
||||
|
||||
class BudgetController:
|
||||
"""强制执行每日和每月预算限制。"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
daily_budget_cny: float = 0.0,
|
||||
monthly_budget_cny: float = 0.0,
|
||||
) -> None:
|
||||
self.daily_budget_cny = daily_budget_cny
|
||||
self.monthly_budget_cny = monthly_budget_cny
|
||||
self._daily_cost_cny: float = 0.0
|
||||
self._monthly_cost_cny: float = 0.0
|
||||
self._current_date: date = date.today()
|
||||
|
||||
def add_cost_cny(self, cost: float) -> None:
|
||||
self._daily_cost_cny += cost
|
||||
self._monthly_cost_cny += cost
|
||||
|
||||
def allow_request(self) -> bool:
|
||||
return not (
|
||||
(self.daily_budget_cny > 0 and self._daily_cost_cny >= self.daily_budget_cny)
|
||||
or (self.monthly_budget_cny > 0 and self._monthly_cost_cny >= self.monthly_budget_cny)
|
||||
)
|
||||
|
||||
def reset_daily(self) -> None:
|
||||
self._daily_cost_cny = 0.0
|
||||
self._current_date = date.today()
|
||||
|
||||
def reset_monthly(self) -> None:
|
||||
self._monthly_cost_cny = 0.0
|
||||
|
||||
|
||||
class BillingFactory:
|
||||
"""计费组件工厂。"""
|
||||
|
||||
@staticmethod
|
||||
def create_tracker() -> BillingTracker:
|
||||
return BillingTracker()
|
||||
|
||||
@staticmethod
|
||||
def create_budget_controller(
|
||||
daily_budget_cny: float = 0.0,
|
||||
monthly_budget_cny: float = 0.0,
|
||||
) -> BudgetController:
|
||||
return BudgetController(
|
||||
daily_budget_cny=daily_budget_cny,
|
||||
monthly_budget_cny=monthly_budget_cny,
|
||||
)
|
||||
|
||||
|
||||
# ── 定价信息 ──────────────────────────────────────────────────────────
|
||||
|
||||
|
||||
class PricingInfo(BaseModel):
|
||||
"""模型定价信息。保持厂商原始单位，不转换不混合。
|
||||
|
||||
各厂商按各自币种/单位记录，展示时原生显示。
|
||||
"""
|
||||
|
||||
currency: str = "CNY"
|
||||
input_per_1m: float | None = None
|
||||
output_per_1m: float | None = None
|
||||
input_cache_read_per_1m: float | None = None
|
||||
|
||||
|
||||
class PricingRegistry:
|
||||
"""模型定价注册表。以模型逻辑名为键查找定价。
|
||||
|
||||
与 ModelRegistry 解耦 —— 改价不影响模型注册，无定价的模型
|
||||
也可以正常使用（不计费）。
|
||||
"""
|
||||
|
||||
def __init__(self, pricing_map: dict[str, PricingInfo] | None = None) -> None:
|
||||
self._pricing: dict[str, PricingInfo] = dict(pricing_map or {})
|
||||
|
||||
def get(self, model_name: str) -> PricingInfo | None:
|
||||
"""获取指定模型的定价信息。无定价时返回 None。"""
|
||||
return self._pricing.get(model_name)
|
||||
|
||||
def list_all(self) -> dict[str, PricingInfo]:
|
||||
"""返回全部定价数据的副本。"""
|
||||
return dict(self._pricing)
|
||||
|
||||
|
||||
class PricingFactory:
|
||||
"""定价注册表工厂。"""
|
||||
|
||||
@staticmethod
|
||||
def from_dict(data: dict[str, dict]) -> PricingRegistry:
|
||||
"""从 YAML 解析后的 dict 构建 PricingRegistry。
|
||||
|
||||
YAML 按厂商嵌套组织（自文档化），工厂将其摊平为按模型逻辑名索引：
|
||||
|
||||
```yaml
|
||||
pricing:
|
||||
deepseek:
|
||||
deepseek-v4-flash:
|
||||
currency: CNY
|
||||
input_per_1m: 1.0
|
||||
anthropic:
|
||||
claude-sonnet:
|
||||
currency: USD
|
||||
input_per_1m: 3.0
|
||||
```
|
||||
"""
|
||||
pricing_map: dict[str, PricingInfo] = {}
|
||||
for _provider_name, models_data in data.items():
|
||||
if not isinstance(models_data, dict):
|
||||
continue
|
||||
for model_name, raw in models_data.items():
|
||||
if raw is None:
|
||||
continue
|
||||
# 跳过所有字段均为 null 的条目（如 mimo-v2.5-asr）
|
||||
values = [v for v in raw.values() if not isinstance(v, bool)]
|
||||
if all(v is None for v in values):
|
||||
continue
|
||||
pricing_map[model_name] = PricingInfo(**raw)
|
||||
return PricingRegistry(pricing_map)
|
||||
@@ -0,0 +1,164 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import random
|
||||
|
||||
from app.ai_provider.config.models import ModelConfig, ProviderConfig
|
||||
from app.ai_provider.core.billing import PricingFactory, PricingRegistry
|
||||
from app.ai_provider.types.errors import AIProviderError
|
||||
|
||||
|
||||
class ModelRegistry:
|
||||
"""内存模型索引。支持逻辑名和别名查找。"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
models: list[ModelConfig],
|
||||
aliases: dict[str, str] | None = None,
|
||||
) -> None:
|
||||
self._models: dict[str, ModelConfig] = {}
|
||||
self._aliases: dict[str, str] = dict(aliases or {})
|
||||
self._default: str | None = None
|
||||
self._fallback: str | None = None
|
||||
self._capability_index: dict[str, list[ModelConfig]] = {}
|
||||
|
||||
for m in models:
|
||||
self._models[m.name] = m
|
||||
if "default" in m.tags and self._default is None:
|
||||
self._default = m.name
|
||||
if "fallback" in m.tags and self._fallback is None:
|
||||
self._fallback = m.name
|
||||
caps = m.resolved_capabilities()
|
||||
for cap, supported in caps.items():
|
||||
if supported:
|
||||
self._capability_index.setdefault(cap, []).append(m)
|
||||
# 索引 tags（如 tts / asr / voice_clone）——这些不在
|
||||
# resolved_capabilities() 中，但 AudioService 等模块依赖它们查找模型。
|
||||
for tag in m.tags:
|
||||
self._capability_index.setdefault(tag, []).append(m)
|
||||
|
||||
if self._default is None and models:
|
||||
self._default = models[0].name
|
||||
|
||||
def get(self, name: str) -> ModelConfig:
|
||||
model = self._models.get(name)
|
||||
if model is None:
|
||||
raise AIProviderError(
|
||||
message=f"Model '{name}' not found in registry",
|
||||
code="MODEL_NOT_FOUND",
|
||||
status_code=404,
|
||||
)
|
||||
return model
|
||||
|
||||
def resolve(self, name_or_alias: str | None = None) -> ModelConfig:
|
||||
"""解析逻辑名或别名为 ModelConfig。
|
||||
|
||||
无参数时返回默认模型。
|
||||
"""
|
||||
if name_or_alias is None:
|
||||
return self._models[self.get_default_name()]
|
||||
if name_or_alias in self._models:
|
||||
return self._models[name_or_alias]
|
||||
if name_or_alias in self._aliases:
|
||||
logical = self._aliases[name_or_alias]
|
||||
return self._models[logical]
|
||||
raise AIProviderError(
|
||||
message=f"Model or alias '{name_or_alias}' not found",
|
||||
code="MODEL_NOT_FOUND",
|
||||
status_code=404,
|
||||
)
|
||||
|
||||
def resolve_name(self, name_or_alias: str) -> str:
|
||||
"""解析为规范逻辑名（不返回完整配置）。"""
|
||||
if name_or_alias in self._models:
|
||||
return name_or_alias
|
||||
if name_or_alias in self._aliases:
|
||||
return self._aliases[name_or_alias]
|
||||
raise AIProviderError(
|
||||
message=f"Model or alias '{name_or_alias}' not found",
|
||||
code="MODEL_NOT_FOUND",
|
||||
status_code=404,
|
||||
)
|
||||
|
||||
def get_default_name(self) -> str:
|
||||
if self._default is None:
|
||||
raise AIProviderError(
|
||||
message="No default model configured",
|
||||
code="NO_DEFAULT_MODEL",
|
||||
status_code=500,
|
||||
)
|
||||
return self._default
|
||||
|
||||
def get_fallback_name(self) -> str | None:
|
||||
return self._fallback
|
||||
|
||||
def contains(self, name: str) -> bool:
|
||||
return name in self._models
|
||||
|
||||
def list_all(self) -> list[ModelConfig]:
|
||||
return list(self._models.values())
|
||||
|
||||
def list_by_capability(self, capability: str) -> list[ModelConfig]:
|
||||
return list(self._capability_index.get(capability, []))
|
||||
|
||||
def resolve_by_capability(self, capability: str) -> str:
|
||||
candidates = self._capability_index.get(capability)
|
||||
if not candidates:
|
||||
raise AIProviderError(
|
||||
message=f"No active models with capability '{capability}'",
|
||||
code="CAPABILITY_NOT_FOUND",
|
||||
status_code=404,
|
||||
)
|
||||
weights = [max(c.weight, 1) for c in candidates]
|
||||
chosen = random.choices(candidates, weights=weights, k=1)[0]
|
||||
return chosen.name
|
||||
|
||||
|
||||
class ProviderRegistry:
|
||||
"""内存提供商配置索引。"""
|
||||
|
||||
def __init__(self, providers: dict[str, ProviderConfig]) -> None:
|
||||
self._providers = dict(providers)
|
||||
|
||||
def get(self, name: str) -> ProviderConfig:
|
||||
prov = self._providers.get(name)
|
||||
if prov is None:
|
||||
raise AIProviderError(
|
||||
message=f"Provider '{name}' not configured",
|
||||
code="PROVIDER_NOT_FOUND",
|
||||
status_code=404,
|
||||
)
|
||||
return prov
|
||||
|
||||
def list_all(self) -> dict[str, ProviderConfig]:
|
||||
return dict(self._providers)
|
||||
|
||||
|
||||
class RegistryFactory:
|
||||
"""工厂：从配置列表创建注册表。"""
|
||||
|
||||
@staticmethod
|
||||
def create(
|
||||
models: list[ModelConfig],
|
||||
providers: dict[str, ProviderConfig],
|
||||
aliases: dict[str, str] | None = None,
|
||||
pricing_data: dict[str, dict] | None = None,
|
||||
) -> RegistryHolder:
|
||||
return RegistryHolder(
|
||||
model_registry=ModelRegistry(models, aliases=aliases),
|
||||
provider_registry=ProviderRegistry(providers),
|
||||
pricing_registry=PricingFactory.from_dict(pricing_data or {}),
|
||||
)
|
||||
|
||||
|
||||
class RegistryHolder:
|
||||
"""同时持有所有注册表。"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
model_registry: ModelRegistry,
|
||||
provider_registry: ProviderRegistry,
|
||||
pricing_registry: PricingRegistry | None = None,
|
||||
) -> None:
|
||||
self.model_registry = model_registry
|
||||
self.provider_registry = provider_registry
|
||||
self.pricing_registry = pricing_registry or PricingRegistry()
|
||||
@@ -0,0 +1,133 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import time
|
||||
from enum import Enum
|
||||
|
||||
|
||||
class CircuitBreakerState(Enum):
|
||||
CLOSED = "closed"
|
||||
OPEN = "open"
|
||||
HALF_OPEN = "half_open"
|
||||
|
||||
|
||||
class CircuitBreaker:
|
||||
"""标准断路器，冷却后自动半开。"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
name: str,
|
||||
failure_threshold: int = 5,
|
||||
cooldown_seconds: float = 30.0,
|
||||
) -> None:
|
||||
self.name = name
|
||||
self.failure_threshold = failure_threshold
|
||||
self.cooldown_seconds = cooldown_seconds
|
||||
self._state = CircuitBreakerState.CLOSED
|
||||
self._failure_count = 0
|
||||
self._last_failure_time: float = 0.0
|
||||
|
||||
@property
|
||||
def state(self) -> CircuitBreakerState:
|
||||
if (
|
||||
self._state == CircuitBreakerState.OPEN
|
||||
and time.monotonic() - self._last_failure_time >= self.cooldown_seconds
|
||||
):
|
||||
self._state = CircuitBreakerState.HALF_OPEN
|
||||
return self._state
|
||||
|
||||
def record_success(self) -> None:
|
||||
self._failure_count = 0
|
||||
self._state = CircuitBreakerState.CLOSED
|
||||
|
||||
def record_failure(self) -> None:
|
||||
self._failure_count += 1
|
||||
self._last_failure_time = time.monotonic()
|
||||
if self._failure_count >= self.failure_threshold:
|
||||
self._state = CircuitBreakerState.OPEN
|
||||
|
||||
def allow_request(self) -> bool:
|
||||
return self.state != CircuitBreakerState.OPEN
|
||||
|
||||
|
||||
class RetryStrategy:
|
||||
"""指数退避重试。"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
max_retries: int = 3,
|
||||
base_delay: float = 1.0,
|
||||
max_delay: float = 60.0,
|
||||
multiplier: float = 2.0,
|
||||
) -> None:
|
||||
self.max_retries = max_retries
|
||||
self.base_delay = base_delay
|
||||
self.max_delay = max_delay
|
||||
self.multiplier = multiplier
|
||||
|
||||
def should_retry(self, attempt_index: int) -> bool:
|
||||
return attempt_index < self.max_retries
|
||||
|
||||
def get_delay(self, attempt_index: int) -> float:
|
||||
delay = self.base_delay * (self.multiplier**attempt_index)
|
||||
return min(delay, self.max_delay)
|
||||
|
||||
|
||||
class DowngradeAction:
|
||||
"""降级链中的一个步骤。"""
|
||||
|
||||
def __init__(self, type: str, target: str) -> None: # noqa: A002
|
||||
self.type = type # "fallback_endpoint", "fallback_model", "cache", "degrade_quality"
|
||||
self.target = target
|
||||
|
||||
|
||||
class DowngradeChain:
|
||||
"""当模型/端点失败时的降级动作链。"""
|
||||
|
||||
def __init__(self, actions: list[DowngradeAction]) -> None:
|
||||
self._actions = actions
|
||||
self._index = 0
|
||||
|
||||
def next_action(self) -> tuple[str, str] | None:
|
||||
if self._index >= len(self._actions):
|
||||
return None
|
||||
action = self._actions[self._index]
|
||||
return (action.type, action.target)
|
||||
|
||||
def advance(self) -> None:
|
||||
self._index += 1
|
||||
|
||||
def is_exhausted(self) -> bool:
|
||||
return self._index >= len(self._actions)
|
||||
|
||||
def reset(self) -> None:
|
||||
self._index = 0
|
||||
|
||||
|
||||
class ResilienceFactory:
|
||||
"""弹性组件工厂。"""
|
||||
|
||||
@staticmethod
|
||||
def create_circuit_breaker(
|
||||
name: str,
|
||||
failure_threshold: int = 5,
|
||||
cooldown_seconds: float = 30.0,
|
||||
) -> CircuitBreaker:
|
||||
return CircuitBreaker(
|
||||
name=name,
|
||||
failure_threshold=failure_threshold,
|
||||
cooldown_seconds=cooldown_seconds,
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def create_retry_strategy(
|
||||
max_retries: int = 3,
|
||||
base_delay: float = 1.0,
|
||||
max_delay: float = 60.0,
|
||||
multiplier: float = 2.0,
|
||||
) -> RetryStrategy:
|
||||
return RetryStrategy(
|
||||
max_retries=max_retries,
|
||||
base_delay=base_delay,
|
||||
max_delay=max_delay,
|
||||
multiplier=multiplier,
|
||||
)
|
||||
@@ -0,0 +1,149 @@
|
||||
"""嵌入服务 — Embedding 适配器外观。
|
||||
|
||||
用法::
|
||||
|
||||
from app.ai_provider.embedding import EmbeddingService
|
||||
|
||||
service = EmbeddingService(manager)
|
||||
result = await service.embed("要嵌入的文本")
|
||||
results = await service.embed(["文本1", "文本2"])
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from app.ai_provider.embedding.base import (
|
||||
EmbeddingProtocol,
|
||||
EmbeddingRequest,
|
||||
EmbeddingResponse,
|
||||
EmbeddingResult,
|
||||
EmbeddingServiceError,
|
||||
)
|
||||
from app.ai_provider.types.errors import AIProviderError
|
||||
|
||||
|
||||
class EmbeddingService:
|
||||
"""外观类，负责解析 Embedding 模型并委托给 ProviderManager。
|
||||
|
||||
接受 ProviderManager 或实现了以下接口的适配器：
|
||||
- ``.model_registry`` → 具有 ``list_by_capability(capability)`` 方法
|
||||
- ``.create_embedding(name)`` → 返回 EmbeddingProtocol 实例
|
||||
"""
|
||||
|
||||
def __init__(self, manager) -> None:
|
||||
self._manager = manager
|
||||
self._adapter: EmbeddingProtocol | None = None
|
||||
|
||||
# ── Registry 兼容适配 ─────────────────────────────────────────
|
||||
|
||||
@property
|
||||
def _registry(self):
|
||||
mgr = self._manager
|
||||
if hasattr(mgr, "model_registry"):
|
||||
return mgr.model_registry
|
||||
return mgr
|
||||
|
||||
@property
|
||||
def _factory(self):
|
||||
mgr = self._manager
|
||||
if hasattr(mgr, "create_embedding"):
|
||||
return mgr
|
||||
return None
|
||||
|
||||
# ── Embedding ──────────────────────────────────────────────────
|
||||
|
||||
@property
|
||||
def adapter(self) -> EmbeddingProtocol:
|
||||
"""解析并缓存 Embedding 适配器。"""
|
||||
if self._adapter is None:
|
||||
self._adapter = self._resolve()
|
||||
return self._adapter
|
||||
|
||||
async def embed(
|
||||
self,
|
||||
input: str | list[str],
|
||||
model_hint: str | None = None,
|
||||
dimensions: int | None = None,
|
||||
) -> EmbeddingResponse:
|
||||
"""快捷方法：构造请求并调用 Embedding 适配器。"""
|
||||
model_name = self._resolve_model_name(model_hint)
|
||||
# 优先从模型配置读 dimensions
|
||||
effective_dimensions = dimensions
|
||||
if effective_dimensions is None:
|
||||
effective_dimensions = self._resolve_model_dimensions(model_hint)
|
||||
request = EmbeddingRequest(
|
||||
input=input if isinstance(input, list) else [input],
|
||||
model=model_name,
|
||||
dimensions=effective_dimensions,
|
||||
)
|
||||
return await self.adapter.embed(request)
|
||||
|
||||
# ── 解析 ───────────────────────────────────────────────────────
|
||||
|
||||
def _resolve(self) -> EmbeddingProtocol:
|
||||
"""查找第一个活跃的 Embedding 模型并通过工厂创建。"""
|
||||
models = self._list_capability_models("embedding")
|
||||
if not models:
|
||||
raise AIProviderError(
|
||||
status_code=503,
|
||||
code="embedding_no_model",
|
||||
message="No embedding model configured in the model registry.",
|
||||
)
|
||||
factory = self._factory
|
||||
if factory is None:
|
||||
raise AIProviderError(
|
||||
status_code=501,
|
||||
code="embedding_no_factory",
|
||||
message="Embedding adapter factory not available (provider manager missing create_embedding).",
|
||||
)
|
||||
return factory.create_embedding(models[0].name)
|
||||
|
||||
def _resolve_model_name(self, model_hint: str | None) -> str:
|
||||
models = self._list_capability_models("embedding")
|
||||
if model_hint:
|
||||
return model_hint
|
||||
if models:
|
||||
return models[0].name
|
||||
return ""
|
||||
|
||||
def _resolve_model_dimensions(self, model_hint: str | None) -> int | None:
|
||||
"""从模型配置读取向量维度。"""
|
||||
models = self._list_capability_models("embedding")
|
||||
if model_hint:
|
||||
for m in models:
|
||||
if m.name == model_hint:
|
||||
return m.dimensions
|
||||
if models:
|
||||
return models[0].dimensions
|
||||
return None
|
||||
|
||||
def _list_capability_models(self, tag: str):
|
||||
reg = self._registry
|
||||
if hasattr(reg, "list_capability_models"):
|
||||
return reg.list_capability_models(tag)
|
||||
if hasattr(reg, "list_by_capability"):
|
||||
return reg.list_by_capability(tag)
|
||||
return []
|
||||
|
||||
# ── 生命周期 ────────────────────────────────────────────────────
|
||||
|
||||
async def close(self) -> None:
|
||||
if self._adapter is not None:
|
||||
close = getattr(self._adapter, "close", None)
|
||||
if close is not None:
|
||||
await close()
|
||||
|
||||
async def __aenter__(self) -> EmbeddingService:
|
||||
return self
|
||||
|
||||
async def __aexit__(self, *exc_info: object) -> None:
|
||||
await self.close()
|
||||
|
||||
|
||||
__all__ = [
|
||||
"EmbeddingProtocol",
|
||||
"EmbeddingRequest",
|
||||
"EmbeddingResponse",
|
||||
"EmbeddingResult",
|
||||
"EmbeddingService",
|
||||
"EmbeddingServiceError",
|
||||
]
|
||||
@@ -0,0 +1,144 @@
|
||||
"""嵌入服务类型 — 请求/响应模型和协议定义。
|
||||
|
||||
Embedding 适配器通过 HTTP 调用嵌入端点（通常是 OpenAI 兼容的 ``POST /v1/embeddings``），
|
||||
经由 ``EmbeddingService`` 或直接从 ``ProviderManager`` 消费。
|
||||
|
||||
用法::
|
||||
|
||||
service = EmbeddingService(manager)
|
||||
result = await service.embed("要嵌入的文本")
|
||||
|
||||
# 或通过工厂直接调用:
|
||||
adapter = manager.create_embedding("text-embedding-3-small")
|
||||
result = await adapter.embed(EmbeddingRequest(input=["文本"]))
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from dataclasses import dataclass
|
||||
from typing import TYPE_CHECKING, Any, Protocol, runtime_checkable
|
||||
|
||||
from app.ai_provider.types.errors import AIProviderError
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from app.ai_provider.config.models import ModelConfig, ModelEndpointConfig
|
||||
|
||||
|
||||
class EmbeddingServiceError(AIProviderError):
|
||||
"""Embedding API 调用失败的基础错误。"""
|
||||
|
||||
|
||||
# ── 请求 / 响应 ────────────────────────────────────────────────────────
|
||||
|
||||
|
||||
@dataclass
|
||||
class EmbeddingRequest:
|
||||
"""将文本转换为向量的请求。
|
||||
|
||||
属性:
|
||||
input: 待嵌入的文本字符串或字符串列表。
|
||||
model: 模型标识符（如 ``"text-embedding-3-small"``）。
|
||||
encoding_format: 返回编码格式（``"float"`` 或 ``"base64"``）。
|
||||
dimensions: 输出向量的维度（仅部分模型支持）。
|
||||
user: 终端用户标识（用于用量追踪）。
|
||||
"""
|
||||
|
||||
input: str | list[str]
|
||||
model: str = ""
|
||||
encoding_format: str = "float"
|
||||
dimensions: int | None = None
|
||||
user: str | None = None
|
||||
|
||||
|
||||
@dataclass
|
||||
class EmbeddingResult:
|
||||
"""单条文本的嵌入结果。"""
|
||||
|
||||
embedding: list[float]
|
||||
index: int = 0
|
||||
|
||||
|
||||
@dataclass
|
||||
class EmbeddingResponse:
|
||||
"""嵌入 API 的完整响应。"""
|
||||
|
||||
data: list[EmbeddingResult]
|
||||
model: str = ""
|
||||
usage: dict[str, int] | None = None
|
||||
|
||||
|
||||
# ── Protocol ────────────────────────────────────────────────────────────
|
||||
|
||||
|
||||
@runtime_checkable
|
||||
class EmbeddingProtocol(Protocol):
|
||||
"""Embedding 适配器的协议。"""
|
||||
|
||||
async def embed(self, request: EmbeddingRequest) -> EmbeddingResponse:
|
||||
"""将文本转换为向量。
|
||||
|
||||
异常:
|
||||
EmbeddingServiceError: API 调用失败或响应异常时抛出。
|
||||
"""
|
||||
...
|
||||
|
||||
async def close(self) -> None:
|
||||
"""释放适配器持有的 HTTP 连接。"""
|
||||
...
|
||||
|
||||
|
||||
# ── 共享工厂基类 ──────────────────────────────────────────────────────
|
||||
|
||||
|
||||
class EmbeddingAdapterBase:
|
||||
"""Embedding 适配器基类，提供共享的工厂方法。
|
||||
|
||||
子类应在 ``__init__`` 中接受 ``base_url``、``api_key``、``model``、
|
||||
``request_timeout_seconds`` 参数。
|
||||
"""
|
||||
|
||||
_adapter_type: str = "embedding"
|
||||
_default_timeout: int = 60
|
||||
|
||||
@classmethod
|
||||
def from_endpoint(
|
||||
cls,
|
||||
endpoint: ModelEndpointConfig,
|
||||
*,
|
||||
timeout_seconds: int | None = None,
|
||||
) -> Any:
|
||||
"""从已解析的端点配置构建适配器。"""
|
||||
base_url = endpoint.base_url
|
||||
api_key = endpoint.api_key
|
||||
if not base_url:
|
||||
raise EmbeddingServiceError(
|
||||
status_code=500,
|
||||
code="embedding_endpoint_invalid",
|
||||
message="Embedding endpoint has no resolved base_url.",
|
||||
details={"endpoint": endpoint.name},
|
||||
)
|
||||
return cls(
|
||||
base_url=base_url,
|
||||
api_key=api_key,
|
||||
model=endpoint.model,
|
||||
request_timeout_seconds=timeout_seconds
|
||||
or endpoint.request_timeout_seconds
|
||||
or cls._default_timeout,
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def from_model_config(
|
||||
cls,
|
||||
config: ModelConfig,
|
||||
*,
|
||||
timeout_seconds: int | None = None,
|
||||
) -> Any:
|
||||
"""从逻辑模型配置构建适配器（使用第一个活跃端点）。"""
|
||||
for ep in config.endpoints:
|
||||
if ep.status == "active":
|
||||
return cls.from_endpoint(ep, timeout_seconds=timeout_seconds)
|
||||
raise EmbeddingServiceError(
|
||||
status_code=500,
|
||||
code="embedding_no_endpoint",
|
||||
message=f"Embedding model '{config.name}' has no active endpoint.",
|
||||
)
|
||||
@@ -0,0 +1,169 @@
|
||||
"""OpenAI 兼容格式的 Embedding 适配器 — ``POST /v1/embeddings``。
|
||||
|
||||
覆盖 OpenAI、DeepSeek、以及任何兼容 OpenAI 嵌入端点的提供商。
|
||||
|
||||
请求::
|
||||
|
||||
POST {base_url}/v1/embeddings
|
||||
Authorization: Bearer {api_key}
|
||||
Content-Type: application/json
|
||||
|
||||
{
|
||||
"model": "text-embedding-3-small",
|
||||
"input": ["要嵌入的文本"],
|
||||
"encoding_format": "float"
|
||||
}
|
||||
|
||||
响应::
|
||||
|
||||
{
|
||||
"object": "list",
|
||||
"data": [
|
||||
{"object": "embedding", "index": 0, "embedding": [0.xxx, ...]}
|
||||
],
|
||||
"model": "text-embedding-3-small",
|
||||
"usage": {"prompt_tokens": 5, "total_tokens": 5}
|
||||
}
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import Any
|
||||
|
||||
import httpx
|
||||
|
||||
from app.ai_provider.embedding.base import (
|
||||
EmbeddingAdapterBase,
|
||||
EmbeddingProtocol,
|
||||
EmbeddingRequest,
|
||||
EmbeddingResponse,
|
||||
EmbeddingResult,
|
||||
EmbeddingServiceError,
|
||||
)
|
||||
from app.core.logging import logger
|
||||
|
||||
|
||||
class OpenAIFormatEmbeddingAdapter(EmbeddingAdapterBase, EmbeddingProtocol):
|
||||
"""OpenAI 兼容嵌入适配器，覆盖 95%+ 的嵌入服务提供商。"""
|
||||
|
||||
_adapter_type = "embedding"
|
||||
_default_timeout = 60
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
*,
|
||||
base_url: str,
|
||||
api_key: str,
|
||||
model: str = "text-embedding-3-small",
|
||||
request_timeout_seconds: int = 60,
|
||||
) -> None:
|
||||
self._base_url = base_url.rstrip("/")
|
||||
self._api_key = api_key
|
||||
self._model = model
|
||||
self._timeout = request_timeout_seconds
|
||||
self._client: httpx.AsyncClient | None = None
|
||||
|
||||
# ── 客户端 ───────────────────────────────────────────────────────────
|
||||
|
||||
@property
|
||||
def _http(self) -> httpx.AsyncClient:
|
||||
if self._client is None or self._client.is_closed:
|
||||
self._client = httpx.AsyncClient(timeout=httpx.Timeout(self._timeout))
|
||||
return self._client
|
||||
|
||||
async def close(self) -> None:
|
||||
if self._client is not None and not self._client.is_closed:
|
||||
await self._client.aclose()
|
||||
|
||||
# ── EmbeddingProtocol ────────────────────────────────────────────────
|
||||
|
||||
async def embed(self, request: EmbeddingRequest) -> EmbeddingResponse:
|
||||
"""通过 ``POST /v1/embeddings`` 将文本转换为向量。"""
|
||||
url = f"{self._base_url}/v1/embeddings"
|
||||
|
||||
payload: dict[str, Any] = {
|
||||
"model": request.model or self._model,
|
||||
"input": request.input,
|
||||
"encoding_format": request.encoding_format,
|
||||
}
|
||||
if request.dimensions is not None:
|
||||
payload["dimensions"] = request.dimensions
|
||||
if request.user is not None:
|
||||
payload["user"] = request.user
|
||||
|
||||
headers = {
|
||||
"Content-Type": "application/json",
|
||||
}
|
||||
if self._api_key:
|
||||
headers["Authorization"] = f"Bearer {self._api_key}"
|
||||
|
||||
logger.info(
|
||||
"embedding.request",
|
||||
model=payload["model"],
|
||||
input_count=len(request.input) if isinstance(request.input, list) else 1,
|
||||
dimensions=request.dimensions,
|
||||
url=url,
|
||||
)
|
||||
|
||||
try:
|
||||
response = await self._http.post(url, json=payload, headers=headers)
|
||||
except httpx.TimeoutException as exc:
|
||||
raise EmbeddingServiceError(
|
||||
status_code=504,
|
||||
code="embedding_timeout",
|
||||
message="Embedding request timed out.",
|
||||
) from exc
|
||||
except httpx.RequestError as exc:
|
||||
raise EmbeddingServiceError(
|
||||
status_code=503,
|
||||
code="embedding_connection_error",
|
||||
message=f"Embedding connection error: {exc}",
|
||||
) from exc
|
||||
|
||||
if response.status_code != 200:
|
||||
body = response.text[:500]
|
||||
raise EmbeddingServiceError(
|
||||
status_code=response.status_code,
|
||||
code="embedding_api_error",
|
||||
message=f"Embedding API returned {response.status_code}: {body}",
|
||||
)
|
||||
|
||||
result = self._parse_response(response.json())
|
||||
|
||||
logger.info(
|
||||
"embedding.success",
|
||||
model=payload["model"],
|
||||
total_vectors=len(result.data),
|
||||
usage=result.usage,
|
||||
)
|
||||
return result
|
||||
|
||||
# ── 辅助方法 ──────────────────────────────────────────────────────────
|
||||
|
||||
@staticmethod
|
||||
def _parse_response(body: dict) -> EmbeddingResponse:
|
||||
"""从 OpenAI 格式的嵌入响应中提取结果。"""
|
||||
try:
|
||||
raw_data = body.get("data", [])
|
||||
results = []
|
||||
for item in raw_data:
|
||||
results.append(
|
||||
EmbeddingResult(
|
||||
embedding=item.get("embedding", []),
|
||||
index=item.get("index", 0),
|
||||
)
|
||||
)
|
||||
# 按 index 排序确保顺序
|
||||
results.sort(key=lambda r: r.index)
|
||||
|
||||
return EmbeddingResponse(
|
||||
data=results,
|
||||
model=body.get("model", ""),
|
||||
usage=body.get("usage"),
|
||||
)
|
||||
except Exception as exc:
|
||||
raise EmbeddingServiceError(
|
||||
status_code=500,
|
||||
code="embedding_response_parse_error",
|
||||
message=f"Failed to parse embedding response: {exc}",
|
||||
) from exc
|
||||
@@ -0,0 +1,18 @@
|
||||
from __future__ import annotations
|
||||
|
||||
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 StreamFactory, StreamHandler, ToolCallAccumulator
|
||||
|
||||
__all__ = [
|
||||
"AnthropicFormatAdapter",
|
||||
"AnthropicFormatAdapterFactory",
|
||||
"OpenAIFormatAdapter",
|
||||
"OpenAIFormatAdapterFactory",
|
||||
"StreamFactory",
|
||||
"StreamHandler",
|
||||
"ToolCallAccumulator",
|
||||
]
|
||||
@@ -0,0 +1,433 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
from collections.abc import AsyncIterator
|
||||
from typing import Any
|
||||
|
||||
import httpx
|
||||
|
||||
from app.ai_provider.config.models import ModelConfig, ProviderConfig
|
||||
from app.ai_provider.pipeline import ParameterPipeline, PipelineContext, PipelineFactory
|
||||
from app.ai_provider.types.chat import ChatChunk, ChatResult, TokenUsage
|
||||
from app.ai_provider.types.errors import AIProviderError
|
||||
from app.ai_provider.types.messages import (
|
||||
AIMessage,
|
||||
BaseMessage,
|
||||
SystemMessage,
|
||||
ToolMessage,
|
||||
)
|
||||
from app.core.logging import logger
|
||||
|
||||
|
||||
class AnthropicFormatAdapter:
|
||||
"""兼容 Anthropic Messages API 格式的适配器。
|
||||
|
||||
覆盖通过 Anthropic API 访问的 Claude 模型，以及任何镜像
|
||||
Anthropic 消息格式的提供商。
|
||||
"""
|
||||
|
||||
API_VERSION = "2023-06-01"
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
provider_config: ProviderConfig,
|
||||
pipeline: ParameterPipeline | None = None,
|
||||
) -> None:
|
||||
self._provider = provider_config
|
||||
self._pipeline = pipeline or PipelineFactory.create(provider_config)
|
||||
self._client: httpx.AsyncClient | None = None
|
||||
|
||||
@property
|
||||
def _http(self) -> httpx.AsyncClient:
|
||||
if self._client is None:
|
||||
self._client = httpx.AsyncClient(
|
||||
base_url=self._provider.base_url or "",
|
||||
timeout=self._provider.timeout_seconds,
|
||||
)
|
||||
return self._client
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# 内部方法：消息转换
|
||||
# ------------------------------------------------------------------
|
||||
|
||||
def _convert_message(self, msg: BaseMessage) -> dict[str, Any]:
|
||||
"""将单条消息转换为 Anthropic 消息格式。
|
||||
|
||||
Anthropic 使用 content blocks 数组而非纯字符串。
|
||||
System 消息**不在此处处理** —— 它们由 :meth:`build_request`
|
||||
提取到顶层的 ``system`` 字段中。
|
||||
"""
|
||||
if isinstance(msg, ToolMessage):
|
||||
return {
|
||||
"role": "user",
|
||||
"content": [
|
||||
{
|
||||
"type": "tool_result",
|
||||
"tool_use_id": msg.tool_call_id,
|
||||
"content": msg.content,
|
||||
}
|
||||
],
|
||||
}
|
||||
|
||||
if isinstance(msg, AIMessage):
|
||||
if msg.raw_content_blocks:
|
||||
return {"role": "assistant", "content": msg.raw_content_blocks}
|
||||
blocks: list[dict[str, Any]] = []
|
||||
if msg.content:
|
||||
blocks.append({"type": "text", "text": msg.content})
|
||||
for tc in msg.tool_calls:
|
||||
blocks.append(
|
||||
{
|
||||
"type": "tool_use",
|
||||
"id": tc.get("id", f"tu_{tc['name']}"),
|
||||
"name": tc["name"],
|
||||
"input": tc.get("args", {}),
|
||||
}
|
||||
)
|
||||
return {"role": "assistant", "content": blocks}
|
||||
|
||||
# HumanMessage
|
||||
return {
|
||||
"role": "user",
|
||||
"content": [{"type": "text", "text": msg.content}],
|
||||
}
|
||||
|
||||
@staticmethod
|
||||
def _convert_tools(tools: list[dict[str, Any]]) -> list[dict[str, Any]]:
|
||||
"""将 OpenAI 兼容格式的工具转换为 Anthropic 格式。
|
||||
|
||||
OpenAI 格式::
|
||||
|
||||
{"type": "function", "function": {"name": ..., "parameters": ...}}
|
||||
|
||||
Anthropic 格式::
|
||||
|
||||
{"name": ..., "input_schema": ...}
|
||||
"""
|
||||
result: list[dict[str, Any]] = []
|
||||
for tool in tools:
|
||||
fn = tool.get("function", tool)
|
||||
result.append(
|
||||
{
|
||||
"name": fn.get("name", ""),
|
||||
"description": fn.get("description", ""),
|
||||
"input_schema": fn.get("parameters", {"type": "object"}),
|
||||
}
|
||||
)
|
||||
return result
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# 请求构建
|
||||
# ------------------------------------------------------------------
|
||||
|
||||
def build_request(
|
||||
self,
|
||||
model: ModelConfig,
|
||||
messages: list[BaseMessage],
|
||||
params: dict[str, Any] | None = None,
|
||||
tools: list[dict[str, Any]] | None = None,
|
||||
tool_choice: str | dict[str, Any] | None = None,
|
||||
stream: bool = False,
|
||||
) -> dict[str, Any]:
|
||||
ctx = PipelineContext(
|
||||
params=params or {},
|
||||
model_config=model,
|
||||
provider_config=self._provider,
|
||||
)
|
||||
pipe_result = self._pipeline.run(ctx)
|
||||
|
||||
# Extract system messages — Anthropic uses a top-level "system" field
|
||||
system_parts: list[str] = []
|
||||
non_system: list[BaseMessage] = []
|
||||
for m in messages:
|
||||
if isinstance(m, SystemMessage):
|
||||
system_parts.append(m.content)
|
||||
else:
|
||||
non_system.append(m)
|
||||
|
||||
body: dict[str, Any] = {
|
||||
"model": model.model,
|
||||
"messages": [self._convert_message(m) for m in non_system],
|
||||
"stream": stream,
|
||||
}
|
||||
|
||||
if system_parts:
|
||||
body["system"] = "\n".join(system_parts)
|
||||
|
||||
body.update(pipe_result.params)
|
||||
|
||||
if tools:
|
||||
body["tools"] = self._convert_tools(tools)
|
||||
if tool_choice:
|
||||
body["tool_choice"] = tool_choice
|
||||
|
||||
# Anthropic requires explicit max_tokens
|
||||
if "max_tokens" not in body:
|
||||
body["max_tokens"] = model.max_output_tokens or 4096
|
||||
|
||||
return body
|
||||
|
||||
def bind_tools(
|
||||
self,
|
||||
request_body: dict[str, Any],
|
||||
tools: list[dict[str, Any]],
|
||||
tool_choice: str | dict[str, Any] | None = None,
|
||||
) -> None:
|
||||
request_body["tools"] = self._convert_tools(tools)
|
||||
if tool_choice:
|
||||
request_body["tool_choice"] = tool_choice
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# 认证头
|
||||
# ------------------------------------------------------------------
|
||||
|
||||
def _headers(self, api_key: str) -> dict[str, str]:
|
||||
return {
|
||||
"x-api-key": api_key,
|
||||
"anthropic-version": self.API_VERSION,
|
||||
"Content-Type": "application/json",
|
||||
}
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# 响应解析 —— 非流式
|
||||
# ------------------------------------------------------------------
|
||||
|
||||
@staticmethod
|
||||
def _parse_content_blocks(raw: dict[str, Any]) -> tuple[str, list[dict[str, Any]], str]:
|
||||
"""解析 Anthropic content blocks 为 (text, tool_calls, reasoning)。"""
|
||||
content_blocks = raw.get("content", [])
|
||||
text_parts: list[str] = []
|
||||
tool_calls: list[dict[str, Any]] = []
|
||||
reasoning = ""
|
||||
|
||||
for block in content_blocks:
|
||||
block_type = block.get("type", "")
|
||||
if block_type == "text":
|
||||
text_parts.append(block.get("text", ""))
|
||||
elif block_type == "tool_use":
|
||||
tool_calls.append(
|
||||
{
|
||||
"name": block.get("name", ""),
|
||||
"args": block.get("input", {}),
|
||||
"id": block.get("id", ""),
|
||||
}
|
||||
)
|
||||
elif block_type == "thinking":
|
||||
reasoning += block.get("thinking", "")
|
||||
# signature 块是 thinking 的加密签名，不是 reasoning 内容，忽略
|
||||
|
||||
return "".join(text_parts), tool_calls, reasoning
|
||||
|
||||
def parse_nonstreaming_response(
|
||||
self,
|
||||
raw: dict[str, Any],
|
||||
model_name: str,
|
||||
) -> ChatResult:
|
||||
content, tool_calls, reasoning = self._parse_content_blocks(raw)
|
||||
|
||||
result = ChatResult(
|
||||
content=content,
|
||||
model_name=model_name,
|
||||
)
|
||||
if tool_calls:
|
||||
result.tool_calls = tool_calls
|
||||
if reasoning:
|
||||
result.reasoning = reasoning
|
||||
|
||||
usage = raw.get("usage")
|
||||
if usage:
|
||||
result.usage = TokenUsage(
|
||||
prompt_tokens=usage.get("input_tokens", 0),
|
||||
prompt_tokens_cached=usage.get("cache_read_input_tokens", 0),
|
||||
completion_tokens=usage.get("output_tokens", 0),
|
||||
total_tokens=usage.get("input_tokens", 0) + usage.get("output_tokens", 0),
|
||||
)
|
||||
|
||||
return result
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# 响应解析 —— 流式（SSE 事件行协议）
|
||||
# ------------------------------------------------------------------
|
||||
|
||||
@staticmethod
|
||||
def _parse_stream_event(event_type: str, data: dict[str, Any]) -> ChatChunk:
|
||||
"""解析单条 SSE 流事件为 ChatChunk。"""
|
||||
chunk = ChatChunk()
|
||||
|
||||
if event_type == "message_start":
|
||||
message = data.get("message", {})
|
||||
usage = message.get("usage")
|
||||
if usage:
|
||||
chunk.usage = TokenUsage(
|
||||
prompt_tokens=usage.get("input_tokens", 0),
|
||||
prompt_tokens_cached=usage.get("cache_read_input_tokens", 0),
|
||||
completion_tokens=usage.get("output_tokens", 0),
|
||||
)
|
||||
|
||||
elif event_type == "content_block_start":
|
||||
block = data.get("content_block", {})
|
||||
chunk.content_block_index = data.get("index", 0)
|
||||
chunk.content_block_start = block
|
||||
block_type = block.get("type", "")
|
||||
if block_type == "text":
|
||||
chunk.delta_content = block.get("text", "")
|
||||
elif block_type == "thinking":
|
||||
chunk.delta_reasoning = block.get("thinking", "")
|
||||
elif block_type == "tool_use":
|
||||
chunk.tool_call_chunks.append(
|
||||
{
|
||||
"index": data.get("index", 0),
|
||||
"id": block.get("id"),
|
||||
"name": block.get("name"),
|
||||
"args": "",
|
||||
}
|
||||
)
|
||||
|
||||
elif event_type == "content_block_delta":
|
||||
delta = data.get("delta", {})
|
||||
chunk.content_block_index = data.get("index", 0)
|
||||
chunk.content_block_delta = delta
|
||||
delta_type = delta.get("type", "")
|
||||
if delta_type == "text_delta":
|
||||
chunk.delta_content = delta.get("text", "")
|
||||
elif delta_type == "thinking_delta":
|
||||
chunk.delta_reasoning = delta.get("thinking", "")
|
||||
elif delta_type == "signature_delta":
|
||||
# signature 不展示，但由 Agent Runtime 原样保留供下一轮回放。
|
||||
pass
|
||||
elif delta_type == "input_json_delta":
|
||||
chunk.tool_call_chunks.append(
|
||||
{
|
||||
"index": data.get("index", 0),
|
||||
"id": None,
|
||||
"name": None,
|
||||
"args": delta.get("partial_json", ""),
|
||||
}
|
||||
)
|
||||
|
||||
elif event_type == "message_delta":
|
||||
delta = data.get("delta", {})
|
||||
finish_reason_map = {
|
||||
"end_turn": "stop",
|
||||
"max_tokens": "length",
|
||||
"stop_sequence": "stop",
|
||||
"tool_use": "tool_calls",
|
||||
}
|
||||
stop_reason = delta.get("stop_reason")
|
||||
if stop_reason:
|
||||
chunk.finish_reason = finish_reason_map.get(stop_reason, stop_reason)
|
||||
|
||||
usage = data.get("usage")
|
||||
if usage:
|
||||
chunk.usage = TokenUsage(
|
||||
prompt_tokens=0,
|
||||
completion_tokens=usage.get("output_tokens", 0),
|
||||
)
|
||||
|
||||
# content_block_stop, message_stop, ping — no content to extract
|
||||
|
||||
return chunk
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# 公开 API：chat
|
||||
# ------------------------------------------------------------------
|
||||
|
||||
async def chat(
|
||||
self,
|
||||
model: ModelConfig,
|
||||
messages: list[BaseMessage],
|
||||
api_key: str,
|
||||
base_url: 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:
|
||||
url = f"{base_url or self._provider.base_url}/v1/messages"
|
||||
body = self.build_request(
|
||||
model,
|
||||
messages,
|
||||
params=params,
|
||||
tools=tools,
|
||||
tool_choice=tool_choice,
|
||||
)
|
||||
|
||||
try:
|
||||
response = await self._http.post(url, json=body, headers=self._headers(api_key))
|
||||
response.raise_for_status()
|
||||
data = response.json()
|
||||
return self.parse_nonstreaming_response(data, model.name)
|
||||
except httpx.HTTPStatusError as exc:
|
||||
raise AIProviderError.from_provider(
|
||||
message=f"Anthropic API error: {exc.response.status_code} {exc.response.text[:200]}",
|
||||
code="PROVIDER_ERROR",
|
||||
provider=self._provider.name,
|
||||
original=exc,
|
||||
) from exc
|
||||
|
||||
async def chat_stream(
|
||||
self,
|
||||
model: ModelConfig,
|
||||
messages: list[BaseMessage],
|
||||
api_key: str,
|
||||
base_url: str | None = None,
|
||||
params: dict[str, Any] | None = None,
|
||||
tools: list[dict[str, Any]] | None = None,
|
||||
tool_choice: str | dict[str, Any] | None = None,
|
||||
) -> AsyncIterator[ChatChunk]:
|
||||
url = f"{base_url or self._provider.base_url}/v1/messages"
|
||||
body = self.build_request(
|
||||
model,
|
||||
messages,
|
||||
params=params,
|
||||
tools=tools,
|
||||
tool_choice=tool_choice,
|
||||
stream=True,
|
||||
)
|
||||
|
||||
async with self._http.stream(
|
||||
"POST",
|
||||
url,
|
||||
json=body,
|
||||
headers=self._headers(api_key),
|
||||
) as response:
|
||||
if response.status_code != 200:
|
||||
error_text = await response.aread()
|
||||
raise AIProviderError.from_provider(
|
||||
message=f"Anthropic stream error: {response.status_code} {error_text[:200]}",
|
||||
code="PROVIDER_ERROR",
|
||||
provider=self._provider.name,
|
||||
)
|
||||
|
||||
current_event = ""
|
||||
async for line in response.aiter_lines():
|
||||
line = line.strip()
|
||||
if not line:
|
||||
continue
|
||||
if line.startswith("event: "):
|
||||
current_event = line[7:]
|
||||
elif line.startswith("data: "):
|
||||
data_str = line[6:]
|
||||
if data_str and current_event:
|
||||
try:
|
||||
raw = json.loads(data_str)
|
||||
yield self._parse_stream_event(current_event, raw)
|
||||
except json.JSONDecodeError:
|
||||
logger.warning(
|
||||
"anthropic_stream_parse_error",
|
||||
event=current_event,
|
||||
raw=data_str[:100],
|
||||
)
|
||||
current_event = ""
|
||||
|
||||
async def close(self) -> None:
|
||||
if self._client:
|
||||
await self._client.aclose()
|
||||
self._client = None
|
||||
|
||||
|
||||
class AnthropicFormatAdapterFactory:
|
||||
"""工厂：创建 Anthropic 格式适配器。"""
|
||||
|
||||
@staticmethod
|
||||
def create(provider_config: ProviderConfig) -> AnthropicFormatAdapter:
|
||||
return AnthropicFormatAdapter(provider_config=provider_config)
|
||||
@@ -0,0 +1,251 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
from collections.abc import AsyncIterator
|
||||
from typing import Any
|
||||
|
||||
import httpx
|
||||
|
||||
from app.ai_provider.config.models import ModelConfig, ProviderConfig
|
||||
from app.ai_provider.pipeline import ParameterPipeline, PipelineContext, PipelineFactory
|
||||
from app.ai_provider.types.chat import ChatChunk, ChatResult, TokenUsage
|
||||
from app.ai_provider.types.errors import AIProviderError
|
||||
from app.ai_provider.types.messages import AIMessage, BaseMessage, ToolMessage
|
||||
from app.core.logging import logger
|
||||
|
||||
|
||||
class OpenAIFormatAdapter:
|
||||
"""兼容 OpenAI API 格式的适配器（覆盖 deepseek、xiaomi 等）。"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
provider_config: ProviderConfig,
|
||||
pipeline: ParameterPipeline | None = None,
|
||||
) -> None:
|
||||
self._provider = provider_config
|
||||
self._pipeline = pipeline or PipelineFactory.create(provider_config)
|
||||
self._client: httpx.AsyncClient | None = None
|
||||
|
||||
@property
|
||||
def _http(self) -> httpx.AsyncClient:
|
||||
if self._client is None:
|
||||
self._client = httpx.AsyncClient(
|
||||
base_url=self._provider.base_url or "",
|
||||
timeout=self._provider.timeout_seconds,
|
||||
)
|
||||
return self._client
|
||||
|
||||
def _convert_message(self, msg: BaseMessage) -> dict[str, Any]:
|
||||
if isinstance(msg, ToolMessage):
|
||||
return {
|
||||
"role": "tool",
|
||||
"content": msg.content,
|
||||
"tool_call_id": msg.tool_call_id,
|
||||
}
|
||||
if isinstance(msg, AIMessage):
|
||||
d: dict[str, Any] = {"role": "assistant", "content": msg.content}
|
||||
if msg.tool_calls:
|
||||
d["tool_calls"] = [
|
||||
{
|
||||
"id": tc.get("id", f"call_{i}"),
|
||||
"type": "function",
|
||||
"function": {
|
||||
"name": tc["name"],
|
||||
"arguments": json.dumps(tc.get("args", {})),
|
||||
},
|
||||
}
|
||||
for i, tc in enumerate(msg.tool_calls)
|
||||
]
|
||||
return d
|
||||
return {"role": msg.role, "content": msg.content}
|
||||
|
||||
def build_request(
|
||||
self,
|
||||
model: ModelConfig,
|
||||
messages: list[BaseMessage],
|
||||
params: dict[str, Any] | None = None,
|
||||
tools: list[dict[str, Any]] | None = None,
|
||||
tool_choice: str | dict[str, Any] | None = None,
|
||||
stream: bool = False,
|
||||
) -> dict[str, Any]:
|
||||
ctx = PipelineContext(
|
||||
params=params or {},
|
||||
model_config=model,
|
||||
provider_config=self._provider,
|
||||
)
|
||||
pipe_result = self._pipeline.run(ctx)
|
||||
|
||||
body: dict[str, Any] = {
|
||||
"model": model.model,
|
||||
"messages": [self._convert_message(m) for m in messages],
|
||||
"stream": stream,
|
||||
}
|
||||
body.update(pipe_result.params)
|
||||
|
||||
if tools:
|
||||
body["tools"] = tools
|
||||
if tool_choice:
|
||||
body["tool_choice"] = tool_choice
|
||||
|
||||
return body
|
||||
|
||||
def bind_tools(
|
||||
self,
|
||||
request_body: dict[str, Any],
|
||||
tools: list[dict[str, Any]],
|
||||
tool_choice: str | dict[str, Any] | None = None,
|
||||
) -> None:
|
||||
request_body["tools"] = tools
|
||||
if tool_choice:
|
||||
request_body["tool_choice"] = tool_choice
|
||||
|
||||
def parse_chunk(self, raw: dict[str, Any]) -> ChatChunk:
|
||||
chunk = ChatChunk()
|
||||
choices = raw.get("choices", [])
|
||||
if choices:
|
||||
delta = choices[0].get("delta", {})
|
||||
chunk.delta_content = delta.get("content") or ""
|
||||
chunk.delta_reasoning = delta.get("reasoning_content") or ""
|
||||
chunk.finish_reason = choices[0].get("finish_reason")
|
||||
|
||||
if delta.get("tool_calls"):
|
||||
for tc in delta["tool_calls"]:
|
||||
fn = tc.get("function", {})
|
||||
chunk.tool_call_chunks.append(
|
||||
{
|
||||
"index": tc.get("index", 0),
|
||||
"id": tc.get("id"),
|
||||
"name": fn.get("name"),
|
||||
"args": fn.get("arguments", ""),
|
||||
}
|
||||
)
|
||||
usage = raw.get("usage")
|
||||
if usage:
|
||||
chunk.usage = TokenUsage(
|
||||
prompt_tokens=usage.get("prompt_tokens", 0),
|
||||
completion_tokens=usage.get("completion_tokens", 0),
|
||||
total_tokens=usage.get("total_tokens", 0),
|
||||
prompt_tokens_cached=usage.get("prompt_tokens_cached", 0),
|
||||
)
|
||||
return chunk
|
||||
|
||||
def parse_nonstreaming_response(
|
||||
self,
|
||||
raw: dict[str, Any],
|
||||
model_name: str,
|
||||
) -> ChatResult:
|
||||
choice = raw.get("choices", [{}])[0]
|
||||
message = choice.get("message", {})
|
||||
content = message.get("content", "")
|
||||
tool_calls_raw = message.get("tool_calls", [])
|
||||
|
||||
result = ChatResult(
|
||||
content=content,
|
||||
reasoning=message.get("reasoning_content", ""),
|
||||
model_name=model_name,
|
||||
)
|
||||
if tool_calls_raw:
|
||||
result.tool_calls = [
|
||||
{
|
||||
"name": tc["function"]["name"],
|
||||
"args": json.loads(tc["function"].get("arguments", "{}")),
|
||||
"id": tc["id"],
|
||||
}
|
||||
for tc in tool_calls_raw
|
||||
]
|
||||
|
||||
usage = raw.get("usage")
|
||||
if usage:
|
||||
result.usage = TokenUsage(
|
||||
prompt_tokens=usage.get("prompt_tokens", 0),
|
||||
completion_tokens=usage.get("completion_tokens", 0),
|
||||
total_tokens=usage.get("total_tokens", 0),
|
||||
prompt_tokens_cached=usage.get("prompt_tokens_cached", 0),
|
||||
)
|
||||
return result
|
||||
|
||||
async def chat(
|
||||
self,
|
||||
model: ModelConfig,
|
||||
messages: list[BaseMessage],
|
||||
api_key: str,
|
||||
base_url: 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:
|
||||
url = f"{base_url or self._provider.base_url}/chat/completions"
|
||||
body = self.build_request(
|
||||
model, messages, params=params, tools=tools, tool_choice=tool_choice
|
||||
)
|
||||
|
||||
headers = {
|
||||
"Authorization": f"Bearer {api_key}",
|
||||
"Content-Type": "application/json",
|
||||
}
|
||||
|
||||
try:
|
||||
response = await self._http.post(url, json=body, headers=headers)
|
||||
response.raise_for_status()
|
||||
data = response.json()
|
||||
return self.parse_nonstreaming_response(data, model.name)
|
||||
except httpx.HTTPStatusError as exc:
|
||||
raise AIProviderError.from_provider(
|
||||
message=f"OpenAI API error: {exc.response.status_code} {exc.response.text[:200]}",
|
||||
code="PROVIDER_ERROR",
|
||||
provider=self._provider.name,
|
||||
original=exc,
|
||||
) from exc
|
||||
|
||||
async def chat_stream(
|
||||
self,
|
||||
model: ModelConfig,
|
||||
messages: list[BaseMessage],
|
||||
api_key: str,
|
||||
base_url: str | None = None,
|
||||
params: dict[str, Any] | None = None,
|
||||
tools: list[dict[str, Any]] | None = None,
|
||||
tool_choice: str | dict[str, Any] | None = None,
|
||||
) -> AsyncIterator[ChatChunk]:
|
||||
url = f"{base_url or self._provider.base_url}/chat/completions"
|
||||
body = self.build_request(
|
||||
model, messages, params=params, tools=tools, tool_choice=tool_choice, stream=True
|
||||
)
|
||||
|
||||
headers = {
|
||||
"Authorization": f"Bearer {api_key}",
|
||||
"Content-Type": "application/json",
|
||||
}
|
||||
|
||||
async with self._http.stream("POST", url, json=body, headers=headers) as response:
|
||||
if response.status_code != 200:
|
||||
error_text = await response.aread()
|
||||
raise AIProviderError.from_provider(
|
||||
message=f"OpenAI stream error: {response.status_code} {error_text[:200]}",
|
||||
code="PROVIDER_ERROR",
|
||||
provider=self._provider.name,
|
||||
)
|
||||
async for line in response.aiter_lines():
|
||||
if line.startswith("data: "):
|
||||
data_str = line[6:].strip()
|
||||
if data_str == "[DONE]":
|
||||
break
|
||||
if data_str:
|
||||
try:
|
||||
raw = json.loads(data_str)
|
||||
yield self.parse_chunk(raw)
|
||||
except json.JSONDecodeError:
|
||||
logger.warning("openai_stream_parse_error", raw=data_str[:100])
|
||||
|
||||
async def close(self) -> None:
|
||||
if self._client:
|
||||
await self._client.aclose()
|
||||
self._client = None
|
||||
|
||||
|
||||
class OpenAIFormatAdapterFactory:
|
||||
"""工厂：创建 OpenAI 格式适配器。"""
|
||||
|
||||
@staticmethod
|
||||
def create(provider_config: ProviderConfig) -> OpenAIFormatAdapter:
|
||||
return OpenAIFormatAdapter(provider_config=provider_config)
|
||||
@@ -0,0 +1,164 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import random
|
||||
from collections.abc import Callable
|
||||
from dataclasses import dataclass, field
|
||||
from typing import Any, Protocol
|
||||
|
||||
|
||||
@dataclass
|
||||
class PoolContext:
|
||||
request_id: str = ""
|
||||
excluded_endpoints: set[str] = field(default_factory=set)
|
||||
|
||||
|
||||
class SchedulingStrategy(Protocol):
|
||||
def select(self, excluded: set[str]) -> dict[str, Any] | None: ...
|
||||
|
||||
|
||||
class WeightedRoundRobinStrategy:
|
||||
"""策略 1：在所有活跃端点间加权轮询。"""
|
||||
|
||||
def __init__(self, endpoints: list[dict[str, Any]]) -> None:
|
||||
self._endpoints = endpoints
|
||||
self._index = 0
|
||||
|
||||
def select(self, excluded: set[str]) -> dict[str, Any] | None:
|
||||
available = [
|
||||
e
|
||||
for e in self._endpoints
|
||||
if e.get("name") not in excluded and e.get("status", "active") == "active"
|
||||
]
|
||||
if not available:
|
||||
return None
|
||||
weights = [e.get("weight", 100) for e in available]
|
||||
self._index = (self._index + 1) % len(available)
|
||||
return random.choices(available, weights=weights, k=1)[0]
|
||||
|
||||
|
||||
class PriorityFailoverStrategy:
|
||||
"""策略 2：选择最高优先级；当排除时顺延。"""
|
||||
|
||||
def __init__(self, endpoints: list[dict[str, Any]]) -> None:
|
||||
self._endpoints = sorted(endpoints, key=lambda e: e.get("priority", 100))
|
||||
|
||||
def select(self, excluded: set[str]) -> dict[str, Any] | None:
|
||||
for ep in self._endpoints:
|
||||
if ep.get("name") not in excluded and ep.get("status", "active") == "active":
|
||||
return ep
|
||||
return None
|
||||
|
||||
|
||||
class LowestLatencyStrategy:
|
||||
"""策略 3：选择最近延迟最低的端点。"""
|
||||
|
||||
def __init__(self, endpoints: list[dict[str, Any]]) -> None:
|
||||
self._endpoints = endpoints
|
||||
|
||||
def select(self, excluded: set[str]) -> dict[str, Any] | None:
|
||||
available = [
|
||||
e
|
||||
for e in self._endpoints
|
||||
if e.get("name") not in excluded and e.get("status", "active") == "active"
|
||||
]
|
||||
if not available:
|
||||
return None
|
||||
return min(available, key=lambda e: e.get("latency_ms", 1000))
|
||||
|
||||
|
||||
class LeastConnectionsStrategy:
|
||||
"""策略 4：选择活跃连接数最少的端点。"""
|
||||
|
||||
def __init__(self, endpoints: list[dict[str, Any]]) -> None:
|
||||
self._endpoints = endpoints
|
||||
|
||||
def select(self, excluded: set[str]) -> dict[str, Any] | None:
|
||||
available = [
|
||||
e
|
||||
for e in self._endpoints
|
||||
if e.get("name") not in excluded and e.get("status", "active") == "active"
|
||||
]
|
||||
if not available:
|
||||
return None
|
||||
return min(available, key=lambda e: e.get("connections", 0))
|
||||
|
||||
|
||||
class CostAwareStrategy:
|
||||
"""策略 5：按每 token 成本选择最便宜的端点。"""
|
||||
|
||||
def __init__(
|
||||
self, endpoints: list[dict[str, Any]], get_cost: Callable[[str], float] | None = None
|
||||
) -> None:
|
||||
self._endpoints = endpoints
|
||||
self._get_cost = get_cost or (lambda _: 0.0)
|
||||
|
||||
def select(self, excluded: set[str]) -> dict[str, Any] | None:
|
||||
available = [
|
||||
e
|
||||
for e in self._endpoints
|
||||
if e.get("name") not in excluded and e.get("status", "active") == "active"
|
||||
]
|
||||
if not available:
|
||||
return None
|
||||
return min(available, key=lambda e: self._get_cost(e.get("name", "")))
|
||||
|
||||
|
||||
_STRATEGIES = {
|
||||
"weighted_round_robin": WeightedRoundRobinStrategy,
|
||||
"priority_failover": PriorityFailoverStrategy,
|
||||
"lowest_latency": LowestLatencyStrategy,
|
||||
"least_connections": LeastConnectionsStrategy,
|
||||
"cost_aware": CostAwareStrategy,
|
||||
}
|
||||
|
||||
|
||||
class EndpointPool:
|
||||
"""模型池，带可插拔的调度策略。"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
logical_model_name: str,
|
||||
endpoints_configs: list[dict[str, Any]],
|
||||
strategy: str = "weighted_round_robin",
|
||||
) -> None:
|
||||
self.logical_model_name = logical_model_name
|
||||
self._endpoints = endpoints_configs
|
||||
strategy_cls = _STRATEGIES.get(strategy, WeightedRoundRobinStrategy)
|
||||
self.strategy = (
|
||||
strategy_cls(endpoints_configs)
|
||||
if strategy == "priority_failover"
|
||||
else strategy_cls(endpoints_configs)
|
||||
)
|
||||
|
||||
def select(self, excluded: set[str]) -> dict[str, Any] | None:
|
||||
"""使用配置的策略选择一个端点。"""
|
||||
return self.strategy.select(excluded)
|
||||
|
||||
def list_endpoints(self) -> list[dict[str, Any]]:
|
||||
return list(self._endpoints)
|
||||
|
||||
|
||||
class PoolFactory:
|
||||
"""工厂：创建端点池。"""
|
||||
|
||||
@staticmethod
|
||||
def create_pool(
|
||||
logical_model_name: str,
|
||||
endpoints: list[dict[str, Any]],
|
||||
strategy_name: str = "weighted_round_robin",
|
||||
) -> EndpointPool:
|
||||
return EndpointPool(
|
||||
logical_model_name=logical_model_name,
|
||||
endpoints_configs=endpoints,
|
||||
strategy=strategy_name,
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def create_strategy(
|
||||
name: str,
|
||||
endpoints: list[dict[str, Any]],
|
||||
) -> SchedulingStrategy:
|
||||
strategy_cls = _STRATEGIES.get(name)
|
||||
if strategy_cls is None:
|
||||
raise ValueError(f"Unknown scheduling strategy: {name}")
|
||||
return strategy_cls(endpoints)
|
||||
@@ -0,0 +1,115 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
from collections import defaultdict
|
||||
from typing import Any
|
||||
|
||||
from app.ai_provider.types.chat import ChatChunk, ChatResult, TokenUsage
|
||||
|
||||
|
||||
class ToolCallAccumulator:
|
||||
"""将流式 tool call 分片累积为完整的 tool call。
|
||||
|
||||
处理流式场景中常见的多分片 tool call 参数，
|
||||
按索引拼接部分 JSON 字符串。
|
||||
"""
|
||||
|
||||
def __init__(self) -> None:
|
||||
self.content: str = ""
|
||||
self.reasoning: str = ""
|
||||
self._tool_call_slots: dict[int, dict[str, Any]] = defaultdict(
|
||||
lambda: {"index": 0, "id": "", "name": "", "args": ""}
|
||||
)
|
||||
self.usage: TokenUsage | None = None
|
||||
|
||||
def add_chunk(self, chunk: ChatChunk) -> None:
|
||||
"""处理一个 ChatChunk，累积 content/reasoning/tool-calls/usage。"""
|
||||
if chunk.delta_content:
|
||||
self.content += chunk.delta_content
|
||||
if chunk.delta_reasoning:
|
||||
self.reasoning += chunk.delta_reasoning
|
||||
if chunk.usage is not None:
|
||||
self.usage = chunk.usage
|
||||
for tc in chunk.tool_call_chunks:
|
||||
idx = tc.get("index", 0)
|
||||
if idx not in self._tool_call_slots:
|
||||
self._tool_call_slots[idx] = {"index": idx, "id": "", "name": "", "args": ""}
|
||||
slot = self._tool_call_slots[idx]
|
||||
# The first chunk for a given index carries id/name; subsequent
|
||||
# arg-only chunks have them as None and should NOT overwrite.
|
||||
if tc.get("id"):
|
||||
slot["id"] = tc["id"]
|
||||
if tc.get("name"):
|
||||
slot["name"] = tc["name"]
|
||||
slot["args"] += tc.get("args", "")
|
||||
|
||||
def drain(self) -> list[dict[str, Any]]:
|
||||
"""返回所有已累积的 tool call 并重置内部槽位。"""
|
||||
result: list[dict[str, Any]] = []
|
||||
for idx in sorted(self._tool_call_slots.keys()):
|
||||
slot = self._tool_call_slots[idx]
|
||||
result.append(
|
||||
{
|
||||
"index": slot["index"],
|
||||
"id": slot["id"] or f"call_{idx}",
|
||||
"name": slot["name"],
|
||||
"args": slot["args"],
|
||||
}
|
||||
)
|
||||
self._tool_call_slots.clear()
|
||||
return result
|
||||
|
||||
|
||||
class StreamHandler:
|
||||
"""收集流式分片并生成最终的 ChatResult。"""
|
||||
|
||||
def __init__(self) -> None:
|
||||
self._accumulator = ToolCallAccumulator()
|
||||
self._chunks: list[ChatChunk] = []
|
||||
|
||||
def add_chunk(self, chunk: ChatChunk) -> None:
|
||||
"""收集一个流式分片。"""
|
||||
self._chunks.append(chunk)
|
||||
self._accumulator.add_chunk(chunk)
|
||||
|
||||
@property
|
||||
def chunks(self) -> list[ChatChunk]:
|
||||
"""返回所有已收集分片的副本。"""
|
||||
return list(self._chunks)
|
||||
|
||||
def to_result(self, model_name: str = "") -> ChatResult:
|
||||
"""从所有累积分片构建最终的 ChatResult。"""
|
||||
raw_tool_calls = self._accumulator.drain()
|
||||
parsed_calls: list[dict[str, Any]] = []
|
||||
for tc in raw_tool_calls:
|
||||
args_raw = tc.get("args", "{}")
|
||||
try:
|
||||
parsed_args = json.loads(args_raw)
|
||||
except json.JSONDecodeError:
|
||||
parsed_args = {"_raw": args_raw}
|
||||
parsed_calls.append(
|
||||
{
|
||||
"name": tc["name"],
|
||||
"args": parsed_args,
|
||||
"id": tc["id"],
|
||||
}
|
||||
)
|
||||
return ChatResult(
|
||||
content=self._accumulator.content,
|
||||
reasoning=self._accumulator.reasoning,
|
||||
tool_calls=parsed_calls,
|
||||
usage=self._accumulator.usage,
|
||||
model_name=model_name,
|
||||
)
|
||||
|
||||
|
||||
class StreamFactory:
|
||||
"""流处理器和累积器的工厂。"""
|
||||
|
||||
@staticmethod
|
||||
def create_handler() -> StreamHandler:
|
||||
return StreamHandler()
|
||||
|
||||
@staticmethod
|
||||
def create_accumulator() -> ToolCallAccumulator:
|
||||
return ToolCallAccumulator()
|
||||
@@ -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,
|
||||
)
|
||||
@@ -0,0 +1,172 @@
|
||||
"""配置驱动的 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(),
|
||||
],
|
||||
)
|
||||
@@ -0,0 +1,153 @@
|
||||
"""参数处理器注册中心与内置 handlers。
|
||||
|
||||
处理跨参数依赖、结构性转换等复杂场景——这些无法通过
|
||||
YAML ParameterMapping 的 min/max/transform/default 表达。
|
||||
|
||||
使用方式（在 providers.yaml 中）：
|
||||
```yaml
|
||||
parameter_mappings:
|
||||
thinking: { provider_param: thinking, handler: anthropic_thinking }
|
||||
|
||||
parameter_handlers:
|
||||
- strip_unsupported_capability_params
|
||||
```
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from collections.abc import Callable
|
||||
|
||||
from app.ai_provider.pipeline import PipelineContext, PipelineResult
|
||||
|
||||
# Type alias for a parameter handler function.
|
||||
# param_key is passed for per-param handlers (from ParameterMapping.handler);
|
||||
# for global handlers (from ProviderConfig.parameter_handlers), param_key is None.
|
||||
ParameterHandler = Callable[[PipelineContext, PipelineResult, str | None], None]
|
||||
|
||||
|
||||
class ParameterHandlerRegistry:
|
||||
"""全局参数处理器注册中心。
|
||||
|
||||
使用 register() 装饰器注册 handler，名称字符串即可在 YAML 中引用。
|
||||
"""
|
||||
|
||||
_handlers: dict[str, ParameterHandler] = {}
|
||||
|
||||
@classmethod
|
||||
def register(cls, name: str) -> Callable[[ParameterHandler], ParameterHandler]:
|
||||
"""装饰器：将函数注册为命名 handler。"""
|
||||
|
||||
def decorator(handler: ParameterHandler) -> ParameterHandler:
|
||||
cls._handlers[name] = handler
|
||||
return handler
|
||||
|
||||
return decorator
|
||||
|
||||
@classmethod
|
||||
def get(cls, name: str) -> ParameterHandler:
|
||||
"""按名称获取 handler。未找到时抛出 KeyError。"""
|
||||
if name not in cls._handlers:
|
||||
msg = f"Unknown parameter handler: '{name}'. "
|
||||
msg += f"Available: {', '.join(sorted(cls._handlers))}"
|
||||
raise KeyError(msg)
|
||||
return cls._handlers[name]
|
||||
|
||||
@classmethod
|
||||
def list(cls) -> list[str]:
|
||||
"""返回所有已注册的 handler 名称。"""
|
||||
return list(cls._handlers.keys())
|
||||
|
||||
|
||||
# ── 内建 Handlers ──────────────────────────────────────────────
|
||||
|
||||
|
||||
@ParameterHandlerRegistry.register("anthropic_thinking")
|
||||
def handle_anthropic_thinking(
|
||||
_ctx: PipelineContext,
|
||||
result: PipelineResult,
|
||||
_param_key: str | None = None,
|
||||
) -> None:
|
||||
"""处理 Anthropic thinking 参数的跨参数约束。
|
||||
|
||||
Anthropic API 要求：
|
||||
- max_tokens 必须大于 thinking.budget_tokens
|
||||
- 开启 thinking 时 temperature 必须为 1（不应传此参数）
|
||||
- thinking 参数自身需转换为 {type, budget_tokens} 结构
|
||||
"""
|
||||
thinking = result.params.get("thinking")
|
||||
if thinking is None:
|
||||
return
|
||||
|
||||
# 将标量 thinking=True 转换为启用 extended thinking
|
||||
if thinking is True:
|
||||
result.params["thinking"] = {"type": "enabled"}
|
||||
|
||||
# 提取 budget_tokens
|
||||
if isinstance(thinking, dict):
|
||||
budget = thinking.get("budget_tokens", 0)
|
||||
# 确保 max_tokens > budget_tokens
|
||||
max_tokens = result.params.get("max_tokens", 0)
|
||||
if budget > 0 and max_tokens <= budget:
|
||||
result.params["max_tokens"] = budget + 1
|
||||
|
||||
# 开启 thinking 时移除 temperature 参数（Anthropic 静默忽略，
|
||||
# 但显式移除更安全干净）
|
||||
result.params.pop("temperature", None)
|
||||
|
||||
|
||||
@ParameterHandlerRegistry.register("validate_max_tokens")
|
||||
def validate_max_tokens(
|
||||
ctx: PipelineContext,
|
||||
result: PipelineResult,
|
||||
_param_key: str | None = None,
|
||||
) -> None:
|
||||
"""校验 max_tokens 不超过模型上限。"""
|
||||
model_config = ctx.model_config
|
||||
if model_config is None:
|
||||
return
|
||||
max_tokens = result.params.get("max_tokens")
|
||||
if max_tokens is None:
|
||||
return
|
||||
limit = model_config.max_output_tokens
|
||||
if limit and max_tokens > limit:
|
||||
result.params["max_tokens"] = limit
|
||||
|
||||
|
||||
@ParameterHandlerRegistry.register("validate_temperature")
|
||||
def validate_temperature(
|
||||
_ctx: PipelineContext,
|
||||
_result: PipelineResult,
|
||||
_param_key: str | None = None,
|
||||
) -> None:
|
||||
"""温度参数跨提供商规范校验。
|
||||
|
||||
非 reasoning/thinking 模式下强制温度 1.0 不是必需的规则，
|
||||
OpenAI/DeepSeek 无此要求，此 handler 留作扩展点。
|
||||
"""
|
||||
# 预留：可为特定 provider 添加温度校验逻辑
|
||||
pass
|
||||
|
||||
|
||||
@ParameterHandlerRegistry.register("strip_unsupported_capability_params")
|
||||
def strip_unsupported_capability_params(
|
||||
ctx: PipelineContext,
|
||||
result: PipelineResult,
|
||||
_param_key: str | None = None,
|
||||
) -> None:
|
||||
"""根据模型能力过滤不支持的参数。
|
||||
|
||||
当模型不支持 vision 时，移除 image_url 等视觉参数。
|
||||
当模型不支持 tool_use 时，移除 tools 相关参数。
|
||||
"""
|
||||
model_config = ctx.model_config
|
||||
if model_config is None:
|
||||
return
|
||||
|
||||
caps = model_config.resolved_capabilities()
|
||||
|
||||
if not caps.get("vision", False):
|
||||
result.params.pop("image_url", None)
|
||||
result.params.pop("image", None)
|
||||
|
||||
if not caps.get("thinking", False):
|
||||
result.params.pop("thinking", None)
|
||||
@@ -0,0 +1,135 @@
|
||||
"""重排序服务 — Reranker 适配器外观。
|
||||
|
||||
用法::
|
||||
|
||||
from app.ai_provider.reranker import RerankerService
|
||||
|
||||
service = RerankerService(manager)
|
||||
result = await service.rerank("查询", ["文档1", "文档2"])
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from app.ai_provider.reranker.base import (
|
||||
RerankerProtocol,
|
||||
RerankerServiceError,
|
||||
RerankRequest,
|
||||
RerankResponse,
|
||||
RerankResult,
|
||||
)
|
||||
from app.ai_provider.types.errors import AIProviderError
|
||||
|
||||
|
||||
class RerankerService:
|
||||
"""外观类，负责解析 Reranker 模型并委托给 ProviderManager。
|
||||
|
||||
接受 ProviderManager 或实现了以下接口的适配器：
|
||||
- ``.model_registry`` → 具有 ``list_by_capability(capability)`` 方法
|
||||
- ``.create_reranker(name)`` → 返回 RerankerProtocol 实例
|
||||
"""
|
||||
|
||||
def __init__(self, manager) -> None:
|
||||
self._manager = manager
|
||||
self._adapter: RerankerProtocol | None = None
|
||||
|
||||
# ── Registry 兼容适配 ─────────────────────────────────────────
|
||||
|
||||
@property
|
||||
def _registry(self):
|
||||
mgr = self._manager
|
||||
if hasattr(mgr, "model_registry"):
|
||||
return mgr.model_registry
|
||||
return mgr
|
||||
|
||||
@property
|
||||
def _factory(self):
|
||||
mgr = self._manager
|
||||
if hasattr(mgr, "create_reranker"):
|
||||
return mgr
|
||||
return None
|
||||
|
||||
# ── Reranker ───────────────────────────────────────────────────
|
||||
|
||||
@property
|
||||
def adapter(self) -> RerankerProtocol:
|
||||
"""解析并缓存 Reranker 适配器。"""
|
||||
if self._adapter is None:
|
||||
self._adapter = self._resolve()
|
||||
return self._adapter
|
||||
|
||||
async def rerank(
|
||||
self,
|
||||
query: str,
|
||||
documents: list[str],
|
||||
model_hint: str | None = None,
|
||||
top_n: int | None = None,
|
||||
) -> RerankResponse:
|
||||
"""快捷方法：构造请求并调用 Reranker 适配器。"""
|
||||
model = self._resolve_model_name(model_hint)
|
||||
request = RerankRequest(
|
||||
query=query,
|
||||
documents=documents,
|
||||
model=model,
|
||||
top_n=top_n,
|
||||
)
|
||||
return await self.adapter.rerank(request)
|
||||
|
||||
# ── 解析 ───────────────────────────────────────────────────────
|
||||
|
||||
def _resolve(self) -> RerankerProtocol:
|
||||
"""查找第一个活跃的 Reranker 模型并通过工厂创建。"""
|
||||
models = self._list_capability_models("reranker")
|
||||
if not models:
|
||||
raise AIProviderError(
|
||||
status_code=503,
|
||||
code="reranker_no_model",
|
||||
message="No reranker model configured in the model registry.",
|
||||
)
|
||||
factory = self._factory
|
||||
if factory is None:
|
||||
raise AIProviderError(
|
||||
status_code=501,
|
||||
code="reranker_no_factory",
|
||||
message="Reranker adapter factory not available (provider manager missing create_reranker).",
|
||||
)
|
||||
return factory.create_reranker(models[0].name)
|
||||
|
||||
def _resolve_model_name(self, model_hint: str | None) -> str:
|
||||
models = self._list_capability_models("reranker")
|
||||
if model_hint:
|
||||
return model_hint
|
||||
if models:
|
||||
return models[0].name
|
||||
return ""
|
||||
|
||||
def _list_capability_models(self, tag: str):
|
||||
reg = self._registry
|
||||
if hasattr(reg, "list_capability_models"):
|
||||
return reg.list_capability_models(tag)
|
||||
if hasattr(reg, "list_by_capability"):
|
||||
return reg.list_by_capability(tag)
|
||||
return []
|
||||
|
||||
# ── 生命周期 ────────────────────────────────────────────────────
|
||||
|
||||
async def close(self) -> None:
|
||||
if self._adapter is not None:
|
||||
close = getattr(self._adapter, "close", None)
|
||||
if close is not None:
|
||||
await close()
|
||||
|
||||
async def __aenter__(self) -> RerankerService:
|
||||
return self
|
||||
|
||||
async def __aexit__(self, *exc_info: object) -> None:
|
||||
await self.close()
|
||||
|
||||
|
||||
__all__ = [
|
||||
"RerankerProtocol",
|
||||
"RerankRequest",
|
||||
"RerankResponse",
|
||||
"RerankResult",
|
||||
"RerankerService",
|
||||
"RerankerServiceError",
|
||||
]
|
||||
@@ -0,0 +1,147 @@
|
||||
"""重排序服务类型 — 请求/响应模型和协议定义。
|
||||
|
||||
Reranker 适配器通过 HTTP 调用重排序端点，经由 ``RerankerService``
|
||||
或直接从 ``ProviderManager`` 消费。
|
||||
|
||||
用法::
|
||||
|
||||
service = RerankerService(manager)
|
||||
result = await service.rerank("查询", ["文档1", "文档2"])
|
||||
|
||||
# 或通过工厂直接调用:
|
||||
adapter = manager.create_reranker("rerank-model")
|
||||
result = await adapter.rerank(RerankRequest(query="...", documents=["..."]))
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from dataclasses import dataclass
|
||||
from typing import TYPE_CHECKING, Any, Protocol, runtime_checkable
|
||||
|
||||
from app.ai_provider.types.errors import AIProviderError
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from app.ai_provider.config.models import ModelConfig, ModelEndpointConfig
|
||||
|
||||
|
||||
class RerankerServiceError(AIProviderError):
|
||||
"""Reranker API 调用失败的基础错误。"""
|
||||
|
||||
|
||||
# ── 请求 / 响应 ────────────────────────────────────────────────────────
|
||||
|
||||
|
||||
@dataclass
|
||||
class RerankRequest:
|
||||
"""重排序请求。
|
||||
|
||||
属性:
|
||||
query: 查询文本。
|
||||
documents: 待重排序的文档列表。
|
||||
model: 模型标识符。
|
||||
top_n: 返回的最相关文档数量（nil 表示全部）。
|
||||
max_chunks_per_doc: 每个文档最多切分为多少块。
|
||||
return_documents: 返回结果中是否包含文档文本。
|
||||
"""
|
||||
|
||||
query: str
|
||||
documents: list[str]
|
||||
model: str = ""
|
||||
top_n: int | None = None
|
||||
max_chunks_per_doc: int | None = None
|
||||
return_documents: bool = False
|
||||
|
||||
|
||||
@dataclass
|
||||
class RerankResult:
|
||||
"""单条重排序结果。"""
|
||||
|
||||
index: int
|
||||
relevance_score: float
|
||||
text: str | None = None
|
||||
|
||||
|
||||
@dataclass
|
||||
class RerankResponse:
|
||||
"""重排序 API 的完整响应。"""
|
||||
|
||||
results: list[RerankResult]
|
||||
model: str = ""
|
||||
usage: dict[str, int] | None = None
|
||||
|
||||
|
||||
# ── Protocol ────────────────────────────────────────────────────────────
|
||||
|
||||
|
||||
@runtime_checkable
|
||||
class RerankerProtocol(Protocol):
|
||||
"""Reranker 适配器的协议。"""
|
||||
|
||||
async def rerank(self, request: RerankRequest) -> RerankResponse:
|
||||
"""对文档进行重排序。
|
||||
|
||||
异常:
|
||||
RerankerServiceError: API 调用失败或响应异常时抛出。
|
||||
"""
|
||||
...
|
||||
|
||||
async def close(self) -> None:
|
||||
"""释放适配器持有的 HTTP 连接。"""
|
||||
...
|
||||
|
||||
|
||||
# ── 共享工厂基类 ──────────────────────────────────────────────────────
|
||||
|
||||
|
||||
class RerankerAdapterBase:
|
||||
"""Reranker 适配器基类，提供共享的工厂方法。
|
||||
|
||||
子类应在 ``__init__`` 中接受 ``base_url``、``api_key``、``model``、
|
||||
``request_timeout_seconds`` 参数。
|
||||
"""
|
||||
|
||||
_adapter_type: str = "reranker"
|
||||
_default_timeout: int = 60
|
||||
|
||||
@classmethod
|
||||
def from_endpoint(
|
||||
cls,
|
||||
endpoint: ModelEndpointConfig,
|
||||
*,
|
||||
timeout_seconds: int | None = None,
|
||||
) -> Any:
|
||||
"""从已解析的端点配置构建适配器。"""
|
||||
base_url = endpoint.base_url
|
||||
api_key = endpoint.api_key
|
||||
if not base_url:
|
||||
raise RerankerServiceError(
|
||||
status_code=500,
|
||||
code="reranker_endpoint_invalid",
|
||||
message="Reranker endpoint has no resolved base_url.",
|
||||
details={"endpoint": endpoint.name},
|
||||
)
|
||||
return cls(
|
||||
base_url=base_url,
|
||||
api_key=api_key,
|
||||
model=endpoint.model,
|
||||
request_timeout_seconds=timeout_seconds
|
||||
or endpoint.request_timeout_seconds
|
||||
or cls._default_timeout,
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def from_model_config(
|
||||
cls,
|
||||
config: ModelConfig,
|
||||
*,
|
||||
timeout_seconds: int | None = None,
|
||||
) -> Any:
|
||||
"""从逻辑模型配置构建适配器（使用第一个活跃端点）。"""
|
||||
for ep in config.endpoints:
|
||||
if ep.status == "active":
|
||||
return cls.from_endpoint(ep, timeout_seconds=timeout_seconds)
|
||||
raise RerankerServiceError(
|
||||
status_code=500,
|
||||
code="reranker_no_endpoint",
|
||||
message=f"Reranker model '{config.name}' has no active endpoint.",
|
||||
)
|
||||
@@ -0,0 +1,186 @@
|
||||
"""Cohere 兼容格式的 Reranker 适配器 — ``POST /v1/rerank``。
|
||||
|
||||
覆盖 Cohere、mixedbread、Jina AI 等采用 Cohere 兼容重排序端点的提供商。
|
||||
|
||||
请求::
|
||||
|
||||
POST {base_url}/v1/rerank
|
||||
Authorization: Bearer {api_key}
|
||||
Content-Type: application/json
|
||||
|
||||
{
|
||||
"model": "rerank-model",
|
||||
"query": "查询文本",
|
||||
"documents": ["文档1", "文档2"],
|
||||
"top_n": 3,
|
||||
"return_documents": false
|
||||
}
|
||||
|
||||
响应::
|
||||
|
||||
{
|
||||
"results": [
|
||||
{"index": 0, "relevance_score": 0.99, "document": {"text": "文档..."}},
|
||||
{"index": 1, "relevance_score": 0.85, "document": {"text": "文档..."}}
|
||||
],
|
||||
"meta": {"billed_units": {"search_units": 2}}
|
||||
}
|
||||
|
||||
另见:
|
||||
|
||||
- Cohere Rerank API: https://docs.cohere.com/reference/rerank
|
||||
- Jina Reranker: https://api.jina.ai/redoc#tag/rerank
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import Any
|
||||
|
||||
import httpx
|
||||
|
||||
from app.ai_provider.reranker.base import (
|
||||
RerankerAdapterBase,
|
||||
RerankerProtocol,
|
||||
RerankerServiceError,
|
||||
RerankRequest,
|
||||
RerankResponse,
|
||||
RerankResult,
|
||||
)
|
||||
from app.core.logging import logger
|
||||
|
||||
|
||||
class CohereFormatRerankerAdapter(RerankerAdapterBase, RerankerProtocol):
|
||||
"""Cohere 兼容重排序适配器，覆盖多数重排序服务提供商。"""
|
||||
|
||||
_adapter_type = "reranker"
|
||||
_default_timeout = 60
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
*,
|
||||
base_url: str,
|
||||
api_key: str,
|
||||
model: str = "rerank-model",
|
||||
request_timeout_seconds: int = 60,
|
||||
) -> None:
|
||||
self._base_url = base_url.rstrip("/")
|
||||
self._api_key = api_key
|
||||
self._model = model
|
||||
self._timeout = request_timeout_seconds
|
||||
self._client: httpx.AsyncClient | None = None
|
||||
|
||||
# ── 客户端 ───────────────────────────────────────────────────────────
|
||||
|
||||
@property
|
||||
def _http(self) -> httpx.AsyncClient:
|
||||
if self._client is None or self._client.is_closed:
|
||||
self._client = httpx.AsyncClient(timeout=httpx.Timeout(self._timeout))
|
||||
return self._client
|
||||
|
||||
async def close(self) -> None:
|
||||
if self._client is not None and not self._client.is_closed:
|
||||
await self._client.aclose()
|
||||
|
||||
# ── RerankerProtocol ────────────────────────────────────────────────
|
||||
|
||||
async def rerank(self, request: RerankRequest) -> RerankResponse:
|
||||
"""通过 ``POST /v1/rerank`` 对文档进行重排序。"""
|
||||
url = f"{self._base_url}/v1/rerank"
|
||||
|
||||
payload: dict[str, Any] = {
|
||||
"model": request.model or self._model,
|
||||
"query": request.query,
|
||||
"documents": request.documents,
|
||||
"return_documents": request.return_documents,
|
||||
}
|
||||
if request.top_n is not None:
|
||||
payload["top_n"] = request.top_n
|
||||
if request.max_chunks_per_doc is not None:
|
||||
payload["max_chunks_per_doc"] = request.max_chunks_per_doc
|
||||
|
||||
headers = {
|
||||
"Authorization": f"Bearer {self._api_key}",
|
||||
"Content-Type": "application/json",
|
||||
}
|
||||
|
||||
logger.info(
|
||||
"reranker.request",
|
||||
model=payload["model"],
|
||||
query_length=len(request.query),
|
||||
document_count=len(request.documents),
|
||||
top_n=request.top_n,
|
||||
url=url,
|
||||
)
|
||||
|
||||
try:
|
||||
response = await self._http.post(url, json=payload, headers=headers)
|
||||
except httpx.TimeoutException as exc:
|
||||
raise RerankerServiceError(
|
||||
status_code=504,
|
||||
code="reranker_timeout",
|
||||
message="Reranker request timed out.",
|
||||
) from exc
|
||||
except httpx.RequestError as exc:
|
||||
raise RerankerServiceError(
|
||||
status_code=503,
|
||||
code="reranker_connection_error",
|
||||
message=f"Reranker connection error: {exc}",
|
||||
) from exc
|
||||
|
||||
if response.status_code != 200:
|
||||
body = response.text[:500]
|
||||
raise RerankerServiceError(
|
||||
status_code=response.status_code,
|
||||
code="reranker_api_error",
|
||||
message=f"Reranker API returned {response.status_code}: {body}",
|
||||
)
|
||||
|
||||
result = self._parse_response(response.json())
|
||||
|
||||
logger.info(
|
||||
"reranker.success",
|
||||
model=payload["model"],
|
||||
total_results=len(result.results),
|
||||
)
|
||||
return result
|
||||
|
||||
# ── 辅助方法 ──────────────────────────────────────────────────────────
|
||||
|
||||
@staticmethod
|
||||
def _parse_response(body: dict) -> RerankResponse:
|
||||
"""从 Cohere 格式的重排序响应中提取结果。"""
|
||||
try:
|
||||
raw_results = body.get("results", [])
|
||||
results = []
|
||||
for item in raw_results:
|
||||
results.append(
|
||||
RerankResult(
|
||||
index=item.get("index", 0),
|
||||
relevance_score=float(item.get("relevance_score", 0.0)),
|
||||
text=item.get("document", {}).get("text")
|
||||
if isinstance(item.get("document"), dict)
|
||||
else None,
|
||||
)
|
||||
)
|
||||
# 按相关性分数降序排列
|
||||
results.sort(key=lambda r: r.relevance_score, reverse=True)
|
||||
|
||||
# 提取用量信息（不同提供商格式各异）
|
||||
usage: dict[str, int] | None = None
|
||||
meta = body.get("meta")
|
||||
if isinstance(meta, dict):
|
||||
billed = meta.get("billed_units")
|
||||
if isinstance(billed, dict):
|
||||
usage = {k: int(v) for k, v in billed.items() if isinstance(v, (int, float))}
|
||||
|
||||
return RerankResponse(
|
||||
results=results,
|
||||
model=body.get("model") or body.get("id", ""),
|
||||
usage=usage,
|
||||
)
|
||||
except Exception as exc:
|
||||
raise RerankerServiceError(
|
||||
status_code=500,
|
||||
code="reranker_response_parse_error",
|
||||
message=f"Failed to parse reranker response: {exc}",
|
||||
) from exc
|
||||
@@ -0,0 +1,14 @@
|
||||
from app.ai_provider.routing.multimodal import MultiModalDecision, MultiModalRouter
|
||||
from app.ai_provider.routing.permissions import (
|
||||
CapabilityPermissionController,
|
||||
PermissionRule,
|
||||
)
|
||||
from app.ai_provider.routing.selector import ModelSelector
|
||||
|
||||
__all__ = [
|
||||
"MultiModalDecision",
|
||||
"MultiModalRouter",
|
||||
"ModelSelector",
|
||||
"PermissionRule",
|
||||
"CapabilityPermissionController",
|
||||
]
|
||||
@@ -0,0 +1,81 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from enum import StrEnum
|
||||
from typing import Any
|
||||
|
||||
|
||||
class MultiModalDecision(StrEnum):
|
||||
"""基于消息内容分析的选路决策。
|
||||
|
||||
每个值代表所选模型必须满足的一项能力需求。
|
||||
PLAIN_TEXT 表示不需要特殊能力。
|
||||
"""
|
||||
|
||||
PLAIN_TEXT = "plain_text"
|
||||
VISION = "vision"
|
||||
THINKING = "thinking"
|
||||
TOOL_USE = "tool_use"
|
||||
STRUCTURED_OUTPUT = "structured_output"
|
||||
|
||||
|
||||
class MultiModalRouter:
|
||||
"""分析消息列表以确定所需的模型能力。
|
||||
|
||||
用法::
|
||||
|
||||
messages = [{"role": "user", "content": "Hello"}]
|
||||
decisions = MultiModalRouter.decide_from_messages(messages)
|
||||
# -> {MultiModalDecision.PLAIN_TEXT}
|
||||
|
||||
messages = [
|
||||
{
|
||||
"role": "user",
|
||||
"content": [
|
||||
{"type": "text", "text": "What's in this image?"},
|
||||
{"type": "image_url", "image_url": {"url": "..."}},
|
||||
],
|
||||
},
|
||||
]
|
||||
decisions = MultiModalRouter.decide_from_messages(messages)
|
||||
# -> {MultiModalDecision.VISION}
|
||||
"""
|
||||
|
||||
@staticmethod
|
||||
def decide_from_messages(messages: list[dict[str, Any]]) -> set[MultiModalDecision]:
|
||||
"""分析消息并返回所需能力的集合。
|
||||
|
||||
参数:
|
||||
messages: 消息字典列表，每个字典至少包含 ``role`` 和 ``content`` 键（兼容 OpenAI 格式）。
|
||||
|
||||
返回:
|
||||
:class:`MultiModalDecision` 值的集合。如果未检测到特殊能力，
|
||||
返回 ``{PLAIN_TEXT}``。
|
||||
"""
|
||||
if not messages:
|
||||
return {MultiModalDecision.PLAIN_TEXT}
|
||||
|
||||
decisions: set[MultiModalDecision] = set()
|
||||
|
||||
for msg in messages:
|
||||
role = msg.get("role", "")
|
||||
content = msg.get("content", "")
|
||||
|
||||
# --- VISION: content is a list of blocks with image type ---
|
||||
if isinstance(content, list):
|
||||
for block in content:
|
||||
if isinstance(block, dict):
|
||||
block_type = block.get("type", "")
|
||||
if block_type in ("image_url", "image"):
|
||||
decisions.add(MultiModalDecision.VISION)
|
||||
|
||||
# --- TOOL_USE: assistant sent tool_calls, or message is tool result ---
|
||||
if msg.get("tool_calls"):
|
||||
decisions.add(MultiModalDecision.TOOL_USE)
|
||||
if role == "tool":
|
||||
decisions.add(MultiModalDecision.TOOL_USE)
|
||||
|
||||
# --- THINKING: assistant included reasoning ---
|
||||
if msg.get("reasoning"):
|
||||
decisions.add(MultiModalDecision.THINKING)
|
||||
|
||||
return decisions if decisions else {MultiModalDecision.PLAIN_TEXT}
|
||||
@@ -0,0 +1,100 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import fnmatch
|
||||
from typing import Any
|
||||
|
||||
|
||||
class PermissionRule:
|
||||
"""模型能力访问控制的单条权限规则。
|
||||
|
||||
规则使用 ``fnmatch`` 风格的通配符模式匹配模型名称。
|
||||
|
||||
参数:
|
||||
model_pattern: ``fnmatch`` 模式，用于匹配模型名（如 ``"deepseek-*"``）。
|
||||
capability: 此规则管控的能力（如 ``"vision"``）。
|
||||
allowed: 能力是否允许（``True`` 允许，``False`` 拒绝）。
|
||||
|
||||
示例::
|
||||
|
||||
# 拒绝 "basic-*" 模型上的 vision 能力
|
||||
PermissionRule(model_pattern="basic-*", capability="vision", allowed=False)
|
||||
|
||||
# 允许所有模型的 thinking 能力
|
||||
PermissionRule(model_pattern="*", capability="thinking", allowed=True)
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
model_pattern: str,
|
||||
capability: str,
|
||||
allowed: bool = True,
|
||||
) -> None:
|
||||
self.model_pattern = model_pattern
|
||||
self.capability = capability
|
||||
self.allowed = allowed
|
||||
|
||||
def match(self, model_name: str) -> bool:
|
||||
"""检查此规则是否适用于给定的模型名。"""
|
||||
return fnmatch.fnmatch(model_name, self.model_pattern)
|
||||
|
||||
def __repr__(self) -> str:
|
||||
verb = "ALLOW" if self.allowed else "DENY"
|
||||
return f"PermissionRule({verb} {self.capability} on {self.model_pattern!r})"
|
||||
|
||||
def __eq__(self, other: Any) -> bool:
|
||||
if not isinstance(other, PermissionRule):
|
||||
return NotImplemented
|
||||
return (
|
||||
self.model_pattern == other.model_pattern
|
||||
and self.capability == other.capability
|
||||
and self.allowed == other.allowed
|
||||
)
|
||||
|
||||
def __hash__(self) -> int:
|
||||
return hash((self.model_pattern, self.capability, self.allowed))
|
||||
|
||||
|
||||
class CapabilityPermissionController:
|
||||
"""控制哪些模型允许哪些能力。
|
||||
|
||||
按顺序评估 :class:`PermissionRule` 列表。
|
||||
**最后匹配的规则** 胜出（最具体的规则应放在较宽泛的规则之后）。
|
||||
如果没有规则匹配某个模型/能力对，该能力**默认允许**。
|
||||
|
||||
用法::
|
||||
|
||||
controller = CapabilityPermissionController([
|
||||
PermissionRule(model_pattern="*", capability="vision", allowed=False),
|
||||
PermissionRule(model_pattern="vision-*", capability="vision", allowed=True),
|
||||
])
|
||||
|
||||
controller.check("fast-model", "vision") # False
|
||||
controller.check("vision-model", "vision") # True
|
||||
"""
|
||||
|
||||
def __init__(self, rules: list[PermissionRule] | None = None) -> None:
|
||||
self._rules = list(rules) if rules else []
|
||||
|
||||
def check(self, model_name: str, capability: str) -> bool:
|
||||
"""检查 *capability* 是否在 *model_name* 上被允许。
|
||||
|
||||
返回 ``True`` 表示允许，``False`` 表示拒绝。
|
||||
"""
|
||||
result: bool | None = None
|
||||
for rule in self._rules:
|
||||
if rule.capability == capability and rule.match(model_name):
|
||||
result = rule.allowed
|
||||
# Default: allow if no rule matched
|
||||
return result if result is not None else True
|
||||
|
||||
def is_allowed(self, model_name: str, capability: str) -> bool:
|
||||
""":meth:`check` 的别名。"""
|
||||
return self.check(model_name, capability)
|
||||
|
||||
def filter_capabilities(
|
||||
self,
|
||||
model_name: str,
|
||||
capabilities: set[str],
|
||||
) -> set[str]:
|
||||
"""返回 *capabilities* 中在 *model_name* 上允许的子集。"""
|
||||
return {c for c in capabilities if self.check(model_name, c)}
|
||||
@@ -0,0 +1,111 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import Any
|
||||
|
||||
from app.ai_provider.config.models import ModelConfig
|
||||
from app.ai_provider.core.registry import ModelRegistry
|
||||
from app.ai_provider.routing.multimodal import MultiModalDecision
|
||||
from app.ai_provider.types.errors import AIProviderError
|
||||
|
||||
_SELECTOR_CAPABILITY_MAP: dict[MultiModalDecision, str] = {
|
||||
MultiModalDecision.VISION: "vision",
|
||||
MultiModalDecision.THINKING: "thinking",
|
||||
MultiModalDecision.TOOL_USE: "tool_use",
|
||||
MultiModalDecision.STRUCTURED_OUTPUT: "structured_output",
|
||||
MultiModalDecision.PLAIN_TEXT: "chat",
|
||||
}
|
||||
|
||||
|
||||
class ModelSelector:
|
||||
"""为给定上下文选择最佳模型。
|
||||
|
||||
用法::
|
||||
|
||||
selector = ModelSelector(registry)
|
||||
model = selector.select({"capabilities": {"vision", "thinking"}})
|
||||
model = selector.select({"preferred": "fast-model"})
|
||||
model = selector.select_from_decisions({MultiModalDecision.VISION})
|
||||
"""
|
||||
|
||||
def __init__(self, registry: ModelRegistry) -> None:
|
||||
self._registry = registry
|
||||
|
||||
def select(self, context: dict[str, Any]) -> ModelConfig:
|
||||
"""根据给定的上下文选择模型。
|
||||
|
||||
上下文键:
|
||||
|
||||
* ``capabilities`` (``set[str]``, 可选) — 必需的能力。
|
||||
* ``preferred`` (``str``, 可选) — 偏好的模型名或别名。
|
||||
* ``avoid`` (``set[str]``, 可选) — 要排除的模型名。
|
||||
|
||||
如果给了偏好模型但不能满足所有必需的能力，此方法会回退到
|
||||
能满足要求的模型。
|
||||
"""
|
||||
capabilities: set[str] = set(context.get("capabilities", set()))
|
||||
preferred: str | None = context.get("preferred")
|
||||
avoid: set[str] = set(context.get("avoid", set()))
|
||||
|
||||
# If preferred is given, resolve it and check capabilities
|
||||
if preferred is not None:
|
||||
model = self._registry.resolve(preferred)
|
||||
if (
|
||||
not capabilities or self._model_satisfies_all(model, capabilities)
|
||||
) and model.name not in avoid:
|
||||
return model
|
||||
|
||||
# Find models that satisfy all required capabilities
|
||||
if capabilities:
|
||||
return self._select_by_capabilities(capabilities, avoid)
|
||||
|
||||
# Fall back to default
|
||||
return self._registry.resolve()
|
||||
|
||||
def select_from_decisions(
|
||||
self,
|
||||
decisions: set[MultiModalDecision],
|
||||
) -> ModelConfig:
|
||||
"""基于 :class:`MultiModalDecision` 值选择模型。
|
||||
|
||||
将决策翻译为能力字符串并委托给 :meth:`select`。
|
||||
"""
|
||||
capabilities: set[str] = set()
|
||||
for d in decisions:
|
||||
cap = _SELECTOR_CAPABILITY_MAP.get(d)
|
||||
if cap and cap != "chat":
|
||||
capabilities.add(cap)
|
||||
return self.select({"capabilities": capabilities})
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
|
||||
def _select_by_capabilities(
|
||||
self,
|
||||
capabilities: set[str],
|
||||
avoid: set[str],
|
||||
) -> ModelConfig:
|
||||
"""找到支持所有必需能力的模型。"""
|
||||
# Narrow down by each capability (ModelConfig is unhashable, use lists)
|
||||
candidates: list[ModelConfig] = list(self._registry.list_all())
|
||||
|
||||
for cap in capabilities:
|
||||
cap_models = list(self._registry.list_by_capability(cap))
|
||||
cap_names = {m.name for m in cap_models}
|
||||
candidates = [m for m in candidates if m.name in cap_names]
|
||||
|
||||
# Remove avoided models
|
||||
candidates = [m for m in candidates if m.name not in avoid]
|
||||
|
||||
if not candidates:
|
||||
raise AIProviderError(
|
||||
message=f"No active model supports all capabilities: {capabilities}",
|
||||
code="CAPABILITY_NOT_FOUND",
|
||||
status_code=404,
|
||||
)
|
||||
|
||||
# Prefer higher-weight models (higher = better)
|
||||
return max(candidates, key=lambda m: m.weight)
|
||||
|
||||
@staticmethod
|
||||
def _model_satisfies_all(model: ModelConfig, capabilities: set[str]) -> bool:
|
||||
caps = model.resolved_capabilities()
|
||||
return all(caps.get(cap, False) for cap in capabilities)
|
||||
@@ -0,0 +1,21 @@
|
||||
from app.ai_provider.types.chat import ChatChunk, ChatResult, TokenUsage
|
||||
from app.ai_provider.types.errors import AIProviderError
|
||||
from app.ai_provider.types.messages import (
|
||||
AIMessage,
|
||||
BaseMessage,
|
||||
HumanMessage,
|
||||
SystemMessage,
|
||||
ToolMessage,
|
||||
)
|
||||
|
||||
__all__ = [
|
||||
"AIMessage",
|
||||
"BaseMessage",
|
||||
"HumanMessage",
|
||||
"SystemMessage",
|
||||
"ToolMessage",
|
||||
"ChatChunk",
|
||||
"ChatResult",
|
||||
"TokenUsage",
|
||||
"AIProviderError",
|
||||
]
|
||||
@@ -0,0 +1,41 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from dataclasses import dataclass, field
|
||||
from typing import Any
|
||||
|
||||
|
||||
@dataclass
|
||||
class TokenUsage:
|
||||
prompt_tokens: int = 0
|
||||
completion_tokens: int = 0
|
||||
total_tokens: int = 0
|
||||
prompt_tokens_cached: int = 0 # 缓存命中的 prompt token 数（按缓存费率计费）
|
||||
|
||||
def __add__(self, other: TokenUsage) -> TokenUsage:
|
||||
return TokenUsage(
|
||||
prompt_tokens=self.prompt_tokens + other.prompt_tokens,
|
||||
completion_tokens=self.completion_tokens + other.completion_tokens,
|
||||
total_tokens=self.total_tokens + other.total_tokens,
|
||||
prompt_tokens_cached=self.prompt_tokens_cached + other.prompt_tokens_cached,
|
||||
)
|
||||
|
||||
|
||||
@dataclass
|
||||
class ChatChunk:
|
||||
delta_content: str = ""
|
||||
delta_reasoning: str = ""
|
||||
tool_call_chunks: list[dict[str, Any]] = field(default_factory=list)
|
||||
usage: TokenUsage | None = None
|
||||
finish_reason: str | None = None
|
||||
content_block_index: int | None = None
|
||||
content_block_start: dict[str, Any] | None = None
|
||||
content_block_delta: dict[str, Any] | None = None
|
||||
|
||||
|
||||
@dataclass
|
||||
class ChatResult:
|
||||
content: str = ""
|
||||
reasoning: str = ""
|
||||
tool_calls: list[dict[str, Any]] = field(default_factory=list)
|
||||
usage: TokenUsage | None = None
|
||||
model_name: str = ""
|
||||
@@ -0,0 +1,33 @@
|
||||
from __future__ import annotations
|
||||
|
||||
|
||||
class AIProviderError(Exception):
|
||||
def __init__(
|
||||
self,
|
||||
message: str,
|
||||
code: str = "UNKNOWN",
|
||||
status_code: int = 500,
|
||||
provider: str | None = None,
|
||||
original: Exception | None = None,
|
||||
) -> None:
|
||||
self.message = message
|
||||
self.code = code
|
||||
self.status_code = status_code
|
||||
self.provider = provider
|
||||
self.original = original
|
||||
super().__init__(message)
|
||||
|
||||
@classmethod
|
||||
def from_provider(
|
||||
cls,
|
||||
message: str,
|
||||
code: str = "PROVIDER_ERROR",
|
||||
provider: str | None = None,
|
||||
original: Exception | None = None,
|
||||
) -> AIProviderError:
|
||||
return cls(
|
||||
message=message, code=code, status_code=502, provider=provider, original=original
|
||||
)
|
||||
|
||||
def __str__(self) -> str:
|
||||
return self.message
|
||||
@@ -0,0 +1,81 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import Any, Literal
|
||||
|
||||
Role = Literal["system", "user", "assistant", "tool"]
|
||||
|
||||
|
||||
class BaseMessage:
|
||||
role: Role
|
||||
content: str
|
||||
|
||||
def __init__(self, *, content: str, role: Role) -> None:
|
||||
self.content = content
|
||||
self.role = role
|
||||
|
||||
def to_dict(self) -> dict[str, Any]:
|
||||
return {"role": self.role, "content": self.content}
|
||||
|
||||
@classmethod
|
||||
def from_dict(cls, d: dict[str, Any]) -> BaseMessage:
|
||||
role = d["role"]
|
||||
content = d.get("content", "")
|
||||
if role == "system":
|
||||
return SystemMessage(content=content)
|
||||
if role == "user":
|
||||
return HumanMessage(content=content)
|
||||
if role == "assistant":
|
||||
return AIMessage(
|
||||
content=content,
|
||||
tool_calls=d.get("tool_calls"),
|
||||
reasoning=d.get("reasoning"),
|
||||
raw_content_blocks=d.get("raw_content_blocks"),
|
||||
)
|
||||
if role == "tool":
|
||||
return ToolMessage(content=content, tool_call_id=d.get("tool_call_id", ""))
|
||||
raise ValueError(f"Unknown role: {role}")
|
||||
|
||||
|
||||
class SystemMessage(BaseMessage):
|
||||
def __init__(self, content: str) -> None:
|
||||
super().__init__(content=content, role="system")
|
||||
|
||||
|
||||
class HumanMessage(BaseMessage):
|
||||
def __init__(self, content: str) -> None:
|
||||
super().__init__(content=content, role="user")
|
||||
|
||||
|
||||
class AIMessage(BaseMessage):
|
||||
def __init__(
|
||||
self,
|
||||
content: str,
|
||||
tool_calls: list[dict[str, Any]] | None = None,
|
||||
reasoning: str | None = None,
|
||||
raw_content_blocks: list[dict[str, Any]] | None = None,
|
||||
) -> None:
|
||||
super().__init__(content=content, role="assistant")
|
||||
self.tool_calls = tool_calls or []
|
||||
self.reasoning = reasoning or ""
|
||||
self.raw_content_blocks = raw_content_blocks or []
|
||||
|
||||
def to_dict(self) -> dict[str, Any]:
|
||||
d = super().to_dict()
|
||||
if self.tool_calls:
|
||||
d["tool_calls"] = self.tool_calls
|
||||
if self.reasoning:
|
||||
d["reasoning"] = self.reasoning
|
||||
if self.raw_content_blocks:
|
||||
d["raw_content_blocks"] = self.raw_content_blocks
|
||||
return d
|
||||
|
||||
|
||||
class ToolMessage(BaseMessage):
|
||||
def __init__(self, content: str, tool_call_id: str) -> None:
|
||||
super().__init__(content=content, role="tool")
|
||||
self.tool_call_id = tool_call_id
|
||||
|
||||
def to_dict(self) -> dict[str, Any]:
|
||||
d = super().to_dict()
|
||||
d["tool_call_id"] = self.tool_call_id
|
||||
return d
|
||||
Reference in new issue
Block a user