RagFlow_Plugs/test/test_es_conn.py

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

import json
from services.utils.es_conn import ESConnection

def test_es_connection():
    """
    测试 Elasticsearch 连接和基本功能
    """
    try:
        # 初始化连接
        print("正在初始化 Elasticsearch 连接...")
        es = ESConnection(hosts=["http://localhost:9200"])
        print("连接成功！")
        
        # 测试索引创建
        index_name = "test_ragflow_index"
        print(f"\n正在创建索引: {index_name}")
        success = es.create_index(index_name)
        if success:
            print(f"索引 {index_name} 创建成功！")
        else:
            print(f"索引 {index_name} 创建失败！")
            return False
        
        # 测试文档插入
        test_doc = {
            "title": "测试文档",
            "content": "这是一个用于测试 Elasticsearch 连接的文档",
            "content_tks": "这 是 一个 用于 测试 Elasticsearch 连接 的 文档",
            "vector_768_vec": [0.1] * 768,
            "created_at": "2024-01-01 00:00:00",
            "count_int": 10,
            "importance_flt": 0.8,
            "tags_kwd": ["测试", "elasticsearch"],
            "kb_id": "test_kb_123"
        }
        
        print("\n正在插入测试文档...")
        insert_success = es.insert(index_name, test_doc)
        if insert_success:
            print("文档插入成功！")
        else:
            print("文档插入失败！")
            return False
        
        # 测试批量插入
        test_docs = []
        for i in range(3):
            doc = {
                "title": f"批量测试文档 {i}",
                "content": f"这是第 {i} 个批量测试文档",
                "content_tks": f"这是 第 {i} 个 批量 测试 文档",
                "vector_768_vec": [0.1] * 768,
                "created_at": "2024-01-01 00:00:00",
                "count_int": i,
                "importance_flt": 0.5 + i * 0.1,
                "tags_kwd": ["批量", "测试"],
                "kb_id": "test_kb_123"
            }
            test_docs.append(doc)
        
        print("\n正在批量插入测试文档...")
        bulk_result = es.bulk_insert(index_name, test_docs)
        print(f"批量插入结果: {bulk_result}")
        
        # 测试全文检索
        print("\n正在测试全文检索...")
        text_query = {
            "match": {
                "content": "测试"
            }
        }
        text_result = es.search(index_name, text_query, size=5)
        print(f"全文检索结果: {text_result['hits']['total']} 个命中")
        
        # 测试向量检索
        print("\n正在测试向量检索...")
        vector = [0.1] * 768
        vector_result = es.knn_search(
            index_name=index_name,
            vector_field="vector_768_vec",
            vector=vector,
            k=3
        )
        print(f"向量检索结果: {vector_result['hits']['total']} 个命中")
        
        # 测试混合检索
        print("\n正在测试混合检索...")
        hybrid_result = es.hybrid_search(
            index_name=index_name,
            text_query="测试",
            vector_field="vector_768_vec",
            vector=vector,
            size=5
        )
        print(f"混合检索结果: {hybrid_result['hits']['total']} 个命中")
        
        # 打印命中的文档
        print("\n混合检索命中的文档:")
        for hit in hybrid_result['hits']['hits']:
            doc = hit['_source']
            print(f"  - 标题: {doc['title']}, 相似度分数: {hit['_score']:.4f}")
        
        # 测试文档删除
        print(f"\n正在删除索引: {index_name}")
        es.es.indices.delete(index=index_name, ignore=[400, 404])
        print(f"索引 {index_name} 删除成功！")
        
        # 关闭连接
        es.close()
        print("\n所有测试完成！")
        return True
        
    except Exception as e:
        print(f"测试失败: {e}")
        return False

if __name__ == "__main__":
    test_es_connection()