feat(logging): support tool logs and per-task log storage (#69)

* feat(logging): support tool logs and per-task log storage

* Update src/logging/logger.py

Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com>

* Update src/logging/logger.py

Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com>

* Update src/logging/logger.py

Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com>

---------

Co-authored-by: Bin Wang <bwang28c@gmail.com>
Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com>

Author: 3 people

common_benchmark.py

@@ -19,13 +19,13 @@
 from omegaconf import DictConfig, OmegaConf
 
 from utils.eval_utils import verify_answer_for_datasets
-from src.logging.logger import bootstrap_logger
+from src.logging.logger import bootstrap_logger, task_logging_context, init_logging_for_benchmark_evaluation
 from config import config_name, config_path
 from src.core.pipeline import (
     create_pipeline_components,
     execute_task_pipeline,
 )
-
+init_logging_for_benchmark_evaluation(print_task_logs=False)
 
 class TaskStatus(StrEnum):
     PENDING = "pending"
@@ -179,7 +179,8 @@ async def run_single_task(self, task: BenchmarkTask) -> BenchmarkResult:
         found_correct_answer = False
 
         # Print debug info about log directory
-        print(f"  Current log directory: {self.output_dir}")
+        print(f"  Current result directory: {self.output_dir}")
+        print(f"  Current task log directory: {self.output_dir}/task_logs")
 
         try:
             # Prepare task
@@ -371,7 +372,9 @@ async def run_parallel_inference(
 
         async def run_with_semaphore(task):
             async with semaphore:
-                return await self.run_single_task(task)
+                with task_logging_context(task.task_id, self.get_log_dir()):
+                        result = await self.run_single_task(task)
+                return result
 
         # Shuffle tasks to avoid order bias and improve balancing
         shuffled_tasks = tasks.copy()

pyproject.toml

@@ -30,6 +30,7 @@ dependencies = [
     "requests>=2.32.4",
     "filetype>=1.2.0",
     "Pillow",
+    "pyzmq"
 ]
 
 [build-system]

src/logging/logger.py

@@ -2,44 +2,201 @@
 #
 # SPDX-License-Identifier: Apache-2.0
 
+import os
+import zmq
+import zmq.asyncio
 import logging
 from functools import lru_cache
-from typing import Literal
-
+from pathlib import Path
+from typing import Literal, Dict
+from contextvars import ContextVar
 import hydra
 from rich.console import Console
 from rich.logging import RichHandler
+import asyncio
+import threading
+from contextlib import contextmanager
+TASK_CONTEXT_VAR: ContextVar[str | None] = ContextVar("CURRENT_TASK_ID", default=None)
+
+class ZMQLogHandler(logging.Handler):
+    def __init__(self, addr="tcp://127.0.0.1:6000", tool_name="unknown_tool"):
+        super().__init__()
+        ctx = zmq.Context()
+        self.sock = ctx.socket(zmq.PUSH)
+        self.sock.connect(addr)
+        self.task_id = os.environ.get("TASK_ID", "0")
+        self.tool_name = tool_name
+
+    def emit(self, record):
+        try:
+            msg = f"{record.getMessage()}"
+            self.sock.send_string(f"{self.task_id}||{self.tool_name}||{msg}")
+        except Exception:
+            self.handleError(record)
+
+async def zmq_log_listener(bind_addr="tcp://127.0.0.1:6000"):
+    ctx = zmq.asyncio.Context()
+    sock = ctx.socket(zmq.PULL)
+    sock.bind(bind_addr)
+
+    root_logger = logging.getLogger()
+
+    while True:
+        raw = await sock.recv_string()
+        if "|" in raw:
+            task_id, tool_name, msg = raw.split("||", 2)
+
+            record = root_logger.makeRecord(
+                name=f'[TOOL] {tool_name}',
+                level=logging.INFO,
+                fn="", lno=0, msg=msg, args=(),
+                exc_info=None
+            )
+            record.task_id = task_id
+
+            root_logger.handle(record)
+        else:
+            root_logger.info(raw)
+
+def start_zmq_listener():
+    loop = asyncio.new_event_loop()
+    asyncio.set_event_loop(loop)
+    loop.run_until_complete(zmq_log_listener())
+
+def setup_mcp_logging(level="INFO", addr="tcp://127.0.0.1:6000", tool_name="unknown_tool"):
+    root = logging.getLogger()
+    root.setLevel(level)
+
+    # Remove root handlers
+    for h in root.handlers[:]:
+        root.removeHandler(h)
+        h.close()
 
+    # Remove all handlers from fastmcp child loggers
+    for name, logger in logging.Logger.manager.loggerDict.items():
+        if isinstance(logger, logging.Logger):
+            for h in logger.handlers[:]:
+                logger.removeHandler(h)
+                h.close()
+            logger.propagate = True  # 确保冒泡到 root
+
+    # Re-add the ZMQ handler
+    handler = ZMQLogHandler(addr=addr, tool_name=tool_name)
+    handler.setFormatter(logging.Formatter("[TOOL] %(asctime)s %(levelname)s: %(message)s"))
+    root.addHandler(handler)
+
+def setup_log_record_factory():
+    old_factory = logging.getLogRecordFactory()
+    def record_factory(*args, **kwargs):
+        record = old_factory(*args, **kwargs)
+        record.task_id = TASK_CONTEXT_VAR.get()
+        return record
+    logging.setLogRecordFactory(record_factory)
+
+class TaskFilter(logging.Filter):
+    def __init__(self, task_id: str):
+        super().__init__()
+        self.task_id = task_id
+    
+    def filter(self, record: logging.LogRecord) -> bool:
+        return getattr(record, "task_id", None) == self.task_id
+
+def make_task_logger(task_id: str, log_dir: Path) -> logging.Handler:
+    log_dir.mkdir(parents=True, exist_ok=True)
+    file_path = log_dir / f"task_{task_id}.log"
+    fh = logging.FileHandler(file_path, encoding="utf-8")
+    fmt = logging.Formatter("%(asctime)s [%(levelname)s] %(name)s: %(message)s")
+    fh.setFormatter(fmt)
+    fh.addFilter(TaskFilter(task_id))
+    logging.getLogger().addHandler(fh)
+    return fh
+
+def remove_all_console_handlers():
+    """
+    移除当前进程中所有 logger 上的 console handler (StreamHandler/RichHandler)。
+    """
+    for name, logger in logging.Logger.manager.loggerDict.items():
+        if isinstance(logger, logging.Logger):
+            handlers_to_remove = []
+            for h in logger.handlers:
+                if isinstance(h, logging.StreamHandler) or isinstance(h, RichHandler):
+                    handlers_to_remove.append(h)
+            for h in handlers_to_remove:
+                logger.removeHandler(h)
+                h.close()
+
+    root_logger = logging.getLogger()
+    handlers_to_remove = []
+    for h in root_logger.handlers:
+        if isinstance(h, logging.StreamHandler):
+            handlers_to_remove.append(h)
+    for h in handlers_to_remove:
+        root_logger.removeHandler(h)
+        h.close()
+
+@contextmanager
+def task_logging_context(task_id: str, log_dir: Path):
+    token = TASK_CONTEXT_VAR.set(task_id)
+    handler = make_task_logger(task_id, log_dir / "task_logs")
+    try:
+        yield
+    finally:
+        TASK_CONTEXT_VAR.reset(token)
+        logging.getLogger().removeHandler(handler)
+        handler.close()
+
+def init_logging_for_benchmark_evaluation(print_task_logs=False):
+    threading.Thread(target=start_zmq_listener, daemon=True).start() #monitoring tool logs
+    logging.basicConfig(handlers=[]) 
+    setup_log_record_factory()
+    if not print_task_logs:
+        remove_all_console_handlers()  
 
 @lru_cache
 def bootstrap_logger(
     level: Literal["DEBUG", "INFO", "WARNING", "ERROR", "CRITICAL"] | int = "INFO",
+    logger_name: str = "miroflow",
     logger: logging.Logger | None = None,
+    log_dir: str | Path | None = None,   # 日志存储目录
+    log_filename: str = "miroflow.log",  # 默认日志文件名
+    to_console: bool = True,             # 是否显示到 console
 ) -> logging.Logger:
     """Configure only this logger, not the root logger"""
     if logger is None:
-        logger = logging.getLogger("miroflow")
+        logger = logging.getLogger(logger_name)
+
     for handler in logger.handlers[:]:
         logger.removeHandler(handler)
 
-    # use rich for better readability of stack trace.
-    handler = RichHandler(
-        console=Console(
-            stderr=True,
-            width=200,
-            color_system=None,  # Disable colors to avoid ANSI escape sequences in log files
-            force_terminal=False,  # Don't force terminal mode
-            legacy_windows=False,
-        ),
-        rich_tracebacks=True,
-        tracebacks_suppress=[hydra],
-        tracebacks_show_locals=True,
-        show_level=False,
-    )
-    formatter = logging.Formatter("[%(levelname)s] %(message)s")
-    handler.setFormatter(formatter)
-    logger.addHandler(handler)
+    if to_console:
+        handler = RichHandler(
+            console=Console(
+                stderr=True,
+                width=200,
+                color_system=None,   
+                force_terminal=False,
+                legacy_windows=False,
+            ),
+            rich_tracebacks=True,
+            tracebacks_suppress=[hydra],
+            tracebacks_show_locals=True,
+            show_level=False,
+        )
+        formatter = logging.Formatter("[%(levelname)s] %(message)s")
+        handler.setFormatter(formatter)
+        logger.addHandler(handler)
+
+    if log_dir is not None:
+        log_dir = Path(log_dir)
+        log_dir.mkdir(parents=True, exist_ok=True)
+        file_path = log_dir / log_filename
+        file_handler = logging.FileHandler(file_path, encoding="utf-8")
+        file_handler.setFormatter(logging.Formatter(
+            "%(asctime)s [%(levelname)s] %(name)s: %(message)s"
+        ))
+        logger.addHandler(file_handler)
+
     logger.setLevel(level)
-    logger.propagate = False
+    logger.propagate = True
 
     return logger

src/tool/manager.py

@@ -20,6 +20,14 @@
 
 R = TypeVar("R")
 
+def update_server_params_with_context_var(server_params: StdioServerParameters) -> StdioServerParameters:
+    """
+    Update the server params with the context var.
+    """
+    from src.logging.logger import TASK_CONTEXT_VAR
+    if TASK_CONTEXT_VAR.get() is not None:
+        server_params.env["TASK_ID"] = TASK_CONTEXT_VAR.get()
+    return server_params
 
 def with_timeout(timeout_s: float = 300.0):
     """
@@ -108,7 +116,7 @@ async def _find_servers_with_tool(self, tool_name):
 
             try:
                 if isinstance(server_params, StdioServerParameters):
-                    async with stdio_client(server_params) as (read, write):
+                    async with stdio_client(update_server_params_with_context_var(server_params)) as (read, write):
                         async with ClientSession(
                             read, write, sampling_callback=None
                         ) as session:
@@ -168,7 +176,7 @@ async def get_all_tool_definitions(self):
 
             try:
                 if isinstance(server_params, StdioServerParameters):
-                    async with stdio_client(server_params) as (read, write):
+                    async with stdio_client(update_server_params_with_context_var(server_params)) as (read, write):
                         async with ClientSession(
                             read, write, sampling_callback=None
                         ) as session:
@@ -342,7 +350,7 @@ async def execute_tool_call(self, server_name, tool_name, arguments) -> Any:
             try:
                 result_content = None
                 if isinstance(server_params, StdioServerParameters):
-                    async with stdio_client(server_params) as (read, write):
+                    async with stdio_client(update_server_params_with_context_var(server_params)) as (read, write):
                         async with ClientSession(
                             read, write, sampling_callback=None
                         ) as session:

src/tool/mcp_servers/audio_mcp_server.py

@@ -25,9 +25,10 @@
 )
 
 # Initialize FastMCP server
+from src.logging.logger import setup_mcp_logging
+setup_mcp_logging(tool_name=os.path.basename(__file__))
 mcp = FastMCP("audio-mcp-server")
 
-
 def _get_audio_extension(url: str, content_type: str = None) -> str:
     """
     Determine the appropriate audio file extension from URL or content type.
@@ -289,4 +290,4 @@ async def audio_question_answering(audio_path_or_url: str, question: str) -> str
 
 
 if __name__ == "__main__":
-    mcp.run(transport="stdio")
+    mcp.run(transport="stdio",show_banner=False)

src/tool/mcp_servers/browser_session.py

@@ -97,4 +97,4 @@ async def test_persistent_session():
 
 
 if __name__ == "__main__":
-    asyncio.run(test_persistent_session())
+    asyncio.run(test_persistent_session(),show_banner=False)

src/tool/mcp_servers/python_server.py

@@ -9,6 +9,8 @@
 from fastmcp import FastMCP
 
 # Initialize FastMCP server
+from src.logging.logger import setup_mcp_logging
+setup_mcp_logging(tool_name=os.path.basename(__file__))
 mcp = FastMCP("e2b-python-interpreter")
 
 # API keys
@@ -411,4 +413,4 @@ async def download_file_from_sandbox_to_local(
 
 
 if __name__ == "__main__":
-    mcp.run(transport="stdio")
+    mcp.run(transport="stdio",show_banner=False)

src/tool/mcp_servers/reading_mcp_server.py

@@ -15,6 +15,8 @@
 from .utils.smart_request import smart_request
 
 # Initialize FastMCP server
+from src.logging.logger import setup_mcp_logging
+setup_mcp_logging(tool_name=os.path.basename(__file__))
 mcp = FastMCP("reading-mcp-server")
 SERPER_API_KEY = os.environ.get("SERPER_API_KEY", "")
 JINA_API_KEY = os.environ.get("JINA_API_KEY", "")
@@ -153,7 +155,7 @@ def _cleanup_tempfile(path):
 
     # Run the server with the specified transport method
     if args.transport == "stdio":
-        mcp.run(transport="stdio")
+        mcp.run(transport="stdio",show_banner=False)
     else:
         # For HTTP transport, include port and path options
-        mcp.run(transport="streamable-http", port=args.port, path=args.path)
+        mcp.run(transport="streamable-http", port=args.port, path=args.path,show_banner=False)

src/tool/mcp_servers/reasoning_mcp_server.py

@@ -19,6 +19,8 @@
 OPENAI_MODEL_NAME = os.environ.get("OPENAI_MODEL_NAME", "o3")
 
 # Initialize FastMCP server
+from src.logging.logger import setup_mcp_logging
+setup_mcp_logging(tool_name=os.path.basename(__file__))
 mcp = FastMCP("reasoning-mcp-server")
 
 
@@ -124,4 +126,4 @@ async def reasoning(question: str) -> str:
 
 
 if __name__ == "__main__":
-    mcp.run(transport="stdio")
+    mcp.run(transport="stdio",show_banner=False)

src/tool/mcp_servers/searching_mcp_server.py

@@ -31,6 +31,8 @@
 )
 
 # Initialize FastMCP server
+from src.logging.logger import setup_mcp_logging
+setup_mcp_logging(tool_name=os.path.basename(__file__))
 mcp = FastMCP("searching-mcp-server")
 
 
@@ -666,4 +668,4 @@ async def scrape_website(url: str) -> str:
 
 
 if __name__ == "__main__":
-    mcp.run(transport="stdio")
+    mcp.run(transport="stdio",show_banner=False)