retrieval.py 96 KB

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