Exposing token usage and other stats

[hellovai]

One of the hardest parts about using BAML is exposing internal workings of BAML to developers using it. Items such as: Token Usage, Raw HTTP Response from the model are all abstracted away for the sake of simplicity. While we do value simplicity, we don’t want it introspection to be a blocker.

A core tenant of BAML is to give the developer complete control.

This document outlines a proposal that will allow BAML developers to:

1. Have access to usage statistics
2. Have access to raw http request / responses made by BAML functions
3. Finely tune where and when BAML logs events to console
4. Aggregate usage statistics

💡 None of the content here changes anything about BAML code. This will only change baml_client code (i.e. how you call BAML functions).

Table of Contents

1. Get token usage statistics
   • Accessing additional metadata
2. Running calls in parallel (id)
3. Reducing repetition of baml_options with with_options
4. Full reference docs

Get usage statistics (logger)

(This document is mostly python, but please see the full reference docs for other languages, they all expose the same capabilities)

Python

from baml_client import b

logger = b.create_logger()

# Call the function, but additionally pass in the logger
invoice = b.ExtractInvoice("..", baml_options={"logger": logger})

# Prints out usage
print(logger.usage)
# Usage(input_tokens=int, output_tokens=int)

# Call another function using the same logger
resume = b.ExtractResume("..", baml_options={"logger": logger})

# Prints out cumulative usage of both functions
print(logger.usage)

# To get usage of just the second function
print(logger.last.usage)

# Print out usage of each call specifically:
for event in logger.logs:
   print(event.function_name, event.id, event.usage)

# Clear the logger of any existing logs and accumulation
logger.clear()

Accessing additional metadata

We also want to expose all metadata such as raw HTTP respones

# Prints out cumulative usage of both functions

# Prints out the HTTP request we made
print(logger.last.selected_call().request)

# Prints out the HTTP response the LLM replied with
print(logger.last.selected_call().response)

Running calls in parallel (id)

When calling in parallel, its logger.last or filtering my function name may not be reliable.

logger = b.create_logger()
bo = { "logger": logger } 
resume1, resume2 = await asyncio.gather(
  b.ExtractResume("text1", baml_options=bo),
  b.ExtractResume("text2", baml_options=bo)
)

# has 2 entries for the ExtractResume function
logger.logs

You can instead use the new b.id scope to generate an (id, T) value from a function.

logger = b.create_logger()
bo = { "logger": logger } 
(id1, resume), (id2, resume2) = await asyncio.gather(
  b.id.ExtractResume("text1", baml_options=bo),
  b.id.ExtractResume("text2", baml_options=bo)
)

# log data about request 1
logger.id(id1).usage

# log data about request 2
logger.id(id1).usage

Reducing repetition of baml_options with with_options

You may find that always passing in baml_options can be cumbersome. We now allow you to construct a BamlClient with default options passed in.

from baml_client import b

my_b = b.with_options(logger=b.create_logger())

# Those options are always forwarded to every call (streamed and otherwise included).
resume = my_b.ExtractResume("...")

# Access the logger for that scoped client.
print(my_b.logger.last.usage)

Full reference docs

Logger Class

https://boundary-preview-12ad00fd-c3dc-43be-ab63-ae17b33338ce.docs.buildwithfern.com/ref/baml-client/logger

id scope

https://boundary-preview-12ad00fd-c3dc-43be-ab63-ae17b33338ce.docs.buildwithfern.com/ref/baml-client/id

with_options

https://boundary-preview-12ad00fd-c3dc-43be-ab63-ae17b33338ce.docs.buildwithfern.com/ref/baml-client/with-options

[hellovai]

Associated discord thread: https://discord.com/channels/1119368998161752075/1119375594984050779/1324870777136353351

[ddehghan]

Maybe further down the road you want to add ability to hand off the transport mechanism to the user. They can make their own http calls and capture what they want and hand the data back to you. People may have special proxy servers they need to set or have to log the calls into ops monitoring pipeline etc. Also this way you also future proof yourself for any additional logs and error handling.

[matthiaswiench]

@hellovai thank you for the update! Is there any workaround to count the tokens already?

[jkonowitch]

This APi makes sense and would address my use cases. I appreciate that you also gave consideration to the streaming API as well.