chat-ai-api/app/utils/embedding_service.py

"""
向量化工具
使用豆包的 doubao-embedding-vision-251215 模型生成向量嵌入
"""
from typing import List
from openai import OpenAI
from ..core.ark_client import get_client


class EmbeddingService:
    def __init__(self):
        self.embedding_model = "doubao-embedding-vision-251215"

    def generate_embeddings(self, texts: List[str], app_name: str) -> List[List[float]]:
        """
        批量生成文本的向量嵌入
        :param texts: 文本列表
        :param app_name: 租户的appName（用于获取对应的API客户端）
        :return: 向量列表
        """
        if not texts:
            return []

        try:
            client = get_client(app_name)

            # 调用豆包的 Embeddings API
            response = client.embeddings.create(
                model=self.embedding_model,
                input=texts
            )

            # 提取向量
            embeddings = [item.embedding for item in response.data]
            return embeddings

        except Exception as e:
            print(f"向量化失败: {str(e)}")
            return []

    def generate_single_embedding(self, text: str, app_name: str) -> List[float]:
        """
        生成单个文本的向量嵌入
        :param text: 文本内容
        :param app_name: 租户的appName
        :return: 向量
        """
        embeddings = self.generate_embeddings([text], app_name)
        return embeddings[0] if embeddings else []


# 全局单例
_embedding_service = None


def get_embedding_service() -> EmbeddingService:
    """获取向量化服务单例"""
    global _embedding_service
    if _embedding_service is None:
        _embedding_service = EmbeddingService()
    return _embedding_service