145 lines
4.5 KiB
Python
145 lines
4.5 KiB
Python
"""嵌入服务类型 — 请求/响应模型和协议定义。
|
|
|
|
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.",
|
|
)
|