from typing import Any
from mlflow.entities.assessment import (
DEFAULT_FEEDBACK_NAME,
Assessment,
AssessmentError,
Expectation,
Feedback,
FeedbackValueType,
IssueReference,
)
from mlflow.entities.assessment_source import AssessmentSource
from mlflow.entities.issue_evidence import IssueEvidence
from mlflow.exceptions import MlflowException
from mlflow.tracing.client import TracingClient
from mlflow.tracing.trace_manager import InMemoryTraceManager
[docs]def get_assessment(trace_id: str, assessment_id: str) -> Assessment:
"""
Get an assessment entity from the backend store.
Args:
trace_id: The ID of the trace.
assessment_id: The ID of the assessment to get.
Returns:
:py:class:`~mlflow.entities.Assessment`: The Assessment object.
"""
return TracingClient().get_assessment(trace_id, assessment_id)
[docs]def log_assessment(trace_id: str, assessment: Assessment) -> Assessment:
"""
Logs an assessment to a Trace. The assessment can be an expectation or a feedback.
- Expectation: A label that represents the expected value for a particular operation.
For example, an expected answer for a user question from a chatbot.
- Feedback: A label that represents the feedback on the quality of the operation.
Feedback can come from different sources, such as human judges, heuristic scorers,
or LLM-as-a-Judge.
The following code annotates a trace with a feedback provided by LLM-as-a-Judge.
.. code-block:: python
import mlflow
from mlflow.entities import Feedback
feedback = Feedback(
name="faithfulness",
value=0.9,
rationale="The model is faithful to the input.",
metadata={"model": "gpt-4o-mini"},
)
mlflow.log_assessment(trace_id="1234", assessment=feedback)
The following code annotates a trace with human-provided ground truth with source information.
When the source is not provided, the default source is set to "default" with type "HUMAN"
.. code-block:: python
import mlflow
from mlflow.entities import AssessmentSource, AssessmentSourceType, Expectation
# Specify the annotator information as a source.
source = AssessmentSource(
source_type=AssessmentSourceType.HUMAN,
source_id="john@example.com",
)
expectation = Expectation(
name="expected_answer",
value=42,
source=source,
)
mlflow.log_assessment(trace_id="1234", assessment=expectation)
The expectation value can be any JSON-serializable value. For example, you may
record the full LLM message as the expectation value.
.. code-block:: python
import mlflow
from mlflow.entities.assessment import Expectation
expectation = Expectation(
name="expected_message",
# Full LLM message including expected tool calls
value={
"role": "assistant",
"content": "The answer is 42.",
"tool_calls": [
{
"id": "1234",
"type": "function",
"function": {"name": "add", "arguments": "40 + 2"},
}
],
},
)
mlflow.log_assessment(trace_id="1234", assessment=expectation)
You can also log an error information during the feedback generation process. To do so,
provide an instance of :py:class:`~mlflow.entities.AssessmentError` to the `error`
parameter, and leave the `value` parameter as `None`.
.. code-block:: python
import mlflow
from mlflow.entities import AssessmentError, Feedback
error = AssessmentError(
error_code="RATE_LIMIT_EXCEEDED",
error_message="Rate limit for the judge exceeded.",
)
feedback = Feedback(
trace_id="1234",
name="faithfulness",
error=error,
)
mlflow.log_assessment(trace_id="1234", assessment=feedback)
"""
return TracingClient().log_assessment(trace_id, assessment)
[docs]def log_expectation(
*,
trace_id: str,
name: str,
value: Any,
source: AssessmentSource | None = None,
metadata: dict[str, Any] | None = None,
span_id: str | None = None,
) -> Assessment:
"""
Logs an expectation (e.g. ground truth label) to a Trace. This API only takes keyword arguments.
Args:
trace_id: The ID of the trace.
name: The name of the expectation assessment e.g., "expected_answer
value: The value of the expectation. It can be any JSON-serializable value.
source: The source of the expectation assessment. Must be an instance of
:py:class:`~mlflow.entities.AssessmentSource`. If not provided,
default to HUMAN source type.
metadata: Additional metadata for the expectation.
span_id: The ID of the span associated with the expectation, if it needs be
associated with a specific span in the trace.
Returns:
:py:class:`~mlflow.entities.Assessment`: The created expectation assessment.
Examples:
.. code-block:: python
import mlflow
from mlflow.entities import AssessmentSource, AssessmentSourceType
# Log simple expected answer
expectation = mlflow.log_expectation(
trace_id="tr-1234567890abcdef",
name="expected_answer",
value="The capital of France is Paris.",
source=AssessmentSource(
source_type=AssessmentSourceType.HUMAN, source_id="annotator@company.com"
),
metadata={"question_type": "factual", "difficulty": "easy"},
)
# Log expected classification label
mlflow.log_expectation(
trace_id="tr-1234567890abcdef",
name="expected_category",
value="positive",
source=AssessmentSource(
source_type=AssessmentSourceType.HUMAN, source_id="data_labeler_001"
),
metadata={"labeling_session": "batch_01", "confidence": 0.95},
)
"""
assessment = Expectation(
name=name,
source=source,
value=value,
metadata=metadata,
span_id=span_id,
)
return TracingClient().log_assessment(trace_id, assessment)
[docs]def update_assessment(
trace_id: str,
assessment_id: str,
assessment: Assessment,
) -> Assessment:
"""
Updates an existing expectation (ground truth) in a Trace.
Args:
trace_id: The ID of the trace.
assessment_id: The ID of the expectation or feedback assessment to update.
assessment: The updated assessment.
Returns:
:py:class:`~mlflow.entities.Assessment`: The updated feedback or expectation assessment.
Example:
The following code updates an existing expectation with a new value.
To update other fields, provide the corresponding parameters.
.. code-block:: python
import mlflow
from mlflow.entities import Expectation, ExpectationValue
# Create an expectation with value 42.
response = mlflow.log_assessment(
trace_id="1234",
assessment=Expectation(name="expected_answer", value=42),
)
assessment_id = response.assessment_id
# Update the expectation with a new value 43.
mlflow.update_assessment(
trace_id="1234",
assessment_id=assessment.assessment_id,
assessment=Expectation(name="expected_answer", value=43),
)
"""
return TracingClient().update_assessment(
assessment_id=assessment_id,
trace_id=trace_id,
assessment=assessment,
)
[docs]def delete_assessment(trace_id: str, assessment_id: str):
"""
Deletes an assessment associated with a trace.
Args:
trace_id: The ID of the trace.
assessment_id: The ID of the assessment to delete.
"""
return TracingClient().delete_assessment(trace_id=trace_id, assessment_id=assessment_id)
[docs]def log_feedback(
*,
trace_id: str,
name: str = DEFAULT_FEEDBACK_NAME,
value: FeedbackValueType | None = None,
source: AssessmentSource | None = None,
error: Exception | AssessmentError | None = None,
rationale: str | None = None,
metadata: dict[str, Any] | None = None,
span_id: str | None = None,
) -> Assessment:
"""
Logs feedback to a Trace. This API only takes keyword arguments.
Args:
trace_id: The ID of the trace.
name: The name of the feedback assessment e.g., "faithfulness". Defaults to
"feedback" if not provided.
value: The value of the feedback. Can be ``None``, or one of the following types:
- float
- int
- str
- bool
- list of values of the same types as above
- dict with string keys and values of the same types as above
source: The source of the feedback assessment. Must be an instance of
:py:class:`~mlflow.entities.AssessmentSource`. If not provided, defaults to
CODE source type
error: An error object representing any issues encountered while computing the
feedback, e.g., a timeout error from an LLM judge. Accepts an exception
object, or an :py:class:`~mlflow.entities.AssessmentError` object. Either
this or `value` must be provided.
rationale: The rationale / justification for the feedback.
metadata: Additional metadata for the feedback.
span_id: The ID of the span associated with the feedback, if it needs be
associated with a specific span in the trace.
Returns:
:py:class:`~mlflow.entities.Assessment`: The created feedback assessment.
Examples:
.. code-block:: python
import mlflow
from mlflow.entities import AssessmentSource, AssessmentSourceType
# Log simple feedback score
feedback = mlflow.log_feedback(
trace_id="tr-1234567890abcdef",
name="relevance",
value=0.9,
source=AssessmentSource(
source_type=AssessmentSourceType.LLM_JUDGE, source_id="gpt-4"
),
rationale="Response directly addresses the user's question",
)
# Log detailed feedback with structured data
mlflow.log_feedback(
trace_id="tr-1234567890abcdef",
name="quality_metrics",
value={"accuracy": 0.95, "completeness": 0.88, "clarity": 0.92, "overall": 0.92},
source=AssessmentSource(
source_type=AssessmentSourceType.HUMAN, source_id="expert_evaluator"
),
rationale="High accuracy and clarity, slightly incomplete coverage",
)
"""
assessment = Feedback(
name=name,
source=source,
value=value,
error=error,
rationale=rationale,
metadata=metadata,
span_id=span_id,
)
return TracingClient().log_assessment(trace_id, assessment)
def _get_trace_span_ids(trace_id: str) -> list[str]:
"""
Every span ID the trace carries, read from the in-progress trace when there is one.
An issue can be logged from inside the traced call, before anything has been exported, which
is why the active trace is consulted first -- the same order ``get_trace_info`` uses. A trace
recorded elsewhere, or already finished, is read from the store.
An in-progress trace is a complete view of itself: spans are only added to it and the whole
trace is popped once its root span ends, so spans that incremental export has already written
to the store are still listed here. The exception is a process that joins a distributed trace,
which registers it locally holding only the spans that process created; evidence logged there
can cite only those spans, the same ones ``get_trace`` returns in that process.
"""
with InMemoryTraceManager.get_instance().get_trace(trace_id) as trace:
if trace is not None:
return list(trace.span_dict)
return [span.span_id for span in TracingClient().get_trace(trace_id).data.spans]
[docs]def log_issue(
*,
trace_id: str,
issue_id: str,
issue_name: str | None = None,
source: AssessmentSource | None = None,
run_id: str | None = None,
rationale: str | None = None,
metadata: dict[str, str] | None = None,
span_id: str | None = None,
evidence: IssueEvidence | dict[str, Any] | None = None,
) -> Assessment:
"""
Logs an issue reference to a Trace. This links a trace to a discovered issue.
This API only takes keyword arguments.
Args:
trace_id: The ID of the trace.
issue_id: The ID of the issue to reference.
issue_name: The name of the issue. If not provided, the issue name will be fetched
using the issue ID.
source: The source represents how this issue was detected on the trace, for example,
human review or automatic AI scan. Must be an instance of
:py:class:`~mlflow.entities.AssessmentSource`. If not provided, defaults to
LLM_JUDGE source type.
run_id: The ID of the run that detect the issue on the given trace.
rationale: The rationale / justification for the issue reference.
metadata: Additional metadata for the issue reference.
span_id: The ID of the span associated with the issue, if it needs to be
associated with a specific span in the trace. Defaults to the evidence's
primary span when ``evidence`` names one.
evidence: Optional :py:class:`~mlflow.entities.IssueEvidence` (or its dictionary form)
explaining the issue on this trace as an ordered walkthrough of its spans. Its
summary is stored as the reference's rationale and its steps as versioned
metadata, so clients that predate structured evidence still read the summary.
Every cited span must belong to the trace.
Returns:
:py:class:`~mlflow.entities.Assessment`: The created issue reference assessment.
Examples:
.. code-block:: python
import mlflow
from mlflow.entities import AssessmentSource, AssessmentSourceType
# Log an issue reference to link a trace to a discovered issue
issue_ref = mlflow.log_issue(
trace_id="tr-1234567890abcdef",
issue_id="iss-abcdef123456",
source=AssessmentSource(
source_type=AssessmentSourceType.LLM_JUDGE, source_id="issue-discovery-agent"
),
run_id="run-123",
rationale="Response time exceeded 2 seconds threshold",
)
"""
fetched_issue_name = TracingClient()._get_issue(issue_id).name
if issue_name is None:
issue_name = fetched_issue_name
elif issue_name != fetched_issue_name:
raise MlflowException.invalid_parameter_value(
f"Provided issue name {issue_name!r} does not match the issue name "
f"{fetched_issue_name!r} with issue ID {issue_id}."
)
if evidence is not None:
evidence = evidence if isinstance(evidence, IssueEvidence) else IssueEvidence(**evidence)
# Validate the cited spans against the trace itself: a fabricated locator would render a
# step that opens nothing, and only the trace can say which spans exist.
evidence.validate_spans(_get_trace_span_ids(trace_id))
# The reference's span is what the UI opens from the card, so it has to be a span the
# walkthrough actually explains. An override that isn't cited would send the reader to a
# span nothing on the card describes.
if span_id is not None and span_id not in evidence.span_ids:
raise MlflowException.invalid_parameter_value(
f"span_id {span_id!r} is not cited by any evidence step. Pass a span the evidence "
"explains, or leave span_id unset to use the evidence's primary span."
)
metadata = evidence.to_metadata(metadata)
# The summary is the rationale, so a reader without this schema still explains the issue.
rationale = rationale or evidence.summary
span_id = span_id or evidence.primary_span_id
assessment = IssueReference(
issue_id=issue_id,
issue_name=issue_name,
source=source,
trace_id=trace_id,
run_id=run_id,
rationale=rationale,
metadata=metadata,
span_id=span_id,
)
return TracingClient().log_assessment(trace_id, assessment)
[docs]def override_feedback(
*,
trace_id: str,
assessment_id: str,
value: FeedbackValueType,
rationale: str | None = None,
source: AssessmentSource | None = None,
metadata: dict[str, Any] | None = None,
) -> Assessment:
"""
Overrides an existing feedback assessment with a new assessment. This API
logs a new assessment with the `overrides` field set to the provided assessment ID.
The original assessment will be marked as invalid, but will otherwise be unchanged.
This is useful when you want to correct an assessment generated by an LLM judge,
but want to preserve the original assessment for future judge fine-tuning.
If you want to mutate an assessment in-place, use :py:func:`update_assessment` instead.
Args:
trace_id: The ID of the trace.
assessment_id: The ID of the assessment to override.
value: The new value of the assessment.
rationale: The rationale of the new assessment.
source: The source of the new assessment.
metadata: Additional metadata for the new assessment.
Returns:
:py:class:`~mlflow.entities.Assessment`: The created assessment.
Examples:
.. code-block:: python
import mlflow
from mlflow.entities import AssessmentSource, AssessmentSourceType
# First, log an initial LLM-generated feedback as a simulation
llm_feedback = mlflow.log_feedback(
trace_id="tr-1234567890abcdef",
name="relevance",
value=0.6,
source=AssessmentSource(
source_type=AssessmentSourceType.LLM_JUDGE, source_id="gpt-4"
),
rationale="Response partially addresses the question",
)
# Later, a human reviewer disagrees and wants to override
corrected_assessment = mlflow.override_feedback(
trace_id="tr-1234567890abcdef",
assessment_id=llm_feedback.assessment_id,
value=0.9,
rationale="Response fully addresses the question with good examples",
source=AssessmentSource(
source_type=AssessmentSourceType.HUMAN, source_id="expert_reviewer@company.com"
),
metadata={
"override_reason": "LLM underestimated relevance",
"review_date": "2024-01-15",
"confidence": "high",
},
)
"""
old_assessment = get_assessment(trace_id, assessment_id)
if not isinstance(old_assessment, Feedback):
raise MlflowException.invalid_parameter_value(
f"The assessment with ID {assessment_id} is not a feedback assessment."
)
new_assessment = Feedback(
name=old_assessment.name,
span_id=old_assessment.span_id,
value=value,
rationale=rationale,
source=source,
metadata=metadata,
overrides=old_assessment.assessment_id,
)
return TracingClient().log_assessment(trace_id, new_assessment)