main.py 55 KB

1234567891011121314151617181920212223242526272829303132333435363738394041424344454647484950515253545556575859606162636465666768697071727374757677787980818283848586878889909192939495969798991001011021031041051061071081091101111121131141151161171181191201211221231241251261271281291301311321331341351361371381391401411421431441451461471481491501511521531541551561571581591601611621631641651661671681691701711721731741751761771781791801811821831841851861871881891901911921931941951961971981992002012022032042052062072082092102112122132142152162172182192202212222232242252262272282292302312322332342352362372382392402412422432442452462472482492502512522532542552562572582592602612622632642652662672682692702712722732742752762772782792802812822832842852862872882892902912922932942952962972982993003013023033043053063073083093103113123133143153163173183193203213223233243253263273283293303313323333343353363373383393403413423433443453463473483493503513523533543553563573583593603613623633643653663673683693703713723733743753763773783793803813823833843853863873883893903913923933943953963973983994004014024034044054064074084094104114124134144154164174184194204214224234244254264274284294304314324334344354364374384394404414424434444454464474484494504514524534544554564574584594604614624634644654664674684694704714724734744754764774784794804814824834844854864874884894904914924934944954964974984995005015025035045055065075085095105115125135145155165175185195205215225235245255265275285295305315325335345355365375385395405415425435445455465475485495505515525535545555565575585595605615625635645655665675685695705715725735745755765775785795805815825835845855865875885895905915925935945955965975985996006016026036046056066076086096106116126136146156166176186196206216226236246256266276286296306316326336346356366376386396406416426436446456466476486496506516526536546556566576586596606616626636646656666676686696706716726736746756766776786796806816826836846856866876886896906916926936946956966976986997007017027037047057067077087097107117127137147157167177187197207217227237247257267277287297307317327337347357367377387397407417427437447457467477487497507517527537547557567577587597607617627637647657667677687697707717727737747757767777787797807817827837847857867877887897907917927937947957967977987998008018028038048058068078088098108118128138148158168178188198208218228238248258268278288298308318328338348358368378388398408418428438448458468478488498508518528538548558568578588598608618628638648658668678688698708718728738748758768778788798808818828838848858868878888898908918928938948958968978988999009019029039049059069079089099109119129139149159169179189199209219229239249259269279289299309319329339349359369379389399409419429439449459469479489499509519529539549559569579589599609619629639649659669679689699709719729739749759769779789799809819829839849859869879889899909919929939949959969979989991000100110021003100410051006100710081009101010111012101310141015101610171018101910201021102210231024102510261027102810291030103110321033103410351036103710381039104010411042104310441045104610471048104910501051105210531054105510561057105810591060106110621063106410651066106710681069107010711072107310741075107610771078107910801081108210831084108510861087108810891090109110921093109410951096109710981099110011011102110311041105110611071108110911101111111211131114111511161117111811191120112111221123112411251126112711281129113011311132113311341135113611371138113911401141114211431144114511461147114811491150115111521153115411551156115711581159116011611162116311641165116611671168116911701171117211731174117511761177117811791180118111821183118411851186118711881189119011911192119311941195119611971198119912001201120212031204120512061207120812091210121112121213121412151216121712181219122012211222122312241225122612271228122912301231123212331234123512361237123812391240124112421243124412451246124712481249125012511252125312541255125612571258125912601261126212631264126512661267126812691270127112721273127412751276127712781279128012811282128312841285128612871288128912901291129212931294129512961297129812991300130113021303130413051306130713081309131013111312131313141315131613171318131913201321132213231324132513261327132813291330133113321333133413351336133713381339134013411342134313441345134613471348134913501351135213531354135513561357135813591360136113621363136413651366136713681369137013711372137313741375137613771378137913801381138213831384138513861387138813891390139113921393139413951396139713981399140014011402140314041405140614071408140914101411141214131414141514161417141814191420142114221423142414251426142714281429143014311432143314341435143614371438143914401441144214431444144514461447144814491450145114521453145414551456145714581459146014611462146314641465146614671468146914701471147214731474147514761477147814791480148114821483148414851486148714881489149014911492149314941495149614971498149915001501150215031504150515061507150815091510151115121513151415151516151715181519152015211522152315241525152615271528152915301531153215331534153515361537153815391540154115421543154415451546154715481549155015511552155315541555155615571558155915601561156215631564156515661567156815691570157115721573157415751576157715781579
  1. import json
  2. import logging
  3. import mimetypes
  4. import os
  5. import shutil
  6. import socket
  7. import urllib.parse
  8. import uuid
  9. from datetime import datetime
  10. from pathlib import Path
  11. from typing import Iterator, Optional, Sequence, Union
  12. import numpy as np
  13. import torch
  14. import requests
  15. import validators
  16. from fastapi import Depends, FastAPI, File, Form, HTTPException, UploadFile, status
  17. from fastapi.middleware.cors import CORSMiddleware
  18. from pydantic import BaseModel
  19. from open_webui.apps.rag.search.main import SearchResult
  20. from open_webui.apps.rag.search.brave import search_brave
  21. from open_webui.apps.rag.search.duckduckgo import search_duckduckgo
  22. from open_webui.apps.rag.search.google_pse import search_google_pse
  23. from open_webui.apps.rag.search.jina_search import search_jina
  24. from open_webui.apps.rag.search.searchapi import search_searchapi
  25. from open_webui.apps.rag.search.searxng import search_searxng
  26. from open_webui.apps.rag.search.serper import search_serper
  27. from open_webui.apps.rag.search.serply import search_serply
  28. from open_webui.apps.rag.search.serpstack import search_serpstack
  29. from open_webui.apps.rag.search.tavily import search_tavily
  30. from open_webui.apps.rag.utils import (
  31. get_embedding_function,
  32. get_model_path,
  33. query_collection,
  34. query_collection_with_hybrid_search,
  35. query_doc,
  36. query_doc_with_hybrid_search,
  37. )
  38. from open_webui.apps.webui.models.documents import DocumentForm, Documents
  39. from open_webui.apps.webui.models.files import Files
  40. from open_webui.config import (
  41. BRAVE_SEARCH_API_KEY,
  42. CHUNK_OVERLAP,
  43. CHUNK_SIZE,
  44. CONTENT_EXTRACTION_ENGINE,
  45. CORS_ALLOW_ORIGIN,
  46. DOCS_DIR,
  47. ENABLE_RAG_HYBRID_SEARCH,
  48. ENABLE_RAG_LOCAL_WEB_FETCH,
  49. ENABLE_RAG_WEB_LOADER_SSL_VERIFICATION,
  50. ENABLE_RAG_WEB_SEARCH,
  51. ENV,
  52. GOOGLE_PSE_API_KEY,
  53. GOOGLE_PSE_ENGINE_ID,
  54. PDF_EXTRACT_IMAGES,
  55. RAG_EMBEDDING_ENGINE,
  56. RAG_EMBEDDING_MODEL,
  57. RAG_EMBEDDING_MODEL_AUTO_UPDATE,
  58. RAG_EMBEDDING_MODEL_TRUST_REMOTE_CODE,
  59. RAG_EMBEDDING_OPENAI_BATCH_SIZE,
  60. RAG_FILE_MAX_COUNT,
  61. RAG_FILE_MAX_SIZE,
  62. RAG_OPENAI_API_BASE_URL,
  63. RAG_OPENAI_API_KEY,
  64. RAG_RELEVANCE_THRESHOLD,
  65. RAG_RERANKING_MODEL,
  66. RAG_RERANKING_MODEL_AUTO_UPDATE,
  67. RAG_RERANKING_MODEL_TRUST_REMOTE_CODE,
  68. DEFAULT_RAG_TEMPLATE,
  69. RAG_TEMPLATE,
  70. RAG_TOP_K,
  71. RAG_WEB_SEARCH_CONCURRENT_REQUESTS,
  72. RAG_WEB_SEARCH_DOMAIN_FILTER_LIST,
  73. RAG_WEB_SEARCH_ENGINE,
  74. RAG_WEB_SEARCH_RESULT_COUNT,
  75. SEARCHAPI_API_KEY,
  76. SEARCHAPI_ENGINE,
  77. SEARXNG_QUERY_URL,
  78. SERPER_API_KEY,
  79. SERPLY_API_KEY,
  80. SERPSTACK_API_KEY,
  81. SERPSTACK_HTTPS,
  82. TAVILY_API_KEY,
  83. TIKA_SERVER_URL,
  84. UPLOAD_DIR,
  85. YOUTUBE_LOADER_LANGUAGE,
  86. AppConfig,
  87. )
  88. from open_webui.constants import ERROR_MESSAGES
  89. from open_webui.env import SRC_LOG_LEVELS, DEVICE_TYPE, DOCKER
  90. from open_webui.utils.misc import (
  91. calculate_sha256,
  92. calculate_sha256_string,
  93. extract_folders_after_data_docs,
  94. sanitize_filename,
  95. )
  96. from open_webui.utils.utils import get_admin_user, get_verified_user
  97. from open_webui.apps.rag.vector.connector import VECTOR_DB_CLIENT
  98. from langchain.text_splitter import RecursiveCharacterTextSplitter
  99. from langchain_community.document_loaders import (
  100. BSHTMLLoader,
  101. CSVLoader,
  102. Docx2txtLoader,
  103. OutlookMessageLoader,
  104. PyPDFLoader,
  105. TextLoader,
  106. UnstructuredEPubLoader,
  107. UnstructuredExcelLoader,
  108. UnstructuredMarkdownLoader,
  109. UnstructuredPowerPointLoader,
  110. UnstructuredRSTLoader,
  111. UnstructuredXMLLoader,
  112. WebBaseLoader,
  113. YoutubeLoader,
  114. )
  115. from langchain_core.documents import Document
  116. from colbert.infra import ColBERTConfig
  117. from colbert.modeling.checkpoint import Checkpoint
  118. log = logging.getLogger(__name__)
  119. log.setLevel(SRC_LOG_LEVELS["RAG"])
  120. app = FastAPI()
  121. app.state.config = AppConfig()
  122. app.state.config.TOP_K = RAG_TOP_K
  123. app.state.config.RELEVANCE_THRESHOLD = RAG_RELEVANCE_THRESHOLD
  124. app.state.config.FILE_MAX_SIZE = RAG_FILE_MAX_SIZE
  125. app.state.config.FILE_MAX_COUNT = RAG_FILE_MAX_COUNT
  126. app.state.config.ENABLE_RAG_HYBRID_SEARCH = ENABLE_RAG_HYBRID_SEARCH
  127. app.state.config.ENABLE_RAG_WEB_LOADER_SSL_VERIFICATION = (
  128. ENABLE_RAG_WEB_LOADER_SSL_VERIFICATION
  129. )
  130. app.state.config.CONTENT_EXTRACTION_ENGINE = CONTENT_EXTRACTION_ENGINE
  131. app.state.config.TIKA_SERVER_URL = TIKA_SERVER_URL
  132. app.state.config.CHUNK_SIZE = CHUNK_SIZE
  133. app.state.config.CHUNK_OVERLAP = CHUNK_OVERLAP
  134. app.state.config.RAG_EMBEDDING_ENGINE = RAG_EMBEDDING_ENGINE
  135. app.state.config.RAG_EMBEDDING_MODEL = RAG_EMBEDDING_MODEL
  136. app.state.config.RAG_EMBEDDING_OPENAI_BATCH_SIZE = RAG_EMBEDDING_OPENAI_BATCH_SIZE
  137. app.state.config.RAG_RERANKING_MODEL = RAG_RERANKING_MODEL
  138. app.state.config.RAG_TEMPLATE = RAG_TEMPLATE
  139. app.state.config.OPENAI_API_BASE_URL = RAG_OPENAI_API_BASE_URL
  140. app.state.config.OPENAI_API_KEY = RAG_OPENAI_API_KEY
  141. app.state.config.PDF_EXTRACT_IMAGES = PDF_EXTRACT_IMAGES
  142. app.state.config.YOUTUBE_LOADER_LANGUAGE = YOUTUBE_LOADER_LANGUAGE
  143. app.state.YOUTUBE_LOADER_TRANSLATION = None
  144. app.state.config.ENABLE_RAG_WEB_SEARCH = ENABLE_RAG_WEB_SEARCH
  145. app.state.config.RAG_WEB_SEARCH_ENGINE = RAG_WEB_SEARCH_ENGINE
  146. app.state.config.RAG_WEB_SEARCH_DOMAIN_FILTER_LIST = RAG_WEB_SEARCH_DOMAIN_FILTER_LIST
  147. app.state.config.SEARXNG_QUERY_URL = SEARXNG_QUERY_URL
  148. app.state.config.GOOGLE_PSE_API_KEY = GOOGLE_PSE_API_KEY
  149. app.state.config.GOOGLE_PSE_ENGINE_ID = GOOGLE_PSE_ENGINE_ID
  150. app.state.config.BRAVE_SEARCH_API_KEY = BRAVE_SEARCH_API_KEY
  151. app.state.config.SERPSTACK_API_KEY = SERPSTACK_API_KEY
  152. app.state.config.SERPSTACK_HTTPS = SERPSTACK_HTTPS
  153. app.state.config.SERPER_API_KEY = SERPER_API_KEY
  154. app.state.config.SERPLY_API_KEY = SERPLY_API_KEY
  155. app.state.config.TAVILY_API_KEY = TAVILY_API_KEY
  156. app.state.config.SEARCHAPI_API_KEY = SEARCHAPI_API_KEY
  157. app.state.config.SEARCHAPI_ENGINE = SEARCHAPI_ENGINE
  158. app.state.config.RAG_WEB_SEARCH_RESULT_COUNT = RAG_WEB_SEARCH_RESULT_COUNT
  159. app.state.config.RAG_WEB_SEARCH_CONCURRENT_REQUESTS = RAG_WEB_SEARCH_CONCURRENT_REQUESTS
  160. def update_embedding_model(
  161. embedding_model: str,
  162. auto_update: bool = False,
  163. ):
  164. if embedding_model and app.state.config.RAG_EMBEDDING_ENGINE == "":
  165. import sentence_transformers
  166. app.state.sentence_transformer_ef = sentence_transformers.SentenceTransformer(
  167. get_model_path(embedding_model, auto_update),
  168. device=DEVICE_TYPE,
  169. trust_remote_code=RAG_EMBEDDING_MODEL_TRUST_REMOTE_CODE,
  170. )
  171. else:
  172. app.state.sentence_transformer_ef = None
  173. def update_reranking_model(
  174. reranking_model: str,
  175. auto_update: bool = False,
  176. ):
  177. if reranking_model:
  178. if any(model in reranking_model for model in ["jinaai/jina-colbert-v2"]):
  179. class ColBERT:
  180. def __init__(self, name) -> None:
  181. print("ColBERT: Loading model", name)
  182. self.device = "cuda" if torch.cuda.is_available() else "cpu"
  183. if DOCKER:
  184. # This is a workaround for the issue with the docker container
  185. # where the torch extension is not loaded properly
  186. # and the following error is thrown:
  187. # /root/.cache/torch_extensions/py311_cpu/segmented_maxsim_cpp/segmented_maxsim_cpp.so: cannot open shared object file: No such file or directory
  188. lock_file = "/root/.cache/torch_extensions/py311_cpu/segmented_maxsim_cpp/lock"
  189. if os.path.exists(lock_file):
  190. os.remove(lock_file)
  191. print("ColBERT: Removed lock file")
  192. self.ckpt = Checkpoint(
  193. name,
  194. colbert_config=ColBERTConfig(model_name=name),
  195. ).to(self.device)
  196. pass
  197. def calculate_similarity_scores(
  198. self, query_embeddings, document_embeddings
  199. ):
  200. query_embeddings = query_embeddings.to(self.device)
  201. document_embeddings = document_embeddings.to(self.device)
  202. # Validate dimensions to ensure compatibility
  203. if query_embeddings.dim() != 3:
  204. raise ValueError(
  205. f"Expected query embeddings to have 3 dimensions, but got {query_embeddings.dim()}."
  206. )
  207. if document_embeddings.dim() != 3:
  208. raise ValueError(
  209. f"Expected document embeddings to have 3 dimensions, but got {document_embeddings.dim()}."
  210. )
  211. if query_embeddings.size(0) not in [1, document_embeddings.size(0)]:
  212. raise ValueError(
  213. "There should be either one query or queries equal to the number of documents."
  214. )
  215. # Transpose the query embeddings to align for matrix multiplication
  216. transposed_query_embeddings = query_embeddings.permute(0, 2, 1)
  217. # Compute similarity scores using batch matrix multiplication
  218. computed_scores = torch.matmul(
  219. document_embeddings, transposed_query_embeddings
  220. )
  221. # Apply max pooling to extract the highest semantic similarity across each document's sequence
  222. maximum_scores = torch.max(computed_scores, dim=1).values
  223. # Sum up the maximum scores across features to get the overall document relevance scores
  224. final_scores = maximum_scores.sum(dim=1)
  225. normalized_scores = torch.softmax(final_scores, dim=0)
  226. return normalized_scores.detach().cpu().numpy().astype(np.float32)
  227. def predict(self, sentences):
  228. query = sentences[0][0]
  229. docs = [i[1] for i in sentences]
  230. # Embedding the documents
  231. embedded_docs = self.ckpt.docFromText(docs, bsize=32)[0]
  232. # Embedding the queries
  233. embedded_queries = self.ckpt.queryFromText([query], bsize=32)
  234. embedded_query = embedded_queries[0]
  235. # Calculate retrieval scores for the query against all documents
  236. scores = self.calculate_similarity_scores(
  237. embedded_query.unsqueeze(0), embedded_docs
  238. )
  239. return scores
  240. try:
  241. app.state.sentence_transformer_rf = ColBERT(
  242. get_model_path(reranking_model, auto_update)
  243. )
  244. except Exception as e:
  245. log.error(f"ColBERT: {e}")
  246. app.state.sentence_transformer_rf = None
  247. app.state.config.ENABLE_RAG_HYBRID_SEARCH = False
  248. else:
  249. import sentence_transformers
  250. try:
  251. app.state.sentence_transformer_rf = sentence_transformers.CrossEncoder(
  252. get_model_path(reranking_model, auto_update),
  253. device=DEVICE_TYPE,
  254. trust_remote_code=RAG_RERANKING_MODEL_TRUST_REMOTE_CODE,
  255. )
  256. except:
  257. log.error("CrossEncoder error")
  258. app.state.sentence_transformer_rf = None
  259. app.state.config.ENABLE_RAG_HYBRID_SEARCH = False
  260. else:
  261. app.state.sentence_transformer_rf = None
  262. update_embedding_model(
  263. app.state.config.RAG_EMBEDDING_MODEL,
  264. RAG_EMBEDDING_MODEL_AUTO_UPDATE,
  265. )
  266. update_reranking_model(
  267. app.state.config.RAG_RERANKING_MODEL,
  268. RAG_RERANKING_MODEL_AUTO_UPDATE,
  269. )
  270. app.state.EMBEDDING_FUNCTION = get_embedding_function(
  271. app.state.config.RAG_EMBEDDING_ENGINE,
  272. app.state.config.RAG_EMBEDDING_MODEL,
  273. app.state.sentence_transformer_ef,
  274. app.state.config.OPENAI_API_KEY,
  275. app.state.config.OPENAI_API_BASE_URL,
  276. app.state.config.RAG_EMBEDDING_OPENAI_BATCH_SIZE,
  277. )
  278. app.add_middleware(
  279. CORSMiddleware,
  280. allow_origins=CORS_ALLOW_ORIGIN,
  281. allow_credentials=True,
  282. allow_methods=["*"],
  283. allow_headers=["*"],
  284. )
  285. class CollectionNameForm(BaseModel):
  286. collection_name: Optional[str] = "test"
  287. class UrlForm(CollectionNameForm):
  288. url: str
  289. class SearchForm(CollectionNameForm):
  290. query: str
  291. @app.get("/")
  292. async def get_status():
  293. return {
  294. "status": True,
  295. "chunk_size": app.state.config.CHUNK_SIZE,
  296. "chunk_overlap": app.state.config.CHUNK_OVERLAP,
  297. "template": app.state.config.RAG_TEMPLATE,
  298. "embedding_engine": app.state.config.RAG_EMBEDDING_ENGINE,
  299. "embedding_model": app.state.config.RAG_EMBEDDING_MODEL,
  300. "reranking_model": app.state.config.RAG_RERANKING_MODEL,
  301. "openai_batch_size": app.state.config.RAG_EMBEDDING_OPENAI_BATCH_SIZE,
  302. }
  303. @app.get("/embedding")
  304. async def get_embedding_config(user=Depends(get_admin_user)):
  305. return {
  306. "status": True,
  307. "embedding_engine": app.state.config.RAG_EMBEDDING_ENGINE,
  308. "embedding_model": app.state.config.RAG_EMBEDDING_MODEL,
  309. "openai_config": {
  310. "url": app.state.config.OPENAI_API_BASE_URL,
  311. "key": app.state.config.OPENAI_API_KEY,
  312. "batch_size": app.state.config.RAG_EMBEDDING_OPENAI_BATCH_SIZE,
  313. },
  314. }
  315. @app.get("/reranking")
  316. async def get_reraanking_config(user=Depends(get_admin_user)):
  317. return {
  318. "status": True,
  319. "reranking_model": app.state.config.RAG_RERANKING_MODEL,
  320. }
  321. class OpenAIConfigForm(BaseModel):
  322. url: str
  323. key: str
  324. batch_size: Optional[int] = None
  325. class EmbeddingModelUpdateForm(BaseModel):
  326. openai_config: Optional[OpenAIConfigForm] = None
  327. embedding_engine: str
  328. embedding_model: str
  329. @app.post("/embedding/update")
  330. async def update_embedding_config(
  331. form_data: EmbeddingModelUpdateForm, user=Depends(get_admin_user)
  332. ):
  333. log.info(
  334. f"Updating embedding model: {app.state.config.RAG_EMBEDDING_MODEL} to {form_data.embedding_model}"
  335. )
  336. try:
  337. app.state.config.RAG_EMBEDDING_ENGINE = form_data.embedding_engine
  338. app.state.config.RAG_EMBEDDING_MODEL = form_data.embedding_model
  339. if app.state.config.RAG_EMBEDDING_ENGINE in ["ollama", "openai"]:
  340. if form_data.openai_config is not None:
  341. app.state.config.OPENAI_API_BASE_URL = form_data.openai_config.url
  342. app.state.config.OPENAI_API_KEY = form_data.openai_config.key
  343. app.state.config.RAG_EMBEDDING_OPENAI_BATCH_SIZE = (
  344. form_data.openai_config.batch_size
  345. if form_data.openai_config.batch_size
  346. else 1
  347. )
  348. update_embedding_model(app.state.config.RAG_EMBEDDING_MODEL)
  349. app.state.EMBEDDING_FUNCTION = get_embedding_function(
  350. app.state.config.RAG_EMBEDDING_ENGINE,
  351. app.state.config.RAG_EMBEDDING_MODEL,
  352. app.state.sentence_transformer_ef,
  353. app.state.config.OPENAI_API_KEY,
  354. app.state.config.OPENAI_API_BASE_URL,
  355. app.state.config.RAG_EMBEDDING_OPENAI_BATCH_SIZE,
  356. )
  357. return {
  358. "status": True,
  359. "embedding_engine": app.state.config.RAG_EMBEDDING_ENGINE,
  360. "embedding_model": app.state.config.RAG_EMBEDDING_MODEL,
  361. "openai_config": {
  362. "url": app.state.config.OPENAI_API_BASE_URL,
  363. "key": app.state.config.OPENAI_API_KEY,
  364. "batch_size": app.state.config.RAG_EMBEDDING_OPENAI_BATCH_SIZE,
  365. },
  366. }
  367. except Exception as e:
  368. log.exception(f"Problem updating embedding model: {e}")
  369. raise HTTPException(
  370. status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
  371. detail=ERROR_MESSAGES.DEFAULT(e),
  372. )
  373. class RerankingModelUpdateForm(BaseModel):
  374. reranking_model: str
  375. @app.post("/reranking/update")
  376. async def update_reranking_config(
  377. form_data: RerankingModelUpdateForm, user=Depends(get_admin_user)
  378. ):
  379. log.info(
  380. f"Updating reranking model: {app.state.config.RAG_RERANKING_MODEL} to {form_data.reranking_model}"
  381. )
  382. try:
  383. app.state.config.RAG_RERANKING_MODEL = form_data.reranking_model
  384. update_reranking_model(app.state.config.RAG_RERANKING_MODEL, True)
  385. return {
  386. "status": True,
  387. "reranking_model": app.state.config.RAG_RERANKING_MODEL,
  388. }
  389. except Exception as e:
  390. log.exception(f"Problem updating reranking model: {e}")
  391. raise HTTPException(
  392. status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
  393. detail=ERROR_MESSAGES.DEFAULT(e),
  394. )
  395. @app.get("/config")
  396. async def get_rag_config(user=Depends(get_admin_user)):
  397. return {
  398. "status": True,
  399. "pdf_extract_images": app.state.config.PDF_EXTRACT_IMAGES,
  400. "file": {
  401. "max_size": app.state.config.FILE_MAX_SIZE,
  402. "max_count": app.state.config.FILE_MAX_COUNT,
  403. },
  404. "content_extraction": {
  405. "engine": app.state.config.CONTENT_EXTRACTION_ENGINE,
  406. "tika_server_url": app.state.config.TIKA_SERVER_URL,
  407. },
  408. "chunk": {
  409. "chunk_size": app.state.config.CHUNK_SIZE,
  410. "chunk_overlap": app.state.config.CHUNK_OVERLAP,
  411. },
  412. "youtube": {
  413. "language": app.state.config.YOUTUBE_LOADER_LANGUAGE,
  414. "translation": app.state.YOUTUBE_LOADER_TRANSLATION,
  415. },
  416. "web": {
  417. "ssl_verification": app.state.config.ENABLE_RAG_WEB_LOADER_SSL_VERIFICATION,
  418. "search": {
  419. "enabled": app.state.config.ENABLE_RAG_WEB_SEARCH,
  420. "engine": app.state.config.RAG_WEB_SEARCH_ENGINE,
  421. "searxng_query_url": app.state.config.SEARXNG_QUERY_URL,
  422. "google_pse_api_key": app.state.config.GOOGLE_PSE_API_KEY,
  423. "google_pse_engine_id": app.state.config.GOOGLE_PSE_ENGINE_ID,
  424. "brave_search_api_key": app.state.config.BRAVE_SEARCH_API_KEY,
  425. "serpstack_api_key": app.state.config.SERPSTACK_API_KEY,
  426. "serpstack_https": app.state.config.SERPSTACK_HTTPS,
  427. "serper_api_key": app.state.config.SERPER_API_KEY,
  428. "serply_api_key": app.state.config.SERPLY_API_KEY,
  429. "tavily_api_key": app.state.config.TAVILY_API_KEY,
  430. "searchapi_api_key": app.state.config.SEARCHAPI_API_KEY,
  431. "seaarchapi_engine": app.state.config.SEARCHAPI_ENGINE,
  432. "result_count": app.state.config.RAG_WEB_SEARCH_RESULT_COUNT,
  433. "concurrent_requests": app.state.config.RAG_WEB_SEARCH_CONCURRENT_REQUESTS,
  434. },
  435. },
  436. }
  437. class FileConfig(BaseModel):
  438. max_size: Optional[int] = None
  439. max_count: Optional[int] = None
  440. class ContentExtractionConfig(BaseModel):
  441. engine: str = ""
  442. tika_server_url: Optional[str] = None
  443. class ChunkParamUpdateForm(BaseModel):
  444. chunk_size: int
  445. chunk_overlap: int
  446. class YoutubeLoaderConfig(BaseModel):
  447. language: list[str]
  448. translation: Optional[str] = None
  449. class WebSearchConfig(BaseModel):
  450. enabled: bool
  451. engine: Optional[str] = None
  452. searxng_query_url: Optional[str] = None
  453. google_pse_api_key: Optional[str] = None
  454. google_pse_engine_id: Optional[str] = None
  455. brave_search_api_key: Optional[str] = None
  456. serpstack_api_key: Optional[str] = None
  457. serpstack_https: Optional[bool] = None
  458. serper_api_key: Optional[str] = None
  459. serply_api_key: Optional[str] = None
  460. tavily_api_key: Optional[str] = None
  461. searchapi_api_key: Optional[str] = None
  462. searchapi_engine: Optional[str] = None
  463. result_count: Optional[int] = None
  464. concurrent_requests: Optional[int] = None
  465. class WebConfig(BaseModel):
  466. search: WebSearchConfig
  467. web_loader_ssl_verification: Optional[bool] = None
  468. class ConfigUpdateForm(BaseModel):
  469. pdf_extract_images: Optional[bool] = None
  470. file: Optional[FileConfig] = None
  471. content_extraction: Optional[ContentExtractionConfig] = None
  472. chunk: Optional[ChunkParamUpdateForm] = None
  473. youtube: Optional[YoutubeLoaderConfig] = None
  474. web: Optional[WebConfig] = None
  475. @app.post("/config/update")
  476. async def update_rag_config(form_data: ConfigUpdateForm, user=Depends(get_admin_user)):
  477. app.state.config.PDF_EXTRACT_IMAGES = (
  478. form_data.pdf_extract_images
  479. if form_data.pdf_extract_images is not None
  480. else app.state.config.PDF_EXTRACT_IMAGES
  481. )
  482. if form_data.file is not None:
  483. app.state.config.FILE_MAX_SIZE = form_data.file.max_size
  484. app.state.config.FILE_MAX_COUNT = form_data.file.max_count
  485. if form_data.content_extraction is not None:
  486. log.info(f"Updating text settings: {form_data.content_extraction}")
  487. app.state.config.CONTENT_EXTRACTION_ENGINE = form_data.content_extraction.engine
  488. app.state.config.TIKA_SERVER_URL = form_data.content_extraction.tika_server_url
  489. if form_data.chunk is not None:
  490. app.state.config.CHUNK_SIZE = form_data.chunk.chunk_size
  491. app.state.config.CHUNK_OVERLAP = form_data.chunk.chunk_overlap
  492. if form_data.youtube is not None:
  493. app.state.config.YOUTUBE_LOADER_LANGUAGE = form_data.youtube.language
  494. app.state.YOUTUBE_LOADER_TRANSLATION = form_data.youtube.translation
  495. if form_data.web is not None:
  496. app.state.config.ENABLE_RAG_WEB_LOADER_SSL_VERIFICATION = (
  497. form_data.web.web_loader_ssl_verification
  498. )
  499. app.state.config.ENABLE_RAG_WEB_SEARCH = form_data.web.search.enabled
  500. app.state.config.RAG_WEB_SEARCH_ENGINE = form_data.web.search.engine
  501. app.state.config.SEARXNG_QUERY_URL = form_data.web.search.searxng_query_url
  502. app.state.config.GOOGLE_PSE_API_KEY = form_data.web.search.google_pse_api_key
  503. app.state.config.GOOGLE_PSE_ENGINE_ID = (
  504. form_data.web.search.google_pse_engine_id
  505. )
  506. app.state.config.BRAVE_SEARCH_API_KEY = (
  507. form_data.web.search.brave_search_api_key
  508. )
  509. app.state.config.SERPSTACK_API_KEY = form_data.web.search.serpstack_api_key
  510. app.state.config.SERPSTACK_HTTPS = form_data.web.search.serpstack_https
  511. app.state.config.SERPER_API_KEY = form_data.web.search.serper_api_key
  512. app.state.config.SERPLY_API_KEY = form_data.web.search.serply_api_key
  513. app.state.config.TAVILY_API_KEY = form_data.web.search.tavily_api_key
  514. app.state.config.SEARCHAPI_API_KEY = form_data.web.search.searchapi_api_key
  515. app.state.config.SEARCHAPI_ENGINE = form_data.web.search.searchapi_engine
  516. app.state.config.RAG_WEB_SEARCH_RESULT_COUNT = form_data.web.search.result_count
  517. app.state.config.RAG_WEB_SEARCH_CONCURRENT_REQUESTS = (
  518. form_data.web.search.concurrent_requests
  519. )
  520. return {
  521. "status": True,
  522. "pdf_extract_images": app.state.config.PDF_EXTRACT_IMAGES,
  523. "file": {
  524. "max_size": app.state.config.FILE_MAX_SIZE,
  525. "max_count": app.state.config.FILE_MAX_COUNT,
  526. },
  527. "content_extraction": {
  528. "engine": app.state.config.CONTENT_EXTRACTION_ENGINE,
  529. "tika_server_url": app.state.config.TIKA_SERVER_URL,
  530. },
  531. "chunk": {
  532. "chunk_size": app.state.config.CHUNK_SIZE,
  533. "chunk_overlap": app.state.config.CHUNK_OVERLAP,
  534. },
  535. "youtube": {
  536. "language": app.state.config.YOUTUBE_LOADER_LANGUAGE,
  537. "translation": app.state.YOUTUBE_LOADER_TRANSLATION,
  538. },
  539. "web": {
  540. "ssl_verification": app.state.config.ENABLE_RAG_WEB_LOADER_SSL_VERIFICATION,
  541. "search": {
  542. "enabled": app.state.config.ENABLE_RAG_WEB_SEARCH,
  543. "engine": app.state.config.RAG_WEB_SEARCH_ENGINE,
  544. "searxng_query_url": app.state.config.SEARXNG_QUERY_URL,
  545. "google_pse_api_key": app.state.config.GOOGLE_PSE_API_KEY,
  546. "google_pse_engine_id": app.state.config.GOOGLE_PSE_ENGINE_ID,
  547. "brave_search_api_key": app.state.config.BRAVE_SEARCH_API_KEY,
  548. "serpstack_api_key": app.state.config.SERPSTACK_API_KEY,
  549. "serpstack_https": app.state.config.SERPSTACK_HTTPS,
  550. "serper_api_key": app.state.config.SERPER_API_KEY,
  551. "serply_api_key": app.state.config.SERPLY_API_KEY,
  552. "serachapi_api_key": app.state.config.SEARCHAPI_API_KEY,
  553. "searchapi_engine": app.state.config.SEARCHAPI_ENGINE,
  554. "tavily_api_key": app.state.config.TAVILY_API_KEY,
  555. "result_count": app.state.config.RAG_WEB_SEARCH_RESULT_COUNT,
  556. "concurrent_requests": app.state.config.RAG_WEB_SEARCH_CONCURRENT_REQUESTS,
  557. },
  558. },
  559. }
  560. @app.get("/template")
  561. async def get_rag_template(user=Depends(get_verified_user)):
  562. return {
  563. "status": True,
  564. "template": app.state.config.RAG_TEMPLATE,
  565. }
  566. @app.get("/query/settings")
  567. async def get_query_settings(user=Depends(get_admin_user)):
  568. return {
  569. "status": True,
  570. "template": app.state.config.RAG_TEMPLATE,
  571. "k": app.state.config.TOP_K,
  572. "r": app.state.config.RELEVANCE_THRESHOLD,
  573. "hybrid": app.state.config.ENABLE_RAG_HYBRID_SEARCH,
  574. }
  575. class QuerySettingsForm(BaseModel):
  576. k: Optional[int] = None
  577. r: Optional[float] = None
  578. template: Optional[str] = None
  579. hybrid: Optional[bool] = None
  580. @app.post("/query/settings/update")
  581. async def update_query_settings(
  582. form_data: QuerySettingsForm, user=Depends(get_admin_user)
  583. ):
  584. app.state.config.RAG_TEMPLATE = (
  585. form_data.template if form_data.template != "" else DEFAULT_RAG_TEMPLATE
  586. )
  587. app.state.config.TOP_K = form_data.k if form_data.k else 4
  588. app.state.config.RELEVANCE_THRESHOLD = form_data.r if form_data.r else 0.0
  589. app.state.config.ENABLE_RAG_HYBRID_SEARCH = (
  590. form_data.hybrid if form_data.hybrid else False
  591. )
  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 QueryDocForm(BaseModel):
  600. collection_name: str
  601. query: str
  602. k: Optional[int] = None
  603. r: Optional[float] = None
  604. hybrid: Optional[bool] = None
  605. @app.post("/query/doc")
  606. def query_doc_handler(
  607. form_data: QueryDocForm,
  608. user=Depends(get_verified_user),
  609. ):
  610. try:
  611. if app.state.config.ENABLE_RAG_HYBRID_SEARCH:
  612. return query_doc_with_hybrid_search(
  613. collection_name=form_data.collection_name,
  614. query=form_data.query,
  615. embedding_function=app.state.EMBEDDING_FUNCTION,
  616. k=form_data.k if form_data.k else app.state.config.TOP_K,
  617. reranking_function=app.state.sentence_transformer_rf,
  618. r=(
  619. form_data.r if form_data.r else app.state.config.RELEVANCE_THRESHOLD
  620. ),
  621. )
  622. else:
  623. return query_doc(
  624. collection_name=form_data.collection_name,
  625. query=form_data.query,
  626. embedding_function=app.state.EMBEDDING_FUNCTION,
  627. k=form_data.k if form_data.k else app.state.config.TOP_K,
  628. )
  629. except Exception as e:
  630. log.exception(e)
  631. raise HTTPException(
  632. status_code=status.HTTP_400_BAD_REQUEST,
  633. detail=ERROR_MESSAGES.DEFAULT(e),
  634. )
  635. class QueryCollectionsForm(BaseModel):
  636. collection_names: list[str]
  637. query: str
  638. k: Optional[int] = None
  639. r: Optional[float] = None
  640. hybrid: Optional[bool] = None
  641. @app.post("/query/collection")
  642. def query_collection_handler(
  643. form_data: QueryCollectionsForm,
  644. user=Depends(get_verified_user),
  645. ):
  646. try:
  647. if app.state.config.ENABLE_RAG_HYBRID_SEARCH:
  648. return query_collection_with_hybrid_search(
  649. collection_names=form_data.collection_names,
  650. query=form_data.query,
  651. embedding_function=app.state.EMBEDDING_FUNCTION,
  652. k=form_data.k if form_data.k else app.state.config.TOP_K,
  653. reranking_function=app.state.sentence_transformer_rf,
  654. r=(
  655. form_data.r if form_data.r else app.state.config.RELEVANCE_THRESHOLD
  656. ),
  657. )
  658. else:
  659. return query_collection(
  660. collection_names=form_data.collection_names,
  661. query=form_data.query,
  662. embedding_function=app.state.EMBEDDING_FUNCTION,
  663. k=form_data.k if form_data.k else app.state.config.TOP_K,
  664. )
  665. except Exception as e:
  666. log.exception(e)
  667. raise HTTPException(
  668. status_code=status.HTTP_400_BAD_REQUEST,
  669. detail=ERROR_MESSAGES.DEFAULT(e),
  670. )
  671. @app.post("/youtube")
  672. def store_youtube_video(form_data: UrlForm, user=Depends(get_verified_user)):
  673. try:
  674. loader = YoutubeLoader.from_youtube_url(
  675. form_data.url,
  676. add_video_info=True,
  677. language=app.state.config.YOUTUBE_LOADER_LANGUAGE,
  678. translation=app.state.YOUTUBE_LOADER_TRANSLATION,
  679. )
  680. data = loader.load()
  681. collection_name = form_data.collection_name
  682. if collection_name == "":
  683. collection_name = calculate_sha256_string(form_data.url)[:63]
  684. store_data_in_vector_db(data, collection_name, overwrite=True)
  685. return {
  686. "status": True,
  687. "collection_name": collection_name,
  688. "filename": form_data.url,
  689. }
  690. except Exception as e:
  691. log.exception(e)
  692. raise HTTPException(
  693. status_code=status.HTTP_400_BAD_REQUEST,
  694. detail=ERROR_MESSAGES.DEFAULT(e),
  695. )
  696. @app.post("/web")
  697. def store_web(form_data: UrlForm, user=Depends(get_verified_user)):
  698. # "https://www.gutenberg.org/files/1727/1727-h/1727-h.htm"
  699. try:
  700. loader = get_web_loader(
  701. form_data.url,
  702. verify_ssl=app.state.config.ENABLE_RAG_WEB_LOADER_SSL_VERIFICATION,
  703. )
  704. data = loader.load()
  705. collection_name = form_data.collection_name
  706. if collection_name == "":
  707. collection_name = calculate_sha256_string(form_data.url)[:63]
  708. store_data_in_vector_db(data, collection_name, overwrite=True)
  709. return {
  710. "status": True,
  711. "collection_name": collection_name,
  712. "filename": form_data.url,
  713. }
  714. except Exception as e:
  715. log.exception(e)
  716. raise HTTPException(
  717. status_code=status.HTTP_400_BAD_REQUEST,
  718. detail=ERROR_MESSAGES.DEFAULT(e),
  719. )
  720. def get_web_loader(url: Union[str, Sequence[str]], verify_ssl: bool = True):
  721. # Check if the URL is valid
  722. if not validate_url(url):
  723. raise ValueError(ERROR_MESSAGES.INVALID_URL)
  724. return SafeWebBaseLoader(
  725. url,
  726. verify_ssl=verify_ssl,
  727. requests_per_second=RAG_WEB_SEARCH_CONCURRENT_REQUESTS,
  728. continue_on_failure=True,
  729. )
  730. def validate_url(url: Union[str, Sequence[str]]):
  731. if isinstance(url, str):
  732. if isinstance(validators.url(url), validators.ValidationError):
  733. raise ValueError(ERROR_MESSAGES.INVALID_URL)
  734. if not ENABLE_RAG_LOCAL_WEB_FETCH:
  735. # Local web fetch is disabled, filter out any URLs that resolve to private IP addresses
  736. parsed_url = urllib.parse.urlparse(url)
  737. # Get IPv4 and IPv6 addresses
  738. ipv4_addresses, ipv6_addresses = resolve_hostname(parsed_url.hostname)
  739. # Check if any of the resolved addresses are private
  740. # This is technically still vulnerable to DNS rebinding attacks, as we don't control WebBaseLoader
  741. for ip in ipv4_addresses:
  742. if validators.ipv4(ip, private=True):
  743. raise ValueError(ERROR_MESSAGES.INVALID_URL)
  744. for ip in ipv6_addresses:
  745. if validators.ipv6(ip, private=True):
  746. raise ValueError(ERROR_MESSAGES.INVALID_URL)
  747. return True
  748. elif isinstance(url, Sequence):
  749. return all(validate_url(u) for u in url)
  750. else:
  751. return False
  752. def resolve_hostname(hostname):
  753. # Get address information
  754. addr_info = socket.getaddrinfo(hostname, None)
  755. # Extract IP addresses from address information
  756. ipv4_addresses = [info[4][0] for info in addr_info if info[0] == socket.AF_INET]
  757. ipv6_addresses = [info[4][0] for info in addr_info if info[0] == socket.AF_INET6]
  758. return ipv4_addresses, ipv6_addresses
  759. def search_web(engine: str, query: str) -> list[SearchResult]:
  760. """Search the web using a search engine and return the results as a list of SearchResult objects.
  761. Will look for a search engine API key in environment variables in the following order:
  762. - SEARXNG_QUERY_URL
  763. - GOOGLE_PSE_API_KEY + GOOGLE_PSE_ENGINE_ID
  764. - BRAVE_SEARCH_API_KEY
  765. - SERPSTACK_API_KEY
  766. - SERPER_API_KEY
  767. - SERPLY_API_KEY
  768. - TAVILY_API_KEY
  769. - SEARCHAPI_API_KEY + SEARCHAPI_ENGINE (by default `google`)
  770. Args:
  771. query (str): The query to search for
  772. """
  773. # TODO: add playwright to search the web
  774. if engine == "searxng":
  775. if app.state.config.SEARXNG_QUERY_URL:
  776. return search_searxng(
  777. app.state.config.SEARXNG_QUERY_URL,
  778. query,
  779. app.state.config.RAG_WEB_SEARCH_RESULT_COUNT,
  780. app.state.config.RAG_WEB_SEARCH_DOMAIN_FILTER_LIST,
  781. )
  782. else:
  783. raise Exception("No SEARXNG_QUERY_URL found in environment variables")
  784. elif engine == "google_pse":
  785. if (
  786. app.state.config.GOOGLE_PSE_API_KEY
  787. and app.state.config.GOOGLE_PSE_ENGINE_ID
  788. ):
  789. return search_google_pse(
  790. app.state.config.GOOGLE_PSE_API_KEY,
  791. app.state.config.GOOGLE_PSE_ENGINE_ID,
  792. query,
  793. app.state.config.RAG_WEB_SEARCH_RESULT_COUNT,
  794. app.state.config.RAG_WEB_SEARCH_DOMAIN_FILTER_LIST,
  795. )
  796. else:
  797. raise Exception(
  798. "No GOOGLE_PSE_API_KEY or GOOGLE_PSE_ENGINE_ID found in environment variables"
  799. )
  800. elif engine == "brave":
  801. if app.state.config.BRAVE_SEARCH_API_KEY:
  802. return search_brave(
  803. app.state.config.BRAVE_SEARCH_API_KEY,
  804. query,
  805. app.state.config.RAG_WEB_SEARCH_RESULT_COUNT,
  806. app.state.config.RAG_WEB_SEARCH_DOMAIN_FILTER_LIST,
  807. )
  808. else:
  809. raise Exception("No BRAVE_SEARCH_API_KEY found in environment variables")
  810. elif engine == "serpstack":
  811. if app.state.config.SERPSTACK_API_KEY:
  812. return search_serpstack(
  813. app.state.config.SERPSTACK_API_KEY,
  814. query,
  815. app.state.config.RAG_WEB_SEARCH_RESULT_COUNT,
  816. app.state.config.RAG_WEB_SEARCH_DOMAIN_FILTER_LIST,
  817. https_enabled=app.state.config.SERPSTACK_HTTPS,
  818. )
  819. else:
  820. raise Exception("No SERPSTACK_API_KEY found in environment variables")
  821. elif engine == "serper":
  822. if app.state.config.SERPER_API_KEY:
  823. return search_serper(
  824. app.state.config.SERPER_API_KEY,
  825. query,
  826. app.state.config.RAG_WEB_SEARCH_RESULT_COUNT,
  827. app.state.config.RAG_WEB_SEARCH_DOMAIN_FILTER_LIST,
  828. )
  829. else:
  830. raise Exception("No SERPER_API_KEY found in environment variables")
  831. elif engine == "serply":
  832. if app.state.config.SERPLY_API_KEY:
  833. return search_serply(
  834. app.state.config.SERPLY_API_KEY,
  835. query,
  836. app.state.config.RAG_WEB_SEARCH_RESULT_COUNT,
  837. app.state.config.RAG_WEB_SEARCH_DOMAIN_FILTER_LIST,
  838. )
  839. else:
  840. raise Exception("No SERPLY_API_KEY found in environment variables")
  841. elif engine == "duckduckgo":
  842. return search_duckduckgo(
  843. query,
  844. app.state.config.RAG_WEB_SEARCH_RESULT_COUNT,
  845. app.state.config.RAG_WEB_SEARCH_DOMAIN_FILTER_LIST,
  846. )
  847. elif engine == "tavily":
  848. if app.state.config.TAVILY_API_KEY:
  849. return search_tavily(
  850. app.state.config.TAVILY_API_KEY,
  851. query,
  852. app.state.config.RAG_WEB_SEARCH_RESULT_COUNT,
  853. )
  854. else:
  855. raise Exception("No TAVILY_API_KEY found in environment variables")
  856. elif engine == "searchapi":
  857. if app.state.config.SEARCHAPI_API_KEY:
  858. return search_searchapi(
  859. app.state.config.SEARCHAPI_API_KEY,
  860. app.state.config.SEARCHAPI_ENGINE,
  861. query,
  862. app.state.config.RAG_WEB_SEARCH_RESULT_COUNT,
  863. app.state.config.RAG_WEB_SEARCH_DOMAIN_FILTER_LIST,
  864. )
  865. else:
  866. raise Exception("No SEARCHAPI_API_KEY found in environment variables")
  867. elif engine == "jina":
  868. return search_jina(query, app.state.config.RAG_WEB_SEARCH_RESULT_COUNT)
  869. else:
  870. raise Exception("No search engine API key found in environment variables")
  871. @app.post("/web/search")
  872. def store_web_search(form_data: SearchForm, user=Depends(get_verified_user)):
  873. try:
  874. logging.info(
  875. f"trying to web search with {app.state.config.RAG_WEB_SEARCH_ENGINE, form_data.query}"
  876. )
  877. web_results = search_web(
  878. app.state.config.RAG_WEB_SEARCH_ENGINE, form_data.query
  879. )
  880. except Exception as e:
  881. log.exception(e)
  882. print(e)
  883. raise HTTPException(
  884. status_code=status.HTTP_400_BAD_REQUEST,
  885. detail=ERROR_MESSAGES.WEB_SEARCH_ERROR(e),
  886. )
  887. try:
  888. urls = [result.link for result in web_results]
  889. loader = get_web_loader(urls)
  890. data = loader.load()
  891. collection_name = form_data.collection_name
  892. if collection_name == "":
  893. collection_name = calculate_sha256_string(form_data.query)[:63]
  894. store_data_in_vector_db(data, collection_name, overwrite=True)
  895. return {
  896. "status": True,
  897. "collection_name": collection_name,
  898. "filenames": urls,
  899. }
  900. except Exception as e:
  901. log.exception(e)
  902. raise HTTPException(
  903. status_code=status.HTTP_400_BAD_REQUEST,
  904. detail=ERROR_MESSAGES.DEFAULT(e),
  905. )
  906. def store_data_in_vector_db(
  907. data, collection_name, metadata: Optional[dict] = None, overwrite: bool = False
  908. ) -> bool:
  909. text_splitter = RecursiveCharacterTextSplitter(
  910. chunk_size=app.state.config.CHUNK_SIZE,
  911. chunk_overlap=app.state.config.CHUNK_OVERLAP,
  912. add_start_index=True,
  913. )
  914. docs = text_splitter.split_documents(data)
  915. if len(docs) > 0:
  916. log.info(f"store_data_in_vector_db {docs}")
  917. return store_docs_in_vector_db(docs, collection_name, metadata, overwrite), None
  918. else:
  919. raise ValueError(ERROR_MESSAGES.EMPTY_CONTENT)
  920. def store_text_in_vector_db(
  921. text, metadata, collection_name, overwrite: bool = False
  922. ) -> bool:
  923. text_splitter = RecursiveCharacterTextSplitter(
  924. chunk_size=app.state.config.CHUNK_SIZE,
  925. chunk_overlap=app.state.config.CHUNK_OVERLAP,
  926. add_start_index=True,
  927. )
  928. docs = text_splitter.create_documents([text], metadatas=[metadata])
  929. return store_docs_in_vector_db(docs, collection_name, overwrite=overwrite)
  930. def store_docs_in_vector_db(
  931. docs, collection_name, metadata: Optional[dict] = None, overwrite: bool = False
  932. ) -> bool:
  933. log.info(f"store_docs_in_vector_db {docs} {collection_name}")
  934. texts = [doc.page_content for doc in docs]
  935. metadatas = [{**doc.metadata, **(metadata if metadata else {})} for doc in docs]
  936. # ChromaDB does not like datetime formats
  937. # for meta-data so convert them to string.
  938. for metadata in metadatas:
  939. for key, value in metadata.items():
  940. if isinstance(value, datetime):
  941. metadata[key] = str(value)
  942. try:
  943. if overwrite:
  944. if VECTOR_DB_CLIENT.has_collection(collection_name=collection_name):
  945. log.info(f"deleting existing collection {collection_name}")
  946. VECTOR_DB_CLIENT.delete_collection(collection_name=collection_name)
  947. embedding_function = get_embedding_function(
  948. app.state.config.RAG_EMBEDDING_ENGINE,
  949. app.state.config.RAG_EMBEDDING_MODEL,
  950. app.state.sentence_transformer_ef,
  951. app.state.config.OPENAI_API_KEY,
  952. app.state.config.OPENAI_API_BASE_URL,
  953. app.state.config.RAG_EMBEDDING_OPENAI_BATCH_SIZE,
  954. )
  955. VECTOR_DB_CLIENT.insert(
  956. collection_name=collection_name,
  957. items=[
  958. {
  959. "id": str(uuid.uuid4()),
  960. "text": text,
  961. "vector": embedding_function(text.replace("\n", " ")),
  962. "metadata": metadatas[idx],
  963. }
  964. for idx, text in enumerate(texts)
  965. ],
  966. )
  967. return True
  968. except Exception as e:
  969. if e.__class__.__name__ == "UniqueConstraintError":
  970. return True
  971. log.exception(e)
  972. return False
  973. class TikaLoader:
  974. def __init__(self, file_path, mime_type=None):
  975. self.file_path = file_path
  976. self.mime_type = mime_type
  977. def load(self) -> list[Document]:
  978. with open(self.file_path, "rb") as f:
  979. data = f.read()
  980. if self.mime_type is not None:
  981. headers = {"Content-Type": self.mime_type}
  982. else:
  983. headers = {}
  984. endpoint = app.state.config.TIKA_SERVER_URL
  985. if not endpoint.endswith("/"):
  986. endpoint += "/"
  987. endpoint += "tika/text"
  988. r = requests.put(endpoint, data=data, headers=headers)
  989. if r.ok:
  990. raw_metadata = r.json()
  991. text = raw_metadata.get("X-TIKA:content", "<No text content found>")
  992. if "Content-Type" in raw_metadata:
  993. headers["Content-Type"] = raw_metadata["Content-Type"]
  994. log.info("Tika extracted text: %s", text)
  995. return [Document(page_content=text, metadata=headers)]
  996. else:
  997. raise Exception(f"Error calling Tika: {r.reason}")
  998. def get_loader(filename: str, file_content_type: str, file_path: str):
  999. file_ext = filename.split(".")[-1].lower()
  1000. known_type = True
  1001. known_source_ext = [
  1002. "go",
  1003. "py",
  1004. "java",
  1005. "sh",
  1006. "bat",
  1007. "ps1",
  1008. "cmd",
  1009. "js",
  1010. "ts",
  1011. "css",
  1012. "cpp",
  1013. "hpp",
  1014. "h",
  1015. "c",
  1016. "cs",
  1017. "sql",
  1018. "log",
  1019. "ini",
  1020. "pl",
  1021. "pm",
  1022. "r",
  1023. "dart",
  1024. "dockerfile",
  1025. "env",
  1026. "php",
  1027. "hs",
  1028. "hsc",
  1029. "lua",
  1030. "nginxconf",
  1031. "conf",
  1032. "m",
  1033. "mm",
  1034. "plsql",
  1035. "perl",
  1036. "rb",
  1037. "rs",
  1038. "db2",
  1039. "scala",
  1040. "bash",
  1041. "swift",
  1042. "vue",
  1043. "svelte",
  1044. "msg",
  1045. "ex",
  1046. "exs",
  1047. "erl",
  1048. "tsx",
  1049. "jsx",
  1050. "hs",
  1051. "lhs",
  1052. ]
  1053. if (
  1054. app.state.config.CONTENT_EXTRACTION_ENGINE == "tika"
  1055. and app.state.config.TIKA_SERVER_URL
  1056. ):
  1057. if file_ext in known_source_ext or (
  1058. file_content_type and file_content_type.find("text/") >= 0
  1059. ):
  1060. loader = TextLoader(file_path, autodetect_encoding=True)
  1061. else:
  1062. loader = TikaLoader(file_path, file_content_type)
  1063. else:
  1064. if file_ext == "pdf":
  1065. loader = PyPDFLoader(
  1066. file_path, extract_images=app.state.config.PDF_EXTRACT_IMAGES
  1067. )
  1068. elif file_ext == "csv":
  1069. loader = CSVLoader(file_path)
  1070. elif file_ext == "rst":
  1071. loader = UnstructuredRSTLoader(file_path, mode="elements")
  1072. elif file_ext == "xml":
  1073. loader = UnstructuredXMLLoader(file_path)
  1074. elif file_ext in ["htm", "html"]:
  1075. loader = BSHTMLLoader(file_path, open_encoding="unicode_escape")
  1076. elif file_ext == "md":
  1077. loader = UnstructuredMarkdownLoader(file_path)
  1078. elif file_content_type == "application/epub+zip":
  1079. loader = UnstructuredEPubLoader(file_path)
  1080. elif (
  1081. file_content_type
  1082. == "application/vnd.openxmlformats-officedocument.wordprocessingml.document"
  1083. or file_ext == "docx"
  1084. ):
  1085. loader = Docx2txtLoader(file_path)
  1086. elif file_content_type in [
  1087. "application/vnd.ms-excel",
  1088. "application/vnd.openxmlformats-officedocument.spreadsheetml.sheet",
  1089. ] or file_ext in ["xls", "xlsx"]:
  1090. loader = UnstructuredExcelLoader(file_path)
  1091. elif file_content_type in [
  1092. "application/vnd.ms-powerpoint",
  1093. "application/vnd.openxmlformats-officedocument.presentationml.presentation",
  1094. ] or file_ext in ["ppt", "pptx"]:
  1095. loader = UnstructuredPowerPointLoader(file_path)
  1096. elif file_ext == "msg":
  1097. loader = OutlookMessageLoader(file_path)
  1098. elif file_ext in known_source_ext or (
  1099. file_content_type and file_content_type.find("text/") >= 0
  1100. ):
  1101. loader = TextLoader(file_path, autodetect_encoding=True)
  1102. else:
  1103. loader = TextLoader(file_path, autodetect_encoding=True)
  1104. known_type = False
  1105. return loader, known_type
  1106. @app.post("/doc")
  1107. def store_doc(
  1108. collection_name: Optional[str] = Form(None),
  1109. file: UploadFile = File(...),
  1110. user=Depends(get_verified_user),
  1111. ):
  1112. # "https://www.gutenberg.org/files/1727/1727-h/1727-h.htm"
  1113. log.info(f"file.content_type: {file.content_type}")
  1114. try:
  1115. unsanitized_filename = file.filename
  1116. filename = os.path.basename(unsanitized_filename)
  1117. file_path = f"{UPLOAD_DIR}/{filename}"
  1118. contents = file.file.read()
  1119. with open(file_path, "wb") as f:
  1120. f.write(contents)
  1121. f.close()
  1122. f = open(file_path, "rb")
  1123. if collection_name is None:
  1124. collection_name = calculate_sha256(f)[:63]
  1125. f.close()
  1126. loader, known_type = get_loader(filename, file.content_type, file_path)
  1127. data = loader.load()
  1128. try:
  1129. result = store_data_in_vector_db(data, collection_name)
  1130. if result:
  1131. return {
  1132. "status": True,
  1133. "collection_name": collection_name,
  1134. "filename": filename,
  1135. "known_type": known_type,
  1136. }
  1137. except Exception as e:
  1138. raise HTTPException(
  1139. status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
  1140. detail=e,
  1141. )
  1142. except Exception as e:
  1143. log.exception(e)
  1144. if "No pandoc was found" in str(e):
  1145. raise HTTPException(
  1146. status_code=status.HTTP_400_BAD_REQUEST,
  1147. detail=ERROR_MESSAGES.PANDOC_NOT_INSTALLED,
  1148. )
  1149. else:
  1150. raise HTTPException(
  1151. status_code=status.HTTP_400_BAD_REQUEST,
  1152. detail=ERROR_MESSAGES.DEFAULT(e),
  1153. )
  1154. class ProcessDocForm(BaseModel):
  1155. file_id: str
  1156. collection_name: Optional[str] = None
  1157. @app.post("/process/doc")
  1158. def process_doc(
  1159. form_data: ProcessDocForm,
  1160. user=Depends(get_verified_user),
  1161. ):
  1162. try:
  1163. file = Files.get_file_by_id(form_data.file_id)
  1164. file_path = file.meta.get("path", f"{UPLOAD_DIR}/{file.filename}")
  1165. f = open(file_path, "rb")
  1166. collection_name = form_data.collection_name
  1167. if collection_name is None:
  1168. collection_name = calculate_sha256(f)[:63]
  1169. f.close()
  1170. loader, known_type = get_loader(
  1171. file.filename, file.meta.get("content_type"), file_path
  1172. )
  1173. data = loader.load()
  1174. try:
  1175. result = store_data_in_vector_db(
  1176. data,
  1177. collection_name,
  1178. {
  1179. "file_id": form_data.file_id,
  1180. "name": file.meta.get("name", file.filename),
  1181. },
  1182. )
  1183. if result:
  1184. return {
  1185. "status": True,
  1186. "collection_name": collection_name,
  1187. "known_type": known_type,
  1188. "filename": file.meta.get("name", file.filename),
  1189. }
  1190. except Exception as e:
  1191. raise HTTPException(
  1192. status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
  1193. detail=e,
  1194. )
  1195. except Exception as e:
  1196. log.exception(e)
  1197. if "No pandoc was found" in str(e):
  1198. raise HTTPException(
  1199. status_code=status.HTTP_400_BAD_REQUEST,
  1200. detail=ERROR_MESSAGES.PANDOC_NOT_INSTALLED,
  1201. )
  1202. else:
  1203. raise HTTPException(
  1204. status_code=status.HTTP_400_BAD_REQUEST,
  1205. detail=ERROR_MESSAGES.DEFAULT(e),
  1206. )
  1207. class TextRAGForm(BaseModel):
  1208. name: str
  1209. content: str
  1210. collection_name: Optional[str] = None
  1211. @app.post("/text")
  1212. def store_text(
  1213. form_data: TextRAGForm,
  1214. user=Depends(get_verified_user),
  1215. ):
  1216. collection_name = form_data.collection_name
  1217. if collection_name is None:
  1218. collection_name = calculate_sha256_string(form_data.content)
  1219. result = store_text_in_vector_db(
  1220. form_data.content,
  1221. metadata={"name": form_data.name, "created_by": user.id},
  1222. collection_name=collection_name,
  1223. )
  1224. if result:
  1225. return {"status": True, "collection_name": collection_name}
  1226. else:
  1227. raise HTTPException(
  1228. status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
  1229. detail=ERROR_MESSAGES.DEFAULT(),
  1230. )
  1231. @app.get("/scan")
  1232. def scan_docs_dir(user=Depends(get_admin_user)):
  1233. for path in Path(DOCS_DIR).rglob("./**/*"):
  1234. try:
  1235. if path.is_file() and not path.name.startswith("."):
  1236. tags = extract_folders_after_data_docs(path)
  1237. filename = path.name
  1238. file_content_type = mimetypes.guess_type(path)
  1239. f = open(path, "rb")
  1240. collection_name = calculate_sha256(f)[:63]
  1241. f.close()
  1242. loader, known_type = get_loader(
  1243. filename, file_content_type[0], str(path)
  1244. )
  1245. data = loader.load()
  1246. try:
  1247. result = store_data_in_vector_db(data, collection_name)
  1248. if result:
  1249. sanitized_filename = sanitize_filename(filename)
  1250. doc = Documents.get_doc_by_name(sanitized_filename)
  1251. if doc is None:
  1252. doc = Documents.insert_new_doc(
  1253. user.id,
  1254. DocumentForm(
  1255. **{
  1256. "name": sanitized_filename,
  1257. "title": filename,
  1258. "collection_name": collection_name,
  1259. "filename": filename,
  1260. "content": (
  1261. json.dumps(
  1262. {
  1263. "tags": list(
  1264. map(
  1265. lambda name: {"name": name},
  1266. tags,
  1267. )
  1268. )
  1269. }
  1270. )
  1271. if len(tags)
  1272. else "{}"
  1273. ),
  1274. }
  1275. ),
  1276. )
  1277. except Exception as e:
  1278. log.exception(e)
  1279. pass
  1280. except Exception as e:
  1281. log.exception(e)
  1282. return True
  1283. @app.post("/reset/db")
  1284. def reset_vector_db(user=Depends(get_admin_user)):
  1285. VECTOR_DB_CLIENT.reset()
  1286. @app.post("/reset/uploads")
  1287. def reset_upload_dir(user=Depends(get_admin_user)) -> bool:
  1288. folder = f"{UPLOAD_DIR}"
  1289. try:
  1290. # Check if the directory exists
  1291. if os.path.exists(folder):
  1292. # Iterate over all the files and directories in the specified directory
  1293. for filename in os.listdir(folder):
  1294. file_path = os.path.join(folder, filename)
  1295. try:
  1296. if os.path.isfile(file_path) or os.path.islink(file_path):
  1297. os.unlink(file_path) # Remove the file or link
  1298. elif os.path.isdir(file_path):
  1299. shutil.rmtree(file_path) # Remove the directory
  1300. except Exception as e:
  1301. print(f"Failed to delete {file_path}. Reason: {e}")
  1302. else:
  1303. print(f"The directory {folder} does not exist")
  1304. except Exception as e:
  1305. print(f"Failed to process the directory {folder}. Reason: {e}")
  1306. return True
  1307. @app.post("/reset")
  1308. def reset(user=Depends(get_admin_user)) -> bool:
  1309. folder = f"{UPLOAD_DIR}"
  1310. for filename in os.listdir(folder):
  1311. file_path = os.path.join(folder, filename)
  1312. try:
  1313. if os.path.isfile(file_path) or os.path.islink(file_path):
  1314. os.unlink(file_path)
  1315. elif os.path.isdir(file_path):
  1316. shutil.rmtree(file_path)
  1317. except Exception as e:
  1318. log.error("Failed to delete %s. Reason: %s" % (file_path, e))
  1319. try:
  1320. VECTOR_DB_CLIENT.reset()
  1321. except Exception as e:
  1322. log.exception(e)
  1323. return True
  1324. class SafeWebBaseLoader(WebBaseLoader):
  1325. """WebBaseLoader with enhanced error handling for URLs."""
  1326. def lazy_load(self) -> Iterator[Document]:
  1327. """Lazy load text from the url(s) in web_path with error handling."""
  1328. for path in self.web_paths:
  1329. try:
  1330. soup = self._scrape(path, bs_kwargs=self.bs_kwargs)
  1331. text = soup.get_text(**self.bs_get_text_kwargs)
  1332. # Build metadata
  1333. metadata = {"source": path}
  1334. if title := soup.find("title"):
  1335. metadata["title"] = title.get_text()
  1336. if description := soup.find("meta", attrs={"name": "description"}):
  1337. metadata["description"] = description.get(
  1338. "content", "No description found."
  1339. )
  1340. if html := soup.find("html"):
  1341. metadata["language"] = html.get("lang", "No language found.")
  1342. yield Document(page_content=text, metadata=metadata)
  1343. except Exception as e:
  1344. # Log the error and continue with the next URL
  1345. log.error(f"Error loading {path}: {e}")
  1346. if ENV == "dev":
  1347. @app.get("/ef")
  1348. async def get_embeddings():
  1349. return {"result": app.state.EMBEDDING_FUNCTION("hello world")}
  1350. @app.get("/ef/{text}")
  1351. async def get_embeddings_text(text: str):
  1352. return {"result": app.state.EMBEDDING_FUNCTION(text)}