[feat] expose tokens/timings through REST API calls

[Elijas]

Use case

Use case use REST APIs to get token counts, timings etc. associated with the request, but keep REST API stateless

Solution Idea

There could be a middle-ground API across all clients, including API clients

pseudocode of /call_with_metrics implementation:

function b.with_stats.Any(...):
    c = Collector()
    r = b.Any(..., collector=c)
    return {result: r, stats: c.stats}

[linear]

BAML-389 [feat] expose tokens/timings through REST API calls

[DomWilliamsEE]

I made an initial PR (#1989) that shows it is technically possible to add a Collector for a /call_with_metrics endpoint.
@hellovai pointed out we should document and decide on what the usage should look like.

I think dumping the raw Collector state in the response is messy and has unwanted fields. I would suggest something like this:

• response: the llm response, i.e. the raw response of /call
• metrics: has input_tokens, output_tokens and duration_ms, the bare minimum of metrics that are useful
• logs: the terminal output of the full function call is useful, but its inaccessible to the caller of the /call endpoint. I would suggest returning it as a free-form string here. Alternatively it could be the full nested structure of all the calls and tool usages.

So as an example:
POST /call_with_metrics/:function_name: Call a BAML function with extended metrics about the function call
returned.

curl \
  -X POST \
  "http://localhost:2024/call_with_metrics/MyFunctionName" \
  -H "Content-Type: application/json" \
  -d '{"arg1": "value1", "arg2": "value2"}'

Example response:

{
  "response": "...",
  "metrics": {
    "input_tokens": 101,
    "output_tokens": 291,
    "duration_ms": 2810
  },
  "logs": "Function MyFunctionName\nClient: GPT4o (gpt-4o-2024-08-06) - 2810ms. StopReason: stop. Tokens(in/out): 101/291\n---PROMPT---..."
}