Use Lakehouse Monitoring for GenAI to monitor your production agent

%pip install -U -qqqq databricks-agents>=0.17.0 databricks-sdk[openai] backoff uv
dbutils.library.restartPython()

%%writefile hello_world_agent.py
from typing import Any, Generator, Optional

import mlflow
from databricks.sdk import WorkspaceClient
from mlflow.entities import SpanType
from mlflow.pyfunc.model import ChatAgent
from mlflow.types.agent import (
    ChatAgentChunk,
    ChatAgentMessage,
    ChatAgentResponse,
    ChatContext,
)

mlflow.openai.autolog()

# Optional: Replace with any model serving endpoint
LLM_ENDPOINT_NAME = "databricks-meta-llama-3-3-70b-instruct"


class SimpleChatAgent(ChatAgent):
    def __init__(self):
        self.workspace_client = WorkspaceClient()
        self.client = self.workspace_client.serving_endpoints.get_open_ai_client()
        self.llm_endpoint = LLM_ENDPOINT_NAME

        # Fake documents to simulate the retriever
        self.documents = [
            mlflow.entities.Document(
                metadata={"doc_uri": "uri1.txt"},
                page_content="""Lakehouse Monitoring for GenAI helps you monitor the quality, cost, and latency of production GenAI apps.  Lakehouse Monitoring for GenAI allows you to:\n- Track quality and operational performance (latency, request volume, errors, etc.).\n- Run LLM-based evaluations on production traffic to detect drift or regressions using Agent Evaluation's LLM judges.\n- Deep dive into individual requests to debug and improve agent responses.\n- Transform real-world logs into evaluation sets to drive continuous improvements.""",
            ),
            # This is a new document about spark.
            mlflow.entities.Document(
                metadata={"doc_uri": "uri2.txt"},
                page_content="The latest spark version in databricks in 3.5.0",
            ),
        ]

        # Tell Agent Evaluation's judges and review app about the schema of your retrieved documents
        mlflow.models.set_retriever_schema(
            name="fake_vector_search",
            primary_key="doc_uri",
            text_column="page_content",
            doc_uri="doc_uri"
            # other_columns=["column1", "column2"],
        )

    @mlflow.trace(span_type=SpanType.RETRIEVER)
    def dummy_retriever(self):
      # Fake retriever
      return self.documents
  

    def prepare_messages_for_llm(
        self, messages: list[ChatAgentMessage]
    ) -> list[dict[str, Any]]:
        """Filter out ChatAgentMessage fields that are not compatible with LLM message formats"""
        compatible_keys = ["role", "content", "name", "tool_calls", "tool_call_id"]
        return [
            {
                k: v
                for k, v in m.model_dump_compat(exclude_none=True).items()
                if k in compatible_keys
            }
            for m in messages
        ]

    @mlflow.trace(span_type=SpanType.PARSER)
    def prepare_rag_prompt(self, messages):

        docs = self.dummy_retriever()

        messages = self.prepare_messages_for_llm(messages)

        messages[-1]['content'] = f"Answer the user's question based on the documents.\nDocuments: <documents>{docs}</documents>.\nUser's question: <user_question>{messages[-1]['content']}</user_question>"

        return messages

    @mlflow.trace(span_type=SpanType.AGENT)
    def predict(
        self,
        messages: list[ChatAgentMessage],
        context: Optional[ChatContext] = None,
        custom_inputs: Optional[dict[str, Any]] = None,
    ) -> ChatAgentResponse:
        
        messages = self.prepare_rag_prompt(messages)

        resp = self.client.chat.completions.create(
            model=self.llm_endpoint,
            messages=messages,
        )

        return ChatAgentResponse(
            messages=[
                ChatAgentMessage(**resp.choices[0].message.to_dict(), id=resp.id)
            ],
        )

    @mlflow.trace(span_type=SpanType.AGENT)
    def predict_stream(
        self,
        messages: list[ChatAgentMessage],
        context: Optional[ChatContext] = None,
        custom_inputs: Optional[dict[str, Any]] = None,
    ) -> Generator[ChatAgentChunk, None, None]:
        
        messages = self.prepare_rag_prompt(messages)

        for chunk in self.client.chat.completions.create(
            model=self.llm_endpoint,
            messages=messages,
            stream=True,
        ):
            if not chunk.choices or not chunk.choices[0].delta.content:
                continue

            yield ChatAgentChunk(
                delta=ChatAgentMessage(
                    **{
                        "role": "assistant",
                        "content": chunk.choices[0].delta.content,
                        "id": chunk.id,
                    }
                )
            )


from mlflow.models import set_model

AGENT = SimpleChatAgent()
set_model(AGENT)

from hello_world_agent import AGENT
AGENT.predict({
        "messages": [{"role": "user", "content": "How do I monitor my genai app?"}]
    })

import mlflow
from mlflow.models.resources import DatabricksServingEndpoint

with mlflow.start_run():
  model_info = mlflow.pyfunc.log_model(
    python_model="hello_world_agent.py",
    artifact_path="agent",
    input_example={
        "messages": [{"role": "user", "content": "How do I monitor my genai app?"}]
    },  
    resources=[DatabricksServingEndpoint(endpoint_name="databricks-meta-llama-3-3-70b-instruct")],
    pip_requirements=["databricks-sdk[openai]", "mlflow", "databricks-agents", "backoff"],    
  )

# Let's validate that the model can be loaded, and try invoking it.
mlflow.models.predict(
    model_uri=model_info.model_uri,
    input_data={"messages": [{"role": "user", "content": "How do I monitor my genai app?"}]},
    env_manager="uv",
)

from databricks.agents import deploy

# Set the name of the model to use in your Unity Catalog schema defined at the top of this notebook

MODEL_NAME = "my_demo_agent"

# Register the model in Unity Catalog and deploy it as a serving endpoint
mlflow.set_registry_uri("databricks-uc")
uc_model_info = mlflow.register_model(
    model_uri=model_info.model_uri, name=f"{UC_CATALOG}.{UC_SCHEMA}.{MODEL_NAME}"
)
deployment = deploy(model_name=uc_model_info.name, model_version=uc_model_info.version)

from databricks.sdk.service.serving import EndpointStateReady, EndpointStateConfigUpdate
from databricks.sdk import WorkspaceClient
import time

print("\nWaiting for endpoint to deploy.  This can take 10 - 20 minutes.", end="")
w = WorkspaceClient()
while w.serving_endpoints.get(deployment.endpoint_name).state.ready == EndpointStateReady.NOT_READY or w.serving_endpoints.get(deployment.endpoint_name).state.config_update == EndpointStateConfigUpdate.IN_PROGRESS:
    print(".", end="")
    time.sleep(30)

print("\nREADY!")

from databricks.agents.evals.monitors import create_monitor, get_monitor, update_monitor, delete_monitor

# Get the current monitor configuration 
monitor = get_monitor(endpoint_name=deployment.endpoint_name)

# Update the monitor to add evaluation metrics
monitor = update_monitor(
    endpoint_name=deployment.endpoint_name,
    monitoring_config={
        "sample": 1,  # Sample 100% of requests - this can be any number from 0 (0%) to 1 (100%).
        # Select 0+ of Agent Evaluation's built-in judges
        "metrics": ['guideline_adherence', 'groundedness', 'safety', 'relevance_to_query', 'chunk_relevance'],
        # Customize these guidelines based on your business requirements.  These guidelines will be analyzed using Agent Evaluation's built in guideline_adherence judge
        "global_guidelines": {
            "english": ["The response must be in English."],
            "clarity": ["The response must be clear, coherent, and concise."],
            "relevant_if_not_refusal": ["Determine if the response provides an answer to the user's request.  A refusal to answer is considered relevant.  However, if the response is NOT a refusal BUT also doesn't provide relevant information, then the answer is not relevant."],
            "no_answer_if_no_docs": ["If the agent can not find a relevant document, it should refuse to answer the question and not discuss the reasons why it could not answer."]
        }
    }
)

from mlflow import deployments

client = deployments.get_deploy_client("databricks")

questions = [
    "What is Mosaic AI Agent Evaluation?",
    "How do you use MLflow with Databricks for experiment tracking?",
    "What should I use Databricks Feature Store for?",
    "How does AutoML work in Databricks?",
    "What is Model Serving in Databricks and what are its deployment options?",
    "How does Databricks handle distributed deep learning training?",
    "Does Unity Catalog support models?",
    "What is the Databricks Lakehouse?",
    "Which Llama models are supported on Databricks?",
    "How does Databricks integrate with popular ML frameworks like PyTorch and TensorFlow?"
]

for i, question in enumerate(questions, 1):
    print(f"\nQuestion {i}: {question}")  
    response = client.predict(
        endpoint=deployment.endpoint_name,
        inputs={
            "messages": [
                {"role": "user", "content": question}
            ]
        }
    )
    print(response)
    

from databricks.agents import review_app

# OPTIONAL: Add a comma separated list of domain experts who will provide feedback/alebsl
# If not provided, only the user running this notebook will be granted access to the review app.
DOMAIN_EXPERT_EMAILS = []

# Get the Review App from the current MLflow Experiment
my_review_app = review_app.get_review_app()

# Optional: Add a custom question for your domain experts
my_review_app.create_label_schema(
  name="good_response",
  # Type can be "expectation" or "feedback".
  type="feedback",
  title="Is this a good response?",
  input=review_app.label_schemas.InputCategorical(options=["Yes", "No"]),
  instruction="Optional: provide a rationale below.",
  enable_comment=True,
  overwrite=True
)

my_session = my_review_app.create_labeling_session(
    name="collect_facts",
    assigned_users=DOMAIN_EXPERT_EMAILS, # If not provided, only the user running this notebook will be granted access
    # Built-in labeling schemas: EXPECTED_FACTS, GUIDELINES, EXPECTED_RESPONSE
    label_schemas=[review_app.label_schemas.GUIDELINES,  "good_response"],
)

# URLs to share with the SME.
print("Review App URL:", my_review_app.url)
print("Labeling session URL: ", my_session.url)