1
0

retrieval.py 102 KB

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