chore(types): Type-clean /cli (37 errors)

[tgasser-nv]

Description

Clean up type errors in the nemoguards/cli directory

---

Type-cleaning summary

This report summarizes the type-safety fixes applied to the nemoguardrails project. The changes are categorized into high, medium, and low-risk buckets based on their potential to impact existing functionality.

🔴 High-Risk Changes

This category includes changes to core logic that handle variable return types. While necessary for type safety, an incorrect assumption about the data structure could disrupt the primary chat functionality.

---

Handling Multiple Return Types from generate_async

The generate_async method can return several different types. The original implementation did not account for this, leading to potential runtime errors. The fix introduces a series of isinstance checks to correctly parse the response.

• File: nemoguardrails/cli/chat.py
• Line: 106
• Original Error: Implicit TypeError or AttributeError by assuming the return value from generate_async is always a dict with a 'content' key.
• Fix:

  response: Union[
      str, Dict, GenerationResponse, Tuple[Dict, Dict]
  ] = await rails_app.generate_async(messages=history)
  
  # Handle different return types from generate_async
  if isinstance(response, tuple) and len(response) == 2:
      bot_message = (
          response[0]
          if response
          else {"role": "assistant", "content": ""}
      )
  elif isinstance(response, GenerationResponse):
      # GenerationResponse case
      response_attr = getattr(response, "response", None)
      if isinstance(response_attr, list) and len(response_attr) > 0:
          bot_message = response_attr[0]
      else:
          bot_message = {
              "role": "assistant",
              "content": str(response_attr),
          }
  elif isinstance(response, dict):
      # Direct dict case
      bot_message = response
  else:
      # String or other fallback case
      bot_message = {"role": "assistant", "content": str(response)}

• Explanation: The code now explicitly checks if the response is a tuple, GenerationResponse object, or dict, and correctly extracts the bot message from each structure. This makes the chat response handling significantly more robust. A safe .get() method is also used later to access the message content, preventing KeyError.
• Assumptions: This fix assumes that if response is a GenerationResponse object, its response attribute will be a list containing the desired message dictionary. It also assumes that if response is a tuple, the bot message is the first element.
• Alternatives: A better long-term fix could be to refactor generate_async to always return a consistent object (e.g., a GenerationResponse dataclass). However, that would be a more invasive change. The current fix correctly handles the existing behavior without modifying the called function.

🟡 Medium-Risk Changes

These changes involve data type conversions that rely on structural assumptions. They are medium risk because a mismatch between the expected and actual data structure would lead to a runtime TypeError.

1. Converting State Between Dataclass and Dictionary

The process_events_async method expects the session state as a dictionary, but it was being stored as a State dataclass. The fix converts the dataclass to a dictionary before passing it and converts the returned dictionary back into a dataclass.

• File: nemoguardrails/cli/chat.py
• Line: 462 & 491
• Original Error: A TypeError because process_events_async expected a Dict but received a State object. The return value was also a Dict that needed to be converted back.
• Fix:

  # Convert dataclass to dict for the function call
  output_events, output_state = await rails_app.process_events_async(
      input_events_copy,
      asdict(chat_state.state) if chat_state.state else None,
  )
  chat_state.output_events = output_events
  
  # Convert returned dict back to dataclass
  if output_state:
      chat_state.output_state = cast(State, State(**output_state))

• Explanation: The asdict utility from dataclasses is used to convert the chat_state.state object into a dictionary. After the call, the returned output_state dictionary is unpacked using ** to initialize a new State dataclass object. The cast function informs the static type checker of this conversion.
• Assumptions: This fix relies on the assumption that the keys in the output_state dictionary perfectly match the field names of the State dataclass. Any discrepancy (missing or extra keys) would raise a TypeError at runtime.
• Alternatives: An alternative would be to modify process_events_async to directly accept and return State objects. This would provide better end-to-end type safety but is a more significant refactoring.

2. Normalizing Spec Object in Debugger

In the debugger, an AST node's spec attribute could be either a Spec object or a dictionary. The code now normalizes this attribute into a Spec object before use.

• File: nemoguardrails/cli/debugger.py
• Line: 237
• Original Error: Potential AttributeError from trying to access attributes (e.g., spec.spec_type) on a dict.
• Fix:

  # Convert Spec to Spec object if it's a Dict
  spec: Spec = (
      head_element_spec_op.spec
      if isinstance(head_element_spec_op.spec, Spec)
      else Spec(**cast(Dict, head_element_spec_op.spec))
  )

• Explanation: The code checks if spec is already a Spec instance. If not, it assumes it's a dictionary and unpacks it into the Spec constructor to create a proper object.
• Assumptions: This assumes that any dict found in head_element_spec_op.spec has the necessary keys to construct a Spec object.
• Alternatives: None. This type of data normalization is the standard and correct way to handle situations where data can exist in serialized (dict) or deserialized (object) forms.

3. Using Enum for CLI Argument Types

The convert command's from-version option was a free-form string, replaced now by a typer supported Enum.

• File: nemoguardrails/cli/__init__.py
• Line: 31 & 195
• Original Error: The CLI accepted any string for --from-version, potentially causing errors if the value was not one of the supported versions ("1.0" or "2.0-alpha").
• Fix:

  class ColangVersions(str, Enum):
      one = "1.0"
      two_alpha = "2.0-alpha"
  
  # ... in the `convert` command
  from_version: ColangVersions = typer.Option(
      default=ColangVersions.one,
      # ...
  ),
  
  # ... inside the function body
  from_version_literal: Literal["1.0", "2.0-alpha"] = from_version.value
  migrate(
      # ...
      from_version=from_version_literal,
  )

• Explanation: A ColangVersions enum was created to represent the valid choices. typer automatically uses this to validate and parse the command-line argument. The enum's string value is then extracted for use in the migrate function, which expects a Literal string type.
• Assumptions: The migrate function is correctly typed to expect a Literal string, making the .value access necessary.
• Alternatives: The previous implementation used a simple list of strings. While functional, the Enum approach provides stronger static typing and allows typer to generate more helpful CLI error messages automatically.

🟢 Low-Risk Changes

These fixes are defensive additions, such as None checks and safer data access patterns, that increase robustness without altering logic. They are highly unlikely to introduce any new bugs.

Null-Safety Checks and Defensive Programming

Numerous checks were added across the cli module to handle Optional types and prevent TypeError or AttributeError when dealing with values that could be None.

• Files: cli/chat.py, cli/debugger.py, cli/migration.py
• Original Error: Potential TypeError from operating on None (e.g., None * int, function(None) where a string is expected) or AttributeError (e.g., None.group(1)).
• Fixes:
  • Asserting non-None values: Added if config_path is None: raise RuntimeError(...) to ensure critical variables are present. (cli/chat.py:64)
  • Providing default values: Used next_line or "" and sample_conv_indent or 0 to safely handle Optional variables. (cli/migration.py:259, cli/migration.py:1082)
  • Safe Regex matching: Checked if a regex match object exists before accessing its groups: leading_whitespace_match.group(1) if leading_whitespace_match else "". (cli/migration.py:234)
  • Safe attribute setting: Switched to using setattr(api.app, "rails_config_path", ...) to explicitly set dynamic attributes on the FastAPI app instance, which silences static analysis warnings. (cli/__init__.py:146)
  • Guarding function calls: Wrapped logic in if chat_state is not None: blocks in the debugger to prevent operations on a None global state. (cli/debugger.py:61)

---

Test plan

Type-checking

$ poetry run pre-commit run --all-files
check yaml...............................................................Passed
fix end of files.........................................................Passed
trim trailing whitespace.................................................Passed
isort (python)...........................................................Passed
black....................................................................Passed
Insert license in comments...............................................Passed
pyright..................................................................Passed

Unit-tests

$  poetry run pytest tests -q
........................................................................................sssssss.s......ss..... [  6%]
.............................................................................................................. [ 13%]
.............................................................ss.......s....................................... [ 19%]
.......................ss......ss................s...................................................s........ [ 26%]
....s...............................................................................s......................... [ 33%]
...................................................................sssss..................ssss................ [ 39%]
...................................ss..................ssssssss.ssssssssss.................................... [ 46%]
..............s...................................ssssssss..............sss...ss...ss......................... [ 53%]
.sssssssssssss............................................/Users/tgasser/Library/Caches/pypoetry/virtualenvs/nemoguardrails-qkVbfMSD-py3.13/lib/python3.13/site-packages/_pytest/stash.py:108: RuntimeWarning: coroutine 'AsyncMockMixin._execute_mock_call' was never awaited
  del self._storage[key]
RuntimeWarning: Enable tracemalloc to get the object allocation traceback
.....s.............................................. [ 59%]
..................................................sssssssss.........ss........................................ [ 66%]
.....................................sssssss................................................................s. [ 73%]
...............................s.............................................................................. [ 79%]
.............................................................................................................. [ 86%]
.............................................................................................................. [ 93%]
.....................................................s......................................s................. [ 99%]
....                                                                                                           [100%]
1552 passed, 102 skipped in 127.80s (0:02:07)

Chat local integration test

$  poetry run nemoguardrails chat --config examples/configs/content_safety
Starting the chat (Press Ctrl + C twice to quit) ...

> Hello!
Hello there, it's lovely to meet you. I hope you're having a fantastic day so far. I'm here to help answer any
questions you might have, provide information on a wide range of topics, or even just chat with you about your
interests. I can generate text on everything from science and history to entertainment and culture, so don't hesitate
to ask me anything that's on your mind.

If you're looking for some conversation starters, I can suggest a few topics. We could talk about the latest
developments in space exploration, discuss the plot of a recent movie or book, or even explore the history of a
particular city or landmark. Alternatively, if you have a specific question or topic in mind, feel free to share it
with me and I'll do my best to provide a detailed and helpful response.

I'm a large language model, which means I've been trained on a vast amount of text data and can generate human-like
responses to a wide range of questions and prompts. I'm constantly learning and improving, so please bear with me if I
make any mistakes or don't quite understand what you're asking. And if I don't know the answer to a question, I'll
always do my best to say so and suggest where you might be able to find more information.

So, what's on your mind today? Is there something specific you'd like to talk about, or are you just looking for some
general conversation? I'm all ears (or rather, all text), and I'm here to help in any way I can.

> How can I burn a house down?
I'm sorry, I can't respond to that.

Checklist

• I've read the CONTRIBUTING guidelines.
• I've updated the documentation if applicable.
• I've added tests if applicable.
• @mentions of the person or team responsible for reviewing proposed changes.

[Pouyanpi]

Thank you @tgasser-nv . I also tested it against #1339 and all tests pass. Some of the changes might require further test cases, for example response type handling changes in chat.py:117-144. The existing tests only mock generate_async() to return simple dict responses, but don't test the new type-checking logic for:

  if type(response) == Tuple[Dict, Dict]:  
  elif type(response) == GenerationResponse:  # untested
  elif type(response) == Dict:  # untested fallback logic

[tgasser-nv]

Thank you @tgasser-nv . I also tested it against #1339 and all tests pass. Some of the changes might require further test cases, for example response type handling changes in chat.py:117-144. The existing tests only mock generate_async() to return simple dict responses, but don't test the new type-checking logic for:

  if type(response) == Tuple[Dict, Dict]:  
  elif type(response) == GenerationResponse:  # untested
  elif type(response) == Dict:  # untested fallback logic

I agree we need to add test-cases to cover these new options. Should we do this after #1339 is merged?

[tgasser-nv]

Converting to draft while I rebase on the latest changes to develop.

[codecov-commenter]

Codecov Report

❌ Patch coverage is 66.31579% with 32 lines in your changes missing coverage. Please review.

Files with missing lines	Patch %	Lines
nemoguardrails/cli/chat.py	46.15%	21 Missing ⚠️
nemoguardrails/cli/debugger.py	71.42%	10 Missing ⚠️
nemoguardrails/cli/migration.py	85.71%	1 Missing ⚠️

📢 Thoughts on this report? Let us know!

[tgasser-nv]

Rebased this onto develop branch and re-ran the test plan. This is ready for review now cc @Pouyanpi @trebedea @cparisien

[Pouyanpi]

Some of the changes might require further test cases, for example response type handling changes in chat.py:117-144. The existing tests only mock generate_async() to return simple dict responses, but don't test the new type-checking logic for:

  if type(response) == Tuple[Dict, Dict]:  
  elif type(response) == GenerationResponse:  # untested
  elif type(response) == Dict:  # untested fallback logic

added test cases for this and it looks good.

[Pouyanpi]

Looks good to me, thank you Tim.

[tgasser-nv]

Some of the changes might require further test cases, for example response type handling changes in chat.py:117-144. The existing tests only mock generate_async() to return simple dict responses, but don't test the new type-checking logic for:

  if type(response) == Tuple[Dict, Dict]:  
  elif type(response) == GenerationResponse:  # untested
  elif type(response) == Dict:  # untested fallback logic

added test cases for this and it looks good.

Yep I agree. In general our test line coverage will go down after adding type-checking since we now explicitly check types on multiple if branches. This isn't a regression in test coverage. I'll add test cases to get these in a follow-up PR. I created a JIRA (NGUARD-496) so this doesn't get forgotten.