Skip to content

Assistants: Count the usage tokens by run and not by assistant instance. #1010

@a-chacon

Description

@a-chacon

I am doing an assistant for a web app using lanchainrb (Thanks for built it). My problem is that I want to register the token usage for every call to run and not by the instance of Assistant. I know this is ok when you set the auto_tool_execution to true because me (as user) dont handle the tool call response, so you plus the tool_call message tokens usage to the last assistant message and return this as the total usage tokens.

I mean, an example of my seudo code:

  def run(user_message)
    assistant = current_assistant_class.new(self)

    # set assistant history

    loop do
      assistant.run

      if assistant.messages.last.tool_calls.any?
        # First register of total tokens, It is OK!
        create_message(tool_calls: assistant.messages.last.tool_calls, role: 'assistant', input_tokens: assistant.total_prompt_tokens, output_tokens: assistant.total_completion_tokens)

        # Handle tool calls....
      end
    end

    # When tool calls finish, I want to be able to register the usage tokens of the "last run"
    # But instead I receive x tool_calls + last_run usage tokens. 
    create_message(content: assistant.messages.last.content, role: 'assistant', input_tokens: assistant.total_prompt_tokens, output_tokens: assistant.total_completion_tokens)
    assistant.messages.last.content
  end

This is an example, My usage case is a bit more complex. But let me know what do you think.

I was taking a look to the class Assistant. With a setter method to the variables :total_prompt_tokens, :total_completion_tokens, :total_tokens I would be able to reset it to 0 before any run. Or add an argument to run for reset these variables. I am able to open a PR with any of this solutions (Or whatever you propose too)

Also would be very nice to get all the metadata of usage of the call ;)

Metadata

Metadata

Assignees

No one assigned

    Labels

    No labels
    No labels

    Type

    No type

    Projects

    No projects

    Milestone

    No milestone

    Relationships

    None yet

    Development

    No branches or pull requests

    Issue actions