Press n or j to go to the next uncovered block, b, p or k for the previous block.
| 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 226 227 228 229 230 231 232 233 234 235 236 237 238 239 240 241 242 243 244 245 246 247 248 249 250 251 252 253 254 255 256 257 258 259 260 261 262 263 264 265 266 267 268 269 270 | 2x 2x 2x 2x 2x 2x 2x 15x 2x 13x 13x 13x 13x 13x 13x 13x 2x 11x 2x 20x 20x 10x 10x 10x 10x 9x 9x 8x 9x 1x 8x 15x 15x 15x 4x 11x 1x 10x 10x 1x 9x 8x 8x 7x 1x 1x 6x 6x 5x 1x 2x 2x 9x 9x 9x 9x 9x 5x 5x 5x 1x 5x 5x 5x 1x 4x 1x 3x 3x 5x 1x 2x 1x 1x 1x 1x | import * as vscode from "vscode";
import { fetchWithRetry } from "../api/client";
import { classifyApiError } from "../api/errors";
import {
BASE_URL,
EXTENSION_VERSION,
MODELS_CACHE_KEY_FINGERPRINT_STATE_KEY,
MODELS_STATE_KEY,
PROVIDER_DISPLAY_NAME,
} from "../shared/constants";
import { ELITE_MODELS_WHITELIST, isNormalizedNvidiaModel } from "../models/catalog";
import { NvidiaApiKeyResolver } from "../api/key-resolver";
/** Maximum accepted image payload (binary bytes) sent to the vision model. */
const MAX_IMAGE_DATA_BYTES = 20 * 1024 * 1024;
/**
* Validate that image_data is a base64 image data URL and return its decoded
* size. Arbitrary remote URLs are rejected so the NIM vision endpoint never
* fetches external resources on the user's behalf.
*/
function measureImageDataUrl(imageData: string): number | undefined {
if (!imageData.startsWith("data:image/")) {
return undefined;
}
const comma = imageData.indexOf(",");
Iif (comma === -1) {
return undefined;
}
const prefix = imageData.slice(5, comma);
const payload = imageData.slice(comma + 1);
Iif (!payload) {
return undefined;
}
const headerFields = prefix.split(";");
if (!headerFields.includes("base64")) {
return undefined;
}
return Math.ceil((payload.length * 3) / 4);
}
/**
* Image-analysis client that uses a cached NVIDIA NIM vision-capable model.
*/
export class NimVisionClient {
constructor(
secrets: vscode.SecretStorage,
private readonly modelStorage?: vscode.Memento,
private readonly keyResolver = new NvidiaApiKeyResolver(secrets),
) {}
private async getApiKey(): Promise<string> {
const storedFingerprint = this.modelStorage?.get<unknown>(
MODELS_CACHE_KEY_FINGERPRINT_STATE_KEY,
);
const cacheKeyFingerprint =
typeof storedFingerprint === "string" ? storedFingerprint : undefined;
const resolved = await this.keyResolver.resolveForTool({ cacheKeyFingerprint });
return resolved?.value ?? "";
}
private getVisionModelId(): string {
const cachedModels = this.modelStorage?.get<unknown>(MODELS_STATE_KEY);
const visionModel = Array.isArray(cachedModels)
? cachedModels.find(
(model) =>
isNormalizedNvidiaModel(model) &&
Object.prototype.hasOwnProperty.call(ELITE_MODELS_WHITELIST, model.id) &&
ELITE_MODELS_WHITELIST[model.id].supportsVision &&
model.supportsVision,
)
: undefined;
if (!visionModel || !isNormalizedNvidiaModel(visionModel)) {
throw new Error(
`No NVIDIA NIM vision model is available. Run "${PROVIDER_DISPLAY_NAME}: Refresh Models" after setting your API key.`,
);
}
return visionModel.id;
}
async analyzeImage(imageData: string, prompt: string, signal?: AbortSignal): Promise<string> {
Iif (signal?.aborted) {
throw createAbortError();
}
const imageByteLength = measureImageDataUrl(imageData);
if (imageByteLength === undefined) {
throw new Error(
`${PROVIDER_DISPLAY_NAME} image analysis requires a base64 image data URL ('data:image/...;base64,...').`,
);
}
if (imageByteLength > MAX_IMAGE_DATA_BYTES) {
throw new Error(
`${PROVIDER_DISPLAY_NAME} image is too large (${Math.ceil(imageByteLength / (1024 * 1024))} MB). Maximum size is ${MAX_IMAGE_DATA_BYTES / (1024 * 1024)} MB.`,
);
}
const apiKey = await this.getApiKey();
if (!apiKey) {
throw new Error(`${PROVIDER_DISPLAY_NAME} API key not found`);
}
const model = this.getVisionModelId();
const ua = `nvidia-nim-provider/${EXTENSION_VERSION} VSCode/${vscode.version}`;
const response = await fetchWithRetry(
`${BASE_URL}/chat/completions`,
{
method: "POST",
headers: {
"Content-Type": "application/json",
Authorization: `Bearer ${apiKey}`,
"User-Agent": ua,
},
body: JSON.stringify({
model,
messages: [
{
role: "user",
content: [
{ type: "text", text: prompt },
{ type: "image_url", image_url: { url: imageData } },
],
},
],
max_tokens: 2000,
}),
signal,
},
3,
{ operation: "vision", model },
);
if (!response.ok) {
const errorText = await response.text();
throw classifyApiError(new Error(`HTTP ${response.status} ${response.statusText}`), {
operation: "vision",
status: response.status,
detail: errorText,
});
}
try {
const data = (await response.json()) as {
choices?: Array<{ message?: { content?: string } }>;
};
return data.choices?.[0]?.message?.content ?? "Failed to analyze image";
} catch (error) {
throw classifyApiError(error, { operation: "vision", model });
}
}
}
/**
* Tool for analyzing images using a cached NVIDIA NIM vision-capable model.
* Non-vision models can delegate image content to this tool for analysis.
*/
export class NimAnalyzeImageTool implements vscode.LanguageModelTool<{
image_data: string;
prompt: string;
}> {
static readonly id = "nvidia_nim_analyze_image";
readonly name = NimAnalyzeImageTool.id;
readonly description =
`Analyze an image using ${PROVIDER_DISPLAY_NAME} Vision. Use this tool when you need to ` +
"understand or describe the content of an image, extract text from images (OCR), " +
"or answer questions about visual content. Returns a detailed analysis of the image.";
readonly tags = ["vision", "image", "ocr", "analysis"];
readonly inputSchema = {
type: "object" as const,
properties: {
image_data: {
type: "string",
description:
"Base64-encoded image data URL (e.g., 'data:image/png;base64,...'). The image to analyze.",
},
prompt: {
type: "string",
description:
"The question or instruction about what to analyze in the image. Be specific about what you want to know.",
},
},
required: ["image_data", "prompt"],
};
private readonly visionClient: NimVisionClient;
constructor(
secrets: vscode.SecretStorage,
modelStorage?: vscode.Memento,
keyResolver?: NvidiaApiKeyResolver,
) {
this.visionClient = new NimVisionClient(secrets, modelStorage, keyResolver);
}
async invoke(
options: vscode.LanguageModelToolInvocationOptions<{ image_data: string; prompt: string }>,
token: vscode.CancellationToken,
): Promise<vscode.LanguageModelToolResult> {
const { image_data, prompt } = options.input;
const abortController = new AbortController();
const cancellationSubscription =
typeof token.onCancellationRequested === "function"
? token.onCancellationRequested(() => abortController.abort())
: undefined;
try {
Iif (token.isCancellationRequested) {
throw createAbortError();
}
const result = await this.visionClient.analyzeImage(
image_data,
prompt,
abortController.signal,
);
return new vscode.LanguageModelToolResult([new vscode.LanguageModelTextPart(result)]);
} catch (error) {
if (
token.isCancellationRequested ||
(error instanceof Error && error.name === "AbortError")
) {
throw createVisionCancellationError();
}
const errorMessage = error instanceof Error ? error.message : "Unknown error";
return new vscode.LanguageModelToolResult([
new vscode.LanguageModelTextPart(`Failed to analyze image: ${errorMessage}`),
]);
} finally {
cancellationSubscription?.dispose();
}
}
prepareInvocation?(
_options: vscode.LanguageModelToolInvocationPrepareOptions<{
image_data: string;
prompt: string;
}>,
_token: vscode.CancellationToken,
): vscode.ProviderResult<vscode.PreparedToolInvocation> {
return { invocationMessage: `Analyzing image with ${PROVIDER_DISPLAY_NAME} Vision...` };
}
}
/**
* Register all NVIDIA NIM tools with the Language Model API.
* @param secrets VS Code secret storage for API key access
* @returns Disposable for the tool registrations
*/
export function registerNimTools(
secrets: vscode.SecretStorage,
modelStorage?: vscode.Memento,
keyResolver?: NvidiaApiKeyResolver,
): vscode.Disposable {
const analyzeImageTool = new NimAnalyzeImageTool(secrets, modelStorage, keyResolver);
return vscode.Disposable.from(vscode.lm.registerTool(NimAnalyzeImageTool.id, analyzeImageTool));
}
function createAbortError(): Error {
const error = new Error("The operation was aborted");
error.name = "AbortError";
return error;
}
function createVisionCancellationError(): Error {
const CancellationError = (vscode as typeof vscode & { CancellationError?: new () => Error })
.CancellationError;
return CancellationError ? new CancellationError() : createAbortError();
}
|