Expose PaLM safety_attributes for each prediction from VertexAI

[wnmurphy]

Not sure of any clean way to do this, since the _VertexAICommon interfaces expect res.text, but it would be really cool to expose the safety_attributes that are also returned in the response (link).

With PaLM text-bison@001 as of google-cloud-aiplatform==1.28.0 you get safety scores for free in each prediction along dimensions like illicit drugs, firearms and weapons, death/harm/tragedy, etc.

Google has some built-in settings that aren't very granular (like "more restrictive"/"less restrictive" etc.) but we're doing mental health crisis detection and escalation, and it would be rad to expose these scores for each prediction so we could trigger monitoring of potential self-harm events.

I see this in _VertexAICommon:

  def _predict(
      self, prompt: str, stop: Optional[List[str]] = None, **kwargs: Any
  ) -> str:
      params = {**self._default_params, **kwargs}
      res = completion_with_retry(self, prompt, **params)  # type: ignore
      if "safety_attributes" in res:
          self.last_prediction
      return self._enforce_stop_words(res.text, stop)

  def _enforce_stop_words(self, text: str, stop: Optional[List[str]] = None) -> str:
      if stop is None and self.stop is not None:
          stop = self.stop
      if stop:
          return enforce_stop_tokens(text, stop)
      return text

One approach might be to add (res.text, res.safety_attributes) to self.last_k_predictions = []. Then these could still be accessible when the same model is used for multiple chains in a single operation.

What do you think?