Skip to content

Log safety attributes from Vertex AI

Tan Le requested to merge 313-track-vertex-blocked-requests into main

What does this merge request do and why?

Vertext LLMs sometime can generate output that violates safety content policy and this results in empty suggestions. We need to log this information to better understand how often this occurs and if any changes to our transformations that would cause this.

https://cloud.google.com/vertex-ai/docs/generative-ai/learn/responsible-ai

How to set up and validate locally

  1. Check out to this merge request's branch.
  2. Ensure a local Docker image built successfully.
    docker buildx build --platform linux/amd64 \
      -t ai-gateway:test .
  3. Run a local service on Docker.
    docker run --platform linux/amd64 --rm \
      -p 5052:5052 \
      -e AUTH_BYPASS_EXTERNAL=true \
      -v $PWD:/app -it ai-gateway:test

Code completions

  1. Send a cURL request to the /v2/completions endpoint

    $ curl --request POST \
      --url http://codesuggestions.gdk.test:5052/v2/completions \
      --header 'Content-Type: application/json' \
      --header 'X-Gitlab-Authentication-Type: oidc' \
      --header 'authorization: Bearer jwt \
      --data '{
      "prompt_version": 1,
      "project_path": "gitlab-org/gitlab",
      "project_id": 278964,
      "current_file": {
        "file_name": "main.rb",
        "content_above_cursor": "# bomberman",
        "content_below_cursor": ""
      }
    }'
  2. Observe the safety attributes in the log (formatted for legibility)

    { 
      ...
      "safety_categories": [],
      "blocked": true,
      ...
    }

Code generations

  1. Send a cURL request to the /v2/code/generations endpoint

    $ curl --request POST \
      --url http://codesuggestions.gdk.test:5052/v2/code/generations \
      --header 'Content-Type: application/json' \
      --header 'X-Gitlab-Authentication-Type: oidc' \
      --header 'authorization: Bearer jwt \
      --data '{
      "prompt_version": 1,
      "project_path": "gitlab-org/gitlab",
      "project_id": 278964,
      "current_file": {
        "file_name": "main.rb",
        "content_above_cursor": "# bomberman ",
        "content_below_cursor": ""
      }
    }'
  2. Observe the safety attributes in the log (formatted for legibility)

    {
      ...
      "safety_categories": ["Firearms & Weapons", "Toxic", "War & Conflict"],
      "blocked": false,
      ...
    }

Merge request checklist

  • Tests added for new functionality. If not, please raise an issue to follow up.
  • Documentation added/updated, if needed.

Relates to #313 (closed)

Edited by Tan Le

Merge request reports