Emil_Shanaty/open-webui-ai4me
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

refac

Browse Source
This commit is contained in:
Timothy Jaeryang Baek
2026-03-17 17:58:01 -05:00
parent fcf7208352
commit de3317e26b
220 changed files with 17200 additions and 22836 deletions
Download Patch File Download Diff File Expand all files Collapse all files
backend/open_webui/retrieval/loaders/tavily.py
+16 -16
View File
@@ -25,7 +25,7 @@ class TavilyLoader(BaseLoader):
self,
urls: Union[str, List[str]],
api_key: str,
extract_depth: Literal["basic", "advanced"] = "basic",
extract_depth: Literal['basic', 'advanced'] = 'basic',
continue_on_failure: bool = True,
) -> None:
"""Initialize Tavily Extract client.
@@ -42,13 +42,13 @@ class TavilyLoader(BaseLoader):
continue_on_failure: Whether to continue if extraction of a URL fails.
"""
if not urls:
raise ValueError("At least one URL must be provided.")
raise ValueError('At least one URL must be provided.')
self.api_key = api_key
self.urls = urls if isinstance(urls, list) else [urls]
self.extract_depth = extract_depth
self.continue_on_failure = continue_on_failure
self.api_url = "https://api.tavily.com/extract"
self.api_url = 'https://api.tavily.com/extract'
def lazy_load(self) -> Iterator[Document]:
"""Extract and yield documents from the URLs using Tavily Extract API."""
@@ -57,35 +57,35 @@ class TavilyLoader(BaseLoader):
batch_urls = self.urls[i : i + batch_size]
try:
headers = {
"Content-Type": "application/json",
"Authorization": f"Bearer {self.api_key}",
'Content-Type': 'application/json',
'Authorization': f'Bearer {self.api_key}',
}
# Use string for single URL, array for multiple URLs
urls_param = batch_urls[0] if len(batch_urls) == 1 else batch_urls
payload = {"urls": urls_param, "extract_depth": self.extract_depth}
payload = {'urls': urls_param, 'extract_depth': self.extract_depth}
# Make the API call
response = requests.post(self.api_url, headers=headers, json=payload)
response.raise_for_status()
response_data = response.json()
# Process successful results
for result in response_data.get("results", []):
url = result.get("url", "")
content = result.get("raw_content", "")
for result in response_data.get('results', []):
url = result.get('url', '')
content = result.get('raw_content', '')
if not content:
log.warning(f"No content extracted from {url}")
log.warning(f'No content extracted from {url}')
continue
# Add URLs as metadata
metadata = {"source": url}
metadata = {'source': url}
yield Document(
page_content=content,
metadata=metadata,
)
for failed in response_data.get("failed_results", []):
url = failed.get("url", "")
error = failed.get("error", "Unknown error")
log.error(f"Failed to extract content from {url}: {error}")
for failed in response_data.get('failed_results', []):
url = failed.get('url', '')
error = failed.get('error', 'Unknown error')
log.error(f'Failed to extract content from {url}: {error}')
except Exception as e:
if self.continue_on_failure:
log.error(f"Error extracting content from batch {batch_urls}: {e}")
log.error(f'Error extracting content from batch {batch_urls}: {e}')
else:
raise e