Global cost tracking #28
Conversation
There was a problem hiding this comment.
This is a great feature. I’ve added a few comments requesting docstrings to ensure future users can easily understand how to implement it. And I’m approving it.
It works well for now, but it could become misleading if we overlook including the cost_tracker in any new features.
llmclient/cost_tracker.py
Outdated
|
|
||
|
|
||
| TRACK_COSTS = contextvars.ContextVar[bool]("track_costs", default=False) | ||
| REPORT_EVERY_USD = 1.0 |
There was a problem hiding this comment.
Few comments here:
- What do you think of making it a
ClassVarofCostTracker? We still get global state, but it's less awkward with theglobaland setters/getters - Can you make this name more intuitive, and add units to it? It's unclear if
- It's a frequency (Hz)
- It's a dollar threshold (USD)
There was a problem hiding this comment.
I'm confused on the last bullet - _USD is the units - how would you describe it?
There was a problem hiding this comment.
I've renamed set_reporting_frequency -> set_reporting_threshold to remove the ambiguity.
There was a problem hiding this comment.
And made both ClassVars.
There was a problem hiding this comment.
Actually I went further and made them instance variables - no reason not to.
| @@ -73,3 +74,10 @@ def fixture_reset_log_levels(caplog) -> Iterator[None]: | |||
| logger = logging.getLogger(name) | |||
| logger.setLevel(logging.NOTSET) | |||
| logger.propagate = True | |||
|
|
|||
|
|
|||
| class CILLMModelNames(StrEnum): | |||
There was a problem hiding this comment.
We ought to just put this in llmclient as one enum:
class CommonLLMNames(StrEnum):
# Use these for model defaults
OPENAI_GENERAL = "gpt-4o-2024-08-06" # Cheap, fast, and decent
# Use these in unit testing
OPENAI_TEST = "gpt-4o-mini-2024-07-18" # Cheap and not OpenAI's cutting edge
ANTHROPIC_TEST = "claude-3-haiku-20240307" # Cheap and not Anthropic's cutting edgeThen both the app and unit tests will just use CommonLLMNames
There was a problem hiding this comment.
I'll leave that for another PR
|
|
||
| if self.lifetime_cost_usd - self.last_report > self.report_every_usd: | ||
| logger.info( | ||
| f"Cumulative llmclient API call cost: ${self.lifetime_cost_usd:.8f}" |
There was a problem hiding this comment.
| f"Cumulative llmclient API call cost: ${self.lifetime_cost_usd:.8f}" | |
| f"Cumulative client API call cost: ${self.lifetime_cost_usd:.8f}" |
We will eventually maybe rename from llmclient, and maybe do things besides just LLMs (e.g. embeddings), so let's just be generally worded here
There was a problem hiding this comment.
This was intentional wording - the cost tracker only tracks llmclient costs right now.
| image = np.zeros((32, 32, 3), dtype=np.uint8) | ||
| image[:] = [255, 0, 0] |
There was a problem hiding this comment.
| image = np.zeros((32, 32, 3), dtype=np.uint8) | |
| image[:] = [255, 0, 0] | |
| image = np.full((32, 32, 3), [255, 0, 0], dtype=np.uint8) |
Please feel free to ignore this too
There was a problem hiding this comment.
Yeah I didn't add this code - just mirroring what's in the other tests.
| text="What color is this square? Show me your chain of reasoning.", | ||
| images=image, | ||
| ), | ||
| ] # TODO: It's not decoding the image. It's trying to guess the color from the encoded image string. |
There was a problem hiding this comment.
Can you:
- Move this
TODOto not be at the end of alist - Clarify what "it" is, when you say "it's not decoding"
There was a problem hiding this comment.
As above - I copied this from test_llms.py
| { | ||
| "model_list": [ | ||
| { | ||
| "model_name": "gpt-4o-mini", |
There was a problem hiding this comment.
Can we use CILLMModelNames.OPENAI here?
There was a problem hiding this comment.
As above - I copied this from test_llms.py
I'd like to have a way to track costs across all LLM API usage - including agents and environments. This PR implements that.
I made a few semi-arbitrary choices here - there are other ways of implementing this and I'm open to suggestions.
First: I decided to capture costs at the lowest level (i.e. all calls to
acompletion,achat, etc.). We could alternatively only support functions that returnLLMResultand useLLMResult.cost. But we'd run the risk of not catching everythingSecond:
litellmasync streaming calls don't returnAsyncGenerators (fixed our type hints), but insteadCustomStreamWrappers, which implement__aiter__and__anext__. I had to create a shim calledTrackedStreamWrapperto handle this - open to alternative suggestions (see comment in code about whyasync for ... yielddoesn't work)