wehub-resource-sync
355 lines
14 KiB
Python
355 lines
14 KiB
Python
import time
|
||
import requests
|
||
from typing import List, Callable, Optional
|
||
from config.settings import settings
|
||
from config.embedding_config import EmbeddingConfig
|
||
import threading
|
||
|
||
|
||
class _OpenAIEmbeddingAPI:
|
||
"""OpenAI Embedding API 客户端 - 内部单例实现"""
|
||
|
||
def __init__(self, base_url: str = None, token: str = None, model: str = None, max_batch_size: int = 100):
|
||
# 初始化embedding配置
|
||
embedding_config = EmbeddingConfig()
|
||
provider_name, model_name = embedding_config.get_default_model()
|
||
model_info = embedding_config.get_model_info(provider_name, model_name)
|
||
|
||
# 使用新配置系统获取配置,允许通过参数覆盖
|
||
self.base_url = (base_url or (model_info.get("api_base_url") if model_info else "")).rstrip('/')
|
||
self.token = token or (model_info.get("api_key") if model_info else "")
|
||
self.model = model or (model_info.get("model") if model_info else model_name)
|
||
self.headers = {
|
||
"Authorization": f"Bearer {self.token}",
|
||
"Content-Type": "application/json"
|
||
}
|
||
self.max_batch_size = max_batch_size
|
||
self.request_timeout = 120
|
||
self.embedding_dim = None
|
||
self._initialized = False
|
||
self._lock = threading.Lock()
|
||
|
||
print(f"🔄 OpenAIEmbeddingAPI 已创建 (model: {self.model}, batch_size: {self.max_batch_size})")
|
||
|
||
def _lazy_init(self):
|
||
"""延迟初始化,在第一次实际调用时初始化"""
|
||
if self._initialized:
|
||
return
|
||
|
||
with self._lock:
|
||
if self._initialized:
|
||
return
|
||
|
||
print(f"🔍 首次调用,正在获取 embedding 维度...")
|
||
# 在第一次调用时获取维度,而不是在初始化时
|
||
if self.embedding_dim is None:
|
||
self.embedding_dim = self._get_actual_embedding_dimension()
|
||
self._initialized = True
|
||
print(f"✅ OpenAIEmbeddingAPI 初始化完成,维度: {self.embedding_dim}")
|
||
|
||
def set_batch_size(self, batch_size: int):
|
||
"""动态设置批处理大小"""
|
||
self._lazy_init()
|
||
self.max_batch_size = batch_size
|
||
print(f"📦 Embedding批处理大小已设置为: {self.max_batch_size}")
|
||
|
||
def test_connection(self):
|
||
"""
|
||
测试API连接(手动调用) - 通过实际调用embedding接口来测试
|
||
|
||
不再依赖/health端点，而是直接测试/v1/embeddings端点
|
||
这样可以同时验证连接性和获取embedding维度
|
||
|
||
Returns:
|
||
dict: {"success": bool, "message": str, "dimension": int}
|
||
"""
|
||
try:
|
||
print("🔍 测试embedding API连接...")
|
||
|
||
# 使用一个简单的测试文本
|
||
test_text = "测试连接"
|
||
|
||
# 直接调用embedding接口进行测试
|
||
test_response = self._encode_single_batch([test_text], get_dimension=True)
|
||
|
||
# 检查返回结果
|
||
if test_response and len(test_response) > 0:
|
||
dimension = len(test_response[0])
|
||
print(f"✅ API连接测试成功，embedding维度: {dimension}")
|
||
return {
|
||
"success": True,
|
||
"message": f"API连接成功，embedding维度: {dimension}",
|
||
"dimension": dimension
|
||
}
|
||
else:
|
||
print("❌ API返回了空结果")
|
||
return {
|
||
"success": False,
|
||
"message": "API返回了空结果，请检查模型配置",
|
||
"dimension": None
|
||
}
|
||
|
||
except requests.exceptions.ConnectionError as e:
|
||
# 连接错误（无法连接到服务器）
|
||
error_msg = f"无法连接到embedding服务 ({self.base_url}): {str(e)}"
|
||
print(f"❌ {error_msg}")
|
||
return {
|
||
"success": False,
|
||
"message": error_msg,
|
||
"dimension": None
|
||
}
|
||
except requests.exceptions.Timeout as e:
|
||
# 超时错误
|
||
error_msg = f"连接超时 (timeout={self.request_timeout}s)，请检查服务是否正常运行: {str(e)}"
|
||
print(f"❌ {error_msg}")
|
||
return {
|
||
"success": False,
|
||
"message": error_msg,
|
||
"dimension": None
|
||
}
|
||
except requests.exceptions.HTTPError as e:
|
||
# HTTP错误（4xx, 5xx）
|
||
error_msg = f"HTTP错误: {str(e)}"
|
||
print(f"❌ {error_msg}")
|
||
return {
|
||
"success": False,
|
||
"message": error_msg,
|
||
"dimension": None
|
||
}
|
||
except Exception as e:
|
||
# 其他未预期的错误
|
||
error_msg = f"测试连接时发生错误: {str(e)}"
|
||
print(f"❌ {error_msg}")
|
||
return {
|
||
"success": False,
|
||
"message": error_msg,
|
||
"dimension": None
|
||
}
|
||
|
||
def _get_actual_embedding_dimension(self) -> int:
|
||
"""通过实际调用API获取embedding维度"""
|
||
try:
|
||
print("🔍 正在检测embedding维度...")
|
||
test_response = self._encode_single_batch(["测试文本"], get_dimension=True)
|
||
if test_response and len(test_response) > 0:
|
||
actual_dim = len(test_response[0])
|
||
print(f"✅ 检测到embedding维度: {actual_dim}")
|
||
return actual_dim
|
||
else:
|
||
print("⚠️ 无法检测embedding维度，使用默认值1024")
|
||
return 1024
|
||
except Exception as e:
|
||
print(f"⚠️ 检测embedding维度时出错: {e}，使用默认值1024")
|
||
return 1024
|
||
|
||
def encode_texts(self, texts: List[str]) -> List[List[float]]:
|
||
"""批量编码文本为向量（带批次级降级）"""
|
||
self._lazy_init()
|
||
if not texts:
|
||
return []
|
||
|
||
all_embeddings = []
|
||
total_batches = (len(texts) + self.max_batch_size - 1) // self.max_batch_size
|
||
|
||
print(f"📦 开始批量向量化: {len(texts)} 个文本，分为 {total_batches} 个批次")
|
||
|
||
for i in range(0, len(texts), self.max_batch_size):
|
||
batch_texts = texts[i:i + self.max_batch_size]
|
||
batch_num = i // self.max_batch_size + 1
|
||
|
||
print(f" 处理批次 {batch_num}/{total_batches} ({len(batch_texts)} 个文本)")
|
||
|
||
try:
|
||
# 尝试批量编码
|
||
batch_embeddings = self._encode_single_batch(batch_texts)
|
||
all_embeddings.extend(batch_embeddings)
|
||
except Exception as batch_error:
|
||
# 批量失败，降级为逐条处理（仅对当前批次）
|
||
print(f" ⚠️ 批次 {batch_num} 批量编码失败: {batch_error}，降级为逐条处理")
|
||
|
||
for j, text in enumerate(batch_texts):
|
||
try:
|
||
single_embedding = self._encode_single_batch([text])
|
||
all_embeddings.append(single_embedding[0])
|
||
except Exception as single_error:
|
||
# 单条也失败，使用 None 占位
|
||
print(f" ❌ 批次 {batch_num} 文本 {j+1}/{len(batch_texts)} 编码失败: {single_error}")
|
||
all_embeddings.append(None)
|
||
|
||
if i + self.max_batch_size < len(texts):
|
||
time.sleep(0.2)
|
||
|
||
success_count = sum(1 for e in all_embeddings if e is not None)
|
||
fail_count = len(all_embeddings) - success_count
|
||
print(f"✅ 批量向量化完成: 成功 {success_count} 个，失败 {fail_count} 个")
|
||
|
||
return all_embeddings
|
||
|
||
def encode_texts_with_progress(self, texts: List[str],
|
||
progress_callback: Optional[Callable[[int, int, str], None]] = None) -> List[
|
||
List[float]]:
|
||
"""批量编码文本为向量（带进度回调和批次级降级）"""
|
||
self._lazy_init()
|
||
if not texts:
|
||
return []
|
||
|
||
all_embeddings = []
|
||
total_batches = (len(texts) + self.max_batch_size - 1) // self.max_batch_size
|
||
|
||
# 初始进度回调
|
||
if progress_callback:
|
||
progress_callback(0, total_batches, f"开始向量化 {len(texts)} 个文本，分为 {total_batches} 个批次")
|
||
|
||
print(f"📦 开始批量向量化: {len(texts)} 个文本，分为 {total_batches} 个批次")
|
||
|
||
for i in range(0, len(texts), self.max_batch_size):
|
||
batch_texts = texts[i:i + self.max_batch_size]
|
||
batch_num = i // self.max_batch_size + 1
|
||
|
||
# 批次开始回调
|
||
if progress_callback:
|
||
progress_callback(batch_num - 1, total_batches,
|
||
f"正在处理第 {batch_num}/{total_batches} 批次 ({len(batch_texts)} 个文本)")
|
||
|
||
print(f" 处理批次 {batch_num}/{total_batches} ({len(batch_texts)} 个文本)")
|
||
|
||
start_time = time.time()
|
||
|
||
try:
|
||
# 尝试批量编码
|
||
batch_embeddings = self._encode_single_batch(batch_texts)
|
||
all_embeddings.extend(batch_embeddings)
|
||
except Exception as batch_error:
|
||
# 批量失败，降级为逐条处理（仅对当前批次）
|
||
print(f" ⚠️ 批次 {batch_num} 批量编码失败: {batch_error}，降级为逐条处理")
|
||
|
||
for j, text in enumerate(batch_texts):
|
||
try:
|
||
single_embedding = self._encode_single_batch([text])
|
||
all_embeddings.append(single_embedding[0])
|
||
except Exception as single_error:
|
||
# 单条也失败，使用 None 占位
|
||
print(f" ❌ 批次 {batch_num} 文本 {j+1}/{len(batch_texts)} 编码失败: {single_error}")
|
||
all_embeddings.append(None)
|
||
|
||
end_time = time.time()
|
||
|
||
# 批次完成回调
|
||
if progress_callback:
|
||
progress_callback(batch_num, total_batches,
|
||
f"第 {batch_num}/{total_batches} 批次完成 ({end_time - start_time:.2f}秒)")
|
||
|
||
if i + self.max_batch_size < len(texts):
|
||
time.sleep(0.2)
|
||
|
||
# 统计成功/失败
|
||
success_count = sum(1 for e in all_embeddings if e is not None)
|
||
fail_count = len(all_embeddings) - success_count
|
||
print(f"✅ 批量向量化完成: 成功 {success_count} 个，失败 {fail_count} 个")
|
||
|
||
# 最终完成回调
|
||
if progress_callback:
|
||
progress_callback(total_batches, total_batches, f"向量化完成，成功 {success_count} 个，失败 {fail_count} 个")
|
||
|
||
return all_embeddings
|
||
|
||
def _encode_single_batch(self, texts: List[str], get_dimension: bool = False) -> List[List[float]]:
|
||
"""
|
||
编码单个批次的文本
|
||
|
||
Args:
|
||
texts: 要编码的文本列表
|
||
get_dimension: 是否用于获取维度（测试连接时使用），为True时失败会抛出异常
|
||
|
||
Returns:
|
||
向量列表
|
||
|
||
Raises:
|
||
当get_dimension=True且失败时抛出异常
|
||
"""
|
||
payload = {
|
||
"model": self.model,
|
||
"input": texts,
|
||
"encoding_format": "float"
|
||
}
|
||
|
||
max_retries = 5
|
||
last_error = None
|
||
|
||
for attempt in range(max_retries):
|
||
try:
|
||
start_time = time.time()
|
||
response = requests.post(
|
||
f"{self.base_url}/v1/embeddings",
|
||
headers=self.headers,
|
||
json=payload,
|
||
timeout=self.request_timeout
|
||
)
|
||
end_time = time.time()
|
||
|
||
if response.status_code == 200:
|
||
data = response.json()
|
||
embeddings = [item["embedding"] for item in data["data"]]
|
||
if not get_dimension:
|
||
print(f" ✅ 批次完成 ({end_time - start_time:.2f}秒)")
|
||
return embeddings
|
||
else:
|
||
error_msg = f"API请求失败: {response.status_code} - {response.text}"
|
||
last_error = Exception(error_msg)
|
||
if attempt == max_retries - 1:
|
||
# 最后一次重试失败，抛出异常
|
||
raise last_error
|
||
else:
|
||
print(f" ⚠️ {error_msg}, 重试中... ({attempt + 1}/{max_retries})")
|
||
time.sleep(2 ** attempt)
|
||
|
||
except requests.exceptions.Timeout as e:
|
||
last_error = e
|
||
if attempt == max_retries - 1:
|
||
# 最后一次重试失败，抛出异常
|
||
raise requests.exceptions.Timeout("API请求超时")
|
||
else:
|
||
print(f" ⚠️ API请求超时，重试中... ({attempt + 1}/{max_retries})")
|
||
time.sleep(2 ** attempt)
|
||
except requests.exceptions.ConnectionError as e:
|
||
# 连接错误通常不需要重试，直接抛出
|
||
last_error = e
|
||
raise
|
||
except requests.exceptions.RequestException as e:
|
||
# 其他requests相关错误
|
||
last_error = e
|
||
if attempt == max_retries - 1:
|
||
raise
|
||
else:
|
||
print(f" ⚠️ 请求异常: {e}, 重试中... ({attempt + 1}/{max_retries})")
|
||
time.sleep(2 ** attempt)
|
||
except Exception as e:
|
||
# 其他未预期的异常（如JSON解析错误）
|
||
last_error = e
|
||
if attempt == max_retries - 1:
|
||
raise Exception(f"API请求异常: {e}")
|
||
else:
|
||
print(f" ⚠️ API请求异常: {e}, 重试中... ({attempt + 1}/{max_retries})")
|
||
time.sleep(2 ** attempt)
|
||
|
||
# 所有重试失败后抛出异常，让调用者决定如何处理
|
||
print(f" ❌ 批次编码失败，已达最大重试次数")
|
||
raise last_error or Exception("批次编码失败，已达最大重试次数")
|
||
|
||
|
||
# 模块级单例实例
|
||
_embedding_instance = None
|
||
_embedding_lock = threading.Lock()
|
||
|
||
|
||
class OpenAIEmbeddingAPI:
|
||
"""OpenAIEmbeddingAPI 的代理类，确保始终返回同一个实例"""
|
||
|
||
def __new__(cls, *args, **kwargs):
|
||
global _embedding_instance
|
||
if _embedding_instance is None:
|
||
with _embedding_lock:
|
||
if _embedding_instance is None:
|
||
_embedding_instance = _OpenAIEmbeddingAPI(*args, **kwargs)
|
||
return _embedding_instance
|