langchain: get_openai_callback not working when using Agent Executor after updating to latest version of Langchain

Having the same issue, thanks for reporting!

It’s a bug that I accidentally introduced while changing how agent executor works to make it easy to stream token by token output from the agent (e.g., https://python.langchain.com/docs/modules/agents/how_to/streaming#custom-streaming-with-events) – didn’t occur to me that callbacks would be affected 😦

I made the agent executor always invoke the open ai model in a streaming mode, so folks can get token by token streaming output from the agent, but because the OpenAI API never returns token counts when the model streams it breaks token counting callbacks.

Likely the best solution is to extend the existing openai callback to calculate tokens with tiktoken when the response from the open ai is missing token counts (as is the case when streaming). This will resolve both this issue as well as issues that folks face when streaming content from the openai api directly.

We just need to verify that the token counts are accurate and figure out mapping between tokenizer and specific models in OpenAI.

I’ll try to pick up work on this soon. For folks looking for a solution, the custom callback that you shared is a good workaround, since that’s what the extension to the existing openai callbacks will do.

@hwchase17 any update?

🤖

I’m sorry to hear that you’re having trouble with the get_openai_callback function after updating to the latest version of LangChain. Based on the code you’ve provided, it seems like you’re not awaiting the agent_executor.ainvoke() function inside the get_openai_callback() context manager. The agent_executor.ainvoke() function is an asynchronous function and needs to be awaited to ensure that it completes execution before the context manager is exited.

Here’s how you can modify your code:

with get_openai_callback() as cb:
    clara_ai_resp = await agent_executor.ainvoke({"input": input})
    clara_ai_output = clara_ai_resp["output"]

    print("CB: ", cb)

    return clara_ai_output, input, cb

Please note that the function ask_assistant() is now an asynchronous function (prefixed with async). This is necessary because we’re using the await keyword inside the function.

Regarding the issue with the ChatOpenAI not calling ainvoke as ChatOpenAI from langchain.chat_models did, it’s possible that there were changes in the latest version of LangChain that affected this behavior. Without more information about the changes made in the latest version, it’s hard to provide a definitive answer. I recommend checking the release notes or the commit history for the latest version of LangChain to see if there were any changes that might have affected this.

I hope this helps! If you have any other questions or need further clarification, please let me know.

Sources

• libs/langchain/tests/integration_tests/callbacks/test_openai_callback.py

---

This response is meant to be useful and save you time. 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.

If you want to continue the conversation, start your reply with @dosu-bot.

🤖

Hello again, @EduardoMM17!

It seems we’re on a coding adventure together, doesn’t it? 😄

The issue you’re experiencing with the get_openai_callback function seems to be related to its usage in the context of the Agent Executor. The get_openai_callback function is a context manager designed to track the usage of tokens in the OpenAI language model. It needs to be active (i.e., within its context) when the language model is called, otherwise it won’t be able to track the tokens.

In your provided code, the get_openai_callback context is not active when the agent_executor.ainvoke method is called. This is why it’s not tracking any tokens and returning zero for everything. To fix this, you need to ensure that the agent_executor.ainvoke method is called within the get_openai_callback context.

Here’s an example of how you can do it:

with get_openai_callback() as cb:
    clara_ai_resp = await agent_executor.ainvoke({"input": input})
    clara_ai_output = clara_ai_resp["output"]

    print("CB: ", cb)

    return clara_ai_output, input, cb

In this corrected code, the agent_executor.ainvoke method is called within the get_openai_callback context, so the callback function will be able to track the tokens used by the language model.

This usage is demonstrated in the LangChain test suite, specifically in the file libs/langchain/tests/integration_tests/callbacks/test_openai_callback.py.

I hope this helps! If you have any other questions or need further clarification, please let me know.

Sources

• libs/langchain/tests/integration_tests/callbacks/test_openai_callback.py

---

This response is meant to be useful and save you time. 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.

If you want to continue the conversation, start your reply with @dosu-bot.

🤖

I’m sorry for the confusion. It seems there might be a misunderstanding about the get_openai_callback function. This function is not present in the AgentExecutor class in the latest version of LangChain. The AgentExecutor class does not have a get_openai_callback method.

The OpenAIFunctionsAgent class, which was deprecated as of version 0.1.0, was replaced by the create_openai_functions_agent function. This function creates an agent that uses OpenAI function calling. It takes as input a language model, a sequence of tools, and a prompt, and returns a Runnable sequence representing an agent.

If you’re trying to use a callback with the AgentExecutor, you might want to use the callback_manager parameter in the from_llm_and_tools class method of the OpenAIFunctionsAgent class. This parameter accepts an instance of a BaseCallbackManager, which can be used to manage callbacks.

Regarding the issue with the ChatOpenAI not calling ainvoke, it’s hard to provide a precise answer without more specific information about the error you’re encountering or the functionality you’re trying to achieve. Could you please provide more details about the issue? For example, any error messages you’re seeing, or a more detailed description of the expected vs. actual behavior would be helpful.

As for the 404 errors you’re encountering, it seems like the links might be broken or the resources have been moved. I recommend checking the LangChain repository directly for the most up-to-date information.

I hope this helps! If you have any more questions or need further clarification, please don’t hesitate to ask.

Sources

• libs/langchain/langchain/agents/openai_functions_agent/base.py

---

This response is meant to be useful and save you time. 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.

If you want to continue the conversation, start your reply with @dosu-bot.

Any help on this issue will be appreciated. Getting the exact token counts isn’t super trivial and will help power the streaming mode as well.

One work-around is to force the model to be invoked in a non streaming mode which will give you token counts from the OpenAI api:

def invoke_model(prompt, config):
    return llm_with_stop.invoke(prompt, config)


    agent = (
        {
            "input": lambda x: x["input"],
            "agent_scratchpad": lambda x: format_log_to_str(x["intermediate_steps"]),
            "chat_history": lambda x: x["chat_history"],
        }
        | prompt
        | invoke_model
        | ReActSingleInputOutputParser()
    )

^ Corrected a typo in the snippet above since it needs to use config rather than callbacks

Thank you both @eyurtsev and @EduardoMM17 for the clarification and for the workaround proposed 😃

🤖

I see that you’re having trouble with the get_openai_callback function when using the Agent Executor in the latest version of LangChain. The issue seems to be that the get_openai_callback function is not being awaited, which is necessary for asynchronous functions in Python.

Here’s a corrected version of your code:

async def ask_assistant(input: str) -> str:
    # ... your code before this ...

    with get_openai_callback() as cb:
        clara_ai_resp = await agent_executor.ainvoke({"input": input})
        clara_ai_output = clara_ai_resp["output"]

        print("CB: ", cb)

        return clara_ai_output, input, cb

In the original code, the agent_executor.ainvoke function was not awaited, so the execution of the code did not wait for the agent_executor.ainvoke function to complete and moved on to the next line. As a result, the get_openai_callback function did not have any data to process and returned zero on everything.

Regarding your observation about ChatOpenAI from langchain_openai not calling ainvoke as ChatOpenAI from langchain.chat_models did, it’s possible that there were changes in the latest version of LangChain that affected this behavior. Unfortunately, I couldn’t find specific changes to the get_openai_callback function in the latest version of LangChain in the repository. It might be helpful to check the release notes or changelog for the latest version to see if there were any changes that could affect this.

I hope this helps! If you have any other questions or need further clarification, feel free to ask.

Sources

• libs/langchain/tests/integration_tests/callbacks/test_openai_callback.py

---

This response is meant to be useful and save you time. 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.

If you want to continue the conversation, start your reply with @dosu-bot.

@danieleratti

I managed to make a workaround using tiktoken inside the CustomAsyncHandler class. Then you can just use a function like get_openai_cost_for_model in line 2 to calculate the costs.

I have also verified that the token count is the same or similar to what I get from Langsmith and to my surprise it is. In the input tokens I see a difference of ~10, which is not that much and in the completion I get a difference of ~1 or no difference. So I will be using this in the meantime until I come up with a more elegant solution or someone fixes the bug. Check the code below:

from langchain.callbacks.base import AsyncCallbackHandler
from langchain_community.callbacks.openai_info import get_openai_token_cost_for_model

from app.services.llm.prompt import prompt_raw
from app.services.llm.tools.build_tools import tools
from app.classes.CustomBuffer import CustomConversationBufferMemory

memory = CustomConversationBufferMemory(memory_key="chat_history", return_messages=True)

# use get_openai_token_cost_for_model to calculate cost
class MyCustomAsyncHandler(AsyncCallbackHandler):
    async def on_llm_start(
        self,
        serialized: dict[str, Any],
        prompts: list[str],
        **kwargs: Any,
    ) -> None:
        encoding = tiktoken.get_encoding("cl100k_base")
        prompts_string = ''.join(prompts)
        num_tokens = len(encoding.encode(prompts_string))
        print("NUM TOKENS: ", num_tokens)


    async def on_llm_end(self, response: LLMResult, **kwargs: Any) -> None:
        """Run when chain ends running."""
        text_response = response.generations[0][0].text
        encoding = tiktoken.get_encoding("cl100k_base")
        response_string = len(encoding.encode(text_response))
        print("NUM TOKENS RESPONSE: ", response_string)

async def ask_assistant(input: str) -> str:
    prompt = PromptTemplate.from_template(prompt_raw)

    prompt = prompt.partial(
        language="Spanish",
        tools=render_text_description(tools),
        tool_names=", ".join([t.name for t in tools]),
    )

    llm = ChatOpenAI(
        temperature=0,
        model_name="gpt-4",
        openai_api_key=os.environ["OPENAI_API_KEY"],
        callbacks=[MyCustomAsyncHandler()],
    )
    llm_with_stop = llm.bind(stop=["\nObservation"])

    agent = (
        {
            "input": lambda x: x["input"],
            "agent_scratchpad": lambda x: format_log_to_str(x["intermediate_steps"]),
            "chat_history": lambda x: x["chat_history"],
        }
        | prompt
        | llm_with_stop
        | ReActSingleInputOutputParser()
    )

    agent_executor = AgentExecutor(
        agent=agent,
        tools=tools,
        verbose=True,
        memory=memory,
        max_execution_time=60,
        handle_parsing_errors=True,
    )

    with get_openai_callback() as cb:
        clara_ai_resp = await agent_executor.ainvoke({"input": input})
        clara_ai_output = clara_ai_resp["output"]

        print("CB: ", cb)

        return clara_ai_output, input, cb

I ran into the same problem and built a context manager that works much in the same way as the original, it just needs the model name as input:

import threading
from contextlib import contextmanager
from typing import Any, Generator

import tiktoken
from langchain_community.callbacks.manager import openai_callback_var
from langchain_community.callbacks.openai_info import standardize_model_name, MODEL_COST_PER_1K_TOKENS, \
    get_openai_token_cost_for_model, OpenAICallbackHandler
from langchain_core.outputs import LLMResult


class CostTrackerCallback(OpenAICallbackHandler):

    def __init__(self, model_name: str) -> None:
        super().__init__()
        self.model_name = model_name
        self._lock = threading.Lock()

    def on_llm_start(
        self,
        serialized: dict[str, Any],
        prompts: list[str],
        **kwargs: Any,
    ) -> None:
        encoding = tiktoken.get_encoding("cl100k_base")
        prompts_string = ''.join(prompts)
        self.prompt_tokens = len(encoding.encode(prompts_string))

    def on_llm_end(self, response: LLMResult, **kwargs: Any) -> None:
        """Run when chain ends running."""
        text_response = response.generations[0][0].text
        encoding = tiktoken.get_encoding("cl100k_base")
        self.completion_tokens = len(encoding.encode(text_response))
        model_name = standardize_model_name(self.model_name)
        if model_name in MODEL_COST_PER_1K_TOKENS:
            completion_cost = get_openai_token_cost_for_model(
                model_name, self.completion_tokens, is_completion=True
            )
            prompt_cost = get_openai_token_cost_for_model(model_name, self.prompt_tokens)
        else:
            completion_cost = 0
            prompt_cost = 0

        # update shared state behind lock
        with self._lock:
            self.total_cost += prompt_cost + completion_cost
            self.total_tokens = self.prompt_tokens + self.completion_tokens
            self.successful_requests += 1


@contextmanager
def get_cost_tracker_callback(model_name) -> Generator[CostTrackerCallback, None, None]:
    cb = CostTrackerCallback(model_name)
    openai_callback_var.set(cb)
    yield cb
    openai_callback_var.set(None)

Same issue here and no solution found

Even if I try to attach a callback to the llm, I get None as response.llm_output (so I can’t access the number of tokens) and if I attach a callback to the AgentExecutor, I only get the on_agent_finish callback (which is useless for getting tokens and cost)…

Hey @EduardoMM17 we are investigating why Dosu linked to a file that doesn’t exist. Sorry for the confusion!