API Reference

This page provides a comprehensive index of MLflow APIs used in GenAI applications, with direct links to the official MLflow documentation.

MLflow features marked as "Databricks only" are only available on Databricks-managed MLflow.

Official documentation links

• All MLflow Python APIs
• MLflow Core APIs
• MLflow GenAI Module
• MLflow Tracing APIs
• MLflow Client APIs
• MLflow Entities

Beta and Experimental Features

Some of the APIs referenced on this page are currently in the Beta or Experimental stages. These APIs are subject to change or removal in future releases. Experimental APIs are available to all customers, and Beta APIs are available to most customers automatically. If you do not have access to a Beta API and need to request access, contact your Databricks support representative.

Experiment management

Manage MLflow experiments and runs for tracking GenAI application development:

SDKs

• mlflow.search_runs() - Search and filter runs by criteria
• mlflow.set_experiment() - Set the active MLflow experiment
• mlflow.start_run() - Start a new MLflow run for tracking

Entities

• mlflow.entities.Experiment - Experiment metadata and configuration
• mlflow.entities.Run - Run metadata, metrics, and parameters

Prompt management

Version control and lifecycle management for prompts used in GenAI applications:

SDKs

• mlflow.genai.load_prompt() - Load a versioned prompt from the registry
• mlflow.genai.optimize_prompt() - Automatically improve prompts using optimization algorithms
• mlflow.genai.register_prompt() - Register a new prompt to the registry
• mlflow.genai.search_prompts() - Search for prompts by name or tags
• mlflow.genai.delete_prompt_alias() - Remove an alias from a prompt version
• mlflow.genai.set_prompt_alias() - Assign an alias to a prompt version

Entities

• mlflow.entities.Prompt - Prompt metadata and version information

Evaluation and monitoring

Scorer lifecycle management (Databricks only)

Beta

This feature is in Beta.

Scorer lifecycle management for continuous quality tracking in production:

Scorer instance methods

• Scorer.register() - Register custom scorer with server
• Scorer.start() - Begin online evaluation with sampling
• Scorer.update() - Modify sampling configuration
• Scorer.stop() - Stop online evaluation
• Scorer.delete() - Remove scorer entirely

Scorer registry functions

• mlflow.genai.scorers.get_scorer() - Retrieve registered scorer by name
• mlflow.genai.scorers.list_scorers() - List all registered scorers
• mlflow.genai.scorers.delete_scorer() - Delete registered scorers by name

Scorer properties

• Scorer.sample_rate - Current sampling rate (0.0-1.0)
• Scorer.filter_string - Current trace filter

Configuration classes

• mlflow.genai.ScorerSamplingConfig - Sampling configuration data class

Core evaluation APIs

Core APIs for offline evaluation and production monitoring:

• mlflow.genai.evaluate() - Evaluation harness to orchestrate offline evaluation with scorers and datasets
• mlflow.genai.to_predict_fn() - Convert model output to standardized prediction function format
• mlflow.genai.Scorer - Custom scorer class for object-oriented implementation with state management
• mlflow.genai.scorer() - Scorer decorator for scorer creation and evaluation logic

Predefined scorers

Quality assessment scorers ready for immediate use:

• mlflow.genai.scorers.Safety - Content safety evaluation
• mlflow.genai.scorers.Correctness - Answer accuracy assessment
• mlflow.genai.scorers.RelevanceToQuery - Query relevance scoring
• mlflow.genai.scorers.Guidelines - Custom guideline compliance
• mlflow.genai.scorers.ExpectationsGuidelines - Guideline evaluation with expectations
• mlflow.genai.scorers.RetrievalGroundedness - RAG grounding assessment
• mlflow.genai.scorers.RetrievalRelevance - Retrieved context relevance
• mlflow.genai.scorers.RetrievalSufficiency - Context sufficiency evaluation

Scorer helpers

• mlflow.genai.scorers.get_all_scorers() - Retrieve all built-in scorers

Judge functions

LLM-based assessment functions for direct use or scorer wrapping:

• mlflow.genai.judges.is_safe() - Safety assessment
• mlflow.genai.judges.is_correct() - Correctness evaluation
• mlflow.genai.judges.is_grounded() - Grounding verification
• mlflow.genai.judges.is_context_relevant() - Context relevance
• mlflow.genai.judges.is_context_sufficient() - Context sufficiency
• mlflow.genai.judges.meets_guidelines() - Custom guideline assessment
• mlflow.genai.make_judge() - Create custom judges (recommended for MLflow 3.4.0 and above)
• mlflow.genai.judges.custom_prompt_judge() - Custom prompt-based evaluation (deprecated in MLflow 3.4.0, use make_judge() instead)

Judge output entities

• mlflow.genai.judges.CategoricalRating - Enum for categorical judge responses
  • mlflow.genai.judges.CategoricalRating.YES - Positive rating
  • mlflow.genai.judges.CategoricalRating.NO - Negative rating
  • mlflow.genai.judges.CategoricalRating.UNKNOWN - Uncertain rating

Evaluation datasets

Create and manage versioned test datasets for systematic evaluation:

SDKs

• mlflow.genai.create_dataset() - Create a new evaluation dataset
• mlflow.genai.delete_dataset() - Delete an evaluation dataset
• mlflow.genai.get_dataset() - Retrieve an existing evaluation dataset

Entities

• mlflow.genai.datasets.EvaluationDataset - Versioned test data container
  • merge_records() - Combine records from multiple sources
  • set_profile() - Configure dataset profile settings
  • to_df() - Convert dataset to pandas DataFrame
  • to_evaluation_dataset() - Convert to evaluation dataset format

Human labeling and review app (Databricks only)

Human feedback collection and review workflows for systematic quality assessment:

Entities

• mlflow.genai.Agent - Agent configuration for review app testing

• mlflow.genai.LabelingSession - Human labeling workflow manager

  • add_dataset() - Add evaluation dataset to labeling session
  • add_traces() - Add traces for human review
  • set_assigned_users() - Assign reviewers to session
  • sync() - Synchronize session state

• mlflow.genai.ReviewApp - Interactive review application

  • add_agent() - Add agent for testing
  • remove_agent() - Remove agent from review app

Labeling session SDKs

• mlflow.genai.create_labeling_session() - Create a new labeling session
• mlflow.genai.delete_labeling_session() - Delete a labeling session
• mlflow.genai.get_labeling_session() - Retrieve labeling session by ID
• mlflow.genai.get_labeling_sessions() - List all labeling sessions
• mlflow.genai.get_review_app() - Retrieve review app instance

Label schema types

• mlflow.genai.label_schemas.InputCategorical - Categorical input field type
• mlflow.genai.label_schemas.InputCategoricalList - Multi-select categorical input
• mlflow.genai.label_schemas.InputNumeric - Numeric input field type
• mlflow.genai.label_schemas.InputText - Text input field type
• mlflow.genai.label_schemas.InputTextList - Multi-text input field type
• mlflow.genai.label_schemas.LabelSchema - Label schema definition
• mlflow.genai.label_schemas.LabelSchemaType - Schema type enum
  • mlflow.genai.label_schemas.LabelSchemaType.EXPECTATION - Expectation schema type
  • mlflow.genai.label_schemas.LabelSchemaType.FEEDBACK - Feedback schema type

Label schema SDKs

• mlflow.genai.label_schemas.create_label_schema() - Create a new label schema
• mlflow.genai.label_schemas.delete_label_schema() - Delete an existing label schema
• mlflow.genai.label_schemas.get_label_schema() - Retrieve label schema by name

Prompt optimization

Beta

This feature is in Beta.

Automated prompt improvement using data-driven optimization algorithms:

Entities

• mlflow.genai.optimize.LLMParams - LLM configuration parameters
• mlflow.genai.optimize.OptimizerConfig - Optimization algorithm configuration
• mlflow.genai.optimize.PromptOptimizationResult - Optimization results and metrics

SDKs

• mlflow.genai.optimize.optimize_prompt() - Run prompt optimization process

Tracing

Instrument and capture execution traces from GenAI applications:

SDKs

• mlflow.delete_trace_tag() - Remove a tag from a trace
• mlflow.get_current_active_span() - Get the currently active span
• mlflow.get_last_active_trace() - Retrieve the most recently completed trace
• mlflow.get_last_active_trace_id() - Get ID of the last active trace
• mlflow.get_trace() - Retrieve a trace by ID
• mlflow.search_traces() - Search and filter traces
• mlflow.set_trace_tag() - Add a tag to a trace
• mlflow.start_span() - Manually start a new span
• mlflow.trace - Decorator to automatically trace function execution
• mlflow.traceName - Context manager to set trace name
• mlflow.traceOutputs - Context manager to set trace outputs
• mlflow.tracing - Tracing module with configuration functions
• mlflow.tracing.disable - Disable tracing globally
• mlflow.tracing.disable_notebook_display() - Disable trace display in notebooks
• mlflow.tracing.enable - Enable tracing globally
• mlflow.tracing.enable_notebook_display() - Enable trace display in notebooks
• mlflow.update_current_trace() - Update metadata for the current trace

Entities

• mlflow.entities.Trace - Complete trace with all spans and metadata
• mlflow.entities.TraceData - Trace execution data
• mlflow.entities.TraceInfo - Trace metadata and summary information
• mlflow.entities.Span - Individual span within a trace
• mlflow.entities.SpanEvent - Event occurring within a span
• mlflow.entities.SpanType - Span type classification enum
• mlflow.entities.Document - Document entity for RAG applications

Assessment entities

Data structures for storing evaluation results and feedback:

• mlflow.entities.Assessment - Evaluation result container
• mlflow.entities.AssessmentError - Assessment error details
• mlflow.entities.AssessmentSource - Source of the assessment
• mlflow.entities.AssessmentSourceType - Assessment source type enum
• mlflow.entities.Expectation - Expected ground truth outcome
• mlflow.entities.Feedback - Scorer output with value and rationale

Tracing integrations

Auto-instrumentation for popular GenAI frameworks and libraries:

• mlflow.anthropic.autolog - Anthropic Claude integration
• mlflow.autogen.autolog - Microsoft AutoGen integration
• mlflow.bedrock.autolog - AWS Bedrock integration
• mlflow.crewai.autolog - CrewAI integration
• mlflow.dspy.autolog - DSPy integration
• mlflow.gemini.autolog - Google Gemini integration
• mlflow.groq.autolog - Groq integration
• mlflow.langchain.autolog - LangChain integration
• mlflow.litellm.autolog - LiteLLM integration
• mlflow.llama_index.autolog - LlamaIndex integration
• mlflow.mistral.autolog - Mistral AI integration
• mlflow.openai.autolog - OpenAI integration

Version tracking

Track and manage GenAI application versions in production:

SDKs

• mlflow.set_active_model() - Set the active model for version tracking
• mlflow.clear_active_model() - Clear the active model context
• mlflow.get_active_model_id() - Get the current active model ID
• mlflow.create_external_model() - Register an external model deployment
• mlflow.delete_logged_model_tag() - Remove a tag from logged model
• mlflow.finalize_logged_model() - Finalize a logged model
• mlflow.get_logged_model() - Retrieve logged model by ID
• mlflow.initialize_logged_model() - Initialize a new logged model
• mlflow.last_logged_model() - Get the most recently logged model
• mlflow.search_logged_models() - Search for logged models
• mlflow.set_logged_model_tags() - Add tags to logged model
• mlflow.log_model_params() - Log parameters for a model

Entities

• mlflow.entities.LoggedModel - Logged model metadata and information
• mlflow.entities.LoggedModelStatus - Logged model status enum
• mlflow.ActiveModel - Active model context manager