mirror of
https://github.com/microsoft/autogen.git
synced 2025-07-14 04:20:52 +00:00
32 lines
940 B
Python
32 lines
940 B
Python
|
from agnext.application_components.logging import LLMUsageTracker, EVENT_LOGGER_NAME, LLMCallEvent
|
||
|
|
|
||
|
|
import logging
|
||
|
|
|
||
|
|
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
|