230 lines
7.5 KiB
Python
230 lines
7.5 KiB
Python
|
"""
|
|||
|
|
Embedding 生成与语义检索服务。
|
|||
|
|
|
|||
|
|
使用 OpenAI text-embedding-3-small 生成文本向量，
|
|||
|
|
在内存中计算余弦相似度实现语义搜索。
|
|||
|
|
|
|||
|
|
如未配置 OpenAI API Key，所有方法静默降级返回空结果。
|
|||
|
|
"""
|
|||
|
|
from __future__ import annotations
|
|||
|
|
|
|||
|
|
import asyncio
|
|||
|
|
import json
|
|||
|
|
import logging
|
|||
|
|
import time
|
|||
|
|
from typing import Any, Dict, List, Optional, TypedDict
|
|||
|
|
|
|||
|
|
from app.core.config import settings
|
|||
|
|
|
|||
|
|
logger = logging.getLogger(__name__)
|
|||
|
|
|
|||
|
|
# 默认 embedding 模型
|
|||
|
|
EMBEDDING_MODEL = "text-embedding-3-small"
|
|||
|
|
EMBEDDING_DIMENSIONS = 1536
|
|||
|
|
|
|||
|
|
|
|||
|
|
class VectorEntry(TypedDict, total=False):
|
|||
|
|
"""向量条目"""
|
|||
|
|
id: str
|
|||
|
|
scope_kind: str
|
|||
|
|
scope_id: str
|
|||
|
|
content_text: str
|
|||
|
|
embedding: List[float]
|
|||
|
|
metadata: Dict[str, Any]
|
|||
|
|
score: float # 余弦相似度，仅检索结果包含
|
|||
|
|
|
|||
|
|
|
|||
|
|
class EmbeddingService:
|
|||
|
|
"""
|
|||
|
|
Embedding 服务。
|
|||
|
|
|
|||
|
|
用法：
|
|||
|
|
svc = EmbeddingService()
|
|||
|
|
emb = await svc.generate_embedding("你好")
|
|||
|
|
results = await svc.similarity_search(query_emb, entries, top_k=5)
|
|||
|
|
"""
|
|||
|
|
|
|||
|
|
def __init__(self):
|
|||
|
|
self._client: Optional[Any] = None
|
|||
|
|
self._client_lock = asyncio.Lock()
|
|||
|
|
|
|||
|
|
async def _get_client(self):
|
|||
|
|
"""延迟初始化 OpenAI 客户端（仅在首次调用时创建）。
|
|||
|
|
|
|||
|
|
优先级：SiliconFlow > OpenAI > DeepSeek。
|
|||
|
|
SiliconFlow 的 netease-youdao/bce-embedding-base_v1 在国内可直连且免费。
|
|||
|
|
"""
|
|||
|
|
if self._client is not None:
|
|||
|
|
return self._client
|
|||
|
|
|
|||
|
|
async with self._client_lock:
|
|||
|
|
if self._client is not None:
|
|||
|
|
return self._client
|
|||
|
|
|
|||
|
|
from openai import AsyncOpenAI
|
|||
|
|
|
|||
|
|
api_key: Optional[str] = None
|
|||
|
|
base_url: Optional[str] = None
|
|||
|
|
backend_label = "none"
|
|||
|
|
|
|||
|
|
# 1) SiliconFlow（国内直连，推荐）
|
|||
|
|
if settings.SILICONFLOW_API_KEY:
|
|||
|
|
api_key = settings.SILICONFLOW_API_KEY
|
|||
|
|
base_url = settings.SILICONFLOW_BASE_URL or "https://api.siliconflow.cn/v1"
|
|||
|
|
backend_label = "siliconflow"
|
|||
|
|
logger.info("Embedding 后端: SiliconFlow (model=%s)", settings.SILICONFLOW_EMBEDDING_MODEL)
|
|||
|
|
|
|||
|
|
# 2) OpenAI
|
|||
|
|
if not api_key:
|
|||
|
|
if settings.OPENAI_API_KEY and settings.OPENAI_API_KEY != "your-openai-api-key":
|
|||
|
|
api_key = settings.OPENAI_API_KEY
|
|||
|
|
base_url = settings.OPENAI_BASE_URL or "https://api.openai.com/v1"
|
|||
|
|
backend_label = "openai"
|
|||
|
|
|
|||
|
|
# 3) DeepSeek（部分代理可能支持 embedding）
|
|||
|
|
if not api_key:
|
|||
|
|
if settings.DEEPSEEK_API_KEY:
|
|||
|
|
api_key = settings.DEEPSEEK_API_KEY
|
|||
|
|
base_url = settings.DEEPSEEK_BASE_URL or "https://api.deepseek.com"
|
|||
|
|
backend_label = "deepseek"
|
|||
|
|
|
|||
|
|
if not api_key:
|
|||
|
|
logger.info("未配置任何 API Key，向量记忆功能已禁用（请配置 SILICONFLOW_API_KEY 或 OPENAI_API_KEY）")
|
|||
|
|
self._client = None
|
|||
|
|
return None
|
|||
|
|
|
|||
|
|
self._client = AsyncOpenAI(
|
|||
|
|
api_key=api_key,
|
|||
|
|
base_url=base_url,
|
|||
|
|
)
|
|||
|
|
logger.info("EmbeddingService 已初始化 (backend=%s)", backend_label)
|
|||
|
|
return self._client
|
|||
|
|
|
|||
|
|
def _get_model(self) -> str:
|
|||
|
|
"""根据后端选择对应的 embedding 模型。"""
|
|||
|
|
if settings.SILICONFLOW_API_KEY:
|
|||
|
|
return settings.SILICONFLOW_EMBEDDING_MODEL
|
|||
|
|
return EMBEDDING_MODEL
|
|||
|
|
|
|||
|
|
async def generate_embedding(self, text: str) -> Optional[List[float]]:
|
|||
|
|
"""
|
|||
|
|
为单段文本生成 embedding 向量。
|
|||
|
|
|
|||
|
|
返回 float 列表，无 API key 或出错时返回 None。
|
|||
|
|
"""
|
|||
|
|
if not text or not text.strip():
|
|||
|
|
return None
|
|||
|
|
|
|||
|
|
client = await self._get_client()
|
|||
|
|
if not client:
|
|||
|
|
return None
|
|||
|
|
|
|||
|
|
model = self._get_model()
|
|||
|
|
try:
|
|||
|
|
kwargs: Dict[str, Any] = {
|
|||
|
|
"model": model,
|
|||
|
|
"input": text.strip()[:8000],
|
|||
|
|
}
|
|||
|
|
# OpenAI text-embedding-3-small 支持指定 dimensions；其它模型可能不支持
|
|||
|
|
if model == EMBEDDING_MODEL:
|
|||
|
|
kwargs["dimensions"] = EMBEDDING_DIMENSIONS
|
|||
|
|
resp = await client.embeddings.create(**kwargs)
|
|||
|
|
return resp.data[0].embedding
|
|||
|
|
except Exception as e:
|
|||
|
|
logger.warning("生成 embedding 失败: %s", e)
|
|||
|
|
return None
|
|||
|
|
|
|||
|
|
async def generate_embeddings(self, texts: List[str]) -> List[Optional[List[float]]]:
|
|||
|
|
"""
|
|||
|
|
批量生成 embedding 向量。
|
|||
|
|
"""
|
|||
|
|
if not texts:
|
|||
|
|
return []
|
|||
|
|
|
|||
|
|
client = await self._get_client()
|
|||
|
|
if not client:
|
|||
|
|
return [None] * len(texts)
|
|||
|
|
|
|||
|
|
# 清理空文本
|
|||
|
|
valid_indices = [i for i, t in enumerate(texts) if t and t.strip()]
|
|||
|
|
if not valid_indices:
|
|||
|
|
return [None] * len(texts)
|
|||
|
|
|
|||
|
|
model = self._get_model()
|
|||
|
|
try:
|
|||
|
|
inputs = [texts[i].strip()[:8000] for i in valid_indices]
|
|||
|
|
kwargs: Dict[str, Any] = {
|
|||
|
|
"model": model,
|
|||
|
|
"input": inputs,
|
|||
|
|
}
|
|||
|
|
if model == EMBEDDING_MODEL:
|
|||
|
|
kwargs["dimensions"] = EMBEDDING_DIMENSIONS
|
|||
|
|
resp = await client.embeddings.create(**kwargs)
|
|||
|
|
embeddings = [r.embedding for r in resp.data]
|
|||
|
|
except Exception as e:
|
|||
|
|
logger.warning("批量生成 embedding 失败: %s", e)
|
|||
|
|
return [None] * len(texts)
|
|||
|
|
|
|||
|
|
# 按原始顺序排列结果
|
|||
|
|
result: List[Optional[List[float]]] = [None] * len(texts)
|
|||
|
|
for idx, emb in zip(valid_indices, embeddings):
|
|||
|
|
result[idx] = emb
|
|||
|
|
return result
|
|||
|
|
|
|||
|
|
@staticmethod
|
|||
|
|
def cosine_similarity(a: List[float], b: List[float]) -> float:
|
|||
|
|
"""计算两个向量的余弦相似度。"""
|
|||
|
|
if not a or not b or len(a) != len(b):
|
|||
|
|
return 0.0
|
|||
|
|
|
|||
|
|
dot = sum(x * y for x, y in zip(a, b))
|
|||
|
|
norm_a = sum(x * x for x in a) ** 0.5
|
|||
|
|
norm_b = sum(y * y for y in b) ** 0.5
|
|||
|
|
|
|||
|
|
if norm_a == 0 or norm_b == 0:
|
|||
|
|
return 0.0
|
|||
|
|
return dot / (norm_a * norm_b)
|
|||
|
|
|
|||
|
|
async def similarity_search(
|
|||
|
|
self,
|
|||
|
|
query_embedding: List[float],
|
|||
|
|
entries: List[VectorEntry],
|
|||
|
|
top_k: int = 5,
|
|||
|
|
min_score: float = 0.3,
|
|||
|
|
) -> List[VectorEntry]:
|
|||
|
|
"""
|
|||
|
|
在内存中对 entries 做余弦相似度搜索，返回 Top-K 结果（按得分降序）。
|
|||
|
|
min_score: 最低相似度阈值，低于该值的结果被过滤。
|
|||
|
|
"""
|
|||
|
|
scored: List[VectorEntry] = []
|
|||
|
|
for entry in entries:
|
|||
|
|
emb = entry.get("embedding")
|
|||
|
|
if not emb:
|
|||
|
|
continue
|
|||
|
|
score = self.cosine_similarity(query_embedding, emb)
|
|||
|
|
if score >= min_score:
|
|||
|
|
entry["score"] = score
|
|||
|
|
scored.append(entry)
|
|||
|
|
|
|||
|
|
scored.sort(key=lambda x: x["score"], reverse=True)
|
|||
|
|
return scored[:top_k]
|
|||
|
|
|
|||
|
|
@staticmethod
|
|||
|
|
def serialize_embedding(embedding: List[float]) -> str:
|
|||
|
|
"""将 embedding 序列化为 JSON 字符串。"""
|
|||
|
|
return json.dumps(embedding, ensure_ascii=False)
|
|||
|
|
|
|||
|
|
@staticmethod
|
|||
|
|
def deserialize_embedding(data: str) -> List[float]:
|
|||
|
|
"""从 JSON 字符串反序列化 embedding。"""
|
|||
|
|
if isinstance(data, list):
|
|||
|
|
return data
|
|||
|
|
try:
|
|||
|
|
return json.loads(data)
|
|||
|
|
except (json.JSONDecodeError, TypeError):
|
|||
|
|
return []
|
|||
|
|
|
|||
|
|
|
|||
|
|
# 全局单例
|
|||
|
|
embedding_service = EmbeddingService()
|