feat: add VLM generation support to the RAG pipeline#2305
feat: add VLM generation support to the RAG pipeline#2305nv-nikkulkarni wants to merge 3 commits into
Conversation
Greptile SummaryThis PR wires multimodal (VLM) answer generation end-to-end through the RAG pipeline — from stored image URIs in LanceDB to an OpenAI vision content array sent to a VLM. Previous-round feedback on SSRF validation, size guards, MIME-type inference, logging level, and constant duplication has been addressed in this revision.
|
| Filename | Overview |
|---|---|
| nemo_retriever/src/nemo_retriever/models/llm/clients/vlm_litellm.py | New VLM client with image loading, SSRF/size guards, and multimodal prompt builder. _DEFAULT_MODEL class attribute on LiteVLMClient is dead code because the inherited from_kwargs bakes in LiteLLMClient._DEFAULT_MODEL at class-definition time; omitting --llm-model silently routes to a text-only model. |
| nemo_retriever/src/nemo_retriever/graph/retriever.py | Merges multimodal support into answer() via multimodal: bool = False; adds _hits_to_multimodal_chunks helper. Both code paths (text-only and multimodal) ultimately call self.query() and construct RetrievalResult equivalently. |
| nemo_retriever/src/nemo_retriever/models/llm/types.py | Adds MultimodalChunk dataclass, MultimodalLLMClient protocol, and VISUAL_CONTENT_TYPES constant — all cleanly defined as the canonical shared location. |
| nemo_retriever/src/nemo_retriever/cli/query/app.py | Adds answer CLI command with --multimodal flag and full LLM config options. The from_kwargs default-model bug (in vlm_litellm.py) means --multimodal without --llm-model silently uses the wrong model. |
| nemo_retriever/src/nemo_retriever/cli/query/options.py | Adds well-typed Annotated option aliases for LLM config and --multimodal flag; straightforward additions. |
| nemo_retriever/src/nemo_retriever/models/llm/init.py | Exports MultimodalChunk, MultimodalLLMClient, and registers LiteVLMClient for lazy import; no issues. |
| nemo_retriever/src/nemo_retriever/models/llm/clients/init.py | Registers LiteVLMClient in the client registry and __all__; straightforward. |
| nemo_retriever/tests/test_vlm_litellm.py | Covers all helper functions with good positive/negative cases; LiteVLMClient.generate_multimodal() itself is not directly exercised. |
Sequence Diagram
%%{init: {'theme': 'neutral'}}%%
sequenceDiagram
participant CLI as retriever query answer
participant Retriever
participant LiteVLMClient
participant LanceDB
participant ImageLoader as _load_image_as_base64
participant VLM as LiteLLM / VLM API
CLI->>Retriever: "answer(query, llm=LiteVLMClient, multimodal=True)"
Retriever->>LanceDB: query(query, top_k)
LanceDB-->>Retriever: hits (text + stored_image_uri)
Retriever->>Retriever: _hits_to_multimodal_chunks(hits)
Retriever->>LiteVLMClient: generate_multimodal(query, mm_chunks)
LiteVLMClient->>LiteVLMClient: _build_multimodal_rag_prompt(query, chunks)
loop per visual chunk
LiteVLMClient->>ImageLoader: _load_image_as_base64(image_uri)
ImageLoader-->>LiteVLMClient: base64 string (or None - text fallback)
end
LiteVLMClient->>VLM: litellm.completion(messages with image_url blocks)
VLM-->>LiteVLMClient: GenerationResult
LiteVLMClient-->>Retriever: GenerationResult
Retriever-->>CLI: AnswerResult (JSON)
%%{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 CLI as retriever query answer
participant Retriever
participant LiteVLMClient
participant LanceDB
participant ImageLoader as _load_image_as_base64
participant VLM as LiteLLM / VLM API
CLI->>Retriever: "answer(query, llm=LiteVLMClient, multimodal=True)"
Retriever->>LanceDB: query(query, top_k)
LanceDB-->>Retriever: hits (text + stored_image_uri)
Retriever->>Retriever: _hits_to_multimodal_chunks(hits)
Retriever->>LiteVLMClient: generate_multimodal(query, mm_chunks)
LiteVLMClient->>LiteVLMClient: _build_multimodal_rag_prompt(query, chunks)
loop per visual chunk
LiteVLMClient->>ImageLoader: _load_image_as_base64(image_uri)
ImageLoader-->>LiteVLMClient: base64 string (or None - text fallback)
end
LiteVLMClient->>VLM: litellm.completion(messages with image_url blocks)
VLM-->>LiteVLMClient: GenerationResult
LiteVLMClient-->>Retriever: GenerationResult
Retriever-->>CLI: AnswerResult (JSON)
Reviews (5): Last reviewed commit: "fix: use explicit MIME table to avoid pl..." | Re-trigger Greptile
5011dec to
62c5869
Compare
|
|
||
| logger = logging.getLogger(__name__) | ||
|
|
||
| _VISUAL_CONTENT_TYPES = frozenset({"image", "chart", "infographic", "table"}) |
There was a problem hiding this comment.
_VISUAL_CONTENT_TYPES duplicated across modules
An identical _VISUAL_CONTENT_TYPES = frozenset({"image", "chart", "infographic", "table"}) constant is defined here and again in graph/retriever.py. If a new visual content type needs to be added (e.g. "diagram"), both sites must be updated in sync. Consider moving it to nemo_retriever.models.llm.types (alongside MultimodalChunk) as the canonical definition and importing from there.
Prompt To Fix With AI
This is a comment left during a code review.
Path: nemo_retriever/src/nemo_retriever/models/llm/clients/vlm_litellm.py
Line: 37
Comment:
**`_VISUAL_CONTENT_TYPES` duplicated across modules**
An identical `_VISUAL_CONTENT_TYPES = frozenset({"image", "chart", "infographic", "table"})` constant is defined here and again in `graph/retriever.py`. If a new visual content type needs to be added (e.g. `"diagram"`), both sites must be updated in sync. Consider moving it to `nemo_retriever.models.llm.types` (alongside `MultimodalChunk`) as the canonical definition and importing from there.
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!
| # SPDX-FileCopyrightText: Copyright (c) 2024-25, NVIDIA CORPORATION & AFFILIATES. | ||
| # All rights reserved. | ||
| # SPDX-License-Identifier: Apache-2.0 | ||
|
|
||
| """ | ||
| VLM-capable LLM generation client for multimodal RAG. | ||
|
|
||
| LiteVLMClient extends LiteLLMClient with generate_multimodal(), which | ||
| accepts MultimodalChunk objects. Chunks that carry a stored_image_uri have | ||
| their image loaded, base64-encoded, and injected into the OpenAI vision | ||
| content array alongside the text caption/description. | ||
|
|
||
| Supported URI schemes for image_uri: | ||
| - Local paths : /abs/path/to/image.png or file:///abs/path/to/image.png | ||
| - HTTP(S) URLs : https://host/path/image.png | ||
| - S3 URIs : s3://bucket/key (requires boto3; skipped with a warning if absent) | ||
| """ | ||
|
|
||
| from __future__ import annotations | ||
|
|
||
| import base64 | ||
| import logging | ||
| from typing import Any, Optional | ||
|
|
||
| from nemo_retriever.models.llm.clients.litellm import ( | ||
| LiteLLMClient, | ||
| _format_rag_system_prompt, | ||
| _with_no_reasoning_controls, | ||
| _NO_REASONING_EXTRA_PARAMS, | ||
| ) | ||
| from nemo_retriever.models.llm.text_utils import strip_think_tags | ||
| from nemo_retriever.models.llm.types import GenerationResult, MultimodalChunk | ||
| from nemo_retriever.common.params.models import LLMInferenceParams, LLMRemoteClientParams | ||
|
|
||
| logger = logging.getLogger(__name__) | ||
|
|
||
| _VISUAL_CONTENT_TYPES = frozenset({"image", "chart", "infographic", "table"}) | ||
|
|
||
| _VLM_RAG_SYSTEM_PROMPT = ( | ||
| "You are a precise question-answering assistant. " | ||
| "Answer the question using ONLY the information provided in the context below, " | ||
| "which may include text passages and images. " | ||
| "If the context does not contain enough information to answer, say so clearly. " | ||
| "Be concise and factual." | ||
| ) | ||
|
|
||
|
|
||
| def _load_image_as_base64(uri: str) -> Optional[str]: | ||
| """Load an image from a local path, file://, http(s)://, or s3:// URI. | ||
|
|
||
| Returns the raw base64 string (no data-URI prefix), or None if loading | ||
| fails so that callers can gracefully degrade to text-only. | ||
| """ | ||
| try: | ||
| if uri.startswith("s3://"): | ||
| try: | ||
| import boto3 | ||
|
|
||
| parts = uri[5:].split("/", 1) | ||
| bucket, key = parts[0], parts[1] if len(parts) > 1 else "" | ||
| s3 = boto3.client("s3") | ||
| response = s3.get_object(Bucket=bucket, Key=key) | ||
| data = response["Body"].read() | ||
| except ImportError: | ||
| logger.warning("boto3 not installed; skipping image at %s", uri) | ||
| return None | ||
| elif uri.startswith("http://") or uri.startswith("https://"): | ||
| import urllib.request | ||
|
|
||
| with urllib.request.urlopen(uri, timeout=10) as resp: # noqa: S310 | ||
| data = resp.read() | ||
| else: | ||
| # Local path — strip optional file:// scheme | ||
| path = uri.removeprefix("file://") | ||
| with open(path, "rb") as fh: | ||
| data = fh.read() | ||
|
|
||
| return base64.b64encode(data).decode("ascii") | ||
| except Exception as exc: | ||
| logger.warning("Failed to load image %s: %s", uri, exc) | ||
| return None | ||
|
|
||
|
|
||
| def _build_multimodal_rag_prompt( | ||
| query: str, | ||
| chunks: list[MultimodalChunk], | ||
| *, | ||
| formatted_rag_system_prompt: str, | ||
| ) -> list[dict]: | ||
| """Build an OpenAI-style vision messages list for multimodal RAG. | ||
|
|
||
| Each chunk becomes one or more content blocks in the user message: | ||
| - Text-only chunks: a single ``{"type": "text", ...}`` block. | ||
| - Visual chunks: a text block (caption/description) followed by an | ||
| ``{"type": "image_url", ...}`` block carrying the base64 data-URI. | ||
| If the image cannot be loaded the chunk falls back to text-only. | ||
| """ | ||
| user_content: list[dict[str, Any]] = [] | ||
|
|
||
| if not chunks: | ||
| user_content.append({"type": "text", "text": "(no context retrieved)"}) | ||
| else: | ||
| user_content.append({"type": "text", "text": "Context:\n"}) | ||
| for i, chunk in enumerate(chunks): | ||
| label = f"[{i + 1}] ({chunk.content_type})" | ||
| if chunk.image_uri and chunk.content_type in _VISUAL_CONTENT_TYPES: | ||
| b64 = _load_image_as_base64(chunk.image_uri) | ||
| if b64: | ||
| if chunk.text: | ||
| user_content.append( | ||
| {"type": "text", "text": f"{label} {chunk.text}\n"} | ||
| ) | ||
| else: | ||
| user_content.append({"type": "text", "text": f"{label}\n"}) | ||
| user_content.append( | ||
| { | ||
| "type": "image_url", | ||
| "image_url": {"url": f"data:image/png;base64,{b64}"}, | ||
| } | ||
| ) | ||
| else: | ||
| # Image load failed — fall back to caption text only | ||
| user_content.append( | ||
| {"type": "text", "text": f"{label} {chunk.text}\n"} | ||
| ) | ||
| else: | ||
| user_content.append( | ||
| {"type": "text", "text": f"{label} {chunk.text}\n"} | ||
| ) | ||
| if i < len(chunks) - 1: | ||
| user_content.append({"type": "text", "text": "\n---\n\n"}) | ||
|
|
||
| user_content.append({"type": "text", "text": f"\nQuestion: {query}\n\nAnswer:"}) | ||
|
|
||
| return [ | ||
| {"role": "system", "content": formatted_rag_system_prompt}, | ||
| {"role": "user", "content": user_content}, | ||
| ] | ||
|
|
||
|
|
||
| class LiteVLMClient(LiteLLMClient): | ||
| """LiteLLM client extended with multimodal (vision) generation. | ||
|
|
||
| Inherits all text-only generation from LiteLLMClient and adds | ||
| generate_multimodal(), which accepts list[MultimodalChunk] and wires | ||
| images into the OpenAI vision content array. | ||
|
|
||
| Usage:: | ||
|
|
||
| client = LiteVLMClient.from_kwargs( | ||
| model="openai/gpt-4o", | ||
| api_base="https://integrate.api.nvidia.com/v1", | ||
| api_key="nvapi-...", | ||
| ) | ||
| result = retriever.answer_multimodal(query, llm=client) | ||
| """ | ||
|
|
||
| _DEFAULT_MODEL: str = "openai/meta/llama-4-scout-17b-16e-instruct" | ||
|
|
||
| def __init__( | ||
| self, | ||
| transport: LLMRemoteClientParams, | ||
| sampling: Optional[LLMInferenceParams] = None, | ||
| ): | ||
| super().__init__(transport=transport, sampling=sampling) | ||
| # Use VLM-specific system prompt when no custom prompt was configured | ||
| if not transport.rag_system_prompt and not transport.rag_system_prompt_prefix: | ||
| self._formatted_rag_system_prompt = _format_rag_system_prompt( | ||
| rag_system_prompt=_VLM_RAG_SYSTEM_PROMPT, | ||
| ) | ||
|
|
||
| def generate_multimodal( | ||
| self, | ||
| query: str, | ||
| chunks: list[MultimodalChunk], | ||
| *, | ||
| reasoning_enabled: Optional[bool] = None, | ||
| ) -> GenerationResult: | ||
| """Generate an answer using both text and image context from chunks.""" | ||
| messages = _build_multimodal_rag_prompt( | ||
| query, | ||
| chunks, | ||
| formatted_rag_system_prompt=self._formatted_rag_system_prompt, | ||
| ) | ||
| request_extra_params: dict[str, Any] | None = None | ||
| effective_reasoning_enabled = ( | ||
| self.transport.reasoning_enabled if reasoning_enabled is None else reasoning_enabled | ||
| ) | ||
| if not effective_reasoning_enabled: | ||
| messages = _with_no_reasoning_controls(messages) | ||
| request_extra_params = _NO_REASONING_EXTRA_PARAMS | ||
| try: | ||
| raw_answer, latency = self.complete(messages, extra_params=request_extra_params) | ||
| answer = strip_think_tags(raw_answer) | ||
| if not answer: | ||
| return GenerationResult( | ||
| answer="", | ||
| latency_s=latency, | ||
| model=self.transport.model, | ||
| error="thinking_truncated", | ||
| ) | ||
| return GenerationResult(answer=answer, latency_s=latency, model=self.transport.model) | ||
| except Exception as exc: | ||
| logger.debug("VLM generation failed for model=%s: %s", self.transport.model, exc) | ||
| return GenerationResult( | ||
| answer="", | ||
| latency_s=0.0, | ||
| model=self.transport.model, | ||
| error=str(exc), | ||
| ) |
There was a problem hiding this comment.
vlm_litellm.py is a new source module with no corresponding test file. The test-mirrors-source-structure and test-coverage-new-code rules require new business logic — especially the multimodal prompt builder, image loading paths (local, HTTP, S3, failure fallback), and generate_multimodal — to have unit tests. The PR checklist also shows this box unchecked.
Rule Used: New source modules should have corresponding test ... (source)
Prompt To Fix With AI
This is a comment left during a code review.
Path: nemo_retriever/src/nemo_retriever/models/llm/clients/vlm_litellm.py
Line: 1-210
Comment:
**No tests for `LiteVLMClient`**
`vlm_litellm.py` is a new source module with no corresponding test file. The `test-mirrors-source-structure` and `test-coverage-new-code` rules require new business logic — especially the multimodal prompt builder, image loading paths (local, HTTP, S3, failure fallback), and `generate_multimodal` — to have unit tests. The PR checklist also shows this box unchecked.
**Rule Used:** New source modules should have corresponding test ... ([source](.greptile))
How can I resolve this? If you propose a fix, please make it concise.There was a problem hiding this comment.
You're right — I missed it. nemo_retriever/tests/test_vlm_litellm.py is included in this PR with 206 lines of tests covering:
_mime_type_from_uri(PNG, JPEG, WebP, unknown fallback)_validate_http_uri(public allowed, AWS/GCP metadata blocked, private/loopback IPs blocked)_validate_local_path(PNG/JPEG allowed, non-image extensions rejected)_load_image_as_base64(local load,file://scheme, missing file, oversized file, SSRF-blocked host, oversized HTTP, missing boto3)_build_multimodal_rag_prompt(empty chunks, text-only, visual with image, MIME inference, image load failure fallback, question appended, system prompt role)
My original comment was incorrect — tests are present and have solid coverage. Sorry for the noise!
| RetrievalResult, | ||
| ) | ||
|
|
||
| _VISUAL_CONTENT_TYPES = frozenset({"image", "chart", "infographic", "table"}) |
There was a problem hiding this comment.
_VISUAL_CONTENT_TYPES is defined identically in both retriever.py and vlm_litellm.py. If one set is updated (e.g. adding "diagram") but not the other, the retriever will attach image_uri for a content type that the VLM prompt builder ignores, silently losing the image. The canonical definition belongs in types.py (or another shared module) and should be imported in both places.
| _VISUAL_CONTENT_TYPES = frozenset({"image", "chart", "infographic", "table"}) | |
| from nemo_retriever.models.llm.types import _VISUAL_CONTENT_TYPES # noqa: PLC2701 |
Prompt To Fix With AI
This is a comment left during a code review.
Path: nemo_retriever/src/nemo_retriever/graph/retriever.py
Line: 41
Comment:
`_VISUAL_CONTENT_TYPES` is defined identically in both `retriever.py` and `vlm_litellm.py`. If one set is updated (e.g. adding `"diagram"`) but not the other, the retriever will attach `image_uri` for a content type that the VLM prompt builder ignores, silently losing the image. The canonical definition belongs in `types.py` (or another shared module) and should be imported in both places.
```suggestion
from nemo_retriever.models.llm.types import _VISUAL_CONTENT_TYPES # noqa: PLC2701
```
How can I resolve this? If you propose a fix, please make it concise.There was a problem hiding this comment.
You're right to push back. Looking at the full PR diff, VISUAL_CONTENT_TYPES is not duplicated — it's defined once in types.py:
# types.py
VISUAL_CONTENT_TYPES: frozenset[str] = frozenset({"image", "chart", "infographic", "table"})And both files import it from there:
retriever.py:from nemo_retriever.models.llm.types import MultimodalChunk, VISUAL_CONTENT_TYPESvlm_litellm.py:from nemo_retriever.models.llm.types import VISUAL_CONTENT_TYPES, GenerationResult, MultimodalChunk
My original comment was incorrect — the PR already follows the canonical single-definition pattern I was recommending. Sorry for the noise!
| # SPDX-FileCopyrightText: Copyright (c) 2024-25, NVIDIA CORPORATION & AFFILIATES. | ||
| # All rights reserved. | ||
| # SPDX-License-Identifier: Apache-2.0 |
There was a problem hiding this comment.
vlm_litellm.py introduces _load_image_as_base64, _build_multimodal_rag_prompt, and LiteVLMClient.generate_multimodal — each with multiple branches (S3, HTTP, local, image load failure, thinking-truncated, generation error). The repo rules require new functionality to have corresponding unit tests covering happy path, error paths, and edge cases, and a test file mirroring the source structure. None are included in this PR.
Prompt To Fix With AI
This is a comment left during a code review.
Path: nemo_retriever/src/nemo_retriever/models/llm/clients/vlm_litellm.py
Line: 1-3
Comment:
**Missing tests for new module**
`vlm_litellm.py` introduces `_load_image_as_base64`, `_build_multimodal_rag_prompt`, and `LiteVLMClient.generate_multimodal` — each with multiple branches (S3, HTTP, local, image load failure, thinking-truncated, generation error). The repo rules require new functionality to have corresponding unit tests covering happy path, error paths, and edge cases, and a test file mirroring the source structure. None are included in this PR.
How can I resolve this? If you propose a fix, please make it concise.There was a problem hiding this comment.
You're right — I missed it. nemo_retriever/tests/test_vlm_litellm.py (206 lines) is included in this PR and covers all the branches I flagged: MIME inference, HTTP/SSRF validation, local path validation, S3 with missing boto3, oversized images, image load failure fallback, and prompt building. My original comment was incorrect. Apologies for the noise.
Tip: You can customize Greptile's behavior for this repo with .greptile/rules.md and .greptile/config.json.
- Add MultimodalChunk dataclass and MultimodalLLMClient protocol - Add LiteVLMClient with generate_multimodal() and image loading from stored_image_uri - Extend Retriever.answer() with multimodal: bool flag; remove separate answer_multimodal() - Add `retriever query answer` CLI command with --multimodal, LLM config, and --reference scoring Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
62c5869 to
0e85d38
Compare
- Fix MIME type: infer from URI extension via mimetypes instead of hardcoding image/png - Fix log level: VLM generation failures now logged at WARNING with exc_info=True - Add size guard: images >50MB are skipped with a warning to prevent OOM - Add SSRF guard: block private IPs, loopback, and cloud metadata endpoints for HTTP URIs - Add path traversal guard: reject local paths with non-image extensions - Fix copyright year: 2026 for new file - Fix stale docstring: answer_multimodal() → answer(..., multimodal=True) - Dedup VISUAL_CONTENT_TYPES: moved to types.py as canonical definition, imported in both retriever.py and vlm_litellm.py - Add tests: 28 unit tests covering MIME inference, SSRF/path guards, size limits, image loading paths, and prompt building Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
mimetypes.guess_type('.webp') returns None on some Linux distros.
Replace with a static _EXTENSION_MIME dict for all supported image types;
mimetypes and image/png remain as fallbacks for unknown extensions.
Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
Introduces multimodal answer generation as a first-class feature of Retriever.answer(), wired end-to-end from ingestion through retrieval to generation.
Key changes:
retriever query answercommand with --multimodal flag, full LLM config options (--llm-model, --llm-invoke-url, --llm-api-key-env, --llm-max-tokens, --llm-temperature, --reasoning), and optional --reference scoringThe ingestion-side image storage (StoreOperator, stored_image_uri in LanceDB) was already in place; this change closes the last-mile gap between stored images and the generation prompt.
Description
Checklist