What problem or use case are you trying to solve?

TASK

We want to enhance our get_token_count() implementation in llm.py, to take advantage of the token counts we are provided from our dependencies, if available, and only fallback to count them if they are not already available.

Read the llm.py file. It uses the litellm library, and you can find some code that uses the Usage object from litellm, too.

This usage data, when available, provides live token counts. It will be available for some LLMs, and it will not be available for others.

If these values could be stored and/or attached to the corresponding messages, these values could be calculated and returned in get_token_count() in llm.py and used to live-compute the token counts.

Read also codeact_agent.py file, to see how the LLM is used, and message.py to understand the messages. Read event.py, there will be multiple files with this name, that's okay, read them all, to see Events.

IMPORTANT:

• the Message objects are computed on the fly at each step
• Messages get their content from Event objects
• it is acceptable to add an event_id attribute to a Message, to link it to the underlying Event.
• it's NOT acceptable to store the event_id_to_tokens or message_to_tokens in the agent. Use message.py class. Or use llm.py, if it's not too crowded. Your choice.

First, read and review the files you need to understand how we can do this.

Then, decide on a solution, and do it.

Unit tests are required.

/TASK

ORIGINAL MESSAGE

The usage data, if filled, provides live token counts. It will work for some LLMs, and it will not be available for others.

If these values could be stored and/or attached to the messages in the codeact_agent.py accordingly, these values could be calculated and returned in get_token_count() in llm.py.

Then this could help a running memory condenser in calculations?

ModelResponse(
    id='chatcmpl-47df6b79-cb8d-4f07-920a-8c5551ea7ee4',
    choices=[
        Choices(
            finish_reason='stop',
            index=0,
            message=Message(
                content='{\n  "action": "run",\n  "args": {\n    "command": "ls"\n  }\n}\n',
                role='assistant',
                tool_calls=None,
                function_call=None
            )
        )
    ],
    created=1721067905,
    model='anthropic/claude-3.5-sonnet:beta',
    object='chat.completion',
    system_fingerprint=None,
    usage=Usage(
        prompt_tokens=10,
        completion_tokens=20,
        total_tokens=30
    )
)