retrieval.py 53 KB

12345678910111213141516171819202122232425262728293031323334353637383940414243444546474849505152535455565758596061626364656667686970717273747576777879808182838485868788899091929394959697989910010110210310410510610710810911011111211311411511611711811912012112212312412512612712812913013113213313413513613713813914014114214314414514614714814915015115215315415515615715815916016116216316416516616716816917017117217317417517617717817918018118218318418518618718818919019119219319419519619719819920020120220320420520620720820921021121221321421521621721821922022122222322422522622722822923023123223323423523623723823924024124224324424524624724824925025125225325425525625725825926026126226326426526626726826927027127227327427527627727827928028128228328428528628728828929029129229329429529629729829930030130230330430530630730830931031131231331431531631731831932032132232332432532632732832933033133233333433533633733833934034134234334434534634734834935035135235335435535635735835936036136236336436536636736836937037137237337437537637737837938038138238338438538638738838939039139239339439539639739839940040140240340440540640740840941041141241341441541641741841942042142242342442542642742842943043143243343443543643743843944044144244344444544644744844945045145245345445545645745845946046146246346446546646746846947047147247347447547647747847948048148248348448548648748848949049149249349449549649749849950050150250350450550650750850951051151251351451551651751851952052152252352452552652752852953053153253353453553653753853954054154254354454554654754854955055155255355455555655755855956056156256356456556656756856957057157257357457557657757857958058158258358458558658758858959059159259359459559659759859960060160260360460560660760860961061161261361461561661761861962062162262362462562662762862963063163263363463563663763863964064164264364464564664764864965065165265365465565665765865966066166266366466566666766866967067167267367467567667767867968068168268368468568668768868969069169269369469569669769869970070170270370470570670770870971071171271371471571671771871972072172272372472572672772872973073173273373473573673773873974074174274374474574674774874975075175275375475575675775875976076176276376476576676776876977077177277377477577677777877978078178278378478578678778878979079179279379479579679779879980080180280380480580680780880981081181281381481581681781881982082182282382482582682782882983083183283383483583683783883984084184284384484584684784884985085185285385485585685785885986086186286386486586686786886987087187287387487587687787887988088188288388488588688788888989089189289389489589689789889990090190290390490590690790890991091191291391491591691791891992092192292392492592692792892993093193293393493593693793893994094194294394494594694794894995095195295395495595695795895996096196296396496596696796896997097197297397497597697797897998098198298398498598698798898999099199299399499599699799899910001001100210031004100510061007100810091010101110121013101410151016101710181019102010211022102310241025102610271028102910301031103210331034103510361037103810391040104110421043104410451046104710481049105010511052105310541055105610571058105910601061106210631064106510661067106810691070107110721073107410751076107710781079108010811082108310841085108610871088108910901091109210931094109510961097109810991100110111021103110411051106110711081109111011111112111311141115111611171118111911201121112211231124112511261127112811291130113111321133113411351136113711381139114011411142114311441145114611471148114911501151115211531154115511561157115811591160116111621163116411651166116711681169117011711172117311741175117611771178117911801181118211831184118511861187118811891190119111921193119411951196119711981199120012011202120312041205120612071208120912101211121212131214121512161217121812191220122112221223122412251226122712281229123012311232123312341235123612371238123912401241124212431244124512461247124812491250125112521253125412551256125712581259126012611262126312641265126612671268126912701271127212731274127512761277127812791280128112821283128412851286128712881289129012911292129312941295129612971298129913001301130213031304130513061307130813091310131113121313131413151316131713181319132013211322132313241325132613271328132913301331133213331334133513361337133813391340134113421343134413451346134713481349135013511352135313541355135613571358135913601361136213631364136513661367136813691370137113721373137413751376137713781379138013811382138313841385138613871388138913901391139213931394139513961397139813991400140114021403140414051406140714081409141014111412141314141415141614171418141914201421142214231424142514261427142814291430143114321433143414351436143714381439144014411442144314441445144614471448144914501451145214531454145514561457145814591460146114621463146414651466146714681469147014711472147314741475147614771478147914801481148214831484148514861487148814891490149114921493149414951496149714981499150015011502150315041505150615071508150915101511151215131514151515161517
  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 pydantic import BaseModel
  23. import tiktoken
  24. from langchain.text_splitter import RecursiveCharacterTextSplitter, TokenTextSplitter
  25. from langchain_core.documents import Document
  26. from open_webui.models.files import FileModel, Files
  27. from open_webui.models.knowledge import Knowledges
  28. from open_webui.storage.provider import Storage
  29. from open_webui.retrieval.vector.connector import VECTOR_DB_CLIENT
  30. # Document loaders
  31. from open_webui.retrieval.loaders.main import Loader
  32. from open_webui.retrieval.loaders.youtube import YoutubeLoader
  33. # Web search engines
  34. from open_webui.retrieval.web.main import SearchResult
  35. from open_webui.retrieval.web.utils import get_web_loader
  36. from open_webui.retrieval.web.brave import search_brave
  37. from open_webui.retrieval.web.kagi import search_kagi
  38. from open_webui.retrieval.web.mojeek import search_mojeek
  39. from open_webui.retrieval.web.duckduckgo import search_duckduckgo
  40. from open_webui.retrieval.web.google_pse import search_google_pse
  41. from open_webui.retrieval.web.jina_search import search_jina
  42. from open_webui.retrieval.web.searchapi import search_searchapi
  43. from open_webui.retrieval.web.searxng import search_searxng
  44. from open_webui.retrieval.web.serper import search_serper
  45. from open_webui.retrieval.web.serply import search_serply
  46. from open_webui.retrieval.web.serpstack import search_serpstack
  47. from open_webui.retrieval.web.tavily import search_tavily
  48. from open_webui.retrieval.web.bing import search_bing
  49. from open_webui.retrieval.utils import (
  50. get_embedding_function,
  51. get_model_path,
  52. query_collection,
  53. query_collection_with_hybrid_search,
  54. query_doc,
  55. query_doc_with_hybrid_search,
  56. )
  57. from open_webui.utils.misc import (
  58. calculate_sha256_string,
  59. )
  60. from open_webui.utils.auth import get_admin_user, get_verified_user
  61. from open_webui.config import (
  62. ENV,
  63. RAG_EMBEDDING_MODEL_AUTO_UPDATE,
  64. RAG_EMBEDDING_MODEL_TRUST_REMOTE_CODE,
  65. RAG_RERANKING_MODEL_AUTO_UPDATE,
  66. RAG_RERANKING_MODEL_TRUST_REMOTE_CODE,
  67. UPLOAD_DIR,
  68. DEFAULT_LOCALE,
  69. )
  70. from open_webui.env import (
  71. SRC_LOG_LEVELS,
  72. DEVICE_TYPE,
  73. DOCKER,
  74. )
  75. from open_webui.constants import ERROR_MESSAGES
  76. log = logging.getLogger(__name__)
  77. log.setLevel(SRC_LOG_LEVELS["RAG"])
  78. ##########################################
  79. #
  80. # Utility functions
  81. #
  82. ##########################################
  83. def get_ef(
  84. engine: str,
  85. embedding_model: str,
  86. auto_update: bool = False,
  87. ):
  88. ef = None
  89. if embedding_model and engine == "":
  90. from sentence_transformers import SentenceTransformer
  91. try:
  92. ef = SentenceTransformer(
  93. get_model_path(embedding_model, auto_update),
  94. device=DEVICE_TYPE,
  95. trust_remote_code=RAG_EMBEDDING_MODEL_TRUST_REMOTE_CODE,
  96. )
  97. except Exception as e:
  98. log.debug(f"Error loading SentenceTransformer: {e}")
  99. return ef
  100. def get_rf(
  101. reranking_model: str,
  102. auto_update: bool = False,
  103. ):
  104. rf = None
  105. if reranking_model:
  106. if any(model in reranking_model for model in ["jinaai/jina-colbert-v2"]):
  107. try:
  108. from open_webui.retrieval.models.colbert import ColBERT
  109. rf = ColBERT(
  110. get_model_path(reranking_model, auto_update),
  111. env="docker" if DOCKER else None,
  112. )
  113. except Exception as e:
  114. log.error(f"ColBERT: {e}")
  115. raise Exception(ERROR_MESSAGES.DEFAULT(e))
  116. else:
  117. import sentence_transformers
  118. try:
  119. rf = sentence_transformers.CrossEncoder(
  120. get_model_path(reranking_model, auto_update),
  121. device=DEVICE_TYPE,
  122. trust_remote_code=RAG_RERANKING_MODEL_TRUST_REMOTE_CODE,
  123. )
  124. except:
  125. log.error("CrossEncoder error")
  126. raise Exception(ERROR_MESSAGES.DEFAULT("CrossEncoder error"))
  127. return rf
  128. ##########################################
  129. #
  130. # API routes
  131. #
  132. ##########################################
  133. router = APIRouter()
  134. class CollectionNameForm(BaseModel):
  135. collection_name: Optional[str] = None
  136. class ProcessUrlForm(CollectionNameForm):
  137. url: str
  138. class SearchForm(CollectionNameForm):
  139. query: str
  140. @router.get("/")
  141. async def get_status(request: Request):
  142. return {
  143. "status": True,
  144. "chunk_size": request.app.state.config.CHUNK_SIZE,
  145. "chunk_overlap": request.app.state.config.CHUNK_OVERLAP,
  146. "template": request.app.state.config.RAG_TEMPLATE,
  147. "embedding_engine": request.app.state.config.RAG_EMBEDDING_ENGINE,
  148. "embedding_model": request.app.state.config.RAG_EMBEDDING_MODEL,
  149. "reranking_model": request.app.state.config.RAG_RERANKING_MODEL,
  150. "embedding_batch_size": request.app.state.config.RAG_EMBEDDING_BATCH_SIZE,
  151. }
  152. @router.get("/embedding")
  153. async def get_embedding_config(request: Request, user=Depends(get_admin_user)):
  154. return {
  155. "status": True,
  156. "embedding_engine": request.app.state.config.RAG_EMBEDDING_ENGINE,
  157. "embedding_model": request.app.state.config.RAG_EMBEDDING_MODEL,
  158. "embedding_batch_size": request.app.state.config.RAG_EMBEDDING_BATCH_SIZE,
  159. "openai_config": {
  160. "url": request.app.state.config.RAG_OPENAI_API_BASE_URL,
  161. "key": request.app.state.config.RAG_OPENAI_API_KEY,
  162. },
  163. "ollama_config": {
  164. "url": request.app.state.config.RAG_OLLAMA_BASE_URL,
  165. "key": request.app.state.config.RAG_OLLAMA_API_KEY,
  166. },
  167. }
  168. @router.get("/reranking")
  169. async def get_reraanking_config(request: Request, user=Depends(get_admin_user)):
  170. return {
  171. "status": True,
  172. "reranking_model": request.app.state.config.RAG_RERANKING_MODEL,
  173. }
  174. class OpenAIConfigForm(BaseModel):
  175. url: str
  176. key: str
  177. class OllamaConfigForm(BaseModel):
  178. url: str
  179. key: str
  180. class EmbeddingModelUpdateForm(BaseModel):
  181. openai_config: Optional[OpenAIConfigForm] = None
  182. ollama_config: Optional[OllamaConfigForm] = None
  183. embedding_engine: str
  184. embedding_model: str
  185. embedding_batch_size: Optional[int] = 1
  186. @router.post("/embedding/update")
  187. async def update_embedding_config(
  188. request: Request, form_data: EmbeddingModelUpdateForm, user=Depends(get_admin_user)
  189. ):
  190. log.info(
  191. f"Updating embedding model: {request.app.state.config.RAG_EMBEDDING_MODEL} to {form_data.embedding_model}"
  192. )
  193. try:
  194. request.app.state.config.RAG_EMBEDDING_ENGINE = form_data.embedding_engine
  195. request.app.state.config.RAG_EMBEDDING_MODEL = form_data.embedding_model
  196. if request.app.state.config.RAG_EMBEDDING_ENGINE in ["ollama", "openai"]:
  197. if form_data.openai_config is not None:
  198. request.app.state.config.RAG_OPENAI_API_BASE_URL = (
  199. form_data.openai_config.url
  200. )
  201. request.app.state.config.RAG_OPENAI_API_KEY = (
  202. form_data.openai_config.key
  203. )
  204. if form_data.ollama_config is not None:
  205. request.app.state.config.RAG_OLLAMA_BASE_URL = (
  206. form_data.ollama_config.url
  207. )
  208. request.app.state.config.RAG_OLLAMA_API_KEY = (
  209. form_data.ollama_config.key
  210. )
  211. request.app.state.config.RAG_EMBEDDING_BATCH_SIZE = (
  212. form_data.embedding_batch_size
  213. )
  214. request.app.state.ef = get_ef(
  215. request.app.state.config.RAG_EMBEDDING_ENGINE,
  216. request.app.state.config.RAG_EMBEDDING_MODEL,
  217. )
  218. request.app.state.EMBEDDING_FUNCTION = get_embedding_function(
  219. request.app.state.config.RAG_EMBEDDING_ENGINE,
  220. request.app.state.config.RAG_EMBEDDING_MODEL,
  221. request.app.state.ef,
  222. (
  223. request.app.state.config.RAG_OPENAI_API_BASE_URL
  224. if request.app.state.config.RAG_EMBEDDING_ENGINE == "openai"
  225. else request.app.state.config.RAG_OLLAMA_BASE_URL
  226. ),
  227. (
  228. request.app.state.config.RAG_OPENAI_API_KEY
  229. if request.app.state.config.RAG_EMBEDDING_ENGINE == "openai"
  230. else request.app.state.config.RAG_OLLAMA_API_KEY
  231. ),
  232. request.app.state.config.RAG_EMBEDDING_BATCH_SIZE,
  233. )
  234. return {
  235. "status": True,
  236. "embedding_engine": request.app.state.config.RAG_EMBEDDING_ENGINE,
  237. "embedding_model": request.app.state.config.RAG_EMBEDDING_MODEL,
  238. "embedding_batch_size": request.app.state.config.RAG_EMBEDDING_BATCH_SIZE,
  239. "openai_config": {
  240. "url": request.app.state.config.RAG_OPENAI_API_BASE_URL,
  241. "key": request.app.state.config.RAG_OPENAI_API_KEY,
  242. },
  243. "ollama_config": {
  244. "url": request.app.state.config.RAG_OLLAMA_BASE_URL,
  245. "key": request.app.state.config.RAG_OLLAMA_API_KEY,
  246. },
  247. }
  248. except Exception as e:
  249. log.exception(f"Problem updating embedding model: {e}")
  250. raise HTTPException(
  251. status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
  252. detail=ERROR_MESSAGES.DEFAULT(e),
  253. )
  254. class RerankingModelUpdateForm(BaseModel):
  255. reranking_model: str
  256. @router.post("/reranking/update")
  257. async def update_reranking_config(
  258. request: Request, form_data: RerankingModelUpdateForm, user=Depends(get_admin_user)
  259. ):
  260. log.info(
  261. f"Updating reranking model: {request.app.state.config.RAG_RERANKING_MODEL} to {form_data.reranking_model}"
  262. )
  263. try:
  264. request.app.state.config.RAG_RERANKING_MODEL = form_data.reranking_model
  265. try:
  266. request.app.state.rf = get_rf(
  267. request.app.state.config.RAG_RERANKING_MODEL,
  268. True,
  269. )
  270. except Exception as e:
  271. log.error(f"Error loading reranking model: {e}")
  272. request.app.state.config.ENABLE_RAG_HYBRID_SEARCH = False
  273. return {
  274. "status": True,
  275. "reranking_model": request.app.state.config.RAG_RERANKING_MODEL,
  276. }
  277. except Exception as e:
  278. log.exception(f"Problem updating reranking 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. "pdf_extract_images": request.app.state.config.PDF_EXTRACT_IMAGES,
  288. "content_extraction": {
  289. "engine": request.app.state.config.CONTENT_EXTRACTION_ENGINE,
  290. "tika_server_url": request.app.state.config.TIKA_SERVER_URL,
  291. },
  292. "chunk": {
  293. "text_splitter": request.app.state.config.TEXT_SPLITTER,
  294. "chunk_size": request.app.state.config.CHUNK_SIZE,
  295. "chunk_overlap": request.app.state.config.CHUNK_OVERLAP,
  296. },
  297. "file": {
  298. "max_size": request.app.state.config.FILE_MAX_SIZE,
  299. "max_count": request.app.state.config.FILE_MAX_COUNT,
  300. },
  301. "youtube": {
  302. "language": request.app.state.config.YOUTUBE_LOADER_LANGUAGE,
  303. "translation": request.app.state.YOUTUBE_LOADER_TRANSLATION,
  304. "proxy_url": request.app.state.config.YOUTUBE_LOADER_PROXY_URL,
  305. },
  306. "web": {
  307. "web_loader_ssl_verification": request.app.state.config.ENABLE_RAG_WEB_LOADER_SSL_VERIFICATION,
  308. "search": {
  309. "enabled": request.app.state.config.ENABLE_RAG_WEB_SEARCH,
  310. "engine": request.app.state.config.RAG_WEB_SEARCH_ENGINE,
  311. "searxng_query_url": request.app.state.config.SEARXNG_QUERY_URL,
  312. "google_pse_api_key": request.app.state.config.GOOGLE_PSE_API_KEY,
  313. "google_pse_engine_id": request.app.state.config.GOOGLE_PSE_ENGINE_ID,
  314. "brave_search_api_key": request.app.state.config.BRAVE_SEARCH_API_KEY,
  315. "kagi_search_api_key": request.app.state.config.KAGI_SEARCH_API_KEY,
  316. "mojeek_search_api_key": request.app.state.config.MOJEEK_SEARCH_API_KEY,
  317. "serpstack_api_key": request.app.state.config.SERPSTACK_API_KEY,
  318. "serpstack_https": request.app.state.config.SERPSTACK_HTTPS,
  319. "serper_api_key": request.app.state.config.SERPER_API_KEY,
  320. "serply_api_key": request.app.state.config.SERPLY_API_KEY,
  321. "tavily_api_key": request.app.state.config.TAVILY_API_KEY,
  322. "searchapi_api_key": request.app.state.config.SEARCHAPI_API_KEY,
  323. "seaarchapi_engine": request.app.state.config.SEARCHAPI_ENGINE,
  324. "jina_api_key": request.app.state.config.JINA_API_KEY,
  325. "bing_search_v7_endpoint": request.app.state.config.BING_SEARCH_V7_ENDPOINT,
  326. "bing_search_v7_subscription_key": request.app.state.config.BING_SEARCH_V7_SUBSCRIPTION_KEY,
  327. "result_count": request.app.state.config.RAG_WEB_SEARCH_RESULT_COUNT,
  328. "concurrent_requests": request.app.state.config.RAG_WEB_SEARCH_CONCURRENT_REQUESTS,
  329. },
  330. },
  331. }
  332. class FileConfig(BaseModel):
  333. max_size: Optional[int] = None
  334. max_count: Optional[int] = None
  335. class ContentExtractionConfig(BaseModel):
  336. engine: str = ""
  337. tika_server_url: Optional[str] = None
  338. class ChunkParamUpdateForm(BaseModel):
  339. text_splitter: Optional[str] = None
  340. chunk_size: int
  341. chunk_overlap: int
  342. class YoutubeLoaderConfig(BaseModel):
  343. language: list[str]
  344. translation: Optional[str] = None
  345. proxy_url: str = ""
  346. class WebSearchConfig(BaseModel):
  347. enabled: bool
  348. engine: Optional[str] = None
  349. searxng_query_url: Optional[str] = None
  350. google_pse_api_key: Optional[str] = None
  351. google_pse_engine_id: Optional[str] = None
  352. brave_search_api_key: Optional[str] = None
  353. kagi_search_api_key: Optional[str] = None
  354. mojeek_search_api_key: Optional[str] = None
  355. serpstack_api_key: Optional[str] = None
  356. serpstack_https: Optional[bool] = None
  357. serper_api_key: Optional[str] = None
  358. serply_api_key: Optional[str] = None
  359. tavily_api_key: Optional[str] = None
  360. searchapi_api_key: Optional[str] = None
  361. searchapi_engine: Optional[str] = None
  362. jina_api_key: Optional[str] = None
  363. bing_search_v7_endpoint: Optional[str] = None
  364. bing_search_v7_subscription_key: Optional[str] = None
  365. result_count: Optional[int] = None
  366. concurrent_requests: Optional[int] = None
  367. class WebConfig(BaseModel):
  368. search: WebSearchConfig
  369. web_loader_ssl_verification: Optional[bool] = None
  370. class ConfigUpdateForm(BaseModel):
  371. pdf_extract_images: Optional[bool] = None
  372. file: Optional[FileConfig] = None
  373. content_extraction: Optional[ContentExtractionConfig] = None
  374. chunk: Optional[ChunkParamUpdateForm] = None
  375. youtube: Optional[YoutubeLoaderConfig] = None
  376. web: Optional[WebConfig] = None
  377. @router.post("/config/update")
  378. async def update_rag_config(
  379. request: Request, form_data: ConfigUpdateForm, user=Depends(get_admin_user)
  380. ):
  381. request.app.state.config.PDF_EXTRACT_IMAGES = (
  382. form_data.pdf_extract_images
  383. if form_data.pdf_extract_images is not None
  384. else request.app.state.config.PDF_EXTRACT_IMAGES
  385. )
  386. if form_data.file is not None:
  387. request.app.state.config.FILE_MAX_SIZE = form_data.file.max_size
  388. request.app.state.config.FILE_MAX_COUNT = form_data.file.max_count
  389. if form_data.content_extraction is not None:
  390. log.info(f"Updating text settings: {form_data.content_extraction}")
  391. request.app.state.config.CONTENT_EXTRACTION_ENGINE = (
  392. form_data.content_extraction.engine
  393. )
  394. request.app.state.config.TIKA_SERVER_URL = (
  395. form_data.content_extraction.tika_server_url
  396. )
  397. if form_data.chunk is not None:
  398. request.app.state.config.TEXT_SPLITTER = form_data.chunk.text_splitter
  399. request.app.state.config.CHUNK_SIZE = form_data.chunk.chunk_size
  400. request.app.state.config.CHUNK_OVERLAP = form_data.chunk.chunk_overlap
  401. if form_data.youtube is not None:
  402. request.app.state.config.YOUTUBE_LOADER_LANGUAGE = form_data.youtube.language
  403. request.app.state.config.YOUTUBE_LOADER_PROXY_URL = form_data.youtube.proxy_url
  404. request.app.state.YOUTUBE_LOADER_TRANSLATION = form_data.youtube.translation
  405. if form_data.web is not None:
  406. request.app.state.config.ENABLE_RAG_WEB_LOADER_SSL_VERIFICATION = (
  407. # Note: When UI "Bypass SSL verification for Websites"=True then ENABLE_RAG_WEB_LOADER_SSL_VERIFICATION=False
  408. form_data.web.web_loader_ssl_verification
  409. )
  410. request.app.state.config.ENABLE_RAG_WEB_SEARCH = form_data.web.search.enabled
  411. request.app.state.config.RAG_WEB_SEARCH_ENGINE = form_data.web.search.engine
  412. request.app.state.config.SEARXNG_QUERY_URL = (
  413. form_data.web.search.searxng_query_url
  414. )
  415. request.app.state.config.GOOGLE_PSE_API_KEY = (
  416. form_data.web.search.google_pse_api_key
  417. )
  418. request.app.state.config.GOOGLE_PSE_ENGINE_ID = (
  419. form_data.web.search.google_pse_engine_id
  420. )
  421. request.app.state.config.BRAVE_SEARCH_API_KEY = (
  422. form_data.web.search.brave_search_api_key
  423. )
  424. request.app.state.config.KAGI_SEARCH_API_KEY = (
  425. form_data.web.search.kagi_search_api_key
  426. )
  427. request.app.state.config.MOJEEK_SEARCH_API_KEY = (
  428. form_data.web.search.mojeek_search_api_key
  429. )
  430. request.app.state.config.SERPSTACK_API_KEY = (
  431. form_data.web.search.serpstack_api_key
  432. )
  433. request.app.state.config.SERPSTACK_HTTPS = form_data.web.search.serpstack_https
  434. request.app.state.config.SERPER_API_KEY = form_data.web.search.serper_api_key
  435. request.app.state.config.SERPLY_API_KEY = form_data.web.search.serply_api_key
  436. request.app.state.config.TAVILY_API_KEY = form_data.web.search.tavily_api_key
  437. request.app.state.config.SEARCHAPI_API_KEY = (
  438. form_data.web.search.searchapi_api_key
  439. )
  440. request.app.state.config.SEARCHAPI_ENGINE = (
  441. form_data.web.search.searchapi_engine
  442. )
  443. request.app.state.config.JINA_API_KEY = form_data.web.search.jina_api_key
  444. request.app.state.config.BING_SEARCH_V7_ENDPOINT = (
  445. form_data.web.search.bing_search_v7_endpoint
  446. )
  447. request.app.state.config.BING_SEARCH_V7_SUBSCRIPTION_KEY = (
  448. form_data.web.search.bing_search_v7_subscription_key
  449. )
  450. request.app.state.config.RAG_WEB_SEARCH_RESULT_COUNT = (
  451. form_data.web.search.result_count
  452. )
  453. request.app.state.config.RAG_WEB_SEARCH_CONCURRENT_REQUESTS = (
  454. form_data.web.search.concurrent_requests
  455. )
  456. return {
  457. "status": True,
  458. "pdf_extract_images": request.app.state.config.PDF_EXTRACT_IMAGES,
  459. "file": {
  460. "max_size": request.app.state.config.FILE_MAX_SIZE,
  461. "max_count": request.app.state.config.FILE_MAX_COUNT,
  462. },
  463. "content_extraction": {
  464. "engine": request.app.state.config.CONTENT_EXTRACTION_ENGINE,
  465. "tika_server_url": request.app.state.config.TIKA_SERVER_URL,
  466. },
  467. "chunk": {
  468. "text_splitter": request.app.state.config.TEXT_SPLITTER,
  469. "chunk_size": request.app.state.config.CHUNK_SIZE,
  470. "chunk_overlap": request.app.state.config.CHUNK_OVERLAP,
  471. },
  472. "youtube": {
  473. "language": request.app.state.config.YOUTUBE_LOADER_LANGUAGE,
  474. "proxy_url": request.app.state.config.YOUTUBE_LOADER_PROXY_URL,
  475. "translation": request.app.state.YOUTUBE_LOADER_TRANSLATION,
  476. },
  477. "web": {
  478. "web_loader_ssl_verification": request.app.state.config.ENABLE_RAG_WEB_LOADER_SSL_VERIFICATION,
  479. "search": {
  480. "enabled": request.app.state.config.ENABLE_RAG_WEB_SEARCH,
  481. "engine": request.app.state.config.RAG_WEB_SEARCH_ENGINE,
  482. "searxng_query_url": request.app.state.config.SEARXNG_QUERY_URL,
  483. "google_pse_api_key": request.app.state.config.GOOGLE_PSE_API_KEY,
  484. "google_pse_engine_id": request.app.state.config.GOOGLE_PSE_ENGINE_ID,
  485. "brave_search_api_key": request.app.state.config.BRAVE_SEARCH_API_KEY,
  486. "kagi_search_api_key": request.app.state.config.KAGI_SEARCH_API_KEY,
  487. "mojeek_search_api_key": request.app.state.config.MOJEEK_SEARCH_API_KEY,
  488. "serpstack_api_key": request.app.state.config.SERPSTACK_API_KEY,
  489. "serpstack_https": request.app.state.config.SERPSTACK_HTTPS,
  490. "serper_api_key": request.app.state.config.SERPER_API_KEY,
  491. "serply_api_key": request.app.state.config.SERPLY_API_KEY,
  492. "serachapi_api_key": request.app.state.config.SEARCHAPI_API_KEY,
  493. "searchapi_engine": request.app.state.config.SEARCHAPI_ENGINE,
  494. "tavily_api_key": request.app.state.config.TAVILY_API_KEY,
  495. "jina_api_key": request.app.state.config.JINA_API_KEY,
  496. "bing_search_v7_endpoint": request.app.state.config.BING_SEARCH_V7_ENDPOINT,
  497. "bing_search_v7_subscription_key": request.app.state.config.BING_SEARCH_V7_SUBSCRIPTION_KEY,
  498. "result_count": request.app.state.config.RAG_WEB_SEARCH_RESULT_COUNT,
  499. "concurrent_requests": request.app.state.config.RAG_WEB_SEARCH_CONCURRENT_REQUESTS,
  500. },
  501. },
  502. }
  503. @router.get("/template")
  504. async def get_rag_template(request: Request, user=Depends(get_verified_user)):
  505. return {
  506. "status": True,
  507. "template": request.app.state.config.RAG_TEMPLATE,
  508. }
  509. @router.get("/query/settings")
  510. async def get_query_settings(request: Request, user=Depends(get_admin_user)):
  511. return {
  512. "status": True,
  513. "template": request.app.state.config.RAG_TEMPLATE,
  514. "k": request.app.state.config.TOP_K,
  515. "r": request.app.state.config.RELEVANCE_THRESHOLD,
  516. "hybrid": request.app.state.config.ENABLE_RAG_HYBRID_SEARCH,
  517. }
  518. class QuerySettingsForm(BaseModel):
  519. k: Optional[int] = None
  520. r: Optional[float] = None
  521. template: Optional[str] = None
  522. hybrid: Optional[bool] = None
  523. @router.post("/query/settings/update")
  524. async def update_query_settings(
  525. request: Request, form_data: QuerySettingsForm, user=Depends(get_admin_user)
  526. ):
  527. request.app.state.config.RAG_TEMPLATE = form_data.template
  528. request.app.state.config.TOP_K = form_data.k if form_data.k else 4
  529. request.app.state.config.RELEVANCE_THRESHOLD = form_data.r if form_data.r else 0.0
  530. request.app.state.config.ENABLE_RAG_HYBRID_SEARCH = (
  531. form_data.hybrid if form_data.hybrid else False
  532. )
  533. return {
  534. "status": True,
  535. "template": request.app.state.config.RAG_TEMPLATE,
  536. "k": request.app.state.config.TOP_K,
  537. "r": request.app.state.config.RELEVANCE_THRESHOLD,
  538. "hybrid": request.app.state.config.ENABLE_RAG_HYBRID_SEARCH,
  539. }
  540. ####################################
  541. #
  542. # Document process and retrieval
  543. #
  544. ####################################
  545. def save_docs_to_vector_db(
  546. request: Request,
  547. docs,
  548. collection_name,
  549. metadata: Optional[dict] = None,
  550. overwrite: bool = False,
  551. split: bool = True,
  552. add: bool = False,
  553. ) -> bool:
  554. def _get_docs_info(docs: list[Document]) -> str:
  555. docs_info = set()
  556. # Trying to select relevant metadata identifying the document.
  557. for doc in docs:
  558. metadata = getattr(doc, "metadata", {})
  559. doc_name = metadata.get("name", "")
  560. if not doc_name:
  561. doc_name = metadata.get("title", "")
  562. if not doc_name:
  563. doc_name = metadata.get("source", "")
  564. if doc_name:
  565. docs_info.add(doc_name)
  566. return ", ".join(docs_info)
  567. log.info(
  568. f"save_docs_to_vector_db: document {_get_docs_info(docs)} {collection_name}"
  569. )
  570. # Check if entries with the same hash (metadata.hash) already exist
  571. if metadata and "hash" in metadata:
  572. result = VECTOR_DB_CLIENT.query(
  573. collection_name=collection_name,
  574. filter={"hash": metadata["hash"]},
  575. )
  576. if result is not None:
  577. existing_doc_ids = result.ids[0]
  578. if existing_doc_ids:
  579. log.info(f"Document with hash {metadata['hash']} already exists")
  580. raise ValueError(ERROR_MESSAGES.DUPLICATE_CONTENT)
  581. if split:
  582. if request.app.state.config.TEXT_SPLITTER in ["", "character"]:
  583. text_splitter = RecursiveCharacterTextSplitter(
  584. chunk_size=request.app.state.config.CHUNK_SIZE,
  585. chunk_overlap=request.app.state.config.CHUNK_OVERLAP,
  586. add_start_index=True,
  587. )
  588. elif request.app.state.config.TEXT_SPLITTER == "token":
  589. log.info(
  590. f"Using token text splitter: {request.app.state.config.TIKTOKEN_ENCODING_NAME}"
  591. )
  592. tiktoken.get_encoding(str(request.app.state.config.TIKTOKEN_ENCODING_NAME))
  593. text_splitter = TokenTextSplitter(
  594. encoding_name=str(request.app.state.config.TIKTOKEN_ENCODING_NAME),
  595. chunk_size=request.app.state.config.CHUNK_SIZE,
  596. chunk_overlap=request.app.state.config.CHUNK_OVERLAP,
  597. add_start_index=True,
  598. )
  599. else:
  600. raise ValueError(ERROR_MESSAGES.DEFAULT("Invalid text splitter"))
  601. docs = text_splitter.split_documents(docs)
  602. if len(docs) == 0:
  603. raise ValueError(ERROR_MESSAGES.EMPTY_CONTENT)
  604. texts = [doc.page_content for doc in docs]
  605. metadatas = [
  606. {
  607. **doc.metadata,
  608. **(metadata if metadata else {}),
  609. "embedding_config": json.dumps(
  610. {
  611. "engine": request.app.state.config.RAG_EMBEDDING_ENGINE,
  612. "model": request.app.state.config.RAG_EMBEDDING_MODEL,
  613. }
  614. ),
  615. }
  616. for doc in docs
  617. ]
  618. # ChromaDB does not like datetime formats
  619. # for meta-data so convert them to string.
  620. for metadata in metadatas:
  621. for key, value in metadata.items():
  622. if isinstance(value, datetime):
  623. metadata[key] = str(value)
  624. try:
  625. if VECTOR_DB_CLIENT.has_collection(collection_name=collection_name):
  626. log.info(f"collection {collection_name} already exists")
  627. if overwrite:
  628. VECTOR_DB_CLIENT.delete_collection(collection_name=collection_name)
  629. log.info(f"deleting existing collection {collection_name}")
  630. elif add is False:
  631. log.info(
  632. f"collection {collection_name} already exists, overwrite is False and add is False"
  633. )
  634. return True
  635. log.info(f"adding to collection {collection_name}")
  636. embedding_function = get_embedding_function(
  637. request.app.state.config.RAG_EMBEDDING_ENGINE,
  638. request.app.state.config.RAG_EMBEDDING_MODEL,
  639. request.app.state.ef,
  640. (
  641. request.app.state.config.RAG_OPENAI_API_BASE_URL
  642. if request.app.state.config.RAG_EMBEDDING_ENGINE == "openai"
  643. else request.app.state.config.RAG_OLLAMA_BASE_URL
  644. ),
  645. (
  646. request.app.state.config.RAG_OPENAI_API_KEY
  647. if request.app.state.config.RAG_EMBEDDING_ENGINE == "openai"
  648. else request.app.state.config.RAG_OLLAMA_API_KEY
  649. ),
  650. request.app.state.config.RAG_EMBEDDING_BATCH_SIZE,
  651. )
  652. embeddings = embedding_function(
  653. list(map(lambda x: x.replace("\n", " "), texts))
  654. )
  655. items = [
  656. {
  657. "id": str(uuid.uuid4()),
  658. "text": text,
  659. "vector": embeddings[idx],
  660. "metadata": metadatas[idx],
  661. }
  662. for idx, text in enumerate(texts)
  663. ]
  664. VECTOR_DB_CLIENT.insert(
  665. collection_name=collection_name,
  666. items=items,
  667. )
  668. return True
  669. except Exception as e:
  670. log.exception(e)
  671. raise e
  672. class ProcessFileForm(BaseModel):
  673. file_id: str
  674. content: Optional[str] = None
  675. collection_name: Optional[str] = None
  676. @router.post("/process/file")
  677. def process_file(
  678. request: Request,
  679. form_data: ProcessFileForm,
  680. user=Depends(get_verified_user),
  681. ):
  682. try:
  683. file = Files.get_file_by_id(form_data.file_id)
  684. collection_name = form_data.collection_name
  685. if collection_name is None:
  686. collection_name = f"file-{file.id}"
  687. if form_data.content:
  688. # Update the content in the file
  689. # Usage: /files/{file_id}/data/content/update
  690. VECTOR_DB_CLIENT.delete_collection(collection_name=f"file-{file.id}")
  691. docs = [
  692. Document(
  693. page_content=form_data.content.replace("<br/>", "\n"),
  694. metadata={
  695. **file.meta,
  696. "name": file.filename,
  697. "created_by": file.user_id,
  698. "file_id": file.id,
  699. "source": file.filename,
  700. },
  701. )
  702. ]
  703. text_content = form_data.content
  704. elif form_data.collection_name:
  705. # Check if the file has already been processed and save the content
  706. # Usage: /knowledge/{id}/file/add, /knowledge/{id}/file/update
  707. result = VECTOR_DB_CLIENT.query(
  708. collection_name=f"file-{file.id}", filter={"file_id": file.id}
  709. )
  710. if result is not None and len(result.ids[0]) > 0:
  711. docs = [
  712. Document(
  713. page_content=result.documents[0][idx],
  714. metadata=result.metadatas[0][idx],
  715. )
  716. for idx, id in enumerate(result.ids[0])
  717. ]
  718. else:
  719. docs = [
  720. Document(
  721. page_content=file.data.get("content", ""),
  722. metadata={
  723. **file.meta,
  724. "name": file.filename,
  725. "created_by": file.user_id,
  726. "file_id": file.id,
  727. "source": file.filename,
  728. },
  729. )
  730. ]
  731. text_content = file.data.get("content", "")
  732. else:
  733. # Process the file and save the content
  734. # Usage: /files/
  735. file_path = file.path
  736. if file_path:
  737. file_path = Storage.get_file(file_path)
  738. loader = Loader(
  739. engine=request.app.state.config.CONTENT_EXTRACTION_ENGINE,
  740. TIKA_SERVER_URL=request.app.state.config.TIKA_SERVER_URL,
  741. PDF_EXTRACT_IMAGES=request.app.state.config.PDF_EXTRACT_IMAGES,
  742. )
  743. docs = loader.load(
  744. file.filename, file.meta.get("content_type"), file_path
  745. )
  746. docs = [
  747. Document(
  748. page_content=doc.page_content,
  749. metadata={
  750. **doc.metadata,
  751. "name": file.filename,
  752. "created_by": file.user_id,
  753. "file_id": file.id,
  754. "source": file.filename,
  755. },
  756. )
  757. for doc in docs
  758. ]
  759. else:
  760. docs = [
  761. Document(
  762. page_content=file.data.get("content", ""),
  763. metadata={
  764. **file.meta,
  765. "name": file.filename,
  766. "created_by": file.user_id,
  767. "file_id": file.id,
  768. "source": file.filename,
  769. },
  770. )
  771. ]
  772. text_content = " ".join([doc.page_content for doc in docs])
  773. log.debug(f"text_content: {text_content}")
  774. Files.update_file_data_by_id(
  775. file.id,
  776. {"content": text_content},
  777. )
  778. hash = calculate_sha256_string(text_content)
  779. Files.update_file_hash_by_id(file.id, hash)
  780. try:
  781. result = save_docs_to_vector_db(
  782. request,
  783. docs=docs,
  784. collection_name=collection_name,
  785. metadata={
  786. "file_id": file.id,
  787. "name": file.filename,
  788. "hash": hash,
  789. },
  790. add=(True if form_data.collection_name else False),
  791. )
  792. if result:
  793. Files.update_file_metadata_by_id(
  794. file.id,
  795. {
  796. "collection_name": collection_name,
  797. },
  798. )
  799. return {
  800. "status": True,
  801. "collection_name": collection_name,
  802. "filename": file.filename,
  803. "content": text_content,
  804. }
  805. except Exception as e:
  806. raise e
  807. except Exception as e:
  808. log.exception(e)
  809. if "No pandoc was found" in str(e):
  810. raise HTTPException(
  811. status_code=status.HTTP_400_BAD_REQUEST,
  812. detail=ERROR_MESSAGES.PANDOC_NOT_INSTALLED,
  813. )
  814. else:
  815. raise HTTPException(
  816. status_code=status.HTTP_400_BAD_REQUEST,
  817. detail=str(e),
  818. )
  819. class ProcessTextForm(BaseModel):
  820. name: str
  821. content: str
  822. collection_name: Optional[str] = None
  823. @router.post("/process/text")
  824. def process_text(
  825. request: Request,
  826. form_data: ProcessTextForm,
  827. user=Depends(get_verified_user),
  828. ):
  829. collection_name = form_data.collection_name
  830. if collection_name is None:
  831. collection_name = calculate_sha256_string(form_data.content)
  832. docs = [
  833. Document(
  834. page_content=form_data.content,
  835. metadata={"name": form_data.name, "created_by": user.id},
  836. )
  837. ]
  838. text_content = form_data.content
  839. log.debug(f"text_content: {text_content}")
  840. result = save_docs_to_vector_db(request, docs, collection_name)
  841. if result:
  842. return {
  843. "status": True,
  844. "collection_name": collection_name,
  845. "content": text_content,
  846. }
  847. else:
  848. raise HTTPException(
  849. status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
  850. detail=ERROR_MESSAGES.DEFAULT(),
  851. )
  852. @router.post("/process/youtube")
  853. def process_youtube_video(
  854. request: Request, form_data: ProcessUrlForm, user=Depends(get_verified_user)
  855. ):
  856. try:
  857. collection_name = form_data.collection_name
  858. if not collection_name:
  859. collection_name = calculate_sha256_string(form_data.url)[:63]
  860. loader = YoutubeLoader(
  861. form_data.url,
  862. language=request.app.state.config.YOUTUBE_LOADER_LANGUAGE,
  863. proxy_url=request.app.state.config.YOUTUBE_LOADER_PROXY_URL,
  864. )
  865. docs = loader.load()
  866. content = " ".join([doc.page_content for doc in docs])
  867. log.debug(f"text_content: {content}")
  868. save_docs_to_vector_db(request, docs, collection_name, overwrite=True)
  869. return {
  870. "status": True,
  871. "collection_name": collection_name,
  872. "filename": form_data.url,
  873. "file": {
  874. "data": {
  875. "content": content,
  876. },
  877. "meta": {
  878. "name": form_data.url,
  879. },
  880. },
  881. }
  882. except Exception as e:
  883. log.exception(e)
  884. raise HTTPException(
  885. status_code=status.HTTP_400_BAD_REQUEST,
  886. detail=ERROR_MESSAGES.DEFAULT(e),
  887. )
  888. @router.post("/process/web")
  889. def process_web(
  890. request: Request, form_data: ProcessUrlForm, user=Depends(get_verified_user)
  891. ):
  892. try:
  893. collection_name = form_data.collection_name
  894. if not collection_name:
  895. collection_name = calculate_sha256_string(form_data.url)[:63]
  896. loader = get_web_loader(
  897. form_data.url,
  898. verify_ssl=request.app.state.config.ENABLE_RAG_WEB_LOADER_SSL_VERIFICATION,
  899. requests_per_second=request.app.state.config.RAG_WEB_SEARCH_CONCURRENT_REQUESTS,
  900. )
  901. docs = loader.load()
  902. content = " ".join([doc.page_content for doc in docs])
  903. log.debug(f"text_content: {content}")
  904. save_docs_to_vector_db(request, docs, collection_name, overwrite=True)
  905. return {
  906. "status": True,
  907. "collection_name": collection_name,
  908. "filename": form_data.url,
  909. "file": {
  910. "data": {
  911. "content": content,
  912. },
  913. "meta": {
  914. "name": form_data.url,
  915. },
  916. },
  917. }
  918. except Exception as e:
  919. log.exception(e)
  920. raise HTTPException(
  921. status_code=status.HTTP_400_BAD_REQUEST,
  922. detail=ERROR_MESSAGES.DEFAULT(e),
  923. )
  924. def search_web(request: Request, engine: str, query: str) -> list[SearchResult]:
  925. """Search the web using a search engine and return the results as a list of SearchResult objects.
  926. Will look for a search engine API key in environment variables in the following order:
  927. - SEARXNG_QUERY_URL
  928. - GOOGLE_PSE_API_KEY + GOOGLE_PSE_ENGINE_ID
  929. - BRAVE_SEARCH_API_KEY
  930. - KAGI_SEARCH_API_KEY
  931. - MOJEEK_SEARCH_API_KEY
  932. - SERPSTACK_API_KEY
  933. - SERPER_API_KEY
  934. - SERPLY_API_KEY
  935. - TAVILY_API_KEY
  936. - SEARCHAPI_API_KEY + SEARCHAPI_ENGINE (by default `google`)
  937. Args:
  938. query (str): The query to search for
  939. """
  940. # TODO: add playwright to search the web
  941. if engine == "searxng":
  942. if request.app.state.config.SEARXNG_QUERY_URL:
  943. return search_searxng(
  944. request.app.state.config.SEARXNG_QUERY_URL,
  945. query,
  946. request.app.state.config.RAG_WEB_SEARCH_RESULT_COUNT,
  947. request.app.state.config.RAG_WEB_SEARCH_DOMAIN_FILTER_LIST,
  948. )
  949. else:
  950. raise Exception("No SEARXNG_QUERY_URL found in environment variables")
  951. elif engine == "google_pse":
  952. if (
  953. request.app.state.config.GOOGLE_PSE_API_KEY
  954. and request.app.state.config.GOOGLE_PSE_ENGINE_ID
  955. ):
  956. return search_google_pse(
  957. request.app.state.config.GOOGLE_PSE_API_KEY,
  958. request.app.state.config.GOOGLE_PSE_ENGINE_ID,
  959. query,
  960. request.app.state.config.RAG_WEB_SEARCH_RESULT_COUNT,
  961. request.app.state.config.RAG_WEB_SEARCH_DOMAIN_FILTER_LIST,
  962. )
  963. else:
  964. raise Exception(
  965. "No GOOGLE_PSE_API_KEY or GOOGLE_PSE_ENGINE_ID found in environment variables"
  966. )
  967. elif engine == "brave":
  968. if request.app.state.config.BRAVE_SEARCH_API_KEY:
  969. return search_brave(
  970. request.app.state.config.BRAVE_SEARCH_API_KEY,
  971. query,
  972. request.app.state.config.RAG_WEB_SEARCH_RESULT_COUNT,
  973. request.app.state.config.RAG_WEB_SEARCH_DOMAIN_FILTER_LIST,
  974. )
  975. else:
  976. raise Exception("No BRAVE_SEARCH_API_KEY found in environment variables")
  977. elif engine == "kagi":
  978. if request.app.state.config.KAGI_SEARCH_API_KEY:
  979. return search_kagi(
  980. request.app.state.config.KAGI_SEARCH_API_KEY,
  981. query,
  982. request.app.state.config.RAG_WEB_SEARCH_RESULT_COUNT,
  983. request.app.state.config.RAG_WEB_SEARCH_DOMAIN_FILTER_LIST,
  984. )
  985. else:
  986. raise Exception("No KAGI_SEARCH_API_KEY found in environment variables")
  987. elif engine == "mojeek":
  988. if request.app.state.config.MOJEEK_SEARCH_API_KEY:
  989. return search_mojeek(
  990. request.app.state.config.MOJEEK_SEARCH_API_KEY,
  991. query,
  992. request.app.state.config.RAG_WEB_SEARCH_RESULT_COUNT,
  993. request.app.state.config.RAG_WEB_SEARCH_DOMAIN_FILTER_LIST,
  994. )
  995. else:
  996. raise Exception("No MOJEEK_SEARCH_API_KEY found in environment variables")
  997. elif engine == "serpstack":
  998. if request.app.state.config.SERPSTACK_API_KEY:
  999. return search_serpstack(
  1000. request.app.state.config.SERPSTACK_API_KEY,
  1001. query,
  1002. request.app.state.config.RAG_WEB_SEARCH_RESULT_COUNT,
  1003. request.app.state.config.RAG_WEB_SEARCH_DOMAIN_FILTER_LIST,
  1004. https_enabled=request.app.state.config.SERPSTACK_HTTPS,
  1005. )
  1006. else:
  1007. raise Exception("No SERPSTACK_API_KEY found in environment variables")
  1008. elif engine == "serper":
  1009. if request.app.state.config.SERPER_API_KEY:
  1010. return search_serper(
  1011. request.app.state.config.SERPER_API_KEY,
  1012. query,
  1013. request.app.state.config.RAG_WEB_SEARCH_RESULT_COUNT,
  1014. request.app.state.config.RAG_WEB_SEARCH_DOMAIN_FILTER_LIST,
  1015. )
  1016. else:
  1017. raise Exception("No SERPER_API_KEY found in environment variables")
  1018. elif engine == "serply":
  1019. if request.app.state.config.SERPLY_API_KEY:
  1020. return search_serply(
  1021. request.app.state.config.SERPLY_API_KEY,
  1022. query,
  1023. request.app.state.config.RAG_WEB_SEARCH_RESULT_COUNT,
  1024. request.app.state.config.RAG_WEB_SEARCH_DOMAIN_FILTER_LIST,
  1025. )
  1026. else:
  1027. raise Exception("No SERPLY_API_KEY found in environment variables")
  1028. elif engine == "duckduckgo":
  1029. return search_duckduckgo(
  1030. query,
  1031. request.app.state.config.RAG_WEB_SEARCH_RESULT_COUNT,
  1032. request.app.state.config.RAG_WEB_SEARCH_DOMAIN_FILTER_LIST,
  1033. )
  1034. elif engine == "tavily":
  1035. if request.app.state.config.TAVILY_API_KEY:
  1036. return search_tavily(
  1037. request.app.state.config.TAVILY_API_KEY,
  1038. query,
  1039. request.app.state.config.RAG_WEB_SEARCH_RESULT_COUNT,
  1040. )
  1041. else:
  1042. raise Exception("No TAVILY_API_KEY found in environment variables")
  1043. elif engine == "searchapi":
  1044. if request.app.state.config.SEARCHAPI_API_KEY:
  1045. return search_searchapi(
  1046. request.app.state.config.SEARCHAPI_API_KEY,
  1047. request.app.state.config.SEARCHAPI_ENGINE,
  1048. query,
  1049. request.app.state.config.RAG_WEB_SEARCH_RESULT_COUNT,
  1050. request.app.state.config.RAG_WEB_SEARCH_DOMAIN_FILTER_LIST,
  1051. )
  1052. else:
  1053. raise Exception("No SEARCHAPI_API_KEY found in environment variables")
  1054. elif engine == "jina":
  1055. return search_jina(
  1056. request.app.state.config.JINA_API_KEY,
  1057. query,
  1058. request.app.state.config.RAG_WEB_SEARCH_RESULT_COUNT,
  1059. )
  1060. elif engine == "bing":
  1061. return search_bing(
  1062. request.app.state.config.BING_SEARCH_V7_SUBSCRIPTION_KEY,
  1063. request.app.state.config.BING_SEARCH_V7_ENDPOINT,
  1064. str(DEFAULT_LOCALE),
  1065. query,
  1066. request.app.state.config.RAG_WEB_SEARCH_RESULT_COUNT,
  1067. request.app.state.config.RAG_WEB_SEARCH_DOMAIN_FILTER_LIST,
  1068. )
  1069. else:
  1070. raise Exception("No search engine API key found in environment variables")
  1071. @router.post("/process/web/search")
  1072. def process_web_search(
  1073. request: Request, form_data: SearchForm, user=Depends(get_verified_user)
  1074. ):
  1075. try:
  1076. logging.info(
  1077. f"trying to web search with {request.app.state.config.RAG_WEB_SEARCH_ENGINE, form_data.query}"
  1078. )
  1079. web_results = search_web(
  1080. request, request.app.state.config.RAG_WEB_SEARCH_ENGINE, form_data.query
  1081. )
  1082. except Exception as e:
  1083. log.exception(e)
  1084. raise HTTPException(
  1085. status_code=status.HTTP_400_BAD_REQUEST,
  1086. detail=ERROR_MESSAGES.WEB_SEARCH_ERROR(e),
  1087. )
  1088. try:
  1089. collection_name = form_data.collection_name
  1090. if collection_name == "":
  1091. collection_name = f"web-search-{calculate_sha256_string(form_data.query)}"[
  1092. :63
  1093. ]
  1094. urls = [result.link for result in web_results]
  1095. loader = get_web_loader(
  1096. urls,
  1097. verify_ssl=request.app.state.config.ENABLE_RAG_WEB_LOADER_SSL_VERIFICATION,
  1098. requests_per_second=request.app.state.config.RAG_WEB_SEARCH_CONCURRENT_REQUESTS,
  1099. )
  1100. docs = loader.aload()
  1101. save_docs_to_vector_db(request, docs, collection_name, overwrite=True)
  1102. return {
  1103. "status": True,
  1104. "collection_name": collection_name,
  1105. "filenames": urls,
  1106. }
  1107. except Exception as e:
  1108. log.exception(e)
  1109. raise HTTPException(
  1110. status_code=status.HTTP_400_BAD_REQUEST,
  1111. detail=ERROR_MESSAGES.DEFAULT(e),
  1112. )
  1113. class QueryDocForm(BaseModel):
  1114. collection_name: str
  1115. query: str
  1116. k: Optional[int] = None
  1117. r: Optional[float] = None
  1118. hybrid: Optional[bool] = None
  1119. @router.post("/query/doc")
  1120. def query_doc_handler(
  1121. request: Request,
  1122. form_data: QueryDocForm,
  1123. user=Depends(get_verified_user),
  1124. ):
  1125. try:
  1126. if request.app.state.config.ENABLE_RAG_HYBRID_SEARCH:
  1127. return query_doc_with_hybrid_search(
  1128. collection_name=form_data.collection_name,
  1129. query=form_data.query,
  1130. embedding_function=request.app.state.EMBEDDING_FUNCTION,
  1131. k=form_data.k if form_data.k else request.app.state.config.TOP_K,
  1132. reranking_function=request.app.state.rf,
  1133. r=(
  1134. form_data.r
  1135. if form_data.r
  1136. else request.app.state.config.RELEVANCE_THRESHOLD
  1137. ),
  1138. )
  1139. else:
  1140. return query_doc(
  1141. collection_name=form_data.collection_name,
  1142. query_embedding=request.app.state.EMBEDDING_FUNCTION(form_data.query),
  1143. k=form_data.k if form_data.k else request.app.state.config.TOP_K,
  1144. )
  1145. except Exception as e:
  1146. log.exception(e)
  1147. raise HTTPException(
  1148. status_code=status.HTTP_400_BAD_REQUEST,
  1149. detail=ERROR_MESSAGES.DEFAULT(e),
  1150. )
  1151. class QueryCollectionsForm(BaseModel):
  1152. collection_names: list[str]
  1153. query: str
  1154. k: Optional[int] = None
  1155. r: Optional[float] = None
  1156. hybrid: Optional[bool] = None
  1157. @router.post("/query/collection")
  1158. def query_collection_handler(
  1159. request: Request,
  1160. form_data: QueryCollectionsForm,
  1161. user=Depends(get_verified_user),
  1162. ):
  1163. try:
  1164. if request.app.state.config.ENABLE_RAG_HYBRID_SEARCH:
  1165. return query_collection_with_hybrid_search(
  1166. collection_names=form_data.collection_names,
  1167. queries=[form_data.query],
  1168. embedding_function=request.app.state.EMBEDDING_FUNCTION,
  1169. k=form_data.k if form_data.k else request.app.state.config.TOP_K,
  1170. reranking_function=request.app.state.rf,
  1171. r=(
  1172. form_data.r
  1173. if form_data.r
  1174. else request.app.state.config.RELEVANCE_THRESHOLD
  1175. ),
  1176. )
  1177. else:
  1178. return query_collection(
  1179. collection_names=form_data.collection_names,
  1180. queries=[form_data.query],
  1181. embedding_function=request.app.state.EMBEDDING_FUNCTION,
  1182. k=form_data.k if form_data.k else request.app.state.config.TOP_K,
  1183. )
  1184. except Exception as e:
  1185. log.exception(e)
  1186. raise HTTPException(
  1187. status_code=status.HTTP_400_BAD_REQUEST,
  1188. detail=ERROR_MESSAGES.DEFAULT(e),
  1189. )
  1190. ####################################
  1191. #
  1192. # Vector DB operations
  1193. #
  1194. ####################################
  1195. class DeleteForm(BaseModel):
  1196. collection_name: str
  1197. file_id: str
  1198. @router.post("/delete")
  1199. def delete_entries_from_collection(form_data: DeleteForm, user=Depends(get_admin_user)):
  1200. try:
  1201. if VECTOR_DB_CLIENT.has_collection(collection_name=form_data.collection_name):
  1202. file = Files.get_file_by_id(form_data.file_id)
  1203. hash = file.hash
  1204. VECTOR_DB_CLIENT.delete(
  1205. collection_name=form_data.collection_name,
  1206. metadata={"hash": hash},
  1207. )
  1208. return {"status": True}
  1209. else:
  1210. return {"status": False}
  1211. except Exception as e:
  1212. log.exception(e)
  1213. return {"status": False}
  1214. @router.post("/reset/db")
  1215. def reset_vector_db(user=Depends(get_admin_user)):
  1216. VECTOR_DB_CLIENT.reset()
  1217. Knowledges.delete_all_knowledge()
  1218. @router.post("/reset/uploads")
  1219. def reset_upload_dir(user=Depends(get_admin_user)) -> bool:
  1220. folder = f"{UPLOAD_DIR}"
  1221. try:
  1222. # Check if the directory exists
  1223. if os.path.exists(folder):
  1224. # Iterate over all the files and directories in the specified directory
  1225. for filename in os.listdir(folder):
  1226. file_path = os.path.join(folder, filename)
  1227. try:
  1228. if os.path.isfile(file_path) or os.path.islink(file_path):
  1229. os.unlink(file_path) # Remove the file or link
  1230. elif os.path.isdir(file_path):
  1231. shutil.rmtree(file_path) # Remove the directory
  1232. except Exception as e:
  1233. print(f"Failed to delete {file_path}. Reason: {e}")
  1234. else:
  1235. print(f"The directory {folder} does not exist")
  1236. except Exception as e:
  1237. print(f"Failed to process the directory {folder}. Reason: {e}")
  1238. return True
  1239. if ENV == "dev":
  1240. @router.get("/ef/{text}")
  1241. async def get_embeddings(request: Request, text: Optional[str] = "Hello World!"):
  1242. return {"result": request.app.state.EMBEDDING_FUNCTION(text)}
  1243. class BatchProcessFilesForm(BaseModel):
  1244. files: List[FileModel]
  1245. collection_name: str
  1246. class BatchProcessFilesResult(BaseModel):
  1247. file_id: str
  1248. status: str
  1249. error: Optional[str] = None
  1250. class BatchProcessFilesResponse(BaseModel):
  1251. results: List[BatchProcessFilesResult]
  1252. errors: List[BatchProcessFilesResult]
  1253. @router.post("/process/files/batch")
  1254. def process_files_batch(
  1255. form_data: BatchProcessFilesForm,
  1256. user=Depends(get_verified_user),
  1257. ) -> BatchProcessFilesResponse:
  1258. """
  1259. Process a batch of files and save them to the vector database.
  1260. """
  1261. results: List[BatchProcessFilesResult] = []
  1262. errors: List[BatchProcessFilesResult] = []
  1263. collection_name = form_data.collection_name
  1264. # Prepare all documents first
  1265. all_docs: List[Document] = []
  1266. for file in form_data.files:
  1267. try:
  1268. text_content = file.data.get("content", "")
  1269. docs: List[Document] = [
  1270. Document(
  1271. page_content=text_content.replace("<br/>", "\n"),
  1272. metadata={
  1273. **file.meta,
  1274. "name": file.filename,
  1275. "created_by": file.user_id,
  1276. "file_id": file.id,
  1277. "source": file.filename,
  1278. },
  1279. )
  1280. ]
  1281. hash = calculate_sha256_string(text_content)
  1282. Files.update_file_hash_by_id(file.id, hash)
  1283. Files.update_file_data_by_id(file.id, {"content": text_content})
  1284. all_docs.extend(docs)
  1285. results.append(BatchProcessFilesResult(file_id=file.id, status="prepared"))
  1286. except Exception as e:
  1287. log.error(f"process_files_batch: Error processing file {file.id}: {str(e)}")
  1288. errors.append(
  1289. BatchProcessFilesResult(file_id=file.id, status="failed", error=str(e))
  1290. )
  1291. # Save all documents in one batch
  1292. if all_docs:
  1293. try:
  1294. save_docs_to_vector_db(
  1295. docs=all_docs, collection_name=collection_name, add=True
  1296. )
  1297. # Update all files with collection name
  1298. for result in results:
  1299. Files.update_file_metadata_by_id(
  1300. result.file_id, {"collection_name": collection_name}
  1301. )
  1302. result.status = "completed"
  1303. except Exception as e:
  1304. log.error(
  1305. f"process_files_batch: Error saving documents to vector DB: {str(e)}"
  1306. )
  1307. for result in results:
  1308. result.status = "failed"
  1309. errors.append(
  1310. BatchProcessFilesResult(file_id=result.file_id, error=str(e))
  1311. )
  1312. return BatchProcessFilesResponse(results=results, errors=errors)