Source code for mlflow.genai.label_schemas

"""
Databricks Agent Label Schemas Python SDK. For more details see Databricks Agent Evaluation:
<https://docs.databricks.com/en/generative-ai/agent-evaluation/index.html>

The API docs can be found here:
<https://api-docs.databricks.com/python/databricks-agents/latest/databricks_agent_eval.html#review-app>
"""

from typing import TYPE_CHECKING, Literal

from mlflow.genai.label_schemas.label_schemas import (
    InputCategorical,
    InputCategoricalList,
    InputNumeric,
    InputText,
    InputTextList,
    LabelSchema,
    LabelSchemaType,
)
from mlflow.genai.labeling import ReviewApp

if TYPE_CHECKING:
    from databricks.agents.review_app import ReviewApp

EXPECTED_FACTS = "expected_facts"
GUIDELINES = "guidelines"
EXPECTED_RESPONSE = "expected_response"


[docs]def create_label_schema( name: str, *, type: Literal["feedback", "expectation"], title: str, input: InputCategorical | InputCategoricalList | InputText | InputTextList | InputNumeric, instruction: str | None = None, enable_comment: bool = False, overwrite: bool = False, ) -> LabelSchema: """Create a new label schema for the review app. A label schema defines the type of input that stakeholders will provide when labeling items in the review app. .. note:: This functionality is only available in Databricks. Please run `pip install mlflow[databricks]` to use it. Args: name: The name of the label schema. Must be unique across the review app. type: The type of the label schema. Either "feedback" or "expectation". title: The title of the label schema shown to stakeholders. input: The input type of the label schema. instruction: Optional. The instruction shown to stakeholders. enable_comment: Optional. Whether to enable comments for the label schema. overwrite: Optional. Whether to overwrite the existing label schema with the same name. Returns: LabelSchema: The created label schema. """ from mlflow.genai.labeling.stores import _get_labeling_store # Nested to avoid circular import store = _get_labeling_store() return store.create_label_schema( name=name, type=type, title=title, input=input, instruction=instruction, enable_comment=enable_comment, overwrite=overwrite, )
[docs]def get_label_schema(name: str) -> LabelSchema: """Get a label schema from the review app. .. note:: This functionality is only available in Databricks. Please run `pip install mlflow[databricks]` to use it. Args: name: The name of the label schema to get. Returns: LabelSchema: The label schema. """ from mlflow.genai.labeling.stores import _get_labeling_store # Nested to avoid circular import store = _get_labeling_store() return store.get_label_schema(name)
[docs]def delete_label_schema(name: str): """Delete a label schema from the review app. .. note:: This functionality is only available in Databricks. Please run `pip install mlflow[databricks]` to use it. Args: name: The name of the label schema to delete. """ # Nested to avoid circular import from mlflow.genai.labeling.databricks_utils import get_databricks_review_app from mlflow.genai.labeling.stores import DatabricksLabelingStore, _get_labeling_store store = _get_labeling_store() store.delete_label_schema(name) # For backwards compatibility, return a ReviewApp instance only if using Databricks store if isinstance(store, DatabricksLabelingStore): return ReviewApp(get_databricks_review_app()) else: # For non-Databricks stores, we can't return a meaningful ReviewApp return None
__all__ = [ "EXPECTED_FACTS", "GUIDELINES", "EXPECTED_RESPONSE", "LabelSchemaType", "LabelSchema", "InputCategorical", "InputCategoricalList", "InputNumeric", "InputText", "InputTextList", "create_label_schema", "get_label_schema", "delete_label_schema", ]