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

1 import os

2 import time

3 import requests

4 from typing import List, Optional, Dict, Any

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

6 from qrclaw.config import TAVILY_API_KEY

7 from qrclaw.logger import get_logger

8

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

10

11 class 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}")