openai.d.ts 14 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332
  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 type { EmbeddingProvider, ProviderEmbedOptions, ProviderEmbedding, ProviderHealth, ProviderKind } from "./provider.js";
  29. /**
  30. * Default batch size — most OpenAI-compatible embedding endpoints accept up to
  31. * 2048 inputs per call but for memory and latency we cap at 64.
  32. */
  33. export declare const DEFAULT_BATCH_SIZE = 64;
  34. /**
  35. * Default in-flight concurrency cap for `embedBatch`. The qmd-embed-worker
  36. * exposes a 4-way semaphore (`MAX_CONCURRENT_REQUESTS=4`) and idles at
  37. * queue-depth 1.0 under sequential clients (i-fkpnar9i baseline). Defaulting
  38. * to 4 matches the worker's advertised concurrency without overshooting the
  39. * GPU. Override per-deploy via `QMD_EMBED_CONCURRENCY`. Setting to 1 reverts
  40. * to the legacy sequential dispatch.
  41. */
  42. export declare const DEFAULT_CONCURRENCY = 4;
  43. /**
  44. * Default per-request timeout (30 s). embeddinggemma-300M on RTX 4090 takes
  45. * <500ms per batch of 64 in practice; 30s is a safe upper bound.
  46. */
  47. export declare const DEFAULT_TIMEOUT_MS = 30000;
  48. /**
  49. * Retry backoff schedule (ms) for 429/503 responses. 3 attempts total
  50. * (initial + 2 retries) — aligns with issue spec "1s/4s/16s".
  51. */
  52. export declare const RETRY_BACKOFFS_MS: readonly number[];
  53. /**
  54. * Rate-limit (429) retry budget, separate from `RETRY_BACKOFFS_MS`.
  55. *
  56. * A 429 is backpressure, not a fault: the shared ai.mm.mk gateway answers
  57. * `{"detail":"Rate limit exceeded (tokens). Retry after 29s."}` with observed
  58. * retry-after values of 7s/29s/30s, while the generic schedule waits at most
  59. * 1+4+16 = 21s in total. Bulk reindex runs therefore burned all three attempts
  60. * inside one bucket refill and reported the chunk as failed (i-yghj098h).
  61. * 429s get their own attempt count and always wait at least as long as the
  62. * server asked for, capped by `RATE_LIMIT_MAX_BACKOFF_MS`.
  63. */
  64. export declare const DEFAULT_RATE_LIMIT_RETRIES = 5;
  65. /** Upper bound on a single rate-limit wait, so a bogus Retry-After can't hang a run. */
  66. export declare const RATE_LIMIT_MAX_BACKOFF_MS = 60000;
  67. /** Fallback wait when a 429 carries no parseable Retry-After (doubles per attempt). */
  68. export declare const RATE_LIMIT_BASE_BACKOFF_MS = 5000;
  69. /**
  70. * Consecutive successful bulk requests before the lane additively recovers one
  71. * unit of concurrency after a 429-triggered halving (AIMD).
  72. */
  73. export declare const LANE_RECOVERY_STREAK = 8;
  74. /**
  75. * Circuit breaker — flips OPEN when error rate exceeds threshold within
  76. * window. While OPEN, every call fails fast so the caller can fall back.
  77. */
  78. export declare const CIRCUIT_WINDOW_MS = 60000;
  79. export declare const CIRCUIT_OPEN_DURATION_MS: number;
  80. export declare const CIRCUIT_FAILURE_RATE_THRESHOLD = 0.5;
  81. export declare const CIRCUIT_MIN_SAMPLES = 4;
  82. export type OpenAIProviderConfig = {
  83. /** Endpoint base URL — e.g. "https://ai.mm.mk" (no trailing slash) */
  84. endpoint: string;
  85. /** Optional bearer token sent as `Authorization: Bearer ...` */
  86. apiKey?: string;
  87. /**
  88. * Stable model identifier to report up via `getModelId()`.
  89. * Defaults to "embeddinggemma" to match qmd's existing DB rows.
  90. */
  91. modelId?: string;
  92. /**
  93. * Upstream model name sent in the HTTP request body. Often differs from
  94. * `modelId` (e.g. modelId="embeddinggemma" but upstream model="embeddinggemma:300m").
  95. */
  96. upstreamModel?: string;
  97. /** Batch size cap (default DEFAULT_BATCH_SIZE = 64) */
  98. batchSize?: number;
  99. /**
  100. * Max in-flight HTTP requests during a single `embedBatch` call. Default
  101. * `DEFAULT_CONCURRENCY=4` matches the worker semaphore. Set to 1 to force
  102. * legacy sequential dispatch (useful for benchmarks / regression bisect).
  103. */
  104. concurrency?: number;
  105. /** Per-request timeout in ms (default DEFAULT_TIMEOUT_MS = 30_000) */
  106. timeoutMs?: number;
  107. /** Custom fetch (for testing). Defaults to global `fetch`. */
  108. fetchImpl?: typeof fetch;
  109. /** Custom retry schedule (for testing). Defaults to RETRY_BACKOFFS_MS. */
  110. retryBackoffsMs?: readonly number[];
  111. /**
  112. * Extra retry attempts granted to 429 responses on top of the generic
  113. * schedule (default DEFAULT_RATE_LIMIT_RETRIES = 5, env
  114. * `QMD_EMBED_RATE_LIMIT_RETRIES`). Rate limits are backpressure, not faults.
  115. */
  116. rateLimitRetries?: number;
  117. /** Custom sleep impl (for testing). Defaults to setTimeout. */
  118. sleep?: (ms: number) => Promise<void>;
  119. /** Custom clock (for testing). Defaults to Date.now. */
  120. now?: () => number;
  121. };
  122. export type OpenAIEmbeddingsResponse = {
  123. object?: string;
  124. model?: string;
  125. data: Array<{
  126. object?: string;
  127. index?: number;
  128. embedding: number[];
  129. }>;
  130. usage?: {
  131. prompt_tokens?: number;
  132. total_tokens?: number;
  133. };
  134. };
  135. /** Map an OpenAI-compat embedding row onto the input slot. Missing `index` = row order. */
  136. export declare function resolveOpenAIEmbeddingIndex(index: unknown, position: number, inputCount: number): number;
  137. /**
  138. * Circuit breaker state — exported for tests
  139. */
  140. export type CircuitState = "closed" | "open" | "half-open";
  141. /**
  142. * Traffic lane for a single `embedBatch` call. `bulk` = reindex traffic
  143. * (multi-input batch), which self-throttles via `BulkLaneGate` and is
  144. * attributed separately in `X-AI-Caller`. `interactive` = query-time embeds,
  145. * which are never made to wait behind a bulk cooldown.
  146. */
  147. export type EmbedLane = "interactive" | "bulk";
  148. /**
  149. * Determine whether an HTTP status is retryable. 429 (Too Many Requests)
  150. * and 503 (Service Unavailable) are retried; 4xx (other than 429) are not.
  151. */
  152. export declare function isRetryableStatus(status: number): boolean;
  153. /**
  154. * Extract the server-advertised cooldown from a rate-limited response.
  155. *
  156. * Two sources, in priority order:
  157. * 1. the standard `Retry-After` header — delta-seconds or an HTTP-date;
  158. * 2. the ai.mm.mk body text, which states the cooldown in prose only:
  159. * `{"detail":"Rate limit exceeded (tokens). Retry after 29s.", ...}`.
  160. *
  161. * Returns `undefined` when neither source yields a sane positive duration, so
  162. * the caller falls back to its own schedule. Values are clamped to
  163. * `RATE_LIMIT_MAX_BACKOFF_MS`.
  164. */
  165. export declare function parseRetryAfterMs(headerValue: string | null | undefined, bodyPreview?: string, now?: () => number): number | undefined;
  166. /**
  167. * Chunk an array into pieces of ≤ size each. `size` MUST be ≥ 1.
  168. */
  169. export declare function chunkArray<T>(items: T[], size: number): T[][];
  170. /**
  171. * Sliding-window circuit breaker. Tracks the last N samples (min 4) over a
  172. * 60-second window; flips OPEN when failure rate exceeds 50%, then auto-
  173. * resets to HALF-OPEN after 5 minutes — at which point the next probe
  174. * decides whether to close (success) or re-open (failure).
  175. */
  176. export declare class CircuitBreaker {
  177. private samples;
  178. private state;
  179. private openedAt;
  180. private readonly windowMs;
  181. private readonly openDurationMs;
  182. private readonly threshold;
  183. private readonly minSamples;
  184. private readonly now;
  185. constructor(opts?: {
  186. windowMs?: number;
  187. openDurationMs?: number;
  188. threshold?: number;
  189. minSamples?: number;
  190. now?: () => number;
  191. });
  192. getState(): CircuitState;
  193. /**
  194. * Returns true when calls should be short-circuited (skip HTTP, fall back).
  195. * Side-effects: may transition OPEN → HALF-OPEN if the open window expired.
  196. */
  197. shouldFailFast(): boolean;
  198. /** Record a successful call. */
  199. recordSuccess(): void;
  200. /** Record a failed call. May trigger OPEN. */
  201. recordFailure(): void;
  202. /** Force-reset the breaker (used by tests / admin) */
  203. reset(): void;
  204. private pushSample;
  205. private evaluate;
  206. private tickAutoReset;
  207. }
  208. /**
  209. * Raised when the circuit breaker is OPEN and a call is short-circuited.
  210. */
  211. export declare class CircuitOpenError extends Error {
  212. constructor(message?: string);
  213. }
  214. /**
  215. * Persistent (non-retryable) HTTP error from upstream. Includes status code.
  216. */
  217. export declare class HttpError extends Error {
  218. readonly status: number;
  219. readonly bodyPreview: string;
  220. /** Server-advertised cooldown for 429s, when the response stated one. */
  221. readonly retryAfterMs?: number;
  222. constructor(status: number, bodyPreview: string, retryAfterMs?: number);
  223. }
  224. /**
  225. * Client-side bulk lane for reindex traffic (i-yghj098h).
  226. *
  227. * qmd's bulk embedding shares the interactive ai.mm.mk token bucket. Without a
  228. * server-side per-caller budget, the only way bulk traffic can stop contending
  229. * head-on with interactive callers is to police itself:
  230. *
  231. * - a 429 on ANY worker pauses the ENTIRE pool for the advertised cooldown
  232. * (one shared promise — concurrent 429s coalesce instead of stacking N
  233. * cooldowns), so the bucket refills for interactive callers rather than
  234. * being re-drained by the remaining workers;
  235. * - the in-flight cap halves on each cooldown (floor 1) and recovers one unit
  236. * per `LANE_RECOVERY_STREAK` successes — classic AIMD, so a run settles at
  237. * whatever share the bucket actually has spare.
  238. *
  239. * Deliberately clock-free: cooldowns are modelled as a promise produced by the
  240. * injected `sleep`, so tests drive them with a fake sleep and no fake clock.
  241. */
  242. export declare class BulkLaneGate {
  243. private readonly maxPermits;
  244. private readonly sleep;
  245. private permits;
  246. private inFlight;
  247. private okStreak;
  248. private cooldown;
  249. private waiters;
  250. constructor(maxPermits: number, sleep: (ms: number) => Promise<void>);
  251. /** Current in-flight cap — exported state for tests/diagnostics. */
  252. get permitCount(): number;
  253. /** True while the lane is serving a rate-limit cooldown. */
  254. get isCoolingDown(): boolean;
  255. /** Take a slot, waiting out any cooldown and respecting the current cap. */
  256. acquire(): Promise<void>;
  257. /** Return a slot. Always call from a `finally`. */
  258. release(): void;
  259. /**
  260. * Enter (or join) a cooldown of `waitMs` and halve the in-flight cap.
  261. * Returns the shared cooldown promise — the caller awaits it INSTEAD of
  262. * sleeping itself, so a burst of 429s costs one cooldown, not one each.
  263. */
  264. penalize(waitMs: number): Promise<void>;
  265. /** Record a successful bulk request; recovers one permit per success streak. */
  266. noteSuccess(): void;
  267. private wakeOne;
  268. private wakeAll;
  269. }
  270. export declare class OpenAIEmbeddingsProvider implements EmbeddingProvider {
  271. readonly kind: ProviderKind;
  272. private readonly endpoint;
  273. private readonly apiKey?;
  274. private readonly modelId;
  275. private readonly upstreamModel;
  276. private readonly batchSize;
  277. private readonly concurrency;
  278. private readonly timeoutMs;
  279. private readonly fetchImpl;
  280. private readonly retryBackoffsMs;
  281. private readonly rateLimitRetries;
  282. private readonly sleep;
  283. private readonly now;
  284. private dimensions;
  285. private lastError;
  286. readonly breaker: CircuitBreaker;
  287. /** Shared bulk-traffic lane — see `BulkLaneGate`. */
  288. readonly lane: BulkLaneGate;
  289. constructor(config: OpenAIProviderConfig);
  290. getModelId(): string;
  291. getDimensions(): number | undefined;
  292. /**
  293. * Most recent per-chunk failure message (HTTP status + body preview, malformed
  294. * JSON, timeout, abort reason). Returns `undefined` after a successful call
  295. * or before the first call. See `EmbeddingProvider.getLastError`.
  296. */
  297. getLastError(): string | undefined;
  298. /** Endpoint URL configured at construction time — used by callers when
  299. * building error messages for failed first-chunk probes. */
  300. getEndpoint(): string;
  301. healthcheck(signal?: AbortSignal): Promise<ProviderHealth>;
  302. embed(text: string, options?: ProviderEmbedOptions): Promise<ProviderEmbedding | null>;
  303. embedBatch(texts: string[], options?: ProviderEmbedOptions): Promise<(ProviderEmbedding | null)[]>;
  304. dispose(): Promise<void>;
  305. /**
  306. * Format a request-failure context string for `lastError`. Includes endpoint
  307. * + HTTP status + body preview when the error was an `HttpError`, otherwise
  308. * falls back to the message of the underlying error (or the value itself
  309. * when not an Error). Kept short — body preview is already capped at 1024
  310. * chars by `HttpError`, but we trim further here for the dimension-probe
  311. * thrown error which surfaces directly to users.
  312. */
  313. private formatErrorContext;
  314. private buildHeaders;
  315. /**
  316. * Single HTTP request with retry on 429/503. Returns embeddings indexed
  317. * the same as `texts`. Throws on non-retryable failure or all attempts
  318. * exhausted.
  319. */
  320. private requestWithRetry;
  321. /**
  322. * How long to wait after a 429. The server's own `Retry-After` wins when it
  323. * gave one; otherwise fall back to the configured schedule (so injected test
  324. * schedules stay authoritative) and then to an exponential 5s/10s/20s… ramp.
  325. * Always clamped to `RATE_LIMIT_MAX_BACKOFF_MS`.
  326. */
  327. private rateLimitWaitMs;
  328. /**
  329. * Issue one HTTP attempt to `POST /v1/embeddings`. Does NOT retry.
  330. */
  331. private requestOnce;
  332. }