LLM Based Task

[KyleZheng1284]

Description

• Adds a reusable TextGenerationOperator built on the existing operator framework.
• Introduces typed tasks for QA, summarization, and generic prompt generation.
• Preserves existing QA APIs, schemas, prompts, and client compatibility.
• Improves secure graph serialization without persisting API credentials.
• Preserves sampling overrides and DataFrame value types across execution.
• Keeps embedding and captioning as separate specialized operator families.
• Adds tests for generation, persistence, validation, and backward compatibility

Checklist

• I am familiar with the Contributing Guidelines.
• New or existing tests cover these changes.
• The documentation is up to date with these changes.

[greptile-apps]

Greptile Summary

This PR introduces a reusable TextGenerationOperator framework with typed task strategies (RagAnswerTask, SummarizeTask, GenericPromptTask), refactors QAGenerationOperator to sit on top of it, adds TextGenerationParams/LLMSamplingOverrides Pydantic models, and significantly hardens graph serialization to prevent credential persistence.

• New task layer (models/llm/tasks/) provides a strict four-phase lifecycle (build → transport → parse → result) with stable error codes and GenerationTaskError, replacing ad-hoc exception swallowing in the old LiteLLMClient.generate().
• New operator layer (operators/generation/) owns DataFrame validation, positional row tracking, ThreadPoolExecutor-based concurrency, and graph-safe constructor-kwarg serialization; QAGenerationOperator is migrated to this base while retaining its flat constructor API.
• Graph serialization in graph_pipeline_registry.py gains credential-aware encoding that blocks literal API keys, supports os.environ/VAR references, and performs lossless Pydantic round-tripping with fields_set preservation.

Confidence Score: 4/5

Safe to merge with one fix: the RAG prompt validator must require at least one {context} or {query} placeholder before the new custom-prompt path goes to production.

The _validate_rag_prompt function accepts prompts with zero placeholders — RagAnswerTask(prompt="Static text") construction succeeds, but build_request silently calls "Static text".format(context=..., query=...) and discards both the retrieved chunks and the user query. GenericPromptTask already enforces an "at least one declared placeholder" check; the RAG path is missing the analogous guard. Everything else — the four-phase task lifecycle, graph-safe serialization, sampling overrides, and QAGenerationOperator migration — looks well-considered and well-tested.

nemo_retriever/src/nemo_retriever/models/llm/tasks/rag_answer.py — _validate_rag_prompt needs a non-empty placeholder check

Important Files Changed

Filename	Overview
nemo_retriever/src/nemo_retriever/models/llm/tasks/rag_answer.py	New RagAnswerTask with prompt validation that permits zero-placeholder templates, silently discarding context and query at inference time.
nemo_retriever/src/nemo_retriever/models/llm/tasks/base.py	New TextGenerationTask/execute lifecycle; well-structured phase separation and error codes. Bare except-blocks in invoke() without exc_info were flagged in prior review cycles.
nemo_retriever/src/nemo_retriever/operators/generation/base.py	New TextGenerationOperator base; positional row tracking, overwrite guards, and graph-safe constructor kwarg validation look solid.
nemo_retriever/src/nemo_retriever/common/params/models.py	Adds TextGenerationParams, LLMSamplingOverrides, environment-reference handling, and secret-redaction helpers. Extra-params protection and _auto_resolve_unset_api_keys split look correct.
nemo_retriever/src/nemo_retriever/models/llm/clients/litellm.py	Refactored to delegate RAG logic to RagAnswerTask and the new TextCompletionClient contract. Response validation is improved. Redundant extra_params validation on transport at call time is minor overhead.
nemo_retriever/src/nemo_retriever/graph/graph_pipeline_registry.py	Significant expansion of graph serialization with credential-aware encoding, env-reference handling, and lossless Pydantic round-tripping. The secret-redaction heuristics look thorough.
nemo_retriever/src/nemo_retriever/tools/evaluation/generation.py	QAGenerationOperator migrated to TextGenerationOperator base with legacy-client adapter; reasoning_enabled now forwarded correctly to legacy generate() callers.
nemo_retriever/tests/test_generation_tasks.py	Comprehensive tests covering lifecycle phases, error codes, prompt validation, and backward compatibility. No test for zero-placeholder RAG prompt behavior.

Sequence Diagram

%%{init: {'theme': 'neutral'}}%%
sequenceDiagram
    participant Caller
    participant Op as TextGenerationOperator
    participant Task as TextGenerationTask
    participant Client as TextCompletionClient

    Caller->>Op: process(DataFrame)
    Op->>Op: _validate_and_resolve_dataframe()
    loop per row (ThreadPoolExecutor)
        Op->>Op: _execute_row(position, inputs)
        Op->>Task: "invoke(client, **inputs)"
        Task->>Task: "_preflight_error(**inputs)"
        Task->>Task: "build_request(**inputs) → GenerationRequest"
        Task->>Client: complete(messages, max_tokens, extra_params)
        Client-->>Task: (raw_text, latency_s)
        Task->>Task: parse(raw_text) → text
        Task-->>Op: GeneratedTextResult
    end
    Op-->>Caller: DataFrame with output/latency/model/error columns

Loading

%%{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"}}}%%
sequenceDiagram
    participant Caller
    participant Op as TextGenerationOperator
    participant Task as TextGenerationTask
    participant Client as TextCompletionClient

    Caller->>Op: process(DataFrame)
    Op->>Op: _validate_and_resolve_dataframe()
    loop per row (ThreadPoolExecutor)
        Op->>Op: _execute_row(position, inputs)
        Op->>Task: "invoke(client, **inputs)"
        Task->>Task: "_preflight_error(**inputs)"
        Task->>Task: "build_request(**inputs) → GenerationRequest"
        Task->>Client: complete(messages, max_tokens, extra_params)
        Client-->>Task: (raw_text, latency_s)
        Task->>Task: parse(raw_text) → text
        Task-->>Op: GeneratedTextResult
    end
    Op-->>Caller: DataFrame with output/latency/model/error columns

Loading

Prompt To Fix All With AI

Fix the following 1 code review issue. Work through them one at a time, proposing concise fixes.

---

### Issue 1 of 1
nemo_retriever/src/nemo_retriever/models/llm/tasks/rag_answer.py:43-47
`_validate_rag_prompt` validates that only `{context}` and `{query}` appear, but doesn't require either to be present. A prompt like `"Here is my answer."` passes validation, and `build_request` then calls `self.prompt.format(context=..., query=...)` which silently returns the static string — both the retrieved context and the user query are discarded. Compare with `GenericPromptTask._validate_prompt_template`, which explicitly raises if `fields` is empty. The fix is to require at least one allowed placeholder.

```suggestion
    used_fields: set[str] = set()
    for _, field_name, format_spec, conversion in parsed_fields:
        if field_name is None:
            continue
        if field_name not in {"context", "query"} or format_spec or conversion:
            raise ValueError("RAG prompt may only use simple {context} and {query} placeholders")
        used_fields.add(field_name)
    if not used_fields:
        raise ValueError("RAG prompt must contain at least one {context} or {query} placeholder")
```

_{Reviews (5): Last reviewed commit: "Merge branch 'main' into feature/llm-tas..." | Re-trigger Greptile}

[mahikaw]

nit: do we want to update it here too for the api_key=None behaviour and gen error codes nemo_retriever/src/nemo_retriever/tools/evaluation/README.md

[mahikaw]

another nit: do you think it might be useful for non-secret generation config like temperature, top_p, max_tokens, and reasoning_enabled be emitted into eval metadata/traces so runs are easier to compare (if that's the goal with this infra)?

[mahikaw]

if this is envisioned to be generation specific base let's rename it to something more reflective like TextGenerationTask

[mahikaw]

looks good, added some comments that might lead to minor refactoring.

[copy-pr-bot]

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.

[greptile-apps]

_validate_rag_prompt validates that only {context} and {query} appear, but doesn't require either to be present. A prompt like "Here is my answer." passes validation, and build_request then calls self.prompt.format(context=..., query=...) which silently returns the static string — both the retrieved context and the user query are discarded. Compare with GenericPromptTask._validate_prompt_template, which explicitly raises if fields is empty. The fix is to require at least one allowed placeholder.

Suggested change

	for _, field_name, format_spec, conversion in parsed_fields:
	if field_name is None:
	continue
	if field_name not in {"context", "query"} or format_spec or conversion:
	raise ValueError("RAG prompt may only use simple {context} and {query} placeholders")
	used_fields: set[str] = set()
	for _, field_name, format_spec, conversion in parsed_fields:
	if field_name is None:
	continue
	if field_name not in {"context", "query"} or format_spec or conversion:
	raise ValueError("RAG prompt may only use simple {context} and {query} placeholders")
	used_fields.add(field_name)
	if not used_fields:
	raise ValueError("RAG prompt must contain at least one {context} or {query} placeholder")

Prompt To Fix With AI

This is a comment left during a code review.
Path: nemo_retriever/src/nemo_retriever/models/llm/tasks/rag_answer.py
Line: 43-47

Comment:
`_validate_rag_prompt` validates that only `{context}` and `{query}` appear, but doesn't require either to be present. A prompt like `"Here is my answer."` passes validation, and `build_request` then calls `self.prompt.format(context=..., query=...)` which silently returns the static string — both the retrieved context and the user query are discarded. Compare with `GenericPromptTask._validate_prompt_template`, which explicitly raises if `fields` is empty. The fix is to require at least one allowed placeholder.

```suggestion
    used_fields: set[str] = set()
    for _, field_name, format_spec, conversion in parsed_fields:
        if field_name is None:
            continue
        if field_name not in {"context", "query"} or format_spec or conversion:
            raise ValueError("RAG prompt may only use simple {context} and {query} placeholders")
        used_fields.add(field_name)
    if not used_fields:
        raise ValueError("RAG prompt must contain at least one {context} or {query} placeholder")
```

How can I resolve this? If you propose a fix, please make it concise.