jina.ts

/**
 * jina.ts - Jina AI remote provider for QMD
 *
 * HTTP clients for the Jina AI embeddings and reranker APIs. Used as optional
 * remote backends instead of the local node-llama-cpp models.
 *
 * Embedding:
 *   QMD_EMBED_PROVIDER=jina
 *   JINA_API_KEY=jina_xxx...
 *   QMD_JINA_MODEL        (default: jina-embeddings-v3)
 *   QMD_JINA_DIMENSION    (default: 1024 — Matryoshka truncation supported on v3)
 *
 * Reranking:
 *   QMD_RERANK_PROVIDER=jina
 *   JINA_API_KEY=jina_xxx...  (shared with embeddings)
 *   QMD_JINA_RERANK_MODEL (default: jina-reranker-v2-base-multilingual)
 *
 * Shared:
 *   QMD_JINA_BASE_URL     (default: https://api.jina.ai/v1)
 *   QMD_JINA_BATCH        (default: 128 — inputs per embed request)
 *   QMD_JINA_CONCURRENCY  (default: 4 — parallel requests in a batch)
 *   QMD_JINA_TIMEOUT_MS   (default: 60000)
 *   QMD_JINA_MAX_RETRIES  (default: 4)
 *
 * Security:
 *   The API key is read from the environment only — never from config files.
 *   Callers must not log or persist it.
 *
 * Usage tracking:
 *   Every successful API call reports token counts via an optional callback
 *   (see `UsageReporter`). QMD wires this to a SQLite `jina_usage` table so
 *   users can monitor their quota with `qmd usage`.
 */

// =============================================================================
// Types
// =============================================================================

/** Jina task types — affect how the model encodes inputs. */
export type JinaTask =
  | "retrieval.query"
  | "retrieval.passage"
  | "text-matching"
  | "classification"
  | "separation";

/**
 * Operation kinds tracked in usage reporting.
 * `embed_query` and `embed_passage` are separated so users can tell which
 * workload consumed the quota (indexing vs querying).
 */
export type JinaOperation = "embed_query" | "embed_passage" | "rerank";

/**
 * Called after every successful Jina API call. Implementations should be
 * cheap and non-throwing — errors in the reporter must not bubble back into
 * the request path. QMD wires this to a SQLite INSERT.
 */
export type UsageReporter = (event: {
  operation: JinaOperation;
  model: string;
  totalTokens: number;
  promptTokens?: number;
  /** ISO 8601 timestamp */
  at: string;
}) => void;

export type JinaEmbedderConfig = {
  apiKey: string;
  model?: string;
  dimensions?: number;
  baseUrl?: string;
  /** Max inputs sent in a single HTTP request. Jina accepts up to 2048 but smaller batches give better retry granularity. */
  batchSize?: number;
  /** Number of in-flight requests when a user batch is split across multiple HTTP calls. */
  concurrency?: number;
  /** Per-request timeout in milliseconds. */
  timeoutMs?: number;
  /** Max retry attempts on 429/5xx/network errors. */
  maxRetries?: number;
  /** Optional callback invoked after each successful API response. */
  usageReporter?: UsageReporter;
};

export type JinaRerankerConfig = {
  apiKey: string;
  model?: string;
  baseUrl?: string;
  /** Per-request timeout in milliseconds. */
  timeoutMs?: number;
  /** Max retry attempts on 429/5xx/network errors. */
  maxRetries?: number;
  usageReporter?: UsageReporter;
};

export type JinaRerankResult = {
  /** Index into the original documents array. */
  index: number;
  /** Relevance score in [0, 1]; higher is more relevant. */
  score: number;
};

type JinaApiResponse = {
  model: string;
  object: string;
  usage?: { total_tokens?: number; prompt_tokens?: number };
  data: Array<{
    index: number;
    embedding: number[];
    object?: string;
  }>;
};

type JinaRerankApiResponse = {
  model: string;
  usage?: { total_tokens?: number; prompt_tokens?: number };
  results: Array<{
    index: number;
    relevance_score: number;
    document?: { text?: string };
  }>;
};

type JinaErrorResponse = {
  detail?: string | { msg?: string }[];
  error?: string | { message?: string };
  message?: string;
};

// =============================================================================
// Defaults
// =============================================================================

const DEFAULT_MODEL = "jina-embeddings-v3";
const DEFAULT_RERANK_MODEL = "jina-reranker-v2-base-multilingual";
const DEFAULT_DIMENSIONS = 1024;
const DEFAULT_BASE_URL = "https://api.jina.ai/v1";
const DEFAULT_BATCH_SIZE = 128;
const DEFAULT_CONCURRENCY = 4;
const DEFAULT_TIMEOUT_MS = 60_000;
const DEFAULT_MAX_RETRIES = 4;

// Jina v3 supports up to 8192 tokens per input. We report this via maxInputTokens().
const JINA_V3_MAX_INPUT_TOKENS = 8192;
// Jina reranker v2 max input tokens per document.
const JINA_RERANKER_V2_MAX_INPUT_TOKENS = 8192;

// =============================================================================
// Helpers
// =============================================================================

function parseIntEnv(name: string, fallback: number): number {
  const raw = process.env[name];
  if (!raw) return fallback;
  const parsed = Number.parseInt(raw.trim(), 10);
  if (!Number.isFinite(parsed) || parsed <= 0) {
    process.stderr.write(`QMD Warning: invalid ${name}="${raw}", using default ${fallback}.\n`);
    return fallback;
  }
  return parsed;
}

function extractErrorMessage(body: unknown): string {
  if (!body || typeof body !== "object") return "";
  const e = body as JinaErrorResponse;
  if (typeof e.message === "string") return e.message;
  if (typeof e.error === "string") return e.error;
  if (e.error && typeof e.error === "object" && typeof e.error.message === "string") {
    return e.error.message;
  }
  if (typeof e.detail === "string") return e.detail;
  if (Array.isArray(e.detail) && e.detail[0]?.msg) return e.detail[0].msg;
  return "";
}

/** Sleep with jitter for backoff. */
function sleep(ms: number): Promise<void> {
  return new Promise((resolve) => setTimeout(resolve, ms));
}

/** Compute exponential backoff with jitter, capped at 30s. */
function backoffMs(attempt: number): number {
  const base = Math.min(30_000, 500 * Math.pow(2, attempt));
  const jitter = Math.random() * 0.3 * base;
  return Math.floor(base + jitter);
}

// =============================================================================
// JinaEmbedder
// =============================================================================

/**
 * Remote embedding client for the Jina AI API.
 *
 * Features:
 *   - Batches inputs into chunks of `batchSize` and runs up to `concurrency` requests in parallel.
 *   - Retries on 429/5xx/network errors with exponential backoff + jitter.
 *   - Respects `Retry-After` header when present.
 *   - Returns a nullable result array, matching the local `embedBatch` contract.
 */
export class JinaEmbedder {
  private readonly apiKey: string;
  readonly model: string;
  readonly dimensions: number;
  private readonly baseUrl: string;
  private readonly batchSize: number;
  private readonly concurrency: number;
  private readonly timeoutMs: number;
  private readonly maxRetries: number;
  private usageReporter: UsageReporter | null;

  constructor(config: JinaEmbedderConfig) {
    if (!config.apiKey) {
      throw new Error(
        "JinaEmbedder: apiKey is required (set JINA_API_KEY env var)",
      );
    }
    this.apiKey = config.apiKey;
    this.model = config.model || DEFAULT_MODEL;
    this.dimensions = config.dimensions || DEFAULT_DIMENSIONS;
    this.baseUrl = (config.baseUrl || DEFAULT_BASE_URL).replace(/\/+$/, "");
    this.batchSize = config.batchSize ?? DEFAULT_BATCH_SIZE;
    this.concurrency = config.concurrency ?? DEFAULT_CONCURRENCY;
    this.timeoutMs = config.timeoutMs ?? DEFAULT_TIMEOUT_MS;
    this.maxRetries = config.maxRetries ?? DEFAULT_MAX_RETRIES;
    this.usageReporter = config.usageReporter ?? null;
  }

  /**
   * Attach or replace the usage reporter after construction. This is used by
   * QMD at startup to wire a SQLite-backed recorder after the store is ready,
   * so the embedder itself never imports the database layer.
   */
  setUsageReporter(reporter: UsageReporter | null): void {
    this.usageReporter = reporter;
  }

  private reportUsage(operation: JinaOperation, usage: JinaApiResponse["usage"] | JinaRerankApiResponse["usage"]): void {
    if (!this.usageReporter) return;
    const total = usage?.total_tokens ?? 0;
    if (total <= 0) return;
    try {
      this.usageReporter({
        operation,
        model: this.model,
        totalTokens: total,
        promptTokens: usage?.prompt_tokens,
        at: new Date().toISOString(),
      });
    } catch {
      // Reporter must not break the request path.
    }
  }

  /**
   * Stable identifier used by QMD to tag stored vectors and to gate format helpers.
   * Kept in sync with isJinaProvider() in llm.ts.
   */
  get modelUri(): string {
    return `jina:${this.model}`;
  }

  /** Max input tokens per text (model-dependent). Used for truncation budget. */
  maxInputTokens(): number {
    return JINA_V3_MAX_INPUT_TOKENS;
  }

  /**
   * Embed a single text. Returns null on failure.
   * `isQuery` selects the task: retrieval.query vs retrieval.passage.
   */
  async embed(
    text: string,
    opts: { isQuery?: boolean } = {},
  ): Promise<number[] | null> {
    const results = await this.embedBatch([text], opts);
    return results[0] ?? null;
  }

  /**
   * Embed multiple texts. Returns an array of the same length, with `null` for
   * any input that failed to embed (preserving index alignment with the input).
   */
  async embedBatch(
    texts: string[],
    opts: { isQuery?: boolean } = {},
  ): Promise<(number[] | null)[]> {
    if (texts.length === 0) return [];

    const task: JinaTask = opts.isQuery ? "retrieval.query" : "retrieval.passage";

    // Split into HTTP request batches.
    const batches: { offset: number; inputs: string[] }[] = [];
    for (let i = 0; i < texts.length; i += this.batchSize) {
      batches.push({
        offset: i,
        inputs: texts.slice(i, i + this.batchSize),
      });
    }

    const out: (number[] | null)[] = new Array(texts.length).fill(null);

    // Run batches with bounded concurrency.
    let cursor = 0;
    const workers: Promise<void>[] = [];
    const workerCount = Math.min(this.concurrency, batches.length);

    for (let w = 0; w < workerCount; w++) {
      workers.push(
        (async () => {
          while (true) {
            const i = cursor++;
            if (i >= batches.length) return;
            const batch = batches[i]!;
            try {
              const vectors = await this.embedBatchRequest(batch.inputs, task);
              for (let j = 0; j < vectors.length; j++) {
                out[batch.offset + j] = vectors[j] ?? null;
              }
            } catch (err) {
              // Individual batch fatal — leave those slots as null, log once.
              const msg = err instanceof Error ? err.message : String(err);
              process.stderr.write(
                `QMD Jina: batch starting at ${batch.offset} failed: ${msg}\n`,
              );
            }
          }
        })(),
      );
    }

    await Promise.all(workers);
    return out;
  }

  /**
   * Low-level: send a single /embeddings request with retries.
   * Returns the vectors in the same order as inputs.
   */
  private async embedBatchRequest(
    inputs: string[],
    task: JinaTask,
  ): Promise<number[][]> {
    const op: JinaOperation =
      task === "retrieval.query" ? "embed_query" : "embed_passage";
    const url = `${this.baseUrl}/embeddings`;
    const body = JSON.stringify({
      model: this.model,
      task,
      dimensions: this.dimensions,
      embedding_type: "float",
      input: inputs,
    });

    let lastError: Error | null = null;

    for (let attempt = 0; attempt <= this.maxRetries; attempt++) {
      const controller = new AbortController();
      const timer = setTimeout(() => controller.abort(), this.timeoutMs);

      try {
        const resp = await fetch(url, {
          method: "POST",
          headers: {
            "Content-Type": "application/json",
            Authorization: `Bearer ${this.apiKey}`,
            Accept: "application/json",
          },
          body,
          signal: controller.signal,
        });

        clearTimeout(timer);

        if (resp.ok) {
          const json = (await resp.json()) as JinaApiResponse;
          if (!json || !Array.isArray(json.data)) {
            throw new Error(`Jina: malformed response (no data array)`);
          }
          // Reorder by .index to be safe.
          const ordered: number[][] = new Array(inputs.length);
          for (const item of json.data) {
            if (
              typeof item.index !== "number" ||
              item.index < 0 ||
              item.index >= inputs.length
            ) {
              throw new Error(
                `Jina: response item has invalid index ${item.index}`,
              );
            }
            if (!Array.isArray(item.embedding)) {
              throw new Error(
                `Jina: response item ${item.index} missing embedding array`,
              );
            }
            ordered[item.index] = item.embedding;
          }
          // Verify all slots filled.
          for (let i = 0; i < inputs.length; i++) {
            if (!ordered[i]) {
              throw new Error(`Jina: response missing embedding for index ${i}`);
            }
          }
          this.reportUsage(op, json.usage);
          return ordered;
        }

        // Non-2xx: decide whether to retry.
        let errText = "";
        try {
          const errJson = await resp.json();
          errText = extractErrorMessage(errJson) || JSON.stringify(errJson);
        } catch {
          try {
            errText = await resp.text();
          } catch {
            errText = resp.statusText;
          }
        }

        const retryable = resp.status === 429 || resp.status >= 500;
        if (!retryable || attempt === this.maxRetries) {
          throw new Error(
            `Jina API ${resp.status} ${resp.statusText}: ${errText || "(no body)"}`,
          );
        }

        // Honor Retry-After if present.
        const retryAfter = resp.headers.get("retry-after");
        let waitMs = backoffMs(attempt);
        if (retryAfter) {
          const n = Number.parseInt(retryAfter, 10);
          if (Number.isFinite(n) && n > 0) waitMs = Math.max(waitMs, n * 1000);
        }
        await sleep(waitMs);
        lastError = new Error(
          `Jina API ${resp.status} (retrying attempt ${attempt + 1}/${this.maxRetries})`,
        );
        continue;
      } catch (err) {
        clearTimeout(timer);
        const isAbort = err instanceof Error && err.name === "AbortError";
        const msg = err instanceof Error ? err.message : String(err);

        // Retry network errors and timeouts; bail on last attempt.
        if (attempt === this.maxRetries) {
          throw new Error(
            `Jina request failed after ${this.maxRetries + 1} attempts: ${msg}`,
          );
        }
        lastError = err instanceof Error ? err : new Error(msg);
        if (isAbort) {
          process.stderr.write(
            `QMD Jina: request timed out after ${this.timeoutMs}ms, retrying...\n`,
          );
        }
        await sleep(backoffMs(attempt));
      }
    }

    // Unreachable, but keeps TS happy.
    throw lastError ?? new Error("Jina: request failed (unknown error)");
  }
}

// =============================================================================
// JinaReranker
// =============================================================================

/**
 * Remote reranker client for the Jina AI /rerank API.
 *
 * Sends (query, documents[]) to the Jina reranker-v2 model and returns a
 * ranked list of {index, score}. Callers are responsible for truncating
 * documents to fit within the provider's max input tokens.
 *
 * Unlike embeddings, the reranker API processes the full document list in a
 * single request. We still apply the same retry logic on 429/5xx.
 */
export class JinaReranker {
  private readonly apiKey: string;
  readonly model: string;
  private readonly baseUrl: string;
  private readonly timeoutMs: number;
  private readonly maxRetries: number;
  private usageReporter: UsageReporter | null;

  constructor(config: JinaRerankerConfig) {
    if (!config.apiKey) {
      throw new Error(
        "JinaReranker: apiKey is required (set JINA_API_KEY env var)",
      );
    }
    this.apiKey = config.apiKey;
    this.model = config.model || DEFAULT_RERANK_MODEL;
    this.baseUrl = (config.baseUrl || DEFAULT_BASE_URL).replace(/\/+$/, "");
    this.timeoutMs = config.timeoutMs ?? DEFAULT_TIMEOUT_MS;
    this.maxRetries = config.maxRetries ?? DEFAULT_MAX_RETRIES;
    this.usageReporter = config.usageReporter ?? null;
  }

  /** Stable identifier for the active rerank backend. */
  get modelUri(): string {
    return `jina:${this.model}`;
  }

  /** Max input tokens per document. Used by callers for truncation budgets. */
  maxInputTokens(): number {
    return JINA_RERANKER_V2_MAX_INPUT_TOKENS;
  }

  setUsageReporter(reporter: UsageReporter | null): void {
    this.usageReporter = reporter;
  }

  private reportUsage(usage: JinaRerankApiResponse["usage"]): void {
    if (!this.usageReporter) return;
    const total = usage?.total_tokens ?? 0;
    if (total <= 0) return;
    try {
      this.usageReporter({
        operation: "rerank",
        model: this.model,
        totalTokens: total,
        promptTokens: usage?.prompt_tokens,
        at: new Date().toISOString(),
      });
    } catch {
      // Never break the request path.
    }
  }

  /**
   * Rerank `documents` by relevance to `query`.
   * Returns results sorted by descending score, with original indices.
   * On error, throws — callers decide how to handle (typically: skip rerank).
   */
  async rerank(query: string, documents: string[]): Promise<JinaRerankResult[]> {
    if (documents.length === 0) return [];

    const url = `${this.baseUrl}/rerank`;
    const body = JSON.stringify({
      model: this.model,
      query,
      documents,
      return_documents: false,
    });

    let lastError: Error | null = null;

    for (let attempt = 0; attempt <= this.maxRetries; attempt++) {
      const controller = new AbortController();
      const timer = setTimeout(() => controller.abort(), this.timeoutMs);

      try {
        const resp = await fetch(url, {
          method: "POST",
          headers: {
            "Content-Type": "application/json",
            Authorization: `Bearer ${this.apiKey}`,
            Accept: "application/json",
          },
          body,
          signal: controller.signal,
        });

        clearTimeout(timer);

        if (resp.ok) {
          const json = (await resp.json()) as JinaRerankApiResponse;
          if (!json || !Array.isArray(json.results)) {
            throw new Error(`Jina rerank: malformed response (no results array)`);
          }
          const results: JinaRerankResult[] = json.results.map((r) => ({
            index: r.index,
            score: r.relevance_score,
          }));
          // Sort by descending score (defensive — Jina already sorts but be explicit).
          results.sort((a, b) => b.score - a.score);
          this.reportUsage(json.usage);
          return results;
        }

        // Non-2xx: decide retry vs fail.
        let errText = "";
        try {
          const errJson = await resp.json();
          errText = extractErrorMessage(errJson) || JSON.stringify(errJson);
        } catch {
          try {
            errText = await resp.text();
          } catch {
            errText = resp.statusText;
          }
        }

        const retryable = resp.status === 429 || resp.status >= 500;
        if (!retryable || attempt === this.maxRetries) {
          throw new Error(
            `Jina rerank API ${resp.status} ${resp.statusText}: ${errText || "(no body)"}`,
          );
        }

        const retryAfter = resp.headers.get("retry-after");
        let waitMs = backoffMs(attempt);
        if (retryAfter) {
          const n = Number.parseInt(retryAfter, 10);
          if (Number.isFinite(n) && n > 0) waitMs = Math.max(waitMs, n * 1000);
        }
        await sleep(waitMs);
        lastError = new Error(
          `Jina rerank API ${resp.status} (retrying attempt ${attempt + 1}/${this.maxRetries})`,
        );
        continue;
      } catch (err) {
        clearTimeout(timer);
        const isAbort = err instanceof Error && err.name === "AbortError";
        const msg = err instanceof Error ? err.message : String(err);

        if (attempt === this.maxRetries) {
          throw new Error(
            `Jina rerank failed after ${this.maxRetries + 1} attempts: ${msg}`,
          );
        }
        lastError = err instanceof Error ? err : new Error(msg);
        if (isAbort) {
          process.stderr.write(
            `QMD Jina: rerank request timed out after ${this.timeoutMs}ms, retrying...\n`,
          );
        }
        await sleep(backoffMs(attempt));
      }
    }

    throw lastError ?? new Error("Jina rerank: failed (unknown error)");
  }
}

// =============================================================================
// Factory from env
// =============================================================================

/** Resolve the shared Jina API key from env, supporting both common names. */
function readJinaApiKey(): string {
  return (process.env.JINA_API_KEY ?? process.env.QMD_JINA_API_KEY ?? "").trim();
}

/**
 * Create a JinaEmbedder from environment variables, or return null if the
 * provider is not selected.
 *
 * Throws if QMD_EMBED_PROVIDER=jina but JINA_API_KEY is missing.
 */
export function jinaEmbedderFromEnv(): JinaEmbedder | null {
  const provider = (process.env.QMD_EMBED_PROVIDER ?? "").trim().toLowerCase();
  if (provider !== "jina") return null;

  const apiKey = readJinaApiKey();
  if (!apiKey) {
    throw new Error(
      "QMD_EMBED_PROVIDER=jina but JINA_API_KEY is not set. " +
        "Get a key at https://jina.ai/ and export JINA_API_KEY=jina_...",
    );
  }

  const model = process.env.QMD_JINA_MODEL?.trim() || DEFAULT_MODEL;
  const dimensions = parseIntEnv("QMD_JINA_DIMENSION", DEFAULT_DIMENSIONS);
  const baseUrl = process.env.QMD_JINA_BASE_URL?.trim() || DEFAULT_BASE_URL;
  const batchSize = parseIntEnv("QMD_JINA_BATCH", DEFAULT_BATCH_SIZE);
  const concurrency = parseIntEnv("QMD_JINA_CONCURRENCY", DEFAULT_CONCURRENCY);
  const timeoutMs = parseIntEnv("QMD_JINA_TIMEOUT_MS", DEFAULT_TIMEOUT_MS);
  const maxRetries = parseIntEnv("QMD_JINA_MAX_RETRIES", DEFAULT_MAX_RETRIES);

  return new JinaEmbedder({
    apiKey,
    model,
    dimensions,
    baseUrl,
    batchSize,
    concurrency,
    timeoutMs,
    maxRetries,
  });
}

/**
 * Create a JinaReranker from environment variables, or return null if the
 * rerank provider is not selected.
 *
 * Throws if QMD_RERANK_PROVIDER=jina but JINA_API_KEY is missing.
 */
export function jinaRerankerFromEnv(): JinaReranker | null {
  const provider = (process.env.QMD_RERANK_PROVIDER ?? "").trim().toLowerCase();
  if (provider !== "jina") return null;

  const apiKey = readJinaApiKey();
  if (!apiKey) {
    throw new Error(
      "QMD_RERANK_PROVIDER=jina but JINA_API_KEY is not set. " +
        "Get a key at https://jina.ai/ and export JINA_API_KEY=jina_...",
    );
  }

  const model = process.env.QMD_JINA_RERANK_MODEL?.trim() || DEFAULT_RERANK_MODEL;
  const baseUrl = process.env.QMD_JINA_BASE_URL?.trim() || DEFAULT_BASE_URL;
  const timeoutMs = parseIntEnv("QMD_JINA_TIMEOUT_MS", DEFAULT_TIMEOUT_MS);
  const maxRetries = parseIntEnv("QMD_JINA_MAX_RETRIES", DEFAULT_MAX_RETRIES);

  return new JinaReranker({
    apiKey,
    model,
    baseUrl,
    timeoutMs,
    maxRetries,
  });
}