app/utils/embedding.py

# -*- coding: utf-8 -*-
"""
Embedding 服务 - 并发优化版

特性:
1. 并发批量请求 - 使用 asyncio.gather 并行处理多个批次
2. 信号量控制 - 限制最大并发数，避免 API 限流
3. 重试机制 - 使用 tenacity 处理临时性错误
4. 智能分批 - 根据 token 数量动态调整批次大小
"""

import asyncio
import logging
from typing import List, Optional
from dataclasses import dataclass

from openai import AsyncOpenAI

from app.core.config import settings
from app.utils.retry import llm_retry, is_retryable_error

logger = logging.getLogger(__name__)


@dataclass
class EmbeddingConfig:
    """Embedding 服务配置"""
    # API 配置
    api_base_url: str = "https://api.siliconflow.cn/v1"
    model_name: str = "BAAI/bge-m3"
    
    # 批处理配置
    batch_size: int = 50              # 每批文本数量
    max_text_length: int = 8000       # 单个文本最大字符数
    
    # 并发控制
    max_concurrent_batches: int = 5   # 最大并发批次数
    
    # 超时配置
    timeout: int = 60                 # 单次请求超时 (秒)


class EmbeddingService:
    """
    高性能 Embedding 服务
    
    使用示例:
    ```python
    service = EmbeddingService()
    
    # 单文本
    embedding = await service.embed_text("Hello world")
    
    # 批量文本 (自动并发优化)
    texts = ["text1", "text2", ..., "text100"]
    embeddings = await service.embed_batch(texts)
    ```
    """
    
    def __init__(self, config: Optional[EmbeddingConfig] = None):
        self.config = config or EmbeddingConfig()
        
        # 初始化 OpenAI 客户端 (SiliconFlow 兼容 OpenAI 协议)
        self._client = AsyncOpenAI(
            api_key=settings.SILICON_API_KEY,
            base_url=self.config.api_base_url,
            timeout=self.config.timeout
        )
        
        # 并发信号量
        self._semaphore = asyncio.Semaphore(self.config.max_concurrent_batches)
        
        # 统计信息
        self._stats = {
            "total_requests": 0,
            "successful_requests": 0,
            "failed_requests": 0,
            "total_texts": 0,
            "retried_requests": 0
        }
    
    def _preprocess_text(self, text: str) -> str:
        """预处理文本: 移除换行、截断长度"""
        text = text.replace("\n", " ").strip()
        if len(text) > self.config.max_text_length:
            text = text[:self.config.max_text_length]
        return text
    
    @llm_retry
    async def _embed_single_batch(self, texts: List[str]) -> List[List[float]]:
        """
        处理单个批次的 Embedding 请求 (带重试)
        
        Args:
            texts: 预处理后的文本列表
            
        Returns:
            embedding 向量列表
        """
        self._stats["total_requests"] += 1
        
        response = await self._client.embeddings.create(
            input=texts,
            model=self.config.model_name
        )
        
        self._stats["successful_requests"] += 1
        return [item.embedding for item in response.data]
    
    async def _embed_batch_with_semaphore(
        self, 
        batch_texts: List[str], 
        batch_index: int
    ) -> tuple[int, List[List[float]]]:
        """
        带信号量控制的批次处理
        
        Returns:
            (batch_index, embeddings) - 返回索引用于结果排序
        """
        async with self._semaphore:
            try:
                embeddings = await self._embed_single_batch(batch_texts)
                logger.debug(f"✅ 批次 {batch_index} 完成: {len(batch_texts)} 文本")
                return (batch_index, embeddings)
            except Exception as e:
                self._stats["failed_requests"] += 1
                logger.error(f"❌ 批次 {batch_index} 失败: {type(e).__name__}: {e}")
                raise
    
    async def embed_text(self, text: str) -> List[float]:
        """
        获取单个文本的 Embedding
        
        Args:
            text: 输入文本
            
        Returns:
            embedding 向量，失败返回空列表
        """
        try:
            processed = self._preprocess_text(text)
            if not processed:
                return []
            
            self._stats["total_texts"] += 1
            embeddings = await self._embed_single_batch([processed])
            return embeddings[0] if embeddings else []
        except Exception as e:
            logger.error(f"embed_text 失败: {e}")
            return []
    
    async def embed_batch(
        self, 
        texts: List[str],
        show_progress: bool = False
    ) -> List[List[float]]:
        """
        批量获取 Embedding (并发优化)
        
        Args:
            texts: 文本列表
            show_progress: 是否显示进度日志
            
        Returns:
            embedding 向量列表 (与输入顺序一致)
            失败的文本对应空列表
        """
        if not texts:
            return []
        
        # 预处理所有文本
        processed_texts = [self._preprocess_text(t) for t in texts]
        self._stats["total_texts"] += len(texts)
        
        # 分批
        batch_size = self.config.batch_size
        batches = [
            processed_texts[i:i + batch_size] 
            for i in range(0, len(processed_texts), batch_size)
        ]
        
        total_batches = len(batches)
        if show_progress:
            logger.info(
                f"📊 Embedding: {len(texts)} 文本 → {total_batches} 批次 "
                f"(并发: {self.config.max_concurrent_batches})"
            )
        
        # 并发执行所有批次
        tasks = [
            self._embed_batch_with_semaphore(batch, idx)
            for idx, batch in enumerate(batches)
        ]
        
        # 收集结果
        results = await asyncio.gather(*tasks, return_exceptions=True)
        
        # 按批次索引排序并合并结果
        embeddings = []
        for result in sorted(results, key=lambda x: x[0] if isinstance(x, tuple) else float('inf')):
            if isinstance(result, tuple):
                batch_idx, batch_embeddings = result
                embeddings.extend(batch_embeddings)
            else:
                # 异常情况: 填充空向量
                # 找出这个批次有多少文本
                failed_batch_size = batch_size  # 保守估计
                embeddings.extend([[] for _ in range(failed_batch_size)])
                logger.warning(f"批次失败，填充 {failed_batch_size} 个空向量")
        
        # 确保返回数量与输入一致
        if len(embeddings) < len(texts):
            embeddings.extend([[] for _ in range(len(texts) - len(embeddings))])
        elif len(embeddings) > len(texts):
            embeddings = embeddings[:len(texts)]
        
        if show_progress:
            success_count = sum(1 for e in embeddings if e)
            logger.info(f"✅ Embedding 完成: {success_count}/{len(texts)} 成功")
        
        return embeddings
    
    def get_stats(self) -> dict:
        """获取统计信息"""
        return self._stats.copy()
    
    def reset_stats(self):
        """重置统计信息"""
        for key in self._stats:
            self._stats[key] = 0


# 全局单例
_embedding_service: Optional[EmbeddingService] = None


def get_embedding_service(config: Optional[EmbeddingConfig] = None) -> EmbeddingService:
    """获取 Embedding 服务单例"""
    global _embedding_service
    if _embedding_service is None:
        _embedding_service = EmbeddingService(config)
    return _embedding_service


# 便捷函数
async def embed_text(text: str) -> List[float]:
    """快捷方式: 获取单个文本的 Embedding"""
    return await get_embedding_service().embed_text(text)


async def embed_batch(texts: List[str], show_progress: bool = False) -> List[List[float]]:
    """快捷方式: 批量获取 Embedding"""
    return await get_embedding_service().embed_batch(texts, show_progress)