|
@@ -20,6 +20,9 @@ from langchain_community.document_loaders import (
|
|
|
YoutubeLoader,
|
|
|
)
|
|
|
from langchain_core.documents import Document
|
|
|
+
|
|
|
+from mistralai import Mistral
|
|
|
+
|
|
|
from open_webui.env import SRC_LOG_LEVELS, GLOBAL_LOG_LEVEL
|
|
|
|
|
|
logging.basicConfig(stream=sys.stdout, level=GLOBAL_LOG_LEVEL)
|
|
@@ -163,6 +166,53 @@ class DoclingLoader:
|
|
|
raise Exception(f"Error calling Docling: {error_msg}")
|
|
|
|
|
|
|
|
|
+class MistralLoader:
|
|
|
+ def __init__(self, api_key: str, file_path: str):
|
|
|
+ self.api_key = api_key
|
|
|
+ self.file_path = file_path
|
|
|
+ self.client = Mistral(api_key=api_key)
|
|
|
+
|
|
|
+ def load(self) -> list[Document]:
|
|
|
+ log.info("Uploading file to Mistral OCR")
|
|
|
+ uploaded_pdf = self.client.files.upload(
|
|
|
+ file={
|
|
|
+ "file_name": self.file_path.split("/")[-1],
|
|
|
+ "content": open(self.file_path, "rb"),
|
|
|
+ },
|
|
|
+ purpose="ocr",
|
|
|
+ )
|
|
|
+ log.info("File uploaded to Mistral OCR, getting signed URL")
|
|
|
+ signed_url = self.client.files.get_signed_url(file_id=uploaded_pdf.id)
|
|
|
+ log.info("Signed URL received, processing OCR")
|
|
|
+ ocr_response = self.client.ocr.process(
|
|
|
+ model="mistral-ocr-latest",
|
|
|
+ document={
|
|
|
+ "type": "document_url",
|
|
|
+ "document_url": signed_url.url,
|
|
|
+ },
|
|
|
+ )
|
|
|
+ log.info("OCR processing done, deleting uploaded file")
|
|
|
+ deleted_pdf = self.client.files.delete(file_id=uploaded_pdf.id)
|
|
|
+ log.info("Uploaded file deleted")
|
|
|
+ log.debug("OCR response: %s", ocr_response)
|
|
|
+ if not hasattr(ocr_response, "pages") or not ocr_response.pages:
|
|
|
+ log.error("No pages found in OCR response")
|
|
|
+ return [Document(page_content="No text content found", metadata={})]
|
|
|
+
|
|
|
+ return [
|
|
|
+ Document(
|
|
|
+ page_content=page.markdown,
|
|
|
+ metadata={
|
|
|
+ "page": page.index,
|
|
|
+ "page_label": page.index + 1,
|
|
|
+ "total_pages": len(ocr_response.pages),
|
|
|
+ },
|
|
|
+ )
|
|
|
+ for page in ocr_response.pages
|
|
|
+ if hasattr(page, "markdown") and hasattr(page, "index")
|
|
|
+ ]
|
|
|
+
|
|
|
+
|
|
|
class Loader:
|
|
|
def __init__(self, engine: str = "", **kwargs):
|
|
|
self.engine = engine
|
|
@@ -222,6 +272,15 @@ class Loader:
|
|
|
api_endpoint=self.kwargs.get("DOCUMENT_INTELLIGENCE_ENDPOINT"),
|
|
|
api_key=self.kwargs.get("DOCUMENT_INTELLIGENCE_KEY"),
|
|
|
)
|
|
|
+ elif (
|
|
|
+ self.engine == "mistral_ocr"
|
|
|
+ and self.kwargs.get("MISTRAL_OCR_API_KEY") != ""
|
|
|
+ and file_ext
|
|
|
+ in ["pdf"] # Mistral OCR currently only supports PDF and images
|
|
|
+ ):
|
|
|
+ loader = MistralLoader(
|
|
|
+ api_key=self.kwargs.get("MISTRAL_OCR_API_KEY"), file_path=file_path
|
|
|
+ )
|
|
|
else:
|
|
|
if file_ext == "pdf":
|
|
|
loader = PyPDFLoader(
|