openai.js 36 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599600601602603604605606607608609610611612613614615616617618619620621622623624625626627628629630631632633634635636637638639640641642643644645646647648649650651652653654655656657658659660661662663664665666667668669670671672673674675676677678679680681682683684685686687688689690691692693694695696697698699700701702703704705706707708709710711712713714715716717718719720721722723724725726727728729730731732733734735736737738739740741742743744745746747748749750751752753754755756757758759760761762763764765766767768769770771772773774775776777778779780781782783784785786787788789790791792793794795796797798799800801802803804805806807808809810811812813814815816817818819820821822823824825826827828829830831832833834835836837838839840841842843844845846847848849850851852853854855856857858859
  1. /**
  2. * openai.ts - OpenAI-compatible HTTP embedding provider
  3. *
  4. * Talks to any endpoint that implements `POST /v1/embeddings` with the OpenAI
  5. * shape: request `{model, input: string|string[]}`, response
  6. * `{data: [{embedding: number[], index?: number}, ...]}`.
  7. * `index` is optional: Gemini's OpenAI-compat layer omits it and returns
  8. * embeddings in input order (session 7fcfc297). Missing index → array position.
  9. *
  10. * Used by qmd to delegate embeddings to an approved commercial HTTPS API.
  11. *
  12. * Features:
  13. * - Batches input in groups of ≤64 (configurable via QMD_EMBED_BATCH_SIZE)
  14. * - Retries 429 / 503 with exponential backoff (1s, 4s, 16s)
  15. * - 429 gets its own, larger retry budget and honours `Retry-After`
  16. * (header or `"Retry after 29s"` body text) — a shared gateway bucket
  17. * refilling in 30s must not exhaust a 1s/4s/16s schedule (i-yghj098h)
  18. * - Bulk lane (`embedBatch` with >1 input): a 429 pauses the WHOLE worker
  19. * pool for the advertised cooldown and halves in-flight concurrency
  20. * (AIMD), so reindex traffic yields the bucket to interactive callers
  21. * instead of contending head-on with them
  22. * - 4xx (non-429) → no retry, count as failure
  23. * - Circuit breaker: >50% failures in a 60s window → OPEN for 5 min;
  24. * callers receive failures and no model fallback is selected
  25. * - Per-call timeout via AbortSignal (default QMD_EMBED_TIMEOUT_MS=30000)
  26. * - Healthcheck via `GET /health` if available, else a probe embed call
  27. */
  28. import os from "node:os";
  29. // ─────────────────────────── Configuration ───────────────────────────────────
  30. /**
  31. * Default batch size — most OpenAI-compatible embedding endpoints accept up to
  32. * 2048 inputs per call but for memory and latency we cap at 64.
  33. */
  34. export const DEFAULT_BATCH_SIZE = 64;
  35. /**
  36. * Default in-flight concurrency cap for `embedBatch`. The qmd-embed-worker
  37. * exposes a 4-way semaphore (`MAX_CONCURRENT_REQUESTS=4`) and idles at
  38. * queue-depth 1.0 under sequential clients (i-fkpnar9i baseline). Defaulting
  39. * to 4 matches the worker's advertised concurrency without overshooting the
  40. * GPU. Override per-deploy via `QMD_EMBED_CONCURRENCY`. Setting to 1 reverts
  41. * to the legacy sequential dispatch.
  42. */
  43. export const DEFAULT_CONCURRENCY = 4;
  44. /**
  45. * Default per-request timeout (30 s). embeddinggemma-300M on RTX 4090 takes
  46. * <500ms per batch of 64 in practice; 30s is a safe upper bound.
  47. */
  48. export const DEFAULT_TIMEOUT_MS = 30_000;
  49. /**
  50. * Retry backoff schedule (ms) for 429/503 responses. 3 attempts total
  51. * (initial + 2 retries) — aligns with issue spec "1s/4s/16s".
  52. */
  53. export const RETRY_BACKOFFS_MS = [1_000, 4_000, 16_000];
  54. /**
  55. * Rate-limit (429) retry budget, separate from `RETRY_BACKOFFS_MS`.
  56. *
  57. * A 429 is backpressure, not a fault: the shared ai.mm.mk gateway answers
  58. * `{"detail":"Rate limit exceeded (tokens). Retry after 29s."}` with observed
  59. * retry-after values of 7s/29s/30s, while the generic schedule waits at most
  60. * 1+4+16 = 21s in total. Bulk reindex runs therefore burned all three attempts
  61. * inside one bucket refill and reported the chunk as failed (i-yghj098h).
  62. * 429s get their own attempt count and always wait at least as long as the
  63. * server asked for, capped by `RATE_LIMIT_MAX_BACKOFF_MS`.
  64. */
  65. export const DEFAULT_RATE_LIMIT_RETRIES = 5;
  66. /** Upper bound on a single rate-limit wait, so a bogus Retry-After can't hang a run. */
  67. export const RATE_LIMIT_MAX_BACKOFF_MS = 60_000;
  68. /** Fallback wait when a 429 carries no parseable Retry-After (doubles per attempt). */
  69. export const RATE_LIMIT_BASE_BACKOFF_MS = 5_000;
  70. /**
  71. * Consecutive successful bulk requests before the lane additively recovers one
  72. * unit of concurrency after a 429-triggered halving (AIMD).
  73. */
  74. export const LANE_RECOVERY_STREAK = 8;
  75. /**
  76. * Circuit breaker — flips OPEN when error rate exceeds threshold within
  77. * window. While OPEN, every call fails fast so the caller can fall back.
  78. */
  79. export const CIRCUIT_WINDOW_MS = 60_000;
  80. export const CIRCUIT_OPEN_DURATION_MS = 5 * 60_000;
  81. export const CIRCUIT_FAILURE_RATE_THRESHOLD = 0.5;
  82. export const CIRCUIT_MIN_SAMPLES = 4;
  83. /** Map an OpenAI-compat embedding row onto the input slot. Missing `index` = row order. */
  84. export function resolveOpenAIEmbeddingIndex(index, position, inputCount) {
  85. const idx = typeof index === "number" ? index : position;
  86. if (!Number.isInteger(idx) || idx < 0 || idx >= inputCount) {
  87. throw new Error(`OpenAIEmbeddingsProvider: data item index out of range (${String(index)}, expected 0..${inputCount - 1})`);
  88. }
  89. return idx;
  90. }
  91. // ─────────────────────────── Helpers ─────────────────────────────────────────
  92. function defaultSleep(ms) {
  93. return new Promise((resolve) => setTimeout(resolve, ms));
  94. }
  95. /** Parse a non-negative integer env value; `undefined` when unset/invalid. */
  96. function parseNonNegativeInt(raw) {
  97. if (raw == null || raw.trim() === "")
  98. return undefined;
  99. const n = Number.parseInt(raw, 10);
  100. return Number.isFinite(n) && n >= 0 ? n : undefined;
  101. }
  102. /**
  103. * Build the advisory `X-AI-Caller` attribution header value (Oivo ai.mm.mk
  104. * "Observability-Driven Fleet Self-Improvement" rollout, Thread 1). Format
  105. * mirrors Oivo's `cli/src/shared/aiCallerHeader.ts` EXACTLY:
  106. * site=<file:line | tool file>; mc=<machine>; sid=<session-8char>;
  107. * comp=<component>; via=<logical-label>
  108. *
  109. * Advisory metadata ONLY — never an auth/authorization input. `mc`/`sid` come
  110. * from the Oivo fleet env when qmd runs as a delegated embedder; both degrade
  111. * to `-`/short-hostname safely when qmd runs standalone. Values are sanitized
  112. * (`;` + control chars stripped, capped at 120) so a pathological env value can
  113. * never break the header grammar.
  114. */
  115. function aiCallerHeaderValue(via) {
  116. const clean = (raw, fallback) => {
  117. if (raw == null)
  118. return fallback;
  119. let out = "";
  120. for (const ch of String(raw)) {
  121. const code = ch.charCodeAt(0);
  122. out += code < 32 || code === 127 || ch === ";" ? " " : ch;
  123. }
  124. const collapsed = out.replace(/\s+/g, " ").trim();
  125. return (collapsed || fallback).slice(0, 120);
  126. };
  127. let mc = (process.env.OIVO_MACHINE_NAME ?? "").trim();
  128. if (!mc) {
  129. try {
  130. mc = os.hostname().split(".")[0] || "unknown";
  131. }
  132. catch {
  133. mc = "unknown";
  134. }
  135. }
  136. const sid = (process.env.OIVO_SESSION_ID || process.env.CLAUDE_CODE_SESSION_ID || "-").slice(0, 8);
  137. return `site=src/embedding/openai.ts; mc=${clean(mc, "unknown")}; sid=${clean(sid, "-")}; comp=qmd; via=${clean(via, "-")}`;
  138. }
  139. /**
  140. * Build the merged AbortSignal for a single HTTP attempt: combines an
  141. * external `userSignal` (from caller / withLLMSession) with a per-attempt
  142. * timeout signal. Returns the merged signal AND the timeout id so the
  143. * caller can `clearTimeout` after the attempt completes (avoids leaks).
  144. */
  145. function buildAttemptSignal(userSignal, timeoutMs) {
  146. const ctrl = new AbortController();
  147. const timeoutId = setTimeout(() => {
  148. ctrl.abort(new Error(`Request timed out after ${timeoutMs}ms`));
  149. }, timeoutMs);
  150. // Don't keep process alive just for this timer
  151. if (typeof timeoutId === "object" && timeoutId !== null && "unref" in timeoutId) {
  152. timeoutId.unref();
  153. }
  154. const onUserAbort = () => ctrl.abort(userSignal?.reason);
  155. if (userSignal) {
  156. if (userSignal.aborted) {
  157. ctrl.abort(userSignal.reason);
  158. }
  159. else {
  160. userSignal.addEventListener("abort", onUserAbort, { once: true });
  161. }
  162. }
  163. const cleanup = () => {
  164. clearTimeout(timeoutId);
  165. if (userSignal)
  166. userSignal.removeEventListener("abort", onUserAbort);
  167. };
  168. return { signal: ctrl.signal, cleanup };
  169. }
  170. /**
  171. * Determine whether an HTTP status is retryable. 429 (Too Many Requests)
  172. * and 503 (Service Unavailable) are retried; 4xx (other than 429) are not.
  173. */
  174. export function isRetryableStatus(status) {
  175. return status === 429 || status === 503;
  176. }
  177. /**
  178. * Extract the server-advertised cooldown from a rate-limited response.
  179. *
  180. * Two sources, in priority order:
  181. * 1. the standard `Retry-After` header — delta-seconds or an HTTP-date;
  182. * 2. the ai.mm.mk body text, which states the cooldown in prose only:
  183. * `{"detail":"Rate limit exceeded (tokens). Retry after 29s.", ...}`.
  184. *
  185. * Returns `undefined` when neither source yields a sane positive duration, so
  186. * the caller falls back to its own schedule. Values are clamped to
  187. * `RATE_LIMIT_MAX_BACKOFF_MS`.
  188. */
  189. export function parseRetryAfterMs(headerValue, bodyPreview, now = Date.now) {
  190. const clamp = (ms) => Number.isFinite(ms) && ms > 0 ? Math.min(ms, RATE_LIMIT_MAX_BACKOFF_MS) : undefined;
  191. const header = headerValue?.trim();
  192. if (header) {
  193. if (/^\d+(\.\d+)?$/.test(header)) {
  194. const fromSeconds = clamp(Number(header) * 1000);
  195. if (fromSeconds !== undefined)
  196. return fromSeconds;
  197. }
  198. const asDate = Date.parse(header);
  199. if (!Number.isNaN(asDate)) {
  200. const fromDate = clamp(asDate - now());
  201. if (fromDate !== undefined)
  202. return fromDate;
  203. }
  204. }
  205. if (bodyPreview) {
  206. const m = /retry\s+after\s+(\d+(?:\.\d+)?)\s*(ms|s|seconds?)?/i.exec(bodyPreview);
  207. if (m) {
  208. const value = Number(m[1]);
  209. const unit = (m[2] ?? "s").toLowerCase();
  210. return clamp(unit === "ms" ? value : value * 1000);
  211. }
  212. }
  213. return undefined;
  214. }
  215. /**
  216. * Chunk an array into pieces of ≤ size each. `size` MUST be ≥ 1.
  217. */
  218. export function chunkArray(items, size) {
  219. if (size < 1)
  220. throw new Error(`chunkArray: size must be ≥ 1, got ${size}`);
  221. if (items.length <= size)
  222. return items.length === 0 ? [] : [items];
  223. const out = [];
  224. for (let i = 0; i < items.length; i += size) {
  225. out.push(items.slice(i, i + size));
  226. }
  227. return out;
  228. }
  229. // ─────────────────────────── Circuit Breaker ─────────────────────────────────
  230. /**
  231. * Sliding-window circuit breaker. Tracks the last N samples (min 4) over a
  232. * 60-second window; flips OPEN when failure rate exceeds 50%, then auto-
  233. * resets to HALF-OPEN after 5 minutes — at which point the next probe
  234. * decides whether to close (success) or re-open (failure).
  235. */
  236. export class CircuitBreaker {
  237. samples = [];
  238. state = "closed";
  239. openedAt = null;
  240. windowMs;
  241. openDurationMs;
  242. threshold;
  243. minSamples;
  244. now;
  245. constructor(opts = {}) {
  246. this.windowMs = opts.windowMs ?? CIRCUIT_WINDOW_MS;
  247. this.openDurationMs = opts.openDurationMs ?? CIRCUIT_OPEN_DURATION_MS;
  248. this.threshold = opts.threshold ?? CIRCUIT_FAILURE_RATE_THRESHOLD;
  249. this.minSamples = opts.minSamples ?? CIRCUIT_MIN_SAMPLES;
  250. this.now = opts.now ?? Date.now;
  251. }
  252. getState() {
  253. this.tickAutoReset();
  254. return this.state;
  255. }
  256. /**
  257. * Returns true when calls should be short-circuited (skip HTTP, fall back).
  258. * Side-effects: may transition OPEN → HALF-OPEN if the open window expired.
  259. */
  260. shouldFailFast() {
  261. return this.getState() === "open";
  262. }
  263. /** Record a successful call. */
  264. recordSuccess() {
  265. // Honor the time-based OPEN→HALF-OPEN transition before deciding what
  266. // to do with this sample. Without this, a success that lands AFTER the
  267. // open window expired would still see state==="open" and never close
  268. // the breaker (a probe call could only flip it via getState()).
  269. this.tickAutoReset();
  270. this.pushSample(true);
  271. if (this.state === "half-open") {
  272. this.state = "closed";
  273. this.openedAt = null;
  274. }
  275. }
  276. /** Record a failed call. May trigger OPEN. */
  277. recordFailure() {
  278. // Same reasoning as recordSuccess — apply lazy auto-reset before
  279. // classifying the sample.
  280. this.tickAutoReset();
  281. this.pushSample(false);
  282. if (this.state === "half-open") {
  283. // Probe failed — re-open
  284. this.state = "open";
  285. this.openedAt = this.now();
  286. return;
  287. }
  288. if (this.state === "closed")
  289. this.evaluate();
  290. }
  291. /** Force-reset the breaker (used by tests / admin) */
  292. reset() {
  293. this.samples = [];
  294. this.state = "closed";
  295. this.openedAt = null;
  296. }
  297. pushSample(ok) {
  298. const ts = this.now();
  299. this.samples.push({ ts, ok });
  300. // Drop samples outside the window
  301. const cutoff = ts - this.windowMs;
  302. while (this.samples.length > 0 && this.samples[0].ts < cutoff) {
  303. this.samples.shift();
  304. }
  305. }
  306. evaluate() {
  307. if (this.samples.length < this.minSamples)
  308. return;
  309. const failures = this.samples.filter((s) => !s.ok).length;
  310. const rate = failures / this.samples.length;
  311. if (rate > this.threshold) {
  312. this.state = "open";
  313. this.openedAt = this.now();
  314. }
  315. }
  316. tickAutoReset() {
  317. if (this.state === "open" && this.openedAt !== null) {
  318. if (this.now() - this.openedAt >= this.openDurationMs) {
  319. this.state = "half-open";
  320. }
  321. }
  322. }
  323. }
  324. // ─────────────────────────── Errors ──────────────────────────────────────────
  325. /**
  326. * Raised when the circuit breaker is OPEN and a call is short-circuited.
  327. */
  328. export class CircuitOpenError extends Error {
  329. constructor(message = "OpenAIEmbeddingsProvider circuit is OPEN") {
  330. super(message);
  331. this.name = "CircuitOpenError";
  332. }
  333. }
  334. /**
  335. * Persistent (non-retryable) HTTP error from upstream. Includes status code.
  336. */
  337. export class HttpError extends Error {
  338. status;
  339. bodyPreview;
  340. /** Server-advertised cooldown for 429s, when the response stated one. */
  341. retryAfterMs;
  342. constructor(status, bodyPreview, retryAfterMs) {
  343. super(`HTTP ${status}: ${bodyPreview.slice(0, 200)}`);
  344. this.name = "HttpError";
  345. this.status = status;
  346. this.bodyPreview = bodyPreview.slice(0, 1024);
  347. this.retryAfterMs = retryAfterMs;
  348. }
  349. }
  350. // ─────────────────────────── Bulk lane ───────────────────────────────────────
  351. /**
  352. * Client-side bulk lane for reindex traffic (i-yghj098h).
  353. *
  354. * qmd's bulk embedding shares the interactive ai.mm.mk token bucket. Without a
  355. * server-side per-caller budget, the only way bulk traffic can stop contending
  356. * head-on with interactive callers is to police itself:
  357. *
  358. * - a 429 on ANY worker pauses the ENTIRE pool for the advertised cooldown
  359. * (one shared promise — concurrent 429s coalesce instead of stacking N
  360. * cooldowns), so the bucket refills for interactive callers rather than
  361. * being re-drained by the remaining workers;
  362. * - the in-flight cap halves on each cooldown (floor 1) and recovers one unit
  363. * per `LANE_RECOVERY_STREAK` successes — classic AIMD, so a run settles at
  364. * whatever share the bucket actually has spare.
  365. *
  366. * Deliberately clock-free: cooldowns are modelled as a promise produced by the
  367. * injected `sleep`, so tests drive them with a fake sleep and no fake clock.
  368. */
  369. export class BulkLaneGate {
  370. maxPermits;
  371. sleep;
  372. permits;
  373. inFlight = 0;
  374. okStreak = 0;
  375. cooldown = null;
  376. waiters = [];
  377. constructor(maxPermits, sleep) {
  378. this.maxPermits = Math.max(1, maxPermits);
  379. this.permits = this.maxPermits;
  380. this.sleep = sleep;
  381. }
  382. /** Current in-flight cap — exported state for tests/diagnostics. */
  383. get permitCount() {
  384. return this.permits;
  385. }
  386. /** True while the lane is serving a rate-limit cooldown. */
  387. get isCoolingDown() {
  388. return this.cooldown !== null;
  389. }
  390. /** Take a slot, waiting out any cooldown and respecting the current cap. */
  391. async acquire() {
  392. while (this.cooldown !== null || this.inFlight >= this.permits) {
  393. if (this.cooldown !== null) {
  394. await this.cooldown;
  395. continue;
  396. }
  397. await new Promise((resolve) => this.waiters.push(resolve));
  398. }
  399. this.inFlight++;
  400. }
  401. /** Return a slot. Always call from a `finally`. */
  402. release() {
  403. if (this.inFlight > 0)
  404. this.inFlight--;
  405. this.wakeOne();
  406. }
  407. /**
  408. * Enter (or join) a cooldown of `waitMs` and halve the in-flight cap.
  409. * Returns the shared cooldown promise — the caller awaits it INSTEAD of
  410. * sleeping itself, so a burst of 429s costs one cooldown, not one each.
  411. */
  412. penalize(waitMs) {
  413. if (this.cooldown !== null)
  414. return this.cooldown;
  415. this.okStreak = 0;
  416. this.permits = Math.max(1, Math.floor(this.permits / 2));
  417. const cooldown = this.sleep(waitMs).then(() => {
  418. if (this.cooldown === cooldown)
  419. this.cooldown = null;
  420. this.wakeAll();
  421. });
  422. this.cooldown = cooldown;
  423. return cooldown;
  424. }
  425. /** Record a successful bulk request; recovers one permit per success streak. */
  426. noteSuccess() {
  427. if (this.permits >= this.maxPermits)
  428. return;
  429. if (++this.okStreak >= LANE_RECOVERY_STREAK) {
  430. this.okStreak = 0;
  431. this.permits++;
  432. this.wakeOne();
  433. }
  434. }
  435. wakeOne() {
  436. this.waiters.shift()?.();
  437. }
  438. wakeAll() {
  439. const pending = this.waiters;
  440. this.waiters = [];
  441. for (const resolve of pending)
  442. resolve();
  443. }
  444. }
  445. // ─────────────────────────── Provider ────────────────────────────────────────
  446. export class OpenAIEmbeddingsProvider {
  447. kind = "openai";
  448. endpoint;
  449. apiKey;
  450. modelId;
  451. upstreamModel;
  452. batchSize;
  453. concurrency;
  454. timeoutMs;
  455. fetchImpl;
  456. retryBackoffsMs;
  457. rateLimitRetries;
  458. sleep;
  459. now;
  460. dimensions = undefined;
  461. lastError = undefined;
  462. breaker;
  463. /** Shared bulk-traffic lane — see `BulkLaneGate`. */
  464. lane;
  465. constructor(config) {
  466. if (!config.endpoint) {
  467. throw new Error("OpenAIEmbeddingsProvider: endpoint is required");
  468. }
  469. this.endpoint = config.endpoint.replace(/\/+$/, "");
  470. this.apiKey = config.apiKey;
  471. this.modelId = config.modelId ?? "embeddinggemma";
  472. this.upstreamModel = config.upstreamModel ?? this.modelId;
  473. this.batchSize = config.batchSize ?? DEFAULT_BATCH_SIZE;
  474. this.concurrency = config.concurrency ?? DEFAULT_CONCURRENCY;
  475. this.timeoutMs = config.timeoutMs ?? DEFAULT_TIMEOUT_MS;
  476. this.fetchImpl = config.fetchImpl ?? globalThis.fetch;
  477. this.retryBackoffsMs = config.retryBackoffsMs ?? RETRY_BACKOFFS_MS;
  478. this.rateLimitRetries =
  479. config.rateLimitRetries ??
  480. parseNonNegativeInt(process.env.QMD_EMBED_RATE_LIMIT_RETRIES) ??
  481. DEFAULT_RATE_LIMIT_RETRIES;
  482. this.sleep = config.sleep ?? defaultSleep;
  483. this.now = config.now ?? Date.now;
  484. this.breaker = new CircuitBreaker({ now: this.now });
  485. this.lane = new BulkLaneGate(this.concurrency, this.sleep);
  486. if (!this.fetchImpl) {
  487. throw new Error("OpenAIEmbeddingsProvider: global fetch is unavailable. " +
  488. "Provide a `fetchImpl` config option (Node ≥18 ships fetch by default).");
  489. }
  490. if (this.batchSize < 1) {
  491. throw new Error(`OpenAIEmbeddingsProvider: batchSize must be ≥ 1, got ${this.batchSize}`);
  492. }
  493. if (this.concurrency < 1) {
  494. throw new Error(`OpenAIEmbeddingsProvider: concurrency must be ≥ 1, got ${this.concurrency}`);
  495. }
  496. }
  497. getModelId() {
  498. return this.modelId;
  499. }
  500. getDimensions() {
  501. return this.dimensions;
  502. }
  503. /**
  504. * Most recent per-chunk failure message (HTTP status + body preview, malformed
  505. * JSON, timeout, abort reason). Returns `undefined` after a successful call
  506. * or before the first call. See `EmbeddingProvider.getLastError`.
  507. */
  508. getLastError() {
  509. return this.lastError;
  510. }
  511. /** Endpoint URL configured at construction time — used by callers when
  512. * building error messages for failed first-chunk probes. */
  513. getEndpoint() {
  514. return this.endpoint;
  515. }
  516. async healthcheck(signal) {
  517. // Try GET /health first (worker exposes it). Fall back to probe embed.
  518. try {
  519. const { signal: attemptSig, cleanup } = buildAttemptSignal(signal, this.timeoutMs);
  520. try {
  521. const resp = await this.fetchImpl(`${this.endpoint}/health`, {
  522. method: "GET",
  523. headers: this.buildHeaders(),
  524. signal: attemptSig,
  525. });
  526. if (resp.ok) {
  527. return {
  528. ok: true,
  529. model: this.modelId,
  530. dimensions: this.dimensions,
  531. detail: `GET /health → ${resp.status}`,
  532. };
  533. }
  534. return {
  535. ok: false,
  536. model: this.modelId,
  537. detail: `GET /health → HTTP ${resp.status}`,
  538. };
  539. }
  540. finally {
  541. cleanup();
  542. }
  543. }
  544. catch (err) {
  545. // Endpoint may not implement /health — try a single embed probe instead.
  546. try {
  547. const probe = await this.embed("healthcheck", { signal });
  548. if (probe) {
  549. return {
  550. ok: true,
  551. model: this.modelId,
  552. dimensions: probe.embedding.length,
  553. detail: "embed probe ok",
  554. };
  555. }
  556. return {
  557. ok: false,
  558. model: this.modelId,
  559. detail: "embed probe returned null",
  560. };
  561. }
  562. catch (probeErr) {
  563. return {
  564. ok: false,
  565. model: this.modelId,
  566. detail: (err instanceof Error ? err.message : String(err)) +
  567. " | probe: " +
  568. (probeErr instanceof Error ? probeErr.message : String(probeErr)),
  569. };
  570. }
  571. }
  572. }
  573. async embed(text, options = {}) {
  574. const batch = await this.embedBatch([text], options);
  575. return batch[0] ?? null;
  576. }
  577. async embedBatch(texts, options = {}) {
  578. if (texts.length === 0)
  579. return [];
  580. if (this.breaker.shouldFailFast()) {
  581. throw new CircuitOpenError();
  582. }
  583. // Bulk = anything that isn't a single query-time embed. Reindex runs go
  584. // through the self-throttling lane and are attributed as `embeddings-bulk`;
  585. // `embed()` (search path) stays interactive and is never parked behind a
  586. // bulk cooldown (i-yghj098h).
  587. const lane = texts.length > 1 ? "bulk" : "interactive";
  588. const chunks = chunkArray(texts, this.batchSize);
  589. const results = new Array(texts.length).fill(null);
  590. // Pre-compute the input-array starting position for each chunk so each
  591. // worker can write its slice of `results` independently — input order is
  592. // preserved end-to-end without a final re-sort step.
  593. const chunkStarts = new Array(chunks.length);
  594. {
  595. let cursor = 0;
  596. for (let i = 0; i < chunks.length; i++) {
  597. chunkStarts[i] = cursor;
  598. cursor += chunks[i].length;
  599. }
  600. }
  601. // Shared state across the worker pool. Each transition is final-write,
  602. // so plain JS scalars are safe — no atomics or locks needed since
  603. // workers only contend on these via cooperative-scheduled awaits.
  604. let nextChunkIdx = 0;
  605. let anySucceeded = false;
  606. let aborted = false;
  607. let circuitTrippedDuringRun = null;
  608. // Workers run as parallel async tasks pulling chunks off `nextChunkIdx`
  609. // until the queue is drained or one of the early-exit flags is set.
  610. // Concurrency is capped at min(this.concurrency, chunks.length) so we
  611. // don't spin up idle workers for tiny inputs.
  612. const workerCount = Math.min(this.concurrency, chunks.length);
  613. const dispatchOne = async () => {
  614. while (true) {
  615. if (aborted || circuitTrippedDuringRun)
  616. return;
  617. const idx = nextChunkIdx++;
  618. if (idx >= chunks.length)
  619. return;
  620. const chunk = chunks[idx];
  621. const start = chunkStarts[idx];
  622. // Honor abort/breaker BEFORE issuing the request so we don't waste
  623. // network for a dispatch we know will be discarded.
  624. if (options.signal?.aborted) {
  625. aborted = true;
  626. this.lastError = `aborted by caller${options.signal.reason ? `: ${String(options.signal.reason)}` : ""}`;
  627. return;
  628. }
  629. if (this.breaker.shouldFailFast()) {
  630. // Capture the breaker-open intent so we throw it AFTER all
  631. // currently in-flight workers settle, instead of leaking
  632. // half-completed results. The thrown error is a fresh instance
  633. // (matching legacy behavior).
  634. circuitTrippedDuringRun = new CircuitOpenError();
  635. return;
  636. }
  637. try {
  638. if (lane === "bulk")
  639. await this.lane.acquire();
  640. let embeddings;
  641. try {
  642. embeddings = await this.requestWithRetry(chunk, options, lane);
  643. }
  644. finally {
  645. if (lane === "bulk")
  646. this.lane.release();
  647. }
  648. if (lane === "bulk")
  649. this.lane.noteSuccess();
  650. for (let i = 0; i < chunk.length; i++) {
  651. const embedding = embeddings[i];
  652. if (embedding) {
  653. results[start + i] = {
  654. embedding,
  655. model: this.modelId,
  656. };
  657. anySucceeded = true;
  658. // Record dimensions on first success. Concurrent workers may
  659. // race on this assignment, but they all observe the same
  660. // length so the race is benign.
  661. if (this.dimensions === undefined) {
  662. this.dimensions = embedding.length;
  663. }
  664. }
  665. }
  666. this.breaker.recordSuccess();
  667. }
  668. catch (err) {
  669. // A rate limit is backpressure, not a fault. Counting 429s toward the
  670. // breaker turns "the bucket is empty for 30s" into a 5-minute hard
  671. // OPEN, which is the outage the retry budget above exists to avoid.
  672. // The lane cooldown is already the correct control response.
  673. const rateLimited = err instanceof HttpError && err.status === 429;
  674. if (!rateLimited)
  675. this.breaker.recordFailure();
  676. if (err instanceof CircuitOpenError) {
  677. circuitTrippedDuringRun = err;
  678. return;
  679. }
  680. // Last-write-wins on lastError matches the legacy semantics — under
  681. // concurrency multiple workers may fail in the same call, but the
  682. // lastError just needs to surface "the most recent cause."
  683. this.lastError = this.formatErrorContext(err);
  684. if (process.env.QMD_EMBED_DEBUG) {
  685. process.stderr.write(`OpenAIEmbeddingsProvider: chunk failed (${err instanceof Error ? err.message : String(err)})\n`);
  686. }
  687. }
  688. }
  689. };
  690. await Promise.all(Array.from({ length: workerCount }, () => dispatchOne()));
  691. // If a worker observed `shouldFailFast()` mid-run, surface the error
  692. // after all in-flight workers have settled.
  693. if (circuitTrippedDuringRun)
  694. throw circuitTrippedDuringRun;
  695. // Clear lastError on a fully-successful sweep (every input got an embedding).
  696. if (anySucceeded && results.every((r) => r !== null)) {
  697. this.lastError = undefined;
  698. }
  699. return results;
  700. }
  701. async dispose() {
  702. // Nothing to release — fetch handles its own connection pooling.
  703. // Reset the breaker so a re-instantiation starts fresh.
  704. this.breaker.reset();
  705. }
  706. // ────────────────────── Internals ──────────────────────
  707. /**
  708. * Format a request-failure context string for `lastError`. Includes endpoint
  709. * + HTTP status + body preview when the error was an `HttpError`, otherwise
  710. * falls back to the message of the underlying error (or the value itself
  711. * when not an Error). Kept short — body preview is already capped at 1024
  712. * chars by `HttpError`, but we trim further here for the dimension-probe
  713. * thrown error which surfaces directly to users.
  714. */
  715. formatErrorContext(err) {
  716. if (err instanceof HttpError) {
  717. const preview = err.bodyPreview.replace(/\s+/g, " ").trim().slice(0, 240);
  718. return `endpoint=${this.endpoint}/v1/embeddings status=${err.status}${preview ? ` body="${preview}"` : ""}`;
  719. }
  720. if (err instanceof Error) {
  721. return `endpoint=${this.endpoint}/v1/embeddings error="${err.message}"`;
  722. }
  723. return `endpoint=${this.endpoint}/v1/embeddings error="${String(err)}"`;
  724. }
  725. buildHeaders(lane = "interactive") {
  726. const headers = {
  727. "Content-Type": "application/json",
  728. "Accept": "application/json",
  729. // Advisory caller attribution for the ai.mm.mk gateway (NULL-safe, never
  730. // auth). Single chokepoint for both /health and /v1/embeddings.
  731. // `via` distinguishes bulk reindex traffic from query-time embeds so the
  732. // gateway can budget the two separately (i-yghj098h part 1b) — the
  733. // gateway-side per-caller bucket keys off this label.
  734. "X-AI-Caller": aiCallerHeaderValue(lane === "bulk" ? "embeddings-bulk" : "embeddings"),
  735. };
  736. if (this.apiKey) {
  737. headers["Authorization"] = `Bearer ${this.apiKey}`;
  738. }
  739. return headers;
  740. }
  741. /**
  742. * Single HTTP request with retry on 429/503. Returns embeddings indexed
  743. * the same as `texts`. Throws on non-retryable failure or all attempts
  744. * exhausted.
  745. */
  746. async requestWithRetry(texts, options, lane = "interactive") {
  747. let lastErr = null;
  748. // 429s draw on their own budget: a token bucket that refills in 30s must
  749. // not be able to exhaust a schedule whose total wait is 21s (i-yghj098h).
  750. let backoffAttempt = 0;
  751. let rateLimitAttempt = 0;
  752. for (;;) {
  753. // Honor user abort BEFORE issuing the call (avoids wasted network)
  754. if (options.signal?.aborted) {
  755. throw new Error("aborted by caller");
  756. }
  757. try {
  758. return await this.requestOnce(texts, options, lane);
  759. }
  760. catch (err) {
  761. lastErr = err;
  762. const status = err instanceof HttpError ? err.status : 0;
  763. if (!isRetryableStatus(status))
  764. throw err;
  765. if (status === 429) {
  766. if (rateLimitAttempt >= this.rateLimitRetries)
  767. throw err;
  768. const waitMs = this.rateLimitWaitMs(err, rateLimitAttempt);
  769. rateLimitAttempt++;
  770. if (lane === "bulk") {
  771. // Pause the whole pool, not just this worker: the other workers
  772. // hammering the same empty bucket is exactly what turns one 429
  773. // into a cascade, and what starves interactive callers.
  774. await this.lane.penalize(waitMs);
  775. }
  776. else {
  777. await this.sleep(waitMs);
  778. }
  779. continue;
  780. }
  781. if (backoffAttempt >= this.retryBackoffsMs.length)
  782. break;
  783. await this.sleep(this.retryBackoffsMs[backoffAttempt]);
  784. backoffAttempt++;
  785. }
  786. }
  787. // Exhausted retries → throw the last error so caller marks the chunk null
  788. throw lastErr ?? new Error("requestWithRetry exhausted");
  789. }
  790. /**
  791. * How long to wait after a 429. The server's own `Retry-After` wins when it
  792. * gave one; otherwise fall back to the configured schedule (so injected test
  793. * schedules stay authoritative) and then to an exponential 5s/10s/20s… ramp.
  794. * Always clamped to `RATE_LIMIT_MAX_BACKOFF_MS`.
  795. */
  796. rateLimitWaitMs(err, rateLimitAttempt) {
  797. const scheduled = this.retryBackoffsMs[Math.min(rateLimitAttempt, this.retryBackoffsMs.length - 1)];
  798. const fallback = scheduled ?? RATE_LIMIT_BASE_BACKOFF_MS * Math.pow(2, rateLimitAttempt);
  799. return Math.min(Math.max(err.retryAfterMs ?? 0, fallback), RATE_LIMIT_MAX_BACKOFF_MS);
  800. }
  801. /**
  802. * Issue one HTTP attempt to `POST /v1/embeddings`. Does NOT retry.
  803. */
  804. async requestOnce(texts, options, lane = "interactive") {
  805. const { signal: attemptSig, cleanup } = buildAttemptSignal(options.signal, this.timeoutMs);
  806. try {
  807. const body = JSON.stringify({
  808. model: options.model ?? this.upstreamModel,
  809. input: texts,
  810. });
  811. const resp = await this.fetchImpl(`${this.endpoint}/v1/embeddings`, {
  812. method: "POST",
  813. headers: this.buildHeaders(lane),
  814. body,
  815. signal: attemptSig,
  816. });
  817. if (!resp.ok) {
  818. const text = await resp.text().catch(() => "");
  819. // ai.mm.mk states the cooldown in the JSON body, not always in a
  820. // Retry-After header — read both (i-yghj098h).
  821. const retryAfterMs = resp.status === 429
  822. ? parseRetryAfterMs(resp.headers?.get?.("retry-after"), text, this.now)
  823. : undefined;
  824. throw new HttpError(resp.status, text, retryAfterMs);
  825. }
  826. let parsed;
  827. try {
  828. parsed = (await resp.json());
  829. }
  830. catch (err) {
  831. throw new Error(`OpenAIEmbeddingsProvider: malformed JSON from ${this.endpoint}/v1/embeddings: ${err instanceof Error ? err.message : String(err)}`);
  832. }
  833. if (!parsed || !Array.isArray(parsed.data)) {
  834. throw new Error(`OpenAIEmbeddingsProvider: response missing "data" array (got ${typeof parsed})`);
  835. }
  836. // Prefer `index` when present (out-of-order servers). Gemini's
  837. // OpenAI-compat layer omits `index` and returns rows in input order.
  838. const out = new Array(texts.length);
  839. for (let i = 0; i < parsed.data.length; i++) {
  840. const item = parsed.data[i];
  841. const idx = resolveOpenAIEmbeddingIndex(item.index, i, texts.length);
  842. if (!Array.isArray(item.embedding)) {
  843. throw new Error(`OpenAIEmbeddingsProvider: data[${idx}].embedding is not an array`);
  844. }
  845. out[idx] = item.embedding;
  846. }
  847. // Sanity check — every slot must be filled
  848. for (let i = 0; i < texts.length; i++) {
  849. if (!out[i]) {
  850. throw new Error(`OpenAIEmbeddingsProvider: response missing embedding for index ${i}`);
  851. }
  852. }
  853. return out;
  854. }
  855. finally {
  856. cleanup();
  857. }
  858. }
  859. }