retrieval.py 99 KB

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