feat: v1.2.0 - 图片去重与管理、微信机器人优化、搜索设置可配置
主要功能: - 图片上传时 OCR 内容去重（3个上传端点统一使用公共函数 _check_ocr_duplicate） - 图片管理 Tab：展示所有图片、手动删除、一键去重 - 搜索结果详情弹窗增加删除按钮（带确认弹窗） - 图片管理卡片点击查看详情（复用 showOcrDetailModal） - 搜索限制和 LLM 批量判断数量可通过网站设置 - MiniMax API 调用添加 reasoning_split=True - 企业微信机器人：WebSocket 长连接、图片搜索、配置化搜索数量 - 版本号升级至 1.2.0
This commit is contained in:
EduBrain Dev
328
app/services/embedding_service.py
Normal file
328
app/services/embedding_service.py
Normal file
@@ -0,0 +1,328 @@
|
||||
"""
|
||||
嵌入服务模块
|
||||
支持多种嵌入模型提供商：OpenAI / 智谱 / 阿里云 DashScope / 本地 BGE
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import abc
|
||||
import logging
|
||||
from typing import List, Optional
|
||||
|
||||
import numpy as np
|
||||
|
||||
from app.config import EmbeddingProvider, settings
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
# ──────────────────────────── 抽象基类 ────────────────────────────
|
||||
|
||||
class EmbeddingProviderBase(abc.ABC):
|
||||
"""嵌入模型提供商抽象基类"""
|
||||
|
||||
@abc.abstractmethod
|
||||
async def embed_batch(self, texts: List[str]) -> List[List[float]]:
|
||||
"""
|
||||
批量生成文本嵌入向量。
|
||||
|
||||
Args:
|
||||
texts: 待嵌入的文本列表
|
||||
|
||||
Returns:
|
||||
嵌入向量列表，每个向量是 float 的列表
|
||||
|
||||
Raises:
|
||||
Exception: 嵌入生成失败时抛出异常
|
||||
"""
|
||||
...
|
||||
|
||||
@abc.abstractmethod
|
||||
async def embed_single(self, text: str) -> List[float]:
|
||||
"""生成单条文本的嵌入向量"""
|
||||
...
|
||||
|
||||
@property
|
||||
@abc.abstractmethod
|
||||
def dimension(self) -> int:
|
||||
"""嵌入向量维度"""
|
||||
...
|
||||
|
||||
|
||||
# ──────────────────────────── OpenAI 实现 ────────────────────────────
|
||||
|
||||
class OpenAIEmbedding(EmbeddingProviderBase):
|
||||
"""OpenAI 嵌入模型（兼容 OpenAI API 格式的服务）"""
|
||||
|
||||
def __init__(self) -> None:
|
||||
from openai import AsyncOpenAI
|
||||
|
||||
self._client = AsyncOpenAI(
|
||||
api_key=settings.OPENAI_API_KEY or "EMPTY",
|
||||
base_url=settings.OPENAI_BASE_URL,
|
||||
)
|
||||
self._model = settings.EMBEDDING_MODEL or "text-embedding-3-small"
|
||||
|
||||
@property
|
||||
def dimension(self) -> int:
|
||||
return settings.EMBEDDING_DIMENSIONS
|
||||
|
||||
async def embed_batch(self, texts: List[str]) -> List[List[float]]:
|
||||
"""调用 OpenAI 兼容接口批量生成嵌入"""
|
||||
# OpenAI 接口单次最多 2048 条，自动分批
|
||||
batch_size = 2048
|
||||
all_embeddings: List[List[float]] = []
|
||||
|
||||
for i in range(0, len(texts), batch_size):
|
||||
batch = texts[i : i + batch_size]
|
||||
response = await self._client.embeddings.create(
|
||||
input=batch,
|
||||
model=self._model,
|
||||
dimensions=self.dimension if self.dimension <= 3072 else None,
|
||||
)
|
||||
batch_embeddings = [item.embedding for item in response.data]
|
||||
all_embeddings.extend(batch_embeddings)
|
||||
|
||||
return all_embeddings
|
||||
|
||||
async def embed_single(self, text: str) -> List[float]:
|
||||
results = await self.embed_batch([text])
|
||||
return results[0]
|
||||
|
||||
|
||||
# ──────────────────────────── 智谱 AI 实现 ────────────────────────────
|
||||
|
||||
class ZhipuEmbedding(EmbeddingProviderBase):
|
||||
"""智谱 AI 嵌入模型"""
|
||||
|
||||
def __init__(self) -> None:
|
||||
from zhipuai import ZhipuAI
|
||||
|
||||
self._client = ZhipuAI(api_key=settings.ZHIPU_API_KEY)
|
||||
self._model = settings.EMBEDDING_MODEL or "embedding-3"
|
||||
|
||||
@property
|
||||
def dimension(self) -> int:
|
||||
return settings.EMBEDDING_DIMENSIONS
|
||||
|
||||
async def embed_batch(self, texts: List[str]) -> List[List[float]]:
|
||||
"""调用智谱 AI 接口批量生成嵌入"""
|
||||
import asyncio
|
||||
|
||||
# 智谱 API 单次最多 16 条
|
||||
batch_size = 16
|
||||
all_embeddings: List[List[float]] = []
|
||||
|
||||
for i in range(0, len(texts), batch_size):
|
||||
batch = texts[i : i + batch_size]
|
||||
# 智谱 SDK 是同步的，用线程池包装
|
||||
loop = asyncio.get_event_loop()
|
||||
response = await loop.run_in_executor(
|
||||
None,
|
||||
lambda: self._client.embeddings.create(
|
||||
input=batch,
|
||||
model=self._model,
|
||||
),
|
||||
)
|
||||
batch_embeddings = [item.embedding for item in response.data]
|
||||
all_embeddings.extend(batch_embeddings)
|
||||
|
||||
return all_embeddings
|
||||
|
||||
async def embed_single(self, text: str) -> List[float]:
|
||||
results = await self.embed_batch([text])
|
||||
return results[0]
|
||||
|
||||
|
||||
# ──────────────────────────── 阿里云 DashScope 实现 ────────────────────────────
|
||||
|
||||
class DashscopeEmbedding(EmbeddingProviderBase):
|
||||
"""阿里云 DashScope 嵌入模型"""
|
||||
|
||||
def __init__(self) -> None:
|
||||
import dashscope
|
||||
dashscope.api_key = settings.DASHSCOPE_API_KEY
|
||||
self._model = settings.EMBEDDING_MODEL or "text-embedding-v3"
|
||||
|
||||
@property
|
||||
def dimension(self) -> int:
|
||||
return settings.EMBEDDING_DIMENSIONS
|
||||
|
||||
async def embed_batch(self, texts: List[str]) -> List[List[float]]:
|
||||
"""调用 DashScope 接口批量生成嵌入"""
|
||||
import asyncio
|
||||
from dashscope import TextEmbedding
|
||||
|
||||
batch_size = 25 # DashScope 单次最多 25 条
|
||||
all_embeddings: List[List[float]] = []
|
||||
|
||||
for i in range(0, len(texts), batch_size):
|
||||
batch = texts[i : i + batch_size]
|
||||
loop = asyncio.get_event_loop()
|
||||
|
||||
def _call() -> list:
|
||||
resp = TextEmbedding.call(
|
||||
model=self._model,
|
||||
input=batch,
|
||||
dimension=self.dimension,
|
||||
)
|
||||
if resp.status_code != 200:
|
||||
raise RuntimeError(f"DashScope 嵌入失败: {resp.code} - {resp.message}")
|
||||
return [item["embedding"] for item in resp.output["embeddings"]]
|
||||
|
||||
batch_embeddings = await loop.run_in_executor(None, _call)
|
||||
all_embeddings.extend(batch_embeddings)
|
||||
|
||||
return all_embeddings
|
||||
|
||||
async def embed_single(self, text: str) -> List[float]:
|
||||
results = await self.embed_batch([text])
|
||||
return results[0]
|
||||
|
||||
|
||||
# ──────────────────────────── 本地 BGE 实现 ────────────────────────────
|
||||
|
||||
class LocalBGEEmbedding(EmbeddingProviderBase):
|
||||
"""
|
||||
本地 BGE 嵌入模型
|
||||
使用 sentence-transformers 加载模型，在 CPU/GPU 上本地推理
|
||||
"""
|
||||
|
||||
def __init__(self) -> None:
|
||||
self._model = None
|
||||
self._dim = settings.EMBEDDING_DIMENSIONS
|
||||
|
||||
def _load_model(self):
|
||||
"""延迟加载模型（首次调用时初始化）"""
|
||||
if self._model is not None:
|
||||
return
|
||||
|
||||
logger.info("正在加载本地 BGE 嵌入模型: %s", settings.EMBEDDING_MODEL)
|
||||
from sentence_transformers import SentenceTransformer
|
||||
|
||||
model_path = settings.LOCAL_BGE_MODEL_PATH or settings.EMBEDDING_MODEL
|
||||
self._model = SentenceTransformer(
|
||||
model_path,
|
||||
trust_remote_code=True,
|
||||
)
|
||||
# 获取实际维度
|
||||
test_embedding = self._model.encode(["测试"])
|
||||
self._dim = len(test_embedding[0])
|
||||
logger.info("BGE 模型加载完成，维度: %d", self._dim)
|
||||
|
||||
@property
|
||||
def dimension(self) -> int:
|
||||
return self._dim
|
||||
|
||||
async def embed_batch(self, texts: List[str]) -> List[List[float]]:
|
||||
"""使用本地模型批量生成嵌入"""
|
||||
import asyncio
|
||||
|
||||
self._load_model()
|
||||
loop = asyncio.get_event_loop()
|
||||
|
||||
def _encode() -> List[List[float]]:
|
||||
embeddings = self._model.encode(texts, normalize_embeddings=True)
|
||||
return embeddings.tolist()
|
||||
|
||||
return await loop.run_in_executor(None, _encode)
|
||||
|
||||
async def embed_single(self, text: str) -> List[float]:
|
||||
results = await self.embed_batch([text])
|
||||
return results[0]
|
||||
|
||||
|
||||
# ──────────────────────────── MiniMax 实现 ────────────────────────────
|
||||
|
||||
class MiniMaxEmbedding(EmbeddingProviderBase):
|
||||
"""
|
||||
MiniMax 嵌入模型
|
||||
通过 MiniMax 的 OpenAI 兼容接口调用嵌入 API
|
||||
"""
|
||||
|
||||
def __init__(self) -> None:
|
||||
from openai import AsyncOpenAI
|
||||
|
||||
self._client = AsyncOpenAI(
|
||||
api_key=settings.MINIMAX_API_KEY or "EMPTY",
|
||||
base_url=settings.MINIMAX_BASE_URL,
|
||||
)
|
||||
self._model = settings.MINIMAX_EMBEDDING_MODEL
|
||||
|
||||
@property
|
||||
def dimension(self) -> int:
|
||||
return settings.EMBEDDING_DIMENSIONS
|
||||
|
||||
async def embed_batch(self, texts: List[str]) -> List[List[float]]:
|
||||
"""调用 MiniMax 兼容接口批量生成嵌入"""
|
||||
batch_size = 64
|
||||
all_embeddings: List[List[float]] = []
|
||||
|
||||
for i in range(0, len(texts), batch_size):
|
||||
batch = texts[i : i + batch_size]
|
||||
response = await self._client.embeddings.create(
|
||||
input=batch,
|
||||
model=self._model,
|
||||
)
|
||||
batch_embeddings = [item.embedding for item in response.data]
|
||||
all_embeddings.extend(batch_embeddings)
|
||||
|
||||
return all_embeddings
|
||||
|
||||
async def embed_single(self, text: str) -> List[float]:
|
||||
results = await self.embed_batch([text])
|
||||
return results[0]
|
||||
|
||||
|
||||
# ──────────────────────────── 工厂类 ────────────────────────────
|
||||
|
||||
class EmbeddingService:
|
||||
"""
|
||||
嵌入服务工厂类
|
||||
根据配置自动选择嵌入模型提供商，提供全局单例访问
|
||||
"""
|
||||
|
||||
_instance: Optional[EmbeddingProviderBase] = None
|
||||
|
||||
@classmethod
|
||||
def get_instance(cls) -> EmbeddingProviderBase:
|
||||
"""获取嵌入服务单例实例"""
|
||||
if cls._instance is not None:
|
||||
return cls._instance
|
||||
|
||||
provider_map = {
|
||||
EmbeddingProvider.OPENAI: OpenAIEmbedding,
|
||||
EmbeddingProvider.ZHIPU: ZhipuEmbedding,
|
||||
EmbeddingProvider.DASHSCOPE: DashscopeEmbedding,
|
||||
EmbeddingProvider.LOCAL_BGE: LocalBGEEmbedding,
|
||||
EmbeddingProvider.MINIMAX: MiniMaxEmbedding,
|
||||
}
|
||||
|
||||
provider_cls = provider_map.get(settings.EMBEDDING_PROVIDER)
|
||||
if provider_cls is None:
|
||||
raise ValueError(f"不支持的嵌入模型提供商: {settings.EMBEDDING_PROVIDER}")
|
||||
|
||||
cls._instance = provider_cls()
|
||||
logger.info(
|
||||
"嵌入服务初始化完成: provider=%s, model=%s",
|
||||
settings.EMBEDDING_PROVIDER.value,
|
||||
settings.EMBEDDING_MODEL,
|
||||
)
|
||||
return cls._instance
|
||||
|
||||
@classmethod
|
||||
async def embed_batch(cls, texts: List[str]) -> List[List[float]]:
|
||||
"""便捷方法：批量生成嵌入"""
|
||||
instance = cls.get_instance()
|
||||
return await instance.embed_batch(texts)
|
||||
|
||||
@classmethod
|
||||
async def embed_single(cls, text: str) -> List[float]:
|
||||
"""便捷方法：生成单条嵌入"""
|
||||
instance = cls.get_instance()
|
||||
return await instance.embed_single(text)
|
||||
|
||||
@classmethod
|
||||
def reset(cls) -> None:
|
||||
"""重置单例（用于测试或切换配置）"""
|
||||
cls._instance = None
|
||||
Reference in New Issue
Block a user