How do OpenAI callback handler work ?

[mspronesti]

I'm specifically looking at this snippet

from langchain.callbacks import get_openai_callback
from langchain.llms import OpenAI

llm = OpenAI(temperature=0)
with get_openai_callback() as cb:
    llm("What is the square root of 4?")
    print(cb)

My understanding is that get_openai_callback() yields an object of type OpenAICallbackHandler here whose on_llm_end method is called in the generate method of the OpenAI class here (from __call__). I don't understand how the an OpenAI object has access to the context var and where the callback is passed.

EDIT:
I guess it all happens inside _configure which appends openai_context_var to the callbacks to be callbed by the LLM here.

[dosubot[bot]]

Answer generated by a 🤖

Answer

You're correct in your understanding that get_openai_callback() yields an OpenAICallbackHandler object. This object's on_llm_end method is indeed called in the generate method of the OpenAI class. Your confusion seems to stem from how the OpenAI object has access to the context variable and where the callback is passed.

In the LangChain framework, the OpenAI model is instantiated with a list of callback handlers passed to the callbacks parameter. This list can include various types of callback handlers, such as StdOutCallbackHandler and ClearMLCallbackHandler. The ClearMLCallbackHandler is specifically used to track and log the model's activities in ClearML, a machine learning experiment management tool. This demonstrates how multiple callback handlers can be used simultaneously with the OpenAI model to customize its behavior and track its activities.

This is one potential explanation based on the code snippets and links you've provided. If you find this explanation accurate and feel it's not covered in the existing documentation, I encourage you to make a pull request to update the documentation. This would greatly benefit other users who might encounter the same confusion.

Sources

• tests/unit_tests/callbacks/test_openai_info.py
• tests/integration_tests/callbacks/test_openai_callback.py
• docs/snippets/modules/callbacks/get_started.mdx
• docs/extras/modules/callbacks/integrations/context.ipynb
• docs/extras/modules/callbacks/how_to/custom_callbacks.ipynb
• docs/snippets/modules/model_io/models/llms/how_to/streaming_llm.mdx
• docs/extras/ecosystem/integrations/aim_tracking.ipynb
• docs/extras/modules/callbacks/how_to/async_callbacks.ipynb
• docs/extras/ecosystem/integrations/clearml_tracking.ipynb

---

This response is meant to be useful, save you time, and share context. It is not meant to be a precise solution, but rather a starting point for your own research.

Help me be more useful! Please leave a 👍 if this is helpful and 👎 if it is irrelevant.

[mspronesti]

The correct answer can be found in the "EDIT" part of my question.
More in detail:

• get_openai_callback sets the value of a context var called openai_context_var to an object of type OpenAICallBackHandler which tracks the token usage and the cost
• the llms derived from BaseOpenAI create a list of callbacks from the configure method of the CallbackManager object
• the on_llm_end method of the OpenAICallBackHandler is applied on the generated tokens once the generation is finished to update its the attributes