From 9cf622d98118e40a92194a4ef76c1cacd92d2640 Mon Sep 17 00:00:00 2001 From: Nicko van Someren Date: Sun, 30 Jun 2024 15:49:15 -0600 Subject: [PATCH 1/2] Added support for using Apache Tika as a document loader. Added persistent configuration options to configure use and location of Tika service. Updated backend.apps.rag.main:get_loader() to make use of Tika document loader. --- .dockerignore | 2 +- backend/apps/rag/main.py | 127 +++++++++++++++++++++++++++------------ backend/config.py | 16 +++++ 3 files changed, 104 insertions(+), 41 deletions(-) diff --git a/.dockerignore b/.dockerignore index e28863bf6..c7f330f4a 100644 --- a/.dockerignore +++ b/.dockerignore @@ -10,7 +10,7 @@ node_modules vite.config.js.timestamp-* vite.config.ts.timestamp-* __pycache__ -.env +.idea _old uploads .ipynb_checkpoints diff --git a/backend/apps/rag/main.py b/backend/apps/rag/main.py index 7c6974535..b6e2ee5e2 100644 --- a/backend/apps/rag/main.py +++ b/backend/apps/rag/main.py @@ -93,6 +93,8 @@ from config import ( SRC_LOG_LEVELS, UPLOAD_DIR, DOCS_DIR, + DOCUMENT_USE_TIKA, + TIKA_SERVER_URL, RAG_TOP_K, RAG_RELEVANCE_THRESHOLD, RAG_EMBEDDING_ENGINE, @@ -985,6 +987,41 @@ def store_docs_in_vector_db(docs, collection_name, overwrite: bool = False) -> b return False +class TikaLoader: + def __init__(self, file_path, mime_type=None): + self.file_path = file_path + self.mime_type = mime_type + + def load(self) -> List[Document]: + with (open(self.file_path, "rb") as f): + data = f.read() + + if self.mime_type is not None: + headers = {"Content-Type": self.mime_type} + else: + headers = {} + + endpoint = str(TIKA_SERVER_URL) + if not endpoint.endswith("/"): + endpoint += "/" + endpoint += "tika/text" + + r = requests.put(endpoint, data=data, headers=headers) + + if r.ok: + raw_metadata = r.json() + text = raw_metadata.get("X-TIKA:content", "") + + if "Content-Type" in raw_metadata: + headers["Content-Type"] = raw_metadata["Content-Type"] + + log.info("Tika extracted text: %s", text) + + return [Document(page_content=text, metadata=headers)] + else: + raise Exception(f"Error calling Tika: {r.reason}") + + def get_loader(filename: str, file_content_type: str, file_path: str): file_ext = filename.split(".")[-1].lower() known_type = True @@ -1035,47 +1072,57 @@ def get_loader(filename: str, file_content_type: str, file_path: str): "msg", ] - if file_ext == "pdf": - loader = PyPDFLoader( - file_path, extract_images=app.state.config.PDF_EXTRACT_IMAGES - ) - elif file_ext == "csv": - loader = CSVLoader(file_path) - elif file_ext == "rst": - loader = UnstructuredRSTLoader(file_path, mode="elements") - elif file_ext == "xml": - loader = UnstructuredXMLLoader(file_path) - elif file_ext in ["htm", "html"]: - loader = BSHTMLLoader(file_path, open_encoding="unicode_escape") - elif file_ext == "md": - loader = UnstructuredMarkdownLoader(file_path) - elif file_content_type == "application/epub+zip": - loader = UnstructuredEPubLoader(file_path) - elif ( - file_content_type - == "application/vnd.openxmlformats-officedocument.wordprocessingml.document" - or file_ext in ["doc", "docx"] - ): - loader = Docx2txtLoader(file_path) - elif file_content_type in [ - "application/vnd.ms-excel", - "application/vnd.openxmlformats-officedocument.spreadsheetml.sheet", - ] or file_ext in ["xls", "xlsx"]: - loader = UnstructuredExcelLoader(file_path) - elif file_content_type in [ - "application/vnd.ms-powerpoint", - "application/vnd.openxmlformats-officedocument.presentationml.presentation", - ] or file_ext in ["ppt", "pptx"]: - loader = UnstructuredPowerPointLoader(file_path) - elif file_ext == "msg": - loader = OutlookMessageLoader(file_path) - elif file_ext in known_source_ext or ( - file_content_type and file_content_type.find("text/") >= 0 - ): - loader = TextLoader(file_path, autodetect_encoding=True) + log.warning("Use tika: %s, server URL: %s", DOCUMENT_USE_TIKA, TIKA_SERVER_URL) + + if DOCUMENT_USE_TIKA and TIKA_SERVER_URL: + if file_ext in known_source_ext or ( + file_content_type and file_content_type.find("text/") >= 0 + ): + loader = TextLoader(file_path, autodetect_encoding=True) + else: + loader = TikaLoader(file_path, file_content_type) else: - loader = TextLoader(file_path, autodetect_encoding=True) - known_type = False + if file_ext == "pdf": + loader = PyPDFLoader( + file_path, extract_images=app.state.config.PDF_EXTRACT_IMAGES + ) + elif file_ext == "csv": + loader = CSVLoader(file_path) + elif file_ext == "rst": + loader = UnstructuredRSTLoader(file_path, mode="elements") + elif file_ext == "xml": + loader = UnstructuredXMLLoader(file_path) + elif file_ext in ["htm", "html"]: + loader = BSHTMLLoader(file_path, open_encoding="unicode_escape") + elif file_ext == "md": + loader = UnstructuredMarkdownLoader(file_path) + elif file_content_type == "application/epub+zip": + loader = UnstructuredEPubLoader(file_path) + elif ( + file_content_type + == "application/vnd.openxmlformats-officedocument.wordprocessingml.document" + or file_ext in ["doc", "docx"] + ): + loader = Docx2txtLoader(file_path) + elif file_content_type in [ + "application/vnd.ms-excel", + "application/vnd.openxmlformats-officedocument.spreadsheetml.sheet", + ] or file_ext in ["xls", "xlsx"]: + loader = UnstructuredExcelLoader(file_path) + elif file_content_type in [ + "application/vnd.ms-powerpoint", + "application/vnd.openxmlformats-officedocument.presentationml.presentation", + ] or file_ext in ["ppt", "pptx"]: + loader = UnstructuredPowerPointLoader(file_path) + elif file_ext == "msg": + loader = OutlookMessageLoader(file_path) + elif file_ext in known_source_ext or ( + file_content_type and file_content_type.find("text/") >= 0 + ): + loader = TextLoader(file_path, autodetect_encoding=True) + else: + loader = TextLoader(file_path, autodetect_encoding=True) + known_type = False return loader, known_type diff --git a/backend/config.py b/backend/config.py index 3a825f53a..fe9d995ef 100644 --- a/backend/config.py +++ b/backend/config.py @@ -878,6 +878,22 @@ WEBUI_SESSION_COOKIE_SECURE = os.environ.get( if WEBUI_AUTH and WEBUI_SECRET_KEY == "": raise ValueError(ERROR_MESSAGES.ENV_VAR_NOT_FOUND) +#################################### +# RAG document text extraction +#################################### + +DOCUMENT_USE_TIKA = PersistentConfig( + "DOCUMENT_USE_TIKA", + "rag.document_use_tika", + os.environ.get("DOCUMENT_USE_TIKA", "false").lower() == "true" +) + +TIKA_SERVER_URL = PersistentConfig( + "TIKA_SERVER_URL", + "rag.tika_server_url", + os.getenv("TIKA_SERVER_URL", "http://tika:9998"), # Default for sidecar deployment +) + #################################### # RAG #################################### From 7aa35a37573c1d0af136176756a16ba73b74f74b Mon Sep 17 00:00:00 2001 From: Nicko van Someren Date: Mon, 1 Jul 2024 12:10:59 -0600 Subject: [PATCH 2/2] Added HTML and Typescript UI components to support configration of text extraction engine. Updated RAG /config and /config/update endpoints to support UI updates. Fixed .dockerignore to prevent Python venv from being copied into Docker image. --- .dockerignore | 1 + backend/apps/rag/main.py | 30 +++++++++-- backend/config.py | 8 +-- src/lib/apis/rag/index.ts | 6 +++ .../admin/Settings/Documents.svelte | 50 +++++++++++++++++++ 5 files changed, 86 insertions(+), 9 deletions(-) diff --git a/.dockerignore b/.dockerignore index c7f330f4a..d7e716758 100644 --- a/.dockerignore +++ b/.dockerignore @@ -11,6 +11,7 @@ vite.config.js.timestamp-* vite.config.ts.timestamp-* __pycache__ .idea +venv _old uploads .ipynb_checkpoints diff --git a/backend/apps/rag/main.py b/backend/apps/rag/main.py index b6e2ee5e2..ce56d4f4d 100644 --- a/backend/apps/rag/main.py +++ b/backend/apps/rag/main.py @@ -93,7 +93,7 @@ from config import ( SRC_LOG_LEVELS, UPLOAD_DIR, DOCS_DIR, - DOCUMENT_USE_TIKA, + TEXT_EXTRACTION_ENGINE, TIKA_SERVER_URL, RAG_TOP_K, RAG_RELEVANCE_THRESHOLD, @@ -150,6 +150,9 @@ app.state.config.ENABLE_RAG_WEB_LOADER_SSL_VERIFICATION = ( ENABLE_RAG_WEB_LOADER_SSL_VERIFICATION ) +app.state.config.TEXT_EXTRACTION_ENGINE = TEXT_EXTRACTION_ENGINE +app.state.config.TIKA_SERVER_URL = TIKA_SERVER_URL + app.state.config.CHUNK_SIZE = CHUNK_SIZE app.state.config.CHUNK_OVERLAP = CHUNK_OVERLAP @@ -390,6 +393,10 @@ async def get_rag_config(user=Depends(get_admin_user)): return { "status": True, "pdf_extract_images": app.state.config.PDF_EXTRACT_IMAGES, + "text_extraction": { + "engine": app.state.config.TEXT_EXTRACTION_ENGINE, + "tika_server_url": app.state.config.TIKA_SERVER_URL, + }, "chunk": { "chunk_size": app.state.config.CHUNK_SIZE, "chunk_overlap": app.state.config.CHUNK_OVERLAP, @@ -419,6 +426,11 @@ async def get_rag_config(user=Depends(get_admin_user)): } +class TextExtractionConfig(BaseModel): + engine: str = "" + tika_server_url: Optional[str] = None + + class ChunkParamUpdateForm(BaseModel): chunk_size: int chunk_overlap: int @@ -452,6 +464,7 @@ class WebConfig(BaseModel): class ConfigUpdateForm(BaseModel): pdf_extract_images: Optional[bool] = None + text_extraction: Optional[TextExtractionConfig] = None chunk: Optional[ChunkParamUpdateForm] = None youtube: Optional[YoutubeLoaderConfig] = None web: Optional[WebConfig] = None @@ -465,6 +478,11 @@ async def update_rag_config(form_data: ConfigUpdateForm, user=Depends(get_admin_ else app.state.config.PDF_EXTRACT_IMAGES ) + if form_data.text_extraction is not None: + log.info(f"Updating text settings: {form_data.text_extraction}") + app.state.config.TEXT_EXTRACTION_ENGINE = form_data.text_extraction.engine + app.state.config.TIKA_SERVER_URL = form_data.text_extraction.tika_server_url + if form_data.chunk is not None: app.state.config.CHUNK_SIZE = form_data.chunk.chunk_size app.state.config.CHUNK_OVERLAP = form_data.chunk.chunk_overlap @@ -501,6 +519,10 @@ async def update_rag_config(form_data: ConfigUpdateForm, user=Depends(get_admin_ return { "status": True, "pdf_extract_images": app.state.config.PDF_EXTRACT_IMAGES, + "text_extraction": { + "engine": app.state.config.TEXT_EXTRACTION_ENGINE, + "tika_server_url": app.state.config.TIKA_SERVER_URL, + }, "chunk": { "chunk_size": app.state.config.CHUNK_SIZE, "chunk_overlap": app.state.config.CHUNK_OVERLAP, @@ -1001,7 +1023,7 @@ class TikaLoader: else: headers = {} - endpoint = str(TIKA_SERVER_URL) + endpoint = app.state.config.TIKA_SERVER_URL if not endpoint.endswith("/"): endpoint += "/" endpoint += "tika/text" @@ -1072,9 +1094,7 @@ def get_loader(filename: str, file_content_type: str, file_path: str): "msg", ] - log.warning("Use tika: %s, server URL: %s", DOCUMENT_USE_TIKA, TIKA_SERVER_URL) - - if DOCUMENT_USE_TIKA and TIKA_SERVER_URL: + if app.state.config.TEXT_EXTRACTION_ENGINE == "tika" and app.state.config.TIKA_SERVER_URL: if file_ext in known_source_ext or ( file_content_type and file_content_type.find("text/") >= 0 ): diff --git a/backend/config.py b/backend/config.py index fe9d995ef..e4097b4cb 100644 --- a/backend/config.py +++ b/backend/config.py @@ -882,10 +882,10 @@ if WEBUI_AUTH and WEBUI_SECRET_KEY == "": # RAG document text extraction #################################### -DOCUMENT_USE_TIKA = PersistentConfig( - "DOCUMENT_USE_TIKA", - "rag.document_use_tika", - os.environ.get("DOCUMENT_USE_TIKA", "false").lower() == "true" +TEXT_EXTRACTION_ENGINE = PersistentConfig( + "TEXT_EXTRACTION_ENGINE", + "rag.text_extraction_engine", + os.environ.get("TEXT_EXTRACTION_ENGINE", "").lower() ) TIKA_SERVER_URL = PersistentConfig( diff --git a/src/lib/apis/rag/index.ts b/src/lib/apis/rag/index.ts index 50f236e06..4047c419a 100644 --- a/src/lib/apis/rag/index.ts +++ b/src/lib/apis/rag/index.ts @@ -32,6 +32,11 @@ type ChunkConfigForm = { chunk_overlap: number; }; +type TextExtractConfigForm = { + engine: string; + tika_server_url: string | null; +}; + type YoutubeConfigForm = { language: string[]; translation?: string | null; @@ -40,6 +45,7 @@ type YoutubeConfigForm = { type RAGConfigForm = { pdf_extract_images?: boolean; chunk?: ChunkConfigForm; + text_extraction?: TextExtractConfigForm; web_loader_ssl_verification?: boolean; youtube?: YoutubeConfigForm; }; diff --git a/src/lib/components/admin/Settings/Documents.svelte b/src/lib/components/admin/Settings/Documents.svelte index 575014b13..1377eb5bb 100644 --- a/src/lib/components/admin/Settings/Documents.svelte +++ b/src/lib/components/admin/Settings/Documents.svelte @@ -37,6 +37,10 @@ let embeddingModel = ''; let rerankingModel = ''; + let textExtractionEngine = 'default'; + let tikaServerUrl = ''; + let showTikaServerUrl = false; + let chunkSize = 0; let chunkOverlap = 0; let pdfExtractImages = true; @@ -163,11 +167,20 @@ rerankingModelUpdateHandler(); } + if (textExtractionEngine === 'tika' && tikaServerUrl === '') { + toast.error($i18n.t('Tika Server URL required.')); + return; + } + const res = await updateRAGConfig(localStorage.token, { pdf_extract_images: pdfExtractImages, chunk: { chunk_overlap: chunkOverlap, chunk_size: chunkSize + }, + text_extraction: { + engine: textExtractionEngine, + tika_server_url: tikaServerUrl } }); @@ -213,6 +226,10 @@ chunkSize = res.chunk.chunk_size; chunkOverlap = res.chunk.chunk_overlap; + + textExtractionEngine = res.text_extraction.engine; + tikaServerUrl = res.text_extraction.tika_server_url; + showTikaServerUrl = textExtractionEngine === 'tika'; } }); @@ -388,6 +405,39 @@ +
+ +
+
{$i18n.t('Text Extraction')}
+ +
+
{$i18n.t('Engine')}
+
+ +
+
+ + {#if showTikaServerUrl} +
+
+ +
+
+ {/if} +