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
    No fields configured for issues without a type.

    Projects

    No projects

    Milestone

    No milestone

    Relationships

    None yet

    Development

    No branches or pull requests

    Issue actions