find-image

/**
 * GGUF multimodal embedding client — talks directly to a native `llama-server`
 * process over its `/embeddings` HTTP route. There is no Python and no
 * intermediate FastAPI server: the chat-template prompt that
 * `llama-vl-embedding` would normally render internally is built here in
 * TypeScript and sent as an already-rendered `prompt_string`.
 */
import { ensureLlamaServerRunning, LLAMA_MEDIA_MARKER, stopLlamaServerEmbeddingServer } from "../llama-server-manager.js";
import { findImageDataPath, legacyPluginDataPath } from "../paths.js";
import { EMBEDDING_INSTRUCTIONS } from "./embeddingInstructions.js";
import type {
  MultimodalEmbeddingBackend,
  MultimodalEmbeddingHealth,
  MultimodalEmbeddingResult,
} from "./multimodalEmbeddingClient.js";

export interface MultimodalEmbeddingGgufClientConfig {
  /** Port the local llama-server embedding process listens on (always bound to 127.0.0.1 — see llama-server-manager.ts). */
  port: number;
  binaryPath: string;
  modelDir: string;
  dimension: number;
  contextSize: number;
  nGpuLayers?: number;
  abortSignal?: AbortSignal;
  onStatus?: (message: string) => void;
}

export const GGUF_DB_PATH = findImageDataPath("multimodal_embeddings.sqlite3");
export const LEGACY_GGUF_DB_PATH = legacyPluginDataPath("multimodal_embeddings.sqlite3");

/**
 * Default system instruction used for image-only Embedding indexing and
 * image-only query embedding — matches `llama-vl-embedding`'s
 * `DEFAULT_INSTRUCTION` exactly (`tools/mtmd/vl-embedding.cpp`).
 */
const DEFAULT_INSTRUCTION = EMBEDDING_INSTRUCTIONS.corpus;

// ASCII punctuation set matching C's ispunct(), used to replicate
// llama-vl-embedding's normalize_instruction()/ends_with_punctuation() 1:1.
const PUNCT_CHARS = "!\"#$%&'()*+,-./:;<=>?@[\\]^_`{|}~";

function endsWithPunctuation(s: string): boolean {
  return s.length > 0 && PUNCT_CHARS.includes(s[s.length - 1]);
}

function normalizeInstruction(instruction: string | undefined): string {
  const trimmed = (instruction ?? DEFAULT_INSTRUCTION).trim();
  if (trimmed && !endsWithPunctuation(trimmed)) {
    return `${trimmed}.`;
  }
  return trimmed;
}

/**
 * Builds the user-message content: the media marker first (if an image is
 * present), then any text, concatenated directly with no separator — or the
 * literal string "NULL" when both are absent. Mirrors
 * `build_conversation()`'s user-content assembly in `vl-embedding.cpp`.
 */
function buildUserContent(hasImage: boolean, text: string | undefined): string {
  const t = text ?? "";
  if (!hasImage && t.length === 0) return "NULL";
  return (hasImage ? LLAMA_MEDIA_MARKER : "") + t;
}

/**
 * Renders the exact Qwen chat-template prompt string
 * (`<|im_start|>role\n...<|im_end|>\n`) that `common_chat_templates_apply()`
 * produces for this model, with `add_generation_prompt=true`. Verified
 * byte-for-byte against `llama-vl-embedding --verbose-prompt` output for
 * text-only, image-only, and image+text cases.
 */
function buildPrompt(instruction: string | undefined, hasImage: boolean, text?: string): string {
  const system = normalizeInstruction(instruction);
  const user = buildUserContent(hasImage, text);
  return `<|im_start|>system\n${system}<|im_end|>\n<|im_start|>user\n${user}<|im_end|>\n<|im_start|>assistant\n`;
}

/**
 * Truncates a raw (already L2-normalized-at-native-dimension) embedding to
 * the requested Matryoshka dimension and re-normalizes the slice — slicing a
 * unit vector does not yield a unit vector. Mirrors the previous Python
 * GGUF/HF backends' `_normalize()`.
 */
function truncateAndNormalize(values: number[], dimension: number): number[] {
  const vec = values.slice(0, dimension);
  let normSq = 0;
  for (const v of vec) normSq += v * v;
  const norm = Math.sqrt(normSq);
  if (norm > 0) {
    for (let i = 0; i < vec.length; i++) vec[i] /= norm;
  }
  return vec;
}

interface EmbeddingsContentItem {
  prompt_string: string;
  multimodal_data?: string[];
}

interface EmbeddingsResultItem {
  index: number;
  embedding: number[][];
}

async function postEmbeddings(port: number, content: EmbeddingsContentItem[], signal?: AbortSignal): Promise<EmbeddingsResultItem[]> {
  const res = await fetch(`http://127.0.0.1:${port}/embeddings`, {
    method: "POST",
    headers: { "content-type": "application/json" },
    body: JSON.stringify({ content }),
    signal,
  });
  if (!res.ok) {
    let detail = `${res.status} ${res.statusText}`;
    try {
      const body = (await res.json()) as { error?: { message?: unknown } };
      detail = String(body.error?.message ?? detail);
    } catch (error) {
      detail = `${detail}; failed to parse error response: ${error instanceof Error ? error.message : String(error)}`;
    }
    throw new Error(detail);
  }
  const results = (await res.json()) as EmbeddingsResultItem[];
  return [...results].sort((a, b) => a.index - b.index);
}

function isRetryableTransportError(error: unknown): boolean {
  if (!(error instanceof TypeError) || error.message !== "fetch failed") return false;
  const cause = (error as TypeError & { cause?: { code?: unknown } }).cause;
  return typeof cause?.code === "string" && ["UND_ERR_SOCKET", "ECONNREFUSED", "ECONNRESET"].includes(cause.code);
}

export class GgufMultimodalEmbeddingBackend implements MultimodalEmbeddingBackend {
  private readyKey: string | null = null;
  private readyPromise: Promise<void> | null = null;

  constructor(private readonly config: MultimodalEmbeddingGgufClientConfig) {}

  async ensureReady(): Promise<void> {
    const port = this.config.port;
    const key = `${port}|${this.config.modelDir}|${this.config.contextSize}|${this.config.nGpuLayers ?? 999}`;
    if (this.readyKey === key) return;
    if (this.readyPromise) return this.readyPromise;

    this.readyPromise = ensureLlamaServerRunning(
      {
        port,
        binaryPath: this.config.binaryPath,
        modelDir: this.config.modelDir,
        contextSize: this.config.contextSize,
        nGpuLayers: this.config.nGpuLayers,
      },
      this.config.onStatus ?? (() => undefined)
    )
      .then(() => {
        this.readyKey = key;
      })
      .finally(() => {
        this.readyPromise = null;
      });

    return this.readyPromise;
  }

  private resetReady(): void {
    this.readyKey = null;
  }

  async health(): Promise<MultimodalEmbeddingHealth> {
    try {
      const res = await fetch(`http://127.0.0.1:${this.config.port}/health`);
      const ready = res.ok;
      return {
        ready,
        backend: "integrated-local-gguf",
        model: this.config.modelDir,
        dimension: this.config.dimension,
        modelLoaded: ready,
      };
    } catch (error: any) {
      return {
        ready: false,
        backend: "integrated-local-gguf",
        model: this.config.modelDir,
        dimension: this.config.dimension,
        error: String(error?.message ?? error),
      };
    }
  }

  private async embedSingle(promptString: string, imageBase64Data: string | undefined, dimension: number, statusLabel?: string): Promise<MultimodalEmbeddingResult> {
    const content: EmbeddingsContentItem[] = [
      { prompt_string: promptString, ...(imageBase64Data ? { multimodal_data: [imageBase64Data] } : {}) },
    ];
    for (let attempt = 0; attempt < 2; attempt++) {
      if (this.config.abortSignal?.aborted) throw new DOMException("Embedding cancelled.", "AbortError");
      await this.ensureReady();
      if (statusLabel) this.config.onStatus?.(statusLabel);
      try {
        const [result] = await postEmbeddings(this.config.port, content, this.config.abortSignal);
        const embedding = truncateAndNormalize(result.embedding[0], dimension);
        return { embedding, model: this.config.modelDir, dimension: embedding.length };
      } catch (error) {
        this.resetReady();
        if (attempt === 0 && isRetryableTransportError(error)) {
          this.config.onStatus?.("Embedding service connection dropped — restarting and retrying...");
          await stopLlamaServerEmbeddingServer(this.config.port);
          continue;
        }
        throw error;
      }
    }
    throw new Error("Embedding request retry loop completed unexpectedly.");
  }

  async embedImagePath(imagePath: string, dimension = this.config.dimension, statusLabel?: string): Promise<MultimodalEmbeddingResult> {
    const fs = await import("fs/promises");
    const bytes = await fs.readFile(imagePath);
    return this.embedImageBytes(new Uint8Array(bytes), dimension, statusLabel);
  }

  async embedImageBytes(bytes: Uint8Array, dimension = this.config.dimension, statusLabel?: string): Promise<MultimodalEmbeddingResult> {
    // Used both for query embedding (find_image.ts's target image, which
    // passes its own statusLabel) and for per-file Embedding indexing
    // (indexer.ts's generateMultimodalEmbeddings(), which reports its own
    // per-file progress and must not have it overwritten with a generic
    // message on every single file — so only announce a status when the
    // caller explicitly supplies one). Each call is its own /embeddings
    // request; the enclosing tool call controls the server lifetime.
    if (statusLabel) this.config.onStatus?.(statusLabel);
    const promptString = buildPrompt(undefined, true, undefined);
    return this.embedSingle(promptString, Buffer.from(bytes).toString("base64"), dimension, statusLabel);
  }

  async embedText(text: string, dimension = this.config.dimension, statusLabel = "Embedding query...", instruction: string = EMBEDDING_INSTRUCTIONS.textRetrieval): Promise<MultimodalEmbeddingResult> {
    this.config.onStatus?.(statusLabel);
    const promptString = buildPrompt(instruction, false, text);
    return this.embedSingle(promptString, undefined, dimension, statusLabel);
  }

  async embedImageBytesAndText(bytes: Uint8Array, text: string, dimension = this.config.dimension, statusLabel?: string, instruction?: string): Promise<MultimodalEmbeddingResult> {
    if (statusLabel) this.config.onStatus?.(statusLabel);
    const promptString = buildPrompt(instruction, true, text);
    return this.embedSingle(promptString, Buffer.from(bytes).toString("base64"), dimension, statusLabel);
  }

  async embedImageAndText(imagePath: string, text: string, dimension = this.config.dimension, statusLabel = "Embedding query...", instruction: string = EMBEDDING_INSTRUCTIONS.fusion): Promise<MultimodalEmbeddingResult> {
    this.config.onStatus?.(statusLabel);
    const fs = await import("fs/promises");
    const bytes = await fs.readFile(imagePath);
    const promptString = buildPrompt(instruction, true, text);
    return this.embedSingle(promptString, Buffer.from(bytes).toString("base64"), dimension, statusLabel);
  }
}