refactor(iorails): Add LLMModel and HTTP Client support to ModelEngine - #2246
Conversation
github-actions
Bot
added
status: needs triage
tgasser-nv
added
status: triaged
Codecov Report✅ All modified and coverable lines are covered by tests. 📢 Thoughts on this report? Let us know! |
Greptile SummaryThe PR adapts IORails model engines to the shared LLM and HTTP-client interfaces.
|
| Filename | Overview |
|---|---|
| nemoguardrails/guardrails/engine_registry.py | Exposes model engines and a shared HTTP client while delegating model calls and telemetry to ModelEngine. |
| nemoguardrails/guardrails/model_engine.py | Implements the LLMModel generation and streaming interfaces, provider metadata, parameter merging, and telemetry. |
| tests/guardrails/test_engine_registry.py | Adds coverage for LLM exposure, shared HTTP-client lifecycle, delegation, telemetry, and streaming cleanup. |
| tests/guardrails/test_model_engine.py | Expands coverage for prompt normalization, model protocol behavior, request parameters, errors, streaming, and telemetry. |
Flowchart
%%{init: {'theme': 'neutral'}}%%
flowchart LR
IORails[IORails callers] --> Registry[EngineRegistry]
Actions[Library rail actions] --> LLMs[EngineRegistry.llms]
Registry --> Model[ModelEngine]
LLMs --> Model
Model --> Adapter[Prompt/message normalization]
Adapter --> HTTP[OpenAI-compatible model endpoint]
Actions --> SharedClient[Managed shared HTTP client]
Registry --> SharedClient
Reviews (4): Last reviewed commit: "Empty commit to trigger new actions CI" | Re-trigger Greptile
|
@coderabbitai Review this PR |
|
✅ Action performedReview finished.
|
📝 WalkthroughWalkthroughEngineRegistry now exposes configured engines as LLMs and manages a shared HTTP client. ModelEngine now handles prompt normalization, parameter merging, telemetry, metrics, content capture, response parsing, and streaming cleanup. Tests cover lifecycle, delegation, protocol behavior, and telemetry. ChangesModel call centralization
Estimated code review effort: 4 (Complex) | ~60 minutes Possibly related PRs
🚥 Pre-merge checks | ✅ 5 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (5 passed)
✨ Finishing Touches📝 Generate docstrings
🧪 Generate unit tests (beta)
Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
🧹 Nitpick comments (3)
tests/guardrails/test_model_engine.py (1)
137-150: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winMove the duplicated
reset_llm_call_contextfixture into a shared conftest. Both test modules define this fixture with identical bodies and identical docstrings. The set of context variables thatllm_callwrites is one fact, so it should have one definition. The docstring already points attests/conftest.pyautouse resets, which shows the project uses conftest-level fixtures for this purpose.
tests/guardrails/test_model_engine.py#L137-L150: delete the local fixture and rely on a newtests/guardrails/conftest.pydefinition.tests/guardrails/test_engine_registry.py#L1777-L1790: delete the local fixture; move this definition intotests/guardrails/conftest.py.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@tests/guardrails/test_model_engine.py` around lines 137 - 150, Move the shared reset_llm_call_context fixture into tests/guardrails/conftest.py, preserving its existing body and docstring. Delete the local definitions from tests/guardrails/test_model_engine.py lines 137-150 and tests/guardrails/test_engine_registry.py lines 1777-1790 so both modules use the shared fixture.nemoguardrails/guardrails/engine_registry.py (2)
275-284: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winAdd a test for early abandonment of
stream_model_call.The explicit
aclose()is the fix for a real cleanup hazard.tests/guardrails/test_engine_registry.pypins that behavior forModelEngine.stream_asynconly, intest_stream_duration_recorded_on_consumer_early_break. No test abandonsstream_model_callmid-stream. Removing thisfinallyblock would therefore not fail any test. Add a case that consumes one chunk fromstream_model_call, callsaclose(), and assertsgen_ai.client.operation.durationwas recorded.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@nemoguardrails/guardrails/engine_registry.py` around lines 275 - 284, Add a regression test in the engine registry test suite for early abandonment of stream_model_call: consume one chunk, explicitly call the returned generator’s aclose(), and assert gen_ai.client.operation.duration was recorded, matching the existing ModelEngine.stream_async coverage and preserving the cleanup behavior in stream_model_call.
178-181: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winLog the swallowed HTTP-client close failure.
The rollback discards any exception from
_close_http_client(). The original engine-start error must still propagate, so swallowing is correct. A silent swallow hides a real cleanup failure during startup. Log it at warning level instead. This also resolves the RuffS110hint.♻️ Proposed change
try: await self._close_http_client() - except Exception: - pass + except Exception as close_error: + log.warning("Error closing the managed HTTP client during start rollback: %s", close_error) raise RuntimeError(f"Failed to start engine: Engine {name}: exception {e}") from e🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@nemoguardrails/guardrails/engine_registry.py` around lines 178 - 181, Update the rollback cleanup around _close_http_client() to catch the exception as a named value and log it at warning level, while preserving propagation of the original engine-start error. Keep the cleanup failure from being re-raised and use the existing logger associated with this flow; this also removes the Ruff S110 violation.Source: Linters/SAST tools
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Nitpick comments:
In `@nemoguardrails/guardrails/engine_registry.py`:
- Around line 275-284: Add a regression test in the engine registry test suite
for early abandonment of stream_model_call: consume one chunk, explicitly call
the returned generator’s aclose(), and assert gen_ai.client.operation.duration
was recorded, matching the existing ModelEngine.stream_async coverage and
preserving the cleanup behavior in stream_model_call.
- Around line 178-181: Update the rollback cleanup around _close_http_client()
to catch the exception as a named value and log it at warning level, while
preserving propagation of the original engine-start error. Keep the cleanup
failure from being re-raised and use the existing logger associated with this
flow; this also removes the Ruff S110 violation.
In `@tests/guardrails/test_model_engine.py`:
- Around line 137-150: Move the shared reset_llm_call_context fixture into
tests/guardrails/conftest.py, preserving its existing body and docstring. Delete
the local definitions from tests/guardrails/test_model_engine.py lines 137-150
and tests/guardrails/test_engine_registry.py lines 1777-1790 so both modules use
the shared fixture.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: 73e4d9d7-3da6-453f-b42e-6380ca567c61
📒 Files selected for processing (4)
nemoguardrails/guardrails/engine_registry.pynemoguardrails/guardrails/model_engine.pytests/guardrails/test_engine_registry.pytests/guardrails/test_model_engine.py
tgasser-nv
changed the title
Pouyanpi
left a comment
Pouyanpi
left a comment
There was a problem hiding this comment.
Thank you @tgasser-nv , I can see the direction, approving 👍🏻
…tion down to match llm_call()
8d3b937 to
2d5fabe
Compare
2 participants
Description
This PR is the second in a series of stacked PRs to allow IORails to run the recently refactored actions in nemoguardrails/library/* directly and avoid duplicating actions for the two engines. The rough PR plan (this may change during implementation) is shown below:
Related Issue(s)
Verification
Pre-commit
Unit-test
Integration test with Chat
AI Assistance
Checklist
Summary by CodeRabbit
New Features
Bug Fixes
Tests