Exposing token usage and other stats

[hellovai]

Update: latest docs link: https://boundary-preview-3c6f166c-5bc3-4b4f-b5b1-7ec76926f143.docs.buildwithfern.com/ref/baml_client/collector (with some api updates)

------ Original post -----
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 (collector)

(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

collector = b.create_collector()

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

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

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

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

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

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

# Clear the collector of any existing logs and accumulation
collector.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(collector.last.selected_call().request)

# Prints out the HTTP response the LLM replied with
print(collector.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_collector()
bo = { "collector": collector } 
resume1, resume2 = await asyncio.gather(
  b.ExtractResume("text1", baml_options=bo),
  b.ExtractResume("text2", baml_options=bo)
)

# has 2 entries for the ExtractResume function
collector.logs

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

collector = b.create_collector()
bo = { "collector": collector } 
(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
collector.id(id1).usage

# log data about request 2
collector.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(collector=b.create_collector())

# 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.collector.last.usage)

Full reference docs

Collector 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.

[aaronvg]

Current ETA: ~Feb 14. Will keep this updated! We are actively working on releasing this (#1 priority)

[ddehghan]

Happy Valentines day ❤️ .... I am really waiting for this feature. 😝

[aaronvg]

still workin on it! Might have something for this monday's release, if not hopefully the next monday.

[phiwi]

Any update on this issue? I am especially looking forward to access to the raw model output (DeepSeek thinking process).

[aaronvg]

Looking like next monday's release.

Making great progress!

[hellovai]

As an update for everyone, it looks like its still 2 releases away (March first week is the target)! We've got python finally working, and so just need to touch up the other languages and do final testing! but we're so excited to finally make this available!! @aaronvg is working hard ⚙️

[aaronvg]

PR to track: #1512

[aaronvg]

We merged the initial python implementation and is in 0.78.0 as a pre-alpha feature (not fully complete). If you use it please be aware the interfaces will change on the next release. https://github.com/BoundaryML/baml/blob/canary/integ-tests/python/tests/test_collector.py

I will keep this updated when it's fully out for all languages (TS and Ruby)

[aaronvg]

updated docs link https://boundary-preview-3c6f166c-5bc3-4b4f-b5b1-7ec76926f143.docs.buildwithfern.com/ref/baml_client/collector