retrieval.py 73 KB

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