CODEWITHSUNDEEP
pythonopenai-apilangchainpy-langchain
wellsinsideout
wellsinsideout

Reputation: 53

Getting token details from Langchain/OpenAI callback handler

I'm currently developing an application that utilizes LangChain and OpenAI's API to handle conversational interactions. The application interacts with users through a chat interface, where both user messages and bot responses are processed using OpenAI's models.

My goal is to track the token costs associated with both the user's messages (as prompt tokens) and the bot's responses (as completion tokens). Right now, it's logging all interactions as 0 tokens & cost, which is incorrect. How can I ensure that the token details are captured accurately? Here's a simplified version of the current setup:

class LangChainService:
    def __init__(self, api_key):
        self.llm = OpenAI(api_key=api_key)
        self.chat_model = ChatOpenAI(api_key=api_key)
        self.memory = ConversationBufferMemory()
        self.prompt_template = ChatPromptTemplate(messages=[
            SystemMessagePromptTemplate.from_template(
            "Prompt here"),
            HumanMessagePromptTemplate.from_template("{user_message}")
        ])

    def generate_response(self, user_message):

        callback_handler = OpenAICallbackHandler()

        conversation = ConversationChain(llm=self.chat_model, memory=self.memory, callbacks=[callback_handler])

        response_data = conversation.invoke(user_message)
        response = response_data.get('response') if response_data and 'response' in response_data else "No response generated."

        print("Total cost:", callback_handler.total_cost)
        print("Prompt tokens:", callback_handler.prompt_tokens)
        print("Completion tokens:", callback_handler.completion_tokens)

        return response

def handle_send_message():

    data = request.json
    message = data.get('message')

    print("Generating bot response for message:", message)

    langchain_service = LangChainService(api_key=['OPENAI_API_KEY'])

    response = langchain_service.generate_response(message)

    print("Generated response:", response)

Upvotes: 2

Views: 1202

Answers (2)

Smit Mehta
Smit Mehta

Reputation: 101

OpenAICallbackHandler Class has utility methods on_llm_end() that can be used for updating the token usages. the method applied on the token once the generation is finished the attributes of the instance gets updated.

I would suggest to call the callback_handler.on_llm_end() explicitly once the conversation response is generated.

    callback_handler = OpenAICallbackHandler()

    conversation = ConversationChain(llm=self.chat_model, memory=self.memory, callbacks=[callback_handler])

    response_data = conversation.invoke(user_message)
   callback_handler.on_llm_end()

    # Access the updated values of the callback_handler
    print("Successful requests:", callback_handler.successful_requests)
    print("Total tokens:", callback_handler.total_tokens)
    print("Prompt tokens:", callback_handler.prompt_tokens)
    print("Completion tokens:", callback_handler.completion_tokens)
    print("Total cost:", callback_handler.total_cost)

Upvotes: 0

user17224039
user17224039

Reputation:

Here’s a simple adjustment you can make in the callback handler:

class OpenAICallbackHandler:
    def __init__(self):
        self.prompt_tokens = 0
        self.completion_tokens = 0
        self.total_cost = 0

    def handle_response(self, response):
        # Assuming 'response' is the JSON response from OpenAI which includes token usage info
        self.prompt_tokens = response.get('usage', {}).get('prompt_tokens', 0)
        self.completion_tokens = response.get('usage', {}).get('completion_tokens', 0)
        self.total_cost = response.get('usage', {}).get('total_cost', 0)

Upvotes: 0

Related Questions