Verfolgung der LLM-Nutzung mit einem Logger

Die in AutoGen enthaltenen Modellclients geben strukturierte Ereignisse aus, die zur Verfolgung der Modellnutzung verwendet werden können. Dieses Notebook zeigt, wie der Logger zur Verfolgung der Modellnutzung verwendet wird.

Diese Ereignisse werden mit dem Namen :py:attr:autogen_core.EVENT_LOGGER_NAME in den Logger geschrieben.

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)

Dieser Logger kann dann wie jeder andere Python-Logger angehängt und die Werte nach Ausführung des Modells ausgelesen werden.

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)