draw-things-index

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

import { tool, type Tool, type ToolsProviderController, type ToolCallContext } from "@lmstudio/sdk";
import { z } from "zod";
import { searchGenerations } from "../search/searchEngine";
import { configSchematics } from "../config";
import { EmbeddingClient } from "../embeddings";
import { checkEmbeddingCapability, type EmbeddingPrimerResult } from "../helpers/embeddingCapabilityPrimer";
import { formatToolMetaBlock } from "../helpers/pluginMeta";
import {
  parseDtcQuery,
  maybeFormatModelForConsumerFromSnapshot,
  maybeGetModelRewriteHintsFromSnapshot,
  type DtcModelMappingSnapshotV1,
} from "../helpers/dtcModelMappingSnapshot";
import type { DrawThingsSearchResult } from "../types";

const IMAGE_FILENAME_EXTS = new Set(["png", "jpg", "jpeg", "webp", "gif", "tif", "tiff", "bmp", "heic", "heif"]);

function looksLikeImageFilenameQuery(q: string): boolean {
  const trimmed = q.trim().replace(/^['"]|['"]$/g, "");
  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 IMAGE_FILENAME_EXTS.has(ext);
}

// Cached embedding client (reused across tool calls)
let embeddingClient: EmbeddingClient | null = null;

// Cached embedding capability check (refreshed periodically)
let cachedCapabilityResult: EmbeddingPrimerResult | null = null;
let lastCapabilityCheckMs = 0;
const CAPABILITY_CHECK_INTERVAL_MS = 30_000; // Re-check every 30s

// Tool parameters schema
const SearchParamsSchema = {
  query: z.string().describe("Search query - filename, prompt text, model name, LoRA name, or keywords"),
};

export function createSearchGenerationsTool(ctl: ToolsProviderController): Tool {
  return tool({
    name: "index_image",
    description: `Search through Draw Things generation history.

The search includes filename/originalName matching.

Use this tool to find metadata about previously generated images:
- File name (Attachments or logged generations)
- Prompt text (supports fuzzy matching for typos)
- Model name (e.g., "flux", "sd3")
- LoRA name
- Keywords from the generation
- Cross-language search (e.g., "Katze" finds "cat" when semantic search enabled)

Returns:
- Exact/fuzzy matches: Direct matches to your query
- Semantic matches: Thematically related results
- Image paths for each matching generation

Example queries:
- "Draw-Things-Generated-Image.png" - finds the image by filename
- "cat in space" - finds generations with similar prompts
- "flux" - finds all FLUX model generations
- "cyberpunk portrait" - finds matching style/subject
- "Sonnenuntergang" - also finds "sunset", "golden hour" via semantic search

${formatToolMetaBlock()}`,
    parameters: SearchParamsSchema,
    implementation: async (args: { query: string }, ctx: ToolCallContext) => {
      try {
        ctx.status("Starting search...");

        const { query, enableModelRewrite, snapshot } = parseDtcQuery(args.query);

        if (enableModelRewrite) {
          ctx.status("Model rewrite enabled; snapshot received.");
        }
        
        // Get indexed generations with progress feedback
        const { indexGenerations } = await import("../indexer");
        const generations = await indexGenerations(
          ctl, 
          false, // use cache if available
          (msg) => ctx.status(msg)
        );
        
        ctx.status(`Searching ${generations.length} generations...`);
        
        if (generations.length === 0) {
          return JSON.stringify({
            type: "draw-things-index-results",
            query: args.query,
            totalFound: 0,
            error: "No generations indexed. Check if JSONL logs directory is configured correctly.",
            searchTimeMs: 0,
            images: [],
          }, null, 2);
        }
        
        // Get config values for search tuning
        const config = ctl.getGlobalPluginConfig(configSchematics);
        const minMatchScore = config.get("minMatchScore");
        const fuzzyTermThreshold = config.get("fuzzyTermThreshold");
        const minTermCoverage = config.get("minTermCoverage");
        
        // Semantic search config (enabled if weight > 0 and model configured)
        const embeddingModel = config.get("embeddingModel");
        const lmStudioUrl = config.get("lmStudioBaseUrl");
        const semanticWeight = config.get("semanticWeight");
        const minSemanticScore = config.get("minSemanticScore");
        const isFilenameQuery = looksLikeImageFilenameQuery(query);
        const userWantsSemanticSearch = !isFilenameQuery && semanticWeight > 0 && !!embeddingModel;
        
        console.info(`[Search Config] minMatchScore=${minMatchScore}, fuzzyTermThreshold=${fuzzyTermThreshold}, minTermCoverage=${minTermCoverage}`);
        console.info(`[Search Config] filenameQuery=${isFilenameQuery}, semantic=${userWantsSemanticSearch}, model=${embeddingModel}, weight=${semanticWeight}`);
        
        // Limit aus Plugin-Config (UI), nicht vom LLM überschreibbar
        const limit = config.get("retrievalLimit");
        const effectiveLimit = limit >= 25 ? 9999 : limit;
        console.info(`[Search Config] retrievalLimit=${limit}, effectiveLimit=${effectiveLimit}`);
        
        // ═══════════════════════════════════════════════════════════════
        // EMBEDDING CAPABILITY CHECK (adapted from vision capability primer)
        // ═══════════════════════════════════════════════════════════════
        let semanticEnabled = false;
        let capabilityMessage: string | undefined;
        let actualEmbeddingModel = embeddingModel;
        
        if (userWantsSemanticSearch) {
          // Check if we need to refresh capability status
          const now = Date.now();
          if (!cachedCapabilityResult || (now - lastCapabilityCheckMs) > CAPABILITY_CHECK_INTERVAL_MS) {
            ctx.status("Checking embedding model availability...");
            cachedCapabilityResult = await checkEmbeddingCapability({
              modelId: embeddingModel,
              baseUrl: lmStudioUrl,
              autoLoad: false, // Don't auto-load during search, just guide user
            });
            lastCapabilityCheckMs = now;
          }
          
          const capResult = cachedCapabilityResult;
          
          if (capResult.ready && capResult.isLoaded) {
            // Great! Semantic search is available
            semanticEnabled = true;
            actualEmbeddingModel = capResult.modelId; // Use actually loaded model
            
            // Update client if model changed
            if (!embeddingClient || embeddingClient.getModelName() !== actualEmbeddingModel) {
              embeddingClient = new EmbeddingClient({
                baseUrl: lmStudioUrl,
                model: actualEmbeddingModel,
              });
            }
          } else {
            // Semantic search not available - graceful degradation
            semanticEnabled = false;
            capabilityMessage = capResult.userMessage;
            console.warn("[index_image] Semantic search disabled:", capResult.error || "No embedding model loaded");
          }
        }
        
        const result = await searchGenerations(query, generations, {
          maxExactMatches: effectiveLimit,
          maxSemanticMatches: Math.floor(effectiveLimit / 2),
          minExactScore: minMatchScore,
          minSemanticScore: minSemanticScore,
          includeSemanticSearch: semanticEnabled,
          embeddingClient: semanticEnabled ? embeddingClient ?? undefined : undefined,
          semanticWeight: semanticWeight,
          fuzzyOptions: {
            fuzzyTermThreshold,
            minTermCoverage,
          },
          snapshot: enableModelRewrite ? snapshot : undefined,
        });
        
        ctx.status(`Found ${result.totalFound} results`);
        
        // Return structured content for draw-things-chat
        return buildToolResponse(result, { capabilityMessage, enableModelRewrite, snapshot });
        
      } catch (e) {
        return JSON.stringify({
          type: "draw-things-index-results",
          query: args.query,
          error: String((e as any)?.message || e),
          totalFound: 0,
          searchTimeMs: 0,
          images: [],
        }, null, 2);
      }
    },
  });
}

/**
 * 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?: { capabilityMessage?: string; enableModelRewrite?: boolean; snapshot?: DtcModelMappingSnapshotV1 }
): string {
  const formatSourceLabel = (sourceInfo: any): string | undefined => {
    if (!sourceInfo || typeof sourceInfo !== "object") return undefined;
    switch (sourceInfo.type) {
      case "attachment":
        return sourceInfo.originalName ? `attachment (${sourceInfo.originalName})` : "attachment";
      case "saved_image":
        return sourceInfo.imageType ? `saved image (${sourceInfo.imageType})` : "saved image";
      case "generate_image_variant":
        return sourceInfo.chatId ? `generate image (${sourceInfo.chatId})` : "generate image";
      case "draw_things_project":
        return sourceInfo.projectFile ? `Draw Things project (${sourceInfo.projectFile})` : "Draw Things project";
      default:
        return undefined;
    }
  };

  const enableModelRewrite = !!opts?.enableModelRewrite;
  const snapshot = opts?.snapshot;
  // Combine all matches into single array
  const allMatches = [
    ...result.exactMatches.map(match => ({
      matchType: match.matchType,
      matchScore: match.matchScore,
      prompt: match.prompt,
      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 || [],
      width: match.width,
      height: match.height,
      numFrames: match.numFrames,
      imagePaths: match.imagePaths,
      httpPreviewUrls: match.httpPreviewUrls || [],
      sourceInfo: match.sourceInfo,
      sourceLabel: formatSourceLabel(match.sourceInfo),
      timestamp: match.timestamp,
    })),
    ...result.semanticMatches.map(match => ({
      matchType: 'semantic' as const,
      matchScore: match.matchScore,
      prompt: match.prompt,
      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 || [],
      width: match.width,
      height: match.height,
      numFrames: match.numFrames,
      imagePaths: match.imagePaths,
      httpPreviewUrls: match.httpPreviewUrls || [],
      sourceInfo: match.sourceInfo,
      sourceLabel: formatSourceLabel(match.sourceInfo),
      timestamp: match.timestamp,
    })),
  ];

  // Sort by matchScore descending
  allMatches.sort((a, b) => b.matchScore - a.matchScore);

  const response: Record<string, any> = {
    type: "draw-things-index-results",
    query: result.query,
    totalFound: allMatches.length,
    searchTimeMs: result.searchTimeMs,
    semanticSearchEnabled: result.semanticSearchEnabled,
    images: allMatches,
  };
  
  // Include capability message for user guidance (displayed by chat model)
  if (opts?.capabilityMessage) {
    response.capabilityNotice = opts.capabilityMessage;
  }

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