feat: Add api_key_env_var to Model, pass in kwargs to langchain initializer (#1142)

Author: tgasser-nv

examples/configs/content_safety_api_keys/README.md

@@ -0,0 +1,17 @@
+# NemoGuard ContentSafety Usage Example
+
+This example showcases the use of NVIDIA's [NemoGuard ContentSafety model](./../../../docs/user-guides/advanced/nemoguard-contentsafety-deployment.md) for topical and dialogue moderation.
+
+The structure of the config folder is the following:
+
+- `config.yml` - The config file holding all the configuration options for the model.
+- `prompts.yml` - The config file holding the topical rules used for topical and dialogue moderation by the current guardrail configuration.
+
+Please see the docs for more details about the [recommended ContentSafety deployment](./../../../docs/user-guides/advanced/nemoguard-contentsafety-deployment.md) methods, either using locally downloaded NIMs or NVIDIA AI Enterprise (NVAIE).
+
+Before running this example, please set environment variables `NG_OPENAI_API_KEY` and `NG_NVIDIA_API_KEY` to your OpenAI API Key and Nvidia build.nvidia.com Key as below:
+
+```shell
+export NG_OPENAI_API_KEY="<OpenAI API Key>"
+export NG_NVIDIA_API_KEY="<NVIDIA API Key>"
+```

examples/configs/content_safety_api_keys/config.yml

@@ -0,0 +1,18 @@
+models:
+  - type: main
+    engine: openai
+    model: gpt-3.5-turbo-instruct
+    api_key_env_var: NG_OPENAI_API_KEY
+
+  - type: content_safety
+    engine: nim
+    model: nvidia/llama-3.1-nemoguard-8b-content-safety
+    api_key_env_var: NG_NVIDIA_API_KEY
+
+rails:
+  input:
+    flows:
+      - content safety check input $model=content_safety
+  output:
+    flows:
+      - content safety check output $model=content_safety

examples/configs/content_safety_api_keys/prompts.yml

@@ -0,0 +1,257 @@
+# These are the default prompts released by Meta, except for policy O7, which was added to address direct insults.
+prompts:
+
+  - task: content_safety_check_input $model=content_safety
+    content: |
+      Task: Check if there is unsafe content in the user or agent messages in the below conversation in accordance with our safety policy with the below categories. Provide a 'safe' or 'unsafe' rating for both the user message and the agent response. If there is unsafe content, please also provide a list of violated categories according to our safety policy below.
+
+      <BEGIN UNSAFE CONTENT CATEGORIES>
+      S1: Violence.
+      S2: Sexual.
+      S3: Criminal Planning/Confessions.
+      S4: Guns and Illegal Weapons.
+      S5: Controlled/Regulated Substances.
+      S6: Suicide and Self Harm.
+      S7: Sexual (minor).
+      S8: Hate/Identity Hate.
+      S9: PII/Privacy.
+      S10: Harassment.
+      S11: Threat.
+      S12: Profanity.
+      S13: Needs Caution.
+      S14: Other.
+      S15: Manipulation.
+      S16: Fraud/Deception.
+      S17: Malware.
+      S18: High Risk Gov Decision Making.
+      S19: Political/Misinformation/Conspiracy.
+      S20: Copyright/Trademark/Plagiarism.
+      S21: Unauthorized Advice.
+      S22: Illegal Activity.
+      S23: Immoral/Unethical.
+      <END UNSAFE CONTENT CATEGORIES>
+
+      <BEGIN CONVERSATION>
+
+      user: {{ user_input }}
+
+      <END CONVERSATION>
+
+      Provide your safety assessment for the above conversation in the below output JSON format:
+      {"User Safety": safe or unsafe rating for the user message, "Response Safety": safe or unsafe rating for the agent response. Omit if no agent response present. "Safety Categories": a comma-separated list of applicable safety categories from the provided taxonomy. Omit if all safe.}
+
+      Do not include anything other than the output JSON in your response.
+      Output JSON:
+    output_parser: nemoguard_parse_prompt_safety
+    max_tokens: 50
+
+  - task: content_safety_check_output $model=content_safety
+    content: |
+      Task: Check if there is unsafe content in the user or agent messages in the below conversation in accordance with our safety policy with the below categories. Provide a 'safe' or 'unsafe' rating for both the user message and the agent response. If there is unsafe content, please also provide a list of violated categories according to our safety policy below.
+
+      <BEGIN UNSAFE CONTENT CATEGORIES>
+      S1: Violence.
+      S2: Sexual.
+      S3: Criminal Planning/Confessions.
+      S4: Guns and Illegal Weapons.
+      S5: Controlled/Regulated Substances.
+      S6: Suicide and Self Harm.
+      S7: Sexual (minor).
+      S8: Hate/Identity Hate.
+      S9: PII/Privacy.
+      S10: Harassment.
+      S11: Threat.
+      S12: Profanity.
+      S13: Needs Caution.
+      S14: Other.
+      S15: Manipulation.
+      S16: Fraud/Deception.
+      S17: Malware.
+      S18: High Risk Gov Decision Making.
+      S19: Political/Misinformation/Conspiracy.
+      S20: Copyright/Trademark/Plagiarism.
+      S21: Unauthorized Advice.
+      S22: Illegal Activity.
+      S23: Immoral/Unethical.
+      <END UNSAFE CONTENT CATEGORIES>
+
+      <BEGIN CONVERSATION>
+
+      user: {{ user_input }}
+
+      response: agent: {{ bot_response }}
+
+      <END CONVERSATION>
+
+      Provide your safety assessment for the above conversation in the below output JSON format:
+      {"User Safety": safe or unsafe rating for the user message, "Response Safety": safe or unsafe rating for the agent response. Omit if no agent response present. "Safety Categories": a comma-separated list of applicable safety categories from the provided taxonomy. Omit if all safe.}
+
+      Do not include anything other than the output JSON in your response.
+      Output JSON:
+    output_parser: nemoguard_parse_response_safety
+    max_tokens: 50
+
+  - task: content_safety_check_input $model=llama_guard
+    content: |
+      <s>[INST] Task: Check if there is unsafe content in 'User' messages in conversations according our safety policy with the below categories.
+      <BEGIN UNSAFE CONTENT CATEGORIES>
+      O1: Violence and Hate.
+      Should not
+      - Help people plan or engage in violence.
+      - Encourage people to engage in violence.
+      - Express hateful, derogatory, or demeaning sentiments against people based on sensitive personal characteristics like their race, color, religion, national origin, sexual orientation, gender, gender identity, or disability.
+      - Encourage discrimination against people based on those characteristics and should not use slurs or create insults based on those characteristics.
+      Can
+      - Provide information on violence and discrimination.
+      - Discuss the topics of hate and violence, and can discuss historical events involving violence.
+      O2: Sexual Content.
+      Should not
+      - Engage in sexually explicit (i.e., erotic) conversations and should not encourage people to engage in sexual activities.
+      Can
+      - Discuss topics related to sexuality and sexual education.
+      - Provide general information regarding sexual activities.
+      - Provide guidance on topics related to sexual health.
+      O3: Criminal Planning.
+      Should not
+      - Help people plan or engage in criminal activities like theft, kidnapping, or financial crimes, especially when the person indicates that they may act on the information provided to them.
+      - Encourage people to engage in criminal activities or condone criminal behavior.
+      Can
+      - Help people understand how crimes are committed, provided that the response is not so detailed that it could be used to provide a roadmap for future crimes.
+      - Help people understand how systems work that are used to prevent crimes, provided those descriptions do not provide detailed information on how to defeat those systems.
+      O4: Guns and Illegal Weapons.
+      Should not
+      - Assist in the planning or execution of firearm or illegal weapons crimes, especially when the person indicates that they may act on the information provided to them.
+      - Encourage people to engage in firearm or illegal weapons crimes or condone such behavior.
+      Can
+      - Discuss firearms and the arguments for and against firearm ownership.
+      O5: Regulated or Controlled Substances.
+      Should not
+      - Encourage or assist people in the trafficking, consumption, or use of regulated or controlled substances, including but not limited to illegal drugs, tobacco, alcohol, and cannabis.
+      - Assist or encourage people to create such substances when it is illegal to do so.
+      Can
+      - Discuss the use of regulated or controlled substances, and can provide general information on those substances (ex: what their chemical composition is and what their effects are).
+      - Provide information on how those substances are manufactured, provided that the information is not so detailed as to provide step-by-step instructions for their creation.
+      O6: Self-Harm.
+      Should not
+      - Encourage people to harm themselves.
+      - Romanticize or condone self-harm or suicide.
+      - Provide information regarding the methods of suicide or self-harm.
+      - Provide any information or advice, other than information regarding appropriate health resources, to people who express or imply that they wish to harm themselves.
+      Can
+      - Provide general information regarding self-harm and suicide (ex: statistics, the mental health challenges faced by those who take these actions).
+      Should
+      - Provide information regarding appropriate health resources to people who imply or express the desire to harm themselves.
+      O7: Offensive Language and Insults.
+      Should not
+      - Address direct insults
+      - Use or encourage the use of swear words, bad language, or vulgar expressions.
+      - Insult, offend, or demean individuals or groups based on any personal characteristics or choices.
+      - Encourage or support the use of offensive or derogatory language.
+      Can
+      - Discuss the impact and societal perceptions of offensive language and insults.
+      - Provide information on communication strategies that avoid the use of offensive language and promote respectful dialogue.
+      - Offer guidance on how to address or respond to offensive language or insults in a constructive manner.
+      <END UNSAFE CONTENT CATEGORIES>
+
+      <BEGIN CONVERSATION>
+      User: {{ user_input }}
+      <END CONVERSATION>
+
+      Provide your safety assessment for User in the above conversation:
+      - First line must read 'safe' or 'unsafe'.
+      - If unsafe, a second line must include a comma-separated list of violated categories. [/INST]
+    output_parser: is_content_safe
+
+  - task: content_safety_check_output $model=llama_guard_2
+    content: |
+      <|begin_of_text|>[INST] Task: Check if there is unsafe content in Agent messages in conversations according our safety policy with the below categories.
+
+      <BEGIN UNSAFE CONTENT CATEGORIES>
+      S1: Violent Crimes.
+      S2: Non-Violent Crimes.
+      S3: Sex Crimes.
+      S4: Child Exploitation.
+      S5: Specialized Advice.
+      S6: Privacy.
+      S7: Intellectual Property.
+      S8: Indiscriminate Weapons.
+      S9: Hate.
+      S10: Self-Harm.
+      S11: Sexual Content.
+      <END UNSAFE CONTENT CATEGORIES>
+
+      <BEGIN CONVERSATION>
+
+      User: {{ user_input }}
+
+      Agent: {{ bot_response }}
+
+      <END CONVERSATION>
+
+      Provide your safety assessment for ONLY THE LAST Agent in the above conversation:
+      - First line must read 'safe' or 'unsafe'.
+      - If unsafe, a second line must include a comma-separated list of violated categories. [/INST]
+    output_parser: is_content_safe
+    max_tokens: 3
+
+  - task: content_safety_check_input $model=shieldgemma
+    content: |
+      You are a policy expert trying to help determine whether a user prompt is in
+      violation of the defined safety policies.
+
+      <start_of_turn>
+      Human Question: {{ user_input }}
+      <end_of_turn>
+
+      Our safety principle is defined in the below:
+
+      * "No Hate Speech": The prompt shall not contain or seek generation of content
+        that expresses, incites, or promotes hate based on race, gender, ethnicity,
+        religion, nationality, sexual orientation, disability status, or caste.
+      * "No Dangerous Content": The prompt shall not contain or seek generation of content that harming oneself
+        and/or others (e.g., accessing or building firearms and explosive devices, promotion of terrorism, instructions for suicide).
+      * "No Sexually Explicit Information": The prompt shall not contain or seek generation of content
+        that contains references to sexual acts or other lewd content (e.g., sexually graphic descriptions,
+        content aimed at causing arousal). Medical or scientific terms related to human anatomy or sex education are permitted.
+      * "No Harassment": The prompt shall not contain or seek generation of content that is malicious,
+        intimidating, bullying, or abusive content targeting another individual (e.g., physical threats,
+        denial of tragic events, disparaging victims of violence).
+      ===
+
+      Does the human prompt/question violate the above principle? Your answer must start with
+      'Yes' or 'No'. And then walk through step by step to be sure we answer
+      correctly.
+    output_parser: is_content_safe
+    max_tokens: 3
+
+  - task: content_safety_check_output $model=shieldgemma
+    content: |
+      You are a policy expert trying to help determine whether a chatbot response to the user is in
+      violation of the defined safety policies:
+
+      <start_of_turn>
+      Human Question: {{ user_input }}
+      Chatbot Response: {{ bot_response }}
+      <end_of_turn>
+
+      Our safety principle is defined in the below:
+
+      * "No Dangerous Content": The chatbot shall not generate content that harming oneself and/or others
+        (e.g., accessing or building firearms and explosive devices, promotion of terrorism, instructions for suicide).
+      * "No Harassment": The chatbot shall not generate content that is malicious, intimidating, bullying, or
+        abusive content targeting another individual (e.g., physical threats, denial of tragic events, disparaging victims of violence).
+      * "No Hate Speech": The chatbot shall not generate content that targeting identity and/or protected
+        attributes (e.g., racial slurs, promotion of discrimination, calls to violence against protected groups) or dehumanizing/belittling/vilifying on the
+        bases of race, ethnicity, religion, disability, age, nationality, veteran status, sexual orientation, gender, gender identity, caste, or any other protected status.
+        We do allow content that includes a reference to a protected individual or group in a positive manner.
+      * "No Sexually Explicit Information": The chatbot shall not generate content that contains references to sexual acts or
+        other lewd content (e.g., sexually graphic descriptions, content aimed at causing arousal).
+        Medical or scientific terms related to human anatomy or sex education are permitted.
+
+      ===
+
+      Does the chatbot response violate the above principle? Your answer must start with
+      'Yes' or 'No'. And then walk through step by step to be sure we answer
+      correctly.
+    output_parser: is_content_safe
+    max_tokens: 3

nemoguardrails/llm/models/initializer.py

@@ -38,6 +38,7 @@ def init_llm_model(
     Args:
         model_name: Name of the model to initialize
         provider_name: Name of the provider to use
+        mode: Literal taking either "chat" or "text" values
         kwargs: Additional arguments to pass to the model initialization
 
     Returns:
@@ -48,7 +49,10 @@ def init_llm_model(
     """
     # currently we only support LangChain models
     return init_langchain_model(
-        model_name=model_name, provider_name=provider_name, mode=mode, kwargs=kwargs
+        model_name=model_name,
+        provider_name=provider_name,
+        mode=mode,
+        kwargs=kwargs,
     )
 
 

nemoguardrails/llm/models/langchain_initializer.py

@@ -103,7 +103,9 @@ def try_initialization_method(
             f"Trying initializer: {initializer.init_method.__name__} for model: {model_name} and provider: {provider_name}"
         )
         result = initializer.execute(
-            model_name=model_name, provider_name=provider_name, kwargs=kwargs
+            model_name=model_name,
+            provider_name=provider_name,
+            kwargs=kwargs,
         )
         log.debug(f"Initializer {initializer.init_method.__name__} returned: {result}")
         if result is not None:
@@ -213,7 +215,7 @@ def _init_chat_completion_model(
 
     # just to document the expected behavior
     # we don't support pre-0.2.7 versions of langchain-core it is in
-    # line wiht our pyproject.toml
+    # line with our pyproject.toml
     package_version = version("langchain-core")
 
     if _parse_version(package_version) < (0, 2, 7):
@@ -225,6 +227,7 @@ def _init_chat_completion_model(
         return init_chat_model(
             model=model_name,
             model_provider=provider_name,
+            **kwargs,
         )
     except ValueError:
         raise
@@ -250,7 +253,6 @@ def _init_text_completion_model(
     if provider_cls is None:
         raise ValueError()
     kwargs = _update_model_kwargs(provider_cls, model_name, kwargs)
-
     return provider_cls(**kwargs)
 
 

nemoguardrails/rails/llm/config.py

@@ -28,6 +28,7 @@
     ValidationError,
     model_validator,
     root_validator,
+    validator,
 )
 from pydantic.fields import Field
 
@@ -101,7 +102,10 @@ class Model(BaseModel):
         default=None,
         description="The name of the model. If not specified, it should be specified through the parameters attribute.",
     )
-
+    api_key_env_var: Optional[str] = Field(
+        default=None,
+        description='Optional environment variable with model\'s API Key. Do not include "$".',
+    )
     reasoning_config: Optional[ReasoningModelConfig] = Field(
         default_factory=ReasoningModelConfig,
         description="Configuration parameters for reasoning LLMs.",
@@ -1352,6 +1356,17 @@ def fill_in_default_values_for_v2_x(cls, values):
 
         return values
 
+    @validator("models")
+    def validate_models_api_key_env_var(cls, models):
+        """Model API Key Env var must be set to make LLM calls"""
+        api_keys = [m.api_key_env_var for m in models]
+        for api_key in api_keys:
+            if api_key and not os.environ.get(api_key):
+                raise ValueError(
+                    f"Model API Key environment variable '{api_key}' not set."
+                )
+        return models
+
     raw_llm_call_action: Optional[str] = Field(
         default="raw llm call",
         description="The name of the action that would execute the original raw LLM call. ",