Token usage and cost tracking in Agentic AI - ML Experiment: Train & Evaluate

Experiment - Token usage and cost tracking

Problem:You are using an AI language model that charges based on the number of tokens processed. Currently, you do not track how many tokens your requests use or the cost incurred. This makes it hard to manage your budget and optimize usage.

Current Metrics:No token usage or cost tracking implemented; total cost unknown.

Issue:Without tracking tokens and cost, you risk overspending and cannot optimize your queries for cost efficiency.

Your Task

Implement a system to count tokens used per request and calculate the total cost based on a fixed price per 1000 tokens. The system should output token counts and cumulative cost after each request.

You must use the provided token counting method or a simple approximation.

Assume a fixed cost rate of $0.002 per 1000 tokens.

Do not change the AI model or request content.

Hint 1

Hint 2

Hint 3

Solution

Agentic AI

import re

class TokenCostTracker:
    def __init__(self, cost_per_1000_tokens=0.002):
        self.total_tokens = 0
        self.total_cost = 0.0
        self.cost_per_1000_tokens = cost_per_1000_tokens

    def count_tokens(self, text):
        # Simple approximation: count words as tokens
        tokens = len(re.findall(r"\w+", text))
        return tokens

    def update(self, prompt, response):
        prompt_tokens = self.count_tokens(prompt)
        response_tokens = self.count_tokens(response)
        used_tokens = prompt_tokens + response_tokens
        self.total_tokens += used_tokens
        cost = (used_tokens / 1000) * self.cost_per_1000_tokens
        self.total_cost += cost
        return {
            "prompt_tokens": prompt_tokens,
            "response_tokens": response_tokens,
            "used_tokens": used_tokens,
            "total_tokens": self.total_tokens,
            "cost": round(cost, 6),
            "total_cost": round(self.total_cost, 6)
        }

# Example usage
tracker = TokenCostTracker()
prompt = "Hello, how are you today?"
response = "I am fine, thank you! How can I help you?"
usage1 = tracker.update(prompt, response)

prompt2 = "Tell me a joke."
response2 = "Why did the chicken cross the road? To get to the other side!"
usage2 = tracker.update(prompt2, response2)

print(usage1)
print(usage2)

Created a TokenCostTracker class to count tokens and track cost.

Implemented a simple token counting method based on word count.

Added methods to update total tokens and cost after each prompt-response pair.

Displayed token usage and cost details after each interaction.

Results Interpretation

Before: No token or cost tracking, total cost unknown.

After: Token usage and cost tracked per request. Example: First request used 15 tokens costing $0.000030, second request used 17 tokens costing $0.000034, total cost $0.000064.

Tracking token usage and cost helps manage AI usage budgets and optimize queries to reduce expenses.

Bonus Experiment

Extend the tracker to visualize token usage and cost over multiple requests using a simple line chart.

💡 Hint

Use matplotlib or any plotting library to plot total tokens and total cost after each request.

Practice

1. What is a token in the context of AI language models?

easy

A. A programming language used for AI

B. A type of AI model architecture

C. A small piece of text like a word or part of a word

D. A hardware component for AI processing

Solution

1. Step 1: Understand the definition of a token

   A token is a small piece of text that AI models use to understand language. It can be a word or part of a word.

2. Step 2: Differentiate tokens from other AI terms

   Tokens are not models, programming languages, or hardware. They are units of text input.

3. Final Answer:

   A small piece of text like a word or part of a word -> Option C

4. Quick Check:

   Token = small text piece [OK]

Hint: Tokens are text pieces, not models or hardware [OK]

Common Mistakes:

• Confusing tokens with AI models
• Thinking tokens are programming languages
• Assuming tokens are hardware parts

2. Which of the following is the correct way to track token usage in a Python script using an AI API?

easy

A. tokens_used = response['total_tokens']

B. tokens_used = response['usage']['total_tokens']

C. tokens_used = response['usage']['tokens']

D. tokens_used = response['token_count']

Solution

1. Step 1: Identify the correct key for token usage in API response

   Most AI APIs return token usage under response['usage']['total_tokens'].

2. Step 2: Compare options with common API response structure

   Only tokens_used = response['usage']['total_tokens'] matches the standard nested key for total tokens used.

3. Final Answer:

   tokens_used = response['usage']['total_tokens'] -> Option B

4. Quick Check:

   Correct key path = response['usage']['total_tokens'] [OK]

Hint: Look for nested 'usage' then 'total_tokens' key [OK]

Common Mistakes:

• Using wrong key names like 'token_count'
• Missing nested 'usage' dictionary
• Assuming flat keys for token counts

3. Given the following code snippet, what will be printed?

response = {'usage': {'prompt_tokens': 50, 'completion_tokens': 30, 'total_tokens': 80}}
print(response['usage']['total_tokens'])

medium

A. 80

B. 30

C. 50

D. Error

Solution

1. Step 1: Access the 'total_tokens' key in the nested dictionary

   The code accesses response['usage']['total_tokens'], which is 80.

2. Step 2: Confirm the print output

   Printing this value outputs 80 without error.

3. Final Answer:

   80 -> Option A

4. Quick Check:

   response['usage']['total_tokens'] = 80 [OK]

Hint: Check nested dictionary keys carefully [OK]

Common Mistakes:

• Confusing prompt_tokens with total_tokens
• Mixing completion_tokens with total_tokens
• Assuming code causes error

4. You wrote this code to track token usage but get a KeyError:

tokens = response['usage']['token_total']
print(tokens)

What is the likely cause?

medium

A. The 'usage' key is missing in the response

B. The response variable is not defined

C. The print statement syntax is incorrect

D. The key 'token_total' does not exist in the response dictionary

Solution

1. Step 1: Check the key names in the response dictionary

   The correct key is 'total_tokens', not 'token_total'.

2. Step 2: Understand KeyError cause

   Using a wrong key name causes KeyError because that key does not exist.

3. Final Answer:

   The key 'token_total' does not exist in the response dictionary -> Option D

4. Quick Check:

   Wrong key name causes KeyError [OK]

Hint: Verify exact key names in API response [OK]

Common Mistakes:

• Assuming similar key names exist
• Ignoring case sensitivity in keys
• Thinking print syntax causes error

5. You want to estimate the cost of an AI request. The model charges $0.002 per 1000 tokens. If your request uses 2500 tokens, what is the total cost?

hard

A. $0.005

B. $0.0025

C. $0.05

D. $0.0005

Solution

1. Step 1: Calculate cost per token

   Cost per token = $0.002 / 1000 = $0.000002 per token.

2. Step 2: Multiply by number of tokens used

   Total cost = 2500 tokens * $0.000002 = $0.005.

3. Final Answer:

   $0.005 -> Option A

4. Quick Check:

   2500 tokens x $0.002/1000 = $0.005 [OK]

Hint: Multiply tokens by cost per token (divide by 1000 first) [OK]

Common Mistakes:

• Multiplying by 0.002 directly without dividing by 1000
• Using wrong token count
• Confusing decimals in cost calculation