main.py 26 KB

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