bhetherman/open-webui
1
0 Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

refac

Browse Source
This commit is contained in:
Timothy Jaeryang Baek committed 2026-03-17 17:58:01 -05:00
1 parent fcf7208352
commit de3317e26b
220 files changed
+17200 -22836
Download Patch File Download Diff File Expand all files Collapse all files

No files matched your search

backend/open_webui/retrieval/loaders/datalab_marker.py
+91 -110
View File
@@ -40,78 +40,76 @@ class DatalabMarkerLoader:
self.output_format = output_format
def _get_mime_type(self, filename: str) -> str:
ext = filename.rsplit(".", 1)[-1].lower()
ext = filename.rsplit('.', 1)[-1].lower()
mime_map = {
"pdf": "application/pdf",
"xls": "application/vnd.ms-excel",
"xlsx": "application/vnd.openxmlformats-officedocument.spreadsheetml.sheet",
"ods": "application/vnd.oasis.opendocument.spreadsheet",
"doc": "application/msword",
"docx": "application/vnd.openxmlformats-officedocument.wordprocessingml.document",
"odt": "application/vnd.oasis.opendocument.text",
"ppt": "application/vnd.ms-powerpoint",
"pptx": "application/vnd.openxmlformats-officedocument.presentationml.presentation",
"odp": "application/vnd.oasis.opendocument.presentation",
"html": "text/html",
"epub": "application/epub+zip",
"png": "image/png",
"jpeg": "image/jpeg",
"jpg": "image/jpeg",
"webp": "image/webp",
"gif": "image/gif",
"tiff": "image/tiff",
'pdf': 'application/pdf',
'xls': 'application/vnd.ms-excel',
'xlsx': 'application/vnd.openxmlformats-officedocument.spreadsheetml.sheet',
'ods': 'application/vnd.oasis.opendocument.spreadsheet',
'doc': 'application/msword',
'docx': 'application/vnd.openxmlformats-officedocument.wordprocessingml.document',
'odt': 'application/vnd.oasis.opendocument.text',
'ppt': 'application/vnd.ms-powerpoint',
'pptx': 'application/vnd.openxmlformats-officedocument.presentationml.presentation',
'odp': 'application/vnd.oasis.opendocument.presentation',
'html': 'text/html',
'epub': 'application/epub+zip',
'png': 'image/png',
'jpeg': 'image/jpeg',
'jpg': 'image/jpeg',
'webp': 'image/webp',
'gif': 'image/gif',
'tiff': 'image/tiff',
}
return mime_map.get(ext, "application/octet-stream")
return mime_map.get(ext, 'application/octet-stream')
def check_marker_request_status(self, request_id: str) -> dict:
url = f"{self.api_base_url}/{request_id}"
headers = {"X-Api-Key": self.api_key}
url = f'{self.api_base_url}/{request_id}'
headers = {'X-Api-Key': self.api_key}
try:
response = requests.get(url, headers=headers)
response.raise_for_status()
result = response.json()
log.info(f"Marker API status check for request {request_id}: {result}")
log.info(f'Marker API status check for request {request_id}: {result}')
return result
except requests.HTTPError as e:
log.error(f"Error checking Marker request status: {e}")
log.error(f'Error checking Marker request status: {e}')
raise HTTPException(
status.HTTP_502_BAD_GATEWAY,
detail=f"Failed to check Marker request: {e}",
detail=f'Failed to check Marker request: {e}',
)
except ValueError as e:
log.error(f"Invalid JSON checking Marker request: {e}")
raise HTTPException(
status.HTTP_502_BAD_GATEWAY, detail=f"Invalid JSON: {e}"
)
log.error(f'Invalid JSON checking Marker request: {e}')
raise HTTPException(status.HTTP_502_BAD_GATEWAY, detail=f'Invalid JSON: {e}')
def load(self) -> List[Document]:
filename = os.path.basename(self.file_path)
mime_type = self._get_mime_type(filename)
headers = {"X-Api-Key": self.api_key}
headers = {'X-Api-Key': self.api_key}
form_data = {
"use_llm": str(self.use_llm).lower(),
"skip_cache": str(self.skip_cache).lower(),
"force_ocr": str(self.force_ocr).lower(),
"paginate": str(self.paginate).lower(),
"strip_existing_ocr": str(self.strip_existing_ocr).lower(),
"disable_image_extraction": str(self.disable_image_extraction).lower(),
"format_lines": str(self.format_lines).lower(),
"output_format": self.output_format,
'use_llm': str(self.use_llm).lower(),
'skip_cache': str(self.skip_cache).lower(),
'force_ocr': str(self.force_ocr).lower(),
'paginate': str(self.paginate).lower(),
'strip_existing_ocr': str(self.strip_existing_ocr).lower(),
'disable_image_extraction': str(self.disable_image_extraction).lower(),
'format_lines': str(self.format_lines).lower(),
'output_format': self.output_format,
}
if self.additional_config and self.additional_config.strip():
form_data["additional_config"] = self.additional_config
form_data['additional_config'] = self.additional_config
log.info(
f"Datalab Marker POST request parameters: {{'filename': '{filename}', 'mime_type': '{mime_type}', **{form_data}}}"
)
try:
with open(self.file_path, "rb") as f:
files = {"file": (filename, f, mime_type)}
with open(self.file_path, 'rb') as f:
files = {'file': (filename, f, mime_type)}
response = requests.post(
f"{self.api_base_url}",
f'{self.api_base_url}',
data=form_data,
files=files,
headers=headers,
@@ -119,29 +117,25 @@ class DatalabMarkerLoader:
response.raise_for_status()
result = response.json()
except FileNotFoundError:
raise HTTPException(
status.HTTP_404_NOT_FOUND, detail=f"File not found: {self.file_path}"
)
raise HTTPException(status.HTTP_404_NOT_FOUND, detail=f'File not found: {self.file_path}')
except requests.HTTPError as e:
raise HTTPException(
status.HTTP_400_BAD_REQUEST,
detail=f"Datalab Marker request failed: {e}",
detail=f'Datalab Marker request failed: {e}',
)
except ValueError as e:
raise HTTPException(
status.HTTP_502_BAD_GATEWAY, detail=f"Invalid JSON response: {e}"
)
raise HTTPException(status.HTTP_502_BAD_GATEWAY, detail=f'Invalid JSON response: {e}')
except Exception as e:
raise HTTPException(status.HTTP_500_INTERNAL_SERVER_ERROR, detail=str(e))
if not result.get("success"):
if not result.get('success'):
raise HTTPException(
status.HTTP_400_BAD_REQUEST,
detail=f"Datalab Marker request failed: {result.get('error', 'Unknown error')}",
detail=f'Datalab Marker request failed: {result.get("error", "Unknown error")}',
)
check_url = result.get("request_check_url")
request_id = result.get("request_id")
check_url = result.get('request_check_url')
request_id = result.get('request_id')
# Check if this is a direct response (self-hosted) or polling response (DataLab)
if check_url:
@@ -154,54 +148,45 @@ class DatalabMarkerLoader:
poll_result = poll_response.json()
except (requests.HTTPError, ValueError) as e:
raw_body = poll_response.text
log.error(f"Polling error: {e}, response body: {raw_body}")
raise HTTPException(
status.HTTP_502_BAD_GATEWAY, detail=f"Polling failed: {e}"
)
log.error(f'Polling error: {e}, response body: {raw_body}')
raise HTTPException(status.HTTP_502_BAD_GATEWAY, detail=f'Polling failed: {e}')
status_val = poll_result.get("status")
success_val = poll_result.get("success")
status_val = poll_result.get('status')
success_val = poll_result.get('success')
if status_val == "complete":
if status_val == 'complete':
summary = {
k: poll_result.get(k)
for k in (
"status",
"output_format",
"success",
"error",
"page_count",
"total_cost",
'status',
'output_format',
'success',
'error',
'page_count',
'total_cost',
)
}
log.info(
f"Marker processing completed successfully: {json.dumps(summary, indent=2)}"
)
log.info(f'Marker processing completed successfully: {json.dumps(summary, indent=2)}')
break
if status_val == "failed" or success_val is False:
log.error(
f"Marker poll failed full response: {json.dumps(poll_result, indent=2)}"
)
error_msg = (
poll_result.get("error")
or "Marker returned failure without error message"
)
if status_val == 'failed' or success_val is False:
log.error(f'Marker poll failed full response: {json.dumps(poll_result, indent=2)}')
error_msg = poll_result.get('error') or 'Marker returned failure without error message'
raise HTTPException(
status.HTTP_400_BAD_REQUEST,
detail=f"Marker processing failed: {error_msg}",
detail=f'Marker processing failed: {error_msg}',
)
else:
raise HTTPException(
status.HTTP_504_GATEWAY_TIMEOUT,
detail="Marker processing timed out",
detail='Marker processing timed out',
)
if not poll_result.get("success", False):
error_msg = poll_result.get("error") or "Unknown processing error"
if not poll_result.get('success', False):
error_msg = poll_result.get('error') or 'Unknown processing error'
raise HTTPException(
status.HTTP_400_BAD_REQUEST,
detail=f"Final processing failed: {error_msg}",
detail=f'Final processing failed: {error_msg}',
)
# DataLab format - content in format-specific fields
@@ -210,69 +195,65 @@ class DatalabMarkerLoader:
final_result = poll_result
else:
# Self-hosted direct response - content in "output" field
if "output" in result:
log.info("Self-hosted Marker returned direct response without polling")
raw_content = result.get("output")
if 'output' in result:
log.info('Self-hosted Marker returned direct response without polling')
raw_content = result.get('output')
final_result = result
else:
available_fields = (
list(result.keys())
if isinstance(result, dict)
else "non-dict response"
)
available_fields = list(result.keys()) if isinstance(result, dict) else 'non-dict response'
raise HTTPException(
status.HTTP_502_BAD_GATEWAY,
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.",
)
if self.output_format.lower() == "json":
if self.output_format.lower() == 'json':
full_text = json.dumps(raw_content, indent=2)
elif self.output_format.lower() in {"markdown", "html"}:
elif self.output_format.lower() in {'markdown', 'html'}:
full_text = str(raw_content).strip()
else:
raise HTTPException(
status.HTTP_400_BAD_REQUEST,
detail=f"Unsupported output format: {self.output_format}",
detail=f'Unsupported output format: {self.output_format}',
)
if not full_text:
raise HTTPException(
status.HTTP_400_BAD_REQUEST,
detail="Marker returned empty content",
detail='Marker returned empty content',
)
marker_output_dir = os.path.join("/app/backend/data/uploads", "marker_output")
marker_output_dir = os.path.join('/app/backend/data/uploads', 'marker_output')
os.makedirs(marker_output_dir, exist_ok=True)
file_ext_map = {"markdown": "md", "json": "json", "html": "html"}
file_ext = file_ext_map.get(self.output_format.lower(), "txt")
output_filename = f"{os.path.splitext(filename)[0]}.{file_ext}"
file_ext_map = {'markdown': 'md', 'json': 'json', 'html': 'html'}
file_ext = file_ext_map.get(self.output_format.lower(), 'txt')
output_filename = f'{os.path.splitext(filename)[0]}.{file_ext}'
output_path = os.path.join(marker_output_dir, output_filename)
try:
with open(output_path, "w", encoding="utf-8") as f:
with open(output_path, 'w', encoding='utf-8') as f:
f.write(full_text)
log.info(f"Saved Marker output to: {output_path}")
log.info(f'Saved Marker output to: {output_path}')
except Exception as e:
log.warning(f"Failed to write marker output to disk: {e}")
log.warning(f'Failed to write marker output to disk: {e}')
metadata = {
"source": filename,
"output_format": final_result.get("output_format", self.output_format),
"page_count": final_result.get("page_count", 0),
"processed_with_llm": self.use_llm,
"request_id": request_id or "",
'source': filename,
'output_format': final_result.get('output_format', self.output_format),
'page_count': final_result.get('page_count', 0),
'processed_with_llm': self.use_llm,
'request_id': request_id or '',
}
images = final_result.get("images", {})
images = final_result.get('images', {})
if images:
metadata["image_count"] = len(images)
metadata["images"] = json.dumps(list(images.keys()))
metadata['image_count'] = len(images)
metadata['images'] = json.dumps(list(images.keys()))
for k, v in metadata.items():
if isinstance(v, (dict, list)):
metadata[k] = json.dumps(v)
elif v is None:
metadata[k] = ""
metadata[k] = ''
return [Document(page_content=full_text, metadata=metadata)]