60 lines
1.8 KiB
Python
60 lines
1.8 KiB
Python
"""Tavily 搜索源"""
|
|
|
|
import time
|
|
import requests
|
|
from providers.base import BaseProvider, SearchResult
|
|
|
|
|
|
class TavilyProvider(BaseProvider):
|
|
name = 'tavily'
|
|
display_name = 'Tavily'
|
|
needs_api_key = True
|
|
enabled = True
|
|
priority = 20
|
|
|
|
def __init__(self, config: dict):
|
|
super().__init__(config)
|
|
tc = config.get('tavily', {})
|
|
self.api_key = tc.get('api_key')
|
|
self.base_url = tc.get('base_url', 'https://api.tavily.com').rstrip('/')
|
|
self.depth = tc.get('depth', 'basic')
|
|
self.max_results = tc.get('max_results', 10)
|
|
|
|
def is_available(self) -> bool:
|
|
return bool(self.api_key)
|
|
|
|
def search(self, query: str, max_results: int = None) -> list:
|
|
if not self.api_key:
|
|
return []
|
|
|
|
url = f'{self.base_url}/search'
|
|
payload = {
|
|
'api_key': self.api_key,
|
|
'query': query,
|
|
'search_depth': self.depth,
|
|
'max_results': max_results or self.max_results,
|
|
'include_answer': False,
|
|
'include_images': False,
|
|
}
|
|
|
|
try:
|
|
resp = requests.post(url, json=payload, timeout=30)
|
|
if resp.status_code != 200:
|
|
return []
|
|
|
|
data = resp.json()
|
|
results = []
|
|
for item in data.get('results', []):
|
|
results.append(SearchResult(
|
|
title=item.get('title', ''),
|
|
url=item.get('url', ''),
|
|
content=item.get('content', ''),
|
|
score=item.get('score', 0),
|
|
source=self.name,
|
|
published_date=item.get('published_date', ''),
|
|
))
|
|
return results
|
|
|
|
except requests.exceptions.RequestException:
|
|
return []
|