RagFlow_Plugs/test/test_vector_db.py

#!/usr/bin/env python3
# -*- coding: utf-8 -*-

"""
向量数据库测试脚本
测试向量数据库工厂类的基本功能和配置切换
"""

from services.utils.vector_db import VectorDBFactory
from conf.config import VectorDBConfig

def test_vector_db_factory():
    """
    测试向量数据库工厂类
    """
    print("=== 测试向量数据库工厂类 ===")
    
    # 获取配置的向量数据库类型
    vector_db_type = VectorDBConfig.get_vector_db_type()
    print(f"当前配置的向量数据库类型: {vector_db_type}")
    
    try:
        # 获取向量数据库实例
        vector_db = VectorDBFactory.get_vector_db()
        print(f"成功获取向量数据库实例: {type(vector_db).__name__}")
        
        # 测试创建索引
        index_name = "test_index"
        print(f"\n测试创建索引: {index_name}")
        result = vector_db.create_index(index_name)
        print(f"创建索引结果: {result}")
        
        # 测试向量检索接口
        print(f"\n测试向量检索接口")
        vector = [0.1] * 768
        result = vector_db.vector_search(index_name, "vector_768_vec", vector, size=5)
        print(f"向量检索结果: {result}")
        
        # 测试混合检索接口
        print(f"\n测试混合检索接口")
        result = vector_db.hybrid_search(
            index_name, 
            text_query="测试", 
            vector_field="vector_768_vec", 
            vector=vector, 
            size=5
        )
        print(f"混合检索结果: {result}")
        
        # 关闭连接
        vector_db.close()
        print(f"\n成功关闭向量数据库连接")
        
        return True
        
    except Exception as e:
        print(f"测试失败: {e}")
        import traceback
        traceback.print_exc()
        return False

def test_vector_db_switch():
    """
    测试向量数据库切换功能
    """
    print("\n=== 测试向量数据库切换功能 ===")
    
    # 测试不同类型的向量数据库
    test_types = ["es", "infinity"]
    
    for db_type in test_types:
        print(f"\n测试向量数据库类型: {db_type}")
        try:
            # 临时修改配置（实际使用时通过环境变量配置）
            from conf.config import VectorDBConfig
            
            # 注意：这里我们不能直接修改配置类的静态方法返回值
            # 所以我们通过工厂类的实现来测试
            
            # 这里只测试工厂类是否能正确创建不同类型的向量数据库
            if db_type == "es":
                from services.utils.vector_db import ElasticsearchVectorDB
                vector_db = ElasticsearchVectorDB()
            else:
                from services.utils.vector_db import InfinityVectorDB
                vector_db = InfinityVectorDB()
            
            print(f"成功创建{db_type}向量数据库实例: {type(vector_db).__name__}")
            vector_db.close()
            print(f"成功关闭{db_type}向量数据库连接")
            
        except Exception as e:
            print(f"测试{db_type}失败: {e}")
            import traceback
            traceback.print_exc()

if __name__ == "__main__":
    # 测试向量数据库工厂类
    test_vector_db_factory()
    
    # 测试向量数据库切换功能
    test_vector_db_switch()
    
    print("\n=== 测试完成 ===")