170 lines
5.5 KiB
Python
170 lines
5.5 KiB
Python
"""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
|