main.py

"""
联网模块 - 入口 & 使用示例

运行方式:
  设置环境变量后:
    python main.py

  或直接在代码中传 config:
    from config import ModuleConfig, LLMConfig, SearchConfig
    from pipeline import run_pipeline
    
    config = ModuleConfig(
        llm=LLMConfig(api_key="sk-xxx"),
        search=SearchConfig(api_key="xxx"),
    )
    result = await run_pipeline("今天比特币多少钱?", config)
    print(result.display())
"""

from __future__ import annotations

import asyncio
import os

from config import ModuleConfig, LLMConfig, SearchConfig
from pipeline import run_pipeline


# ── 示例 1: 完整流程 ────────────────────────────────────────────

async def example_basic():
    """基础示例: 问一个需要联网的问题"""
    
    config = ModuleConfig(
        llm=LLMConfig(
            api_base=os.getenv("LLM_API_BASE", "https://api.openai.com/v1"),
            api_key=os.getenv("LLM_API_KEY", "sk-your-key-here"),
            model="gpt-4o-mini",
            summary_model="gpt-4o",
        ),
        search=SearchConfig(
            provider="serpapi",
            api_key=os.getenv("SEARCH_API_KEY", "your-serpapi-key"),
            max_results=5,
            fetch_content=True,
            fetch_timeout=10,
        ),
        verbose=True,
    )
    
    query = "2024年诺贝尔物理学奖得主是谁? 他们的主要贡献是什么?"
    
    print("=" * 60)
    print(f"用户问题: {query}")
    print("=" * 60)
    
    result = await run_pipeline(query, config)
    
    print("\n" + "=" * 60)
    print(result.display())
    print("=" * 60)
    print(f"\n总耗时: {result.elapsed_ms:.0f}ms")


# ── 示例 2: 不需要搜索的问题 ───────────────────────────────────

async def example_no_search():
    """测试联网判断: 常识问题应该跳过搜索"""
    
    config = ModuleConfig(
        llm=LLMConfig(
            api_base=os.getenv("LLM_API_BASE", "https://api.openai.com/v1"),
            api_key=os.getenv("LLM_API_KEY", "sk-your-key-here"),
            model="gpt-4o-mini",
        ),
        verbose=True,
    )
    
    queries = [
        "1+1等于几?",
        "帮我写一个 Python 冒泡排序",
        "今天天气怎么样?",
        "昨天晚上那场足球比赛谁赢了?",
        "最新的 iPhone 什么时候发布?",
    ]
    
    for query in queries:
        result = await run_pipeline(query, config)
        print(f"\nQ: {query}")
        print(f"  搜索? {result.searched} | {result.answer[:100]}...")


# ── 示例 3: 多轮对话上下文 ─────────────────────────────────────

async def example_multi_turn():
    """多轮对话: query 改写需要结合上下文"""
    
    config = ModuleConfig(
        llm=LLMConfig(
            api_base=os.getenv("LLM_API_BASE", "https://api.openai.com/v1"),
            api_key=os.getenv("LLM_API_KEY", "sk-your-key-here"),
            model="gpt-4o-mini",
        ),
        search=SearchConfig(
            provider="serpapi",
            api_key=os.getenv("SEARCH_API_KEY", "your-serpapi-key"),
            max_results=3,
            fetch_content=False,  # 多轮示例不抓取详情
        ),
        verbose=True,
    )
    
    history = [
        {"role": "user", "content": "OpenAI 最近有什么新消息?"},
        {"role": "assistant", "content": "OpenAI 最近发布了 GPT-4o 模型..."},
        {"role": "user", "content": "它的价格是多少?"},  # "它" 需要从上下文推断
    ]
    
    # 取最后一轮的用户问题
    result = await run_pipeline(history[-1]["content"], config, history=history)
    
    print(f"\n上下文: OpenAI 相关对话...")
    print(f"当前问题: {history[-1]['content']}")
    print(f"改写 query: {result.rewritten_query}")
    print(f"答案: {result.answer[:300]}...")


# ── 示例 4: 不使用 LLM 的纯规则判断 ──────────────────────────

async def example_rule_only():
    """
    演示: 规则引擎的覆盖率。
    很多常见问题用规则就能判断, 不需要调 LLM。
    """
    from need_search import _rule_based_check
    
    test_cases = [
        ("帮我搜一下最新的人工智能新闻", True),
        ("你好啊", False),
        ("写一个 Python 函数", False),
        ("今天的比特币价格是多少", True),
        ("帮我翻译这段话", False),
        ("苹果刚刚发布了什么新产品", True),
        ("1+1等于几", False),
        ("百度搜索一下气候变化", True),
    ]
    
    print("\n规则引擎测试:")
    print("-" * 50)
    
    hits = 0
    for query, expected in test_cases:
        decision = _rule_based_check(query)
        if decision is not None:
            match = "[OK]" if decision.need_search == expected else "[FAIL]"
            hits += 1
            print(f"{match} \"{query}\" -> need_search={decision.need_search} [{decision.source}]")
        else:
            print(f"[??] \"{query}\" -> 规则无法判断 (需要 LLM)")
    
    print(f"\n规则命中率: {hits}/{len(test_cases)}")


# ── Main ────────────────────────────────────────────────────────

async def main():
    """运行所有示例"""
    import sys
    
    examples = {
        "basic": ("完整联网流程", example_basic),
        "no-search": ("联网判断测试", example_no_search),
        "multi-turn": ("多轮对话上下文", example_multi_turn),
        "rule": ("规则引擎演示", example_rule_only),
    }
    
    if len(sys.argv) > 1 and sys.argv[1] in examples:
        name, fn = examples[sys.argv[1]]
        print(f"\n>> 运行示例: {name}\n")
        await fn()
    else:
        print("用法: python main.py [basic|no-search|multi-turn|rule]")
        print(f"\n可用示例: {', '.join(examples.keys())}")
        
        # 默认跑不需要 API key 的示例
        print("\n>> 运行规则引擎演示 (无需 API key)...\n")
        await example_rule_only()


if __name__ == "__main__":
    asyncio.run(main())