LLM Based Task#2283
Draft
KyleZheng1284 wants to merge 1 commit into
Draft
Conversation
Signed-off-by: Kyle Zheng <126034466+KyleZheng1284@users.noreply.github.com>
Contributor
Greptile SummaryThis PR introduces a reusable
|
| Filename | Overview |
|---|---|
| nemo_retriever/src/nemo_retriever/models/llm/tasks/base.py | New abstract base class for generation tasks; execute() swallows all exceptions silently into an error result without any logging, violating the no-bare-except rule. |
| nemo_retriever/src/nemo_retriever/operators/generation/base.py | Core DataFrame operator base; _failure_model() has a bare except Exception without logging, and the overall design for bounded threaded execution and positional result ordering is sound. |
| nemo_retriever/src/nemo_retriever/tools/evaluation/generation.py | Refactored QAGenerationOperator to extend TextGenerationOperator; legacy LLMClient adapter path drops reasoning_enabled when forwarding to generate(), silently ignoring the operator-level override. |
| nemo_retriever/src/nemo_retriever/common/params/models.py | Adds LLMSamplingOverrides and TextGenerationParams; the custom model_serializer preserves omit-vs-null distinction correctly; _no_api_key_fields PrivateAttr mechanism is well-designed for no-auth graph persistence. |
| nemo_retriever/src/nemo_retriever/graph/graph_pipeline_registry.py | Major expansion of graph serialization to v2 format with full type-safe encoding, secret redaction, and backward-compatible v1 loading; logic is thorough and the test suite validates edge cases well. |
| nemo_retriever/src/nemo_retriever/models/llm/tasks/rag_answer.py | New RagAnswerTask correctly moves RAG prompt construction and reasoning controls out of LiteLLMClient; backward-compatible re-export of _build_rag_prompt is properly handled. |
| nemo_retriever/src/nemo_retriever/models/llm/tasks/generic.py | GenericPromptTask with strict template validation; correctly rejects attribute access, indexing, format specs, and undeclared/missing placeholders at construction time. |
| nemo_retriever/src/nemo_retriever/models/llm/tasks/summarize.py | SummarizeTask with empty-input short-circuit and prompt validation; clean and well-tested. |
| nemo_retriever/src/nemo_retriever/models/llm/clients/litellm.py | Substantial cleanup; RAG prompt logic correctly extracted to RagAnswerTask; backward-compatible _build_rag_prompt re-export retained; generate() now delegates to RagAnswerTask.execute(). |
| nemo_retriever/tests/test_generation_tasks.py | Comprehensive 925-line test suite covering tasks, operators, sampling round-trips, graph persistence, secret redaction, and legacy compatibility; good coverage of error paths and edge cases. |
Class Diagram
%%{init: {'theme': 'neutral'}}%%
classDiagram
class GenerationTask {
<<abstract>>
+required_inputs tuple
+_default_sampling dict
+build_request(**inputs) GenerationRequest
+parse(raw_text) str
+execute(client, **inputs) GeneratedTextResult
}
class RagAnswerTask {
+prompt Optional[str]
+system_prompt Optional[str]
+reasoning_enabled Optional[bool]
+build_request(**inputs) GenerationRequest
+parse(raw_text) str
}
class SummarizeTask {
+prompt Optional[str]
+system_prompt Optional[str]
+reasoning_enabled Optional[bool]
+build_request(**inputs) GenerationRequest
+parse(raw_text) str
}
class GenericPromptTask {
+prompt str
+input_names tuple
+system_prompt Optional[str]
+reasoning_enabled Optional[bool]
+build_request(**inputs) GenerationRequest
+parse(raw_text) str
}
class TextGenerationOperator {
<<abstract>>
+params TextGenerationParams
+_client CompletionClient
+_create_task() GenerationTask
+_get_generation_constructor_kwargs() dict
+process(data) DataFrame
}
class SummarizationOperator {
+input_column str
+output_column str
+_create_task() GenerationTask
}
class GenericGenerationOperator {
+input_columns Mapping
+output_column str
+_create_task() GenerationTask
}
class QAGenerationOperator {
+model str
+_execute_task(inputs) GeneratedTextResult
+_create_task() GenerationTask
}
class TextGenerationParams {
+transport LLMRemoteClientParams
+sampling LLMSamplingOverrides
+prompt Optional[str]
+reasoning_enabled Optional[bool]
+resolve_sampling(defaults) LLMInferenceParams
+from_kwargs(...) TextGenerationParams
}
GenerationTask <|-- RagAnswerTask
GenerationTask <|-- SummarizeTask
GenerationTask <|-- GenericPromptTask
TextGenerationOperator <|-- SummarizationOperator
TextGenerationOperator <|-- GenericGenerationOperator
TextGenerationOperator <|-- QAGenerationOperator
TextGenerationOperator --> GenerationTask : creates via _create_task
TextGenerationOperator --> TextGenerationParams
SummarizationOperator --> SummarizeTask
GenericGenerationOperator --> GenericPromptTask
QAGenerationOperator --> RagAnswerTask
%%{init: {'theme': 'base', 'themeVariables': {"darkMode": true, "background": "#0d1117", "primaryColor": "#21262d", "primaryTextColor": "#e6edf3", "primaryBorderColor": "#8b949e", "lineColor": "#8b949e", "textColor": "#e6edf3", "edgeLabelBackground": "#161b22", "actorBkg": "#21262d", "actorBorder": "#8b949e", "actorTextColor": "#e6edf3", "actorLineColor": "#8b949e", "signalColor": "#8b949e", "signalTextColor": "#e6edf3", "noteBkgColor": "#373320", "noteBorderColor": "#d4a72c", "noteTextColor": "#f0e6c0", "labelBoxBkgColor": "#21262d", "labelBoxBorderColor": "#8b949e", "labelTextColor": "#e6edf3", "loopTextColor": "#e6edf3", "activationBkgColor": "#30363d", "activationBorderColor": "#8b949e"}}}%%
classDiagram
class GenerationTask {
<<abstract>>
+required_inputs tuple
+_default_sampling dict
+build_request(**inputs) GenerationRequest
+parse(raw_text) str
+execute(client, **inputs) GeneratedTextResult
}
class RagAnswerTask {
+prompt Optional[str]
+system_prompt Optional[str]
+reasoning_enabled Optional[bool]
+build_request(**inputs) GenerationRequest
+parse(raw_text) str
}
class SummarizeTask {
+prompt Optional[str]
+system_prompt Optional[str]
+reasoning_enabled Optional[bool]
+build_request(**inputs) GenerationRequest
+parse(raw_text) str
}
class GenericPromptTask {
+prompt str
+input_names tuple
+system_prompt Optional[str]
+reasoning_enabled Optional[bool]
+build_request(**inputs) GenerationRequest
+parse(raw_text) str
}
class TextGenerationOperator {
<<abstract>>
+params TextGenerationParams
+_client CompletionClient
+_create_task() GenerationTask
+_get_generation_constructor_kwargs() dict
+process(data) DataFrame
}
class SummarizationOperator {
+input_column str
+output_column str
+_create_task() GenerationTask
}
class GenericGenerationOperator {
+input_columns Mapping
+output_column str
+_create_task() GenerationTask
}
class QAGenerationOperator {
+model str
+_execute_task(inputs) GeneratedTextResult
+_create_task() GenerationTask
}
class TextGenerationParams {
+transport LLMRemoteClientParams
+sampling LLMSamplingOverrides
+prompt Optional[str]
+reasoning_enabled Optional[bool]
+resolve_sampling(defaults) LLMInferenceParams
+from_kwargs(...) TextGenerationParams
}
GenerationTask <|-- RagAnswerTask
GenerationTask <|-- SummarizeTask
GenerationTask <|-- GenericPromptTask
TextGenerationOperator <|-- SummarizationOperator
TextGenerationOperator <|-- GenericGenerationOperator
TextGenerationOperator <|-- QAGenerationOperator
TextGenerationOperator --> GenerationTask : creates via _create_task
TextGenerationOperator --> TextGenerationParams
SummarizationOperator --> SummarizeTask
GenericGenerationOperator --> GenericPromptTask
QAGenerationOperator --> RagAnswerTask
Comments Outside Diff (1)
-
nemo_retriever/src/nemo_retriever/models/llm/tasks/base.py, line 1448-1454 (link)Exception swallowed silently — no logging
executeis the pipeline row-level error boundary, and the customno-bare-exceptrule requires that everyexcept Exceptionboundary logs with full context (exc_info=True) before converting to a failure state. As written, transport failures (network timeouts, model errors, etc.) are silently absorbed into theerrorfield with no trace in the log. Becausetask.execute()returns aGeneratedTextResultrather than raising, the operator-levellogger.warning("Row %d generation failed …")inTextGenerationOperator.processnever fires for these failures, so the exception disappears entirely.Prompt To Fix With AI
This is a comment left during a code review. Path: nemo_retriever/src/nemo_retriever/models/llm/tasks/base.py Line: 1448-1454 Comment: **Exception swallowed silently — no logging** `execute` is the pipeline row-level error boundary, and the custom `no-bare-except` rule requires that every `except Exception` boundary logs with full context (`exc_info=True`) before converting to a failure state. As written, transport failures (network timeouts, model errors, etc.) are silently absorbed into the `error` field with no trace in the log. Because `task.execute()` returns a `GeneratedTextResult` rather than raising, the operator-level `logger.warning("Row %d generation failed …")` in `TextGenerationOperator.process` never fires for these failures, so the exception disappears entirely. How can I resolve this? If you propose a fix, please make it concise.
Prompt To Fix All With AI
Fix the following 4 code review issues. Work through them one at a time, proposing concise fixes.
---
### Issue 1 of 4
nemo_retriever/src/nemo_retriever/models/llm/tasks/base.py:1448-1454
**Exception swallowed silently — no logging**
`execute` is the pipeline row-level error boundary, and the custom `no-bare-except` rule requires that every `except Exception` boundary logs with full context (`exc_info=True`) before converting to a failure state. As written, transport failures (network timeouts, model errors, etc.) are silently absorbed into the `error` field with no trace in the log. Because `task.execute()` returns a `GeneratedTextResult` rather than raising, the operator-level `logger.warning("Row %d generation failed …")` in `TextGenerationOperator.process` never fires for these failures, so the exception disappears entirely.
### Issue 2 of 4
nemo_retriever/src/nemo_retriever/operators/generation/base.py:238-243
**Bare `except Exception` without logging**
`_failure_model` is not a defined boundary (not a FastAPI handler, CLI entry point, Ray actor loop, or top-level pipeline handler), so the `no-bare-except` rule applies strictly. Any exception that occurs while reading `self._client.model` is silently discarded, which hides misconfigured clients and makes failures hard to diagnose.
```suggestion
def _failure_model(self) -> str:
try:
model = self._client.model
except Exception:
logger.debug("Could not read client model name, falling back to configured model", exc_info=True)
return self._configured_model
return model if isinstance(model, str) and model else self._configured_model
```
### Issue 3 of 4
nemo_retriever/src/nemo_retriever/tools/evaluation/generation.py:95-97
**`reasoning_enabled` is not forwarded to legacy `generate()` callers**
When a legacy `LLMClient` (one that only exposes `generate`, not `complete`) is injected, the adapter silently drops the operator-level `reasoning_enabled`. A caller who constructs `QAGenerationOperator(reasoning_enabled=False)` and injects a legacy client will find that reasoning is *not* disabled — the legacy client receives `reasoning_enabled=None` and falls back to its own transport default (`True`). The `LLMClient.generate` protocol already accepts `reasoning_enabled` as a keyword argument, so it can be forwarded without breaking anything.
```suggestion
if isinstance(client, LLMClient):
task_reasoning = getattr(self._task, "reasoning_enabled", None)
result = client.generate(inputs["query"], inputs["chunks"], reasoning_enabled=task_reasoning)
return GeneratedTextResult(
```
### Issue 4 of 4
nemo_retriever/src/nemo_retriever/models/llm/tasks/__init__.py:1-3
**SPDX copyright year format inconsistency across new files**
The `tasks/` package and operator files use three different year formats: `2024-26` (non-standard shorthand), `2024-25`, and `2026` (standalone). The `spdx-license-header` rule requires the current year. New files in this PR use `2024-26` (`tasks/__init__.py`, `tasks/base.py`, `tasks/generic.py`, `tasks/rag_answer.py`, `tasks/summarize.py`) while `test_generation_tasks.py` uses `2026`. The standard form in the rest of the repo is the four-digit year; the abbreviated `2024-26` is not conventional. All new files should use the same format as the existing codebase.
Reviews (1): Last reviewed commit: "Add abstraction for LLM-based tasks" | Re-trigger Greptile
![greptile-apps[bot]](https://avatars.githubusercontent.com/in/867647?s=60&v=4){kind=small}
Comment on lines
+238
to
+243
| def _failure_model(self) -> str: | ||
| try: | ||
| model = self._client.model | ||
| except Exception: | ||
| return self._configured_model | ||
| return model if isinstance(model, str) and model else self._configured_model |
Contributor
There was a problem hiding this comment.
except Exception without logging
_failure_model is not a defined boundary (not a FastAPI handler, CLI entry point, Ray actor loop, or top-level pipeline handler), so the no-bare-except rule applies strictly. Any exception that occurs while reading self._client.model is silently discarded, which hides misconfigured clients and makes failures hard to diagnose.
Suggested change
| def _failure_model(self) -> str: | |
| try: | |
| model = self._client.model | |
| except Exception: | |
| return self._configured_model | |
| return model if isinstance(model, str) and model else self._configured_model | |
| def _failure_model(self) -> str: | |
| try: | |
| model = self._client.model | |
| except Exception: | |
| logger.debug("Could not read client model name, falling back to configured model", exc_info=True) | |
| return self._configured_model | |
| return model if isinstance(model, str) and model else self._configured_model |
Prompt To Fix With AI
This is a comment left during a code review.
Path: nemo_retriever/src/nemo_retriever/operators/generation/base.py
Line: 238-243
Comment:
**Bare `except Exception` without logging**
`_failure_model` is not a defined boundary (not a FastAPI handler, CLI entry point, Ray actor loop, or top-level pipeline handler), so the `no-bare-except` rule applies strictly. Any exception that occurs while reading `self._client.model` is silently discarded, which hides misconfigured clients and makes failures hard to diagnose.
```suggestion
def _failure_model(self) -> str:
try:
model = self._client.model
except Exception:
logger.debug("Could not read client model name, falling back to configured model", exc_info=True)
return self._configured_model
return model if isinstance(model, str) and model else self._configured_model
```
How can I resolve this? If you propose a fix, please make it concise.
Comment on lines
+95
to
+97
| if isinstance(client, LLMClient): | ||
| result = client.generate(inputs["query"], inputs["chunks"]) | ||
| return GeneratedTextResult( |
Contributor
There was a problem hiding this comment.
reasoning_enabled is not forwarded to legacy generate() callers
When a legacy LLMClient (one that only exposes generate, not complete) is injected, the adapter silently drops the operator-level reasoning_enabled. A caller who constructs QAGenerationOperator(reasoning_enabled=False) and injects a legacy client will find that reasoning is not disabled — the legacy client receives reasoning_enabled=None and falls back to its own transport default (True). The LLMClient.generate protocol already accepts reasoning_enabled as a keyword argument, so it can be forwarded without breaking anything.
Suggested change
| if isinstance(client, LLMClient): | |
| result = client.generate(inputs["query"], inputs["chunks"]) | |
| return GeneratedTextResult( | |
| if isinstance(client, LLMClient): | |
| task_reasoning = getattr(self._task, "reasoning_enabled", None) | |
| result = client.generate(inputs["query"], inputs["chunks"], reasoning_enabled=task_reasoning) | |
| return GeneratedTextResult( |
Prompt To Fix With AI
This is a comment left during a code review.
Path: nemo_retriever/src/nemo_retriever/tools/evaluation/generation.py
Line: 95-97
Comment:
**`reasoning_enabled` is not forwarded to legacy `generate()` callers**
When a legacy `LLMClient` (one that only exposes `generate`, not `complete`) is injected, the adapter silently drops the operator-level `reasoning_enabled`. A caller who constructs `QAGenerationOperator(reasoning_enabled=False)` and injects a legacy client will find that reasoning is *not* disabled — the legacy client receives `reasoning_enabled=None` and falls back to its own transport default (`True`). The `LLMClient.generate` protocol already accepts `reasoning_enabled` as a keyword argument, so it can be forwarded without breaking anything.
```suggestion
if isinstance(client, LLMClient):
task_reasoning = getattr(self._task, "reasoning_enabled", None)
result = client.generate(inputs["query"], inputs["chunks"], reasoning_enabled=task_reasoning)
return GeneratedTextResult(
```
How can I resolve this? If you propose a fix, please make it concise.
Comment on lines
+1
to
+3
| # SPDX-FileCopyrightText: Copyright (c) 2024-26, NVIDIA CORPORATION & AFFILIATES. | ||
| # All rights reserved. | ||
| # SPDX-License-Identifier: Apache-2.0 |
Contributor
There was a problem hiding this comment.
The tasks/ package and operator files use three different year formats: 2024-26 (non-standard shorthand), 2024-25, and 2026 (standalone). The spdx-license-header rule requires the current year. New files in this PR use 2024-26 (tasks/__init__.py, tasks/base.py, tasks/generic.py, tasks/rag_answer.py, tasks/summarize.py) while test_generation_tasks.py uses 2026. The standard form in the rest of the repo is the four-digit year; the abbreviated 2024-26 is not conventional. All new files should use the same format as the existing codebase.
Prompt To Fix With AI
This is a comment left during a code review.
Path: nemo_retriever/src/nemo_retriever/models/llm/tasks/__init__.py
Line: 1-3
Comment:
**SPDX copyright year format inconsistency across new files**
The `tasks/` package and operator files use three different year formats: `2024-26` (non-standard shorthand), `2024-25`, and `2026` (standalone). The `spdx-license-header` rule requires the current year. New files in this PR use `2024-26` (`tasks/__init__.py`, `tasks/base.py`, `tasks/generic.py`, `tasks/rag_answer.py`, `tasks/summarize.py`) while `test_generation_tasks.py` uses `2026`. The standard form in the rest of the repo is the four-digit year; the abbreviated `2024-26` is not conventional. All new files should use the same format as the existing codebase.
How can I resolve this? If you propose a fix, please make it concise.Note: If this suggestion doesn't match your team's coding style, reply to this and let me know. I'll remember it for next time!
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Description
Checklist