four-level-schedule/test_search_final.py

#!/usr/bin/env python3
"""
测试向量检索功能
"""
import requests
import json

# 配置
EMBEDDING_API_URL = "http://localhost:8084"


def test_search():
    """测试搜索功能"""
    print("=" * 60)
    print("测试向量检索功能")
    print("=" * 60)
    print()

    # 测试查询
    queries = [
        "数学",
        "口算题",
        "一年级",
    ]

    for query in queries:
        print("-" * 60)
        print(f"查询: {query}")
        print("-" * 60)

        url = f"{EMBEDDING_API_URL}/api/v1/search"

        try:
            response = requests.post(
                url,
                json={"query": query, "topK": 5},
                headers={"Content-Type": "application/json"},
                timeout=30,
            )

            if response.status_code != 200:
                print(f"❌ HTTP错误: {response.status_code}")
                print(f"   {response.text}")
                continue

            results = response.json()
            print(f"✅ 找到 {len(results)} 个结果:\n")

            for i, result in enumerate(results, 1):
                print(f"  结果 {i}:")
                print(f"    文档ID: {result.get('docId')}")
                print(f"    文件名: {result.get('metadata', {}).get('file_name')}")
                print(f"    相似度: {result.get('score'):.4f}")
                print(f"    文件类型: {result.get('metadata', {}).get('file_type')}")
                content = result.get('content', '')
                if len(content) > 100:
                    content = content[:100] + "..."
                print(f"    内容片段: {content}")
                print()

        except Exception as e:
            print(f"❌ 调用异常: {e}")
            import traceback
            traceback.print_exc()
        print()

    print("=" * 60)


if __name__ == "__main__":
    test_search()