mirror of
https://gitee.com/wanwujie/deer-flow
synced 2026-04-03 14:22:13 +08:00
3b4e993531
2. use tavily from`langchain-tavily` rather than the older one from `langchain-community` Co-authored-by: Willem Jiang <willem.jiang@gmail.com>
82 lines
2.4 KiB
Python
82 lines
2.4 KiB
Python
# Copyright (c) 2025 Bytedance Ltd. and/or its affiliates
|
|
# SPDX-License-Identifier: MIT
|
|
|
|
import functools
|
|
import logging
|
|
from typing import Any, Callable, Type, TypeVar
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
T = TypeVar("T")
|
|
|
|
|
|
def log_io(func: Callable) -> Callable:
|
|
"""
|
|
A decorator that logs the input parameters and output of a tool function.
|
|
|
|
Args:
|
|
func: The tool function to be decorated
|
|
|
|
Returns:
|
|
The wrapped function with input/output logging
|
|
"""
|
|
|
|
@functools.wraps(func)
|
|
def wrapper(*args: Any, **kwargs: Any) -> Any:
|
|
# Log input parameters
|
|
func_name = func.__name__
|
|
params = ", ".join(
|
|
[*(str(arg) for arg in args), *(f"{k}={v}" for k, v in kwargs.items())]
|
|
)
|
|
logger.info(f"Tool {func_name} called with parameters: {params}")
|
|
|
|
# Execute the function
|
|
result = func(*args, **kwargs)
|
|
|
|
# Log the output
|
|
logger.info(f"Tool {func_name} returned: {result}")
|
|
|
|
return result
|
|
|
|
return wrapper
|
|
|
|
|
|
class LoggedToolMixin:
|
|
"""A mixin class that adds logging functionality to any tool."""
|
|
|
|
def _log_operation(self, method_name: str, *args: Any, **kwargs: Any) -> None:
|
|
"""Helper method to log tool operations."""
|
|
tool_name = self.__class__.__name__.replace("Logged", "")
|
|
params = ", ".join(
|
|
[*(str(arg) for arg in args), *(f"{k}={v}" for k, v in kwargs.items())]
|
|
)
|
|
logger.debug(f"Tool {tool_name}.{method_name} called with parameters: {params}")
|
|
|
|
def _run(self, *args: Any, **kwargs: Any) -> Any:
|
|
"""Override _run method to add logging."""
|
|
self._log_operation("_run", *args, **kwargs)
|
|
result = super()._run(*args, **kwargs)
|
|
logger.debug(
|
|
f"Tool {self.__class__.__name__.replace('Logged', '')} returned: {result}"
|
|
)
|
|
return result
|
|
|
|
|
|
def create_logged_tool(base_tool_class: Type[T]) -> Type[T]:
|
|
"""
|
|
Factory function to create a logged version of any tool class.
|
|
|
|
Args:
|
|
base_tool_class: The original tool class to be enhanced with logging
|
|
|
|
Returns:
|
|
A new class that inherits from both LoggedToolMixin and the base tool class
|
|
"""
|
|
|
|
class LoggedTool(LoggedToolMixin, base_tool_class):
|
|
pass
|
|
|
|
# Set a more descriptive name for the class
|
|
LoggedTool.__name__ = f"Logged{base_tool_class.__name__}"
|
|
return LoggedTool
|