retrieval.py 79 KB

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