feat(TokenTracker): Add context manager support to simplify token tracking

This commit is contained in:
choizhang
2025-03-30 00:59:23 +08:00
parent 6eea8bdf5d
commit 164faf94e2
3 changed files with 31 additions and 30 deletions

View File

@@ -44,14 +44,10 @@ async def embedding_func(texts: list[str]) -> np.ndarray:
# function test
async def test_funcs():
# Reset tracker before processing queries
token_tracker.reset()
result = await llm_model_func("How are you?")
print("llm_model_func: ", result)
# Display final token usage after main query
print("Token usage:", token_tracker.get_usage())
# Context Manager Method
with token_tracker:
result = await llm_model_func("How are you?")
print("llm_model_func: ", result)
asyncio.run(test_funcs())