analyze-images-advanced

import { tool, Tool, ToolsProviderController } from "@lmstudio/sdk";
import { readdir, stat, readFile } from "fs/promises";
import { basename, dirname, extname, isAbsolute, join, normalize, relative } from "path";
import { z } from "zod";

const IMAGE_EXTENSIONS = new Set([
	".jpg",
	".jpeg",
	".png",
	".webp",
	".gif",
	".bmp",
	".tiff",
	".tif",
	".avif",
]);
const MAX_ANALYSIS_TOKENS_HARD_CAP = 2048;
const ANALYSIS_TOKENS_SOFT_TARGET = 512;

type FoundImage = {
	relativePath: string;
	absolutePath: string;
	sizeBytes: number;
	modifiedAt: string;
	metadataPath?: string;
	caption?: string;
	sourcePageUrl?: string;
	downloadedAt?: string;
};

export async function toolsProvider(ctl: ToolsProviderController): Promise<Tool[]> {
	const listLocalImagesTool = tool({
		name: "List Local Images Advanced",
		description: "Lists local image files in the working directory, including sidecar metadata such as captions and source URLs when available. Use this before choosing images to analyze; do not infer image identity from local file numbering.",
		parameters: {
			recursive: z.boolean().optional().default(false).describe("Whether to search subdirectories."),
			maxResults: z.number().int().min(1).max(200).optional().default(50).describe("Maximum number of image files to return."),
		},
		implementation: async ({ recursive, maxResults }, { status, warn }) => {
			status("Listing local images...");

			const workingDirectory = ctl.getWorkingDirectory();
			const images = await collectImages(
				workingDirectory,
				recursive ?? false,
				maxResults ?? 50,
				warn,
			);

			return {
				count: images.length,
				images: images.map((image) => ({
					name: image.relativePath,
					sizeBytes: image.sizeBytes,
					modifiedAt: image.modifiedAt,
					...(image.metadataPath ? { metadataPath: image.metadataPath } : {}),
					...(image.caption ? { caption: image.caption } : {}),
					...(image.sourcePageUrl ? { sourcePageUrl: image.sourcePageUrl } : {}),
					...(image.downloadedAt ? { downloadedAt: image.downloadedAt } : {}),
				})),
				hint: images.length
					? "If the user is asking for a described image, such as a caption or subject, use Search Local Images Advanced before Analyze Local Image Advanced. Use exact image names from the search result. Do not infer image identity from local file numbering or JSON-LD ordering."
					: "No local images were found in the working directory.",
			};
		},
	});
	const searchLocalImagesTool = tool({
		name: "Search Local Images Advanced",
		description:
			"Searches local images by sidecar metadata such as caption, alt text, source URL, original image URL, and file name. Use this tool before Analyze Local Image Advanced whenever the user refers to an image by description, caption, topic, source, or page context instead of an exact local file name. This prevents confusing local file numbering with JSON-LD image ordering.",
		parameters: {
			query: z
				.string()
				.describe("Search text such as a caption, subject, source page, or phrase to find matching local images."),
			recursive: z
				.boolean()
				.optional()
				.default(false)
				.describe("Whether to search subdirectories."),
			maxResults: z
				.number()
				.int()
				.min(1)
				.max(50)
				.optional()
				.default(10)
				.describe("Maximum number of matching images to return."),
		},
		implementation: async ({ query, recursive, maxResults }, { status, warn }) => {
			status("Searching local images...");

			const workingDirectory = ctl.getWorkingDirectory();
			const images = await collectImages(
				workingDirectory,
				recursive ?? false,
				200,
				warn,
			);

			const terms = query
				.toLowerCase()
				.split(/\s+/)
				.map((term) => term.trim())
				.filter(Boolean);

			if (terms.length === 0) {
				return {
					count: 0,
					images: [],
					hint: "Search query was empty.",
				};
			}

			const scored = images
				.map((image) => {
					const haystack = [
						image.relativePath,
						image.metadataPath,
						image.caption,
						image.sourcePageUrl,
						image.downloadedAt,
					]
						.filter(Boolean)
						.join("\n")
						.toLowerCase();

					const score = terms.reduce((total, term) => {
						if (!haystack.includes(term)) return total;

						let bonus = 1;
						if (image.caption?.toLowerCase().includes(term)) bonus += 5;
						if (image.sourcePageUrl?.toLowerCase().includes(term)) bonus += 2;
						if (image.relativePath.toLowerCase().includes(term)) bonus += 1;

						return total + bonus;
					}, 0);

					return { image, score };
				})
				.filter((item) => item.score > 0)
				.sort((a, b) => b.score - a.score)
				.slice(0, maxResults ?? 10);

			return {
				count: scored.length,
				images: scored.map(({ image, score }) => ({
					name: image.relativePath,
					score,
					sizeBytes: image.sizeBytes,
					modifiedAt: image.modifiedAt,
					...(image.metadataPath ? { metadataPath: image.metadataPath } : {}),
					...(image.caption ? { caption: image.caption } : {}),
					...(image.sourcePageUrl ? { sourcePageUrl: image.sourcePageUrl } : {}),
					...(image.downloadedAt ? { downloadedAt: image.downloadedAt } : {}),
				})),
				hint: scored.length
					? "Use these exact image names with Analyze Local Image Advanced. Do not infer image identity from local file numbering or JSON-LD ordering."
					: "No matching local images were found. Try List Local Images Advanced to inspect available files.",
			};
		},
	});
	const analyzeLocalImageTool = tool({
		name: "Analyze Local Image Advanced",
		description:
			"Analyzes one exact local image from the working directory using the currently loaded vision model. If the user describes an image by caption, subject, source page, or text such as 'the snack mama image', first use Search Local Images Advanced to find the correct image name. Do not guess image names from local numbering or JSON-LD ordering.",
		parameters: {
			imageName: z
				.string()
				.describe("Exact image file name or relative path. If you do not already know the exact local file name, use Search Local Images Advanced first. Do not guess from numbering."),
			prompt: z
				.string()
				.describe("Required analysis task/question. Keep it clear and specific."),
			context: z
				.string()
				.optional()
				.default("")
				.describe("Optional known context for this image and task, such as source, intent, constraints, known facts, or prior findings."),
		},
		implementation: async ({ imageName, prompt, context }, { status, warn }) => {
			const workingDirectory = ctl.getWorkingDirectory();
			const safeImageName = sanitizeRelativeInput(imageName);
			if (!safeImageName) {
				return "Error: imageName is empty or invalid.";
			}

			const resolvedImagePath = await resolveImagePathByName(workingDirectory, safeImageName, warn);
			if (!resolvedImagePath) {
				return `Error: image not found: ${safeImageName}`;
			}

			status("Preparing image for multimodal model...");
			const model = await ctl.client.llm.model();
			if (!model.vision) {
				return "Error: currently loaded model does not support vision. Load a vision model and retry.";
			}

			const fileHandle = await ctl.client.files.prepareImage(resolvedImagePath.absolutePath);
			const sidecarMetadata = await readImageSidecarMetadata(
				resolvedImagePath.absolutePath,
				resolvedImagePath.relativePath,
			);

			if (sidecarMetadata?.localMetadata) {
				status(`Loaded image metadata: ${sidecarMetadata.localMetadata}`);
			}

			const userPrompt = prompt.trim();
			const userContext = (context || "").trim();

			const metadataContext = sidecarMetadata
				? buildImageMetadataContext(sidecarMetadata)
				: "";

			const contextParts = [
				metadataContext,
				userContext ? `User-provided context:\n${userContext}` : "",
			].filter(Boolean);

			const contextBlock = contextParts.length
				? `Known context:\n${contextParts.join("\n\n")}\n\n`
				: "";

			const analysisPrompt =
				`You are a vision assistant. Analyze the provided image and give a concise final answer. ` +
				`Do not provide hidden reasoning or step-by-step chain-of-thought. ` +
				`If uncertain, state uncertainty briefly. ` +
				`Target up to ${ANALYSIS_TOKENS_SOFT_TARGET} tokens in the final answer.\n\n` +
				contextBlock +
				`User request:\n${userPrompt}`;
			const effectiveMaxTokens = MAX_ANALYSIS_TOKENS_HARD_CAP;

			status("Running multimodal analysis...");
			const result = await model.respond(
				[
					{
						role: "user",
						content: analysisPrompt,
						images: [fileHandle],
					},
				],
				{
					maxTokens: effectiveMaxTokens,
				},
			);

			return result.content;
		},
	});

	return [listLocalImagesTool, searchLocalImagesTool, analyzeLocalImageTool];
}

async function collectImages(
	directoryPath: string,
	recursive: boolean,
	maxResults: number,
	warn: (text: string) => void,
): Promise<FoundImage[]> {
	const found: FoundImage[] = [];
	const queue: string[] = [directoryPath];
	const root = directoryPath;

	while (queue.length > 0 && found.length < maxResults) {
		const current = queue.shift() as string;
		let entries: Array<{ name: string; isFile: () => boolean; isDirectory: () => boolean }>;
		try {
			entries = await readdir(current, { withFileTypes: true });
		} catch (error: any) {
			warn(`Cannot read directory '${current}': ${error?.message || String(error)}`);
			continue;
		}

		for (const entry of entries) {
			if (found.length >= maxResults) break;
			const absolutePath = join(current, entry.name);

			if (entry.isDirectory()) {
				if (recursive) {
					queue.push(absolutePath);
				}
				continue;
			}

			if (!entry.isFile()) continue;
			if (!isImagePath(entry.name)) continue;

			try {
				const metadata = await stat(absolutePath);
				const relativePath = normalize(relative(root, absolutePath)).replace(/\\/g, "/");
				const sidecarMetadata = await readImageSidecarMetadata(absolutePath, relativePath);

				found.push({
					absolutePath,
					relativePath,
					sizeBytes: metadata.size,
					modifiedAt: metadata.mtime.toISOString(),
					...(sidecarMetadata?.localMetadata ? { metadataPath: sidecarMetadata.localMetadata } : {}),
					...(sidecarMetadata?.caption ? { caption: sidecarMetadata.caption } : {}),
					...(sidecarMetadata?.alt && !sidecarMetadata?.caption ? { caption: sidecarMetadata.alt } : {}),
					...(sidecarMetadata?.sourcePageUrl ? { sourcePageUrl: sidecarMetadata.sourcePageUrl } : {}),
					...(sidecarMetadata?.downloadedAt ? { downloadedAt: sidecarMetadata.downloadedAt } : {}),
				});
			} catch (error: any) {
				warn(`Cannot stat file '${absolutePath}': ${error?.message || String(error)}`);
			}
		}
	}

	return found;
}

function isImagePath(value: string): boolean {
	return IMAGE_EXTENSIONS.has(extname(value).toLowerCase());
}

function sanitizeRelativeInput(input?: string): string | null {
	if (!input) return null;
	const trimmed = input.trim();
	if (!trimmed) return null;
	if (isAbsolute(trimmed)) return null;
	const normalized = normalize(trimmed).replace(/\\/g, "/").replace(/^\.\/+/, "");
	if (!normalized || normalized.startsWith("../") || normalized.includes("/../")) return null;
	return normalized;
}

async function resolveImagePathByName(
	workingDirectory: string,
	imageName: string,
	warn: (text: string) => void,
): Promise<{ absolutePath: string; relativePath: string } | null> {
	// 1) Try working-directory root first: <workingDirectory>/<imageName>
	const directAbsolutePath = join(workingDirectory, imageName);
	const directStats = await stat(directAbsolutePath).catch(() => null);
	if (directStats?.isFile() && isImagePath(imageName)) {
		return await preferFullImageIfThumb(directAbsolutePath, imageName);
	}

	// 2) Fallback: recursive basename match across all images
	const allImages = await collectImages(workingDirectory, true, 1000, warn);
	const targetBasename = basename(imageName).toLowerCase();
	const matched = allImages.find((item) => basename(item.relativePath).toLowerCase() === targetBasename);
	if (!matched) return null;

	return await preferFullImageIfThumb(matched.absolutePath, matched.relativePath);
}

async function preferFullImageIfThumb(
	absolutePath: string,
	relativePath: string,
): Promise<{ absolutePath: string; relativePath: string }> {
	const fileName = basename(relativePath);
	const thumbMatch = fileName.match(/^(.*)-thumb\.webp$/i);
	if (!thumbMatch) {
		return { absolutePath, relativePath };
	}

	const baseNameWithoutThumb = thumbMatch[1];
	const parentDirAbsolute = dirname(absolutePath);
	const parentDirRelative = dirname(relativePath).replace(/\\/g, "/");
	const candidateExtensions = [".png", ".jpg", ".jpeg", ".webp", ".gif", ".bmp", ".tiff", ".tif", ".avif"];

	for (const extension of candidateExtensions) {
		const candidateFileName = `${baseNameWithoutThumb}${extension}`;
		const candidateAbsolutePath = join(parentDirAbsolute, candidateFileName);
		const candidateStats = await stat(candidateAbsolutePath).catch(() => null);
		if (!candidateStats?.isFile()) continue;
		if (!isImagePath(candidateFileName)) continue;
		const candidateRelativePath =
			parentDirRelative === "." ? candidateFileName : `${parentDirRelative}/${candidateFileName}`;
		return {
			absolutePath: candidateAbsolutePath,
			relativePath: candidateRelativePath,
		};
	}

	return { absolutePath, relativePath };
}

async function readImageSidecarMetadata(
	absoluteImagePath: string,
	relativeImagePath: string,
): Promise<Record<string, any> | null> {
	const candidates = getSidecarMetadataCandidates(absoluteImagePath, relativeImagePath);

	for (const candidate of candidates) {
		try {
			const raw = await readFile(candidate.absolutePath, "utf8");
			const parsed = JSON.parse(raw);
			if (!parsed || typeof parsed !== "object") continue;
			return parsed;
		} catch {
			// Try next candidate
		}
	}

	return null;
}

function getSidecarMetadataCandidates(
	absoluteImagePath: string,
	relativeImagePath: string,
): Array<{ absolutePath: string; relativePath: string }> {
	const absoluteDir = dirname(absoluteImagePath);
	const relativeDir = dirname(relativeImagePath).replace(/\\/g, "/");

	const imageFileName = basename(relativeImagePath);
	const withoutThumb = imageFileName.replace(/-thumb\.webp$/i, "");
	const withoutExtension = withoutThumb.replace(/\.[^.]+$/i, "");

	const metadataFileName = `${withoutExtension}.json`;

	const relativeMetadataPath =
		relativeDir === "." ? metadataFileName : `${relativeDir}/${metadataFileName}`;

	return [
		{
			absolutePath: join(absoluteDir, metadataFileName),
			relativePath: relativeMetadataPath,
		},
	];
}

function buildImageMetadataContext(metadata: Record<string, any>): string {
	const lines: string[] = [];

	lines.push("Image metadata from sidecar JSON:");

	if (typeof metadata.caption === "string" && metadata.caption.trim()) {
		lines.push(`- caption: ${metadata.caption.trim()}`);
	} else if (typeof metadata.alt === "string" && metadata.alt.trim()) {
		lines.push(`- alt: ${metadata.alt.trim()}`);
	}

	if (typeof metadata.sourcePageUrl === "string" && metadata.sourcePageUrl.trim()) {
		lines.push(`- source page: ${metadata.sourcePageUrl.trim()}`);
	}

	if (typeof metadata.url === "string" && metadata.url.trim()) {
		lines.push(`- original image URL: ${metadata.url.trim()}`);
	}

	if (typeof metadata.imageModifiedAt === "string" && metadata.imageModifiedAt.trim()) {
		lines.push(`- image modified at: ${metadata.imageModifiedAt.trim()}`);
	}

	if (typeof metadata.imagePublishedAt === "string" && metadata.imagePublishedAt.trim()) {
		lines.push(`- image published at: ${metadata.imagePublishedAt.trim()}`);
	}

	if (typeof metadata.downloadedAt === "string" && metadata.downloadedAt.trim()) {
		lines.push(`- downloaded at: ${metadata.downloadedAt.trim()}`);
	}

	if (typeof metadata.source === "string" && metadata.source.trim()) {
		lines.push(`- metadata source: ${metadata.source.trim()}`);
	}

	lines.push(
		"Use this metadata only as contextual information. Do not treat alt text, captions, or timestamps as visual proof unless the image itself supports it."
	);

	return lines.join("\n");
}