langfuse._task_manager.ingestion_consumer

  1import json
  2import logging
  3import threading
  4import time
  5
  6from queue import Empty, Queue
  7from typing import Any, List, Optional
  8
  9import backoff
 10
 11try:
 12    import pydantic.v1 as pydantic
 13except ImportError:
 14    import pydantic
 15
 16from langfuse.parse_error import handle_exception
 17from langfuse.request import APIError, LangfuseClient
 18from langfuse.Sampler import Sampler
 19from langfuse.serializer import EventSerializer
 20from langfuse.types import MaskFunction
 21
 22from .media_manager import MediaManager
 23
 24MAX_EVENT_SIZE_BYTES = 1_000_000
 25MAX_BATCH_SIZE_BYTES = 2_500_000
 26
 27
 28class IngestionMetadata(pydantic.BaseModel):
 29    batch_size: int
 30    sdk_integration: str
 31    sdk_name: str
 32    sdk_version: str
 33    public_key: str
 34
 35
 36class IngestionConsumer(threading.Thread):
 37    _log = logging.getLogger("langfuse")
 38    _ingestion_queue: Queue
 39    _identifier: int
 40    _client: LangfuseClient
 41    _flush_at: int
 42    _flush_interval: float
 43    _max_retries: int
 44    _public_key: str
 45    _sdk_name: str
 46    _sdk_version: str
 47    _sdk_integration: str
 48    _mask: Optional[MaskFunction]
 49    _sampler: Sampler
 50    _media_manager: MediaManager
 51
 52    def __init__(
 53        self,
 54        *,
 55        ingestion_queue: Queue,
 56        identifier: int,
 57        client: LangfuseClient,
 58        flush_at: int,
 59        flush_interval: float,
 60        max_retries: int,
 61        public_key: str,
 62        media_manager: MediaManager,
 63        sdk_name: str,
 64        sdk_version: str,
 65        sdk_integration: str,
 66        sample_rate: float,
 67        mask: Optional[MaskFunction] = None,
 68    ):
 69        """Create a consumer thread."""
 70        super().__init__()
 71        # It's important to set running in the constructor: if we are asked to
 72        # pause immediately after construction, we might set running to True in
 73        # run() *after* we set it to False in pause... and keep running
 74        # forever.
 75        self.running = True
 76        # Make consumer a daemon thread so that it doesn't block program exit
 77        self.daemon = True
 78        self._ingestion_queue = ingestion_queue
 79        self._identifier = identifier
 80        self._client = client
 81        self._flush_at = flush_at
 82        self._flush_interval = flush_interval
 83        self._max_retries = max_retries
 84        self._public_key = public_key
 85        self._sdk_name = sdk_name
 86        self._sdk_version = sdk_version
 87        self._sdk_integration = sdk_integration
 88        self._mask = mask
 89        self._sampler = Sampler(sample_rate)
 90        self._media_manager = media_manager
 91
 92    def _next(self):
 93        """Return the next batch of items to upload."""
 94        events = []
 95
 96        start_time = time.monotonic()
 97        total_size = 0
 98
 99        while len(events) < self._flush_at:
100            elapsed = time.monotonic() - start_time
101            if elapsed >= self._flush_interval:
102                break
103            try:
104                event = self._ingestion_queue.get(
105                    block=True, timeout=self._flush_interval - elapsed
106                )
107
108                # convert pydantic models to dicts
109                if "body" in event and isinstance(event["body"], pydantic.BaseModel):
110                    event["body"] = event["body"].dict(exclude_none=True)
111
112                # sample event
113                if not self._sampler.sample_event(event):
114                    self._ingestion_queue.task_done()
115
116                    continue
117
118                # handle multimodal data
119                self._media_manager.process_media_in_event(event)
120
121                # truncate item if it exceeds size limit
122                item_size = self._truncate_item_in_place(
123                    event=event,
124                    max_size=MAX_EVENT_SIZE_BYTES,
125                    log_message="<truncated due to size exceeding limit>",
126                )
127
128                # apply mask
129                self._apply_mask_in_place(event)
130
131                # check for serialization errors
132                try:
133                    json.dumps(event, cls=EventSerializer)
134                except Exception as e:
135                    self._log.error(f"Error serializing item, skipping: {e}")
136                    self._ingestion_queue.task_done()
137
138                    continue
139
140                events.append(event)
141
142                total_size += item_size
143                if total_size >= MAX_BATCH_SIZE_BYTES:
144                    self._log.debug("hit batch size limit (size: %d)", total_size)
145                    break
146
147            except Empty:
148                break
149
150            except Exception as e:
151                self._log.warning(
152                    "Failed to process event in IngestionConsumer, skipping",
153                    exc_info=e,
154                )
155                self._ingestion_queue.task_done()
156
157        self._log.debug(
158            "~%d items in the Langfuse queue", self._ingestion_queue.qsize()
159        )
160
161        return events
162
163    def _truncate_item_in_place(
164        self,
165        *,
166        event: Any,
167        max_size: int,
168        log_message: Optional[str] = None,
169    ) -> int:
170        """Truncate the item in place to fit within the size limit."""
171        item_size = self._get_item_size(event)
172        self._log.debug(f"item size {item_size}")
173
174        if item_size > max_size:
175            self._log.warning(
176                "Item exceeds size limit (size: %s), dropping input / output / metadata of item until it fits.",
177                item_size,
178            )
179
180            if "body" in event:
181                drop_candidates = ["input", "output", "metadata"]
182                sorted_field_sizes = sorted(
183                    [
184                        (
185                            field,
186                            self._get_item_size((event["body"][field]))
187                            if field in event["body"]
188                            else 0,
189                        )
190                        for field in drop_candidates
191                    ],
192                    key=lambda x: x[1],
193                )
194
195                # drop the largest field until the item size is within the limit
196                for _ in range(len(sorted_field_sizes)):
197                    field_to_drop, size_to_drop = sorted_field_sizes.pop()
198
199                    if field_to_drop not in event["body"]:
200                        continue
201
202                    event["body"][field_to_drop] = log_message
203                    item_size -= size_to_drop
204
205                    self._log.debug(
206                        f"Dropped field {field_to_drop}, new item size {item_size}"
207                    )
208
209                    if item_size <= max_size:
210                        break
211
212            # if item does not have body or input/output fields, drop the event
213            if "body" not in event or (
214                "input" not in event["body"] and "output" not in event["body"]
215            ):
216                self._log.warning(
217                    "Item does not have body or input/output fields, dropping item."
218                )
219                self._ingestion_queue.task_done()
220                return 0
221
222        return self._get_item_size(event)
223
224    def _get_item_size(self, item: Any) -> int:
225        """Return the size of the item in bytes."""
226        return len(json.dumps(item, cls=EventSerializer).encode())
227
228    def _apply_mask_in_place(self, event: dict):
229        """Apply the mask function to the event. This is done in place."""
230        if not self._mask:
231            return
232
233        body = event["body"] if "body" in event else {}
234        for key in ("input", "output"):
235            if key in body:
236                try:
237                    body[key] = self._mask(data=body[key])
238                except Exception as e:
239                    self._log.error(f"Mask function failed with error: {e}")
240                    body[key] = "<fully masked due to failed mask function>"
241
242    def run(self):
243        """Run the consumer."""
244        self._log.debug("consumer is running...")
245        while self.running:
246            self.upload()
247
248    def upload(self):
249        """Upload the next batch of items, return whether successful."""
250        batch = self._next()
251        if len(batch) == 0:
252            return
253
254        try:
255            self._upload_batch(batch)
256        except Exception as e:
257            handle_exception(e)
258        finally:
259            # mark items as acknowledged from queue
260            for _ in batch:
261                self._ingestion_queue.task_done()
262
263    def pause(self):
264        """Pause the consumer."""
265        self.running = False
266
267    def _upload_batch(self, batch: List[Any]):
268        self._log.debug("uploading batch of %d items", len(batch))
269
270        metadata = IngestionMetadata(
271            batch_size=len(batch),
272            sdk_integration=self._sdk_integration,
273            sdk_name=self._sdk_name,
274            sdk_version=self._sdk_version,
275            public_key=self._public_key,
276        ).dict()
277
278        @backoff.on_exception(
279            backoff.expo, Exception, max_tries=self._max_retries, logger=None
280        )
281        def execute_task_with_backoff(batch: List[Any]):
282            try:
283                self._client.batch_post(batch=batch, metadata=metadata)
284            except Exception as e:
285                if (
286                    isinstance(e, APIError)
287                    and 400 <= int(e.status) < 500
288                    and int(e.status) != 429  # retry if rate-limited
289                ):
290                    return
291
292                raise e
293
294        execute_task_with_backoff(batch)
295        self._log.debug("successfully uploaded batch of %d events", len(batch))