luojiehua
/
BidiRag


			
							123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276
							# -*- coding: utf-8 -*-
"""
BidiRag 使用示例 - 演示如何使用 BidiRag 类进行 HTML 文档检索

功能：
1. 加载 HTML 文档
2. 使用不同的 RAG 方法进行检索
3. 支持关键词检索
4. 支持自然语言查询
"""

import os
from loguru import logger

from bdirag.bidi_rag import BidiRag


def demo_basic_usage():
    """基础使用示例"""
    print("=" * 80)
    print("示例 1: 基础使用 - BM25 HTML Tree 方法")
    print("=" * 80)
    
    # 1. 初始化 BidiRag（使用 BM25 HTML Tree 方法，适合 HTML 文档）
    rag = BidiRag(
        rag_method='bm25_html_tree',  # 使用 BM25 HTML Tree 方法
        chunk_size=512,
        chunk_overlap=50
    )
    
    # 2. 添加 HTML 文档
    html_files = [
        "examples/sample_data.html",  # 假设的 HTML 文件
        # "data/documents/announcement1.html",
        # "data/documents/announcement2.html",
    ]
    
    # 如果示例文件不存在，使用文本替代
    if not any(os.path.exists(f) for f in html_files):
        print("未找到 HTML 文件，使用示例文本...")
        sample_texts = [
            """
            <html>
            <body>
            <h1>招标公告</h1>
            <p>项目名称：大连长兴岛经济技术开发区交流岛街道办事处2026年5月至6月政府采购意向</p>
            <p>采购单位：大连长兴岛经济技术开发区交流岛街道办事处</p>
            <p>采购项目名称：交流岛滨海路夜间出行照明提升工程</p>
            <p>预算金额：147.060000万元(人民币)</p>
            <p>采购需求：在滨海路安装太阳能路灯200盏（单排）</p>
            <p>预计采购时间：2026-05</p>
            </body>
            </html>
            """,
            """
            <html>
            <body>
            <h1>中标公告</h1>
            <p>项目名称：XX市智慧交通系统建设项目</p>
            <p>采购人：XX市交通运输局</p>
            <p>中标人：XX科技有限公司</p>
            <p>中标金额：5000万元</p>
            <p>项目内容：交通信号控制系统、视频监控系统、交通流量监测系统</p>
            </body>
            </html>
            """
        ]
        rag.add_texts(sample_texts)
    else:
        rag.add_documents(html_files)
    
    print(f"\n已加载 {rag.get_document_count()} 个文档\n")
    
    # 3. 检索文档 - 使用关键词
    print("-" * 80)
    print("检索 1: 查找'招标人'和'中标人'相关信息")
    print("-" * 80)
    
    results = rag.retrieve(
        query="招标人和中标人",
        top_k=3,
        keywords=["招标人", "中标人", "采购人", "中标"]
    )
    
    for i, (doc, score) in enumerate(results, 1):
        print(f"\n结果 {i} (相似度: {score:.3f}):")
        print(doc.page_content[:200])
        print("...")
    
    print("\n" + "=" * 80)


def demo_keyword_search():
    """关键词搜索示例"""
    print("\n示例 2: 关键词搜索")
    print("=" * 80)
    
    # 初始化
    rag = BidiRag(rag_method='bm25')
    
    # 添加示例数据
    sample_docs = [
        "招标公告：本项目招标人为XX市财政局，项目预算100万元",
        "中标公告：中标人为XX建设有限公司，中标金额98万元",
        "采购公告：采购单位XX医院，采购医疗设备一批",
        "中标公告：本项目中标人为XX科技公司，中标价格50万元",
    ]
    rag.add_texts(sample_docs)
    
    # 搜索特定关键词
    print("\n搜索关键词：'招标人'")
    results = rag.search_keywords(["招标人"], top_k=5)
    for i, doc in enumerate(results, 1):
        print(f"{i}. {doc.page_content}")
    
    print("\n搜索关键词：'中标人'")
    results = rag.search_keywords(["中标人"], top_k=5)
    for i, doc in enumerate(results, 1):
        print(f"{i}. {doc.page_content}")
    
    print("\n搜索关键词：'招标人' AND '中标人'")
    results = rag.search_keywords(["招标人", "中标人"], top_k=5)
    for i, doc in enumerate(results, 1):
        print(f"{i}. {doc.page_content}")
    
    print("\n" + "=" * 80)


def demo_different_methods():
    """不同 RAG 方法对比"""
    print("\n示例 3: 不同 RAG 方法对比")
    print("=" * 80)
    
    # 准备相同的测试数据
    test_texts = [
        "招标公告：XX市政府采购项目，招标人：XX局，预算200万",
        "中标公告：中标人XX公司，中标金额180万，项目：信息化建设",
        "采购意向：XX学校设备采购，采购单位：XX学校，预计采购时间2026年",
    ]
    
    # 测试不同的方法
    methods = ['bm25', 'tfidf', 'keyword', 'naive']
    
    for method in methods:
        print(f"\n测试方法: {method}")
        print("-" * 40)
        
        try:
            rag = BidiRag(rag_method=method)
            rag.add_texts(test_texts)
            
            # 检索
            results = rag.retrieve(
                query="招标人信息",
                top_k=2,
                keywords=["招标人", "采购人"]
            )
            
            for i, (doc, score) in enumerate(results, 1):
                print(f"  {i}. (score: {score:.3f}) {doc.page_content[:100]}...")
            
        except Exception as e:
            print(f"  方法 {method} 失败: {e}")
    
    print("\n" + "=" * 80)


def demo_query_with_answer():
    """完整问答示例"""
    print("\n示例 4: 完整问答（需要 LLM）")
    print("=" * 80)
    
    # 注意：这个方法需要配置 LLM
    print("注意：此示例需要配置 LLM 客户端")
    print("如果您有 OpenAI API key，可以这样使用：\n")
    
    example_code = '''
from openai import OpenAI

# 初始化 LLM 客户端
llm_client = OpenAI(
    api_key="your-api-key",
    base_url="https://api.openai.com/v1"
)

# 使用 BidiRag
rag = BidiRag(
    rag_method='naive',
    llm_client=llm_client,
    llm_model='gpt-4o'
)

rag.add_texts([
    "招标公告：招标人XX局，项目预算100万",
    "中标公告：中标人XX公司，金额98万"
])

# 完整问答
result = rag.query(
    query="谁是招标人？",
    keywords=["招标人", "采购人"]
)

print("答案:", result.answer)
print("检索到的文档数:", len(result.retrieved_docs))
'''
    
    print(example_code)
    print("=" * 80)


def demo_html_structure_aware():
    """HTML 结构感知检索示例"""
    print("\n示例 5: HTML 结构感知检索（推荐用于招投标公告）")
    print("=" * 80)
    
    # BM25 HTML Tree 方法会保留 HTML 结构信息
    rag = BidiRag(rag_method='bm25_html_tree')
    
    html_content = """
    <html>
    <body>
        <div class="announcement">
            <h1>政府采购中标公告</h1>
            <table>
                <tr><td>项目名称</td><td>办公设备采购项目</td></tr>
                <tr><td>采购人</td><td>XX市财政局</td></tr>
                <tr><td>中标人</td><td>XX办公设备有限公司</td></tr>
                <tr><td>中标金额</td><td>50万元</td></tr>
                <tr><td>采购内容</td><td>电脑、打印机、复印机等办公设备</td></tr>
            </table>
        </div>
    </body>
    </html>
    """
    
    rag.add_texts([html_content])
    
    # 检索表格中的结构化信息
    results = rag.retrieve(
        query="中标人信息",
        top_k=3,
        keywords=["中标人", "采购人", "中标金额"]
    )
    
    print(f"\n找到 {len(results)} 个相关结果:")
    for i, (doc, score) in enumerate(results, 1):
        print(f"\n结果 {i} (score: {score:.3f}):")
        print(doc.page_content)
    
    print("\n" + "=" * 80)


def main():
    """运行所有示例"""
    print("\n" + "🚀 " * 20)
    print("BidiRag 使用示例")
    print(" " * 20 + "\n")
    
    # 运行示例
    demo_basic_usage()
    demo_keyword_search()
    demo_different_methods()
    demo_query_with_answer()
    demo_html_structure_aware()
    
    print("\n✅ 所有示例运行完成！")
    print("\n提示：")
    print("1. 对于招投标公告，推荐使用 'bm25_html_tree' 方法")
    print("2. 可以使用 keywords 参数精化检索结果")
    print("3. 如果需要生成答案，需要配置 LLM 客户端")
    print("4. 可用的 RAG 方法:", BidiRag(rag_method='bm25').list_available_methods())


if __name__ == "__main__":
    main()