retrieval.py 90 KB

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