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
This commit is contained in:
Nuno Campos 2023-06-22 19:36:19 +01:00 committed by GitHub
parent a9108c1809
commit 74ac6fb6b9
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
2 changed files with 13 additions and 3 deletions

View File

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

View File

@ -223,6 +223,9 @@ async def _ahandle_event_for_handler(
event = getattr(handler, event_name) event = getattr(handler, event_name)
if asyncio.iscoroutinefunction(event): if asyncio.iscoroutinefunction(event):
await event(*args, **kwargs) await event(*args, **kwargs)
else:
if handler.run_inline:
event(*args, **kwargs)
else: else:
await asyncio.get_event_loop().run_in_executor( await asyncio.get_event_loop().run_in_executor(
None, functools.partial(event, *args, **kwargs) None, functools.partial(event, *args, **kwargs)
@ -259,12 +262,17 @@ async def _ahandle_event(
**kwargs: Any, **kwargs: Any,
) -> None: ) -> None:
"""Generic event handler for AsyncCallbackManager.""" """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( await asyncio.gather(
*( *(
_ahandle_event_for_handler( _ahandle_event_for_handler(
handler, event_name, ignore_condition_name, *args, **kwargs handler, event_name, ignore_condition_name, *args, **kwargs
) )
for handler in handlers for handler in handlers
if not handler.run_inline
) )
) )