find-image

import { accessSync, constants, existsSync, readFileSync } from "fs";
import { spawn } from "child_process";
import path from "path";
import { resolveGgufModelFiles } from "../llama-server-manager.js";

export interface RerankingQuery {
  text: string;
  imagePath?: string;
  imageBytes?: Uint8Array;
}

export interface RerankingDocument {
  text: string;
  imagePath?: string;
  imageBytes?: Uint8Array;
}

export interface LlamaServerRerankerConfig {
  binaryPath: string;
  modelPath: string;
  port: number;
  contextSize: number;
  batchSize: number;
  ubatchSize: number;
  gpuLayers: number;
}

const HEALTH_TIMEOUT_MS = 90_000;
const REQUEST_TIMEOUT_MS = 240_000;
const STOP_TIMEOUT_MS = 5_000;

function imageMimeType(imagePath: string | undefined, bytes: Uint8Array): string {
  if (bytes.length >= 3 && bytes[0] === 0xff && bytes[1] === 0xd8 && bytes[2] === 0xff) return "jpeg";
  if (bytes.length >= 8 && bytes[0] === 0x89 && bytes[1] === 0x50 && bytes[2] === 0x4e && bytes[3] === 0x47) return "png";
  const extension = imagePath?.split(".").pop()?.toLowerCase() || "png";
  return extension === "jpg" ? "jpeg" : extension;
}

function dataUrlForImage(imagePath?: string, imageBytes?: Uint8Array): string {
  const bytes = imageBytes ?? (imagePath ? readFileSync(imagePath) : undefined);
  if (!bytes) throw new Error("Reranking image has no readable path or image bytes.");
  const mimeType = imageMimeType(imagePath, bytes);
  return `data:image/${mimeType};base64,${Buffer.from(bytes).toString("base64")}`;
}

async function waitForHealth(port: number): Promise<void> {
  const deadline = Date.now() + HEALTH_TIMEOUT_MS;
  while (Date.now() < deadline) {
    try {
      const response = await fetch(`http://127.0.0.1:${port}/health`);
      if (response.ok) return;
    } catch {
      // The server is still loading.
    }
    await new Promise<void>((resolve) => setTimeout(resolve, 500));
  }
  throw new Error(`Reranking service did not become healthy within ${HEALTH_TIMEOUT_MS / 1000}s.`);
}

async function assertPortIsAvailable(port: number): Promise<void> {
  try {
    const response = await fetch(`http://127.0.0.1:${port}/health`);
    if (response.ok) {
      throw new Error(`Reranking service port ${port} is already occupied by another healthy server.`);
    }
  } catch (error) {
    if (error instanceof Error && error.message.includes("already occupied")) throw error;
  }
}

async function waitForExit(child: ReturnType<typeof spawn>, timeoutMs: number): Promise<boolean> {
  if (child.exitCode !== null) return true;
  return new Promise<boolean>((resolve) => {
    const timeout = setTimeout(() => resolve(child.exitCode !== null), timeoutMs);
    child.once("exit", () => {
      clearTimeout(timeout);
      resolve(true);
    });
  });
}

async function waitForPortRelease(port: number): Promise<boolean> {
  const deadline = Date.now() + STOP_TIMEOUT_MS;
  while (Date.now() < deadline) {
    try {
      await fetch(`http://127.0.0.1:${port}/health`);
    } catch {
      return true;
    }
    await new Promise<void>((resolve) => setTimeout(resolve, 100));
  }
  return false;
}

async function stopProcess(child: ReturnType<typeof spawn>, port: number): Promise<void> {
  if (!child.pid) return;
  if (child.exitCode === null) {
    child.kill("SIGTERM");
    if (!(await waitForExit(child, STOP_TIMEOUT_MS))) {
      child.kill("SIGKILL");
      if (!(await waitForExit(child, STOP_TIMEOUT_MS))) {
        throw new Error(`Reranking service PID ${child.pid} did not exit after SIGKILL.`);
      }
    }
  }
  if (!(await waitForPortRelease(port))) {
    throw new Error(`Reranking service port ${port} remained occupied after stopping PID ${child.pid}.`);
  }
}

function requestValue(value: RerankingQuery | RerankingDocument): string | { text: string; image: string } {
  if (!value.imagePath && !value.imageBytes) return value.text;
  return { text: value.text, image: dataUrlForImage(value.imagePath, value.imageBytes) };
}

async function rerankDocument(port: number, query: RerankingQuery, instruction: string, document: RerankingDocument): Promise<number> {
  const controller = new AbortController();
  const timeout = setTimeout(() => controller.abort(), REQUEST_TIMEOUT_MS);
  const response = await fetch(`http://127.0.0.1:${port}/rerank`, {
    method: "POST",
    headers: { "content-type": "application/json" },
    signal: controller.signal,
    body: JSON.stringify({
      instruction,
      query: requestValue(query),
      documents: [requestValue(document)],
    }),
  }).finally(() => clearTimeout(timeout));
  if (!response.ok) throw new Error(`Reranking request failed: ${response.status} ${await response.text()}`);
  const body = await response.json() as { results?: Array<{ relevance_score?: number }>; data?: Array<{ relevance_score?: number }> };
  const score = Number((body.results ?? body.data)?.[0]?.relevance_score);
  if (!Number.isFinite(score)) throw new Error("Reranking response had an invalid result.");
  return score;
}

export async function rerankWithLlamaServer(
  config: LlamaServerRerankerConfig,
  query: RerankingQuery,
  instruction: string,
  documents: RerankingDocument[],
  onStatus: (message: string) => void,
): Promise<Array<number | undefined>> {
  if (documents.length === 0) return [];
  if (!config.modelPath.trim()) throw new Error("Reranking model path is not configured.");
  if (!existsSync(config.binaryPath)) throw new Error(`llama-server binary does not exist: ${config.binaryPath}`);
  accessSync(config.binaryPath, constants.X_OK);

  const { modelFile, mmprojFile } = resolveGgufModelFiles(config.modelPath);
  if (!/q8[_-]?0/i.test(path.basename(modelFile))) {
    throw new Error(`Reranking model must be a Q8_0 GGUF: ${modelFile}`);
  }
  await assertPortIsAvailable(config.port);
  const child = spawn(config.binaryPath, [
    "-m", modelFile,
    "--mmproj", mmprojFile,
    "--host", "127.0.0.1",
    "--port", String(config.port),
    "--ctx-size", String(config.contextSize),
    "--batch-size", String(config.batchSize),
    "--ubatch-size", String(config.ubatchSize),
    "-ngl", String(config.gpuLayers),
    "--reranking",
    "--no-webui",
  ], { stdio: ["ignore", "ignore", "ignore"] });

  if (!child.pid) throw new Error("Failed to start the reranking service.");
  try {
    onStatus("Loading reranking model...");
    await waitForHealth(config.port);
    const scores: Array<number | undefined> = [];
    for (let index = 0; index < documents.length; index++) {
      const completed = index + 1;
      onStatus(`Reranking candidates ${completed}/${documents.length} (${Math.round(completed / documents.length * 100)}%)`);
      try {
        scores.push(await rerankDocument(config.port, query, instruction, documents[index]));
      } catch (error) {
        console.warn(`[reranker] Skipping candidate ${completed}/${documents.length}:`, error);
        scores.push(undefined);
      }
    }
    return scores;
  } finally {
    await stopProcess(child, config.port);
  }
}