citu_vanna/embedding_function.py

import requests
import time
import numpy as np
from typing import List, Callable

class EmbeddingFunction:
    def __init__(self, model_name: str, api_key: str, base_url: str, embedding_dimension: int):
        self.model_name = model_name
        self.api_key = api_key
        self.base_url = base_url
        self.embedding_dimension = embedding_dimension
        self.headers = {
            "Authorization": f"Bearer {api_key}",
            "Content-Type": "application/json"
        }
        self.max_retries = 2  # 设置默认的最大重试次数
        self.retry_interval = 2  # 设置默认的重试间隔秒数
        self.normalize_embeddings = True # 设置默认是否归一化

    def _normalize_vector(self, vector: List[float]) -> List[float]:
        """
        对向量进行L2归一化
        Args:
            vector: 输入向量   
        Returns:
            List[float]: 归一化后的向量
        """

        if not vector:
            return []
        norm = np.linalg.norm(vector)
        if norm == 0:
            return vector
        return (np.array(vector) / norm).tolist()
    

    def __call__(self, input) -> List[List[float]]:
        """
        为文本列表生成嵌入向量
        
        Args:
            input: 要嵌入的文本或文本列表
            
        Returns:
            List[List[float]]: 嵌入向量列表
        """
        if not isinstance(input, list):
            input = [input]
            
        embeddings = []
        for text in input:
            # 直接调用generate_embedding，让它处理异常
            try:
                vector = self.generate_embedding(text)
                embeddings.append(vector)
            except Exception as e:
                print(f"为文本 '{text}' 生成embedding失败: {e}")
                # 重新抛出异常，不返回零向量
                raise e
                
        return embeddings
    
    def embed_documents(self, texts: List[str]) -> List[List[float]]:
        """
        为文档列表生成嵌入向量 (LangChain 接口)
        
        Args:
            texts: 要嵌入的文档列表
            
        Returns:
            List[List[float]]: 嵌入向量列表
        """
        return self.__call__(texts)
    
    def embed_query(self, text: str) -> List[float]:
        """
        为查询文本生成嵌入向量 (LangChain 接口)
        
        Args:
            text: 要嵌入的查询文本
            
        Returns:
            List[float]: 嵌入向量
        """
        return self.generate_embedding(text)
    
    def generate_embedding(self, text: str) -> List[float]:
        """
        为单个文本生成嵌入向量
        
        Args:
            text (str): 要嵌入的文本
            
        Returns:
            List[float]: 嵌入向量
        """
        # 处理空文本
        if not text or len(text.strip()) == 0:
            # 空文本返回零向量是合理的行为
            if self.embedding_dimension is None:
                raise ValueError("Embedding dimension (self.embedding_dimension) 未被正确初始化。")
            return [0.0] * self.embedding_dimension
        
        # 准备请求体
        payload = {
            "model": self.model_name,
            "input": text,
            "encoding_format": "float"
        }
        
        # 添加重试机制
        retries = 0
        while retries <= self.max_retries:
            try:
                # 发送API请求
                url = self.base_url
                if not url.endswith("/embeddings"):
                    url = url.rstrip("/")  # 移除尾部斜杠，避免双斜杠
                    if not url.endswith("/v1/embeddings"):
                        url = f"{url}/embeddings"
                
                response = requests.post(
                    url, 
                    json=payload, 
                    headers=self.headers,
                    timeout=30  # 设置超时时间
                )
                
                # 检查响应状态
                if response.status_code != 200:
                    error_msg = f"API请求错误: {response.status_code}, {response.text}"
                    
                    # 根据错误码判断是否需要重试
                    if response.status_code in (429, 500, 502, 503, 504):
                        retries += 1
                        if retries <= self.max_retries:
                            wait_time = self.retry_interval * (2 ** (retries - 1))  # 指数退避
                            print(f"API请求失败，等待 {wait_time} 秒后重试 ({retries}/{self.max_retries})")
                            time.sleep(wait_time)
                            continue
                    
                    raise ValueError(error_msg)
                
                # 解析响应
                result = response.json()
                
                # 提取embedding向量
                if "data" in result and len(result["data"]) > 0 and "embedding" in result["data"][0]:
                    vector = result["data"][0]["embedding"]
                    
                    # 如果是首次调用且未提供维度，则自动设置
                    if self.embedding_dimension is None:
                        self.embedding_dimension = len(vector)
                    else:
                        # 验证向量维度
                        actual_dim = len(vector)
                        if actual_dim != self.embedding_dimension:
                            print(f"警告: 向量维度不匹配: 期望 {self.embedding_dimension}, 实际 {actual_dim}")
                    
                    # 如果需要归一化
                    if self.normalize_embeddings:
                        vector = self._normalize_vector(vector)
                    
                    return vector
                else:
                    error_msg = f"API返回格式异常: {result}"
                    raise ValueError(error_msg)
                
            except Exception as e:
                retries += 1
                
                if retries <= self.max_retries:
                    wait_time = self.retry_interval * (2 ** (retries - 1))  # 指数退避
                    print(f"生成embedding时出错: {str(e)}, 等待 {wait_time} 秒后重试 ({retries}/{self.max_retries})")
                    time.sleep(wait_time)
                else:
                    # 抛出异常而不是返回零向量，确保问题不被掩盖
                    raise RuntimeError(f"生成embedding失败，已重试{self.max_retries}次: {str(e)}")
        
        # 这里不应该到达，但为了完整性添加
        raise RuntimeError("生成embedding失败")

    def test_connection(self, test_text="测试文本") -> dict:
        """
        测试嵌入模型的连接和功能
        
        Args:
            test_text (str): 用于测试的文本
            
        Returns:
            dict: 包含测试结果的字典，包括是否成功、维度信息等
        """
        result = {
            "success": False,
            "model": self.model_name,
            "base_url": self.base_url,
            "message": "",
            "actual_dimension": None,
            "expected_dimension": self.embedding_dimension
        }
        
        try:
            print(f"测试嵌入模型连接 - 模型: {self.model_name}")
            print(f"API服务地址: {self.base_url}")
            
            # 验证配置
            if not self.api_key:
                result["message"] = "API密钥未设置或为空"
                return result
                
            if not self.base_url:
                result["message"] = "API服务地址未设置或为空"
                return result
                
            # 测试生成向量
            vector = self.generate_embedding(test_text)
            actual_dimension = len(vector)
            
            result["success"] = True
            result["actual_dimension"] = actual_dimension
            
            # 检查维度是否一致
            if actual_dimension != self.embedding_dimension:
                result["message"] = f"警告: 模型实际生成的向量维度({actual_dimension})与配置维度({self.embedding_dimension})不一致"
            else:
                result["message"] = f"连接测试成功，向量维度: {actual_dimension}"
                
            return result
            
        except Exception as e:
            result["message"] = f"连接测试失败: {str(e)}"
            return result

def test_embedding_connection() -> dict:
    """
    测试嵌入模型连接和配置是否正确
    
    Returns:
        dict: 测试结果，包括成功/失败状态、错误消息等
    """
    try:
        # 获取嵌入函数实例
        embedding_function = get_embedding_function()
        
        # 测试连接
        test_result = embedding_function.test_connection()
        
        if test_result["success"]:
            print(f"嵌入模型连接测试成功!")
            if "警告" in test_result["message"]:
                print(test_result["message"])
                print(f"建议将app_config.py中的EMBEDDING_CONFIG['embedding_dimension']修改为{test_result['actual_dimension']}")
        else:
            print(f"嵌入模型连接测试失败: {test_result['message']}")
            
        return test_result
        
    except Exception as e:
        error_message = f"无法测试嵌入模型连接: {str(e)}"
        print(error_message)
        return {
            "success": False,
            "message": error_message
        }

def get_embedding_function():
    """
    根据当前配置创建合适的EmbeddingFunction实例
    支持API和Ollama两种提供商
    
    Returns:
        EmbeddingFunction或OllamaEmbeddingFunction: 根据配置类型返回相应的实例
        
    Raises:
        ImportError: 如果无法导入必要的模块
        ValueError: 如果配置无效
    """
    try:
        from common.utils import get_current_embedding_config, is_using_ollama_embedding
    except ImportError:
        raise ImportError("无法导入 common.utils，请确保该文件存在")
    
    # 获取当前embedding配置
    embedding_config = get_current_embedding_config()
    
    if is_using_ollama_embedding():
        # 使用Ollama Embedding
        try:
            from customembedding.ollama_embedding import OllamaEmbeddingFunction
        except ImportError:
            raise ImportError("无法导入 OllamaEmbeddingFunction，请确保 customembedding 包存在")
            
        return OllamaEmbeddingFunction(
            model_name=embedding_config["model_name"],
            base_url=embedding_config["base_url"],
            embedding_dimension=embedding_config["embedding_dimension"]
        )
    else:
        # 使用API Embedding
        try:
            api_key = embedding_config["api_key"]
            model_name = embedding_config["model_name"]
            base_url = embedding_config["base_url"]
            embedding_dimension = embedding_config["embedding_dimension"]
            
            if api_key is None:
                raise KeyError("API模式下 'api_key' 未设置 (可能环境变量 EMBEDDING_API_KEY 未定义)")
                
        except KeyError as e:
            raise KeyError(f"API Embedding配置中缺少必要的键或值无效：{e}")

        return EmbeddingFunction(
            model_name=model_name,
            api_key=api_key,
            base_url=base_url,
            embedding_dimension=embedding_dimension
        )