88import logging
99from functools import lru_cache
1010from pathlib import Path
11- from typing import Literal , Dict
11+ from typing import Literal
1212from contextvars import ContextVar
1313import hydra
1414from rich .console import Console
1515from rich .logging import RichHandler
1616import asyncio
1717import threading
1818from contextlib import contextmanager
19+
1920TASK_CONTEXT_VAR : ContextVar [str | None ] = ContextVar ("CURRENT_TASK_ID" , default = None )
2021
22+
2123class ZMQLogHandler (logging .Handler ):
2224 def __init__ (self , addr = "tcp://127.0.0.1:6000" , tool_name = "unknown_tool" ):
2325 super ().__init__ ()
@@ -34,6 +36,7 @@ def emit(self, record):
3436 except Exception :
3537 self .handleError (record )
3638
39+
3740async def zmq_log_listener (bind_addr = "tcp://127.0.0.1:6000" ):
3841 ctx = zmq .asyncio .Context ()
3942 sock = ctx .socket (zmq .PULL )
@@ -47,23 +50,30 @@ async def zmq_log_listener(bind_addr="tcp://127.0.0.1:6000"):
4750 task_id , tool_name , msg = raw .split ("||" , 2 )
4851
4952 record = root_logger .makeRecord (
50- name = f' [TOOL] { tool_name } ' ,
53+ name = f" [TOOL] { tool_name } " ,
5154 level = logging .INFO ,
52- fn = "" , lno = 0 , msg = msg , args = (),
53- exc_info = None
55+ fn = "" ,
56+ lno = 0 ,
57+ msg = msg ,
58+ args = (),
59+ exc_info = None ,
5460 )
5561 record .task_id = task_id
5662
5763 root_logger .handle (record )
5864 else :
5965 root_logger .info (raw )
6066
67+
6168def start_zmq_listener ():
6269 loop = asyncio .new_event_loop ()
6370 asyncio .set_event_loop (loop )
6471 loop .run_until_complete (zmq_log_listener ())
6572
66- def setup_mcp_logging (level = "INFO" , addr = "tcp://127.0.0.1:6000" , tool_name = "unknown_tool" ):
73+
74+ def setup_mcp_logging (
75+ level = "INFO" , addr = "tcp://127.0.0.1:6000" , tool_name = "unknown_tool"
76+ ):
6777 root = logging .getLogger ()
6878 root .setLevel (level )
6979
@@ -78,29 +88,36 @@ def setup_mcp_logging(level="INFO", addr="tcp://127.0.0.1:6000", tool_name="unkn
7888 for h in logger .handlers [:]:
7989 logger .removeHandler (h )
8090 h .close ()
81- logger .propagate = True # 确保冒泡到 root
91+ logger .propagate = True # Ensure bubbling to root
8292
8393 # Re-add the ZMQ handler
8494 handler = ZMQLogHandler (addr = addr , tool_name = tool_name )
85- handler .setFormatter (logging .Formatter ("[TOOL] %(asctime)s %(levelname)s: %(message)s" ))
95+ handler .setFormatter (
96+ logging .Formatter ("[TOOL] %(asctime)s %(levelname)s: %(message)s" )
97+ )
8698 root .addHandler (handler )
8799
100+
88101def setup_log_record_factory ():
89102 old_factory = logging .getLogRecordFactory ()
103+
90104 def record_factory (* args , ** kwargs ):
91105 record = old_factory (* args , ** kwargs )
92106 record .task_id = TASK_CONTEXT_VAR .get ()
93107 return record
108+
94109 logging .setLogRecordFactory (record_factory )
95110
111+
96112class TaskFilter (logging .Filter ):
97113 def __init__ (self , task_id : str ):
98114 super ().__init__ ()
99115 self .task_id = task_id
100-
116+
101117 def filter (self , record : logging .LogRecord ) -> bool :
102118 return getattr (record , "task_id" , None ) == self .task_id
103119
120+
104121def make_task_logger (task_id : str , log_dir : Path ) -> logging .Handler :
105122 log_dir .mkdir (parents = True , exist_ok = True )
106123 file_path = log_dir / f"task_{ task_id } .log"
@@ -111,9 +128,10 @@ def make_task_logger(task_id: str, log_dir: Path) -> logging.Handler:
111128 logging .getLogger ().addHandler (fh )
112129 return fh
113130
131+
114132def remove_all_console_handlers ():
115133 """
116- 移除当前进程中所有 logger 上的 console handler (StreamHandler/RichHandler)。
134+ Remove all console handlers (StreamHandler/RichHandler) from all loggers in the current process.
117135 """
118136 for name , logger in logging .Logger .manager .loggerDict .items ():
119137 if isinstance (logger , logging .Logger ):
@@ -134,6 +152,7 @@ def remove_all_console_handlers():
134152 root_logger .removeHandler (h )
135153 h .close ()
136154
155+
137156@contextmanager
138157def task_logging_context (task_id : str , log_dir : Path ):
139158 token = TASK_CONTEXT_VAR .set (task_id )
@@ -145,21 +164,25 @@ def task_logging_context(task_id: str, log_dir: Path):
145164 logging .getLogger ().removeHandler (handler )
146165 handler .close ()
147166
167+
148168def init_logging_for_benchmark_evaluation (print_task_logs = False ):
149- threading .Thread (target = start_zmq_listener , daemon = True ).start () #monitoring tool logs
150- logging .basicConfig (handlers = [])
169+ threading .Thread (
170+ target = start_zmq_listener , daemon = True
171+ ).start () # monitoring tool logs
172+ logging .basicConfig (handlers = [])
151173 setup_log_record_factory ()
152174 if not print_task_logs :
153- remove_all_console_handlers ()
175+ remove_all_console_handlers ()
176+
154177
155178@lru_cache
156179def bootstrap_logger (
157180 level : Literal ["DEBUG" , "INFO" , "WARNING" , "ERROR" , "CRITICAL" ] | int = "INFO" ,
158181 logger_name : str = "miroflow" ,
159182 logger : logging .Logger | None = None ,
160- log_dir : str | Path | None = None , # 日志存储目录
161- log_filename : str = "miroflow.log" , # 默认日志文件名
162- to_console : bool = True , # 是否显示到 console
183+ log_dir : str | Path | None = None , # Log storage directory
184+ log_filename : str = "miroflow.log" , # Default log filename
185+ to_console : bool = True , # Whether to display to console
163186) -> logging .Logger :
164187 """Configure only this logger, not the root logger"""
165188 if logger is None :
@@ -173,7 +196,7 @@ def bootstrap_logger(
173196 console = Console (
174197 stderr = True ,
175198 width = 200 ,
176- color_system = None ,
199+ color_system = None ,
177200 force_terminal = False ,
178201 legacy_windows = False ,
179202 ),
@@ -191,9 +214,9 @@ def bootstrap_logger(
191214 log_dir .mkdir (parents = True , exist_ok = True )
192215 file_path = log_dir / log_filename
193216 file_handler = logging .FileHandler (file_path , encoding = "utf-8" )
194- file_handler .setFormatter (logging . Formatter (
195- "%(asctime)s [%(levelname)s] %(name)s: %(message)s"
196- ))
217+ file_handler .setFormatter (
218+ logging . Formatter ( "%(asctime)s [%(levelname)s] %(name)s: %(message)s" )
219+ )
197220 logger .addHandler (file_handler )
198221
199222 logger .setLevel (level )
0 commit comments