From 74ac6fb6b9792dfe89ce9255b3320f1c86b9b927 Mon Sep 17 00:00:00 2001
From: Nuno Campos <nuno@boringbits.io>
Date: Thu, 22 Jun 2023 19:36:19 +0100
Subject: [PATCH] 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
---
 langchain/callbacks/base.py    |  2 ++
 langchain/callbacks/manager.py | 14 +++++++++++---
 2 files changed, 13 insertions(+), 3 deletions(-)

diff --git a/langchain/callbacks/base.py b/langchain/callbacks/base.py
index 2bdf22ff6f3..2f4ed83dc97 100644
--- a/langchain/callbacks/base.py
+++ b/langchain/callbacks/base.py
@@ -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."""
diff --git a/langchain/callbacks/manager.py b/langchain/callbacks/manager.py
index b438550b692..1f73c029c0d 100644
--- a/langchain/callbacks/manager.py
+++ b/langchain/callbacks/manager.py
@@ -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
         )
     )