chat.py
1 import logging 2 from typing import TYPE_CHECKING, Union 3 4 from opentelemetry.sdk.trace import Span 5 6 from mlflow.tracing.utils import set_span_chat_tools 7 from mlflow.types.chat import ChatTool 8 9 if TYPE_CHECKING: 10 from autogen_core.tools import BaseTool, ToolSchema 11 12 _logger = logging.getLogger(__name__) 13 14 15 def log_tools(span: Span, tools: list[Union["BaseTool", "ToolSchema"]]): 16 """ 17 Log Autogen tool definitions into the passed in span. 18 19 Ref: https://microsoft.github.io/autogen/stable/user-guide/core-user-guide/components/tools.html 20 21 Args: 22 span: The span to log the tools into. 23 tools: A list of Autogen BaseTool. 24 """ 25 from autogen_core.tools import BaseTool 26 27 try: 28 tools = [ 29 ChatTool( 30 type="function", 31 function=tool.schema if isinstance(tool, BaseTool) else tool, 32 ) 33 for tool in tools 34 ] 35 set_span_chat_tools(span, tools) 36 except Exception: 37 _logger.debug(f"Failed to log tools to Span {span}.", exc_info=True)