'HarmCategory' object has no attribute 'lower' error when used with google-cloud-aiplatform

[michael-erasmus]

I'm using google-cloud-aiplatform with weave and running into the error 'HarmCategory' object has no attribute 'lower when specifying safety_settings

Here's a minimal reproducible example:

Install dependencies

$ pip install weave google-cloud-aiplatform

Then create a run.py file like this:

import vertexai
from vertexai.generative_models import GenerativeModel, HarmBlockThreshold, HarmCategory

import weave

@weave.op
def run():
    vertexai.init(project='<PROJECT_NAME>', location='<LOCATION>')
    model = GenerativeModel('gemini-1.5-pro-002')

    generation_config = {
            "temperature": 0.7,
            "max_output_tokens": 100,
    }

    safety_config = {
        HarmCategory.HARM_CATEGORY_UNSPECIFIED: HarmBlockThreshold.BLOCK_NONE,
        HarmCategory.HARM_CATEGORY_DANGEROUS_CONTENT: HarmBlockThreshold.BLOCK_NONE,
        HarmCategory.HARM_CATEGORY_HARASSMENT: HarmBlockThreshold.BLOCK_NONE,
        HarmCategory.HARM_CATEGORY_HATE_SPEECH: HarmBlockThreshold.BLOCK_NONE,
        HarmCategory.HARM_CATEGORY_SEXUALLY_EXPLICIT: HarmBlockThreshold.BLOCK_NONE,
    }

    response = model.generate_content(
        "Hello Gemini!", generation_config=generation_config, safety_settings=safety_config
    )

    print(response.candidates[0].content.parts[0].text)


if __name__ == "__main__":
    weave.init('repro-error')
    run()

Then run

$ python run.py

This results in the error:

Logged in as Weights & Biases user: michael-erasmus.
View Weave data at https://wandb.ai/chesscom/repro-error/weave
Error creating call:
Traceback (most recent call last):
  File "/Users/michaelerasmus/.pyenv/versions/3.10.13/envs/venv-chess.com-weave-repro/lib/python3.10/site-packages/weave/trace/op.py", line 444, in _do_call
    call = _create_call(op, *args, __weave=__weave, **kwargs)
  File "/Users/michaelerasmus/.pyenv/versions/3.10.13/envs/venv-chess.com-weave-repro/lib/python3.10/site-packages/weave/trace/op.py", line 259, in _create_call
    return client.create_call(
  File "/Users/michaelerasmus/.pyenv/versions/3.10.13/envs/venv-chess.com-weave-repro/lib/python3.10/site-packages/weave/trace/trace_sentry.py", line 204, in wrapper
    return func(*args, **kwargs)
  File "/Users/michaelerasmus/.pyenv/versions/3.10.13/envs/venv-chess.com-weave-repro/lib/python3.10/site-packages/weave/trace/weave_client.py", line 766, in create_call
    inputs_redacted = redact_sensitive_keys(inputs)
  File "/Users/michaelerasmus/.pyenv/versions/3.10.13/envs/venv-chess.com-weave-repro/lib/python3.10/site-packages/weave/trace/weave_client.py", line 1654, in redact_sensitive_keys
    dict_res[k] = redact_sensitive_keys(v)
  File "/Users/michaelerasmus/.pyenv/versions/3.10.13/envs/venv-chess.com-weave-repro/lib/python3.10/site-packages/weave/trace/weave_client.py", line 1651, in redact_sensitive_keys
    if should_redact(k):
  File "/Users/michaelerasmus/.pyenv/versions/3.10.13/envs/venv-chess.com-weave-repro/lib/python3.10/site-packages/weave/trace/sanitize.py", line 11, in should_redact
    return key.lower() in REDACT_KEYS
AttributeError: 'HarmCategory' object has no attribute 'lower'
 (subsequent messages of this type will be suppressed)
Hello! How can I help you today?

🍩 https://wandb.ai/chesscom/repro-error/r/call/0194526d-983e-7691-bc38-2850e3d2fc94

It seems like the code assumes that the keys will be strings when they're not.