RagFlow_Plugs/test_infinity_encapsulation.py

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

"""
测试Infinity向量数据库封装
"""

from services.utils.infinity import InfinityVectorDB
from conf.config import VectorDBConfig

def test_infinity_connection():
    """
    测试Infinity连接
    """
    print("=== 测试Infinity连接 ===")
    
    try:
        # 初始化InfinityVectorDB
        infinity_db = InfinityVectorDB()
        print("✅ InfinityVectorDB初始化成功")
        
        # 测试索引创建
        index_name = "test_collection"
        print(f"\n测试创建索引: {index_name}")
        result = infinity_db.create_index(index_name)
        if result:
            print(f"✅ 索引 {index_name} 创建成功")
        else:
            print(f"❌ 索引 {index_name} 创建失败")
            return False
        
        # 测试索引存在检查
        print(f"\n测试检查索引存在: {index_name}")
        exists = infinity_db.index_exists(index_name)
        if exists:
            print(f"✅ 索引 {index_name} 存在")
        else:
            print(f"❌ 索引 {index_name} 不存在")
            return False
        
        # 测试插入文档
        print(f"\n测试插入文档")
        document = {
            "file_name": "test.pdf",
            "file_page_count": 10,
            "page_number": 1,
            "text": "这是一个测试文档",
            "image_path": "test.png",
            "sparse_vector": [],
            "dense_vector_1024": [0.1] * 1024,
            "dataset_id": "test_dataset",
            "document_id": "test_doc_id"
        }
        
        insert_result = infinity_db.insert_document(index_name, document)
        if insert_result:
            print(f"✅ 文档插入成功")
        else:
            print(f"❌ 文档插入失败")
            return False
        
        # 测试批量插入
        print(f"\n测试批量插入文档")
        documents = []
        for i in range(2, 5):
            doc = {
                "file_name": "test.pdf",
                "file_page_count": 10,
                "page_number": i,
                "text": f"这是第 {i} 页",
                "image_path": f"test_{i}.png",
                "sparse_vector": [],
                "dense_vector_1024": [0.1] * 1024,
                "dataset_id": "test_dataset",
                "document_id": "test_doc_id"
            }
            documents.append(doc)
        
        bulk_result = infinity_db.bulk_insert(index_name, documents)
        if bulk_result["success"] == len(documents):
            print(f"✅ 批量插入成功，共插入 {bulk_result['success']} 个文档")
        else:
            print(f"❌ 批量插入失败，成功 {bulk_result['success']} 个，失败 {bulk_result['failed']} 个")
            return False
        
        # 测试向量检索
        print(f"\n测试向量检索")
        vector = [0.1] * 1024
        search_result = infinity_db.vector_search(index_name, "dense_vector_1024", vector, size=5)
        if search_result["hits"]["total"] > 0:
            print(f"✅ 向量检索成功，找到 {search_result['hits']['total']} 个结果")
        else:
            print(f"❌ 向量检索失败，未找到结果")
        
        # 测试混合检索
        print(f"\n测试混合检索")
        hybrid_result = infinity_db.hybrid_search(
            index_name,
            text_query="测试",
            vector_field="dense_vector_1024",
            vector=vector,
            size=5
        )
        if hybrid_result["hits"]["total"] > 0:
            print(f"✅ 混合检索成功，找到 {hybrid_result['hits']['total']} 个结果")
        else:
            print(f"❌ 混合检索失败，未找到结果")
        
        # 测试删除索引
        print(f"\n测试删除索引: {index_name}")
        delete_result = infinity_db.delete_index(index_name)
        if delete_result:
            print(f"✅ 索引 {index_name} 删除成功")
        else:
            print(f"❌ 索引 {index_name} 删除失败")
            return False
        
        # 关闭连接
        infinity_db.close()
        print(f"\n✅ 成功关闭连接")
        
        return True
        
    except Exception as e:
        print(f"\n❌ 测试失败: {e}")
        import traceback
        traceback.print_exc()
        return False


def test_vector_db_factory():
    """
    测试VectorDBFactory
    """
    print("\n=== 测试VectorDBFactory ===")
    
    try:
        from services.utils.vector_db import VectorDBFactory
        
        # 获取向量数据库实例
        vector_db = VectorDBFactory.get_vector_db()
        print(f"✅ 成功获取向量数据库实例: {type(vector_db).__name__}")
        
        # 测试创建索引
        index_name = "test_factory_collection"
        result = vector_db.create_index(index_name)
        if result:
            print(f"✅ 通过工厂创建索引 {index_name} 成功")
        else:
            print(f"❌ 通过工厂创建索引 {index_name} 失败")
        
        vector_db.close()
        print(f"✅ 成功关闭通过工厂获取的连接")
        
        return True
        
    except Exception as e:
        print(f"\n❌ 工厂测试失败: {e}")
        import traceback
        traceback.print_exc()
        return False


if __name__ == "__main__":
    print("开始测试Infinity向量数据库封装...")
    
    # 测试Infinity连接
    connection_result = test_infinity_connection()
    
    # 测试VectorDBFactory
    factory_result = test_vector_db_factory()
    
    # 总结
    print("\n=== 测试总结 ===")
    if connection_result and factory_result:
        print("✅ 所有测试通过！")
        exit(0)
    else:
        print("❌ 部分测试失败！")
        exit(1)