md-semantic-search 1.0.3 → 2.0.0

package/src/models.ts

@@ -3,7 +3,7 @@
  */
 import crypto from 'node:crypto';
 
-export type Pooling = 'mean' | 'last_token' | 'none';
+export type Pooling = 'mean' | 'last_token' | 'none' | 'cls';
 export type DType = 'q8' | 'q4' | 'fp16' | 'fp32' | 'int8' | 'uint8';
 
 export interface ModelAdapter {
@@ -19,6 +19,12 @@ export interface ModelAdapter {
   dtype?: DType | string;
   maxTokens?: number;
   sessionOptions?: Record<string, unknown>;
+  /**
+   * issue #200: the ONNX graph already emits a pooled `sentence_embedding`
+   * (e.g. EmbeddingGemma) — embedding bypasses the feature-extraction
+   * pipeline and reads that output directly; `pooling: 'none'` is implied.
+   */
+  pooledOutput?: boolean;
   family?: string;
   note?: string;
   unknownAdapter?: boolean;
@@ -68,6 +74,7 @@ export const MODELS: Record<string, ModelAdapter> = {
     id: 'Xenova/bge-m3',
     nativeDim: 1024,
     dim: 1024,
+    dimensions: [64, 128, 256, 512, 1024],
     queryPrefix: '',
     passagePrefix: '',
     pooling: 'mean',
@@ -82,6 +89,7 @@ export const MODELS: Record<string, ModelAdapter> = {
     revision: 'c25a394dd583836952667c12f008335071b3f43d',
     nativeDim: 1024,
     dim: 1024,
+    dimensions: [32, 64, 128, 256, 512, 768, 1024],
     queryPrefix: 'Instruct: Given a web search query, retrieve relevant passages that answer the query\nQuery:',
     passagePrefix: '',
     pooling: 'last_token',
@@ -89,12 +97,42 @@ export const MODELS: Record<string, ModelAdapter> = {
     dtype: 'q8',
     maxTokens: 32768,
     family: 'qwen3',
-    note: '~613 MB q8. Opt-in Qwen3 embedding model.',
+    note: '~613 MB q8. Opt-in Qwen3 embedding model with Matryoshka/MRL support.',
+  },
+  'granite-multilingual-311m': {
+    id: 'onnx-community/granite-embedding-311m-multilingual-r2-ONNX',
+    revision: '8f039f21d4181327268271bea4b11ddcc7eef88d',
+    nativeDim: 768,
+    dim: 768,
+    queryPrefix: '',
+    passagePrefix: '',
+    pooling: 'cls',
+    normalize: true,
+    dtype: 'q8',
+    maxTokens: 8192,
+    family: 'granite',
+    note: '~298 MB q8. IBM Granite r2, ModernBERT (8k context), CLS pooling, Apache-2.0. Issue #200.',
+  },
+  'embeddinggemma-300m': {
+    id: 'onnx-community/embeddinggemma-300m-ONNX',
+    revision: '5090578d9565bb06545b4552f76e6bc2c93e4a66',
+    nativeDim: 768,
+    dim: 768,
+    dimensions: [64, 128, 256, 512, 768],
+    queryPrefix: 'task: search result | query: ',
+    passagePrefix: 'title: none | text: ',
+    pooling: 'none',
+    normalize: true,
+    dtype: 'q8',
+    maxTokens: 2048,
+    pooledOutput: true,
+    family: 'gemma',
+    note: '~300 MB q8. Google EmbeddingGemma, pre-pooled sentence_embedding output, MRL. Gemma license (not MIT/Apache). Issue #200.',
   },
 };
 
 export const DEFAULT_MODEL = 'e5-base';
-export const POOLING: Pooling[] = ['mean', 'last_token', 'none'];
+export const POOLING: Pooling[] = ['mean', 'last_token', 'none', 'cls'];
 const POOLING_SET = new Set<string>(POOLING);
 
 const aliasByRepoId = new Map<string, string>();
@@ -139,6 +177,7 @@ export function normalizeAdapter(raw: Partial<ModelAdapter> | ModelAdapter): Mod
     normalize: raw.normalize !== false,
     dtype: typeof raw.dtype === 'string' && raw.dtype.length > 0 ? raw.dtype : 'q8',
     ...(Number.isSafeInteger(raw.maxTokens) && raw.maxTokens && raw.maxTokens > 0 ? { maxTokens: raw.maxTokens } : {}),
+    ...(raw.pooledOutput === true ? { pooledOutput: true } : {}),
     ...(typeof raw.family === 'string' && raw.family.length > 0 ? { family: raw.family } : {}),
     ...(typeof raw.note === 'string' ? { note: raw.note } : {}),
   });
@@ -148,42 +187,65 @@ export function normalizeAdapter(raw: Partial<ModelAdapter> | ModelAdapter): Mod
   return adapter;
 }
 
-export function resolveModel(name?: string | Partial<ModelAdapter> | ModelAdapter | null): ModelAdapter {
+export function resolveModel(
+  name?: string | Partial<ModelAdapter> | ModelAdapter | null,
+  chosenDim?: number
+): ModelAdapter {
+  let base: ModelAdapter;
   if (name === undefined || name === null || name === '') {
-    return withDim(MODELS[DEFAULT_MODEL]);
-  }
-  if (typeof name === 'object' && !Array.isArray(name)) {
-    return normalizeAdapter(name);
-  }
-  if (typeof name !== 'string') {
+    base = withDim(MODELS[DEFAULT_MODEL]);
+  } else if (typeof name === 'object' && !Array.isArray(name)) {
+    base = normalizeAdapter(name);
+  } else if (typeof name !== 'string') {
     throw new Error(`model must be an alias, id, or adapter object — got ${typeof name}`);
+  } else if (MODELS[name]) {
+    base = withDim(MODELS[name]);
+  } else {
+    let id = name;
+    let revision = 'main';
+    const at = name.indexOf('@');
+    const hasExplicitRevision = at > 0;
+    if (hasExplicitRevision) {
+      id = name.slice(0, at);
+      const requested = name.slice(at + 1);
+      revision = requested === '' ? 'main' : requested;
+    }
+    const alias = aliasByRepoId.get(id);
+    if (alias !== undefined) {
+      const registered = MODELS[alias];
+      const resolved = hasExplicitRevision ? { ...registered, revision } : registered;
+      base = withDim(resolved);
+    } else {
+      base = normalizeAdapter({
+        id,
+        ...(revision !== 'main' ? { revision } : {}),
+        pooling: 'none',
+        normalize: false,
+        nativeDim: 0,
+        dim: 0,
+        unknownAdapter: true,
+      });
+    }
   }
-  if (MODELS[name]) return withDim(MODELS[name]);
 
-  let id = name;
-  let revision = 'main';
-  const at = name.indexOf('@');
-  const hasExplicitRevision = at > 0;
-  if (hasExplicitRevision) {
-    id = name.slice(0, at);
-    const requested = name.slice(at + 1);
-    revision = requested === '' ? 'main' : requested;
+  if (chosenDim !== undefined) {
+    if (typeof chosenDim !== 'number' || !Number.isSafeInteger(chosenDim) || chosenDim <= 0) {
+      throw new Error(`--dimensions must be a positive integer, got "${chosenDim}"`);
+    }
+    if (!base.dimensions || base.dimensions.length === 0) {
+      throw new Error(
+        `model "${base.id}" does not declare Matryoshka/MRL dimensions support (supported MRL models: qwen3-embedding-0.6b, bge-m3)`
+      );
+    }
+    if (!base.dimensions.includes(chosenDim)) {
+      throw new Error(
+        `invalid dimension ${chosenDim} for model "${base.id}" — supported MRL dimensions: ${base.dimensions.join(', ')}`
+      );
+    }
+    return { ...base, dim: chosenDim, nativeDim: base.nativeDim || base.dim };
   }
-  const alias = aliasByRepoId.get(id);
-  if (alias !== undefined) {
-    const registered = MODELS[alias];
-    const resolved = hasExplicitRevision ? { ...registered, revision } : registered;
-    return withDim(resolved);
-  }
-  return normalizeAdapter({
-    id,
-    ...(revision !== 'main' ? { revision } : {}),
-    pooling: 'none',
-    normalize: false,
-    nativeDim: 0,
-    dim: 0,
-    unknownAdapter: true,
-  });
+
+  return base;
 }
 
 export interface FingerprintInput {
@@ -202,7 +264,7 @@ export function embeddingAdapterFingerprint(model: FingerprintInput): string {
     `passage-prefix:${model.passagePrefix ?? ''}`,
     `pooling:${model.pooling ?? 'mean'}`,
     `normalize:${model.normalize !== false}`,
-    `dimension:${model.nativeDim ?? model.dim ?? 0}`,
+    `dimension:${model.dim ?? model.nativeDim ?? 0}`,
   ].join('\u0000');
   return `adapter-v1:${crypto.createHash('sha256').update(semantics).digest('hex')}`;
 }

package/src/providers.ts

@@ -13,6 +13,7 @@
  */
 import { loadApiKeyFile } from './serve.js';
 import { ModelAdapter } from './models.js';
+import { fetchTimeoutMs } from './core.js';
 
 export type EmbedderRole = 'query' | 'passage';
 export type ExternalEmbedFn = (
@@ -59,6 +60,9 @@ async function postJson(fetchImpl: typeof fetch, url: string, headers: Record<st
     method: 'POST',
     headers: { 'content-type': 'application/json', ...headers },
     body: JSON.stringify(body),
+    // issue #184: a hung embedder endpoint must fail into the retry path,
+    // not hold an index run / serve slot for undici's ~300s default
+    signal: AbortSignal.timeout(fetchTimeoutMs()),
   });
   if (!res.ok) {
     const text = await res.text().catch(() => '');
@@ -124,6 +128,10 @@ export function resolveExternalEmbedder(cfg: ProviderConfig): ResolvedEmbedder {
       family: cfg.embedder,
     };
   })();
+  // The probe starts eagerly; callers that only ever use embedFn never await
+  // it, so its rejection (e.g. the issue #184 fetch timeout) must not become
+  // an unhandled one. descriptor() re-surfaces the same rejection to awaiters.
+  descriptorPromise.catch(() => {});
 
   return {
     id: cfg.embedder,
@@ -132,3 +140,340 @@ export function resolveExternalEmbedder(cfg: ProviderConfig): ResolvedEmbedder {
     descriptor: () => descriptorPromise,
   };
 }
+
+/**
+ * Chat provider (issue #199): one OpenAI-compatible `/chat/completions` path
+ * for answer synthesis — the same lingua franca HyDE already speaks — with an
+ * Ollama-native `/api/generate` legacy path (`--llm-provider ollama`) that
+ * preserves the pre-#199 wire format for existing setups.
+ */
+export type ChatProviderId = 'openai' | 'ollama';
+
+export interface ChatProviderConfig {
+  provider: ChatProviderId;
+  /** Provider-side model name; omitted bodies let the server pick its default. */
+  model?: string;
+  baseUrl?: string;
+  keyFile?: string;
+  /** Injected fetch for tests; defaults to globalThis.fetch. */
+  fetchImpl?: typeof fetch;
+}
+
+export interface ChatRequest {
+  system: string;
+  user: string;
+  /** Streaming render hook; fires once per token chunk when the wire streams. */
+  onDelta?: (delta: string) => void;
+  /** Cooperative cancel (serve: client disconnected mid-stream). */
+  signal?: AbortSignal;
+}
+
+export interface ResolvedChat {
+  id: ChatProviderId;
+  model: string;
+  /** Streams one completion; resolves to the assembled full text. */
+  complete(req: ChatRequest): Promise<string>;
+}
+
+/**
+ * Incremental SSE frame parser (issue #199): feed it raw network chunks, it
+ * emits one `data:` payload per complete event. Handles CRLF, LF-only streams,
+ * multi-line `data:` fields (joined with \n per the WHATWG spec), comment
+ * keep-alive lines (`: ping`), and chunks that split a line or a CRLF pair
+ * across the boundary. Non-`data` fields (event/id/retry) are irrelevant to
+ * the chat delta stream and dropped.
+ */
+export function createSseParser(onData: (data: string) => void): { push(chunk: string): void; end(): void } {
+  let buf = '';
+  let dataLines: string[] = [];
+  const dispatch = () => {
+    if (dataLines.length > 0) {
+      const payload = dataLines.join('\n');
+      dataLines = [];
+      onData(payload);
+    }
+  };
+  const consumeLine = (rawLine: string) => {
+    const line = rawLine.endsWith('\r') ? rawLine.slice(0, -1) : rawLine;
+    if (line === '') {
+      dispatch();
+      return;
+    }
+    if (line.startsWith(':')) return; // comment / keep-alive
+    if (line.startsWith('data:')) dataLines.push(line.slice(5).replace(/^ /, ''));
+  };
+  return {
+    push(chunk: string) {
+      buf += chunk;
+      let nl: number;
+      while ((nl = buf.indexOf('\n')) >= 0) {
+        const line = buf.slice(0, nl);
+        buf = buf.slice(nl + 1);
+        consumeLine(line);
+      }
+    },
+    end() {
+      if (buf !== '') {
+        consumeLine(buf);
+        buf = '';
+      }
+      dispatch();
+    },
+  };
+}
+
+/** `data:` payload of an OpenAI streaming chunk → delta text ('' when none). */
+export function openAiDeltaText(payload: string): string {
+  if (payload === '[DONE]') return '';
+  try {
+    const json: any = JSON.parse(payload);
+    const text = json?.choices?.[0]?.delta?.content;
+    return typeof text === 'string' ? text : '';
+  } catch {
+    return '';
+  }
+}
+
+/** One Ollama NDJSON streaming line → delta text ('' when none). */
+export function ollamaDeltaText(line: string): string {
+  const trimmed = line.trim();
+  if (!trimmed) return '';
+  try {
+    const json: any = JSON.parse(trimmed);
+    const text = json?.response;
+    return typeof text === 'string' ? text : '';
+  } catch {
+    return '';
+  }
+}
+
+function chatBaseOf(cfg: ChatProviderConfig): string {
+  const raw =
+    cfg.baseUrl ||
+    (cfg.provider === 'ollama' ? process.env.OLLAMA_HOST || 'http://127.0.0.1:11434' : 'https://api.openai.com/v1');
+  return raw.replace(/\/+$/, '');
+}
+
+function chatBearerKey(cfg: ChatProviderConfig): string | undefined {
+  if (cfg.provider !== 'openai') return undefined;
+  if (cfg.keyFile) return loadApiKeyFile(cfg.keyFile);
+  return process.env.OPENAI_API_KEY || undefined;
+}
+
+const IDLE_ABORT_REASON = () =>
+  new DOMException('The operation was aborted due to timeout', 'TimeoutError');
+
+/**
+ * Streams one completion with an IDLE timeout (issue #199 × #184): the
+ * abort timer arms before connect and resets on every received chunk, so a
+ * hung endpoint fails fast while a long streaming generation is never killed
+ * mid-answer the way a whole-request `AbortSignal.timeout` would. An external
+ * `signal` (serve: client disconnected mid-stream) aborts immediately.
+ */
+function makeIdleGuard(external?: AbortSignal) {
+  const controller = new AbortController();
+  const idleMs = fetchTimeoutMs();
+  let timer: ReturnType<typeof setTimeout> | null = setTimeout(() => controller.abort(IDLE_ABORT_REASON()), idleMs);
+  (timer as any)?.unref?.();
+  const onExternalAbort = () => controller.abort(external!.reason);
+  if (external) {
+    if (external.aborted) onExternalAbort();
+    else external.addEventListener('abort', onExternalAbort, { once: true });
+  }
+  return {
+    signal: controller.signal,
+    /** Reset the idle window (after connect, after every chunk). */
+    arm() {
+      if (timer) clearTimeout(timer);
+      timer = setTimeout(() => controller.abort(IDLE_ABORT_REASON()), idleMs);
+      (timer as any)?.unref?.();
+    },
+    done() {
+      if (timer) {
+        clearTimeout(timer);
+        timer = null;
+      }
+      external?.removeEventListener('abort', onExternalAbort);
+    },
+  };
+}
+
+/** Line-splitting over arbitrary network chunks (Ollama NDJSON). */
+function createLineParser(onLine: (line: string) => void): { push(chunk: string): void; end(): void } {
+  let buf = '';
+  return {
+    push(chunk: string) {
+      buf += chunk;
+      let nl: number;
+      while ((nl = buf.indexOf('\n')) >= 0) {
+        const line = buf.slice(0, nl).replace(/\r$/, '');
+        buf = buf.slice(nl + 1);
+        onLine(line);
+      }
+    },
+    end() {
+      if (buf !== '') {
+        onLine(buf.replace(/\r$/, ''));
+        buf = '';
+      }
+    },
+  };
+}
+
+function isLegacyFallbackEligible(err: unknown): boolean {
+  // Only a clearly OpenAI-incompatible shape (no such route) justifies trying
+  // the Ollama-native endpoint; auth/4xx/5xx errors must surface as-is.
+  return /HTTP 40[45]\b/.test(String((err as any)?.message || ''));
+}
+
+/**
+ * Build a chat completion provider over HTTP streaming (issue #199).
+ * `openai` → POST {base}/chat/completions with `stream:true`, SSE deltas.
+ * `ollama` → POST {base}/api/generate with `stream:true`, NDJSON deltas —
+ * the pre-#199 endpoint, upgraded from `stream:false` to streaming. The
+ * openai path additionally falls back to /api/generate once when the endpoint
+ * answers 404/405 (not OpenAI-shaped), so an Ollama root URL keeps working
+ * with default flags.
+ */
+export function resolveChatProvider(cfg: ChatProviderConfig): ResolvedChat {
+  const fetchImpl = cfg.fetchImpl ?? globalThis.fetch;
+  if (typeof fetchImpl !== 'function') throw new Error('chat: no fetch implementation available');
+  const base = chatBaseOf(cfg);
+
+  const streamOllamaGenerate = async (
+    system: string,
+    user: string,
+    onDelta: ((d: string) => void) | undefined,
+    signal: AbortSignal | undefined,
+  ): Promise<string> => {
+    const guard = makeIdleGuard(signal);
+    let res: Response;
+    try {
+      res = await fetchImpl(`${base}/api/generate`, {
+        method: 'POST',
+        headers: { 'content-type': 'application/json' },
+        signal: guard.signal,
+        body: JSON.stringify({
+          model: cfg.model || 'llama3',
+          // legacy wire format: one prompt string, byte-identical to the
+          // pre-#199 single-string synthesis prompt
+          prompt: system + user,
+          stream: true,
+        }),
+      });
+    } catch (e) {
+      guard.done();
+      throw e;
+    }
+    if (!res.ok) {
+      guard.done();
+      const text = await res.text().catch(() => '');
+      throw new Error(`chat ${base}/api/generate -> HTTP ${res.status}${text ? `: ${text.slice(0, 200)}` : ''}`);
+    }
+    let full = '';
+    const parser = createLineParser((line) => {
+      const delta = ollamaDeltaText(line);
+      if (delta) {
+        full += delta;
+        onDelta?.(delta);
+      }
+    });
+    const decoder = new TextDecoder();
+    try {
+      for await (const chunk of res.body as unknown as AsyncIterable<Uint8Array | string>) {
+        parser.push(typeof chunk === 'string' ? chunk : decoder.decode(chunk, { stream: true }));
+        guard.arm();
+      }
+      parser.push(decoder.decode());
+      parser.end();
+    } finally {
+      guard.done();
+    }
+    return full;
+  };
+
+  const streamOpenAiChat = async (
+    system: string,
+    user: string,
+    onDelta: ((d: string) => void) | undefined,
+    signal: AbortSignal | undefined,
+  ): Promise<string> => {
+    const headers: Record<string, string> = { 'content-type': 'application/json' };
+    const key = chatBearerKey(cfg);
+    if (key) headers.authorization = `Bearer ${key}`;
+    const guard = makeIdleGuard(signal);
+    let res: Response;
+    try {
+      res = await fetchImpl(`${base}/chat/completions`, {
+        method: 'POST',
+        headers,
+        signal: guard.signal,
+        body: JSON.stringify({
+          ...(cfg.model ? { model: cfg.model } : {}),
+          messages: [
+            { role: 'system', content: system },
+            { role: 'user', content: user },
+          ],
+          stream: true,
+        }),
+      });
+    } catch (e) {
+      guard.done();
+      throw e;
+    }
+    if (!res.ok) {
+      guard.done();
+      const text = await res.text().catch(() => '');
+      throw new Error(`chat ${base}/chat/completions -> HTTP ${res.status}${text ? `: ${text.slice(0, 200)}` : ''}`);
+    }
+    let full = '';
+    let wireDone = false;
+    const parser = createSseParser((payload) => {
+      if (payload === '[DONE]') {
+        wireDone = true;
+        return;
+      }
+      const delta = openAiDeltaText(payload);
+      if (delta) {
+        full += delta;
+        onDelta?.(delta);
+      }
+    });
+    const decoder = new TextDecoder();
+    try {
+      for await (const chunk of res.body as unknown as AsyncIterable<Uint8Array | string>) {
+        parser.push(typeof chunk === 'string' ? chunk : decoder.decode(chunk, { stream: true }));
+        guard.arm();
+        // leaving the for-await on [DONE] cancels the stream through the
+        // iterator's return() — a manual cancel() here would just reject
+        // with "ReadableStream is locked"
+        if (wireDone) break;
+      }
+      if (!wireDone) {
+        parser.push(decoder.decode());
+        parser.end();
+      }
+    } finally {
+      guard.done();
+    }
+    return full;
+  };
+
+  return {
+    id: cfg.provider,
+    model: cfg.model || '',
+    async complete(req) {
+      if (cfg.provider === 'ollama') {
+        return streamOllamaGenerate(req.system, req.user, req.onDelta, req.signal);
+      }
+      try {
+        return await streamOpenAiChat(req.system, req.user, req.onDelta, req.signal);
+      } catch (e) {
+        if (isLegacyFallbackEligible(e)) {
+          return streamOllamaGenerate(req.system, req.user, req.onDelta, req.signal);
+        }
+        throw e;
+      }
+    },
+  };
+}

package/src/quantization.ts

@@ -91,3 +91,64 @@ export function asymmetricCosineInt8(queryFp32: Float32Array, chunkInt8: Int8Arr
   if (normSq <= 0) return 0.0;
   return dot / Math.sqrt(normSq);
 }
+
+/**
+ * Precomputed int8 chunk norm = sqrt(sum(c_i^2)) (issue #190): constant per
+ * chunk, so caching it at load time removes half of the asymmetric sweep.
+ * Accumulates in the same order as `asymmetricCosineInt8` so the cached-norm
+ * fast path returns bitwise-equal scores.
+ */
+export function int8Norm(chunkInt8: Int8Array): number {
+  const dim = chunkInt8.length;
+  let normSq = 0.0;
+
+  let i = 0;
+  const limit = dim - (dim % 8);
+  for (; i < limit; i += 8) {
+    normSq +=
+      chunkInt8[i] * chunkInt8[i] +
+      chunkInt8[i + 1] * chunkInt8[i + 1] +
+      chunkInt8[i + 2] * chunkInt8[i + 2] +
+      chunkInt8[i + 3] * chunkInt8[i + 3] +
+      chunkInt8[i + 4] * chunkInt8[i + 4] +
+      chunkInt8[i + 5] * chunkInt8[i + 5] +
+      chunkInt8[i + 6] * chunkInt8[i + 6] +
+      chunkInt8[i + 7] * chunkInt8[i + 7];
+  }
+  for (; i < dim; i++) {
+    normSq += chunkInt8[i] * chunkInt8[i];
+  }
+  return Math.sqrt(normSq);
+}
+
+/**
+ * Dot-only sweep fast path (issue #190): identical to `asymmetricCosineInt8`
+ * when `chunkNorm` came from `int8Norm` on the same vector.
+ */
+export function asymmetricCosineInt8CachedNorm(
+  queryFp32: Float32Array,
+  chunkInt8: Int8Array,
+  chunkNorm: number,
+): number {
+  if (!(chunkNorm > 0)) return 0.0;
+  const dim = queryFp32.length;
+  let dot = 0.0;
+
+  let i = 0;
+  const limit = dim - (dim % 8);
+  for (; i < limit; i += 8) {
+    dot +=
+      queryFp32[i] * chunkInt8[i] +
+      queryFp32[i + 1] * chunkInt8[i + 1] +
+      queryFp32[i + 2] * chunkInt8[i + 2] +
+      queryFp32[i + 3] * chunkInt8[i + 3] +
+      queryFp32[i + 4] * chunkInt8[i + 4] +
+      queryFp32[i + 5] * chunkInt8[i + 5] +
+      queryFp32[i + 6] * chunkInt8[i + 6] +
+      queryFp32[i + 7] * chunkInt8[i + 7];
+  }
+  for (; i < dim; i++) {
+    dot += queryFp32[i] * chunkInt8[i];
+  }
+  return dot / chunkNorm;
+}