mirror of
https://github.com/open-webui/open-webui
synced 2025-04-29 10:42:55 +00:00
54 lines
1.8 KiB
Python
54 lines
1.8 KiB
Python
import requests
|
|
import logging
|
|
from typing import Iterator, List, 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 ExternalLoader(BaseLoader):
|
|
def __init__(
|
|
self,
|
|
web_paths: Union[str, List[str]],
|
|
external_url: str,
|
|
external_api_key: str,
|
|
continue_on_failure: bool = True,
|
|
**kwargs,
|
|
) -> None:
|
|
self.external_url = external_url
|
|
self.external_api_key = external_api_key
|
|
self.urls = web_paths if isinstance(web_paths, list) else [web_paths]
|
|
self.continue_on_failure = continue_on_failure
|
|
|
|
def lazy_load(self) -> Iterator[Document]:
|
|
batch_size = 20
|
|
for i in range(0, len(self.urls), batch_size):
|
|
urls = self.urls[i : i + batch_size]
|
|
try:
|
|
response = requests.post(
|
|
self.external_url,
|
|
headers={
|
|
"User-Agent": "Open WebUI (https://github.com/open-webui/open-webui) RAG Bot",
|
|
"Authorization": f"Bearer {self.external_api_key}",
|
|
},
|
|
json={
|
|
"urls": urls,
|
|
},
|
|
)
|
|
response.raise_for_status()
|
|
results = response.json()
|
|
for result in results:
|
|
yield Document(
|
|
page_content=result.get("page_content", ""),
|
|
metadata=result.get("metadata", {}),
|
|
)
|
|
except Exception as e:
|
|
if self.continue_on_failure:
|
|
log.error(f"Error extracting content from batch {urls}: {e}")
|
|
else:
|
|
raise e
|