diff --git a/Dockerfile b/Dockerfile index 8fc3fe0..9b02061 100644 --- a/Dockerfile +++ b/Dockerfile @@ -15,10 +15,7 @@ ENV GITHUB_BUILD=${GITHUB_BUILD}\ UV_LINK_MODE=copy \ PORT=8191 \ XDG_CACHE_HOME=/cache \ - HOME=/tmp \ - # Optional: set to require bearer token on /load requests. - # Must match EXTERNAL_WEB_LOADER_API_KEY in Open WebUI. - OWUI_API_KEY="" + HOME=/tmp RUN apt-get update &&\ apt-get install -y --no-install-recommends curl ca-certificates git tini &&\ diff --git a/src/consts.py b/src/consts.py index 249d4ec..207139b 100644 --- a/src/consts.py +++ b/src/consts.py @@ -23,6 +23,7 @@ class Settings(BaseSettings): block_media: bool = False return_only_cookies: bool = False + owui_api_key: str | None = None settings = Settings() @@ -42,6 +43,8 @@ PORT = settings.port BLOCK_MEDIA = settings.block_media RETURN_ONLY_COOKIES = settings.return_only_cookies +OWUI_API_KEY = settings.owui_api_key + CHALLENGE_TITLES_MAP: dict[CaptchaType, list[str]] = { # Cloudflare CaptchaType.CLOUDFLARE_INTERSTITIAL: ["Just a moment..."], diff --git a/src/owui.py b/src/owui.py index 55e660b..293e760 100644 --- a/src/owui.py +++ b/src/owui.py @@ -1,114 +1,72 @@ -""" -src/owui.py — Open WebUI external web loader endpoint. - -Implements the contract expected by Open WebUI's WEB_LOADER_ENGINE=external: - POST /load - Request: {"urls": ["https://..."]} - Response: [{"page_content": str, "metadata": {"source": str}}] - -Reuses Byparr's existing browser dependency injection so each request -gets a properly-initialised browser context with challenge solving. -Uses document.body.innerText for content extraction. - -Configure in Open WebUI: - WEB_LOADER_ENGINE=external - EXTERNAL_WEB_LOADER_URL=http://byparr:8191/load - EXTERNAL_WEB_LOADER_API_KEY= -""" +"""Open WebUI external web loader endpoint: POST /load.""" from __future__ import annotations -import os -from typing import Annotated, Any +from hmac import compare_digest +from typing import Annotated from fastapi import APIRouter, Depends, Header, HTTPException +from playwright.async_api import Page +from playwright.async_api import TimeoutError as PlaywrightTimeoutError from pydantic import BaseModel +from src.consts import OWUI_API_KEY from src.utils import BrowserDepClass, get_browser, logger -_API_KEY: str | None = os.getenv("OWUI_API_KEY") or None - router = APIRouter(tags=["Open WebUI"]) BrowserDep = Annotated[BrowserDepClass, Depends(get_browser)] -# --------------------------------------------------------------------------- -# Models -# --------------------------------------------------------------------------- - - class LoadRequest(BaseModel): urls: list[str] class LoadResult(BaseModel): page_content: str - metadata: dict[str, Any] + metadata: dict[str, str] -# --------------------------------------------------------------------------- -# Content extraction -# --------------------------------------------------------------------------- - - -async def _extract_content(page: Any) -> str: - """ - Extract text content from the page using document.body.innerText. - Whitespace is collapsed for clean RAG chunking. - """ - result: str = await page.evaluate(""" - () => { - return document.body ? document.body.innerText : ""; - } - """) - lines = [line.strip() for line in result.splitlines() if line.strip()] - return "\n".join(lines) - - -# --------------------------------------------------------------------------- -# Auth -# --------------------------------------------------------------------------- - - -def _check_auth(authorization: str | None) -> None: - if not _API_KEY: +def require_auth(authorization: Annotated[str | None, Header()] = None) -> None: + """Enforce a bearer token on /load when OWUI_API_KEY is set.""" + if not OWUI_API_KEY: return - if authorization != f"Bearer {_API_KEY}": + if authorization is None or not compare_digest( + authorization.encode(), f"Bearer {OWUI_API_KEY}".encode() + ): raise HTTPException(status_code=401, detail="Unauthorized") -# --------------------------------------------------------------------------- -# Endpoint -# --------------------------------------------------------------------------- +async def _extract_content(page: Page) -> str: + """Return the page's visible text, with blank lines removed.""" + result = await page.evaluate("() => document.body ? document.body.innerText : ''") + return "\n".join(line.strip() for line in result.splitlines() if line.strip()) @router.post("/load", response_model=list[LoadResult]) async def load_urls( request: LoadRequest, - authorization: Annotated[str | None, Header()] = None, - dep: BrowserDep = None, # type: ignore[assignment] + _auth: Annotated[None, Depends(require_auth)], + dep: BrowserDep, ) -> list[LoadResult]: """ - Fetch URLs through InvisiblePlaywright, bypassing Cloudflare challenges. - Returns plain-text content for Open WebUI's RAG pipeline. - Failed URLs return empty page_content so the search degrades gracefully. + Fetch URLs through the anti-bot browser and return their text content. + + Each URL is fetched sequentially; a failing URL yields empty + page_content so Open WebUI's RAG pipeline degrades gracefully. """ - _check_auth(authorization) - results: list[LoadResult] = [] - for url in request.urls: try: - logger.info("OWUI loader fetching: %s", url) await dep.page.goto(url, timeout=60_000) await dep.page.wait_for_load_state("domcontentloaded", timeout=30_000) - await dep.page.wait_for_load_state("networkidle", timeout=15_000) + try: + await dep.page.wait_for_load_state("networkidle", timeout=15_000) + except PlaywrightTimeoutError: + logger.debug("networkidle timed out for %s; extracting anyway", url) content = await _extract_content(dep.page) - logger.debug("OWUI loader: %s -> %d chars", url, len(content)) - results.append(LoadResult(page_content=content, metadata={"source": url})) except Exception as exc: # noqa: BLE001 - logger.warning("OWUI loader failed for %s: %s", url, exc) - results.append(LoadResult(page_content="", metadata={"source": url})) - + logger.warning("Failed to load %s: %s", url, exc) + content = "" + results.append(LoadResult(page_content=content, metadata={"source": url})) return results diff --git a/tests/owui_test.py b/tests/owui_test.py index a51b9b7..1b79fb1 100644 --- a/tests/owui_test.py +++ b/tests/owui_test.py @@ -1,66 +1,63 @@ from http import HTTPStatus -from unittest.mock import patch +from unittest.mock import AsyncMock, MagicMock, patch import pytest +from playwright.async_api import TimeoutError as PlaywrightTimeoutError from starlette.testclient import TestClient from main import app +from src.owui import LoadRequest, load_urls +from src.utils import BrowserDepClass client = TestClient(app) def test_owui_load_basic(): - """Test /load endpoint returns correct structure.""" + """/load returns one result per URL with the expected shape.""" response = client.post("/load", json={"urls": ["https://example.com"]}) assert response.status_code == HTTPStatus.OK results = response.json() assert len(results) == 1 - assert "page_content" in results[0] - assert "metadata" in results[0] - assert results[0]["metadata"]["source"] == "https://example.com" + assert results[0]["page_content"] + assert results[0]["metadata"] == {"source": "https://example.com"} def test_owui_load_multiple_urls(): - """Test /load endpoint handles multiple URLs.""" + """/load returns one result per URL, in order.""" urls = ["https://example.com", "https://example.org"] response = client.post("/load", json={"urls": urls}) assert response.status_code == HTTPStatus.OK results = response.json() - assert len(results) == 2 - assert results[0]["metadata"]["source"] == urls[0] - assert results[1]["metadata"]["source"] == urls[1] + assert [r["metadata"]["source"] for r in results] == urls -def test_owui_load_extracts_content(): - """Test /load endpoint extracts non-empty content from a real page.""" - response = client.post("/load", json={"urls": ["https://example.com"]}) +def test_owui_load_invalid_url_graceful(): + """Unreachable URLs yield empty page_content instead of an error.""" + response = client.post( + "/load", json={"urls": ["https://this-domain-does-not-exist-12345.invalid"]} + ) assert response.status_code == HTTPStatus.OK results = response.json() - # example.com has minimal content but should have something - assert len(results[0]["page_content"]) > 0 + assert len(results) == 1 + assert results[0]["page_content"] == "" -def test_owui_load_auth_required(): - """Test /load returns 401 when API key is set but not provided.""" - with patch("src.owui._API_KEY", "test-secret-key"): - response = client.post("/load", json={"urls": ["https://example.com"]}) - assert response.status_code == HTTPStatus.UNAUTHORIZED - - -def test_owui_load_auth_invalid(): - """Test /load returns 401 when API key is wrong.""" - with patch("src.owui._API_KEY", "test-secret-key"): +@pytest.mark.parametrize( + "headers", + [None, {"Authorization": "Bearer wrong-key"}], +) +def test_owui_load_rejects_missing_or_wrong_key(headers): + """/load returns 401 without a valid bearer token when a key is set.""" + with patch("src.owui.OWUI_API_KEY", "test-secret-key"): response = client.post( - "/load", - json={"urls": ["https://example.com"]}, - headers={"Authorization": "Bearer wrong-key"}, + "/load", json={"urls": ["https://example.com"]}, headers=headers ) assert response.status_code == HTTPStatus.UNAUTHORIZED -def test_owui_load_auth_valid(): - """Test /load succeeds with correct API key.""" - with patch("src.owui._API_KEY", "test-secret-key"): +def test_owui_load_accepts_valid_key(): + """/load succeeds with the configured bearer token.""" + with patch("src.owui.OWUI_API_KEY", "test-secret-key"): response = client.post( "/load", json={"urls": ["https://example.com"]}, @@ -69,13 +66,25 @@ def test_owui_load_auth_valid(): assert response.status_code == HTTPStatus.OK -def test_owui_load_invalid_url_graceful(): - """Test /load returns empty content for unreachable URLs.""" - response = client.post( - "/load", json={"urls": ["https://this-domain-does-not-exist-12345.invalid"]} +def fake_dep() -> BrowserDepClass: + """Browser dependency whose page loads text but never reaches networkidle.""" + page = AsyncMock() + page.goto.return_value = MagicMock() + page.evaluate.return_value = "line one\n\nline two" + + def wait_for_load_state(state: str, **_kwargs: object) -> None: + if state == "networkidle": + message = "load state wait timed out" + raise PlaywrightTimeoutError(message) + + page.wait_for_load_state.side_effect = wait_for_load_state + return BrowserDepClass(page=page, solver=AsyncMock(), context=AsyncMock()) + + +@pytest.mark.asyncio +async def test_networkidle_timeout_still_extracts_content(): + """A page that never reaches networkidle still yields its text.""" + results = await load_urls( + LoadRequest(urls=["https://example.test"]), None, fake_dep() ) - assert response.status_code == HTTPStatus.OK - results = response.json() - assert len(results) == 1 - assert results[0]["page_content"] == "" - assert "this-domain-does-not-exist" in results[0]["metadata"]["source"] + assert results[0].page_content == "line one\nline two"