diff --git a/.env.example b/.env.example index 79c2c2be0..aad7b0267 100644 --- a/.env.example +++ b/.env.example @@ -54,6 +54,12 @@ INDEXERUI_PORT=8060 # Port to expose the Indexer UI INDEXERUI_URL='http://X.X.X.X:INDEXERUI_PORT' API_BASE_URL='http://X.X.X.X:APP_PORT' # Base URL of your FastAPI backend. +# Web Search +# WEBSEARCH_API_TOKEN= # Web search provider API token. If unset, web search is silently disabled. +# WEBSEARCH_BASE_URL=https://api.staan.ai/search/web # Web search provider endpoint +# WEBSEARCH_TOP_K=5 # Number of web results to include (default: 5) +# WEBSEARCH_LANG=fr-FR # Search language/market (default: fr-FR) + # LOGGING LOG_LEVEL=DEBUG # See possible values https://loguru.readthedocs.io/en/stable/api/logger.html diff --git a/.gitignore b/.gitignore index 3bf1ad2ee..41d956636 100644 --- a/.gitignore +++ b/.gitignore @@ -69,6 +69,9 @@ services/* #helm charts/openrag-stack/charts/*.tgz +# Planning +.planning/ + # Astro / Starlight .astro/ dist/ diff --git a/.hydra_config/config.yaml b/.hydra_config/config.yaml index 1cb693e5e..45f370b60 100644 --- a/.hydra_config/config.yaml +++ b/.hydra_config/config.yaml @@ -3,6 +3,7 @@ defaults: - chunker: ${oc.env:CHUNKER, recursive_splitter} # recursive_splitter - retriever: ${oc.env:RETRIEVER_TYPE, single} # single # multiQuery # hyde - rag: ChatBotRag + - websearch: ${oc.env:WEBSEARCH_PROVIDER, staan} llm_params: &llm_params temperature: 0.1 @@ -67,6 +68,7 @@ map_reduce: # Enable debug logging for map & reduce debug: ${oc.decode:${oc.env:MAP_REDUCE_DEBUG, false}} + verbose: level: ${oc.env:LOG_LEVEL, DEBUG} diff --git a/.hydra_config/websearch/base.yaml b/.hydra_config/websearch/base.yaml new file mode 100644 index 000000000..83381dd81 --- /dev/null +++ b/.hydra_config/websearch/base.yaml @@ -0,0 +1,11 @@ +provider: '' +api_token: ${oc.env:WEBSEARCH_API_TOKEN, ""} +base_url: '' +top_k: ${oc.decode:${oc.env:WEBSEARCH_TOP_K, 5}} +lang: ${oc.env:WEBSEARCH_LANG, fr-FR} +max_tokens: ${oc.decode:${oc.env:WEBSEARCH_MAX_TOKENS, 2000}} +fetch_content: ${oc.decode:${oc.env:WEBSEARCH_FETCH_CONTENT, true}} +fetch_max_results: ${oc.decode:${oc.env:WEBSEARCH_FETCH_MAX_RESULTS, 3}} +fetch_timeout: ${oc.decode:${oc.env:WEBSEARCH_FETCH_TIMEOUT, 1.0}} +fetch_max_tokens: ${oc.decode:${oc.env:WEBSEARCH_FETCH_MAX_TOKENS, 500}} +fetch_verify_ssl: ${oc.decode:${oc.env:WEBSEARCH_FETCH_VERIFY_SSL, false}} diff --git a/.hydra_config/websearch/staan.yaml b/.hydra_config/websearch/staan.yaml new file mode 100644 index 000000000..de14bc331 --- /dev/null +++ b/.hydra_config/websearch/staan.yaml @@ -0,0 +1,5 @@ +defaults: + - base + +provider: staan +base_url: ${oc.env:WEBSEARCH_BASE_URL, "https://api.staan.ai/search/web"} diff --git a/CLAUDE.md b/CLAUDE.md index 2bda129f7..96dad6ced 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -189,6 +189,28 @@ await vectordb.list_partition_members.remote(partition) - For admins with `SUPER_ADMIN_MODE=true`, `all` resolves to all system partitions - Model prefix is `openrag-` (legacy: `ragondin-`) +### Web Search Integration + +Optional web search augmentation via the Staan API, allowing the LLM to combine RAG document context with live web results. + +**Configuration** (`.hydra_config/config.yaml` → `websearch:` block, env vars): +- `WEBSEARCH_API_TOKEN` — provider API token; if unset, web search is silently disabled +- `WEBSEARCH_BASE_URL` — provider endpoint (default: Staan API) +- `WEBSEARCH_TOP_K` — number of web results (default: 5) +- `WEBSEARCH_LANG` — search language/market (default: `fr-FR`) + +**How it works:** +- Client sends `metadata: {"websearch": true}` in the chat completion request +- **Combined mode** (partition + websearch): RAG retrieval and web search run concurrently via `asyncio.gather()`; web results are appended after document sources with continuous `[Source N]` numbering +- **Web-only mode** (no partition + websearch): skips RAG retrieval entirely, uses web results as sole context; if no results (token unset / search fails), falls back to plain direct LLM mode +- Source entries include `source_type: "document"` or `source_type: "web"` in the `extra.sources` response + +**Key files:** +- `openrag/components/websearch/` — `WebSearchService`, `BaseWebSearchProvider`, `StaanProvider` +- `openrag/components/utils.py` — `format_web_context()` formats web results as numbered source blocks +- `openrag/components/pipeline.py` — `_prepare_for_web_only()`, web search logic in `_prepare_for_chat_completion()` +- `openrag/routers/openai.py` — `__prepare_sources()` merges document and web sources + ### File Quota System Per-user file quota enforcement tracked via the `file_count` and `file_quota` columns on `users`, and `created_by` on `files`. @@ -236,7 +258,7 @@ Environment variables override config values (see `.env.example`). act -j api-tests -W .github/workflows/api_tests.yml --bind ``` -**Mock VLLM for CI:** `.github/workflows/api_tests/mock_vllm.py` provides fake embeddings and completions endpoints (streaming and non-streaming) for testing without a real LLM. Pydantic request models use `ConfigDict(extra="allow")` to accept vendor-specific fields like `extra_body`. +**Mock VLLM for CI:** `tests/api_tests/api_run/mock_vllm.py` provides fake embeddings and completions endpoints (streaming and non-streaming) for testing without a real LLM. Pydantic request models use `ConfigDict(extra="allow")` to accept vendor-specific fields like `extra_body`. ## Key Patterns diff --git a/docs/content/docs/documentation/API.mdx b/docs/content/docs/documentation/API.mdx index 7174560a2..c2fc43466 100644 --- a/docs/content/docs/documentation/API.mdx +++ b/docs/content/docs/documentation/API.mdx @@ -400,17 +400,16 @@ OpenAI-compatible text completion endpoint. #### Extra arguments -* When using the openai endpoint /v1/chat/completions, one can provide extra arguments in the request body to customize the RAG behavior: -- `spoken_style_answer`: boolean (default: false) - If true, the model will generate a succint spoken style conversational answer based on the retrieved documents. -- `use_map_reduce`: boolean (default: false) - If true, the model will use a map-reduce strategy to aggregate information from multiple documents. For more information see the [map-reduce documentation](/openrag/documentation/env_vars/#map--reduce-configuration). -- `llm_override`: object (optional) - Route the request to a different LLM endpoint while still using OpenRAG's RAG pipeline (retrieval, reranking, prompt construction). Accepts the following fields: - - `base_url`: string - Base URL of the target LLM API (e.g. `https://api.openai.com/v1`) - - `api_key`: string - API key for the target LLM - - `model`: string - Model name to use on the target endpoint +* When using the openai endpoint /v1/chat/completions, you can pass extra arguments **via the `metadata` field** of the request body to customize the RAG behavior: - Any field not provided falls back to the default OpenRAG LLM configuration. +| Option | Type | Default | Description | +|--------|------|---------|-------------| +| `websearch` | `bool` | `false` | Augments the RAG context with live web search results. When used with a partition (`openrag-{partition}`), document and web results are combined. When used without a partition (direct LLM mode), web results are the sole context. Requires `WEBSEARCH_API_TOKEN` to be configured. See [web search configuration](/openrag/documentation/env_vars/#web-search-configuration). | +| `spoken_style_answer` | `bool` | `false` | Generates a succinct spoken-style conversational answer based on the retrieved documents. | +| `use_map_reduce` | `bool` | `false` | Uses a map-reduce strategy to aggregate information from multiple documents. See [map-reduce configuration](/openrag/documentation/env_vars/#map--reduce-configuration). | +| `llm_override` | `object` | `null` | Routes the request to a different LLM endpoint while still using OpenRAG's RAG pipeline (retrieval, reranking, prompt construction). Accepts: `base_url` (string), `api_key` (string), `model` (string). Any field not provided falls back to the default OpenRAG LLM configuration. | -These arguments are supplied via the metadata field of the OpenAI request body. Example: +Examples: ```bash title="Enabling conversational answer with openai chat completions endpoint" curl -X 'POST' 'http://localhost:8080/v1/chat/completions' \ @@ -433,6 +432,46 @@ curl -X 'POST' 'http://localhost:8080/v1/chat/completions' \ }' ``` +```bash title="Enabling web search with RAG documents" +curl -X 'POST' 'http://localhost:8080/v1/chat/completions' \ + -H 'accept: application/json' \ + -H 'Authorization: Bearer YOUR_AUTH_TOKEN' \ + -H 'Content-Type: application/json' \ + -d '{ + "model": "openrag-{partition_name}", + "messages": [ + { + "role": "user", + "content": "your_query" + } + ], + "stream": false, + "metadata": { + "websearch": true + } +}' +``` + +```bash title="Web search only (no RAG partition)" +curl -X 'POST' 'http://localhost:8080/v1/chat/completions' \ + -H 'accept: application/json' \ + -H 'Authorization: Bearer YOUR_AUTH_TOKEN' \ + -H 'Content-Type: application/json' \ + -d '{ + "model": "", + "messages": [ + { + "role": "user", + "content": "your_query" + } + ], + "stream": false, + "metadata": { + "websearch": true + } +}' +``` + ```bash title="Using a custom LLM endpoint with OpenRAG's RAG pipeline" curl -X 'POST' 'http://localhost:8080/v1/chat/completions' \ -H 'accept: application/json' \ diff --git a/docs/content/docs/documentation/env_vars.md b/docs/content/docs/documentation/env_vars.md index 8d012a034..32534c3b1 100644 --- a/docs/content/docs/documentation/env_vars.md +++ b/docs/content/docs/documentation/env_vars.md @@ -383,6 +383,28 @@ Ray Serve enables deployment of the FastAPI as a scalable service. For simple de | `CHAINLIT_PORT` | int | 8090 | Port for the Chainlit UI interface if ray serve is enable `ENABLE_RAY_SERVE`. If not chainlit UI is simply a subroute (`/chainlit` [see this](/openrag/getting_started/usage/#default-ports)) of the FastAPI **`base_url`**| +### Web Search Configuration + +Web search allows the LLM to augment RAG document context with live web results. It is disabled by default — set `WEBSEARCH_API_TOKEN` to enable it. + +| Variable | Type | Default | Description | +|----------|------|---------|-------------| +| `WEBSEARCH_PROVIDER` | `str` | `staan` | Web search provider to use. Currently supported: `staan`. | +| `WEBSEARCH_API_TOKEN` | `str` | `""` | API token for the web search provider. If empty, web search is disabled. | +| `WEBSEARCH_BASE_URL` | `str` | (provider default) | Base URL of the web search provider API. | +| `WEBSEARCH_TOP_K` | `int` | `5` | Number of web search results to return. | +| `WEBSEARCH_LANG` | `str` | `fr-FR` | Language/market code for web search queries. | +| `WEBSEARCH_MAX_TOKENS` | `int` | `2000` | Maximum token budget for all web sources combined in the LLM context. This budget is reserved from the global context window when web results are present. | +| `WEBSEARCH_FETCH_CONTENT` | `bool` | `true` | When enabled, fetches actual page content from the top URLs instead of relying on short search snippets. | +| `WEBSEARCH_FETCH_MAX_RESULTS` | `int` | `3` | Number of top URLs to fetch content from (the remaining results use their search snippet). | +| `WEBSEARCH_FETCH_TIMEOUT` | `float` | `1.0` | Per-URL timeout in seconds for content fetching. URLs that don't respond within this time fall back to their snippet. | +| `WEBSEARCH_FETCH_MAX_TOKENS` | `int` | `500` | Maximum approximate tokens of content to extract per page. Content is truncated at word boundaries. | +| `WEBSEARCH_FETCH_VERIFY_SSL` | `bool` | `false` | Whether to verify SSL certificates when fetching page content. | + +:::tip[How to Enable Web Search?] +When chatting, you can enable web search through the OpenAI-compatible API by setting `"websearch": true` in the `metadata` field of the request body. See the [API documentation](/openrag/documentation/api/#extra-arguments) for examples. +::: + ### Map & Reduce Configuration The map & reduce mechanism processes documents by fetching chunks (map phase), filtering out irrelevant ones and summarizing relevant content (reduce phase) with respect to the user's query. The algorithm works as follows: diff --git a/docs/content/docs/documentation/features_in_details.md b/docs/content/docs/documentation/features_in_details.md index 5cce7a411..bc65fc886 100644 --- a/docs/content/docs/documentation/features_in_details.md +++ b/docs/content/docs/documentation/features_in_details.md @@ -84,6 +84,22 @@ See the section on [distributed deployment in a ray cluster](#5-distributed-depl +### 🌐 Web Search Augmentation +Enhance RAG responses with live web search results. When enabled, the LLM can combine document context with up-to-date information from the web. + +
+ +Web Search Features + +* **Combined mode** — RAG retrieval and web search run concurrently; web results are appended as additional sources alongside document sources +* **Web-only mode** — Skip RAG entirely by omitting the partition; uses web results as the sole context +* **Content fetching** — The top 3 URLs are fetched in parallel (1s timeout) and their main content is extracted, providing richer context than search snippets alone +* **Boilerplate filtering** — Navigation, footers, headers, and other non-content HTML elements are stripped before extraction +* **Graceful fallback** — If web search fails or returns no results, the pipeline continues with document context only (or falls back to direct LLM mode) +* **Source attribution** — Web sources are tagged with `source_type: "web"` in the response, distinct from `source_type: "document"` + +
+ ### 🔍 Advanced Retrieval & Reranking [OpenRag](https://open-rag.ai/) Leverages state-of-the-art retrieval techniques for superior accuracy. diff --git a/openrag/app_front.py b/openrag/app_front.py index 47960aa12..3999e3646 100644 --- a/openrag/app_front.py +++ b/openrag/app_front.py @@ -36,6 +36,11 @@ "description": "Get a conversational text answer suitable for voice assistants.\nThe answer is concise, clear, and factual.", "persistent": True, }, + { + "id": "WebSearch", + "icon": "globe", + "description": "Augment the RAG context with live web search results.\nCombines document and web sources for more comprehensive answers.", + }, ] @@ -162,6 +167,17 @@ async def _format_sources(metadata_sources, only_txt=False, api_key=None): d = {} headers = get_headers(api_key) for i, s in enumerate(metadata_sources): + if s.get("source_type") == "web": + title = s.get("title") or s.get("url", f"Web source {i + 1}") + url = s.get("url", "") + snippet = s.get("snippet", "") + content = f"**[{title}]({url})**\n\n{snippet}" + source_name = title + if source_name in d: + source_name = f"{title} ({i})" + d[source_name] = cl.Text(content=content, name=source_name, display="side") + continue + filename = Path(s["filename"]) file_url = s["file_url"] file_url = file_url.replace(INTERNAL_BASE_URL, external_url) # put the correct base url @@ -222,6 +238,7 @@ async def on_message(message: cl.Message): "metadata": { "use_map_reduce": message.command == "DeepSearch", "spoken_style_answer": message.command == "SpokenStyleAnswer", + "websearch": message.command == "WebSearch", }, } diff --git a/openrag/components/pipeline.py b/openrag/components/pipeline.py index f349438fa..421d834d6 100644 --- a/openrag/components/pipeline.py +++ b/openrag/components/pipeline.py @@ -1,3 +1,4 @@ +import asyncio import copy from enum import Enum @@ -6,6 +7,7 @@ SPOKEN_STYLE_ANSWER_PROMPT, SYS_PROMPT_TMPLT, ) +from components.websearch import WebSearchFactory from config import load_config from langchain_core.documents.base import Document from openai import AsyncOpenAI @@ -15,7 +17,7 @@ from .map_reduce import RAGMapReduce from .reranker import Reranker from .retriever import BaseRetriever, RetrieverFactory -from .utils import format_context +from .utils import SOURCE_SEPARATOR, format_context, format_web_context logger = get_logger() config = load_config() @@ -89,6 +91,13 @@ def __init__(self) -> None: # map reduce self.map_reduce: RAGMapReduce = RAGMapReduce(config=config) + # Web search + self.web_search_service = WebSearchFactory.create_service(config) + if self.web_search_service.provider: + logger.info("Web search enabled", provider=config.websearch.get("provider")) + else: + logger.info("Web search disabled (WEBSEARCH_API_TOKEN not set)") + async def generate_query(self, messages: list[dict]) -> str: match RAGMODE(self.rag_mode): case RAGMODE.SIMPLERAG: @@ -121,7 +130,7 @@ async def generate_query(self, messages: list[dict]) -> str: contextualized_query = response.choices[0].message.content return contextualized_query - async def _prepare_for_chat_completion(self, partition: list[str], payload: dict): + async def _prepare_for_chat_completion(self, partition: list[str] | None, payload: dict): messages = payload["messages"] messages = messages[-self.chat_history_depth :] # limit history depth @@ -129,28 +138,68 @@ async def _prepare_for_chat_completion(self, partition: list[str], payload: dict query = await self.generate_query(messages) logger.debug("Prepared query for chat completion", query=query) - metadata = payload.get("metadata", {}) + metadata = payload.get("metadata") or {} use_map_reduce = metadata.get("use_map_reduce", False) spoken_style_answer = metadata.get("spoken_style_answer", False) + use_websearch = metadata.get("websearch", False) logger.debug( "Metadata parameters", use_map_reduce=use_map_reduce, spoken_style_answer=spoken_style_answer, + use_websearch=use_websearch, ) - # 2. get docs + # 2. get docs and/or web results concurrently top_k = config.map_reduce["max_total_documents"] if use_map_reduce else None - docs = await self.retriever_pipeline.retrieve_docs(partition=partition, query=query, top_k=top_k) + if partition is not None and use_websearch: + docs, web_results = await asyncio.gather( + self.retriever_pipeline.retrieve_docs(partition=partition, query=query, top_k=top_k), + self.web_search_service.search(query), + ) + elif partition is not None: + docs = await self.retriever_pipeline.retrieve_docs(partition=partition, query=query, top_k=top_k) + web_results = [] + else: + # Web-only mode (partition is None): no RAG retrieval + docs = [] + web_results = await self.web_search_service.search(query) + + # Web-only with no results: fall back to plain direct LLM mode + if not docs and not web_results and partition is None: + return payload, [], [] if use_map_reduce and docs: docs = await self.map_reduce.map(query=query, chunks=docs) - # 3. Format the retrieved docs - context, included_indices = format_context(docs, max_context_tokens=self.max_context_tokens) + # 3. Format web results first to know actual token usage, then allocate remaining budget to RAG + web_formatted = "" + web_tokens_used = 0 + if web_results: + web_formatted, _, web_tokens_used = format_web_context( + web_results, start_index=1, max_tokens=self.web_search_service.max_tokens + ) + + rag_max_tokens = self.max_context_tokens - web_tokens_used + context, included_indices = format_context(docs, max_context_tokens=rag_max_tokens) docs = [docs[i] for i in included_indices] + # Re-number web sources after RAG sources and rebuild if needed + if web_results: + n_rag_sources = len(docs) + if n_rag_sources > 0: + # Re-format with correct start_index now that we know RAG source count + web_formatted, _, _ = format_web_context( + web_results, start_index=n_rag_sources + 1, max_tokens=self.web_search_service.max_tokens + ) + + # Avoid misleading "No document found" when web results provide context + if not docs: + context = "" + + context = f"{context}{SOURCE_SEPARATOR}{web_formatted}" if context else web_formatted + # 4. prepare the output messages: list = copy.deepcopy(messages) @@ -165,7 +214,7 @@ async def _prepare_for_chat_completion(self, partition: list[str], payload: dict }, ) payload["messages"] = messages - return payload, docs + return payload, docs, web_results async def _prepare_for_completions(self, partition: list[str], payload: dict): prompt = payload["prompt"] @@ -199,9 +248,14 @@ async def completions(self, partition: list[str], payload: dict): return llm_output, docs async def chat_completion(self, partition: list[str] | None, payload: dict): - if partition is None: + metadata = payload.get("metadata") or {} + use_websearch = metadata.get("websearch", False) + + if partition is None and not use_websearch: + # Direct LLM mode: no RAG, no web search docs = [] + web_results = [] else: - payload, docs = await self._prepare_for_chat_completion(partition=partition, payload=payload) + payload, docs, web_results = await self._prepare_for_chat_completion(partition=partition, payload=payload) llm_output = self.llm_client.chat_completion(request=payload) - return llm_output, docs + return llm_output, docs, web_results diff --git a/openrag/components/utils.py b/openrag/components/utils.py index 1cdcb0c93..cd871ae24 100644 --- a/openrag/components/utils.py +++ b/openrag/components/utils.py @@ -7,12 +7,15 @@ from typing import ClassVar import ray +from components.indexer.utils.text_sanitizer import sanitize_text from config import load_config from fast_langdetect import LangDetectConfig, LangDetector from langchain_core.documents.base import Document from langchain_openai import ChatOpenAI from utils.logger import get_logger +SOURCE_SEPARATOR = "-" * 10 + "\n\n" + # Global variables config = load_config() logger = get_logger() @@ -115,9 +118,50 @@ def format_context( included_indices.append(i) total_tokens += n_tokens - sep = "-" * 10 + "\n\n" logger.debug("Context formatted", total_tokens=total_tokens, doc_count=len(reduced_docs)) - return f"{sep}".join(reduced_docs), included_indices + return SOURCE_SEPARATOR.join(reduced_docs), included_indices + + +def format_web_context( + web_results: list, + start_index: int = 1, + max_tokens: int = 2000, +) -> tuple[str, list[int], int]: + """Format web results as numbered [Source N] blocks within a token budget. + + Uses fetched page content when available, falling back to the search snippet. + + Args: + web_results: Results from web search provider (list of WebResult) + start_index: First source number (continues numbering after RAG sources) + max_tokens: Maximum token budget for all web sources combined + + Returns: + (formatted_string, list_of_source_numbers_used, total_tokens_used) + """ + if not web_results: + return "", [], 0 + + _length_function = get_num_tokens() + + parts = [] + source_numbers = [] + total_tokens = 0 + + for i, result in enumerate(web_results): + n = start_index + i + title = sanitize_text(result.title) + body = sanitize_text(result.content) if result.content else sanitize_text(result.snippet) + block = f"[Source {n}]\n{title}\n{body}" + block_tokens = _length_function(block) + if total_tokens + block_tokens > max_tokens and parts: + break + parts.append(block) + source_numbers.append(n) + total_tokens += block_tokens + + logger.debug("Web context formatted", total_tokens=total_tokens, source_count=len(parts)) + return SOURCE_SEPARATOR.join(parts), source_numbers, total_tokens _SOURCES_NONE_RE = re.compile(r"\n?\[?Sources?\]?\s*:\s*\[?\s*none\s*\]?\s*$", re.IGNORECASE) diff --git a/openrag/components/websearch/__init__.py b/openrag/components/websearch/__init__.py new file mode 100644 index 000000000..45df91ac5 --- /dev/null +++ b/openrag/components/websearch/__init__.py @@ -0,0 +1,48 @@ +from .base import BaseWebSearchProvider as BaseWebSearchProvider +from .base import WebResult as WebResult +from .content_fetcher import ContentFetcher as ContentFetcher +from .providers.staan import StaanProvider +from .service import WebSearchService as WebSearchService + +PROVIDER_MAPPING = { + "staan": StaanProvider, +} + + +class WebSearchFactory: + @staticmethod + def create_service(config) -> WebSearchService: + """Create a WebSearchService from Hydra config, following the embedder/retriever pattern.""" + ws_config = config.websearch + api_token = ws_config.get("api_token", "") + max_tokens = ws_config.get("max_tokens") + + if not api_token: + return WebSearchService(provider=None, max_tokens=max_tokens) + + provider_name = ws_config.get("provider", "") + provider_cls = PROVIDER_MAPPING.get(provider_name) + if provider_cls is None: + raise ValueError(f"Unsupported web search provider: {provider_name}") + + provider = provider_cls( + api_token=api_token, + base_url=ws_config.get("base_url"), + top_k=ws_config.get("top_k"), + lang=ws_config.get("lang"), + ) + + content_fetcher = None + if ws_config.get("fetch_content"): + content_fetcher = ContentFetcher( + max_results=ws_config.get("fetch_max_results"), + timeout=ws_config.get("fetch_timeout"), + max_tokens_per_page=ws_config.get("fetch_max_tokens"), + verify_ssl=ws_config.get("fetch_verify_ssl"), + ) + + return WebSearchService( + provider=provider, + content_fetcher=content_fetcher, + max_tokens=max_tokens, + ) diff --git a/openrag/components/websearch/base.py b/openrag/components/websearch/base.py new file mode 100644 index 000000000..3d740396c --- /dev/null +++ b/openrag/components/websearch/base.py @@ -0,0 +1,17 @@ +from abc import ABC, abstractmethod +from dataclasses import dataclass + + +@dataclass +class WebResult: + title: str + url: str + snippet: str + content: str | None = None + + +class BaseWebSearchProvider(ABC): + @abstractmethod + async def search(self, query: str) -> list[WebResult]: + """Return web results for the query. May raise on failure.""" + ... diff --git a/openrag/components/websearch/content_fetcher.py b/openrag/components/websearch/content_fetcher.py new file mode 100644 index 000000000..3e3ffc10a --- /dev/null +++ b/openrag/components/websearch/content_fetcher.py @@ -0,0 +1,135 @@ +import asyncio +import ipaddress +from urllib.parse import urlparse + +import httpx +import lxml.html +from components.websearch.base import WebResult +from html_to_markdown import convert +from utils.logger import get_logger + +logger = get_logger() + +# Rough chars-per-token estimate for truncation (conservative) +_CHARS_PER_TOKEN = 4 + +# HTML tags that typically contain boilerplate, not main content +_BOILERPLATE_TAGS = {"nav", "footer", "header", "aside", "script", "style", "noscript"} + +_USER_AGENT = "Mozilla/5.0 (compatible; OpenRAG/1.0; +https://github.com/linagora/openrag)" + + +class ContentFetcher: + """Fetch and extract text content from web search result URLs.""" + + def __init__( + self, + max_results: int = 3, + timeout: float = 1.0, + max_tokens_per_page: int = 500, + verify_ssl: bool = False, + ): + self.max_results = max_results + self.timeout = timeout + self.max_tokens_per_page = max_tokens_per_page + self.verify_ssl = verify_ssl + self._client_override: httpx.AsyncClient | None = None # For testing + + def _truncate(self, text: str) -> str: + """Truncate text to approximately max_tokens_per_page tokens.""" + max_chars = self.max_tokens_per_page * _CHARS_PER_TOKEN + if len(text) <= max_chars: + return text + truncated = text[:max_chars] + last_space = truncated.rfind(" ") + if last_space > max_chars * 0.8: + truncated = truncated[:last_space] + return truncated.rstrip() + " [...]" + + @staticmethod + def _is_loopback_url(url: str) -> bool: + host = urlparse(url).hostname or "" + if host == "localhost": + return True + try: + return not ipaddress.ip_address(host).is_global + except ValueError: + return False # Regular hostname, let it through + + async def _fetch_single(self, client: httpx.AsyncClient, url: str) -> str | None: + """Fetch a single URL and extract text. Returns None on any failure.""" + # Guard against SSRF: URLs come from the search provider, but a compromised + # or misbehaving provider could return loopback addresses targeting internal services. + if self._is_loopback_url(url): + logger.warning("Blocked loopback URL in web search results", url=url) + return None + try: + response = await asyncio.wait_for( + client.get(url, follow_redirects=True), + timeout=self.timeout, + ) + response.raise_for_status() + + content_type = response.headers.get("content-type", "") + if "text/html" not in content_type and "text/plain" not in content_type: + return None + + html = response.text + if not html.strip(): + return None + + # Strip boilerplate elements (nav, footer, etc.) before conversion + try: + tree = lxml.html.fromstring(html) + for tag in _BOILERPLATE_TAGS: + for el in tree.iter(tag): + el.getparent().remove(el) + html = lxml.html.tostring(tree, encoding="unicode") + except Exception: + pass # If lxml parsing fails, convert the raw HTML + + text = convert(html) + text = text.strip() + if not text: + return None + + return self._truncate(text) + + except TimeoutError: + logger.debug("Content fetch timed out", url=url) + return None + except Exception as e: + logger.debug("Content fetch failed", url=url, error=str(e)) + return None + + async def enrich(self, results: list[WebResult]) -> list[WebResult]: + """Fetch content for the top N results in parallel. Mutates results in place.""" + if not results: + return results + + to_fetch = results[: self.max_results] + + client = self._client_override + if client is not None: + tasks = [self._fetch_single(client, r.url) for r in to_fetch] + contents = await asyncio.gather(*tasks) + for result, content in zip(to_fetch, contents): + result.content = content + else: + timeout = httpx.Timeout( + connect=self.timeout, + read=self.timeout, + write=self.timeout, + pool=self.timeout, + ) + async with httpx.AsyncClient( + timeout=timeout, verify=self.verify_ssl, headers={"User-Agent": _USER_AGENT} + ) as client: + tasks = [self._fetch_single(client, r.url) for r in to_fetch] + contents = await asyncio.gather(*tasks) + for result, content in zip(to_fetch, contents): + result.content = content + + n_enriched = sum(1 for c in contents if c is not None) + logger.debug("Content fetching done", enriched=n_enriched, total=len(to_fetch)) + return results diff --git a/openrag/components/websearch/providers/__init__.py b/openrag/components/websearch/providers/__init__.py new file mode 100644 index 000000000..dd4920add --- /dev/null +++ b/openrag/components/websearch/providers/__init__.py @@ -0,0 +1 @@ +from .staan import StaanProvider as StaanProvider diff --git a/openrag/components/websearch/providers/staan.py b/openrag/components/websearch/providers/staan.py new file mode 100644 index 000000000..6fedaf3ee --- /dev/null +++ b/openrag/components/websearch/providers/staan.py @@ -0,0 +1,33 @@ +import httpx +from components.websearch.base import BaseWebSearchProvider, WebResult +from utils.logger import get_logger + +logger = get_logger() + + +class StaanProvider(BaseWebSearchProvider): + def __init__(self, api_token: str, base_url: str, top_k: int = 5, lang: str = "fr-FR"): + self.api_token = api_token + self.base_url = base_url + self.top_k = top_k + self.lang = lang + + async def search(self, query: str) -> list[WebResult]: + headers = {"Authorization": f"Bearer {self.api_token}"} + params = {"q": query, "market": self.lang, "offset": 0} + timeout = httpx.Timeout(connect=5.0, read=10.0, write=5.0, pool=2.0) + + async with httpx.AsyncClient(timeout=timeout) as client: + response = await client.get(self.base_url, headers=headers, params=params) + response.raise_for_status() + data = response.json() + + results = data if isinstance(data, list) else data.get("web", {}).get("results", []) + return [ + WebResult( + title=r.get("title", ""), + url=r.get("url", ""), + snippet=r.get("snippet", ""), + ) + for r in results[: self.top_k] + ] diff --git a/openrag/components/websearch/service.py b/openrag/components/websearch/service.py new file mode 100644 index 000000000..8ef0e93bd --- /dev/null +++ b/openrag/components/websearch/service.py @@ -0,0 +1,35 @@ +from components.websearch.base import BaseWebSearchProvider, WebResult +from components.websearch.content_fetcher import ContentFetcher +from utils.logger import get_logger + +logger = get_logger() + + +class WebSearchService: + def __init__( + self, + provider: BaseWebSearchProvider | None, + content_fetcher: ContentFetcher | None = None, + max_tokens: int = 2000, + ): + self.provider = provider # None when WEBSEARCH_API_TOKEN is not set + self.content_fetcher = content_fetcher + self.max_tokens = max_tokens + + async def search(self, query: str) -> list[WebResult]: + if self.provider is None: + logger.warning("Web search requested but no provider configured — ignoring websearch flag") + return [] + try: + results = await self.provider.search(query) + if not results: + logger.warning("Web search returned zero results", query=query) + return results + + if self.content_fetcher: + results = await self.content_fetcher.enrich(results) + + return results + except Exception as e: + logger.warning("Web search failed, continuing without web context", error=str(e)) + return [] diff --git a/openrag/components/websearch/test_content_fetcher.py b/openrag/components/websearch/test_content_fetcher.py new file mode 100644 index 000000000..4ce23bb1b --- /dev/null +++ b/openrag/components/websearch/test_content_fetcher.py @@ -0,0 +1,165 @@ +import asyncio + +import httpx +import pytest +from components.websearch.base import WebResult +from components.websearch.content_fetcher import ContentFetcher + + +@pytest.fixture +def fetcher(): + return ContentFetcher(max_results=3, timeout=1.0, max_tokens_per_page=500) + + +def _make_result(url="https://example.com", snippet="short snippet"): + return WebResult(title="Test", url=url, snippet=snippet) + + +class TestFetchSingleURL: + @pytest.mark.asyncio + async def test_extracts_text_from_html(self, fetcher): + html = "

Hello

World paragraph content here.

" + + async def mock_handler(request): + return httpx.Response(200, text=html) + + transport = httpx.MockTransport(mock_handler) + async with httpx.AsyncClient(transport=transport) as client: + text = await fetcher._fetch_single(client, "https://example.com") + + assert text is not None + assert "Hello" in text + assert "World paragraph content" in text + + @pytest.mark.asyncio + async def test_returns_none_on_timeout(self, fetcher): + async def slow_handler(request): + await asyncio.sleep(5) + return httpx.Response(200, text="too late") + + transport = httpx.MockTransport(slow_handler) + async with httpx.AsyncClient(transport=transport) as client: + text = await fetcher._fetch_single(client, "https://slow.example.com") + + assert text is None + + @pytest.mark.asyncio + async def test_returns_none_on_http_error(self, fetcher): + async def error_handler(request): + return httpx.Response(500, text="error") + + transport = httpx.MockTransport(error_handler) + async with httpx.AsyncClient(transport=transport) as client: + text = await fetcher._fetch_single(client, "https://error.example.com") + + assert text is None + + @pytest.mark.asyncio + @pytest.mark.parametrize( + "url", + [ + "http://localhost/secret", + "http://127.0.0.1/admin", + "http://127.0.0.42/x", + "http://[::1]/admin", + "http://10.0.0.1/internal", + "http://192.168.1.1/router", + "http://169.254.169.254/metadata", + "http://0.0.0.0/x", + ], + ) + async def test_skips_loopback_urls(self, fetcher, url): + async def mock_handler(request): + return httpx.Response(200, text="secret") + + transport = httpx.MockTransport(mock_handler) + async with httpx.AsyncClient(transport=transport) as client: + text = await fetcher._fetch_single(client, url) + + assert text is None + + @pytest.mark.asyncio + async def test_strips_boilerplate_html(self, fetcher): + html = """ + +

Site Header

+

This is the actual article content.

+ + + """ + + async def mock_handler(request): + return httpx.Response(200, text=html) + + transport = httpx.MockTransport(mock_handler) + async with httpx.AsyncClient(transport=transport) as client: + text = await fetcher._fetch_single(client, "https://example.com") + + assert text is not None + assert "actual article content" in text + assert "Home" not in text + assert "Site Header" not in text + assert "Sidebar ad" not in text + assert "Copyright" not in text + + @pytest.mark.asyncio + async def test_returns_none_for_non_html(self, fetcher): + async def pdf_handler(request): + return httpx.Response( + 200, + content=b"%PDF-1.4", + headers={"content-type": "application/pdf"}, + ) + + transport = httpx.MockTransport(pdf_handler) + async with httpx.AsyncClient(transport=transport) as client: + text = await fetcher._fetch_single(client, "https://example.com/file.pdf") + + assert text is None + + +class TestTruncation: + def test_truncates_long_content(self, fetcher): + long_text = "word " * 2000 + result = fetcher._truncate(long_text) + assert len(result) < len(long_text) + + def test_preserves_short_content(self, fetcher): + short_text = "A brief sentence." + result = fetcher._truncate(short_text) + assert result == short_text + + +class TestEnrichResults: + @pytest.mark.asyncio + async def test_enriches_only_top_n_results(self, fetcher): + html = "

Page content for testing.

" + + async def mock_handler(request): + return httpx.Response(200, text=html) + + transport = httpx.MockTransport(mock_handler) + results = [_make_result(f"https://example.com/{i}") for i in range(5)] + + async with httpx.AsyncClient(transport=transport) as client: + fetcher._client_override = client + enriched = await fetcher.enrich(results) + + for r in enriched[:3]: + assert r.content is not None + for r in enriched[3:]: + assert r.content is None + + @pytest.mark.asyncio + async def test_failed_fetch_keeps_none_content(self, fetcher): + async def error_handler(request): + return httpx.Response(500, text="error") + + transport = httpx.MockTransport(error_handler) + results = [_make_result()] + + async with httpx.AsyncClient(transport=transport) as client: + fetcher._client_override = client + enriched = await fetcher.enrich(results) + + assert enriched[0].content is None diff --git a/openrag/models/openai.py b/openrag/models/openai.py index 7f233bdec..323e44d64 100644 --- a/openrag/models/openai.py +++ b/openrag/models/openai.py @@ -29,6 +29,8 @@ class OpenAIChatCompletionRequest(BaseModel): { "use_map_reduce": False, "spoken_style_answer": False, + "websearch": False, + "llm_override": None, }, description="Extra custom parameters. Supports 'llm_override' object with optional 'base_url', 'api_key', and 'model' to override the downstream LLM endpoint.", ) diff --git a/openrag/routers/openai.py b/openrag/routers/openai.py index 263dbdadf..9eb686fb7 100644 --- a/openrag/routers/openai.py +++ b/openrag/routers/openai.py @@ -1,9 +1,10 @@ import asyncio import json from pathlib import Path -from urllib.parse import quote +from urllib.parse import quote, urlparse import consts +from components.indexer.utils.text_sanitizer import sanitize_text from components.pipeline import RagPipeline from components.utils import ( extract_and_strip_sources_block, @@ -110,7 +111,7 @@ async def list_models( return JSONResponse(content={"object": "list", "data": models}) -def __prepare_sources(request: Request, docs: list[Document]): +def __prepare_sources(request: Request, docs: list[Document], web_results: list | None = None): links = [] for doc in docs: doc_metadata = dict(doc.metadata) @@ -119,11 +120,24 @@ def __prepare_sources(request: Request, docs: list[Document]): encoded_url = quote(file_url, safe=":/") links.append( { + "source_type": "document", "file_url": encoded_url, "chunk_url": str(request.url_for("get_extract", extract_id=doc_metadata["_id"])), **doc_metadata, } ) + for result in web_results or []: + url = sanitize_text(result.url or "") + if not url or urlparse(url).scheme not in ("http", "https"): + continue + links.append( + { + "source_type": "web", + "url": url, + "title": sanitize_text(result.title), + "snippet": sanitize_text(result.snippet), + } + ) return links @@ -319,10 +333,10 @@ async def openai_chat_completion( partitions = await get_partition_name(model_name, user_partitions, is_admin=user["is_admin"]) log.debug(f"Using partitions: {partitions}") - llm_output, docs = await ragpipe.chat_completion(partition=partitions, payload=request.model_dump()) + llm_output, docs, web_results = await ragpipe.chat_completion(partition=partitions, payload=request.model_dump()) log.debug("RAG chat completion pipeline executed.") - sources = __prepare_sources(request2, docs) + sources = __prepare_sources(request2, docs, web_results=web_results) if request.stream: diff --git a/pyproject.toml b/pyproject.toml index 7990c9870..648e51878 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -43,6 +43,7 @@ dependencies = [ "pytest-env>=1.1.5", "markitdown[docx]>=0.1.3", "html-to-markdown>=2.4.0", + "lxml>=5.0.0", "alembic>=1.17.0", "fast-langdetect>=1.0.0", "ruff>=0.14.1", diff --git a/tests/api_tests/test_openai_compat.py b/tests/api_tests/test_openai_compat.py index 8c1050203..b0446da3a 100644 --- a/tests/api_tests/test_openai_compat.py +++ b/tests/api_tests/test_openai_compat.py @@ -234,6 +234,52 @@ def test_streaming_has_finish_reason(self, api_client, indexed_partition): assert isinstance(extra["sources"], list) +class TestWebOnlyMode: + """Test web-only mode: metadata.websearch=true with no partition.""" + + def test_web_only_mode_no_partition(self, api_client): + """Web-only mode returns 200 with valid response when websearch=true and no partition.""" + response = api_client.post( + "/v1/chat/completions", + json={ + "model": "", # empty string → is_direct_llm_model()=True → partition=None + "messages": [{"role": "user", "content": "What is the capital of France?"}], + "metadata": {"websearch": True}, + "stream": False, + }, + ) + assert response.status_code == 200 + data = response.json() + assert "choices" in data + assert len(data["choices"]) > 0 + assert data["choices"][0]["message"]["content"] # non-empty string + + # If sources present, all must be web type (no document sources in web-only mode) + extra = json.loads(data["extra"]) if data.get("extra") else {} + sources = extra.get("sources", []) + for source in sources: + assert source.get("source_type") == "web" + + def test_web_only_mode_graceful_degradation(self, api_client): + """Web-only mode with no web results still returns 200 (plain LLM answer).""" + # In CI with mock VLLM and no Staan, web results will be empty. + # The LLM should still answer — graceful degradation, not an error. + response = api_client.post( + "/v1/chat/completions", + json={ + "model": "", + "messages": [{"role": "user", "content": "Hello"}], + "metadata": {"websearch": True}, + "stream": False, + }, + ) + assert response.status_code == 200 + data = response.json() + assert "choices" in data + assert len(data["choices"]) > 0 + assert data["choices"][0]["message"]["content"] # non-empty + + class TestChatCompletionsMultiPartition: """Test chat completions with multi-partition access."""