mirror of
https://github.com/garrytan/gbrain.git
synced 2026-07-28 14:59:47 +00:00
Garry TanandClaude Fable 5
595eeb7d6f
Four embedding-throughput/correctness fixes: - #970: google recipe now declares max_batch_tokens (204,800 — derived from Gemini's real limits: 100 inputs per batchEmbedContents × 2048 tokens per input) + max_batch_count 100 + chars_per_token, silencing the missing-cap startup warning and enabling the gateway pre-split. Deliberately NOT the 2048 per-input limit, which would over-split 50x. - #1199: new optional EmbeddingTouchpoint.max_batch_count enforced in splitByTokenBudget (flush at N inputs even when the token budget has room); dashscope sets 10 (provider hard-caps embeddings at 10 inputs per request). isTokenLimitError also learns DashScope's "batch size is invalid" message so recursive halving backstops it. - #1207: gbrain import without --workers now resolves through the shared autoConcurrency policy (PGLite → 1, >100 files on Postgres → 4) instead of hardcoding serial; explicit --workers still wins. - #1818: embedBatch dispatches its 100-input sub-batches through a bounded worker pool (default 4; EmbedBatchOptions.concurrency / GBRAIN_EMBED_BATCH_CONCURRENCY override) with index-addressed results so output order is preserved; single-batch fast path unchanged. Also: listRecipes() now reads the exported RECIPES map instead of the private ALL array (one source of truth; lets tests inject a synthetic capless recipe to keep the startup-warning path covered now that every real recipe declares a cap). Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
70 lines
2.9 KiB
TypeScript
70 lines
2.9 KiB
TypeScript
/**
|
|
* #1207: `gbrain import` without `--workers` used to hardcode workerCount=1,
|
|
* so a large Postgres import paid one serial embedding round-trip per file.
|
|
* runImport now routes the default through the shared autoConcurrency policy
|
|
* (PGLite → 1, >100 files on Postgres → DEFAULT_PARALLEL_WORKERS), while an
|
|
* explicit `--workers N` still wins.
|
|
*
|
|
* The engine here is a minimal postgres-kind stub with no database_url in
|
|
* config — runImport's parallel branch then falls back to serial processing
|
|
* (its PR #490 guard) but the WORKER-COUNT DECISION (the thing #1207 fixes)
|
|
* is still observable via the "Using N parallel workers" log line. Per-file
|
|
* imports fail against the stub engine and are swallowed by runImport's
|
|
* per-file catch; that's fine — this test pins the policy, not the import.
|
|
*/
|
|
|
|
import { afterEach, beforeEach, describe, expect, test } from 'bun:test';
|
|
import { mkdtempSync, writeFileSync, mkdirSync, rmSync, realpathSync } from 'fs';
|
|
import { tmpdir } from 'os';
|
|
import { join } from 'path';
|
|
import { withEnv } from './helpers/with-env.ts';
|
|
import { runImport } from '../src/commands/import.ts';
|
|
|
|
const fakePostgresEngine = {
|
|
kind: 'postgres',
|
|
executeRaw: async () => [],
|
|
logIngest: async () => {},
|
|
setConfig: async () => {},
|
|
getConfig: async () => null,
|
|
} as any;
|
|
|
|
let workspace: string;
|
|
let brainDir: string;
|
|
let logs: string[];
|
|
const realLog = console.log;
|
|
|
|
beforeEach(() => {
|
|
workspace = mkdtempSync(join(tmpdir(), 'gbrain-import-workers-home-'));
|
|
mkdirSync(join(workspace, '.gbrain'), { recursive: true });
|
|
brainDir = realpathSync(mkdtempSync(join(tmpdir(), 'gbrain-import-workers-brain-')));
|
|
// 101 files: one past AUTO_CONCURRENCY_FILE_THRESHOLD (100).
|
|
for (let i = 0; i < 101; i++) {
|
|
writeFileSync(join(brainDir, `page-${i}.md`), `# Page ${i}\n\nbody ${i}\n`);
|
|
}
|
|
logs = [];
|
|
console.log = (msg?: unknown) => logs.push(String(msg));
|
|
});
|
|
|
|
afterEach(() => {
|
|
console.log = realLog;
|
|
rmSync(workspace, { recursive: true, force: true });
|
|
rmSync(brainDir, { recursive: true, force: true });
|
|
});
|
|
|
|
describe('import default worker count (#1207)', () => {
|
|
test('no --workers flag → autoConcurrency picks 4 for >100 files on Postgres', async () => {
|
|
await withEnv({ GBRAIN_HOME: join(workspace, '.gbrain'), GBRAIN_SOURCE: undefined }, async () => {
|
|
await runImport(fakePostgresEngine, [brainDir, '--no-embed'], { sourceId: 'default' });
|
|
});
|
|
expect(logs.some(l => l.includes('Using 4 parallel workers'))).toBe(true);
|
|
});
|
|
|
|
test('explicit --workers 2 still wins over the auto policy', async () => {
|
|
await withEnv({ GBRAIN_HOME: join(workspace, '.gbrain'), GBRAIN_SOURCE: undefined }, async () => {
|
|
await runImport(fakePostgresEngine, [brainDir, '--no-embed', '--workers', '2'], { sourceId: 'default' });
|
|
});
|
|
expect(logs.some(l => l.includes('Using 2 parallel workers'))).toBe(true);
|
|
expect(logs.some(l => l.includes('Using 4 parallel workers'))).toBe(false);
|
|
});
|
|
});
|