LQAgentPlatform/test/milvus入库脚本.py

#!/usr/bin/env python3
"""
测试修复后的 Milvus 向量实现
"""

import sys
import os

# 添加项目根目录到路径
sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))

print("测试修复后的 Milvus 向量实现")
print("=" * 50)

def test_basic_functionality():
    """测试基本功能"""
    try:
        # 导入并初始化 MilvusVectorManager
        from foundation.database.base.vector.milvus_vector import MilvusVectorManager
        print("成功导入 MilvusVectorManager")

        # 初始化管理器
        manager = MilvusVectorManager()
        print("MilvusVectorManager 初始化成功")

        # 测试 text_to_vector 方法
        test_text = "桥梁建设技术"
        vector = manager.text_to_vector(test_text)
        print(f"text_to_vector 测试成功，向量维度: {len(vector)}")

        # 简单测试文档
        test_documents = [
            {
                'content': '四川路桥建设集团专注于桥梁和隧道工程建设',
                'metadata': {'category': 'company', 'type': 'construction'}
            },
            {
                'content': '高速公路桥梁建设技术包括预应力混凝土和钢结构',
                'metadata': {'category': 'technology', 'type': 'highway'}
            }
        ]

        collection_name = "test_fix_validation"

        print(f"\n测试 create_hybrid_collection 方法...")
        vectorstore = manager.create_hybrid_collection(
            collection_name=collection_name,
            documents=test_documents
        )
        print("create_hybrid_collection 执行成功!")
        print(f"返回的 vectorstore 类型: {type(vectorstore)}")

        # 等待索引创建完成
        import time
        time.sleep(5)

        print(f"\n测试 hybrid_search 方法...")
        param = {'collection_name': collection_name}

        # 测试加权搜索
        results = manager.hybrid_search(
            param=param,
            query_text="桥梁建设",
            top_k=2,
            ranker_type="weighted",
            dense_weight=0.7,
            sparse_weight=0.3
        )
        print(f"Hybrid search 执行成功，返回 {len(results)} 个结果")

        for i, result in enumerate(results):
            content = result.get('text_content', '')[:50]
            print(f"  {i+1}. {content}...")

        # 清理测试集合
        print(f"\n清理测试集合...")
        try:
            from pymilvus import utility
            if utility.has_collection(collection_name):
                utility.drop_collection(collection_name)
                print(f"成功清理集合: {collection_name}")
        except Exception as e:
            print(f"清理集合失败: {e}")

        return True

    except Exception as e:
        print(f"测试失败: {e}")
        import traceback
        traceback.print_exc()
        return False

if __name__ == "__main__":
    success = test_basic_functionality()
    print("\n" + "=" * 50)
    print(f"测试结果: {'成功' if success else '失败'}")

    if success:
        print("修复验证成功！")
        print("- text_to_vector 方法正常工作")
        print("- create_hybrid_collection 方法正常工作")
        print("- hybrid_search 方法正常工作")