使用日志记录器跟踪 LLM 使用情况#

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:
            # Use the StructuredMessage if the message is an instance of it
            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

# 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]

# client.create(...)

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