{
  "cells": [
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "7d9bbf86da5e"
      },
      "outputs": [],
      "source": [
        "# Copyright 2023 Google LLC\n",
        "#\n",
        "# Licensed under the Apache License, Version 2.0 (the \"License\");\n",
        "# you may not use this file except in compliance with the License.\n",
        "# You may obtain a copy of the License at\n",
        "#\n",
        "#     https://www.apache.org/licenses/LICENSE-2.0\n",
        "#\n",
        "# Unless required by applicable law or agreed to in writing, software\n",
        "# distributed under the License is distributed on an \"AS IS\" BASIS,\n",
        "# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n",
        "# See the License for the specific language governing permissions and\n",
        "# limitations under the License."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "e95c369b3b9d"
      },
      "source": [
        "# Log Anomaly Detection & Investigation with Text Embeddings + BigQuery Vector Search\n",
        "\n",
        "\n",
        "<table align=\"left\">\n",
        "  <td style=\"text-align: center\">\n",
        "    <a href=\"https://colab.research.google.com/github/GoogleCloudPlatform/generative-ai/blob/main/embeddings/use-cases/outlier-detection/bq-vector-search-outlier-detection-audit-logs.ipynb\">\n",
        "      <img width=\"32px\" src=\"https://www.gstatic.com/pantheon/images/bigquery/welcome_page/colab-logo.svg\" alt=\"Google Colaboratory logo\"><br> Open in Colab\n",
        "    </a>\n",
        "  </td>\n",
        "  <td style=\"text-align: center\">\n",
        "    <a href=\"https://console.cloud.google.com/vertex-ai/colab/import/https:%2F%2Fraw.githubusercontent.com%2FGoogleCloudPlatform%2Fgenerative-ai%2Fmain%2Fembeddings%2Fuse-cases%2Foutlier-detection%2Fbq-vector-search-outlier-detection-audit-logs.ipynb\">\n",
        "      <img width=\"32px\" src=\"https://lh3.googleusercontent.com/JmcxdQi-qOpctIvWKgPtrzZdJJK-J3sWE1RsfjZNwshCFgE_9fULcNpuXYTilIR2hjwN\" alt=\"Google Cloud Colab Enterprise logo\"><br> Open in Colab Enterprise\n",
        "    </a>\n",
        "  </td>\n",
        "  <td style=\"text-align: center\">\n",
        "    <a href=\"https://console.cloud.google.com/vertex-ai/workbench/deploy-notebook?download_url=https://raw.githubusercontent.com/GoogleCloudPlatform/generative-ai/main/embeddings/use-cases/outlier-detection/bq-vector-search-outlier-detection-audit-logs.ipynb\">\n",
        "      <img src=\"https://www.gstatic.com/images/branding/gcpiconscolors/vertexai/v1/32px.svg\" alt=\"Vertex AI logo\"><br> Open in Vertex AI Workbench\n",
        "    </a>\n",
        "  </td>\n",
        "  <td style=\"text-align: center\">\n",
        "    <a href=\"https://console.cloud.google.com/bigquery/import?url=https://github.com/GoogleCloudPlatform/generative-ai/blob/main/embeddings/use-cases/outlier-detection/bq-vector-search-outlier-detection-audit-logs.ipynb\">\n",
        "      <img src=\"https://www.gstatic.com/images/branding/gcpiconscolors/bigquery/v1/32px.svg\" alt=\"BigQuery Studio logo\"><br> Open in BigQuery Studio\n",
        "    </a>\n",
        "  </td>\n",
        "  <td style=\"text-align: center\">\n",
        "    <a href=\"https://github.com/GoogleCloudPlatform/generative-ai/blob/main/embeddings/use-cases/outlier-detection/bq-vector-search-outlier-detection-audit-logs.ipynb\">\n",
        "      <img width=\"32px\" src=\"https://raw.githubusercontent.com/primer/octicons/refs/heads/main/icons/mark-github-24.svg\" alt=\"GitHub logo\"><br> View on GitHub\n",
        "    </a>\n",
        "  </td>\n",
        "</table>\n",
        "\n",
        "<div style=\"clear: both;\"></div>\n",
        "\n",
        "<b>Share to:</b>\n",
        "\n",
        "<a href=\"https://www.linkedin.com/sharing/share-offsite/?url=https%3A//github.com/GoogleCloudPlatform/generative-ai/blob/main/embeddings/use-cases/outlier-detection/bq-vector-search-outlier-detection-audit-logs.ipynb\" target=\"_blank\">\n",
        "  <img width=\"20px\" src=\"https://upload.wikimedia.org/wikipedia/commons/8/81/LinkedIn_icon.svg\" alt=\"LinkedIn logo\">\n",
        "</a>\n",
        "\n",
        "<a href=\"https://bsky.app/intent/compose?text=https%3A//github.com/GoogleCloudPlatform/generative-ai/blob/main/embeddings/use-cases/outlier-detection/bq-vector-search-outlier-detection-audit-logs.ipynb\" target=\"_blank\">\n",
        "  <img width=\"20px\" src=\"https://upload.wikimedia.org/wikipedia/commons/7/7a/Bluesky_Logo.svg\" alt=\"Bluesky logo\">\n",
        "</a>\n",
        "\n",
        "<a href=\"https://twitter.com/intent/tweet?url=https%3A//github.com/GoogleCloudPlatform/generative-ai/blob/main/embeddings/use-cases/outlier-detection/bq-vector-search-outlier-detection-audit-logs.ipynb\" target=\"_blank\">\n",
        "  <img width=\"20px\" src=\"https://upload.wikimedia.org/wikipedia/commons/5/5a/X_icon_2.svg\" alt=\"X logo\">\n",
        "</a>\n",
        "\n",
        "<a href=\"https://reddit.com/submit?url=https%3A//github.com/GoogleCloudPlatform/generative-ai/blob/main/embeddings/use-cases/outlier-detection/bq-vector-search-outlier-detection-audit-logs.ipynb\" target=\"_blank\">\n",
        "  <img width=\"20px\" src=\"https://redditinc.com/hubfs/Reddit%20Inc/Brand/Reddit_Logo.png\" alt=\"Reddit logo\">\n",
        "</a>\n",
        "\n",
        "<a href=\"https://www.facebook.com/sharer/sharer.php?u=https%3A//github.com/GoogleCloudPlatform/generative-ai/blob/main/embeddings/use-cases/outlier-detection/bq-vector-search-outlier-detection-audit-logs.ipynb\" target=\"_blank\">\n",
        "  <img width=\"20px\" src=\"https://upload.wikimedia.org/wikipedia/commons/5/51/Facebook_f_logo_%282019%29.svg\" alt=\"Facebook logo\">\n",
        "</a>            "
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "61bf5ba500f8"
      },
      "source": [
        "| | |\n",
        "|-|-|\n",
        "| Author: | [Roy Arsan](https://github.com/rarsan) |"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "TGE5SpdQrFUD"
      },
      "source": [
        "## Overview\n",
        "Security and IT operations teams are tasked with analyzing an increasing amount of logs such as [Cloud Audit Logs](https://cloud.google.com/logging/docs/audit). At the same time, they are faced with a number of challenges, including shortage of talent, increased toil, and increased cyber threats. By leveraging AI for log analytics, cloud administrators, security analysts,  and threat hunters can keep up with these challenges:\n",
        "- Proactively and intelligently identify cloud threats or misconfigurations\n",
        "- Reduce the time and effort during a threat or incident investigation.\n",
        "\n",
        "In this notebook, we will demonstrate **Generative AI-enabled outliers detection and incident investigation** using your Cloud Audit logs in your own project.\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "-kHvlYPt9OsY"
      },
      "source": [
        "![log_anomaly_detection_with_bq_vector_search.png](https://storage.googleapis.com/github-repo/generative-ai/embeddings/use-cases/outlier-detection/bq-vector-search-outlier-detection-audit-logs/log_anomaly_detection_with_bq_vector_search.png)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "38b2358a2636"
      },
      "source": [
        "### Objective\n",
        "\n",
        "Using [Vertex AI](https://cloud.google.com/vertex-ai/docs) and [BigQuery](https://cloud.google.com/bigquery/docs), you will:\n",
        "\n",
        "*   Summarize and transform your logs from a verbose deeply nested complex structured payload into a natural language text summary.\n",
        "*   Create embeddings for each textual log summary directly in BigQuery using [`text-embedding-005` Generative AI model](https://cloud.google.com/vertex-ai/docs/generative-ai/embeddings/get-text-embeddings) from Vertex AI.\n",
        "*   Create a fully-managed BigQuery vector index to search all logs summaries without operational overhead.\n",
        "*   Search directly in BigQuery to find any past actions that are semantically similar to suspected current actions.\n",
        "*   Visualize findings in an interative table and timechart for rapid investigation & confirmation of such potential outliers."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "3b1ffd5ab768"
      },
      "source": [
        "## Getting Started"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "hsX82WHAR0Gk"
      },
      "source": [
        "### Prerequisite\n",
        "If you haven't already done so, the only requirement is to [upgrade your existing log bucket](https://cloud.google.com/logging/docs/buckets#upgrade-bucket) to use Log Analytics which provides you with a linked BigQuery dataset with your own queryable logs data. This is a **one-click step without incurring any additional costs**. By default, Cloud Audit Admin Activity logs are enabled, ingested and stored in every project's `_Required` bucket without any charges.\n",
        "\n",
        "![one click prerequisite](https://services.google.com/fh/files/misc/upgrade_log_bucket.png)\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Y0PJAOju-yxx"
      },
      "source": [
        "### Install packages and SDKs"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "CuonDAxt-ewW"
      },
      "outputs": [],
      "source": [
        "# Install mpld3 package for interactive rendering of matplotlib\n",
        "%pip install mpld3\n",
        "# Install Vertex AI and BigQuery SDKs\n",
        "%pip install google-cloud-aiplatform google-cloud-bigquery --upgrade --user"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "RFcQnSibyn4F"
      },
      "source": [
        "### Restart runtime\n",
        "\n",
        "To use the newly installed packages in this Jupyter runtime, you must restart the runtime. You can do this by running the cell below, which will restart the current kernel."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "rGUqrj9jyxOd"
      },
      "outputs": [],
      "source": [
        "# Restart kernel after installs so that your environment can access the new packages\n",
        "import IPython\n",
        "\n",
        "app = IPython.Application.instance()\n",
        "app.kernel.do_shutdown(True)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "02d6dfc513c3"
      },
      "source": [
        "<div class=\"alert alert-block alert-warning\">\n",
        "<b>⚠️ The kernel is going to restart. Please wait until it is finished before continuing to the next step. ⚠️</b>\n",
        "</div>"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "7Eg9Mmp1xEAK"
      },
      "source": [
        "### Authenticate your notebook environment\n",
        "\n",
        "Authenticating your notebook environment\n",
        "- If you are using **Colab** to run this notebook, run the cell below and continue.\n",
        "- If you are using **Vertex AI Workbench**, check out the setup instructions [here](https://github.com/GoogleCloudPlatform/generative-ai/tree/main/setup-env)."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "RjHTFGJbxXXN"
      },
      "outputs": [],
      "source": [
        "import sys\n",
        "\n",
        "if \"google.colab\" in sys.modules:\n",
        "    from google.colab import auth\n",
        "\n",
        "    auth.authenticate_user()"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "I0DdmYrZ-lsY"
      },
      "outputs": [],
      "source": [
        "# For debug only\n",
        "!gcloud config list --format 'value(core.account)'"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "fKvFT-9d_GlP"
      },
      "source": [
        "### Import libraries"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "Xrz9wjvP-mE-"
      },
      "outputs": [],
      "source": [
        "from google.cloud import bigquery\n",
        "import pandas as pd"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "M_-tnDXOwY3i"
      },
      "outputs": [],
      "source": [
        "from google.colab import data_table\n",
        "\n",
        "data_table.enable_dataframe_formatter()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "RgTWVU2z_9KS"
      },
      "source": [
        "### Set Google Cloud region and project ID"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "wH0Qwzzt_Wq0"
      },
      "outputs": [],
      "source": [
        "PROJECT_ID = \"[project-id]\"  # @param {type:\"string\"}\n",
        "LOCATION = \"US\"  # @param {type:\"string\"}\n",
        "REGION = \"us-central1\"  # @param {type:\"string\"}\n",
        "\n",
        "# Set the project id for gcloud CLI and BigQuery client libraries\n",
        "!gcloud config set project {PROJECT_ID}\n",
        "bq = bigquery.Client(project=PROJECT_ID)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "BFwOxRrMA4vc"
      },
      "source": [
        "## Prepare log data and convert to text"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ShreV-cE0X53"
      },
      "source": [
        "You'll preprocess the raw logs that reside in your BigQuery linked dataset (or `SOURCE_DATASET`) into a summary table in your new BigQuery dataset (`PROCESSED_DATASET`). This new dataset will contain the aggregated logs (daily user actions are coalesced) which will also be converted into a simple natural language text. It will also contain the text embeddings and vector index for semantic search.\n",
        "\n",
        "Make sure you have **BigQuery Data Viewer** role over `SOURCE_DATASET` dataset"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "NcQguptKCPba"
      },
      "outputs": [],
      "source": [
        "SOURCE_DATASET = \"[linked-dataset]\"  # @param {type:\"string\"}\n",
        "PROCESSED_DATASET = \"[new-dataset]\"  # @param {type:\"string\"}\n",
        "\n",
        "TABLE_NAME = \"nb_admin_actions_summary\"  # @param {type:\"string\"}"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "_Zb53qoUBf-K"
      },
      "source": [
        "### Create BigQuery dataset"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "tqpblzhADUn_"
      },
      "outputs": [],
      "source": [
        "!bq --location={LOCATION} mk --dataset {PROJECT_ID}:{PROCESSED_DATASET}"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "pgnAhPHwCu1S"
      },
      "source": [
        "### Define BigQuery User-Defined Functions (UDFs)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "b9jABImgMQht"
      },
      "source": [
        "The following UDF helper function converts a structured log entry payload into a natural language text explaining the log entry activity of who did what, when, on which resource and from which ip and user agent. All of these dimensions are important facets to include in the text content and as a result, to consider in the embeddings vector space."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "SWSjMkDkDx3D"
      },
      "outputs": [],
      "source": [
        "# Produce user-friendly text content from an audit log entry's parameters\n",
        "UDF_NAME = \"stringifyAdminLogEntry\"\n",
        "\n",
        "sql = f\"\"\"\n",
        "CREATE OR REPLACE FUNCTION `{PROJECT_ID}.{PROCESSED_DATASET}.{UDF_NAME}`(\n",
        "  day DATE,\n",
        "  principal_email STRING,\n",
        "  action STRING,\n",
        "  resource_type STRING,\n",
        "  resource_id STRING,\n",
        "  container_type STRING,\n",
        "  container_id STRING,\n",
        "  channel STRING,\n",
        "  ip STRING,\n",
        "  counter INT64\n",
        ")\n",
        "RETURNS STRING\n",
        "AS (\n",
        "  \"On \" || CAST(day AS STRING) || \", principal \" || principal_email ||\n",
        "  \"  ran operation \" || action || \" over \" || resource_type || \" \" ||\n",
        "  resource_id || \" in \" || container_type || \" \" || container_id ||\n",
        "  \" using \" || channel || \" from ip \" || ip || \" \" || counter || \" time(s)\"\n",
        ");\n",
        "\"\"\"\n",
        "\n",
        "query_job = bq.query(sql)\n",
        "print(query_job.result())  # Wait for the job to complete.\n",
        "print(f\"Created UDF {PROJECT_ID}.{PROCESSED_DATASET}.{UDF_NAME}\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ATM61oqTMmfe"
      },
      "source": [
        "The following UDF extracts the resource ID that was acted on per the audit log entry. In the audit log entry, The resource ID is specified in a different resource label field depending on the resource type. That's why this UDF is needed to normalize that resource ID field."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "NYQsa3Lo2-qn"
      },
      "outputs": [],
      "source": [
        "# Deduce resource ID from a log entry resource field\n",
        "UDF_NAME = \"getResourceId\"\n",
        "\n",
        "sql = f\"\"\"\n",
        "CREATE OR REPLACE FUNCTION `{PROJECT_ID}.{PROCESSED_DATASET}.{UDF_NAME}`(\n",
        "  type STRING,\n",
        "  labels JSON\n",
        ")\n",
        "RETURNS STRING\n",
        "AS (\n",
        " COALESCE(\n",
        "  JSON_VALUE(labels.email_id),     # service_account\n",
        "  JSON_VALUE(labels.pod_id),       # container\n",
        "  JSON_VALUE(labels.instance_id),  # gce_instance, spanner_instance, redis_instance, ...\n",
        "  JSON_VALUE(labels.subnetwork_id),# gce_subnetwork,\n",
        "  JSON_VALUE(labels.network_id),   # gce_network, gce_network_region, ...\n",
        "  JSON_VALUE(labels.topic_id),     # pubsub_topic\n",
        "  JSON_VALUE(labels.subscription_id), # pubsub_subscription\n",
        "  JSON_VALUE(labels.endpoint_id),  # aiplatform.googleapis.com/Endpoint\n",
        "  JSON_VALUE(labels.job_id),       # dataflow_step\n",
        "  JSON_VALUE(labels.dataset_id),   # bigquery_dataset\n",
        "  JSON_VALUE(labels.project_id),\n",
        "  JSON_VALUE(labels.organization_id),\n",
        "  JSON_VALUE(labels.id),\n",
        "  \"other\")\n",
        ");\n",
        "\"\"\"\n",
        "\n",
        "query_job = bq.query(sql)\n",
        "print(query_job.result())  # Wait for the job to complete.\n",
        "print(f\"Created UDF {PROJECT_ID}.{PROCESSED_DATASET}.{UDF_NAME}\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "YqGoF8lfNEmL"
      },
      "source": [
        "The following UDF deduces where a user or system action occurred from per the audit log entry. For example, an action may have occurred through the Cloud Console, or using gcloud CLI, or via Terraform script or another unknown client or channel."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "PfRcocm65-rk"
      },
      "outputs": [],
      "source": [
        "# Deduce channel from a log entry request user agent\n",
        "UDF_NAME = \"getChannelType\"\n",
        "\n",
        "sql = f\"\"\"\n",
        "CREATE OR REPLACE FUNCTION `{PROJECT_ID}.{PROCESSED_DATASET}.{UDF_NAME}`(\n",
        "  caller_supplied_user_agent STRING\n",
        ")\n",
        "RETURNS STRING\n",
        "AS (\n",
        "  CASE\n",
        "    WHEN caller_supplied_user_agent LIKE \"Mozilla/%\" THEN 'Cloud Console'\n",
        "    WHEN caller_supplied_user_agent LIKE \"google-cloud-sdk gcloud/%\" THEN 'gcloud CLI'\n",
        "    WHEN caller_supplied_user_agent LIKE \"google-api-go-client/% Terraform/%\" THEN 'Terraform'\n",
        "    ELSE 'other'\n",
        "  END\n",
        ");\n",
        "\"\"\"\n",
        "\n",
        "query_job = bq.query(sql)\n",
        "print(query_job.result())  # Wait for the job to complete.\n",
        "print(f\"Created UDF {PROJECT_ID}.{PROCESSED_DATASET}.{UDF_NAME}\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "hWCB3WeUBiDN"
      },
      "source": [
        "### Build summary table of admin actions"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "5lUGTp-MhqoJ"
      },
      "outputs": [],
      "source": [
        "LOOKBACK_WINDOW_DAYS = 360  # @param {type:\"slider\", min:30, max:1080, step:30}"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "GPhiu82VBx6P"
      },
      "outputs": [],
      "source": [
        "TABLE_ID = f\"{PROJECT_ID}.{PROCESSED_DATASET}.{TABLE_NAME}\"\n",
        "\n",
        "job_config = bigquery.QueryJobConfig(\n",
        "    destination=TABLE_ID, write_disposition=\"WRITE_TRUNCATE\"\n",
        ")\n",
        "\n",
        "sql = f\"\"\"\n",
        "SELECT\n",
        "  {PROCESSED_DATASET}.stringifyAdminLogEntry(\n",
        "    day, principal_email, action, resource_type, resource_id, container_type,\n",
        "    container_id, channel, ip, counter) AS content,\n",
        "  *\n",
        "FROM (\n",
        "  SELECT\n",
        "    EXTRACT(DATE FROM timestamp) AS day,\n",
        "    IFNULL(proto_payload.audit_log.authentication_info.principal_email, \"unknown\") as principal_email,\n",
        "    IFNULL(proto_payload.audit_log.method_name, \"unknown\") as action,\n",
        "    IFNULL(resource.type, \"unknown\") as resource_type,\n",
        "    {PROCESSED_DATASET}.getResourceId(resource.type, resource.labels) AS resource_id,\n",
        "    -- proto_payload.audit_log.resource_name as resource_name,\n",
        "    SPLIT(log_name, '/')[SAFE_OFFSET(0)] as container_type,\n",
        "    SPLIT(log_name, '/')[SAFE_OFFSET(1)] as container_id,\n",
        "    {PROCESSED_DATASET}.getChannelType(proto_payload.audit_log.request_metadata.caller_supplied_user_agent) AS channel,\n",
        "    IFNULL(proto_payload.audit_log.request_metadata.caller_ip, \"unknown\") as ip,\n",
        "    COUNT(*) counter,\n",
        "    -- ANY_VALUE(resource) as resource,           -- for debugging\n",
        "    -- ANY_VALUE(proto_payload) as proto_payload  -- for debugging\n",
        "  FROM  `{PROJECT_ID}.{SOURCE_DATASET}._AllLogs`\n",
        "  WHERE\n",
        "    -- log_id = \"cloudaudit.googleapis.com/activity\" AND\n",
        "    timestamp > TIMESTAMP_SUB(CURRENT_TIMESTAMP(), INTERVAL {LOOKBACK_WINDOW_DAYS} DAY)\n",
        "  GROUP BY\n",
        "    day, principal_email, action, resource_type, resource_id, container_type, container_id, channel, ip, log_name\n",
        "  ORDER BY\n",
        "    day DESC, principal_email, action\n",
        ")\n",
        "\"\"\"\n",
        "\n",
        "# Start the query and save results in new table\n",
        "query_job = bq.query(sql, job_config=job_config)\n",
        "result = query_job.result()  # Wait for the job to complete.\n",
        "\n",
        "print(f\"{result.total_rows} admin action records loaded to table {TABLE_ID}\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "PXSx722dNf4x"
      },
      "source": [
        "Let's peak into some of these summarized daily user actions:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "Ao2gLpW6NoCI"
      },
      "outputs": [],
      "source": [
        "result.to_dataframe().head(10)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "1194baac6a52"
      },
      "source": [
        "Below are example daily user actions summarizing what action was taken by whom, on what resource, and from where:"
      ]
    },
    {
      "attachments": {
        "4c5c1883-5ab6-4fca-9b4f-c7b1584d50e3.png": {
          "image/png": "[encoded data, 147120 characters]"
        }
      },
      "cell_type": "markdown",
      "metadata": {
        "id": "ce36b3ed8e7d"
      },
      "source": [
        "![image.png](attachment:4c5c1883-5ab6-4fca-9b4f-c7b1584d50e3.png)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "pxhwaND4_IvS"
      },
      "source": [
        "## Generate text embeddings for logs"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "gjy7jR0T_eAL"
      },
      "source": [
        "### Create BigQuery Cloud resource connection"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "hpQqsysu_zCr"
      },
      "source": [
        "You need to create a Cloud resource connection to enable BigQuery Dataframes to interact with Vertex AI services"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "dbWYeyEzii71"
      },
      "outputs": [],
      "source": [
        "CONN_NAME = \"bqdf-llm-embeddings\""
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "5_ko_ngACnPd"
      },
      "outputs": [],
      "source": [
        "!bq mk --connection --location={LOCATION} --project_id={PROJECT_ID} \\\n",
        "    --connection_type=CLOUD_RESOURCE {CONN_NAME}"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "hS0kzr2gD0-h"
      },
      "source": [
        "Retrieve the connection service account ID"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "IG3KaOQb_doE"
      },
      "outputs": [],
      "source": [
        "!bq show --connection {PROJECT_ID}.{LOCATION}.{CONN_NAME}"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "D143kdK7Hjpf"
      },
      "source": [
        "Copy the service account ID in the following parameter:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "ryD1ufggF6Fj"
      },
      "outputs": [],
      "source": [
        "# Copy the value of `serviceAccountId` field in last cell's output\n",
        "CONN_SA_ID = \"[bq-connection-service-account]\"  # @param {type:\"string\"}"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "pa3lCS2hHqv-"
      },
      "source": [
        "Give that service account access to Vertex AI"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "h-SoTNgKHpvg"
      },
      "outputs": [],
      "source": [
        "!gcloud projects add-iam-policy-binding {PROJECT_ID} --member='serviceAccount:{CONN_SA_ID}' --role='roles/aiplatform.user' --condition=None --no-user-output-enabled"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "8HJgiz91IBHz"
      },
      "source": [
        "### Define the LLM model"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "b32791b4d884"
      },
      "source": [
        "Let's create a remote model in BigQuery using the above BigQuery connection. The remote endpoint is set as `textembeddings-gecko` LLM model in Vertex AI."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "4ff3c629535e"
      },
      "outputs": [],
      "source": [
        "LLM_ENDPOINT = \"text-embedding-005\"\n",
        "MODEL_NAME = \"embedding_model\"\n",
        "\n",
        "sql = f\"\"\"\n",
        "CREATE OR REPLACE MODEL `{PROJECT_ID}.{PROCESSED_DATASET}.{MODEL_NAME}`\n",
        "REMOTE WITH CONNECTION `{PROJECT_ID}.{LOCATION}.{CONN_NAME}`\n",
        "OPTIONS (ENDPOINT = '{LLM_ENDPOINT}');\n",
        "\"\"\"\n",
        "\n",
        "# Start the query\n",
        "query_job = bq.query(sql)\n",
        "print(query_job.result())  # Wait for the job to complete.\n",
        "print(f\"Created remote model {PROJECT_ID}.{PROCESSED_DATASET}.{MODEL_NAME}\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "AXw5FKWu9WoJ"
      },
      "source": [
        "### Generate embedding for each log summary"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "AUeODs3ejZpO"
      },
      "source": [
        "The following code generates text embedding for each past action in your admin actions summary table. All records including their embeddings are saved into a new table with suffix `_embeddings`.\n",
        "\n",
        "  CAUTION: This may take several minutes and incur Vertex AI costs depending on how many rows in your table."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "936b956c5a56"
      },
      "outputs": [],
      "source": [
        "TABLE_ID = f\"{PROJECT_ID}.{PROCESSED_DATASET}.{TABLE_NAME}_embeddings\"\n",
        "\n",
        "job_config = bigquery.QueryJobConfig(\n",
        "    destination=TABLE_ID, write_disposition=\"WRITE_TRUNCATE\"\n",
        ")\n",
        "\n",
        "sql = f\"\"\"\n",
        "SELECT\n",
        "  * EXCEPT (ml_embed_text_status, statistics),\n",
        "FROM ML.GENERATE_TEXT_EMBEDDING(\n",
        "  MODEL `{PROJECT_ID}.{PROCESSED_DATASET}.{MODEL_NAME}`,\n",
        "  TABLE `{PROJECT_ID}.{PROCESSED_DATASET}.{TABLE_NAME}`,\n",
        "  STRUCT(TRUE AS flatten_json_output)\n",
        ")\n",
        "\"\"\"\n",
        "\n",
        "# Start the query and save results in new table\n",
        "query_job = bq.query(sql, job_config=job_config)\n",
        "results = query_job.result()  # Wait for the job to complete.\n",
        "\n",
        "print(f\"{results.total_rows} embedded admin actions loaded to table {TABLE_ID}\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "6ef250380afb"
      },
      "source": [
        "Let's peak into some of the embedded admin actions:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "f7117e41dc5e"
      },
      "outputs": [],
      "source": [
        "results.to_dataframe().head(2)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "da26c0b996bf"
      },
      "source": [
        "## Search logs with BigQuery Vector Search"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "qGvFqZd6q9pG"
      },
      "source": [
        "### Create a vector index"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "M1TENz6mrAuN"
      },
      "outputs": [],
      "source": [
        "sql = f\"\"\"\n",
        "CREATE OR REPLACE VECTOR INDEX `my_vector_index`\n",
        "ON `{PROJECT_ID}.{PROCESSED_DATASET}.{TABLE_NAME}_embeddings`(text_embedding)\n",
        "OPTIONS(distance_type='COSINE', index_type='IVF');\n",
        "\"\"\"\n",
        "\n",
        "query_job = bq.query(sql)\n",
        "print(query_job.result())  # Wait for the job to complete.\n",
        "print(\n",
        "    f\"Created Vector index for {PROJECT_ID}.{PROCESSED_DATASET}.{TABLE_NAME}_embeddings\"\n",
        ")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ullFyiSloDu5"
      },
      "source": [
        "Inspect vector index progress, and total storage (chargeable) from `INFORMATION_SCHEMA.VECTOR_INDEXES` view:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "M3Qi1hVsoHf7"
      },
      "outputs": [],
      "source": [
        "sql = f\"\"\"\n",
        "SELECT\n",
        "  table_name,\n",
        "  index_name,\n",
        "  coverage_percentage,\n",
        "  unindexed_row_count,\n",
        "  total_logical_bytes,\n",
        "  total_storage_bytes\n",
        "FROM `{PROJECT_ID}.{PROCESSED_DATASET}.INFORMATION_SCHEMA.VECTOR_INDEXES`\n",
        "WHERE index_status = 'ACTIVE';\n",
        "\"\"\"\n",
        "query_job = bq.query(sql)\n",
        "results = query_job.result()  # Wait for the job to complete\n",
        "results.to_dataframe()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "gR6FcfBEfZFt"
      },
      "source": [
        "### Vector search across existing embeddings"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "pRF7JGg_H2Xg"
      },
      "source": [
        "Pick a couple of recent suspicious actions from the summary table with embeddings, and save into a new table of actions to investigate:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "782xJrFxh3fy"
      },
      "outputs": [],
      "source": [
        "TABLE_ID = f\"{PROJECT_ID}.{PROCESSED_DATASET}.{TABLE_NAME}_test_actions\"\n",
        "\n",
        "job_config = bigquery.QueryJobConfig(\n",
        "    destination=TABLE_ID, write_disposition=\"WRITE_TRUNCATE\"\n",
        ")\n",
        "\n",
        "# Search for any N destructive actions (containing 'delete') over the last week\n",
        "sql = f\"\"\"\n",
        "SELECT\n",
        "  *\n",
        "FROM  `{PROJECT_ID}.{PROCESSED_DATASET}.{TABLE_NAME}_embeddings`\n",
        "WHERE\n",
        "  day > DATE_SUB(CURRENT_DATE(), INTERVAL 7 DAY) AND\n",
        "  lower(action) LIKE \"%delete%\"\n",
        "LIMIT 2\n",
        "\"\"\"\n",
        "\n",
        "# Start the query and save results in new table\n",
        "query_job = bq.query(sql, job_config=job_config)\n",
        "results = query_job.result()  # Wait for the job to complete.\n",
        "\n",
        "print(\n",
        "    f\"{results.total_rows} suspicious admin action records loaded to table {TABLE_ID}\"\n",
        ")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "dS340NF0gSfa"
      },
      "source": [
        "As part of your investigation, let's confirm whether each of these actions is an anomaly or not. We can do so using a semantic search across all past actions to find out if a similar activity has occurred in the past."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "XkQkZLnMxIaK"
      },
      "outputs": [],
      "source": [
        "sql = f\"\"\"\n",
        "SELECT\n",
        "  query.content as suspicious_action,\n",
        "  base.content as past_similar_action,\n",
        "  distance,\n",
        "  base.day as day,\n",
        "  base.counter as counter,\n",
        "FROM VECTOR_SEARCH(\n",
        "  TABLE `{PROJECT_ID}.{PROCESSED_DATASET}.{TABLE_NAME}_embeddings`, 'text_embedding',\n",
        "  TABLE `{PROJECT_ID}.{PROCESSED_DATASET}.{TABLE_NAME}_test_actions`,\n",
        "  top_k => 5\n",
        ")\n",
        "WHERE query.content != base.content -- remove exact dups\n",
        "\"\"\"\n",
        "\n",
        "# Start the query and save results in an in-memory dataframe\n",
        "query_job = bq.query(sql)\n",
        "results = query_job.result()  # Wait for the job to complete\n",
        "df = results.to_dataframe()\n",
        "df"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "zgJ9IDVHp8vX"
      },
      "source": [
        "#### Visualize results"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "pNMlbSqu4PoO"
      },
      "source": [
        "<ins>In a table</ins>:"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "5c29f5876f95"
      },
      "source": [
        "First, create helper functions to:\n",
        "- Reshape the dataframe to group by suspicious action and list associated similar actions,\n",
        "- Infer likelihood of equivalency (`likelihood`) for each pair given nearest neighbor distance (`distance`),\n",
        "- Apply custom table styling to highlight \"very similar\" actions (d < 0.01) that are deemed equivalent"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "oaF4oLxI8v3h"
      },
      "outputs": [],
      "source": [
        "def similarity(d):\n",
        "    if d < 0.01:\n",
        "        return \"Very similar\"\n",
        "    elif d < 0.04:\n",
        "        return \"Similar\"\n",
        "    return \"Somewhat similar\"\n",
        "\n",
        "\n",
        "def reshape(df):\n",
        "    df2 = df.copy().query(\"distance > 0.000000000001\")\n",
        "    df2[\"suspected\"] = True\n",
        "    df2[\"likelihood\"] = df2[\"distance\"].apply(similarity)\n",
        "    df2.sort_values(by=\"day\", inplace=True)\n",
        "\n",
        "    pivoted = df2.pivot(\n",
        "        index=[\"suspicious_action\", \"past_similar_action\"],\n",
        "        columns=\"suspected\",\n",
        "        values=[\"day\", \"distance\", \"likelihood\"],\n",
        "    )\n",
        "    return pivoted\n",
        "\n",
        "\n",
        "def highlight_similar_actions(styler):\n",
        "    styler.highlight_between(\n",
        "        left=0,\n",
        "        right=0.01,\n",
        "        axis=1,\n",
        "        subset=[\"distance\"],\n",
        "        props=\"color:white; font-weight:bold; background-color:darkblue;\",\n",
        "    )\n",
        "    return styler"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "f052e822aa1d"
      },
      "outputs": [],
      "source": [
        "reshape(df).style.pipe(highlight_similar_actions)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "4c6a1269aeb2"
      },
      "source": [
        "The results will be rendered in a table as follows:"
      ]
    },
    {
      "attachments": {
        "7a699ddf-ba3c-436d-a938-33505258c2c9.png": {
          "image/png": "[encoded data, 246096 characters]"
        }
      },
      "cell_type": "markdown",
      "metadata": {
        "id": "aabbd61af7df"
      },
      "source": [
        "![3TAsfLVJs93D4eF.png](attachment:7a699ddf-ba3c-436d-a938-33505258c2c9.png)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "673408ff9b26"
      },
      "source": [
        "In this example, the search returned similar past action by the same service account on several occasions in the past, so the suspected action is deemed **not an anomaly**"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "NfLG816P4bxA"
      },
      "source": [
        "<ins>In a timechart</ins>:"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "5FKlNvhU33or"
      },
      "source": [
        "First, create helper function to plot the dataframe of past actions grouped by a given suspicious action:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "YHJmyNqGh958"
      },
      "outputs": [],
      "source": [
        "import matplotlib.pyplot as plt\n",
        "import mpld3\n",
        "from mpld3 import plugins\n",
        "\n",
        "\n",
        "def plot_actions_over_time(df):\n",
        "    mpld3.enable_notebook()\n",
        "\n",
        "    # Set CSS to style our custom labels\n",
        "    css = \"\"\"\n",
        "  p {\n",
        "    font-size: 14px;\n",
        "    color: blue;\n",
        "    background-color: white;\n",
        "  }\n",
        "  \"\"\"\n",
        "\n",
        "    fig, ax = plt.subplots(figsize=(12, 8))\n",
        "\n",
        "    df[\"day\"] = pd.to_datetime(df[\"day\"])\n",
        "    for i, ts in df.groupby(\"suspicious_action\"):\n",
        "        points = ax.plot(\n",
        "            ts.day, ts.distance, \"o\", label=str(i), mec=\"k\", ms=15, mew=1, alpha=0.6\n",
        "        )\n",
        "        labels = list(ts.past_similar_action.apply(lambda x: f\"<p>{x}</p>\"))\n",
        "        tooltip = plugins.PointHTMLTooltip(\n",
        "            points[0], labels, voffset=10, hoffset=10, css=css\n",
        "        )\n",
        "        plugins.connect(fig, tooltip)\n",
        "\n",
        "    ax.legend(loc=\"upper left\", bbox_to_anchor=(-0.1, 1.1))\n",
        "    ax.set_xlabel(\"time\")\n",
        "    ax.set_ylabel(\"distance\")\n",
        "    plt.show()\n",
        "\n",
        "    mpld3.disable_notebook()"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "R7ABHBfO3OJG"
      },
      "outputs": [],
      "source": [
        "plot_actions_over_time(df)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ea954aabc228"
      },
      "source": [
        "The results will be rendered in a timechart as follows, where past closest matches of each suspected action are displayed in a separate line chart. The y-axis represents the distance of these nearest neighbors to the suspected action, where 0 is exact match."
      ]
    },
    {
      "attachments": {
        "cbe85182-03c5-4aa0-936b-a94252b91d64.png": {
          "image/png": "[encoded data, 91292 characters]"
        }
      },
      "cell_type": "markdown",
      "metadata": {
        "id": "afee4a455a64"
      },
      "source": [
        "![AyVjx6h7w6MxLAN.png](attachment:cbe85182-03c5-4aa0-936b-a94252b91d64.png)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "f293b5bcedec"
      },
      "source": [
        "In this example, the 'orange' dots are not close enough matches and therefore the suspected action of 'DeleteRoutine' is likely an anomaly that needs to be further investigated."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "3Egnvumyot88"
      },
      "source": [
        "### Vector search of a new log entry"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "WOWRvRf4iT2F"
      },
      "source": [
        "Let's pick a seemingly suspicious destructive action from your recent admin activity log stream"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "Iz9R9gUTnJ-c"
      },
      "outputs": [],
      "source": [
        "inner_sql = f\"\"\"\n",
        "SELECT\n",
        "  EXTRACT(DATE FROM timestamp) AS day,\n",
        "  IFNULL(proto_payload.audit_log.authentication_info.principal_email, \"unknown\") as principal_email,\n",
        "  IFNULL(proto_payload.audit_log.method_name, \"unknown\") as action,\n",
        "  IFNULL(resource.type, \"unknown\") as resource_type,\n",
        "  {PROCESSED_DATASET}.getResourceId(resource.type, resource.labels) AS resource_id,\n",
        "  -- proto_payload.audit_log.resource_name as resource_name,\n",
        "  SPLIT(log_name, '/')[SAFE_OFFSET(0)] as container_type,\n",
        "  SPLIT(log_name, '/')[SAFE_OFFSET(1)] as container_id,\n",
        "  {PROCESSED_DATASET}.getChannelType(proto_payload.audit_log.request_metadata.caller_supplied_user_agent) AS channel,\n",
        "  IFNULL(proto_payload.audit_log.request_metadata.caller_ip, \"unknown\") as ip,\n",
        "  COUNT(*) counter,\n",
        "FROM  `{PROJECT_ID}.{SOURCE_DATASET}._AllLogs`\n",
        "WHERE\n",
        "  log_id = \"cloudaudit.googleapis.com/activity\" AND\n",
        "  timestamp > TIMESTAMP_SUB(CURRENT_TIMESTAMP(), INTERVAL 1 DAY) AND\n",
        "  lower(proto_payload.audit_log.method_name) LIKE \"%delete%\"\n",
        "GROUP BY\n",
        "  day, principal_email, action, resource_type, resource_id, container_type, container_id, channel, ip, log_name\n",
        "ORDER BY\n",
        "  day DESC,\n",
        "  counter DESC\n",
        "LIMIT 2\n",
        "\"\"\"\n",
        "\n",
        "# Start the query and save results in an in-memory dataframe\n",
        "query_job = bq.query(inner_sql)\n",
        "results = query_job.result()  # Wait for the job to complete\n",
        "results.to_dataframe()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "029c09ceb983"
      },
      "source": [
        "In the following singular SQL command, 1) convert the suspected log into a text string summary, 2) generate the text embeddings 2) and do vector search across all past actions to find out similar activity (if any) has occurred in the past:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "CEqc2MSSpkv6"
      },
      "outputs": [],
      "source": [
        "sql = f\"\"\"\n",
        "SELECT\n",
        "  query.content as suspicious_action,\n",
        "  base.content as past_similar_action,\n",
        "  distance,\n",
        "  base.day as day,\n",
        "  base.counter as counter\n",
        "FROM VECTOR_SEARCH(\n",
        "  TABLE `{PROJECT_ID}.{PROCESSED_DATASET}.{TABLE_NAME}_embeddings`,\n",
        "  'text_embedding',\n",
        "  (SELECT\n",
        "    text_embedding, content\n",
        "    FROM ML.GENERATE_TEXT_EMBEDDING(\n",
        "      MODEL `{PROJECT_ID}.{PROCESSED_DATASET}.embedding_model`,\n",
        "      (SELECT\n",
        "        {PROCESSED_DATASET}.stringifyAdminLogEntry(\n",
        "          day, principal_email, action, resource_type, resource_id, container_type,\n",
        "          container_id, channel, ip, counter) AS content\n",
        "        FROM ({inner_sql})),\n",
        "      STRUCT(TRUE AS flatten_json_output)\n",
        "  )),\n",
        "  top_k => 5\n",
        ")\n",
        "WHERE query.content != base.content -- remove potential dups\n",
        "\"\"\"\n",
        "\n",
        "# Start the query and save results in an in-memory dataframe\n",
        "query_job = bq.query(sql)\n",
        "results = query_job.result()  # Wait for the job to complete\n",
        "df = results.to_dataframe()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "VrX4CSU57F-b"
      },
      "source": [
        "#### Visualize results"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Q4mphXYM7JoF"
      },
      "source": [
        "<ins>In a table<ins>:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "8_jWjpjP7JBH"
      },
      "outputs": [],
      "source": [
        "reshape(df).style.pipe(highlight_similar_actions)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ouYmslmU8SFV"
      },
      "source": [
        "<ins>In a timechart</ins>:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "0XT6wkTK8gs1"
      },
      "outputs": [],
      "source": [
        "plot_actions_over_time(df)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "959b92998af7"
      },
      "source": [
        "## Summary\n",
        "\n",
        "In this notebook, you were able to:\n",
        "- Process your Google Cloud audit logs into a summary table for easier, faster and cost-effective log analysis\n",
        "- Convert logs from semi-structured deeply nested JSON payload into a concise textual log summary \n",
        "- Generate embeddings for each log summary using Vertex AI, and store them in a BigQuery table alongside the log content\n",
        "- Create a vector index in BigQuery to semantically search all historical log summaries using embeddings column\n",
        "- Search logs using BigQuery `VECTOR_SEARCH` SQL command to detect & investigate suspicious or anomalous activity\n",
        "- Visualize findings in a table and a timechart to highlight past similar or equivalent actions, if any"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "9633a0a9a970"
      },
      "source": [
        "## Cleanup\n",
        "\n",
        "To clean up all Google Cloud resources used in this notebook, you can delete the Google Cloud project you used for the tutorial.\n",
        "\n",
        "Otherwise, you can delete the individual resources you created in this tutorial, namely the BigQuery dataset `PROCESSED_DATASET` with the processed data:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "554bae6af364"
      },
      "outputs": [],
      "source": [
        "# Delete the created BigQuery dataset\n",
        "!bq rm -r -f {PROJECT_ID}:{PROCESSED_DATASET}"
      ]
    }
  ],
  "metadata": {
    "colab": {
      "collapsed_sections": [
        "WN2LVDh0-c_l"
      ],
      "name": "bq-vector-search-outlier-detection-audit-logs.ipynb",
      "toc_visible": true
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}
