使用日志记录器跟踪大语言模型使用情况

AutoGen中包含的模型客户端会发出结构化事件，可用于跟踪模型使用情况。本笔记本演示如何使用日志记录器来跟踪模型使用。

这些事件会被记录到名为 :py:attr:autogen_core.EVENT_LOGGER_NAME 的日志记录器中。

import logging

from autogen_core.logging import LLMCallEvent


class LLMUsageTracker(logging.Handler):
    def __init__(self) -> None:
        """Logging handler that tracks the number of tokens used in the prompt and completion."""
        super().__init__()
        self._prompt_tokens = 0
        self._completion_tokens = 0

    @property
    def tokens(self) -> int:
        return self._prompt_tokens + self._completion_tokens

    @property
    def prompt_tokens(self) -> int:
        return self._prompt_tokens

    @property
    def completion_tokens(self) -> int:
        return self._completion_tokens

    def reset(self) -> None:
        self._prompt_tokens = 0
        self._completion_tokens = 0

    def emit(self, record: logging.LogRecord) -> None:
        """Emit the log record. To be used by the logging module."""
        try:
            # 如果消息是StructuredMessage的实例，则使用该类型
            if isinstance(record.msg, LLMCallEvent):
                event = record.msg
                self._prompt_tokens += event.prompt_tokens
                self._completion_tokens += event.completion_tokens
        except Exception:
            self.handleError(record)

然后，可以像其他Python日志记录器一样附加此记录器，并在模型运行后读取记录值

from autogen_core import EVENT_LOGGER_NAME

# 设置日志配置以使用自定义处理器
logger = logging.getLogger(EVENT_LOGGER_NAME)
logger.setLevel(logging.INFO)
llm_usage = LLMUsageTracker()
logger.handlers = [llm_usage]

# client.create(...)

print(llm_usage.prompt_tokens)
print(llm_usage.completion_tokens)