"""
通过 API 测试知识库检索性能
"""

import requests
import json
import time

API_BASE = "http://localhost:8000"
API_KEY = "oQfkSAbL8iafzyHxqb--G7zRWSOYJHvlzQxia2KpYms"
TENANT_ID = "szmp@ash@2026"

def test_kb_search():
    """测试知识库搜索 API"""
    print("=" * 80)
    print("测试知识库检索 API")
    print("=" * 80)
    
    headers = {
        "Content-Type": "application/json",
        "X-API-Key": API_KEY,
        "X-Tenant-Id": TENANT_ID,
    }
    
    # 测试数据
    test_cases = [
        {
            "name": "完整参数（含context过滤）",
            "data": {
                "query": "三年级语文学习",
                "scene": "学习方案",
                "top_k": 5,
                "context": {"grade": "三年级", "subject": "语文"},
            }
        },
        {
            "name": "简化参数（无context）",
            "data": {
                "query": "三年级语文学习",
                "scene": "学习方案",
                "top_k": 5,
            }
        },
    ]
    
    for test_case in test_cases:
        print(f"\n{'='*80}")
        print(f"测试: {test_case['name']}")
        print(f"{'='*80}")
        print(f"请求数据: {json.dumps(test_case['data'], ensure_ascii=False)}")
        
        try:
            start = time.time()
            response = requests.post(
                f"{API_BASE}/api/v1/mid/kb-search-dynamic",
                headers=headers,
                json=test_case['data'],
                timeout=30,
            )
            elapsed = (time.time() - start) * 1000
            
            print(f"\n响应状态: {response.status_code}")
            print(f"总耗时: {elapsed:.2f} ms")
            
            if response.status_code == 200:
                result = response.json()
                print(f"API 结果:")
                print(f"  success: {result.get('success')}")
                print(f"  hits count: {len(result.get('hits', []))}")
                print(f"  duration_ms: {result.get('duration_ms')}")
                print(f"  applied_filter: {result.get('applied_filter')}")
            else:
                print(f"错误: {response.text}")
                
        except Exception as e:
            print(f"请求失败: {e}")
    
    print("\n" + "=" * 80)


if __name__ == "__main__":
    test_kb_search()