retrieval.py 94 KB

1234567891011121314151617181920212223242526272829303132333435363738394041424344454647484950515253545556575859606162636465666768697071727374757677787980818283848586878889909192939495969798991001011021031041051061071081091101111121131141151161171181191201211221231241251261271281291301311321331341351361371381391401411421431441451461471481491501511521531541551561571581591601611621631641651661671681691701711721731741751761771781791801811821831841851861871881891901911921931941951961971981992002012022032042052062072082092102112122132142152162172182192202212222232242252262272282292302312322332342352362372382392402412422432442452462472482492502512522532542552562572582592602612622632642652662672682692702712722732742752762772782792802812822832842852862872882892902912922932942952962972982993003013023033043053063073083093103113123133143153163173183193203213223233243253263273283293303313323333343353363373383393403413423433443453463473483493503513523533543553563573583593603613623633643653663673683693703713723733743753763773783793803813823833843853863873883893903913923933943953963973983994004014024034044054064074084094104114124134144154164174184194204214224234244254264274284294304314324334344354364374384394404414424434444454464474484494504514524534544554564574584594604614624634644654664674684694704714724734744754764774784794804814824834844854864874884894904914924934944954964974984995005015025035045055065075085095105115125135145155165175185195205215225235245255265275285295305315325335345355365375385395405415425435445455465475485495505515525535545555565575585595605615625635645655665675685695705715725735745755765775785795805815825835845855865875885895905915925935945955965975985996006016026036046056066076086096106116126136146156166176186196206216226236246256266276286296306316326336346356366376386396406416426436446456466476486496506516526536546556566576586596606616626636646656666676686696706716726736746756766776786796806816826836846856866876886896906916926936946956966976986997007017027037047057067077087097107117127137147157167177187197207217227237247257267277287297307317327337347357367377387397407417427437447457467477487497507517527537547557567577587597607617627637647657667677687697707717727737747757767777787797807817827837847857867877887897907917927937947957967977987998008018028038048058068078088098108118128138148158168178188198208218228238248258268278288298308318328338348358368378388398408418428438448458468478488498508518528538548558568578588598608618628638648658668678688698708718728738748758768778788798808818828838848858868878888898908918928938948958968978988999009019029039049059069079089099109119129139149159169179189199209219229239249259269279289299309319329339349359369379389399409419429439449459469479489499509519529539549559569579589599609619629639649659669679689699709719729739749759769779789799809819829839849859869879889899909919929939949959969979989991000100110021003100410051006100710081009101010111012101310141015101610171018101910201021102210231024102510261027102810291030103110321033103410351036103710381039104010411042104310441045104610471048104910501051105210531054105510561057105810591060106110621063106410651066106710681069107010711072107310741075107610771078107910801081108210831084108510861087108810891090109110921093109410951096109710981099110011011102110311041105110611071108110911101111111211131114111511161117111811191120112111221123112411251126112711281129113011311132113311341135113611371138113911401141114211431144114511461147114811491150115111521153115411551156115711581159116011611162116311641165116611671168116911701171117211731174117511761177117811791180118111821183118411851186118711881189119011911192119311941195119611971198119912001201120212031204120512061207120812091210121112121213121412151216121712181219122012211222122312241225122612271228122912301231123212331234123512361237123812391240124112421243124412451246124712481249125012511252125312541255125612571258125912601261126212631264126512661267126812691270127112721273127412751276127712781279128012811282128312841285128612871288128912901291129212931294129512961297129812991300130113021303130413051306130713081309131013111312131313141315131613171318131913201321132213231324132513261327132813291330133113321333133413351336133713381339134013411342134313441345134613471348134913501351135213531354135513561357135813591360136113621363136413651366136713681369137013711372137313741375137613771378137913801381138213831384138513861387138813891390139113921393139413951396139713981399140014011402140314041405140614071408140914101411141214131414141514161417141814191420142114221423142414251426142714281429143014311432143314341435143614371438143914401441144214431444144514461447144814491450145114521453145414551456145714581459146014611462146314641465146614671468146914701471147214731474147514761477147814791480148114821483148414851486148714881489149014911492149314941495149614971498149915001501150215031504150515061507150815091510151115121513151415151516151715181519152015211522152315241525152615271528152915301531153215331534153515361537153815391540154115421543154415451546154715481549155015511552155315541555155615571558155915601561156215631564156515661567156815691570157115721573157415751576157715781579158015811582158315841585158615871588158915901591159215931594159515961597159815991600160116021603160416051606160716081609161016111612161316141615161616171618161916201621162216231624162516261627162816291630163116321633163416351636163716381639164016411642164316441645164616471648164916501651165216531654165516561657165816591660166116621663166416651666166716681669167016711672167316741675167616771678167916801681168216831684168516861687168816891690169116921693169416951696169716981699170017011702170317041705170617071708170917101711171217131714171517161717171817191720172117221723172417251726172717281729173017311732173317341735173617371738173917401741174217431744174517461747174817491750175117521753175417551756175717581759176017611762176317641765176617671768176917701771177217731774177517761777177817791780178117821783178417851786178717881789179017911792179317941795179617971798179918001801180218031804180518061807180818091810181118121813181418151816181718181819182018211822182318241825182618271828182918301831183218331834183518361837183818391840184118421843184418451846184718481849185018511852185318541855185618571858185918601861186218631864186518661867186818691870187118721873187418751876187718781879188018811882188318841885188618871888188918901891189218931894189518961897189818991900190119021903190419051906190719081909191019111912191319141915191619171918191919201921192219231924192519261927192819291930193119321933193419351936193719381939194019411942194319441945194619471948194919501951195219531954195519561957195819591960196119621963196419651966196719681969197019711972197319741975197619771978197919801981198219831984198519861987198819891990199119921993199419951996199719981999200020012002200320042005200620072008200920102011201220132014201520162017201820192020202120222023202420252026202720282029203020312032203320342035203620372038203920402041204220432044204520462047204820492050205120522053205420552056205720582059206020612062206320642065206620672068206920702071207220732074207520762077207820792080208120822083208420852086208720882089209020912092209320942095209620972098209921002101210221032104210521062107210821092110211121122113211421152116211721182119212021212122212321242125212621272128212921302131213221332134213521362137213821392140214121422143214421452146214721482149215021512152215321542155215621572158215921602161216221632164216521662167216821692170217121722173217421752176217721782179218021812182218321842185218621872188218921902191219221932194219521962197219821992200220122022203220422052206220722082209221022112212221322142215221622172218221922202221222222232224222522262227222822292230223122322233223422352236223722382239224022412242224322442245224622472248224922502251225222532254225522562257225822592260226122622263226422652266226722682269227022712272227322742275227622772278227922802281228222832284228522862287228822892290229122922293229422952296229722982299230023012302230323042305230623072308
  1. import json
  2. import logging
  3. import mimetypes
  4. import os
  5. import shutil
  6. import asyncio
  7. import uuid
  8. from datetime import datetime
  9. from pathlib import Path
  10. from typing import Iterator, List, Optional, Sequence, Union
  11. from fastapi import (
  12. Depends,
  13. FastAPI,
  14. File,
  15. Form,
  16. HTTPException,
  17. UploadFile,
  18. Request,
  19. status,
  20. APIRouter,
  21. )
  22. from fastapi.middleware.cors import CORSMiddleware
  23. from fastapi.concurrency import run_in_threadpool
  24. from pydantic import BaseModel
  25. import tiktoken
  26. from langchain.text_splitter import RecursiveCharacterTextSplitter, TokenTextSplitter
  27. from langchain_text_splitters import MarkdownHeaderTextSplitter
  28. from langchain_core.documents import Document
  29. from open_webui.models.files import FileModel, Files
  30. from open_webui.models.knowledge import Knowledges
  31. from open_webui.storage.provider import Storage
  32. from open_webui.retrieval.vector.factory import VECTOR_DB_CLIENT
  33. # Document loaders
  34. from open_webui.retrieval.loaders.main import Loader
  35. from open_webui.retrieval.loaders.youtube import YoutubeLoader
  36. # Web search engines
  37. from open_webui.retrieval.web.main import SearchResult
  38. from open_webui.retrieval.web.utils import get_web_loader
  39. from open_webui.retrieval.web.brave import search_brave
  40. from open_webui.retrieval.web.kagi import search_kagi
  41. from open_webui.retrieval.web.mojeek import search_mojeek
  42. from open_webui.retrieval.web.bocha import search_bocha
  43. from open_webui.retrieval.web.duckduckgo import search_duckduckgo
  44. from open_webui.retrieval.web.google_pse import search_google_pse
  45. from open_webui.retrieval.web.jina_search import search_jina
  46. from open_webui.retrieval.web.searchapi import search_searchapi
  47. from open_webui.retrieval.web.serpapi import search_serpapi
  48. from open_webui.retrieval.web.searxng import search_searxng
  49. from open_webui.retrieval.web.yacy import search_yacy
  50. from open_webui.retrieval.web.serper import search_serper
  51. from open_webui.retrieval.web.serply import search_serply
  52. from open_webui.retrieval.web.serpstack import search_serpstack
  53. from open_webui.retrieval.web.tavily import search_tavily
  54. from open_webui.retrieval.web.bing import search_bing
  55. from open_webui.retrieval.web.exa import search_exa
  56. from open_webui.retrieval.web.perplexity import search_perplexity
  57. from open_webui.retrieval.web.sougou import search_sougou
  58. from open_webui.retrieval.web.firecrawl import search_firecrawl
  59. from open_webui.retrieval.web.external import search_external
  60. from open_webui.retrieval.utils import (
  61. get_embedding_function,
  62. get_reranking_function,
  63. get_model_path,
  64. query_collection,
  65. query_collection_with_hybrid_search,
  66. query_doc,
  67. query_doc_with_hybrid_search,
  68. )
  69. from open_webui.utils.misc import (
  70. calculate_sha256_string,
  71. )
  72. from open_webui.utils.auth import get_admin_user, get_verified_user
  73. from open_webui.config import (
  74. ENV,
  75. RAG_EMBEDDING_MODEL_AUTO_UPDATE,
  76. RAG_EMBEDDING_MODEL_TRUST_REMOTE_CODE,
  77. RAG_RERANKING_MODEL_AUTO_UPDATE,
  78. RAG_RERANKING_MODEL_TRUST_REMOTE_CODE,
  79. UPLOAD_DIR,
  80. DEFAULT_LOCALE,
  81. RAG_EMBEDDING_CONTENT_PREFIX,
  82. RAG_EMBEDDING_QUERY_PREFIX,
  83. )
  84. from open_webui.env import (
  85. SRC_LOG_LEVELS,
  86. DEVICE_TYPE,
  87. DOCKER,
  88. SENTENCE_TRANSFORMERS_BACKEND,
  89. SENTENCE_TRANSFORMERS_MODEL_KWARGS,
  90. SENTENCE_TRANSFORMERS_CROSS_ENCODER_BACKEND,
  91. SENTENCE_TRANSFORMERS_CROSS_ENCODER_MODEL_KWARGS,
  92. )
  93. from open_webui.constants import ERROR_MESSAGES
  94. log = logging.getLogger(__name__)
  95. log.setLevel(SRC_LOG_LEVELS["RAG"])
  96. ##########################################
  97. #
  98. # Utility functions
  99. #
  100. ##########################################
  101. def get_ef(
  102. engine: str,
  103. embedding_model: str,
  104. auto_update: bool = False,
  105. ):
  106. ef = None
  107. if embedding_model and engine == "":
  108. from sentence_transformers import SentenceTransformer
  109. try:
  110. ef = SentenceTransformer(
  111. get_model_path(embedding_model, auto_update),
  112. device=DEVICE_TYPE,
  113. trust_remote_code=RAG_EMBEDDING_MODEL_TRUST_REMOTE_CODE,
  114. backend=SENTENCE_TRANSFORMERS_BACKEND,
  115. model_kwargs=SENTENCE_TRANSFORMERS_MODEL_KWARGS,
  116. )
  117. except Exception as e:
  118. log.debug(f"Error loading SentenceTransformer: {e}")
  119. return ef
  120. def get_rf(
  121. engine: str = "",
  122. reranking_model: Optional[str] = None,
  123. external_reranker_url: str = "",
  124. external_reranker_api_key: str = "",
  125. auto_update: bool = False,
  126. ):
  127. rf = None
  128. if reranking_model:
  129. if any(model in reranking_model for model in ["jinaai/jina-colbert-v2"]):
  130. try:
  131. from open_webui.retrieval.models.colbert import ColBERT
  132. rf = ColBERT(
  133. get_model_path(reranking_model, auto_update),
  134. env="docker" if DOCKER else None,
  135. )
  136. except Exception as e:
  137. log.error(f"ColBERT: {e}")
  138. raise Exception(ERROR_MESSAGES.DEFAULT(e))
  139. else:
  140. if engine == "external":
  141. try:
  142. from open_webui.retrieval.models.external import ExternalReranker
  143. rf = ExternalReranker(
  144. url=external_reranker_url,
  145. api_key=external_reranker_api_key,
  146. model=reranking_model,
  147. )
  148. except Exception as e:
  149. log.error(f"ExternalReranking: {e}")
  150. raise Exception(ERROR_MESSAGES.DEFAULT(e))
  151. else:
  152. import sentence_transformers
  153. try:
  154. rf = sentence_transformers.CrossEncoder(
  155. get_model_path(reranking_model, auto_update),
  156. device=DEVICE_TYPE,
  157. trust_remote_code=RAG_RERANKING_MODEL_TRUST_REMOTE_CODE,
  158. backend=SENTENCE_TRANSFORMERS_CROSS_ENCODER_BACKEND,
  159. model_kwargs=SENTENCE_TRANSFORMERS_CROSS_ENCODER_MODEL_KWARGS,
  160. )
  161. except Exception as e:
  162. log.error(f"CrossEncoder: {e}")
  163. raise Exception(ERROR_MESSAGES.DEFAULT("CrossEncoder error"))
  164. return rf
  165. ##########################################
  166. #
  167. # API routes
  168. #
  169. ##########################################
  170. router = APIRouter()
  171. class CollectionNameForm(BaseModel):
  172. collection_name: Optional[str] = None
  173. class ProcessUrlForm(CollectionNameForm):
  174. url: str
  175. class SearchForm(BaseModel):
  176. queries: List[str]
  177. @router.get("/")
  178. async def get_status(request: Request):
  179. return {
  180. "status": True,
  181. "chunk_size": request.app.state.config.CHUNK_SIZE,
  182. "chunk_overlap": request.app.state.config.CHUNK_OVERLAP,
  183. "template": request.app.state.config.RAG_TEMPLATE,
  184. "embedding_engine": request.app.state.config.RAG_EMBEDDING_ENGINE,
  185. "embedding_model": request.app.state.config.RAG_EMBEDDING_MODEL,
  186. "reranking_model": request.app.state.config.RAG_RERANKING_MODEL,
  187. "embedding_batch_size": request.app.state.config.RAG_EMBEDDING_BATCH_SIZE,
  188. }
  189. @router.get("/embedding")
  190. async def get_embedding_config(request: Request, user=Depends(get_admin_user)):
  191. return {
  192. "status": True,
  193. "embedding_engine": request.app.state.config.RAG_EMBEDDING_ENGINE,
  194. "embedding_model": request.app.state.config.RAG_EMBEDDING_MODEL,
  195. "embedding_batch_size": request.app.state.config.RAG_EMBEDDING_BATCH_SIZE,
  196. "openai_config": {
  197. "url": request.app.state.config.RAG_OPENAI_API_BASE_URL,
  198. "key": request.app.state.config.RAG_OPENAI_API_KEY,
  199. },
  200. "ollama_config": {
  201. "url": request.app.state.config.RAG_OLLAMA_BASE_URL,
  202. "key": request.app.state.config.RAG_OLLAMA_API_KEY,
  203. },
  204. "azure_openai_config": {
  205. "url": request.app.state.config.RAG_AZURE_OPENAI_BASE_URL,
  206. "key": request.app.state.config.RAG_AZURE_OPENAI_API_KEY,
  207. "version": request.app.state.config.RAG_AZURE_OPENAI_API_VERSION,
  208. },
  209. }
  210. class OpenAIConfigForm(BaseModel):
  211. url: str
  212. key: str
  213. class OllamaConfigForm(BaseModel):
  214. url: str
  215. key: str
  216. class AzureOpenAIConfigForm(BaseModel):
  217. url: str
  218. key: str
  219. version: str
  220. class EmbeddingModelUpdateForm(BaseModel):
  221. openai_config: Optional[OpenAIConfigForm] = None
  222. ollama_config: Optional[OllamaConfigForm] = None
  223. azure_openai_config: Optional[AzureOpenAIConfigForm] = None
  224. embedding_engine: str
  225. embedding_model: str
  226. embedding_batch_size: Optional[int] = 1
  227. @router.post("/embedding/update")
  228. async def update_embedding_config(
  229. request: Request, form_data: EmbeddingModelUpdateForm, user=Depends(get_admin_user)
  230. ):
  231. log.info(
  232. f"Updating embedding model: {request.app.state.config.RAG_EMBEDDING_MODEL} to {form_data.embedding_model}"
  233. )
  234. try:
  235. request.app.state.config.RAG_EMBEDDING_ENGINE = form_data.embedding_engine
  236. request.app.state.config.RAG_EMBEDDING_MODEL = form_data.embedding_model
  237. if request.app.state.config.RAG_EMBEDDING_ENGINE in [
  238. "ollama",
  239. "openai",
  240. "azure_openai",
  241. ]:
  242. if form_data.openai_config is not None:
  243. request.app.state.config.RAG_OPENAI_API_BASE_URL = (
  244. form_data.openai_config.url
  245. )
  246. request.app.state.config.RAG_OPENAI_API_KEY = (
  247. form_data.openai_config.key
  248. )
  249. if form_data.ollama_config is not None:
  250. request.app.state.config.RAG_OLLAMA_BASE_URL = (
  251. form_data.ollama_config.url
  252. )
  253. request.app.state.config.RAG_OLLAMA_API_KEY = (
  254. form_data.ollama_config.key
  255. )
  256. if form_data.azure_openai_config is not None:
  257. request.app.state.config.RAG_AZURE_OPENAI_BASE_URL = (
  258. form_data.azure_openai_config.url
  259. )
  260. request.app.state.config.RAG_AZURE_OPENAI_API_KEY = (
  261. form_data.azure_openai_config.key
  262. )
  263. request.app.state.config.RAG_AZURE_OPENAI_API_VERSION = (
  264. form_data.azure_openai_config.version
  265. )
  266. request.app.state.config.RAG_EMBEDDING_BATCH_SIZE = (
  267. form_data.embedding_batch_size
  268. )
  269. request.app.state.ef = get_ef(
  270. request.app.state.config.RAG_EMBEDDING_ENGINE,
  271. request.app.state.config.RAG_EMBEDDING_MODEL,
  272. )
  273. request.app.state.EMBEDDING_FUNCTION = get_embedding_function(
  274. request.app.state.config.RAG_EMBEDDING_ENGINE,
  275. request.app.state.config.RAG_EMBEDDING_MODEL,
  276. request.app.state.ef,
  277. (
  278. request.app.state.config.RAG_OPENAI_API_BASE_URL
  279. if request.app.state.config.RAG_EMBEDDING_ENGINE == "openai"
  280. else (
  281. request.app.state.config.RAG_OLLAMA_BASE_URL
  282. if request.app.state.config.RAG_EMBEDDING_ENGINE == "ollama"
  283. else request.app.state.config.RAG_AZURE_OPENAI_BASE_URL
  284. )
  285. ),
  286. (
  287. request.app.state.config.RAG_OPENAI_API_KEY
  288. if request.app.state.config.RAG_EMBEDDING_ENGINE == "openai"
  289. else (
  290. request.app.state.config.RAG_OLLAMA_API_KEY
  291. if request.app.state.config.RAG_EMBEDDING_ENGINE == "ollama"
  292. else request.app.state.config.RAG_AZURE_OPENAI_API_KEY
  293. )
  294. ),
  295. request.app.state.config.RAG_EMBEDDING_BATCH_SIZE,
  296. azure_api_version=(
  297. request.app.state.config.RAG_AZURE_OPENAI_API_VERSION
  298. if request.app.state.config.RAG_EMBEDDING_ENGINE == "azure_openai"
  299. else None
  300. ),
  301. )
  302. return {
  303. "status": True,
  304. "embedding_engine": request.app.state.config.RAG_EMBEDDING_ENGINE,
  305. "embedding_model": request.app.state.config.RAG_EMBEDDING_MODEL,
  306. "embedding_batch_size": request.app.state.config.RAG_EMBEDDING_BATCH_SIZE,
  307. "openai_config": {
  308. "url": request.app.state.config.RAG_OPENAI_API_BASE_URL,
  309. "key": request.app.state.config.RAG_OPENAI_API_KEY,
  310. },
  311. "ollama_config": {
  312. "url": request.app.state.config.RAG_OLLAMA_BASE_URL,
  313. "key": request.app.state.config.RAG_OLLAMA_API_KEY,
  314. },
  315. "azure_openai_config": {
  316. "url": request.app.state.config.RAG_AZURE_OPENAI_BASE_URL,
  317. "key": request.app.state.config.RAG_AZURE_OPENAI_API_KEY,
  318. "version": request.app.state.config.RAG_AZURE_OPENAI_API_VERSION,
  319. },
  320. }
  321. except Exception as e:
  322. log.exception(f"Problem updating embedding model: {e}")
  323. raise HTTPException(
  324. status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
  325. detail=ERROR_MESSAGES.DEFAULT(e),
  326. )
  327. @router.get("/config")
  328. async def get_rag_config(request: Request, user=Depends(get_admin_user)):
  329. return {
  330. "status": True,
  331. # RAG settings
  332. "RAG_TEMPLATE": request.app.state.config.RAG_TEMPLATE,
  333. "TOP_K": request.app.state.config.TOP_K,
  334. "BYPASS_EMBEDDING_AND_RETRIEVAL": request.app.state.config.BYPASS_EMBEDDING_AND_RETRIEVAL,
  335. "RAG_FULL_CONTEXT": request.app.state.config.RAG_FULL_CONTEXT,
  336. # Hybrid search settings
  337. "ENABLE_RAG_HYBRID_SEARCH": request.app.state.config.ENABLE_RAG_HYBRID_SEARCH,
  338. "TOP_K_RERANKER": request.app.state.config.TOP_K_RERANKER,
  339. "RELEVANCE_THRESHOLD": request.app.state.config.RELEVANCE_THRESHOLD,
  340. "HYBRID_BM25_WEIGHT": request.app.state.config.HYBRID_BM25_WEIGHT,
  341. # Content extraction settings
  342. "CONTENT_EXTRACTION_ENGINE": request.app.state.config.CONTENT_EXTRACTION_ENGINE,
  343. "PDF_EXTRACT_IMAGES": request.app.state.config.PDF_EXTRACT_IMAGES,
  344. "DATALAB_MARKER_API_KEY": request.app.state.config.DATALAB_MARKER_API_KEY,
  345. "DATALAB_MARKER_LANGS": request.app.state.config.DATALAB_MARKER_LANGS,
  346. "DATALAB_MARKER_SKIP_CACHE": request.app.state.config.DATALAB_MARKER_SKIP_CACHE,
  347. "DATALAB_MARKER_FORCE_OCR": request.app.state.config.DATALAB_MARKER_FORCE_OCR,
  348. "DATALAB_MARKER_PAGINATE": request.app.state.config.DATALAB_MARKER_PAGINATE,
  349. "DATALAB_MARKER_STRIP_EXISTING_OCR": request.app.state.config.DATALAB_MARKER_STRIP_EXISTING_OCR,
  350. "DATALAB_MARKER_DISABLE_IMAGE_EXTRACTION": request.app.state.config.DATALAB_MARKER_DISABLE_IMAGE_EXTRACTION,
  351. "DATALAB_MARKER_USE_LLM": request.app.state.config.DATALAB_MARKER_USE_LLM,
  352. "DATALAB_MARKER_OUTPUT_FORMAT": request.app.state.config.DATALAB_MARKER_OUTPUT_FORMAT,
  353. "EXTERNAL_DOCUMENT_LOADER_URL": request.app.state.config.EXTERNAL_DOCUMENT_LOADER_URL,
  354. "EXTERNAL_DOCUMENT_LOADER_API_KEY": request.app.state.config.EXTERNAL_DOCUMENT_LOADER_API_KEY,
  355. "TIKA_SERVER_URL": request.app.state.config.TIKA_SERVER_URL,
  356. "DOCLING_SERVER_URL": request.app.state.config.DOCLING_SERVER_URL,
  357. "DOCLING_OCR_ENGINE": request.app.state.config.DOCLING_OCR_ENGINE,
  358. "DOCLING_OCR_LANG": request.app.state.config.DOCLING_OCR_LANG,
  359. "DOCLING_DO_PICTURE_DESCRIPTION": request.app.state.config.DOCLING_DO_PICTURE_DESCRIPTION,
  360. "DOCLING_PICTURE_DESCRIPTION_MODE": request.app.state.config.DOCLING_PICTURE_DESCRIPTION_MODE,
  361. "DOCLING_PICTURE_DESCRIPTION_LOCAL": request.app.state.config.DOCLING_PICTURE_DESCRIPTION_LOCAL,
  362. "DOCLING_PICTURE_DESCRIPTION_API": request.app.state.config.DOCLING_PICTURE_DESCRIPTION_API,
  363. "DOCUMENT_INTELLIGENCE_ENDPOINT": request.app.state.config.DOCUMENT_INTELLIGENCE_ENDPOINT,
  364. "DOCUMENT_INTELLIGENCE_KEY": request.app.state.config.DOCUMENT_INTELLIGENCE_KEY,
  365. "MISTRAL_OCR_API_KEY": request.app.state.config.MISTRAL_OCR_API_KEY,
  366. # Reranking settings
  367. "RAG_RERANKING_MODEL": request.app.state.config.RAG_RERANKING_MODEL,
  368. "RAG_RERANKING_ENGINE": request.app.state.config.RAG_RERANKING_ENGINE,
  369. "RAG_EXTERNAL_RERANKER_URL": request.app.state.config.RAG_EXTERNAL_RERANKER_URL,
  370. "RAG_EXTERNAL_RERANKER_API_KEY": request.app.state.config.RAG_EXTERNAL_RERANKER_API_KEY,
  371. # Chunking settings
  372. "TEXT_SPLITTER": request.app.state.config.TEXT_SPLITTER,
  373. "CHUNK_SIZE": request.app.state.config.CHUNK_SIZE,
  374. "CHUNK_OVERLAP": request.app.state.config.CHUNK_OVERLAP,
  375. # File upload settings
  376. "FILE_MAX_SIZE": request.app.state.config.FILE_MAX_SIZE,
  377. "FILE_MAX_COUNT": request.app.state.config.FILE_MAX_COUNT,
  378. "FILE_IMAGE_COMPRESSION_WIDTH": request.app.state.config.FILE_IMAGE_COMPRESSION_WIDTH,
  379. "FILE_IMAGE_COMPRESSION_HEIGHT": request.app.state.config.FILE_IMAGE_COMPRESSION_HEIGHT,
  380. "ALLOWED_FILE_EXTENSIONS": request.app.state.config.ALLOWED_FILE_EXTENSIONS,
  381. # Integration settings
  382. "ENABLE_GOOGLE_DRIVE_INTEGRATION": request.app.state.config.ENABLE_GOOGLE_DRIVE_INTEGRATION,
  383. "ENABLE_ONEDRIVE_INTEGRATION": request.app.state.config.ENABLE_ONEDRIVE_INTEGRATION,
  384. # Web search settings
  385. "web": {
  386. "ENABLE_WEB_SEARCH": request.app.state.config.ENABLE_WEB_SEARCH,
  387. "WEB_SEARCH_ENGINE": request.app.state.config.WEB_SEARCH_ENGINE,
  388. "WEB_SEARCH_TRUST_ENV": request.app.state.config.WEB_SEARCH_TRUST_ENV,
  389. "WEB_SEARCH_RESULT_COUNT": request.app.state.config.WEB_SEARCH_RESULT_COUNT,
  390. "WEB_SEARCH_CONCURRENT_REQUESTS": request.app.state.config.WEB_SEARCH_CONCURRENT_REQUESTS,
  391. "WEB_SEARCH_DOMAIN_FILTER_LIST": request.app.state.config.WEB_SEARCH_DOMAIN_FILTER_LIST,
  392. "BYPASS_WEB_SEARCH_EMBEDDING_AND_RETRIEVAL": request.app.state.config.BYPASS_WEB_SEARCH_EMBEDDING_AND_RETRIEVAL,
  393. "BYPASS_WEB_SEARCH_WEB_LOADER": request.app.state.config.BYPASS_WEB_SEARCH_WEB_LOADER,
  394. "SEARXNG_QUERY_URL": request.app.state.config.SEARXNG_QUERY_URL,
  395. "YACY_QUERY_URL": request.app.state.config.YACY_QUERY_URL,
  396. "YACY_USERNAME": request.app.state.config.YACY_USERNAME,
  397. "YACY_PASSWORD": request.app.state.config.YACY_PASSWORD,
  398. "GOOGLE_PSE_API_KEY": request.app.state.config.GOOGLE_PSE_API_KEY,
  399. "GOOGLE_PSE_ENGINE_ID": request.app.state.config.GOOGLE_PSE_ENGINE_ID,
  400. "BRAVE_SEARCH_API_KEY": request.app.state.config.BRAVE_SEARCH_API_KEY,
  401. "KAGI_SEARCH_API_KEY": request.app.state.config.KAGI_SEARCH_API_KEY,
  402. "MOJEEK_SEARCH_API_KEY": request.app.state.config.MOJEEK_SEARCH_API_KEY,
  403. "BOCHA_SEARCH_API_KEY": request.app.state.config.BOCHA_SEARCH_API_KEY,
  404. "SERPSTACK_API_KEY": request.app.state.config.SERPSTACK_API_KEY,
  405. "SERPSTACK_HTTPS": request.app.state.config.SERPSTACK_HTTPS,
  406. "SERPER_API_KEY": request.app.state.config.SERPER_API_KEY,
  407. "SERPLY_API_KEY": request.app.state.config.SERPLY_API_KEY,
  408. "TAVILY_API_KEY": request.app.state.config.TAVILY_API_KEY,
  409. "SEARCHAPI_API_KEY": request.app.state.config.SEARCHAPI_API_KEY,
  410. "SEARCHAPI_ENGINE": request.app.state.config.SEARCHAPI_ENGINE,
  411. "SERPAPI_API_KEY": request.app.state.config.SERPAPI_API_KEY,
  412. "SERPAPI_ENGINE": request.app.state.config.SERPAPI_ENGINE,
  413. "JINA_API_KEY": request.app.state.config.JINA_API_KEY,
  414. "BING_SEARCH_V7_ENDPOINT": request.app.state.config.BING_SEARCH_V7_ENDPOINT,
  415. "BING_SEARCH_V7_SUBSCRIPTION_KEY": request.app.state.config.BING_SEARCH_V7_SUBSCRIPTION_KEY,
  416. "EXA_API_KEY": request.app.state.config.EXA_API_KEY,
  417. "PERPLEXITY_API_KEY": request.app.state.config.PERPLEXITY_API_KEY,
  418. "PERPLEXITY_MODEL": request.app.state.config.PERPLEXITY_MODEL,
  419. "PERPLEXITY_SEARCH_CONTEXT_USAGE": request.app.state.config.PERPLEXITY_SEARCH_CONTEXT_USAGE,
  420. "SOUGOU_API_SID": request.app.state.config.SOUGOU_API_SID,
  421. "SOUGOU_API_SK": request.app.state.config.SOUGOU_API_SK,
  422. "WEB_LOADER_ENGINE": request.app.state.config.WEB_LOADER_ENGINE,
  423. "ENABLE_WEB_LOADER_SSL_VERIFICATION": request.app.state.config.ENABLE_WEB_LOADER_SSL_VERIFICATION,
  424. "PLAYWRIGHT_WS_URL": request.app.state.config.PLAYWRIGHT_WS_URL,
  425. "PLAYWRIGHT_TIMEOUT": request.app.state.config.PLAYWRIGHT_TIMEOUT,
  426. "FIRECRAWL_API_KEY": request.app.state.config.FIRECRAWL_API_KEY,
  427. "FIRECRAWL_API_BASE_URL": request.app.state.config.FIRECRAWL_API_BASE_URL,
  428. "TAVILY_EXTRACT_DEPTH": request.app.state.config.TAVILY_EXTRACT_DEPTH,
  429. "EXTERNAL_WEB_SEARCH_URL": request.app.state.config.EXTERNAL_WEB_SEARCH_URL,
  430. "EXTERNAL_WEB_SEARCH_API_KEY": request.app.state.config.EXTERNAL_WEB_SEARCH_API_KEY,
  431. "EXTERNAL_WEB_LOADER_URL": request.app.state.config.EXTERNAL_WEB_LOADER_URL,
  432. "EXTERNAL_WEB_LOADER_API_KEY": request.app.state.config.EXTERNAL_WEB_LOADER_API_KEY,
  433. "YOUTUBE_LOADER_LANGUAGE": request.app.state.config.YOUTUBE_LOADER_LANGUAGE,
  434. "YOUTUBE_LOADER_PROXY_URL": request.app.state.config.YOUTUBE_LOADER_PROXY_URL,
  435. "YOUTUBE_LOADER_TRANSLATION": request.app.state.YOUTUBE_LOADER_TRANSLATION,
  436. },
  437. }
  438. class WebConfig(BaseModel):
  439. ENABLE_WEB_SEARCH: Optional[bool] = None
  440. WEB_SEARCH_ENGINE: Optional[str] = None
  441. WEB_SEARCH_TRUST_ENV: Optional[bool] = None
  442. WEB_SEARCH_RESULT_COUNT: Optional[int] = None
  443. WEB_SEARCH_CONCURRENT_REQUESTS: Optional[int] = None
  444. WEB_SEARCH_DOMAIN_FILTER_LIST: Optional[List[str]] = []
  445. BYPASS_WEB_SEARCH_EMBEDDING_AND_RETRIEVAL: Optional[bool] = None
  446. BYPASS_WEB_SEARCH_WEB_LOADER: Optional[bool] = None
  447. SEARXNG_QUERY_URL: Optional[str] = None
  448. YACY_QUERY_URL: Optional[str] = None
  449. YACY_USERNAME: Optional[str] = None
  450. YACY_PASSWORD: Optional[str] = None
  451. GOOGLE_PSE_API_KEY: Optional[str] = None
  452. GOOGLE_PSE_ENGINE_ID: Optional[str] = None
  453. BRAVE_SEARCH_API_KEY: Optional[str] = None
  454. KAGI_SEARCH_API_KEY: Optional[str] = None
  455. MOJEEK_SEARCH_API_KEY: Optional[str] = None
  456. BOCHA_SEARCH_API_KEY: Optional[str] = None
  457. SERPSTACK_API_KEY: Optional[str] = None
  458. SERPSTACK_HTTPS: Optional[bool] = None
  459. SERPER_API_KEY: Optional[str] = None
  460. SERPLY_API_KEY: Optional[str] = None
  461. TAVILY_API_KEY: Optional[str] = None
  462. SEARCHAPI_API_KEY: Optional[str] = None
  463. SEARCHAPI_ENGINE: Optional[str] = None
  464. SERPAPI_API_KEY: Optional[str] = None
  465. SERPAPI_ENGINE: Optional[str] = None
  466. JINA_API_KEY: Optional[str] = None
  467. BING_SEARCH_V7_ENDPOINT: Optional[str] = None
  468. BING_SEARCH_V7_SUBSCRIPTION_KEY: Optional[str] = None
  469. EXA_API_KEY: Optional[str] = None
  470. PERPLEXITY_API_KEY: Optional[str] = None
  471. PERPLEXITY_MODEL: Optional[str] = None
  472. PERPLEXITY_SEARCH_CONTEXT_USAGE: Optional[str] = None
  473. SOUGOU_API_SID: Optional[str] = None
  474. SOUGOU_API_SK: Optional[str] = None
  475. WEB_LOADER_ENGINE: Optional[str] = None
  476. ENABLE_WEB_LOADER_SSL_VERIFICATION: Optional[bool] = None
  477. PLAYWRIGHT_WS_URL: Optional[str] = None
  478. PLAYWRIGHT_TIMEOUT: Optional[int] = None
  479. FIRECRAWL_API_KEY: Optional[str] = None
  480. FIRECRAWL_API_BASE_URL: Optional[str] = None
  481. TAVILY_EXTRACT_DEPTH: Optional[str] = None
  482. EXTERNAL_WEB_SEARCH_URL: Optional[str] = None
  483. EXTERNAL_WEB_SEARCH_API_KEY: Optional[str] = None
  484. EXTERNAL_WEB_LOADER_URL: Optional[str] = None
  485. EXTERNAL_WEB_LOADER_API_KEY: Optional[str] = None
  486. YOUTUBE_LOADER_LANGUAGE: Optional[List[str]] = None
  487. YOUTUBE_LOADER_PROXY_URL: Optional[str] = None
  488. YOUTUBE_LOADER_TRANSLATION: Optional[str] = None
  489. class ConfigForm(BaseModel):
  490. # RAG settings
  491. RAG_TEMPLATE: Optional[str] = None
  492. TOP_K: Optional[int] = None
  493. BYPASS_EMBEDDING_AND_RETRIEVAL: Optional[bool] = None
  494. RAG_FULL_CONTEXT: Optional[bool] = None
  495. # Hybrid search settings
  496. ENABLE_RAG_HYBRID_SEARCH: Optional[bool] = None
  497. TOP_K_RERANKER: Optional[int] = None
  498. RELEVANCE_THRESHOLD: Optional[float] = None
  499. HYBRID_BM25_WEIGHT: Optional[float] = None
  500. # Content extraction settings
  501. CONTENT_EXTRACTION_ENGINE: Optional[str] = None
  502. PDF_EXTRACT_IMAGES: Optional[bool] = None
  503. DATALAB_MARKER_API_KEY: Optional[str] = None
  504. DATALAB_MARKER_LANGS: Optional[str] = None
  505. DATALAB_MARKER_SKIP_CACHE: Optional[bool] = None
  506. DATALAB_MARKER_FORCE_OCR: Optional[bool] = None
  507. DATALAB_MARKER_PAGINATE: Optional[bool] = None
  508. DATALAB_MARKER_STRIP_EXISTING_OCR: Optional[bool] = None
  509. DATALAB_MARKER_DISABLE_IMAGE_EXTRACTION: Optional[bool] = None
  510. DATALAB_MARKER_USE_LLM: Optional[bool] = None
  511. DATALAB_MARKER_OUTPUT_FORMAT: Optional[str] = None
  512. EXTERNAL_DOCUMENT_LOADER_URL: Optional[str] = None
  513. EXTERNAL_DOCUMENT_LOADER_API_KEY: Optional[str] = None
  514. TIKA_SERVER_URL: Optional[str] = None
  515. DOCLING_SERVER_URL: Optional[str] = None
  516. DOCLING_OCR_ENGINE: Optional[str] = None
  517. DOCLING_OCR_LANG: Optional[str] = None
  518. DOCLING_DO_PICTURE_DESCRIPTION: Optional[bool] = None
  519. DOCLING_PICTURE_DESCRIPTION_MODE: Optional[str] = None
  520. DOCLING_PICTURE_DESCRIPTION_LOCAL: Optional[dict] = None
  521. DOCLING_PICTURE_DESCRIPTION_API: Optional[dict] = None
  522. DOCUMENT_INTELLIGENCE_ENDPOINT: Optional[str] = None
  523. DOCUMENT_INTELLIGENCE_KEY: Optional[str] = None
  524. MISTRAL_OCR_API_KEY: Optional[str] = None
  525. # Reranking settings
  526. RAG_RERANKING_MODEL: Optional[str] = None
  527. RAG_RERANKING_ENGINE: Optional[str] = None
  528. RAG_EXTERNAL_RERANKER_URL: Optional[str] = None
  529. RAG_EXTERNAL_RERANKER_API_KEY: Optional[str] = None
  530. # Chunking settings
  531. TEXT_SPLITTER: Optional[str] = None
  532. CHUNK_SIZE: Optional[int] = None
  533. CHUNK_OVERLAP: Optional[int] = None
  534. # File upload settings
  535. FILE_MAX_SIZE: Optional[int] = None
  536. FILE_MAX_COUNT: Optional[int] = None
  537. FILE_IMAGE_COMPRESSION_WIDTH: Optional[int] = None
  538. FILE_IMAGE_COMPRESSION_HEIGHT: Optional[int] = None
  539. ALLOWED_FILE_EXTENSIONS: Optional[List[str]] = None
  540. # Integration settings
  541. ENABLE_GOOGLE_DRIVE_INTEGRATION: Optional[bool] = None
  542. ENABLE_ONEDRIVE_INTEGRATION: Optional[bool] = None
  543. # Web search settings
  544. web: Optional[WebConfig] = None
  545. @router.post("/config/update")
  546. async def update_rag_config(
  547. request: Request, form_data: ConfigForm, user=Depends(get_admin_user)
  548. ):
  549. # RAG settings
  550. request.app.state.config.RAG_TEMPLATE = (
  551. form_data.RAG_TEMPLATE
  552. if form_data.RAG_TEMPLATE is not None
  553. else request.app.state.config.RAG_TEMPLATE
  554. )
  555. request.app.state.config.TOP_K = (
  556. form_data.TOP_K
  557. if form_data.TOP_K is not None
  558. else request.app.state.config.TOP_K
  559. )
  560. request.app.state.config.BYPASS_EMBEDDING_AND_RETRIEVAL = (
  561. form_data.BYPASS_EMBEDDING_AND_RETRIEVAL
  562. if form_data.BYPASS_EMBEDDING_AND_RETRIEVAL is not None
  563. else request.app.state.config.BYPASS_EMBEDDING_AND_RETRIEVAL
  564. )
  565. request.app.state.config.RAG_FULL_CONTEXT = (
  566. form_data.RAG_FULL_CONTEXT
  567. if form_data.RAG_FULL_CONTEXT is not None
  568. else request.app.state.config.RAG_FULL_CONTEXT
  569. )
  570. # Hybrid search settings
  571. request.app.state.config.ENABLE_RAG_HYBRID_SEARCH = (
  572. form_data.ENABLE_RAG_HYBRID_SEARCH
  573. if form_data.ENABLE_RAG_HYBRID_SEARCH is not None
  574. else request.app.state.config.ENABLE_RAG_HYBRID_SEARCH
  575. )
  576. # Free up memory if hybrid search is disabled
  577. if not request.app.state.config.ENABLE_RAG_HYBRID_SEARCH:
  578. request.app.state.rf = None
  579. request.app.state.config.TOP_K_RERANKER = (
  580. form_data.TOP_K_RERANKER
  581. if form_data.TOP_K_RERANKER is not None
  582. else request.app.state.config.TOP_K_RERANKER
  583. )
  584. request.app.state.config.RELEVANCE_THRESHOLD = (
  585. form_data.RELEVANCE_THRESHOLD
  586. if form_data.RELEVANCE_THRESHOLD is not None
  587. else request.app.state.config.RELEVANCE_THRESHOLD
  588. )
  589. request.app.state.config.HYBRID_BM25_WEIGHT = (
  590. form_data.HYBRID_BM25_WEIGHT
  591. if form_data.HYBRID_BM25_WEIGHT is not None
  592. else request.app.state.config.HYBRID_BM25_WEIGHT
  593. )
  594. # Content extraction settings
  595. request.app.state.config.CONTENT_EXTRACTION_ENGINE = (
  596. form_data.CONTENT_EXTRACTION_ENGINE
  597. if form_data.CONTENT_EXTRACTION_ENGINE is not None
  598. else request.app.state.config.CONTENT_EXTRACTION_ENGINE
  599. )
  600. request.app.state.config.PDF_EXTRACT_IMAGES = (
  601. form_data.PDF_EXTRACT_IMAGES
  602. if form_data.PDF_EXTRACT_IMAGES is not None
  603. else request.app.state.config.PDF_EXTRACT_IMAGES
  604. )
  605. request.app.state.config.DATALAB_MARKER_API_KEY = (
  606. form_data.DATALAB_MARKER_API_KEY
  607. if form_data.DATALAB_MARKER_API_KEY is not None
  608. else request.app.state.config.DATALAB_MARKER_API_KEY
  609. )
  610. request.app.state.config.DATALAB_MARKER_LANGS = (
  611. form_data.DATALAB_MARKER_LANGS
  612. if form_data.DATALAB_MARKER_LANGS is not None
  613. else request.app.state.config.DATALAB_MARKER_LANGS
  614. )
  615. request.app.state.config.DATALAB_MARKER_SKIP_CACHE = (
  616. form_data.DATALAB_MARKER_SKIP_CACHE
  617. if form_data.DATALAB_MARKER_SKIP_CACHE is not None
  618. else request.app.state.config.DATALAB_MARKER_SKIP_CACHE
  619. )
  620. request.app.state.config.DATALAB_MARKER_FORCE_OCR = (
  621. form_data.DATALAB_MARKER_FORCE_OCR
  622. if form_data.DATALAB_MARKER_FORCE_OCR is not None
  623. else request.app.state.config.DATALAB_MARKER_FORCE_OCR
  624. )
  625. request.app.state.config.DATALAB_MARKER_PAGINATE = (
  626. form_data.DATALAB_MARKER_PAGINATE
  627. if form_data.DATALAB_MARKER_PAGINATE is not None
  628. else request.app.state.config.DATALAB_MARKER_PAGINATE
  629. )
  630. request.app.state.config.DATALAB_MARKER_STRIP_EXISTING_OCR = (
  631. form_data.DATALAB_MARKER_STRIP_EXISTING_OCR
  632. if form_data.DATALAB_MARKER_STRIP_EXISTING_OCR is not None
  633. else request.app.state.config.DATALAB_MARKER_STRIP_EXISTING_OCR
  634. )
  635. request.app.state.config.DATALAB_MARKER_DISABLE_IMAGE_EXTRACTION = (
  636. form_data.DATALAB_MARKER_DISABLE_IMAGE_EXTRACTION
  637. if form_data.DATALAB_MARKER_DISABLE_IMAGE_EXTRACTION is not None
  638. else request.app.state.config.DATALAB_MARKER_DISABLE_IMAGE_EXTRACTION
  639. )
  640. request.app.state.config.DATALAB_MARKER_OUTPUT_FORMAT = (
  641. form_data.DATALAB_MARKER_OUTPUT_FORMAT
  642. if form_data.DATALAB_MARKER_OUTPUT_FORMAT is not None
  643. else request.app.state.config.DATALAB_MARKER_OUTPUT_FORMAT
  644. )
  645. request.app.state.config.DATALAB_MARKER_USE_LLM = (
  646. form_data.DATALAB_MARKER_USE_LLM
  647. if form_data.DATALAB_MARKER_USE_LLM is not None
  648. else request.app.state.config.DATALAB_MARKER_USE_LLM
  649. )
  650. request.app.state.config.EXTERNAL_DOCUMENT_LOADER_URL = (
  651. form_data.EXTERNAL_DOCUMENT_LOADER_URL
  652. if form_data.EXTERNAL_DOCUMENT_LOADER_URL is not None
  653. else request.app.state.config.EXTERNAL_DOCUMENT_LOADER_URL
  654. )
  655. request.app.state.config.EXTERNAL_DOCUMENT_LOADER_API_KEY = (
  656. form_data.EXTERNAL_DOCUMENT_LOADER_API_KEY
  657. if form_data.EXTERNAL_DOCUMENT_LOADER_API_KEY is not None
  658. else request.app.state.config.EXTERNAL_DOCUMENT_LOADER_API_KEY
  659. )
  660. request.app.state.config.TIKA_SERVER_URL = (
  661. form_data.TIKA_SERVER_URL
  662. if form_data.TIKA_SERVER_URL is not None
  663. else request.app.state.config.TIKA_SERVER_URL
  664. )
  665. request.app.state.config.DOCLING_SERVER_URL = (
  666. form_data.DOCLING_SERVER_URL
  667. if form_data.DOCLING_SERVER_URL is not None
  668. else request.app.state.config.DOCLING_SERVER_URL
  669. )
  670. request.app.state.config.DOCLING_OCR_ENGINE = (
  671. form_data.DOCLING_OCR_ENGINE
  672. if form_data.DOCLING_OCR_ENGINE is not None
  673. else request.app.state.config.DOCLING_OCR_ENGINE
  674. )
  675. request.app.state.config.DOCLING_OCR_LANG = (
  676. form_data.DOCLING_OCR_LANG
  677. if form_data.DOCLING_OCR_LANG is not None
  678. else request.app.state.config.DOCLING_OCR_LANG
  679. )
  680. request.app.state.config.DOCLING_DO_PICTURE_DESCRIPTION = (
  681. form_data.DOCLING_DO_PICTURE_DESCRIPTION
  682. if form_data.DOCLING_DO_PICTURE_DESCRIPTION is not None
  683. else request.app.state.config.DOCLING_DO_PICTURE_DESCRIPTION
  684. )
  685. request.app.state.config.DOCLING_PICTURE_DESCRIPTION_MODE = (
  686. form_data.DOCLING_PICTURE_DESCRIPTION_MODE
  687. if form_data.DOCLING_PICTURE_DESCRIPTION_MODE is not None
  688. else request.app.state.config.DOCLING_PICTURE_DESCRIPTION_MODE
  689. )
  690. request.app.state.config.DOCLING_PICTURE_DESCRIPTION_LOCAL = (
  691. form_data.DOCLING_PICTURE_DESCRIPTION_LOCAL
  692. if form_data.DOCLING_PICTURE_DESCRIPTION_LOCAL is not None
  693. else request.app.state.config.DOCLING_PICTURE_DESCRIPTION_LOCAL
  694. )
  695. request.app.state.config.DOCLING_PICTURE_DESCRIPTION_API = (
  696. form_data.DOCLING_PICTURE_DESCRIPTION_API
  697. if form_data.DOCLING_PICTURE_DESCRIPTION_API is not None
  698. else request.app.state.config.DOCLING_PICTURE_DESCRIPTION_API
  699. )
  700. request.app.state.config.DOCUMENT_INTELLIGENCE_ENDPOINT = (
  701. form_data.DOCUMENT_INTELLIGENCE_ENDPOINT
  702. if form_data.DOCUMENT_INTELLIGENCE_ENDPOINT is not None
  703. else request.app.state.config.DOCUMENT_INTELLIGENCE_ENDPOINT
  704. )
  705. request.app.state.config.DOCUMENT_INTELLIGENCE_KEY = (
  706. form_data.DOCUMENT_INTELLIGENCE_KEY
  707. if form_data.DOCUMENT_INTELLIGENCE_KEY is not None
  708. else request.app.state.config.DOCUMENT_INTELLIGENCE_KEY
  709. )
  710. request.app.state.config.MISTRAL_OCR_API_KEY = (
  711. form_data.MISTRAL_OCR_API_KEY
  712. if form_data.MISTRAL_OCR_API_KEY is not None
  713. else request.app.state.config.MISTRAL_OCR_API_KEY
  714. )
  715. # Reranking settings
  716. request.app.state.config.RAG_RERANKING_ENGINE = (
  717. form_data.RAG_RERANKING_ENGINE
  718. if form_data.RAG_RERANKING_ENGINE is not None
  719. else request.app.state.config.RAG_RERANKING_ENGINE
  720. )
  721. request.app.state.config.RAG_EXTERNAL_RERANKER_URL = (
  722. form_data.RAG_EXTERNAL_RERANKER_URL
  723. if form_data.RAG_EXTERNAL_RERANKER_URL is not None
  724. else request.app.state.config.RAG_EXTERNAL_RERANKER_URL
  725. )
  726. request.app.state.config.RAG_EXTERNAL_RERANKER_API_KEY = (
  727. form_data.RAG_EXTERNAL_RERANKER_API_KEY
  728. if form_data.RAG_EXTERNAL_RERANKER_API_KEY is not None
  729. else request.app.state.config.RAG_EXTERNAL_RERANKER_API_KEY
  730. )
  731. log.info(
  732. f"Updating reranking model: {request.app.state.config.RAG_RERANKING_MODEL} to {form_data.RAG_RERANKING_MODEL}"
  733. )
  734. try:
  735. request.app.state.config.RAG_RERANKING_MODEL = (
  736. form_data.RAG_RERANKING_MODEL
  737. if form_data.RAG_RERANKING_MODEL is not None
  738. else request.app.state.config.RAG_RERANKING_MODEL
  739. )
  740. try:
  741. request.app.state.rf = get_rf(
  742. request.app.state.config.RAG_RERANKING_ENGINE,
  743. request.app.state.config.RAG_RERANKING_MODEL,
  744. request.app.state.config.RAG_EXTERNAL_RERANKER_URL,
  745. request.app.state.config.RAG_EXTERNAL_RERANKER_API_KEY,
  746. True,
  747. )
  748. request.app.state.RERANKING_FUNCTION = get_reranking_function(
  749. request.app.state.config.RAG_RERANKING_ENGINE,
  750. request.app.state.config.RAG_RERANKING_MODEL,
  751. request.app.state.rf,
  752. )
  753. except Exception as e:
  754. log.error(f"Error loading reranking model: {e}")
  755. request.app.state.config.ENABLE_RAG_HYBRID_SEARCH = False
  756. except Exception as e:
  757. log.exception(f"Problem updating reranking model: {e}")
  758. raise HTTPException(
  759. status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
  760. detail=ERROR_MESSAGES.DEFAULT(e),
  761. )
  762. # Chunking settings
  763. request.app.state.config.TEXT_SPLITTER = (
  764. form_data.TEXT_SPLITTER
  765. if form_data.TEXT_SPLITTER is not None
  766. else request.app.state.config.TEXT_SPLITTER
  767. )
  768. request.app.state.config.CHUNK_SIZE = (
  769. form_data.CHUNK_SIZE
  770. if form_data.CHUNK_SIZE is not None
  771. else request.app.state.config.CHUNK_SIZE
  772. )
  773. request.app.state.config.CHUNK_OVERLAP = (
  774. form_data.CHUNK_OVERLAP
  775. if form_data.CHUNK_OVERLAP is not None
  776. else request.app.state.config.CHUNK_OVERLAP
  777. )
  778. # File upload settings
  779. request.app.state.config.FILE_MAX_SIZE = form_data.FILE_MAX_SIZE
  780. request.app.state.config.FILE_MAX_COUNT = form_data.FILE_MAX_COUNT
  781. request.app.state.config.FILE_IMAGE_COMPRESSION_WIDTH = (
  782. form_data.FILE_IMAGE_COMPRESSION_WIDTH
  783. )
  784. request.app.state.config.FILE_IMAGE_COMPRESSION_HEIGHT = (
  785. form_data.FILE_IMAGE_COMPRESSION_HEIGHT
  786. )
  787. request.app.state.config.ALLOWED_FILE_EXTENSIONS = (
  788. form_data.ALLOWED_FILE_EXTENSIONS
  789. if form_data.ALLOWED_FILE_EXTENSIONS is not None
  790. else request.app.state.config.ALLOWED_FILE_EXTENSIONS
  791. )
  792. # Integration settings
  793. request.app.state.config.ENABLE_GOOGLE_DRIVE_INTEGRATION = (
  794. form_data.ENABLE_GOOGLE_DRIVE_INTEGRATION
  795. if form_data.ENABLE_GOOGLE_DRIVE_INTEGRATION is not None
  796. else request.app.state.config.ENABLE_GOOGLE_DRIVE_INTEGRATION
  797. )
  798. request.app.state.config.ENABLE_ONEDRIVE_INTEGRATION = (
  799. form_data.ENABLE_ONEDRIVE_INTEGRATION
  800. if form_data.ENABLE_ONEDRIVE_INTEGRATION is not None
  801. else request.app.state.config.ENABLE_ONEDRIVE_INTEGRATION
  802. )
  803. if form_data.web is not None:
  804. # Web search settings
  805. request.app.state.config.ENABLE_WEB_SEARCH = form_data.web.ENABLE_WEB_SEARCH
  806. request.app.state.config.WEB_SEARCH_ENGINE = form_data.web.WEB_SEARCH_ENGINE
  807. request.app.state.config.WEB_SEARCH_TRUST_ENV = (
  808. form_data.web.WEB_SEARCH_TRUST_ENV
  809. )
  810. request.app.state.config.WEB_SEARCH_RESULT_COUNT = (
  811. form_data.web.WEB_SEARCH_RESULT_COUNT
  812. )
  813. request.app.state.config.WEB_SEARCH_CONCURRENT_REQUESTS = (
  814. form_data.web.WEB_SEARCH_CONCURRENT_REQUESTS
  815. )
  816. request.app.state.config.WEB_SEARCH_DOMAIN_FILTER_LIST = (
  817. form_data.web.WEB_SEARCH_DOMAIN_FILTER_LIST
  818. )
  819. request.app.state.config.BYPASS_WEB_SEARCH_EMBEDDING_AND_RETRIEVAL = (
  820. form_data.web.BYPASS_WEB_SEARCH_EMBEDDING_AND_RETRIEVAL
  821. )
  822. request.app.state.config.BYPASS_WEB_SEARCH_WEB_LOADER = (
  823. form_data.web.BYPASS_WEB_SEARCH_WEB_LOADER
  824. )
  825. request.app.state.config.SEARXNG_QUERY_URL = form_data.web.SEARXNG_QUERY_URL
  826. request.app.state.config.YACY_QUERY_URL = form_data.web.YACY_QUERY_URL
  827. request.app.state.config.YACY_USERNAME = form_data.web.YACY_USERNAME
  828. request.app.state.config.YACY_PASSWORD = form_data.web.YACY_PASSWORD
  829. request.app.state.config.GOOGLE_PSE_API_KEY = form_data.web.GOOGLE_PSE_API_KEY
  830. request.app.state.config.GOOGLE_PSE_ENGINE_ID = (
  831. form_data.web.GOOGLE_PSE_ENGINE_ID
  832. )
  833. request.app.state.config.BRAVE_SEARCH_API_KEY = (
  834. form_data.web.BRAVE_SEARCH_API_KEY
  835. )
  836. request.app.state.config.KAGI_SEARCH_API_KEY = form_data.web.KAGI_SEARCH_API_KEY
  837. request.app.state.config.MOJEEK_SEARCH_API_KEY = (
  838. form_data.web.MOJEEK_SEARCH_API_KEY
  839. )
  840. request.app.state.config.BOCHA_SEARCH_API_KEY = (
  841. form_data.web.BOCHA_SEARCH_API_KEY
  842. )
  843. request.app.state.config.SERPSTACK_API_KEY = form_data.web.SERPSTACK_API_KEY
  844. request.app.state.config.SERPSTACK_HTTPS = form_data.web.SERPSTACK_HTTPS
  845. request.app.state.config.SERPER_API_KEY = form_data.web.SERPER_API_KEY
  846. request.app.state.config.SERPLY_API_KEY = form_data.web.SERPLY_API_KEY
  847. request.app.state.config.TAVILY_API_KEY = form_data.web.TAVILY_API_KEY
  848. request.app.state.config.SEARCHAPI_API_KEY = form_data.web.SEARCHAPI_API_KEY
  849. request.app.state.config.SEARCHAPI_ENGINE = form_data.web.SEARCHAPI_ENGINE
  850. request.app.state.config.SERPAPI_API_KEY = form_data.web.SERPAPI_API_KEY
  851. request.app.state.config.SERPAPI_ENGINE = form_data.web.SERPAPI_ENGINE
  852. request.app.state.config.JINA_API_KEY = form_data.web.JINA_API_KEY
  853. request.app.state.config.BING_SEARCH_V7_ENDPOINT = (
  854. form_data.web.BING_SEARCH_V7_ENDPOINT
  855. )
  856. request.app.state.config.BING_SEARCH_V7_SUBSCRIPTION_KEY = (
  857. form_data.web.BING_SEARCH_V7_SUBSCRIPTION_KEY
  858. )
  859. request.app.state.config.EXA_API_KEY = form_data.web.EXA_API_KEY
  860. request.app.state.config.PERPLEXITY_API_KEY = form_data.web.PERPLEXITY_API_KEY
  861. request.app.state.config.PERPLEXITY_MODEL = form_data.web.PERPLEXITY_MODEL
  862. request.app.state.config.PERPLEXITY_SEARCH_CONTEXT_USAGE = (
  863. form_data.web.PERPLEXITY_SEARCH_CONTEXT_USAGE
  864. )
  865. request.app.state.config.SOUGOU_API_SID = form_data.web.SOUGOU_API_SID
  866. request.app.state.config.SOUGOU_API_SK = form_data.web.SOUGOU_API_SK
  867. # Web loader settings
  868. request.app.state.config.WEB_LOADER_ENGINE = form_data.web.WEB_LOADER_ENGINE
  869. request.app.state.config.ENABLE_WEB_LOADER_SSL_VERIFICATION = (
  870. form_data.web.ENABLE_WEB_LOADER_SSL_VERIFICATION
  871. )
  872. request.app.state.config.PLAYWRIGHT_WS_URL = form_data.web.PLAYWRIGHT_WS_URL
  873. request.app.state.config.PLAYWRIGHT_TIMEOUT = form_data.web.PLAYWRIGHT_TIMEOUT
  874. request.app.state.config.FIRECRAWL_API_KEY = form_data.web.FIRECRAWL_API_KEY
  875. request.app.state.config.FIRECRAWL_API_BASE_URL = (
  876. form_data.web.FIRECRAWL_API_BASE_URL
  877. )
  878. request.app.state.config.EXTERNAL_WEB_SEARCH_URL = (
  879. form_data.web.EXTERNAL_WEB_SEARCH_URL
  880. )
  881. request.app.state.config.EXTERNAL_WEB_SEARCH_API_KEY = (
  882. form_data.web.EXTERNAL_WEB_SEARCH_API_KEY
  883. )
  884. request.app.state.config.EXTERNAL_WEB_LOADER_URL = (
  885. form_data.web.EXTERNAL_WEB_LOADER_URL
  886. )
  887. request.app.state.config.EXTERNAL_WEB_LOADER_API_KEY = (
  888. form_data.web.EXTERNAL_WEB_LOADER_API_KEY
  889. )
  890. request.app.state.config.TAVILY_EXTRACT_DEPTH = (
  891. form_data.web.TAVILY_EXTRACT_DEPTH
  892. )
  893. request.app.state.config.YOUTUBE_LOADER_LANGUAGE = (
  894. form_data.web.YOUTUBE_LOADER_LANGUAGE
  895. )
  896. request.app.state.config.YOUTUBE_LOADER_PROXY_URL = (
  897. form_data.web.YOUTUBE_LOADER_PROXY_URL
  898. )
  899. request.app.state.YOUTUBE_LOADER_TRANSLATION = (
  900. form_data.web.YOUTUBE_LOADER_TRANSLATION
  901. )
  902. return {
  903. "status": True,
  904. # RAG settings
  905. "RAG_TEMPLATE": request.app.state.config.RAG_TEMPLATE,
  906. "TOP_K": request.app.state.config.TOP_K,
  907. "BYPASS_EMBEDDING_AND_RETRIEVAL": request.app.state.config.BYPASS_EMBEDDING_AND_RETRIEVAL,
  908. "RAG_FULL_CONTEXT": request.app.state.config.RAG_FULL_CONTEXT,
  909. # Hybrid search settings
  910. "ENABLE_RAG_HYBRID_SEARCH": request.app.state.config.ENABLE_RAG_HYBRID_SEARCH,
  911. "TOP_K_RERANKER": request.app.state.config.TOP_K_RERANKER,
  912. "RELEVANCE_THRESHOLD": request.app.state.config.RELEVANCE_THRESHOLD,
  913. "HYBRID_BM25_WEIGHT": request.app.state.config.HYBRID_BM25_WEIGHT,
  914. # Content extraction settings
  915. "CONTENT_EXTRACTION_ENGINE": request.app.state.config.CONTENT_EXTRACTION_ENGINE,
  916. "PDF_EXTRACT_IMAGES": request.app.state.config.PDF_EXTRACT_IMAGES,
  917. "DATALAB_MARKER_API_KEY": request.app.state.config.DATALAB_MARKER_API_KEY,
  918. "DATALAB_MARKER_LANGS": request.app.state.config.DATALAB_MARKER_LANGS,
  919. "DATALAB_MARKER_SKIP_CACHE": request.app.state.config.DATALAB_MARKER_SKIP_CACHE,
  920. "DATALAB_MARKER_FORCE_OCR": request.app.state.config.DATALAB_MARKER_FORCE_OCR,
  921. "DATALAB_MARKER_PAGINATE": request.app.state.config.DATALAB_MARKER_PAGINATE,
  922. "DATALAB_MARKER_STRIP_EXISTING_OCR": request.app.state.config.DATALAB_MARKER_STRIP_EXISTING_OCR,
  923. "DATALAB_MARKER_DISABLE_IMAGE_EXTRACTION": request.app.state.config.DATALAB_MARKER_DISABLE_IMAGE_EXTRACTION,
  924. "DATALAB_MARKER_USE_LLM": request.app.state.config.DATALAB_MARKER_USE_LLM,
  925. "DATALAB_MARKER_OUTPUT_FORMAT": request.app.state.config.DATALAB_MARKER_OUTPUT_FORMAT,
  926. "EXTERNAL_DOCUMENT_LOADER_URL": request.app.state.config.EXTERNAL_DOCUMENT_LOADER_URL,
  927. "EXTERNAL_DOCUMENT_LOADER_API_KEY": request.app.state.config.EXTERNAL_DOCUMENT_LOADER_API_KEY,
  928. "TIKA_SERVER_URL": request.app.state.config.TIKA_SERVER_URL,
  929. "DOCLING_SERVER_URL": request.app.state.config.DOCLING_SERVER_URL,
  930. "DOCLING_OCR_ENGINE": request.app.state.config.DOCLING_OCR_ENGINE,
  931. "DOCLING_OCR_LANG": request.app.state.config.DOCLING_OCR_LANG,
  932. "DOCLING_DO_PICTURE_DESCRIPTION": request.app.state.config.DOCLING_DO_PICTURE_DESCRIPTION,
  933. "DOCLING_PICTURE_DESCRIPTION_MODE": request.app.state.config.DOCLING_PICTURE_DESCRIPTION_MODE,
  934. "DOCLING_PICTURE_DESCRIPTION_LOCAL": request.app.state.config.DOCLING_PICTURE_DESCRIPTION_LOCAL,
  935. "DOCLING_PICTURE_DESCRIPTION_API": request.app.state.config.DOCLING_PICTURE_DESCRIPTION_API,
  936. "DOCUMENT_INTELLIGENCE_ENDPOINT": request.app.state.config.DOCUMENT_INTELLIGENCE_ENDPOINT,
  937. "DOCUMENT_INTELLIGENCE_KEY": request.app.state.config.DOCUMENT_INTELLIGENCE_KEY,
  938. "MISTRAL_OCR_API_KEY": request.app.state.config.MISTRAL_OCR_API_KEY,
  939. # Reranking settings
  940. "RAG_RERANKING_MODEL": request.app.state.config.RAG_RERANKING_MODEL,
  941. "RAG_RERANKING_ENGINE": request.app.state.config.RAG_RERANKING_ENGINE,
  942. "RAG_EXTERNAL_RERANKER_URL": request.app.state.config.RAG_EXTERNAL_RERANKER_URL,
  943. "RAG_EXTERNAL_RERANKER_API_KEY": request.app.state.config.RAG_EXTERNAL_RERANKER_API_KEY,
  944. # Chunking settings
  945. "TEXT_SPLITTER": request.app.state.config.TEXT_SPLITTER,
  946. "CHUNK_SIZE": request.app.state.config.CHUNK_SIZE,
  947. "CHUNK_OVERLAP": request.app.state.config.CHUNK_OVERLAP,
  948. # File upload settings
  949. "FILE_MAX_SIZE": request.app.state.config.FILE_MAX_SIZE,
  950. "FILE_MAX_COUNT": request.app.state.config.FILE_MAX_COUNT,
  951. "FILE_IMAGE_COMPRESSION_WIDTH": request.app.state.config.FILE_IMAGE_COMPRESSION_WIDTH,
  952. "FILE_IMAGE_COMPRESSION_HEIGHT": request.app.state.config.FILE_IMAGE_COMPRESSION_HEIGHT,
  953. "ALLOWED_FILE_EXTENSIONS": request.app.state.config.ALLOWED_FILE_EXTENSIONS,
  954. # Integration settings
  955. "ENABLE_GOOGLE_DRIVE_INTEGRATION": request.app.state.config.ENABLE_GOOGLE_DRIVE_INTEGRATION,
  956. "ENABLE_ONEDRIVE_INTEGRATION": request.app.state.config.ENABLE_ONEDRIVE_INTEGRATION,
  957. # Web search settings
  958. "web": {
  959. "ENABLE_WEB_SEARCH": request.app.state.config.ENABLE_WEB_SEARCH,
  960. "WEB_SEARCH_ENGINE": request.app.state.config.WEB_SEARCH_ENGINE,
  961. "WEB_SEARCH_TRUST_ENV": request.app.state.config.WEB_SEARCH_TRUST_ENV,
  962. "WEB_SEARCH_RESULT_COUNT": request.app.state.config.WEB_SEARCH_RESULT_COUNT,
  963. "WEB_SEARCH_CONCURRENT_REQUESTS": request.app.state.config.WEB_SEARCH_CONCURRENT_REQUESTS,
  964. "WEB_SEARCH_DOMAIN_FILTER_LIST": request.app.state.config.WEB_SEARCH_DOMAIN_FILTER_LIST,
  965. "BYPASS_WEB_SEARCH_EMBEDDING_AND_RETRIEVAL": request.app.state.config.BYPASS_WEB_SEARCH_EMBEDDING_AND_RETRIEVAL,
  966. "BYPASS_WEB_SEARCH_WEB_LOADER": request.app.state.config.BYPASS_WEB_SEARCH_WEB_LOADER,
  967. "SEARXNG_QUERY_URL": request.app.state.config.SEARXNG_QUERY_URL,
  968. "YACY_QUERY_URL": request.app.state.config.YACY_QUERY_URL,
  969. "YACY_USERNAME": request.app.state.config.YACY_USERNAME,
  970. "YACY_PASSWORD": request.app.state.config.YACY_PASSWORD,
  971. "GOOGLE_PSE_API_KEY": request.app.state.config.GOOGLE_PSE_API_KEY,
  972. "GOOGLE_PSE_ENGINE_ID": request.app.state.config.GOOGLE_PSE_ENGINE_ID,
  973. "BRAVE_SEARCH_API_KEY": request.app.state.config.BRAVE_SEARCH_API_KEY,
  974. "KAGI_SEARCH_API_KEY": request.app.state.config.KAGI_SEARCH_API_KEY,
  975. "MOJEEK_SEARCH_API_KEY": request.app.state.config.MOJEEK_SEARCH_API_KEY,
  976. "BOCHA_SEARCH_API_KEY": request.app.state.config.BOCHA_SEARCH_API_KEY,
  977. "SERPSTACK_API_KEY": request.app.state.config.SERPSTACK_API_KEY,
  978. "SERPSTACK_HTTPS": request.app.state.config.SERPSTACK_HTTPS,
  979. "SERPER_API_KEY": request.app.state.config.SERPER_API_KEY,
  980. "SERPLY_API_KEY": request.app.state.config.SERPLY_API_KEY,
  981. "TAVILY_API_KEY": request.app.state.config.TAVILY_API_KEY,
  982. "SEARCHAPI_API_KEY": request.app.state.config.SEARCHAPI_API_KEY,
  983. "SEARCHAPI_ENGINE": request.app.state.config.SEARCHAPI_ENGINE,
  984. "SERPAPI_API_KEY": request.app.state.config.SERPAPI_API_KEY,
  985. "SERPAPI_ENGINE": request.app.state.config.SERPAPI_ENGINE,
  986. "JINA_API_KEY": request.app.state.config.JINA_API_KEY,
  987. "BING_SEARCH_V7_ENDPOINT": request.app.state.config.BING_SEARCH_V7_ENDPOINT,
  988. "BING_SEARCH_V7_SUBSCRIPTION_KEY": request.app.state.config.BING_SEARCH_V7_SUBSCRIPTION_KEY,
  989. "EXA_API_KEY": request.app.state.config.EXA_API_KEY,
  990. "PERPLEXITY_API_KEY": request.app.state.config.PERPLEXITY_API_KEY,
  991. "PERPLEXITY_MODEL": request.app.state.config.PERPLEXITY_MODEL,
  992. "PERPLEXITY_SEARCH_CONTEXT_USAGE": request.app.state.config.PERPLEXITY_SEARCH_CONTEXT_USAGE,
  993. "SOUGOU_API_SID": request.app.state.config.SOUGOU_API_SID,
  994. "SOUGOU_API_SK": request.app.state.config.SOUGOU_API_SK,
  995. "WEB_LOADER_ENGINE": request.app.state.config.WEB_LOADER_ENGINE,
  996. "ENABLE_WEB_LOADER_SSL_VERIFICATION": request.app.state.config.ENABLE_WEB_LOADER_SSL_VERIFICATION,
  997. "PLAYWRIGHT_WS_URL": request.app.state.config.PLAYWRIGHT_WS_URL,
  998. "PLAYWRIGHT_TIMEOUT": request.app.state.config.PLAYWRIGHT_TIMEOUT,
  999. "FIRECRAWL_API_KEY": request.app.state.config.FIRECRAWL_API_KEY,
  1000. "FIRECRAWL_API_BASE_URL": request.app.state.config.FIRECRAWL_API_BASE_URL,
  1001. "TAVILY_EXTRACT_DEPTH": request.app.state.config.TAVILY_EXTRACT_DEPTH,
  1002. "EXTERNAL_WEB_SEARCH_URL": request.app.state.config.EXTERNAL_WEB_SEARCH_URL,
  1003. "EXTERNAL_WEB_SEARCH_API_KEY": request.app.state.config.EXTERNAL_WEB_SEARCH_API_KEY,
  1004. "EXTERNAL_WEB_LOADER_URL": request.app.state.config.EXTERNAL_WEB_LOADER_URL,
  1005. "EXTERNAL_WEB_LOADER_API_KEY": request.app.state.config.EXTERNAL_WEB_LOADER_API_KEY,
  1006. "YOUTUBE_LOADER_LANGUAGE": request.app.state.config.YOUTUBE_LOADER_LANGUAGE,
  1007. "YOUTUBE_LOADER_PROXY_URL": request.app.state.config.YOUTUBE_LOADER_PROXY_URL,
  1008. "YOUTUBE_LOADER_TRANSLATION": request.app.state.YOUTUBE_LOADER_TRANSLATION,
  1009. },
  1010. }
  1011. ####################################
  1012. #
  1013. # Document process and retrieval
  1014. #
  1015. ####################################
  1016. def save_docs_to_vector_db(
  1017. request: Request,
  1018. docs,
  1019. collection_name,
  1020. metadata: Optional[dict] = None,
  1021. overwrite: bool = False,
  1022. split: bool = True,
  1023. add: bool = False,
  1024. user=None,
  1025. ) -> bool:
  1026. def _get_docs_info(docs: list[Document]) -> str:
  1027. docs_info = set()
  1028. # Trying to select relevant metadata identifying the document.
  1029. for doc in docs:
  1030. metadata = getattr(doc, "metadata", {})
  1031. doc_name = metadata.get("name", "")
  1032. if not doc_name:
  1033. doc_name = metadata.get("title", "")
  1034. if not doc_name:
  1035. doc_name = metadata.get("source", "")
  1036. if doc_name:
  1037. docs_info.add(doc_name)
  1038. return ", ".join(docs_info)
  1039. log.info(
  1040. f"save_docs_to_vector_db: document {_get_docs_info(docs)} {collection_name}"
  1041. )
  1042. # Check if entries with the same hash (metadata.hash) already exist
  1043. if metadata and "hash" in metadata:
  1044. result = VECTOR_DB_CLIENT.query(
  1045. collection_name=collection_name,
  1046. filter={"hash": metadata["hash"]},
  1047. )
  1048. if result is not None:
  1049. existing_doc_ids = result.ids[0]
  1050. if existing_doc_ids:
  1051. log.info(f"Document with hash {metadata['hash']} already exists")
  1052. raise ValueError(ERROR_MESSAGES.DUPLICATE_CONTENT)
  1053. if split:
  1054. if request.app.state.config.TEXT_SPLITTER in ["", "character"]:
  1055. text_splitter = RecursiveCharacterTextSplitter(
  1056. chunk_size=request.app.state.config.CHUNK_SIZE,
  1057. chunk_overlap=request.app.state.config.CHUNK_OVERLAP,
  1058. add_start_index=True,
  1059. )
  1060. docs = text_splitter.split_documents(docs)
  1061. elif request.app.state.config.TEXT_SPLITTER == "token":
  1062. log.info(
  1063. f"Using token text splitter: {request.app.state.config.TIKTOKEN_ENCODING_NAME}"
  1064. )
  1065. tiktoken.get_encoding(str(request.app.state.config.TIKTOKEN_ENCODING_NAME))
  1066. text_splitter = TokenTextSplitter(
  1067. encoding_name=str(request.app.state.config.TIKTOKEN_ENCODING_NAME),
  1068. chunk_size=request.app.state.config.CHUNK_SIZE,
  1069. chunk_overlap=request.app.state.config.CHUNK_OVERLAP,
  1070. add_start_index=True,
  1071. )
  1072. docs = text_splitter.split_documents(docs)
  1073. elif request.app.state.config.TEXT_SPLITTER == "markdown_header":
  1074. log.info("Using markdown header text splitter")
  1075. # Define headers to split on - covering most common markdown header levels
  1076. headers_to_split_on = [
  1077. ("#", "Header 1"),
  1078. ("##", "Header 2"),
  1079. ("###", "Header 3"),
  1080. ("####", "Header 4"),
  1081. ("#####", "Header 5"),
  1082. ("######", "Header 6"),
  1083. ]
  1084. markdown_splitter = MarkdownHeaderTextSplitter(
  1085. headers_to_split_on=headers_to_split_on,
  1086. strip_headers=False, # Keep headers in content for context
  1087. )
  1088. md_split_docs = []
  1089. for doc in docs:
  1090. md_header_splits = markdown_splitter.split_text(doc.page_content)
  1091. text_splitter = RecursiveCharacterTextSplitter(
  1092. chunk_size=request.app.state.config.CHUNK_SIZE,
  1093. chunk_overlap=request.app.state.config.CHUNK_OVERLAP,
  1094. add_start_index=True,
  1095. )
  1096. md_header_splits = text_splitter.split_documents(md_header_splits)
  1097. # Convert back to Document objects, preserving original metadata
  1098. for split_chunk in md_header_splits:
  1099. headings_list = []
  1100. # Extract header values in order based on headers_to_split_on
  1101. for _, header_meta_key_name in headers_to_split_on:
  1102. if header_meta_key_name in split_chunk.metadata:
  1103. headings_list.append(
  1104. split_chunk.metadata[header_meta_key_name]
  1105. )
  1106. md_split_docs.append(
  1107. Document(
  1108. page_content=split_chunk.page_content,
  1109. metadata={**doc.metadata, "headings": headings_list},
  1110. )
  1111. )
  1112. docs = md_split_docs
  1113. else:
  1114. raise ValueError(ERROR_MESSAGES.DEFAULT("Invalid text splitter"))
  1115. if len(docs) == 0:
  1116. raise ValueError(ERROR_MESSAGES.EMPTY_CONTENT)
  1117. texts = [doc.page_content for doc in docs]
  1118. metadatas = [
  1119. {
  1120. **doc.metadata,
  1121. **(metadata if metadata else {}),
  1122. "embedding_config": {
  1123. "engine": request.app.state.config.RAG_EMBEDDING_ENGINE,
  1124. "model": request.app.state.config.RAG_EMBEDDING_MODEL,
  1125. },
  1126. }
  1127. for doc in docs
  1128. ]
  1129. try:
  1130. if VECTOR_DB_CLIENT.has_collection(collection_name=collection_name):
  1131. log.info(f"collection {collection_name} already exists")
  1132. if overwrite:
  1133. VECTOR_DB_CLIENT.delete_collection(collection_name=collection_name)
  1134. log.info(f"deleting existing collection {collection_name}")
  1135. elif add is False:
  1136. log.info(
  1137. f"collection {collection_name} already exists, overwrite is False and add is False"
  1138. )
  1139. return True
  1140. log.info(f"adding to collection {collection_name}")
  1141. embedding_function = get_embedding_function(
  1142. request.app.state.config.RAG_EMBEDDING_ENGINE,
  1143. request.app.state.config.RAG_EMBEDDING_MODEL,
  1144. request.app.state.ef,
  1145. (
  1146. request.app.state.config.RAG_OPENAI_API_BASE_URL
  1147. if request.app.state.config.RAG_EMBEDDING_ENGINE == "openai"
  1148. else (
  1149. request.app.state.config.RAG_OLLAMA_BASE_URL
  1150. if request.app.state.config.RAG_EMBEDDING_ENGINE == "ollama"
  1151. else request.app.state.config.RAG_AZURE_OPENAI_BASE_URL
  1152. )
  1153. ),
  1154. (
  1155. request.app.state.config.RAG_OPENAI_API_KEY
  1156. if request.app.state.config.RAG_EMBEDDING_ENGINE == "openai"
  1157. else (
  1158. request.app.state.config.RAG_OLLAMA_API_KEY
  1159. if request.app.state.config.RAG_EMBEDDING_ENGINE == "ollama"
  1160. else request.app.state.config.RAG_AZURE_OPENAI_API_KEY
  1161. )
  1162. ),
  1163. request.app.state.config.RAG_EMBEDDING_BATCH_SIZE,
  1164. azure_api_version=(
  1165. request.app.state.config.RAG_AZURE_OPENAI_API_VERSION
  1166. if request.app.state.config.RAG_EMBEDDING_ENGINE == "azure_openai"
  1167. else None
  1168. ),
  1169. )
  1170. embeddings = embedding_function(
  1171. list(map(lambda x: x.replace("\n", " "), texts)),
  1172. prefix=RAG_EMBEDDING_CONTENT_PREFIX,
  1173. user=user,
  1174. )
  1175. items = [
  1176. {
  1177. "id": str(uuid.uuid4()),
  1178. "text": text,
  1179. "vector": embeddings[idx],
  1180. "metadata": metadatas[idx],
  1181. }
  1182. for idx, text in enumerate(texts)
  1183. ]
  1184. VECTOR_DB_CLIENT.insert(
  1185. collection_name=collection_name,
  1186. items=items,
  1187. )
  1188. return True
  1189. except Exception as e:
  1190. log.exception(e)
  1191. raise e
  1192. class ProcessFileForm(BaseModel):
  1193. file_id: str
  1194. content: Optional[str] = None
  1195. collection_name: Optional[str] = None
  1196. @router.post("/process/file")
  1197. def process_file(
  1198. request: Request,
  1199. form_data: ProcessFileForm,
  1200. user=Depends(get_verified_user),
  1201. ):
  1202. try:
  1203. file = Files.get_file_by_id(form_data.file_id)
  1204. collection_name = form_data.collection_name
  1205. if collection_name is None:
  1206. collection_name = f"file-{file.id}"
  1207. if form_data.content:
  1208. # Update the content in the file
  1209. # Usage: /files/{file_id}/data/content/update, /files/ (audio file upload pipeline)
  1210. try:
  1211. # /files/{file_id}/data/content/update
  1212. VECTOR_DB_CLIENT.delete_collection(collection_name=f"file-{file.id}")
  1213. except:
  1214. # Audio file upload pipeline
  1215. pass
  1216. docs = [
  1217. Document(
  1218. page_content=form_data.content.replace("<br/>", "\n"),
  1219. metadata={
  1220. **file.meta,
  1221. "name": file.filename,
  1222. "created_by": file.user_id,
  1223. "file_id": file.id,
  1224. "source": file.filename,
  1225. },
  1226. )
  1227. ]
  1228. text_content = form_data.content
  1229. elif form_data.collection_name:
  1230. # Check if the file has already been processed and save the content
  1231. # Usage: /knowledge/{id}/file/add, /knowledge/{id}/file/update
  1232. result = VECTOR_DB_CLIENT.query(
  1233. collection_name=f"file-{file.id}", filter={"file_id": file.id}
  1234. )
  1235. if result is not None and len(result.ids[0]) > 0:
  1236. docs = [
  1237. Document(
  1238. page_content=result.documents[0][idx],
  1239. metadata=result.metadatas[0][idx],
  1240. )
  1241. for idx, id in enumerate(result.ids[0])
  1242. ]
  1243. else:
  1244. docs = [
  1245. Document(
  1246. page_content=file.data.get("content", ""),
  1247. metadata={
  1248. **file.meta,
  1249. "name": file.filename,
  1250. "created_by": file.user_id,
  1251. "file_id": file.id,
  1252. "source": file.filename,
  1253. },
  1254. )
  1255. ]
  1256. text_content = file.data.get("content", "")
  1257. else:
  1258. # Process the file and save the content
  1259. # Usage: /files/
  1260. file_path = file.path
  1261. if file_path:
  1262. file_path = Storage.get_file(file_path)
  1263. loader = Loader(
  1264. engine=request.app.state.config.CONTENT_EXTRACTION_ENGINE,
  1265. DATALAB_MARKER_API_KEY=request.app.state.config.DATALAB_MARKER_API_KEY,
  1266. DATALAB_MARKER_LANGS=request.app.state.config.DATALAB_MARKER_LANGS,
  1267. DATALAB_MARKER_SKIP_CACHE=request.app.state.config.DATALAB_MARKER_SKIP_CACHE,
  1268. DATALAB_MARKER_FORCE_OCR=request.app.state.config.DATALAB_MARKER_FORCE_OCR,
  1269. DATALAB_MARKER_PAGINATE=request.app.state.config.DATALAB_MARKER_PAGINATE,
  1270. DATALAB_MARKER_STRIP_EXISTING_OCR=request.app.state.config.DATALAB_MARKER_STRIP_EXISTING_OCR,
  1271. DATALAB_MARKER_DISABLE_IMAGE_EXTRACTION=request.app.state.config.DATALAB_MARKER_DISABLE_IMAGE_EXTRACTION,
  1272. DATALAB_MARKER_USE_LLM=request.app.state.config.DATALAB_MARKER_USE_LLM,
  1273. DATALAB_MARKER_OUTPUT_FORMAT=request.app.state.config.DATALAB_MARKER_OUTPUT_FORMAT,
  1274. EXTERNAL_DOCUMENT_LOADER_URL=request.app.state.config.EXTERNAL_DOCUMENT_LOADER_URL,
  1275. EXTERNAL_DOCUMENT_LOADER_API_KEY=request.app.state.config.EXTERNAL_DOCUMENT_LOADER_API_KEY,
  1276. TIKA_SERVER_URL=request.app.state.config.TIKA_SERVER_URL,
  1277. DOCLING_SERVER_URL=request.app.state.config.DOCLING_SERVER_URL,
  1278. DOCLING_PARAMS={
  1279. "ocr_engine": request.app.state.config.DOCLING_OCR_ENGINE,
  1280. "ocr_lang": request.app.state.config.DOCLING_OCR_LANG,
  1281. "do_picture_description": request.app.state.config.DOCLING_DO_PICTURE_DESCRIPTION,
  1282. "picture_description_mode": request.app.state.config.DOCLING_PICTURE_DESCRIPTION_MODE,
  1283. "picture_description_local": request.app.state.config.DOCLING_PICTURE_DESCRIPTION_LOCAL,
  1284. "picture_description_api": request.app.state.config.DOCLING_PICTURE_DESCRIPTION_API,
  1285. },
  1286. PDF_EXTRACT_IMAGES=request.app.state.config.PDF_EXTRACT_IMAGES,
  1287. DOCUMENT_INTELLIGENCE_ENDPOINT=request.app.state.config.DOCUMENT_INTELLIGENCE_ENDPOINT,
  1288. DOCUMENT_INTELLIGENCE_KEY=request.app.state.config.DOCUMENT_INTELLIGENCE_KEY,
  1289. MISTRAL_OCR_API_KEY=request.app.state.config.MISTRAL_OCR_API_KEY,
  1290. )
  1291. docs = loader.load(
  1292. file.filename, file.meta.get("content_type"), file_path
  1293. )
  1294. docs = [
  1295. Document(
  1296. page_content=doc.page_content,
  1297. metadata={
  1298. **doc.metadata,
  1299. "name": file.filename,
  1300. "created_by": file.user_id,
  1301. "file_id": file.id,
  1302. "source": file.filename,
  1303. },
  1304. )
  1305. for doc in docs
  1306. ]
  1307. else:
  1308. docs = [
  1309. Document(
  1310. page_content=file.data.get("content", ""),
  1311. metadata={
  1312. **file.meta,
  1313. "name": file.filename,
  1314. "created_by": file.user_id,
  1315. "file_id": file.id,
  1316. "source": file.filename,
  1317. },
  1318. )
  1319. ]
  1320. text_content = " ".join([doc.page_content for doc in docs])
  1321. log.debug(f"text_content: {text_content}")
  1322. Files.update_file_data_by_id(
  1323. file.id,
  1324. {"content": text_content},
  1325. )
  1326. hash = calculate_sha256_string(text_content)
  1327. Files.update_file_hash_by_id(file.id, hash)
  1328. if not request.app.state.config.BYPASS_EMBEDDING_AND_RETRIEVAL:
  1329. try:
  1330. result = save_docs_to_vector_db(
  1331. request,
  1332. docs=docs,
  1333. collection_name=collection_name,
  1334. metadata={
  1335. "file_id": file.id,
  1336. "name": file.filename,
  1337. "hash": hash,
  1338. },
  1339. add=(True if form_data.collection_name else False),
  1340. user=user,
  1341. )
  1342. if result:
  1343. Files.update_file_metadata_by_id(
  1344. file.id,
  1345. {
  1346. "collection_name": collection_name,
  1347. },
  1348. )
  1349. return {
  1350. "status": True,
  1351. "collection_name": collection_name,
  1352. "filename": file.filename,
  1353. "content": text_content,
  1354. }
  1355. except Exception as e:
  1356. raise e
  1357. else:
  1358. return {
  1359. "status": True,
  1360. "collection_name": None,
  1361. "filename": file.filename,
  1362. "content": text_content,
  1363. }
  1364. except Exception as e:
  1365. log.exception(e)
  1366. if "No pandoc was found" in str(e):
  1367. raise HTTPException(
  1368. status_code=status.HTTP_400_BAD_REQUEST,
  1369. detail=ERROR_MESSAGES.PANDOC_NOT_INSTALLED,
  1370. )
  1371. else:
  1372. raise HTTPException(
  1373. status_code=status.HTTP_400_BAD_REQUEST,
  1374. detail=str(e),
  1375. )
  1376. class ProcessTextForm(BaseModel):
  1377. name: str
  1378. content: str
  1379. collection_name: Optional[str] = None
  1380. @router.post("/process/text")
  1381. def process_text(
  1382. request: Request,
  1383. form_data: ProcessTextForm,
  1384. user=Depends(get_verified_user),
  1385. ):
  1386. collection_name = form_data.collection_name
  1387. if collection_name is None:
  1388. collection_name = calculate_sha256_string(form_data.content)
  1389. docs = [
  1390. Document(
  1391. page_content=form_data.content,
  1392. metadata={"name": form_data.name, "created_by": user.id},
  1393. )
  1394. ]
  1395. text_content = form_data.content
  1396. log.debug(f"text_content: {text_content}")
  1397. result = save_docs_to_vector_db(request, docs, collection_name, user=user)
  1398. if result:
  1399. return {
  1400. "status": True,
  1401. "collection_name": collection_name,
  1402. "content": text_content,
  1403. }
  1404. else:
  1405. raise HTTPException(
  1406. status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
  1407. detail=ERROR_MESSAGES.DEFAULT(),
  1408. )
  1409. @router.post("/process/youtube")
  1410. def process_youtube_video(
  1411. request: Request, form_data: ProcessUrlForm, user=Depends(get_verified_user)
  1412. ):
  1413. try:
  1414. collection_name = form_data.collection_name
  1415. if not collection_name:
  1416. collection_name = calculate_sha256_string(form_data.url)[:63]
  1417. loader = YoutubeLoader(
  1418. form_data.url,
  1419. language=request.app.state.config.YOUTUBE_LOADER_LANGUAGE,
  1420. proxy_url=request.app.state.config.YOUTUBE_LOADER_PROXY_URL,
  1421. )
  1422. docs = loader.load()
  1423. content = " ".join([doc.page_content for doc in docs])
  1424. log.debug(f"text_content: {content}")
  1425. save_docs_to_vector_db(
  1426. request, docs, collection_name, overwrite=True, user=user
  1427. )
  1428. return {
  1429. "status": True,
  1430. "collection_name": collection_name,
  1431. "filename": form_data.url,
  1432. "file": {
  1433. "data": {
  1434. "content": content,
  1435. },
  1436. "meta": {
  1437. "name": form_data.url,
  1438. },
  1439. },
  1440. }
  1441. except Exception as e:
  1442. log.exception(e)
  1443. raise HTTPException(
  1444. status_code=status.HTTP_400_BAD_REQUEST,
  1445. detail=ERROR_MESSAGES.DEFAULT(e),
  1446. )
  1447. @router.post("/process/web")
  1448. def process_web(
  1449. request: Request, form_data: ProcessUrlForm, user=Depends(get_verified_user)
  1450. ):
  1451. try:
  1452. collection_name = form_data.collection_name
  1453. if not collection_name:
  1454. collection_name = calculate_sha256_string(form_data.url)[:63]
  1455. loader = get_web_loader(
  1456. form_data.url,
  1457. verify_ssl=request.app.state.config.ENABLE_WEB_LOADER_SSL_VERIFICATION,
  1458. requests_per_second=request.app.state.config.WEB_SEARCH_CONCURRENT_REQUESTS,
  1459. )
  1460. docs = loader.load()
  1461. content = " ".join([doc.page_content for doc in docs])
  1462. log.debug(f"text_content: {content}")
  1463. if not request.app.state.config.BYPASS_WEB_SEARCH_EMBEDDING_AND_RETRIEVAL:
  1464. save_docs_to_vector_db(
  1465. request, docs, collection_name, overwrite=True, user=user
  1466. )
  1467. else:
  1468. collection_name = None
  1469. return {
  1470. "status": True,
  1471. "collection_name": collection_name,
  1472. "filename": form_data.url,
  1473. "file": {
  1474. "data": {
  1475. "content": content,
  1476. },
  1477. "meta": {
  1478. "name": form_data.url,
  1479. "source": form_data.url,
  1480. },
  1481. },
  1482. }
  1483. except Exception as e:
  1484. log.exception(e)
  1485. raise HTTPException(
  1486. status_code=status.HTTP_400_BAD_REQUEST,
  1487. detail=ERROR_MESSAGES.DEFAULT(e),
  1488. )
  1489. def search_web(request: Request, engine: str, query: str) -> list[SearchResult]:
  1490. """Search the web using a search engine and return the results as a list of SearchResult objects.
  1491. Will look for a search engine API key in environment variables in the following order:
  1492. - SEARXNG_QUERY_URL
  1493. - YACY_QUERY_URL + YACY_USERNAME + YACY_PASSWORD
  1494. - GOOGLE_PSE_API_KEY + GOOGLE_PSE_ENGINE_ID
  1495. - BRAVE_SEARCH_API_KEY
  1496. - KAGI_SEARCH_API_KEY
  1497. - MOJEEK_SEARCH_API_KEY
  1498. - BOCHA_SEARCH_API_KEY
  1499. - SERPSTACK_API_KEY
  1500. - SERPER_API_KEY
  1501. - SERPLY_API_KEY
  1502. - TAVILY_API_KEY
  1503. - EXA_API_KEY
  1504. - PERPLEXITY_API_KEY
  1505. - SOUGOU_API_SID + SOUGOU_API_SK
  1506. - SEARCHAPI_API_KEY + SEARCHAPI_ENGINE (by default `google`)
  1507. - SERPAPI_API_KEY + SERPAPI_ENGINE (by default `google`)
  1508. Args:
  1509. query (str): The query to search for
  1510. """
  1511. # TODO: add playwright to search the web
  1512. if engine == "searxng":
  1513. if request.app.state.config.SEARXNG_QUERY_URL:
  1514. return search_searxng(
  1515. request.app.state.config.SEARXNG_QUERY_URL,
  1516. query,
  1517. request.app.state.config.WEB_SEARCH_RESULT_COUNT,
  1518. request.app.state.config.WEB_SEARCH_DOMAIN_FILTER_LIST,
  1519. )
  1520. else:
  1521. raise Exception("No SEARXNG_QUERY_URL found in environment variables")
  1522. elif engine == "yacy":
  1523. if request.app.state.config.YACY_QUERY_URL:
  1524. return search_yacy(
  1525. request.app.state.config.YACY_QUERY_URL,
  1526. request.app.state.config.YACY_USERNAME,
  1527. request.app.state.config.YACY_PASSWORD,
  1528. query,
  1529. request.app.state.config.WEB_SEARCH_RESULT_COUNT,
  1530. request.app.state.config.WEB_SEARCH_DOMAIN_FILTER_LIST,
  1531. )
  1532. else:
  1533. raise Exception("No YACY_QUERY_URL found in environment variables")
  1534. elif engine == "google_pse":
  1535. if (
  1536. request.app.state.config.GOOGLE_PSE_API_KEY
  1537. and request.app.state.config.GOOGLE_PSE_ENGINE_ID
  1538. ):
  1539. return search_google_pse(
  1540. request.app.state.config.GOOGLE_PSE_API_KEY,
  1541. request.app.state.config.GOOGLE_PSE_ENGINE_ID,
  1542. query,
  1543. request.app.state.config.WEB_SEARCH_RESULT_COUNT,
  1544. request.app.state.config.WEB_SEARCH_DOMAIN_FILTER_LIST,
  1545. )
  1546. else:
  1547. raise Exception(
  1548. "No GOOGLE_PSE_API_KEY or GOOGLE_PSE_ENGINE_ID found in environment variables"
  1549. )
  1550. elif engine == "brave":
  1551. if request.app.state.config.BRAVE_SEARCH_API_KEY:
  1552. return search_brave(
  1553. request.app.state.config.BRAVE_SEARCH_API_KEY,
  1554. query,
  1555. request.app.state.config.WEB_SEARCH_RESULT_COUNT,
  1556. request.app.state.config.WEB_SEARCH_DOMAIN_FILTER_LIST,
  1557. )
  1558. else:
  1559. raise Exception("No BRAVE_SEARCH_API_KEY found in environment variables")
  1560. elif engine == "kagi":
  1561. if request.app.state.config.KAGI_SEARCH_API_KEY:
  1562. return search_kagi(
  1563. request.app.state.config.KAGI_SEARCH_API_KEY,
  1564. query,
  1565. request.app.state.config.WEB_SEARCH_RESULT_COUNT,
  1566. request.app.state.config.WEB_SEARCH_DOMAIN_FILTER_LIST,
  1567. )
  1568. else:
  1569. raise Exception("No KAGI_SEARCH_API_KEY found in environment variables")
  1570. elif engine == "mojeek":
  1571. if request.app.state.config.MOJEEK_SEARCH_API_KEY:
  1572. return search_mojeek(
  1573. request.app.state.config.MOJEEK_SEARCH_API_KEY,
  1574. query,
  1575. request.app.state.config.WEB_SEARCH_RESULT_COUNT,
  1576. request.app.state.config.WEB_SEARCH_DOMAIN_FILTER_LIST,
  1577. )
  1578. else:
  1579. raise Exception("No MOJEEK_SEARCH_API_KEY found in environment variables")
  1580. elif engine == "bocha":
  1581. if request.app.state.config.BOCHA_SEARCH_API_KEY:
  1582. return search_bocha(
  1583. request.app.state.config.BOCHA_SEARCH_API_KEY,
  1584. query,
  1585. request.app.state.config.WEB_SEARCH_RESULT_COUNT,
  1586. request.app.state.config.WEB_SEARCH_DOMAIN_FILTER_LIST,
  1587. )
  1588. else:
  1589. raise Exception("No BOCHA_SEARCH_API_KEY found in environment variables")
  1590. elif engine == "serpstack":
  1591. if request.app.state.config.SERPSTACK_API_KEY:
  1592. return search_serpstack(
  1593. request.app.state.config.SERPSTACK_API_KEY,
  1594. query,
  1595. request.app.state.config.WEB_SEARCH_RESULT_COUNT,
  1596. request.app.state.config.WEB_SEARCH_DOMAIN_FILTER_LIST,
  1597. https_enabled=request.app.state.config.SERPSTACK_HTTPS,
  1598. )
  1599. else:
  1600. raise Exception("No SERPSTACK_API_KEY found in environment variables")
  1601. elif engine == "serper":
  1602. if request.app.state.config.SERPER_API_KEY:
  1603. return search_serper(
  1604. request.app.state.config.SERPER_API_KEY,
  1605. query,
  1606. request.app.state.config.WEB_SEARCH_RESULT_COUNT,
  1607. request.app.state.config.WEB_SEARCH_DOMAIN_FILTER_LIST,
  1608. )
  1609. else:
  1610. raise Exception("No SERPER_API_KEY found in environment variables")
  1611. elif engine == "serply":
  1612. if request.app.state.config.SERPLY_API_KEY:
  1613. return search_serply(
  1614. request.app.state.config.SERPLY_API_KEY,
  1615. query,
  1616. request.app.state.config.WEB_SEARCH_RESULT_COUNT,
  1617. request.app.state.config.WEB_SEARCH_DOMAIN_FILTER_LIST,
  1618. )
  1619. else:
  1620. raise Exception("No SERPLY_API_KEY found in environment variables")
  1621. elif engine == "duckduckgo":
  1622. return search_duckduckgo(
  1623. query,
  1624. request.app.state.config.WEB_SEARCH_RESULT_COUNT,
  1625. request.app.state.config.WEB_SEARCH_DOMAIN_FILTER_LIST,
  1626. )
  1627. elif engine == "tavily":
  1628. if request.app.state.config.TAVILY_API_KEY:
  1629. return search_tavily(
  1630. request.app.state.config.TAVILY_API_KEY,
  1631. query,
  1632. request.app.state.config.WEB_SEARCH_RESULT_COUNT,
  1633. request.app.state.config.WEB_SEARCH_DOMAIN_FILTER_LIST,
  1634. )
  1635. else:
  1636. raise Exception("No TAVILY_API_KEY found in environment variables")
  1637. elif engine == "exa":
  1638. if request.app.state.config.EXA_API_KEY:
  1639. return search_exa(
  1640. request.app.state.config.EXA_API_KEY,
  1641. query,
  1642. request.app.state.config.WEB_SEARCH_RESULT_COUNT,
  1643. request.app.state.config.WEB_SEARCH_DOMAIN_FILTER_LIST,
  1644. )
  1645. else:
  1646. raise Exception("No EXA_API_KEY found in environment variables")
  1647. elif engine == "searchapi":
  1648. if request.app.state.config.SEARCHAPI_API_KEY:
  1649. return search_searchapi(
  1650. request.app.state.config.SEARCHAPI_API_KEY,
  1651. request.app.state.config.SEARCHAPI_ENGINE,
  1652. query,
  1653. request.app.state.config.WEB_SEARCH_RESULT_COUNT,
  1654. request.app.state.config.WEB_SEARCH_DOMAIN_FILTER_LIST,
  1655. )
  1656. else:
  1657. raise Exception("No SEARCHAPI_API_KEY found in environment variables")
  1658. elif engine == "serpapi":
  1659. if request.app.state.config.SERPAPI_API_KEY:
  1660. return search_serpapi(
  1661. request.app.state.config.SERPAPI_API_KEY,
  1662. request.app.state.config.SERPAPI_ENGINE,
  1663. query,
  1664. request.app.state.config.WEB_SEARCH_RESULT_COUNT,
  1665. request.app.state.config.WEB_SEARCH_DOMAIN_FILTER_LIST,
  1666. )
  1667. else:
  1668. raise Exception("No SERPAPI_API_KEY found in environment variables")
  1669. elif engine == "jina":
  1670. return search_jina(
  1671. request.app.state.config.JINA_API_KEY,
  1672. query,
  1673. request.app.state.config.WEB_SEARCH_RESULT_COUNT,
  1674. )
  1675. elif engine == "bing":
  1676. return search_bing(
  1677. request.app.state.config.BING_SEARCH_V7_SUBSCRIPTION_KEY,
  1678. request.app.state.config.BING_SEARCH_V7_ENDPOINT,
  1679. str(DEFAULT_LOCALE),
  1680. query,
  1681. request.app.state.config.WEB_SEARCH_RESULT_COUNT,
  1682. request.app.state.config.WEB_SEARCH_DOMAIN_FILTER_LIST,
  1683. )
  1684. elif engine == "exa":
  1685. return search_exa(
  1686. request.app.state.config.EXA_API_KEY,
  1687. query,
  1688. request.app.state.config.WEB_SEARCH_RESULT_COUNT,
  1689. request.app.state.config.WEB_SEARCH_DOMAIN_FILTER_LIST,
  1690. )
  1691. elif engine == "perplexity":
  1692. return search_perplexity(
  1693. request.app.state.config.PERPLEXITY_API_KEY,
  1694. query,
  1695. request.app.state.config.WEB_SEARCH_RESULT_COUNT,
  1696. request.app.state.config.WEB_SEARCH_DOMAIN_FILTER_LIST,
  1697. model=request.app.state.config.PERPLEXITY_MODEL,
  1698. search_context_usage=request.app.state.config.PERPLEXITY_SEARCH_CONTEXT_USAGE,
  1699. )
  1700. elif engine == "sougou":
  1701. if (
  1702. request.app.state.config.SOUGOU_API_SID
  1703. and request.app.state.config.SOUGOU_API_SK
  1704. ):
  1705. return search_sougou(
  1706. request.app.state.config.SOUGOU_API_SID,
  1707. request.app.state.config.SOUGOU_API_SK,
  1708. query,
  1709. request.app.state.config.WEB_SEARCH_RESULT_COUNT,
  1710. request.app.state.config.WEB_SEARCH_DOMAIN_FILTER_LIST,
  1711. )
  1712. else:
  1713. raise Exception(
  1714. "No SOUGOU_API_SID or SOUGOU_API_SK found in environment variables"
  1715. )
  1716. elif engine == "firecrawl":
  1717. return search_firecrawl(
  1718. request.app.state.config.FIRECRAWL_API_BASE_URL,
  1719. request.app.state.config.FIRECRAWL_API_KEY,
  1720. query,
  1721. request.app.state.config.WEB_SEARCH_RESULT_COUNT,
  1722. request.app.state.config.WEB_SEARCH_DOMAIN_FILTER_LIST,
  1723. )
  1724. elif engine == "external":
  1725. return search_external(
  1726. request.app.state.config.EXTERNAL_WEB_SEARCH_URL,
  1727. request.app.state.config.EXTERNAL_WEB_SEARCH_API_KEY,
  1728. query,
  1729. request.app.state.config.WEB_SEARCH_RESULT_COUNT,
  1730. request.app.state.config.WEB_SEARCH_DOMAIN_FILTER_LIST,
  1731. )
  1732. else:
  1733. raise Exception("No search engine API key found in environment variables")
  1734. @router.post("/process/web/search")
  1735. async def process_web_search(
  1736. request: Request, form_data: SearchForm, user=Depends(get_verified_user)
  1737. ):
  1738. urls = []
  1739. try:
  1740. logging.info(
  1741. f"trying to web search with {request.app.state.config.WEB_SEARCH_ENGINE, form_data.queries}"
  1742. )
  1743. search_tasks = [
  1744. run_in_threadpool(
  1745. search_web,
  1746. request,
  1747. request.app.state.config.WEB_SEARCH_ENGINE,
  1748. query,
  1749. )
  1750. for query in form_data.queries
  1751. ]
  1752. search_results = await asyncio.gather(*search_tasks)
  1753. for result in search_results:
  1754. if result:
  1755. for item in result:
  1756. if item and item.link:
  1757. urls.append(item.link)
  1758. urls = list(dict.fromkeys(urls))
  1759. log.debug(f"urls: {urls}")
  1760. except Exception as e:
  1761. log.exception(e)
  1762. raise HTTPException(
  1763. status_code=status.HTTP_400_BAD_REQUEST,
  1764. detail=ERROR_MESSAGES.WEB_SEARCH_ERROR(e),
  1765. )
  1766. try:
  1767. if request.app.state.config.BYPASS_WEB_SEARCH_WEB_LOADER:
  1768. search_results = [
  1769. item for result in search_results for item in result if result
  1770. ]
  1771. docs = [
  1772. Document(
  1773. page_content=result.snippet,
  1774. metadata={
  1775. "source": result.link,
  1776. "title": result.title,
  1777. "snippet": result.snippet,
  1778. "link": result.link,
  1779. },
  1780. )
  1781. for result in search_results
  1782. if hasattr(result, "snippet")
  1783. ]
  1784. else:
  1785. loader = get_web_loader(
  1786. urls,
  1787. verify_ssl=request.app.state.config.ENABLE_WEB_LOADER_SSL_VERIFICATION,
  1788. requests_per_second=request.app.state.config.WEB_SEARCH_CONCURRENT_REQUESTS,
  1789. trust_env=request.app.state.config.WEB_SEARCH_TRUST_ENV,
  1790. )
  1791. docs = await loader.aload()
  1792. urls = [
  1793. doc.metadata.get("source") for doc in docs if doc.metadata.get("source")
  1794. ] # only keep the urls returned by the loader
  1795. if request.app.state.config.BYPASS_WEB_SEARCH_EMBEDDING_AND_RETRIEVAL:
  1796. return {
  1797. "status": True,
  1798. "collection_name": None,
  1799. "filenames": urls,
  1800. "docs": [
  1801. {
  1802. "content": doc.page_content,
  1803. "metadata": doc.metadata,
  1804. }
  1805. for doc in docs
  1806. ],
  1807. "loaded_count": len(docs),
  1808. }
  1809. else:
  1810. # Create a single collection for all documents
  1811. collection_name = (
  1812. f"web-search-{calculate_sha256_string('-'.join(form_data.queries))}"[
  1813. :63
  1814. ]
  1815. )
  1816. try:
  1817. await run_in_threadpool(
  1818. save_docs_to_vector_db,
  1819. request,
  1820. docs,
  1821. collection_name,
  1822. overwrite=True,
  1823. user=user,
  1824. )
  1825. except Exception as e:
  1826. log.debug(f"error saving docs: {e}")
  1827. return {
  1828. "status": True,
  1829. "collection_names": [collection_name],
  1830. "filenames": urls,
  1831. "loaded_count": len(docs),
  1832. }
  1833. except Exception as e:
  1834. log.exception(e)
  1835. raise HTTPException(
  1836. status_code=status.HTTP_400_BAD_REQUEST,
  1837. detail=ERROR_MESSAGES.DEFAULT(e),
  1838. )
  1839. class QueryDocForm(BaseModel):
  1840. collection_name: str
  1841. query: str
  1842. k: Optional[int] = None
  1843. k_reranker: Optional[int] = None
  1844. r: Optional[float] = None
  1845. hybrid: Optional[bool] = None
  1846. @router.post("/query/doc")
  1847. def query_doc_handler(
  1848. request: Request,
  1849. form_data: QueryDocForm,
  1850. user=Depends(get_verified_user),
  1851. ):
  1852. try:
  1853. if request.app.state.config.ENABLE_RAG_HYBRID_SEARCH:
  1854. collection_results = {}
  1855. collection_results[form_data.collection_name] = VECTOR_DB_CLIENT.get(
  1856. collection_name=form_data.collection_name
  1857. )
  1858. return query_doc_with_hybrid_search(
  1859. collection_name=form_data.collection_name,
  1860. collection_result=collection_results[form_data.collection_name],
  1861. query=form_data.query,
  1862. embedding_function=lambda query, prefix: request.app.state.EMBEDDING_FUNCTION(
  1863. query, prefix=prefix, user=user
  1864. ),
  1865. k=form_data.k if form_data.k else request.app.state.config.TOP_K,
  1866. reranking_function=(
  1867. (
  1868. lambda sentences: request.app.state.RERANKING_FUNCTION(
  1869. sentences, user=user
  1870. )
  1871. )
  1872. if request.app.state.RERANKING_FUNCTION
  1873. else None
  1874. ),
  1875. k_reranker=form_data.k_reranker
  1876. or request.app.state.config.TOP_K_RERANKER,
  1877. r=(
  1878. form_data.r
  1879. if form_data.r
  1880. else request.app.state.config.RELEVANCE_THRESHOLD
  1881. ),
  1882. hybrid_bm25_weight=(
  1883. form_data.hybrid_bm25_weight
  1884. if form_data.hybrid_bm25_weight
  1885. else request.app.state.config.HYBRID_BM25_WEIGHT
  1886. ),
  1887. user=user,
  1888. )
  1889. else:
  1890. return query_doc(
  1891. collection_name=form_data.collection_name,
  1892. query_embedding=request.app.state.EMBEDDING_FUNCTION(
  1893. form_data.query, prefix=RAG_EMBEDDING_QUERY_PREFIX, user=user
  1894. ),
  1895. k=form_data.k if form_data.k else request.app.state.config.TOP_K,
  1896. user=user,
  1897. )
  1898. except Exception as e:
  1899. log.exception(e)
  1900. raise HTTPException(
  1901. status_code=status.HTTP_400_BAD_REQUEST,
  1902. detail=ERROR_MESSAGES.DEFAULT(e),
  1903. )
  1904. class QueryCollectionsForm(BaseModel):
  1905. collection_names: list[str]
  1906. query: str
  1907. k: Optional[int] = None
  1908. k_reranker: Optional[int] = None
  1909. r: Optional[float] = None
  1910. hybrid: Optional[bool] = None
  1911. hybrid_bm25_weight: Optional[float] = None
  1912. @router.post("/query/collection")
  1913. def query_collection_handler(
  1914. request: Request,
  1915. form_data: QueryCollectionsForm,
  1916. user=Depends(get_verified_user),
  1917. ):
  1918. try:
  1919. if request.app.state.config.ENABLE_RAG_HYBRID_SEARCH:
  1920. return query_collection_with_hybrid_search(
  1921. collection_names=form_data.collection_names,
  1922. queries=[form_data.query],
  1923. embedding_function=lambda query, prefix: request.app.state.EMBEDDING_FUNCTION(
  1924. query, prefix=prefix, user=user
  1925. ),
  1926. k=form_data.k if form_data.k else request.app.state.config.TOP_K,
  1927. reranking_function=(
  1928. (
  1929. lambda sentences: request.app.state.RERANKING_FUNCTION(
  1930. sentences, user=user
  1931. )
  1932. )
  1933. if request.app.state.RERANKING_FUNCTION
  1934. else None
  1935. ),
  1936. k_reranker=form_data.k_reranker
  1937. or request.app.state.config.TOP_K_RERANKER,
  1938. r=(
  1939. form_data.r
  1940. if form_data.r
  1941. else request.app.state.config.RELEVANCE_THRESHOLD
  1942. ),
  1943. hybrid_bm25_weight=(
  1944. form_data.hybrid_bm25_weight
  1945. if form_data.hybrid_bm25_weight
  1946. else request.app.state.config.HYBRID_BM25_WEIGHT
  1947. ),
  1948. )
  1949. else:
  1950. return query_collection(
  1951. collection_names=form_data.collection_names,
  1952. queries=[form_data.query],
  1953. embedding_function=lambda query, prefix: request.app.state.EMBEDDING_FUNCTION(
  1954. query, prefix=prefix, user=user
  1955. ),
  1956. k=form_data.k if form_data.k else request.app.state.config.TOP_K,
  1957. )
  1958. except Exception as e:
  1959. log.exception(e)
  1960. raise HTTPException(
  1961. status_code=status.HTTP_400_BAD_REQUEST,
  1962. detail=ERROR_MESSAGES.DEFAULT(e),
  1963. )
  1964. ####################################
  1965. #
  1966. # Vector DB operations
  1967. #
  1968. ####################################
  1969. class DeleteForm(BaseModel):
  1970. collection_name: str
  1971. file_id: str
  1972. @router.post("/delete")
  1973. def delete_entries_from_collection(form_data: DeleteForm, user=Depends(get_admin_user)):
  1974. try:
  1975. if VECTOR_DB_CLIENT.has_collection(collection_name=form_data.collection_name):
  1976. file = Files.get_file_by_id(form_data.file_id)
  1977. hash = file.hash
  1978. VECTOR_DB_CLIENT.delete(
  1979. collection_name=form_data.collection_name,
  1980. metadata={"hash": hash},
  1981. )
  1982. return {"status": True}
  1983. else:
  1984. return {"status": False}
  1985. except Exception as e:
  1986. log.exception(e)
  1987. return {"status": False}
  1988. @router.post("/reset/db")
  1989. def reset_vector_db(user=Depends(get_admin_user)):
  1990. VECTOR_DB_CLIENT.reset()
  1991. Knowledges.delete_all_knowledge()
  1992. @router.post("/reset/uploads")
  1993. def reset_upload_dir(user=Depends(get_admin_user)) -> bool:
  1994. folder = f"{UPLOAD_DIR}"
  1995. try:
  1996. # Check if the directory exists
  1997. if os.path.exists(folder):
  1998. # Iterate over all the files and directories in the specified directory
  1999. for filename in os.listdir(folder):
  2000. file_path = os.path.join(folder, filename)
  2001. try:
  2002. if os.path.isfile(file_path) or os.path.islink(file_path):
  2003. os.unlink(file_path) # Remove the file or link
  2004. elif os.path.isdir(file_path):
  2005. shutil.rmtree(file_path) # Remove the directory
  2006. except Exception as e:
  2007. log.exception(f"Failed to delete {file_path}. Reason: {e}")
  2008. else:
  2009. log.warning(f"The directory {folder} does not exist")
  2010. except Exception as e:
  2011. log.exception(f"Failed to process the directory {folder}. Reason: {e}")
  2012. return True
  2013. if ENV == "dev":
  2014. @router.get("/ef/{text}")
  2015. async def get_embeddings(request: Request, text: Optional[str] = "Hello World!"):
  2016. return {
  2017. "result": request.app.state.EMBEDDING_FUNCTION(
  2018. text, prefix=RAG_EMBEDDING_QUERY_PREFIX
  2019. )
  2020. }
  2021. class BatchProcessFilesForm(BaseModel):
  2022. files: List[FileModel]
  2023. collection_name: str
  2024. class BatchProcessFilesResult(BaseModel):
  2025. file_id: str
  2026. status: str
  2027. error: Optional[str] = None
  2028. class BatchProcessFilesResponse(BaseModel):
  2029. results: List[BatchProcessFilesResult]
  2030. errors: List[BatchProcessFilesResult]
  2031. @router.post("/process/files/batch")
  2032. def process_files_batch(
  2033. request: Request,
  2034. form_data: BatchProcessFilesForm,
  2035. user=Depends(get_verified_user),
  2036. ) -> BatchProcessFilesResponse:
  2037. """
  2038. Process a batch of files and save them to the vector database.
  2039. """
  2040. results: List[BatchProcessFilesResult] = []
  2041. errors: List[BatchProcessFilesResult] = []
  2042. collection_name = form_data.collection_name
  2043. # Prepare all documents first
  2044. all_docs: List[Document] = []
  2045. for file in form_data.files:
  2046. try:
  2047. text_content = file.data.get("content", "")
  2048. docs: List[Document] = [
  2049. Document(
  2050. page_content=text_content.replace("<br/>", "\n"),
  2051. metadata={
  2052. **file.meta,
  2053. "name": file.filename,
  2054. "created_by": file.user_id,
  2055. "file_id": file.id,
  2056. "source": file.filename,
  2057. },
  2058. )
  2059. ]
  2060. hash = calculate_sha256_string(text_content)
  2061. Files.update_file_hash_by_id(file.id, hash)
  2062. Files.update_file_data_by_id(file.id, {"content": text_content})
  2063. all_docs.extend(docs)
  2064. results.append(BatchProcessFilesResult(file_id=file.id, status="prepared"))
  2065. except Exception as e:
  2066. log.error(f"process_files_batch: Error processing file {file.id}: {str(e)}")
  2067. errors.append(
  2068. BatchProcessFilesResult(file_id=file.id, status="failed", error=str(e))
  2069. )
  2070. # Save all documents in one batch
  2071. if all_docs:
  2072. try:
  2073. save_docs_to_vector_db(
  2074. request=request,
  2075. docs=all_docs,
  2076. collection_name=collection_name,
  2077. add=True,
  2078. user=user,
  2079. )
  2080. # Update all files with collection name
  2081. for result in results:
  2082. Files.update_file_metadata_by_id(
  2083. result.file_id, {"collection_name": collection_name}
  2084. )
  2085. result.status = "completed"
  2086. except Exception as e:
  2087. log.error(
  2088. f"process_files_batch: Error saving documents to vector DB: {str(e)}"
  2089. )
  2090. for result in results:
  2091. result.status = "failed"
  2092. errors.append(
  2093. BatchProcessFilesResult(file_id=result.file_id, error=str(e))
  2094. )
  2095. return BatchProcessFilesResponse(results=results, errors=errors)