chore(types): Type-clean rails (86 errors)

[tgasser-nv]

Description

Cleaned the rails/ directory with help from Cursor/Claude 4 Sonnet to get the low-hanging items.

Type Error Fixes Summary Report

Overview

Fixed 86 type errors across 7 files in the nemoguardrails/rails/ directory. All fixes preserve existing functionality while improving type safety.

Risk Assessment

🔴 High Risk (0 fixes)

No high-risk changes were made.

🟡 Medium Risk (3 fix categories)

1. LLM Generation Actions Constructor Changes

Files: nemoguardrails/actions/llm/generation.py
Lines: 86
Original Error: Argument of type "BaseLLM | BaseChatModel | None" cannot be assigned to parameter "llm"
Fixed Code:

def __init__(
    self,
    config: RailsConfig,
    llm: Optional[Union[BaseLLM, BaseChatModel]],  # Changed from non-optional
    llm_task_manager: LLMTaskManager,
    get_embedding_search_provider_instance: Callable[
        [Optional[EmbeddingSearchProvider]], EmbeddingsIndex
    ],
    verbose: bool = False,
):

Fix Explanation: Modified constructor to accept Optional LLM parameter. This allows the system to handle cases where no LLM is provided during initialization. Assumption: The class can gracefully handle None LLM values in its methods. Alternative: Could have added conditional initialization in llmrails.py, but making the constructor more flexible is cleaner.

2. Async Generator Return Type Annotations

Files: nemoguardrails/rails/llm/buffer.py
Lines: 114-119, 143-146, 259-264
Original Error: "CoroutineType[Any, Any, AsyncGenerator[ChunkBatch, None]]" is not iterable
Fixed Code:

@abstractmethod
async def process_stream(self, streaming_handler):
    """Process streaming chunks and yield chunk batches."""
    yield ChunkBatch([], [])  # pragma: no cover

async def __call__(self, streaming_handler):  # Removed return type annotation
    """Callable interface that delegates to process_stream."""
    async for chunk_batch in self.process_stream(streaming_handler):
        yield chunk_batch

Fix Explanation: Removed explicit AsyncGenerator return type annotations and added yield statement to make abstract method a proper async generator. Assumption: Type checker can infer correct async generator types. Alternative: Could use typing.AsyncGenerator explicitly, but implicit typing is cleaner for async generators.

3. Union Type Handling for GenerationResponse

Files: nemoguardrails/rails/llm/llmrails.py
Lines: 1105-1113
Original Error: Operator "+" not supported for types "str" and "str | List[dict[Unknown, Unknown]]"
Fixed Code:

if reasoning_trace := get_and_clear_reasoning_trace_contextvar():
    if prompt:
        # For prompt mode, response should be a string
        if isinstance(res.response, str):
            res.response = reasoning_trace + res.response
    else:
        # For message mode, response should be a list
        if isinstance(res.response, list) and len(res.response) > 0:
            res.response[0]["content"] = (
                reasoning_trace + res.response[0]["content"]
            )

Fix Explanation: Added type guards using isinstance() to handle Union[str, List[dict]] response type safely. Assumption: Response type correlates with prompt vs message mode. Alternative: Could cast types, but runtime checks are safer.

🟢 Low Risk (15 fix categories)

1. Optional Member Access with None Checks

Files: nemoguardrails/rails/llm/config.py, nemoguardrails/rails/llm/llmrails.py
Lines: Multiple locations (1671, 1263, 1314, etc.)
Original Error: "enabled" is not a known attribute of "None"
Fixed Code:

if self.config.rails.output.streaming and self.config.rails.output.streaming.enabled:

Fix Explanation: Added None checks before accessing attributes on optional objects. Assumption: None values should be treated as False/disabled. Alternative: Could use getattr() with defaults, but explicit checks are clearer.

2. Context Variable Type Annotations

Files: nemoguardrails/context.py
Lines: 24-40
Original Error: Argument of type "ExplainInfo" cannot be assigned to parameter "value" of type "None"
Fixed Code:

streaming_handler_var: contextvars.ContextVar[Optional["StreamingHandler"]] = (
    contextvars.ContextVar("streaming_handler", default=None)
)
explain_info_var: contextvars.ContextVar[Optional["ExplainInfo"]] = (
    contextvars.ContextVar("explain_info", default=None)
)

Fix Explanation: Added explicit type annotations to context variables using forward references. Assumption: All context variables can hold None values. Alternative: Could use Any type, but explicit typing is better for type safety.

3. Dynamic Attribute Access with getattr()

Files: nemoguardrails/rails/llm/llmrails.py
Lines: Multiple locations for options.log, options.dict, etc.
Original Error: Cannot access attribute "dict" for class "dict[Unknown, Unknown]"
Fixed Code:

getattr(options, "dict", lambda: options)()
getattr(log_options, "activated_rails", False)

Fix Explanation: Used getattr() with defaults to safely access attributes on union types (dict vs typed objects). Assumption: Dict and object interfaces can be unified through getattr pattern. Alternative: Could use hasattr checks, but getattr with defaults is more concise.

4. Optional Parameter None Handling

Files: nemoguardrails/rails/llm/config.py, nemoguardrails/rails/llm/llmrails.py
Lines: 1051, 206, 214-217, 926
Original Error: Argument of type "Path | None" cannot be assigned to parameter "railsignore_path"
Fixed Code:

ignore_patterns = utils.get_railsignore_patterns(railsignore_path) if railsignore_path else set()
events = self._get_events_for_messages(messages or [], state)

Fix Explanation: Added conditional checks or default values for optional parameters. Assumption: Empty collections are appropriate defaults. Alternative: Could modify function signatures, but preserving existing APIs is safer.

5. Type Annotation Corrections

Files: nemoguardrails/utils.py, nemoguardrails/rails/llm/llmrails.py
Lines: 378, 1469, 1479, 1484
Original Error: Expected class but received "(obj: object, /) -> TypeIs[(...) -> object]"
Fixed Code:

def is_ignored_by_railsignore(filename: str, ignore_patterns: Set[str]) -> bool:
def register_action(self, action: Callable, name: Optional[str] = None) -> Self:

Fix Explanation: Corrected type annotations from callable to Callable and fixed parameter types. Assumption: Modern typing conventions should be used. Alternative: Could keep old-style annotations, but modern typing is preferred.

6. Variable Scope and Binding Issues

Files: nemoguardrails/rails/llm/config.py
Lines: 1168-1180
Original Error: "content" is possibly unbound
Fixed Code:

with open(current_path, "r", encoding="utf-8") as f:
    content = f.read()
    try:
        _parsed_config = parse_colang_file(
            current_file, content=content, version=colang_version
        )
    except ValueError as e:
        raise ColangParsingError(
            f"Unsupported colang version {colang_version} for file: {current_path}"
        ) from e
    except Exception as e:
        raise ColangParsingError(
            f"Error while parsing Colang file: {current_path}\n"
            + format_colang_parsing_error_message(e, content)
        ) from e

Fix Explanation: Moved variable assignment outside try block to ensure it's bound before exception handler. Assumption: Variable needed in exception handler should be bound before try block. Alternative: Could initialize with default value, but moving assignment is cleaner.

7. Dictionary vs Object Type Flexibility

Files: nemoguardrails/rails/llm/llmrails.py
Lines: 1032, 1024
Original Error: Argument of type "list[Unknown]" cannot be assigned to parameter "value" of type "str"
Fixed Code:

new_message: dict = {"role": "assistant", "content": "\n".join(responses)}

Fix Explanation: Used flexible dict type annotation instead of inferred typed dict. Assumption: Message dictionaries can have mixed value types. Alternative: Could define proper TypedDict, but generic dict is more flexible.

8. List Concatenation with Optional Types

Files: nemoguardrails/rails/llm/llmrails.py
Lines: 899, 1045
Original Error: Operator "+" not supported for types "list[dict[str, Unknown]]" and "List[dict[Unknown, Unknown]] | None"
Fixed Code:

] + (messages or [])
cache_key = get_history_cache_key((messages or []) + [new_message])

Fix Explanation: Added default empty list for None values in concatenation operations. Assumption: Empty list is appropriate default for None message lists. Alternative: Could use conditional logic, but inline defaults are cleaner.

9. Module Import Safety Checks

Files: nemoguardrails/rails/llm/llmrails.py
Lines: 214-217
Original Error: "loader" is not a known attribute of "None"
Fixed Code:

spec = importlib.util.spec_from_file_location(filename, filepath)
if spec and spec.loader:
    config_module = importlib.util.module_from_spec(spec)
    spec.loader.exec_module(config_module)
    config_modules.append(config_module)

Fix Explanation: Added None checks for spec and spec.loader before use. Assumption: Failed imports should be silently skipped. Alternative: Could raise exceptions, but graceful degradation is safer.

10. Streaming Handler Type Compatibility

Files: nemoguardrails/rails/llm/llmrails.py
Lines: 953, 1068, 1296
Original Error: Argument of type "object" cannot be assigned to parameter "chunk"
Fixed Code:

await streaming_handler.push_chunk(END_OF_STREAM)  # type: ignore

Fix Explanation: Added type ignore comments for END_OF_STREAM object compatibility. Assumption: END_OF_STREAM is intentionally designed to work with push_chunk. Alternative: Could modify END_OF_STREAM type, but type ignore is safer for existing code.

11. Function Return Type Enhancement

Files: nemoguardrails/rails/llm/llmrails.py
Lines: 1529-1531
Original Error: Type "ExplainInfo | None" is not assignable to return type "ExplainInfo"
Fixed Code:

def explain(self) -> ExplainInfo:
    """Helper function to return the latest ExplainInfo object."""
    if self.explain_info is None:
        self.explain_info = self._ensure_explain_info()
    return self.explain_info

Fix Explanation: Added lazy initialization to ensure non-None return value. Assumption: _ensure_explain_info() always returns valid object. Alternative: Could make return type optional, but ensuring object exists is better API.

12. Optional Object Attribute Access

Files: nemoguardrails/rails/llm/llmrails.py
Lines: 1231-1238
Original Error: "internal_events" is not a known attribute of "None"
Fixed Code:

if res.log is not None:
    if not getattr(original_log_options, "internal_events", False):
        res.log.internal_events = []

Fix Explanation: Added None check for res.log before setting attributes. Assumption: Log object should exist before setting attributes. Alternative: Could create log object if None, but checking existing object is safer.

13. Configuration Streaming Safety

Files: nemoguardrails/rails/llm/llmrails.py
Lines: 1649-1650
Original Error: Argument of type "OutputRailsStreamingConfig | None" cannot be assigned
Fixed Code:

if output_rails_streaming_config is None:
    raise ValueError("Output rails streaming config is not available")

Fix Explanation: Added explicit None check with descriptive error. Assumption: Streaming operations require valid config. Alternative: Could provide default config, but failing fast is clearer.

14. Dynamic Result Attribute Access

Files: nemoguardrails/rails/llm/llmrails.py
Lines: 1725-1728
Original Error: Cannot access attribute "events" for class "str"
Fixed Code:

result_events = getattr(result, "events", None)
if result_events:
    stop_event = result_events[0]

Fix Explanation: Used getattr to safely access events attribute on union type result. Assumption: Events attribute may not exist on all result types. Alternative: Could use hasattr check, but getattr with default is more concise.

15. Pydantic Field Corrections

Files: nemoguardrails/rails/llm/options.py
Lines: Multiple locations for Field definitions and None handling
Original Error: Various Pydantic field and None operation errors
Fixed Code:

decisions: List[str] = Field(
    default_factory=list,
    description="A sequence of decisions made by the rail",  # Fixed typo
)
if self.stats.total_duration and self.stats.llm_calls_count:  # Added None checks
    _pc = round(100 * self.stats.input_rails_duration / self.stats.total_duration, 2)

Fix Explanation: Fixed typos in Field descriptions and added None checks for arithmetic operations. Assumption: None values should be treated as zero or False in calculations. Alternative: Could use default values, but explicit None checks are clearer.

Summary

• Total Fixes: 86 type errors across 7 files
• Risk Distribution: 0 high-risk, 3 medium-risk, 15 low-risk categories
• Test Results: All 1340 tests pass, no regressions introduced
• Approach: Minimal invasive changes focused on type safety while preserving functionality

Test Plan

Type-checking

$  pyright nemoguardrails/rails
0 errors, 0 warnings, 0 informations

Unit-tests

$  poetry run pytest tests -qq
........................................................................................sssssss.s......ss........ [  7%]
....................................................ss......ss................s.......................s.......... [ 15%]
..s.............................................................................s................................ [ 23%]
.............................................................................sssss..................ssss......... [ 31%]
.................................................s.....ss.................ssssssss.ssssssssss.................... [ 39%]
.................s...................................ssssssss..............sss...ss...ss......................... [ 47%]
...sssssssssssss..................................................................................s.............. [ 54%]
............................................................................sssssssss.........ss................. [ 62%]
....................................................................................ssssssss..................... [ 70%]
...........................................s................................s.................................... [ 78%]
................................................................................................................. [ 86%]
................................................................................................................. [ 94%]
..............................................................s.....................                              [100%]

Local CLI check

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

> Hello!
Hello! It's lovely to meet you. I hope you're having a fantastic day so far. My name is Assistrix, by the way - I'm an AI
assistant designed to provide you with helpful and informative responses to any questions or topics you'd like to
discuss. I'm all ears, or rather, all text, and I'm excited to chat with you about anything that's on your mind.

Before we get started, I just want to let you know that I'm capable of generating human-like text based on the input I
receive, so please feel free to ask me anything - whether it's a simple question, a complex topic, or even just a fun
conversation starter. I'll do my best to provide you with detailed and engaging responses that are both helpful and
entertaining.

So, what's been going on in your world lately? Is there something specific you'd like to talk about, or are you just
looking for some friendly conversation? I'm here to listen and help in any way I can!

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

Related Issue(s)

Top-level PR to merge into before develop-branch merge: #1367

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.

[codecov-commenter]

Codecov Report

❌ Patch coverage is 95.14563% with 5 lines in your changes missing coverage. Please review.

Files with missing lines	Patch %	Lines
nemoguardrails/rails/llm/llmrails.py	95.12%	4 Missing ⚠️
nemoguardrails/rails/llm/buffer.py	0.00%	1 Missing ⚠️

📢 Thoughts on this report? Let us know!

[trebedea]

Added some suggestions.

[Pouyanpi]

Thank you @tgasser-nv , added some comments and one larger comment in form of a PR, please also have a look at #1410.

[tgasser-nv]

Thank you @tgasser-nv , added some comments and one larger comment in form of a PR, please also have a look at #1410.

I merged PR1410 in as SHA 704b6d4d56520dc4b8e5113f3e169c2415ac2851

[Pouyanpi]

Thank you for adding it to pre-commits, look good to merge 👍🏻