langfuse._task_manager.media_manager

  1import logging
  2import time
  3from queue import Empty, Queue
  4from typing import Any, Callable, Optional, TypeVar
  5
  6import backoff
  7import requests
  8from typing_extensions import ParamSpec
  9
 10from langfuse.api import GetMediaUploadUrlRequest, PatchMediaBody
 11from langfuse.api.client import FernLangfuse
 12from langfuse.api.core import ApiError
 13from langfuse.media import LangfuseMedia
 14from langfuse.utils import _get_timestamp
 15
 16from .media_upload_queue import UploadMediaJob
 17
 18T = TypeVar("T")
 19P = ParamSpec("P")
 20
 21
 22class MediaManager:
 23    _log = logging.getLogger(__name__)
 24
 25    def __init__(
 26        self,
 27        *,
 28        api_client: FernLangfuse,
 29        media_upload_queue: Queue,
 30        max_retries: Optional[int] = 3,
 31    ):
 32        self._api_client = api_client
 33        self._queue = media_upload_queue
 34        self._max_retries = max_retries
 35
 36    def process_next_media_upload(self):
 37        try:
 38            upload_job = self._queue.get(block=True, timeout=1)
 39            self._log.debug(f"Processing upload for {upload_job['media_id']}")
 40            self._process_upload_media_job(data=upload_job)
 41
 42            self._queue.task_done()
 43        except Empty:
 44            self._log.debug("Media upload queue is empty")
 45            pass
 46        except Exception as e:
 47            self._log.error(f"Error uploading media: {e}")
 48            self._queue.task_done()
 49
 50    def process_media_in_event(self, event: dict):
 51        try:
 52            if "body" not in event:
 53                return
 54
 55            body = event["body"]
 56            trace_id = body.get("traceId", None) or (
 57                body.get("id", None)
 58                if "type" in event and "trace" in event["type"]
 59                else None
 60            )
 61
 62            if trace_id is None:
 63                raise ValueError("trace_id is required for media upload")
 64
 65            observation_id = (
 66                body.get("id", None)
 67                if "type" in event
 68                and ("generation" in event["type"] or "span" in event["type"])
 69                else None
 70            )
 71
 72            multimodal_fields = ["input", "output", "metadata"]
 73
 74            for field in multimodal_fields:
 75                if field in body:
 76                    processed_data = self._find_and_process_media(
 77                        data=body[field],
 78                        trace_id=trace_id,
 79                        observation_id=observation_id,
 80                        field=field,
 81                    )
 82
 83                    body[field] = processed_data
 84
 85        except Exception as e:
 86            self._log.error(f"Error processing multimodal event: {e}")
 87
 88    def _find_and_process_media(
 89        self,
 90        *,
 91        data: Any,
 92        trace_id: str,
 93        observation_id: Optional[str],
 94        field: str,
 95    ):
 96        seen = set()
 97        max_levels = 10
 98
 99        def _process_data_recursively(data: Any, level: int):
100            if id(data) in seen or level > max_levels:
101                return data
102
103            seen.add(id(data))
104
105            if isinstance(data, LangfuseMedia):
106                self._process_media(
107                    media=data,
108                    trace_id=trace_id,
109                    observation_id=observation_id,
110                    field=field,
111                )
112
113                return data
114
115            if isinstance(data, str) and data.startswith("data:"):
116                media = LangfuseMedia(
117                    obj=data,
118                    base64_data_uri=data,
119                )
120
121                self._process_media(
122                    media=media,
123                    trace_id=trace_id,
124                    observation_id=observation_id,
125                    field=field,
126                )
127
128                return media
129
130            # Anthropic
131            if (
132                isinstance(data, dict)
133                and "type" in data
134                and data["type"] == "base64"
135                and "media_type" in data
136                and "data" in data
137            ):
138                media = LangfuseMedia(
139                    base64_data_uri=f"data:{data['media_type']};base64," + data["data"],
140                )
141
142                self._process_media(
143                    media=media,
144                    trace_id=trace_id,
145                    observation_id=observation_id,
146                    field=field,
147                )
148
149                data["data"] = media
150
151                return data
152
153            # Vertex
154            if (
155                isinstance(data, dict)
156                and "type" in data
157                and data["type"] == "media"
158                and "mime_type" in data
159                and "data" in data
160            ):
161                media = LangfuseMedia(
162                    base64_data_uri=f"data:{data['mime_type']};base64," + data["data"],
163                )
164
165                self._process_media(
166                    media=media,
167                    trace_id=trace_id,
168                    observation_id=observation_id,
169                    field=field,
170                )
171
172                data["data"] = media
173
174                return data
175
176            if isinstance(data, list):
177                return [_process_data_recursively(item, level + 1) for item in data]
178
179            if isinstance(data, dict):
180                return {
181                    key: _process_data_recursively(value, level + 1)
182                    for key, value in data.items()
183                }
184
185            return data
186
187        return _process_data_recursively(data, 1)
188
189    def _process_media(
190        self,
191        *,
192        media: LangfuseMedia,
193        trace_id: str,
194        observation_id: Optional[str],
195        field: str,
196    ):
197        if (
198            media._content_length is None
199            or media._content_type is None
200            or media._content_sha256_hash is None
201            or media._content_bytes is None
202        ):
203            return
204
205        upload_url_response = self._request_with_backoff(
206            self._api_client.media.get_upload_url,
207            request=GetMediaUploadUrlRequest(
208                contentLength=media._content_length,
209                contentType=media._content_type,
210                sha256Hash=media._content_sha256_hash,
211                field=field,
212                traceId=trace_id,
213                observationId=observation_id,
214            ),
215        )
216
217        upload_url = upload_url_response.upload_url
218        media._media_id = upload_url_response.media_id  # Important as this is will be used in the media reference string in serializer
219
220        if upload_url is not None:
221            self._log.debug(f"Scheduling upload for {media._media_id}")
222            self._queue.put(
223                item={
224                    "upload_url": upload_url,
225                    "media_id": media._media_id,
226                    "content_bytes": media._content_bytes,
227                    "content_type": media._content_type,
228                    "content_sha256_hash": media._content_sha256_hash,
229                },
230                block=True,
231                timeout=1,
232            )
233
234        else:
235            self._log.debug(f"Media {media._media_id} already uploaded")
236
237    def _process_upload_media_job(
238        self,
239        *,
240        data: UploadMediaJob,
241    ):
242        upload_start_time = time.time()
243        upload_response = self._request_with_backoff(
244            requests.put,
245            data["upload_url"],
246            headers={
247                "Content-Type": data["content_type"],
248                "x-amz-checksum-sha256": data["content_sha256_hash"],
249                "x-ms-blob-type": "BlockBlob",
250            },
251            data=data["content_bytes"],
252        )
253        upload_time_ms = int((time.time() - upload_start_time) * 1000)
254
255        self._request_with_backoff(
256            self._api_client.media.patch,
257            media_id=data["media_id"],
258            request=PatchMediaBody(
259                uploadedAt=_get_timestamp(),
260                uploadHttpStatus=upload_response.status_code,
261                uploadHttpError=upload_response.text,
262                uploadTimeMs=upload_time_ms,
263            ),
264        )
265
266        self._log.debug(
267            f"Media upload completed for {data['media_id']} in {upload_time_ms}ms"
268        )
269
270    def _request_with_backoff(
271        self, func: Callable[P, T], *args: P.args, **kwargs: P.kwargs
272    ) -> T:
273        def _should_give_up(e: Exception) -> bool:
274            if isinstance(e, ApiError):
275                return (
276                    e.status_code is not None
277                    and 400 <= e.status_code < 500
278                    and e.status_code != 429
279                )
280            if isinstance(e, requests.exceptions.RequestException):
281                return (
282                    e.response is not None
283                    and e.response.status_code < 500
284                    and e.response.status_code != 429
285                )
286            return False
287
288        @backoff.on_exception(
289            backoff.expo,
290            Exception,
291            max_tries=self._max_retries,
292            giveup=_should_give_up,
293            logger=None,
294        )
295        def execute_task_with_backoff() -> T:
296            return func(*args, **kwargs)
297
298        return execute_task_with_backoff()