Collibra AI model Python integration

An AI Model is a statistical model that is trained on datasets in order to make predictions on new data. You can integrate AI/ML models in Collibra or manually create AI Model assets and link them to the relevant AI use cases.

To help you transfer the metadata related to your AI models to Collibra, adapt and use the following code in your integration applications. We use Python and the Collibra REST Import API for this example.

Prerequisites

  • You have Python version 3.
  • You have installed the requests module.
  • You have created a programmatic user for API requests that has at least the AI Business User global role.
  • You have obtained the ID of the AI Governance domain.

    The ID of the AI Governance domain should be 00000000-0000-0000-0000-000000006023.

Code sample

import json
import requests

# Update the following variables with your own values
DOMAIN_UUID = "YOUR_DOMAIN_UUID"
USERNAME = "YOUR_USERNAME"
PASSWORD = "YOUR_PASSWORD"

# COLLIBRA / API variables
COLLIBRA_API_URL = "YOUR_COLLIBRA_ENVIRONMENT_URL"
MODEL_ASSET_TYPE = "00000000-0000-0000-0000-000000031402"
MODEL_ACCURACY_ATTRIBUTE_ID = "00000000-0000-0000-0000-000000000328"
MODEL_PRECISION_ATTRIBUTE_ID = "00000000-0000-0000-0000-000000000329"
MODEL_TYPE_ATTRIBUTE_ID = "00000000-0000-0000-0000-000000000334"

class CollibraAIGov:
  def __init__(self):
    # Unique identifier for the domain where all AI Model assets are stored
    self.domain = DOMAIN_UUID
    # Base URL for the Collibra API
    self.COLLIBRA_API_URL = COLLIBRA_API_URL
    # Retrieve password from secrets manager
    self.PASSWORD = PASSWORD
    # Username for API authentication
    self.USERNAME = USERNAME

  def model_update_creation(
    self,
    model_name,
    model_precision="0",
    model_accuracy="0",
    model_type="Generative AI",
  ):
    """
    Create a model and add all relevant information in the AI Governance Module.

    Parameters:
    - model_name: Name of the model to be created.
    - model_precision: Precision metric of the model (optional).
    - model_accuracy: Accuracy metric of the model (optional).
    - model_type: Type of the model (optional). Allowed values are "Generative AI", "Classification", "Regression", "Computer Vision", "Reinforcement Learning" or "Image Classification"
    """

    def builder(model_name):
      """
      Constructs the query payload for model creation.

      Parameters:
      - model_name: Name of the model to be created.

      Returns:
      - A dictionary representing the query payload.
      """
      query = {
        "attributes": {
          MODEL_ACCURACY_ATTRIBUTE_ID: [{ "value": model_accuracy }],
          MODEL_PRECISION_ATTRIBUTE_ID: [{ "value": model_precision }],
          MODEL_TYPE_ATTRIBUTE_ID: [{ "value": model_type }],
        },
        "displayName": model_name,
        "identifier": {
          "domain": {
            "id": self.domain,
          },
          "name": model_name,
        },
        "relations": {},
        "resourceType": "Asset",
        "type": {
          "id": MODEL_ASSET_TYPE,
        },
      }
      return query

    # Serialize the query payload to JSON
    query_payload = json.dumps([builder(model_name)])

    # Construct the file object to be sent in the request
    file = {"file": ("data.json", query_payload)}

    # Make a POST request to the Collibra import API
	# Increase the timeout if the request takes longer
    response = requests.post(
      f"{self.COLLIBRA_API_URL}/rest/2.0/import/json-job",
      files=file,
      auth=(self.USERNAME, self.PASSWORD),
      timeout=30
    )

    # Print the response for debugging purposes
    # Consider handling the response more gracefully in a real application
    print(response)

# Example usage
# collibra = CollibraAIGov()
# collibra.model_update_creation("Example 1 Model")
# collibra.model_update_creation("Example 2 Model", "0.95", "0.90", "Generative AI")

Metadata from your ML models are ingested as assets in Collibra. All AI Model assets, and the assets of child asset types, appear in the Assets drop-down list.

If you need to add more information that is important for your model registry, contact your Collibra administrator to expand the number of characteristics that you can add to the model asset type.

Additional resources