Codex
207 lines
7.1 KiB
Python
207 lines
7.1 KiB
Python
"""
|
||
Embedding 服务
|
||
|
||
通过 OpenAI 兼容的 Embeddings API 将文本转换为向量
|
||
支持所有兼容 OpenAI 接口的提供商（OpenAI、Qwen、Ollama 等）
|
||
"""
|
||
import logging
|
||
import math
|
||
from typing import List, Optional
|
||
|
||
from sqlalchemy import select
|
||
from sqlalchemy.ext.asyncio import AsyncSession
|
||
|
||
logger = logging.getLogger(__name__)
|
||
|
||
# 单次批量请求的默认最大文本数
|
||
DEFAULT_BATCH_SIZE = 50
|
||
|
||
# 不同提供商的批次大小限制
|
||
PROVIDER_BATCH_SIZE = {
|
||
'qwen': 10, # 阿里云 DashScope 限制 10
|
||
'dashscope': 10,
|
||
'siliconflow': 10, # 硅基流动限制较小
|
||
'ollama': 1, # Ollama 通常逐条处理
|
||
}
|
||
|
||
|
||
|
||
class EmbeddingService:
|
||
"""
|
||
Embedding 服务
|
||
|
||
通过模型 ID 获取对应的提供商，调用 OpenAI 兼容的 Embeddings API
|
||
"""
|
||
|
||
def __init__(self, db: AsyncSession):
|
||
self._db = db
|
||
self._client_cache = {}
|
||
|
||
async def _get_client_and_model(self, model_id: str):
|
||
"""
|
||
根据模型 ID 获取异步客户端和模型名称
|
||
|
||
Returns:
|
||
(async_client, model_name, max_tokens, provider_type)
|
||
"""
|
||
from ai_platform.models import LLMModel, LLMProvider
|
||
|
||
result = await self._db.execute(
|
||
select(LLMModel).where(
|
||
LLMModel.id == model_id,
|
||
LLMModel.is_active == True,
|
||
LLMModel.is_deleted == False
|
||
)
|
||
)
|
||
model = result.scalar_one_or_none()
|
||
if not model:
|
||
raise ValueError(f'Embedding 模型不存在或已禁用: {model_id}')
|
||
if model.model_type != 'embedding':
|
||
raise ValueError(f'模型 {model.display_name} 不是 Embedding 类型')
|
||
|
||
provider_result = await self._db.execute(
|
||
select(LLMProvider).where(
|
||
LLMProvider.id == model.provider_id,
|
||
LLMProvider.is_active == True,
|
||
LLMProvider.is_deleted == False
|
||
)
|
||
)
|
||
provider = provider_result.scalar_one_or_none()
|
||
if not provider:
|
||
raise ValueError('Embedding 模型对应的提供商不存在或已禁用')
|
||
|
||
cache_key = str(provider.id)
|
||
if cache_key not in self._client_cache:
|
||
import httpx
|
||
from openai import AsyncOpenAI
|
||
|
||
# 根据提供商类型确定 base_url
|
||
if provider.provider_type == 'ollama':
|
||
base_url = (provider.ollama_host or 'http://localhost:11434').rstrip('/') + '/v1'
|
||
else:
|
||
from ai_platform.providers.registry import ProviderRegistry
|
||
provider_class = ProviderRegistry.get(provider.provider_type)
|
||
default_base = getattr(provider_class, 'DEFAULT_API_BASE', 'https://api.openai.com/v1') if provider_class else 'https://api.openai.com/v1'
|
||
base_url = provider.api_base or default_base
|
||
|
||
self._client_cache[cache_key] = AsyncOpenAI(
|
||
api_key=provider.api_key or 'ollama',
|
||
base_url=base_url,
|
||
timeout=httpx.Timeout(120.0, connect=30.0),
|
||
max_retries=5,
|
||
)
|
||
|
||
return self._client_cache[cache_key], model.model_name, model.context_window or 8191, provider.provider_type
|
||
|
||
async def embed_text(
|
||
self,
|
||
model_id: str,
|
||
text: str,
|
||
dimensions: Optional[int] = None,
|
||
) -> List[float]:
|
||
"""
|
||
将单个文本转换为向量
|
||
|
||
Args:
|
||
model_id: Embedding 模型 ID
|
||
text: 文本内容
|
||
dimensions: 向量维度（可选，部分模型支持）
|
||
|
||
Returns:
|
||
向量列表 List[float]
|
||
"""
|
||
results = await self.embed_texts(model_id, [text], dimensions)
|
||
return results[0]
|
||
|
||
async def embed_texts(
|
||
self,
|
||
model_id: str,
|
||
texts: List[str],
|
||
dimensions: Optional[int] = None,
|
||
) -> List[List[float]]:
|
||
"""
|
||
批量将文本转换为向量
|
||
|
||
Args:
|
||
model_id: Embedding 模型 ID
|
||
texts: 文本列表
|
||
dimensions: 向量维度（可选）
|
||
|
||
Returns:
|
||
向量列表 List[List[float]]
|
||
"""
|
||
if not texts:
|
||
return []
|
||
|
||
client, model_name, max_tokens, provider_type = await self._get_client_and_model(model_id)
|
||
|
||
# 根据提供商类型确定批次大小
|
||
batch_size = PROVIDER_BATCH_SIZE.get(provider_type, DEFAULT_BATCH_SIZE)
|
||
|
||
# 预处理：截断超长文本
|
||
processed_texts = []
|
||
for text in texts:
|
||
if not text or not text.strip():
|
||
processed_texts.append(" ")
|
||
else:
|
||
# 粗略估算 token 数（中文约 1 字 = 1.5 token，英文约 4 字符 = 1 token）
|
||
# 保守截断到 max_tokens * 2 个字符
|
||
max_chars = max_tokens * 2
|
||
if len(text) > max_chars:
|
||
processed_texts.append(text[:max_chars])
|
||
else:
|
||
processed_texts.append(text)
|
||
|
||
# 分批处理
|
||
all_embeddings = [None] * len(processed_texts)
|
||
total_batches = math.ceil(len(processed_texts) / batch_size)
|
||
|
||
for batch_idx in range(total_batches):
|
||
start = batch_idx * batch_size
|
||
end = min(start + batch_size, len(processed_texts))
|
||
batch_texts = processed_texts[start:end]
|
||
|
||
try:
|
||
kwargs = {
|
||
'model': model_name,
|
||
'input': batch_texts,
|
||
}
|
||
# 仅对明确支持 dimensions 参数的模型传递该参数
|
||
if dimensions:
|
||
model_lower = model_name.lower()
|
||
# OpenAI text-embedding-3 系列原生支持任意 dimensions
|
||
if 'text-embedding-3' in model_lower:
|
||
kwargs['dimensions'] = dimensions
|
||
# DashScope text-embedding-v3 只接受 [64,128,256,512,768,1024]
|
||
elif 'text-embedding-v3' in model_lower and dimensions in (64, 128, 256, 512, 768, 1024):
|
||
kwargs['dimensions'] = dimensions
|
||
|
||
response = await client.embeddings.create(**kwargs)
|
||
|
||
for item in response.data:
|
||
all_embeddings[start + item.index] = item.embedding
|
||
|
||
except Exception as e:
|
||
logger.exception(f'Embedding 批次 {batch_idx + 1}/{total_batches} 失败: {e}')
|
||
raise ValueError(f'Embedding 调用失败: {str(e)}')
|
||
|
||
# 检查是否所有向量都已生成
|
||
for i, emb in enumerate(all_embeddings):
|
||
if emb is None:
|
||
raise ValueError(f'第 {i} 个文本的向量未生成')
|
||
|
||
return all_embeddings
|
||
|
||
async def get_embedding_dimensions(self, model_id: str) -> int:
|
||
"""
|
||
获取模型的向量维度（通过嵌入一个测试文本来检测）
|
||
|
||
Args:
|
||
model_id: Embedding 模型 ID
|
||
|
||
Returns:
|
||
向量维度
|
||
"""
|
||
test_embedding = await self.embed_text(model_id, "test")
|
||
return len(test_embedding)
|