From a515a1d40d4adf625380f5e5c900d6bb4934ccea Mon Sep 17 00:00:00 2001
From: Aafimalek <91587101+Aafimalek@users.noreply.github.com>
Date: Fri, 6 Nov 2026 18:51:18 +0530
Subject: [PATCH] fix(datasets): hydrate media under numeric object keys

---
 langfuse/_utils/json_path.py | 21 ++++++++++++++-------
 tests/unit/test_datasets.py  | 10 ++++++++++
 tests/unit/test_json_path.py | 21 +++++++++++++++++----
 3 files changed, 41 insertions(+), 11 deletions(-)

diff --git a/langfuse/_utils/json_path.py b/langfuse/_utils/json_path.py
index 8a8beba87..b3d4eb5b5 100644
--- a/langfuse/_utils/json_path.py
+++ b/langfuse/_utils/json_path.py
@@ -16,7 +16,8 @@
 def parse_path(json_path: str) -> List[Union[str, int]]:
     """Parse a jsonpath-plus normalized path into ordered segments.
 
-    Object keys become ``str`` segments, array indices become ``int`` segments.
+    Quoted keys become ``str`` segments and canonical numeric segments become
+    ``int``. Numeric segments with leading zeroes remain strings.
     Returns an empty list for the root ``$``.
     """
     if not json_path.startswith("$"):
@@ -43,7 +44,10 @@ def parse_path(json_path: str) -> List[Union[str, int]]:
                 i += 1
             if i == start or i >= n or json_path[i] != "]":
                 raise ValueError(json_path)
-            segments.append(int(json_path[start:i]))
+            digits = json_path[start:i]
+            segments.append(
+                digits if len(digits) > 1 and digits[0] == "0" else int(digits)
+            )
             i += 1
 
     return segments
@@ -61,14 +65,17 @@ def set_value_at_path(value: Any, json_path: str, replacement: Any) -> Any:
 
     target = value
     for segment in segments[:-1]:
+        if isinstance(target, dict) and isinstance(segment, int):
+            segment = str(segment)
         target = target[segment]
 
     leaf = segments[-1]
-    # JSON object keys are always strings, so an int leaf on a dict is an
-    # all-digit key that jsonpath-plus rendered ambiguously as "[0]". We can't
-    # tell it from a list index, so raise rather than add a bogus int key.
-    if isinstance(leaf, int) and isinstance(target, dict):
-        raise KeyError(json_path)
+    # jsonpath-plus uses unquoted numeric segments for both dict keys and indices.
+    if isinstance(target, dict):
+        if isinstance(leaf, int):
+            leaf = str(leaf)
+        if leaf not in target:
+            raise KeyError(json_path)
     target[leaf] = replacement
 
     return value
diff --git a/tests/unit/test_datasets.py b/tests/unit/test_datasets.py
index f1ca3fca3..98854ab8f 100644
--- a/tests/unit/test_datasets.py
+++ b/tests/unit/test_datasets.py
@@ -30,6 +30,16 @@
             "$['image']",
             lambda item: isinstance(item.input["image"], LangfuseMediaReference),
         ),
+        (
+            DatasetItemMediaReferenceField.INPUT,
+            {
+                "00": [
+                    {"2": "@@@langfuseMedia:type=image/png|id=media-id|source=bytes@@@"}
+                ]
+            },
+            "$[00][0][2]",
+            lambda item: isinstance(item.input["00"][0]["2"], LangfuseMediaReference),
+        ),
         (
             DatasetItemMediaReferenceField.EXPECTED_OUTPUT,
             ["@@@langfuseMedia:type=image/png|id=media-id|source=bytes@@@"],
diff --git a/tests/unit/test_json_path.py b/tests/unit/test_json_path.py
index 676bde096..39e15710e 100644
--- a/tests/unit/test_json_path.py
+++ b/tests/unit/test_json_path.py
@@ -18,6 +18,19 @@
     ("bracket key", {"arr[0]": _REF}, "$['arr[0]']", {"arr[0]": "__MEDIA__"}),
     ("dot key", {"a.b": _REF}, "$['a.b']", {"a.b": "__MEDIA__"}),
     ("list root", [_REF], "$[0]", ["__MEDIA__"]),
+    ("numeric key", {"0": _REF}, "$[0]", {"0": "__MEDIA__"}),
+    (
+        "leading zero key",
+        {"01": _REF, "1": "keep"},
+        "$[01]",
+        {"01": "__MEDIA__", "1": "keep"},
+    ),
+    (
+        "nested numeric key",
+        {"00": [{"2": _REF}]},
+        "$[00][0][2]",
+        {"00": [{"2": "__MEDIA__"}]},
+    ),
     (
         "list element",
         {"items": [0, _REF, 2]},
@@ -72,11 +85,11 @@ def test_set_value_at_path_matches_jsonpath_plus(value, json_path, expected):
 @pytest.mark.parametrize(
     ("value", "json_path"),
     [
-        # All-digit keys and keys containing "']" are indistinguishable / broken in
-        # jsonpath-plus' output, so they cannot be resolved and must raise (the
-        # caller leaves the value unchanged rather than guessing).
-        ({"0": _REF}, "$[0]"),
+        # Keys containing "']" are broken in jsonpath-plus' output.
         ({"a']b": _REF}, "$['a']b']"),
+        ({"1": _REF}, "$[0]"),
+        ({"01": _REF}, "$[00]"),
+        ([_REF], "$[00]"),
         # Malformed paths the API should never emit.
         ("x", "image"),
         ({"a": _REF}, "$['a'"),