rwandanmtgorilla--kumi/services/embedding_service.py

"""
Embedding服务 - 处理多供应商embedding逻辑
"""
from typing import List, Optional, Callable
from config.settings import settings
from config.embedding_config import EmbeddingConfig
from vector_db.embedding_client import _OpenAIEmbeddingAPI
import logging

logger = logging.getLogger(__name__)


class EmbeddingService:
    """Embedding服务,支持多供应商配置"""

    def __init__(self):
        """初始化Embedding服务"""
        self.config: Optional[EmbeddingConfig] = None
        self._clients_cache = {}  # 缓存客户端实例

    def _ensure_config(self):
        """确保配置已加载"""
        if self.config is None:
            self.config = settings.get_embedding_config()

    def get_client(self, provider_name: str = None, model_name: str = None) -> _OpenAIEmbeddingAPI:
        """
        获取embedding客户端

        Args:
            provider_name: 供应商名称,None表示使用默认
            model_name: 模型名称,None表示使用默认

        Returns:
            _OpenAIEmbeddingAPI实例
        """
        self._ensure_config()

        # 使用默认值
        if provider_name is None or model_name is None:
            provider_name, model_name = self.config.get_default_model()

        # 检查缓存
        cache_key = f"{provider_name},{model_name}"
        if cache_key in self._clients_cache:
            return self._clients_cache[cache_key]

        # 获取模型信息
        model_info = self.config.get_model_info(provider_name, model_name)
        if not model_info:
            raise ValueError(f"未找到模型配置: provider={provider_name}, model={model_name}")

        # 创建客户端
        client = _OpenAIEmbeddingAPI(
            base_url=model_info['api_base_url'],
            token=model_info['api_key'],
            model=model_info['model']
        )

        # 缓存
        self._clients_cache[cache_key] = client
        logger.info(f"创建embedding客户端: {cache_key}")

        return client

    def encode_texts(
            self,
            texts: List[str],
            provider_name: str = None,
            model_name: str = None,
            batch_size: int = 20
    ) -> List[List[float]]:
        """
        批量编码文本为向量

        Args:
            texts: 文本列表
            provider_name: 供应商名称
            model_name: 模型名称
            batch_size: 批处理大小

        Returns:
            向量列表
        """
        client = self.get_client(provider_name, model_name)
        client.set_batch_size(batch_size)
        return client.encode_texts(texts)

    def encode_texts_with_progress(
            self,
            texts: List[str],
            progress_callback: Optional[Callable[[int, int, str], None]],
            provider_name: str = None,
            model_name: str = None,
            batch_size: int = 20
    ) -> List[List[float]]:
        """
        批量编码文本为向量(带进度回调)

        Args:
            texts: 文本列表
            progress_callback: 进度回调函数
            provider_name: 供应商名称
            model_name: 模型名称
            batch_size: 批处理大小

        Returns:
            向量列表
        """
        client = self.get_client(provider_name, model_name)
        client.set_batch_size(batch_size)
        return client.encode_texts_with_progress(texts, progress_callback)

    def encode_texts_with_progress_concurrent(
            self,
            texts: List[str],
            progress_callback: Optional[Callable[[int, int, str], None]],
            provider_name: str = None,
            model_name: str = None,
            batch_size: int = 20
    ) -> List[List[float]]:
        """
        并发安全的批量编码文本为向量(带进度回调)

        与 encode_texts_with_progress 不同，此方法每次创建独立的客户端实例，
        避免在并发场景下共享状态导致的问题。

        Args:
            texts: 文本列表
            progress_callback: 进度回调函数
            provider_name: 供应商名称
            model_name: 模型名称
            batch_size: 批处理大小

        Returns:
            向量列表
        """
        self._ensure_config()

        # 使用默认值
        if provider_name is None or model_name is None:
            provider_name, model_name = self.config.get_default_model()

        # 获取模型信息
        model_info = self.config.get_model_info(provider_name, model_name)
        if not model_info:
            raise ValueError(f"未找到模型配置: provider={provider_name}, model={model_name}")

        # 为并发调用创建独立的客户端实例（不使用缓存）
        client = _OpenAIEmbeddingAPI(
            base_url=model_info['api_base_url'],
            token=model_info['api_key'],
            model=model_info['model'],
            max_batch_size=batch_size
        )

        return client.encode_texts_with_progress(texts, progress_callback)

    def test_connection(
            self,
            provider_name: str = None,
            model_name: str = None
    ) -> dict:
        """
        测试embedding API连接

        Args:
            provider_name: 供应商名称
            model_name: 模型名称

        Returns:
            {"success": bool, "message": str, "dimension": int}
        """
        try:
            logger.info(f"开始测试embedding连接: provider={provider_name}, model={model_name}")
            client = self.get_client(provider_name, model_name)
            result = client.test_connection()

            if result["success"]:
                logger.info(f"Embedding连接测试成功: provider={provider_name}, model={model_name}, dimension={result['dimension']}")
            else:
                logger.warning(f"Embedding连接测试失败: provider={provider_name}, model={model_name}, reason={result['message']}")

            return result
        except Exception as e:
            logger.error(f"测试连接时发生异常: provider={provider_name}, model={model_name}, error={e}", exc_info=True)
            return {
                "success": False,
                "message": f"测试连接失败: {str(e)}",
                "dimension": None
            }

    def get_available_models(self) -> List[dict]:
        """
        获取所有可用模型

        Returns:
            [
                {
                    "provider": "openrouter",
                    "model": "google/gemini-embedding-001",
                    "display_name": "openrouter,google/gemini-embedding-001"
                },
                ...
            ]
        """
        self._ensure_config()
        return self.config.get_all_models()

    def get_default_model(self) -> dict:
        """
        获取默认模型信息

        Returns:
            {
                "provider": "openrouter",
                "model": "google/gemini-embedding-001",
                "display_name": "openrouter,google/gemini-embedding-001"
            }
        """
        self._ensure_config()
        provider, model = self.config.get_default_model()
        return {
            "provider": provider,
            "model": model,
            "display_name": f"{provider},{model}"
        }

    def parse_model_identifier(self, model_identifier: str) -> tuple:
        """
        解析模型标识符

        Args:
            model_identifier: "provider,model" 格式

        Returns:
            (provider_name, model_name)
        """
        self._ensure_config()
        return self.config.parse_model_identifier(model_identifier)

    def generate_model_abbreviation(self, model_name: str) -> str:
        """
        生成模型缩写用于collection名称

        Args:
            model_name: 模型名称

        Returns:
            缩写字符串,如 "gem-1", "Qwe-6"
        """
        self._ensure_config()
        return EmbeddingConfig.generate_model_abbreviation(model_name)