mirror of
https://github.com/microsoft/autogen.git
synced 2025-07-13 12:01:04 +00:00
33 lines
930 B
Python
33 lines
930 B
Python
import logging
|
|
|
|
from agnext.application.logging import EVENT_LOGGER_NAME, LLMCallEvent, LLMUsageTracker
|
|
|
|
|
|
def test_llm_usage() -> None:
|
|
|
|
# Set up the logging configuration to use the custom handler
|
|
logger = logging.getLogger(EVENT_LOGGER_NAME)
|
|
logger.setLevel(logging.INFO)
|
|
llm_usage = LLMUsageTracker()
|
|
logger.handlers = [llm_usage]
|
|
|
|
logger.info(LLMCallEvent(prompt_tokens=10, completion_tokens=20))
|
|
|
|
assert llm_usage.prompt_tokens == 10
|
|
assert llm_usage.completion_tokens == 20
|
|
|
|
logger.info(LLMCallEvent(prompt_tokens=1, completion_tokens=1))
|
|
|
|
assert llm_usage.prompt_tokens == 11
|
|
assert llm_usage.completion_tokens == 21
|
|
|
|
llm_usage.reset()
|
|
|
|
assert llm_usage.prompt_tokens == 0
|
|
assert llm_usage.completion_tokens == 0
|
|
|
|
logger.info(LLMCallEvent(prompt_tokens=1, completion_tokens=1))
|
|
|
|
assert llm_usage.prompt_tokens == 1
|
|
assert llm_usage.completion_tokens == 1
|