refac

Author: tjbck

backend/open_webui/retrieval/loaders/external_document.py

@@ -6,7 +6,7 @@
 import requests
 from langchain_core.document_loaders import BaseLoader
 from langchain_core.documents import Document
-from open_webui.utils.headers import include_user_info_headers
+from open_webui.utils.headers import get_custom_headers, include_user_info_headers
 
 log = logging.getLogger(__name__)
 
@@ -19,6 +19,8 @@ def __init__(
         api_key: str,
         mime_type=None,
         user=None,
+        headers=None,
+        metadata=None,
         **kwargs,
     ) -> None:
         self.url = url
@@ -28,6 +30,8 @@ def __init__(
         self.mime_type = mime_type
 
         self.user = user
+        self.headers = headers
+        self.metadata = metadata
 
     def load(self) -> List[Document]:
         with open(self.file_path, 'rb') as f:
@@ -45,6 +49,8 @@ def load(self) -> List[Document]:
         except Exception:
             pass
 
+        headers.update(get_custom_headers(self.headers, self.user, self.metadata))
+
         if self.user is not None:
             headers = include_user_info_headers(headers, self.user)
 

backend/open_webui/retrieval/loaders/main.py

@@ -229,6 +229,7 @@ class Loader:
     def __init__(self, engine: str = '', **kwargs):
         self.engine = engine
         self.user = kwargs.get('user', None)
+        self.metadata = kwargs.get('metadata', {})
         self.kwargs = kwargs
 
     def load(self, filename: str, file_content_type: str, file_path: str) -> list[Document]:
@@ -404,6 +405,12 @@ def _get_loader(self, filename: str, file_content_type: str, file_path: str):
                 api_key=self.kwargs.get('EXTERNAL_DOCUMENT_LOADER_API_KEY'),
                 mime_type=file_content_type,
                 user=self.user,
+                headers=self.kwargs.get('EXTERNAL_DOCUMENT_LOADER_HEADERS'),
+                metadata={
+                    **self.metadata,
+                    'file_name': filename,
+                    'file_content_type': file_content_type,
+                },
             )
         elif self.engine == 'tika' and self.kwargs.get('TIKA_SERVER_URL'):
             if self._is_text_file(file_ext, file_content_type):

backend/open_webui/utils/headers.py

@@ -75,6 +75,9 @@ def get_custom_headers(custom_headers: dict, user=None, metadata: dict = None) -
         '{{MESSAGE_ID}}': metadata.get('message_id', '') or '',
         '{{USER_MESSAGE_ID}}': user_message_id or '',
         '{{USER_MESSAGE_PARENT_ID}}': user_message_parent_id or '',
+        '{{FILE_ID}}': metadata.get('file_id', '') or '',
+        '{{FILE_NAME}}': metadata.get('file_name', '') or '',
+        '{{FILE_CONTENT_TYPE}}': metadata.get('file_content_type', '') or '',
         '{{TASK}}': metadata.get('task', '') or '',
         '{{USER_ID}}': (user.id if user else '') or '',
         '{{USER_NAME}}': (user.name if user else '') or '',