main.py 49 KB

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