draw-things-index

/**
 * Draw Things Index - Prompt Preprocessor
 * Searches generation history and provides structured results
 */

import {
  type ChatMessage,
  type PromptPreprocessorController,
} from "@lmstudio/sdk";
import { configSchematics } from "./config";
import { indexGenerations } from "./indexer";
import { searchGenerations } from "./search/searchEngine";
import type { DrawThingsSearchResult, GenerationMatch } from "./types";
import path from "path";

// ═══════════════════════════════════════════════════════════════
// Main Preprocessor
// ═══════════════════════════════════════════════════════════════

export async function preprocess(
  ctl: PromptPreprocessorController,
  userMessage: ChatMessage
): Promise<ChatMessage | string> {
  const userPrompt = userMessage.getText();
  const config = ctl.getGlobalPluginConfig(configSchematics);

  // Check if any search source is enabled
  const contentDirectories = config.get("contentDirectories");
  const anySourceEnabled = 
    config.get("searchJsonlLogs") ||
    contentDirectories.length > 0;

  if (!anySourceEnabled) {
    return userMessage;
  }

  try {
    // ─────────────────────────────────────────────────────────────
    // Index sources (uses shared cache from indexer.ts)
    // ─────────────────────────────────────────────────────────────
    const indexStatus = ctl.createStatus({
      status: "loading",
      text: "Indexing Draw Things data...",
    });

    const generations = await indexGenerations(ctl);
    
    indexStatus.setState({
      status: "done",
      text: `Indexed ${generations.length} generations`,
    });

    // ─────────────────────────────────────────────────────────────
    // Search
    // ─────────────────────────────────────────────────────────────
    const searchStatus = ctl.createStatus({
      status: "loading",
      text: "Searching...",
    });

    // 25 = "all" (unlimited)
    const configLimit = config.get("retrievalLimit");
    const limit = configLimit >= 25 ? 9999 : configLimit;
    
    const searchResult = await searchGenerations(userPrompt, generations, {
      maxExactMatches: limit,
      maxSemanticMatches: Math.floor(limit / 2),
      minExactScore: 25,
      includeSemanticSearch: false,
    });

    if (searchResult.totalFound === 0) {
      searchStatus.setState({
        status: "canceled",
        text: "No matching generations found",
      });

      const note = buildNoResultsContext(userPrompt, generations.length);
      userMessage.replaceText(note);
      return userMessage;
    }

    searchStatus.setState({
      status: "done",
      text: `Found ${searchResult.totalFound} results (${searchResult.searchTimeMs}ms)`,
    });

    // ─────────────────────────────────────────────────────────────
    // Build context for LLM
    // ─────────────────────────────────────────────────────────────
    const context = buildSearchContext(userPrompt, searchResult);
    userMessage.replaceText(context);

    return userMessage;

  } catch (error) {
    ctl.createStatus({
      status: "canceled",
      text: `Error: ${error instanceof Error ? error.message : 'Unknown error'}`,
    });
    return userMessage;
  }
}

// ═══════════════════════════════════════════════════════════════
// Context Building
// ═══════════════════════════════════════════════════════════════

/**
 * Build context when no results found
 */
function buildNoResultsContext(query: string, totalIndexed: number): string {
  return `## 🎨 Draw Things Index

**Query:** "${query}"
**Result:** No matching generations found in ${totalIndexed} indexed entries.

### Suggestions:
- Try broader search terms
- Check if the generation logs directory is configured correctly
- The prompt might have been worded differently

---

**User Question:** ${query}`;
}

/**
 * Build search context for LLM
 */
function buildSearchContext(query: string, result: DrawThingsSearchResult): string {
  const parts: string[] = [];
  
  parts.push('## 🎨 Draw Things Generation Search Results\n');
  parts.push(`**Query:** "${query}"`);
  parts.push(`**Found:** ${result.totalFound} results in ${result.searchTimeMs}ms\n`);

  // Volltreffer (Exact Matches)
  if (result.exactMatches.length > 0) {
    parts.push('### 🎯 Volltreffer (Direct Matches)\n');
    
    for (let i = 0; i < result.exactMatches.length; i++) {
      const match = result.exactMatches[i];
      parts.push(formatGenerationMatch(i + 1, match));
    }
  }

  // Beifang (Semantic Matches)
  if (result.semanticMatches.length > 0) {
    parts.push('### 🔮 Beifang (Related Results)\n');
    parts.push('*These results are thematically related but may not contain exact query terms.*\n');
    
    for (let i = 0; i < result.semanticMatches.length; i++) {
      const match = result.semanticMatches[i];
      parts.push(formatGenerationMatch(i + 1, match, true));
    }
  }

  // Image List for Reference
  if (result.imageResults.length > 0) {
    parts.push('### 📸 Image Paths\n');
    parts.push('```');
    for (const img of result.imageResults) {
      const matchIcon = img.matchType === 'exact' ? '🎯' : 
                        img.matchType === 'fuzzy' ? '≈' : 
                        img.matchType === 'semantic' ? '🔮' : '○';
      parts.push(`${matchIcon} ${img.path}`);
    }
    parts.push('```\n');
  }

  // Instructions for LLM
  parts.push('---\n');
  parts.push('## Instructions\n');
  parts.push('Use the generation history above to answer the user\'s question.');
  parts.push('- Reference specific prompts, models, LoRAs, and settings when relevant.');
  parts.push('- If asked about images, provide the file paths.');
  parts.push('- Distinguish between direct matches (🎯) and related results (🔮).\n');
  
  parts.push(`## User Question\n${query}`);

  return parts.join('\n');
}

/**
 * Format a single generation match
 */
function formatGenerationMatch(
  index: number, 
  match: GenerationMatch,
  isSemanticMatch = false
): string {
  const lines: string[] = [];
  
  const matchTypeIcon = match.matchType === 'exact' ? '✓ Exact' :
                        match.matchType === 'fuzzy' ? '≈ Fuzzy' :
                        match.matchType === 'partial' ? '◐ Partial' :
                        '🔮 Semantic';
  
  lines.push(`**${index}. [${matchTypeIcon}] Score: ${match.matchScore}**\n`);
  
  // Prompt (truncate if very long)
  const promptDisplay = match.prompt.length > 200 
    ? match.prompt.slice(0, 200) + '...' 
    : match.prompt;
  lines.push(`**Prompt:** ${promptDisplay}\n`);
  
  // Negative prompt (if present)
  if (match.negativePrompt) {
    const negDisplay = match.negativePrompt.length > 100
      ? match.negativePrompt.slice(0, 100) + '...'
      : match.negativePrompt;
    lines.push(`**Negative:** ${negDisplay}\n`);
  }
  
  // Model and LoRAs
  lines.push(`**Model:** ${match.model}`);
  if (match.loras && match.loras.length > 0) {
    lines.push(`**LoRAs:** ${match.loras.join(', ')}`);
  }
  
  // Generation settings
  const settings: string[] = [];
  if (match.seed !== undefined) settings.push(`Seed: ${match.seed}`);
  if (match.steps !== undefined) settings.push(`Steps: ${match.steps}`);
  if (match.cfgScale !== undefined) settings.push(`CFG: ${match.cfgScale}`);
  if (match.width && match.height) settings.push(`Size: ${match.width}×${match.height}`);
  
  if (settings.length > 0) {
    lines.push(`**Settings:** ${settings.join(' | ')}`);
  }
  
  // Images
  if (match.imagePaths.length > 0) {
    lines.push(`**Images:** ${match.imagePaths.length} file(s)`);
    if (match.imagePaths.length <= 3) {
      for (const p of match.imagePaths) {
        lines.push(`  • \`${path.basename(p)}\``);
      }
    }
  }
  
  // Timestamp
  if (match.timestamp) {
    lines.push(`**Generated:** ${match.timestamp}`);
  }
  
  // Matched terms (for debugging/transparency)
  if (match.matchedTerms && match.matchedTerms.length > 0 && !isSemanticMatch) {
    lines.push(`*Matched: ${match.matchedTerms.join(', ')}*`);
  }
  
  lines.push('');
  return lines.join('\n');
}