ailine/tools/test/test_tavily_search.py

#!/usr/bin/env python3
"""测试 Tavily 搜索功能 - 直接调用 API"""
import sys
from pathlib import Path
from dataclasses import dataclass
from datetime import datetime
from typing import List, Optional
from dotenv import load_dotenv

# 路径设置
project_root = Path(__file__).resolve().parent.parent.parent
sys.path.insert(0, str(project_root))
load_dotenv(project_root / ".env")

import os

TAVILY_API_KEY = os.getenv("TAVILY_API_KEY")
TAVILY_MAX_RESULTS = int(os.getenv("TAVILY_MAX_RESULTS") or "5")


@dataclass
class SearchResult:
    """搜索结果数据类"""
    title: str
    url: str
    snippet: str
    source: str = "DuckDuckGo"
    timestamp: datetime = None

    def __post_init__(self):
        if self.timestamp is None:
            self.timestamp = datetime.now()


def test_tavily_api_key():
    """测试 API Key 配置"""
    print("=" * 60)
    print("测试 1: 检查 Tavily API Key")
    print("=" * 60)

    if TAVILY_API_KEY:
        print(f"✓ TAVILY_API_KEY 已配置: {TAVILY_API_KEY[:15]}...")
    else:
        print("✗ TAVILY_API_KEY 未配置")
    print()


def test_tavily_search_direct():
    """直接测试 Tavily API"""
    print("=" * 60)
    print("测试 2: 直接调用 Tavily API")
    print("=" * 60)

    if not TAVILY_API_KEY:
        print("✗ 未配置 API Key，跳过测试")
        return

    from tavily import TavilyClient

    client = TavilyClient(api_key=TAVILY_API_KEY)

    test_queries = [
        "Python 编程语言最新版本",
        "LangGraph AI 框架",
    ]

    for query in test_queries:
        print(f"\n搜索: {query}")
        print("-" * 40)
        try:
            response = client.search(
                query=query,
                max_results=3,
                include_answer=True,
                include_raw_content=False
            )

            print(f"✓ 搜索成功")
            print(f"  - 结果数量: {len(response.get('results', []))}")

            # 打印结果
            for i, item in enumerate(response.get("results", []), 1):
                print(f"\n  [{i}] {item.get('title', '')}")
                print(f"      URL: {item.get('url', '')}")
                print(f"      摘要: {item.get('content', '')[:100]}...")

            # 如果有 answer
            if response.get("answer"):
                print(f"\n  🤖 AI 摘要: {response['answer'][:200]}...")

        except Exception as e:
            print(f"✗ 搜索失败: {e}")
    print()


def test_web_search_integration():
    """测试 web_search 模块集成"""
    print("=" * 60)
    print("测试 3: 测试 web_search 模块集成")
    print("=" * 60)

    # 直接导入 web_search 模块（避免循环依赖）
    web_search_path = project_root / "backend" / "app" / "core" / "web_search.py"
    if not web_search_path.exists():
        print(f"✗ web_search.py 不存在于 {web_search_path}")
        return

    print(f"✓ 找到 web_search.py: {web_search_path}")

    # 使用 exec 动态加载模块
    import importlib.util
    spec = importlib.util.spec_from_file_location("web_search_module", web_search_path)
    web_search_module = importlib.util.module_from_spec(spec)

    try:
        spec.loader.exec_module(web_search_module)
        print("✓ web_search 模块加载成功")
    except Exception as e:
        print(f"✗ 模块加载失败: {e}")
        return

    # 测试搜索
    print("\n执行搜索测试:")
    try:
        result = web_search_module.web_search("今天天气怎么样", max_results=3)
        print(f"✓ 搜索成功，返回 {len(result)} 字符")
        print("-" * 40)
        print(result[:800] + "..." if len(result) > 800 else result)
    except Exception as e:
        print(f"✗ 搜索失败: {e}")
    print()


def main():
    print("\n" + "=" * 60)
    print("🚀 Tavily 搜索功能测试")
    print("=" * 60 + "\n")

    test_tavily_api_key()
    test_tavily_search_direct()
    test_web_search_integration()

    print("=" * 60)
    print("✅ 测试完成")
    print("=" * 60)


if __name__ == "__main__":
    main()