learningcircuit--local-deep-research/examples/elasticsearch/search_example.py

"""
使用 Elasticsearch 搜索引擎的示例脚本。(Example script for using Elasticsearch search engine.)
展示如何索引文档和搜索数据。(Demonstrates how to index documents and search data.)
"""

import sys
from pathlib import Path

from loguru import logger

# 添加项目根目录到 Python 路径 (Add project root directory to Python path)
sys.path.append(str(Path(__file__).parent.parent))

# Import after adding project root to path
from src.local_deep_research.utilities.es_utils import (
    ElasticsearchManager,
)
from src.local_deep_research.web_search_engines.engines.search_engine_elasticsearch import (
    ElasticsearchSearchEngine,
)

# 配置日志 (Configure logging)
# Loguru automatically handles logging configuration


def index_sample_documents():
    """索引示例文档到 Elasticsearch。(Index sample documents to Elasticsearch.)"""

    # 创建 Elasticsearch 管理器 (Create Elasticsearch manager)
    es_manager = ElasticsearchManager(
        hosts=["http://localhost:9200"],
        # 如果需要可以提供认证信息 (Authentication credentials can be provided if needed)
        # username="elastic",
        # password="password",
    )

    # 创建索引 (Create index)
    index_name = "documents"
    es_manager.create_index(index_name)

    # 准备示例文档 (Prepare sample documents)
    documents = [
        {
            "title": "Elasticsearch 简介",
            "content": "Elasticsearch 是一个分布式、开源的搜索和分析引擎，适用于所有类型的数据。",
            "tags": ["搜索引擎", "数据库", "全文搜索"],
            "category": "技术",
        },
        {
            "title": "Python 编程基础",
            "content": "Python 是一种解释型、高级、通用型编程语言。Python 的设计强调代码的可读性，使用缩进表示代码块。",
            "tags": ["编程语言", "脚本语言", "开发"],
            "category": "编程",
        },
        {
            "title": "自然语言处理介绍",
            "content": "自然语言处理(NLP)是人工智能的一个子领域，专注于计算机与人类语言之间的交互。",
            "tags": ["NLP", "AI", "机器学习"],
            "category": "人工智能",
        },
        {
            "title": "深度学习基础知识",
            "content": "深度学习是机器学习的一个分支，它使用多层神经网络来模拟人脑的学习过程。",
            "tags": ["深度学习", "神经网络", "AI"],
            "category": "人工智能",
        },
        {
            "title": "向量数据库比较",
            "content": "向量数据库是专门为存储和检索高维向量而设计的数据库。常见的向量数据库包括Elasticsearch、Pinecone、Milvus等。",
            "tags": ["向量数据库", "embeddings", "相似性搜索"],
            "category": "数据库",
        },
    ]

    # 批量索引文档 (Bulk index documents)
    success_count = es_manager.bulk_index_documents(
        index_name=index_name,
        documents=documents,
        refresh=True,  # 立即刷新索引使文档可搜索 (Immediately refresh index to make documents searchable)
    )

    logger.info(
        f"成功索引了 {success_count} 个文档到 '{index_name}' 索引"
    )  # Successfully indexed {success_count} documents to '{index_name}' index
    return index_name


def search_documents(index_name, query):
    """使用 Elasticsearch 搜索引擎搜索文档。(Search documents using Elasticsearch search engine.)"""

    # 创建 Elasticsearch 搜索引擎 (Create Elasticsearch search engine)
    search_engine = ElasticsearchSearchEngine(
        hosts=["http://localhost:9200"],
        index_name=index_name,
        max_results=10,
        # 如果需要可以提供认证信息 (Authentication credentials can be provided if needed)
        # username="elastic",
        # password="password",
    )

    # 执行搜索 (Execute search)
    logger.info(f"搜索查询: '{query}'")  # Search query: '{query}'
    results = search_engine.run(query)

    # 显示搜索结果 (Display search results)
    logger.info(f"找到 {len(results)} 个结果:")  # Found {len(results)} results:
    for i, result in enumerate(results, 1):
        print(f"\n结果 {i}:")  # Result {i}:
        print(
            f"标题: {result.get('title', '无标题')}"
        )  # Title: {result.get('title', 'No title')}
        print(
            f"片段: {result.get('snippet', '无摘要')[:100]}..."
        )  # Snippet: {result.get('snippet', 'No summary')[:100]}...
        if "score" in result:
            print(
                f"相关性分数: {result.get('score')}"
            )  # Relevance score: {result.get('score')}
        print("-" * 50)

    return results


def advanced_search_examples(index_name):
    """展示高级搜索功能的示例。(Demonstrate examples of advanced search features.)"""

    # 创建 Elasticsearch 搜索引擎 (Create Elasticsearch search engine)
    search_engine = ElasticsearchSearchEngine(
        hosts=["http://localhost:9200"],
        index_name=index_name,
    )

    # 1. 使用查询字符串语法 (1. Using query string syntax)
    print("\n=== 使用查询字符串语法 ===")  # === Using query string syntax ===
    query_string = "content:深度学习 OR title:elasticsearch"
    print(f"查询字符串: '{query_string}'")  # Query string: '{query_string}'
    results = search_engine.search_by_query_string(query_string)
    print(f"找到 {len(results)} 个结果")  # Found {len(results)} results

    # 2. 使用 DSL 查询 (2. Using DSL query)
    print("\n=== 使用 DSL 查询 ===")  # === Using DSL query ===
    query_dsl = {
        "query": {
            "bool": {
                "must": {"match": {"content": "人工智能"}},
                "filter": {"term": {"category.keyword": "人工智能"}},
            }
        }
    }
    print(f"DSL 查询: {query_dsl}")  # DSL query: {query_dsl}
    results = search_engine.search_by_dsl(query_dsl)
    print(f"找到 {len(results)} 个结果")  # Found {len(results)} results


def main():
    """主函数，运行示例。(Main function, run examples.)"""
    try:
        # 索引示例文档 (Index sample documents)
        index_name = index_sample_documents()

        # 执行基本搜索 (Execute basic searches)
        search_documents(index_name, "elasticsearch")
        search_documents(index_name, "深度学习")

        # 展示高级搜索功能 (Demonstrate advanced search features)
        advanced_search_examples(index_name)

    except Exception:
        logger.exception("运行示例时出错")  # Error running example
        logger.info(
            "请确保 Elasticsearch 正在运行，默认地址为 http://localhost:9200"
        )  # Make sure Elasticsearch is running, default address is http://localhost:9200


if __name__ == "__main__":
    main()