bella-openapi 1.0.2__py3-none-any.whl → 1.0.2.2__py3-none-any.whl

bella_openapi/__init__.py

@@ -1,22 +1,26 @@
-from .authorize import validate_token, support_model, account_balance_enough, check_configuration
-from .log import operation_log, submit_log
-from .openapi_contexvar import trace_id_context, caller_id_context, request_url_context
-from .auth_billing import ErrorInfo, async_authenticate_decorator_args, authenticate_user, print_context, \
-get_context, set_context, clean_context, report
-__all__ = ["validate_token", "operation_log",
-           "support_model",
-           "account_balance_enough",
-           "check_configuration",
-           "trace_id_context",
-           "caller_id_context",
-           "request_url_context",
-           "submit_log",
-           "ErrorInfo",
-           "async_authenticate_decorator_args",
-           "authenticate_user",
-           "print_context",
-           "get_context",
-           "set_context",
-           "clean_context",
-           "report"
-           ]
+from .authorize import validate_token, support_model, account_balance_enough, check_configuration
+from .log import operation_log, submit_log
+from .openapi_contexvar import trace_id_context, caller_id_context, request_url_context
+from .auth_billing import ErrorInfo, async_authenticate_decorator_args, authenticate_user, print_context, \
+get_context, set_context, clean_context, report
+from .domtree import StandardDomTree, StandardNode
+
+__all__ = ["validate_token", "operation_log",
+           "support_model",
+           "account_balance_enough",
+           "check_configuration",
+           "trace_id_context",
+           "caller_id_context",
+           "request_url_context",
+           "submit_log",
+           "ErrorInfo",
+           "async_authenticate_decorator_args",
+           "authenticate_user",
+           "print_context",
+           "get_context",
+           "set_context",
+           "clean_context",
+           "report",
+           "StandardDomTree",
+           "StandardNode"
+           ]

bella_openapi/auth_billing.py

@@ -1,91 +1,91 @@
-from .log import operation_log
-from .authorize import validate_token, check_configuration, account_balance_enough
-from .openapi_contexvar import trace_id_context, caller_id_context, request_url_context
-from pydantic import BaseModel
-import uuid
-from fastapi import Request
-import logging
-# 创建一个日志记录器
-logger = logging.getLogger(__name__)
-logger.setLevel(logging.INFO)  # 设置日志级别为INFO
-
-# 创建一个控制台处理器，并设置其级别和格式
-console_handler = logging.StreamHandler()
-console_handler.setLevel(logging.INFO)
-formatter = logging.Formatter('%(asctime)s - %(levelname)s - %(message)s')
-console_handler.setFormatter(formatter)
-
-# 将处理器添加到日志记录器
-logger.addHandler(console_handler)
-
-class ErrorInfo(BaseModel):
-    task_id: str = ""
-    result: str = ""
-    status: int = 40000001
-    message: str = ""
-
-def async_authenticate_decorator_args(end_point):
-    def async_authenticate_decorator(func):
-        async def wrapper(*args, **kwargs):
-            request_arg = None
-            for arg in args:
-                if type(arg) == Request:
-                    request_arg = arg
-                    break
-            if request_arg is not None:
-                task_id = str(uuid.uuid4())
-                supported, error_json, caller_id = authenticate_user(request_arg.headers.get("Authorization"), task_id)
-                if not supported:
-                    return error_json
-                t_token, c_token, r_token = set_context(task_id, caller_id, end_point)
-                result = await func(*args, **kwargs)
-                clean_context(t_token, c_token, r_token)
-                return result
-            else:
-                logger.warn("please check your request param,have not a param's type is Request of fastapi!")
-            return func(*args, **kwargs)
-        return wrapper
-    return async_authenticate_decorator
-
-
-def authenticate_user(token, task_id):
-    if check_configuration():
-        if token is None:
-            return False, ErrorInfo(task_id=task_id, result="", status=40000001, message="token is missing").dict(), ""
-        try:
-            caller_id = validate_token(token)
-            balance = account_balance_enough(token, cost=1.0)
-
-        except Exception as e:
-            return False, ErrorInfo(task_id=task_id, result="", status=40000001, message=str(e)[:100]).dict(), ""
-        return balance, ErrorInfo(task_id=task_id, result="", status=40000001, message="Insufficient balance").dict(), caller_id
-    else:
-        return True, "", ""
-
-
-@operation_log(op_type='upload_cost_log', is_cost_log=True, ucid_key="ucid")
-def upload_cost_log(result_obejct, ucid):
-    response = result_obejct.dict()
-    return response
-
-def print_context(log):
-    logger.info(f"{log} trace_id:{trace_id_context.get()}, caller_id:{caller_id_context.get()}, end_point:{request_url_context.get()}")
-
-def get_context():
-    return trace_id_context.get(), caller_id_context.get(), request_url_context.get()
-
-def set_context(trace_id, caller_id, end_point):
-    t_token = trace_id_context.set(trace_id)
-    c_token = caller_id_context.set(caller_id)
-    r_token = request_url_context.set(end_point)
-    return t_token, c_token, r_token
-
-def clean_context(t_token, c_token, r_token):
-    trace_id_context.reset(t_token)
-    caller_id_context.reset(c_token)
-    request_url_context.reset(r_token)
-
-def report(result_obejct, ucid = ""):
-    if not check_configuration():
-        return
-    upload_cost_log(result_obejct, ucid)
+from .log import operation_log
+from .authorize import validate_token, check_configuration, account_balance_enough
+from .openapi_contexvar import trace_id_context, caller_id_context, request_url_context
+from pydantic import BaseModel
+import uuid
+from fastapi import Request
+import logging
+# 创建一个日志记录器
+logger = logging.getLogger(__name__)
+logger.setLevel(logging.INFO)  # 设置日志级别为INFO
+
+# 创建一个控制台处理器，并设置其级别和格式
+console_handler = logging.StreamHandler()
+console_handler.setLevel(logging.INFO)
+formatter = logging.Formatter('%(asctime)s - %(levelname)s - %(message)s')
+console_handler.setFormatter(formatter)
+
+# 将处理器添加到日志记录器
+logger.addHandler(console_handler)
+
+class ErrorInfo(BaseModel):
+    task_id: str = ""
+    result: str = ""
+    status: int = 40000001
+    message: str = ""
+
+def async_authenticate_decorator_args(end_point):
+    def async_authenticate_decorator(func):
+        async def wrapper(*args, **kwargs):
+            request_arg = None
+            for arg in args:
+                if type(arg) == Request:
+                    request_arg = arg
+                    break
+            if request_arg is not None:
+                task_id = str(uuid.uuid4())
+                supported, error_json, caller_id = authenticate_user(request_arg.headers.get("Authorization"), task_id)
+                if not supported:
+                    return error_json
+                t_token, c_token, r_token = set_context(task_id, caller_id, end_point)
+                result = await func(*args, **kwargs)
+                clean_context(t_token, c_token, r_token)
+                return result
+            else:
+                logger.warn("please check your request param,have not a param's type is Request of fastapi!")
+            return func(*args, **kwargs)
+        return wrapper
+    return async_authenticate_decorator
+
+
+def authenticate_user(token, task_id):
+    if check_configuration():
+        if token is None:
+            return False, ErrorInfo(task_id=task_id, result="", status=40000001, message="token is missing").dict(), ""
+        try:
+            caller_id = validate_token(token)
+            balance = account_balance_enough(token, cost=1.0)
+
+        except Exception as e:
+            return False, ErrorInfo(task_id=task_id, result="", status=40000001, message=str(e)[:100]).dict(), ""
+        return balance, ErrorInfo(task_id=task_id, result="", status=40000001, message="Insufficient balance").dict(), caller_id
+    else:
+        return True, "", ""
+
+
+@operation_log(op_type='upload_cost_log', is_cost_log=True, ucid_key="ucid")
+def upload_cost_log(result_obejct, ucid):
+    response = result_obejct.dict()
+    return response
+
+def print_context(log):
+    logger.info(f"{log} trace_id:{trace_id_context.get()}, caller_id:{caller_id_context.get()}, end_point:{request_url_context.get()}")
+
+def get_context():
+    return trace_id_context.get(), caller_id_context.get(), request_url_context.get()
+
+def set_context(trace_id, caller_id, end_point):
+    t_token = trace_id_context.set(trace_id)
+    c_token = caller_id_context.set(caller_id)
+    r_token = request_url_context.set(end_point)
+    return t_token, c_token, r_token
+
+def clean_context(t_token, c_token, r_token):
+    trace_id_context.reset(t_token)
+    caller_id_context.reset(c_token)
+    request_url_context.reset(r_token)
+
+def report(result_obejct, ucid = ""):
+    if not check_configuration():
+        return
+    upload_cost_log(result_obejct, ucid)

bella_openapi/authorize.py

@@ -1,61 +1,61 @@
-import httpx
-import logging
-from .config import openapi_config
-from .exception import AuthorizationException
-
-
-def check_configuration() -> bool:
-    res = openapi_config.OPENAPI_HOST is not None
-    return res
-
-
-def validate_token(token: str) -> str:
-    """
-    根据传入token， 解析用户身份
-    :param token:
-    :return: 用户身份，返回账户id
-    :raises: AuthorizationException 如果token无效， 抛出异常
-    """
-    return _validate_token_request(token)
-
-
-def support_model(token: str, model: str) -> bool:
-    """
-    根据传入token，判断用户是否有对应的模型权限
-    :param token: 用户token
-    :param model: 模型名
-    :return: bool
-    """
-    url = openapi_config.OPENAPI_HOST + "/v1/openapi/support/model"
-    # 使用httpx发送get请求
-    response = httpx.get(url, headers={"Authorization": token}, params={"model": model})
-    if response.status_code == 200:
-        return response.json()['data']
-    else:
-        raise AuthorizationException(response.text, response.status_code)
-
-
-def account_balance_enough(token: str, cost: float = 0) -> bool:
-    """
-    账户余额判断，支持传入指定判断阈值，服务方可以在每次用户请求前，根据本次预估花费，对用户余额进行校验，若余额不足，则可拒绝请求
-    :param token: 用户token
-    :param cost: 消耗金额, 单位：元/RMB
-    :return: bool
-    """
-    url = openapi_config.OPENAPI_HOST + "/v1/openapi/check/account/balance"
-    # 使用httpx发送get请求
-    response = httpx.get(url, headers={"Authorization": token}, params={"cost": cost})
-    if response.status_code == 200:
-        return response.json()['data']
-    else:
-        raise AuthorizationException(response.text, response.status_code)
-
-
-def _validate_token_request(token: str) -> str:
-    url = openapi_config.OPENAPI_HOST + "/v1/openapi/validate/tokens"
-    # 使用httpx发送get请求
-    response = httpx.get(url, headers={"Authorization": token})
-    if response.status_code == 200:
-        return response.json()['data']
-    else:
-        raise AuthorizationException(response.text, response.status_code)
+import httpx
+import logging
+from .config import openapi_config
+from .exception import AuthorizationException
+
+
+def check_configuration() -> bool:
+    res = openapi_config.OPENAPI_HOST is not None
+    return res
+
+
+def validate_token(token: str) -> str:
+    """
+    根据传入token， 解析用户身份
+    :param token:
+    :return: 用户身份，返回账户id
+    :raises: AuthorizationException 如果token无效， 抛出异常
+    """
+    return _validate_token_request(token)
+
+
+def support_model(token: str, model: str) -> bool:
+    """
+    根据传入token，判断用户是否有对应的模型权限
+    :param token: 用户token
+    :param model: 模型名
+    :return: bool
+    """
+    url = openapi_config.OPENAPI_HOST + "/v1/openapi/support/model"
+    # 使用httpx发送get请求
+    response = httpx.get(url, headers={"Authorization": token}, params={"model": model})
+    if response.status_code == 200:
+        return response.json()['data']
+    else:
+        raise AuthorizationException(response.text, response.status_code)
+
+
+def account_balance_enough(token: str, cost: float = 0) -> bool:
+    """
+    账户余额判断，支持传入指定判断阈值，服务方可以在每次用户请求前，根据本次预估花费，对用户余额进行校验，若余额不足，则可拒绝请求
+    :param token: 用户token
+    :param cost: 消耗金额, 单位：元/RMB
+    :return: bool
+    """
+    url = openapi_config.OPENAPI_HOST + "/v1/openapi/check/account/balance"
+    # 使用httpx发送get请求
+    response = httpx.get(url, headers={"Authorization": token}, params={"cost": cost})
+    if response.status_code == 200:
+        return response.json()['data']
+    else:
+        raise AuthorizationException(response.text, response.status_code)
+
+
+def _validate_token_request(token: str) -> str:
+    url = openapi_config.OPENAPI_HOST + "/v1/openapi/validate/tokens"
+    # 使用httpx发送get请求
+    response = httpx.get(url, headers={"Authorization": token})
+    if response.status_code == 200:
+        return response.json()['data']
+    else:
+        raise AuthorizationException(response.text, response.status_code)

bella_openapi/bella_trace/__init__.py

@@ -1,13 +1,13 @@
-# -*- coding: utf-8 -*-
-# ======================
-# Date    : 2024/12/30
-# Author  : Liu Yuchen
-# Content : 
-# 协议规范：https://doc.weixin.qq.com/doc/w3_AagAxwZdAD4dsCIEHU3RL26Knh1x8?scode=AJMA1Qc4AAwYUI6MJrAAEASgZXANE
-# ======================
-from ._context import TraceContext, TRACE_ID
-from .fastapi_interceptor import FastapiBellaTraceMiddleware
-import bella_openapi.bella_trace.trace_requests as requests
-from .record_log import trace, BellaTraceHandler
-
-__all__ = ["TraceContext", "TRACE_ID", "FastapiBellaTraceMiddleware", "requests", "trace", "BellaTraceHandler"]
+# -*- coding: utf-8 -*-
+# ======================
+# Date    : 2024/12/30
+# Author  : Liu Yuchen
+# Content : 
+# 协议规范：https://doc.weixin.qq.com/doc/w3_AagAxwZdAD4dsCIEHU3RL26Knh1x8?scode=AJMA1Qc4AAwYUI6MJrAAEASgZXANE
+# ======================
+from ._context import TraceContext, TRACE_ID
+from .fastapi_interceptor import FastapiBellaTraceMiddleware
+import bella_openapi.bella_trace.trace_requests as requests
+from .record_log import trace, BellaTraceHandler
+
+__all__ = ["TraceContext", "TRACE_ID", "FastapiBellaTraceMiddleware", "requests", "trace", "BellaTraceHandler"]

bella_openapi/bella_trace/_context.py

@@ -1,61 +1,61 @@
-# -*- coding: utf-8 -*-
-# ======================
-# Date    : 2024/12/30
-# Author  : Liu Yuchen
-# Content : 
-# 
-# ======================
-import os
-import uuid
-from contextvars import ContextVar
-
-__all__ = ["TraceContext", "TRACE_ID"]
-
-
-_trace_id = ContextVar("bella_trace_id", default="")
-_mock_request = ContextVar("mock_request", default="false")
-
-
-TRACE_ID = "X-BELLA-TRACE-ID"
-MOCK_REQUEST = "X-BELLA-MOCK-REQUEST"
-
-
-class _TraceContext(object):
-    @property
-    def trace_id(self) -> str:
-        return _trace_id.get()
-    
-    @trace_id.setter
-    def trace_id(self, value):
-        _trace_id.set(value)
-    
-    @property
-    def service_id(self) -> str:
-        return _get_service_id()
-    
-    @staticmethod
-    def generate_trace_id() -> str:
-        return f"{_get_service_id()}-{uuid.uuid4().hex}"
-    
-    @property
-    def mock_request(self):
-        return _mock_request.get()
-    
-    @mock_request.setter
-    def mock_request(self, value: str):
-        _mock_request.set(value)
-    
-    @property
-    def is_mock_request(self) -> bool:
-        return self.mock_request.lower() == "true"
-    
-    @property
-    def headers(self) -> dict:
-        return {TRACE_ID: self.trace_id, MOCK_REQUEST: self.mock_request}
-
-
-TraceContext = _TraceContext()
-
-
-def _get_service_id() -> str:
-    return os.environ.get("SERVICE_ID", "")
+# -*- coding: utf-8 -*-
+# ======================
+# Date    : 2024/12/30
+# Author  : Liu Yuchen
+# Content : 
+# 
+# ======================
+import os
+import uuid
+from contextvars import ContextVar
+
+__all__ = ["TraceContext", "TRACE_ID"]
+
+
+_trace_id = ContextVar("bella_trace_id", default="")
+_mock_request = ContextVar("mock_request", default="false")
+
+
+TRACE_ID = "X-BELLA-TRACE-ID"
+MOCK_REQUEST = "X-BELLA-MOCK-REQUEST"
+
+
+class _TraceContext(object):
+    @property
+    def trace_id(self) -> str:
+        return _trace_id.get()
+    
+    @trace_id.setter
+    def trace_id(self, value):
+        _trace_id.set(value)
+    
+    @property
+    def service_id(self) -> str:
+        return _get_service_id()
+    
+    @staticmethod
+    def generate_trace_id() -> str:
+        return f"{_get_service_id()}-{uuid.uuid4().hex}"
+    
+    @property
+    def mock_request(self):
+        return _mock_request.get()
+    
+    @mock_request.setter
+    def mock_request(self, value: str):
+        _mock_request.set(value)
+    
+    @property
+    def is_mock_request(self) -> bool:
+        return self.mock_request.lower() == "true"
+    
+    @property
+    def headers(self) -> dict:
+        return {TRACE_ID: self.trace_id, MOCK_REQUEST: self.mock_request}
+
+
+TraceContext = _TraceContext()
+
+
+def _get_service_id() -> str:
+    return os.environ.get("SERVICE_ID", "")

bella_openapi/bella_trace/fastapi_interceptor.py

@@ -1,28 +1,28 @@
-# -*- coding: utf-8 -*-
-# ======================
-# Date    : 2024/12/30
-# Author  : Liu Yuchen
-# Content : 
-# 
-# ======================
-from starlette.middleware.base import BaseHTTPMiddleware
-
-from ._context import TraceContext, TRACE_ID, MOCK_REQUEST
-
-__all__ = ["FastapiBellaTraceMiddleware"]
-
-
-class FastapiBellaTraceMiddleware(BaseHTTPMiddleware):
-    async def dispatch(self, request, call_next):
-        # 设置 trace_id
-        if trace_id_h := request.headers.get(TRACE_ID):
-            TraceContext.trace_id = trace_id_h
-        else:
-            TraceContext.trace_id = TraceContext.generate_trace_id()
-        
-        if mock_request_h := request.headers.get(MOCK_REQUEST):
-            TraceContext.mock_request = mock_request_h
-        
-        # 继续处理请求
-        return await call_next(request)
-
+# -*- coding: utf-8 -*-
+# ======================
+# Date    : 2024/12/30
+# Author  : Liu Yuchen
+# Content : 
+# 
+# ======================
+from starlette.middleware.base import BaseHTTPMiddleware
+
+from ._context import TraceContext, TRACE_ID, MOCK_REQUEST
+
+__all__ = ["FastapiBellaTraceMiddleware"]
+
+
+class FastapiBellaTraceMiddleware(BaseHTTPMiddleware):
+    async def dispatch(self, request, call_next):
+        # 设置 trace_id
+        if trace_id_h := request.headers.get(TRACE_ID):
+            TraceContext.trace_id = trace_id_h
+        else:
+            TraceContext.trace_id = TraceContext.generate_trace_id()
+        
+        if mock_request_h := request.headers.get(MOCK_REQUEST):
+            TraceContext.mock_request = mock_request_h
+        
+        # 继续处理请求
+        return await call_next(request)
+