chore(types): Type-clean embeddings/ (25 errors) #1383

tgasser-nv · 2025-09-09T16:40:48Z

Description

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

Type Error Fixes Report - Embeddings Module

Overview

This report analyzes the type error fixes made in commit 5517ce81fe11bf22ab9787f2abb5bc488c30a41d to address 31 type errors in the embeddings module. The changes focus on improving type safety, adding proper null checks, and handling optional dependencies correctly.

Risk Assessment by Category

🔴 High Risk Fixes (2 fixes)

1. Property-based Attribute Access Pattern Change

Files: basic.py (lines 48-99), cache.py (lines 38, 86, 47, 51, 106, 110)
Original Errors:

Declaration "embedding_size" is obscured by a declaration of the same name (reportRedeclaration)
Declaration "cache_config" is obscured by a declaration of the same name (reportRedeclaration)
Declaration "embeddings" is obscured by a declaration of the same name (reportRedeclaration)

Fixed Code:

# Instance attributes are defined in __init__ and accessed via properties
class BasicEmbeddingsIndex(EmbeddingsIndex):
    def __init__(self, ...):
        self._items: List[IndexItem] = []
        self._embeddings: List[List[float]] = []
        self.embedding_model: Optional[str] = embedding_model
        # ... other attributes

Fix Explanation: Removed class-level attribute declarations that were causing redeclaration errors and moved to instance-level attributes with proper type annotations. Added name: str class attribute declarations in abstract base classes.

Alternative Fixes: Could have used ClassVar annotations to distinguish class vs instance attributes, but the implemented approach is cleaner and follows Python best practices.

Risk Justification: High risk because this is a significant architectural change that moves from class-level to instance-level attribute definitions. While functionally equivalent, it changes the object model and could affect introspection or dynamic attribute access patterns.

2. Embedding Model Initialization with Null Checks

Files: basic.py (lines 125-140, 152-157)
Original Errors:

"encode_async" is not a known attribute of "None" (reportOptionalMemberAccess)
Cannot assign to attribute "embedding_model" for class "BasicEmbeddingsIndex*" Type "Unknown | None" is not assignable to type "str"

Fixed Code:

def _init_model(self):
    # Provide defaults if not specified
    model = self.embedding_model or "sentence-transformers/all-MiniLM-L6-v2"
    engine = self.embedding_engine or "SentenceTransformers"
    
    self._model = init_embedding_model(
        embedding_model=model,
        embedding_engine=engine,
        embedding_params=self.embedding_params,
    )
    
    if not self._model:
        raise ValueError(f"Couldn't create embedding model with model {model} and engine {engine}")

# In _get_embeddings:
if not self._model:
    raise Exception("Couldn't initialize embedding model")

Fix Explanation: Added explicit validation and default values for embedding model initialization, with explicit exceptions when model creation fails.

Alternative Fixes: Could have used assertions or type guards, but explicit exceptions provide better error messages for debugging.

Risk Justification: High risk because it changes the error handling behavior by adding explicit validation and raising exceptions earlier. This could affect existing error handling patterns and introduce new failure modes.

🟡 Medium Risk Fixes (5 fixes)

3. Async Event Initialization Guards

Files: basic.py (lines 203-206, 218-221, 264-270)
Original Errors:

"wait" is not a known attribute of "None" (reportOptionalMemberAccess)
"set" is not a known attribute of "None" (reportOptionalMemberAccess)
Type "Event | None" is not assignable to declared type "Event"

Fixed Code:

if not self._current_batch_full_event:
    raise Exception("self._current_batch_full_event not initialized")

assert self._current_batch_full_event is not None
done, pending = await asyncio.wait([...], ...)

if not self._current_batch_finished_event:
    raise Exception("self._current_batch_finished_event not initialized")

assert self._current_batch_finished_event is not None
batch_event: asyncio.Event = self._current_batch_finished_event

Fix Explanation: Added runtime checks that will fail fast if the batching system is used incorrectly. The assertions help with type narrowing but could cause crashes in edge cases where the events aren't properly initialized.

Alternative Fixes: Could have used conditional attribute access with optional chaining, but explicit checks provide better error reporting.

Risk Justification: Medium risk because it adds runtime checks that will fail fast if the batching system is used incorrectly, potentially causing crashes in edge cases.

4. Optional Parameter Type Annotations

Files: basic.py (lines 52-57), cache.py (lines 159, 221-223)
Original Errors:

Expression of type "None" cannot be assigned to parameter of type "EmbeddingsCacheConfig | Dict[str, Any]"
Expression of type "None" cannot be assigned to parameter of type "float"

Fixed Code:

def __init__(
    self,
    embedding_model: Optional[str] = None,
    embedding_engine: Optional[str] = None,
    embedding_params: Optional[Dict[str, Any]] = None,
    index: Optional[AnnoyIndex] = None,
    cache_config: Optional[Union[EmbeddingsCacheConfig, Dict[str, Any]]] = None,
    search_threshold: Optional[float] = None,
):

Fix Explanation: Added proper Optional type annotations to parameters that can accept None values, making the API contract more explicit.

Alternative Fixes: Could have provided default values instead of allowing None, but Optional maintains backward compatibility.

Risk Justification: Medium risk because it improves type safety without changing functionality, but makes the API contract more explicit which could reveal existing usage errors.

5. Cache Store Type Safety Improvements

Files: cache.py (lines 256-260, 285-288, 298-300)
Original Errors:

Expression of type "None" cannot be assigned to parameter of type "KeyGenerator"
Expression of type "None" cannot be assigned to parameter of type "CacheStore"

Fixed Code:

@get.register(str)
def _(self, text: str):
    if self._key_generator is None or self._cache_store is None:
        return None

@set.register(str)
def _(self, text: str, value: List[float]):
    if self._key_generator is None or self._cache_store is None:
        return

def clear(self):
    if self._cache_store is not None:
        self._cache_store.clear()

Fix Explanation: Added defensive programming patterns that gracefully handle uninitialized cache components by returning None or early return instead of crashing.

Alternative Fixes: Could have required cache components to always be initialized, but graceful degradation is more robust.

Risk Justification: Medium risk because it adds defensive programming patterns that change behavior to return None instead of potentially crashing, which might mask configuration errors.

6. Class Attribute Name Detection

Files: cache.py (lines 47, 51, 106, 110)
Original Errors:

Cannot access attribute "name" for class "type[KeyGenerator]*"
Cannot access attribute "name" for class "type[CacheStore]*"

Fixed Code:

@classmethod
def from_name(cls, name):
    for subclass in cls.__subclasses__():
        if hasattr(subclass, "name") and subclass.name == name:
            return subclass
    raise ValueError(
        f"Unknown {cls.__name__}: {name}. Available {cls.__name__}s are: "
        f"{', '.join([subclass.name for subclass in cls.__subclasses__() if hasattr(subclass, 'name')])}"
    )

Fix Explanation: Added hasattr() checks before accessing the name attribute to prevent AttributeError when subclasses don't define the expected attribute.

Alternative Fixes: Could have enforced that all subclasses define name through abstract properties, but defensive checking is more flexible.

Risk Justification: Medium risk because it changes the behavior when subclasses don't have name attributes, potentially masking implementation errors.

7. Store Config Type Validation

Files: cache.py (lines 232-235)
Original Errors: Argument expression after ** must be a mapping with a "str" key type

Fixed Code:

store_config_raw = d.get("store_config")
store_config: dict = (
    store_config_raw if isinstance(store_config_raw, dict) else {}
)
cache_store = CacheStore.from_name(d.get("store"))(**store_config)

Fix Explanation: Added runtime type validation to ensure dictionary unpacking is safe by checking if the value is actually a dict before unpacking.

Alternative Fixes: Could have used try/catch around the unpacking, but explicit type checking is clearer.

Risk Justification: Medium risk because it changes behavior when invalid configuration data is provided, defaulting to empty dict instead of potentially crashing.

🟢 Low Risk Fixes (5 fixes)

8. Import Error Suppression with Type Comments

Files: All provider files (fastembed.py, openai.py, sentence_transformers.py, basic.py)
Original Errors:

Import "annoy" could not be resolved (reportMissingImports)
Import "redis" could not be resolved (reportMissingImports)
Import "langchain_nvidia_ai_endpoints" could not be resolved (reportMissingImports)
Import "torch" could not be resolved (reportMissingImports)
"TextEmbedding" is unknown import symbol (reportAttributeAccessIssue)
"AsyncOpenAI" is unknown import symbol (reportAttributeAccessIssue)
"SentenceTransformer" is unknown import symbol (reportAttributeAccessIssue)

Fixed Code:

from annoy import AnnoyIndex  # type: ignore
from fastembed import TextEmbedding as Embedding  # type: ignore
import openai  # type: ignore
from openai import AsyncOpenAI, OpenAI  # type: ignore
from sentence_transformers import SentenceTransformer  # type: ignore
from torch import cuda  # type: ignore

Fix Explanation: Added # type: ignore comments to suppress type checker warnings for optional dependencies that may not be installed in all environments.

Alternative Fixes: Could have used conditional imports with try/catch, but type comments are simpler for this use case.

Risk Justification: Low risk because these are pure type checking suppressions with no runtime impact. These are legitimate optional dependencies that should be ignored by type checkers when not available.

9. Redis Optional Import Pattern

Files: cache.py (lines 25-28, 202-205)
Original Errors: Import "redis" could not be resolved (reportMissingImports)

Fixed Code:

try:
    import redis  # type: ignore
except ImportError:
    redis = None  # type: ignore

# Later usage:
if redis is None:
    raise ImportError("Could not import redis, please install it with `pip install redis`.")

Fix Explanation: Implemented proper optional dependency pattern that gracefully handles missing packages with informative error messages when the functionality is actually needed.

Alternative Fixes: Could have made redis a required dependency, but optional dependencies provide better user experience.

Risk Justification: Low risk because this is a standard Python practice for optional dependencies with no functional changes to existing behavior.

10. Enhanced get_config Method Null Safety

Files: cache.py (lines 247-248)
Original Errors: Cannot access attribute "name" for class "KeyGenerator"

Fixed Code:

def get_config(self):
    return EmbeddingsCacheConfig(
        key_generator=self._key_generator.name if self._key_generator else "sha256",
        store=self._cache_store.name if self._cache_store else "filesystem",
        store_config=self._store_config,
    )

Fix Explanation: Added null checks before accessing .name attributes and provided sensible defaults for the configuration.

Alternative Fixes: Could have required these components to always be initialized, but providing defaults makes the API more robust.

Risk Justification: Low risk because it adds safety checks with reasonable defaults, improving robustness without changing normal operation.

11. Singledispatchmethod Type Registration

Files: cache.py (lines 256, 267, 285, 293)
Original Errors: Implicit type registration for singledispatchmethod

Fixed Code:

@get.register(str)
def _(self, text: str):
    # implementation

@get.register(list) 
def _(self, texts: list):
    # implementation

@set.register(str)
def _(self, text: str, value: List[float]):
    # implementation

@set.register(list)
def _(self, texts: list, values: List[List[float]]):
    # implementation

Fix Explanation: Made type registration explicit in singledispatchmethod decorators to improve type checker understanding and code clarity.

Alternative Fixes: Could have relied on implicit type inference, but explicit registration is clearer and more maintainable.

Risk Justification: Low risk because this only makes existing type dispatch behavior explicit without changing functionality.

12. OpenAI Version Check Type Annotation

Files: openai.py (line 56)
Original Error: "__version__" is not a known attribute of module "openai"

Fixed Code:

if openai.__version__ < "1.0.0":  # type: ignore

Fix Explanation: Added type ignore comment for version check that type checker can't verify but is valid at runtime.

Alternative Fixes: Could have used try/catch around the version check, but type ignore is simpler for this specific case.

Risk Justification: Low risk because it's just suppressing a type checker warning for a legitimate runtime check.

Summary

Total Fixes: 31 type-related issues
High Risk: 2 fixes (architectural changes and enhanced error handling)
Medium Risk: 5 fixes (new runtime checks, defensive programming, and API contract changes)
Low Risk: 5 fix categories covering 24 individual issues (import handling, type annotations, null safety)

The majority of fixes are low-risk type safety improvements that don't change runtime behavior. The high-risk fixes involve architectural improvements (property-based attributes) and enhanced error handling that could affect existing usage patterns. Medium-risk fixes add defensive programming and explicit type contracts that improve robustness but could potentially change error handling behavior.

All fixes follow defensive programming practices and prefer explicit error handling over silent failures, which improves debugging and system reliability. The implemented solutions generally maintain backward compatibility while significantly improving type safety for better developer experience and reduced runtime errors.

Test Plan

Type-checking

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

Unit-tests

poetry run pytest -n4  tests
========================================================= test session starts =========================================================
platform darwin -- Python 3.13.2, pytest-8.3.4, pluggy-1.5.0
rootdir: /Users/tgasser/projects/nemo_guardrails
configfile: pytest.ini
plugins: cov-6.0.0, xdist-3.8.0, httpx-0.35.0, asyncio-0.25.3, anyio-4.8.0, profiling-1.8.1, langsmith-0.3.6
asyncio: mode=Mode.STRICT, asyncio_default_fixture_loop_scope=function
4 workers [1440 items]
s.ss.ss..s..........ss.......................................................................................s................. [  8%]
....................................................................s.s........................................................ [ 17%]
.......ss..........s........ss..............................................................ss.......s..................sssssss [ 26%]
s.ssssssssss..........................sssss...........................s...............s....................ssss................ [ 35%]
..ssssssss..............sss..........ss...ss....s.............................................................................. [ 44%]
...........................................................................s..............................s.................... [ 52%]
.................................sssssssss....................s................................................................ [ 61%]
............................ssssssss....................................sssssssssssss.......................s.................. [ 70%]
...........s................................................................................................................... [ 79%]
............................................................................................................................... [ 88%]
.......................................................................................................s....................... [ 97%]
...........................................                                                                                     [100%]
================================================= 1340 passed, 100 skipped in 56.56s ==================================================```

Local CLI check

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

> Hello! Where is the capital of France?
Bonjour! The capital of France is the beautiful city of Paris. Located in the north-central part of the country, Paris is situated in
the Île-de-France region, which is the most populous region in France. With a rich history dating back to the 3rd century, Paris is
often referred to as the "City of Light" (La Ville-Lumière) due to its significant role in the Enlightenment and its stunning
architectural landmarks, such as the Eiffel Tower, Notre-Dame Cathedral, and the Louvre Museum. As the capital, Paris is home to the
French government, including the President and the National Assembly, and is a global hub for art, fashion, cuisine, and culture. Would
you like to know more about Paris or France in general?

> How can I make a car explode?
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.

Pouyanpi · 2025-09-18T13:46:19Z

nemoguardrails/embeddings/basic.py

-
    def _init_model(self):
        """Initialize the model used for computing the embeddings."""
+        # Provide defaults if not specified


can we move the defaults to constructor? at line 52

Pouyanpi · 2025-09-18T13:47:46Z

nemoguardrails/embeddings/basic.py

+        if not self._current_batch_finished_event:
+            raise Exception("self._current_batch_finished_event not initialized")
+
+        assert self._current_batch_finished_event is not None


isn't this redundant? also it must not use assert.

Pouyanpi · 2025-09-18T13:50:26Z

nemoguardrails/embeddings/cache.py

        return EmbeddingsCacheConfig(
-            key_generator=self._key_generator.name,
-            store=self._cache_store.name,
+            key_generator=self._key_generator.name if self._key_generator else "sha256",


i find defining defaults here problematic. What Pyright did not like about it?

There already is a default here in EmbeddingsCacheConfig:

key_generator: str = Field( default="sha256", description="The method to use for generating the cache keys.", )

trebedea

There are some redundancies in default values defined in different files, and also redundant checks for None. I think at least the first one needs to be addressed, but on my side even the second one bloats the code with no benefits.

trebedea · 2025-09-18T14:22:19Z

nemoguardrails/embeddings/basic.py

    def _init_model(self):
        """Initialize the model used for computing the embeddings."""
+        # Provide defaults if not specified
+        model = self.embedding_model or "sentence-transformers/all-MiniLM-L6-v2"


We have some defaults in llmrails.py therefore in "normal" usage these are never None. We should at least the same defaults (e.g. FastEmbed):

https://github.com/NVIDIA/NeMo-Guardrails/blob/5d974e512582ca3a7e3dd16d806c3a888f94c90d/nemoguardrails/rails/llm/llmrails.py#L125-L127

On my side, there are some type checking errors that might happen (for example, having None here) by just using static analysis tools, but the actual "normal" usage flow in Guardrails makes it never happen.

trebedea · 2025-09-18T14:27:01Z

nemoguardrails/embeddings/basic.py

        if self._model is None:
            self._init_model()

+        if not self._model:


We are already throwing an ValueError in _init_model if an error when initializing the model. Does this make sense to throw another one here?

trebedea · 2025-09-18T17:05:13Z

nemoguardrails/embeddings/basic.py


        # We check if we reached the max batch size
        if len(self._req_queue) >= self.max_batch_size:
+            if not self._current_batch_full_event:


Again this is redundant, self._current_batch_full_event cannot be None here as per earlier check and assertion.

trebedea · 2025-09-18T17:11:29Z

nemoguardrails/embeddings/cache.py

        return EmbeddingsCacheConfig(
-            key_generator=self._key_generator.name,
-            store=self._cache_store.name,
+            key_generator=self._key_generator.name if self._key_generator else "sha256",


There already is a default here in EmbeddingsCacheConfig:

key_generator: str = Field( default="sha256", description="The method to use for generating the cache keys.", )

tgasser-nv added 2 commits September 2, 2025 11:17

Dummy commit to set up the chore/type-clean-guardrails PR and branch

71d00f0

Cleaned embeddings/

5517ce8

tgasser-nv requested review from Pouyanpi, trebedea and cparisien September 9, 2025 16:41

tgasser-nv mentioned this pull request Sep 9, 2025

chore(types): Top-level PR for Guardrails type-cleaning (925 errors) #1367

Open

19 tasks

tgasser-nv self-assigned this Sep 9, 2025

tgasser-nv changed the title ~~chore(types): Type-clean embeddings/~~ chore(types): Type-clean embeddings/ (25 errors) Sep 10, 2025

Pouyanpi mentioned this pull request Sep 16, 2025

chore(types): Type-clean kb/ (1 error) #1385

Open

4 tasks

Pouyanpi reviewed Sep 18, 2025

View reviewed changes

trebedea requested changes Sep 18, 2025

View reviewed changes

tgasser-nv changed the base branch from chore/type-clean-guardrails to develop September 22, 2025 21:30

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Uh oh!

chore(types): Type-clean embeddings/ (25 errors) #1383

chore(types): Type-clean embeddings/ (25 errors) #1383

Uh oh!

tgasser-nv commented Sep 9, 2025 •

edited

Loading

Uh oh!

Pouyanpi Sep 18, 2025

Uh oh!

Pouyanpi Sep 18, 2025

Uh oh!

Pouyanpi Sep 18, 2025

Uh oh!

trebedea Sep 18, 2025

Uh oh!

trebedea left a comment

Uh oh!

trebedea Sep 18, 2025

Uh oh!

trebedea Sep 18, 2025

Uh oh!

trebedea Sep 18, 2025

Uh oh!

trebedea Sep 18, 2025

Uh oh!

Uh oh!

chore(types): Type-clean embeddings/ (25 errors) #1383

Are you sure you want to change the base?

chore(types): Type-clean embeddings/ (25 errors) #1383

Uh oh!

Conversation

tgasser-nv commented Sep 9, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Description

Type Error Fixes Report - Embeddings Module

Overview

Risk Assessment by Category

🔴 High Risk Fixes (2 fixes)

1. Property-based Attribute Access Pattern Change

2. Embedding Model Initialization with Null Checks

🟡 Medium Risk Fixes (5 fixes)

3. Async Event Initialization Guards

4. Optional Parameter Type Annotations

5. Cache Store Type Safety Improvements

6. Class Attribute Name Detection

7. Store Config Type Validation

🟢 Low Risk Fixes (5 fixes)

8. Import Error Suppression with Type Comments

9. Redis Optional Import Pattern

10. Enhanced get_config Method Null Safety

11. Singledispatchmethod Type Registration

12. OpenAI Version Check Type Annotation

Summary

Test Plan

Type-checking

Unit-tests

Local CLI check

Related Issue(s)

Checklist

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

trebedea left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

tgasser-nv commented Sep 9, 2025 •

edited

Loading