{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "# 使用日志记录器跟踪大语言模型使用情况\n\nAutoGen中包含的模型客户端会发出结构化事件，可用于跟踪模型使用情况。本笔记本演示如何使用日志记录器来跟踪模型使用。\n\n这些事件会被记录到名为 :py:attr:`autogen_core.EVENT_LOGGER_NAME` 的日志记录器中。\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "import logging\n",
        "\n",
        "from autogen_core.logging import LLMCallEvent\n",
        "\n",
        "\n",
        "class LLMUsageTracker(logging.Handler):\n",
        "    def __init__(self) -> None:\n",
        "        \"\"\"Logging handler that tracks the number of tokens used in the prompt and completion.\"\"\"\n",
        "        super().__init__()\n",
        "        self._prompt_tokens = 0\n",
        "        self._completion_tokens = 0\n",
        "\n",
        "    @property\n",
        "    def tokens(self) -> int:\n",
        "        return self._prompt_tokens + self._completion_tokens\n",
        "\n",
        "    @property\n",
        "    def prompt_tokens(self) -> int:\n",
        "        return self._prompt_tokens\n",
        "\n",
        "    @property\n",
        "    def completion_tokens(self) -> int:\n",
        "        return self._completion_tokens\n",
        "\n",
        "    def reset(self) -> None:\n",
        "        self._prompt_tokens = 0\n",
        "        self._completion_tokens = 0\n",
        "\n",
        "    def emit(self, record: logging.LogRecord) -> None:\n",
        "        \"\"\"Emit the log record. To be used by the logging module.\"\"\"\n",
        "        try:\n",
        "            # 如果消息是StructuredMessage的实例，则使用该类型\n",
        "            if isinstance(record.msg, LLMCallEvent):\n",
        "                event = record.msg\n",
        "                self._prompt_tokens += event.prompt_tokens\n",
        "                self._completion_tokens += event.completion_tokens\n",
        "        except Exception:\n",
        "            self.handleError(record)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "然后，可以像其他Python日志记录器一样附加此记录器，并在模型运行后读取记录值\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "from autogen_core import EVENT_LOGGER_NAME\n",
        "\n",
        "# 设置日志配置以使用自定义处理器\n",
        "logger = logging.getLogger(EVENT_LOGGER_NAME)\n",
        "logger.setLevel(logging.INFO)\n",
        "llm_usage = LLMUsageTracker()\n",
        "logger.handlers = [llm_usage]\n",
        "\n",
        "# client.create(...)\n",
        "\n",
        "print(llm_usage.prompt_tokens)\n",
        "print(llm_usage.completion_tokens)"
      ]
    }
  ],
  "metadata": {
    "kernelspec": {
      "display_name": ".venv",
      "language": "python",
      "name": "python3"
    },
    "language_info": {
      "codemirror_mode": {
        "name": "ipython",
        "version": 3
      },
      "file_extension": ".py",
      "mimetype": "text/x-python",
      "name": "python",
      "nbconvert_exporter": "python",
      "pygments_lexer": "ipython3",
      "version": "3.11.5"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 2
}