retrieval.py 105 KB

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