import os import mlflow from rich import print from guardrails import Guard from guardrails.hub import RestrictToTopic, ValidLength from guardrails.integrations.databricks import MlFlowInstrumentor # Setup some environment variables for the LLM os.environ["DATABRICKS_API_KEY"] = os.environ.get("DATABRICKS_TOKEN", "your-databricks-key") os.environ["DATABRICKS_API_BASE"] = os.environ.get("DATABRICKS_HOST", "https://abc-123ab12a-1234.cloud.databricks.com") # Run this in the terminal # ! mlflow server --host localhost --port 8080 mlflow.set_tracking_uri(uri="http://localhost:8080") MlFlowInstrumentor(experiment_name="MLFlow x Guardrails Async").instrument() guard = Guard(name='content-guard').use_many( RestrictToTopic(valid_topics=["computer programming", "computer science", "algorithms"], disable_llm=True, on_fail="exception"), ValidLength(min=1, max=150, on_fail="exception") ) instructions = { "role": "system", "content": "You are a helpful assistant that gives advice about writing clean code and other programming practices." } prompt = "Write a short summary about recursion in less than 100 characters." try: response = guard( model="databricks/databricks-dbrx-instruct", messages=[instructions, { "role":"user", "content": prompt }], ) print(response) except Exception as e: print(e) # Go to http://localhost:8080 # Select the MLFlow x Guardrails Async experiment # Go to Traces tab and see that there are 3 traces not just one, the latter two of which are from the validators and are unattached from the first trace which holds the Guard hierarchy.