269 lines
10 KiB
Python
269 lines
10 KiB
Python
"""
|
||
阿里云DashScope嵌入模型实现
|
||
直接调用阿里云文本嵌入API，不通过OneAPI兼容层
|
||
"""
|
||
|
||
import os
|
||
import sys
|
||
import json
|
||
import time
|
||
import requests
|
||
import numpy as np
|
||
from typing import List, Union, Optional
|
||
from dotenv import load_dotenv
|
||
|
||
# 添加项目根目录到路径
|
||
project_root = os.path.join(os.path.dirname(__file__), '..', '..')
|
||
sys.path.append(project_root)
|
||
|
||
from atlas_rag.vectorstore.embedding_model import BaseEmbeddingModel
|
||
|
||
# 加载环境变量
|
||
load_dotenv(os.path.join(os.path.dirname(__file__), '..', '..', '.env'))
|
||
|
||
|
||
class DashScopeEmbeddingModel(BaseEmbeddingModel):
|
||
"""
|
||
阿里云DashScope嵌入模型实现
|
||
直接调用DashScope文本嵌入API
|
||
"""
|
||
|
||
def __init__(
|
||
self,
|
||
api_key: Optional[str] = None,
|
||
model_name: Optional[str] = None,
|
||
max_retries: int = 3,
|
||
retry_delay: float = 1.0
|
||
):
|
||
"""
|
||
初始化DashScope嵌入模型
|
||
|
||
Args:
|
||
api_key: 阿里云DashScope API Key
|
||
model_name: 嵌入模型名称，如 text-embedding-v3
|
||
max_retries: 最大重试次数
|
||
retry_delay: 重试延迟时间（秒）
|
||
"""
|
||
# 初始化父类，传入None作为sentence_encoder（我们不使用它）
|
||
super().__init__(sentence_encoder=None)
|
||
|
||
self.api_key = api_key or os.getenv('ONEAPI_KEY') # 复用现有环境变量
|
||
self.model_name = model_name or os.getenv('ONEAPI_MODEL_EMBED', 'text-embedding-v3')
|
||
self.max_retries = max_retries
|
||
self.retry_delay = retry_delay
|
||
|
||
if not self.api_key:
|
||
raise ValueError("DashScope API Key未设置，请在.env文件中设置ONEAPI_KEY")
|
||
|
||
# DashScope文本嵌入API端点
|
||
self.api_url = "https://dashscope.aliyuncs.com/api/v1/services/embeddings/text-embedding/text-embedding"
|
||
|
||
# 设置请求头
|
||
self.headers = {
|
||
'Authorization': f'Bearer {self.api_key}',
|
||
'Content-Type': 'application/json'
|
||
}
|
||
|
||
print(f"[OK] DashScope嵌入模型初始化完成: {self.model_name}")
|
||
self._test_connection()
|
||
|
||
def encode(self, texts, **kwargs):
|
||
"""
|
||
实现基类的抽象方法
|
||
对文本进行编码，返回嵌入向量
|
||
|
||
Args:
|
||
texts: 文本或文本列表
|
||
**kwargs: 其他参数（batch_size, show_progress_bar, query_type等）
|
||
|
||
Returns:
|
||
嵌入向量或嵌入向量列表（numpy数组格式）
|
||
"""
|
||
import numpy as np
|
||
|
||
if isinstance(texts, str):
|
||
# 单个文本
|
||
result = self.embed_query(texts)
|
||
return np.array(result)
|
||
else:
|
||
# 文本列表
|
||
results = self.embed_texts(texts)
|
||
return np.array(results)
|
||
|
||
def _test_connection(self):
|
||
"""测试DashScope嵌入API连接"""
|
||
try:
|
||
# 发送一个简单的测试请求
|
||
test_payload = {
|
||
"model": self.model_name,
|
||
"input": {
|
||
"texts": ["test"]
|
||
}
|
||
}
|
||
|
||
response = requests.post(
|
||
self.api_url,
|
||
headers=self.headers,
|
||
json=test_payload,
|
||
timeout=(5, 15) # 短超时进行连接测试
|
||
)
|
||
|
||
if response.status_code == 200:
|
||
result = response.json()
|
||
if result.get('output') and result['output'].get('embeddings'):
|
||
print("[OK] DashScope嵌入API连接测试成功")
|
||
else:
|
||
print(f"[WARNING] DashScope嵌入API响应格式异常: {result}")
|
||
else:
|
||
print(f"[WARNING] DashScope嵌入API连接异常，状态码: {response.status_code}")
|
||
print(f" 响应: {response.text[:200]}")
|
||
|
||
except requests.exceptions.Timeout:
|
||
print("[WARNING] DashScope嵌入API连接超时，但将在实际请求时重试")
|
||
except requests.exceptions.ConnectionError:
|
||
print("[WARNING] DashScope嵌入API连接失败，请检查网络状态")
|
||
except Exception as e:
|
||
print(f"[WARNING] DashScope嵌入API连接测试出错: {e}")
|
||
|
||
def embed_texts(self, texts: List[str]) -> List[List[float]]:
|
||
"""
|
||
对文本列表进行嵌入
|
||
|
||
Args:
|
||
texts: 文本列表
|
||
|
||
Returns:
|
||
嵌入向量列表
|
||
"""
|
||
import numpy as np
|
||
if not texts:
|
||
return []
|
||
|
||
# DashScope API支持批量处理，但需要分批处理大量文本
|
||
batch_size = 10 # 每批处理10个文本
|
||
all_embeddings = []
|
||
|
||
for i in range(0, len(texts), batch_size):
|
||
batch_texts = texts[i:i + batch_size]
|
||
batch_embeddings = self._embed_batch(batch_texts)
|
||
all_embeddings.extend(batch_embeddings)
|
||
|
||
return all_embeddings
|
||
|
||
def _embed_batch(self, texts: List[str]) -> List[List[float]]:
|
||
"""对一批文本进行嵌入"""
|
||
payload = {
|
||
"model": self.model_name,
|
||
"input": {
|
||
"texts": texts
|
||
}
|
||
}
|
||
|
||
for attempt in range(self.max_retries):
|
||
try:
|
||
response = requests.post(
|
||
self.api_url,
|
||
headers=self.headers,
|
||
json=payload,
|
||
timeout=(30, 120) # 连接30秒，读取120秒
|
||
)
|
||
|
||
if response.status_code == 200:
|
||
result = response.json()
|
||
|
||
# 检查是否有错误
|
||
if result.get('code'):
|
||
error_msg = result.get('message', f'API错误代码: {result["code"]}')
|
||
if attempt == self.max_retries - 1:
|
||
raise RuntimeError(f"DashScope嵌入API错误: {error_msg}")
|
||
else:
|
||
print(f"API错误，正在重试 ({attempt + 2}/{self.max_retries}): {error_msg}")
|
||
time.sleep(self.retry_delay * (attempt + 1))
|
||
continue
|
||
|
||
# 提取嵌入向量
|
||
if result.get('output') and result['output'].get('embeddings'):
|
||
embeddings_data = result['output']['embeddings']
|
||
|
||
# 提取向量数据
|
||
embeddings = []
|
||
for embedding_item in embeddings_data:
|
||
if isinstance(embedding_item, dict) and 'embedding' in embedding_item:
|
||
embeddings.append(embedding_item['embedding'])
|
||
elif isinstance(embedding_item, list):
|
||
embeddings.append(embedding_item)
|
||
else:
|
||
print(f"[WARNING] 未知的嵌入格式: {type(embedding_item)}")
|
||
|
||
return embeddings
|
||
else:
|
||
error_msg = f"API响应格式错误: {result}"
|
||
if attempt == self.max_retries - 1:
|
||
raise RuntimeError(error_msg)
|
||
else:
|
||
print(f"响应格式错误，正在重试 ({attempt + 2}/{self.max_retries})")
|
||
time.sleep(self.retry_delay * (attempt + 1))
|
||
else:
|
||
error_text = response.text[:500] if response.text else "无响应内容"
|
||
error_msg = f"API请求失败，状态码: {response.status_code}, 响应: {error_text}"
|
||
if attempt == self.max_retries - 1:
|
||
raise RuntimeError(error_msg)
|
||
else:
|
||
print(f"请求失败，正在重试 ({attempt + 2}/{self.max_retries}): 状态码 {response.status_code}")
|
||
time.sleep(self.retry_delay * (attempt + 1))
|
||
|
||
except KeyboardInterrupt:
|
||
print(f"\n[WARNING] 用户中断请求")
|
||
raise KeyboardInterrupt("用户中断请求")
|
||
|
||
except requests.exceptions.Timeout as e:
|
||
error_msg = f"请求超时: {str(e)}"
|
||
if attempt == self.max_retries - 1:
|
||
raise RuntimeError(f"经过 {self.max_retries} 次重试后仍超时: {error_msg}")
|
||
else:
|
||
print(f"请求超时，正在重试 ({attempt + 2}/{self.max_retries}): {error_msg}")
|
||
time.sleep(self.retry_delay * (attempt + 1))
|
||
|
||
except requests.exceptions.ConnectionError as e:
|
||
error_msg = f"连接错误: {str(e)}"
|
||
if attempt == self.max_retries - 1:
|
||
raise RuntimeError(f"经过 {self.max_retries} 次重试后仍无法连接: {error_msg}")
|
||
else:
|
||
print(f"连接错误，正在重试 ({attempt + 2}/{self.max_retries}): {error_msg}")
|
||
time.sleep(self.retry_delay * (attempt + 1))
|
||
|
||
except requests.RequestException as e:
|
||
error_msg = f"网络请求异常: {str(e)}"
|
||
if attempt == self.max_retries - 1:
|
||
raise RuntimeError(f"经过 {self.max_retries} 次重试后仍失败: {error_msg}")
|
||
else:
|
||
print(f"网络异常，正在重试 ({attempt + 2}/{self.max_retries}): {str(e)[:100]}")
|
||
time.sleep(self.retry_delay * (attempt + 1))
|
||
|
||
raise RuntimeError("所有重试都失败了")
|
||
|
||
def embed_query(self, text: str) -> List[float]:
|
||
"""
|
||
对单个查询文本进行嵌入
|
||
|
||
Args:
|
||
text: 查询文本
|
||
|
||
Returns:
|
||
嵌入向量
|
||
"""
|
||
embeddings = self.embed_texts([text])
|
||
return embeddings[0] if embeddings else []
|
||
|
||
|
||
def create_dashscope_embedding_model(
|
||
api_key: Optional[str] = None,
|
||
model_name: Optional[str] = None,
|
||
**kwargs
|
||
) -> DashScopeEmbeddingModel:
|
||
"""创建DashScope嵌入模型实例的便捷函数"""
|
||
return DashScopeEmbeddingModel(
|
||
api_key=api_key,
|
||
model_name=model_name,
|
||
**kwargs
|
||
) |