/ mlflow / autogen / chat.py
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)