retrieval.py 102 KB

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