-
Notifications
You must be signed in to change notification settings - Fork 559
chore(types): Type-clean embeddings/ (25 errors) #1383
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
Conversation
trebedea
left a comment
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
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.
|
Converting to draft while I rebase on the latest changes to develop. |
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 |
trebedea
left a comment
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Just a minor comment.
b78de2c to
994cc14
Compare
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
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
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
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
Pouyanpi
left a comment
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Thank you @tgasser-nv . Please see the review PR: #1473. After addressing those we should be good to merge 👍🏻
* first round of review * redundant pyright checks * fix cohere type errors * remove redundant model validation
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Greptile Overview
Greptile Summary
This review covers only the changes made since the last review, not the entire PR.
This update reverts the problematic constructor signature changes to EmbeddingsCache that were flagged in the previous review. The developer has restored key_generator and cache_store as optional parameters in the __init__ method of EmbeddingsCache (in cache.py), fixing the breaking change that would have caused runtime failures in code paths where the cache is disabled or uninitialized. Additionally, the defensive hasattr checks in the from_name factory methods have been restored, and the None-guard clauses in get and set methods remain in place to gracefully degrade when cache components are not configured. The configuration schema changes in config.py remain unchanged (fields still non-optional with defaults), which aligns with the internal usage pattern where these fields are always populated by the LLMRails constructor. The Cohere provider fix in cohere.py remains as-is with the clarifying comment and type ignore directive.
These changes address the critical feedback from the previous review that identified the removal of Optional types as a breaking change. The revised approach maintains backward compatibility while still satisfying Pyright's type checker through the use of targeted # type: ignore directives and defensive runtime checks.
Important Files Changed
| Filename | Score | Overview |
|---|---|---|
| nemoguardrails/embeddings/cache.py | 4/5 | Reverted breaking constructor signature changes; restored optional parameters and defensive hasattr checks in factory methods while keeping None-guards in get/set methods |
| nemoguardrails/rails/llm/config.py | 5/5 | Configuration schema fields remain non-optional with defaults; aligns with internal constructor usage pattern |
| nemoguardrails/embeddings/providers/cohere.py | 5/5 | Added clarifying comment and type ignore directive for Cohere SDK union return type; no runtime behavior change |
| nemoguardrails/embeddings/basic.py | 4/5 | Refactored BasicEmbeddingsIndex initialization with explicit type annotations and added runtime checks for async event objects; defaults moved to constructor |
Confidence score: 4/5
- This PR is safe to merge with minimal risk after the critical cache.py breaking changes were reverted in response to previous review feedback
- Score reflects that the remaining changes are defensive improvements (explicit initialization, None-guards, type ignores for third-party stubs) that enhance type safety without altering core logic, though the BasicEmbeddingsIndex refactoring introduces new default values that warrant verification in edge cases
- Pay close attention to nemoguardrails/embeddings/basic.py to ensure the new default values for
embedding_modelandembedding_enginealign with all call sites, and verify that the async event initialization checks do not introduce unexpected exceptions in legitimate usage patterns
Additional Comments (2)
-
nemoguardrails/embeddings/cache.py, line 47-48 (link)logic: Removing
hasattrcheck will cause AttributeError if any subclass does not definename. The old code was defensive; this assumes all subclasses are correct. -
nemoguardrails/embeddings/cache.py, line 106-107 (link)logic: Same as line 47—removing
hasattrassumes all subclasses definename. If any don't, this will raise AttributeError at runtime.
4 files reviewed, 5 comments
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Greptile Overview
Greptile Summary
This review covers only the changes made since the last review, not the entire PR. The current diff makes two minor changes: (1) simplifies type-ignore comments in azureopenai.py by replacing verbose specifiers ([attr-defined], [arg-type]) with a blanket # type: ignore, and (2) removes a clarifying comment from line 86 of cache.py. These are cosmetic adjustments aligned with the broader Pyright type-cleaning effort. The changes do not introduce new logic or alter behavior. However, the azureopenai.py file now masks the fact that os.getenv(...) may return None, which could cause runtime failures if the AzureOpenAI constructor requires non-None arguments. The cache.py change is purely stylistic and does not address the pre-existing critical issues previously identified in the file (redundant checks, indentation causing potential deadlocks, insufficient validation, missing synchronization, breaking constructor changes).
Important Files Changed
| Filename | Score | Overview |
|---|---|---|
| nemoguardrails/embeddings/providers/azureopenai.py | 4/5 | Simplified type-ignore comments; may hide runtime errors if environment variables are not set |
| nemoguardrails/embeddings/cache.py | 5/5 | Removed a clarifying comment on line 86; no logic changes or new issues introduced |
Confidence score: 4/5
- This PR is safe to merge with minimal risk, as the changes are cosmetic and do not introduce new functionality or breaking changes.
- Score reflects that the changes are purely stylistic and pass all tests, but the blanket
# type: ignoreinazureopenai.pycould mask runtime issues if environment variables are missing, and thecache.pyfile retains pre-existing issues from earlier commits. - Pay close attention to
nemoguardrails/embeddings/providers/azureopenai.pyto ensure environment variables are properly validated at runtime before instantiating AzureOpenAI.
Additional Comments (1)
-
nemoguardrails/embeddings/cache.py, line 345 (link)logic: : Line 345re-fetches uncached_texts from cache immediately after storing them (line 343). This redundant round-trip can be removed by directly using
uncached_resultsinstead of callingembeddings_cache.get(uncached_texts)again.
2 files reviewed, 2 comments
Pouyanpi
left a comment
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Thank you @tgasser-nv , it looks good 🚀
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Greptile Overview
Greptile Summary
This review covers only the changes made since the last review, not the entire PR. The most recent change removes a single explanatory comment (line 48) from nemoguardrails/embeddings/basic.py that documented instance attributes being defined in __init__ and accessed via properties. This cosmetic cleanup is part of the broader type-cleaning effort across the embeddings/ directory. The comment removal has no runtime impact and aligns with the explicit property decorators already present in the codebase (lines 100–123). The change integrates cleanly with the existing BasicEmbeddingsIndex class structure, which implements asynchronous batch processing for embeddings with caching support. All type-checking (pyright) and unit tests (1932 passed) continue to pass.
Important Files Changed
| Filename | Score | Overview |
|---|---|---|
| nemoguardrails/embeddings/basic.py | 5/5 | Removed explanatory comment about instance attributes; purely cosmetic change with no functional impact. |
Confidence score: 5/5
- This change is safe to merge with minimal risk—it only removes a comment.
- Score reflects that the modification is purely cosmetic documentation cleanup with no logic or behavior changes; passing tests confirm no regression.
- No files require special attention; the single-line comment removal introduces zero risk.
1 file reviewed, no comments
f84713f to
61ac1c9
Compare
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Greptile Overview
Greptile Summary
This review covers only the changes made since the last review, not the entire PR. The latest changes continue the type-cleaning effort in the embeddings directory, focusing on suppressing type-checker errors. The key modifications are: (1) a # type: ignore comment added to azureopenai.py to suppress the type error where os.getenv() returns str | None but AzureOpenAI expects a non-None value for azure_endpoint, and (2) removal of redundant comments in basic.py. These changes align with the PR's goal of achieving a clean Pyright run across the embeddings module, and they preserve existing runtime behavior without adding validation logic.
Important Files Changed
| Filename | Score | Overview |
|---|---|---|
| nemoguardrails/embeddings/providers/azureopenai.py | 3/5 | Added # type: ignore to suppress type error on azure_endpoint parameter without validating the environment variable at runtime |
| nemoguardrails/embeddings/basic.py | 4/5 | Removed redundant comment; no logic changes |
Confidence score: 3/5
- This PR introduces minimal risk but does not address potential runtime errors; safe to merge with caveats
- Score reflects that while type-checking now passes and no logic changed, the
# type: ignorein azureopenai.py masks a real risk: ifAZURE_OPENAI_ENDPOINTis unset, the code will pass type-checking but may crash at runtime depending on the AzureOpenAI SDK's behavior - Pay close attention to nemoguardrails/embeddings/providers/azureopenai.py to ensure the suppressed type error does not hide a configuration issue that could cause runtime failures in production
2 files reviewed, no comments
* Initial checkin * Add nemoguardrails/server to pyright type-checking * chore(types): Type-clean embeddings/ (25 errors) (#1383) * test: restore test that was skipped due to Colang 2.0 serialization issue (#1449) * fix(llm): add fallback extraction for reasoning traces from <think> tags (#1474) Adds a compatibility layer for LLM providers that don't properly populate reasoning_content in additional_kwargs. When reasoning_content is missing, the system now falls back to extracting reasoning traces from <think>...</think> tags in the response content and removes the tags from the final output. This fixes compatibility with certain NVIDIA models (e.g., nvidia/llama-3.3-nemotron-super-49b-v1.5) in langchain-nvidia-ai-endpoints that include reasoning traces in <think> tags but fail to populate the reasoning_content field. All reasoning models using ChatNVIDIA should expose reasoning content consistently through the same interface * Clean up the config_id logic based on Traian and Greptile feedback --------- Co-authored-by: Pouyan <[email protected]>
* Initial checkin * Add nemoguardrails/server to pyright type-checking * chore(types): Type-clean embeddings/ (25 errors) (#1383) * test: restore test that was skipped due to Colang 2.0 serialization issue (#1449) * fix(llm): add fallback extraction for reasoning traces from <think> tags (#1474) Adds a compatibility layer for LLM providers that don't properly populate reasoning_content in additional_kwargs. When reasoning_content is missing, the system now falls back to extracting reasoning traces from <think>...</think> tags in the response content and removes the tags from the final output. This fixes compatibility with certain NVIDIA models (e.g., nvidia/llama-3.3-nemotron-super-49b-v1.5) in langchain-nvidia-ai-endpoints that include reasoning traces in <think> tags but fail to populate the reasoning_content field. All reasoning models using ChatNVIDIA should expose reasoning content consistently through the same interface * Clean up the config_id logic based on Traian and Greptile feedback --------- Co-authored-by: Pouyan <[email protected]>
* Initial checkin * Add nemoguardrails/server to pyright type-checking * chore(types): Type-clean embeddings/ (25 errors) (#1383) * test: restore test that was skipped due to Colang 2.0 serialization issue (#1449) * fix(llm): add fallback extraction for reasoning traces from <think> tags (#1474) Adds a compatibility layer for LLM providers that don't properly populate reasoning_content in additional_kwargs. When reasoning_content is missing, the system now falls back to extracting reasoning traces from <think>...</think> tags in the response content and removes the tags from the final output. This fixes compatibility with certain NVIDIA models (e.g., nvidia/llama-3.3-nemotron-super-49b-v1.5) in langchain-nvidia-ai-endpoints that include reasoning traces in <think> tags but fail to populate the reasoning_content field. All reasoning models using ChatNVIDIA should expose reasoning content consistently through the same interface * Clean up the config_id logic based on Traian and Greptile feedback --------- Co-authored-by: Pouyan <[email protected]>
Description
Cleaned the nemoguardrails/embeddings directory using Pyright.
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