Add ABAC Agent into uc-quickstart/utils #611
Conversation
There was a problem hiding this comment.
Pull request overview
This PR introduces an ABAC (Attribute-Based Access Control) Policy Assistant agent into the uc-quickstart/utils directory. The agent is designed to analyze Unity Catalog tables and generate ABAC policy recommendations using AI-powered analysis.
Key changes:
- Adds a complete Streamlit-based chat application for ABAC policy generation
- Implements an MLflow ResponsesAgent with Unity Catalog function tools
- Provides comprehensive documentation and deployment configuration
Reviewed changes
Copilot reviewed 7 out of 7 changed files in this pull request and generated 10 comments.
Show a summary per file
| File | Description |
|---|---|
| requirements.txt | Defines Python dependencies for the application (mlflow, streamlit) |
| model_serving_utils.py | Provides utility functions for querying serving endpoints and handling different agent response formats |
| messages.py | Defines message classes for the chat interface with rendering logic |
| driver.py | Databricks notebook implementing the agent with MLflow logging and deployment |
| app.yaml | Configuration file for Databricks Apps deployment |
| app.py | Main Streamlit application implementing the chat interface |
| README.md | Comprehensive documentation covering features, setup, and usage |
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
| # MAGIC | ||
| # MAGIC Find additional examples in the documentation - https://docs.databricks.com/aws/en/data-governance/unity-catalog/abac/policies?language=SQL | ||
| # MAGIC | ||
| # MAGIC Usually table name is given as catalog_name.schem_name.table_name. |
There was a problem hiding this comment.
Corrected spelling of 'schem_name' to 'schema_name'.
| # MAGIC Usually table name is given as catalog_name.schem_name.table_name. | |
| # MAGIC Usually table name is given as catalog_name.schema_name.table_name. |
| scorers=[RelevanceToQuery(), Safety()], # add more scorers here if they're applicable | ||
| ) | ||
|
|
||
| # Review the evaluation results in the MLfLow UI (see console output) |
There was a problem hiding this comment.
Corrected spelling of 'MLfLow' to 'MLflow'.
| # Review the evaluation results in the MLfLow UI (see console output) | |
| # Review the evaluation results in the MLflow UI (see console output) |
| secret_scope = 'david_scope' | ||
| client_secret_key = 'DATABRICKS_CLIENT_SECRET' | ||
| client_id_key = 'DATABRICKS_CLIENT_ID' |
There was a problem hiding this comment.
The hardcoded secret scope name 'david_scope' appears to be user-specific and should not be committed to the repository. This should be parameterized or documented as requiring user configuration.
| deployment_info = agents.deploy( | ||
| UC_MODEL_NAME, | ||
| uc_registered_model_info.version, | ||
| environment_vars={ | ||
| "DATABRICKS_HOST": "https://dbc-a612b3a4-f0ff.cloud.databricks.com", |
There was a problem hiding this comment.
The hardcoded Databricks host URL appears to be workspace-specific and should not be committed. This should be parameterized or retrieved from the environment/workspace context.
| deployment_info = agents.deploy( | |
| UC_MODEL_NAME, | |
| uc_registered_model_info.version, | |
| environment_vars={ | |
| "DATABRICKS_HOST": "https://dbc-a612b3a4-f0ff.cloud.databricks.com", | |
| workspace_host = dbutils.notebook.entry_point.getDbutils().notebook().getContext().apiUrl().get() | |
| deployment_info = agents.deploy( | |
| UC_MODEL_NAME, | |
| uc_registered_model_info.version, | |
| environment_vars={ | |
| "DATABRICKS_HOST": workspace_host, |
| 1. **Clone the repository** | ||
| ```bash | ||
| git clone <repository-url> | ||
| cd e2e-chatbot-app |
There was a problem hiding this comment.
The directory name 'e2e-chatbot-app' in the installation instructions doesn't match the actual directory structure 'uc-quickstart/utils/abac-agent'. This should be updated to reflect the correct path.
| cd e2e-chatbot-app | |
| cd uc-quickstart/utils/abac-agent |
| except: | ||
| pass |
There was a problem hiding this comment.
Bare except clause catches all exceptions including system exits. Specify the expected exception type (e.g., except json.JSONDecodeError:) or at minimum use except Exception:.
| except: | ||
| # If not JSON, show as code |
There was a problem hiding this comment.
Bare except clause catches all exceptions including system exits. Specify the expected exception type (e.g., except json.JSONDecodeError:) or at minimum use except Exception:.
| messages=[{"role": "assistant", "content": accumulated_content}], | ||
| request_id=request_id | ||
| ) | ||
| except Exception: |
There was a problem hiding this comment.
The error message doesn't provide any information about what went wrong. Consider logging the exception details and providing a more descriptive error message to help with debugging.
| except Exception: | |
| except Exception as e: | |
| logging.exception("Error during streaming query to serving endpoint") |
| except Exception: | ||
| response_area.markdown("_Ran into an error. Retrying without streaming..._") |
There was a problem hiding this comment.
The error message doesn't provide any information about what went wrong. Consider logging the exception details and providing a more descriptive error message to help with debugging.
| except Exception: | |
| response_area.markdown("_Ran into an error. Retrying without streaming..._") | |
| except Exception as exc: | |
| logging.exception( | |
| "Error while streaming response from endpoint '%s'. Falling back to non-streaming query.", | |
| SERVING_ENDPOINT, | |
| ) | |
| response_area.markdown( | |
| "_Ran into an error while streaming from the endpoint. " | |
| "Retrying without streaming (see logs for details)..._" | |
| ) |
| except Exception: | ||
| response_area.markdown("_Ran into an error. Retrying without streaming..._") |
There was a problem hiding this comment.
The error message doesn't provide any information about what went wrong. Consider logging the exception details and providing a more descriptive error message to help with debugging.
| except Exception: | |
| response_area.markdown("_Ran into an error. Retrying without streaming..._") | |
| except Exception as e: | |
| logging.exception("Error during streaming response in query_responses_endpoint_and_render") | |
| response_area.markdown("_Ran into an error while streaming the response. Retrying without streaming..._") |
|
@dlevy324 please fix Copilot suggestions |
2 participants
No description provided.