From ee49ce7e650aadd03f4ff72d586ed9afec2f97d2 Mon Sep 17 00:00:00 2001 From: Apeksha_2022 <70309060+Apeksha9389@users.noreply.github.com> Date: Mon, 19 Oct 2026 23:07:05 +0530 Subject: [PATCH] perf(media): reuse successful resolutions within a call --- langfuse/media.py | 21 ++++---- tests/unit/test_media.py | 113 +++++++++++++++++++++++++++++++++++++++ 2 files changed, 124 insertions(+), 10 deletions(-) diff --git a/langfuse/media.py b/langfuse/media.py index 6410079ab..d332103a7 100644 --- a/langfuse/media.py +++ b/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,6 +352,8 @@ def resolve_media_references( # } """ + resolved_media: dict[str, str] = {} + def traverse(obj: Any, depth: int) -> Any: if depth > max_depth: return obj @@ -357,12 +361,11 @@ def traverse(obj: Any, depth: int) -> Any: # 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 @@ -370,6 +373,8 @@ def traverse(obj: Any, depth: int) -> Any: ) 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 diff --git a/tests/unit/test_media.py b/tests/unit/test_media.py index f222df576..410d9ff19 100644 --- a/tests/unit/test_media.py +++ b/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}, + }