find-image

/**
 * find_image Tool
 * Searches Draw Things generation history and returns structured results
 */

import { tool, type Tool, type ToolsProviderController, type ToolCallContext } from "@lmstudio/sdk";
// @ts-ignore — zod/lib re-export chain breaks with NodeNext moduleResolution; runtime is fine
import { z } from "zod";
import path from "node:path";
import fs from "node:fs";
import crypto from "node:crypto";
import { getThumbnailFromProject, syncAttachmentsToState } from "../core-bundle.mjs";
import { queryEmbeddingPayloadPath, resolveMediaQueries, type ResolvedMediaQuery } from "../media/mediaResolver.js";
import { previewPayloadBytes } from "../media/previewPayload.js";
import { parseStructuredQuery } from "../search/queryParser.js";
import { searchGenerations } from "../search/searchEngine.js";
import { configSchematics, globalConfigSchematics } from "../config.js";
import { MultimodalEmbeddingClient, MultimodalEmbeddingStore, multimodalModelIdentity, resolveMultimodalEmbeddingBackend } from "../embeddings/index.js";
import type { MultimodalEmbeddingBackendKind } from "../embeddings/index.js";
import {
  EMBEDDING_INSTRUCTIONS,
  RERANKING_INSTRUCTIONS,
  selectQueryEmbeddingInstruction,
  selectRerankingInstruction,
} from "../embeddings/embeddingInstructions.js";
import { renderEmbeddingMetadataText } from "../embeddings/embeddingMetadata.js";
import { rerankWithLlamaServer } from "../reranking/llamaServerReranker.js";
import { releaseOwnedLlamaServerEmbeddingServer } from "../llama-server-manager.js";
import { GGUF_DB_PATH, LEGACY_GGUF_DB_PATH } from "../embeddings/multimodalEmbeddingClientGguf.js";
import { migrateLegacyFileIfMissing, sanitizeConfigPath } from "../paths.js";
import { resolvePluginStringSetting } from "../pluginSettings.js";
import { PngMetadataParser, type ParsedImageMetadata } from "../documents/parsers/pngMetadataParser.js";
import { parseProjectUri } from "../documents/parsers/projectFileParser.js";
import { localTimestamp } from "../helpers/localTimestamp.js";
import { formatToolMetaBlock } from "../helpers/pluginMeta.js";
import { SUPPORTED_IMAGE_FORMAT_SET } from "../documents/imageFormats.js";
import {
  parseDtcQuery,
  maybeFormatModelForConsumerFromSnapshot,
  maybeGetModelRewriteHintsFromSnapshot,
  type DtcModelMappingSnapshotV1,
} from "../helpers/dtcModelMappingSnapshot.js";
import type { DrawThingsSearchResult, IndexedGeneration } from "../types.js";

const EMBEDDING_RERANKING_DIAGNOSTIC_LOG_FILE = "embedding-reranking-diagnostics.jsonl";

type RerankingDiagnosticMatch = {
  recallRank?: number;
  rerankRequestOrder?: number;
  finalRank?: number;
  imagePath?: string;
  imagePaths: string[];
  prompt: string;
  model: string;
  matchType: string;
  embeddingScore?: number;
  embeddingSimilarity?: number;
  rerankingScore?: number;
  displayedScore: number;
};

type RerankingDiagnosticPositions = Pick<
  RerankingDiagnosticMatch,
  "recallRank" | "rerankRequestOrder" | "finalRank"
>;

function diagnosticMatch(match: {
  imagePaths: string[];
  prompt: string;
  model: string;
  matchType: string;
  multimodalScore?: number;
  rankingSimilarity?: number;
  rerankingScore?: number;
  matchScore: number;
}, positions: RerankingDiagnosticPositions, imagePath?: string): RerankingDiagnosticMatch {
  return {
    ...positions,
    imagePath,
    imagePaths: match.imagePaths,
    prompt: match.prompt,
    model: match.model,
    matchType: match.matchType,
    embeddingScore: match.multimodalScore,
    embeddingSimilarity: match.rankingSimilarity,
    rerankingScore: match.rerankingScore,
    displayedScore: match.matchScore,
  };
}

function appendEmbeddingRerankingDiagnostic(record: Record<string, unknown>): void {
  try {
    const logsDir = path.join(process.cwd(), "logs");
    fs.mkdirSync(logsDir, { recursive: true });
    fs.appendFileSync(
      path.join(logsDir, EMBEDDING_RERANKING_DIAGNOSTIC_LOG_FILE),
      `${JSON.stringify({ timestamp: localTimestamp(), ...record }, null, 2)}\n\n`,
      "utf8",
    );
  } catch (error) {
    console.warn("[find_image] Failed to write Embedding reranking diagnostic:", error);
  }
}

async function rerankingImageForMatch(
  match: { multimodalImageRef?: string },
  usePreviews: boolean,
): Promise<{ imagePath?: string; imageBytes?: Uint8Array } | undefined> {
  const imageRef = match.multimodalImageRef;
  if (!imageRef) return undefined;
  if (!imageRef.startsWith("project://")) {
    if (!fs.existsSync(imageRef)) return undefined;
    const originalBytes = new Uint8Array(await fs.promises.readFile(imageRef));
    return { imagePath: imageRef, imageBytes: await previewPayloadBytes(originalBytes, usePreviews) };
  }

  const projectImage = parseProjectUri(imageRef);
  if (!projectImage) return undefined;

  const originalBytes = await getThumbnailFromProject(projectImage.projectPath, projectImage.thumbnailId);
  if (!originalBytes) return undefined;
  return { imageBytes: await previewPayloadBytes(originalBytes, usePreviews) };
}

function looksLikeImageFilenameQuery(q: string): boolean {
  const trimmed = q.trim().replace(/^['"]/,'').replace(/['"']$/,'');
  const normalized = trimmed.replace(/\\/g, "/");
  const base = normalized.split("/").pop() ?? normalized;
  const dot = base.lastIndexOf(".");
  if (dot <= 0 || dot === base.length - 1) return false;
  const ext = base.slice(dot + 1).toLowerCase();
  return SUPPORTED_IMAGE_FORMAT_SET.has(ext);
}

function looksLikeProjectFilenameQuery(q: string): boolean {
  const trimmed = q.trim().replace(/^['"]/,'').replace(/['"']$/,'');
  const normalized = trimmed.replace(/\\/g, "/");
  const base = normalized.split("/").pop() ?? normalized;
  return base.toLowerCase().endsWith('.sqlite3');
}

function unquoteQueryPath(q: string): string {
  return q.trim().replace(/^['"]/, '').replace(/['"]$/, '');
}

function metadataQueryFromImage(metadata: ParsedImageMetadata, fallbackName: string): string {
  const parts: string[] = [];
  const prompt = metadata.prompt?.trim();
  if (prompt) {
    parts.push(prompt);
  } else if (fallbackName.trim()) {
    parts.push(fallbackName.trim());
  }

  if (metadata.model?.trim()) {
    parts.push(`Model: ${metadata.model.trim()}`);
  }

  if (metadata.loras?.length) {
    const loras = metadata.loras.map((lora) => lora.model).filter(Boolean);
    if (loras.length) {
      parts.push(`LoRAs: ${loras.join(', ')}`);
    }
  }

  if (metadata.width > 0 && metadata.height > 0) {
    parts.push(`Size: ${metadata.width}x${metadata.height}`);
  }

  if (metadata.rawType === 'comfyui') {
    parts.push('Origin: ComfyUI');
  } else if (metadata.rawType === 'drawthings') {
    parts.push('Origin: Draw Things');
  }

  // This metadata text is used only for embedding and reranking. Deterministic
  // search continues to receive the user's original query.
  return parts.join('\n').trim() || fallbackName.trim();
}

async function queryFromReferencedImage(filePath: string, fallbackName: string): Promise<string> {
  try {
    const metadata = await PngMetadataParser.extractMetadata(filePath);
    if (metadata) {
      return metadataQueryFromImage(metadata, fallbackName);
    }
  } catch (error) {
    console.warn('[find_image] Query image metadata extraction failed:', error);
  }
  return fallbackName.trim();
}

function metadataQueryFromGeneration(gen: IndexedGeneration, fallbackName: string): string {
  return renderEmbeddingMetadataText(gen, fallbackName);
}

function findGenerationForImageRef(generations: IndexedGeneration[], imageRef: string): IndexedGeneration | undefined {
  const normalized = imageRef.trim();
  if (!normalized) return undefined;
  return generations.find((gen) => (gen.imagePaths ?? []).includes(normalized));
}

async function metadataQueryForTargetImage(
  filePath: string,
  fallbackName: string,
  generations: IndexedGeneration[],
  useIndexedGeneration: boolean,
): Promise<{ query: string; targetGeneration?: IndexedGeneration }> {
  if (useIndexedGeneration) {
    const indexed = findGenerationForImageRef(generations, filePath);
    if (indexed) return { query: metadataQueryFromGeneration(indexed, fallbackName), targetGeneration: indexed };
  }
  return { query: await queryFromReferencedImage(filePath, fallbackName) };
}

function appendQueryPart(base: string, addition: string): string {
  const left = base.trim();
  const right = addition.trim();
  if (!left) return right;
  if (!right) return left;
  return `${left}\n${right}`;
}

export function rankingTextForParsedQuery(rawQuery: string, parsed: { hasFilters: boolean; promptQuery: string }): string {
  return parsed.hasFilters ? parsed.promptQuery.trim() : rawQuery.trim();
}

let multimodalEmbeddingStore: MultimodalEmbeddingStore | null = null;
let multimodalEmbeddingStoreDbPath: string | null = null;

// Tool parameters schema
const SearchParamsSchema = {
  target: z.string().optional().describe("Reference image: aN, vN, iN, pN, or an absolute path. Set it for 'like this image', 'same style as a1', or changes to a shown image."),
  query: z.string().optional().describe("Search instruction. With target, positively describe additions or changes; without target, write a complete description. Model:, LoRAs:, Tag:, Tags:, Size:, Source:, Origin:, and Timestamp: are hard AND filters."),
  includeMetadata: z.boolean().optional().describe("Add target generation metadata as a ranking signal across the full corpus, not an identical-metadata filter."),
  excludeImage: z.boolean().optional().describe("Ignore target pixels. Use only for prompt/metadata similarity, not 'same style' or other image-guided changes."),
};

interface FindImageArgs {
  target?: string;
  query?: string;
  includeMetadata?: boolean;
  excludeImage?: boolean;
}

export function createSearchGenerationsTool(ctl: ToolsProviderController): Tool {
  return tool({
    name: "find_image",
    description: `Find visually similar and metadata-related images.

Use this tool for cross-modal retrieval over indexed images and Draw Things projects:
- target: image references (aN, vN, iN, pN) or an absolute image file path
- query: a positive search instruction, structured metadata filters, or a filename
- includeMetadata: add target generation metadata as a similarity signal
- excludeImage: ignore target image pixels; use query text and optional target metadata only
- query: File/project names: indexed image filenames or Draw Things .sqlite3 project filenames

Image-reference rules:
- If the user refers to an image, set target; do not replace it with a text-only reconstruction.
- With target, query positively describes changes: "people in front of a brick wall", not "no cats". Do not use excludeImage for image-guided changes.
- Target and includeMetadata rank the complete corpus. Structured query metadata, such as Model: or Size:, are hard AND filters.
- Tag: and Tags: apply hard AND filters to user-managed tags.

Target searches use visual similarity by default; when a query is given alongside target, the two are fused into one combined image+text cross-modal search. Project filename queries return generations from that project in chronological order, limited by the configured retrievalLimit like any other search.

Returns:
- Content matches: images retrieved from visual, text, or combined image-and-text similarity
- Exact matches: filename, project filename, and structured metadata filter results
- Image paths for each matching generation

Example queries:
- { "target": "p3" } - finds images visually similar to picture p3
- { "target": "a1", "includeMetadata": true } - searches from attachment a1 and its generation metadata
- { "target": "a1", "query": "night city", "excludeImage": true } - uses the text without target pixels
- { "target": "a1", "query": "Same chibi cartoon style as the reference image, but with people instead of cats: thick outlines, flat colors, big eyes, cute characters." } - keeps a1 as the visual anchor and describes the desired change
- { "query": "A standing person in front of a brick wall, chibi cartoon style, thick outlines, flat colors, large expressive eyes." } - text-only search, used when no reference image should guide results
- { "query": "portrait Model: z_image_1.0_q8p.ckpt" } - searches portraits only among images made with that model
- { "query": "portrait in studio lighting Tags: favorite, reviewed" } - searches studio portraits tagged favorite and reviewed
- { "query": "my-project.sqlite3" } - returns generations from that Draw Things project in chronological order (subject to retrievalLimit)

${formatToolMetaBlock()}`,
    parameters: SearchParamsSchema,
  implementation: async (args: FindImageArgs, ctx: ToolCallContext) => {
      let releaseEmbeddingServer: (() => Promise<void>) | undefined;
      try {
        // Sync attachments from conversation.json → chat_media_state.json (non-fatal).
        // Must run before query rewrite so tryBuildMediaIndex can resolve a1 → originalName.
        try {
          const workingDir = ctl.getWorkingDirectory();
          if (typeof workingDir === "string" && workingDir.trim().length > 0) {
            await syncAttachmentsToState(workingDir, false, Number.MAX_SAFE_INTEGER);
          }
        } catch (syncErr: any) {
          console.warn(
            "[find_image] attachment sync failed (non-fatal):",
            syncErr?.message ?? syncErr,
          );
        }

        ctx.status("Starting search...");

        const targetInput = String(args.target ?? "").trim();
        const queryInput = String(args.query ?? "").trim();
        if (!targetInput && !queryInput) {
          throw new Error("find_image requires target, query, or both.");
        }
        if (args.includeMetadata === true && !targetInput) {
          throw new Error("includeMetadata requires a target image.");
        }
        if (args.excludeImage === true && targetInput && !queryInput && args.includeMetadata !== true) {
          throw new Error("excludeImage removes the target image; provide query or set includeMetadata to keep a search signal.");
        }

        const { query: parsedQuery, enableModelRewrite, snapshot } = parseDtcQuery(queryInput);

        if (enableModelRewrite) {
          ctx.status("Model rewrite enabled; snapshot received.");
        }

        let query = parsedQuery;
        let targetMedia: ResolvedMediaQuery[] = [];
        let directReferencedImagePath: string | undefined;
        if (targetInput) {
          if (/^\s*[avip]\d+\s*$/i.test(targetInput)) {
            const workingDir = ctl.getWorkingDirectory();
            if (typeof workingDir === "string" && workingDir.trim()) {
              const resolved = await resolveMediaQueries(workingDir, targetInput);
              targetMedia = resolved.mediaQueries;
            }
          } else {
            const candidatePath = unquoteQueryPath(targetInput);
            if (path.isAbsolute(candidatePath) && fs.existsSync(candidatePath) && looksLikeImageFilenameQuery(candidatePath)) {
              directReferencedImagePath = candidatePath;
            } else {
              throw new Error(`Invalid target: ${targetInput}. Use aN/vN/iN/pN or an absolute image file path.`);
            }
          }
        }
        
        const config = ctl.getGlobalPluginConfig(globalConfigSchematics);
        const multimodalModelPath = sanitizeConfigPath(config.get("multimodalEmbeddingModelPath"));
        const multimodalBackend = config.get("multimodalEmbeddingBackend") as MultimodalEmbeddingBackendKind;
        const resolvedMultimodalBackend = resolveMultimodalEmbeddingBackend(multimodalBackend, multimodalModelPath);
        const multimodalEmbeddingServicePort = config.get("multimodalEmbeddingServicePort");
        const llamaServerTTL = config.get("llamaServerTTL");
        let embeddingServerReleased = false;
        releaseEmbeddingServer = async (): Promise<void> => {
          if (embeddingServerReleased || resolvedMultimodalBackend !== "integrated-local-gguf") return;
          embeddingServerReleased = true;
          await releaseOwnedLlamaServerEmbeddingServer(multimodalEmbeddingServicePort, llamaServerTTL);
        };

        // Get indexed generations with progress feedback
        const { indexGenerations } = await import("../indexer.js");
        const generations = await indexGenerations(
          ctl, 
          false, // use cache if available
          (msg: string) => ctx.status(msg)
        );
        
        ctx.status(`Searching ${generations.length} generations...`);
        
        if (generations.length === 0) {
          await releaseEmbeddingServer();
          return JSON.stringify({
            type: "find-image-results",
            query: queryInput,
            target: targetInput || undefined,
            totalFound: 0,
            error: "No images indexed. Check image directories, LM Studio user-files, working-directories, or Draw Things project settings.",
            searchTimeMs: 0,
            images: [],
          }, null, 2);
        }

        const firstTargetImage = targetMedia.find((media) => media.filePath);
        const targetImagePath = firstTargetImage?.filePath ?? directReferencedImagePath;
        const targetFallbackName = firstTargetImage?.displayName
          ?? firstTargetImage?.originalName
          ?? firstTargetImage?.notation
          ?? (targetImagePath ? path.basename(targetImagePath) : targetInput);
        if (targetInput && !targetImagePath) {
          throw new Error(`Target did not resolve to an image file: ${targetInput}`);
        }
        const usePreviewsForQueryEmbedding = config.get("usePreviewsForQueryEmbedding");
        const queryEmbeddingImagePath = firstTargetImage
          ? await queryEmbeddingPayloadPath(firstTargetImage, usePreviewsForQueryEmbedding)
          : targetImagePath;
        const projectTarget = firstTargetImage?.sourceReference
          ? parseProjectUri(firstTargetImage.sourceReference)
          : null;
        const projectTargetBytes = projectTarget
          ? await getThumbnailFromProject(projectTarget.projectPath, projectTarget.thumbnailId)
          : null;
        if (projectTarget && !projectTargetBytes) {
          throw new Error(`Target project thumbnail could not be read: ${firstTargetImage?.sourceReference}`);
        }
        const rerankingEnabled = config.get("rerankingEnabled");
        let embeddingTargetMetadataText = "";
        let targetGenerationMetadata: IndexedGeneration | undefined;
        if (args.includeMetadata && targetImagePath) {
          const { query: metadataQuery, targetGeneration } = await metadataQueryForTargetImage(
            firstTargetImage?.sourceReference ?? targetImagePath,
            targetFallbackName,
            generations,
            true,
          );
          targetGenerationMetadata = targetGeneration;
          embeddingTargetMetadataText = metadataQuery;
          console.info(`[find_image] Prepared target metadata for Embedding: ${metadataQuery.replace(/\n/g, ' • ')}`);
        }
        
        const chatConfig = ctl.getPluginConfig(configSchematics);
        
        const isFilenameQuery = looksLikeImageFilenameQuery(query);
        const isProjectFilenameQuery = looksLikeProjectFilenameQuery(query);
        const multimodalSearchEnabled = config.get("multimodalSearchEnabled");
        const multimodalDimension = config.get("multimodalEmbeddingDimension");
        const multimodalContextSize = config.get("multimodalEmbeddingContextSize");
        const multimodalModel = config.get("multimodalEmbeddingModel");
        const multimodalGgufBinaryPath = sanitizeConfigPath(resolvePluginStringSetting(
          config.get("multimodalEmbeddingGgufBinaryPath"),
          "multimodalEmbeddingGgufBinaryPath",
        ));
        const multimodalGgufGpuLayers = config.get("multimodalEmbeddingGgufGpuLayers");
        const multimodalModelKey = multimodalModelIdentity(multimodalModelPath, multimodalModel);
        const multimodalEmbeddingDataStorePath = sanitizeConfigPath(config.get("multimodalEmbeddingDataStorePath"));
        const embeddingMinScore = config.get("embeddingMinScore");
        const embeddingDisplayMultiplier = config.get("embeddingDisplayMultiplier");
        const rerankingModelPath = sanitizeConfigPath(config.get("rerankingModelPath"));
        const rerankingServicePort = config.get("rerankingServicePort");
        const rerankingContextSize = config.get("rerankingContextSize");
        const rerankingBatchSize = config.get("rerankingBatchSize");
        const rerankingUbatchSize = config.get("rerankingUbatchSize");
        const rerankingDisplayMultiplier = config.get("rerankingDisplayMultiplier");
        const rerankingMinScore = config.get("rerankingMinScore");
        const rerankingCandidateMultiplier = config.get("rerankingCandidateMultiplier");
        const rerankingMaxCandidates = config.get("rerankingMaxCandidates");
        
        console.info(`[Search Config] filenameQuery=${isFilenameQuery}, multimodal=${multimodalSearchEnabled}, model=${multimodalModelKey}`);
        
        // Every result path, including filename/project fast paths, is capped
        // by this per-chat setting. The reranker recall expansion applies only
        // to non-fast-path content candidates.
        const allowByteIdenticalResults = chatConfig.get("allowByteIdenticalResults");
        const limit = chatConfig.get("retrievalLimit");
        const structuredFilterTolerances = {
          gpsRadiusMeters: chatConfig.get("gpsToleranceMeters"),
          timestampToleranceMinutes: chatConfig.get("createdToleranceMinutes"),
        };
        const finalLimit = limit >= 25 ? undefined : Math.max(1, Math.floor(limit));
        console.info(`[Search Config] allowByteIdenticalResults=${allowByteIdenticalResults}, retrievalLimit=${limit}, finalLimit=${finalLimit ?? 'all'}`);
        
        const knownProjects = [...new Set(generations.flatMap((generation) =>
          generation.sourceInfo?.type === "draw_things_project" ? [generation.sourceInfo.projectFile] : [],
        ))];
        const parsedStructuredQuery = parseStructuredQuery(
          query,
          enableModelRewrite ? snapshot : undefined,
          knownProjects,
          structuredFilterTolerances,
        );
        const rankingText = rankingTextForParsedQuery(query, parsedStructuredQuery);
        const embeddingQueryText = appendQueryPart(rankingText, embeddingTargetMetadataText);
        const excludeImage = args.excludeImage === true;
        const queryEmbeddingInstruction = EMBEDDING_INSTRUCTIONS[selectQueryEmbeddingInstruction({
          hasTargetImage: !!queryEmbeddingImagePath,
          hasFreeText: Boolean(rankingText),
          includeMetadata: args.includeMetadata === true,
          excludeImage,
        })];
        const hasExplicitRerankingText = Boolean(rankingText);
        const hasTargetMetadataBeyondFallback = Boolean(
          embeddingTargetMetadataText.trim()
          && embeddingTargetMetadataText.trim() !== targetFallbackName.trim(),
        );
        const rerankingMetadataOnly = !hasExplicitRerankingText
          && !!args.includeMetadata
          && hasTargetMetadataBeyondFallback;
        const useReranking = rerankingEnabled
          && !!rerankingModelPath
          && fs.existsSync(rerankingModelPath)
          && fs.existsSync(multimodalGgufBinaryPath)
          && (hasExplicitRerankingText || rerankingMetadataOnly);
        const rerankingQuery = hasExplicitRerankingText
          ? { text: excludeImage ? embeddingQueryText : rankingText, imagePath: excludeImage ? undefined : queryEmbeddingImagePath }
          : { text: embeddingTargetMetadataText };
        const rerankingInstruction = RERANKING_INSTRUCTIONS[selectRerankingInstruction({
          hasTargetImage: !!rerankingQuery.imagePath,
          hasFreeText: hasExplicitRerankingText,
          includeMetadata: rerankingMetadataOnly || (excludeImage && args.includeMetadata === true),
          excludeImage,
        })];
        let multimodalQueryEmbedding: number[] | undefined;
        let multimodalQueryMode: "text" | "image" | "mixed" | undefined;
        // True only for a genuinely fused image+text query (embedImageAndText,
        // Use Case 4/5) — distinguishes it from a plain target-image-only
        // query (embedImagePath, Use Case 2/3), which also uses mode "mixed"
        // but must NOT be treated as hybrid on the merge/ranking side (see
        // SearchOptions.fusedImageTextQuery docs in searchEngine.ts).
        const directQueryPath = targetImagePath;
        // multimodalSearchEnabled is a single, consistent killswitch: when
        // false, ensureMultimodalEmbeddingsForGenerations() never indexes any
        // image (indexer.ts), so there is nothing for an image-target query
        // to search against either — no override for hasImageQuery here.
        const shouldRunMultimodalQuery = !!multimodalModelKey && !isProjectFilenameQuery && multimodalSearchEnabled;
        const targetContentHash = projectTargetBytes
          ? crypto.createHash("sha256").update(projectTargetBytes).digest("hex")
          : targetImagePath
            ? crypto.createHash("sha256").update(await fs.promises.readFile(targetImagePath)).digest("hex")
            : undefined;

        if (shouldRunMultimodalQuery) {
          try {
            const isGgufBackend = resolvedMultimodalBackend === "integrated-local-gguf";
            if (isGgufBackend) {
              migrateLegacyFileIfMissing(GGUF_DB_PATH, LEGACY_GGUF_DB_PATH);
            }
            const expectedDbPath = isGgufBackend ? multimodalEmbeddingDataStorePath : null; // null = MultimodalEmbeddingStore's own default (unused by any implemented backend)
            // The indexer owns a separate sql.js store and can write a newer
            // database between searches. Drop this read-side snapshot so a
            // model switch is visible without restarting the host process.
            multimodalEmbeddingStore?.close();
            multimodalEmbeddingStore = new MultimodalEmbeddingStore(isGgufBackend ? { dbPath: multimodalEmbeddingDataStorePath } : {});
            multimodalEmbeddingStoreDbPath = expectedDbPath;
            await multimodalEmbeddingStore.init();
            const multimodalClient = new MultimodalEmbeddingClient({
              backend: resolvedMultimodalBackend,
              port: multimodalEmbeddingServicePort,
              lmStudioUrl: config.get("multimodalEmbeddingLmStudioUrl"),
              modelPath: multimodalModelPath,
              ggufBinaryPath: multimodalGgufBinaryPath,
              ggufGpuLayers: multimodalGgufGpuLayers,
              model: multimodalModel,
              dimension: multimodalDimension,
              contextSize: multimodalContextSize,
              onStatus: (message) => ctx.status(message),
            });

            if (projectTargetBytes && !excludeImage && embeddingQueryText.trim()) {
              const targetImageStatus = firstTargetImage?.notation ? `Embedding target image ${firstTargetImage.notation}...` : "Embedding target image...";
              ctx.status(targetImageStatus);
              multimodalQueryEmbedding = (await multimodalClient.embedImageBytesAndText(
                new Uint8Array(projectTargetBytes),
                embeddingQueryText,
                multimodalDimension,
                targetImageStatus,
                queryEmbeddingInstruction,
              )).embedding;
              multimodalQueryMode = "mixed";
            } else if (projectTargetBytes && !excludeImage) {
              const targetImageStatus = firstTargetImage?.notation ? `Embedding target image ${firstTargetImage.notation}...` : "Embedding target image...";
              ctx.status(targetImageStatus);
              multimodalQueryEmbedding = (await multimodalClient.embedImageBytes(
                new Uint8Array(projectTargetBytes),
                multimodalDimension,
                targetImageStatus,
              )).embedding;
              multimodalQueryMode = "image";
            } else if (queryEmbeddingImagePath && !excludeImage && embeddingQueryText.trim()) {
              const targetImageStatus = firstTargetImage?.notation ? `Embedding target image ${firstTargetImage.notation}...` : "Embedding target image...";
              ctx.status(targetImageStatus);
              multimodalQueryEmbedding = (await multimodalClient.embedImageAndText(
                queryEmbeddingImagePath,
                embeddingQueryText,
                multimodalDimension,
                targetImageStatus,
                queryEmbeddingInstruction,
              )).embedding;
              multimodalQueryMode = "mixed";
            } else if (queryEmbeddingImagePath && !excludeImage) {
              const targetImageStatus = firstTargetImage?.notation ? `Embedding target image ${firstTargetImage.notation}...` : "Embedding target image...";
              ctx.status(targetImageStatus);
              multimodalQueryEmbedding = (await multimodalClient.embedImagePath(
                queryEmbeddingImagePath,
                multimodalDimension,
                targetImageStatus,
              )).embedding;
              multimodalQueryMode = "image";
            } else if (embeddingQueryText.trim()) {
              ctx.status("Embedding query...");
              multimodalQueryEmbedding = (await multimodalClient.embedText(
                embeddingQueryText,
                multimodalDimension,
                "Embedding query...",
                queryEmbeddingInstruction,
              )).embedding;
              multimodalQueryMode = "text";
            }
          } catch (error: any) {
            const message = String(error?.message ?? error);
            throw new Error(`Multimodal query embedding failed: ${message}`);
          } finally {
            await releaseEmbeddingServer();
          }
        } else {
          await releaseEmbeddingServer();
        }
        
        const result = await searchGenerations(query, generations, {
          ...(() => {
            const isFastPath = isFilenameQuery || isProjectFilenameQuery;
            const maxResults = !isFastPath && useReranking
              ? Math.min(rerankingCandidateMultiplier * limit, rerankingMaxCandidates)
              : finalLimit;
            return maxResults !== undefined ? { maxResults } : {};
          })(),
          snapshot: enableModelRewrite ? snapshot : undefined,
          includeMultimodalSearch: !!multimodalQueryEmbedding && !!multimodalEmbeddingStore,
          multimodalQueryEmbedding,
          multimodalStore: multimodalEmbeddingStore ?? undefined,
          multimodalModel: multimodalModelKey,
          multimodalDimension,
          minMultimodalScore: embeddingMinScore,
          embeddingDisplayMultiplier,
          multimodalQueryMode,
          contentOnly: true,
          targetContentHash,
          targetGenerationMetadata,
          metadataExactMatches: args.includeMetadata === true && excludeImage,
          allowByteIdenticalResults,
          structuredFilterTolerances,
        });

        if (useReranking) {
          const recalled = [...result.exactMatches, ...result.semanticMatches];
          const recalledContent = recalled.filter((match) => match.matchType !== "exact");
          const recallRanks = new Map(recalledContent.map((match, index) => [match.imagePaths.join("\u0000"), index + 1]));
          const recallRankFor = (match: typeof recalled[number]) => recallRanks.get(match.imagePaths.join("\u0000"));
          const rerankable = (await Promise.all(recalledContent.map(async (match) => {
            const image = await rerankingImageForMatch(match, usePreviewsForQueryEmbedding);
            return image ? { match, ...image } : undefined;
          }))).filter((entry): entry is { match: typeof recalled[number]; imagePath?: string; imageBytes?: Uint8Array } => !!entry);
          if (rerankable.length > 0) {
            const scores = await rerankWithLlamaServer(
              {
                binaryPath: multimodalGgufBinaryPath,
                modelPath: rerankingModelPath,
                port: rerankingServicePort,
                contextSize: rerankingContextSize,
                batchSize: rerankingBatchSize,
                ubatchSize: rerankingUbatchSize,
                gpuLayers: multimodalGgufGpuLayers,
              },
              rerankingQuery,
              rerankingInstruction,
              rerankable.map(({ imagePath, imageBytes }) => ({
                text: "",
                imagePath,
                imageBytes,
              })),
              (message) => ctx.status(message),
            );
            const rerankedCandidates = rerankable.flatMap(({ match, imagePath }, index) => {
              const score = scores[index];
              if (typeof score !== "number" || !Number.isFinite(score)) return [];
              return [{
                match: {
                  ...match,
                  matchType: "reranked" as const,
                  rerankingScore: score,
                  matchScore: Math.min(100, Math.round(score * 100 * rerankingDisplayMultiplier)),
                },
                imagePath,
                rerankRequestOrder: index + 1,
              }];
            });
            const reranked = rerankedCandidates
              .map(({ match }) => match)
              .sort((left, right) => (right.rerankingScore ?? 0) - (left.rerankingScore ?? 0))
              .filter((match) => (match.rerankingScore ?? 0) >= rerankingMinScore);
            const rerankedOriginalMatches = new Set(rerankedCandidates.map(({ match }) => match.imagePaths.join("\u0000")));
            const fallbackMatches = recalledContent.filter((match) => !rerankedOriginalMatches.has(match.imagePaths.join("\u0000")));
            const remainingLimit = finalLimit === undefined
              ? undefined
              : Math.max(0, finalLimit - result.exactMatches.length);
            result.semanticMatches = [...reranked, ...fallbackMatches].slice(0, remainingLimit);
            const ranked = [...result.exactMatches, ...reranked];
            result.totalFound = ranked.length;
            result.imageResults = ranked.flatMap((match) => match.imagePaths.map((imagePath) => ({
              path: imagePath,
              prompt: match.prompt,
              model: match.model,
              matchType: match.matchType,
              matchScore: match.matchScore,
              score: match.matchScore,
            })));
            appendEmbeddingRerankingDiagnostic({
              mode: "embedding-plus-reranking",
              query: parsedQuery,
              targetImagePath,
              includeMetadata: !!args.includeMetadata,
              retrievalLimit: limit,
              recallCandidates: recalledContent.map((match, index) => diagnosticMatch(match, { recallRank: index + 1 })),
              skippedRecallCandidates: recalledContent
                .filter((match) => !rerankable.some((entry) => entry.match === match))
                .map((match) => diagnosticMatch(match, { recallRank: recallRankFor(match) })),
              rerankedCandidates: rerankedCandidates.map(({ match, imagePath, rerankRequestOrder }) => diagnosticMatch(match, {
                recallRank: recallRankFor(match),
                rerankRequestOrder,
              }, imagePath)),
              finalResults: ranked.map((match, index) => diagnosticMatch(match, {
                recallRank: recallRankFor(match),
                finalRank: index + 1,
              })),
              distinctVisibleImages: new Set(result.imageResults.map((image) => image.path)).size,
            });
          } else {
            appendEmbeddingRerankingDiagnostic({
              mode: "embedding-plus-reranking-no-rerankable-candidates",
              query: parsedQuery,
              targetImagePath,
              includeMetadata: !!args.includeMetadata,
              retrievalLimit: limit,
              recallCandidates: recalledContent.map((match, index) => diagnosticMatch(match, { recallRank: index + 1 })),
              distinctVisibleImages: new Set(result.imageResults.map((image) => image.path)).size,
            });
          }
        } else if (multimodalQueryEmbedding) {
          const ranked = [...result.exactMatches, ...result.semanticMatches];
          const recallRanks = new Map(result.semanticMatches.map((match, index) => [match, index + 1]));
          appendEmbeddingRerankingDiagnostic({
            mode: "embedding-only",
            query: parsedQuery,
            targetImagePath,
            includeMetadata: !!args.includeMetadata,
            retrievalLimit: limit,
            recallCandidates: result.semanticMatches.map((match, index) => diagnosticMatch(match, { recallRank: index + 1 })),
            finalResults: ranked.map((match, index) => diagnosticMatch(match, {
              recallRank: recallRanks.get(match),
              finalRank: index + 1,
            })),
            distinctVisibleImages: new Set(result.imageResults.map((image) => image.path)).size,
          });
        }
        
        ctx.status(`Found ${result.totalFound} results`);
        
        // Return structured content for draw-things-chat
        return buildToolResponse(result, {
          enableModelRewrite,
          snapshot,
          rankByVisualScore: multimodalQueryMode === "mixed" || multimodalQueryMode === "text",
          target: targetInput || undefined,
          includeMetadata: !!args.includeMetadata,
        });
        
      } catch (e) {
        return JSON.stringify({
          type: "find-image-results",
          query: String(args.query ?? ""),
          target: String(args.target ?? "") || undefined,
          includeMetadata: !!args.includeMetadata,
          error: String((e as any)?.message || e),
          totalFound: 0,
          searchTimeMs: 0,
          images: [],
        }, null, 2);
      } finally {
        await releaseEmbeddingServer?.();
      }
    },
  });
}

/**
 * Build structured tool response
 * Returns a JSON string that can be parsed by the model
 * 
 * @param result - Search results
 * @param capabilityMessage - Optional message about embedding capability (for user guidance)
 */
function buildToolResponse(
  result: DrawThingsSearchResult,
  opts?: {
    enableModelRewrite?: boolean;
    snapshot?: DtcModelMappingSnapshotV1;
    rankByVisualScore?: boolean;
    target?: string;
    includeMetadata?: boolean;
  }
): string {
  const enableModelRewrite = !!opts?.enableModelRewrite;
  const snapshot = opts?.snapshot;
  const rankByVisualScore = !!opts?.rankByVisualScore;
  // Combine all matches into single array
  const matches = [...result.exactMatches, ...result.semanticMatches];
  matches.sort((a, b) => {
    const aIsExact = a.matchType === "exact";
    const bIsExact = b.matchType === "exact";
    if (aIsExact !== bIsExact) return aIsExact ? -1 : 1;
    const arerank = a.rerankingScore;
    const brerank = b.rerankingScore;
    if (arerank !== undefined || brerank !== undefined) {
      return (brerank ?? -1) - (arerank ?? -1);
    }
    if (rankByVisualScore) {
      const ar = a.rankingSimilarity ?? -1;
      const br = b.rankingSimilarity ?? -1;
      if (ar !== br) return br - ar;
      const av = a.visualScore ?? a.multimodalScore ?? -1;
      const bv = b.visualScore ?? b.multimodalScore ?? -1;
      if (av !== bv) return bv - av;
    }
    const scoreDifference = b.matchScore - a.matchScore;
    if (scoreDifference !== 0) return scoreDifference;
    if (aIsExact && bIsExact) {
      const aPaths = a.imagePaths.join("\u0000");
      const bPaths = b.imagePaths.join("\u0000");
      return aPaths < bPaths ? -1 : aPaths > bPaths ? 1 : 0;
    }
    return 0;
  });

  const allMatches = matches.map(match => ({
      matchType: match.matchType,
      matchScore: match.matchScore,
      prompt: match.prompt,
      negativePrompt: match.negativePrompt,
      model: match.model,
      ...(enableModelRewrite
        ? { model_display: maybeFormatModelForConsumerFromSnapshot(match.model, true, snapshot) }
        : {}),
      ...(enableModelRewrite
        ? (() => {
            const hints = maybeGetModelRewriteHintsFromSnapshot(match.model, true, snapshot);
            return hints ? { model_use_hints: hints } : {};
          })()
        : {}),
      loras: match.loras || [],
      tags: match.tags || [],
      sampler: match.sampler,
      steps: match.steps,
      cfgScale: match.cfgScale,
      seed: match.seed,
      seedMode: match.seedMode,
      shift: match.shift,
      strength: match.strength,
      mode: match.mode,
      width: match.width,
      height: match.height,
      numFrames: match.numFrames,
      imagePaths: match.imagePaths,
      httpPreviewUrls: match.httpPreviewUrls || [],
      sourceInfo: match.sourceInfo,
      timestamp: match.timestamp,
      latitude: match.latitude,
      longitude: match.longitude,
      lensModel: match.lensModel,
      exposureTime: match.exposureTime,
      fNumber: match.fNumber,
      iso: match.iso,
      exposureCompensation: match.exposureCompensation,
      focalLength: match.focalLength,
      focalLength35mm: match.focalLength35mm,
      exposureProgram: match.exposureProgram,
      meteringMode: match.meteringMode,
      whiteBalance: match.whiteBalance,
      flash: match.flash,
      orientation: match.orientation,
      exposureMode: match.exposureMode,
    }));

  const response: Record<string, any> = {
    type: "find-image-results",
    query: result.query,
    target: opts?.target,
    includeMetadata: !!opts?.includeMetadata,
    totalFound: allMatches.length,
    searchTimeMs: result.searchTimeMs,
    semanticSearchEnabled: result.semanticSearchEnabled,
    multimodalSearchEnabled: result.multimodalSearchEnabled ?? false,
    images: allMatches,
  };

  return JSON.stringify(response, null, 2);
}