rbm_tools 0.0.0__py3-none-any.whl

rbm_tools/__init__.py

@@ -0,0 +1,5 @@
+"""
+定义 sdk 为 python 的一个包，注意导包顺序，必须保证 logger 最先导包，否则会导致 opentelemetry 注册失效
+"""
+
+from rbm_tools import logger

rbm_tools/client/rbm_tools_cli.py

@@ -0,0 +1,51 @@
+"""
+命令行工具类
+"""
+import sys
+import os
+
+from jsonargparse import auto_cli
+from rbm_tools.client.sort_class_methods import reorder_class_methods
+
+class RBMToolsCLI:
+    """
+    命令行工具类
+    """
+    def __init__(self):
+        pass
+
+    def check_code_style(self):
+        """
+        检查代码风格
+        """
+        success_flag = True
+        directory = os.getcwd()
+        entries = os.listdir(directory)
+        # 只看 py 文件
+        filename_list = [
+            os.path.join(directory, i) for i in entries
+            if os.path.isfile(os.path.join(directory, i)) and i.endswith(".py")
+        ]
+        for filename in filename_list:
+            # 检查对象的成员方法排序
+            old_code = ""
+            with open(filename, "r", encoding="utf-8") as f:
+                old_code = f.read()
+            new_code = reorder_class_methods(old_code)
+            with open(filename, "w", encoding="utf-8") as f:
+                f.write(new_code)
+            if old_code != new_code:
+                success_flag = False
+        if success_flag:
+            sys.exit(0)
+        else:
+            sys.exit(1)
+
+def main():
+    """
+    入口函数
+    """
+    auto_cli(RBMToolsCLI)
+
+if __name__ == "__main__":
+    main()

rbm_tools/client/sort_class_methods.py

@@ -0,0 +1,78 @@
+"""
+自动整理 Python 类中的方法顺序
+
+依赖：
+    需要安装 libcst 库（pip install libcst）
+
+适用场景：
+    适用于需要规范化类方法顺序、提升代码可读性和维护性的场景
+"""
+
+import libcst as cst
+import libcst.matchers as m
+
+
+class ClassMethodSorter(cst.CSTTransformer):
+    """
+    一个 CSTTransformer，用于对每个类中的方法按如下规则排序：
+    1. 魔法方法（即名称前后都有双下划线，如 __init__、__str__ 等），按名称排序，排在最前面；
+    2. 公共方法（不以下划线开头），按名称排序，排在中间；
+    3. 私有方法（以下划线开头，但不是魔法方法），按名称排序，排在最后；
+    其他类成员（如属性、docstring、嵌套类等）保持原有顺序，并始终排在方法之前。
+    """
+
+    def leave_ClassDef(self, original_node, updated_node):
+        """
+        在离开类定义节点时，对类体中的方法按指定规则排序。
+        非方法成员保持原有顺序，并排在方法之前。
+
+        :param original_node: 原始的 ClassDef 节点
+        :param updated_node: 可能已被修改的 ClassDef 节点
+        :return: 排序后的新的 ClassDef 节点
+        """
+        # 分离方法和其他成员
+        magic_methods = []
+        public_methods = []
+        private_methods = []
+        others = []
+        for stmt in updated_node.body.body:
+            if m.matches(stmt, m.FunctionDef()):
+                name = stmt.name.value  # type: ignore
+                # 魔法方法：前后都有双下划线
+                if name.startswith("__") and name.endswith("__"):
+                    magic_methods.append(stmt)
+                # 私有方法：以单下划线或双下划线开头，但不是魔法方法
+                elif name.startswith("_"):
+                    private_methods.append(stmt)
+                # 公共方法：不以下划线开头
+                else:
+                    public_methods.append(stmt)
+            else:
+                others.append(stmt)
+        # 分别排序
+        magic_methods_sorted = sorted(magic_methods, key=lambda f: f.name.value)
+        public_methods_sorted = sorted(public_methods, key=lambda f: f.name.value)
+        private_methods_sorted = sorted(private_methods, key=lambda f: f.name.value)
+        # 合并：其他成员 + 魔法方法 + 公共方法 + 私有方法
+        new_body = (
+            list(others)
+            + magic_methods_sorted
+            + public_methods_sorted
+            + private_methods_sorted
+        )
+        # 返回新的类定义节点
+        return updated_node.with_changes(
+            body=updated_node.body.with_changes(body=new_body)
+        )
+
+
+def reorder_class_methods(source_code):
+    """
+    对给定 Python 源代码中的每个类的方法和成员按照 ClassMethodSorter 定义的规则排序
+
+    :param source_code: 需要处理的 Python 源代码字符串
+    :return: 方法已按类型排序后的新源代码字符串
+    """
+    tree = cst.parse_module(source_code)
+    new_tree = tree.visit(ClassMethodSorter())
+    return new_tree.code

rbm_tools/env.py

@@ -0,0 +1,11 @@
+"""
+环境变量注册
+"""
+
+import os
+
+CONF_ENV = os.getenv("CONF_ENV", "test")
+HOSTNAME = os.getenv("HOSTNAME", "localhost")
+MODULE = os.getenv("MODULE", "defaultmodule")
+POD_IP = os.getenv("POD_IP", "127.0.0.1")
+POD_NAME = os.getenv("POD_NAME", "defaultpodname")

rbm_tools/logger.py

@@ -0,0 +1,278 @@
+"""
+日志类
+"""
+
+import json
+import logging
+import os
+from copy import copy
+from logging.handlers import QueueHandler, QueueListener, RotatingFileHandler
+from multiprocessing import Queue
+
+import click
+from opentelemetry import trace
+from opentelemetry.exporter.otlp.proto.http.trace_exporter import OTLPSpanExporter
+from opentelemetry.instrumentation.fastapi import FastAPIInstrumentor
+from opentelemetry.instrumentation.logging import LoggingInstrumentor
+from opentelemetry.instrumentation.openai import OpenAIInstrumentor
+from opentelemetry.instrumentation.openai.shared import chat_wrappers
+from opentelemetry.sdk.resources import Resource
+from opentelemetry.sdk.trace import TracerProvider
+from opentelemetry.sdk.trace.export import BatchSpanProcessor
+from pydantic import BaseModel, Field
+
+from rbm_tools.env import CONF_ENV, HOSTNAME, MODULE, POD_IP
+
+
+class ColourizedFormatter(logging.Formatter):
+    """
+    日志内容上颜色模块
+    """
+
+    level_name_colors = {
+        logging.DEBUG: lambda field: click.style(str(field), fg="bright_cyan"),
+        logging.WARNING: lambda field: click.style(str(field), fg="bright_yellow"),
+        logging.ERROR: lambda field: click.style(str(field), fg="bright_red"),
+        logging.CRITICAL: lambda field: click.style(str(field), fg="bright_red"),
+    }
+
+    def formatMessage(self, record: logging.LogRecord) -> str:
+        recordcopy = copy(record)
+        # message
+        if "color_message" in recordcopy.__dict__:
+            recordcopy.msg = recordcopy.__dict__["color_message"]
+            recordcopy.__dict__["message"] = recordcopy.getMessage()
+        recordcopy.__dict__["message"] = self.__color_field_by_level_no(
+            recordcopy.message, recordcopy.levelno
+        )
+        # asctime
+        recordcopy.__dict__["asctime"] = self.__color_field_by_fg(
+            recordcopy.asctime, "bright_green"
+        )
+        # levelname
+        seperator = " " * (8 - len(recordcopy.levelname))
+        levelname = self.__color_field_by_level_no(
+            recordcopy.levelname, recordcopy.levelno
+        )
+        recordcopy.__dict__["levelname"] = levelname + seperator
+        # otelTraceID
+        otel_trace_id = recordcopy.otelTraceID if hasattr(recordcopy, "otelTraceID") else "null"  # type: ignore # pylint: disable=C0301
+        seperator = " " * (32 - len(otel_trace_id))
+        recordcopy.__dict__["otelTraceID"] = otel_trace_id + seperator
+        # otelSpanID
+        otel_span_id = recordcopy.otelSpanID if hasattr(recordcopy, "otelSpanID") else "null"  # type: ignore # pylint: disable=C0301
+        seperator = " " * (16 - len(otel_span_id))
+        recordcopy.__dict__["otelSpanID"] = otel_span_id + seperator
+        # otelTraceSampled
+        otel_trace_sampled = recordcopy.otelTraceSampled if hasattr(recordcopy, "otelTraceSampled") else "null"  # type: ignore # pylint: disable=C0301
+        seperator = " " * (5 - len(str(otel_trace_sampled)))
+        recordcopy.__dict__["otelTraceSampled"] = str(otel_trace_sampled) + seperator
+        # name
+        recordcopy.__dict__["name"] = self.__color_field_by_fg(
+            recordcopy.name, "bright_magenta"
+        )
+        # module
+        recordcopy.__dict__["module"] = self.__color_field_by_fg(
+            recordcopy.module, "bright_blue"
+        )
+        # funcName
+        recordcopy.__dict__["funcName"] = self.__color_field_by_fg(
+            recordcopy.funcName, "bright_blue"
+        )
+        # lineno
+        recordcopy.__dict__["lineno"] = self.__color_field_by_fg(
+            str(recordcopy.lineno), "bright_blue"
+        )
+        # customized_data_info
+        customized_data_info = "{}"
+        if hasattr(recordcopy, "customized_data_info"):
+            try:
+                customized_data_info = json.dumps(
+                    recordcopy.customized_data_info, cls=CustomJSONEncoder, ensure_ascii=False  # type: ignore # pylint: disable=C0301
+                )
+            except Exception:  # pylint: disable=W0718
+                customized_data_info = recordcopy.customized_data_info  # type: ignore
+        recordcopy.__dict__["customized_data_info"] = self.__color_field_by_level_no(
+            field=customized_data_info,
+            level_no=recordcopy.levelno,
+            default_fg="bright_green",
+        )
+        return super().formatMessage(recordcopy)
+
+    def __color_field_by_fg(self, field: str, fg: str) -> str:
+        return click.style(str(field), fg=fg)
+
+    def __color_field_by_level_no(
+        self, field: str, level_no: int, default_fg: str = "reset"
+    ) -> str:
+        def default(field: str) -> str:
+            return click.style(str(field), fg=default_fg)
+
+        func = self.level_name_colors.get(level_no, default)
+        return func(field)
+
+
+class CustomJSONEncoder(json.JSONEncoder):
+    """
+    自定义的 JSONEncoder
+    """
+
+    def default(self, o):
+        if isinstance(o, BaseModel):
+            return o.model_dump()
+        return super().default(o)
+
+
+class NoneRootWarningFilter(logging.Filter):
+    """
+    非 root 的 logger 只允许 WARNING 级别的日志（除 uvicorn 相关的日志外）
+    """
+
+    def filter(self, record):
+        return (
+            "uvicorn" in record.name
+            or record.name == "root"
+            or record.levelno >= logging.WARNING
+        )
+
+
+def process_image_item(item: dict) -> dict:
+    """
+    自定义大模型调用中图片 base64 字符串的处理
+    """
+    customized_url_for_trace = item.get("image_url", {}).get(
+        "customized_url_for_trace", ""
+    )
+    if customized_url_for_trace:
+        return {"type": "image_url", "image_url": {"url": customized_url_for_trace}}
+    else:
+        return item
+
+
+async def _process_image_item(
+    item, trace_id, span_id, message_index, content_index  # pylint: disable=W0613
+) -> dict:
+    return process_image_item(item)
+
+
+# 不优雅的解决方案，如果要优雅的话需要开源社区支持
+chat_wrappers._process_image_item = _process_image_item  # pylint: disable=W0212
+
+# 获取环境变量
+app_name = MODULE
+app_env = CONF_ENV
+hostname = HOSTNAME
+ip = POD_IP
+# 日志配置
+format_str = "%(asctime)s | %(levelname)-8s | trace_id=%(otelTraceID)-32s span_id=%(otelSpanID)-16s sampled=%(otelTraceSampled)-5s | %(name)s | %(module)s:%(funcName)s:%(lineno)s - %(message)s | %(customized_data_info)s"  # pylint: disable=C0103,C0301
+log_config = {
+    "version": 1,
+    "loggers": {
+        "uvicorn": {"level": "INFO", "propagate": True},
+        "uvicorn.error": {"level": "INFO", "propagate": True},
+        "uvicorn.access": {"level": "INFO", "propagate": True},
+    },
+}
+resource = Resource.create(
+    attributes={
+        "service_name": app_name,
+        "service_env": app_env,
+        "service.name": app_name,
+        "service.env": app_env,
+    }
+)
+# trace Tempo 配置
+# os.environ["OTEL_EXPORTER_OTLP_TRACES_PROTOCOL"] = "http/protobuf"
+# os.environ["OTEL_EXPORTER_OTLP_ENDPOINT"] = tempo_url
+# os.environ["OTEL_EXPORTER_OTLP_HEADERS"] = f"Authorization=Basic%20{tempo_credentials}"
+# os.environ["OTEL_PYTHON_LOG_CORRELATION"] = "true"
+# os.environ["OTEL_TRACES_SAMPLER"] = "traceidratio"
+# os.environ["OTEL_TRACES_SAMPLER_ARG"] = "1.00"
+tracer_provider = TracerProvider(resource=resource)
+tracer_provider.add_span_processor(BatchSpanProcessor(OTLPSpanExporter()))
+# 设置 trace
+trace.set_tracer_provider(tracer_provider)
+# 初始化 root logger
+logger = logging.getLogger()
+logger.setLevel(logging.DEBUG)
+logger.addFilter(NoneRootWarningFilter())
+logger.handlers.clear()
+formatter = ColourizedFormatter(fmt=format_str)
+# 设置 logger 的控制台输出
+console_handler = logging.StreamHandler()
+console_handler.setLevel(logging.INFO)
+console_handler.addFilter(NoneRootWarningFilter())
+console_handler.setFormatter(formatter)
+logger.addHandler(console_handler)
+# 设置 logger 的文件输出（最多保留 3 + 1 个文件）
+file_handler = RotatingFileHandler(
+    filename=os.getcwd() + "/logger.log",
+    maxBytes=500 * 1024 * 1024,  # 每个文件最大 500 MB
+    backupCount=3,
+    encoding="utf-8"
+)
+file_handler.setLevel(logging.DEBUG)
+file_handler.addFilter(NoneRootWarningFilter())
+file_handler.setFormatter(formatter)
+# 使用 Python 官方推荐的异步日志方案，注意程序结束前需要调用 queue_listener.stop()
+report_queue = Queue(-1)
+queue_handler = QueueHandler(report_queue)
+queue_handler.setLevel(logging.DEBUG)
+queue_handler.addFilter(NoneRootWarningFilter())
+queue_listener = QueueListener(report_queue, file_handler, respect_handler_level=True)
+queue_listener.start()
+logger.addHandler(queue_handler)
+# 设置 Instrumentor
+FastAPIInstrumentor().instrument(
+    tracer_provider=tracer_provider, excluded_urls="metrics,health"
+)
+LoggingInstrumentor().instrument(set_logging_format=True, log_format=format_str)
+OpenAIInstrumentor().instrument(tracer_provider=tracer_provider)
+
+
+def close_logger():
+    """
+    确保日志和 trace 都写入完成
+    """
+    logger.info("close_logger start")
+    queue_listener.stop()
+    tracer_provider.shutdown()
+    logger.info("close_logger success")  # 这条日志永远写不了
+
+
+if __name__ == "__main__":
+    from typing import Generic, TypeVar
+
+    T = TypeVar("T")
+
+    class TestClass(BaseModel, Generic[T]):
+        """
+        用来测试日志输出的类
+        """
+
+        code: int = Field(..., title="业务相关的 code")
+        message: str = Field(..., title="业务相关的信息")
+        data: T = Field(..., title="数据")
+
+    a = TestClass[str](code=0, message="testa", data="testa")
+    b = TestClass[str](code=0, message="testb", data="testb")
+    logger.info(
+        "logger init success",
+        extra={
+            "customized_data_info": {
+                "test_list": [1, 2, 3],
+                # "test_object_list": [a, b],  # bad
+                "test_object_str_list": [a.model_dump_json(), b.model_dump_json()],
+                "test_dict": {"a": 1, "b": 2, "c": 3},
+                # "test_object_dict": {"a": a, "b": b},  # bad
+                "test_object_str_dict": {
+                    "a": a.model_dump_json(),
+                    "b": b.model_dump_json(),
+                },
+                "test_none": None,
+                # "test_bytes": "test_bytes".encode(),  # bad
+                "test_bytes_str": str("test_bytes".encode()),
+            }
+        },
+    )
+    close_logger()