main.py 48 KB

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