Using LLMs hosted on Vertex AI#

This guide teaches you how to use NeMo Guardrails with LLMs hosted on Vertex AI. It uses the ABC Bot configuration and changes the model to gemini-1.0-pro.

This guide assumes you have configured and tested working with Vertex AI models. If not, refer to this guide.

Prerequisites#

You need to install the following Python libraries:

  1. Install the google-cloud-aiplatform and langchain-google-vertexai packages:

pip install --quiet "google-cloud-aiplatform>=1.38.0" langchain-google-vertexai==0.1.0

  1. Set the GOOGLE_APPLICATION_CREDENTIALS environment variable:

export GOOGLE_APPLICATION_CREDENTIALS=$GOOGLE_APPLICATION_CREDENTIALS # Replace with your own key

  1. If you’re running this inside a notebook, patch the AsyncIO loop.

import nest_asyncio
nest_asyncio.apply()

Configuration#

To get started, copy the ABC bot configuration into a subdirectory called config:

cp -r ../../../../examples/bots/abc config

Update the config/config.yml file to use the gemini-1.0-pro model with the vertexai provider:

...

models:
  - type: main
    engine: vertexai
    model: gemini-1.0-pro

...

Load the guardrails configuration:

from nemoguardrails import RailsConfig
from nemoguardrails import LLMRails

config = RailsConfig.from_path("./config")
rails = LLMRails(config)

Test that it works:

response = rails.generate(messages=[{
    "role": "user",
    "content": "Hi! How are you?"
}])
print(response)

{'role': 'assistant', 'content': "I'm doing great! Thank you for asking. I'm here to help you with any questions you may have about the ABC Company."}

You can see that the bot responds correctly. To see in more detail what LLM calls have been made, you can use the print_llm_calls_summary method as follows:

info = rails.explain()
info.print_llm_calls_summary()

Summary: 5 LLM call(s) took 3.99 seconds .

1. Task `self_check_input` took 0.58 seconds .
2. Task `generate_user_intent` took 1.19 seconds .
3. Task `generate_next_steps` took 0.71 seconds .
4. Task `generate_bot_message` took 0.88 seconds .
5. Task `self_check_output` took 0.63 seconds .

Evaluation#

The gemini-1.0-pro and text-bison models have been evaluated for topical rails, and gemini-1.0-pro has also been evaluated as a self-checking model for hallucination and content moderation. Evaluation results can be found here.

Conclusion#

In this guide, you learned how to connect a NeMo Guardrails configuration to a Vertex AI LLM model. This guide uses gemini-1.0-pro, however, you can connect any other model following the same steps.