wehub-resource-sync
1283 lines
46 KiB
Python
1283 lines
46 KiB
Python
import asyncio
|
|
import functools
|
|
import inspect
|
|
import json
|
|
import logging
|
|
import math
|
|
import threading
|
|
from collections import OrderedDict
|
|
from typing import TYPE_CHECKING, Any, Callable
|
|
|
|
from cachetools.func import cached
|
|
from opentelemetry.trace import NoOpTracer
|
|
from pydantic import BaseModel, Field
|
|
|
|
import mlflow
|
|
from mlflow.entities.assessment_source import AssessmentSourceType
|
|
from mlflow.entities.span import Span, SpanType
|
|
from mlflow.entities.trace import Trace
|
|
from mlflow.environment_variables import (
|
|
MLFLOW_GENAI_EVAL_ASYNC_TIMEOUT,
|
|
MLFLOW_GENAI_EVAL_ENABLE_SCORER_TRACING,
|
|
MLFLOW_GENAI_EVAL_SKIP_TRACE_VALIDATION,
|
|
)
|
|
from mlflow.exceptions import MlflowException
|
|
from mlflow.genai.discovery.constants import DEFAULT_TOP_N_SLOWEST_SPANS
|
|
from mlflow.genai.judges.utils import get_chat_completions_with_structured_output, get_default_model
|
|
from mlflow.genai.utils.data_validation import check_model_prediction
|
|
from mlflow.genai.utils.prompts.available_tools_extraction import (
|
|
get_available_tools_extraction_prompts,
|
|
)
|
|
from mlflow.models.evaluation.utils.trace import configure_autologging_for_evaluation
|
|
from mlflow.protos.databricks_pb2 import INVALID_PARAMETER_VALUE
|
|
from mlflow.tracing.constant import (
|
|
AssessmentMetadataKey,
|
|
SpanAttributeKey,
|
|
TraceMetadataKey,
|
|
TraceTagKey,
|
|
)
|
|
from mlflow.tracing.display import IPythonTraceDisplayHandler
|
|
from mlflow.tracing.utils import TraceJSONEncoder
|
|
from mlflow.tracing.utils.search import traces_to_df
|
|
from mlflow.tracking.client import MlflowClient
|
|
from mlflow.utils.uri import is_databricks_uri
|
|
|
|
if TYPE_CHECKING:
|
|
import pandas as pd
|
|
|
|
from mlflow.genai.evaluation.entities import EvalItem, EvalResult
|
|
from mlflow.genai.utils.type import FunctionCall
|
|
from mlflow.types.chat import ChatTool
|
|
|
|
_logger = logging.getLogger(__name__)
|
|
|
|
_MESSAGE_KEY = "message"
|
|
_MESSAGES_KEY = "messages"
|
|
_CHOICES_KEY = "choices"
|
|
_CONTENT_KEY = "content"
|
|
_OUTPUT_KEY = "output"
|
|
|
|
|
|
def extract_request_from_trace(trace: Trace) -> str | None:
|
|
"""
|
|
Extract request text from an MLflow trace object.
|
|
|
|
Args:
|
|
trace: MLflow trace object
|
|
|
|
Returns:
|
|
Extracted request text as string, or None if no root span
|
|
"""
|
|
root_span = trace.data._get_root_span()
|
|
if root_span is None:
|
|
return None
|
|
return parse_inputs_to_str(root_span.inputs)
|
|
|
|
|
|
def extract_response_from_trace(trace: Trace) -> str | None:
|
|
"""
|
|
Extract response text from an MLflow trace object.
|
|
|
|
Args:
|
|
trace: MLflow trace object
|
|
|
|
Returns:
|
|
Extracted response text as string, or None if no root span
|
|
"""
|
|
root_span = trace.data._get_root_span()
|
|
if root_span is None:
|
|
return None
|
|
return parse_outputs_to_str(root_span.outputs)
|
|
|
|
|
|
def extract_inputs_from_trace(trace: Trace) -> Any:
|
|
"""
|
|
Extract inputs from the root span of an MLflow trace.
|
|
|
|
Args:
|
|
trace: MLflow trace object
|
|
|
|
Returns:
|
|
Inputs from the root span, or None if no root span or inputs
|
|
"""
|
|
root_span = trace.data._get_root_span()
|
|
if root_span and root_span.inputs is not None:
|
|
return root_span.inputs
|
|
return None
|
|
|
|
|
|
def extract_outputs_from_trace(trace: Trace) -> Any:
|
|
"""
|
|
Extract outputs from the root span of an MLflow trace.
|
|
|
|
Args:
|
|
trace: MLflow trace object
|
|
|
|
Returns:
|
|
Outputs from the root span, or None if no root span or outputs
|
|
"""
|
|
root_span = trace.data._get_root_span()
|
|
if root_span and root_span.outputs is not None:
|
|
return root_span.outputs
|
|
return None
|
|
|
|
|
|
def resolve_inputs_from_trace(
|
|
inputs: Any | None, trace: Trace, *, extract_if_none: bool = True
|
|
) -> Any | None:
|
|
"""
|
|
Extract inputs from trace if not provided.
|
|
|
|
Args:
|
|
inputs: Input data to evaluate. If None, will be extracted from trace.
|
|
trace: MLflow trace object containing the execution to evaluate.
|
|
extract_if_none: If True, extract from trace when inputs is None. If False, only
|
|
return the provided inputs value. Defaults to True.
|
|
|
|
Returns:
|
|
The provided inputs if not None, otherwise extracted inputs from trace,
|
|
or None if extraction fails.
|
|
"""
|
|
if inputs is None and trace is not None and extract_if_none:
|
|
try:
|
|
return extract_inputs_from_trace(trace)
|
|
except Exception as e:
|
|
_logger.debug(f"Could not extract inputs from trace: {e}")
|
|
return inputs
|
|
|
|
|
|
def resolve_outputs_from_trace(
|
|
outputs: Any | None, trace: Trace, *, extract_if_none: bool = True
|
|
) -> Any | None:
|
|
"""
|
|
Extract outputs from trace if not provided.
|
|
|
|
Args:
|
|
outputs: Output data to evaluate. If None, will be extracted from trace.
|
|
trace: MLflow trace object containing the execution to evaluate.
|
|
extract_if_none: If True, extract from trace when outputs is None. If False, only
|
|
return the provided outputs value. Defaults to True.
|
|
|
|
Returns:
|
|
The provided outputs if not None, otherwise extracted outputs from trace,
|
|
or None if extraction fails.
|
|
"""
|
|
if outputs is None and trace is not None and extract_if_none:
|
|
try:
|
|
return extract_outputs_from_trace(trace)
|
|
except Exception as e:
|
|
_logger.debug(f"Could not extract outputs from trace: {e}")
|
|
return outputs
|
|
|
|
|
|
def _get_exception_from_span(span: Span) -> str | None:
|
|
"""
|
|
Extract exception information from span events.
|
|
|
|
Args:
|
|
span: The span to check for exception events.
|
|
|
|
Returns:
|
|
A formatted string containing exception information if found, None otherwise.
|
|
"""
|
|
exception_events = [event for event in span.events if event.name == "exception"]
|
|
if not exception_events:
|
|
return None
|
|
|
|
exception_event = exception_events[0]
|
|
attrs = exception_event.attributes
|
|
|
|
exception_type = attrs.get("exception.type", "Exception")
|
|
|
|
if exception_message := attrs.get("exception.message"):
|
|
return f"{exception_type}: {exception_message}"
|
|
return exception_type
|
|
|
|
|
|
def _extract_tool_name_from_span(span: Span) -> str:
|
|
inputs = span.attributes.get(SpanAttributeKey.INPUTS)
|
|
if isinstance(inputs, dict):
|
|
call_data = inputs.get("call")
|
|
if isinstance(call_data, dict) and "tool_name" in call_data:
|
|
return call_data["tool_name"]
|
|
return span.name
|
|
|
|
|
|
def extract_tools_called_from_trace(trace: Trace) -> list["FunctionCall"]:
|
|
"""
|
|
Extract tool call information from TOOL type spans in a trace.
|
|
|
|
This function extracts tool spans (spans with span_type==SpanType.TOOL) from a trace
|
|
and returns them as a list of FunctionCall objects containing the tool name, inputs,
|
|
and outputs.
|
|
|
|
Args:
|
|
trace: A single Trace object to extract tool calls from.
|
|
|
|
Returns:
|
|
List of FunctionCall objects.
|
|
Returns empty list if no tool spans are found.
|
|
|
|
Example:
|
|
>>> trace = mlflow.get_trace(trace_id)
|
|
>>> tools = extract_tools_called_from_trace(trace)
|
|
>>> # Returns: [FunctionCall(name="tool_name", arguments={...}, outputs={...})]
|
|
"""
|
|
from mlflow.genai.utils.type import FunctionCall
|
|
|
|
tools_called = []
|
|
tool_spans = trace.search_spans(span_type=SpanType.TOOL)
|
|
|
|
for tool_span in sorted(tool_spans, key=lambda s: s.start_time_ns or 0):
|
|
tool_info = FunctionCall(
|
|
name=_extract_tool_name_from_span(tool_span),
|
|
arguments=tool_span.inputs or None,
|
|
outputs=tool_span.outputs or None,
|
|
exception=_get_exception_from_span(tool_span),
|
|
)
|
|
tools_called.append(tool_info)
|
|
|
|
return tools_called
|
|
|
|
|
|
def parse_tool_call_messages_from_trace(trace: Trace) -> list[dict[str, str]]:
|
|
"""
|
|
Extract and format tool call information from TOOL type spans in a trace.
|
|
|
|
This function extracts tool spans (spans with span_type==SpanType.TOOL) from a trace
|
|
and formats them as conversation messages with role='tool'. Each tool message includes
|
|
the tool name, inputs, and outputs.
|
|
|
|
Args:
|
|
trace: A single Trace object to extract tool calls from.
|
|
|
|
Returns:
|
|
List of tool call messages in the format [{"role": "tool", "content": str}].
|
|
Tool content includes the tool name, inputs, and outputs formatted as a string.
|
|
Returns empty list if no tool spans are found.
|
|
|
|
Example:
|
|
>>> trace = mlflow.get_trace(trace_id)
|
|
>>> tool_messages = parse_tool_call_messages_from_trace(trace)
|
|
>>> # Returns: [{"role": "tool", "content": "Tool: name\\nInputs: ...\\nOutputs: ..."}]
|
|
"""
|
|
tools_called = extract_tools_called_from_trace(trace)
|
|
|
|
tool_messages = []
|
|
for tool in tools_called:
|
|
tool_info = f"Tool: {tool.name}"
|
|
if tool.arguments is not None:
|
|
tool_info += f"\nInputs: {tool.arguments}"
|
|
if tool.outputs is not None:
|
|
tool_info += f"\nOutputs: {tool.outputs}"
|
|
if tool.exception is not None:
|
|
tool_info += f"\nException: {tool.exception}"
|
|
tool_messages.append({"role": "tool", "content": tool_info})
|
|
|
|
return tool_messages
|
|
|
|
|
|
def validate_session(session: list[Trace]) -> None:
|
|
"""
|
|
Validate that all traces in session belong to the same session.
|
|
|
|
Args:
|
|
session: List of traces to validate.
|
|
|
|
Raises:
|
|
MlflowException: If traces are missing session_id or belong to different sessions.
|
|
"""
|
|
session_id_to_trace_ids: dict[str, list[str]] = {}
|
|
for trace in session:
|
|
session_id = trace.info.trace_metadata.get(TraceMetadataKey.TRACE_SESSION)
|
|
if session_id is None:
|
|
raise MlflowException(
|
|
f"All traces in 'session' must have a session_id. "
|
|
f"Trace {trace.info.trace_id} is missing session_id. "
|
|
f"See https://mlflow.org/docs/latest/genai/tracing/track-users-sessions/ "
|
|
f"for information on how to set session_id on traces.",
|
|
error_code=INVALID_PARAMETER_VALUE,
|
|
)
|
|
if session_id not in session_id_to_trace_ids:
|
|
session_id_to_trace_ids[session_id] = []
|
|
session_id_to_trace_ids[session_id].append(trace.info.trace_id)
|
|
|
|
if len(session_id_to_trace_ids) != 1:
|
|
session_details = "\n".join(
|
|
f"session_id '{sid}': trace_ids {trace_ids[:3]}"
|
|
+ (
|
|
f" and {len(trace_ids) - 3} more trace{'s' if len(trace_ids) - 3 != 1 else ''}"
|
|
if len(trace_ids) > 3
|
|
else ""
|
|
)
|
|
for sid, trace_ids in session_id_to_trace_ids.items()
|
|
)
|
|
raise MlflowException.invalid_parameter_value(
|
|
f"All traces in 'session' must belong to the same session. "
|
|
f"Found {len(session_id_to_trace_ids)} different session(s):\n{session_details}"
|
|
)
|
|
|
|
|
|
def _extract_trace_timing_info(
|
|
trace: Trace, *, top_n_slowest_spans: int = DEFAULT_TOP_N_SLOWEST_SPANS
|
|
) -> dict[str, Any] | None:
|
|
"""
|
|
Extract timing information from a trace for display in evaluations.
|
|
|
|
Args:
|
|
trace: The trace to extract timing from.
|
|
top_n_slowest_spans: Number of slowest spans to include in the output.
|
|
|
|
Returns:
|
|
Dict containing 'duration_s' (float) and 'slowest_spans_formatted' (str | None),
|
|
or None if the trace has no execution duration.
|
|
"""
|
|
if trace.info.execution_duration is None:
|
|
return None
|
|
|
|
duration_s = trace.info.execution_duration / 1000
|
|
slowest_spans_formatted = None
|
|
|
|
# Extract top N slowest spans for context on bottlenecks
|
|
if trace.data.spans:
|
|
# Filter out spans that do not have an end time to avoid None arithmetic
|
|
if completed_spans := [span for span in trace.data.spans if span.end_time_ns is not None]:
|
|
if sorted_spans := sorted(
|
|
completed_spans, key=lambda s: s.end_time_ns - s.start_time_ns, reverse=True
|
|
)[:top_n_slowest_spans]:
|
|
slow_spans = [
|
|
f"{span.name} ({(span.end_time_ns - span.start_time_ns) / 1_000_000_000:.2f}s)"
|
|
for span in sorted_spans
|
|
]
|
|
slowest_spans_formatted = ", ".join(slow_spans)
|
|
|
|
return {
|
|
"duration_s": duration_s,
|
|
"slowest_spans_formatted": slowest_spans_formatted,
|
|
}
|
|
|
|
|
|
def resolve_conversation_from_session(
|
|
session: list[Trace],
|
|
*,
|
|
include_tool_calls: bool = False,
|
|
include_timing: bool = False,
|
|
) -> list[dict[str, str]]:
|
|
"""
|
|
Extract conversation history from traces in session.
|
|
|
|
Args:
|
|
session: List of traces from the same session.
|
|
include_tool_calls: If True, include tool call information from TOOL type spans
|
|
in the conversation. Default is False for backward compatibility.
|
|
include_timing: If True, append timing information to assistant responses.
|
|
This includes total duration and slowest spans for latency analysis.
|
|
|
|
Returns:
|
|
List of conversation messages in the format:
|
|
[{"role": "user"|"assistant"|"tool", "content": str}].
|
|
Each trace contributes user input and assistant output messages.
|
|
If include_tool_calls is True, tool call messages (with inputs/outputs)
|
|
are also included in chronological order.
|
|
If include_timing is True, assistant messages include performance metadata.
|
|
"""
|
|
# Sort traces by creation time (timestamp_ms)
|
|
sorted_traces = sorted(session, key=lambda t: t.info.timestamp_ms)
|
|
|
|
conversation = []
|
|
for trace in sorted_traces:
|
|
# Extract and parse input (user message)
|
|
if inputs := extract_inputs_from_trace(trace):
|
|
user_content = parse_inputs_to_str(inputs)
|
|
if user_content and user_content.strip():
|
|
conversation.append({"role": "user", "content": user_content})
|
|
|
|
# Extract tool calls from TOOL type spans (if requested)
|
|
if include_tool_calls:
|
|
tool_messages = parse_tool_call_messages_from_trace(trace)
|
|
conversation.extend(tool_messages)
|
|
|
|
# Extract and parse output (assistant message)
|
|
if outputs := extract_outputs_from_trace(trace):
|
|
assistant_content = parse_outputs_to_str(outputs)
|
|
if assistant_content and assistant_content.strip():
|
|
if include_timing:
|
|
if timing_info := _extract_trace_timing_info(trace):
|
|
timing_parts = [f"\n[Response duration: {timing_info['duration_s']:.2f}s"]
|
|
if slowest_spans_formatted := timing_info["slowest_spans_formatted"]:
|
|
timing_parts.append(f", slowest spans: {slowest_spans_formatted}")
|
|
timing_parts.append("]")
|
|
assistant_content += "".join(timing_parts)
|
|
|
|
conversation.append({"role": "assistant", "content": assistant_content})
|
|
|
|
return conversation
|
|
|
|
|
|
def resolve_expectations_from_trace(
|
|
expectations: dict[str, Any] | None,
|
|
trace: Trace,
|
|
source_type: AssessmentSourceType = AssessmentSourceType.HUMAN,
|
|
*,
|
|
extract_if_none: bool = True,
|
|
) -> dict[str, Any] | None:
|
|
"""
|
|
Extract expectations from trace if not provided.
|
|
|
|
Args:
|
|
expectations: Dictionary of expected outcomes. If None, will be extracted from trace.
|
|
trace: MLflow trace object containing the execution to evaluate.
|
|
source_type: Assessment source type to filter expectations by. Defaults to HUMAN.
|
|
extract_if_none: If True, extract from trace when expectations is None. If False, only
|
|
return the provided expectations value. Defaults to True.
|
|
|
|
Returns:
|
|
The provided expectations if not None, otherwise extracted expectations from trace,
|
|
or None if extraction fails.
|
|
"""
|
|
if expectations is None and trace is not None and extract_if_none:
|
|
try:
|
|
return extract_expectations_from_trace(trace, source_type=source_type)
|
|
except Exception as e:
|
|
_logger.debug(f"Could not extract expectations from trace: {e}")
|
|
return expectations
|
|
|
|
|
|
def resolve_expectations_from_session(
|
|
expectations: dict[str, Any] | None,
|
|
session: list[Trace],
|
|
source_type: AssessmentSourceType = AssessmentSourceType.HUMAN,
|
|
) -> dict[str, Any] | None:
|
|
"""
|
|
Extract session-level expectations from the first trace in a session if not provided.
|
|
|
|
Args:
|
|
expectations: Dictionary of expected outcomes. If provided, this is returned as-is
|
|
(ground truth). If None, will be extracted from session.
|
|
session: List of traces from the same session.
|
|
source_type: Assessment source type to filter expectations by. Defaults to HUMAN.
|
|
|
|
Returns:
|
|
The provided expectations if not None (ground truth), otherwise extracted
|
|
session-level expectations from the first trace, or None if extraction fails.
|
|
"""
|
|
if expectations is None and session:
|
|
try:
|
|
sorted_traces = sorted(session, key=lambda t: t.info.timestamp_ms)
|
|
first_trace = sorted_traces[0]
|
|
|
|
expectation_assessments = first_trace.search_assessments(type="expectation")
|
|
|
|
expectation_assessments = [
|
|
exp
|
|
for exp in expectation_assessments
|
|
if exp.source
|
|
and exp.source.source_type == source_type
|
|
and exp.metadata
|
|
and TraceMetadataKey.TRACE_SESSION in exp.metadata
|
|
]
|
|
|
|
return {exp.name: exp.expectation.value for exp in expectation_assessments} or None
|
|
except Exception as e:
|
|
_logger.debug(f"Could not extract expectations from session: {e}")
|
|
return expectations
|
|
|
|
|
|
def extract_expectations_from_trace(
|
|
trace: Trace,
|
|
source_type: str | None = None,
|
|
) -> dict[str, Any] | None:
|
|
"""
|
|
Extract expectations from trace assessments.
|
|
|
|
Args:
|
|
trace: MLflow trace object
|
|
source_type: If specified, only extract expectations from the given source type.
|
|
Must be one of the valid AssessmentSourceType values
|
|
If None, extract all expectations regardless of source.
|
|
|
|
Returns:
|
|
Dictionary of expectations, or None if no expectations found
|
|
"""
|
|
validated_source = (
|
|
AssessmentSourceType._standardize(source_type) if source_type is not None else None
|
|
)
|
|
|
|
expectation_assessments = trace.search_assessments(type="expectation")
|
|
|
|
if validated_source is not None:
|
|
expectation_assessments = [
|
|
exp
|
|
for exp in expectation_assessments
|
|
if exp.source and exp.source.source_type == validated_source
|
|
]
|
|
|
|
if not expectation_assessments:
|
|
return None
|
|
|
|
return {exp.name: exp.expectation.value for exp in expectation_assessments}
|
|
|
|
|
|
def _wrap_async_predict_fn(async_fn: Callable[..., Any]) -> Callable[..., Any]:
|
|
"""
|
|
Wrap an async function to make it synchronous using asyncio.run with timeout.
|
|
|
|
Args:
|
|
async_fn: The async function to wrap
|
|
|
|
Returns:
|
|
A synchronous wrapper function that calls the async function with timeout
|
|
"""
|
|
timeout = MLFLOW_GENAI_EVAL_ASYNC_TIMEOUT.get()
|
|
|
|
@functools.wraps(async_fn)
|
|
def sync_wrapper(*args, **kwargs):
|
|
try:
|
|
loop = asyncio.get_running_loop()
|
|
except RuntimeError:
|
|
loop = None
|
|
|
|
if loop is not None:
|
|
try:
|
|
import nest_asyncio
|
|
|
|
nest_asyncio.apply()
|
|
except ImportError:
|
|
raise MlflowException(
|
|
"Detected a running event loop (e.g., in Jupyter notebook). "
|
|
"To use async predict functions in notebook environments, "
|
|
"install nest-asyncio: pip install nest-asyncio"
|
|
)
|
|
|
|
return asyncio.run(asyncio.wait_for(async_fn(*args, **kwargs), timeout=timeout))
|
|
|
|
return sync_wrapper
|
|
|
|
|
|
def convert_predict_fn(predict_fn: Callable[..., Any], sample_input: Any) -> Callable[..., Any]:
|
|
"""
|
|
Check the predict_fn is callable and add trace decorator if it is not already traced.
|
|
If the predict_fn is an async function, wrap it to make it synchronous.
|
|
"""
|
|
# Detect if predict_fn is an async function and wrap it
|
|
if inspect.iscoroutinefunction(predict_fn):
|
|
_logger.debug(
|
|
f"Detected async predict_fn. Wrapping with asyncio.run() with timeout of "
|
|
f"{MLFLOW_GENAI_EVAL_ASYNC_TIMEOUT.get()} seconds."
|
|
)
|
|
predict_fn = _wrap_async_predict_fn(predict_fn)
|
|
if not MLFLOW_GENAI_EVAL_SKIP_TRACE_VALIDATION.get() and sample_input:
|
|
with (
|
|
NoOpTracerPatcher() as counter,
|
|
# Enable auto-tracing before checking if the predict_fn produces traces, so that
|
|
# functions using auto-traceable libraries (OpenAI, LangChain, etc.) are correctly
|
|
# identified as traced functions
|
|
configure_autologging_for_evaluation(enable_tracing=True),
|
|
):
|
|
check_model_prediction(predict_fn, sample_input)
|
|
|
|
if counter.count == 0:
|
|
predict_fn = mlflow.trace(predict_fn)
|
|
|
|
# Wrap the prediction function to unwrap the inputs dictionary into keyword arguments.
|
|
return lambda request: predict_fn(**request)
|
|
|
|
|
|
class NoOpTracerPatcher:
|
|
"""
|
|
A context manager to count the number of times NoOpTracer's start_span is called.
|
|
|
|
The check is done in the following steps so it doesn't have any side effects:
|
|
1. Disable tracing.
|
|
2. Patch the NoOpTracer.start_span method to count the number of times it is called.
|
|
NoOpTracer is used when tracing is disabled.
|
|
3. Call the predict function with the sample input.
|
|
4. Restore the original NoOpTracer.start_span method and re-enable tracing.
|
|
|
|
|
|
WARNING: This function is not thread-safe. We do not provide support for running
|
|
`mlflow.genai.evaluate` in multi-threaded environments.`
|
|
"""
|
|
|
|
def __init__(self):
|
|
self.count = 0
|
|
|
|
def __enter__(self):
|
|
self.original = NoOpTracer.start_span
|
|
|
|
def _patched_start_span(_self, *args, **kwargs):
|
|
self.count += 1
|
|
return self.original(_self, *args, **kwargs)
|
|
|
|
NoOpTracer.start_span = _patched_start_span
|
|
return self
|
|
|
|
def __exit__(self, exc_type, exc_value, traceback):
|
|
NoOpTracer.start_span = self.original
|
|
|
|
|
|
def is_none_or_nan(value: Any) -> bool:
|
|
"""
|
|
Checks whether a value is None or NaN.
|
|
|
|
NB: This function does not handle pandas.NA.
|
|
"""
|
|
# isinstance(value, float) check is needed to ensure that math.isnan is not called on an array.
|
|
return value is None or (isinstance(value, float) and math.isnan(value))
|
|
|
|
|
|
def _is_empty(value: Any) -> bool:
|
|
"""
|
|
Check if a value is empty (None, empty dict, empty list, empty string, etc.).
|
|
"""
|
|
if value is None:
|
|
return True
|
|
if isinstance(value, (dict, list, str)):
|
|
return len(value) == 0
|
|
return False
|
|
|
|
|
|
def parse_inputs_to_str(value: Any) -> str:
|
|
"""Parse the inputs to a string compatible with the judges API"""
|
|
if is_none_or_nan(value):
|
|
# The DBX managed backend doesn't allow empty inputs. This is
|
|
# a temporary workaround to bypass the validation.
|
|
return " "
|
|
if isinstance(value, str):
|
|
return value
|
|
|
|
value = _to_dict(value)
|
|
|
|
# Handle case where _to_dict returns a non-dict (e.g., a list that gets serialized
|
|
# and remains a list)
|
|
if not isinstance(value, dict):
|
|
return json.dumps(value, cls=TraceJSONEncoder)
|
|
|
|
if (messages := value.get(_MESSAGES_KEY)) and len(messages) > 0:
|
|
contents = [m.get(_CONTENT_KEY) for m in messages]
|
|
if len(contents) > 1 and all(isinstance(c, str) for c in contents):
|
|
return json.dumps(messages)
|
|
elif isinstance(contents[-1], str):
|
|
return contents[-1]
|
|
return str(value)
|
|
|
|
|
|
def parse_outputs_to_str(value: Any) -> str:
|
|
"""Parse the outputs to a string compatible with the judges API"""
|
|
if is_none_or_nan(value):
|
|
return " "
|
|
if isinstance(value, str):
|
|
return value
|
|
|
|
# PyFuncModel.predict wraps the output in a list
|
|
if isinstance(value, list) and len(value) > 0:
|
|
return parse_outputs_to_str(value[0])
|
|
|
|
value = _to_dict(value)
|
|
if _is_chat_choices(value.get(_CHOICES_KEY)):
|
|
content = value[_CHOICES_KEY][0][_MESSAGE_KEY][_CONTENT_KEY]
|
|
elif _is_chat_messages(value.get(_MESSAGES_KEY)):
|
|
content = value[_MESSAGES_KEY][-1][_CONTENT_KEY]
|
|
elif _is_responses_api_output(value.get(_OUTPUT_KEY)):
|
|
content = _extract_responses_api_content(value[_OUTPUT_KEY])
|
|
else:
|
|
content = json.dumps(value, cls=TraceJSONEncoder)
|
|
return content
|
|
|
|
|
|
def _is_responses_api_output(maybe_output: Any) -> bool:
|
|
"""Check if the value is an OpenAI Responses API output format."""
|
|
if not maybe_output or not isinstance(maybe_output, list) or len(maybe_output) == 0:
|
|
return False
|
|
last_item = maybe_output[-1]
|
|
return (
|
|
isinstance(last_item, dict)
|
|
and last_item.get("type") == "message"
|
|
and "content" in last_item
|
|
)
|
|
|
|
|
|
def _extract_responses_api_content(output: list[dict[str, Any]]) -> str:
|
|
"""Extract text content from OpenAI Responses API output format."""
|
|
for item in reversed(output):
|
|
if item.get("role") == "assistant" and "content" in item:
|
|
content = item["content"]
|
|
if isinstance(content, str):
|
|
return content
|
|
if isinstance(content, list):
|
|
for part in content:
|
|
if isinstance(part, dict) and part.get("type") in ("text", "output_text"):
|
|
return part.get("text", json.dumps(output))
|
|
return json.dumps(output)
|
|
|
|
|
|
def _is_chat_choices(maybe_choices: Any) -> bool:
|
|
if (
|
|
not maybe_choices
|
|
or not isinstance(maybe_choices, list)
|
|
or not isinstance(maybe_choices[0], dict)
|
|
):
|
|
return False
|
|
|
|
message = maybe_choices[0].get(_MESSAGE_KEY)
|
|
return _is_chat_messages([message])
|
|
|
|
|
|
def _is_chat_messages(maybe_messages: Any) -> bool:
|
|
return (
|
|
maybe_messages
|
|
and len(maybe_messages) > 0
|
|
and isinstance(maybe_messages[-1], dict)
|
|
and isinstance(maybe_messages[-1].get(_CONTENT_KEY), str)
|
|
)
|
|
|
|
|
|
def _to_dict(obj: Any) -> dict[str, Any]:
|
|
if hasattr(obj, "to_dict"):
|
|
return obj.to_dict()
|
|
|
|
if isinstance(obj, BaseModel):
|
|
return obj.model_dump()
|
|
|
|
# Convert to JSON string and then back to dictionary to handle nested objects
|
|
json_str = json.dumps(obj, cls=TraceJSONEncoder)
|
|
return json.loads(json_str)
|
|
|
|
|
|
def extract_retrieval_context_from_trace(trace: Trace | None) -> dict[str, list[Any]]:
|
|
"""
|
|
Extract the retrieval context from the trace.
|
|
Extracts all top-level retrieval spans from the trace if there are multiple retrieval spans.
|
|
If the trace does not have a retrieval span, return an empty dictionary.
|
|
⚠️ Warning: Please make sure to not throw exception. If fails, return an empty dictionary.
|
|
"""
|
|
if trace is None or trace.data is None:
|
|
return {}
|
|
|
|
top_level_retrieval_spans = _get_top_level_retrieval_spans(trace)
|
|
if len(top_level_retrieval_spans) == 0:
|
|
return {}
|
|
|
|
retrieved = {}
|
|
|
|
for retrieval_span in top_level_retrieval_spans:
|
|
try:
|
|
outputs = retrieval_span.outputs
|
|
outputs = json.loads(outputs) if isinstance(outputs, str) else outputs
|
|
contexts = [_parse_chunk(chunk) for chunk in outputs or []]
|
|
retrieved[retrieval_span.span_id] = [c for c in contexts if c is not None]
|
|
except Exception as e:
|
|
_logger.debug(
|
|
f"Fail to get retrieval context from span: {retrieval_span}. Error: {e!r}"
|
|
)
|
|
|
|
return retrieved
|
|
|
|
|
|
def _get_top_level_retrieval_spans(trace: Trace) -> list[Span]:
|
|
"""
|
|
Get the top-level retrieval spans in the trace.
|
|
Top-level retrieval spans are retrieval spans that are not children of other retrieval spans.
|
|
For example, given the following spans:
|
|
- Span A (Chain)
|
|
- Span B (Retriever)
|
|
- Span C (Retriever)
|
|
- Span D (Retriever)
|
|
- Span E (LLM)
|
|
- Span F (Retriever)
|
|
Span B and Span D are top-level retrieval spans.
|
|
Span C and Span F are NOT top-level because they are children of other retrieval spans.
|
|
"""
|
|
top_level_retrieval_spans = []
|
|
# Cache span_id -> span mapping for fast lookup
|
|
all_spans = {span.span_id: span for span in trace.data.spans}
|
|
for span in trace.search_spans(span_type=SpanType.RETRIEVER):
|
|
# Check if this span is a child of another retrieval span
|
|
parent_id = span.parent_id
|
|
while parent_id:
|
|
parent_span = all_spans.get(parent_id)
|
|
if not parent_span:
|
|
# Malformed trace
|
|
_logger.debug(
|
|
f"Malformed trace: span {span} has parent span ID {parent_id}, "
|
|
"but the parent span is not found in the trace."
|
|
)
|
|
break
|
|
|
|
if parent_span.span_type == SpanType.RETRIEVER:
|
|
# This span is a child of another retrieval span
|
|
break
|
|
|
|
parent_id = parent_span.parent_id
|
|
else:
|
|
top_level_retrieval_spans.append(span)
|
|
|
|
return top_level_retrieval_spans
|
|
|
|
|
|
_RETRIEVER_DOCUMENT_CONTENT_KEYS = ("page_content", "content", "text")
|
|
_RETRIEVER_DOCUMENT_METADATA_KEYS = ("metadata",)
|
|
_MAX_RETRIEVER_DOCUMENT_WARNING_KEY_SETS = 128
|
|
_WARNED_RETRIEVER_DOCUMENT_KEY_SETS: OrderedDict[frozenset[str], None] = OrderedDict()
|
|
_WARNED_RETRIEVER_DOCUMENT_KEY_SETS_LOCK = threading.Lock()
|
|
|
|
|
|
def _should_warn_for_retriever_document_key_set(key_set: frozenset[str]) -> bool:
|
|
with _WARNED_RETRIEVER_DOCUMENT_KEY_SETS_LOCK:
|
|
if key_set in _WARNED_RETRIEVER_DOCUMENT_KEY_SETS:
|
|
_WARNED_RETRIEVER_DOCUMENT_KEY_SETS.move_to_end(key_set)
|
|
return False
|
|
|
|
_WARNED_RETRIEVER_DOCUMENT_KEY_SETS[key_set] = None
|
|
if len(_WARNED_RETRIEVER_DOCUMENT_KEY_SETS) > _MAX_RETRIEVER_DOCUMENT_WARNING_KEY_SETS:
|
|
_WARNED_RETRIEVER_DOCUMENT_KEY_SETS.popitem(last=False)
|
|
|
|
return True
|
|
|
|
|
|
def _parse_chunk(chunk: Any) -> dict[str, Any] | None:
|
|
if not isinstance(chunk, dict):
|
|
return None
|
|
|
|
content_key = next(
|
|
(key for key in _RETRIEVER_DOCUMENT_CONTENT_KEYS if key in chunk),
|
|
None,
|
|
)
|
|
content = chunk.get(content_key) if content_key is not None else None
|
|
|
|
if content_key is None:
|
|
# Many retriever libraries store source/citation details under metadata.
|
|
# Avoid warning for metadata-only chunks, but warn when other fields are
|
|
# present because they may contain text under an unsupported key.
|
|
non_metadata_keys = set(chunk) - set(_RETRIEVER_DOCUMENT_METADATA_KEYS)
|
|
if non_metadata_keys:
|
|
key_set = frozenset(map(str, chunk.keys()))
|
|
if _should_warn_for_retriever_document_key_set(key_set):
|
|
_logger.warning(
|
|
"RETRIEVER span document does not contain any recognized text field. "
|
|
"Expected one of %s. Found fields: %s",
|
|
list(_RETRIEVER_DOCUMENT_CONTENT_KEYS),
|
|
sorted(key_set),
|
|
)
|
|
|
|
metadata = chunk.get("metadata")
|
|
metadata = metadata if isinstance(metadata, dict) else {}
|
|
doc = {"content": content}
|
|
if doc_uri := metadata.get("doc_uri"):
|
|
doc["doc_uri"] = doc_uri
|
|
return doc
|
|
|
|
|
|
def clean_up_extra_traces(
|
|
traces: list[Trace],
|
|
eval_start_time: int,
|
|
experiment_id: str,
|
|
input_trace_ids: set[str] | None = None,
|
|
) -> None:
|
|
"""
|
|
Clean up noisy traces generated outside predict function.
|
|
|
|
Evaluation run should only contain traces that is being evaluated or generated by the predict
|
|
function. If not, the result will not show the correct list of traces.
|
|
Sometimes, there are extra traces generated during the evaluation, for example, custom scorer
|
|
code might generate traces. This function cleans up those noisy traces.
|
|
|
|
Args:
|
|
traces: List of traces to clean up.
|
|
eval_start_time: The start time of the evaluation run.
|
|
experiment_id: The experiment ID of the evaluation run.
|
|
input_trace_ids: Set of trace IDs that were passed in the input DataFrame.
|
|
These traces should never be deleted.
|
|
"""
|
|
try:
|
|
extra_trace_ids = [
|
|
trace.info.trace_id
|
|
for trace in traces
|
|
if not _should_keep_trace(trace, eval_start_time, input_trace_ids)
|
|
]
|
|
if extra_trace_ids:
|
|
_logger.debug(
|
|
f"Found {len(extra_trace_ids)} extra traces generated during evaluation run. "
|
|
"Deleting them."
|
|
)
|
|
# Import MlflowClient locally to avoid issues with tracing-only SDK
|
|
from mlflow.tracking.client import MlflowClient
|
|
|
|
MlflowClient().delete_traces(experiment_id=experiment_id, trace_ids=extra_trace_ids)
|
|
for trace_id in extra_trace_ids:
|
|
IPythonTraceDisplayHandler.get_instance().traces_to_display.pop(trace_id, None)
|
|
else:
|
|
_logger.debug("No extra traces found during evaluation run.")
|
|
except Exception as e:
|
|
_logger.debug(
|
|
f"Failed to clean up extra traces generated during evaluation. The "
|
|
f"result page might not show the correct list of traces. Error: {e}"
|
|
)
|
|
|
|
|
|
def _should_keep_trace(
|
|
trace: Trace,
|
|
eval_start_time: int,
|
|
input_trace_ids: set[str] | None = None,
|
|
) -> bool:
|
|
# Never delete traces that were explicitly passed in the input DataFrame.
|
|
if input_trace_ids and trace.info.trace_id in input_trace_ids:
|
|
return True
|
|
|
|
# We should not delete traces that are generated before the evaluation run started.
|
|
if trace.info.timestamp_ms < eval_start_time:
|
|
return True
|
|
|
|
# If the scorer tracing is enabled, keep traces generated by scorers.
|
|
if (
|
|
MLFLOW_GENAI_EVAL_ENABLE_SCORER_TRACING.get()
|
|
and TraceTagKey.SOURCE_SCORER_NAME in trace.info.tags
|
|
):
|
|
return True
|
|
|
|
# Otherwise, only keep traces from the prediction function.
|
|
return TraceTagKey.EVAL_REQUEST_ID in trace.info.tags
|
|
|
|
|
|
def construct_eval_result_df(
|
|
run_id: str,
|
|
traces: list[Trace],
|
|
eval_results: list["EvalResult"],
|
|
) -> "pd.DataFrame | None":
|
|
"""
|
|
Construct a pandas DataFrame from the traces and eval results.
|
|
|
|
Args:
|
|
run_id: The MLflow run ID of the evaluation run.
|
|
traces: List of traces. Only TraceInfo is used here, and **spans are ignored&**.
|
|
The expected input to this function is the result of
|
|
`mlflow.search_traces(include_spans=False, return_type="list")`.
|
|
eval_results: List of eval results containing the full spans.
|
|
|
|
Returns:
|
|
A pandas DataFrame with the eval results.
|
|
"""
|
|
import pandas as pd
|
|
|
|
if not traces:
|
|
return None
|
|
|
|
try:
|
|
trace_id_to_info = {t.info.trace_id: t.info for t in traces}
|
|
# Skip results whose trace could not be materialized (e.g. a clone read-back miss
|
|
# nulled eval_item.trace); otherwise a single missing trace would collapse the entire
|
|
# result DataFrame to None via the except below.
|
|
traces = [
|
|
Trace(
|
|
info=trace_id_to_info[eval_result.eval_item.trace.info.trace_id],
|
|
data=eval_result.eval_item.trace.data,
|
|
)
|
|
for eval_result in eval_results
|
|
if eval_result.eval_item.trace is not None
|
|
]
|
|
df = traces_to_df(traces)
|
|
# Add unpacked assessment columns. The result df should look like:
|
|
# [trace_id, score_1/value, score_2/value, trace, state, ...]
|
|
assessments = (
|
|
df["assessments"].apply(lambda x: _get_assessment_values(x, run_id)).apply(pd.Series)
|
|
)
|
|
trace_id_column = df.pop("trace_id")
|
|
return pd.concat([trace_id_column, assessments, df], axis=1)
|
|
except Exception as e:
|
|
_logger.debug(f"Failed to construct eval result DataFrame: {e}", exc_info=True)
|
|
|
|
|
|
def _get_assessment_values(assessments: list[dict[str, Any]], run_id: str) -> dict[str, Any]:
|
|
result = {}
|
|
for a in assessments:
|
|
if (
|
|
# Exclude feedbacks from other evaluation runs
|
|
(source_run_id := a.get("metadata", {}).get(AssessmentMetadataKey.SOURCE_RUN_ID))
|
|
and source_run_id != run_id
|
|
):
|
|
continue
|
|
name = a["assessment_name"]
|
|
if feedback := a.get("feedback"):
|
|
result[f"{name}/value"] = feedback.get("value")
|
|
# Carry the rationale and any scorer error so downstream consumers (e.g.
|
|
# EvaluationResult.passed/reason) can surface them. Emitted only when
|
|
# present to keep the result DataFrame compact.
|
|
if (rationale := a.get("rationale")) is not None:
|
|
result[f"{name}/rationale"] = rationale
|
|
if (error := feedback.get("error")) and (msg := error.get("error_message")):
|
|
result[f"{name}/error_message"] = msg
|
|
elif expectation := a.get("expectation"):
|
|
result[f"{name}/value"] = expectation.get("value")
|
|
|
|
return result
|
|
|
|
|
|
def create_minimal_trace(eval_item: "EvalItem") -> Trace:
|
|
"""
|
|
Create a minimal trace object with a single span, based on given inputs/outputs.
|
|
|
|
If the eval_item has a source with session metadata (from a dataset created from traces),
|
|
the session metadata will be restored on the newly created trace. This enables session-level
|
|
scorers to identify which traces belong to the same session.
|
|
"""
|
|
from mlflow.pyfunc.context import Context, set_prediction_context
|
|
|
|
# Extract session metadata from source if available
|
|
session_metadata = {}
|
|
if eval_item.source and hasattr(eval_item.source, "source_data"):
|
|
source_data = eval_item.source.source_data
|
|
if session_id := source_data.get("session_id"):
|
|
session_metadata[TraceMetadataKey.TRACE_SESSION] = session_id
|
|
|
|
context = Context(request_id=eval_item.request_id, is_evaluate=True)
|
|
with set_prediction_context(context):
|
|
with mlflow.start_span(name="root_span", span_type=SpanType.CHAIN) as root_span:
|
|
root_span.set_inputs(eval_item.inputs)
|
|
root_span.set_outputs(eval_item.outputs)
|
|
|
|
# Set session metadata on the trace while it's still active
|
|
if session_metadata:
|
|
mlflow.update_current_trace(metadata=session_metadata)
|
|
|
|
return mlflow.get_trace(root_span.trace_id)
|
|
|
|
|
|
# MB: Caching on tracking URI level to avoid unnecessary checks for each trace.
|
|
@cached(cache={}, key=lambda **kwargs: kwargs["tracking_uri"])
|
|
def _does_store_support_trace_linking(*, tracking_uri: str, trace: Trace, run_id: str) -> bool:
|
|
# Databricks backend is guaranteed to support trace linking
|
|
if is_databricks_uri(tracking_uri):
|
|
return True
|
|
|
|
try:
|
|
MlflowClient(tracking_uri).link_traces_to_run([trace.info.trace_id], run_id=run_id)
|
|
return True
|
|
except Exception:
|
|
return False
|
|
|
|
|
|
def batch_link_traces_to_run(
|
|
run_id: str | None, eval_results: list["EvalResult"], max_batch_size: int = 100
|
|
) -> None:
|
|
"""
|
|
Batch link traces to a run to avoid rate limits.
|
|
|
|
Args:
|
|
run_id: The MLflow run ID to link traces to
|
|
eval_results: List of evaluation results containing traces
|
|
max_batch_size: Maximum number of traces to link per batch call
|
|
"""
|
|
trace_ids = [
|
|
eval_result.eval_item.trace.info.trace_id
|
|
for eval_result in eval_results
|
|
if eval_result.eval_item.trace is not None
|
|
]
|
|
# Batch the trace IDs to avoid overwhelming the MLflow backend
|
|
for i in range(0, len(trace_ids), max_batch_size):
|
|
batch = trace_ids[i : i + max_batch_size]
|
|
try:
|
|
MlflowClient().link_traces_to_run(run_id=run_id, trace_ids=batch)
|
|
except Exception as e:
|
|
# FileStore doesn't support trace linking, so we skip it
|
|
if "Linking traces to runs is not supported in FileStore." in str(e):
|
|
return
|
|
|
|
_logger.warning(f"Failed to link batch of traces to run: {e}")
|
|
|
|
|
|
class ExtractedToolsFromTrace(BaseModel):
|
|
tools: list["ChatTool"] = Field(
|
|
default_factory=list,
|
|
description="List of all available tools found in the trace",
|
|
)
|
|
|
|
model_config = {"extra": "forbid"}
|
|
|
|
|
|
def extract_available_tools_from_trace(trace: Trace, model: str | None = None) -> list["ChatTool"]:
|
|
"""
|
|
Extract available tools from a trace by checking all LLM spans.
|
|
|
|
This function uses a two-stage approach:
|
|
1. Programmatic extraction: Checks all LLM and CHAT_MODEL spans for tools in
|
|
attributes (mlflow.chat.tools) and inputs (inputs.tools field).
|
|
2. LLM fallback: If no tools are found programmatically, uses an LLM to analyze
|
|
the trace and identify tool definitions.
|
|
|
|
The programmatic approach mirrors the frontend's getChatToolsFromSpan logic in
|
|
ModelTraceExplorer.utils.tsx, which extracts tools per-span and returns a
|
|
deduplicated list of all unique tools found across the trace.
|
|
|
|
Args:
|
|
trace: MLflow trace object
|
|
model: Optional model URI to use for LLM-based fallback extraction
|
|
(e.g., "openai:/gpt-4"). If None, uses a default model.
|
|
|
|
Returns:
|
|
List of unique ChatTool objects, or an empty list if no valid tools are found.
|
|
"""
|
|
# Stage 1: Programmatic extraction from span attributes and inputs
|
|
all_tools = []
|
|
seen_tool_signatures = set()
|
|
|
|
relevant_span_types = [SpanType.LLM, SpanType.CHAT_MODEL]
|
|
|
|
for span in trace.data.spans:
|
|
span_type = span.get_attribute(SpanAttributeKey.SPAN_TYPE)
|
|
if span_type not in relevant_span_types:
|
|
continue
|
|
|
|
span_tools = _extract_tools_from_span(span)
|
|
|
|
for tool in span_tools:
|
|
if tool.function:
|
|
tool_signature = _get_tool_signature(tool)
|
|
if tool_signature not in seen_tool_signatures:
|
|
seen_tool_signatures.add(tool_signature)
|
|
all_tools.append(tool)
|
|
|
|
if all_tools:
|
|
return all_tools
|
|
|
|
# Stage 2: LLM fallback when programmatic extraction yields no results
|
|
return _try_extract_available_tools_with_llm(trace, model)
|
|
|
|
|
|
def _get_tool_signature(tool: "ChatTool") -> str:
|
|
if not tool.function:
|
|
return ""
|
|
|
|
try:
|
|
tool_dict = tool.function.model_dump()
|
|
except AttributeError:
|
|
tool_dict = tool.function.dict()
|
|
|
|
return json.dumps(tool_dict, sort_keys=True)
|
|
|
|
|
|
def _extract_tools_from_span(span: Span) -> list["ChatTool"]:
|
|
"""
|
|
Extract tools from a single LLM or CHAT_MODEL span, checking attribute first, then inputs.
|
|
|
|
This mirrors the frontend's getChatToolsFromSpan logic exactly, but returns
|
|
validated ChatTool objects using Pydantic validation.
|
|
|
|
Args:
|
|
span: MLflow span object
|
|
|
|
Returns:
|
|
List of ChatTool objects for this span
|
|
"""
|
|
tools_attribute = span.get_attribute(SpanAttributeKey.CHAT_TOOLS)
|
|
if tools_attribute is not None:
|
|
try:
|
|
if isinstance(tools_attribute, str):
|
|
tools_attribute = json.loads(tools_attribute)
|
|
return _parse_tools_to_chat_tool(tools_attribute)
|
|
except Exception as e:
|
|
_logger.debug(f"Failed to parse tools from attribute in span {span.span_id}: {e}")
|
|
|
|
if span.inputs is not None:
|
|
try:
|
|
inputs = _to_dict(span.inputs)
|
|
if "tools" in inputs:
|
|
return _parse_tools_to_chat_tool(inputs["tools"])
|
|
except Exception as e:
|
|
_logger.debug(f"Failed to parse tools from inputs in span {span.span_id}: {e}")
|
|
|
|
return []
|
|
|
|
|
|
def _parse_tools_to_chat_tool(tools_data: list[dict[str, Any]]) -> list["ChatTool"]:
|
|
"""
|
|
Parse a list of tool dictionaries into ChatTool objects using Pydantic validation.
|
|
|
|
Args:
|
|
tools_data: List of tool dictionaries
|
|
|
|
Returns:
|
|
List of validated ChatTool objects. Invalid tools are skipped with debug logging.
|
|
"""
|
|
from mlflow.types.chat import ChatTool
|
|
|
|
validated_tools = []
|
|
for data in tools_data:
|
|
try:
|
|
tool = ChatTool(**data)
|
|
validated_tools.append(tool)
|
|
except Exception as e:
|
|
_logger.debug(f"Skipping invalid tool {data}: {e}")
|
|
|
|
return validated_tools
|
|
|
|
|
|
def _try_extract_available_tools_with_llm(
|
|
trace: Trace, model: str | None = None
|
|
) -> list["ChatTool"]:
|
|
"""
|
|
Attempt to extract available tools from trace using LLM with structured output.
|
|
|
|
This is a fallback method when programmatic extraction fails. It uses an LLM to
|
|
analyze the trace and identify tool definitions that were available to the agent.
|
|
|
|
Args:
|
|
trace: MLflow trace object to analyze
|
|
model: Optional model URI to use for extraction (e.g., "openai:/gpt-4").
|
|
If None, uses a default model.
|
|
|
|
Returns:
|
|
List of ChatTool objects extracted by the LLM, or empty list if extraction fails.
|
|
"""
|
|
if model is None:
|
|
model = get_default_model()
|
|
|
|
try:
|
|
from mlflow.types.chat import (
|
|
ChatTool,
|
|
FunctionParams,
|
|
FunctionToolDefinition,
|
|
ParamProperty,
|
|
)
|
|
|
|
output_example = json.dumps(
|
|
ExtractedToolsFromTrace(
|
|
tools=[
|
|
ChatTool(
|
|
type="function",
|
|
function=FunctionToolDefinition(
|
|
name="example_tool",
|
|
description="Description of what the tool does",
|
|
parameters=FunctionParams(
|
|
type="object",
|
|
properties={
|
|
"param1": ParamProperty(
|
|
type="string",
|
|
description="A parameter",
|
|
)
|
|
},
|
|
required=["param1"],
|
|
),
|
|
),
|
|
)
|
|
]
|
|
).model_dump(),
|
|
indent=2,
|
|
)
|
|
|
|
messages = get_available_tools_extraction_prompts(output_example)
|
|
|
|
result = get_chat_completions_with_structured_output(
|
|
model_uri=model,
|
|
messages=messages,
|
|
output_schema=ExtractedToolsFromTrace,
|
|
trace=trace,
|
|
)
|
|
|
|
return result.tools
|
|
|
|
except Exception as e:
|
|
_logger.warning(
|
|
f"Failed to extract tools from trace using LLM. Returning empty list. Error: {e!r}"
|
|
)
|
|
return []
|