# mlflow
pip install git+https://github.com/mlflow/mlflow.git@refs/pull/19654/merge
# mlflow-skinny
pip install git+https://github.com/mlflow/mlflow.git@refs/pull/19654/merge#subdirectory=libs/skinny

%sh curl -LsSf https://raw.githubusercontent.com/mlflow/mlflow/HEAD/dev/install-skinny.sh | sh -s pull/19654/merge

What changes are proposed in this pull request?

Adds support to extract request and response using the Databricks ChatCompletions endpoint. This functionality is already supported for custom judge models.

How is this PR tested?

• Existing unit/integration tests
• New unit/integration tests
• Manual tests

Validation script:

import os

os.environ["MLFLOW_ENABLE_ASYNC_TRACE_LOGGING"] = "false"

import mlflow
from mlflow.entities.span import SpanType
from mlflow.genai.scorers import Safety

mlflow.set_tracking_uri("databricks")
mlflow.set_experiment(f"/Users/{os.environ.get('USER')}@databricks.com/demo_llm_extraction")

# Create trace with empty root span but populated child span
with mlflow.start_span(name="agent_call", span_type=SpanType.CHAIN) as root:
    with mlflow.start_span(name="chat_completion", span_type=SpanType.CHAT_MODEL) as child:
        child.set_inputs({"messages": [{"role": "user", "content": "What is the capital of France?"}]})
        child.set_outputs({"choices": [{"message": {"role": "assistant", "content": "The capital of France is Paris."}}]})

trace = mlflow.get_trace(root.trace_id)
print(f"Root span inputs: {trace.data.spans[0].inputs}, outputs: {trace.data.spans[0].outputs}")

# Call Safety scorer - will use LLM to extract fields from child spans
result = Safety()(trace=trace)
print(f"Score: {result.value}, Rationale: {result.rationale}")

Succeeds on the feature branch:

$ python tmp/scripts/demo_llm_extraction.py
Root span inputs: None, outputs: None
Score: yes, Rationale: The text simply states a factual piece of information about the capital of France, which does not contain any elements of hate speech, harassment, incitement of violence, or promotion of illegal or severely harmful acts. Therefore, it is safe.

Fails on master:

$ python tmp/scripts/demo_llm_extraction.py
Root span inputs: None, outputs: None
2025/12/27 00:46:44 WARNING mlflow.genai.scorers.builtin_scorers: Failed to extract required fields from trace using LLM: Malformed model uri 'databricks'. The URI must be in the format of <provider>:/<model-name>, e.g., 'openai:/gpt-4.1-mini'.
Traceback (most recent call last):
  File "/home/avesh.singh/mlflow/tmp/scripts/demo_llm_extraction.py", line 34, in <module>
    result = Safety()(trace=trace)
  File "/home/avesh.singh/mlflow/mlflow/genai/scorers/base.py", line 149, in wrapper
    return telemetry_wrapped(*args, **kwargs)
  File "/home/avesh.singh/mlflow/mlflow/telemetry/track.py", line 24, in wrapper
    return func(*args, **kwargs)
  File "/home/avesh.singh/mlflow/mlflow/genai/scorers/builtin_scorers.py", line 1386, in __call__
    _validate_required_fields(fields, self, "Safety scorer")
  File "/home/avesh.singh/mlflow/mlflow/genai/scorers/builtin_scorers.py", line 214, in _validate_required_fields
    raise MlflowException(
mlflow.exceptions.MlflowException: Safety scorer requires the following fields: outputs. Provide them directly or pass a trace containing them.

Does this PR require documentation update?

• No. You can skip the rest of this section.
• Yes. I've updated:
  • Examples
  • API references
  • Instructions

Release Notes

Is this a user-facing change?

• No. You can skip the rest of this section.
• Yes. Give a description of this change to be included in the release notes for MLflow users.

Adds Databricks support to automatically extract request and response from traces which do not contain inputs and outputs in the root span. This is already supported for non-Databricks judge models.

What component(s), interfaces, languages, and integrations does this PR affect?

Components

• area/tracking: Tracking Service, tracking client APIs, autologging
• area/models: MLmodel format, model serialization/deserialization, flavors
• area/model-registry: Model Registry service, APIs, and the fluent client calls for Model Registry
• area/scoring: MLflow Model server, model deployment tools, Spark UDFs
• area/evaluation: MLflow model evaluation features, evaluation metrics, and evaluation workflows
• area/gateway: MLflow AI Gateway client APIs, server, and third-party integrations
• area/prompts: MLflow prompt engineering features, prompt templates, and prompt management
• area/tracing: MLflow Tracing features, tracing APIs, and LLM tracing functionality
• area/projects: MLproject format, project running backends
• area/uiux: Front-end, user experience, plotting, JavaScript, JavaScript dev server
• area/build: Build and test infrastructure for MLflow
• area/docs: MLflow documentation pages

How should the PR be classified in the release notes? Choose one:

• rn/none - No description will be included. The PR will be mentioned only by the PR number in the "Small Bugfixes and Documentation Updates" section
• rn/breaking-change - The PR will be mentioned in the "Breaking Changes" section
• rn/feature - A new user-facing feature worth mentioning in the release notes
• rn/bug-fix - A user-facing bug fix worth mentioning in the release notes
• rn/documentation - A user-facing documentation change worth mentioning in the release notes

Should this PR be included in the next patch release?

Yes should be selected for bug fixes, documentation updates, and other small changes. No should be selected for new features and larger changes. If you're unsure about the release classification of this PR, leave this unchecked to let the maintainers decide.

• Yes (this PR will be cherry-picked and included in the next patch release)
• No (this PR will be included in the next minor release)