Coverage for open_webui/retrieval/loaders/datalab_marker.py: 9%
118 statements
« prev ^ index » next coverage.py v7.15.2, created at 2026-10-07 05:07 +0000
1import json
2import logging
3import os
4import time
5from typing import List, Optional
7import requests
8from fastapi import HTTPException, status
9from langchain_core.documents import Document
10from open_webui.utils.json_codec import JSONCodec
12log = logging.getLogger(__name__)
15class DatalabMarkerLoader:
16 def __init__(
17 self,
18 file_path: str,
19 api_key: str,
20 api_base_url: str,
21 additional_config: Optional[str] = None,
22 use_llm: bool = False,
23 skip_cache: bool = False,
24 force_ocr: bool = False,
25 paginate: bool = False,
26 strip_existing_ocr: bool = False,
27 disable_image_extraction: bool = False,
28 format_lines: bool = False,
29 output_format: str = None,
30 ):
31 self.file_path = file_path
32 self.api_key = api_key
33 self.api_base_url = api_base_url
34 self.additional_config = additional_config
35 self.use_llm = use_llm
36 self.skip_cache = skip_cache
37 self.force_ocr = force_ocr
38 self.paginate = paginate
39 self.strip_existing_ocr = strip_existing_ocr
40 self.disable_image_extraction = disable_image_extraction
41 self.format_lines = format_lines
42 self.output_format = output_format
44 def _get_mime_type(self, filename: str) -> str:
45 ext = filename.rsplit('.', 1)[-1].lower()
46 mime_map = {
47 'pdf': 'application/pdf',
48 'xls': 'application/vnd.ms-excel',
49 'xlsx': 'application/vnd.openxmlformats-officedocument.spreadsheetml.sheet',
50 'ods': 'application/vnd.oasis.opendocument.spreadsheet',
51 'doc': 'application/msword',
52 'docx': 'application/vnd.openxmlformats-officedocument.wordprocessingml.document',
53 'odt': 'application/vnd.oasis.opendocument.text',
54 'ppt': 'application/vnd.ms-powerpoint',
55 'pptx': 'application/vnd.openxmlformats-officedocument.presentationml.presentation',
56 'odp': 'application/vnd.oasis.opendocument.presentation',
57 'html': 'text/html',
58 'epub': 'application/epub+zip',
59 'png': 'image/png',
60 'jpeg': 'image/jpeg',
61 'jpg': 'image/jpeg',
62 'webp': 'image/webp',
63 'gif': 'image/gif',
64 'tiff': 'image/tiff',
65 }
66 return mime_map.get(ext, 'application/octet-stream')
68 def load(self) -> List[Document]:
69 filename = os.path.basename(self.file_path)
70 mime_type = self._get_mime_type(filename)
71 headers = {'X-Api-Key': self.api_key}
73 form_data = {
74 'use_llm': str(self.use_llm).lower(),
75 'skip_cache': str(self.skip_cache).lower(),
76 'force_ocr': str(self.force_ocr).lower(),
77 'paginate': str(self.paginate).lower(),
78 'strip_existing_ocr': str(self.strip_existing_ocr).lower(),
79 'disable_image_extraction': str(self.disable_image_extraction).lower(),
80 'format_lines': str(self.format_lines).lower(),
81 'output_format': self.output_format,
82 }
84 if self.additional_config and self.additional_config.strip():
85 form_data['additional_config'] = self.additional_config
87 log.info(
88 "Datalab Marker POST request parameters: {'filename': '%s', 'mime_type': '%s', **%s}",
89 filename,
90 mime_type,
91 form_data,
92 )
94 try:
95 with open(self.file_path, 'rb') as f:
96 files = {'file': (filename, f, mime_type)}
97 response = requests.post(
98 f'{self.api_base_url}',
99 data=form_data,
100 files=files,
101 headers=headers,
102 )
103 response.raise_for_status()
104 result = response.json()
105 except FileNotFoundError:
106 raise HTTPException(status.HTTP_404_NOT_FOUND, detail=f'File not found: {self.file_path}')
107 except requests.HTTPError as e:
108 raise HTTPException(
109 status.HTTP_400_BAD_REQUEST,
110 detail=f'Datalab Marker request failed: {e}',
111 )
112 except ValueError as e:
113 raise HTTPException(status.HTTP_502_BAD_GATEWAY, detail=f'Invalid JSON response: {e}')
114 except Exception as e:
115 raise HTTPException(status.HTTP_500_INTERNAL_SERVER_ERROR, detail=str(e))
117 if not result.get('success'):
118 raise HTTPException(
119 status.HTTP_400_BAD_REQUEST,
120 detail=f'Datalab Marker request failed: {result.get("error", "Unknown error")}',
121 )
123 check_url = result.get('request_check_url')
124 request_id = result.get('request_id')
126 # Check if this is a direct response (self-hosted) or polling response (DataLab)
127 if check_url:
128 # DataLab polling pattern
129 for _ in range(300): # Up to 10 minutes
130 time.sleep(2)
131 try:
132 poll_response = requests.get(check_url, headers=headers)
133 poll_response.raise_for_status()
134 poll_result = poll_response.json()
135 except (requests.HTTPError, ValueError) as e:
136 raw_body = poll_response.text
137 log.error(f'Polling error: {e}, response body: {raw_body}')
138 raise HTTPException(status.HTTP_502_BAD_GATEWAY, detail=f'Polling failed: {e}')
140 status_val = poll_result.get('status')
141 success_val = poll_result.get('success')
143 if status_val == 'complete':
144 summary = {
145 k: poll_result.get(k)
146 for k in (
147 'status',
148 'output_format',
149 'success',
150 'error',
151 'page_count',
152 'total_cost',
153 )
154 }
155 log.info('Marker processing completed successfully: %s', json.dumps(summary, indent=2))
156 break
158 if status_val == 'failed' or success_val is False:
159 log.error(f'Marker poll failed full response: {json.dumps(poll_result, indent=2)}')
160 error_msg = poll_result.get('error') or 'Marker returned failure without error message'
161 raise HTTPException(
162 status.HTTP_400_BAD_REQUEST,
163 detail=f'Marker processing failed: {error_msg}',
164 )
165 else:
166 raise HTTPException(
167 status.HTTP_504_GATEWAY_TIMEOUT,
168 detail='Marker processing timed out',
169 )
171 if not poll_result.get('success', False):
172 error_msg = poll_result.get('error') or 'Unknown processing error'
173 raise HTTPException(
174 status.HTTP_400_BAD_REQUEST,
175 detail=f'Final processing failed: {error_msg}',
176 )
178 # DataLab format - content in format-specific fields
179 content_key = self.output_format.lower()
180 raw_content = poll_result.get(content_key)
181 final_result = poll_result
182 else:
183 # Self-hosted direct response - content in "output" field
184 if 'output' in result:
185 log.info('Self-hosted Marker returned direct response without polling')
186 raw_content = result.get('output')
187 final_result = result
188 else:
189 available_fields = list(result.keys()) if isinstance(result, dict) else 'non-dict response'
190 raise HTTPException(
191 status.HTTP_502_BAD_GATEWAY,
192 detail=f"Custom Marker endpoint returned success but no 'output' field found. Available fields: {available_fields}. Expected either 'request_check_url' for polling or 'output' field for direct response.",
193 )
195 if self.output_format.lower() == 'json':
196 full_text = json.dumps(raw_content, indent=2)
197 elif self.output_format.lower() in {'markdown', 'html'}:
198 full_text = str(raw_content).strip()
199 else:
200 raise HTTPException(
201 status.HTTP_400_BAD_REQUEST,
202 detail=f'Unsupported output format: {self.output_format}',
203 )
205 if not full_text:
206 raise HTTPException(
207 status.HTTP_400_BAD_REQUEST,
208 detail='Marker returned empty content',
209 )
211 marker_output_dir = os.path.join('/app/backend/data/uploads', 'marker_output')
212 os.makedirs(marker_output_dir, exist_ok=True)
214 file_ext_map = {'markdown': 'md', 'json': 'json', 'html': 'html'}
215 file_ext = file_ext_map.get(self.output_format.lower(), 'txt')
216 output_filename = f'{os.path.splitext(filename)[0]}.{file_ext}'
217 output_path = os.path.join(marker_output_dir, output_filename)
219 try:
220 with open(output_path, 'w', encoding='utf-8') as f:
221 f.write(full_text)
222 log.info('Saved Marker output to: %s', output_path)
223 except Exception as e:
224 log.warning(f'Failed to write marker output to disk: {e}')
226 metadata = {
227 'source': filename,
228 'output_format': final_result.get('output_format', self.output_format),
229 'page_count': final_result.get('page_count', 0),
230 'processed_with_llm': self.use_llm,
231 'request_id': request_id or '',
232 }
234 images = final_result.get('images', {})
235 if images:
236 metadata['image_count'] = len(images)
237 metadata['images'] = JSONCodec.dumps(list(images.keys()))
239 for k, v in metadata.items():
240 if isinstance(v, (dict, list)):
241 metadata[k] = JSONCodec.dumps(v)
242 elif v is None:
243 metadata[k] = ''
245 return [Document(page_content=full_text, metadata=metadata)]