Class PromptfooTracingProcessor

examples/openai-agents/promptfoo_tracing.py:303–360

Buffer spans per trace and export them once the workflow finishes.

class PromptfooTracingProcessor(TracingProcessor):
    """Buffer spans per trace and export them once the workflow finishes."""

    def __init__(self, exporter: PromptfooOTLPExporter) -> None:
        self._exporter = exporter
        self._lock = threading.Lock()
        self._traces: dict[str, Trace] = {}
        self._spans_by_trace: dict[str, list[Span[Any]]] = {}

    def on_trace_start(self, trace: Trace) -> None:
        with self._lock:
            self._traces[trace.trace_id] = trace

    def on_trace_end(self, trace: Trace) -> None:
        with self._lock:
            spans = list(self._spans_by_trace.pop(trace.trace_id, []))
            self._traces.pop(trace.trace_id, None)

        try:
            self._exporter.export([trace, *spans])
        except Exception as exc:
            print(
                f"[promptfoo_tracing] Failed to export trace {trace.trace_id}: {exc}",
                file=sys.stderr,
            )

    def on_span_start(self, span: Span[Any]) -> None:
        return None

    def on_span_end(self, span: Span[Any]) -> None:
        with self._lock:
            self._spans_by_trace.setdefault(span.trace_id, []).append(span)

    def shutdown(self) -> None:
        self.force_flush()

    def force_flush(self) -> None:
        with self._lock:
            pending_trace_ids = set(self._traces) | set(self._spans_by_trace)
            pending_batches = []
            for trace_id in pending_trace_ids:
                trace = self._traces.pop(trace_id, None)
                spans = list(self._spans_by_trace.pop(trace_id, []))
                pending_batches.append((trace, spans))

        for trace, spans in pending_batches:
            if trace is not None or spans:
                items: list[Trace | Span[Any]] = [*spans]
                if trace is not None:
                    items.insert(0, trace)
                try:
                    self._exporter.export(items)
                except Exception as exc:
                    print(
                        f"[promptfoo_tracing] Failed to flush "
                        f"{len(spans)} span(s): {exc}",
                        file=sys.stderr,
                    )