diff --git a/openrag/api/routers/user/chat.py b/openrag/api/routers/user/chat.py index 9e315d110..81f93eeb7 100644 --- a/openrag/api/routers/user/chat.py +++ b/openrag/api/routers/user/chat.py @@ -14,7 +14,6 @@ import asyncio import json from typing import TYPE_CHECKING -from urllib.parse import urlparse import consts from api.dependencies.auth import ( @@ -34,6 +33,7 @@ from core.utils.exceptions import OpenRAGError from core.utils.logging import get_logger from core.utils.text import get_num_tokens, sanitize_text +from core.utils.web_url import normalize_web_url from di.providers import get_config, get_partition_service, get_query_service from fastapi import APIRouter, Body, Depends, HTTPException, Request, status from fastapi.responses import JSONResponse, StreamingResponse @@ -225,8 +225,8 @@ def chunk_url(extract_id) -> str: doc_metadata = dict(doc.metadata) links.append(build_document_source_link(doc_metadata, static_url, chunk_url)) for result in web_results or []: - url = sanitize_text(result.url or "") - if not url or urlparse(url).scheme not in ("http", "https"): + url = normalize_web_url(result.url) + if url is None: continue links.append( { diff --git a/openrag/api/routers/user/source_links.py b/openrag/api/routers/user/source_links.py index 7ace374d9..cb17289f5 100644 --- a/openrag/api/routers/user/source_links.py +++ b/openrag/api/routers/user/source_links.py @@ -33,9 +33,11 @@ def build_document_source_link( encoded_url = None if filename: encoded_url = quote(static_url_builder(doc_metadata["_id"]), safe=":/") - return { - "source_type": "document", - **({"file_url": encoded_url} if encoded_url else {}), - "chunk_url": chunk_url_builder(doc_metadata["_id"]), - **doc_metadata, - } + link = dict(doc_metadata) + link["source_type"] = "document" + link["chunk_url"] = chunk_url_builder(doc_metadata["_id"]) + if encoded_url: + link["file_url"] = encoded_url + else: + link.pop("file_url", None) + return link diff --git a/openrag/app_front.py b/openrag/app_front.py index fd5be7434..7b8efd788 100644 --- a/openrag/app_front.py +++ b/openrag/app_front.py @@ -1,10 +1,13 @@ import json import os +import re import secrets +import string import time +import unicodedata from functools import lru_cache from pathlib import Path -from urllib.parse import urlparse +from urllib.parse import quote, urlparse import chainlit as cl import httpx @@ -19,6 +22,7 @@ CHAINLIT_TOKEN_COOKIE_PATH, ) from core.utils.logging import get_logger, mask_email +from core.utils.web_url import normalize_web_url from dotenv import load_dotenv from openai import AsyncOpenAI @@ -42,12 +46,94 @@ OPENRAG_CHAT_PROFILES_METADATA_KEY = "openrag_chat_profiles" OPENRAG_SESSION_COOKIE_NAME = "openrag_session" _OPENRAG_TOKEN_STORE: dict[str, tuple[str, float]] = {} +_MARKDOWN_ESCAPE_TABLE = str.maketrans({char: f"\\{char}" for char in string.punctuation}) +_MARKDOWN_URL_SAFE_CHARS = ":/?#[]@!$&'+,;=%" +# Chainlit inserts source names into Markdown link labels. Strip only characters +# that can break or restyle that label so ordinary filenames stay recognizable. +_MARKDOWN_UNSAFE_SOURCE_NAME_CHARS = str.maketrans(dict.fromkeys("[]*`\\<>", " ")) +_MARKDOWN_BLOCK_PREFIX_RE = re.compile(r"^(?:(?:#{1,6}|[-+]|\d{1,9}[.)])\s+)+") +_MARKDOWN_THEMATIC_BREAK_RE = re.compile(r"^(?:(?:-\s*){3,}|(?:_\s*){3,}|(?:\*\s*){3,})$") class MissingOpenRAGCredentialError(RuntimeError): pass +def _escape_markdown_text(value: str) -> str: + """Render untrusted source metadata as literal Markdown text.""" + return value.translate(_MARKDOWN_ESCAPE_TABLE) + + +def _neutralize_source_name_delimiters(value: str) -> str: + """Remove active Markdown delimiters while preserving safe filename punctuation.""" + characters = list(value) + underscore_openers: list[tuple[int, int]] = [] + index = 0 + + while index < len(value): + if value[index] != "_": + index += 1 + continue + + end = index + 1 + while end < len(value) and value[end] == "_": + end += 1 + + previous = value[index - 1] if index else None + following = value[end] if end < len(value) else None + previous_whitespace = previous is None or previous.isspace() + following_whitespace = following is None or following.isspace() + previous_punctuation = previous is not None and unicodedata.category(previous)[0] in {"P", "S"} + following_punctuation = following is not None and unicodedata.category(following)[0] in {"P", "S"} + left_flanking = not following_whitespace and ( + not following_punctuation or previous_whitespace or previous_punctuation + ) + right_flanking = not previous_whitespace and ( + not previous_punctuation or following_whitespace or following_punctuation + ) + can_open = left_flanking and (not right_flanking or previous_punctuation) + can_close = right_flanking and (not left_flanking or following_punctuation) + + if can_close and underscore_openers: + opener_start, opener_end = underscore_openers.pop() + characters[opener_start:opener_end] = " " * (opener_end - opener_start) + characters[index:end] = " " * (end - index) + elif can_open: + underscore_openers.append((index, end)) + index = end + + index = 0 + while index < len(value): + if value[index] != "~": + index += 1 + continue + end = index + 1 + while end < len(value) and value[end] == "~": + end += 1 + if end - index >= 2: + characters[index:end] = " " * (end - index) + index = end + + return "".join(characters) + + +def _safe_source_name(value: str, existing: dict) -> str: + """Build a Markdown-inert name that Chainlit can match to its element.""" + value = _neutralize_source_name_delimiters(value) + value = value.lstrip() + if _MARKDOWN_THEMATIC_BREAK_RE.fullmatch(value): + value = "" + else: + value = _MARKDOWN_BLOCK_PREFIX_RE.sub("", value) + base = " ".join(value.translate(_MARKDOWN_UNSAFE_SOURCE_NAME_CHARS).split()) or "source" + candidate = base + suffix = 2 + while candidate in existing: + candidate = f"{base} {suffix}" + suffix += 1 + return candidate + + def get_user_language() -> str: """Return the active language: env override if set, otherwise browser's Accept-Language.""" if DEFAULT_LANGUAGE: @@ -465,26 +551,48 @@ 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 + if not isinstance(metadata_sources, list) or not metadata_sources: + 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}") - url = s.get("url", "") + title = s.get("title", "") snippet = s.get("snippet", "") - content = f"**[{title}]({url})**\n\n{snippet}" - source_name = title - if source_name in d: - source_name = f"{title} ({i})" + title = title.strip() if isinstance(title, str) else "" + snippet = snippet.strip() if isinstance(snippet, str) else "" + url = normalize_web_url(s.get("url")) + if url is None: + continue + markdown_url = quote(url, safe=_MARKDOWN_URL_SAFE_CHARS) + + source_label = title or url + source_name = _safe_source_name(source_label, d) + content = f"**[{_escape_markdown_text(source_label)}]({markdown_url})**" + if snippet: + content += f"\n\n{_escape_markdown_text(snippet)}" d[source_name] = cl.Text(content=content, name=source_name, display="side") continue - 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() + ): + continue + + filename = Path(filename_value.strip()) + suffix = filename.suffix.lower() + 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 @@ -495,36 +603,51 @@ 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 "" + page_label = str(page).strip() if page is not None else "" + source_label = f"{filename}" + ( + f" (page: {page_label})" if suffix in [".pdf", ".pptx", ".docx", ".doc"] and page_label else "" ) + source_name = _safe_source_name(source_label, d) - 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 suffix: + case ".pdf": + elem = cl.Pdf( + name=source_name, + url=file_url, + page=int(page) if page_label else None, + 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, httpx.InvalidURL, TypeError, ValueError, AttributeError): + logger.warning("Skipping an unavailable source", source_index=i) + continue d[source_name] = elem - source_names = list(d.keys()) + source_names = list(d) elements = list(d.values()) return elements, source_names @@ -582,7 +705,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() diff --git a/openrag/core/models/query.py b/openrag/core/models/query.py index b5dad680d..86c512de5 100644 --- a/openrag/core/models/query.py +++ b/openrag/core/models/query.py @@ -103,6 +103,10 @@ class SearchQueries(BaseModel): """Collection of sub-queries produced by query decomposition.""" query_list: list[Query] = Field(..., description="Search sub-queries to retrieve relevant documents.") + requires_retrieval: bool = Field( + default=True, + description="Whether the user's request needs document retrieval.", + ) def __str__(self) -> str: return " --- ".join(str(q) for q in self.query_list) diff --git a/openrag/core/utils/source_filtering.py b/openrag/core/utils/source_filtering.py index e64418bda..ff7b52833 100644 --- a/openrag/core/utils/source_filtering.py +++ b/openrag/core/utils/source_filtering.py @@ -17,6 +17,15 @@ re.IGNORECASE, ) _SOURCES_NUMS_RE = re.compile(r"\n?[ \t]*\[?Sources?\]?\s*:\s*\[?([\d,\s]+)\]?[.\s]*?(?=\n|$)", re.IGNORECASE) +_INLINE_SOURCE_NUMS_RE = re.compile( + r"[ \t]*\[\s*Sources?\s+(\d+(?:\s*,\s*\d+)*)\s*\]", + re.IGNORECASE, +) +_UNCLOSED_SOURCE_NUMS_RE = re.compile( + r"[ \t]*\[\s*Sources?\s+(\d+(?:\s*,\s*\d+)*)\s*(?=\n|$)", + re.IGNORECASE, +) +_DANGLING_SOURCE_RE = re.compile(r"[ \t]*\[\s*Sources?\s*(?=\n|$)", re.IGNORECASE) def _sanitize_log_preview(text: str, max_length: int = 150) -> str: @@ -26,22 +35,37 @@ def _sanitize_log_preview(text: str, max_length: int = 150) -> str: return preview -def _strip_sources_tags(text: str) -> tuple[str, set[int], bool]: - """Strip line-terminal source tags and return citations found.""" +def _strip_sources_tags(text: str, *, include_inline_markers: bool = True) -> tuple[str, set[int], bool]: + """Strip source tags and return citations found.""" cited: set[int] = set() - for match in _SOURCES_NUMS_RE.finditer(text): - cited.update(int(n.strip()) for n in match.group(1).split(",") if n.strip().isdigit()) + patterns = [_SOURCES_NUMS_RE] + if include_inline_markers: + patterns.extend((_INLINE_SOURCE_NUMS_RE, _UNCLOSED_SOURCE_NUMS_RE)) + for pattern in patterns: + for match in pattern.finditer(text): + cited.update(int(n.strip()) for n in match.group(1).split(",") if n.strip().isdigit()) saw_none = bool(_SOURCES_NONE_RE.search(text)) cleaned = _SOURCES_NUMS_RE.sub("", text) cleaned = _SOURCES_NONE_RE.sub("", cleaned) + if include_inline_markers: + cleaned = _INLINE_SOURCE_NUMS_RE.sub("", cleaned) + cleaned = _UNCLOSED_SOURCE_NUMS_RE.sub("", cleaned) + cleaned = _DANGLING_SOURCE_RE.sub("", cleaned) return cleaned, cited, saw_none -def extract_and_strip_sources_block(text: str) -> tuple[str, set[int] | None]: - """Strip line-terminal source tags and return merged citations.""" - cleaned, citations, saw_none = _strip_sources_tags(text) +def extract_and_strip_sources_block( + text: str, + *, + include_inline_markers: bool = True, +) -> tuple[str, set[int] | None]: + """Strip source tags and return merged citations.""" + cleaned, citations, saw_none = _strip_sources_tags(text, include_inline_markers=include_inline_markers) if not citations and not saw_none: + if cleaned != text: + logger.debug("Removed incomplete source marker from LLM response") + return cleaned.rstrip(), None tail = text[-150:] if len(text) > 150 else text logger.debug("No [Sources: ...] tag found in LLM response", tail=repr(_sanitize_log_preview(tail))) return text, None @@ -55,14 +79,18 @@ def extract_and_strip_sources_block(text: str) -> tuple[str, set[int] | None]: return cleaned, set() -def filter_sources_by_citations(sources: list, citations: set[int] | None) -> list: +def filter_sources_by_citations( + sources: list, + citations: set[int] | None, + *, + allow_uncited: bool = False, +) -> list: """Keep only sources whose 1-based index was cited.""" if citations is None: - return sources + return sources if allow_uncited else [] if not citations: return [] - filtered = [source for i, source in enumerate(sources, start=1) if i in citations] - return filtered if filtered else sources + return [source for i, source in enumerate(sources, start=1) if i in citations] def _min_sources_tag_buffer_size(n_sources: int) -> int: @@ -83,8 +111,11 @@ async def stream_with_source_filtering( sources: list, model_name: str, buffer_size: int | None = None, + *, + allow_uncited_sources: bool = False, + citation_protocol_active: bool = True, ): - """Process an LLM SSE stream, stripping line-terminal source tags. + """Process an LLM SSE stream and, when active, strip source tags. The terminal flush (tail content + ``extra.sources``) runs exactly once after the loop on *every* termination path — a clean ``data: [DONE]``, the @@ -99,6 +130,7 @@ async def stream_with_source_filtering( """ if buffer_size is None: buffer_size = max(_MIN_STREAM_LOOKAHEAD, _min_sources_tag_buffer_size(len(sources))) + include_inline_markers = citation_protocol_active and bool(sources) pending = "" emitted_len = 0 chunk_template = None @@ -155,7 +187,13 @@ async def stream_with_source_filtering( if len(pending) <= buffer_size: continue - cleaned, _, _ = _strip_sources_tags(pending) + if citation_protocol_active: + cleaned, _, _ = _strip_sources_tags( + pending, + include_inline_markers=include_inline_markers, + ) + else: + cleaned = pending safe_end = max(0, len(cleaned) - buffer_size) if safe_end > emitted_len: out = { @@ -216,10 +254,16 @@ async def stream_with_source_filtering( logger.warning("Upstream stream raised before any content; surfacing error", error=str(stream_error)) raise stream_error - final_clean, citations = extract_and_strip_sources_block(pending) - final_clean = final_clean.rstrip() + if citation_protocol_active: + final_clean, citations = extract_and_strip_sources_block( + pending, + include_inline_markers=include_inline_markers, + ) + final_clean = final_clean.rstrip() + else: + final_clean, citations = pending, None - filtered = filter_sources_by_citations(sources, citations) + filtered = filter_sources_by_citations(sources, citations, allow_uncited=allow_uncited_sources) extra_payload = {"sources": filtered} if not saw_done: extra_payload["truncated"] = True diff --git a/openrag/core/utils/web_url.py b/openrag/core/utils/web_url.py new file mode 100644 index 000000000..8662e8706 --- /dev/null +++ b/openrag/core/utils/web_url.py @@ -0,0 +1,25 @@ +"""Shared validation and normalization for displayable web URLs.""" + +from urllib.parse import urlparse + +import httpx +from core.utils.text import sanitize_text + + +def normalize_web_url(value: object) -> str | None: + """Return a renderable HTTP(S) URL, or None when the value is invalid.""" + if not isinstance(value, str): + return None + url = sanitize_text(value) + if not url: + return None + try: + parsed = urlparse(url) + if parsed.scheme not in {"http", "https"} or not parsed.netloc: + return None + return str(httpx.URL(url)) + except (ValueError, httpx.InvalidURL): + return None + + +__all__ = ["normalize_web_url"] diff --git a/openrag/prompts/templates/query_contextualizer_tmpl.txt b/openrag/prompts/templates/query_contextualizer_tmpl.txt index 67a134d3a..a2458377d 100644 --- a/openrag/prompts/templates/query_contextualizer_tmpl.txt +++ b/openrag/prompts/templates/query_contextualizer_tmpl.txt @@ -1,17 +1,18 @@ Produce a JSON object listing sub-queries derived from the user's last message. Output shape (return this JSON object, nothing else): -{{"query_list": [ {{"query": "", "temporal_filters": }} ]}} +{{"requires_retrieval": , "query_list": [ {{"query": "", "temporal_filters": }} ]}} Current date: {current_date} Language for `query` field: {query_language} Timestamps: UTC (`+00:00`). Week starts Monday. # Rewrite rules (`query` field) +- Set `requires_retrieval: false` and return an empty `query_list` when the entire last user message is only a greeting, thanks, casual conversation, or a question about the assistant's general capabilities. +- Set `requires_retrieval: true` for any request that asks for factual or document-backed information. A greeting combined with a factual question still requires retrieval. - Rewrite the last `user:` line as one standalone descriptive sentence. - Use earlier turns only to resolve pronouns or add context directly relevant to the query; do not inject unrelated history. - For independent questions: minimal changes (grammar, missing keywords). -- Greetings / thanks: copy verbatim, `temporal_filters: null`. - Do not answer. Only reformulate. # Sub-queries — when to split @@ -46,19 +47,25 @@ For exclusions, split into two sub-queries covering each remaining range. Never # Examples (Current date = Wednesday, April 15, 2026) User: "Summary of meeting notes uploaded in the past month" -{{"query_list":[{{"query":"Summary of meeting notes uploaded in the past month","temporal_filters":[{{"field":"created_at","operator":">=","value":"2026-03-15T00:00:00+00:00"}},{{"field":"created_at","operator":"<","value":"2026-04-16T00:00:00+00:00"}}]}}]}} +{{"requires_retrieval":true,"query_list":[{{"query":"Summary of meeting notes uploaded in the past month","temporal_filters":[{{"field":"created_at","operator":">=","value":"2026-03-15T00:00:00+00:00"}},{{"field":"created_at","operator":"<","value":"2026-04-16T00:00:00+00:00"}}]}}]}} User: "Sales figures for Product A and Product B" -{{"query_list":[{{"query":"Sales figures for Product A","temporal_filters":null}},{{"query":"Sales figures for Product B","temporal_filters":null}}]}} +{{"requires_retrieval":true,"query_list":[{{"query":"Sales figures for Product A","temporal_filters":null}},{{"query":"Sales figures for Product B","temporal_filters":null}}]}} User: "Documents from last year except March" -{{"query_list":[{{"query":"Documents from January or February 2025","temporal_filters":[{{"field":"created_at","operator":">=","value":"2025-01-01T00:00:00+00:00"}},{{"field":"created_at","operator":"<","value":"2025-03-01T00:00:00+00:00"}}]}},{{"query":"Documents from April to December 2025","temporal_filters":[{{"field":"created_at","operator":">=","value":"2025-04-01T00:00:00+00:00"}},{{"field":"created_at","operator":"<","value":"2026-01-01T00:00:00+00:00"}}]}}]}} +{{"requires_retrieval":true,"query_list":[{{"query":"Documents from January or February 2025","temporal_filters":[{{"field":"created_at","operator":">=","value":"2025-01-01T00:00:00+00:00"}},{{"field":"created_at","operator":"<","value":"2025-03-01T00:00:00+00:00"}}]}},{{"query":"Documents from April to December 2025","temporal_filters":[{{"field":"created_at","operator":">=","value":"2025-04-01T00:00:00+00:00"}},{{"field":"created_at","operator":"<","value":"2026-01-01T00:00:00+00:00"}}]}}]}} User: "Q3 2024 reporting template" -{{"query_list":[{{"query":"Q3 2024 reporting template","temporal_filters":null}}]}} +{{"requires_retrieval":true,"query_list":[{{"query":"Q3 2024 reporting template","temporal_filters":null}}]}} User: "Evolution of the Department of Justice budget between 2020 and 2022" -{{"query_list":[{{"query":"Department of Justice budget in 2020","temporal_filters":null}},{{"query":"Department of Justice budget in 2021","temporal_filters":null}},{{"query":"Department of Justice budget in 2022","temporal_filters":null}}]}} +{{"requires_retrieval":true,"query_list":[{{"query":"Department of Justice budget in 2020","temporal_filters":null}},{{"query":"Department of Justice budget in 2021","temporal_filters":null}},{{"query":"Department of Justice budget in 2022","temporal_filters":null}}]}} User: "Latest safety bulletins" -{{"query_list":[{{"query":"Latest safety bulletins","temporal_filters":[{{"field":"created_at","operator":">=","value":"2026-01-15T00:00:00+00:00"}},{{"field":"created_at","operator":"<","value":"2026-04-16T00:00:00+00:00"}}]}}]}} \ No newline at end of file +{{"requires_retrieval":true,"query_list":[{{"query":"Latest safety bulletins","temporal_filters":[{{"field":"created_at","operator":">=","value":"2026-01-15T00:00:00+00:00"}},{{"field":"created_at","operator":"<","value":"2026-04-16T00:00:00+00:00"}}]}}]}} + +User: "How can you help me?" +{{"requires_retrieval":false,"query_list":[]}} + +User: "Hello, what was Product A revenue in Q1?" +{{"requires_retrieval":true,"query_list":[{{"query":"Product A revenue in Q1","temporal_filters":null}}]}} diff --git a/openrag/prompts/templates/spoken_style_answer_tmpl.txt b/openrag/prompts/templates/spoken_style_answer_tmpl.txt index a2fe6c311..32534293d 100644 --- a/openrag/prompts/templates/spoken_style_answer_tmpl.txt +++ b/openrag/prompts/templates/spoken_style_answer_tmpl.txt @@ -1,4 +1,4 @@ -You are an AI assistant designed for **spoken, conversational answers**. +You are **OpenRAG**, a retrieval-augmented generation system built by **LINAGORA**, designed for spoken, conversational answers. Your goal is to give short (1-2 sentences), clear, and accurate explanations, based only on the retrieved documents in `Context`. # Context @@ -9,10 +9,15 @@ Your goal is to give short (1-2 sentences), clear, and accurate explanations, ba 1. Use only the provided Context * Answer strictly from the information in `Context`. * Do not guess, infer, or use outside knowledge. + * For greetings, thanks, casual conversation, identity questions, or questions about your capabilities, answer briefly without using the Context and end with `[Sources: none]`. + * For identity or capability questions, explain that you are OpenRAG, built by LINAGORA, and that you retrieve and synthesize information from indexed documents with supporting sources. + * Do not present yourself as a general-purpose assistant or list unrelated abilities. + * If a message also asks a factual question, answer that part from the Context as usual. * If the Context lacks enough information, say so briefly and ask the user for more details. 2. Citations * Never place citations, source numbers, or references **inside** your answer text: citation is only at the end + * Do not copy the Context markers such as `[Source 1]` into the answer body. Use source numbers only in the final `[Sources: ...]` line. * Cite sources **only once**, on a **single line** that is the **very last line** of your response, separated from the body by a blank line. * The final line MUST match **exactly one** of these two formats (no other text on that line): - `[Sources: 1, 3]` — when one or more numbered sources from the Context contributed to your answer (comma-separated, ascending order, no duplicates) diff --git a/openrag/prompts/templates/sys_prompt_tmpl.txt b/openrag/prompts/templates/sys_prompt_tmpl.txt index 5e3598e5e..e82a9b1f9 100644 --- a/openrag/prompts/templates/sys_prompt_tmpl.txt +++ b/openrag/prompts/templates/sys_prompt_tmpl.txt @@ -1,4 +1,4 @@ -You are an AI conversational assistant specialized in **information retrieval and synthesis**. +You are **OpenRAG**, a retrieval-augmented generation system built by **LINAGORA**. Your goal is to provide **precise, reliable, and well-structured answers** using **only the retrieved documents** (`Context`). Prioritize **clarity, accuracy, and completeness** in your responses. @@ -10,11 +10,17 @@ Prioritize **clarity, accuracy, and completeness** in your responses. 1. Use only the provided Context * Base your answer **exclusively** on the information contained in the `Context`. * **Never infer**, assume, or rely on any external knowledge. + * For greetings, thanks, casual conversation, identity questions, or questions about your capabilities, answer briefly without using the Context and end with `[Sources: none]`. + * For identity or capability questions, explain that you are OpenRAG, built by LINAGORA, and that you search, retrieve, and synthesize information from indexed documents while citing the supporting sources. + * Present yourself as a document-grounded RAG system, not as a general-purpose assistant. Do not advertise unrelated abilities such as general knowledge, travel advice, creative writing, or coding unless the user asks about indexed documents covering those topics. + * Keep conversational, identity, and capability answers concise: normally 1-3 sentences. + * If a message combines conversation with a factual question, answer the factual part from the Context as usual. * If the context is **insufficient**, **invite the user** to clarify their query or provide additional keywords. * **Always answer with at least one sentence of text.** Your response body must **never be empty**, even when no source is relevant — in that case briefly explain (in the user's language) that the documents do not cover the question, then end with `[Sources: none]`. 2. Citations * Never place citations, source numbers, or references **inside** your answer text: citation is only at the end + * Do not copy the Context markers such as `[Source 1]` into the answer body. Use source numbers only in the final `[Sources: ...]` line. * Cite sources **only once**, on a **single line** that is the **very last line** of your response, separated from the body by a blank line. * The final line MUST match **exactly one** of these two formats (no other text on that line): - `[Sources: 1, 3, 5]` — when one or more numbered sources from the Context contributed to your answer (comma-separated, ascending order, no duplicates) @@ -27,4 +33,4 @@ Prioritize **clarity, accuracy, and completeness** in your responses. * Use **headings**, **bullet points**, **numbered lists**, or **tables** to organize information clearly. * Ensure responses are **concise yet complete**, avoiding omission of key details. -Here are the retrieved documents: `{context}` \ No newline at end of file +Here are the retrieved documents: `{context}` diff --git a/openrag/services/orchestrators/query_service.py b/openrag/services/orchestrators/query_service.py index 17054fc52..152cebf7e 100644 --- a/openrag/services/orchestrators/query_service.py +++ b/openrag/services/orchestrators/query_service.py @@ -49,6 +49,7 @@ format_context, format_web_context, load_template_by_key, + prepend_system_prompt, ) from core.utils.exceptions import WorkspaceNotFoundError from core.utils.logging import get_logger @@ -58,6 +59,7 @@ stream_with_source_filtering, ) from core.utils.text import get_num_tokens +from core.utils.web_url import normalize_web_url from services.inference.runtime import detect_language, get_llm_semaphore if TYPE_CHECKING: @@ -88,8 +90,10 @@ {query}""" _QUERY_JSON_HINT = ( - "\n\nRespond ONLY with a JSON object of the form " - '{"query_list": [{"query": "", "temporal_filters": null}]}.' + "\n\nRespond ONLY with one of these JSON forms: " + '{"requires_retrieval": false, "query_list": []} when retrieval is not needed, or ' + '{"requires_retrieval": true, ' + '"query_list": [{"query": "", "temporal_filters": null}]} when it is.' ) @@ -412,6 +416,19 @@ async def _prepare_chat(self, partition: list[str] | None, payload: dict, llm: L partition = [scope.partition] filter_params = {"file_id": scope.file_ids} + force_retrieval = use_websearch or use_map_reduce + if not queries.query_list: + if not queries.requires_retrieval and not force_retrieval: + tmpl = self._spoken_style_answer_prompt if spoken_style else self._sys_prompt_tmplt + payload["messages"] = prepend_system_prompt( + messages, + tmpl, + context="", + current_date=datetime.now().strftime("%A, %B %d, %Y, %H:%M:%S"), + ) + return payload, [], [], True + queries = SearchQueries(query_list=[Query(query=messages[-1]["content"])]) + web_results: list = [] if partition is not None and use_websearch: chunks, web_lists = await self._gather_rag_and_web(queries, partition, top_k, filter_params) @@ -425,18 +442,19 @@ async def _prepare_chat(self, partition: list[str] | None, payload: dict, llm: L chunks = [] if not chunks and not web_results and partition is None: - return payload, [], [] + return payload, [], [], False docs = [c.to_langchain() for c in chunks] if use_map_reduce and docs: docs = await self._map_reduce(" ".join(q.query for q in queries.query_list), docs) - web_formatted, web_tokens = "", 0 + web_formatted, web_source_numbers, web_tokens = "", [], 0 + web_start_index = 1 if web_results: - web_formatted, _, web_tokens = format_web_context( + web_formatted, web_source_numbers, web_tokens = format_web_context( web_results, length_function=get_num_tokens(), - start_index=1, + start_index=web_start_index, max_tokens=self._web.max_tokens, ) context, included = format_context( @@ -448,15 +466,17 @@ async def _prepare_chat(self, partition: list[str] | None, payload: dict, llm: L if web_results: if docs: - web_formatted, _, _ = format_web_context( + web_start_index = len(docs) + 1 + web_formatted, web_source_numbers, _ = format_web_context( web_results, length_function=get_num_tokens(), - start_index=len(docs) + 1, + start_index=web_start_index, max_tokens=self._web.max_tokens, ) else: context = "" context = f"{context}{SOURCE_SEPARATOR}{web_formatted}" if context else web_formatted + web_results = [web_results[number - web_start_index] for number in web_source_numbers] new_messages = copy.deepcopy(messages) tmpl = self._spoken_style_answer_prompt if spoken_style else self._sys_prompt_tmplt @@ -470,7 +490,7 @@ async def _prepare_chat(self, partition: list[str] | None, payload: dict, llm: L }, ) payload["messages"] = new_messages - return payload, docs, web_results + return payload, docs, web_results, True async def _gather_rag_and_web(self, queries, partition, top_k, filter_params): # Fuse the doc branch through retrieve_multi so a partition's rrf_k drives @@ -488,20 +508,30 @@ async def _gather_rag_and_web(self, queries, partition, top_k, filter_params): async def _prepare_completions(self, partition: list[str], payload: dict, llm: LLM | None = None): prompt = payload["prompt"] queries = await self.generate_query([{"role": "user", "content": prompt}], llm=llm) - chunks = await self._retrieval.retrieve_multi(partitions=partition, search_queries=queries) - docs = [c.to_langchain() for c in chunks] - context, included = format_context( - [doc.page_content for doc in docs], - max_context_tokens=self._max_context_tokens, - length_function=get_num_tokens(), - ) - docs = [docs[i] for i in included] - if docs: - payload["prompt"] = ( - f"Given the content\n{context}\nComplete the following prompt: {prompt}\n" - "At the very end of your response, on a new line, list which source numbers " - "you used: [Sources: 1, 3]" + if not queries.query_list: + if not queries.requires_retrieval: + docs, context = [], "" + else: + queries = SearchQueries(query_list=[Query(query=prompt)]) + if queries.query_list: + chunks = await self._retrieval.retrieve_multi(partitions=partition, search_queries=queries) + docs = [c.to_langchain() for c in chunks] + context, included = format_context( + [doc.page_content for doc in docs], + max_context_tokens=self._max_context_tokens, + length_function=get_num_tokens(), ) + docs = [docs[i] for i in included] + + metadata = payload.get("metadata") or {} + tmpl = ( + self._spoken_style_answer_prompt if metadata.get("spoken_style_answer", False) else self._sys_prompt_tmplt + ) + instructions = tmpl.format( + context=context, + current_date=datetime.now().strftime("%A, %B %d, %Y, %H:%M:%S"), + ) + payload["prompt"] = f"{instructions}\n\n# User request\n{prompt}" return payload, docs # ------------------------------------------------------------------ @@ -568,19 +598,35 @@ async def chat( """Non-streaming chat completion → finalized OpenAI dict.""" metadata = payload.get("metadata") or {} llm = self._resolve_llm(partitions) + citation_protocol_active = False if partitions is None and not metadata.get("websearch", False): docs, web_results = [], [] else: - payload, docs, web_results = await self._prepare_chat(partitions, payload, llm) + payload, docs, web_results, citation_protocol_active = await self._prepare_chat(partitions, payload, llm) sources = prepare_sources(docs, web_results) + structured_output = _allows_uncited_sources(payload) payload["messages"] = self._sanitize_messages(payload["messages"]) chunk = await llm.chat(payload["messages"], **_sampling(payload)) chunk["model"] = model_name content = chunk.get("choices", [{}])[0].get("message", {}).get("content", "") or "" - clean, citations = extract_and_strip_sources_block(content) + if citation_protocol_active and not structured_output: + clean, citations = extract_and_strip_sources_block( + content, + include_inline_markers=bool(sources), + ) + else: + clean, citations = content, None chunk["choices"][0]["message"]["content"] = clean - chunk["extra"] = json.dumps({"sources": filter_sources_by_citations(sources, citations)}) + chunk["extra"] = json.dumps( + { + "sources": filter_sources_by_citations( + sources, + citations, + allow_uncited=structured_output, + ) + } + ) return chunk async def chat_stream( @@ -594,15 +640,23 @@ async def chat_stream( """Streaming chat completion → SSE strings with filtered sources.""" metadata = payload.get("metadata") or {} llm = self._resolve_llm(partitions) + citation_protocol_active = False if partitions is None and not metadata.get("websearch", False): docs, web_results = [], [] else: - payload, docs, web_results = await self._prepare_chat(partitions, payload, llm) + payload, docs, web_results, citation_protocol_active = await self._prepare_chat(partitions, payload, llm) sources = prepare_sources(docs, web_results) + structured_output = _allows_uncited_sources(payload) payload["messages"] = self._sanitize_messages(payload["messages"]) llm_stream = llm.stream_chat(payload["messages"], **_sampling(payload)) - async for sse_line in stream_with_source_filtering(llm_stream, sources, model_name): + async for sse_line in stream_with_source_filtering( + llm_stream, + sources, + model_name, + allow_uncited_sources=structured_output, + citation_protocol_active=citation_protocol_active and not structured_output, + ): yield sse_line async def complete( @@ -614,17 +668,33 @@ async def complete( ) -> dict: """Non-streaming text completion → finalized OpenAI dict.""" llm = self._resolve_llm(partitions) + citation_protocol_active = partitions is not None if partitions is None: docs = [] else: payload, docs = await self._prepare_completions(partitions, payload, llm) sources = prepare_sources(docs, []) + structured_output = _allows_uncited_sources(payload) resp = await llm.generate(payload["prompt"], **_sampling(payload, key="prompt")) text = resp.get("choices", [{}])[0].get("text", "") or "" - clean, citations = extract_and_strip_sources_block(text) + if citation_protocol_active and not structured_output: + clean, citations = extract_and_strip_sources_block( + text, + include_inline_markers=bool(sources), + ) + else: + clean, citations = text, None resp["choices"][0]["text"] = clean - resp["extra"] = json.dumps({"sources": filter_sources_by_citations(sources, citations)}) + resp["extra"] = json.dumps( + { + "sources": filter_sources_by_citations( + sources, + citations, + allow_uncited=structured_output, + ) + } + ) return resp @@ -649,8 +719,10 @@ def _dedupe_web(web_lists: list[list]) -> list: seen: set[str] = set() out: list = [] for r in (r for lst in web_lists for r in lst): - if r.url not in seen: - seen.add(r.url) + url = normalize_web_url(r.url) + if url is not None and url not in seen: + r.url = url + seen.add(url) out.append(r) return out @@ -666,4 +738,10 @@ def _sampling(payload: dict, key: str = "messages") -> dict: return {k: v for k, v in payload.items() if k not in drop} +def _allows_uncited_sources(payload: dict) -> bool: + """Structured output cannot carry the plain-text citation marker.""" + response_format = payload.get("response_format") + return isinstance(response_format, dict) and response_format.get("type") in {"json_object", "json_schema"} + + __all__ = ["QueryService", "RAGMODE"] diff --git a/tests/unit/api/routers/user/test_source_links.py b/tests/unit/api/routers/user/test_source_links.py index ed20af42f..9127837c1 100644 --- a/tests/unit/api/routers/user/test_source_links.py +++ b/tests/unit/api/routers/user/test_source_links.py @@ -61,3 +61,25 @@ def test_metadata_is_passed_through(): link = _build({"_id": "x", "source": "doc.pdf", "author": "alice"}) assert link["author"] == "alice" assert link["chunk_url"] == "https://host/extract/x" + + +def test_metadata_cannot_override_authoritative_source_fields(): + link = _build( + { + "_id": "42", + "source": "diagram.png", + "file_url": "https://attacker.example/file", + "chunk_url": "https://attacker.example/chunk", + "source_type": "web", + } + ) + + assert link["source_type"] == "document" + assert link["file_url"] == "https://host/static/42" + assert link["chunk_url"] == "https://host/extract/42" + + +def test_metadata_file_url_is_removed_when_source_is_missing(): + link = _build({"_id": "42", "file_url": "https://attacker.example/file"}) + + assert "file_url" not in link diff --git a/tests/unit/core/utils/test_source_filtering.py b/tests/unit/core/utils/test_source_filtering.py index 629c0c3eb..877587542 100644 --- a/tests/unit/core/utils/test_source_filtering.py +++ b/tests/unit/core/utils/test_source_filtering.py @@ -135,6 +135,24 @@ def test_tag_inline_in_prose_preserved(self): assert clean == text assert citations is None + def test_context_source_markers_are_recovered_and_stripped(self): + text = "The footprint fell by 28% [Source 7].\nThe partners include Flexis [Source 8][Source 9]." + clean, citations = extract_and_strip_sources_block(text) + assert clean == "The footprint fell by 28%.\nThe partners include Flexis." + assert citations == {7, 8, 9} + + def test_unclosed_numbered_source_marker_is_recovered(self): + text = "The target is 2040 [Source 2" + clean, citations = extract_and_strip_sources_block(text) + assert clean == "The target is 2040" + assert citations == {2} + + def test_dangling_source_marker_is_removed_without_a_citation(self): + text = "Logistics emissions fell by 30% [Source" + clean, citations = extract_and_strip_sources_block(text) + assert clean == "Logistics emissions fell by 30%" + assert citations is None + class TestFilterSourcesByCitations: def test_basic_filtering(self): @@ -142,9 +160,14 @@ def test_basic_filtering(self): result = filter_sources_by_citations(sources, {1, 3, 5}) assert result == ["a", "c", "e"] - def test_none_citations_returns_all(self): + def test_none_citations_returns_empty(self): sources = ["a", "b", "c"] result = filter_sources_by_citations(sources, None) + assert result == [] + + def test_none_citations_can_be_allowed_for_structured_output(self): + sources = ["a", "b", "c"] + result = filter_sources_by_citations(sources, None, allow_uncited=True) assert result == ["a", "b", "c"] def test_empty_citations_returns_empty(self): @@ -152,10 +175,10 @@ def test_empty_citations_returns_empty(self): result = filter_sources_by_citations(sources, set()) assert result == [] - def test_out_of_range_citations_fallback(self): + def test_out_of_range_citations_returns_empty(self): sources = ["a", "b", "c"] result = filter_sources_by_citations(sources, {99}) - assert result == ["a", "b", "c"] + assert result == [] def test_partial_out_of_range(self): sources = ["a", "b", "c"] @@ -370,8 +393,8 @@ async def test_case2_llm_says_sources_none(self): assert _parse_finish_sources(result) == [] @pytest.mark.asyncio - async def test_case3_llm_no_tag_fallback_all(self): - """Case 3: LLM omits tag entirely → fallback to all sources.""" + async def test_case3_llm_no_tag_returns_no_sources(self): + """Case 3: LLM omits tag entirely → no source is attributed.""" lines = [ _make_chunk("Answer without any sources tag."), _make_finish(), @@ -379,8 +402,64 @@ async def test_case3_llm_no_tag_fallback_all(self): ] result = await _collect(stream_with_source_filtering(_fake_stream(lines), self.SOURCES, "test-model")) assert _collect_content(result) == "Answer without any sources tag." + assert _parse_finish_sources(result) == [] + + @pytest.mark.asyncio + async def test_no_tag_keeps_sources_when_uncited_output_is_allowed(self): + lines = [ + _make_chunk('{"answer": "structured"}'), + _make_finish(), + DONE_LINE, + ] + result = await _collect( + stream_with_source_filtering( + _fake_stream(lines), + self.SOURCES, + "test-model", + allow_uncited_sources=True, + ) + ) assert _parse_finish_sources(result) == self.SOURCES + @pytest.mark.asyncio + async def test_structured_output_preserves_source_like_json_values(self): + structured = '{"answer":"Use [Source 1]","literal_format":"[Sources: 1]"}' + lines = [ + _make_chunk(structured), + _make_finish(), + DONE_LINE, + ] + result = await _collect( + stream_with_source_filtering( + _fake_stream(lines), + self.SOURCES, + "test-model", + allow_uncited_sources=True, + citation_protocol_active=False, + ) + ) + assert _collect_content(result) == structured + assert _parse_finish_sources(result) == self.SOURCES + + @pytest.mark.asyncio + async def test_direct_output_preserves_terminal_source_marker(self): + answer = "The requested literal notation is:\n[Sources: 1]" + lines = [ + _make_chunk(answer), + _make_finish(), + DONE_LINE, + ] + result = await _collect( + stream_with_source_filtering( + _fake_stream(lines), + [], + "test-model", + citation_protocol_active=False, + ) + ) + assert _collect_content(result) == answer + assert _parse_finish_sources(result) == [] + @pytest.mark.asyncio async def test_multiple_inline_tags_stripped_from_stream(self): """Bullet-leak: LLM emits [Sources: X] per bullet. All inline tags must be stripped.""" @@ -399,6 +478,30 @@ async def test_multiple_inline_tags_stripped_from_stream(self): assert "Claim two about APEX." in content assert _parse_finish_sources(result) == [{"file": "a.pdf"}, {"file": "c.pdf"}] + @pytest.mark.asyncio + async def test_context_source_markers_are_stripped_and_rendered_as_sources(self): + lines = [ + _make_chunk("First claim [Sour"), + _make_chunk("ce 1]. Second claim [Source 2][Source 3]."), + _make_finish(), + DONE_LINE, + ] + result = await _collect(stream_with_source_filtering(_fake_stream(lines), self.SOURCES, "test-model")) + assert _collect_content(result) == "First claim. Second claim." + assert _parse_finish_sources(result) == self.SOURCES + + @pytest.mark.asyncio + async def test_literal_source_marker_is_preserved_without_sources(self): + lines = [ + _make_chunk("The literal notation [Sour"), + _make_chunk("ce 1] identifies the first source."), + _make_finish(), + DONE_LINE, + ] + result = await _collect(stream_with_source_filtering(_fake_stream(lines), [], "test-model")) + assert _collect_content(result) == "The literal notation [Source 1] identifies the first source." + assert _parse_finish_sources(result) == [] + @pytest.mark.asyncio async def test_inline_prose_tag_preserved_in_stream(self): """Meta-discussion: a [Sources: 1, 3] inside a sentence must NOT be stripped.""" @@ -411,8 +514,8 @@ async def test_inline_prose_tag_preserved_in_stream(self): result = await _collect(stream_with_source_filtering(_fake_stream(lines), self.SOURCES, "test-model")) content = _collect_content(result) assert content == "Use the format [Sources: 1, 3] at the very end of your response." - # No line-terminal tag → fallback to all sources - assert _parse_finish_sources(result) == self.SOURCES + # No line-terminal tag means no source was actually cited. + assert _parse_finish_sources(result) == [] @pytest.mark.asyncio async def test_mid_response_tag_stripped_plus_trailing_tag(self): diff --git a/tests/unit/core/utils/test_web_url.py b/tests/unit/core/utils/test_web_url.py new file mode 100644 index 000000000..f87357128 --- /dev/null +++ b/tests/unit/core/utils/test_web_url.py @@ -0,0 +1,21 @@ +import pytest +from core.utils.web_url import normalize_web_url + + +@pytest.mark.parametrize( + "value", + [ + None, + 42, + "", + "javascript:alert(1)", + "https://", + "http://[::1", + ], +) +def test_normalize_web_url_rejects_unrenderable_values(value): + assert normalize_web_url(value) is None + + +def test_normalize_web_url_returns_canonical_http_url(): + assert normalize_web_url(" https://example.com/a path ") == "https://example.com/a%20path" diff --git a/tests/unit/services/orchestrators/test_query_service.py b/tests/unit/services/orchestrators/test_query_service.py index 62fa1b27d..742e8e5f9 100644 --- a/tests/unit/services/orchestrators/test_query_service.py +++ b/tests/unit/services/orchestrators/test_query_service.py @@ -42,6 +42,7 @@ def __init__(self, *, chat_responses=None, gen_text="answer", stream_lines=None) self._gen_text = gen_text self._stream_lines = stream_lines or ['data: {"choices":[{"delta":{"content":"hi"}}]}\n\n', "data: [DONE]\n\n"] self.chat_calls: list = [] + self.generate_calls: list = [] async def chat(self, messages, **kwargs): self.chat_calls.append((messages, kwargs)) @@ -52,6 +53,7 @@ async def chat(self, messages, **kwargs): return {"choices": [{"message": {"content": content}}]} async def generate(self, prompt, **kwargs): + self.generate_calls.append((prompt, kwargs)) return {"choices": [{"text": self._gen_text}]} async def stream_chat(self, messages, **kwargs): @@ -83,8 +85,10 @@ class FakeWeb: def __init__(self, results=None): self._results = results or [] + self.calls: list[str] = [] async def search(self, query): + self.calls.append(query) return list(self._results) @@ -346,6 +350,16 @@ async def test_generate_query_chatbotrag_parses_json(): svc = _svc(llm=FakeLLM(chat_responses=[payload]), mode="ChatBotRag") sq = await svc.generate_query([{"role": "user", "content": "hi"}]) assert sq.query_list[0].query == "rewritten" + assert sq.requires_retrieval is True + + +@pytest.mark.asyncio +async def test_generate_query_chatbotrag_can_skip_retrieval(): + payload = json.dumps({"requires_retrieval": False, "query_list": []}) + svc = _svc(llm=FakeLLM(chat_responses=[payload]), mode="ChatBotRag") + sq = await svc.generate_query([{"role": "user", "content": "How can you help me?"}]) + assert sq.requires_retrieval is False + assert sq.query_list == [] @pytest.mark.asyncio @@ -379,8 +393,40 @@ async def _spy(**kwargs): ) assert called["n"] == 0 # no retrieval in direct mode assert out["model"] == "m1" - assert out["choices"][0]["message"]["content"] == "hello" # sources tag stripped - assert json.loads(out["extra"])["sources"] == [] # [Sources: none] → no sources + assert out["choices"][0]["message"]["content"] == "hello [Sources: none]" + assert json.loads(out["extra"])["sources"] == [] + + +@pytest.mark.asyncio +async def test_chat_direct_mode_preserves_literal_source_marker(): + answer = "The literal notation [Source 1] identifies the first source." + svc = _svc(llm=FakeLLM(chat_responses=[answer])) + + out = await svc.chat( + partitions=None, + payload={"messages": [{"role": "user", "content": "Explain [Source 1]"}], "metadata": {}}, + prepare_sources=lambda d, w: [], + model_name="m1", + ) + + assert out["choices"][0]["message"]["content"] == answer + assert json.loads(out["extra"])["sources"] == [] + + +@pytest.mark.asyncio +async def test_chat_direct_mode_preserves_literal_terminal_sources_marker(): + answer = "The requested literal notation is:\n[Sources: 1]" + svc = _svc(llm=FakeLLM(chat_responses=[answer])) + + out = await svc.chat( + partitions=None, + payload={"messages": [{"role": "user", "content": "Repeat [Sources: 1]"}], "metadata": {}}, + prepare_sources=lambda d, w: [], + model_name="m1", + ) + + assert out["choices"][0]["message"]["content"] == answer + assert json.loads(out["extra"])["sources"] == [] @pytest.mark.asyncio @@ -397,6 +443,268 @@ async def test_chat_with_partition_retrieves_and_filters_sources(): assert filtered == [{"source_type": "document", "n": 1}] # only cited source 1 +@pytest.mark.asyncio +async def test_chat_recovers_context_markers_as_citations(): + svc = _svc(llm=FakeLLM(chat_responses=["First claim [Source 2]. Second claim [Source 1][Source 2]."])) + sources = [{"source_type": "document", "n": 1}, {"source_type": "document", "n": 2}] + out = await svc.chat( + partitions=["p"], + payload={"messages": [{"role": "user", "content": "q"}], "metadata": {}}, + prepare_sources=lambda d, w: sources, + model_name="m", + ) + + assert out["choices"][0]["message"]["content"] == "First claim. Second claim." + assert json.loads(out["extra"])["sources"] == sources + + +@pytest.mark.asyncio +async def test_chat_conversational_request_skips_partition_retrieval(): + query_json = json.dumps({"requires_retrieval": False, "query_list": []}) + llm = FakeLLM(chat_responses=[query_json, "I can help you search and summarize documents."]) + retrieval = FakeRetrieval() + svc = _svc(mode="ChatBotRag", llm=llm, retrieval=retrieval) + + out = await svc.chat( + partitions=["p"], + payload={"messages": [{"role": "user", "content": "How can you help me?"}], "metadata": {}}, + prepare_sources=lambda d, w: [{"source_type": "document"}], + model_name="m", + ) + + assert retrieval.retrieve_multi_calls == [] + assert out["choices"][0]["message"]["content"] == "I can help you search and summarize documents." + assert json.loads(out["extra"])["sources"] == [] + answer_messages = llm.chat_calls[1][0] + assert answer_messages[0]["role"] == "system" + assert "OpenRAG" in answer_messages[0]["content"] + assert "LINAGORA" in answer_messages[0]["content"] + assert "document-grounded RAG system" in answer_messages[0]["content"] + + +@pytest.mark.asyncio +async def test_chat_conversational_request_keeps_spoken_style_prompt(): + query_json = json.dumps({"requires_retrieval": False, "query_list": []}) + llm = FakeLLM(chat_responses=[query_json, "I'm OpenRAG, built by LINAGORA."]) + svc = _svc(mode="ChatBotRag", llm=llm) + + await svc.chat( + partitions=["p"], + payload={ + "messages": [{"role": "user", "content": "Who are you?"}], + "metadata": {"spoken_style_answer": True}, + }, + prepare_sources=lambda d, w: [], + model_name="m", + ) + + answer_system_prompt = llm.chat_calls[1][0][0]["content"] + assert "OpenRAG" in answer_system_prompt + assert "LINAGORA" in answer_system_prompt + assert "short (1-2 sentences)" in answer_system_prompt + + +@pytest.mark.asyncio +async def test_chat_mixed_request_still_retrieves_documents(): + query_json = json.dumps( + { + "requires_retrieval": True, + "query_list": [{"query": "Product A revenue in Q1", "temporal_filters": None}], + } + ) + llm = FakeLLM(chat_responses=[query_json, "Revenue was 10 million. [Sources: 1]"]) + retrieval = FakeRetrieval() + svc = _svc(mode="ChatBotRag", llm=llm, retrieval=retrieval) + + out = await svc.chat( + partitions=["p"], + payload={"messages": [{"role": "user", "content": "Hello, what was Product A revenue in Q1?"}]}, + prepare_sources=lambda d, w: [{"source_type": "document", "filename": "report.pdf"}], + model_name="m", + ) + + assert len(retrieval.retrieve_multi_calls) == 1 + assert json.loads(out["extra"])["sources"] == [{"source_type": "document", "filename": "report.pdf"}] + + +@pytest.mark.asyncio +async def test_chat_inconsistent_classifier_result_prefers_supplied_query(): + query_json = json.dumps( + { + "requires_retrieval": False, + "query_list": [{"query": "Product A revenue", "temporal_filters": None}], + } + ) + llm = FakeLLM(chat_responses=[query_json, "Revenue was 10 million. [Sources: 1]"]) + retrieval = FakeRetrieval() + svc = _svc(mode="ChatBotRag", llm=llm, retrieval=retrieval) + + await svc.chat( + partitions=["p"], + payload={"messages": [{"role": "user", "content": "What was Product A revenue?"}]}, + prepare_sources=lambda d, w: [{"source_type": "document"}], + model_name="m", + ) + + assert len(retrieval.retrieve_multi_calls) == 1 + + +@pytest.mark.asyncio +async def test_chat_without_citation_does_not_attribute_retrieved_sources(): + svc = _svc(llm=FakeLLM(chat_responses=["A general answer with no citation marker."])) + sources = [{"source_type": "document", "filename": "unrelated.pdf"}] + + out = await svc.chat( + partitions=["p"], + payload={"messages": [{"role": "user", "content": "How can you help me?"}], "metadata": {}}, + prepare_sources=lambda d, w: sources, + model_name="m", + ) + + assert json.loads(out["extra"])["sources"] == [] + + +@pytest.mark.asyncio +async def test_chat_invalid_citation_does_not_fallback_to_unrelated_sources(): + svc = _svc(llm=FakeLLM(chat_responses=["Answer. [Sources: 99]"])) + sources = [{"source_type": "document", "filename": "unrelated.pdf"}] + + out = await svc.chat( + partitions=["p"], + payload={"messages": [{"role": "user", "content": "Question"}], "metadata": {}}, + prepare_sources=lambda d, w: sources, + model_name="m", + ) + + assert json.loads(out["extra"])["sources"] == [] + + +@pytest.mark.asyncio +async def test_chat_structured_output_keeps_retrieved_sources_without_citation_marker(): + structured_answer = '{"answer": "Use [Source 1]", "literal_format": "[Sources: 1]"}' + svc = _svc(llm=FakeLLM(chat_responses=[structured_answer])) + sources = [{"source_type": "document", "filename": "report.pdf"}] + + out = await svc.chat( + partitions=["p"], + payload={ + "messages": [{"role": "user", "content": "Question"}], + "metadata": {}, + "response_format": {"type": "json_object"}, + }, + prepare_sources=lambda d, w: sources, + model_name="m", + ) + + assert out["choices"][0]["message"]["content"] == structured_answer + assert json.loads(out["extra"])["sources"] == sources + + +@pytest.mark.asyncio +async def test_chat_stream_structured_output_preserves_source_like_json_values(): + structured_answer = '{"answer":"Use [Source 1]","literal_format":"[Sources: 1]"}' + stream_lines = [ + "data: " + + json.dumps( + { + "choices": [ + { + "delta": {"content": structured_answer}, + "finish_reason": None, + } + ] + } + ) + + "\n\n", + 'data: {"choices":[{"delta":{},"finish_reason":"stop"}]}\n\n', + "data: [DONE]\n\n", + ] + svc = _svc(llm=FakeLLM(stream_lines=stream_lines)) + sources = [{"source_type": "document", "filename": "report.pdf"}] + + lines = [ + line + async for line in svc.chat_stream( + partitions=["p"], + payload={ + "messages": [{"role": "user", "content": "Question"}], + "metadata": {}, + "response_format": {"type": "json_object"}, + }, + prepare_sources=lambda d, w: sources, + model_name="m", + ) + ] + chunks = [ + json.loads(line[len("data: ") :]) + for line in lines + if line.startswith("data: ") and line.strip() != "data: [DONE]" + ] + content = "".join( + choice.get("delta", {}).get("content", "") for chunk in chunks for choice in chunk.get("choices", []) + ) + extra = next(json.loads(chunk["extra"]) for chunk in reversed(chunks) if chunk.get("extra") not in (None, "{}")) + + assert content == structured_answer + assert extra["sources"] == sources + + +@pytest.mark.asyncio +async def test_structured_websearch_returns_only_sources_included_in_context(): + first = SimpleNamespace( + url="https://example.test/included", + title="Included", + content="short evidence", + snippet="", + ) + excluded = SimpleNamespace( + url="https://example.test/excluded", + title="Excluded", + content="long evidence that does not fit", + snippet="", + ) + web = FakeWeb(results=[first, excluded]) + web.max_tokens = qs.get_num_tokens()("[Source 1]\nIncluded\nshort evidence") + svc = _svc( + llm=FakeLLM(chat_responses=['{"answer": "structured"}']), + retrieval=FakeRetrieval(chunks=[]), + web=web, + ) + + out = await svc.chat( + partitions=None, + payload={ + "messages": [{"role": "user", "content": "Question"}], + "metadata": {"websearch": True}, + "response_format": {"type": "json_object"}, + }, + prepare_sources=lambda _docs, results: [{"url": result.url} for result in results], + model_name="m", + ) + + assert json.loads(out["extra"])["sources"] == [{"url": "https://example.test/included"}] + + +@pytest.mark.asyncio +async def test_explicit_websearch_forces_retrieval_for_conversational_classifier_result(): + query_json = json.dumps({"requires_retrieval": False, "query_list": []}) + llm = FakeLLM(chat_responses=[query_json]) + retrieval = FakeRetrieval(chunks=[]) + web = FakeWeb() + svc = _svc(mode="ChatBotRag", llm=llm, retrieval=retrieval, web=web) + + await svc._prepare_chat( + ["p"], + { + "messages": [{"role": "user", "content": "What is happening today?"}], + "metadata": {"websearch": True}, + }, + ) + + assert len(retrieval.retrieve_multi_calls) == 1 + assert web.calls == ["What is happening today?"] + + @pytest.mark.asyncio async def test_websearch_with_partition_fuses_docs_via_retrieve_multi(): # #707/#740: with a partition AND websearch enabled, the document branch must @@ -406,7 +714,7 @@ async def test_websearch_with_partition_fuses_docs_via_retrieve_multi(): retrieval = FakeRetrieval() web_result = SimpleNamespace(url="https://ex.com", title="T", content="web body", snippet="") svc = _svc(retrieval=retrieval, web=FakeWeb(results=[web_result])) - _payload, _docs, web = await svc._prepare_chat( + _payload, _docs, web, _citation_protocol_active = await svc._prepare_chat( ["p"], {"messages": [{"role": "user", "content": "q"}], "metadata": {"websearch": True}} ) assert len(retrieval.retrieve_multi_calls) == 1 # doc branch fused via the rrf_k-aware retrieve_multi @@ -513,17 +821,59 @@ async def test_chat_without_workspace_unaffected(): @pytest.mark.asyncio -async def test_complete_strips_and_filters(): - svc = _svc(llm=FakeLLM(gen_text="text body [Sources: none]")) +async def test_complete_direct_mode_preserves_literal_source_marker(): + answer = "text body\n[Sources: none]" + svc = _svc(llm=FakeLLM(gen_text=answer)) out = await svc.complete( partitions=None, payload={"prompt": "do x"}, prepare_sources=lambda d, w: [{"x": 1}], ) - assert out["choices"][0]["text"] == "text body" + assert out["choices"][0]["text"] == answer assert json.loads(out["extra"])["sources"] == [] +@pytest.mark.asyncio +async def test_complete_conversational_request_uses_openrag_prompt_without_retrieval(): + query_json = json.dumps({"requires_retrieval": False, "query_list": []}) + llm = FakeLLM(chat_responses=[query_json], gen_text="I am OpenRAG.\n[Sources: none]") + retrieval = FakeRetrieval() + svc = _svc(mode="ChatBotRag", llm=llm, retrieval=retrieval) + + out = await svc.complete( + partitions=["p"], + payload={"prompt": "Who are you?"}, + prepare_sources=lambda d, w: [], + ) + + assert retrieval.retrieve_multi_calls == [] + assert out["choices"][0]["text"] == "I am OpenRAG." + answer_prompt = llm.generate_calls[0][0] + assert "OpenRAG" in answer_prompt + assert "LINAGORA" in answer_prompt + assert "document-grounded RAG system" in answer_prompt + assert "Who are you?" in answer_prompt + + +@pytest.mark.asyncio +async def test_complete_partition_request_keeps_context_and_filters_citations(): + llm = FakeLLM(gen_text="The answer is grounded.\n[Sources: 1]") + svc = _svc(llm=llm) + sources = [{"source_type": "document", "filename": "report.pdf"}] + + out = await svc.complete( + partitions=["p"], + payload={"prompt": "What does the report say?"}, + prepare_sources=lambda d, w: sources, + ) + + assert out["choices"][0]["text"] == "The answer is grounded." + assert json.loads(out["extra"])["sources"] == sources + answer_prompt = llm.generate_calls[0][0] + assert "ctx" in answer_prompt + assert "What does the report say?" in answer_prompt + + @pytest.mark.asyncio async def test_chat_stream_yields_sse_and_done(): svc = _svc(llm=FakeLLM()) @@ -567,12 +917,19 @@ def test_json_slice_extracts_object(): def test_dedupe_web_preserves_first_seen(): - a = SimpleNamespace(url="u1") - b = SimpleNamespace(url="u1") - c = SimpleNamespace(url="u2") + a = SimpleNamespace(url="https://example.test/one") + b = SimpleNamespace(url="https://example.test/one") + c = SimpleNamespace(url="https://example.test/two") assert qs._dedupe_web([[a, b], [c]]) == [a, c] +def test_dedupe_web_drops_invalid_urls_before_source_numbering(): + invalid = SimpleNamespace(url="javascript:alert(1)") + valid = SimpleNamespace(url="https://example.test/evidence") + + assert qs._dedupe_web([[invalid, valid]]) == [valid] + + def test_sampling_strips_transport_keys(): out = qs._sampling({"messages": [], "stream": True, "model": "m", "temperature": 0.5}) assert out == {"temperature": 0.5} diff --git a/tests/unit/test_app_front_secret.py b/tests/unit/test_app_front_secret.py index b0a9a5e37..49565bd88 100644 --- a/tests/unit/test_app_front_secret.py +++ b/tests/unit/test_app_front_secret.py @@ -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. @@ -398,6 +408,303 @@ async def fake_load_model_ids(_client, api_key): assert module._OPENRAG_TOKEN_STORE[auth_handle][0] == "handoff-token" +@pytest.mark.parametrize( + "sources", + [ + None, + [], + 42, + "not-a-source-list", + {"filename": "notes.txt"}, + [{}], + [None], + [{"source_type": "web", "title": "", "url": "", "snippet": ""}], + [{"source_type": "web", "title": "Invalid URL", "url": "http://[invalid"}], + [{"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_page_less_non_pdf_sources(monkeypatch): + module = _load_app_front(monkeypatch, auth_mode="token", module_name="app_front_page_less_sources_test") + monkeypatch.setattr(module, "get_external_url", lambda: "https://openrag.example") + _stub_chainlit_elements(module) + + async def available_chunk(*_args, **_kwargs): + return "Page-less text content" + + monkeypatch.setattr(module, "__fetch_page_content", available_chunk) + + elements, source_names = await module._format_sources( + [ + { + "filename": "diagram.png", + "file_url": "https://openrag.example/static/image-id", + }, + { + "filename": "demo.mp4", + "file_url": "https://openrag.example/static/video-id", + }, + { + "filename": "recording.mp3", + "file_url": "https://openrag.example/static/audio-id", + }, + { + "filename": "notes.txt", + "file_url": "https://openrag.example/static/text-id", + "chunk_url": "https://openrag.example/chunks/text-id", + }, + ] + ) + + assert source_names == ["diagram.png", "demo.mp4", "recording.mp3", "notes.txt"] + assert [element.name for element in elements] == ["diagram.png", "demo.mp4", "recording.mp3", "notes.txt"] + assert source_names == [element.name for element in elements] + assert elements[-1].content == "Page-less text content" + + +@pytest.mark.asyncio +async def test_chainlit_keeps_page_less_pdf_sources(monkeypatch): + module = _load_app_front(monkeypatch, auth_mode="token", module_name="app_front_page_less_pdf_test") + monkeypatch.setattr(module, "get_external_url", lambda: "https://openrag.example") + _stub_chainlit_elements(module) + + elements, source_names = await module._format_sources( + [ + { + "filename": "report_[draft].pdf", + "file_url": "https://openrag.example/static/pdf-id", + } + ] + ) + + assert source_names == ["report_ draft .pdf"] + assert source_names == [element.name for element in elements] + assert elements[0].page is None + + +@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_escapes_untrusted_web_source_markdown(monkeypatch): + module = _load_app_front(monkeypatch, auth_mode="token", module_name="app_front_web_source_escaping_test") + monkeypatch.setattr(module, "get_external_url", lambda: "https://openrag.example") + _stub_chainlit_elements(module) + + title = "Reference ](https://spoof.test) **trusted**" + snippet = "Evidence [click here](https://spoof.test) or *ignore this*." + url = "https://example.test/reference_(draft)" + elements, source_names = await module._format_sources( + [ + { + "source_type": "web", + "title": title, + "url": url, + "snippet": snippet, + } + ] + ) + + assert source_names == ["Reference (https://spoof.test) trusted"] + assert elements[0].name == source_names[0] + assert elements[0].content == ( + r"**[Reference \]\(https\:\/\/spoof\.test\) \*\*trusted\*\*]" + "(https://example.test/reference_%28draft%29)**\n\n" + r"Evidence \[click here\]\(https\:\/\/spoof\.test\) or \*ignore this\*\." + ) + + +@pytest.mark.asyncio +async def test_chainlit_keeps_source_names_unique_after_sanitizing(monkeypatch): + module = _load_app_front(monkeypatch, auth_mode="token", module_name="app_front_source_name_collision_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": "Reference [draft]", "url": "https://example.test/one"}, + {"source_type": "web", "title": "Reference *draft*", "url": "https://example.test/two"}, + ] + ) + + assert source_names == ["Reference draft", "Reference draft 2"] + assert source_names == [element.name for element in elements] + + +@pytest.mark.parametrize( + ("source_name", "expected"), + [ + ("rapport_annuel_2026.pdf", "rapport_annuel_2026.pdf"), + ("report.pdf (page: 3)", "report.pdf (page: 3)"), + ("C++_style_guide.pdf", "C++_style_guide.pdf"), + ], +) +def test_chainlit_preserves_safe_source_name_punctuation(monkeypatch, source_name, expected): + module = _load_app_front(monkeypatch, auth_mode="token", module_name="app_front_source_name_punctuation_test") + + assert module._safe_source_name(source_name, {}) == expected + + +@pytest.mark.parametrize( + ("source_name", "expected"), + [ + ("_x_", "x"), + ("_trusted_", "trusted"), + ("__trusted__", "trusted"), + ("foo _trusted_ bar", "foo trusted bar"), + ("~~irrelevant~~", "irrelevant"), + ("# Trusted source", "Trusted source"), + ("1. Official result", "Official result"), + ("- Search result", "Search result"), + ("+ Search result", "Search result"), + ("# 1. Official result", "Official result"), + ("___", "source"), + ], +) +def test_chainlit_neutralizes_active_source_name_markdown(monkeypatch, source_name, expected): + module = _load_app_front(monkeypatch, auth_mode="token", module_name="app_front_source_name_markdown_test") + + assert module._safe_source_name(source_name, {}) == expected + + +@pytest.mark.parametrize("source_name", ["#report.pdf", "1.report.pdf"]) +def test_chainlit_preserves_non_block_source_name_prefixes(monkeypatch, source_name): + module = _load_app_front(monkeypatch, auth_mode="token", module_name="app_front_source_name_prefix_test") + + assert module._safe_source_name(source_name, {}) == source_name + + +def test_chainlit_handles_long_source_name_delimiter_runs(monkeypatch): + module = _load_app_front(monkeypatch, auth_mode="token", module_name="app_front_source_name_long_run_test") + source_name = "_" * 10_000 + + assert module._safe_source_name(source_name, {}) == "source" + + +@pytest.mark.parametrize( + "hostile_name", + [ + "x](https://evil.test)", + "nested[a](b)c", + "trailing\\", + "**bold**", + "back`tick`", + ">quote", + ], +) +def test_chainlit_source_name_cannot_break_markdown_link(monkeypatch, hostile_name): + module = _load_app_front(monkeypatch, auth_mode="token", module_name="app_front_source_name_safety_test") + + safe_name = module._safe_source_name(hostile_name, {}) + + assert not set(safe_name) & set("[]*`\\>") + assert safe_name.strip() == safe_name + + +@pytest.mark.parametrize( + "source_error", + [ + pytest.param(httpx.ConnectError("source unavailable"), id="connection-error"), + pytest.param(httpx.InvalidURL("invalid source URL"), id="invalid-url"), + ], +) +@pytest.mark.asyncio +async def test_chainlit_skips_unavailable_text_sources(monkeypatch, source_error): + 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 source_error + + 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_chainlit_skips_text_source_with_non_object_json(monkeypatch): + module = _load_app_front(monkeypatch, auth_mode="token", module_name="app_front_malformed_source_test") + monkeypatch.setattr(module, "get_external_url", lambda: "https://openrag.example") + _stub_chainlit_elements(module) + + class FakeResponse: + def raise_for_status(self): + return None + + def json(self): + return [{"page_content": "unexpected list response"}] + + class FakeAsyncClient: + async def __aenter__(self): + return self + + async def __aexit__(self, *_args): + return None + + async def get(self, *_args, **_kwargs): + return FakeResponse() + + monkeypatch.setattr(module.httpx, "AsyncClient", FakeAsyncClient) + + 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") @@ -410,14 +717,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( [ @@ -445,14 +746,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( [