b17786b57b
主要功能: - 图片上传时 OCR 内容去重（3个上传端点统一使用公共函数 _check_ocr_duplicate） - 图片管理 Tab：展示所有图片、手动删除、一键去重 - 搜索结果详情弹窗增加删除按钮（带确认弹窗） - 图片管理卡片点击查看详情（复用 showOcrDetailModal） - 搜索限制和 LLM 批量判断数量可通过网站设置 - MiniMax API 调用添加 reasoning_split=True - 企业微信机器人：WebSocket 长连接、图片搜索、配置化搜索数量 - 版本号升级至 1.2.0
329 lines
11 KiB
Python
329 lines
11 KiB
Python
"""
|
||
嵌入服务模块
|
||
支持多种嵌入模型提供商：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
|