Coverage for open_webui/routers/retrieval.py: 54%
1268 statements
« prev ^ index » next coverage.py v7.15.2, created at 2026-10-07 05:07 +0000
1from __future__ import annotations
3import asyncio
4import io
5import logging
6import mimetypes
7import os
8import re
9import shutil
10import uuid
11from datetime import datetime
12from pathlib import Path
13from types import SimpleNamespace
14from typing import Callable, Iterator, Optional, Sequence, Union
15from urllib.parse import unquote, urlparse
17import tiktoken
18from fastapi import (
19 APIRouter,
20 Depends,
21 FastAPI,
22 File,
23 Form,
24 HTTPException,
25 Query,
26 Request,
27 UploadFile,
28 status,
29)
30from fastapi.concurrency import run_in_threadpool
31from fastapi.middleware.cors import CORSMiddleware
32from langchain_core.documents import Document
33from langchain_text_splitters import (
34 MarkdownHeaderTextSplitter,
35 RecursiveCharacterTextSplitter,
36 TokenTextSplitter,
37)
38from open_webui.config import (
39 DEFAULT_LOCALE,
40 ENV,
41 RAG_EMBEDDING_CONTENT_PREFIX,
42 RAG_EMBEDDING_MODEL_AUTO_UPDATE,
43 RAG_EMBEDDING_MODEL_TRUST_REMOTE_CODE,
44 RAG_EMBEDDING_QUERY_PREFIX,
45 RAG_RERANKING_MODEL_AUTO_UPDATE,
46 RAG_RERANKING_MODEL_TRUST_REMOTE_CODE,
47 UPLOAD_DIR,
48)
49from open_webui.constants import ERROR_MESSAGES
50from open_webui.env import (
51 AIOHTTP_CLIENT_ALLOW_REDIRECTS,
52 AIOHTTP_CLIENT_SESSION_SSL,
53 DEVICE_TYPE,
54 DOCKER,
55 RAG_EMBEDDING_TIMEOUT,
56 SENTENCE_TRANSFORMERS_BACKEND,
57 SENTENCE_TRANSFORMERS_CROSS_ENCODER_BACKEND,
58 SENTENCE_TRANSFORMERS_CROSS_ENCODER_MODEL_KWARGS,
59 SENTENCE_TRANSFORMERS_CROSS_ENCODER_SIGMOID_ACTIVATION_FUNCTION,
60 SENTENCE_TRANSFORMERS_MODEL_KWARGS,
61 USE_SLIM,
62)
63from open_webui.events import EVENTS, publish_event
64from open_webui.internal.db import get_async_db, get_async_session
65from open_webui.models.files import FileModel, Files, FileUpdateForm
66from open_webui.models.knowledge import Knowledges
67from open_webui.models.config import Config
69# Document loaders
70from open_webui.retrieval.loaders.youtube import YoutubeLoader, YoutubeTranscriptError
71from open_webui.retrieval.utils import (
72 build_loader_from_config,
73 get_loader_config,
74 filter_accessible_collections,
75 get_content_from_url,
76 get_embedding_function,
77 get_model_path,
78 get_reranking_function,
79 is_youtube_url,
80 query_collection,
81 query_collection_with_hybrid_search,
82 query_doc,
83 query_doc_with_hybrid_search,
84)
85from open_webui.retrieval.vector.async_client import ASYNC_VECTOR_DB_CLIENT
86from open_webui.retrieval.vector.factory import get_vector_db_client
87from open_webui.retrieval.vector.utils import filter_metadata
88from open_webui.retrieval.web.azure import search_azure
89from open_webui.retrieval.web.bing import search_bing
90from open_webui.retrieval.web.bocha import search_bocha
91from open_webui.retrieval.web.brave import search_brave
92from open_webui.retrieval.web.brave_llm_context import search_brave_llm_context
93from open_webui.retrieval.web.duckduckgo import search_duckduckgo
94from open_webui.retrieval.web.exa import search_exa
95from open_webui.retrieval.web.external import search_external
96from open_webui.retrieval.web.firecrawl import search_firecrawl
97from open_webui.retrieval.web.google_pse import search_google_pse
98from open_webui.retrieval.web.jina_search import search_jina
99from open_webui.retrieval.web.kagi import search_kagi
100from open_webui.retrieval.web.utils import get_ssrf_safe_session, validate_url
102# Web search engines
103from open_webui.retrieval.web.main import SearchResult
104from open_webui.retrieval.web.microsoft_web_iq import search_microsoft_web_iq
105from open_webui.retrieval.web.mojeek import search_mojeek
106from open_webui.retrieval.web.ollama import search_ollama_cloud
107from open_webui.retrieval.web.perplexity import search_perplexity
108from open_webui.retrieval.web.perplexity_search import search_perplexity_search
109from open_webui.retrieval.web.searchapi import search_searchapi
110from open_webui.retrieval.web.openserp import search_openserp
111from open_webui.retrieval.web.searxng import search_searxng
112from open_webui.retrieval.web.serpapi import search_serpapi
113from open_webui.retrieval.web.serper import search_serper
114from open_webui.retrieval.web.serphouse import search_serphouse
115from open_webui.retrieval.web.serply import search_serply
116from open_webui.retrieval.web.serpstack import search_serpstack
117from open_webui.retrieval.web.sougou import search_sougou
118from open_webui.retrieval.web.staan import search_staan
119from open_webui.retrieval.web.tavily import search_tavily
120from open_webui.retrieval.web.utils import get_web_loader
121from open_webui.retrieval.web.yacy import search_yacy
122from open_webui.retrieval.web.yandex import search_yandex
123from open_webui.retrieval.web.ydc import search_youcom
124from open_webui.retrieval.web.linkup import search_linkup
125from open_webui.storage.provider import Storage
126from open_webui.utils.access_control import has_permission
127from open_webui.utils.access_control.files import has_access_to_file
128from open_webui.utils.auth import get_admin_user, get_verified_user
129from open_webui.utils.misc import (
130 calculate_sha256_string,
131 sanitize_text_for_db,
132)
133from pydantic import BaseModel, Field
134from sqlalchemy.ext.asyncio import AsyncSession
136log = logging.getLogger(__name__)
138TIKTOKEN_DISALLOWED_SPECIAL = ()
140##########################################
141#
142# Utility functions
143# Give us this day our relevant chunks, and lead us
144# not into hallucination, but deliver us from noise.
145#
146##########################################
149def get_ef(
150 engine: str,
151 embedding_model: str,
152 auto_update: bool = RAG_EMBEDDING_MODEL_AUTO_UPDATE,
153):
154 ef = None
155 if embedding_model and engine == '' and not USE_SLIM:
156 from sentence_transformers import SentenceTransformer
158 try:
159 ef = SentenceTransformer(
160 get_model_path(embedding_model, auto_update),
161 device=DEVICE_TYPE,
162 trust_remote_code=RAG_EMBEDDING_MODEL_TRUST_REMOTE_CODE,
163 backend=SENTENCE_TRANSFORMERS_BACKEND,
164 model_kwargs=SENTENCE_TRANSFORMERS_MODEL_KWARGS,
165 )
166 except Exception as e:
167 log.error(f'Error loading SentenceTransformer: {e}')
169 return ef
172def get_rf(
173 engine: str = '',
174 reranking_model: str | None = None,
175 external_reranker_url: str = '',
176 external_reranker_api_key: str = '',
177 external_reranker_timeout: str = '',
178 auto_update: bool = RAG_RERANKING_MODEL_AUTO_UPDATE,
179):
180 rf = None
181 # Convert timeout string to int or None (system default)
182 timeout_value = int(external_reranker_timeout) if external_reranker_timeout else None
183 if reranking_model and engine == 'external':
184 from open_webui.retrieval.models.external import ExternalReranker
186 return ExternalReranker(
187 url=external_reranker_url,
188 api_key=external_reranker_api_key,
189 model=reranking_model,
190 timeout=timeout_value,
191 )
192 if USE_SLIM:
193 return None
194 if reranking_model:
195 if any(model in reranking_model for model in ['jinaai/jina-colbert-v2']):
196 try:
197 from open_webui.retrieval.models.colbert import ColBERT
199 rf = ColBERT(
200 get_model_path(reranking_model, auto_update),
201 env='docker' if DOCKER else None,
202 )
204 except Exception as e:
205 log.error(f'ColBERT: {e}')
206 raise Exception(ERROR_MESSAGES.DEFAULT(e, 'Error loading reranking model'))
207 else:
208 import sentence_transformers
209 import torch
211 try:
212 rf = sentence_transformers.CrossEncoder(
213 get_model_path(reranking_model, auto_update),
214 device=DEVICE_TYPE,
215 trust_remote_code=RAG_RERANKING_MODEL_TRUST_REMOTE_CODE,
216 backend=SENTENCE_TRANSFORMERS_CROSS_ENCODER_BACKEND,
217 model_kwargs=SENTENCE_TRANSFORMERS_CROSS_ENCODER_MODEL_KWARGS,
218 activation_fn=(
219 torch.nn.Sigmoid() if SENTENCE_TRANSFORMERS_CROSS_ENCODER_SIGMOID_ACTIVATION_FUNCTION else None
220 ),
221 )
222 except Exception as e:
223 log.error(f'CrossEncoder: {e}')
224 raise Exception(ERROR_MESSAGES.DEFAULT(e, 'CrossEncoder error'))
226 # Safely adjust pad_token_id if missing as some models do not have this in config
227 try:
228 model_cfg = getattr(rf, 'model', None)
229 if model_cfg and hasattr(model_cfg, 'config'):
230 cfg = model_cfg.config
231 if getattr(cfg, 'pad_token_id', None) is None:
232 # Fallback to eos_token_id when available
233 eos = getattr(cfg, 'eos_token_id', None)
234 if eos is not None:
235 cfg.pad_token_id = eos
236 log.debug('Missing pad_token_id detected; set to eos_token_id=%s', eos)
237 else:
238 log.warning('Neither pad_token_id nor eos_token_id present in model config')
239 except Exception as e2:
240 log.warning(f'Failed to adjust pad_token_id on CrossEncoder: {e2}')
242 return rf
245##########################################
246#
247# API routes
248#
249##########################################
252router = APIRouter()
254RETRIEVAL_CONFIG_KEYS = {
255 'ALLOWED_FILE_EXTENSIONS': 'rag.file.allowed_extensions',
256 'AZURE_AI_SEARCH_API_KEY': 'web.search.azure_ai_search_api_key',
257 'AZURE_AI_SEARCH_ENDPOINT': 'web.search.azure_ai_search_endpoint',
258 'AZURE_AI_SEARCH_INDEX_NAME': 'web.search.azure_ai_search_index_name',
259 'BING_SEARCH_V7_ENDPOINT': 'web.search.bing_search_v7_endpoint',
260 'BING_SEARCH_V7_SUBSCRIPTION_KEY': 'web.search.bing_search_v7_subscription_key',
261 'BOCHA_SEARCH_API_KEY': 'web.search.bocha_search_api_key',
262 'BRAVE_SEARCH_API_KEY': 'web.search.brave_search_api_key',
263 'BRAVE_SEARCH_CONTEXT_TOKENS': 'web.search.brave_search_context_tokens',
264 'BYPASS_EMBEDDING_AND_RETRIEVAL': 'rag.bypass_embedding_and_retrieval',
265 'BYPASS_WEB_SEARCH_EMBEDDING_AND_RETRIEVAL': 'web.search.bypass_embedding_and_retrieval',
266 'BYPASS_WEB_SEARCH_WEB_LOADER': 'web.search.bypass_web_loader',
267 'CHUNK_MIN_SIZE_TARGET': 'rag.chunk_min_size_target',
268 'CHUNK_OVERLAP': 'rag.chunk_overlap',
269 'CHUNK_SIZE': 'rag.chunk_size',
270 'CONTENT_EXTRACTION_SUPPORTED_MEDIA_MIME_TYPES': 'rag.content_extraction.supported_media_mime_types',
271 'CONTENT_EXTRACTION_ENGINE': 'rag.content_extraction_engine',
272 'DATALAB_MARKER_ADDITIONAL_CONFIG': 'rag.datalab_marker_additional_config',
273 'DATALAB_MARKER_API_BASE_URL': 'rag.datalab_marker_api_base_url',
274 'DATALAB_MARKER_API_KEY': 'rag.datalab_marker_api_key',
275 'DATALAB_MARKER_DISABLE_IMAGE_EXTRACTION': 'rag.datalab_marker_disable_image_extraction',
276 'DATALAB_MARKER_FORCE_OCR': 'rag.datalab_marker_force_ocr',
277 'DATALAB_MARKER_FORMAT_LINES': 'rag.datalab_marker_format_lines',
278 'DATALAB_MARKER_OUTPUT_FORMAT': 'rag.datalab_marker_output_format',
279 'DATALAB_MARKER_PAGINATE': 'rag.datalab_marker_paginate',
280 'DATALAB_MARKER_SKIP_CACHE': 'rag.datalab_marker_skip_cache',
281 'DATALAB_MARKER_STRIP_EXISTING_OCR': 'rag.datalab_marker_strip_existing_ocr',
282 'DATALAB_MARKER_USE_LLM': 'rag.datalab_marker_use_llm',
283 'DDGS_BACKEND': 'web.search.ddgs_backend',
284 'DOCLING_API_KEY': 'rag.docling_api_key',
285 'DOCLING_PARAMS': 'rag.docling_params',
286 'DOCLING_SERVER_URL': 'rag.docling_server_url',
287 'DOCUMENT_INTELLIGENCE_ENDPOINT': 'rag.document_intelligence_endpoint',
288 'DOCUMENT_INTELLIGENCE_KEY': 'rag.document_intelligence_key',
289 'DOCUMENT_INTELLIGENCE_MODEL': 'rag.document_intelligence_model',
290 'ENABLE_ASYNC_EMBEDDING': 'rag.enable_async_embedding',
291 'ENABLE_GOOGLE_DRIVE_INTEGRATION': 'google_drive.enable',
292 'ENABLE_MARKDOWN_HEADER_TEXT_SPLITTER': 'rag.enable_markdown_header_text_splitter',
293 'ENABLE_ONEDRIVE_INTEGRATION': 'onedrive.enable',
294 'ENABLE_RAG_HYBRID_SEARCH': 'rag.enable_hybrid_search',
295 'ENABLE_RAG_HYBRID_SEARCH_ENRICHED_TEXTS': 'rag.enable_hybrid_search_enriched_texts',
296 'ENABLE_WEB_LOADER_SSL_VERIFICATION': 'web.loader.ssl_verification',
297 'ENABLE_WEB_SEARCH': 'web.search.enable',
298 'ENABLE_WEB_SEARCH_CONFIRMATION': 'web.search.confirmation.enable',
299 'WEB_SEARCH_CONFIRMATION_CONTENT': 'web.search.confirmation.content',
300 'EXA_API_KEY': 'web.search.exa_api_key',
301 'EXA_MAX_CONTENT_LENGTH': 'web.search.exa_max_content_length',
302 'EXTERNAL_DOCUMENT_LOADER_API_KEY': 'rag.external_document_loader_api_key',
303 'EXTERNAL_DOCUMENT_LOADER_HEADERS': 'rag.external_document_loader_headers',
304 'EXTERNAL_DOCUMENT_LOADER_URL': 'rag.external_document_loader_url',
305 'EXTERNAL_WEB_LOADER_API_KEY': 'web.loader.external_web_loader_api_key',
306 'EXTERNAL_WEB_LOADER_URL': 'web.loader.external_web_loader_url',
307 'EXTERNAL_WEB_SEARCH_API_KEY': 'web.search.external_web_search_api_key',
308 'EXTERNAL_WEB_SEARCH_URL': 'web.search.external_web_search_url',
309 'FILE_IMAGE_COMPRESSION_HEIGHT': 'file.image_compression_height',
310 'FILE_IMAGE_COMPRESSION_WIDTH': 'file.image_compression_width',
311 'FILE_MAX_COUNT': 'rag.file.max_count',
312 'FILE_MAX_SIZE': 'rag.file.max_size',
313 'FIRECRAWL_API_BASE_URL': 'web.loader.firecrawl_api_url',
314 'FIRECRAWL_API_KEY': 'web.loader.firecrawl_api_key',
315 'FIRECRAWL_TIMEOUT': 'web.loader.firecrawl_timeout',
316 'GOOGLE_PSE_API_KEY': 'web.search.google_pse_api_key',
317 'GOOGLE_PSE_ENGINE_ID': 'web.search.google_pse_engine_id',
318 'HYBRID_BM25_WEIGHT': 'rag.hybrid_bm25_weight',
319 'JINA_API_BASE_URL': 'web.search.jina_api_base_url',
320 'JINA_API_KEY': 'web.search.jina_api_key',
321 'KAGI_SEARCH_API_KEY': 'web.search.kagi_search_api_key',
322 'LINKUP_API_KEY': 'web.search.linkup_api_key',
323 'LINKUP_SEARCH_PARAMS': 'web.search.linkup_search_params',
324 'MINERU_API_KEY': 'rag.mineru_api_key',
325 'MINERU_API_MODE': 'rag.mineru_api_mode',
326 'MINERU_API_TIMEOUT': 'rag.mineru_api_timeout',
327 'MINERU_API_URL': 'rag.mineru_api_url',
328 'MINERU_FILE_EXTENSIONS': 'rag.mineru_file_extensions',
329 'MINERU_PARAMS': 'rag.mineru_params',
330 'MICROSOFT_WEB_IQ_API_BASE_URL': 'web.search.microsoft_web_iq_api_base_url',
331 'MICROSOFT_WEB_IQ_API_KEY': 'web.search.microsoft_web_iq_api_key',
332 'MICROSOFT_WEB_IQ_LANGUAGE': 'web.search.microsoft_web_iq_language',
333 'MISTRAL_OCR_API_BASE_URL': 'rag.mistral_ocr_api_base_url',
334 'MISTRAL_OCR_API_KEY': 'rag.mistral_ocr_api_key',
335 'MISTRAL_OCR_USE_BASE64': 'rag.mistral_ocr_use_base64',
336 'MOJEEK_SEARCH_API_KEY': 'web.search.mojeek_search_api_key',
337 'OLLAMA_CLOUD_WEB_SEARCH_API_KEY': 'web.search.ollama_cloud_api_key',
338 'PADDLEOCR_VL_BASE_URL': 'rag.paddleocr_vl_base_url',
339 'PADDLEOCR_VL_TOKEN': 'rag.paddleocr_vl_token',
340 'PDF_EXTRACT_IMAGES': 'rag.pdf_extract_images',
341 'PDF_LOADER_MODE': 'rag.pdf_loader_mode',
342 'PERPLEXITY_API_KEY': 'web.search.perplexity_api_key',
343 'PERPLEXITY_MODEL': 'web.search.perplexity_model',
344 'PERPLEXITY_SEARCH_API_URL': 'web.search.perplexity_search_api_url',
345 'PERPLEXITY_SEARCH_CONTEXT_USAGE': 'web.search.perplexity_search_context_usage',
346 'PLAYWRIGHT_TIMEOUT': 'web.loader.playwright_timeout',
347 'PLAYWRIGHT_WS_URL': 'web.loader.playwright_ws_url',
348 'RAG_AZURE_OPENAI_API_KEY': 'rag.azure_openai.api_key',
349 'RAG_AZURE_OPENAI_API_VERSION': 'rag.azure_openai.api_version',
350 'RAG_AZURE_OPENAI_BASE_URL': 'rag.azure_openai.base_url',
351 'RAG_EMBEDDING_BATCH_SIZE': 'rag.embedding_batch_size',
352 'RAG_EMBEDDING_CONCURRENT_REQUESTS': 'rag.embedding_concurrent_requests',
353 'RAG_EMBEDDING_ENGINE': 'rag.embedding_engine',
354 'RAG_EMBEDDING_MODEL': 'rag.embedding_model',
355 'RAG_TOKENIZER_MODEL': 'rag.tokenizer_model',
356 'RAG_EXTERNAL_RERANKER_API_KEY': 'rag.external_reranker_api_key',
357 'RAG_EXTERNAL_RERANKER_TIMEOUT': 'rag.external_reranker_timeout',
358 'RAG_EXTERNAL_RERANKER_URL': 'rag.external_reranker_url',
359 'RAG_FULL_CONTEXT': 'rag.full_context',
360 'RAG_OLLAMA_API_KEY': 'rag.ollama.api_key',
361 'RAG_OLLAMA_BASE_URL': 'rag.ollama.base_url',
362 'RAG_OPENAI_API_BASE_URL': 'rag.openai.api_base_url',
363 'RAG_OPENAI_API_KEY': 'rag.openai.api_key',
364 'RAG_RERANKING_BATCH_SIZE': 'rag.reranking_batch_size',
365 'RAG_RERANKING_ENGINE': 'rag.reranking_engine',
366 'RAG_RERANKING_MODEL': 'rag.reranking_model',
367 'RAG_TEMPLATE': 'rag.template',
368 'RELEVANCE_THRESHOLD': 'rag.relevance_threshold',
369 'SEARCHAPI_API_KEY': 'web.search.searchapi_api_key',
370 'SEARCHAPI_ENGINE': 'web.search.searchapi_engine',
371 'SEARXNG_LANGUAGE': 'web.search.searxng_language',
372 'SEARXNG_QUERY_URL': 'web.search.searxng_query_url',
373 'OPENSERP_BASE_URL': 'web.search.openserp_base_url',
374 'SERPAPI_API_KEY': 'web.search.serpapi_api_key',
375 'SERPAPI_ENGINE': 'web.search.serpapi_engine',
376 'SERPER_API_KEY': 'web.search.serper_api_key',
377 'SERPHOUSE_API_KEY': 'web.search.serphouse_api_key',
378 'SERPHOUSE_DOMAIN': 'web.search.serphouse_domain',
379 'SERPLY_API_KEY': 'web.search.serply_api_key',
380 'SERPSTACK_API_KEY': 'web.search.serpstack_api_key',
381 'SERPSTACK_HTTPS': 'web.search.serpstack_https',
382 'SOUGOU_API_SID': 'web.search.sougou_api_sid',
383 'SOUGOU_API_SK': 'web.search.sougou_api_sk',
384 'STAAN_API_KEY': 'web.search.staan_api_key',
385 'STAAN_MARKET': 'web.search.staan_market',
386 'STAAN_MAX_SNIPPETS': 'web.search.staan_max_snippets',
387 'TAVILY_API_KEY': 'web.search.tavily_api_key',
388 'TAVILY_EXTRACT_DEPTH': 'web.search.tavily_extract_depth',
389 'TEXT_SPLITTER': 'rag.text_splitter',
390 'TIKA_SERVER_URL': 'rag.tika_server_url',
391 'TIKA_SERVER_VERSION': 'rag.tika_server_version',
392 'TIKTOKEN_ENCODING_NAME': 'rag.tiktoken_encoding_name',
393 'TOP_K': 'rag.top_k',
394 'TOP_K_RERANKER': 'rag.top_k_reranker',
395 'USER_PERMISSIONS': 'user.permissions',
396 'WEBUI_URL': 'webui.url',
397 'WEB_FETCH_MAX_CONTENT_LENGTH': 'web.fetch.max_content_length',
398 'WEB_LOADER_CONCURRENT_REQUESTS': 'web.loader.concurrent_requests',
399 'WEB_LOADER_ENGINE': 'web.loader.engine',
400 'WEB_LOADER_TIMEOUT': 'web.loader.timeout',
401 'WEB_SEARCH_CONCURRENT_REQUESTS': 'web.search.concurrent_requests',
402 'WEB_SEARCH_DOMAIN_FILTER_LIST': 'web.search.domain.filter_list',
403 'WEB_SEARCH_ENGINE': 'web.search.engine',
404 'WEB_SEARCH_RESULT_COUNT': 'web.search.result_count',
405 'WEB_SEARCH_TRUST_ENV': 'web.search.trust_env',
406 'YACY_PASSWORD': 'web.search.yacy_password',
407 'YACY_QUERY_URL': 'web.search.yacy_query_url',
408 'YACY_USERNAME': 'web.search.yacy_username',
409 'YANDEX_WEB_SEARCH_API_KEY': 'web.search.yandex_web_search_api_key',
410 'YANDEX_WEB_SEARCH_CONFIG': 'web.search.yandex_web_search_config',
411 'YANDEX_WEB_SEARCH_URL': 'web.search.yandex_web_search_url',
412 'YOUCOM_API_KEY': 'web.search.youcom_api_key',
413 'YOUTUBE_LOADER_LANGUAGE': 'rag.youtube_loader_language',
414 'YOUTUBE_LOADER_PROXY_URL': 'rag.youtube_loader_proxy_url',
415}
418class RetrievalConfig(SimpleNamespace):
419 def __init__(self, values: dict):
420 super().__init__(**values)
421 object.__setattr__(self, '_updates', {})
423 def __setattr__(self, key: str, value):
424 if key.startswith('_'): 424 ↛ 425line 424 didn't jump to line 425 because the condition on line 424 was never true
425 object.__setattr__(self, key, value)
426 return
427 object.__setattr__(self, key, value)
428 if key in RETRIEVAL_CONFIG_KEYS: 428 ↛ exitline 428 didn't return from function '__setattr__' because the condition on line 428 was always true
429 self._updates[RETRIEVAL_CONFIG_KEYS[key]] = value
431 async def save(self) -> None:
432 if self._updates:
433 await Config.upsert(dict(self._updates))
434 self._updates.clear()
437async def get_config_values(key_map: dict[str, str]) -> dict:
438 values = await Config.get_many(*key_map.values())
439 return {field: values[storage_key] for field, storage_key in key_map.items() if storage_key in values}
442async def get_retrieval_config() -> RetrievalConfig:
443 return RetrievalConfig(await get_config_values(RETRIEVAL_CONFIG_KEYS))
446class CollectionNameForm(BaseModel):
447 collection_name: str | None = None
450class ProcessUrlForm(CollectionNameForm):
451 url: str
454class ProcessUrlResponse(BaseModel):
455 status: bool
456 type: str
457 name: str
458 url: str
459 collection_name: str | None = None
460 content: str | None = None
461 file: dict | None = None
464class SearchForm(BaseModel):
465 queries: list[str]
468@router.get('/embedding')
469async def get_embedding_config(request: Request, user=Depends(get_admin_user)):
470 config = await get_retrieval_config()
471 return {
472 'status': True,
473 'RAG_EMBEDDING_ENGINE': config.RAG_EMBEDDING_ENGINE,
474 'RAG_EMBEDDING_MODEL': config.RAG_EMBEDDING_MODEL,
475 'RAG_EMBEDDING_BATCH_SIZE': config.RAG_EMBEDDING_BATCH_SIZE,
476 'ENABLE_ASYNC_EMBEDDING': config.ENABLE_ASYNC_EMBEDDING,
477 'RAG_EMBEDDING_CONCURRENT_REQUESTS': config.RAG_EMBEDDING_CONCURRENT_REQUESTS,
478 'openai_config': {
479 'url': config.RAG_OPENAI_API_BASE_URL,
480 'key': config.RAG_OPENAI_API_KEY,
481 },
482 'ollama_config': {
483 'url': config.RAG_OLLAMA_BASE_URL,
484 'key': config.RAG_OLLAMA_API_KEY,
485 },
486 'azure_openai_config': {
487 'url': config.RAG_AZURE_OPENAI_BASE_URL,
488 'key': config.RAG_AZURE_OPENAI_API_KEY,
489 'version': config.RAG_AZURE_OPENAI_API_VERSION,
490 },
491 }
494class OpenAIConfigForm(BaseModel):
495 url: str | None = None
496 key: str | None = None
499class OllamaConfigForm(BaseModel):
500 url: str | None = None
501 key: str | None = None
504class AzureOpenAIConfigForm(BaseModel):
505 url: str | None = None
506 key: str | None = None
507 version: str | None = None
510class EmbeddingModelUpdateForm(BaseModel):
511 openai_config: OpenAIConfigForm | None = None
512 ollama_config: OllamaConfigForm | None = None
513 azure_openai_config: AzureOpenAIConfigForm | None = None
514 RAG_EMBEDDING_ENGINE: str
515 RAG_EMBEDDING_MODEL: str
516 RAG_EMBEDDING_BATCH_SIZE: int | None = 1
517 ENABLE_ASYNC_EMBEDDING: bool | None = True
518 RAG_EMBEDDING_CONCURRENT_REQUESTS: int | None = 0
521async def unload_embedding_model(request: Request):
522 config = await get_retrieval_config()
523 if config.RAG_EMBEDDING_ENGINE == '': 523 ↛ exitline 523 didn't return from function 'unload_embedding_model' because the condition on line 523 was always true
524 # unloads current internal embedding model and clears VRAM cache
525 request.app.state.ef = None
526 request.app.state.EMBEDDING_FUNCTION = None
527 import gc
529 gc.collect()
530 if DEVICE_TYPE == 'cuda': 530 ↛ 531line 530 didn't jump to line 531 because the condition on line 530 was never true
531 import torch
533 if torch.cuda.is_available():
534 torch.cuda.empty_cache()
537@router.post('/embedding/update')
538async def update_embedding_config(request: Request, form_data: EmbeddingModelUpdateForm, user=Depends(get_admin_user)):
539 if USE_SLIM and form_data.RAG_EMBEDDING_ENGINE == '': 539 ↛ 540line 539 didn't jump to line 540 because the condition on line 539 was never true
540 raise HTTPException(400, 'Slim requires an external embedding engine (openai, ollama, azure_openai).')
541 config = await get_retrieval_config()
542 log.info('Updating embedding model: %s to %s', config.RAG_EMBEDDING_MODEL, form_data.RAG_EMBEDDING_MODEL)
543 await unload_embedding_model(request)
544 try:
545 config.RAG_EMBEDDING_ENGINE = form_data.RAG_EMBEDDING_ENGINE
546 config.RAG_EMBEDDING_MODEL = form_data.RAG_EMBEDDING_MODEL.strip()
547 config.RAG_EMBEDDING_BATCH_SIZE = form_data.RAG_EMBEDDING_BATCH_SIZE
548 config.ENABLE_ASYNC_EMBEDDING = form_data.ENABLE_ASYNC_EMBEDDING
549 config.RAG_EMBEDDING_CONCURRENT_REQUESTS = form_data.RAG_EMBEDDING_CONCURRENT_REQUESTS
551 if config.RAG_EMBEDDING_ENGINE == 'openai' and form_data.openai_config is not None: 551 ↛ 552line 551 didn't jump to line 552 because the condition on line 551 was never true
552 config.RAG_OPENAI_API_BASE_URL = form_data.openai_config.url or ''
553 config.RAG_OPENAI_API_KEY = form_data.openai_config.key or ''
555 if config.RAG_EMBEDDING_ENGINE == 'ollama' and form_data.ollama_config is not None: 555 ↛ 556line 555 didn't jump to line 556 because the condition on line 555 was never true
556 config.RAG_OLLAMA_BASE_URL = form_data.ollama_config.url or ''
557 config.RAG_OLLAMA_API_KEY = form_data.ollama_config.key or ''
559 if config.RAG_EMBEDDING_ENGINE == 'azure_openai' and form_data.azure_openai_config is not None: 559 ↛ 560line 559 didn't jump to line 560 because the condition on line 559 was never true
560 config.RAG_AZURE_OPENAI_BASE_URL = form_data.azure_openai_config.url or ''
561 config.RAG_AZURE_OPENAI_API_KEY = form_data.azure_openai_config.key or ''
562 config.RAG_AZURE_OPENAI_API_VERSION = form_data.azure_openai_config.version or ''
564 request.app.state.ef = get_ef(
565 config.RAG_EMBEDDING_ENGINE,
566 config.RAG_EMBEDDING_MODEL,
567 )
569 request.app.state.EMBEDDING_FUNCTION = get_embedding_function(
570 config.RAG_EMBEDDING_ENGINE,
571 config.RAG_EMBEDDING_MODEL,
572 request.app.state.ef,
573 (
574 config.RAG_OPENAI_API_BASE_URL
575 if config.RAG_EMBEDDING_ENGINE == 'openai'
576 else (
577 config.RAG_OLLAMA_BASE_URL
578 if config.RAG_EMBEDDING_ENGINE == 'ollama'
579 else config.RAG_AZURE_OPENAI_BASE_URL
580 )
581 ),
582 (
583 config.RAG_OPENAI_API_KEY
584 if config.RAG_EMBEDDING_ENGINE == 'openai'
585 else (
586 config.RAG_OLLAMA_API_KEY
587 if config.RAG_EMBEDDING_ENGINE == 'ollama'
588 else config.RAG_AZURE_OPENAI_API_KEY
589 )
590 ),
591 config.RAG_EMBEDDING_BATCH_SIZE,
592 azure_api_version=(
593 config.RAG_AZURE_OPENAI_API_VERSION if config.RAG_EMBEDDING_ENGINE == 'azure_openai' else None
594 ),
595 enable_async=config.ENABLE_ASYNC_EMBEDDING,
596 concurrent_requests=config.RAG_EMBEDDING_CONCURRENT_REQUESTS,
597 )
599 await config.save()
600 return {
601 'status': True,
602 'RAG_EMBEDDING_ENGINE': config.RAG_EMBEDDING_ENGINE,
603 'RAG_EMBEDDING_MODEL': config.RAG_EMBEDDING_MODEL,
604 'RAG_EMBEDDING_BATCH_SIZE': config.RAG_EMBEDDING_BATCH_SIZE,
605 'ENABLE_ASYNC_EMBEDDING': config.ENABLE_ASYNC_EMBEDDING,
606 'RAG_EMBEDDING_CONCURRENT_REQUESTS': config.RAG_EMBEDDING_CONCURRENT_REQUESTS,
607 'openai_config': {
608 'url': config.RAG_OPENAI_API_BASE_URL,
609 'key': config.RAG_OPENAI_API_KEY,
610 },
611 'ollama_config': {
612 'url': config.RAG_OLLAMA_BASE_URL,
613 'key': config.RAG_OLLAMA_API_KEY,
614 },
615 'azure_openai_config': {
616 'url': config.RAG_AZURE_OPENAI_BASE_URL,
617 'key': config.RAG_AZURE_OPENAI_API_KEY,
618 'version': config.RAG_AZURE_OPENAI_API_VERSION,
619 },
620 }
621 except Exception as e:
622 log.exception(f'Problem updating embedding model: {e}')
623 raise HTTPException(
624 status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
625 detail=ERROR_MESSAGES.DEFAULT(e, 'Error updating embedding configuration'),
626 )
629@router.get('/config')
630async def get_rag_config(request: Request, user=Depends(get_admin_user)):
631 config = await get_retrieval_config()
632 await config.save()
633 return {
634 'status': True,
635 # RAG settings
636 'RAG_TEMPLATE': config.RAG_TEMPLATE,
637 'TOP_K': config.TOP_K,
638 'BYPASS_EMBEDDING_AND_RETRIEVAL': config.BYPASS_EMBEDDING_AND_RETRIEVAL,
639 'RAG_FULL_CONTEXT': config.RAG_FULL_CONTEXT,
640 # Hybrid search settings
641 'ENABLE_RAG_HYBRID_SEARCH': config.ENABLE_RAG_HYBRID_SEARCH,
642 'ENABLE_RAG_HYBRID_SEARCH_ENRICHED_TEXTS': config.ENABLE_RAG_HYBRID_SEARCH_ENRICHED_TEXTS,
643 'TOP_K_RERANKER': config.TOP_K_RERANKER,
644 'RELEVANCE_THRESHOLD': config.RELEVANCE_THRESHOLD,
645 'HYBRID_BM25_WEIGHT': config.HYBRID_BM25_WEIGHT,
646 # Content extraction settings
647 'CONTENT_EXTRACTION_ENGINE': config.CONTENT_EXTRACTION_ENGINE,
648 'CONTENT_EXTRACTION_SUPPORTED_MEDIA_MIME_TYPES': config.CONTENT_EXTRACTION_SUPPORTED_MEDIA_MIME_TYPES,
649 'PDF_EXTRACT_IMAGES': config.PDF_EXTRACT_IMAGES,
650 'PDF_LOADER_MODE': config.PDF_LOADER_MODE,
651 'DATALAB_MARKER_API_KEY': config.DATALAB_MARKER_API_KEY,
652 'DATALAB_MARKER_API_BASE_URL': config.DATALAB_MARKER_API_BASE_URL,
653 'DATALAB_MARKER_ADDITIONAL_CONFIG': config.DATALAB_MARKER_ADDITIONAL_CONFIG,
654 'DATALAB_MARKER_SKIP_CACHE': config.DATALAB_MARKER_SKIP_CACHE,
655 'DATALAB_MARKER_FORCE_OCR': config.DATALAB_MARKER_FORCE_OCR,
656 'DATALAB_MARKER_PAGINATE': config.DATALAB_MARKER_PAGINATE,
657 'DATALAB_MARKER_STRIP_EXISTING_OCR': config.DATALAB_MARKER_STRIP_EXISTING_OCR,
658 'DATALAB_MARKER_DISABLE_IMAGE_EXTRACTION': config.DATALAB_MARKER_DISABLE_IMAGE_EXTRACTION,
659 'DATALAB_MARKER_FORMAT_LINES': config.DATALAB_MARKER_FORMAT_LINES,
660 'DATALAB_MARKER_USE_LLM': config.DATALAB_MARKER_USE_LLM,
661 'DATALAB_MARKER_OUTPUT_FORMAT': config.DATALAB_MARKER_OUTPUT_FORMAT,
662 'EXTERNAL_DOCUMENT_LOADER_URL': config.EXTERNAL_DOCUMENT_LOADER_URL,
663 'EXTERNAL_DOCUMENT_LOADER_API_KEY': config.EXTERNAL_DOCUMENT_LOADER_API_KEY,
664 'EXTERNAL_DOCUMENT_LOADER_HEADERS': config.EXTERNAL_DOCUMENT_LOADER_HEADERS,
665 'TIKA_SERVER_URL': config.TIKA_SERVER_URL,
666 'TIKA_SERVER_VERSION': config.TIKA_SERVER_VERSION,
667 'DOCLING_SERVER_URL': config.DOCLING_SERVER_URL,
668 'DOCLING_API_KEY': config.DOCLING_API_KEY,
669 'DOCLING_PARAMS': config.DOCLING_PARAMS,
670 'DOCUMENT_INTELLIGENCE_ENDPOINT': config.DOCUMENT_INTELLIGENCE_ENDPOINT,
671 'DOCUMENT_INTELLIGENCE_KEY': config.DOCUMENT_INTELLIGENCE_KEY,
672 'DOCUMENT_INTELLIGENCE_MODEL': config.DOCUMENT_INTELLIGENCE_MODEL,
673 'MISTRAL_OCR_API_BASE_URL': config.MISTRAL_OCR_API_BASE_URL,
674 'MISTRAL_OCR_API_KEY': config.MISTRAL_OCR_API_KEY,
675 'MISTRAL_OCR_USE_BASE64': config.MISTRAL_OCR_USE_BASE64,
676 'PADDLEOCR_VL_BASE_URL': config.PADDLEOCR_VL_BASE_URL,
677 'PADDLEOCR_VL_TOKEN': config.PADDLEOCR_VL_TOKEN,
678 # MinerU settings
679 'MINERU_API_MODE': config.MINERU_API_MODE,
680 'MINERU_API_URL': config.MINERU_API_URL,
681 'MINERU_API_KEY': config.MINERU_API_KEY,
682 'MINERU_API_TIMEOUT': config.MINERU_API_TIMEOUT,
683 'MINERU_PARAMS': config.MINERU_PARAMS,
684 'MINERU_FILE_EXTENSIONS': config.MINERU_FILE_EXTENSIONS,
685 # Reranking settings
686 'RAG_RERANKING_MODEL': config.RAG_RERANKING_MODEL,
687 'RAG_RERANKING_ENGINE': config.RAG_RERANKING_ENGINE,
688 'RAG_RERANKING_BATCH_SIZE': config.RAG_RERANKING_BATCH_SIZE,
689 'RAG_EXTERNAL_RERANKER_URL': config.RAG_EXTERNAL_RERANKER_URL,
690 'RAG_EXTERNAL_RERANKER_API_KEY': config.RAG_EXTERNAL_RERANKER_API_KEY,
691 'RAG_EXTERNAL_RERANKER_TIMEOUT': config.RAG_EXTERNAL_RERANKER_TIMEOUT,
692 # Chunking settings
693 'TEXT_SPLITTER': config.TEXT_SPLITTER,
694 'RAG_TOKENIZER_MODEL': config.RAG_TOKENIZER_MODEL,
695 'ENABLE_MARKDOWN_HEADER_TEXT_SPLITTER': config.ENABLE_MARKDOWN_HEADER_TEXT_SPLITTER,
696 'CHUNK_SIZE': config.CHUNK_SIZE,
697 'CHUNK_MIN_SIZE_TARGET': config.CHUNK_MIN_SIZE_TARGET,
698 'CHUNK_OVERLAP': config.CHUNK_OVERLAP,
699 # File upload settings
700 'FILE_MAX_SIZE': config.FILE_MAX_SIZE,
701 'FILE_MAX_COUNT': config.FILE_MAX_COUNT,
702 'FILE_IMAGE_COMPRESSION_WIDTH': config.FILE_IMAGE_COMPRESSION_WIDTH,
703 'FILE_IMAGE_COMPRESSION_HEIGHT': config.FILE_IMAGE_COMPRESSION_HEIGHT,
704 'ALLOWED_FILE_EXTENSIONS': config.ALLOWED_FILE_EXTENSIONS,
705 # Integration settings
706 'ENABLE_GOOGLE_DRIVE_INTEGRATION': config.ENABLE_GOOGLE_DRIVE_INTEGRATION,
707 'ENABLE_ONEDRIVE_INTEGRATION': config.ENABLE_ONEDRIVE_INTEGRATION,
708 # Web search settings
709 'web': {
710 'ENABLE_WEB_SEARCH': config.ENABLE_WEB_SEARCH,
711 'ENABLE_WEB_SEARCH_CONFIRMATION': config.ENABLE_WEB_SEARCH_CONFIRMATION,
712 'WEB_SEARCH_CONFIRMATION_CONTENT': config.WEB_SEARCH_CONFIRMATION_CONTENT,
713 'WEB_SEARCH_ENGINE': config.WEB_SEARCH_ENGINE,
714 'WEB_SEARCH_TRUST_ENV': config.WEB_SEARCH_TRUST_ENV,
715 'WEB_SEARCH_RESULT_COUNT': config.WEB_SEARCH_RESULT_COUNT,
716 'WEB_SEARCH_CONCURRENT_REQUESTS': config.WEB_SEARCH_CONCURRENT_REQUESTS,
717 'WEB_FETCH_MAX_CONTENT_LENGTH': config.WEB_FETCH_MAX_CONTENT_LENGTH,
718 'WEB_LOADER_CONCURRENT_REQUESTS': config.WEB_LOADER_CONCURRENT_REQUESTS,
719 'WEB_SEARCH_DOMAIN_FILTER_LIST': config.WEB_SEARCH_DOMAIN_FILTER_LIST,
720 'BYPASS_WEB_SEARCH_EMBEDDING_AND_RETRIEVAL': config.BYPASS_WEB_SEARCH_EMBEDDING_AND_RETRIEVAL,
721 'BYPASS_WEB_SEARCH_WEB_LOADER': config.BYPASS_WEB_SEARCH_WEB_LOADER,
722 'OLLAMA_CLOUD_WEB_SEARCH_API_KEY': config.OLLAMA_CLOUD_WEB_SEARCH_API_KEY,
723 'SEARXNG_QUERY_URL': config.SEARXNG_QUERY_URL,
724 'SEARXNG_LANGUAGE': config.SEARXNG_LANGUAGE,
725 'OPENSERP_BASE_URL': config.OPENSERP_BASE_URL,
726 'YACY_QUERY_URL': config.YACY_QUERY_URL,
727 'YACY_USERNAME': config.YACY_USERNAME,
728 'YACY_PASSWORD': config.YACY_PASSWORD,
729 'GOOGLE_PSE_API_KEY': config.GOOGLE_PSE_API_KEY,
730 'GOOGLE_PSE_ENGINE_ID': config.GOOGLE_PSE_ENGINE_ID,
731 'BRAVE_SEARCH_API_KEY': config.BRAVE_SEARCH_API_KEY,
732 'BRAVE_SEARCH_CONTEXT_TOKENS': config.BRAVE_SEARCH_CONTEXT_TOKENS,
733 'KAGI_SEARCH_API_KEY': config.KAGI_SEARCH_API_KEY,
734 'MOJEEK_SEARCH_API_KEY': config.MOJEEK_SEARCH_API_KEY,
735 'BOCHA_SEARCH_API_KEY': config.BOCHA_SEARCH_API_KEY,
736 'SERPSTACK_API_KEY': config.SERPSTACK_API_KEY,
737 'SERPSTACK_HTTPS': config.SERPSTACK_HTTPS,
738 'SERPER_API_KEY': config.SERPER_API_KEY,
739 'SERPHOUSE_API_KEY': config.SERPHOUSE_API_KEY,
740 'SERPHOUSE_DOMAIN': config.SERPHOUSE_DOMAIN,
741 'SERPLY_API_KEY': config.SERPLY_API_KEY,
742 'DDGS_BACKEND': config.DDGS_BACKEND,
743 'TAVILY_API_KEY': config.TAVILY_API_KEY,
744 'STAAN_API_KEY': config.STAAN_API_KEY,
745 'STAAN_MARKET': config.STAAN_MARKET,
746 'STAAN_MAX_SNIPPETS': config.STAAN_MAX_SNIPPETS,
747 'SEARCHAPI_API_KEY': config.SEARCHAPI_API_KEY,
748 'SEARCHAPI_ENGINE': config.SEARCHAPI_ENGINE,
749 'SERPAPI_API_KEY': config.SERPAPI_API_KEY,
750 'SERPAPI_ENGINE': config.SERPAPI_ENGINE,
751 'JINA_API_KEY': config.JINA_API_KEY,
752 'JINA_API_BASE_URL': config.JINA_API_BASE_URL,
753 'BING_SEARCH_V7_ENDPOINT': config.BING_SEARCH_V7_ENDPOINT,
754 'BING_SEARCH_V7_SUBSCRIPTION_KEY': config.BING_SEARCH_V7_SUBSCRIPTION_KEY,
755 'EXA_API_KEY': config.EXA_API_KEY,
756 'EXA_MAX_CONTENT_LENGTH': config.EXA_MAX_CONTENT_LENGTH,
757 'PERPLEXITY_API_KEY': config.PERPLEXITY_API_KEY,
758 'PERPLEXITY_MODEL': config.PERPLEXITY_MODEL,
759 'PERPLEXITY_SEARCH_CONTEXT_USAGE': config.PERPLEXITY_SEARCH_CONTEXT_USAGE,
760 'PERPLEXITY_SEARCH_API_URL': config.PERPLEXITY_SEARCH_API_URL,
761 'MICROSOFT_WEB_IQ_API_BASE_URL': config.MICROSOFT_WEB_IQ_API_BASE_URL,
762 'MICROSOFT_WEB_IQ_API_KEY': config.MICROSOFT_WEB_IQ_API_KEY,
763 'MICROSOFT_WEB_IQ_LANGUAGE': config.MICROSOFT_WEB_IQ_LANGUAGE,
764 'SOUGOU_API_SID': config.SOUGOU_API_SID,
765 'SOUGOU_API_SK': config.SOUGOU_API_SK,
766 'WEB_LOADER_ENGINE': config.WEB_LOADER_ENGINE,
767 'WEB_LOADER_TIMEOUT': config.WEB_LOADER_TIMEOUT,
768 'ENABLE_WEB_LOADER_SSL_VERIFICATION': config.ENABLE_WEB_LOADER_SSL_VERIFICATION,
769 'PLAYWRIGHT_WS_URL': config.PLAYWRIGHT_WS_URL,
770 'PLAYWRIGHT_TIMEOUT': config.PLAYWRIGHT_TIMEOUT,
771 'FIRECRAWL_API_KEY': config.FIRECRAWL_API_KEY,
772 'FIRECRAWL_API_BASE_URL': config.FIRECRAWL_API_BASE_URL,
773 'FIRECRAWL_TIMEOUT': config.FIRECRAWL_TIMEOUT,
774 'TAVILY_EXTRACT_DEPTH': config.TAVILY_EXTRACT_DEPTH,
775 'EXTERNAL_WEB_SEARCH_URL': config.EXTERNAL_WEB_SEARCH_URL,
776 'EXTERNAL_WEB_SEARCH_API_KEY': config.EXTERNAL_WEB_SEARCH_API_KEY,
777 'EXTERNAL_WEB_LOADER_URL': config.EXTERNAL_WEB_LOADER_URL,
778 'EXTERNAL_WEB_LOADER_API_KEY': config.EXTERNAL_WEB_LOADER_API_KEY,
779 'YOUTUBE_LOADER_LANGUAGE': config.YOUTUBE_LOADER_LANGUAGE,
780 'YOUTUBE_LOADER_PROXY_URL': config.YOUTUBE_LOADER_PROXY_URL,
781 'YOUTUBE_LOADER_TRANSLATION': request.app.state.YOUTUBE_LOADER_TRANSLATION,
782 'YANDEX_WEB_SEARCH_URL': config.YANDEX_WEB_SEARCH_URL,
783 'YANDEX_WEB_SEARCH_API_KEY': config.YANDEX_WEB_SEARCH_API_KEY,
784 'YANDEX_WEB_SEARCH_CONFIG': config.YANDEX_WEB_SEARCH_CONFIG,
785 'YOUCOM_API_KEY': config.YOUCOM_API_KEY,
786 'LINKUP_API_KEY': config.LINKUP_API_KEY,
787 'LINKUP_SEARCH_PARAMS': config.LINKUP_SEARCH_PARAMS,
788 },
789 }
792class WebConfig(BaseModel):
793 ENABLE_WEB_SEARCH: bool | None = None
794 ENABLE_WEB_SEARCH_CONFIRMATION: bool | None = None
795 WEB_SEARCH_CONFIRMATION_CONTENT: str | None = None
796 WEB_SEARCH_ENGINE: str | None = None
797 WEB_SEARCH_TRUST_ENV: bool | None = None
798 WEB_SEARCH_RESULT_COUNT: int | None = None
799 WEB_SEARCH_CONCURRENT_REQUESTS: int | None = None
800 WEB_SEARCH_DOMAIN_FILTER_LIST: list[str] | None = []
801 WEB_FETCH_MAX_CONTENT_LENGTH: int | None = None
802 WEB_LOADER_CONCURRENT_REQUESTS: int | None = None
803 BYPASS_WEB_SEARCH_EMBEDDING_AND_RETRIEVAL: bool | None = None
804 BYPASS_WEB_SEARCH_WEB_LOADER: bool | None = None
805 OLLAMA_CLOUD_WEB_SEARCH_API_KEY: str | None = None
806 SEARXNG_QUERY_URL: str | None = None
807 SEARXNG_LANGUAGE: str | None = None
808 OPENSERP_BASE_URL: str | None = None
809 YACY_QUERY_URL: str | None = None
810 YACY_USERNAME: str | None = None
811 YACY_PASSWORD: str | None = None
812 GOOGLE_PSE_API_KEY: str | None = None
813 GOOGLE_PSE_ENGINE_ID: str | None = None
814 BRAVE_SEARCH_API_KEY: str | None = None
815 BRAVE_SEARCH_CONTEXT_TOKENS: int | None = None
816 KAGI_SEARCH_API_KEY: str | None = None
817 MOJEEK_SEARCH_API_KEY: str | None = None
818 BOCHA_SEARCH_API_KEY: str | None = None
819 SERPSTACK_API_KEY: str | None = None
820 SERPSTACK_HTTPS: bool | None = None
821 SERPER_API_KEY: str | None = None
822 SERPHOUSE_API_KEY: str | None = None
823 SERPHOUSE_DOMAIN: str | None = None
824 SERPLY_API_KEY: str | None = None
825 DDGS_BACKEND: str | None = None
826 TAVILY_API_KEY: str | None = None
827 STAAN_API_KEY: str | None = None
828 STAAN_MARKET: str | None = None
829 STAAN_MAX_SNIPPETS: int | None = None
830 SEARCHAPI_API_KEY: str | None = None
831 SEARCHAPI_ENGINE: str | None = None
832 SERPAPI_API_KEY: str | None = None
833 SERPAPI_ENGINE: str | None = None
834 JINA_API_KEY: str | None = None
835 JINA_API_BASE_URL: str | None = None
836 BING_SEARCH_V7_ENDPOINT: str | None = None
837 BING_SEARCH_V7_SUBSCRIPTION_KEY: str | None = None
838 EXA_API_KEY: str | None = None
839 EXA_MAX_CONTENT_LENGTH: int | None = Field(default=None, gt=0, strict=True)
840 PERPLEXITY_API_KEY: str | None = None
841 PERPLEXITY_MODEL: str | None = None
842 PERPLEXITY_SEARCH_CONTEXT_USAGE: str | None = None
843 PERPLEXITY_SEARCH_API_URL: str | None = None
844 MICROSOFT_WEB_IQ_API_BASE_URL: str | None = None
845 MICROSOFT_WEB_IQ_API_KEY: str | None = None
846 MICROSOFT_WEB_IQ_LANGUAGE: str | None = None
847 SOUGOU_API_SID: str | None = None
848 SOUGOU_API_SK: str | None = None
849 WEB_LOADER_ENGINE: str | None = None
850 WEB_LOADER_TIMEOUT: str | None = None
851 ENABLE_WEB_LOADER_SSL_VERIFICATION: bool | None = None
852 PLAYWRIGHT_WS_URL: str | None = None
853 PLAYWRIGHT_TIMEOUT: int | None = None
854 FIRECRAWL_API_KEY: str | None = None
855 FIRECRAWL_API_BASE_URL: str | None = None
856 FIRECRAWL_TIMEOUT: str | None = None
857 TAVILY_EXTRACT_DEPTH: str | None = None
858 EXTERNAL_WEB_SEARCH_URL: str | None = None
859 EXTERNAL_WEB_SEARCH_API_KEY: str | None = None
860 EXTERNAL_WEB_LOADER_URL: str | None = None
861 EXTERNAL_WEB_LOADER_API_KEY: str | None = None
862 YOUTUBE_LOADER_LANGUAGE: list[str] | None = None
863 YOUTUBE_LOADER_PROXY_URL: str | None = None
864 YOUTUBE_LOADER_TRANSLATION: str | None = None
865 YANDEX_WEB_SEARCH_URL: str | None = None
866 YANDEX_WEB_SEARCH_API_KEY: str | None = None
867 YANDEX_WEB_SEARCH_CONFIG: str | None = None
868 YOUCOM_API_KEY: str | None = None
869 LINKUP_API_KEY: str | None = None
870 LINKUP_SEARCH_PARAMS: dict | None = None
873class ConfigForm(BaseModel):
874 # RAG settings
875 RAG_TEMPLATE: str | None = None
876 TOP_K: int | None = None
877 BYPASS_EMBEDDING_AND_RETRIEVAL: bool | None = None
878 RAG_FULL_CONTEXT: bool | None = None
880 # Hybrid search settings
881 ENABLE_RAG_HYBRID_SEARCH: bool | None = None
882 ENABLE_RAG_HYBRID_SEARCH_ENRICHED_TEXTS: bool | None = None
883 TOP_K_RERANKER: int | None = None
884 RELEVANCE_THRESHOLD: float | None = None
885 HYBRID_BM25_WEIGHT: float | None = None
887 # Content extraction settings
888 CONTENT_EXTRACTION_ENGINE: str | None = None
889 CONTENT_EXTRACTION_SUPPORTED_MEDIA_MIME_TYPES: list[str] | None = None
890 PDF_EXTRACT_IMAGES: bool | None = None
891 PDF_LOADER_MODE: str | None = None
893 DATALAB_MARKER_API_KEY: str | None = None
894 DATALAB_MARKER_API_BASE_URL: str | None = None
895 DATALAB_MARKER_ADDITIONAL_CONFIG: str | None = None
896 DATALAB_MARKER_SKIP_CACHE: bool | None = None
897 DATALAB_MARKER_FORCE_OCR: bool | None = None
898 DATALAB_MARKER_PAGINATE: bool | None = None
899 DATALAB_MARKER_STRIP_EXISTING_OCR: bool | None = None
900 DATALAB_MARKER_DISABLE_IMAGE_EXTRACTION: bool | None = None
901 DATALAB_MARKER_FORMAT_LINES: bool | None = None
902 DATALAB_MARKER_USE_LLM: bool | None = None
903 DATALAB_MARKER_OUTPUT_FORMAT: str | None = None
905 EXTERNAL_DOCUMENT_LOADER_URL: str | None = None
906 EXTERNAL_DOCUMENT_LOADER_API_KEY: str | None = None
907 EXTERNAL_DOCUMENT_LOADER_HEADERS: dict | None = None
909 TIKA_SERVER_URL: str | None = None
910 TIKA_SERVER_VERSION: str | None = None
911 DOCLING_SERVER_URL: str | None = None
912 DOCLING_API_KEY: str | None = None
913 DOCLING_PARAMS: dict | None = None
914 DOCUMENT_INTELLIGENCE_ENDPOINT: str | None = None
915 DOCUMENT_INTELLIGENCE_KEY: str | None = None
916 DOCUMENT_INTELLIGENCE_MODEL: str | None = None
917 MISTRAL_OCR_API_BASE_URL: str | None = None
918 MISTRAL_OCR_API_KEY: str | None = None
919 MISTRAL_OCR_USE_BASE64: bool | None = None
920 PADDLEOCR_VL_BASE_URL: str | None = None
921 PADDLEOCR_VL_TOKEN: str | None = None
923 # MinerU settings
924 MINERU_API_MODE: str | None = None
925 MINERU_API_URL: str | None = None
926 MINERU_API_KEY: str | None = None
927 MINERU_API_TIMEOUT: int | None = None
928 MINERU_PARAMS: dict | None = None
929 MINERU_FILE_EXTENSIONS: list[str] | None = None
931 # Reranking settings
932 RAG_RERANKING_MODEL: str | None = None
933 RAG_RERANKING_ENGINE: str | None = None
934 RAG_RERANKING_BATCH_SIZE: int | None = None
935 RAG_EXTERNAL_RERANKER_URL: str | None = None
936 RAG_EXTERNAL_RERANKER_API_KEY: str | None = None
937 RAG_EXTERNAL_RERANKER_TIMEOUT: str | None = None
939 # Chunking settings
940 TEXT_SPLITTER: str | None = None
941 RAG_TOKENIZER_MODEL: str | None = None
942 ENABLE_MARKDOWN_HEADER_TEXT_SPLITTER: bool | None = None
943 CHUNK_SIZE: int | None = None
944 CHUNK_MIN_SIZE_TARGET: int | None = None
945 CHUNK_OVERLAP: int | None = None
947 # File upload settings
948 FILE_MAX_SIZE: Union[int, str | None] = None
949 FILE_MAX_COUNT: Union[int, str | None] = None
950 FILE_IMAGE_COMPRESSION_WIDTH: Union[int, str | None] = None
951 FILE_IMAGE_COMPRESSION_HEIGHT: Union[int, str | None] = None
952 ALLOWED_FILE_EXTENSIONS: list[str] | None = None
954 # Integration settings
955 ENABLE_GOOGLE_DRIVE_INTEGRATION: bool | None = None
956 ENABLE_ONEDRIVE_INTEGRATION: bool | None = None
958 # Web search settings
959 web: WebConfig | None = None
962@router.post('/config/update')
963async def update_rag_config(request: Request, form_data: ConfigForm, user=Depends(get_admin_user)):
964 # RAG settings
965 config = await get_retrieval_config()
966 if USE_SLIM: 966 ↛ 967line 966 didn't jump to line 967 because the condition on line 966 was never true
967 if (
968 form_data.web
969 and form_data.web.WEB_SEARCH_ENGINE == 'duckduckgo'
970 and config.WEB_SEARCH_ENGINE != 'duckduckgo'
971 ):
972 raise HTTPException(
973 400,
974 'DDGS is unavailable in slim. Configure another web search provider in Admin Settings > Web Search.',
975 )
976 if (
977 form_data.web
978 and form_data.web.WEB_LOADER_ENGINE == 'playwright'
979 and config.WEB_LOADER_ENGINE != 'playwright'
980 ):
981 raise HTTPException(
982 400, 'Playwright is unavailable in slim. Use basic HTTP fetching or an external web loader.'
983 )
984 if form_data.TEXT_SPLITTER == 'token_transformers' and config.TEXT_SPLITTER != 'token_transformers':
985 raise HTTPException(
986 400, 'Transformers tokenization is unavailable in slim. Use character or token splitting.'
987 )
988 reranker_engine = (
989 form_data.RAG_RERANKING_ENGINE
990 if form_data.RAG_RERANKING_ENGINE is not None
991 else config.RAG_RERANKING_ENGINE
992 )
993 reranker_model = (
994 form_data.RAG_RERANKING_MODEL if form_data.RAG_RERANKING_MODEL is not None else config.RAG_RERANKING_MODEL
995 )
996 if (
997 reranker_engine != 'external'
998 and reranker_model
999 and (reranker_engine != config.RAG_RERANKING_ENGINE or reranker_model != config.RAG_RERANKING_MODEL)
1000 ):
1001 raise HTTPException(
1002 400, 'Slim requires an external reranker, or an empty reranking model for cosine scoring.'
1003 )
1004 config.RAG_TEMPLATE = form_data.RAG_TEMPLATE if form_data.RAG_TEMPLATE is not None else config.RAG_TEMPLATE
1005 config.TOP_K = form_data.TOP_K if form_data.TOP_K is not None else config.TOP_K
1006 config.BYPASS_EMBEDDING_AND_RETRIEVAL = (
1007 form_data.BYPASS_EMBEDDING_AND_RETRIEVAL
1008 if form_data.BYPASS_EMBEDDING_AND_RETRIEVAL is not None
1009 else config.BYPASS_EMBEDDING_AND_RETRIEVAL
1010 )
1011 config.RAG_FULL_CONTEXT = (
1012 form_data.RAG_FULL_CONTEXT if form_data.RAG_FULL_CONTEXT is not None else config.RAG_FULL_CONTEXT
1013 )
1015 # Hybrid search settings
1016 config.ENABLE_RAG_HYBRID_SEARCH = (
1017 form_data.ENABLE_RAG_HYBRID_SEARCH
1018 if form_data.ENABLE_RAG_HYBRID_SEARCH is not None
1019 else config.ENABLE_RAG_HYBRID_SEARCH
1020 )
1021 config.ENABLE_RAG_HYBRID_SEARCH_ENRICHED_TEXTS = (
1022 form_data.ENABLE_RAG_HYBRID_SEARCH_ENRICHED_TEXTS
1023 if form_data.ENABLE_RAG_HYBRID_SEARCH_ENRICHED_TEXTS is not None
1024 else config.ENABLE_RAG_HYBRID_SEARCH_ENRICHED_TEXTS
1025 )
1027 config.TOP_K_RERANKER = form_data.TOP_K_RERANKER if form_data.TOP_K_RERANKER is not None else config.TOP_K_RERANKER
1028 config.RELEVANCE_THRESHOLD = (
1029 form_data.RELEVANCE_THRESHOLD if form_data.RELEVANCE_THRESHOLD is not None else config.RELEVANCE_THRESHOLD
1030 )
1031 config.HYBRID_BM25_WEIGHT = (
1032 form_data.HYBRID_BM25_WEIGHT if form_data.HYBRID_BM25_WEIGHT is not None else config.HYBRID_BM25_WEIGHT
1033 )
1035 # Content extraction settings
1036 config.CONTENT_EXTRACTION_ENGINE = (
1037 form_data.CONTENT_EXTRACTION_ENGINE
1038 if form_data.CONTENT_EXTRACTION_ENGINE is not None
1039 else config.CONTENT_EXTRACTION_ENGINE
1040 )
1041 config.CONTENT_EXTRACTION_SUPPORTED_MEDIA_MIME_TYPES = (
1042 form_data.CONTENT_EXTRACTION_SUPPORTED_MEDIA_MIME_TYPES
1043 if form_data.CONTENT_EXTRACTION_SUPPORTED_MEDIA_MIME_TYPES is not None
1044 else config.CONTENT_EXTRACTION_SUPPORTED_MEDIA_MIME_TYPES
1045 )
1046 config.PDF_EXTRACT_IMAGES = (
1047 form_data.PDF_EXTRACT_IMAGES if form_data.PDF_EXTRACT_IMAGES is not None else config.PDF_EXTRACT_IMAGES
1048 )
1049 config.PDF_LOADER_MODE = (
1050 form_data.PDF_LOADER_MODE if form_data.PDF_LOADER_MODE is not None else config.PDF_LOADER_MODE
1051 )
1052 config.DATALAB_MARKER_API_KEY = (
1053 form_data.DATALAB_MARKER_API_KEY
1054 if form_data.DATALAB_MARKER_API_KEY is not None
1055 else config.DATALAB_MARKER_API_KEY
1056 )
1057 config.DATALAB_MARKER_API_BASE_URL = (
1058 form_data.DATALAB_MARKER_API_BASE_URL
1059 if form_data.DATALAB_MARKER_API_BASE_URL is not None
1060 else config.DATALAB_MARKER_API_BASE_URL
1061 )
1062 config.DATALAB_MARKER_ADDITIONAL_CONFIG = (
1063 form_data.DATALAB_MARKER_ADDITIONAL_CONFIG
1064 if form_data.DATALAB_MARKER_ADDITIONAL_CONFIG is not None
1065 else config.DATALAB_MARKER_ADDITIONAL_CONFIG
1066 )
1067 config.DATALAB_MARKER_SKIP_CACHE = (
1068 form_data.DATALAB_MARKER_SKIP_CACHE
1069 if form_data.DATALAB_MARKER_SKIP_CACHE is not None
1070 else config.DATALAB_MARKER_SKIP_CACHE
1071 )
1072 config.DATALAB_MARKER_FORCE_OCR = (
1073 form_data.DATALAB_MARKER_FORCE_OCR
1074 if form_data.DATALAB_MARKER_FORCE_OCR is not None
1075 else config.DATALAB_MARKER_FORCE_OCR
1076 )
1077 config.DATALAB_MARKER_PAGINATE = (
1078 form_data.DATALAB_MARKER_PAGINATE
1079 if form_data.DATALAB_MARKER_PAGINATE is not None
1080 else config.DATALAB_MARKER_PAGINATE
1081 )
1082 config.DATALAB_MARKER_STRIP_EXISTING_OCR = (
1083 form_data.DATALAB_MARKER_STRIP_EXISTING_OCR
1084 if form_data.DATALAB_MARKER_STRIP_EXISTING_OCR is not None
1085 else config.DATALAB_MARKER_STRIP_EXISTING_OCR
1086 )
1087 config.DATALAB_MARKER_DISABLE_IMAGE_EXTRACTION = (
1088 form_data.DATALAB_MARKER_DISABLE_IMAGE_EXTRACTION
1089 if form_data.DATALAB_MARKER_DISABLE_IMAGE_EXTRACTION is not None
1090 else config.DATALAB_MARKER_DISABLE_IMAGE_EXTRACTION
1091 )
1092 config.DATALAB_MARKER_FORMAT_LINES = (
1093 form_data.DATALAB_MARKER_FORMAT_LINES
1094 if form_data.DATALAB_MARKER_FORMAT_LINES is not None
1095 else config.DATALAB_MARKER_FORMAT_LINES
1096 )
1097 config.DATALAB_MARKER_OUTPUT_FORMAT = (
1098 form_data.DATALAB_MARKER_OUTPUT_FORMAT
1099 if form_data.DATALAB_MARKER_OUTPUT_FORMAT is not None
1100 else config.DATALAB_MARKER_OUTPUT_FORMAT
1101 )
1102 config.DATALAB_MARKER_USE_LLM = (
1103 form_data.DATALAB_MARKER_USE_LLM
1104 if form_data.DATALAB_MARKER_USE_LLM is not None
1105 else config.DATALAB_MARKER_USE_LLM
1106 )
1107 config.EXTERNAL_DOCUMENT_LOADER_URL = (
1108 form_data.EXTERNAL_DOCUMENT_LOADER_URL
1109 if form_data.EXTERNAL_DOCUMENT_LOADER_URL is not None
1110 else config.EXTERNAL_DOCUMENT_LOADER_URL
1111 )
1112 config.EXTERNAL_DOCUMENT_LOADER_API_KEY = (
1113 form_data.EXTERNAL_DOCUMENT_LOADER_API_KEY
1114 if form_data.EXTERNAL_DOCUMENT_LOADER_API_KEY is not None
1115 else config.EXTERNAL_DOCUMENT_LOADER_API_KEY
1116 )
1117 config.EXTERNAL_DOCUMENT_LOADER_HEADERS = (
1118 form_data.EXTERNAL_DOCUMENT_LOADER_HEADERS
1119 if form_data.EXTERNAL_DOCUMENT_LOADER_HEADERS is not None
1120 else config.EXTERNAL_DOCUMENT_LOADER_HEADERS
1121 )
1122 config.TIKA_SERVER_URL = (
1123 form_data.TIKA_SERVER_URL if form_data.TIKA_SERVER_URL is not None else config.TIKA_SERVER_URL
1124 )
1125 config.TIKA_SERVER_VERSION = (
1126 form_data.TIKA_SERVER_VERSION if form_data.TIKA_SERVER_VERSION is not None else config.TIKA_SERVER_VERSION
1127 )
1128 config.DOCLING_SERVER_URL = (
1129 form_data.DOCLING_SERVER_URL if form_data.DOCLING_SERVER_URL is not None else config.DOCLING_SERVER_URL
1130 )
1131 config.DOCLING_API_KEY = (
1132 form_data.DOCLING_API_KEY if form_data.DOCLING_API_KEY is not None else config.DOCLING_API_KEY
1133 )
1134 config.DOCLING_PARAMS = form_data.DOCLING_PARAMS if form_data.DOCLING_PARAMS is not None else config.DOCLING_PARAMS
1135 config.DOCUMENT_INTELLIGENCE_ENDPOINT = (
1136 form_data.DOCUMENT_INTELLIGENCE_ENDPOINT
1137 if form_data.DOCUMENT_INTELLIGENCE_ENDPOINT is not None
1138 else config.DOCUMENT_INTELLIGENCE_ENDPOINT
1139 )
1140 config.DOCUMENT_INTELLIGENCE_KEY = (
1141 form_data.DOCUMENT_INTELLIGENCE_KEY
1142 if form_data.DOCUMENT_INTELLIGENCE_KEY is not None
1143 else config.DOCUMENT_INTELLIGENCE_KEY
1144 )
1145 config.DOCUMENT_INTELLIGENCE_MODEL = (
1146 form_data.DOCUMENT_INTELLIGENCE_MODEL
1147 if form_data.DOCUMENT_INTELLIGENCE_MODEL is not None
1148 else config.DOCUMENT_INTELLIGENCE_MODEL
1149 )
1151 config.MISTRAL_OCR_API_BASE_URL = (
1152 form_data.MISTRAL_OCR_API_BASE_URL
1153 if form_data.MISTRAL_OCR_API_BASE_URL is not None
1154 else config.MISTRAL_OCR_API_BASE_URL
1155 )
1156 config.MISTRAL_OCR_API_KEY = (
1157 form_data.MISTRAL_OCR_API_KEY if form_data.MISTRAL_OCR_API_KEY is not None else config.MISTRAL_OCR_API_KEY
1158 )
1159 config.MISTRAL_OCR_USE_BASE64 = (
1160 form_data.MISTRAL_OCR_USE_BASE64
1161 if form_data.MISTRAL_OCR_USE_BASE64 is not None
1162 else config.MISTRAL_OCR_USE_BASE64
1163 )
1164 config.PADDLEOCR_VL_BASE_URL = (
1165 form_data.PADDLEOCR_VL_BASE_URL if form_data.PADDLEOCR_VL_BASE_URL is not None else config.PADDLEOCR_VL_BASE_URL
1166 )
1167 config.PADDLEOCR_VL_TOKEN = (
1168 form_data.PADDLEOCR_VL_TOKEN if form_data.PADDLEOCR_VL_TOKEN is not None else config.PADDLEOCR_VL_TOKEN
1169 )
1171 # MinerU settings
1172 config.MINERU_API_MODE = (
1173 form_data.MINERU_API_MODE if form_data.MINERU_API_MODE is not None else config.MINERU_API_MODE
1174 )
1175 config.MINERU_API_URL = form_data.MINERU_API_URL if form_data.MINERU_API_URL is not None else config.MINERU_API_URL
1176 config.MINERU_API_KEY = form_data.MINERU_API_KEY if form_data.MINERU_API_KEY is not None else config.MINERU_API_KEY
1177 config.MINERU_API_TIMEOUT = (
1178 form_data.MINERU_API_TIMEOUT if form_data.MINERU_API_TIMEOUT is not None else config.MINERU_API_TIMEOUT
1179 )
1180 config.MINERU_PARAMS = form_data.MINERU_PARAMS if form_data.MINERU_PARAMS is not None else config.MINERU_PARAMS
1181 config.MINERU_FILE_EXTENSIONS = (
1182 form_data.MINERU_FILE_EXTENSIONS
1183 if form_data.MINERU_FILE_EXTENSIONS is not None
1184 else config.MINERU_FILE_EXTENSIONS
1185 )
1187 # Reranking settings
1188 if config.RAG_RERANKING_ENGINE == '':
1189 # Unloading the internal reranker and clear VRAM memory
1190 request.app.state.rf = None
1191 request.app.state.RERANKING_FUNCTION = None
1192 import gc
1194 gc.collect()
1195 if DEVICE_TYPE == 'cuda': 1195 ↛ 1196line 1195 didn't jump to line 1196 because the condition on line 1195 was never true
1196 import torch
1198 if torch.cuda.is_available():
1199 torch.cuda.empty_cache()
1200 config.RAG_RERANKING_ENGINE = (
1201 form_data.RAG_RERANKING_ENGINE if form_data.RAG_RERANKING_ENGINE is not None else config.RAG_RERANKING_ENGINE
1202 )
1204 config.RAG_EXTERNAL_RERANKER_URL = (
1205 form_data.RAG_EXTERNAL_RERANKER_URL
1206 if form_data.RAG_EXTERNAL_RERANKER_URL is not None
1207 else config.RAG_EXTERNAL_RERANKER_URL
1208 )
1210 config.RAG_EXTERNAL_RERANKER_API_KEY = (
1211 form_data.RAG_EXTERNAL_RERANKER_API_KEY
1212 if form_data.RAG_EXTERNAL_RERANKER_API_KEY is not None
1213 else config.RAG_EXTERNAL_RERANKER_API_KEY
1214 )
1216 config.RAG_EXTERNAL_RERANKER_TIMEOUT = (
1217 form_data.RAG_EXTERNAL_RERANKER_TIMEOUT
1218 if form_data.RAG_EXTERNAL_RERANKER_TIMEOUT is not None
1219 else config.RAG_EXTERNAL_RERANKER_TIMEOUT
1220 )
1222 config.RAG_RERANKING_BATCH_SIZE = (
1223 form_data.RAG_RERANKING_BATCH_SIZE
1224 if form_data.RAG_RERANKING_BATCH_SIZE is not None
1225 else config.RAG_RERANKING_BATCH_SIZE
1226 )
1228 if form_data.RAG_RERANKING_MODEL is not None:
1229 log.info('Updating reranking model: %s to %s', config.RAG_RERANKING_MODEL, form_data.RAG_RERANKING_MODEL)
1230 try:
1231 config.RAG_RERANKING_MODEL = (
1232 form_data.RAG_RERANKING_MODEL if form_data.RAG_RERANKING_MODEL is not None else config.RAG_RERANKING_MODEL
1233 )
1235 try:
1236 if config.ENABLE_RAG_HYBRID_SEARCH and not config.BYPASS_EMBEDDING_AND_RETRIEVAL:
1237 request.app.state.rf = get_rf(
1238 config.RAG_RERANKING_ENGINE,
1239 config.RAG_RERANKING_MODEL,
1240 config.RAG_EXTERNAL_RERANKER_URL,
1241 config.RAG_EXTERNAL_RERANKER_API_KEY,
1242 config.RAG_EXTERNAL_RERANKER_TIMEOUT,
1243 )
1245 request.app.state.RERANKING_FUNCTION = get_reranking_function(
1246 config.RAG_RERANKING_ENGINE,
1247 config.RAG_RERANKING_MODEL,
1248 request.app.state.rf,
1249 reranking_batch_size=config.RAG_RERANKING_BATCH_SIZE,
1250 )
1251 except Exception as e:
1252 log.error(f'Error loading reranking model: {e}')
1253 config.ENABLE_RAG_HYBRID_SEARCH = False
1254 except Exception as e:
1255 log.exception(f'Problem updating reranking model: {e}')
1256 raise HTTPException(
1257 status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
1258 detail=ERROR_MESSAGES.DEFAULT(e, 'Error updating reranking configuration'),
1259 )
1261 # Chunking settings
1262 config.TEXT_SPLITTER = form_data.TEXT_SPLITTER if form_data.TEXT_SPLITTER is not None else config.TEXT_SPLITTER
1263 config.ENABLE_MARKDOWN_HEADER_TEXT_SPLITTER = (
1264 form_data.ENABLE_MARKDOWN_HEADER_TEXT_SPLITTER
1265 if form_data.ENABLE_MARKDOWN_HEADER_TEXT_SPLITTER is not None
1266 else config.ENABLE_MARKDOWN_HEADER_TEXT_SPLITTER
1267 )
1268 config.CHUNK_SIZE = form_data.CHUNK_SIZE if form_data.CHUNK_SIZE is not None else config.CHUNK_SIZE
1269 config.CHUNK_MIN_SIZE_TARGET = (
1270 form_data.CHUNK_MIN_SIZE_TARGET if form_data.CHUNK_MIN_SIZE_TARGET is not None else config.CHUNK_MIN_SIZE_TARGET
1271 )
1272 config.CHUNK_OVERLAP = form_data.CHUNK_OVERLAP if form_data.CHUNK_OVERLAP is not None else config.CHUNK_OVERLAP
1273 config.RAG_TOKENIZER_MODEL = (
1274 form_data.RAG_TOKENIZER_MODEL.strip()
1275 if form_data.RAG_TOKENIZER_MODEL is not None
1276 else config.RAG_TOKENIZER_MODEL
1277 )
1279 # File upload settings
1280 # Empty string means "clear to None" (unlimited/no compression),
1281 # None means "don't change", int means "set to this value"
1282 if form_data.FILE_MAX_SIZE is not None:
1283 config.FILE_MAX_SIZE = None if form_data.FILE_MAX_SIZE == '' else form_data.FILE_MAX_SIZE
1284 if form_data.FILE_MAX_COUNT is not None:
1285 config.FILE_MAX_COUNT = None if form_data.FILE_MAX_COUNT == '' else form_data.FILE_MAX_COUNT
1286 if form_data.FILE_IMAGE_COMPRESSION_WIDTH is not None:
1287 config.FILE_IMAGE_COMPRESSION_WIDTH = (
1288 None if form_data.FILE_IMAGE_COMPRESSION_WIDTH == '' else form_data.FILE_IMAGE_COMPRESSION_WIDTH
1289 )
1290 if form_data.FILE_IMAGE_COMPRESSION_HEIGHT is not None:
1291 config.FILE_IMAGE_COMPRESSION_HEIGHT = (
1292 None if form_data.FILE_IMAGE_COMPRESSION_HEIGHT == '' else form_data.FILE_IMAGE_COMPRESSION_HEIGHT
1293 )
1295 config.ALLOWED_FILE_EXTENSIONS = (
1296 form_data.ALLOWED_FILE_EXTENSIONS
1297 if form_data.ALLOWED_FILE_EXTENSIONS is not None
1298 else config.ALLOWED_FILE_EXTENSIONS
1299 )
1301 # Integration settings
1302 config.ENABLE_GOOGLE_DRIVE_INTEGRATION = (
1303 form_data.ENABLE_GOOGLE_DRIVE_INTEGRATION
1304 if form_data.ENABLE_GOOGLE_DRIVE_INTEGRATION is not None
1305 else config.ENABLE_GOOGLE_DRIVE_INTEGRATION
1306 )
1307 config.ENABLE_ONEDRIVE_INTEGRATION = (
1308 form_data.ENABLE_ONEDRIVE_INTEGRATION
1309 if form_data.ENABLE_ONEDRIVE_INTEGRATION is not None
1310 else config.ENABLE_ONEDRIVE_INTEGRATION
1311 )
1313 if form_data.web is not None:
1314 # Web search settings
1315 config.ENABLE_WEB_SEARCH = form_data.web.ENABLE_WEB_SEARCH
1316 config.ENABLE_WEB_SEARCH_CONFIRMATION = form_data.web.ENABLE_WEB_SEARCH_CONFIRMATION
1317 config.WEB_SEARCH_CONFIRMATION_CONTENT = form_data.web.WEB_SEARCH_CONFIRMATION_CONTENT
1318 config.WEB_SEARCH_ENGINE = form_data.web.WEB_SEARCH_ENGINE
1319 config.WEB_SEARCH_TRUST_ENV = form_data.web.WEB_SEARCH_TRUST_ENV
1320 config.WEB_SEARCH_RESULT_COUNT = form_data.web.WEB_SEARCH_RESULT_COUNT
1321 config.WEB_SEARCH_CONCURRENT_REQUESTS = form_data.web.WEB_SEARCH_CONCURRENT_REQUESTS
1322 config.WEB_FETCH_MAX_CONTENT_LENGTH = form_data.web.WEB_FETCH_MAX_CONTENT_LENGTH
1323 config.WEB_LOADER_CONCURRENT_REQUESTS = form_data.web.WEB_LOADER_CONCURRENT_REQUESTS
1324 config.WEB_SEARCH_DOMAIN_FILTER_LIST = form_data.web.WEB_SEARCH_DOMAIN_FILTER_LIST
1325 config.BYPASS_WEB_SEARCH_EMBEDDING_AND_RETRIEVAL = form_data.web.BYPASS_WEB_SEARCH_EMBEDDING_AND_RETRIEVAL
1326 config.BYPASS_WEB_SEARCH_WEB_LOADER = form_data.web.BYPASS_WEB_SEARCH_WEB_LOADER
1327 config.OLLAMA_CLOUD_WEB_SEARCH_API_KEY = form_data.web.OLLAMA_CLOUD_WEB_SEARCH_API_KEY
1328 config.SEARXNG_QUERY_URL = form_data.web.SEARXNG_QUERY_URL
1329 config.SEARXNG_LANGUAGE = form_data.web.SEARXNG_LANGUAGE
1330 config.OPENSERP_BASE_URL = form_data.web.OPENSERP_BASE_URL
1331 config.YACY_QUERY_URL = form_data.web.YACY_QUERY_URL
1332 config.YACY_USERNAME = form_data.web.YACY_USERNAME
1333 config.YACY_PASSWORD = form_data.web.YACY_PASSWORD
1334 config.GOOGLE_PSE_API_KEY = form_data.web.GOOGLE_PSE_API_KEY
1335 config.GOOGLE_PSE_ENGINE_ID = form_data.web.GOOGLE_PSE_ENGINE_ID
1336 config.BRAVE_SEARCH_API_KEY = form_data.web.BRAVE_SEARCH_API_KEY
1337 if form_data.web.BRAVE_SEARCH_CONTEXT_TOKENS is not None: 1337 ↛ 1338line 1337 didn't jump to line 1338 because the condition on line 1337 was never true
1338 config.BRAVE_SEARCH_CONTEXT_TOKENS = form_data.web.BRAVE_SEARCH_CONTEXT_TOKENS
1339 config.KAGI_SEARCH_API_KEY = form_data.web.KAGI_SEARCH_API_KEY
1340 config.MOJEEK_SEARCH_API_KEY = form_data.web.MOJEEK_SEARCH_API_KEY
1341 config.BOCHA_SEARCH_API_KEY = form_data.web.BOCHA_SEARCH_API_KEY
1342 config.SERPSTACK_API_KEY = form_data.web.SERPSTACK_API_KEY
1343 config.SERPSTACK_HTTPS = form_data.web.SERPSTACK_HTTPS
1344 config.SERPER_API_KEY = form_data.web.SERPER_API_KEY
1345 config.SERPHOUSE_API_KEY = form_data.web.SERPHOUSE_API_KEY
1346 config.SERPHOUSE_DOMAIN = form_data.web.SERPHOUSE_DOMAIN
1347 config.SERPLY_API_KEY = form_data.web.SERPLY_API_KEY
1348 config.DDGS_BACKEND = form_data.web.DDGS_BACKEND
1349 config.TAVILY_API_KEY = form_data.web.TAVILY_API_KEY
1350 config.STAAN_API_KEY = form_data.web.STAAN_API_KEY
1351 config.STAAN_MARKET = form_data.web.STAAN_MARKET
1352 config.STAAN_MAX_SNIPPETS = form_data.web.STAAN_MAX_SNIPPETS
1353 config.SEARCHAPI_API_KEY = form_data.web.SEARCHAPI_API_KEY
1354 config.SEARCHAPI_ENGINE = form_data.web.SEARCHAPI_ENGINE
1355 config.SERPAPI_API_KEY = form_data.web.SERPAPI_API_KEY
1356 config.SERPAPI_ENGINE = form_data.web.SERPAPI_ENGINE
1357 config.JINA_API_KEY = form_data.web.JINA_API_KEY
1358 config.JINA_API_BASE_URL = form_data.web.JINA_API_BASE_URL
1359 config.BING_SEARCH_V7_ENDPOINT = form_data.web.BING_SEARCH_V7_ENDPOINT
1360 config.BING_SEARCH_V7_SUBSCRIPTION_KEY = form_data.web.BING_SEARCH_V7_SUBSCRIPTION_KEY
1361 config.EXA_API_KEY = form_data.web.EXA_API_KEY
1362 config.EXA_MAX_CONTENT_LENGTH = form_data.web.EXA_MAX_CONTENT_LENGTH
1363 config.PERPLEXITY_API_KEY = form_data.web.PERPLEXITY_API_KEY
1364 config.PERPLEXITY_MODEL = form_data.web.PERPLEXITY_MODEL
1365 config.PERPLEXITY_SEARCH_CONTEXT_USAGE = form_data.web.PERPLEXITY_SEARCH_CONTEXT_USAGE
1366 config.PERPLEXITY_SEARCH_API_URL = form_data.web.PERPLEXITY_SEARCH_API_URL
1367 config.MICROSOFT_WEB_IQ_API_BASE_URL = form_data.web.MICROSOFT_WEB_IQ_API_BASE_URL
1368 config.MICROSOFT_WEB_IQ_API_KEY = form_data.web.MICROSOFT_WEB_IQ_API_KEY
1369 config.MICROSOFT_WEB_IQ_LANGUAGE = form_data.web.MICROSOFT_WEB_IQ_LANGUAGE
1370 config.SOUGOU_API_SID = form_data.web.SOUGOU_API_SID
1371 config.SOUGOU_API_SK = form_data.web.SOUGOU_API_SK
1373 # Web loader settings
1374 config.WEB_LOADER_ENGINE = form_data.web.WEB_LOADER_ENGINE
1375 config.WEB_LOADER_TIMEOUT = form_data.web.WEB_LOADER_TIMEOUT
1377 config.ENABLE_WEB_LOADER_SSL_VERIFICATION = form_data.web.ENABLE_WEB_LOADER_SSL_VERIFICATION
1378 config.PLAYWRIGHT_WS_URL = form_data.web.PLAYWRIGHT_WS_URL
1379 config.PLAYWRIGHT_TIMEOUT = form_data.web.PLAYWRIGHT_TIMEOUT
1380 config.FIRECRAWL_API_KEY = form_data.web.FIRECRAWL_API_KEY
1381 config.FIRECRAWL_API_BASE_URL = form_data.web.FIRECRAWL_API_BASE_URL
1382 config.FIRECRAWL_TIMEOUT = form_data.web.FIRECRAWL_TIMEOUT
1383 config.EXTERNAL_WEB_SEARCH_URL = form_data.web.EXTERNAL_WEB_SEARCH_URL
1384 config.EXTERNAL_WEB_SEARCH_API_KEY = form_data.web.EXTERNAL_WEB_SEARCH_API_KEY
1385 config.EXTERNAL_WEB_LOADER_URL = form_data.web.EXTERNAL_WEB_LOADER_URL
1386 config.EXTERNAL_WEB_LOADER_API_KEY = form_data.web.EXTERNAL_WEB_LOADER_API_KEY
1387 config.TAVILY_EXTRACT_DEPTH = form_data.web.TAVILY_EXTRACT_DEPTH
1388 config.YOUTUBE_LOADER_LANGUAGE = form_data.web.YOUTUBE_LOADER_LANGUAGE
1389 config.YOUTUBE_LOADER_PROXY_URL = form_data.web.YOUTUBE_LOADER_PROXY_URL
1390 request.app.state.YOUTUBE_LOADER_TRANSLATION = form_data.web.YOUTUBE_LOADER_TRANSLATION
1391 config.YANDEX_WEB_SEARCH_URL = form_data.web.YANDEX_WEB_SEARCH_URL
1392 config.YANDEX_WEB_SEARCH_API_KEY = form_data.web.YANDEX_WEB_SEARCH_API_KEY
1393 config.YANDEX_WEB_SEARCH_CONFIG = form_data.web.YANDEX_WEB_SEARCH_CONFIG
1394 config.YOUCOM_API_KEY = form_data.web.YOUCOM_API_KEY
1395 config.LINKUP_API_KEY = form_data.web.LINKUP_API_KEY
1396 config.LINKUP_SEARCH_PARAMS = form_data.web.LINKUP_SEARCH_PARAMS
1398 await config.save()
1400 return {
1401 'status': True,
1402 # RAG settings
1403 'RAG_TEMPLATE': config.RAG_TEMPLATE,
1404 'TOP_K': config.TOP_K,
1405 'BYPASS_EMBEDDING_AND_RETRIEVAL': config.BYPASS_EMBEDDING_AND_RETRIEVAL,
1406 'RAG_FULL_CONTEXT': config.RAG_FULL_CONTEXT,
1407 # Hybrid search settings
1408 'ENABLE_RAG_HYBRID_SEARCH': config.ENABLE_RAG_HYBRID_SEARCH,
1409 'TOP_K_RERANKER': config.TOP_K_RERANKER,
1410 'RELEVANCE_THRESHOLD': config.RELEVANCE_THRESHOLD,
1411 'HYBRID_BM25_WEIGHT': config.HYBRID_BM25_WEIGHT,
1412 # Content extraction settings
1413 'CONTENT_EXTRACTION_ENGINE': config.CONTENT_EXTRACTION_ENGINE,
1414 'CONTENT_EXTRACTION_SUPPORTED_MEDIA_MIME_TYPES': config.CONTENT_EXTRACTION_SUPPORTED_MEDIA_MIME_TYPES,
1415 'PDF_EXTRACT_IMAGES': config.PDF_EXTRACT_IMAGES,
1416 'PDF_LOADER_MODE': config.PDF_LOADER_MODE,
1417 'DATALAB_MARKER_API_KEY': config.DATALAB_MARKER_API_KEY,
1418 'DATALAB_MARKER_API_BASE_URL': config.DATALAB_MARKER_API_BASE_URL,
1419 'DATALAB_MARKER_ADDITIONAL_CONFIG': config.DATALAB_MARKER_ADDITIONAL_CONFIG,
1420 'DATALAB_MARKER_SKIP_CACHE': config.DATALAB_MARKER_SKIP_CACHE,
1421 'DATALAB_MARKER_FORCE_OCR': config.DATALAB_MARKER_FORCE_OCR,
1422 'DATALAB_MARKER_PAGINATE': config.DATALAB_MARKER_PAGINATE,
1423 'DATALAB_MARKER_STRIP_EXISTING_OCR': config.DATALAB_MARKER_STRIP_EXISTING_OCR,
1424 'DATALAB_MARKER_DISABLE_IMAGE_EXTRACTION': config.DATALAB_MARKER_DISABLE_IMAGE_EXTRACTION,
1425 'DATALAB_MARKER_USE_LLM': config.DATALAB_MARKER_USE_LLM,
1426 'DATALAB_MARKER_OUTPUT_FORMAT': config.DATALAB_MARKER_OUTPUT_FORMAT,
1427 'EXTERNAL_DOCUMENT_LOADER_URL': config.EXTERNAL_DOCUMENT_LOADER_URL,
1428 'EXTERNAL_DOCUMENT_LOADER_API_KEY': config.EXTERNAL_DOCUMENT_LOADER_API_KEY,
1429 'EXTERNAL_DOCUMENT_LOADER_HEADERS': config.EXTERNAL_DOCUMENT_LOADER_HEADERS,
1430 'TIKA_SERVER_URL': config.TIKA_SERVER_URL,
1431 'TIKA_SERVER_VERSION': config.TIKA_SERVER_VERSION,
1432 'DOCLING_SERVER_URL': config.DOCLING_SERVER_URL,
1433 'DOCLING_API_KEY': config.DOCLING_API_KEY,
1434 'DOCLING_PARAMS': config.DOCLING_PARAMS,
1435 'DOCUMENT_INTELLIGENCE_ENDPOINT': config.DOCUMENT_INTELLIGENCE_ENDPOINT,
1436 'DOCUMENT_INTELLIGENCE_KEY': config.DOCUMENT_INTELLIGENCE_KEY,
1437 'DOCUMENT_INTELLIGENCE_MODEL': config.DOCUMENT_INTELLIGENCE_MODEL,
1438 'MISTRAL_OCR_API_BASE_URL': config.MISTRAL_OCR_API_BASE_URL,
1439 'MISTRAL_OCR_API_KEY': config.MISTRAL_OCR_API_KEY,
1440 'MISTRAL_OCR_USE_BASE64': config.MISTRAL_OCR_USE_BASE64,
1441 'PADDLEOCR_VL_BASE_URL': config.PADDLEOCR_VL_BASE_URL,
1442 'PADDLEOCR_VL_TOKEN': config.PADDLEOCR_VL_TOKEN,
1443 # MinerU settings
1444 'MINERU_API_MODE': config.MINERU_API_MODE,
1445 'MINERU_API_URL': config.MINERU_API_URL,
1446 'MINERU_API_KEY': config.MINERU_API_KEY,
1447 'MINERU_API_TIMEOUT': config.MINERU_API_TIMEOUT,
1448 'MINERU_PARAMS': config.MINERU_PARAMS,
1449 # Reranking settings
1450 'RAG_RERANKING_MODEL': config.RAG_RERANKING_MODEL,
1451 'RAG_RERANKING_ENGINE': config.RAG_RERANKING_ENGINE,
1452 'RAG_EXTERNAL_RERANKER_URL': config.RAG_EXTERNAL_RERANKER_URL,
1453 'RAG_EXTERNAL_RERANKER_API_KEY': config.RAG_EXTERNAL_RERANKER_API_KEY,
1454 'RAG_EXTERNAL_RERANKER_TIMEOUT': config.RAG_EXTERNAL_RERANKER_TIMEOUT,
1455 # Chunking settings
1456 'TEXT_SPLITTER': config.TEXT_SPLITTER,
1457 'RAG_TOKENIZER_MODEL': config.RAG_TOKENIZER_MODEL,
1458 'CHUNK_SIZE': config.CHUNK_SIZE,
1459 'CHUNK_MIN_SIZE_TARGET': config.CHUNK_MIN_SIZE_TARGET,
1460 'ENABLE_MARKDOWN_HEADER_TEXT_SPLITTER': config.ENABLE_MARKDOWN_HEADER_TEXT_SPLITTER,
1461 'CHUNK_OVERLAP': config.CHUNK_OVERLAP,
1462 # File upload settings
1463 'FILE_MAX_SIZE': config.FILE_MAX_SIZE,
1464 'FILE_MAX_COUNT': config.FILE_MAX_COUNT,
1465 'FILE_IMAGE_COMPRESSION_WIDTH': config.FILE_IMAGE_COMPRESSION_WIDTH,
1466 'FILE_IMAGE_COMPRESSION_HEIGHT': config.FILE_IMAGE_COMPRESSION_HEIGHT,
1467 'ALLOWED_FILE_EXTENSIONS': config.ALLOWED_FILE_EXTENSIONS,
1468 # Integration settings
1469 'ENABLE_GOOGLE_DRIVE_INTEGRATION': config.ENABLE_GOOGLE_DRIVE_INTEGRATION,
1470 'ENABLE_ONEDRIVE_INTEGRATION': config.ENABLE_ONEDRIVE_INTEGRATION,
1471 # Web search settings
1472 'web': {
1473 'ENABLE_WEB_SEARCH': config.ENABLE_WEB_SEARCH,
1474 'ENABLE_WEB_SEARCH_CONFIRMATION': config.ENABLE_WEB_SEARCH_CONFIRMATION,
1475 'WEB_SEARCH_CONFIRMATION_CONTENT': config.WEB_SEARCH_CONFIRMATION_CONTENT,
1476 'WEB_SEARCH_ENGINE': config.WEB_SEARCH_ENGINE,
1477 'WEB_SEARCH_TRUST_ENV': config.WEB_SEARCH_TRUST_ENV,
1478 'WEB_SEARCH_RESULT_COUNT': config.WEB_SEARCH_RESULT_COUNT,
1479 'WEB_SEARCH_CONCURRENT_REQUESTS': config.WEB_SEARCH_CONCURRENT_REQUESTS,
1480 'WEB_FETCH_MAX_CONTENT_LENGTH': config.WEB_FETCH_MAX_CONTENT_LENGTH,
1481 'WEB_LOADER_CONCURRENT_REQUESTS': config.WEB_LOADER_CONCURRENT_REQUESTS,
1482 'WEB_SEARCH_DOMAIN_FILTER_LIST': config.WEB_SEARCH_DOMAIN_FILTER_LIST,
1483 'BYPASS_WEB_SEARCH_EMBEDDING_AND_RETRIEVAL': config.BYPASS_WEB_SEARCH_EMBEDDING_AND_RETRIEVAL,
1484 'BYPASS_WEB_SEARCH_WEB_LOADER': config.BYPASS_WEB_SEARCH_WEB_LOADER,
1485 'OLLAMA_CLOUD_WEB_SEARCH_API_KEY': config.OLLAMA_CLOUD_WEB_SEARCH_API_KEY,
1486 'SEARXNG_QUERY_URL': config.SEARXNG_QUERY_URL,
1487 'SEARXNG_LANGUAGE': config.SEARXNG_LANGUAGE,
1488 'OPENSERP_BASE_URL': config.OPENSERP_BASE_URL,
1489 'YACY_QUERY_URL': config.YACY_QUERY_URL,
1490 'YACY_USERNAME': config.YACY_USERNAME,
1491 'YACY_PASSWORD': config.YACY_PASSWORD,
1492 'GOOGLE_PSE_API_KEY': config.GOOGLE_PSE_API_KEY,
1493 'GOOGLE_PSE_ENGINE_ID': config.GOOGLE_PSE_ENGINE_ID,
1494 'BRAVE_SEARCH_API_KEY': config.BRAVE_SEARCH_API_KEY,
1495 'BRAVE_SEARCH_CONTEXT_TOKENS': config.BRAVE_SEARCH_CONTEXT_TOKENS,
1496 'KAGI_SEARCH_API_KEY': config.KAGI_SEARCH_API_KEY,
1497 'MOJEEK_SEARCH_API_KEY': config.MOJEEK_SEARCH_API_KEY,
1498 'BOCHA_SEARCH_API_KEY': config.BOCHA_SEARCH_API_KEY,
1499 'SERPSTACK_API_KEY': config.SERPSTACK_API_KEY,
1500 'SERPSTACK_HTTPS': config.SERPSTACK_HTTPS,
1501 'SERPER_API_KEY': config.SERPER_API_KEY,
1502 'SERPHOUSE_API_KEY': config.SERPHOUSE_API_KEY,
1503 'SERPHOUSE_DOMAIN': config.SERPHOUSE_DOMAIN,
1504 'SERPLY_API_KEY': config.SERPLY_API_KEY,
1505 'TAVILY_API_KEY': config.TAVILY_API_KEY,
1506 'STAAN_API_KEY': config.STAAN_API_KEY,
1507 'STAAN_MARKET': config.STAAN_MARKET,
1508 'STAAN_MAX_SNIPPETS': config.STAAN_MAX_SNIPPETS,
1509 'SEARCHAPI_API_KEY': config.SEARCHAPI_API_KEY,
1510 'SEARCHAPI_ENGINE': config.SEARCHAPI_ENGINE,
1511 'SERPAPI_API_KEY': config.SERPAPI_API_KEY,
1512 'SERPAPI_ENGINE': config.SERPAPI_ENGINE,
1513 'JINA_API_KEY': config.JINA_API_KEY,
1514 'JINA_API_BASE_URL': config.JINA_API_BASE_URL,
1515 'BING_SEARCH_V7_ENDPOINT': config.BING_SEARCH_V7_ENDPOINT,
1516 'BING_SEARCH_V7_SUBSCRIPTION_KEY': config.BING_SEARCH_V7_SUBSCRIPTION_KEY,
1517 'EXA_API_KEY': config.EXA_API_KEY,
1518 'EXA_MAX_CONTENT_LENGTH': config.EXA_MAX_CONTENT_LENGTH,
1519 'PERPLEXITY_API_KEY': config.PERPLEXITY_API_KEY,
1520 'PERPLEXITY_MODEL': config.PERPLEXITY_MODEL,
1521 'PERPLEXITY_SEARCH_CONTEXT_USAGE': config.PERPLEXITY_SEARCH_CONTEXT_USAGE,
1522 'PERPLEXITY_SEARCH_API_URL': config.PERPLEXITY_SEARCH_API_URL,
1523 'MICROSOFT_WEB_IQ_API_BASE_URL': config.MICROSOFT_WEB_IQ_API_BASE_URL,
1524 'MICROSOFT_WEB_IQ_API_KEY': config.MICROSOFT_WEB_IQ_API_KEY,
1525 'MICROSOFT_WEB_IQ_LANGUAGE': config.MICROSOFT_WEB_IQ_LANGUAGE,
1526 'SOUGOU_API_SID': config.SOUGOU_API_SID,
1527 'SOUGOU_API_SK': config.SOUGOU_API_SK,
1528 'WEB_LOADER_ENGINE': config.WEB_LOADER_ENGINE,
1529 'WEB_LOADER_TIMEOUT': config.WEB_LOADER_TIMEOUT,
1530 'ENABLE_WEB_LOADER_SSL_VERIFICATION': config.ENABLE_WEB_LOADER_SSL_VERIFICATION,
1531 'PLAYWRIGHT_WS_URL': config.PLAYWRIGHT_WS_URL,
1532 'PLAYWRIGHT_TIMEOUT': config.PLAYWRIGHT_TIMEOUT,
1533 'FIRECRAWL_API_KEY': config.FIRECRAWL_API_KEY,
1534 'FIRECRAWL_API_BASE_URL': config.FIRECRAWL_API_BASE_URL,
1535 'FIRECRAWL_TIMEOUT': config.FIRECRAWL_TIMEOUT,
1536 'TAVILY_EXTRACT_DEPTH': config.TAVILY_EXTRACT_DEPTH,
1537 'EXTERNAL_WEB_SEARCH_URL': config.EXTERNAL_WEB_SEARCH_URL,
1538 'EXTERNAL_WEB_SEARCH_API_KEY': config.EXTERNAL_WEB_SEARCH_API_KEY,
1539 'EXTERNAL_WEB_LOADER_URL': config.EXTERNAL_WEB_LOADER_URL,
1540 'EXTERNAL_WEB_LOADER_API_KEY': config.EXTERNAL_WEB_LOADER_API_KEY,
1541 'YOUTUBE_LOADER_LANGUAGE': config.YOUTUBE_LOADER_LANGUAGE,
1542 'YOUTUBE_LOADER_PROXY_URL': config.YOUTUBE_LOADER_PROXY_URL,
1543 'YOUTUBE_LOADER_TRANSLATION': request.app.state.YOUTUBE_LOADER_TRANSLATION,
1544 'YANDEX_WEB_SEARCH_URL': config.YANDEX_WEB_SEARCH_URL,
1545 'YANDEX_WEB_SEARCH_API_KEY': config.YANDEX_WEB_SEARCH_API_KEY,
1546 'YANDEX_WEB_SEARCH_CONFIG': config.YANDEX_WEB_SEARCH_CONFIG,
1547 'YOUCOM_API_KEY': config.YOUCOM_API_KEY,
1548 'LINKUP_API_KEY': config.LINKUP_API_KEY,
1549 'LINKUP_SEARCH_PARAMS': config.LINKUP_SEARCH_PARAMS,
1550 },
1551 }
1554####################################
1555#
1556# Document process and retrieval
1557#
1558####################################
1561def can_merge_chunks(a: Document, b: Document) -> bool:
1562 if a.metadata.get('source') != b.metadata.get('source'):
1563 return False
1565 a_file_id = a.metadata.get('file_id')
1566 b_file_id = b.metadata.get('file_id')
1568 if a_file_id is not None and b_file_id is not None:
1569 return a_file_id == b_file_id
1571 return True
1574def merge_docs_to_target_size(
1575 request: Request,
1576 chunks: list[Document],
1577 config: RetrievalConfig,
1578) -> list[Document]:
1579 """
1580 Best-effort normalization of chunk sizes.
1582 Attempts to grow small chunks up to a desired minimum size,
1583 without exceeding the maximum size or crossing source/file
1584 boundaries.
1586 Uses forward merging first (absorb the next chunk), then
1587 backward merging (append into the previous emitted chunk)
1588 for undersized chunks that can't grow forward.
1589 """
1590 min_size = config.CHUNK_MIN_SIZE_TARGET
1591 max_size = config.CHUNK_SIZE
1593 if min_size <= 0: 1593 ↛ 1594line 1593 didn't jump to line 1594 because the condition on line 1593 was never true
1594 return chunks
1596 measure = get_splitter_length_function(request, config)
1598 def _merge_backward(result: list[Document], content: str, chunk: Document) -> bool:
1599 """Try to append content into the last emitted chunk. Returns True on success."""
1600 if not result: 1600 ↛ 1602line 1600 didn't jump to line 1602 because the condition on line 1600 was always true
1601 return False
1602 prev = result[-1]
1603 if not can_merge_chunks(prev, chunk):
1604 return False
1605 merged = f'{prev.page_content}\n\n{content}'
1606 if measure(merged) > max_size:
1607 return False
1608 result[-1] = Document(page_content=merged, metadata={**prev.metadata})
1609 return True
1611 def _emit(result: list[Document], content: str, chunk: Document) -> None:
1612 """Emit a chunk, trying backward merge first if it's undersized."""
1613 is_undersized = measure(content) < min_size
1614 if is_undersized and _merge_backward(result, content, chunk): 1614 ↛ 1615line 1614 didn't jump to line 1615 because the condition on line 1614 was never true
1615 return
1616 result.append(Document(page_content=content, metadata={**chunk.metadata}))
1618 result: list[Document] = []
1619 current_chunk: Document | None = None
1620 current_content: str = ''
1622 for next_chunk in chunks:
1623 if current_chunk is None: 1623 ↛ 1629line 1623 didn't jump to line 1629 because the condition on line 1623 was always true
1624 current_chunk = next_chunk
1625 current_content = next_chunk.page_content
1626 continue
1628 # Forward merge: absorb next chunk into current if undersized and fits
1629 merged_content = f'{current_content}\n\n{next_chunk.page_content}'
1630 can_merge_forward = (
1631 can_merge_chunks(current_chunk, next_chunk)
1632 and measure(current_content) < min_size
1633 and measure(merged_content) <= max_size
1634 )
1636 if can_merge_forward:
1637 current_content = merged_content
1638 else:
1639 _emit(result, current_content, current_chunk)
1640 current_chunk = next_chunk
1641 current_content = next_chunk.page_content
1643 if current_chunk is not None:
1644 _emit(result, current_content, current_chunk)
1646 return result
1649def get_transformers_tokenizer(request: Request, config: RetrievalConfig):
1650 if USE_SLIM:
1651 raise HTTPException(503, 'Transformers tokenization is unavailable in slim. Use character or token splitting.')
1652 if config.RAG_TOKENIZER_MODEL:
1653 from transformers import AutoTokenizer
1655 tokenizer_model = config.RAG_TOKENIZER_MODEL
1656 if not os.path.exists(tokenizer_model) and '/' not in tokenizer_model:
1657 tokenizer_model = f'sentence-transformers/{tokenizer_model}'
1659 cache_dir = os.getenv('SENTENCE_TRANSFORMERS_HOME') or os.getenv('HF_HUB_CACHE')
1660 local_files_only = not RAG_EMBEDDING_MODEL_AUTO_UPDATE
1661 tokenizer_key = (tokenizer_model, cache_dir, local_files_only)
1662 cached_tokenizer = getattr(request.app.state, 'transformers_tokenizer', None)
1663 if cached_tokenizer and cached_tokenizer[0] == tokenizer_key:
1664 return cached_tokenizer[1]
1666 tokenizer = AutoTokenizer.from_pretrained(
1667 tokenizer_model,
1668 cache_dir=cache_dir,
1669 trust_remote_code=RAG_EMBEDDING_MODEL_TRUST_REMOTE_CODE,
1670 local_files_only=local_files_only,
1671 )
1672 request.app.state.transformers_tokenizer = (tokenizer_key, tokenizer)
1673 return tokenizer
1675 tokenizer = getattr(getattr(request.app.state, 'ef', None), 'tokenizer', None)
1676 if tokenizer is not None:
1677 return tokenizer
1679 raise ValueError('Tokenizer model required for Token (Transformers) text splitter')
1682def get_splitter_length_function(
1683 request: Request,
1684 config: RetrievalConfig,
1685) -> Callable[[str], int]:
1686 if config.TEXT_SPLITTER == 'token': 1686 ↛ 1687line 1686 didn't jump to line 1687 because the condition on line 1686 was never true
1687 encoding = tiktoken.get_encoding(str(config.TIKTOKEN_ENCODING_NAME))
1688 return lambda text: len(encoding.encode(text, disallowed_special=TIKTOKEN_DISALLOWED_SPECIAL))
1690 if config.TEXT_SPLITTER == 'token_transformers': 1690 ↛ 1691line 1690 didn't jump to line 1691 because the condition on line 1690 was never true
1691 tokenizer = get_transformers_tokenizer(request, config)
1692 return lambda text: len(tokenizer.encode(text))
1694 return len
1697def filter_file_metadata(metadata: dict | None) -> dict:
1698 metadata = dict(metadata or {})
1699 data = metadata.pop('data', None)
1700 if isinstance(data, dict): 1700 ↛ 1702line 1700 didn't jump to line 1702 because the condition on line 1700 was always true
1701 metadata = {**filter_metadata(data), **metadata}
1702 return filter_metadata(metadata)
1705def save_docs_to_vector_db(
1706 request: Request,
1707 docs,
1708 collection_name,
1709 config: RetrievalConfig,
1710 metadata: dict | None = None,
1711 overwrite: bool = False,
1712 split: bool = True,
1713 add: bool = False,
1714 user=None,
1715) -> bool:
1716 def _get_docs_info(docs: list[Document]) -> str:
1717 docs_info = set()
1719 # Trying to select relevant metadata identifying the document.
1720 for doc in docs:
1721 metadata = getattr(doc, 'metadata', {})
1722 doc_name = metadata.get('name', '')
1723 if not doc_name:
1724 doc_name = metadata.get('title', '')
1725 if not doc_name:
1726 doc_name = metadata.get('source', '')
1727 if doc_name:
1728 docs_info.add(doc_name)
1730 return ', '.join(docs_info)
1732 log.debug('save_docs_to_vector_db: document %s %s', _get_docs_info(docs), collection_name)
1734 # Check if entries with the same hash (metadata.hash) already exist
1735 if metadata and 'hash' in metadata:
1736 result = get_vector_db_client().query(
1737 collection_name=collection_name,
1738 filter={'hash': metadata['hash']},
1739 )
1741 if result is not None and result.ids and len(result.ids) > 0: 1741 ↛ 1742line 1741 didn't jump to line 1742 because the condition on line 1741 was never true
1742 existing_doc_ids = result.ids[0]
1743 if existing_doc_ids:
1744 # Check if the existing document belongs to the same file
1745 # If same file_id, this is a re-add/reindex - allow it
1746 # If different file_id, this is a duplicate - block it
1747 existing_file_id = None
1748 if result.metadatas and result.metadatas[0]:
1749 existing_file_id = result.metadatas[0][0].get('file_id')
1751 if existing_file_id != metadata.get('file_id'):
1752 log.info('Document with hash %s already exists', metadata['hash'])
1753 raise ValueError(ERROR_MESSAGES.DUPLICATE_CONTENT)
1755 if split: 1755 ↛ 1819line 1755 didn't jump to line 1819 because the condition on line 1755 was always true
1756 if config.ENABLE_MARKDOWN_HEADER_TEXT_SPLITTER:
1757 log.info('Using markdown header text splitter')
1758 # Define headers to split on - covering most common markdown header levels
1759 markdown_splitter = MarkdownHeaderTextSplitter(
1760 headers_to_split_on=[
1761 ('#', 'Header 1'),
1762 ('##', 'Header 2'),
1763 ('###', 'Header 3'),
1764 ('####', 'Header 4'),
1765 ('#####', 'Header 5'),
1766 ('######', 'Header 6'),
1767 ],
1768 strip_headers=False, # Keep headers in content for context
1769 )
1771 split_docs = []
1772 for doc in docs:
1773 split_docs.extend(
1774 [
1775 Document(
1776 page_content=split_chunk.page_content,
1777 metadata={**doc.metadata},
1778 )
1779 for split_chunk in markdown_splitter.split_text(doc.page_content)
1780 ]
1781 )
1783 docs = split_docs
1784 if config.CHUNK_MIN_SIZE_TARGET > 0:
1785 docs = merge_docs_to_target_size(request, docs, config)
1787 if config.TEXT_SPLITTER in ['', 'character']:
1788 text_splitter = RecursiveCharacterTextSplitter(
1789 chunk_size=config.CHUNK_SIZE,
1790 chunk_overlap=config.CHUNK_OVERLAP,
1791 add_start_index=True,
1792 )
1793 docs = text_splitter.split_documents(docs)
1794 elif config.TEXT_SPLITTER == 'token': 1794 ↛ 1795line 1794 didn't jump to line 1795 because the condition on line 1794 was never true
1795 log.info('Using token text splitter: %s', config.TIKTOKEN_ENCODING_NAME)
1797 tiktoken.get_encoding(str(config.TIKTOKEN_ENCODING_NAME))
1798 text_splitter = TokenTextSplitter(
1799 encoding_name=str(config.TIKTOKEN_ENCODING_NAME),
1800 chunk_size=config.CHUNK_SIZE,
1801 chunk_overlap=config.CHUNK_OVERLAP,
1802 add_start_index=True,
1803 disallowed_special=TIKTOKEN_DISALLOWED_SPECIAL,
1804 )
1805 docs = text_splitter.split_documents(docs)
1806 elif config.TEXT_SPLITTER == 'token_transformers': 1806 ↛ 1807line 1806 didn't jump to line 1807 because the condition on line 1806 was never true
1807 log.info('Using transformers token text splitter')
1809 text_splitter = RecursiveCharacterTextSplitter(
1810 chunk_size=config.CHUNK_SIZE,
1811 chunk_overlap=config.CHUNK_OVERLAP,
1812 length_function=get_splitter_length_function(request, config),
1813 add_start_index=True,
1814 )
1815 docs = text_splitter.split_documents(docs)
1816 else:
1817 raise ValueError(ERROR_MESSAGES.DEFAULT('Invalid text splitter'))
1819 if len(docs) == 0:
1820 raise ValueError(ERROR_MESSAGES.EMPTY_CONTENT)
1822 texts = [sanitize_text_for_db(doc.page_content) for doc in docs]
1823 metadatas = [
1824 {
1825 **doc.metadata,
1826 **(metadata if metadata else {}),
1827 'embedding_config': {
1828 'engine': config.RAG_EMBEDDING_ENGINE,
1829 'model': config.RAG_EMBEDDING_MODEL,
1830 },
1831 }
1832 for doc in docs
1833 ]
1835 try:
1836 if get_vector_db_client().has_collection(collection_name=collection_name): 1836 ↛ 1837line 1836 didn't jump to line 1837 because the condition on line 1836 was never true
1837 log.info('collection %s already exists', collection_name)
1839 if overwrite:
1840 get_vector_db_client().delete_collection(collection_name=collection_name)
1841 log.info('deleting existing collection %s', collection_name)
1842 elif add is False:
1843 log.info('collection %s already exists, overwrite is False and add is False', collection_name)
1844 return True
1846 log.info('generating embeddings for %s', collection_name)
1847 embedding_function = get_embedding_function(
1848 config.RAG_EMBEDDING_ENGINE,
1849 config.RAG_EMBEDDING_MODEL,
1850 request.app.state.ef,
1851 (
1852 config.RAG_OPENAI_API_BASE_URL
1853 if config.RAG_EMBEDDING_ENGINE == 'openai'
1854 else (
1855 config.RAG_OLLAMA_BASE_URL
1856 if config.RAG_EMBEDDING_ENGINE == 'ollama'
1857 else config.RAG_AZURE_OPENAI_BASE_URL
1858 )
1859 ),
1860 (
1861 config.RAG_OPENAI_API_KEY
1862 if config.RAG_EMBEDDING_ENGINE == 'openai'
1863 else (
1864 config.RAG_OLLAMA_API_KEY
1865 if config.RAG_EMBEDDING_ENGINE == 'ollama'
1866 else config.RAG_AZURE_OPENAI_API_KEY
1867 )
1868 ),
1869 config.RAG_EMBEDDING_BATCH_SIZE,
1870 azure_api_version=(
1871 config.RAG_AZURE_OPENAI_API_VERSION if config.RAG_EMBEDDING_ENGINE == 'azure_openai' else None
1872 ),
1873 enable_async=config.ENABLE_ASYNC_EMBEDDING,
1874 concurrent_requests=config.RAG_EMBEDDING_CONCURRENT_REQUESTS,
1875 )
1877 # Run async embedding in sync context using the main event loop
1878 # This allows the main loop to stay responsive to health checks during long operations
1879 embedding_timeout = RAG_EMBEDDING_TIMEOUT
1881 future = asyncio.run_coroutine_threadsafe(
1882 embedding_function(
1883 list(map(lambda x: x.replace('\n', ' '), texts)),
1884 prefix=RAG_EMBEDDING_CONTENT_PREFIX,
1885 user=user,
1886 ),
1887 request.app.state.main_loop,
1888 )
1889 embeddings = future.result(timeout=embedding_timeout)
1890 log.info('embeddings generated %s for %s items', len(embeddings), len(texts))
1892 items = [
1893 {
1894 'id': str(uuid.uuid4()),
1895 'text': text,
1896 'vector': embeddings[idx],
1897 'metadata': metadatas[idx],
1898 }
1899 for idx, text in enumerate(texts)
1900 ]
1902 log.info('adding to collection %s', collection_name)
1903 get_vector_db_client().insert(
1904 collection_name=collection_name,
1905 items=items,
1906 )
1908 log.info('added %s items to collection %s', len(items), collection_name)
1909 return True
1910 except Exception as e:
1911 log.exception(e)
1912 raise e
1915class ProcessFileForm(BaseModel):
1916 file_id: str
1917 content: str | None = None
1918 collection_name: str | None = None
1921def has_vector_results(result) -> bool:
1922 return bool(result and result.ids and result.ids[0])
1925@router.post('/process/file')
1926async def process_file(
1927 request: Request,
1928 form_data: ProcessFileForm,
1929 user=Depends(get_verified_user),
1930 db: AsyncSession = Depends(get_async_session),
1931):
1932 """
1933 Process a file and save its content to the vector database.
1934 Note: granular session management is used to prevent connection pool exhaustion.
1935 The session is committed before external API calls, and updates use a fresh session.
1936 """
1937 config = await get_retrieval_config()
1938 if user.role == 'admin': 1938 ↛ 1941line 1938 didn't jump to line 1941 because the condition on line 1938 was always true
1939 file = await Files.get_file_by_id(form_data.file_id, db=db)
1940 else:
1941 file = await Files.get_file_by_id_and_user_id(form_data.file_id, user.id, db=db)
1943 if file:
1944 try:
1945 collection_name = form_data.collection_name
1946 file_collection_name = f'file-{file.id}'
1948 if collection_name is None: 1948 ↛ 1951line 1948 didn't jump to line 1951 because the condition on line 1948 was always true
1949 collection_name = file_collection_name
1950 else:
1951 await _validate_collection_access([collection_name], user, access_type='write')
1952 collection_names = [collection_name]
1954 if form_data.content:
1955 # Update the content in the file
1956 # Usage: /files/{file_id}/data/content/update, /files/ (audio file upload pipeline)
1958 try:
1959 # /files/{file_id}/data/content/update
1960 await ASYNC_VECTOR_DB_CLIENT.delete_collection(collection_name=file_collection_name)
1961 except Exception:
1962 # Audio file upload pipeline
1963 pass
1965 docs = [
1966 Document(
1967 page_content=form_data.content.replace('<br/>', '\n'),
1968 metadata={
1969 **filter_file_metadata(file.meta),
1970 'name': file.filename,
1971 'created_by': file.user_id,
1972 'file_id': file.id,
1973 'source': file.filename,
1974 },
1975 )
1976 ]
1978 text_content = form_data.content
1979 elif form_data.collection_name: 1979 ↛ 1985line 1979 didn't jump to line 1985 because the condition on line 1979 was never true
1980 # Add this file to a knowledge collection.
1981 # Usage: /knowledge/{id}/file/add, /knowledge/{id}/file/update
1982 # Reuse file-{id} chunks when they exist; otherwise restore file-{id}
1983 # from stored file content while adding the file to the knowledge collection.
1985 file_result = await ASYNC_VECTOR_DB_CLIENT.query(
1986 collection_name=file_collection_name, filter={'file_id': file.id}
1987 )
1988 stored_content = (file.data or {}).get('content')
1990 if has_vector_results(file_result):
1991 # Normal path: reuse the already-processed per-file chunks.
1992 docs = [
1993 Document(
1994 page_content=file_result.documents[0][idx],
1995 metadata=file_result.metadatas[0][idx],
1996 )
1997 for idx, id in enumerate(file_result.ids[0])
1998 ]
1999 elif stored_content is not None:
2000 # Repair path: vector chunks are missing, but SQL still has the file text.
2001 docs = [
2002 Document(
2003 page_content=stored_content,
2004 metadata={
2005 **filter_file_metadata(file.meta),
2006 'name': file.filename,
2007 'created_by': file.user_id,
2008 'file_id': file.id,
2009 'source': file.filename,
2010 },
2011 )
2012 ]
2013 collection_names.append(file_collection_name)
2014 else:
2015 raise ValueError(ERROR_MESSAGES.EMPTY_CONTENT)
2017 text_content = stored_content or ''
2018 else:
2019 # Process the file and save the content
2020 # Usage: /files/
2021 file_path = file.path
2022 if file_path: 2022 ↛ 2049line 2022 didn't jump to line 2049 because the condition on line 2022 was always true
2023 file_path = await asyncio.to_thread(Storage.get_file, file_path)
2024 loader_config = await get_loader_config()
2025 loader = build_loader_from_config(request, loader_config)
2026 loader.user = user
2027 loader.metadata = {
2028 'file_id': file.id,
2029 'file_name': file.filename,
2030 'file_content_type': file.meta.get('content_type'),
2031 }
2032 docs = await loader.aload(file.filename, file.meta.get('content_type'), file_path)
2034 docs = [
2035 Document(
2036 page_content=doc.page_content,
2037 metadata={
2038 **filter_file_metadata(file.meta),
2039 **filter_metadata(doc.metadata),
2040 'name': file.filename,
2041 'created_by': file.user_id,
2042 'file_id': file.id,
2043 'source': file.filename,
2044 },
2045 )
2046 for doc in docs
2047 ]
2048 else:
2049 docs = [
2050 Document(
2051 page_content=file.data.get('content', ''),
2052 metadata={
2053 **filter_file_metadata(file.meta),
2054 'name': file.filename,
2055 'created_by': file.user_id,
2056 'file_id': file.id,
2057 'source': file.filename,
2058 },
2059 )
2060 ]
2061 text_content = ' '.join([doc.page_content for doc in docs])
2063 log.debug('text_content: %s', text_content)
2064 await Files.update_file_data_by_id(
2065 file.id,
2066 {'content': text_content},
2067 db=db,
2068 )
2069 hash = calculate_sha256_string(text_content)
2071 if config.BYPASS_EMBEDDING_AND_RETRIEVAL: 2071 ↛ 2072line 2071 didn't jump to line 2072 because the condition on line 2071 was never true
2072 await Files.update_file_data_by_id(file.id, {'status': 'completed', 'error': None}, db=db)
2073 await Files.update_file_hash_by_id(file.id, hash, db=db)
2074 await publish_event(
2075 request,
2076 EVENTS.RETRIEVAL_CONTENT_PROCESSED,
2077 actor=user,
2078 subject_id=file.id,
2079 subject_type='file',
2080 data={'collection_name': None, 'filename': file.filename},
2081 )
2082 return {
2083 'status': True,
2084 'collection_name': None,
2085 'filename': file.filename,
2086 'content': text_content,
2087 }
2088 else:
2089 try:
2090 # Commit any pending changes before the slow embedding step.
2091 # Note: file is already a Pydantic model (not ORM), so no expunge needed.
2092 await db.commit()
2094 # External embedding API takes time (5-60s+).
2095 # Subsequent updates use fresh async sessions.
2096 # NOTE: save_docs_to_vector_db is a sync function that
2097 # calls asyncio.run_coroutine_threadsafe(..., main_loop).result()
2098 # which blocks the calling thread. We MUST run it in a
2099 # worker thread to avoid deadlocking the event loop.
2100 result = True
2101 for name in collection_names: 2101 ↛ 2116line 2101 didn't jump to line 2116 because the loop on line 2101 didn't complete
2102 result = await run_in_threadpool(
2103 save_docs_to_vector_db,
2104 request,
2105 docs=docs,
2106 collection_name=name,
2107 config=config,
2108 metadata={
2109 'file_id': file.id,
2110 'name': file.filename,
2111 'hash': hash,
2112 },
2113 add=(True if form_data.collection_name else False),
2114 user=user,
2115 )
2116 log.info('added %s items to collection %s', len(docs), collection_name)
2118 if result:
2119 # Fresh session for the final update.
2120 async with get_async_db() as session:
2121 await Files.update_file_metadata_by_id(
2122 file.id,
2123 {
2124 'collection_name': collection_name,
2125 },
2126 db=session,
2127 )
2129 await Files.update_file_data_by_id(
2130 file.id,
2131 {'status': 'completed', 'error': None},
2132 db=session,
2133 )
2134 await Files.update_file_hash_by_id(file.id, hash, db=session)
2136 await publish_event(
2137 request,
2138 EVENTS.RETRIEVAL_CONTENT_PROCESSED,
2139 actor=user,
2140 subject_id=file.id,
2141 subject_type='file',
2142 data={'collection_name': collection_name, 'filename': file.filename},
2143 )
2144 return {
2145 'status': True,
2146 'collection_name': collection_name,
2147 'filename': file.filename,
2148 'content': text_content,
2149 }
2150 else:
2151 raise Exception('Error saving document to vector database')
2152 except Exception as e:
2153 raise e
2155 except Exception as e:
2156 log.exception(e)
2157 # Fresh session for error status update.
2158 async with get_async_db() as session:
2159 await Files.update_file_data_by_id(
2160 file.id,
2161 {'status': 'failed', 'error': str(e)},
2162 db=session,
2163 )
2164 # Clear the hash so the file can be re-uploaded after fixing the issue
2165 await Files.update_file_hash_by_id(file.id, None, db=session)
2167 await publish_event(
2168 request,
2169 EVENTS.RETRIEVAL_CONTENT_PROCESS_FAILED,
2170 actor=user,
2171 subject_id=file.id,
2172 subject_type='file',
2173 data={
2174 'collection_name': collection_name,
2175 'filename': file.filename,
2176 'message': f'{file.filename}: {e}',
2177 },
2178 )
2180 if 'No pandoc was found' in str(e): 2180 ↛ 2181line 2180 didn't jump to line 2181 because the condition on line 2180 was never true
2181 raise HTTPException(
2182 status_code=status.HTTP_400_BAD_REQUEST,
2183 detail=ERROR_MESSAGES.PANDOC_NOT_INSTALLED,
2184 )
2185 else:
2186 raise HTTPException(
2187 status_code=status.HTTP_400_BAD_REQUEST,
2188 detail=str(e),
2189 )
2191 else:
2192 raise HTTPException(status_code=status.HTTP_404_NOT_FOUND, detail=ERROR_MESSAGES.NOT_FOUND)
2195class ProcessTextForm(BaseModel):
2196 name: str
2197 content: str
2198 collection_name: str | None = None
2201@router.post('/process/text')
2202async def process_text(
2203 request: Request,
2204 form_data: ProcessTextForm,
2205 user=Depends(get_verified_user),
2206):
2207 collection_name = form_data.collection_name
2208 if collection_name is None:
2209 collection_name = calculate_sha256_string(form_data.content)
2210 else:
2211 await _validate_collection_access([collection_name], user, access_type='write')
2213 docs = [
2214 Document(
2215 page_content=form_data.content,
2216 metadata={'name': form_data.name, 'created_by': user.id},
2217 )
2218 ]
2219 text_content = form_data.content
2220 log.debug('text_content: %s', text_content)
2222 config = await get_retrieval_config()
2223 result = await run_in_threadpool(save_docs_to_vector_db, request, docs, collection_name, config, user=user)
2224 if result:
2225 await publish_event(
2226 request,
2227 EVENTS.RETRIEVAL_CONTENT_PROCESSED,
2228 actor=user,
2229 subject_id=collection_name,
2230 subject_type='retrieval.collection',
2231 data={'name': form_data.name, 'content_preview': text_content[:300]},
2232 )
2233 return {
2234 'status': True,
2235 'collection_name': collection_name,
2236 'content': text_content,
2237 }
2238 else:
2239 raise HTTPException(
2240 status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
2241 detail=ERROR_MESSAGES.DEFAULT(),
2242 )
2245async def _fetch_url(url: str, max_size_mb: int | str | None) -> dict:
2246 await asyncio.to_thread(validate_url, url)
2247 max_bytes = None
2248 if max_size_mb:
2249 try:
2250 max_bytes = int(max_size_mb) * 1024 * 1024
2251 except (TypeError, ValueError):
2252 max_bytes = None
2254 async with get_ssrf_safe_session() as session:
2255 async with session.get(
2256 url, ssl=AIOHTTP_CLIENT_SESSION_SSL, allow_redirects=AIOHTTP_CLIENT_ALLOW_REDIRECTS
2257 ) as response:
2258 response.raise_for_status()
2260 content_type = response.headers.get('Content-Type', '')
2261 content_disposition = response.headers.get('Content-Disposition', '')
2262 content_length = response.headers.get('Content-Length')
2263 base_content_type = content_type.split(';')[0].strip().lower()
2264 is_attachment = content_disposition.split(';')[0].strip().lower() == 'attachment'
2266 chunks = []
2267 total = 0
2269 iterator = response.content.iter_chunked(64 * 1024)
2270 first_chunk = await anext(iterator, b'')
2272 if not is_attachment and base_content_type in {'text/html', 'application/xhtml+xml'}:
2273 return {'kind': 'web'}
2275 if not is_attachment and base_content_type in {'', 'application/octet-stream', 'binary/octet-stream'}:
2276 sample = first_chunk[:4096].lstrip().lower()
2277 if (
2278 sample.startswith((b'<!doctype html', b'<html', b'<head', b'<body', b'<?xml'))
2279 or b'<html' in sample[:1024]
2280 ):
2281 return {'kind': 'web'}
2283 if max_bytes and content_length:
2284 try:
2285 if int(content_length) > max_bytes:
2286 raise HTTPException(
2287 status_code=status.HTTP_413_REQUEST_ENTITY_TOO_LARGE,
2288 detail=ERROR_MESSAGES.FILE_TOO_LARGE(size=f'{max_size_mb} MB'),
2289 )
2290 except ValueError:
2291 pass
2293 if first_chunk:
2294 chunks.append(first_chunk)
2295 total += len(first_chunk)
2296 if max_bytes and total > max_bytes:
2297 raise HTTPException(
2298 status_code=status.HTTP_413_REQUEST_ENTITY_TOO_LARGE,
2299 detail=ERROR_MESSAGES.FILE_TOO_LARGE(size=f'{max_size_mb} MB'),
2300 )
2302 async for chunk in iterator:
2303 if not chunk:
2304 continue
2305 chunks.append(chunk)
2306 total += len(chunk)
2307 if max_bytes and total > max_bytes:
2308 raise HTTPException(
2309 status_code=status.HTTP_413_REQUEST_ENTITY_TOO_LARGE,
2310 detail=ERROR_MESSAGES.FILE_TOO_LARGE(size=f'{max_size_mb} MB'),
2311 )
2313 data = b''.join(chunks)
2315 image_mime = None
2316 try:
2317 from PIL import Image
2319 image = Image.open(io.BytesIO(data))
2320 image.verify()
2321 image_mime = Image.MIME.get(image.format) if image.format else None
2322 except Exception:
2323 image_mime = None
2325 if base_content_type.startswith('image/') and image_mime is None:
2326 raise HTTPException(
2327 status_code=status.HTTP_400_BAD_REQUEST,
2328 detail=ERROR_MESSAGES.DEFAULT('Invalid image content'),
2329 )
2331 filename = ''
2332 filename_star = re.search(r"filename\*=UTF-8''([^;]+)", content_disposition, re.IGNORECASE)
2333 filename_plain = re.search(r'filename="?([^";]+)"?', content_disposition, re.IGNORECASE)
2334 if filename_star:
2335 filename = unquote(filename_star.group(1))
2336 elif filename_plain:
2337 filename = filename_plain.group(1)
2338 if not filename:
2339 filename = os.path.basename(urlparse(url).path)
2340 filename = os.path.basename(filename or 'download')
2342 resolved_content_type = (
2343 image_mime or base_content_type or mimetypes.guess_type(filename)[0] or 'application/octet-stream'
2344 )
2345 if not os.path.splitext(filename)[1]:
2346 filename = f'{filename}{mimetypes.guess_extension(resolved_content_type) or ".bin"}'
2348 return {
2349 'kind': 'file',
2350 'data': data,
2351 'filename': filename,
2352 'content_type': resolved_content_type,
2353 }
2356@router.post('/process/url', response_model=ProcessUrlResponse)
2357async def process_url(
2358 request: Request,
2359 form_data: ProcessUrlForm,
2360 process: bool = Query(True, description='Whether to process and save the content'),
2361 user=Depends(get_verified_user),
2362):
2363 try:
2364 if is_youtube_url(form_data.url): 2364 ↛ 2365line 2364 didn't jump to line 2365 because the condition on line 2364 was never true
2365 result = await process_web(request, form_data, process=process, user=user)
2366 return {
2367 'status': True,
2368 'type': 'youtube',
2369 'name': form_data.url,
2370 'url': form_data.url,
2371 'collection_name': result.get('collection_name'),
2372 'content': result.get('content'),
2373 }
2375 config = await get_retrieval_config()
2376 url_result = await _fetch_url(form_data.url, config.FILE_MAX_SIZE)
2378 if url_result['kind'] == 'web':
2379 result = await process_web(request, form_data, process=process, user=user)
2380 return {
2381 'status': True,
2382 'type': 'web',
2383 'name': form_data.url,
2384 'url': form_data.url,
2385 'collection_name': result.get('collection_name'),
2386 'content': result.get('content'),
2387 }
2389 from open_webui.routers.files import upload_file_handler
2391 is_image = url_result['content_type'].startswith('image/')
2392 file = UploadFile(
2393 file=io.BytesIO(url_result['data']),
2394 filename=url_result['filename'],
2395 headers={'content-type': url_result['content_type']},
2396 )
2397 uploaded_file = await upload_file_handler(
2398 request,
2399 file=file,
2400 metadata={'source_url': form_data.url},
2401 process=process and not is_image,
2402 process_in_background=False,
2403 user=user,
2404 )
2405 file_data = uploaded_file.model_dump() if hasattr(uploaded_file, 'model_dump') else uploaded_file
2406 file_id = file_data.get('id') if isinstance(file_data, dict) else None
2407 if file_id:
2408 refreshed_file = await Files.get_file_by_id(file_id)
2409 if refreshed_file:
2410 file_data = refreshed_file.model_dump()
2411 return {
2412 'status': True,
2413 'type': 'image' if is_image else 'file',
2414 'name': url_result['filename'],
2415 'url': form_data.url,
2416 'collection_name': (file_data.get('meta') or {}).get('collection_name'),
2417 'file': file_data,
2418 }
2419 except HTTPException:
2420 raise
2421 except Exception as e:
2422 log.exception(e)
2423 raise HTTPException(
2424 status_code=status.HTTP_400_BAD_REQUEST,
2425 detail=ERROR_MESSAGES.DEFAULT(e, 'Error processing URL'),
2426 )
2429@router.post('/process/youtube')
2430@router.post('/process/web')
2431async def process_web(
2432 request: Request,
2433 form_data: ProcessUrlForm,
2434 process: bool = Query(True, description='Whether to process and save the content'),
2435 overwrite: bool = Query(True, description='Whether to overwrite existing collection'),
2436 user=Depends(get_verified_user),
2437):
2438 config = await get_retrieval_config()
2440 try:
2441 content, docs = await get_content_from_url(request, form_data.url)
2442 except HTTPException:
2443 raise
2444 except YoutubeTranscriptError as e:
2445 log.warning('YouTube transcript unavailable for %s: %s', form_data.url, e)
2446 raise HTTPException(
2447 status_code=status.HTTP_400_BAD_REQUEST,
2448 detail=str(e),
2449 )
2450 except Exception as e:
2451 log.exception(e)
2452 raise HTTPException(
2453 status_code=status.HTTP_400_BAD_REQUEST,
2454 detail=ERROR_MESSAGES.DEFAULT(e, f'Could not read content from {form_data.url}'),
2455 )
2457 try:
2458 log.debug('text_content: %s', content)
2460 if process:
2461 collection_name = form_data.collection_name
2462 if not collection_name:
2463 collection_name = calculate_sha256_string(form_data.url)[:63]
2464 else:
2465 await _validate_collection_access([collection_name], user, access_type='write')
2467 if not config.BYPASS_WEB_SEARCH_EMBEDDING_AND_RETRIEVAL:
2468 await run_in_threadpool(
2469 save_docs_to_vector_db,
2470 request,
2471 docs,
2472 collection_name,
2473 config,
2474 overwrite=overwrite,
2475 add=(not overwrite),
2476 user=user,
2477 )
2478 else:
2479 collection_name = None
2481 return {
2482 'status': True,
2483 'collection_name': collection_name,
2484 'filename': form_data.url,
2485 'content': content,
2486 'file': {
2487 'data': {
2488 'content': content,
2489 },
2490 'meta': {
2491 'name': form_data.url,
2492 'source': form_data.url,
2493 },
2494 },
2495 }
2496 else:
2497 return {
2498 'status': True,
2499 'content': content,
2500 }
2501 except HTTPException:
2502 raise
2503 except Exception as e:
2504 log.exception(e)
2505 raise HTTPException(
2506 status_code=status.HTTP_400_BAD_REQUEST,
2507 detail=ERROR_MESSAGES.DEFAULT(e, 'Error querying knowledge base'),
2508 )
2511async def search_web(request: Request, engine: str, query: str, user=None) -> list[SearchResult]:
2512 """Dispatch a web search query to the configured engine and return results.
2514 Providers that have been migrated to async (aiohttp) are awaited natively.
2515 Legacy sync providers are offloaded via ``asyncio.to_thread`` to avoid
2516 blocking the event loop.
2517 """
2519 # TODO: add playwright to search the web
2520 config = await get_retrieval_config()
2521 if engine == 'ollama_cloud':
2522 return await asyncio.to_thread(
2523 search_ollama_cloud,
2524 'https://ollama.com',
2525 config.OLLAMA_CLOUD_WEB_SEARCH_API_KEY,
2526 query,
2527 config.WEB_SEARCH_RESULT_COUNT,
2528 config.WEB_SEARCH_DOMAIN_FILTER_LIST,
2529 )
2530 elif engine == 'perplexity_search':
2531 if config.PERPLEXITY_API_KEY:
2532 return await asyncio.to_thread(
2533 search_perplexity_search,
2534 config.PERPLEXITY_API_KEY,
2535 query,
2536 config.WEB_SEARCH_RESULT_COUNT,
2537 config.WEB_SEARCH_DOMAIN_FILTER_LIST,
2538 config.PERPLEXITY_SEARCH_API_URL,
2539 user,
2540 )
2541 else:
2542 raise Exception('No PERPLEXITY_API_KEY found in environment variables')
2543 elif engine == 'searxng':
2544 if config.SEARXNG_QUERY_URL:
2545 searxng_kwargs = {'language': config.SEARXNG_LANGUAGE}
2546 return await search_searxng(
2547 config.SEARXNG_QUERY_URL,
2548 query,
2549 config.WEB_SEARCH_RESULT_COUNT,
2550 config.WEB_SEARCH_DOMAIN_FILTER_LIST,
2551 **searxng_kwargs,
2552 )
2553 else:
2554 raise Exception('No SEARXNG_QUERY_URL found in environment variables')
2555 elif engine == 'openserp':
2556 if config.OPENSERP_BASE_URL:
2557 return await search_openserp(
2558 config.OPENSERP_BASE_URL,
2559 query,
2560 config.WEB_SEARCH_RESULT_COUNT,
2561 config.WEB_SEARCH_DOMAIN_FILTER_LIST,
2562 )
2563 else:
2564 raise Exception('No OPENSERP_BASE_URL found in environment variables')
2565 elif engine == 'yacy':
2566 if config.YACY_QUERY_URL:
2567 return await asyncio.to_thread(
2568 search_yacy,
2569 config.YACY_QUERY_URL,
2570 config.YACY_USERNAME,
2571 config.YACY_PASSWORD,
2572 query,
2573 config.WEB_SEARCH_RESULT_COUNT,
2574 config.WEB_SEARCH_DOMAIN_FILTER_LIST,
2575 )
2576 else:
2577 raise Exception('No YACY_QUERY_URL found in environment variables')
2578 elif engine == 'google_pse':
2579 if config.GOOGLE_PSE_API_KEY and config.GOOGLE_PSE_ENGINE_ID:
2580 return await search_google_pse(
2581 config.GOOGLE_PSE_API_KEY,
2582 config.GOOGLE_PSE_ENGINE_ID,
2583 query,
2584 config.WEB_SEARCH_RESULT_COUNT,
2585 config.WEB_SEARCH_DOMAIN_FILTER_LIST,
2586 referer=config.WEBUI_URL,
2587 )
2588 else:
2589 raise Exception('No GOOGLE_PSE_API_KEY or GOOGLE_PSE_ENGINE_ID found in environment variables')
2590 elif engine == 'brave':
2591 if config.BRAVE_SEARCH_API_KEY:
2592 return await search_brave(
2593 config.BRAVE_SEARCH_API_KEY,
2594 query,
2595 config.WEB_SEARCH_RESULT_COUNT,
2596 config.WEB_SEARCH_DOMAIN_FILTER_LIST,
2597 )
2598 else:
2599 raise Exception('No BRAVE_SEARCH_API_KEY found in environment variables')
2600 elif engine == 'brave_llm_context':
2601 if config.BRAVE_SEARCH_API_KEY: 2601 ↛ anywhereline 2601 didn't jump anywhere: it always raised an exception.
2602 return await asyncio.to_thread(
2603 search_brave_llm_context,
2604 config.BRAVE_SEARCH_API_KEY,
2605 query,
2606 config.WEB_SEARCH_RESULT_COUNT,
2607 config.WEB_SEARCH_DOMAIN_FILTER_LIST,
2608 config.BRAVE_SEARCH_CONTEXT_TOKENS,
2609 )
2610 else:
2611 raise Exception('No BRAVE_SEARCH_API_KEY found in environment variables')
2612 elif engine == 'kagi':
2613 if config.KAGI_SEARCH_API_KEY:
2614 return await asyncio.to_thread(
2615 search_kagi,
2616 config.KAGI_SEARCH_API_KEY,
2617 query,
2618 config.WEB_SEARCH_RESULT_COUNT,
2619 config.WEB_SEARCH_DOMAIN_FILTER_LIST,
2620 )
2621 else:
2622 raise Exception('No KAGI_SEARCH_API_KEY found in environment variables')
2623 elif engine == 'mojeek':
2624 if config.MOJEEK_SEARCH_API_KEY:
2625 return await asyncio.to_thread(
2626 search_mojeek,
2627 config.MOJEEK_SEARCH_API_KEY,
2628 query,
2629 config.WEB_SEARCH_RESULT_COUNT,
2630 config.WEB_SEARCH_DOMAIN_FILTER_LIST,
2631 )
2632 else:
2633 raise Exception('No MOJEEK_SEARCH_API_KEY found in environment variables')
2634 elif engine == 'bocha':
2635 if config.BOCHA_SEARCH_API_KEY:
2636 return await asyncio.to_thread(
2637 search_bocha,
2638 config.BOCHA_SEARCH_API_KEY,
2639 query,
2640 config.WEB_SEARCH_RESULT_COUNT,
2641 config.WEB_SEARCH_DOMAIN_FILTER_LIST,
2642 )
2643 else:
2644 raise Exception('No BOCHA_SEARCH_API_KEY found in environment variables')
2645 elif engine == 'serpstack':
2646 if config.SERPSTACK_API_KEY:
2647 return await search_serpstack(
2648 config.SERPSTACK_API_KEY,
2649 query,
2650 config.WEB_SEARCH_RESULT_COUNT,
2651 config.WEB_SEARCH_DOMAIN_FILTER_LIST,
2652 https_enabled=config.SERPSTACK_HTTPS,
2653 )
2654 else:
2655 raise Exception('No SERPSTACK_API_KEY found in environment variables')
2656 elif engine == 'serper':
2657 if config.SERPER_API_KEY:
2658 return await search_serper(
2659 config.SERPER_API_KEY,
2660 query,
2661 config.WEB_SEARCH_RESULT_COUNT,
2662 config.WEB_SEARCH_DOMAIN_FILTER_LIST,
2663 )
2664 else:
2665 raise Exception('No SERPER_API_KEY found in environment variables')
2666 elif engine == 'serphouse':
2667 if config.SERPHOUSE_API_KEY:
2668 return await search_serphouse(
2669 config.SERPHOUSE_API_KEY,
2670 config.SERPHOUSE_DOMAIN,
2671 query,
2672 config.WEB_SEARCH_RESULT_COUNT,
2673 config.WEB_SEARCH_DOMAIN_FILTER_LIST,
2674 )
2675 else:
2676 raise Exception('No SERPHOUSE_API_KEY found in environment variables')
2677 elif engine == 'serply':
2678 if config.SERPLY_API_KEY:
2679 return await asyncio.to_thread(
2680 search_serply,
2681 config.SERPLY_API_KEY,
2682 query,
2683 config.WEB_SEARCH_RESULT_COUNT,
2684 filter_list=config.WEB_SEARCH_DOMAIN_FILTER_LIST,
2685 )
2686 else:
2687 raise Exception('No SERPLY_API_KEY found in environment variables')
2688 elif engine == 'duckduckgo':
2689 return await asyncio.to_thread(
2690 search_duckduckgo,
2691 query,
2692 config.WEB_SEARCH_RESULT_COUNT,
2693 config.WEB_SEARCH_DOMAIN_FILTER_LIST,
2694 concurrent_requests=config.WEB_SEARCH_CONCURRENT_REQUESTS,
2695 backend=config.DDGS_BACKEND,
2696 )
2697 elif engine == 'tavily':
2698 if config.TAVILY_API_KEY:
2699 return await asyncio.to_thread(
2700 search_tavily,
2701 config.TAVILY_API_KEY,
2702 query,
2703 config.WEB_SEARCH_RESULT_COUNT,
2704 config.WEB_SEARCH_DOMAIN_FILTER_LIST,
2705 )
2706 else:
2707 raise Exception('No TAVILY_API_KEY found in environment variables')
2708 elif engine == 'staan':
2709 if config.STAAN_API_KEY:
2710 return await asyncio.to_thread(
2711 search_staan,
2712 config.STAAN_API_KEY,
2713 query,
2714 config.WEB_SEARCH_RESULT_COUNT,
2715 config.WEB_SEARCH_DOMAIN_FILTER_LIST,
2716 market=config.STAAN_MARKET,
2717 max_snippets=config.STAAN_MAX_SNIPPETS,
2718 )
2719 else:
2720 raise Exception('No STAAN_API_KEY found in environment variables')
2721 elif engine == 'exa':
2722 if config.EXA_API_KEY:
2723 return await asyncio.to_thread(
2724 search_exa,
2725 config.EXA_API_KEY,
2726 query,
2727 config.WEB_SEARCH_RESULT_COUNT,
2728 config.WEB_SEARCH_DOMAIN_FILTER_LIST,
2729 max_content_length=config.EXA_MAX_CONTENT_LENGTH,
2730 )
2731 else:
2732 raise Exception('No EXA_API_KEY found in environment variables')
2733 elif engine == 'searchapi': 2733 ↛ 2745line 2733 didn't jump to line 2745 because the condition on line 2733 was always true
2734 if config.SEARCHAPI_API_KEY: 2734 ↛ anywhereline 2734 didn't jump anywhere: it always raised an exception.
2735 return await asyncio.to_thread(
2736 search_searchapi,
2737 config.SEARCHAPI_API_KEY,
2738 config.SEARCHAPI_ENGINE,
2739 query,
2740 config.WEB_SEARCH_RESULT_COUNT,
2741 config.WEB_SEARCH_DOMAIN_FILTER_LIST,
2742 )
2743 else:
2744 raise Exception('No SEARCHAPI_API_KEY found in environment variables')
2745 elif engine == 'serpapi':
2746 if config.SERPAPI_API_KEY:
2747 return await asyncio.to_thread(
2748 search_serpapi,
2749 config.SERPAPI_API_KEY,
2750 config.SERPAPI_ENGINE,
2751 query,
2752 config.WEB_SEARCH_RESULT_COUNT,
2753 config.WEB_SEARCH_DOMAIN_FILTER_LIST,
2754 )
2755 else:
2756 raise Exception('No SERPAPI_API_KEY found in environment variables')
2757 elif engine == 'jina':
2758 return await asyncio.to_thread(
2759 search_jina,
2760 config.JINA_API_KEY,
2761 query,
2762 config.WEB_SEARCH_RESULT_COUNT,
2763 config.JINA_API_BASE_URL,
2764 )
2765 elif engine == 'bing':
2766 return await asyncio.to_thread(
2767 search_bing,
2768 config.BING_SEARCH_V7_SUBSCRIPTION_KEY,
2769 config.BING_SEARCH_V7_ENDPOINT,
2770 str(DEFAULT_LOCALE),
2771 query,
2772 config.WEB_SEARCH_RESULT_COUNT,
2773 config.WEB_SEARCH_DOMAIN_FILTER_LIST,
2774 )
2775 elif engine == 'azure':
2776 if config.AZURE_AI_SEARCH_API_KEY and config.AZURE_AI_SEARCH_ENDPOINT and config.AZURE_AI_SEARCH_INDEX_NAME:
2777 return await asyncio.to_thread(
2778 search_azure,
2779 config.AZURE_AI_SEARCH_API_KEY,
2780 config.AZURE_AI_SEARCH_ENDPOINT,
2781 config.AZURE_AI_SEARCH_INDEX_NAME,
2782 query,
2783 config.WEB_SEARCH_RESULT_COUNT,
2784 config.WEB_SEARCH_DOMAIN_FILTER_LIST,
2785 )
2786 else:
2787 raise Exception(
2788 'AZURE_AI_SEARCH_API_KEY, AZURE_AI_SEARCH_ENDPOINT, and AZURE_AI_SEARCH_INDEX_NAME are required for Azure AI Search'
2789 )
2790 elif engine == 'perplexity':
2791 return await asyncio.to_thread(
2792 search_perplexity,
2793 config.PERPLEXITY_API_KEY,
2794 query,
2795 config.WEB_SEARCH_RESULT_COUNT,
2796 config.WEB_SEARCH_DOMAIN_FILTER_LIST,
2797 model=config.PERPLEXITY_MODEL,
2798 search_context_usage=config.PERPLEXITY_SEARCH_CONTEXT_USAGE,
2799 )
2800 elif engine == 'microsoft_web_iq':
2801 if config.MICROSOFT_WEB_IQ_API_KEY:
2802 return await asyncio.to_thread(
2803 search_microsoft_web_iq,
2804 config.MICROSOFT_WEB_IQ_API_BASE_URL,
2805 config.MICROSOFT_WEB_IQ_API_KEY,
2806 query,
2807 config.WEB_SEARCH_RESULT_COUNT,
2808 config.WEB_SEARCH_DOMAIN_FILTER_LIST,
2809 config.MICROSOFT_WEB_IQ_LANGUAGE,
2810 user,
2811 )
2812 else:
2813 raise Exception('No MICROSOFT_WEB_IQ_API_KEY found in environment variables')
2814 elif engine == 'sougou':
2815 if config.SOUGOU_API_SID and config.SOUGOU_API_SK:
2816 return await asyncio.to_thread(
2817 search_sougou,
2818 config.SOUGOU_API_SID,
2819 config.SOUGOU_API_SK,
2820 query,
2821 config.WEB_SEARCH_RESULT_COUNT,
2822 config.WEB_SEARCH_DOMAIN_FILTER_LIST,
2823 )
2824 else:
2825 raise Exception('No SOUGOU_API_SID or SOUGOU_API_SK found in environment variables')
2826 elif engine == 'firecrawl':
2827 return await asyncio.to_thread(
2828 search_firecrawl,
2829 config.FIRECRAWL_API_BASE_URL,
2830 config.FIRECRAWL_API_KEY,
2831 query,
2832 config.WEB_SEARCH_RESULT_COUNT,
2833 config.WEB_SEARCH_DOMAIN_FILTER_LIST,
2834 )
2835 elif engine == 'external':
2836 return await asyncio.to_thread(
2837 search_external,
2838 request,
2839 config.EXTERNAL_WEB_SEARCH_URL,
2840 config.EXTERNAL_WEB_SEARCH_API_KEY,
2841 query,
2842 config.WEB_SEARCH_RESULT_COUNT,
2843 config.WEB_SEARCH_DOMAIN_FILTER_LIST,
2844 user=user,
2845 )
2846 elif engine == 'yandex':
2847 return await asyncio.to_thread(
2848 search_yandex,
2849 request,
2850 config.YANDEX_WEB_SEARCH_URL,
2851 config.YANDEX_WEB_SEARCH_API_KEY,
2852 config.YANDEX_WEB_SEARCH_CONFIG,
2853 query,
2854 config.WEB_SEARCH_RESULT_COUNT,
2855 config.WEB_SEARCH_DOMAIN_FILTER_LIST,
2856 user=user,
2857 )
2858 elif engine == 'youcom':
2859 return await asyncio.to_thread(
2860 search_youcom,
2861 config.YOUCOM_API_KEY,
2862 query,
2863 config.WEB_SEARCH_RESULT_COUNT,
2864 config.WEB_SEARCH_DOMAIN_FILTER_LIST,
2865 )
2866 elif engine == 'linkup':
2867 if config.LINKUP_API_KEY:
2868 return await asyncio.to_thread(
2869 search_linkup,
2870 api_key=config.LINKUP_API_KEY,
2871 query=query,
2872 count=config.WEB_SEARCH_RESULT_COUNT,
2873 filter_list=config.WEB_SEARCH_DOMAIN_FILTER_LIST,
2874 params=config.LINKUP_SEARCH_PARAMS,
2875 )
2876 else:
2877 raise Exception('No LINKUP_API_KEY found in environment variables')
2878 else:
2879 raise Exception('No search engine API key found in environment variables')
2882@router.post('/process/web/search')
2883async def process_web_search(request: Request, form_data: SearchForm, user=Depends(get_verified_user)):
2884 config = await get_retrieval_config()
2885 if not config.ENABLE_WEB_SEARCH: 2885 ↛ 2891line 2885 didn't jump to line 2891 because the condition on line 2885 was always true
2886 raise HTTPException(
2887 status_code=status.HTTP_403_FORBIDDEN,
2888 detail=ERROR_MESSAGES.ACCESS_PROHIBITED,
2889 )
2891 if user.role != 'admin' and not await has_permission(user.id, 'features.web_search', config.USER_PERMISSIONS):
2892 raise HTTPException(
2893 status_code=status.HTTP_403_FORBIDDEN,
2894 detail=ERROR_MESSAGES.ACCESS_PROHIBITED,
2895 )
2897 urls = []
2898 result_items = []
2900 try:
2901 logging.debug('trying to web search with %s', (config.WEB_SEARCH_ENGINE, form_data.queries))
2903 # Use semaphore to limit concurrent requests based on WEB_SEARCH_CONCURRENT_REQUESTS
2904 # 0 or None = unlimited (previous behavior), positive number = limited concurrency
2905 # Set to 1 for sequential execution (rate-limited APIs like Brave free tier)
2906 concurrent_limit = config.WEB_SEARCH_CONCURRENT_REQUESTS
2908 if concurrent_limit:
2909 # Limited concurrency with semaphore
2910 semaphore = asyncio.Semaphore(concurrent_limit)
2912 async def search_query_with_semaphore(query):
2913 async with semaphore:
2914 return await search_web(
2915 request,
2916 config.WEB_SEARCH_ENGINE,
2917 query,
2918 user,
2919 )
2921 search_tasks = [search_query_with_semaphore(query) for query in form_data.queries]
2922 else:
2923 # Unlimited parallel execution
2924 search_tasks = [
2925 search_web(
2926 request,
2927 config.WEB_SEARCH_ENGINE,
2928 query,
2929 user,
2930 )
2931 for query in form_data.queries
2932 ]
2934 search_results = await asyncio.gather(*search_tasks)
2936 for result in search_results:
2937 if result:
2938 for item in result:
2939 if item and item.link:
2940 result_items.append(item)
2941 urls.append(item.link)
2943 urls = list(dict.fromkeys(urls))
2944 log.debug('urls: %s', urls)
2946 except Exception as e:
2947 log.exception('Web search failed')
2948 raise HTTPException(
2949 status.HTTP_400_BAD_REQUEST,
2950 detail=ERROR_MESSAGES.DEFAULT(e, ERROR_MESSAGES.WEB_SEARCH_ERROR),
2951 )
2953 if len(urls) == 0:
2954 raise HTTPException(
2955 status_code=status.HTTP_404_NOT_FOUND,
2956 detail=ERROR_MESSAGES.DEFAULT('No results found from web search'),
2957 )
2959 try:
2960 if config.BYPASS_WEB_SEARCH_WEB_LOADER:
2961 search_results = [item for result in search_results for item in result if result]
2963 docs = [
2964 Document(
2965 page_content=result.snippet,
2966 metadata={
2967 'source': result.link,
2968 'title': result.title,
2969 'snippet': result.snippet,
2970 'link': result.link,
2971 },
2972 )
2973 for result in search_results
2974 if hasattr(result, 'snippet') and result.snippet is not None
2975 ]
2976 else:
2977 loader_config = await get_loader_config()
2978 loader = get_web_loader(
2979 urls,
2980 verify_ssl=loader_config.get('web_loader_ssl_verification'),
2981 requests_per_second=loader_config.get('web_loader_concurrent_requests'),
2982 trust_env=loader_config.get('web_search_trust_env'),
2983 loader_config=loader_config,
2984 )
2985 docs = await loader.aload()
2987 urls = [
2988 doc.metadata.get('source') for doc in docs if doc.metadata.get('source')
2989 ] # only keep the urls returned by the loader
2990 url_set = set(urls)
2991 result_items = [
2992 dict(item) for item in result_items if item.link in url_set
2993 ] # only keep the search results that have been loaded
2995 if config.BYPASS_WEB_SEARCH_EMBEDDING_AND_RETRIEVAL:
2996 return {
2997 'status': True,
2998 'collection_name': None,
2999 'filenames': urls,
3000 'items': result_items,
3001 'docs': [
3002 {
3003 'content': doc.page_content,
3004 'metadata': doc.metadata,
3005 }
3006 for doc in docs
3007 ],
3008 'loaded_count': len(docs),
3009 }
3010 else:
3011 # Create a single collection for all documents
3012 # Bind the ephemeral collection to its owner so filter_accessible_collections can scope it per-user.
3013 collection_name = f'web-search-{user.id}-{calculate_sha256_string("-".join(form_data.queries))}'[:63]
3015 try:
3016 await run_in_threadpool(
3017 save_docs_to_vector_db,
3018 request,
3019 docs,
3020 collection_name,
3021 config,
3022 overwrite=True,
3023 user=user,
3024 )
3025 except Exception as e:
3026 # Surface the failure instead of returning an unusable collection
3027 log.exception(f'Error saving web search results to vector DB: {e}')
3028 raise HTTPException(
3029 status.HTTP_500_INTERNAL_SERVER_ERROR,
3030 detail='Failed to embed and store the retrieved web pages. Check the embedding configuration in Admin Settings > Documents.',
3031 )
3033 return {
3034 'status': True,
3035 'collection_names': [collection_name],
3036 'items': result_items,
3037 'filenames': urls,
3038 'loaded_count': len(docs),
3039 }
3040 except HTTPException:
3041 raise
3042 except Exception as e:
3043 log.exception('Web search content loading failed')
3044 raise HTTPException(
3045 status.HTTP_400_BAD_REQUEST,
3046 detail=ERROR_MESSAGES.DEFAULT(e, ERROR_MESSAGES.WEB_SEARCH_ERROR),
3047 )
3050async def _validate_collection_access(collection_names: list[str], user, access_type: str = 'read') -> None:
3051 """
3052 Raise 403 if the user lacks access to any of the requested collections.
3053 Delegates to the shared filter_accessible_collections utility so the
3054 access rules stay in one place.
3055 """
3056 requested = set(collection_names)
3057 allowed = await filter_accessible_collections(requested, user, access_type=access_type)
3058 denied = requested - allowed
3059 if denied:
3060 raise HTTPException(
3061 status_code=status.HTTP_403_FORBIDDEN,
3062 detail=ERROR_MESSAGES.ACCESS_PROHIBITED,
3063 )
3066class QueryDocForm(BaseModel):
3067 collection_name: str
3068 query: str
3069 k: int | None = None
3070 k_reranker: int | None = None
3071 r: float | None = None
3072 hybrid: bool | None = None
3073 hybrid_bm25_weight: float | None = None
3076@router.post('/query/doc')
3077async def query_doc_handler(
3078 request: Request,
3079 form_data: QueryDocForm,
3080 user=Depends(get_verified_user),
3081):
3082 config = await get_retrieval_config()
3083 await _validate_collection_access([form_data.collection_name], user)
3085 try:
3086 if config.ENABLE_RAG_HYBRID_SEARCH and (form_data.hybrid is None or form_data.hybrid): 3086 ↛ 3087line 3086 didn't jump to line 3087 because the condition on line 3086 was never true
3087 return await query_doc_with_hybrid_search(
3088 collection_name=form_data.collection_name,
3089 collection_result=None,
3090 query=form_data.query,
3091 embedding_function=lambda query, prefix: request.app.state.EMBEDDING_FUNCTION(
3092 query, prefix=prefix, user=user
3093 ),
3094 k=form_data.k if form_data.k else config.TOP_K,
3095 reranking_function=(
3096 (lambda query, documents: request.app.state.RERANKING_FUNCTION(query, documents, user=user))
3097 if request.app.state.RERANKING_FUNCTION
3098 else None
3099 ),
3100 k_reranker=form_data.k_reranker or config.TOP_K_RERANKER,
3101 r=(form_data.r if form_data.r else config.RELEVANCE_THRESHOLD),
3102 hybrid_bm25_weight=(
3103 form_data.hybrid_bm25_weight
3104 if form_data.hybrid_bm25_weight is not None
3105 else config.HYBRID_BM25_WEIGHT
3106 ),
3107 )
3108 else:
3109 query_embedding = await request.app.state.EMBEDDING_FUNCTION(
3110 form_data.query, prefix=RAG_EMBEDDING_QUERY_PREFIX, user=user
3111 )
3112 # query_doc wraps a blocking get_vector_db_client().search call;
3113 # offload so the request's event loop stays responsive.
3114 return await asyncio.to_thread(
3115 query_doc,
3116 collection_name=form_data.collection_name,
3117 query_embedding=query_embedding,
3118 k=form_data.k if form_data.k else config.TOP_K,
3119 user=user,
3120 )
3121 except HTTPException:
3122 raise
3123 except Exception as e:
3124 log.exception(e)
3125 raise HTTPException(
3126 status_code=status.HTTP_400_BAD_REQUEST,
3127 detail=ERROR_MESSAGES.DEFAULT(e, 'Error querying knowledge base'),
3128 )
3131class QueryCollectionsForm(BaseModel):
3132 collection_names: list[str]
3133 query: str
3134 k: int | None = None
3135 k_reranker: int | None = None
3136 r: float | None = None
3137 hybrid: bool | None = None
3138 hybrid_bm25_weight: float | None = None
3139 enable_enriched_texts: bool | None = None
3142@router.post('/query/collection')
3143async def query_collection_handler(
3144 request: Request,
3145 form_data: QueryCollectionsForm,
3146 user=Depends(get_verified_user),
3147):
3148 config = await get_retrieval_config()
3149 await _validate_collection_access(form_data.collection_names, user)
3151 try:
3152 if config.ENABLE_RAG_HYBRID_SEARCH and (form_data.hybrid is None or form_data.hybrid): 3152 ↛ 3153line 3152 didn't jump to line 3153 because the condition on line 3152 was never true
3153 return await query_collection_with_hybrid_search(
3154 collection_names=form_data.collection_names,
3155 queries=[form_data.query],
3156 embedding_function=lambda query, prefix: request.app.state.EMBEDDING_FUNCTION(
3157 query, prefix=prefix, user=user
3158 ),
3159 k=form_data.k if form_data.k else config.TOP_K,
3160 reranking_function=(
3161 (lambda query, documents: request.app.state.RERANKING_FUNCTION(query, documents, user=user))
3162 if request.app.state.RERANKING_FUNCTION
3163 else None
3164 ),
3165 k_reranker=form_data.k_reranker or config.TOP_K_RERANKER,
3166 r=(form_data.r if form_data.r else config.RELEVANCE_THRESHOLD),
3167 hybrid_bm25_weight=(
3168 form_data.hybrid_bm25_weight
3169 if form_data.hybrid_bm25_weight is not None
3170 else config.HYBRID_BM25_WEIGHT
3171 ),
3172 enable_enriched_texts=(
3173 form_data.enable_enriched_texts
3174 if form_data.enable_enriched_texts is not None
3175 else config.ENABLE_RAG_HYBRID_SEARCH_ENRICHED_TEXTS
3176 ),
3177 )
3178 else:
3179 return await query_collection(
3180 request,
3181 collection_names=form_data.collection_names,
3182 queries=[form_data.query],
3183 embedding_function=lambda query, prefix: request.app.state.EMBEDDING_FUNCTION(
3184 query, prefix=prefix, user=user
3185 ),
3186 k=form_data.k if form_data.k else config.TOP_K,
3187 )
3189 except HTTPException:
3190 raise
3191 except Exception as e:
3192 log.exception(e)
3193 raise HTTPException(
3194 status_code=status.HTTP_400_BAD_REQUEST,
3195 detail=ERROR_MESSAGES.DEFAULT(e, 'Error querying knowledge base'),
3196 )
3199####################################
3200#
3201# Vector DB operations
3202#
3203####################################
3206class DeleteForm(BaseModel):
3207 collection_name: str
3208 file_id: str
3211@router.post('/delete')
3212async def delete_entries_from_collection(
3213 request: Request,
3214 form_data: DeleteForm,
3215 user=Depends(get_admin_user),
3216 db: AsyncSession = Depends(get_async_session),
3217):
3218 try:
3219 if await ASYNC_VECTOR_DB_CLIENT.has_collection(collection_name=form_data.collection_name): 3219 ↛ 3220line 3219 didn't jump to line 3220 because the condition on line 3219 was never true
3220 file = await Files.get_file_by_id(form_data.file_id, db=db)
3221 if not file:
3222 raise HTTPException(
3223 status_code=status.HTTP_404_NOT_FOUND,
3224 detail=ERROR_MESSAGES.NOT_FOUND,
3225 )
3226 hash = file.hash
3228 # Refuse to issue a `filter={'hash': None}` query — the
3229 # match semantics of a null filter value are
3230 # backend-dependent (some backends ignore the key, some
3231 # match every row whose metadata lacks `hash`) and risk
3232 # deleting unrelated entries. Files without a hash are
3233 # typically unprocessed / failed / legacy records that
3234 # can't be targeted by hash anyway.
3235 if hash is None:
3236 raise HTTPException(
3237 status_code=status.HTTP_400_BAD_REQUEST,
3238 detail=ERROR_MESSAGES.DEFAULT('File has no hash; cannot delete vector entries by hash.'),
3239 )
3241 # Pre-existing bug: this used `metadata=` which is not a
3242 # parameter on `VectorDBBase.delete` nor on any backend
3243 # implementation, so the call always raised TypeError that
3244 # was silently swallowed by the surrounding `except
3245 # Exception` and the endpoint reported `{'status': False}`
3246 # for every request. Use `filter` to actually do what the
3247 # endpoint name promises.
3248 await ASYNC_VECTOR_DB_CLIENT.delete(
3249 collection_name=form_data.collection_name,
3250 filter={'hash': hash},
3251 )
3252 await publish_event(
3253 request,
3254 EVENTS.RETRIEVAL_COLLECTION_DELETED,
3255 actor=user,
3256 subject_id=form_data.collection_name,
3257 data={'file_id': form_data.file_id},
3258 )
3259 return {'status': True}
3260 else:
3261 return {'status': False}
3262 except HTTPException:
3263 # Caller-meaningful errors (404/400 above) must not be
3264 # swallowed and re-shaped as `{'status': False}`.
3265 raise
3266 except Exception as e:
3267 log.exception(e)
3268 return {'status': False}
3271@router.post('/reset/db')
3272async def reset_vector_db(
3273 request: Request,
3274 user=Depends(get_admin_user),
3275 db: AsyncSession = Depends(get_async_session),
3276):
3277 await ASYNC_VECTOR_DB_CLIENT.reset()
3278 await Knowledges.delete_all_knowledge(db=db)
3279 await publish_event(
3280 request,
3281 EVENTS.RETRIEVAL_VECTOR_DB_RESET,
3282 actor=user,
3283 subject_id='default',
3284 )
3287@router.post('/reset/uploads')
3288async def reset_upload_dir(request: Request, user=Depends(get_admin_user)) -> bool:
3289 folder = f'{UPLOAD_DIR}'
3290 try:
3291 # Check if the directory exists
3292 if await asyncio.to_thread(os.path.exists, folder): 3292 ↛ 3306line 3292 didn't jump to line 3306 because the condition on line 3292 was always true
3293 # Iterate over all the files and directories in the specified directory
3294 for filename in await asyncio.to_thread(os.listdir, folder):
3295 file_path = os.path.join(folder, filename)
3296 try:
3297 if await asyncio.to_thread(os.path.isfile, file_path) or await asyncio.to_thread( 3297 ↛ 3301line 3297 didn't jump to line 3301 because the condition on line 3297 was always true
3298 os.path.islink, file_path
3299 ):
3300 await asyncio.to_thread(os.unlink, file_path) # Remove the file or link
3301 elif await asyncio.to_thread(os.path.isdir, file_path):
3302 await asyncio.to_thread(shutil.rmtree, file_path) # Remove the directory
3303 except Exception as e:
3304 log.exception(f'Failed to delete {file_path}. Reason: {e}')
3305 else:
3306 log.warning(f'The directory {folder} does not exist')
3307 except Exception as e:
3308 log.exception(f'Failed to process the directory {folder}. Reason: {e}')
3310 await publish_event(
3311 request,
3312 EVENTS.RETRIEVAL_UPLOADS_RESET,
3313 actor=user,
3314 subject_id='all',
3315 subject_type='file.uploads',
3316 )
3317 return True
3320if ENV == 'dev': 3320 ↛ 3327line 3320 didn't jump to line 3327 because the condition on line 3320 was always true
3322 @router.get('/ef/{text}')
3323 async def get_embeddings(request: Request, text: str | None = 'Hello World!'):
3324 return {'result': await request.app.state.EMBEDDING_FUNCTION(text, prefix=RAG_EMBEDDING_QUERY_PREFIX)}
3327class BatchProcessFilesForm(BaseModel):
3328 files: list[FileModel]
3329 collection_name: str
3332class BatchProcessFilesResult(BaseModel):
3333 file_id: str
3334 status: str
3335 error: str | None = None
3338class BatchProcessFilesResponse(BaseModel):
3339 results: list[BatchProcessFilesResult]
3340 errors: list[BatchProcessFilesResult]
3343@router.post('/process/files/batch')
3344async def process_files_batch(
3345 request: Request,
3346 form_data: BatchProcessFilesForm,
3347 user=Depends(get_verified_user),
3348 db=None,
3349) -> BatchProcessFilesResponse:
3350 """
3351 Process a batch of files and save them to the vector database.
3353 NOTE: We intentionally do NOT use Depends(get_async_session) here.
3354 The save_docs_to_vector_db() call makes external embedding API calls which
3355 can take 5-60+ seconds for batch operations. Database operations after
3356 embedding (Files.update_file_by_id) manage their own short-lived sessions.
3357 """
3359 config = await get_retrieval_config()
3360 collection_name = form_data.collection_name
3362 if collection_name:
3363 await _validate_collection_access([collection_name], user, access_type='write')
3365 file_results: list[BatchProcessFilesResult] = []
3366 file_errors: list[BatchProcessFilesResult] = []
3367 file_updates: list[FileUpdateForm] = []
3369 # Prepare all documents first
3370 all_docs: list[Document] = []
3372 for file in form_data.files:
3373 try:
3374 # Ownership check: verify the requesting user owns the file or is an admin
3375 db_file = await Files.get_file_by_id(file.id, db=db)
3376 if not db_file: 3376 ↛ 3385line 3376 didn't jump to line 3385 because the condition on line 3376 was always true
3377 file_errors.append(
3378 BatchProcessFilesResult(
3379 file_id=file.id,
3380 status='failed',
3381 error='File not found',
3382 )
3383 )
3384 continue
3385 if db_file.user_id != user.id and user.role != 'admin':
3386 file_errors.append(
3387 BatchProcessFilesResult(
3388 file_id=file.id,
3389 status='failed',
3390 error='Permission denied: not file owner',
3391 )
3392 )
3393 continue
3395 text_content = file.data.get('content', '')
3396 docs: list[Document] = [
3397 Document(
3398 page_content=text_content.replace('<br/>', '\n'),
3399 metadata={
3400 **filter_file_metadata(file.meta),
3401 'name': file.filename,
3402 'created_by': file.user_id,
3403 'file_id': file.id,
3404 'source': file.filename,
3405 },
3406 )
3407 ]
3409 all_docs.extend(docs)
3411 file_updates.append(
3412 FileUpdateForm(
3413 hash=calculate_sha256_string(text_content),
3414 data={'content': text_content},
3415 )
3416 )
3417 file_results.append(BatchProcessFilesResult(file_id=file.id, status='prepared'))
3419 except Exception as e:
3420 log.error(f'process_files_batch: Error processing file {file.id}: {str(e)}')
3421 file_errors.append(BatchProcessFilesResult(file_id=file.id, status='failed', error=str(e)))
3423 # Save all documents in one batch
3424 if all_docs: 3424 ↛ 3425line 3424 didn't jump to line 3425 because the condition on line 3424 was never true
3425 try:
3426 await run_in_threadpool(
3427 save_docs_to_vector_db,
3428 request,
3429 all_docs,
3430 collection_name,
3431 config,
3432 add=True,
3433 user=user,
3434 )
3436 # Update all files with collection name
3437 for file_update, file_result in zip(file_updates, file_results):
3438 await Files.update_file_by_id(id=file_result.file_id, form_data=file_update, db=db)
3439 file_result.status = 'completed'
3441 except Exception as e:
3442 log.error(f'process_files_batch: Error saving documents to vector DB: {str(e)}')
3443 for file_result in file_results:
3444 file_result.status = 'failed'
3445 file_errors.append(BatchProcessFilesResult(file_id=file_result.file_id, status='failed', error=str(e)))
3447 response = BatchProcessFilesResponse(results=file_results, errors=file_errors)
3448 await publish_event(
3449 request,
3450 EVENTS.RETRIEVAL_CONTENT_PROCESSED,
3451 actor=user,
3452 subject_id=collection_name,
3453 subject_type='retrieval.collection',
3454 data={
3455 'count': len([item for item in file_results if item.status == 'completed']),
3456 'errors': len(file_errors),
3457 },
3458 )
3459 return response