|
| 1 | +"""CrewAI patchers — one ``CallbackPatcher`` that registers the listener.""" |
| 2 | + |
| 3 | +from typing import Any, ClassVar |
| 4 | + |
| 5 | +from braintrust.integrations.base import CallbackPatcher |
| 6 | + |
| 7 | + |
| 8 | +# Module-level cache for the single registered listener instance. Keeping |
| 9 | +# it at module scope (rather than on the integration class) means |
| 10 | +# ``setup_crewai`` / ``patch_crewai`` / ``auto_instrument`` all see the |
| 11 | +# same ``BraintrustCrewAIListener`` regardless of entry point. |
| 12 | +_LISTENER: Any | None = None |
| 13 | + |
| 14 | + |
| 15 | +def _unregister_event_handler(event_bus: Any, event_type: Any, handler: Any) -> None: |
| 16 | + """Best-effort wrapper around ``CrewAIEventsBus.off``. |
| 17 | +
|
| 18 | + Pylint cannot currently infer the dynamically-generated event-bus API in |
| 19 | + CrewAI, so we use ``getattr`` here instead of calling ``off`` directly. |
| 20 | + """ |
| 21 | + unregister = getattr(event_bus, "off", None) |
| 22 | + if callable(unregister): |
| 23 | + unregister(event_type, handler) |
| 24 | + |
| 25 | + |
| 26 | +def _register_braintrust_listener() -> bool: |
| 27 | + """Idempotently create and register the Braintrust listener. |
| 28 | +
|
| 29 | + The listener subclasses :class:`crewai.events.BaseEventListener`, whose |
| 30 | + ``__init__`` registers handlers on the process-singleton |
| 31 | + ``crewai_event_bus``. We cache the instance at module scope so repeat |
| 32 | + calls (e.g. ``setup_crewai()`` followed by ``auto_instrument()``) do |
| 33 | + not register a second listener. |
| 34 | + """ |
| 35 | + global _LISTENER # noqa: PLW0603 |
| 36 | + if _LISTENER is not None: |
| 37 | + return True |
| 38 | + |
| 39 | + # Lazy import: CrewAI may not be installed in the environment. |
| 40 | + from .tracing import BraintrustCrewAIListener |
| 41 | + |
| 42 | + _LISTENER = BraintrustCrewAIListener() |
| 43 | + return True |
| 44 | + |
| 45 | + |
| 46 | +def _listener_registered() -> bool: |
| 47 | + """Return whether the Braintrust listener is currently registered.""" |
| 48 | + return _LISTENER is not None |
| 49 | + |
| 50 | + |
| 51 | +def _get_registered_listener() -> Any | None: |
| 52 | + """Return the registered listener, or ``None`` when setup has not run.""" |
| 53 | + return _LISTENER |
| 54 | + |
| 55 | + |
| 56 | +def _reset_for_testing() -> None: |
| 57 | + """Unregister the Braintrust listener and forget cached runtime state. |
| 58 | +
|
| 59 | + Intended for pytest fixtures that need to restart from a clean slate. |
| 60 | + Safe to call when CrewAI is not importable or nothing has been |
| 61 | + registered. Not part of the public API. |
| 62 | + """ |
| 63 | + global _LISTENER # noqa: PLW0603 |
| 64 | + if _LISTENER is None: |
| 65 | + return |
| 66 | + |
| 67 | + try: |
| 68 | + from crewai.events.event_bus import crewai_event_bus |
| 69 | + except ImportError: |
| 70 | + _LISTENER = None |
| 71 | + return |
| 72 | + |
| 73 | + for event_type, handlers in list(crewai_event_bus._sync_handlers.items()): |
| 74 | + for handler in list(handlers): |
| 75 | + handler_mod = getattr(handler, "__module__", "") |
| 76 | + if "braintrust" in handler_mod and "crewai" in handler_mod: |
| 77 | + _unregister_event_handler(crewai_event_bus, event_type, handler) |
| 78 | + |
| 79 | + _LISTENER = None |
| 80 | + |
| 81 | + # Clear the runtime subclass cache so the next setup rebuilds it; this |
| 82 | + # matters for tests that monkey-patch the listener base class. |
| 83 | + from .tracing import BraintrustCrewAIListener |
| 84 | + |
| 85 | + BraintrustCrewAIListener._cls = None |
| 86 | + |
| 87 | + |
| 88 | +class EventBusPatcher(CallbackPatcher): |
| 89 | + """Register :class:`BraintrustCrewAIListener` on ``crewai_event_bus``. |
| 90 | +
|
| 91 | + The target module check gates this patcher on the event-bus module being |
| 92 | + importable, not the top-level ``crewai`` package. That means users who |
| 93 | + install only a CrewAI fork missing the event-bus surface get a clean |
| 94 | + skip rather than an import error during setup. |
| 95 | + """ |
| 96 | + |
| 97 | + name: ClassVar[str] = "crewai.event_bus" |
| 98 | + target_module: ClassVar[str] = "crewai.events.event_bus" |
| 99 | + callback: ClassVar[Any] = staticmethod(_register_braintrust_listener) |
| 100 | + state_getter: ClassVar[Any] = staticmethod(_listener_registered) |
0 commit comments