|
import requests |
|
import logging |
|
from typing import Iterator, List, Literal, Union |
|
|
|
from langchain_core.document_loaders import BaseLoader |
|
from langchain_core.documents import Document |
|
from open_webui.env import SRC_LOG_LEVELS |
|
|
|
log = logging.getLogger(__name__) |
|
log.setLevel(SRC_LOG_LEVELS["RAG"]) |
|
|
|
|
|
class TavilyLoader(BaseLoader): |
|
"""Extract web page content from URLs using Tavily Extract API. |
|
|
|
This is a LangChain document loader that uses Tavily's Extract API to |
|
retrieve content from web pages and return it as Document objects. |
|
|
|
Args: |
|
urls: URL or list of URLs to extract content from. |
|
api_key: The Tavily API key. |
|
extract_depth: Depth of extraction, either "basic" or "advanced". |
|
continue_on_failure: Whether to continue if extraction of a URL fails. |
|
""" |
|
|
|
def __init__( |
|
self, |
|
urls: Union[str, List[str]], |
|
api_key: str, |
|
extract_depth: Literal["basic", "advanced"] = "basic", |
|
continue_on_failure: bool = True, |
|
) -> None: |
|
"""Initialize Tavily Extract client. |
|
|
|
Args: |
|
urls: URL or list of URLs to extract content from. |
|
api_key: The Tavily API key. |
|
include_images: Whether to include images in the extraction. |
|
extract_depth: Depth of extraction, either "basic" or "advanced". |
|
advanced extraction retrieves more data, including tables and |
|
embedded content, with higher success but may increase latency. |
|
basic costs 1 credit per 5 successful URL extractions, |
|
advanced costs 2 credits per 5 successful URL extractions. |
|
continue_on_failure: Whether to continue if extraction of a URL fails. |
|
""" |
|
if not urls: |
|
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" |
|
|
|
def lazy_load(self) -> Iterator[Document]: |
|
"""Extract and yield documents from the URLs using Tavily Extract API.""" |
|
batch_size = 20 |
|
for i in range(0, len(self.urls), batch_size): |
|
batch_urls = self.urls[i : i + batch_size] |
|
try: |
|
headers = { |
|
"Content-Type": "application/json", |
|
"Authorization": f"Bearer {self.api_key}", |
|
} |
|
|
|
urls_param = batch_urls[0] if len(batch_urls) == 1 else batch_urls |
|
payload = {"urls": urls_param, "extract_depth": self.extract_depth} |
|
|
|
response = requests.post(self.api_url, headers=headers, json=payload) |
|
response.raise_for_status() |
|
response_data = response.json() |
|
|
|
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}") |
|
continue |
|
|
|
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}") |
|
except Exception as e: |
|
if self.continue_on_failure: |
|
log.error(f"Error extracting content from batch {batch_urls}: {e}") |
|
else: |
|
raise e |
|
|
