main.py 52 KB

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