添加 print 函数

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"]
 

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:

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!")