External monitor setup

%pip install -U "databricks-agents>=0.18.1" "mlflow>=2.21.2" "databricks-sdk[openai]"
dbutils.library.restartPython()

import mlflow

# Get the current user name
user_email = spark.sql("SELECT current_user() as username").collect()[0].username
user = user_email.split("@")[0].replace(".", "").lower()[:35]

# For convinience, we auto-infer these from the current notebook.
# TODO: In prod and outside notebooks, set these explicitly.

MLFLOW_EXPERIMENT_NAME = "my_app_monitor"
MLFLOW_EXPERIMENT_PATH = f"/Users/{user_email}/{MLFLOW_EXPERIMENT_NAME}"
mlflow.set_experiment(MLFLOW_EXPERIMENT_PATH)
MLFLOW_EXPERIMENT_ID = mlflow.tracking.fluent._get_experiment_id() 

# Get the current Databricks workspace URL
DATABRICKS_HOST = mlflow.utils.databricks_utils.get_workspace_url() # https://<workspace-url>.databricks.com

print(f"MLFLOW_EXPERIMENT_ID: {MLFLOW_EXPERIMENT_ID}")
print(f"MLFLOW_EXPERIMENT_PATH: {MLFLOW_EXPERIMENT_PATH}")
print(f"DATABRICKS_HOST: {DATABRICKS_HOST}")

print(f"\nView the monitoring UI in the MLflow experiment here: {DATABRICKS_HOST}/ml/experiments/{MLFLOW_EXPERIMENT_ID}/evaluation-monitoring")

# Get the workspace default UC catalog / schema
uc_default_location = spark.sql("select current_catalog() as current_catalog, current_schema() as current_schema").collect()[0]
current_catalog = uc_default_location["current_catalog"]
current_schema = uc_default_location["current_schema"]


# Modify the UC catalog / schema here or at the top of the notebook in the widget editor
dbutils.widgets.text("uc_catalog", current_catalog)
dbutils.widgets.text("uc_schema", current_schema)
UC_CATALOG = dbutils.widgets.get("uc_catalog")
UC_SCHEMA = dbutils.widgets.get("uc_schema")

from databricks.agents.monitoring import create_external_monitor, AssessmentsSuiteConfig, BuiltinJudge, GuidelinesJudge, get_external_monitor
from requests.exceptions import HTTPError

try:
  external_monitor = create_external_monitor(
    catalog_name=UC_CATALOG,
    schema_name=UC_SCHEMA,
    experiment_id=MLFLOW_EXPERIMENT_ID,
    assessments_config=AssessmentsSuiteConfig(
      # The % of traces on which the `assessments` are run; a number between 0 and 1.
      sample=1.0,
      assessments=[
        # Builtin judges: "safety", "groundedness", "relevance_to_query", "chunk_relevance"
        BuiltinJudge(name='safety'),
        BuiltinJudge(name='groundedness'),
        BuiltinJudge(name='relevance_to_query'),
        BuiltinJudge(name='chunk_relevance'),
        # Create custom LLM judges with the guidelines judge.
        GuidelinesJudge(guidelines={
          "pii": ["The response must not contain personal information."],
          "english": ["The response must be in English"]
        }),
      ]
    ))
except HTTPError as e:
    if e.response.status_code == 409 and "ALREADY_EXISTS" in e.response.text:
        print("A monitor for this experiment already exists. Retrieving that monitor...\n\n")
        external_monitor = get_external_monitor(experiment_id=MLFLOW_EXPERIMENT_ID)
        print(external_monitor)
    else:
        raise

import mlflow
from mlflow.tracing.destination import Databricks
import os

# In your production app, set these environment variables to enable MLflow Tracing to connect to your monitor.
# os.environ["DATABRICKS_HOST"] = DATABRICKS_HOST  # Your workspace's URL e.g., https://<workspace-url>.databricks.com
# os.environ["DATABRICKS_TOKEN"] = DATABRICKS_TOKEN  # A PAT token

# Use set_destination to configure MLflow tracing to log to your monitor.
mlflow.tracing.set_destination(Databricks(experiment_id=MLFLOW_EXPERIMENT_ID))

# Here, we use the databricks-sdk's convenience method to get an OpenAI client authenticated to your workspace.  For details of this convenience method, see https://docs.databricks.com/aws/en/machine-learning/model-serving/score-foundation-models#install-packages
from databricks.sdk import WorkspaceClient

w = WorkspaceClient()
openai_client = w.serving_endpoints.get_open_ai_client()

# In your app, you can use any OpenAI client.
# import openai
# openai_client = openai.OpenAI(api_key=os.environ.get("OPENAI_API_KEY"))


# Instrument your Agent using MLflow auto-logging or manual instrumentation.  Here we use a combination of auto-logging and manual instrumentation.
mlflow.openai.autolog()

# Manual instrumentation
@mlflow.trace(span_type="AGENT")
def openai_agent(user_input: str):
    return openai_client.chat.completions.create(
        model="databricks-meta-llama-3-3-70b-instruct",
        messages=[
            {
                "role": "system",
                "content": "You are a helpful assistant that always responds in CAPS!",
            },
            {"role": "user", "content": user_input},
        ],
    )

openai_agent("What is GenAI observability?")