Request: tracer equivalent of AsyncCallbackHandler

[jamesbraza]

Checked

• I searched existing ideas and did not find a similar one
• I added a very descriptive title
• I've clearly described the feature request and motivation for it

Feature request

A BaseTracer subclass that works with with asyncio, in langchain-core.

Think AsyncCallbackHandler converted to a tracer, named perhaps AsyncTracer.

Motivation

I am migrating my code base from BaseCallbackHandler to BaseTracer. I have an AsyncCallbackHandler I would like to migrate, but there is not an async subclass for BaseTracer.

Proposal (If applicable)

No response

[tgram-3D]

I second this FR. One potential use case is with RunnableWithMessageHistory. I have a RAG app that uses Beanie so I got it in my head it'd be a good idea to make an async version of MongoDBChatMessageHistory just for consistency's sake. It works fine if you handle messages manually but if you try to plug in an async ChatMessageHistory to RunnableWithMessageHistory it doesn't work since the _exit_history method callback is attached with the Runnable with_listeners method, which binds a RootListenersTracer as the callback handler:

 bound = (
            history_chain | runnable.with_listeners(on_end=self._exit_history)
        ).with_config(run_name="RunnableWithMessageHistory")

I'm sure there's an easier way, but one workaround for now is to rewrite the entire BaseTracer class as async then make a corresponding async RootListenersTracer.

class AsyncBaseTracer(AsyncCallbackHandler, ABC):
    """Async Base interface for tracers."""
    def __init__(
        self,
        *,
        _schema_format: Literal["original", "streaming_events"] = "original",
        **kwargs: Any,
    ) -> None:
       
        super().__init__(**kwargs)
        self._schema_format = _schema_format  # For internal use only API will change.
        self.run_map: Dict[str, Run] = {}

Example method:

async def _end_trace(self, run: Run) -> None:
        """End a trace for a run."""
        if not run.parent_run_id:
            self._persist_run(run)
        else:
            parent_run = self.run_map.get(str(run.parent_run_id))
            if parent_run is None:
                logger.debug(f"Parent run with UUID {run.parent_run_id} not found.")
            elif (
                run.child_execution_order is not None
                and parent_run.child_execution_order is not None
                and run.child_execution_order > parent_run.child_execution_order
            ):
                parent_run.child_execution_order = run.child_execution_order
        self.run_map.pop(str(run.id))
        await self._on_run_update(run)

And for the RootListenersTracer same thing just subclass your AsyncBaseTracer:

class AsyncRootListenersTracer(AsyncBaseTracer):
  
     # Just copy the init method of RootListenersTracer

Example method:

async def _on_run_update(self, run: Run) -> None:
        if run.id != self.root_id:
            return
        if run.error is None:
            if self._arg_on_end is not None:
                await call_func_with_variable_args(self._arg_on_end, run, self.config)
        else:
            if self._arg_on_error is not None:
                await call_func_with_variable_args(self._arg_on_error, run, self.config)

This lets you pass async Listener functions as your args. To use this with RunnableWithMessageHistory, I had to rewrite that entire class as async too, and change the init method to use the AsyncRootListenersTracer instead of the Runnable with_listeners method:

class AsyncRunnableWithMessageHistory(RunnableBindingBase):

...

modified_runnable = RunnableBinding(
            bound=runnable,
            config_factories=[
                lambda config: {
                    "callbacks": [
                        AsyncRootListenersTracer(
                            config=config,
                            on_end=self._aexit_history
                        )
                    ],
                }
            ],
        )

bound = (
            history_chain | modified_runnable
  ).with_config(run_name="RunnableWithMessageHistory")

Then finally, you can use an async add_messages function in your exit method:

async def _aexit_history(self, run: Run, config: RunnableConfig) -> None:
        hist: BaseChatMessageHistory = config["configurable"]["message_history"]
        inputs = load(run.inputs)
        input_val = inputs[self.input_messages_key or "input"]
        input_messages = self._get_input_messages(input_val)

        if not self.history_messages_key:
            historic_messages = config["configurable"]["message_history"].messages
            input_messages = input_messages[len(historic_messages):]

        output_val = load(run.outputs)
        output_messages = self._get_output_messages(output_val)
        # Await the asynchronous add_messages method
        await hist.add_messages(input_messages + output_messages)

This approach does run at least. The callbacks are triggered successfully and the chat history is added to Mongo with Beanie/motor. I figured since BaseTracer inherits from BaseCallbackHandler an async version should work the same as a standard AsyncCallbackHandler, with the coros getting run by the manager with ahandle_event, but I'm not sure if there are any undesirable side-effects from doing this sort of thing. I was just curious if I could get it to work.