add media

Author: hassiebp

langfuse/_task_manager/media_manager.py

@@ -1,7 +1,9 @@
+import base64
+import hashlib
 import logging
 import time
-from queue import Empty, Queue
-from typing import Any, Callable, Optional, TypeVar
+from queue import Empty, Full, Queue
+from typing import Any, Callable, Optional, TypeVar, cast
 
 import backoff
 import requests
@@ -10,6 +12,7 @@
 from langfuse.api import GetMediaUploadUrlRequest, PatchMediaBody
 from langfuse.api.client import FernLangfuse
 from langfuse.api.core import ApiError
+from langfuse.api.resources.media.types.media_content_type import MediaContentType
 from langfuse.media import LangfuseMedia
 from langfuse.utils import _get_timestamp
 
@@ -20,7 +23,7 @@
 
 
 class MediaManager:
-    _log = logging.getLogger(__name__)
+    _log = logging.getLogger("langfuse")
 
     def __init__(
         self,
@@ -92,7 +95,11 @@ def _find_and_process_media(
         trace_id: str,
         observation_id: Optional[str],
         field: str,
+        project_id: Optional[str],
     ):
+        if not project_id:
+            return data
+
         seen = set()
         max_levels = 10
 
@@ -108,6 +115,7 @@ def _process_data_recursively(data: Any, level: int):
                     trace_id=trace_id,
                     observation_id=observation_id,
                     field=field,
+                    project_id=project_id,
                 )
 
                 return data
@@ -123,6 +131,7 @@ def _process_data_recursively(data: Any, level: int):
                     trace_id=trace_id,
                     observation_id=observation_id,
                     field=field,
+                    project_id=project_id,
                 )
 
                 return media
@@ -144,6 +153,7 @@ def _process_data_recursively(data: Any, level: int):
                     trace_id=trace_id,
                     observation_id=observation_id,
                     field=field,
+                    project_id=project_id,
                 )
 
                 data["data"] = media
@@ -167,6 +177,7 @@ def _process_data_recursively(data: Any, level: int):
                     trace_id=trace_id,
                     observation_id=observation_id,
                     field=field,
+                    project_id=project_id,
                 )
 
                 data["data"] = media
@@ -193,6 +204,7 @@ def _process_media(
         trace_id: str,
         observation_id: Optional[str],
         field: str,
+        project_id: str,
     ):
         if (
             media._content_length is None
@@ -202,47 +214,83 @@ def _process_media(
         ):
             return
 
-        upload_url_response = self._request_with_backoff(
-            self._api_client.media.get_upload_url,
-            request=GetMediaUploadUrlRequest(
-                contentLength=media._content_length,
-                contentType=media._content_type,
-                sha256Hash=media._content_sha256_hash,
-                field=field,
-                traceId=trace_id,
-                observationId=observation_id,
-            ),
+        # Important as this is will be used in the media reference string in serializer
+        media._media_id = self._get_media_id(
+            project_id=project_id, content_sha256_hash=media._content_sha256_hash
         )
 
-        upload_url = upload_url_response.upload_url
-        media._media_id = upload_url_response.media_id  # Important as this is will be used in the media reference string in serializer
+        try:
+            upload_media_job = UploadMediaJob(
+                media_id=media._media_id,
+                content_bytes=media._content_bytes,
+                content_type=media._content_type,
+                content_length=media._content_length,
+                content_sha256_hash=media._content_sha256_hash,
+                trace_id=trace_id,
+                observation_id=observation_id,
+                field=field,
+            )
 
-        if upload_url is not None:
-            self._log.debug(f"Scheduling upload for {media._media_id}")
             self._queue.put(
-                item={
-                    "upload_url": upload_url,
-                    "media_id": media._media_id,
-                    "content_bytes": media._content_bytes,
-                    "content_type": media._content_type,
-                    "content_sha256_hash": media._content_sha256_hash,
-                },
-                block=True,
-                timeout=1,
+                item=upload_media_job,
+                block=False,
+            )
+            self._log.debug(
+                f"Enqueued media ID {media._media_id} for upload processing"
+            )
+
+        except Full:
+            self._log.debug(
+                f"Media queue is full. Failed to process media id {media._media_id}"
+            )
+
+        except Exception as e:
+            self._log.debug(
+                f"Failed to process media with id {media._media_id}: {str(e)}"
             )
 
-        else:
-            self._log.debug(f"Media {media._media_id} already uploaded")
+    def _get_media_id(self, *, project_id: str, content_sha256_hash) -> str:
+        hash_obj = hashlib.sha256()
+        hash_obj.update((project_id + content_sha256_hash).encode("utf-8"))
+        media_id = base64.urlsafe_b64encode(hash_obj.digest()).decode("utf-8")[:22]
+
+        return media_id
 
     def _process_upload_media_job(
         self,
         *,
         data: UploadMediaJob,
     ):
+        upload_url_response = self._request_with_backoff(
+            self._api_client.media.get_upload_url,
+            request=GetMediaUploadUrlRequest(
+                contentLength=data["content_length"],
+                contentType=cast(MediaContentType, data["content_type"]),
+                sha256Hash=data["content_sha256_hash"],
+                field=data["field"],
+                traceId=data["trace_id"],
+                observationId=data["observation_id"],
+            ),
+        )
+
+        upload_url = upload_url_response.upload_url
+
+        if not upload_url:
+            self._log.debug(f"Media with ID {data['media_id']} already uploaded.")
+
+            return
+
+        if upload_url_response.media_id != data["media_id"]:
+            self._log.error(
+                f"Media ID mismatch: SDK {data['media_id']} vs Server {upload_url_response.media_id}. Upload cancelled."
+            )
+
+            return
+
         upload_start_time = time.time()
         upload_response = self._request_with_backoff(
             requests.put,
-            data["upload_url"],
+            upload_url,
             headers={
                 "Content-Type": data["content_type"],
                 "x-amz-checksum-sha256": data["content_sha256_hash"],

langfuse/_task_manager/media_upload_queue.py

@@ -1,9 +1,12 @@
-from typing import TypedDict
+from typing import Optional, TypedDict
 
 
 class UploadMediaJob(TypedDict):
-    upload_url: str
     media_id: str
     content_type: str
+    content_length: int
     content_bytes: bytes
     content_sha256_hash: str
+    trace_id: str
+    observation_id: Optional[str]
+    field: str

langfuse/media.py

@@ -5,8 +5,9 @@
 import logging
 import os
 import re
+from typing import Any, Literal, Optional, Tuple, TypeVar, cast
+
 import requests
-from typing import Optional, cast, Tuple, Any, TypeVar, Literal
 
 from langfuse.api import MediaContentType
 from langfuse.types import ParsedMediaReference