Skip to content
Merged
Show file tree
Hide file tree
Changes from 1 commit
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
101 changes: 70 additions & 31 deletions openrag/app_front.py
Original file line number Diff line number Diff line change
Expand Up @@ -465,26 +465,52 @@ async def __fetch_page_content(chunk_url, headers=None):


async def _format_sources(metadata_sources, only_txt=False, api_key=None):
external_url = get_external_url() # used to override the base URL when the front-end requests a file resource
if not metadata_sources:
return None, None
return [], []

d = {}
headers = get_headers(api_key)
external_url = get_external_url() # used to override the base URL when the front-end requests a file resource
for i, s in enumerate(metadata_sources):
if not isinstance(s, dict):
continue

if s.get("source_type") == "web":
title = s.get("title") or s.get("url", f"Web source {i + 1}")
title = s.get("title", "")
url = s.get("url", "")
snippet = s.get("snippet", "")
content = f"**[{title}]({url})**\n\n{snippet}"
source_name = title
title = title.strip() if isinstance(title, str) else ""
url = url.strip() if isinstance(url, str) else ""
snippet = snippet.strip() if isinstance(snippet, str) else ""
parsed_url = urlparse(url)
if parsed_url.scheme not in {"http", "https"} or not parsed_url.netloc:
continue

source_label = title or url
source_name = source_label
if source_name in d:
source_name = f"{title} ({i})"
source_name = f"{source_name} ({i})"
content = f"**[{source_label}]({url})**"
if snippet:
content += f"\n\n{snippet}"
d[source_name] = cl.Text(content=content, name=source_name, display="side")
continue
Comment thread
coderabbitai[bot] marked this conversation as resolved.

filename = Path(s["filename"])
file_url = s["file_url"]
filename_value = s.get("filename")
file_url = s.get("file_url")
page = s.get("page")
if (
not isinstance(filename_value, str)
or not filename_value.strip()
or not isinstance(file_url, str)
or not file_url.strip()
or page is None
or not str(page).strip()
Comment thread
hedhoud marked this conversation as resolved.
Outdated
):
continue

filename = Path(filename_value.strip())
file_url = file_url.strip()
file_url = file_url.replace(INTERNAL_BASE_URL, external_url) # put the correct base url
# Avoid leaking the credential in the URL (browser history, proxy logs,
# Referer headers). In OIDC mode the browser already sends the
Expand All @@ -495,32 +521,45 @@ async def _format_sources(metadata_sources, only_txt=False, api_key=None):
# authenticate the fetch.
if api_key and (AUTH_MODE != "oidc" or _current_openrag_auth_provider() == "credentials"):
file_url = f"{file_url}?token={api_key}"
page = s["page"]
source_name = f"{filename}" + (
f" (page: {page})" if filename.suffix in [".pdf", ".pptx", ".docx", ".doc"] else ""
)

if only_txt:
chunk_content = await __fetch_page_content(chunk_url=s["chunk_url"], headers=headers)
elem = cl.Text(content=chunk_content, name=source_name, display="side")
else:
match filename.suffix.lower():
case ".pdf":
elem = cl.Pdf(
name=source_name,
url=file_url,
page=int(s["page"]),
display="side",
)
case suffix if suffix in [".png", ".jpg", ".jpeg"]:
elem = cl.Image(name=source_name, url=file_url, display="side")
case ".mp4":
elem = cl.Video(name=source_name, url=file_url, display="side")
case ".mp3":
elem = cl.Audio(name=source_name, url=file_url, display="side")
case _:
chunk_content = await __fetch_page_content(chunk_url=s["chunk_url"], headers=headers)
elem = cl.Text(content=chunk_content, name=source_name, display="side")
try:
if only_txt:
chunk_url = s.get("chunk_url")
if not isinstance(chunk_url, str) or not chunk_url.strip():
continue
chunk_content = await __fetch_page_content(chunk_url=chunk_url, headers=headers)
if not isinstance(chunk_content, str) or not chunk_content.strip():
continue
elem = cl.Text(content=chunk_content, name=source_name, display="side")
else:
match filename.suffix.lower():
case ".pdf":
elem = cl.Pdf(
name=source_name,
url=file_url,
page=int(page),
display="side",
)
case suffix if suffix in [".png", ".jpg", ".jpeg"]:
elem = cl.Image(name=source_name, url=file_url, display="side")
case ".mp4":
elem = cl.Video(name=source_name, url=file_url, display="side")
case ".mp3":
elem = cl.Audio(name=source_name, url=file_url, display="side")
case _:
chunk_url = s.get("chunk_url")
if not isinstance(chunk_url, str) or not chunk_url.strip():
continue
chunk_content = await __fetch_page_content(chunk_url=chunk_url, headers=headers)
if not isinstance(chunk_content, str) or not chunk_content.strip():
continue
elem = cl.Text(content=chunk_content, name=source_name, display="side")
except (httpx.HTTPError, TypeError, ValueError):
logger.warning("Skipping an unavailable source", source_index=i)
continue
Comment thread
coderabbitai[bot] marked this conversation as resolved.

d[source_name] = elem

Expand Down Expand Up @@ -582,7 +621,7 @@ async def on_message(message: cl.Message):
# Show sources
elements, source_names = await _format_sources(sources, api_key=api_key, only_txt=False)
msg.elements = elements if elements else []
if source_names:
if elements and source_names:
s = "\n\n" + "-" * 50 + f"\n\n{t('sources_label')}: \n" + "\n".join(source_names)
await msg.stream_token(s)
await msg.update()
Expand Down
101 changes: 85 additions & 16 deletions tests/unit/test_app_front_secret.py
Original file line number Diff line number Diff line change
Expand Up @@ -29,6 +29,16 @@ def _load_app_front(monkeypatch, *, auth_mode: str, module_name: str):
return module


def _stub_chainlit_elements(module):
module.cl = SimpleNamespace(
Pdf=lambda **kwargs: SimpleNamespace(**kwargs),
Text=lambda **kwargs: SimpleNamespace(**kwargs),
Image=lambda **kwargs: SimpleNamespace(**kwargs),
Video=lambda **kwargs: SimpleNamespace(**kwargs),
Audio=lambda **kwargs: SimpleNamespace(**kwargs),
)


def test_no_hardcoded_default_secret_assignment_in_source():
"""The fall-through to a literal default secret must be gone.

Expand Down Expand Up @@ -398,6 +408,77 @@ async def fake_load_model_ids(_client, api_key):
assert module._OPENRAG_TOKEN_STORE[auth_handle][0] == "handoff-token"


@pytest.mark.parametrize(
"sources",
[
None,
[],
[{}],
[None],
[{"source_type": "web", "title": "", "url": "", "snippet": ""}],
[{"filename": "", "file_url": "", "page": ""}],
],
)
@pytest.mark.asyncio
async def test_chainlit_hides_sources_when_none_are_displayable(monkeypatch, sources):
module = _load_app_front(monkeypatch, auth_mode="token", module_name="app_front_empty_sources_test")
monkeypatch.setattr(module, "get_external_url", lambda: "https://openrag.example")
_stub_chainlit_elements(module)

elements, source_names = await module._format_sources(sources)

assert elements == []
assert source_names == []


@pytest.mark.asyncio
async def test_chainlit_keeps_valid_web_sources(monkeypatch):
module = _load_app_front(monkeypatch, auth_mode="token", module_name="app_front_web_source_test")
monkeypatch.setattr(module, "get_external_url", lambda: "https://openrag.example")
_stub_chainlit_elements(module)

elements, source_names = await module._format_sources(
[
{
"source_type": "web",
"title": "Example reference",
"url": "https://example.test/reference",
"snippet": "Supporting evidence",
}
]
)

assert source_names == ["Example reference"]
assert elements[0].name == "Example reference"
assert elements[0].content == ("**[Example reference](https://example.test/reference)**\n\nSupporting evidence")


@pytest.mark.asyncio
async def test_chainlit_skips_unavailable_text_sources(monkeypatch):
module = _load_app_front(monkeypatch, auth_mode="token", module_name="app_front_unavailable_source_test")
monkeypatch.setattr(module, "get_external_url", lambda: "https://openrag.example")
_stub_chainlit_elements(module)

async def unavailable_chunk(*_args, **_kwargs):
raise httpx.ConnectError("source unavailable")

monkeypatch.setattr(module, "__fetch_page_content", unavailable_chunk)

elements, source_names = await module._format_sources(
[
{
"filename": "notes.txt",
"file_url": "http://internal:8080/static/source-id",
"page": "1",
"chunk_url": "http://internal:8080/chunks/source-id",
}
]
)

assert elements == []
assert source_names == []


@pytest.mark.asyncio
async def test_oidc_token_handoff_keeps_bearer_on_static_source_urls(monkeypatch):
module = _load_app_front(monkeypatch, auth_mode="oidc", module_name="app_front_source_token_test")
Expand All @@ -410,14 +491,8 @@ def get(self, key):
return SimpleNamespace(metadata={"provider": "credentials"})
return None

module.cl = SimpleNamespace(
user_session=UserSession(),
Pdf=lambda **kwargs: SimpleNamespace(**kwargs),
Text=lambda **kwargs: SimpleNamespace(**kwargs),
Image=lambda **kwargs: SimpleNamespace(**kwargs),
Video=lambda **kwargs: SimpleNamespace(**kwargs),
Audio=lambda **kwargs: SimpleNamespace(**kwargs),
)
_stub_chainlit_elements(module)
module.cl.user_session = UserSession()

elements, _ = await module._format_sources(
[
Expand Down Expand Up @@ -445,14 +520,8 @@ def get(self, key):
return SimpleNamespace(metadata={"provider": "oidc"})
return None

module.cl = SimpleNamespace(
user_session=UserSession(),
Pdf=lambda **kwargs: SimpleNamespace(**kwargs),
Text=lambda **kwargs: SimpleNamespace(**kwargs),
Image=lambda **kwargs: SimpleNamespace(**kwargs),
Video=lambda **kwargs: SimpleNamespace(**kwargs),
Audio=lambda **kwargs: SimpleNamespace(**kwargs),
)
_stub_chainlit_elements(module)
module.cl.user_session = UserSession()

elements, _ = await module._format_sources(
[
Expand Down
Loading