main.py 27 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599600601602603604605606607608609610611612613614615616617618619620621622623624625626627628629630631632633634635636637638639640641642643644645646647648649650651652653654655656657658659660661662663664665666667668669670671672673674675676677678679680681682683684685686687688689690691692693694695696697698699700701702703704705706707708709710711712713714715716717718719720721722723724725726727728729730731732733734735736737738739740741742743744745746747748749750751752753754755756757758759760761762763764765766767768769770771772773774775776777778779780781782783784785786787788789790791792793794795796797798799800801802803804805806807808809810811812813814815816817818819820821822823824825826827828829830831832833834835836837838839840841842843844845846847848849850851852853854855856857858859860861862863864865866867868869870871872873874875876877878879880881882883884885886887888889890891892893894895896897898899900901902903904905906907908909910911912
  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 os, shutil, logging, re
  12. from pathlib import Path
  13. from typing import List, Union, Sequence
  14. from chromadb.utils.batch_utils import create_batches
  15. from langchain_community.document_loaders import (
  16. WebBaseLoader,
  17. TextLoader,
  18. PyPDFLoader,
  19. CSVLoader,
  20. BSHTMLLoader,
  21. Docx2txtLoader,
  22. UnstructuredEPubLoader,
  23. UnstructuredWordDocumentLoader,
  24. UnstructuredMarkdownLoader,
  25. UnstructuredXMLLoader,
  26. UnstructuredRSTLoader,
  27. UnstructuredExcelLoader,
  28. YoutubeLoader,
  29. )
  30. from langchain.text_splitter import RecursiveCharacterTextSplitter
  31. import validators
  32. import urllib.parse
  33. import socket
  34. from pydantic import BaseModel
  35. from typing import Optional
  36. import mimetypes
  37. import uuid
  38. import json
  39. import sentence_transformers
  40. from apps.web.models.documents import (
  41. Documents,
  42. DocumentForm,
  43. DocumentResponse,
  44. )
  45. from apps.rag.utils import (
  46. get_model_path,
  47. get_embedding_function,
  48. query_doc,
  49. query_doc_with_hybrid_search,
  50. query_collection,
  51. query_collection_with_hybrid_search,
  52. search_web,
  53. )
  54. from utils.misc import (
  55. calculate_sha256,
  56. calculate_sha256_string,
  57. sanitize_filename,
  58. extract_folders_after_data_docs,
  59. )
  60. from utils.utils import get_current_user, get_admin_user
  61. from config import (
  62. SRC_LOG_LEVELS,
  63. UPLOAD_DIR,
  64. DOCS_DIR,
  65. RAG_TOP_K,
  66. RAG_RELEVANCE_THRESHOLD,
  67. RAG_EMBEDDING_ENGINE,
  68. RAG_EMBEDDING_MODEL,
  69. RAG_EMBEDDING_MODEL_AUTO_UPDATE,
  70. RAG_EMBEDDING_MODEL_TRUST_REMOTE_CODE,
  71. ENABLE_RAG_HYBRID_SEARCH,
  72. RAG_RERANKING_MODEL,
  73. PDF_EXTRACT_IMAGES,
  74. RAG_RERANKING_MODEL_AUTO_UPDATE,
  75. RAG_RERANKING_MODEL_TRUST_REMOTE_CODE,
  76. RAG_OPENAI_API_BASE_URL,
  77. RAG_OPENAI_API_KEY,
  78. DEVICE_TYPE,
  79. CHROMA_CLIENT,
  80. CHUNK_SIZE,
  81. CHUNK_OVERLAP,
  82. RAG_TEMPLATE,
  83. ENABLE_LOCAL_WEB_FETCH,
  84. )
  85. from constants import ERROR_MESSAGES
  86. log = logging.getLogger(__name__)
  87. log.setLevel(SRC_LOG_LEVELS["RAG"])
  88. app = FastAPI()
  89. app.state.TOP_K = RAG_TOP_K
  90. app.state.RELEVANCE_THRESHOLD = RAG_RELEVANCE_THRESHOLD
  91. app.state.ENABLE_RAG_HYBRID_SEARCH = ENABLE_RAG_HYBRID_SEARCH
  92. app.state.CHUNK_SIZE = CHUNK_SIZE
  93. app.state.CHUNK_OVERLAP = CHUNK_OVERLAP
  94. app.state.RAG_EMBEDDING_ENGINE = RAG_EMBEDDING_ENGINE
  95. app.state.RAG_EMBEDDING_MODEL = RAG_EMBEDDING_MODEL
  96. app.state.RAG_RERANKING_MODEL = RAG_RERANKING_MODEL
  97. app.state.RAG_TEMPLATE = RAG_TEMPLATE
  98. app.state.OPENAI_API_BASE_URL = RAG_OPENAI_API_BASE_URL
  99. app.state.OPENAI_API_KEY = RAG_OPENAI_API_KEY
  100. app.state.PDF_EXTRACT_IMAGES = PDF_EXTRACT_IMAGES
  101. def update_embedding_model(
  102. embedding_model: str,
  103. update_model: bool = False,
  104. ):
  105. if embedding_model and app.state.RAG_EMBEDDING_ENGINE == "":
  106. app.state.sentence_transformer_ef = sentence_transformers.SentenceTransformer(
  107. get_model_path(embedding_model, update_model),
  108. device=DEVICE_TYPE,
  109. trust_remote_code=RAG_EMBEDDING_MODEL_TRUST_REMOTE_CODE,
  110. )
  111. else:
  112. app.state.sentence_transformer_ef = None
  113. def update_reranking_model(
  114. reranking_model: str,
  115. update_model: bool = False,
  116. ):
  117. if reranking_model:
  118. app.state.sentence_transformer_rf = sentence_transformers.CrossEncoder(
  119. get_model_path(reranking_model, update_model),
  120. device=DEVICE_TYPE,
  121. trust_remote_code=RAG_RERANKING_MODEL_TRUST_REMOTE_CODE,
  122. )
  123. else:
  124. app.state.sentence_transformer_rf = None
  125. update_embedding_model(
  126. app.state.RAG_EMBEDDING_MODEL,
  127. RAG_EMBEDDING_MODEL_AUTO_UPDATE,
  128. )
  129. update_reranking_model(
  130. app.state.RAG_RERANKING_MODEL,
  131. RAG_RERANKING_MODEL_AUTO_UPDATE,
  132. )
  133. app.state.EMBEDDING_FUNCTION = get_embedding_function(
  134. app.state.RAG_EMBEDDING_ENGINE,
  135. app.state.RAG_EMBEDDING_MODEL,
  136. app.state.sentence_transformer_ef,
  137. app.state.OPENAI_API_KEY,
  138. app.state.OPENAI_API_BASE_URL,
  139. )
  140. origins = ["*"]
  141. app.add_middleware(
  142. CORSMiddleware,
  143. allow_origins=origins,
  144. allow_credentials=True,
  145. allow_methods=["*"],
  146. allow_headers=["*"],
  147. )
  148. class CollectionNameForm(BaseModel):
  149. collection_name: Optional[str] = "test"
  150. class UrlForm(CollectionNameForm):
  151. url: str
  152. class SearchForm(CollectionNameForm):
  153. query: str
  154. @app.get("/")
  155. async def get_status():
  156. return {
  157. "status": True,
  158. "chunk_size": app.state.CHUNK_SIZE,
  159. "chunk_overlap": app.state.CHUNK_OVERLAP,
  160. "template": app.state.RAG_TEMPLATE,
  161. "embedding_engine": app.state.RAG_EMBEDDING_ENGINE,
  162. "embedding_model": app.state.RAG_EMBEDDING_MODEL,
  163. "reranking_model": app.state.RAG_RERANKING_MODEL,
  164. }
  165. @app.get("/embedding")
  166. async def get_embedding_config(user=Depends(get_admin_user)):
  167. return {
  168. "status": True,
  169. "embedding_engine": app.state.RAG_EMBEDDING_ENGINE,
  170. "embedding_model": app.state.RAG_EMBEDDING_MODEL,
  171. "openai_config": {
  172. "url": app.state.OPENAI_API_BASE_URL,
  173. "key": app.state.OPENAI_API_KEY,
  174. },
  175. }
  176. @app.get("/reranking")
  177. async def get_reraanking_config(user=Depends(get_admin_user)):
  178. return {"status": True, "reranking_model": app.state.RAG_RERANKING_MODEL}
  179. class OpenAIConfigForm(BaseModel):
  180. url: str
  181. key: str
  182. class EmbeddingModelUpdateForm(BaseModel):
  183. openai_config: Optional[OpenAIConfigForm] = None
  184. embedding_engine: str
  185. embedding_model: str
  186. @app.post("/embedding/update")
  187. async def update_embedding_config(
  188. form_data: EmbeddingModelUpdateForm, user=Depends(get_admin_user)
  189. ):
  190. log.info(
  191. f"Updating embedding model: {app.state.RAG_EMBEDDING_MODEL} to {form_data.embedding_model}"
  192. )
  193. try:
  194. app.state.RAG_EMBEDDING_ENGINE = form_data.embedding_engine
  195. app.state.RAG_EMBEDDING_MODEL = form_data.embedding_model
  196. if app.state.RAG_EMBEDDING_ENGINE in ["ollama", "openai"]:
  197. if form_data.openai_config != None:
  198. app.state.OPENAI_API_BASE_URL = form_data.openai_config.url
  199. app.state.OPENAI_API_KEY = form_data.openai_config.key
  200. update_embedding_model(app.state.RAG_EMBEDDING_MODEL, True)
  201. app.state.EMBEDDING_FUNCTION = get_embedding_function(
  202. app.state.RAG_EMBEDDING_ENGINE,
  203. app.state.RAG_EMBEDDING_MODEL,
  204. app.state.sentence_transformer_ef,
  205. app.state.OPENAI_API_KEY,
  206. app.state.OPENAI_API_BASE_URL,
  207. )
  208. return {
  209. "status": True,
  210. "embedding_engine": app.state.RAG_EMBEDDING_ENGINE,
  211. "embedding_model": app.state.RAG_EMBEDDING_MODEL,
  212. "openai_config": {
  213. "url": app.state.OPENAI_API_BASE_URL,
  214. "key": app.state.OPENAI_API_KEY,
  215. },
  216. }
  217. except Exception as e:
  218. log.exception(f"Problem updating embedding model: {e}")
  219. raise HTTPException(
  220. status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
  221. detail=ERROR_MESSAGES.DEFAULT(e),
  222. )
  223. class RerankingModelUpdateForm(BaseModel):
  224. reranking_model: str
  225. @app.post("/reranking/update")
  226. async def update_reranking_config(
  227. form_data: RerankingModelUpdateForm, user=Depends(get_admin_user)
  228. ):
  229. log.info(
  230. f"Updating reranking model: {app.state.RAG_RERANKING_MODEL} to {form_data.reranking_model}"
  231. )
  232. try:
  233. app.state.RAG_RERANKING_MODEL = form_data.reranking_model
  234. update_reranking_model(app.state.RAG_RERANKING_MODEL, True)
  235. return {
  236. "status": True,
  237. "reranking_model": app.state.RAG_RERANKING_MODEL,
  238. }
  239. except Exception as e:
  240. log.exception(f"Problem updating reranking model: {e}")
  241. raise HTTPException(
  242. status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
  243. detail=ERROR_MESSAGES.DEFAULT(e),
  244. )
  245. @app.get("/config")
  246. async def get_rag_config(user=Depends(get_admin_user)):
  247. return {
  248. "status": True,
  249. "pdf_extract_images": app.state.PDF_EXTRACT_IMAGES,
  250. "chunk": {
  251. "chunk_size": app.state.CHUNK_SIZE,
  252. "chunk_overlap": app.state.CHUNK_OVERLAP,
  253. },
  254. }
  255. class ChunkParamUpdateForm(BaseModel):
  256. chunk_size: int
  257. chunk_overlap: int
  258. class ConfigUpdateForm(BaseModel):
  259. pdf_extract_images: bool
  260. chunk: ChunkParamUpdateForm
  261. @app.post("/config/update")
  262. async def update_rag_config(form_data: ConfigUpdateForm, user=Depends(get_admin_user)):
  263. app.state.PDF_EXTRACT_IMAGES = form_data.pdf_extract_images
  264. app.state.CHUNK_SIZE = form_data.chunk.chunk_size
  265. app.state.CHUNK_OVERLAP = form_data.chunk.chunk_overlap
  266. return {
  267. "status": True,
  268. "pdf_extract_images": app.state.PDF_EXTRACT_IMAGES,
  269. "chunk": {
  270. "chunk_size": app.state.CHUNK_SIZE,
  271. "chunk_overlap": app.state.CHUNK_OVERLAP,
  272. },
  273. }
  274. @app.get("/template")
  275. async def get_rag_template(user=Depends(get_current_user)):
  276. return {
  277. "status": True,
  278. "template": app.state.RAG_TEMPLATE,
  279. }
  280. @app.get("/query/settings")
  281. async def get_query_settings(user=Depends(get_admin_user)):
  282. return {
  283. "status": True,
  284. "template": app.state.RAG_TEMPLATE,
  285. "k": app.state.TOP_K,
  286. "r": app.state.RELEVANCE_THRESHOLD,
  287. "hybrid": app.state.ENABLE_RAG_HYBRID_SEARCH,
  288. }
  289. class QuerySettingsForm(BaseModel):
  290. k: Optional[int] = None
  291. r: Optional[float] = None
  292. template: Optional[str] = None
  293. hybrid: Optional[bool] = None
  294. @app.post("/query/settings/update")
  295. async def update_query_settings(
  296. form_data: QuerySettingsForm, user=Depends(get_admin_user)
  297. ):
  298. app.state.RAG_TEMPLATE = form_data.template if form_data.template else RAG_TEMPLATE
  299. app.state.TOP_K = form_data.k if form_data.k else 4
  300. app.state.RELEVANCE_THRESHOLD = form_data.r if form_data.r else 0.0
  301. app.state.ENABLE_RAG_HYBRID_SEARCH = form_data.hybrid if form_data.hybrid else False
  302. return {
  303. "status": True,
  304. "template": app.state.RAG_TEMPLATE,
  305. "k": app.state.TOP_K,
  306. "r": app.state.RELEVANCE_THRESHOLD,
  307. "hybrid": app.state.ENABLE_RAG_HYBRID_SEARCH,
  308. }
  309. class QueryDocForm(BaseModel):
  310. collection_name: str
  311. query: str
  312. k: Optional[int] = None
  313. r: Optional[float] = None
  314. hybrid: Optional[bool] = None
  315. @app.post("/query/doc")
  316. def query_doc_handler(
  317. form_data: QueryDocForm,
  318. user=Depends(get_current_user),
  319. ):
  320. try:
  321. if app.state.ENABLE_RAG_HYBRID_SEARCH:
  322. return query_doc_with_hybrid_search(
  323. collection_name=form_data.collection_name,
  324. query=form_data.query,
  325. embedding_function=app.state.EMBEDDING_FUNCTION,
  326. k=form_data.k if form_data.k else app.state.TOP_K,
  327. reranking_function=app.state.sentence_transformer_rf,
  328. r=form_data.r if form_data.r else app.state.RELEVANCE_THRESHOLD,
  329. )
  330. else:
  331. return query_doc(
  332. collection_name=form_data.collection_name,
  333. query=form_data.query,
  334. embedding_function=app.state.EMBEDDING_FUNCTION,
  335. k=form_data.k if form_data.k else app.state.TOP_K,
  336. )
  337. except Exception as e:
  338. log.exception(e)
  339. raise HTTPException(
  340. status_code=status.HTTP_400_BAD_REQUEST,
  341. detail=ERROR_MESSAGES.DEFAULT(e),
  342. )
  343. class QueryCollectionsForm(BaseModel):
  344. collection_names: List[str]
  345. query: str
  346. k: Optional[int] = None
  347. r: Optional[float] = None
  348. hybrid: Optional[bool] = None
  349. @app.post("/query/collection")
  350. def query_collection_handler(
  351. form_data: QueryCollectionsForm,
  352. user=Depends(get_current_user),
  353. ):
  354. try:
  355. if app.state.ENABLE_RAG_HYBRID_SEARCH:
  356. return query_collection_with_hybrid_search(
  357. collection_names=form_data.collection_names,
  358. query=form_data.query,
  359. embedding_function=app.state.EMBEDDING_FUNCTION,
  360. k=form_data.k if form_data.k else app.state.TOP_K,
  361. reranking_function=app.state.sentence_transformer_rf,
  362. r=form_data.r if form_data.r else app.state.RELEVANCE_THRESHOLD,
  363. )
  364. else:
  365. return query_collection(
  366. collection_names=form_data.collection_names,
  367. query=form_data.query,
  368. embedding_function=app.state.EMBEDDING_FUNCTION,
  369. k=form_data.k if form_data.k else app.state.TOP_K,
  370. )
  371. except Exception as e:
  372. log.exception(e)
  373. raise HTTPException(
  374. status_code=status.HTTP_400_BAD_REQUEST,
  375. detail=ERROR_MESSAGES.DEFAULT(e),
  376. )
  377. @app.post("/youtube")
  378. def store_youtube_video(form_data: UrlForm, user=Depends(get_current_user)):
  379. try:
  380. loader = YoutubeLoader.from_youtube_url(form_data.url, add_video_info=False)
  381. data = loader.load()
  382. collection_name = form_data.collection_name
  383. if collection_name == "":
  384. collection_name = calculate_sha256_string(form_data.url)[:63]
  385. store_data_in_vector_db(data, collection_name, overwrite=True)
  386. return {
  387. "status": True,
  388. "collection_name": collection_name,
  389. "filename": form_data.url,
  390. }
  391. except Exception as e:
  392. log.exception(e)
  393. raise HTTPException(
  394. status_code=status.HTTP_400_BAD_REQUEST,
  395. detail=ERROR_MESSAGES.DEFAULT(e),
  396. )
  397. @app.post("/web")
  398. def store_web(form_data: UrlForm, user=Depends(get_current_user)):
  399. # "https://www.gutenberg.org/files/1727/1727-h/1727-h.htm"
  400. try:
  401. loader = get_web_loader(form_data.url)
  402. data = loader.load()
  403. collection_name = form_data.collection_name
  404. if collection_name == "":
  405. collection_name = calculate_sha256_string(form_data.url)[:63]
  406. store_data_in_vector_db(data, collection_name, overwrite=True)
  407. return {
  408. "status": True,
  409. "collection_name": collection_name,
  410. "filename": form_data.url,
  411. }
  412. except Exception as e:
  413. log.exception(e)
  414. raise HTTPException(
  415. status_code=status.HTTP_400_BAD_REQUEST,
  416. detail=ERROR_MESSAGES.DEFAULT(e),
  417. )
  418. def get_web_loader(url: Union[str, Sequence[str]]):
  419. # Check if the URL is valid
  420. if not validate_url(url):
  421. raise ValueError(ERROR_MESSAGES.INVALID_URL)
  422. return WebBaseLoader(url)
  423. def validate_url(url: Union[str, Sequence[str]]):
  424. if isinstance(url, str):
  425. if isinstance(validators.url(url), validators.ValidationError):
  426. raise ValueError(ERROR_MESSAGES.INVALID_URL)
  427. if not ENABLE_LOCAL_WEB_FETCH:
  428. # Local web fetch is disabled, filter out any URLs that resolve to private IP addresses
  429. parsed_url = urllib.parse.urlparse(url)
  430. # Get IPv4 and IPv6 addresses
  431. ipv4_addresses, ipv6_addresses = resolve_hostname(parsed_url.hostname)
  432. # Check if any of the resolved addresses are private
  433. # This is technically still vulnerable to DNS rebinding attacks, as we don't control WebBaseLoader
  434. for ip in ipv4_addresses:
  435. if validators.ipv4(ip, private=True):
  436. raise ValueError(ERROR_MESSAGES.INVALID_URL)
  437. for ip in ipv6_addresses:
  438. if validators.ipv6(ip, private=True):
  439. raise ValueError(ERROR_MESSAGES.INVALID_URL)
  440. return True
  441. elif isinstance(url, Sequence):
  442. return all(validate_url(u) for u in url)
  443. else:
  444. return False
  445. def resolve_hostname(hostname):
  446. # Get address information
  447. addr_info = socket.getaddrinfo(hostname, None)
  448. # Extract IP addresses from address information
  449. ipv4_addresses = [info[4][0] for info in addr_info if info[0] == socket.AF_INET]
  450. ipv6_addresses = [info[4][0] for info in addr_info if info[0] == socket.AF_INET6]
  451. return ipv4_addresses, ipv6_addresses
  452. @app.post("/websearch")
  453. def store_websearch(form_data: SearchForm, user=Depends(get_current_user)):
  454. try:
  455. web_results = search_web(form_data.query)
  456. urls = [result.link for result in web_results]
  457. loader = get_web_loader(urls)
  458. data = loader.load()
  459. collection_name = form_data.collection_name
  460. if collection_name == "":
  461. collection_name = calculate_sha256_string(form_data.query)[:63]
  462. store_data_in_vector_db(data, collection_name, overwrite=True)
  463. return {
  464. "status": True,
  465. "collection_name": collection_name,
  466. "filenames": urls,
  467. }
  468. except Exception as e:
  469. log.exception(e)
  470. raise HTTPException(
  471. status_code=status.HTTP_400_BAD_REQUEST,
  472. detail=ERROR_MESSAGES.DEFAULT(e),
  473. )
  474. def store_data_in_vector_db(data, collection_name, overwrite: bool = False) -> bool:
  475. text_splitter = RecursiveCharacterTextSplitter(
  476. chunk_size=app.state.CHUNK_SIZE,
  477. chunk_overlap=app.state.CHUNK_OVERLAP,
  478. add_start_index=True,
  479. )
  480. docs = text_splitter.split_documents(data)
  481. if len(docs) > 0:
  482. log.info(f"store_data_in_vector_db {docs}")
  483. return store_docs_in_vector_db(docs, collection_name, overwrite), None
  484. else:
  485. raise ValueError(ERROR_MESSAGES.EMPTY_CONTENT)
  486. def store_text_in_vector_db(
  487. text, metadata, collection_name, overwrite: bool = False
  488. ) -> bool:
  489. text_splitter = RecursiveCharacterTextSplitter(
  490. chunk_size=app.state.CHUNK_SIZE,
  491. chunk_overlap=app.state.CHUNK_OVERLAP,
  492. add_start_index=True,
  493. )
  494. docs = text_splitter.create_documents([text], metadatas=[metadata])
  495. return store_docs_in_vector_db(docs, collection_name, overwrite)
  496. def store_docs_in_vector_db(docs, collection_name, overwrite: bool = False) -> bool:
  497. log.info(f"store_docs_in_vector_db {docs} {collection_name}")
  498. texts = [doc.page_content for doc in docs]
  499. metadatas = [doc.metadata for doc in docs]
  500. try:
  501. if overwrite:
  502. for collection in CHROMA_CLIENT.list_collections():
  503. if collection_name == collection.name:
  504. log.info(f"deleting existing collection {collection_name}")
  505. CHROMA_CLIENT.delete_collection(name=collection_name)
  506. collection = CHROMA_CLIENT.create_collection(name=collection_name)
  507. embedding_func = get_embedding_function(
  508. app.state.RAG_EMBEDDING_ENGINE,
  509. app.state.RAG_EMBEDDING_MODEL,
  510. app.state.sentence_transformer_ef,
  511. app.state.OPENAI_API_KEY,
  512. app.state.OPENAI_API_BASE_URL,
  513. )
  514. embedding_texts = list(map(lambda x: x.replace("\n", " "), texts))
  515. embeddings = embedding_func(embedding_texts)
  516. for batch in create_batches(
  517. api=CHROMA_CLIENT,
  518. ids=[str(uuid.uuid1()) for _ in texts],
  519. metadatas=metadatas,
  520. embeddings=embeddings,
  521. documents=texts,
  522. ):
  523. collection.add(*batch)
  524. return True
  525. except Exception as e:
  526. log.exception(e)
  527. if e.__class__.__name__ == "UniqueConstraintError":
  528. return True
  529. return False
  530. def get_loader(filename: str, file_content_type: str, file_path: str):
  531. file_ext = filename.split(".")[-1].lower()
  532. known_type = True
  533. known_source_ext = [
  534. "go",
  535. "py",
  536. "java",
  537. "sh",
  538. "bat",
  539. "ps1",
  540. "cmd",
  541. "js",
  542. "ts",
  543. "css",
  544. "cpp",
  545. "hpp",
  546. "h",
  547. "c",
  548. "cs",
  549. "sql",
  550. "log",
  551. "ini",
  552. "pl",
  553. "pm",
  554. "r",
  555. "dart",
  556. "dockerfile",
  557. "env",
  558. "php",
  559. "hs",
  560. "hsc",
  561. "lua",
  562. "nginxconf",
  563. "conf",
  564. "m",
  565. "mm",
  566. "plsql",
  567. "perl",
  568. "rb",
  569. "rs",
  570. "db2",
  571. "scala",
  572. "bash",
  573. "swift",
  574. "vue",
  575. "svelte",
  576. ]
  577. if file_ext == "pdf":
  578. loader = PyPDFLoader(file_path, extract_images=app.state.PDF_EXTRACT_IMAGES)
  579. elif file_ext == "csv":
  580. loader = CSVLoader(file_path)
  581. elif file_ext == "rst":
  582. loader = UnstructuredRSTLoader(file_path, mode="elements")
  583. elif file_ext == "xml":
  584. loader = UnstructuredXMLLoader(file_path)
  585. elif file_ext in ["htm", "html"]:
  586. loader = BSHTMLLoader(file_path, open_encoding="unicode_escape")
  587. elif file_ext == "md":
  588. loader = UnstructuredMarkdownLoader(file_path)
  589. elif file_content_type == "application/epub+zip":
  590. loader = UnstructuredEPubLoader(file_path)
  591. elif (
  592. file_content_type
  593. == "application/vnd.openxmlformats-officedocument.wordprocessingml.document"
  594. or file_ext in ["doc", "docx"]
  595. ):
  596. loader = Docx2txtLoader(file_path)
  597. elif file_content_type in [
  598. "application/vnd.ms-excel",
  599. "application/vnd.openxmlformats-officedocument.spreadsheetml.sheet",
  600. ] or file_ext in ["xls", "xlsx"]:
  601. loader = UnstructuredExcelLoader(file_path)
  602. elif file_ext in known_source_ext or (
  603. file_content_type and file_content_type.find("text/") >= 0
  604. ):
  605. loader = TextLoader(file_path, autodetect_encoding=True)
  606. else:
  607. loader = TextLoader(file_path, autodetect_encoding=True)
  608. known_type = False
  609. return loader, known_type
  610. @app.post("/doc")
  611. def store_doc(
  612. collection_name: Optional[str] = Form(None),
  613. file: UploadFile = File(...),
  614. user=Depends(get_current_user),
  615. ):
  616. # "https://www.gutenberg.org/files/1727/1727-h/1727-h.htm"
  617. log.info(f"file.content_type: {file.content_type}")
  618. try:
  619. unsanitized_filename = file.filename
  620. filename = os.path.basename(unsanitized_filename)
  621. file_path = f"{UPLOAD_DIR}/{filename}"
  622. contents = file.file.read()
  623. with open(file_path, "wb") as f:
  624. f.write(contents)
  625. f.close()
  626. f = open(file_path, "rb")
  627. if collection_name == None:
  628. collection_name = calculate_sha256(f)[:63]
  629. f.close()
  630. loader, known_type = get_loader(filename, file.content_type, file_path)
  631. data = loader.load()
  632. try:
  633. result = store_data_in_vector_db(data, collection_name)
  634. if result:
  635. return {
  636. "status": True,
  637. "collection_name": collection_name,
  638. "filename": filename,
  639. "known_type": known_type,
  640. }
  641. except Exception as e:
  642. raise HTTPException(
  643. status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
  644. detail=e,
  645. )
  646. except Exception as e:
  647. log.exception(e)
  648. if "No pandoc was found" in str(e):
  649. raise HTTPException(
  650. status_code=status.HTTP_400_BAD_REQUEST,
  651. detail=ERROR_MESSAGES.PANDOC_NOT_INSTALLED,
  652. )
  653. else:
  654. raise HTTPException(
  655. status_code=status.HTTP_400_BAD_REQUEST,
  656. detail=ERROR_MESSAGES.DEFAULT(e),
  657. )
  658. class TextRAGForm(BaseModel):
  659. name: str
  660. content: str
  661. collection_name: Optional[str] = None
  662. @app.post("/text")
  663. def store_text(
  664. form_data: TextRAGForm,
  665. user=Depends(get_current_user),
  666. ):
  667. collection_name = form_data.collection_name
  668. if collection_name == None:
  669. collection_name = calculate_sha256_string(form_data.content)
  670. result = store_text_in_vector_db(
  671. form_data.content,
  672. metadata={"name": form_data.name, "created_by": user.id},
  673. collection_name=collection_name,
  674. )
  675. if result:
  676. return {"status": True, "collection_name": collection_name}
  677. else:
  678. raise HTTPException(
  679. status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
  680. detail=ERROR_MESSAGES.DEFAULT(),
  681. )
  682. @app.get("/scan")
  683. def scan_docs_dir(user=Depends(get_admin_user)):
  684. for path in Path(DOCS_DIR).rglob("./**/*"):
  685. try:
  686. if path.is_file() and not path.name.startswith("."):
  687. tags = extract_folders_after_data_docs(path)
  688. filename = path.name
  689. file_content_type = mimetypes.guess_type(path)
  690. f = open(path, "rb")
  691. collection_name = calculate_sha256(f)[:63]
  692. f.close()
  693. loader, known_type = get_loader(
  694. filename, file_content_type[0], str(path)
  695. )
  696. data = loader.load()
  697. try:
  698. result = store_data_in_vector_db(data, collection_name)
  699. if result:
  700. sanitized_filename = sanitize_filename(filename)
  701. doc = Documents.get_doc_by_name(sanitized_filename)
  702. if doc == None:
  703. doc = Documents.insert_new_doc(
  704. user.id,
  705. DocumentForm(
  706. **{
  707. "name": sanitized_filename,
  708. "title": filename,
  709. "collection_name": collection_name,
  710. "filename": filename,
  711. "content": (
  712. json.dumps(
  713. {
  714. "tags": list(
  715. map(
  716. lambda name: {"name": name},
  717. tags,
  718. )
  719. )
  720. }
  721. )
  722. if len(tags)
  723. else "{}"
  724. ),
  725. }
  726. ),
  727. )
  728. except Exception as e:
  729. log.exception(e)
  730. pass
  731. except Exception as e:
  732. log.exception(e)
  733. return True
  734. @app.get("/reset/db")
  735. def reset_vector_db(user=Depends(get_admin_user)):
  736. CHROMA_CLIENT.reset()
  737. @app.get("/reset")
  738. def reset(user=Depends(get_admin_user)) -> bool:
  739. folder = f"{UPLOAD_DIR}"
  740. for filename in os.listdir(folder):
  741. file_path = os.path.join(folder, filename)
  742. try:
  743. if os.path.isfile(file_path) or os.path.islink(file_path):
  744. os.unlink(file_path)
  745. elif os.path.isdir(file_path):
  746. shutil.rmtree(file_path)
  747. except Exception as e:
  748. log.error("Failed to delete %s. Reason: %s" % (file_path, e))
  749. try:
  750. CHROMA_CLIENT.reset()
  751. except Exception as e:
  752. log.exception(e)
  753. return True