Coverage for web_search / providers / tavily.py: 40%

1import os 

2import time 

3import requests 

4from typing import List, Optional, Dict, Any 

5from qrclaw.web_search.types import WebSearchProvider, SearchResult, WebSearchResponse 

6from qrclaw.config import TAVILY_API_KEY 

7from qrclaw.logger import get_logger 

8 

9logger = get_logger("qrclaw.web_search.providers.tavily") 

10 

11class TavilySearchProvider(WebSearchProvider): 

12 @property 

13 def id(self) -> str: 

14 return "tavily" 

15 

16 @property 

17 def name(self) -> str: 

18 return "Tavily Search" 

19 

20 def is_available(self) -> bool: 

21 return bool(TAVILY_API_KEY) 

22 

23 def search(self, query: str, max_results: int = 5, **kwargs) -> WebSearchResponse: 

24 """ 

25 执行 Tavily 搜索 

26  

27 Args: 

28 query: 搜索词 

29 max_results: 返回结果数量 (1-20) 

30 **kwargs: 其他参数，如 search_depth, include_answer 

31 """ 

32 if not self.is_available(): 

33 raise ValueError("TAVILY_API_KEY is not configured") 

34 

35 start_time = time.time() 

36 

37 # 参数处理 

38 count = max(1, min(20, int(max_results))) 

39 search_depth = kwargs.get("search_depth", "basic") 

40 include_answer = kwargs.get("include_answer", False) 

41 

42 url = "https://api.tavily.com/search" 

43 payload = { 

44 "api_key": TAVILY_API_KEY, 

45 "query": query, 

46 "max_results": count, 

47 "search_depth": search_depth, 

48 "include_answer": include_answer, 

49 "include_images": False, 

50 "include_raw_content": False 

51 } 

52 

53 try: 

54 logger.debug(f"Tavily search request: {query} (count={count})") 

55 response = requests.post(url, json=payload, timeout=30) 

56 response.raise_for_status() 

57 data = response.json() 

58 

59 results: List[SearchResult] = [] 

60 raw_results = data.get("results", []) 

61 

62 for item in raw_results: 

63 results.append(SearchResult( 

64 title=item.get("title", ""), 

65 url=item.get("url", ""), 

66 snippet=item.get("content", "")[:500], # 限制摘要长度 

67 score=item.get("score"), 

68 published_date=item.get("published_date") 

69 )) 

70 

71 took_ms = (time.time() - start_time) * 1000 

72 

73 return WebSearchResponse( 

74 query=query, 

75 provider=self.id, 

76 count=len(results), 

77 took_ms=took_ms, 

78 results=results, 

79 answer=data.get("answer"), 

80 raw_response=data 

81 ) 

82 

83 except requests.exceptions.RequestException as e: 

84 logger.error(f"Tavily search failed: {e}") 

85 raise RuntimeError(f"Tavily API error: {e}")