chore(types): Type-clean embeddings/ (25 errors) #1383
Conversation
tgasser-nv
changed the title
tgasser-nv
changed the base branch from
chore/type-clean-guardrails
to
develop
|
Converting to draft while I rebase on the latest changes to develop. |
tgasser-nv
force-pushed
the
chore/type-clean-embeddings
branch
from
7fa0c6f to
1b6eaab
Compare
Documentation preview |
Codecov Report❌ Patch coverage is 📢 Thoughts on this report? Let us know! |
|
Requesting review after rebasing on top of develop and addressing feedback in comments from Pouyan and Traian. cc @Pouyanpi , @trebedea , @cparisien |
tgasser-nv
force-pushed
the
chore/type-clean-embeddings
branch
from
b78de2c to
994cc14
Compare
![greptile-apps[bot]](https://avatars.githubusercontent.com/in/867647?s=60&v=4){kind=small}
There was a problem hiding this comment.
Greptile Overview
Greptile Summary
This PR eliminates 25 Pyright type-checking errors in the nemoguardrails/embeddings/ module by adding explicit type annotations, defensive null-checks, and # type: ignore directives for third-party libraries without type stubs. The changes transform implicitly-typed class attributes in BasicEmbeddingsIndex into explicitly-typed instance attributes initialized in __init__, add graceful-degradation guards to EmbeddingsCache when components are uninitialized, mark optional configuration fields in EmbeddingsCacheConfig, and enable Pyright checking for the entire embeddings directory. These fixes integrate with the existing NeMo-Guardrails architecture by preserving the optional nature of caching and embedding model selection while providing clearer error messages when components fail to initialize, aligning with the project's broader effort to incrementally add static type safety (similar to existing Pyright coverage for rails/, actions/, cli/, kb/, logging/, tracing/).
Important Files Changed
| Filename | Score | Overview |
|---|---|---|
nemoguardrails/embeddings/basic.py |
1/5 | Critical indentation bug on line 275 will cause unconditional await deadlock in batching code path; needs immediate fix |
nemoguardrails/embeddings/cache.py |
4/5 | Adds null-safety guards, explicit type hints, and graceful degradation when cache components are uninitialized; silent failure may mask configuration issues |
nemoguardrails/rails/llm/config.py |
5/5 | Marks three EmbeddingsCacheConfig fields as Optional to align type signatures with runtime null-handling |
pyproject.toml |
5/5 | Enables Pyright type-checking for nemoguardrails/embeddings/** directory |
nemoguardrails/embeddings/providers/fastembed.py |
5/5 | Adds # type: ignore to suppress missing type stub warning for fastembed import |
nemoguardrails/embeddings/providers/nim.py |
5/5 | Adds # type: ignore to suppress missing type stub warning for langchain_nvidia_ai_endpoints import |
nemoguardrails/embeddings/providers/sentence_transformers.py |
5/5 | Adds # type: ignore comments for sentence_transformers and torch imports |
nemoguardrails/embeddings/providers/openai.py |
5/5 | Adds three # type: ignore comments for openai library imports and version checks |
Confidence score: 1/5
- This PR contains a critical indentation bug that will cause production deadlocks when batching is enabled, making it unsafe to merge in its current state.
- Score reflects one blocking issue in
basic.pyline 275 where anawaitstatement is incorrectly unindented, causing unconditional execution regardless of the preceding batch-full check; the remaining changes are sound but overshadowed by this regression. nemoguardrails/embeddings/basic.pylines 267-275 require immediate attention—theawait self._current_batch_finished_event.wait()must be re-indented to be inside theif len(self._req_queue) >= self.max_batch_size:block.
Sequence Diagram
sequenceDiagram
participant User
participant BasicEmbeddingsIndex
participant EmbeddingsCache
participant EmbeddingModel
participant CacheStore
participant KeyGenerator
Note over User,KeyGenerator: Type-Clean Embeddings Flow
User->>BasicEmbeddingsIndex: __init__(embedding_model, embedding_engine, ...)
BasicEmbeddingsIndex->>BasicEmbeddingsIndex: Initialize typed instance attributes
Note right of BasicEmbeddingsIndex: _items: List[IndexItem]<br/>_embeddings: List[List[float]]<br/>_req_queue: Dict[int, str]<br/>_current_batch_full_event: Optional[Event]
User->>BasicEmbeddingsIndex: add_item(item)
BasicEmbeddingsIndex->>BasicEmbeddingsIndex: _init_model()
BasicEmbeddingsIndex->>BasicEmbeddingsIndex: Check if self._model exists
alt _model is None
BasicEmbeddingsIndex->>EmbeddingModel: init_embedding_model(model, engine, params)
EmbeddingModel-->>BasicEmbeddingsIndex: Return model instance or None
BasicEmbeddingsIndex->>BasicEmbeddingsIndex: Validate model is not None
alt model is None
BasicEmbeddingsIndex-->>User: raise ValueError("Couldn't create embedding model")
end
end
BasicEmbeddingsIndex->>BasicEmbeddingsIndex: _get_embeddings([item.text])
BasicEmbeddingsIndex->>EmbeddingsCache: @cache_embeddings decorator
EmbeddingsCache->>EmbeddingsCache: Check if cache enabled
alt cache enabled
EmbeddingsCache->>EmbeddingsCache: get(texts)
EmbeddingsCache->>EmbeddingsCache: Check _key_generator and _cache_store not None
alt key_generator or cache_store is None
EmbeddingsCache-->>BasicEmbeddingsIndex: return None (graceful degradation)
else both initialized
EmbeddingsCache->>KeyGenerator: generate_key(text)
KeyGenerator-->>EmbeddingsCache: key
EmbeddingsCache->>CacheStore: get(key)
CacheStore-->>EmbeddingsCache: cached_value or None
end
end
alt cache miss or disabled
BasicEmbeddingsIndex->>EmbeddingModel: encode_async(texts)
EmbeddingModel-->>BasicEmbeddingsIndex: embeddings
BasicEmbeddingsIndex->>EmbeddingsCache: set(texts, embeddings)
EmbeddingsCache->>EmbeddingsCache: Check _key_generator and _cache_store not None
alt both initialized
EmbeddingsCache->>KeyGenerator: generate_key(text)
KeyGenerator-->>EmbeddingsCache: key
EmbeddingsCache->>CacheStore: set(key, value)
else either is None
Note right of EmbeddingsCache: Silently skip caching
end
end
BasicEmbeddingsIndex-->>User: Item added successfully
User->>BasicEmbeddingsIndex: _run_batch()
BasicEmbeddingsIndex->>BasicEmbeddingsIndex: Check _current_batch_full_event not None
alt event is None
BasicEmbeddingsIndex-->>User: raise Exception("event not initialized")
else event exists
BasicEmbeddingsIndex->>BasicEmbeddingsIndex: Process batch with event synchronization
end
8 files reviewed, 4 comments
| 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 |
There was a problem hiding this comment.
style: Redundant check; line 208–209 already raise if _current_batch_full_event is None
| if len(self._req_queue) >= self.max_batch_size: | ||
| if not self._current_batch_full_event: | ||
| raise Exception("self._current_batch_full_event not initialized") | ||
| self._current_batch_full_event.set() | ||
|
|
||
| # Wait for the batch to finish | ||
| # Wait for the batch to finish | ||
| if not self._current_batch_finished_event: | ||
| raise Exception("self._current_batch_finished_event not initialized") | ||
| await self._current_batch_finished_event.wait() |
There was a problem hiding this comment.
logic: Indentation issue: line 275 is outside the if block on lines 267–270, so the await always executes regardless of batch size. This will likely deadlock.
| store_config_raw = d.get("store_config") | ||
| store_config: dict = ( | ||
| store_config_raw if isinstance(store_config_raw, dict) else {} | ||
| ) |
There was a problem hiding this comment.
style: Type narrowing with isinstance is good, but d.get('store_config') could be any type from untrusted input. Consider validating the structure beyond just dict type checking.
| @get.register(str) | ||
| def _(self, text: str): | ||
| if self._key_generator is None or self._cache_store is None: | ||
| return None |
There was a problem hiding this comment.
style: Silently returning None when cache components are uninitialized may hide configuration errors. Consider logging a warning on first occurrence.
There was a problem hiding this comment.
Greptile Overview
Greptile Summary
This review covers only the changes made since the last review, not the entire PR. The latest changes address three minor type-checking issues in provider files for Google, Cohere, and AzureOpenAI embedding services. Specifically, the developer added # type: ignore directives to suppress Pyright warnings where third-party libraries lack complete type stubs. In google.py, the google.genai import is marked with # type: ignore[import]. In cohere.py, imports are moved into a TYPE_CHECKING block (allowing static analyzers to resolve types without altering runtime behavior), and a # type: ignore[reportCallIssue] is added to the cohere.Client constructor call. In azureopenai.py, a # type: ignore[arg-type] is added where os.getenv() returns str | None but the Azure SDK expects str. All three changes are documentation-only annotations for the type checker and have zero impact on runtime execution. They fit into the broader type-safety cleanup of the nemoguardrails/embeddings/ directory, ensuring the module passes strict Pyright checks while maintaining backward compatibility with the existing lazy-import pattern for optional dependencies.
Important Files Changed
| Filename | Score | Overview |
|---|---|---|
| nemoguardrails/embeddings/providers/google.py | 5/5 | Added # type: ignore[import] to suppress missing type stubs warning for google.genai import. |
| nemoguardrails/embeddings/providers/cohere.py | 5/5 | Moved imports into TYPE_CHECKING block and suppressed type-checker warning on cohere.Client call; no runtime logic changed. |
| nemoguardrails/embeddings/providers/azureopenai.py | 4/5 | Silenced type-checker warning on azure_endpoint parameter that expects str but receives `str |
Confidence score: 4/5
- This incremental update is safe to merge with very low risk; all changes are type-annotation directives that do not alter runtime behavior.
- Score is 4/5 instead of 5/5 solely because the
azureopenai.pychange silences a type mismatch whereos.getenv()may returnNone, and the Azure SDK will fail at runtime if the environment variable is missing—no validation logic is added to catch this earlier. - Pay close attention to
nemoguardrails/embeddings/providers/azureopenai.pyto ensure environment variables are documented as required in deployment guides; no other files require special scrutiny.
3 files reviewed, no comments
There was a problem hiding this comment.
Greptile Overview
Greptile Summary
This review covers only the changes made since the last review, not the entire PR. The developer has addressed previous feedback by removing redundant None checks that were unreachable after earlier guard clauses, relocating default values from _init_model back into the __init__ constructor to centralize initialization, and eliminating an assert statement in favor of an explicit ValueError. These refinements clean up the defensive programming approach while maintaining the type-safety improvements. However, the critical concurrency bug in _batch_get_embeddings (lines 260–278 of basic.py) remains unresolved: multiple concurrent callers can still create separate batch-event instances, causing some coroutines to await the wrong event and deadlock.
PR Description Notes:
- Line 208: References to line numbers in the PR body do not match the latest code after refactoring; consider updating the description to reflect the final state.
Important Files Changed
| Filename | Score | Overview |
|---|---|---|
nemoguardrails/embeddings/basic.py |
2/5 | Moved defaults to constructor; removed redundant checks; concurrency race in _batch_get_embeddings (lines 260–278) remains—multiple coroutines can create separate batch events and deadlock. |
nemoguardrails/embeddings/providers/azureopenai.py |
4/5 | Cleaned up # type: ignore directives; no functional changes. |
Confidence score: 2/5
- This PR remains risky due to an unresolved concurrency bug that can cause deadlocks under concurrent embedding requests.
- Score reflects the same critical issue as the initial review: the batch-event initialization in
_batch_get_embeddingsis not thread-safe and can create orphaned events. While the recent changes improved code clarity and removed redundant checks, they did not address the race condition. - Pay close attention to
nemoguardrails/embeddings/basic.pylines 260–278—this section requires a lock or single-initialization pattern to prevent multiple coroutines from creating conflicting batch events.
2 files reviewed, 3 comments
| if ( | ||
| self._current_batch_finished_event is None | ||
| or self._current_batch_full_event is None | ||
| ): | ||
| self._current_batch_finished_event = asyncio.Event() | ||
| self._current_batch_full_event = asyncio.Event() | ||
| self._current_batch_submitted.clear() | ||
| asyncio.ensure_future(self._run_batch()) | ||
|
|
||
| # 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: | ||
| raise Exception("self._current_batch_full_event not initialized") | ||
| self._current_batch_full_event.set() | ||
|
|
||
| # Wait for the batch to finish | ||
| # Wait for the batch to finish | ||
| if not self._current_batch_finished_event: | ||
| raise Exception("self._current_batch_finished_event not initialized") | ||
| await self._current_batch_finished_event.wait() |
There was a problem hiding this comment.
logic: Lines 260–267 create new Event objects only if both are None, but line 278 always awaits _current_batch_finished_event regardless of whether this coroutine created the event. If another coroutine set the event between lines 260–267, this await will hang indefinitely waiting for its own request to complete. Is there a lock or other synchronization mechanism ensuring only one coroutine creates the batch events at a time?
| try: | ||
| from openai import AzureOpenAI | ||
| from openai import ( | ||
| AzureOpenAI, # type: ignore[attr-defined] (Assume this is installed) |
There was a problem hiding this comment.
style: # type: ignore[attr-defined] is incorrect for this import. AzureOpenAI is a class, not an attribute. If the OpenAI library is installed, this import succeeds without issues. Use # type: ignore without a specifier, or remove it entirely if type stubs are available.
| api_key=os.getenv("AZURE_OPENAI_API_KEY"), | ||
| api_version=os.getenv("AZURE_OPENAI_API_VERSION"), | ||
| azure_endpoint=os.getenv("AZURE_OPENAI_ENDPOINT"), | ||
| azure_endpoint=os.getenv("AZURE_OPENAI_ENDPOINT"), # type: ignore[arg-type] (comes from `$AZURE_OPENAI_ENDPOINT`) |
There was a problem hiding this comment.
style: The # type: ignore[arg-type] is redundant; os.getenv returns str | None, which AzureOpenAI's constructor should accept. If the constructor requires a non-None value, validate it explicitly rather than silencing the type error. Does the AzureOpenAI constructor accept None for azure_endpoint, or does it require a valid string?
4 participants
Description
Cleaned the nemoguardrails/embeddings directory using Pyright.
This report summarizes the type-safety fixes applied to the
nemoguardrailsembeddings module. The changes are categorized into medium and low-risk buckets based on their potential to impact existing functionality.🟡 Medium-Risk Changes
These changes involve refactoring class initialization and adding explicit null-safety checks in critical code paths. While they improve robustness, they alter error handling and make assumptions about default behavior, which introduces a moderate risk of unintended side effects.
1. Refactoring
BasicEmbeddingsIndexfor Type SafetyThe
BasicEmbeddingsIndexclass was refactored to properly initialize instance attributes and handle potentially uninitialized async components, preventingAttributeErrorat runtime.nemoguardrails/embeddings/basic.pyAttributeErrorif async methods like_run_batchwere called before certain event objects were created. Class attributes were also implicitly treated as instance attributes, leading to poor type inference.__init__. In critical async methods, explicit checks were added to ensure event loop objects (_current_batch_full_event, etc.) are notNonebefore they are accessed, raising a descriptiveExceptioninstead of a genericAttributeError. The_init_modelmethod now provides sensible defaults if a model or engine is not specified.Nonewhen accessed, it represents an unrecoverable state error, justifying an exception.Exception, a custom exception type could have been used for more specific error handling. However, the current fix is sufficient to prevent the runtime crash and clearly signals the programming error.2. Making
EmbeddingsCacheRobust AgainstNoneThe
EmbeddingsCacheclass was hardened to preventAttributeErrorwhen its core components (_key_generator,_cache_store) are not configured.nemoguardrails/embeddings/cache.pyTypeErrororAttributeErrorwould occur ifgetorsetwere called on a cache instance that was not fully configured (e.g.,_cache_storewasNone).getandsetmethods now begin with a guard clause that checks if the key generator and cache store have been initialized. If not, they exit early, withgetreturningNoneandsetdoing nothing. This makes the cache's behavior more predictable when it's disabled or misconfigured.ConfigurationErrorif the methods are called on an uninitialized cache. The current implementation prioritizes graceful degradation over strictness, which is a reasonable choice for an optional component like a cache.🟢 Low-Risk Changes
These fixes are minor, defensive additions that silence type-checker warnings and prevent simple errors without altering program logic. They are highly unlikely to introduce any new bugs.
1. Ignoring Missing Type Stubs in Third-Party Libraries
Type checkers were reporting errors for popular libraries that do not ship with type stubs. These have been silenced using
# type: ignore.embeddings/basic.py,embeddings/providers/fastembed.py,embeddings/providers/openai.py,embeddings/providers/sentence_transformers.pyannoy,fastembed,openai,torch, andsentence_transformersas missing type information.# type: ignoreto the import statements instructs static type checkers to skip validation for these lines. This has no effect on the runtime behavior of the code and is the standard way to handle dependencies that are not yet fully typed..pyi) for these libraries, which would be a significant and unnecessary effort.2. Adding
OptionalTypes and DefensiveNoneChecksFunction signatures were updated with
Optionalto accurately reflect thatNoneis a valid input. Defensivehasattrandis Nonechecks were added to prevent errors.nemoguardrails/embeddings/cache.pyAttributeErrorwhen accessing the.nameattribute on a subclass that might not have defined it, orTypeErrorwhen passingNoneto functions not expecting it.from_namefactory methods now usehasattr(subclass, 'name')before attempting to access thenameattribute, making them more robust to incorrectly defined subclasses. Additionally, theredisimport is now handled lazily, with a check inRedisCacheStorethat provides a clear error message if the library is not installed.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..................................................................PassedUnit-tests
Local CLI check
Checklist