Allow callback handlers to opt into being run inline (#6424)

This is useful eg for callback handlers that use context vars (like open
telemetry)

See https://github.com/hwchase17/langchain/pull/6095
pull/6604/head
Nuno Campos 1 year ago committed by GitHub
parent a9108c1809
commit 74ac6fb6b9
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23

@ -194,6 +194,8 @@ class BaseCallbackHandler(
raise_error: bool = False
run_inline: bool = False
@property
def ignore_llm(self) -> bool:
"""Whether to ignore LLM callbacks."""

@ -224,9 +224,12 @@ async def _ahandle_event_for_handler(
if asyncio.iscoroutinefunction(event):
await event(*args, **kwargs)
else:
await asyncio.get_event_loop().run_in_executor(
None, functools.partial(event, *args, **kwargs)
)
if handler.run_inline:
event(*args, **kwargs)
else:
await asyncio.get_event_loop().run_in_executor(
None, functools.partial(event, *args, **kwargs)
)
except NotImplementedError as e:
if event_name == "on_chat_model_start":
message_strings = [get_buffer_string(m) for m in args[1]]
@ -259,12 +262,17 @@ async def _ahandle_event(
**kwargs: Any,
) -> None:
"""Generic event handler for AsyncCallbackManager."""
for handler in [h for h in handlers if h.run_inline]:
await _ahandle_event_for_handler(
handler, event_name, ignore_condition_name, *args, **kwargs
)
await asyncio.gather(
*(
_ahandle_event_for_handler(
handler, event_name, ignore_condition_name, *args, **kwargs
)
for handler in handlers
if not handler.run_inline
)
)

Loading…
Cancel
Save