Return dict from ai_parse_document and keep metadata fields

Author: bohutang

databend_aiserver/udfs/docparse.py

@@ -256,29 +256,23 @@ def ai_parse_document(stage_location: StageLocation, path: str) -> Dict[str, Any
             full_path = resolved_path or path
 
         # Keep metadata first for predictable JSON ordering.
-        payload: Dict[str, Any] = OrderedDict(
-            [
-                (
-                    "metadata",
-                    {
-                        "chunk_count": chunk_count,
-                        "chunk_size": DEFAULT_CHUNK_SIZE,
-                        "duration_ms": duration_ms,
-                        "file_size": file_size if file_size is not None else 0,
-                        "filename": Path(path).name,
-                        "path": full_path or path,
-                        "timings_ms": {
-                            "convert": (t_convert_end_ns - t_convert_start_ns)
-                            / 1_000_000.0,
-                            "chunk": (t_chunk_end_ns - t_convert_end_ns) / 1_000_000.0,
-                            "total": duration_ms,
-                        },
-                        "version": 1,
-                    },
-                ),
-                ("chunks", pages),
-            ]
-        )
+        payload: Dict[str, Any] = {
+            "metadata": {
+                "chunk_count": chunk_count,
+                "chunk_size": DEFAULT_CHUNK_SIZE,
+                "duration_ms": duration_ms,
+                "file_size": file_size if file_size is not None else 0,
+                "filename": Path(path).name,
+                "path": full_path or path,
+                "timings_ms": {
+                    "convert": (t_convert_end_ns - t_convert_start_ns) / 1_000_000.0,
+                    "chunk": (t_chunk_end_ns - t_convert_end_ns) / 1_000_000.0,
+                    "total": duration_ms,
+                },
+                "version": 1,
+            },
+            "chunks": pages,
+        }
         if fallback:
             payload["error_information"] = [
                 {
@@ -296,19 +290,11 @@ def ai_parse_document(stage_location: StageLocation, path: str) -> Dict[str, Any
         )
         return payload
     except Exception as exc:  # pragma: no cover - defensive for unexpected docling errors
-        return OrderedDict(
-            [
-                (
-                    "metadata",
-                    {
-                        "path": path,
-                        "filename": Path(path).name,
-                    },
-                ),
-                ("chunks", []),
-                (
-                    "error_information",
-                    [{"message": str(exc), "type": exc.__class__.__name__}],
-                ),
-            ]
-        )
+        return {
+            "metadata": {
+                "path": path,
+                "filename": Path(path).name,
+            },
+            "chunks": [],
+            "error_information": [{"message": str(exc), "type": exc.__class__.__name__}],
+        }

tests/integration/test_docparse_integration.py

@@ -34,13 +34,13 @@ def _call_docparse(client: UDFClient, path: str, memory_stage):
     )
     assert len(result) == 1
     payload_raw = result[0]
+    if hasattr(payload_raw, "as_py"):
+        payload_raw = payload_raw.as_py()
     if isinstance(payload_raw, (bytes, bytearray)):
         payload_raw = payload_raw.decode("utf-8")
     if isinstance(payload_raw, str):
-        payload = json.loads(payload_raw)
-    else:
-        payload = payload_raw
-    return payload
+        return json.loads(payload_raw)
+    return payload_raw
 
 
 def _normalize_payload(payload):

tests/unit/test_docparse_path.py

@@ -13,10 +13,12 @@
 # limitations under the License.
 
 from databend_aiserver.udfs.docparse import ai_parse_document
+import json
 
 
 def test_docparse_metadata_path_uses_root(memory_stage_with_root):
-    payload = ai_parse_document(memory_stage_with_root, "2206.01062.pdf")
+    raw = ai_parse_document(memory_stage_with_root, "2206.01062.pdf")
+    payload = json.loads(raw) if isinstance(raw, str) else raw
     meta = payload.get("metadata", {})
     assert meta["path"] == "s3://wizardbend/dataset/data/2206.01062.pdf"
     assert meta["filename"] == "2206.01062.pdf"