|
@@ -15,24 +15,28 @@ class DatalabMarkerLoader:
|
|
|
self,
|
|
|
file_path: str,
|
|
|
api_key: str,
|
|
|
- langs: Optional[str] = None,
|
|
|
+ api_base_url: str,
|
|
|
+ additional_config: Optional[str] = None,
|
|
|
use_llm: bool = False,
|
|
|
skip_cache: bool = False,
|
|
|
force_ocr: bool = False,
|
|
|
paginate: bool = False,
|
|
|
strip_existing_ocr: bool = False,
|
|
|
disable_image_extraction: bool = False,
|
|
|
+ format_lines: bool = False,
|
|
|
output_format: str = None,
|
|
|
):
|
|
|
self.file_path = file_path
|
|
|
self.api_key = api_key
|
|
|
- self.langs = langs
|
|
|
+ self.api_base_url = api_base_url
|
|
|
+ self.additional_config = additional_config
|
|
|
self.use_llm = use_llm
|
|
|
self.skip_cache = skip_cache
|
|
|
self.force_ocr = force_ocr
|
|
|
self.paginate = paginate
|
|
|
self.strip_existing_ocr = strip_existing_ocr
|
|
|
self.disable_image_extraction = disable_image_extraction
|
|
|
+ self.format_lines = format_lines
|
|
|
self.output_format = output_format
|
|
|
|
|
|
def _get_mime_type(self, filename: str) -> str:
|
|
@@ -60,7 +64,7 @@ class DatalabMarkerLoader:
|
|
|
return mime_map.get(ext, "application/octet-stream")
|
|
|
|
|
|
def check_marker_request_status(self, request_id: str) -> dict:
|
|
|
- url = f"https://www.datalab.to/api/v1/marker/{request_id}"
|
|
|
+ url = f"{self.api_base_url}/{request_id}"
|
|
|
headers = {"X-Api-Key": self.api_key}
|
|
|
try:
|
|
|
response = requests.get(url, headers=headers)
|
|
@@ -81,22 +85,25 @@ class DatalabMarkerLoader:
|
|
|
)
|
|
|
|
|
|
def load(self) -> List[Document]:
|
|
|
- url = "https://www.datalab.to/api/v1/marker"
|
|
|
+ url = self.api_base_url
|
|
|
filename = os.path.basename(self.file_path)
|
|
|
mime_type = self._get_mime_type(filename)
|
|
|
headers = {"X-Api-Key": self.api_key}
|
|
|
|
|
|
form_data = {
|
|
|
- "langs": self.langs,
|
|
|
"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
|
|
|
+
|
|
|
log.info(
|
|
|
f"Datalab Marker POST request parameters: {{'filename': '{filename}', 'mime_type': '{mime_type}', **{form_data}}}"
|
|
|
)
|
|
@@ -133,74 +140,92 @@ class DatalabMarkerLoader:
|
|
|
|
|
|
check_url = result.get("request_check_url")
|
|
|
request_id = result.get("request_id")
|
|
|
- if not check_url:
|
|
|
- raise HTTPException(
|
|
|
- status.HTTP_502_BAD_GATEWAY, detail="No request_check_url returned."
|
|
|
- )
|
|
|
|
|
|
- for _ in range(300): # Up to 10 minutes
|
|
|
- time.sleep(2)
|
|
|
- try:
|
|
|
- poll_response = requests.get(check_url, headers=headers)
|
|
|
- poll_response.raise_for_status()
|
|
|
- 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}"
|
|
|
- )
|
|
|
+ # Check if this is a direct response (self-hosted) or polling response (DataLab)
|
|
|
+ if check_url:
|
|
|
+ # DataLab polling pattern
|
|
|
+ for _ in range(300): # Up to 10 minutes
|
|
|
+ time.sleep(2)
|
|
|
+ try:
|
|
|
+ poll_response = requests.get(check_url, headers=headers)
|
|
|
+ poll_response.raise_for_status()
|
|
|
+ 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}"
|
|
|
+ )
|
|
|
+
|
|
|
+ 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":
|
|
|
- summary = {
|
|
|
- k: poll_result.get(k)
|
|
|
- for k in (
|
|
|
- "status",
|
|
|
- "output_format",
|
|
|
- "success",
|
|
|
- "error",
|
|
|
- "page_count",
|
|
|
- "total_cost",
|
|
|
+ if status_val == "complete":
|
|
|
+ summary = {
|
|
|
+ k: poll_result.get(k)
|
|
|
+ for k in (
|
|
|
+ "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
|
|
|
+ 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}",
|
|
|
+ )
|
|
|
+ else:
|
|
|
+ raise HTTPException(
|
|
|
+ status.HTTP_504_GATEWAY_TIMEOUT,
|
|
|
+ detail="Marker processing timed out",
|
|
|
)
|
|
|
+
|
|
|
+ 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"Marker processing failed: {error_msg}",
|
|
|
+ detail=f"Final processing failed: {error_msg}",
|
|
|
)
|
|
|
- else:
|
|
|
- raise HTTPException(
|
|
|
- status.HTTP_504_GATEWAY_TIMEOUT, detail="Marker processing timed out"
|
|
|
- )
|
|
|
-
|
|
|
- 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}",
|
|
|
- )
|
|
|
|
|
|
- content_key = self.output_format.lower()
|
|
|
- raw_content = poll_result.get(content_key)
|
|
|
+ # DataLab format - content in format-specific fields
|
|
|
+ content_key = self.output_format.lower()
|
|
|
+ raw_content = poll_result.get(content_key)
|
|
|
+ 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")
|
|
|
+ final_result = result
|
|
|
+ else:
|
|
|
+ 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 content_key == "json":
|
|
|
+ if self.output_format.lower() == "json":
|
|
|
full_text = json.dumps(raw_content, indent=2)
|
|
|
- elif content_key in {"markdown", "html"}:
|
|
|
+ elif self.output_format.lower() in {"markdown", "html"}:
|
|
|
full_text = str(raw_content).strip()
|
|
|
else:
|
|
|
raise HTTPException(
|
|
@@ -211,14 +236,14 @@ class DatalabMarkerLoader:
|
|
|
if not full_text:
|
|
|
raise HTTPException(
|
|
|
status.HTTP_400_BAD_REQUEST,
|
|
|
- detail="Datalab Marker returned empty content",
|
|
|
+ detail="Marker returned empty content",
|
|
|
)
|
|
|
|
|
|
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(content_key, "txt")
|
|
|
+ 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)
|
|
|
|
|
@@ -231,13 +256,13 @@ class DatalabMarkerLoader:
|
|
|
|
|
|
metadata = {
|
|
|
"source": filename,
|
|
|
- "output_format": poll_result.get("output_format", self.output_format),
|
|
|
- "page_count": poll_result.get("page_count", 0),
|
|
|
+ "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 = poll_result.get("images", {})
|
|
|
+ images = final_result.get("images", {})
|
|
|
if images:
|
|
|
metadata["image_count"] = len(images)
|
|
|
metadata["images"] = json.dumps(list(images.keys()))
|