diff --git a/cucyuqing/cmd/embedding.py b/cucyuqing/cmd/embedding.py index 0cb5bdf..ac8b7ec 100644 --- a/cucyuqing/cmd/embedding.py +++ b/cucyuqing/cmd/embedding.py @@ -9,6 +9,7 @@ from pydantic import BaseModel from typing import Any, Literal from cucyuqing.pg import pool, get_cur from cucyuqing.config import OPENAI_API_KEY, OPENAI_BASE_URL +from cucyuqing.utils import print EmbeddingModel = Literal["acge-large-zh", "text-embedding-3-large"] diff --git a/cucyuqing/cmd/es-sync.py b/cucyuqing/cmd/es-sync.py index dd0de20..f3b5a2a 100644 --- a/cucyuqing/cmd/es-sync.py +++ b/cucyuqing/cmd/es-sync.py @@ -12,6 +12,7 @@ from typing import AsyncIterable from cucyuqing.config import ES_API from cucyuqing.pg import get_cur, pool from cucyuqing.mysql import mysql +from cucyuqing.utils import print async def post(url: str, json: dict) -> dict: diff --git a/cucyuqing/utils.py b/cucyuqing/utils.py new file mode 100644 index 0000000..61a2d42 --- /dev/null +++ b/cucyuqing/utils.py @@ -0,0 +1,29 @@ +import sys +import inspect + +# 保存原始的 print 函数 +_print = print + +def print(*args, **kwargs): + # 获取调用者的栈帧信息 + frame = inspect.currentframe().f_back + module = inspect.getmodule(frame) + filename = inspect.getsourcefile(frame) + lineno = frame.f_lineno + # 获取调用者的函数名 + caller_frame = frame.f_back + caller_name = inspect.getframeinfo(caller_frame).function + + # 如果未指定输出文件,将输出重定向到 stderr + if 'file' not in kwargs or kwargs['file'] is sys.stdout: + kwargs['file'] = sys.stderr + + # 输出模块名、文件名和行号 + _print( + f"\033[92m[{module.__name__}.{caller_name}] {filename}:{lineno}\033[0m:", + *args, + **kwargs, + ) + +if __name__ == "__main__": + print("Hello, world!")