|
@@ -91,6 +91,8 @@ from config import (
|
|
SRC_LOG_LEVELS,
|
|
SRC_LOG_LEVELS,
|
|
UPLOAD_DIR,
|
|
UPLOAD_DIR,
|
|
DOCS_DIR,
|
|
DOCS_DIR,
|
|
|
|
+ TEXT_EXTRACTION_ENGINE,
|
|
|
|
+ TIKA_SERVER_URL,
|
|
RAG_TOP_K,
|
|
RAG_TOP_K,
|
|
RAG_RELEVANCE_THRESHOLD,
|
|
RAG_RELEVANCE_THRESHOLD,
|
|
RAG_EMBEDDING_ENGINE,
|
|
RAG_EMBEDDING_ENGINE,
|
|
@@ -146,6 +148,9 @@ app.state.config.ENABLE_RAG_WEB_LOADER_SSL_VERIFICATION = (
|
|
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_SIZE = CHUNK_SIZE
|
|
app.state.config.CHUNK_OVERLAP = CHUNK_OVERLAP
|
|
app.state.config.CHUNK_OVERLAP = CHUNK_OVERLAP
|
|
|
|
|
|
@@ -390,6 +395,10 @@ async def get_rag_config(user=Depends(get_admin_user)):
|
|
return {
|
|
return {
|
|
"status": True,
|
|
"status": True,
|
|
"pdf_extract_images": app.state.config.PDF_EXTRACT_IMAGES,
|
|
"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": {
|
|
"chunk_size": app.state.config.CHUNK_SIZE,
|
|
"chunk_size": app.state.config.CHUNK_SIZE,
|
|
"chunk_overlap": app.state.config.CHUNK_OVERLAP,
|
|
"chunk_overlap": app.state.config.CHUNK_OVERLAP,
|
|
@@ -419,6 +428,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):
|
|
class ChunkParamUpdateForm(BaseModel):
|
|
chunk_size: int
|
|
chunk_size: int
|
|
chunk_overlap: int
|
|
chunk_overlap: int
|
|
@@ -452,6 +466,7 @@ class WebConfig(BaseModel):
|
|
|
|
|
|
class ConfigUpdateForm(BaseModel):
|
|
class ConfigUpdateForm(BaseModel):
|
|
pdf_extract_images: Optional[bool] = None
|
|
pdf_extract_images: Optional[bool] = None
|
|
|
|
+ text_extraction: Optional[TextExtractionConfig] = None
|
|
chunk: Optional[ChunkParamUpdateForm] = None
|
|
chunk: Optional[ChunkParamUpdateForm] = None
|
|
youtube: Optional[YoutubeLoaderConfig] = None
|
|
youtube: Optional[YoutubeLoaderConfig] = None
|
|
web: Optional[WebConfig] = None
|
|
web: Optional[WebConfig] = None
|
|
@@ -465,6 +480,11 @@ async def update_rag_config(form_data: ConfigUpdateForm, user=Depends(get_admin_
|
|
else app.state.config.PDF_EXTRACT_IMAGES
|
|
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:
|
|
if form_data.chunk is not None:
|
|
app.state.config.CHUNK_SIZE = form_data.chunk.chunk_size
|
|
app.state.config.CHUNK_SIZE = form_data.chunk.chunk_size
|
|
app.state.config.CHUNK_OVERLAP = form_data.chunk.chunk_overlap
|
|
app.state.config.CHUNK_OVERLAP = form_data.chunk.chunk_overlap
|
|
@@ -501,6 +521,10 @@ async def update_rag_config(form_data: ConfigUpdateForm, user=Depends(get_admin_
|
|
return {
|
|
return {
|
|
"status": True,
|
|
"status": True,
|
|
"pdf_extract_images": app.state.config.PDF_EXTRACT_IMAGES,
|
|
"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": {
|
|
"chunk_size": app.state.config.CHUNK_SIZE,
|
|
"chunk_size": app.state.config.CHUNK_SIZE,
|
|
"chunk_overlap": app.state.config.CHUNK_OVERLAP,
|
|
"chunk_overlap": app.state.config.CHUNK_OVERLAP,
|
|
@@ -987,6 +1011,41 @@ def store_docs_in_vector_db(docs, collection_name, overwrite: bool = False) -> b
|
|
return False
|
|
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 = app.state.config.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", "<No text content found>")
|
|
|
|
+
|
|
|
|
+ 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):
|
|
def get_loader(filename: str, file_content_type: str, file_path: str):
|
|
file_ext = filename.split(".")[-1].lower()
|
|
file_ext = filename.split(".")[-1].lower()
|
|
known_type = True
|
|
known_type = True
|
|
@@ -1037,47 +1096,55 @@ def get_loader(filename: str, file_content_type: str, file_path: str):
|
|
"msg",
|
|
"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)
|
|
|
|
|
|
+ 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
|
|
|
|
+ ):
|
|
|
|
+ loader = TextLoader(file_path, autodetect_encoding=True)
|
|
|
|
+ else:
|
|
|
|
+ loader = TikaLoader(file_path, file_content_type)
|
|
else:
|
|
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
|
|
return loader, known_type
|
|
|
|
|