wehub-resource-sync
278 lines
8.9 KiB
Python
278 lines
8.9 KiB
Python
"""
|
|
Utility functions for trace evaluation output formatting.
|
|
"""
|
|
|
|
from dataclasses import dataclass
|
|
from typing import Any
|
|
|
|
import click
|
|
import pandas as pd
|
|
|
|
from mlflow.exceptions import MlflowException
|
|
from mlflow.genai.scorers import Scorer, get_all_scorers, get_scorer
|
|
from mlflow.tracing.constant import AssessmentMetadataKey
|
|
|
|
# Represents the absence of a value for an assessment
|
|
NA_VALUE = "N/A"
|
|
|
|
|
|
@dataclass
|
|
class Assessment:
|
|
"""
|
|
Structured assessment data for a trace evaluation.
|
|
"""
|
|
|
|
name: str | None
|
|
"""The name of the assessment"""
|
|
|
|
result: Any | None = None
|
|
"""The result value from the assessment"""
|
|
|
|
rationale: str | None = None
|
|
"""The rationale text explaining the assessment"""
|
|
|
|
error: str | None = None
|
|
"""Error message if the assessment failed"""
|
|
|
|
|
|
@dataclass
|
|
class Cell:
|
|
"""
|
|
Structured cell data for table display with metadata.
|
|
"""
|
|
|
|
value: str
|
|
"""The formatted display value for the cell"""
|
|
|
|
assessment: Assessment | None = None
|
|
"""The assessment data for this cell, if it represents an assessment"""
|
|
|
|
|
|
@dataclass
|
|
class EvalResult:
|
|
"""
|
|
Container for evaluation results for a single trace.
|
|
|
|
This dataclass provides structured access to trace evaluation data,
|
|
replacing dict-based access for better type safety.
|
|
"""
|
|
|
|
trace_id: str
|
|
"""The trace ID"""
|
|
|
|
assessments: list[Assessment]
|
|
"""List of Assessment objects for this trace"""
|
|
|
|
|
|
@dataclass
|
|
class TableOutput:
|
|
"""Container for formatted table data."""
|
|
|
|
headers: list[str]
|
|
rows: list[list[Cell]]
|
|
|
|
|
|
def _format_assessment_cell(assessment: Assessment | None) -> Cell:
|
|
"""
|
|
Format a single assessment cell for table display.
|
|
|
|
Args:
|
|
assessment: Assessment object with result, rationale, and error fields
|
|
|
|
Returns:
|
|
Cell object with formatted value and assessment metadata
|
|
"""
|
|
if not assessment:
|
|
return Cell(value=NA_VALUE)
|
|
|
|
if assessment.error:
|
|
display_value = f"error: {assessment.error}"
|
|
elif assessment.result is not None and assessment.rationale:
|
|
display_value = f"value: {assessment.result}, rationale: {assessment.rationale}"
|
|
elif assessment.result is not None:
|
|
display_value = f"value: {assessment.result}"
|
|
elif assessment.rationale:
|
|
display_value = f"rationale: {assessment.rationale}"
|
|
else:
|
|
display_value = NA_VALUE
|
|
|
|
return Cell(value=display_value, assessment=assessment)
|
|
|
|
|
|
def resolve_scorers(scorer_names: list[str], experiment_id: str) -> list[Scorer]:
|
|
"""
|
|
Resolve scorer names to scorer objects.
|
|
|
|
Checks built-in scorers first, then registered scorers.
|
|
Supports both class names (e.g., "RelevanceToQuery") and snake_case
|
|
scorer names (e.g., "relevance_to_query").
|
|
|
|
Args:
|
|
scorer_names: List of scorer names to resolve
|
|
experiment_id: Experiment ID for looking up registered scorers
|
|
|
|
Returns:
|
|
List of resolved scorer objects
|
|
|
|
Raises:
|
|
click.UsageError: If a scorer is not found or no valid scorers specified
|
|
"""
|
|
resolved_scorers = []
|
|
builtin_scorers = get_all_scorers()
|
|
# Build map with both class name and snake_case name for lookup
|
|
builtin_scorer_map = {}
|
|
for scorer in builtin_scorers:
|
|
# Map by class name (e.g., "RelevanceToQuery")
|
|
builtin_scorer_map[scorer.__class__.__name__] = scorer
|
|
# Map by scorer.name (snake_case, e.g., "relevance_to_query")
|
|
if scorer.name is not None:
|
|
builtin_scorer_map[scorer.name] = scorer
|
|
|
|
for scorer_name in scorer_names:
|
|
if scorer_name in builtin_scorer_map:
|
|
resolved_scorers.append(builtin_scorer_map[scorer_name])
|
|
else:
|
|
# Try to get it as a registered scorer
|
|
try:
|
|
registered_scorer = get_scorer(name=scorer_name, experiment_id=experiment_id)
|
|
resolved_scorers.append(registered_scorer)
|
|
except MlflowException as e:
|
|
error_message = str(e)
|
|
if "not found" in error_message.lower():
|
|
available_builtin = ", ".join(
|
|
sorted({scorer.__class__.__name__ for scorer in builtin_scorers})
|
|
)
|
|
raise click.UsageError(
|
|
f"Could not identify Scorer '{scorer_name}'. "
|
|
f"Only built-in or registered scorers can be resolved. "
|
|
f"Available built-in scorers: {available_builtin}. "
|
|
f"To use a custom scorer, register it first in experiment {experiment_id} "
|
|
f"using the register_scorer() API."
|
|
)
|
|
else:
|
|
raise click.UsageError(
|
|
f"An error occurred when retrieving information for Scorer "
|
|
f"`{scorer_name}`: {error_message}"
|
|
)
|
|
|
|
if not resolved_scorers:
|
|
raise click.UsageError("No valid scorers specified")
|
|
|
|
return resolved_scorers
|
|
|
|
|
|
def extract_assessments_from_results(
|
|
results_df: pd.DataFrame, evaluation_run_id: str
|
|
) -> list[EvalResult]:
|
|
"""
|
|
Extract assessments from evaluation results DataFrame.
|
|
|
|
The evaluate() function returns results with a DataFrame that contains
|
|
an 'assessments' column. Each row has a list of assessment dictionaries
|
|
with metadata including AssessmentMetadataKey.SOURCE_RUN_ID that we use to
|
|
filter assessments from this specific evaluation run.
|
|
|
|
Args:
|
|
results_df: DataFrame from evaluate() results containing assessments column
|
|
evaluation_run_id: The MLflow run ID from the evaluation that generated the assessments
|
|
|
|
Returns:
|
|
List of EvalResult objects with trace_id and assessments
|
|
"""
|
|
output_data = []
|
|
|
|
for _, row in results_df.iterrows():
|
|
trace_id = row.get("trace_id", "unknown")
|
|
assessments_list = []
|
|
|
|
for assessment_dict in row.get("assessments", []):
|
|
# Only consider assessments from the evaluation run
|
|
metadata = assessment_dict.get("metadata", {})
|
|
source_run_id = metadata.get(AssessmentMetadataKey.SOURCE_RUN_ID)
|
|
|
|
if source_run_id != evaluation_run_id:
|
|
continue
|
|
|
|
assessment_name = assessment_dict.get("assessment_name")
|
|
assessment_result = None
|
|
assessment_rationale = None
|
|
assessment_error = None
|
|
|
|
if (feedback := assessment_dict.get("feedback")) and isinstance(feedback, dict):
|
|
assessment_result = feedback.get("value")
|
|
|
|
if rationale := assessment_dict.get("rationale"):
|
|
assessment_rationale = rationale
|
|
|
|
if error := assessment_dict.get("error"):
|
|
assessment_error = str(error)
|
|
|
|
assessments_list.append(
|
|
Assessment(
|
|
name=assessment_name,
|
|
result=assessment_result,
|
|
rationale=assessment_rationale,
|
|
error=assessment_error,
|
|
)
|
|
)
|
|
|
|
# If no assessments were found for this trace, add error markers
|
|
if not assessments_list:
|
|
assessments_list.append(
|
|
Assessment(
|
|
name=NA_VALUE,
|
|
result=None,
|
|
rationale=None,
|
|
error="No assessments found on trace",
|
|
)
|
|
)
|
|
|
|
output_data.append(EvalResult(trace_id=trace_id, assessments=assessments_list))
|
|
|
|
return output_data
|
|
|
|
|
|
def format_table_output(output_data: list[EvalResult]) -> TableOutput:
|
|
"""
|
|
Format evaluation results as table data.
|
|
|
|
Args:
|
|
output_data: List of EvalResult objects with assessments
|
|
|
|
Returns:
|
|
TableOutput dataclass containing headers and rows
|
|
"""
|
|
# Extract unique assessment names from output_data to use as column headers
|
|
# Note: assessment name can be None, so we filter it out
|
|
assessment_names_set = set()
|
|
for trace_result in output_data:
|
|
for assessment in trace_result.assessments:
|
|
if assessment.name and assessment.name != NA_VALUE:
|
|
assessment_names_set.add(assessment.name)
|
|
|
|
# Sort for consistent ordering
|
|
assessment_names = sorted(assessment_names_set)
|
|
|
|
headers = ["trace_id"] + assessment_names
|
|
table_data = []
|
|
|
|
for trace_result in output_data:
|
|
# Create Cell for trace_id column
|
|
row = [Cell(value=trace_result.trace_id)]
|
|
|
|
# Build a map of assessment name -> assessment for this trace
|
|
assessment_map = {
|
|
assessment.name: assessment
|
|
for assessment in trace_result.assessments
|
|
if assessment.name and assessment.name != NA_VALUE
|
|
}
|
|
|
|
# For each assessment name in headers, get the corresponding assessment
|
|
for assessment_name in assessment_names:
|
|
cell_content = _format_assessment_cell(assessment_map.get(assessment_name))
|
|
row.append(cell_content)
|
|
|
|
table_data.append(row)
|
|
|
|
return TableOutput(headers=headers, rows=table_data)
|