perf(media): reuse downloads for repeated references within a call

langfuse/media.py

@@ -319,6 +319,8 @@ def resolve_media_references(
         If fetching media content fails for a reference string, a warning is logged and the reference
         string is left unchanged.
 
+        Successful downloads are reused for repeated references within this call.
+
         Args:
             obj: The object to process. Can be a primitive value, array, or nested object.
                 If the object has a __dict__ attribute, a dict will be returned instead of the original object type.
@@ -350,26 +352,29 @@ def resolve_media_references(
             # }
         """
 
+        resolved_media: dict[str, str] = {}
+
         def traverse(obj: Any, depth: int) -> Any:
             if depth > max_depth:
                 return obj
 
             # Handle string
             if isinstance(obj, str):
                 regex = r"@@@langfuseMedia:.+?@@@"
-                reference_string_matches = re.findall(regex, obj)
+                reference_string_matches = list(dict.fromkeys(re.findall(regex, obj)))
                 if len(reference_string_matches) == 0:
                     return obj
 
                 result = obj
-                reference_string_to_media_content = {}
                 httpx_client = (
                     langfuse_client._resources.httpx_client
                     if langfuse_client._resources is not None
                     else None
                 )
 
                 for reference_string in reference_string_matches:
+                    if reference_string in resolved_media:
+                        continue
                     try:
                         parsed_media_reference = LangfuseMedia.parse_reference_string(
                             reference_string
@@ -394,9 +399,7 @@ def traverse(obj: Any, depth: int) -> Any:
                         ).decode()
                         base64_data_uri = f"data:{media_data.content_type};base64,{base64_media_content}"
 
-                        reference_string_to_media_content[reference_string] = (
-                            base64_data_uri
-                        )
+                        resolved_media[reference_string] = base64_data_uri
                     except Exception as e:
                         logger.warning(
                             "Error fetching media content for reference string %s: %s",
@@ -406,11 +409,9 @@ def traverse(obj: Any, depth: int) -> Any:
                         # Do not replace the reference string if there's an error
                         continue
 
-                for (
-                    ref_str,
-                    media_content_str,
-                ) in reference_string_to_media_content.items():
-                    result = result.replace(ref_str, media_content_str)
+                for ref_str in reference_string_matches:
+                    if ref_str in resolved_media:
+                        result = result.replace(ref_str, resolved_media[ref_str])
 
                 return result
 

tests/unit/test_media.py

@@ -275,3 +275,116 @@ def test_init_with_urlsafe_base64_data_uri():
     assert media._source == "base64_data_uri"
     assert media._content_type == "application/octet-stream"
     assert media._content_bytes == original_bytes
+
+
+@pytest.fixture
+def media_resolution_client():
+    media_api = Mock()
+    media_api.get.return_value = SimpleNamespace(
+        url="https://example.com/image.jpg", content_type="image/jpeg"
+    )
+    response = Mock(content=b"image")
+    httpx_client = Mock()
+    httpx_client.get.return_value = response
+    return SimpleNamespace(
+        api=SimpleNamespace(media=media_api),
+        _resources=SimpleNamespace(httpx_client=httpx_client),
+    )
+
+
+def test_resolve_media_references_downloads_repeated_media_once(
+    media_resolution_client,
+):
+    client = media_resolution_client
+    reference = "@@@langfuseMedia:type=image/jpeg|id=repeated|source=bytes@@@"
+    original = {"text": f"{reference} {reference}", "messages": [reference]}
+
+    result = LangfuseMedia.resolve_media_references(
+        obj=original, langfuse_client=client, resolve_with="base64_data_uri"
+    )
+
+    uri = "data:image/jpeg;base64,aW1hZ2U="
+    assert result == {"text": f"{uri} {uri}", "messages": [uri]}
+    assert original == {"text": f"{reference} {reference}", "messages": [reference]}
+    client.api.media.get.assert_called_once_with("repeated")
+    client._resources.httpx_client.get.assert_called_once_with(
+        "https://example.com/image.jpg", timeout=10
+    )
+
+    LangfuseMedia.resolve_media_references(
+        obj=original, langfuse_client=client, resolve_with="base64_data_uri"
+    )
+    assert client.api.media.get.call_count == 2
+    assert client._resources.httpx_client.get.call_count == 2
+
+
+def test_resolve_media_references_keeps_distinct_media_separate(
+    media_resolution_client,
+):
+    client = media_resolution_client
+    references = [
+        f"@@@langfuseMedia:type=image/jpeg|id={media_id}|source=bytes@@@"
+        for media_id in ("first", "second")
+    ]
+    client._resources.httpx_client.get.side_effect = [
+        Mock(content=b"first"),
+        Mock(content=b"second"),
+    ]
+
+    result = LangfuseMedia.resolve_media_references(
+        obj=references + references,
+        langfuse_client=client,
+        resolve_with="base64_data_uri",
+    )
+
+    assert (
+        result
+        == [
+            "data:image/jpeg;base64,Zmlyc3Q=",
+            "data:image/jpeg;base64,c2Vjb25k",
+        ]
+        * 2
+    )
+    assert client.api.media.get.call_count == 2
+    assert client._resources.httpx_client.get.call_count == 2
+
+
+@pytest.mark.parametrize("failure_stage", ["metadata", "download"])
+def test_resolve_media_references_retries_failed_media(
+    media_resolution_client, failure_stage
+):
+    client = media_resolution_client
+    reference = "@@@langfuseMedia:type=image/jpeg|id=repeated|source=bytes@@@"
+    if failure_stage == "metadata":
+        client.api.media.get.side_effect = [
+            RuntimeError("temporary failure"),
+            client.api.media.get.return_value,
+        ]
+    else:
+        client._resources.httpx_client.get.side_effect = [
+            RuntimeError("temporary failure"),
+            client._resources.httpx_client.get.return_value,
+        ]
+
+    result = LangfuseMedia.resolve_media_references(
+        obj=[reference] * 3,
+        langfuse_client=client,
+        resolve_with="base64_data_uri",
+    )
+
+    assert result == [reference] + ["data:image/jpeg;base64,aW1hZ2U="] * 2
+    assert client.api.media.get.call_count == 2
+
+
+def test_resolve_media_references_retains_depth_limit(media_resolution_client):
+    reference = "@@@langfuseMedia:type=image/jpeg|id=repeated|source=bytes@@@"
+    result = LangfuseMedia.resolve_media_references(
+        obj={"image": reference, "nested": {"image": reference}},
+        langfuse_client=media_resolution_client,
+        resolve_with="base64_data_uri",
+        max_depth=1,
+    )
+    assert result == {
+        "image": "data:image/jpeg;base64,aW1hZ2U=",
+        "nested": {"image": reference},
+    }