diff --git a/.github/workflows/publish-preview.yml b/.github/workflows/publish-preview.yml index da48932ae..1f4f3e834 100644 --- a/.github/workflows/publish-preview.yml +++ b/.github/workflows/publish-preview.yml @@ -22,4 +22,9 @@ jobs: bun-version: latest - run: bun i - run: bun run build + # Publish every workspace under packages/* and providers/* as a PR- + # preview package, plus the cli example. `pkg-pr-new` honors the + # `"private": true` flag on `@workglow/test`, so internal-only + # workspaces are silently skipped. The shell expands the globs into + # one positional argument per matched directory. - run: bunx pkg-pr-new publish ./packages/* ./providers/* ./examples/cli diff --git a/bun.lock b/bun.lock index fe1b429f1..68c47b7e6 100644 --- a/bun.lock +++ b/bun.lock @@ -46,7 +46,7 @@ }, "examples/cli": { "name": "@workglow/cli", - "version": "0.2.35", + "version": "0.2.36", "bin": "./dist/workglow.js", "dependencies": { "@anthropic-ai/sdk": "catalog:", @@ -87,7 +87,7 @@ }, "examples/web": { "name": "@workglow/web", - "version": "0.2.35", + "version": "0.2.36", "dependencies": { "@codemirror/lang-json": "^6.0.2", "@huggingface/transformers": "catalog:", @@ -129,7 +129,7 @@ }, "packages/ai": { "name": "@workglow/ai", - "version": "0.2.35", + "version": "0.2.36", "devDependencies": { "@workglow/job-queue": "workspace:*", "@workglow/knowledge-base": "workspace:*", @@ -147,7 +147,7 @@ }, "packages/browser-control": { "name": "@workglow/browser-control", - "version": "0.2.35", + "version": "0.2.36", "devDependencies": { "@workglow/task-graph": "workspace:*", "@workglow/util": "workspace:*", @@ -159,7 +159,7 @@ }, "packages/indexeddb": { "name": "@workglow/indexeddb", - "version": "0.2.35", + "version": "0.2.36", "devDependencies": { "@workglow/job-queue": "workspace:*", "@workglow/storage": "workspace:*", @@ -174,7 +174,7 @@ }, "packages/javascript": { "name": "@workglow/javascript", - "version": "0.2.35", + "version": "0.2.36", "devDependencies": { "@workglow/task-graph": "workspace:*", "@workglow/util": "workspace:*", @@ -186,7 +186,7 @@ }, "packages/job-queue": { "name": "@workglow/job-queue", - "version": "0.2.35", + "version": "0.2.36", "devDependencies": { "@workglow/util": "workspace:*", }, @@ -196,19 +196,21 @@ }, "packages/knowledge-base": { "name": "@workglow/knowledge-base", - "version": "0.2.35", + "version": "0.2.36", "devDependencies": { "@workglow/storage": "workspace:*", + "@workglow/task-graph": "workspace:*", "@workglow/util": "workspace:*", }, "peerDependencies": { "@workglow/storage": "workspace:*", + "@workglow/task-graph": "workspace:*", "@workglow/util": "workspace:*", }, }, "packages/mcp": { "name": "@workglow/mcp", - "version": "0.2.35", + "version": "0.2.36", "dependencies": { "@modelcontextprotocol/sdk": "catalog:", }, @@ -225,7 +227,7 @@ }, "packages/storage": { "name": "@workglow/storage", - "version": "0.2.35", + "version": "0.2.36", "devDependencies": { "@workglow/util": "workspace:*", }, @@ -235,7 +237,7 @@ }, "packages/task-graph": { "name": "@workglow/task-graph", - "version": "0.2.35", + "version": "0.2.36", "devDependencies": { "@workglow/job-queue": "workspace:*", "@workglow/storage": "workspace:*", @@ -249,7 +251,7 @@ }, "packages/tasks": { "name": "@workglow/tasks", - "version": "0.2.35", + "version": "0.2.36", "dependencies": { "ipaddr.js": "^2.4.0", "undici": "^8.3.0", @@ -273,7 +275,7 @@ }, "packages/test": { "name": "@workglow/test", - "version": "0.2.35", + "version": "0.2.36", "devDependencies": { "@electric-sql/pglite": "catalog:", "@modelcontextprotocol/sdk": "catalog:", @@ -315,7 +317,7 @@ }, "packages/util": { "name": "@workglow/util", - "version": "0.2.35", + "version": "0.2.36", "dependencies": { "@sroussey/json-schema-library": "^11.4.0", "@sroussey/json-schema-to-ts": "3.1.4", @@ -326,7 +328,7 @@ }, "packages/workglow": { "name": "workglow", - "version": "0.2.35", + "version": "0.2.36", "dependencies": { "@workglow/ai": "workspace:*", "@workglow/anthropic": "workspace:*", @@ -353,7 +355,7 @@ }, "providers/anthropic": { "name": "@workglow/anthropic", - "version": "0.2.35", + "version": "0.2.36", "dependencies": { "@anthropic-ai/sdk": "catalog:", }, @@ -374,7 +376,7 @@ }, "providers/bun-webview": { "name": "@workglow/bun-webview", - "version": "0.2.35", + "version": "0.2.36", "devDependencies": { "@workglow/browser-control": "workspace:*", "@workglow/util": "workspace:*", @@ -386,7 +388,7 @@ }, "providers/chrome-ai": { "name": "@workglow/chrome-ai", - "version": "0.2.35", + "version": "0.2.36", "devDependencies": { "@workglow/ai": "workspace:*", "@workglow/job-queue": "workspace:*", @@ -404,7 +406,7 @@ }, "providers/electron": { "name": "@workglow/electron", - "version": "0.2.35", + "version": "0.2.36", "dependencies": { "electron": "catalog:", }, @@ -419,7 +421,7 @@ }, "providers/google-gemini": { "name": "@workglow/google-gemini", - "version": "0.2.35", + "version": "0.2.36", "dependencies": { "@google/generative-ai": "catalog:", }, @@ -440,7 +442,7 @@ }, "providers/huggingface-inference": { "name": "@workglow/huggingface-inference", - "version": "0.2.35", + "version": "0.2.36", "dependencies": { "@huggingface/inference": "catalog:", }, @@ -461,7 +463,7 @@ }, "providers/huggingface-transformers": { "name": "@workglow/huggingface-transformers", - "version": "0.2.35", + "version": "0.2.36", "dependencies": { "@huggingface/transformers": "catalog:", }, @@ -482,7 +484,7 @@ }, "providers/node-llama-cpp": { "name": "@workglow/node-llama-cpp", - "version": "0.2.35", + "version": "0.2.36", "dependencies": { "node-llama-cpp": "catalog:", }, @@ -503,7 +505,7 @@ }, "providers/ollama": { "name": "@workglow/ollama", - "version": "0.2.35", + "version": "0.2.36", "dependencies": { "ollama": "catalog:", }, @@ -524,7 +526,7 @@ }, "providers/openai": { "name": "@workglow/openai", - "version": "0.2.35", + "version": "0.2.36", "dependencies": { "js-tiktoken": "catalog:", "openai": "catalog:", @@ -547,7 +549,7 @@ }, "providers/playwright": { "name": "@workglow/playwright", - "version": "0.2.35", + "version": "0.2.36", "dependencies": { "playwright": "catalog:", }, @@ -560,7 +562,7 @@ }, "providers/postgres": { "name": "@workglow/postgres", - "version": "0.2.35", + "version": "0.2.36", "devDependencies": { "@types/pg": "^8.18.0", "@workglow/job-queue": "workspace:*", @@ -581,7 +583,7 @@ }, "providers/sqlite": { "name": "@workglow/sqlite", - "version": "0.2.35", + "version": "0.2.36", "devDependencies": { "@sqlite.org/sqlite-wasm": "catalog:", "@sqliteai/sqlite-vector": "catalog:", @@ -607,7 +609,7 @@ }, "providers/supabase": { "name": "@workglow/supabase", - "version": "0.2.35", + "version": "0.2.36", "dependencies": { "@supabase/supabase-js": "catalog:", }, @@ -624,7 +626,7 @@ }, "providers/tf-mediapipe": { "name": "@workglow/tf-mediapipe", - "version": "0.2.35", + "version": "0.2.36", "dependencies": { "@mediapipe/tasks-audio": "catalog:", "@mediapipe/tasks-genai": "catalog:", diff --git a/examples/cli/CHANGELOG.md b/examples/cli/CHANGELOG.md index 2b6b825fd..724196e89 100644 --- a/examples/cli/CHANGELOG.md +++ b/examples/cli/CHANGELOG.md @@ -1,5 +1,7 @@ # @workglow/cli +## 0.2.36 + ## 0.2.35 ### Refactors diff --git a/examples/cli/package.json b/examples/cli/package.json index 62ba61811..a739e0bf8 100644 --- a/examples/cli/package.json +++ b/examples/cli/package.json @@ -1,7 +1,7 @@ { "name": "@workglow/cli", "type": "module", - "version": "0.2.35", + "version": "0.2.36", "repository": { "type": "git", "url": "https://github.com/workglow-dev/libs.git", diff --git a/examples/web/CHANGELOG.md b/examples/web/CHANGELOG.md index e5e034f68..86dbcf676 100644 --- a/examples/web/CHANGELOG.md +++ b/examples/web/CHANGELOG.md @@ -1,5 +1,20 @@ # @workglow/web +## 0.2.36 + +### Chores + +- update deps + +#### deps-dev + +- bump vite from 8.0.12 to 8.0.13 + +### Updated Dependencies + +- `@vitejs/plugin-react`: ^6.0.2 +- `vite`: ^8.0.13 + ## 0.2.35 ### Refactors diff --git a/examples/web/package.json b/examples/web/package.json index e8dfd2fa9..6de1a88a7 100644 --- a/examples/web/package.json +++ b/examples/web/package.json @@ -1,6 +1,6 @@ { "name": "@workglow/web", - "version": "0.2.35", + "version": "0.2.36", "type": "module", "description": "Web interface example for Workglow, showcasing a React-based UI for building and visualizing AI task pipelines.", "scripts": { diff --git a/package.json b/package.json index 5476d2793..133020a8b 100644 --- a/package.json +++ b/package.json @@ -1,7 +1,7 @@ { "name": "@workglow-dev/libs", "type": "module", - "version": "0.2.35", + "version": "0.2.36", "repository": { "type": "git", "url": "https://github.com/workglow-dev/libs.git" @@ -130,4 +130,4 @@ ], "{packages,providers,examples}/*/src/**/*.json": "prettier --write" } -} \ No newline at end of file +} diff --git a/packages/ai/CHANGELOG.md b/packages/ai/CHANGELOG.md index 5b3fccf40..86daf1d99 100644 --- a/packages/ai/CHANGELOG.md +++ b/packages/ai/CHANGELOG.md @@ -1,5 +1,26 @@ # @workglow/ai +## 0.2.36 + +### Features + +#### kb + +- pluggable strategy with model config, IRunConfig threading, and document tasks + +### Bug Fixes + +#### util/worker, ai/task + +- TTL-based pendingAborts eviction; clarify runWithIterable bond (#500) + +#### ai + +- avoid NaN reranker scores for empty queries +- escape regex metacharacters in RerankerTask.simpleRerank +- classify provider-error vs no-finish in AiTask.execute +- replace runWithIterable Proxy with shallow clone + ## 0.2.35 ### Features diff --git a/packages/ai/package.json b/packages/ai/package.json index fafab1e84..cd49507a7 100644 --- a/packages/ai/package.json +++ b/packages/ai/package.json @@ -1,7 +1,7 @@ { "name": "@workglow/ai", "type": "module", - "version": "0.2.35", + "version": "0.2.36", "repository": { "type": "git", "url": "https://github.com/workglow-dev/libs.git", diff --git a/packages/ai/src/common.ts b/packages/ai/src/common.ts index fc5136f3c..2bad1c0aa 100644 --- a/packages/ai/src/common.ts +++ b/packages/ai/src/common.ts @@ -22,6 +22,7 @@ export * from "./provider/AiProviderRegistry"; export * from "./provider/QueuedAiProvider"; export * from "./capability"; +export * from "./kb/createStandardKbStrategy"; export * from "./task"; import { AiVisionTask } from "./task/base/AiVisionTask"; diff --git a/packages/ai/src/kb/createStandardKbStrategy.ts b/packages/ai/src/kb/createStandardKbStrategy.ts new file mode 100644 index 000000000..591d09227 --- /dev/null +++ b/packages/ai/src/kb/createStandardKbStrategy.ts @@ -0,0 +1,327 @@ +/** + * @license + * Copyright 2026 Steven Roussey + * SPDX-License-Identifier: Apache-2.0 + */ + +import type { + ChunkSearchResult, + ChunkStrategy, + Document, + IKbAiStrategy, + IKbStrategyTarget, + ISearchOptions, + SearchMode, +} from "@workglow/knowledge-base"; +import { chunkText, toInsertChunkEntities } from "@workglow/knowledge-base"; +import type { IRunConfig } from "@workglow/task-graph"; +import type { TypedArray } from "@workglow/util/schema"; + +import { HierarchicalChunkerTask } from "../task/HierarchicalChunkerTask"; +import { RerankerTask } from "../task/RerankerTask"; +import { TextEmbeddingTask } from "../task/TextEmbeddingTask"; +import { TextRerankerTask } from "../task/TextRerankerTask"; + +/** + * Tuning knobs for the standard strategy. Most defaults come straight from + * the KB (model IDs, chunkStrategy, searchMode); these overrides exist for + * callers that want different chunker token budgets than the built-in + * defaults or that need to pin a search mode different from what the KB + * has stored. + */ +export interface CreateStandardKbStrategyOptions { + readonly chunker?: { + readonly maxTokens?: number; + readonly overlap?: number; + readonly reservedTokens?: number; + }; + /** Override KB's chunkStrategy at strategy-build time. */ + readonly chunkStrategy?: ChunkStrategy; + /** Override KB's searchMode at strategy-build time. */ + readonly searchMode?: SearchMode; + /** + * Multiplier applied to `topK` to size the first-stage candidate pool + * when `searchMode === "rerank"`. The reranker then narrows the pool + * back down to `topK`. Defaults to `5`, i.e. first stage fetches + * `topK * 5` candidates. Used together with `firstStageMinimum` — + * the actual first-stage size is `max(topK * firstStageMultiplier, + * firstStageMinimum)`, so a tiny `topK` (e.g. `1`) still yields a + * meaningful candidate pool for the reranker to choose from instead of + * collapsing to a single candidate. + */ + readonly firstStageMultiplier?: number; + /** + * Minimum first-stage candidate pool size when `searchMode === "rerank"`. + * Defaults to `20`. Prevents the rerank pool from collapsing to + * `topK` for very small `topK` values where `topK * firstStageMultiplier` + * would still be too few candidates for the reranker to do useful work. + * The effective first-stage size is + * `max(topK * firstStageMultiplier, firstStageMinimum)`. + */ + readonly firstStageMinimum?: number; +} + +/** + * The standard KB strategy: hierarchical-by-default chunking + embedding + * during ingest, and a single search mode for retrieval. Search and ingest + * read the KB's stored model IDs (`docEmbeddingModel` / + * `queryEmbeddingModel` / `rerankerModel`) and config fields + * (`chunkStrategy` / `searchMode`) on every call, so updates to the KB + * record take effect immediately on the next op. + * + * Score semantics: results carry `scoreType` matching the retrieval + * mode — `"cosine"` for similarity, `"rrf"` for hybrid, `"rerank"` for + * both reranker-model and heuristic fallback paths. **Cross-encoder + * rerank scores are raw logits**, not probabilities or similarities, and + * they are NOT comparable to cosine / BM25 / RRF scores. Always check + * `scoreType` before applying a score threshold; the strategy itself + * ignores `ISearchOptions.scoreThreshold` in the rerank branch. + * + * For custom RAG flows (per-tenant scoping, alternative chunkers, etc.) + * write your own `IKbAiStrategy` — this factory is the "good defaults" + * path, not the only path. + */ +export function createStandardKbStrategy( + options: CreateStandardKbStrategyOptions = {} +): IKbAiStrategy { + const chunkerDefaults = { + maxTokens: options.chunker?.maxTokens ?? 512, + overlap: options.chunker?.overlap ?? 50, + reservedTokens: options.chunker?.reservedTokens ?? 10, + } as const; + const firstStageMultiplier = options.firstStageMultiplier ?? 5; + const firstStageMinimum = options.firstStageMinimum ?? 20; + + const resolveSearchMode = (kb: IKbStrategyTarget): SearchMode => { + if (options.searchMode) return options.searchMode; + if (kb.searchMode) return kb.searchMode; + if (kb.rerankerModel) return "rerank"; + if (kb.supportsHybridSearch()) return "hybrid"; + return "similarity"; + }; + + const resolveChunkStrategy = (kb: IKbStrategyTarget): ChunkStrategy => + options.chunkStrategy ?? kb.chunkStrategy ?? "hierarchical"; + + const requireQueryEmbedModel = (kb: IKbStrategyTarget): string => { + const m = kb.queryEmbeddingModel ?? kb.docEmbeddingModel; + if (!m) { + throw new Error( + `KnowledgeBase "${kb.name}": no queryEmbeddingModel or docEmbeddingModel configured.` + ); + } + return m; + }; + + const requireDocEmbedModel = (kb: IKbStrategyTarget): string => { + const m = kb.docEmbeddingModel; + if (!m) { + throw new Error(`KnowledgeBase "${kb.name}": no docEmbeddingModel configured.`); + } + return m; + }; + + const embedTexts = async ( + texts: readonly string[], + modelId: string, + runConfig?: Partial + ): Promise => { + if (texts.length === 0) return []; + const result = await new TextEmbeddingTask().run( + { text: texts as string[], model: modelId }, + runConfig + ); + const vector = result.vector; + return Array.isArray(vector) ? (vector as TypedArray[]) : [vector as TypedArray]; + }; + + return { + async ingest(kb, doc, runConfig): Promise { + // Order matters: delete old chunks BEFORE rewriting the document. + // If upsertDocument or any later step fails partway through, the + // worst the KB can be left in is "doc row preserved, chunks + // removed" rather than "new doc row pointing at old stale chunks" + // — chunks always reflect the in-flight ingest, never a previous + // version. The text-index removal piggy-backs on + // deleteChunksForDocument, so RRF rankings can't end up surfacing + // chunks that no longer exist either. + const initialDocId = doc.doc_id; + if (initialDocId) { + await kb.deleteChunksForDocument(initialDocId); + } + const stored = await kb.upsertDocument(doc); + const docId = stored.doc_id!; + if (!initialDocId) { + // Fresh-id case: chunks under this new id can't pre-exist in a + // well-behaved storage backend, but call delete unconditionally + // so the post-condition ("after ingest returns, the doc owns + // exactly the newly-embedded chunks") holds even if a backend + // recycles ids or a stale row survived a prior aborted run. + await kb.deleteChunksForDocument(docId); + } + + const chunker = new HierarchicalChunkerTask(); + const chunkResult = await chunker.run( + { + doc_id: docId, + documentTree: stored.root as never, + strategy: resolveChunkStrategy(kb), + ...chunkerDefaults, + }, + runConfig + ); + const chunks = chunkResult.chunks ?? []; + if (chunks.length === 0) return stored; + + const vectors = await embedTexts( + chunks.map((c) => c.text), + requireDocEmbedModel(kb), + runConfig + ); + const inserts = toInsertChunkEntities( + { chunks, vectors }, + { doc_id: docId, doc_title: stored.metadata.title } + ); + await kb.upsertChunksBulk(inserts); + return stored; + }, + + async delete(kb, doc_id): Promise { + await kb.deleteDocument(doc_id); + }, + + async search( + kb, + query, + options?: ISearchOptions, + runConfig?: Partial + ): Promise { + const mode = resolveSearchMode(kb); + const topK = options?.topK ?? 5; + const filter = options?.filter; + const scoreThreshold = options?.scoreThreshold; + + if (mode === "text") { + // Pure FTS via hybridSearch with vectorWeight=0 — works on backends + // that support hybrid (Postgres, in-memory). For storage that doesn't + // support hybrid this falls back to a zero-vector similarity search, + // which is mostly useless; callers should pick a different mode. + if (!kb.supportsHybridSearch()) { + throw new Error( + `searchMode "text" needs hybrid-capable storage; install a backend with hybridSearch.` + ); + } + const dummy = new Float32Array(kb.getVectorDimensions()); + return kb.hybridSearch(dummy, { + textQuery: query, + topK, + filter, + scoreThreshold, + vectorWeight: 0, + }); + } + + const queryVec = await embedTexts([query], requireQueryEmbedModel(kb), runConfig); + const vector = queryVec[0]; + + if (mode === "similarity") { + return kb.similaritySearch(vector, { topK, filter, scoreThreshold }); + } + + if (mode === "hybrid") { + if (!kb.supportsHybridSearch()) { + // Graceful fallback — hybrid requested but backend doesn't have it. + return kb.similaritySearch(vector, { topK, filter, scoreThreshold }); + } + return kb.hybridSearch(vector, { + textQuery: query, + topK, + filter, + scoreThreshold, + }); + } + + // mode === "rerank" + // First-stage pool is `topK * firstStageMultiplier`, but never + // smaller than `firstStageMinimum`. The floor matters for small + // `topK`: with `topK=1, multiplier=5` the raw product is 5, which + // robs the reranker of any real choice. The minimum keeps the + // candidate pool meaningful regardless of how small `topK` is. + const firstStageTopK = Math.max(topK * firstStageMultiplier, firstStageMinimum); + const firstStage: ChunkSearchResult[] = kb.supportsHybridSearch() + ? await kb.hybridSearch(vector, { + textQuery: query, + topK: firstStageTopK, + filter, + scoreThreshold, + }) + : await kb.similaritySearch(vector, { + topK: firstStageTopK, + filter, + scoreThreshold, + }); + if (firstStage.length === 0) return []; + + // `chunkText` enforces the metadata.text contract — chunks missing + // text throw with the offending chunk_id rather than silently + // feeding `JSON.stringify(metadata)` to the reranker, which would + // produce meaningless relevance scores. + const docs = firstStage.map(chunkText); + + // Note: `scoreThreshold` is intentionally NOT honored in the rerank + // branch. The first stage already filtered by score; cross-encoder + // logits live on a completely different scale (often negative) and + // a cosine-style threshold would either drop everything or nothing. + // Callers wanting a rerank-relative cutoff should clip on the + // returned `score` themselves after inspecting `scoreType`. + if (kb.rerankerModel) { + const result = await new TextRerankerTask().run( + { + query, + documents: docs, + model: kb.rerankerModel, + topK, + }, + runConfig + ); + const indices = (result.indices as number[]) ?? []; + const scores = (result.scores as number[]) ?? []; + return indices.map((idx) => { + const candidate = firstStage[idx]; + const newScore = scores[idx]; + return { + ...candidate, + score: typeof newScore === "number" ? newScore : candidate.score, + scoreType: "rerank" as const, + }; + }); + } + + // No reranker model configured but mode is "rerank" — fall back to a + // local heuristic so callers still get a usable ordering. We still + // tag the result with scoreType: "rerank" because callers asked for + // rerank semantics; the score scale isn't comparable to cosine/RRF. + const heuristic = await new RerankerTask().run( + { + query, + chunks: docs, + scores: firstStage.map((c) => c.score), + metadata: firstStage.map((c) => c.metadata as Record), + topK, + method: "simple", + }, + runConfig + ); + const indices = (heuristic.originalIndices as number[]) ?? []; + const newScores = (heuristic.scores as number[]) ?? []; + return indices.map((idx, rank) => { + const candidate = firstStage[idx]; + return { + ...candidate, + score: newScores[rank] ?? candidate.score, + scoreType: "rerank" as const, + }; + }); + }, + }; +} diff --git a/packages/ai/src/task/ChunkRetrievalTask.ts b/packages/ai/src/task/ChunkRetrievalTask.ts index 4e4617946..7bdf03fef 100644 --- a/packages/ai/src/task/ChunkRetrievalTask.ts +++ b/packages/ai/src/task/ChunkRetrievalTask.ts @@ -141,12 +141,14 @@ const outputSchema = { }, scoreType: { type: "string", - enum: ["cosine", "bm25", "rrf"], + enum: ["cosine", "bm25", "rrf", "rerank"], title: "Score Type", description: "Discriminator naming the scorer used for `scores`: 'cosine' for similarity search " + "and for hybrid fallback when the text query is empty/whitespace; 'rrf' for hybrid " + - "fusion. ('bm25' is reserved for direct text search and is not produced by this task.)", + "fusion. ('bm25' is reserved for direct text search and is not produced by this task. " + + "'rerank' is produced by the standard KB strategy after cross-encoder reranking and " + + "is not produced by this task either.)", }, vectors: { type: "array", @@ -289,9 +291,14 @@ export class ChunkRetrievalTask extends Task< // want to surface that to callers even when the result set is empty. const hybridFallsBackToCosine = method === "hybrid" && (queryText === undefined || queryText.trim().length === 0); - const defaultScoreType: "cosine" | "bm25" | "rrf" = + // `ChunkRetrievalTask` itself only produces cosine or RRF scores; it + // can't emit "bm25" (no text-only path) or "rerank" (that comes from + // a downstream reranker, not this task). The output-schema enum + // includes them so the field is consistent with the canonical + // `ScoreType` union, but they won't appear in `defaultScoreType`. + const defaultScoreType: "cosine" | "rrf" = method === "hybrid" && !hybridFallsBackToCosine ? "rrf" : "cosine"; - const scoreType = + const scoreType: "cosine" | "bm25" | "rrf" | "rerank" = results.length > 0 ? (results[0].scoreType ?? defaultScoreType) : defaultScoreType; const output: ChunkRetrievalTaskOutput = { diff --git a/packages/ai/src/task/KbAddDocumentTask.ts b/packages/ai/src/task/KbAddDocumentTask.ts new file mode 100644 index 000000000..203b68c6f --- /dev/null +++ b/packages/ai/src/task/KbAddDocumentTask.ts @@ -0,0 +1,108 @@ +/** + * @license + * Copyright 2026 Steven Roussey + * SPDX-License-Identifier: Apache-2.0 + */ + +import type { KnowledgeBase } from "@workglow/knowledge-base"; +import { Document, TypeKnowledgeBase } from "@workglow/knowledge-base"; +import type { IRunConfig, TaskConfig } from "@workglow/task-graph"; +import { CreateWorkflow, IExecuteContext, Task, Workflow } from "@workglow/task-graph"; +import type { DataPortSchema, FromSchema } from "@workglow/util/schema"; +import type { Capability } from "../capability/Capabilities"; + +const inputSchema = { + type: "object", + properties: { + knowledgeBase: TypeKnowledgeBase({ + title: "Knowledge Base", + description: "Knowledge base to add the document to.", + }), + document: { + title: "Document", + description: "The Document instance to chunk, embed, and store.", + additionalProperties: true, + }, + }, + required: ["knowledgeBase", "document"], + additionalProperties: false, +} as const satisfies DataPortSchema; + +const outputSchema = { + type: "object", + properties: { + doc_id: { + type: "string", + title: "Document ID", + description: "The stored document ID.", + }, + }, + required: ["doc_id"], + additionalProperties: false, +} as const satisfies DataPortSchema; + +export type KbAddDocumentTaskInput = Omit, "document"> & { + readonly document: Document; +}; +export type KbAddDocumentTaskOutput = FromSchema; +export type KbAddDocumentTaskConfig = TaskConfig; + +/** + * Ingest a document into a knowledge base end-to-end: chunk, embed, and + * store via the KB's installed strategy. Threads the task's execution + * context (signal, resourceScope, registry) into the KB call so model + * resources are shared and abort signals propagate. + */ +export class KbAddDocumentTask extends Task< + KbAddDocumentTaskInput, + KbAddDocumentTaskOutput, + KbAddDocumentTaskConfig +> { + public static override type = "KbAddDocumentTask"; + public static readonly requires: readonly Capability[] = [] as const satisfies Capability[]; + public static override category = "RAG"; + public static override title = "KB Add Document"; + public static override description = + "Ingest a document into a knowledge base: chunk, embed, and store via the KB's installed strategy."; + public static override cacheable = false; + + public static override inputSchema(): DataPortSchema { + return inputSchema as DataPortSchema; + } + public static override outputSchema(): DataPortSchema { + return outputSchema as DataPortSchema; + } + + override async execute( + input: KbAddDocumentTaskInput, + context: IExecuteContext + ): Promise { + const kb = input.knowledgeBase as KnowledgeBase; + const stored = await kb.upsert(input.document, { + signal: context.signal, + resourceScope: context.resourceScope, + registry: context.registry, + }); + return { doc_id: stored.doc_id! }; + } +} + +export const kbAddDocument = ( + input: KbAddDocumentTaskInput, + config?: KbAddDocumentTaskConfig, + runConfig?: Partial +) => { + return new KbAddDocumentTask(config).run(input, runConfig); +}; + +declare module "@workglow/task-graph" { + interface Workflow { + kbAddDocument: CreateWorkflow< + KbAddDocumentTaskInput, + KbAddDocumentTaskOutput, + KbAddDocumentTaskConfig + >; + } +} + +Workflow.prototype.kbAddDocument = CreateWorkflow(KbAddDocumentTask); diff --git a/packages/ai/src/task/KbDeleteTask.ts b/packages/ai/src/task/KbDeleteTask.ts new file mode 100644 index 000000000..fef1d1621 --- /dev/null +++ b/packages/ai/src/task/KbDeleteTask.ts @@ -0,0 +1,96 @@ +/** + * @license + * Copyright 2026 Steven Roussey + * SPDX-License-Identifier: Apache-2.0 + */ + +import type { KnowledgeBase } from "@workglow/knowledge-base"; +import { TypeKnowledgeBase } from "@workglow/knowledge-base"; +import type { IRunConfig, TaskConfig } from "@workglow/task-graph"; +import { CreateWorkflow, IExecuteContext, Task, Workflow } from "@workglow/task-graph"; +import type { DataPortSchema, FromSchema } from "@workglow/util/schema"; +import type { Capability } from "../capability/Capabilities"; + +const inputSchema = { + type: "object", + properties: { + knowledgeBase: TypeKnowledgeBase({ + title: "Knowledge Base", + description: "Knowledge base to delete the document from.", + }), + doc_id: { + type: "string", + title: "Document ID", + description: "ID of the document to delete.", + }, + }, + required: ["knowledgeBase", "doc_id"], + additionalProperties: false, +} as const satisfies DataPortSchema; + +const outputSchema = { + type: "object", + properties: { + doc_id: { + type: "string", + title: "Document ID", + description: "ID of the deleted document (echoed for pipeline composability).", + }, + }, + required: ["doc_id"], + additionalProperties: false, +} as const satisfies DataPortSchema; + +export type KbDeleteTaskInput = FromSchema; +export type KbDeleteTaskOutput = FromSchema; +export type KbDeleteTaskConfig = TaskConfig; + +/** + * Delete a document and its chunks from a knowledge base via the KB's + * installed strategy. Echoes `doc_id` so the task is composable in + * pipelines that need to pass the id to a downstream step. + */ +export class KbDeleteTask extends Task { + public static override type = "KbDeleteTask"; + public static readonly requires: readonly Capability[] = [] as const satisfies Capability[]; + public static override category = "RAG"; + public static override title = "KB Delete Document"; + public static override description = "Delete a document and its chunks from a knowledge base."; + public static override cacheable = false; + + public static override inputSchema(): DataPortSchema { + return inputSchema as DataPortSchema; + } + public static override outputSchema(): DataPortSchema { + return outputSchema as DataPortSchema; + } + + override async execute( + input: KbDeleteTaskInput, + context: IExecuteContext + ): Promise { + const kb = input.knowledgeBase as KnowledgeBase; + await kb.delete(input.doc_id, { + signal: context.signal, + resourceScope: context.resourceScope, + registry: context.registry, + }); + return { doc_id: input.doc_id }; + } +} + +export const kbDelete = ( + input: KbDeleteTaskInput, + config?: KbDeleteTaskConfig, + runConfig?: Partial +) => { + return new KbDeleteTask(config).run(input, runConfig); +}; + +declare module "@workglow/task-graph" { + interface Workflow { + kbDelete: CreateWorkflow; + } +} + +Workflow.prototype.kbDelete = CreateWorkflow(KbDeleteTask); diff --git a/packages/ai/src/task/KbReindexTask.ts b/packages/ai/src/task/KbReindexTask.ts new file mode 100644 index 000000000..ec8f00f50 --- /dev/null +++ b/packages/ai/src/task/KbReindexTask.ts @@ -0,0 +1,96 @@ +/** + * @license + * Copyright 2026 Steven Roussey + * SPDX-License-Identifier: Apache-2.0 + */ + +import type { KnowledgeBase } from "@workglow/knowledge-base"; +import { TypeKnowledgeBase } from "@workglow/knowledge-base"; +import type { IRunConfig, TaskConfig } from "@workglow/task-graph"; +import { CreateWorkflow, IExecuteContext, Task, Workflow } from "@workglow/task-graph"; +import type { DataPortSchema, FromSchema } from "@workglow/util/schema"; +import type { Capability } from "../capability/Capabilities"; + +const inputSchema = { + type: "object", + properties: { + knowledgeBase: TypeKnowledgeBase({ + title: "Knowledge Base", + description: "Knowledge base to re-index", + }), + }, + required: ["knowledgeBase"], + additionalProperties: false, +} as const satisfies DataPortSchema; + +const outputSchema = { + type: "object", + properties: { + count: { + type: "number", + title: "Documents Re-indexed", + description: "Number of documents re-indexed", + }, + }, + required: ["count"], + additionalProperties: false, +} as const satisfies DataPortSchema; + +export type KbReindexTaskInput = FromSchema; +export type KbReindexTaskOutput = FromSchema; +export type KbReindexTaskConfig = TaskConfig; + +/** + * Re-index every document in a knowledge base via its installed AI strategy. + * The strategy handles chunking + embedding using the KB's configured + * `docEmbeddingModel`. Replaces the multi-task embed workflow pattern. + */ +export class KbReindexTask extends Task< + KbReindexTaskInput, + KbReindexTaskOutput, + KbReindexTaskConfig +> { + public static override type = "KbReindexTask"; + public static override category = "RAG"; + public static override title = "KB Reindex"; + public static override description = + "Re-chunk and re-embed every document in a knowledge base using its configured models."; + public static readonly requires: readonly Capability[] = [] as const satisfies Capability[]; + public static override cacheable = false; + + public static override inputSchema(): DataPortSchema { + return inputSchema as DataPortSchema; + } + public static override outputSchema(): DataPortSchema { + return outputSchema as DataPortSchema; + } + + override async execute( + input: KbReindexTaskInput, + context: IExecuteContext + ): Promise { + const kb = input.knowledgeBase as KnowledgeBase; + const count = await kb.reindex({ + signal: context.signal, + resourceScope: context.resourceScope, + registry: context.registry, + }); + return { count }; + } +} + +export const kbReindex = async ( + input: KbReindexTaskInput, + config?: KbReindexTaskConfig, + runConfig?: Partial +) => { + return new KbReindexTask(config).run(input, runConfig); +}; + +declare module "@workglow/task-graph" { + interface Workflow { + kbReindex: CreateWorkflow; + } +} + +Workflow.prototype.kbReindex = CreateWorkflow(KbReindexTask); diff --git a/packages/ai/src/task/KbSearchTask.ts b/packages/ai/src/task/KbSearchTask.ts index 2297e96d2..571351ff1 100644 --- a/packages/ai/src/task/KbSearchTask.ts +++ b/packages/ai/src/task/KbSearchTask.ts @@ -5,7 +5,7 @@ */ import type { ChunkSearchResult, KnowledgeBase } from "@workglow/knowledge-base"; -import { TypeKnowledgeBase } from "@workglow/knowledge-base"; +import { chunkText, TypeKnowledgeBase } from "@workglow/knowledge-base"; import type { IRunConfig, TaskConfig } from "@workglow/task-graph"; import { CreateWorkflow, IExecuteContext, Task, Workflow } from "@workglow/task-graph"; import type { DataPortSchema, FromSchema } from "@workglow/util/schema"; @@ -16,24 +16,30 @@ const inputSchema = { properties: { knowledgeBase: TypeKnowledgeBase({ title: "Knowledge Base", - description: "The knowledge base instance to search in", + description: "Knowledge base to search.", }), query: { type: "string", title: "Query", - description: "Search query (the KB's onSearch handles embedding internally)", + description: "Search query text.", }, topK: { type: "number", title: "Top K", - description: "Number of top results to return", + description: "Number of top results to return.", minimum: 1, default: 5, }, filter: { type: "object", title: "Metadata Filter", - description: "Filter results by metadata fields", + description: "Filter results by chunk metadata fields.", + }, + scoreThreshold: { + type: "number", + title: "Score Threshold", + description: + "Minimum score to include a result. Honored by similarity and hybrid search modes; ignored by the strategy in rerank mode (cross-encoder logits aren't comparable to cosine/RRF scores). Not constrained to a range — some embedding models (un-normalized dot product) produce scores outside [0, 1].", }, }, required: ["knowledgeBase", "query"], @@ -53,28 +59,49 @@ const outputSchema = { title: "Results", description: "Matching chunks in score-desc order", }, + chunks: { + type: "array", + items: { type: "string" }, + title: "Chunks", + description: "The chunk text content, parallel to `results`", + }, + chunk_ids: { + type: "array", + items: { type: "string" }, + title: "Chunk IDs", + description: "The chunk ids, parallel to `results`", + }, + scores: { + type: "array", + items: { type: "number" }, + title: "Scores", + description: "Scores parallel to `results`", + }, count: { type: "number", title: "Count", description: "Number of results returned", }, }, - required: ["results", "count"], + required: ["results", "chunks", "chunk_ids", "scores", "count"], additionalProperties: false, } as const satisfies DataPortSchema; export type KbSearchTaskInput = FromSchema; export type KbSearchTaskOutput = { readonly results: ChunkSearchResult[]; + readonly chunks: string[]; + readonly chunk_ids: string[]; + readonly scores: number[]; readonly count: number; }; export type KbSearchTaskConfig = TaskConfig; /** - * Observable wrapper around `kb.search(text, opts)` — the KB's `onSearch` - * callback handles embedding and any custom retrieval logic. Distinct from - * `ChunkRetrievalTask`, which embeds via an explicit model and calls - * `kb.similaritySearch(vector)` (bypassing `onSearch`). + * High-level KB search task. Delegates to `kb.search(query, options)`; the + * KB and its installed strategy decide everything else (embedding model, + * retrieval mode, rerank). No model or kind input here — that's the + * point of moving the config onto the KB itself. */ export class KbSearchTask extends Task { public static override type = "KbSearchTask"; @@ -88,7 +115,7 @@ export class KbSearchTask extends Task { - const { knowledgeBase, query, topK = 5, filter } = input; + const { knowledgeBase, query, topK = 5, filter, scoreThreshold } = input; const kb = knowledgeBase as KnowledgeBase; - const results = await kb.search(query, { topK, filter }); - return { results, count: results.length }; + const results = await kb.search( + query, + { topK, filter, scoreThreshold }, + { + signal: context.signal, + resourceScope: context.resourceScope, + registry: context.registry, + } + ); + return { + results, + chunks: results.map(chunkText), + chunk_ids: results.map((r) => r.chunk_id), + scores: results.map((r) => r.score), + count: results.length, + }; } } diff --git a/packages/ai/src/task/RerankerTask.ts b/packages/ai/src/task/RerankerTask.ts index d5ceabd8d..c2a3f5708 100644 --- a/packages/ai/src/task/RerankerTask.ts +++ b/packages/ai/src/task/RerankerTask.ts @@ -50,7 +50,7 @@ const inputSchema = { type: "string", enum: ["reciprocal-rank-fusion", "simple"], title: "Reranking Method", - description: "Method to use for reranking", + description: "Heuristic reranking method", default: "simple", }, }, @@ -144,7 +144,7 @@ export class RerankerTask extends Task { const { query, chunks, scores = [], metadata = [], topK, method = "simple" } = input; @@ -217,7 +217,9 @@ export class RerankerTask extends Task + * SPDX-License-Identifier: Apache-2.0 + */ + +import type { TaskConfig } from "@workglow/task-graph"; +import { CreateWorkflow, Workflow } from "@workglow/task-graph"; +import type { DataPortSchema, FromSchema } from "@workglow/util/schema"; +import type { Capability } from "../capability/Capabilities"; +import { AiTask } from "./base/AiTask"; +import { TypeModel } from "./base/AiTaskSchemas"; + +const inputSchema = { + type: "object", + properties: { + query: { + type: "string", + title: "Query", + description: "The query to score documents against", + }, + documents: { + type: "array", + items: { type: "string" }, + title: "Documents", + description: "Candidate documents to score", + }, + topK: { + type: "number", + title: "Top K", + description: "Return at most this many results (default: all)", + minimum: 1, + }, + model: TypeModel("model:TextRerankerTask", { + title: "Reranker Model", + description: "Cross-encoder reranker model (e.g. bge-reranker, Cohere rerank). Required.", + }), + }, + required: ["query", "documents", "model"], + additionalProperties: false, +} as const satisfies DataPortSchema; + +const outputSchema = { + type: "object", + properties: { + scores: { + type: "array", + items: { type: "number" }, + title: "Scores", + description: "Relevance score for each document, in the original order", + }, + indices: { + type: "array", + items: { type: "number" }, + title: "Indices", + description: "Indices of documents sorted best-first (length = topK if set)", + }, + }, + required: ["scores", "indices"], + additionalProperties: false, +} as const satisfies DataPortSchema; + +export type TextRerankerTaskInput = FromSchema; +export type TextRerankerTaskOutput = FromSchema; +export type TextRerankerTaskConfig = TaskConfig; + +/** + * Thrown by reranker provider run-fns when the underlying ML pipeline + * returns output that doesn't match the expected `{ label, score }` + * shape (or array thereof when `top_k > 1`). Co-located with the task + * definition so callers can `instanceof`-test against a single import + * regardless of which provider is installed. + * + * `actualShape` is a truncated, JSON-stringified snippet of the offending + * entry — enough to point an operator at the misconfigured model without + * dumping arbitrary tensors into logs. + */ +export class KbRerankerOutputError extends Error { + public readonly actualShape: unknown; + constructor(message: string, actualShape: unknown) { + super(message); + this.name = "KbRerankerOutputError"; + this.actualShape = actualShape; + } +} + +/** + * AiTask for cross-encoder reranking. Providers register a run-fn for this + * task type (e.g. HuggingFace Transformers using a `text-classification` + * cross-encoder pipeline on `[query, doc]` pairs). `createStandardKbStrategy` + * invokes this task as the rerank stage of `kb.search()` when the KB is + * configured with `searchMode: "rerank"` and has a `rerankerModel` set. + */ +export class TextRerankerTask extends AiTask< + TextRerankerTaskInput, + TextRerankerTaskOutput, + TextRerankerTaskConfig +> { + public static override type = "TextRerankerTask"; + public static override category = "RAG"; + public static override title = "Text Reranker"; + public static override description = + "Score documents against a query using a cross-encoder reranker model"; + public static override readonly requires = ["text.reranking"] as const satisfies Capability[]; + + public static override inputSchema(): DataPortSchema { + return inputSchema as DataPortSchema; + } + public static override outputSchema(): DataPortSchema { + return outputSchema as DataPortSchema; + } +} + +export const textReranker = async ( + input: TextRerankerTaskInput, + config?: TextRerankerTaskConfig +) => { + return new TextRerankerTask(config).run(input); +}; + +declare module "@workglow/task-graph" { + interface Workflow { + textReranker: CreateWorkflow< + TextRerankerTaskInput, + TextRerankerTaskOutput, + TextRerankerTaskConfig + >; + } +} + +Workflow.prototype.textReranker = CreateWorkflow(TextRerankerTask); diff --git a/packages/ai/src/task/index.ts b/packages/ai/src/task/index.ts index e4f315aee..4dc2438dc 100644 --- a/packages/ai/src/task/index.ts +++ b/packages/ai/src/task/index.ts @@ -36,6 +36,9 @@ export * from "./ImageClassificationTask"; export * from "./ImageEmbeddingTask"; export * from "./ImageSegmentationTask"; export * from "./ImageToTextTask"; +export * from "./KbAddDocumentTask"; +export * from "./KbDeleteTask"; +export * from "./KbReindexTask"; export * from "./KbSearchTask"; export * from "./KbToDocumentsTask"; export * from "./MessageConversion"; @@ -57,6 +60,7 @@ export * from "./TextGenerationTask"; export * from "./TextLanguageDetectionTask"; export * from "./TextNamedEntityRecognitionTask"; export * from "./TextQuestionAnswerTask"; +export * from "./TextRerankerTask"; export * from "./TextRewriterTask"; export * from "./TextSummaryTask"; export * from "./TextTranslationTask"; diff --git a/packages/ai/src/task/registerAiTasks.ts b/packages/ai/src/task/registerAiTasks.ts index d602de8da..7c469e2ec 100644 --- a/packages/ai/src/task/registerAiTasks.ts +++ b/packages/ai/src/task/registerAiTasks.ts @@ -26,6 +26,9 @@ import { ImageClassificationTask } from "./ImageClassificationTask"; import { ImageEmbeddingTask } from "./ImageEmbeddingTask"; import { ImageSegmentationTask } from "./ImageSegmentationTask"; import { ImageToTextTask } from "./ImageToTextTask"; +import { KbAddDocumentTask } from "./KbAddDocumentTask"; +import { KbDeleteTask } from "./KbDeleteTask"; +import { KbReindexTask } from "./KbReindexTask"; import { KbSearchTask } from "./KbSearchTask"; import { KbToDocumentsTask } from "./KbToDocumentsTask"; import { ModelDownloadRemoveTask } from "./ModelDownloadRemoveTask"; @@ -46,6 +49,7 @@ import { TextGenerationTask } from "./TextGenerationTask"; import { TextLanguageDetectionTask } from "./TextLanguageDetectionTask"; import { TextNamedEntityRecognitionTask } from "./TextNamedEntityRecognitionTask"; import { TextQuestionAnswerTask } from "./TextQuestionAnswerTask"; +import { TextRerankerTask } from "./TextRerankerTask"; import { TextRewriterTask } from "./TextRewriterTask"; import { TextSummaryTask } from "./TextSummaryTask"; import { TextTranslationTask } from "./TextTranslationTask"; @@ -77,12 +81,15 @@ export const registerAiTasks = () => { HandLandmarkerTask, HierarchicalChunkerTask, HierarchyJoinTask, + KbAddDocumentTask, + KbDeleteTask, KbSearchTask, KbToDocumentsTask, ImageClassificationTask, ImageEmbeddingTask, ImageSegmentationTask, ImageToTextTask, + KbReindexTask, ModelInfoTask, ModelSearchTask, ObjectDetectionTask, @@ -99,6 +106,7 @@ export const registerAiTasks = () => { TextLanguageDetectionTask, TextNamedEntityRecognitionTask, TextQuestionAnswerTask, + TextRerankerTask, TextRewriterTask, TextSummaryTask, TextTranslationTask, diff --git a/packages/browser-control/CHANGELOG.md b/packages/browser-control/CHANGELOG.md index 9966d1c74..6efecd297 100644 --- a/packages/browser-control/CHANGELOG.md +++ b/packages/browser-control/CHANGELOG.md @@ -1,5 +1,7 @@ # Changelog +## 0.2.36 + ## 0.2.35 ### Chores diff --git a/packages/browser-control/package.json b/packages/browser-control/package.json index aa35fe950..38fa3f3b4 100644 --- a/packages/browser-control/package.json +++ b/packages/browser-control/package.json @@ -1,7 +1,7 @@ { "name": "@workglow/browser-control", "type": "module", - "version": "0.2.35", + "version": "0.2.36", "repository": { "type": "git", "url": "https://github.com/workglow-dev/libs.git", diff --git a/packages/indexeddb/CHANGELOG.md b/packages/indexeddb/CHANGELOG.md index 19dfb7867..9d34bcfbc 100644 --- a/packages/indexeddb/CHANGELOG.md +++ b/packages/indexeddb/CHANGELOG.md @@ -1,5 +1,7 @@ # Changelog +## 0.2.36 + ## 0.2.35 ### Features diff --git a/packages/indexeddb/package.json b/packages/indexeddb/package.json index d61150f1d..7884de4ea 100644 --- a/packages/indexeddb/package.json +++ b/packages/indexeddb/package.json @@ -1,7 +1,7 @@ { "name": "@workglow/indexeddb", "type": "module", - "version": "0.2.35", + "version": "0.2.36", "repository": { "type": "git", "url": "https://github.com/workglow-dev/libs.git", diff --git a/packages/javascript/CHANGELOG.md b/packages/javascript/CHANGELOG.md index 6c04250ac..9e855980d 100644 --- a/packages/javascript/CHANGELOG.md +++ b/packages/javascript/CHANGELOG.md @@ -1,5 +1,7 @@ # Changelog +## 0.2.36 + ## 0.2.35 ### Features diff --git a/packages/javascript/package.json b/packages/javascript/package.json index bcfcc70fe..491fafc7e 100644 --- a/packages/javascript/package.json +++ b/packages/javascript/package.json @@ -2,7 +2,7 @@ "name": "@workglow/javascript", "type": "module", "sideEffects": false, - "version": "0.2.35", + "version": "0.2.36", "repository": { "type": "git", "url": "https://github.com/workglow-dev/libs.git", diff --git a/packages/job-queue/CHANGELOG.md b/packages/job-queue/CHANGELOG.md index da26569dc..b4bb43c70 100644 --- a/packages/job-queue/CHANGELOG.md +++ b/packages/job-queue/CHANGELOG.md @@ -1,5 +1,7 @@ # @workglow/job-queue +## 0.2.36 + ## 0.2.35 ### Performance diff --git a/packages/job-queue/package.json b/packages/job-queue/package.json index a29e7704e..53799f390 100644 --- a/packages/job-queue/package.json +++ b/packages/job-queue/package.json @@ -2,7 +2,7 @@ "name": "@workglow/job-queue", "type": "module", "sideEffects": false, - "version": "0.2.35", + "version": "0.2.36", "repository": { "type": "git", "url": "https://github.com/workglow-dev/libs.git", diff --git a/packages/knowledge-base/CHANGELOG.md b/packages/knowledge-base/CHANGELOG.md index f2fe73c7c..06f34b929 100644 --- a/packages/knowledge-base/CHANGELOG.md +++ b/packages/knowledge-base/CHANGELOG.md @@ -1,5 +1,13 @@ # @workglow/knowledge-base +## 0.2.36 + +### Features + +#### kb + +- pluggable strategy with model config, IRunConfig threading, and document tasks + ## 0.2.35 ### Features diff --git a/packages/knowledge-base/package.json b/packages/knowledge-base/package.json index d7d2ade48..6268cdf32 100644 --- a/packages/knowledge-base/package.json +++ b/packages/knowledge-base/package.json @@ -1,7 +1,7 @@ { "name": "@workglow/knowledge-base", "type": "module", - "version": "0.2.35", + "version": "0.2.36", "repository": { "type": "git", "url": "https://github.com/workglow-dev/libs.git", @@ -27,18 +27,23 @@ }, "peerDependencies": { "@workglow/storage": "workspace:*", + "@workglow/task-graph": "workspace:*", "@workglow/util": "workspace:*" }, "peerDependenciesMeta": { "@workglow/storage": { "optional": false }, + "@workglow/task-graph": { + "optional": false + }, "@workglow/util": { "optional": false } }, "devDependencies": { "@workglow/storage": "workspace:*", + "@workglow/task-graph": "workspace:*", "@workglow/util": "workspace:*" }, "exports": { diff --git a/packages/knowledge-base/src/chunk/ChunkVectorStorageSchema.ts b/packages/knowledge-base/src/chunk/ChunkVectorStorageSchema.ts index ba20e5989..ba6da21b8 100644 --- a/packages/knowledge-base/src/chunk/ChunkVectorStorageSchema.ts +++ b/packages/knowledge-base/src/chunk/ChunkVectorStorageSchema.ts @@ -5,8 +5,8 @@ */ import { IVectorStorage } from "@workglow/storage"; -import { TypedArraySchema } from "@workglow/util/schema"; import type { DataPortSchemaObject, TypedArray } from "@workglow/util/schema"; +import { TypedArraySchema } from "@workglow/util/schema"; import type { ChunkRecord } from "./ChunkSchema"; /** @@ -41,6 +41,18 @@ export interface ChunkVectorEntity< /** * Type for inserting chunk vectors - chunk_id is optional (auto-generated) + * + * @remarks + * `metadata.text` is a load-bearing field — it carries the chunk's + * canonical text (the same string that was embedded to produce + * {@link ChunkVectorEntity.vector}). Downstream callers — notably + * cross-encoder rerankers and any UI that displays the chunk — read + * `metadata.text` directly via {@link chunkText}. Strategies that build + * `InsertChunkVectorEntity` from custom chunkers MUST populate + * `metadata.text` or rerank/display paths will throw. The standard + * strategy populates it via `toInsertChunkEntities` from + * `HierarchicalChunkerTask` output, which always emits `text` on each + * chunk. */ export type InsertChunkVectorEntity< Metadata extends ChunkRecord = ChunkRecord, @@ -63,7 +75,7 @@ export type ChunkVectorStorage = IVectorStorage< /** * Discriminator for the scoring function used to produce a * {@link ChunkSearchResult.score}. Callers (typically UI) use this to render - * the score appropriately, since the three scorers live on different scales: + * the score appropriately, since the scorers live on different scales: * * - `"cosine"`: cosine similarity in `[-1, 1]`, typically `[0, 1]` for text * embeddings. Absolute — higher means more similar. @@ -73,8 +85,14 @@ export type ChunkVectorStorage = IVectorStorage< * `2 / (rrfK + 1)` (~`0.033` with the default `rrfK=60`). Rank-based, not * absolute — the magnitude is not a similarity, only an ordering signal. * Not comparable across queries. + * - `"rerank"`: cross-encoder reranker output (e.g. bge-reranker, Cohere + * rerank). Raw logit, not a probability and not comparable to cosine / + * BM25 / RRF scores. Callers MUST inspect `scoreType` before applying + * any score-threshold gate; cross-encoder scores often span wide negative + * ranges that look invalid under a cosine-style threshold but are + * perfectly normal here. */ -export type ScoreType = "cosine" | "bm25" | "rrf"; +export type ScoreType = "cosine" | "bm25" | "rrf" | "rerank"; /** * Search result with score @@ -83,3 +101,26 @@ export type ChunkSearchResult = ChunkVectorEntity & { score: number; scoreType?: ScoreType; }; + +/** + * Extract the canonical chunk text from a search result. + * + * Reads `metadata.text` directly. Throws (with the offending chunk_id) if + * the field is missing — chunks without text can't be reranked, displayed, + * or fed into downstream NLP tasks. Use this helper everywhere a chunk's + * text is needed instead of inlining `metadata.text` access; it keeps the + * contract — "every chunk in the KB owns its source text in + * `metadata.text`" — enforced at exactly one place. See + * {@link InsertChunkVectorEntity} for the writer-side requirement. + */ +export function chunkText(c: { chunk_id: string; metadata?: ChunkRecord }): string { + const text = c.metadata?.text; + if (typeof text !== "string") { + throw new Error( + `chunkText: chunk ${c.chunk_id} is missing metadata.text. ` + + `Every chunk in a KnowledgeBase must carry its source text on metadata.text — ` + + `update the chunker / strategy that produced this chunk to populate it.` + ); + } + return text; +} diff --git a/packages/knowledge-base/src/common.ts b/packages/knowledge-base/src/common.ts index 9d9466df9..f76d6558f 100644 --- a/packages/knowledge-base/src/common.ts +++ b/packages/knowledge-base/src/common.ts @@ -8,6 +8,7 @@ export * from "./chunk/ChunkSchema"; export * from "./chunk/ChunkVectorStorageSchema"; +export * from "./knowledge-base/IKbAiStrategy"; export * from "./knowledge-base/KnowledgeBase"; export * from "./knowledge-base/KnowledgeBaseSchema"; export * from "./knowledge-base/KnowledgeBaseRepository"; diff --git a/packages/knowledge-base/src/knowledge-base/IKbAiStrategy.ts b/packages/knowledge-base/src/knowledge-base/IKbAiStrategy.ts new file mode 100644 index 000000000..71e43f479 --- /dev/null +++ b/packages/knowledge-base/src/knowledge-base/IKbAiStrategy.ts @@ -0,0 +1,167 @@ +/** + * @license + * Copyright 2026 Steven Roussey + * SPDX-License-Identifier: Apache-2.0 + */ + +import type { IRunConfig } from "@workglow/task-graph"; +import type { TypedArray } from "@workglow/util/schema"; +import type { ChunkRecord } from "../chunk/ChunkSchema"; +import type { + ChunkSearchResult, + ChunkVectorEntity, + InsertChunkVectorEntity, +} from "../chunk/ChunkVectorStorageSchema"; +import type { Document } from "../document/Document"; +import type { ISearchOptions } from "./KnowledgeBase"; + +/** + * Strategy that bridges a {@link KnowledgeBase} to an AI runtime. The + * strategy is the single extension point: a KB has exactly one installed + * strategy, and `kb.upsert` / `kb.delete` / `kb.search` delegate to it. + * + * Two flavors ship in `@workglow/ai`: + * - `createStandardKbStrategy(...)` — defaults parameterized by chunker + * strategy and search mode; reads the KB's model IDs at op time. + * - Custom — write your own to add scoping, alternative chunkers, or + * unusual retrieval flows. The builder's KBs use a custom strategy on + * top of `ScopedKnowledgeBase` so user/project ids ride along. + * + * Strategies receive the KB instance on every call. Calls into + * `kb.upsertDocument` / `kb.upsertChunksBulk` / `kb.similaritySearch` etc. + * go through virtual dispatch — subclasses (e.g. `ScopedKnowledgeBase`) + * can intercept the low-level ops without the strategy knowing. + * + * **Trust model.** Strategies are TRUSTED CODE. An installed strategy + * receives an {@link IKbStrategyTarget} that exposes the KB's full + * low-level storage surface (`upsertDocument`, `deleteChunksForDocument`, + * `upsertChunksBulk`, `similaritySearch`, `hybridSearch`). These + * operations bypass any application-level access control (e.g. + * user/project scoping enforced by `ScopedKnowledgeBase`) because + * scoping is implemented via virtual dispatch on the *target instance*, + * and a malicious or buggy strategy can violate the contract by routing + * data through a different KB. Do NOT load strategies from untrusted + * sources, user input, or remote registries. Install only strategies you + * ship in trusted application code. + */ +export interface IKbAiStrategy { + /** + * Ingest a single document: chunk + embed + write document + write + * chunks. The strategy decides chunker strategy, dedup behavior, + * embedding model, etc. Returns the stored document (possibly with a + * newly-assigned doc_id). + */ + ingest(kb: IKbStrategyTarget, doc: Document, runConfig?: Partial): Promise; + + /** + * Remove a document and its chunks. The default cascading delete works + * for most cases; override to add audit logging, soft delete, etc. + */ + delete(kb: IKbStrategyTarget, doc_id: string, runConfig?: Partial): Promise; + + /** + * Run a text query and return matching chunks. The strategy picks the + * retrieval flavor (similarity, hybrid, reranker, plain text) — callers + * don't choose per-call. + * + * The returned `score` is only comparable within a single result list, + * and only when results share a `scoreType`. The standard strategy + * tags rerank results with `scoreType: "rerank"` — cross-encoder + * logits are NOT comparable to cosine/BM25/RRF scores, so callers + * MUST inspect `scoreType` before applying any score threshold. In + * particular, `ISearchOptions.scoreThreshold` is not honored under + * `searchMode === "rerank"` because there's no meaningful default + * threshold across rerankers. + */ + search( + kb: IKbStrategyTarget, + query: string, + options?: ISearchOptions, + runConfig?: Partial + ): Promise; +} + +/** + * The narrow KB surface strategies operate against. Spells out exactly the + * building-block methods strategies need so the public KB API + * (`upsert`/`delete`/`search`) stays the only surface callers see. + * + * **Trust model.** This interface is the full low-level storage surface a + * strategy can reach: `upsertDocument`, `deleteChunksForDocument`, + * `upsertChunksBulk`, `similaritySearch`, `hybridSearch`. These + * operations bypass any application-level access control (e.g. + * user/project scoping enforced by `ScopedKnowledgeBase`) because + * scoping is implemented via virtual dispatch on the *target instance*, + * and a malicious or buggy strategy can violate the contract by routing + * data through a different KB. Treat installed strategies as TRUSTED + * CODE — do NOT load them from untrusted sources, user input, or remote + * registries. Install only strategies you ship in trusted application + * code. See {@link IKbAiStrategy} for the full trust model. + */ +export interface IKbStrategyTarget { + readonly name: string; + readonly docEmbeddingModel: string | undefined; + readonly queryEmbeddingModel: string | undefined; + readonly rerankerModel: string | undefined; + readonly chunkStrategy: ChunkStrategy | undefined; + readonly searchMode: SearchMode | undefined; + getVectorDimensions(): number; + supportsHybridSearch(): boolean; + /** Low-level: store a document JSON record without chunking. */ + upsertDocument(doc: Document): Promise; + /** Low-level: cascade delete a document + its chunks. */ + deleteDocument(doc_id: string): Promise; + /** Low-level: drop every chunk row for the given doc_id. */ + deleteChunksForDocument(doc_id: string): Promise; + /** Low-level: bulk-write chunk vectors. */ + upsertChunksBulk(chunks: InsertChunkVectorEntity[]): Promise; + /** Low-level: pure-vector retrieval. */ + similaritySearch( + query: TypedArray, + options?: { topK?: number; filter?: Readonly>; scoreThreshold?: number } + ): Promise; + /** Low-level: vector + full-text retrieval. */ + hybridSearch( + query: TypedArray, + options: { + readonly textQuery: string; + readonly topK?: number; + readonly filter?: Readonly>; + readonly scoreThreshold?: number; + readonly vectorWeight?: number; + } + ): Promise; +} + +/** Document-chunker strategy registered on the KB; consumed by ingest. */ +export type ChunkStrategy = "hierarchical" | "flat" | "sentence"; + +/** + * Retrieval mode registered on the KB; consumed by search. `text` is pure + * full-text (FTS) and bypasses embedding; the others require an embedding + * model (and `rerank` also requires `rerankerModel`). + */ +export type SearchMode = "text" | "similarity" | "hybrid" | "rerank"; + +/** + * Convert chunker output (chunks + parallel vectors) into the + * `InsertChunkVectorEntity` records that `kb.upsertChunksBulk()` expects. + * Shared helper so every strategy uses identical key derivation. + */ +export function toInsertChunkEntities( + result: { readonly chunks: ChunkRecord[]; readonly vectors: TypedArray[] }, + context: { readonly doc_id: string; readonly doc_title?: string } +): InsertChunkVectorEntity[] { + const { chunks, vectors } = result; + if (chunks.length !== vectors.length) { + throw new Error( + `Chunk/vector length mismatch: ${chunks.length} chunks but ${vectors.length} vectors` + ); + } + return chunks.map((chunk, i) => ({ + chunk_id: chunk.chunk_id, + doc_id: context.doc_id, + vector: vectors[i], + metadata: { ...chunk, doc_title: context.doc_title }, + })) as InsertChunkVectorEntity[]; +} diff --git a/packages/knowledge-base/src/knowledge-base/KnowledgeBase.ts b/packages/knowledge-base/src/knowledge-base/KnowledgeBase.ts index 49978820f..633e625b6 100644 --- a/packages/knowledge-base/src/knowledge-base/KnowledgeBase.ts +++ b/packages/knowledge-base/src/knowledge-base/KnowledgeBase.ts @@ -5,6 +5,7 @@ */ import type { ITextIndex, TextFields, VectorSearchOptions } from "@workglow/storage"; +import type { IRunConfig } from "@workglow/task-graph"; import type { TypedArray } from "@workglow/util/schema"; import type { ChunkRecord } from "../chunk/ChunkSchema"; import type { @@ -20,17 +21,26 @@ import type { DocumentTabularStorage, InsertDocumentStorageEntity, } from "../document/DocumentStorageSchema"; +import type { ChunkStrategy, IKbAiStrategy, SearchMode } from "./IKbAiStrategy"; /** - * Options passed through `kb.search()` to the `onSearch` callback. - * The callback decides how to interpret them (similarity vs hybrid, etc.). - * `filter` is intentionally a loose record — the callback and its backing - * vector storage define the allowed keys. + * Options passed through `kb.search()`. `filter` is a loose record; allowed + * keys are defined by the underlying vector storage. */ export interface ISearchOptions { readonly topK?: number; readonly filter?: Readonly>; readonly scoreThreshold?: number; + /** + * For hybrid retrieval and the first stage of rerank retrieval: vector + * vs. text weighting in [0, 1]. Defaults to the storage backend's default. + */ + readonly vectorWeight?: number; + /** + * For `kind: "rerank"`: how many candidates to retrieve before reranking. + * Defaults to `max(topK * 5, 20)`. + */ + readonly firstStageTopK?: number; } /** @@ -52,6 +62,7 @@ export interface HybridSearchOptions< readonly textQuery: string; readonly topK?: number; readonly filter?: Partial; + readonly scoreThreshold?: number; readonly vectorWeight?: number; /** RRF saturation constant; standard value is 60. */ readonly rrfK?: number; @@ -126,34 +137,9 @@ function matchesFilter>( return true; } -/** - * Callback invoked after a document is upserted. - * Receives the KB instance and the upserted document. - */ -export type OnDocumentUpsertCallback = (kb: KnowledgeBase, doc: Document) => Promise; - -/** - * Callback invoked after a document (and its chunks) are deleted. - * Receives the KB instance and the deleted document's ID. - */ -export type OnDocumentDeleteCallback = (kb: KnowledgeBase, doc_id: string) => Promise; - -/** - * Callback invoked by `search()` to handle text-to-vector conversion - * and the actual search. Returns search results. - */ -export type OnSearchCallback = ( - kb: KnowledgeBase, - query: string, - options?: ISearchOptions -) => Promise; - export interface KnowledgeBaseOptions { readonly title?: string; readonly description?: string; - readonly onDocumentUpsert?: OnDocumentUpsertCallback; - readonly onDocumentDelete?: OnDocumentDeleteCallback; - readonly onSearch?: OnSearchCallback; /** * Optional text index. When installed, chunk upserts auto-write to it and * {@link KnowledgeBase.hybridSearch} fuses vector + text rankings via RRF. @@ -161,35 +147,71 @@ export interface KnowledgeBaseOptions { * {@link KnowledgeBase.installTextIndex} after. */ readonly textIndex?: ITextIndex; + /** + * Model ID used to embed document chunks during ingest. Consumed by the + * installed {@link IKbAiStrategy} — the KB itself doesn't run AI. + */ + readonly docEmbeddingModel?: string; + /** + * Model ID used to embed search queries. Defaults to `docEmbeddingModel` + * if absent (the common case — symmetric embedding). + */ + readonly queryEmbeddingModel?: string; + /** + * Optional cross-encoder reranker model ID. Required when `searchMode` + * is `"rerank"`. + */ + readonly rerankerModel?: string; + /** Chunker mode used by ingest. Defaults to `"hierarchical"`. */ + readonly chunkStrategy?: ChunkStrategy; + /** + * Retrieval mode used by search. Defaults to `"rerank"` when a reranker + * model is configured, `"hybrid"` when the storage supports it, + * otherwise `"similarity"`. + */ + readonly searchMode?: SearchMode; + /** + * The AI strategy used by `upsert`, `delete`, and `search`. Installable + * post-construction via {@link KnowledgeBase.setAiStrategy}. + */ + readonly aiStrategy?: IKbAiStrategy; } /** - * Unified KnowledgeBase that owns both document and vector storage, - * providing lifecycle management and cascading deletes. + * Unified KnowledgeBase that owns both document and vector storage. + * + * The public API is intentionally tiny: `upsert`, `delete`, `search`, plus + * lifecycle and inspection helpers. RAG behavior (chunking, embedding, + * retrieval flavor) is fully delegated to an installed + * {@link IKbAiStrategy}. Two flavors ship: + * - `createStandardKbStrategy(...)` from `@workglow/ai` — picks chunker + * mode and search mode from this KB's `chunkStrategy` / `searchMode` + * fields. Uses the registered model IDs. + * - Custom strategies — write your own when you need scoping or unusual + * retrieval; the builder ships one for per-project KBs. + * + * Storage access methods (`upsertDocument`, `upsertChunksBulk`, + * `similaritySearch`, `hybridSearch`, etc.) remain on the class as + * building blocks that strategies and subclasses use. Every one of + * these goes through virtual dispatch, so subclasses (e.g. a + * tenant-scoped KB) can intercept any of them without the strategy + * knowing. + * + * See {@link IKbAiStrategy} for the strategy trust model — installed + * strategies are TRUSTED CODE and must not come from untrusted sources. */ export class KnowledgeBase { readonly name: string; readonly title: string = ""; readonly description: string = ""; + readonly docEmbeddingModel: string | undefined = undefined; + readonly queryEmbeddingModel: string | undefined = undefined; + readonly rerankerModel: string | undefined = undefined; + readonly chunkStrategy: ChunkStrategy | undefined = undefined; + readonly searchMode: SearchMode | undefined = undefined; private readonly tabularStorage: DocumentTabularStorage; private readonly chunkStorage: ChunkVectorStorage; - - /** - * Called after `upsertDocument` successfully writes to storage. - * Awaited — throwing rejects the upsert call, but storage is already committed. - * Use for chunk re-indexing, audit logging, etc. - */ - onDocumentUpsert: OnDocumentUpsertCallback | undefined; - /** - * Called after `deleteDocument` successfully deletes the document and its chunks. - * Awaited — throwing rejects the delete call, but storage is already committed. - */ - onDocumentDelete: OnDocumentDeleteCallback | undefined; - /** - * Called by `search()` to embed the query and execute the search. - * Required if you intend to call `kb.search()`. - */ - onSearch: OnSearchCallback | undefined; + private aiStrategy: IKbAiStrategy | undefined; /** * Optional full-text index. When installed, chunk upserts auto-write to it @@ -212,12 +234,15 @@ export class KnowledgeBase { if (typeof options === "object" && options !== null) { this.title = options.title ?? name; this.description = options.description ?? ""; - this.onDocumentUpsert = options.onDocumentUpsert; - this.onDocumentDelete = options.onDocumentDelete; - this.onSearch = options.onSearch; if (options.textIndex) { this.textIndex = options.textIndex; } + this.docEmbeddingModel = options.docEmbeddingModel; + this.queryEmbeddingModel = options.queryEmbeddingModel ?? options.docEmbeddingModel; + this.rerankerModel = options.rerankerModel; + this.chunkStrategy = options.chunkStrategy; + this.searchMode = options.searchMode; + this.aiStrategy = options.aiStrategy; } } @@ -337,155 +362,149 @@ export class KnowledgeBase { } // =========================================================================== - // Document CRUD + // Strategy installation // =========================================================================== /** - * Upsert a document. - * @returns The document with the generated doc_id if it was auto-generated + * Install (or replace) the AI strategy used by `upsert`/`delete`/`search`. + * + * See {@link IKbAiStrategy} for the strategy trust model — strategies + * receive the KB's full low-level storage surface and are TRUSTED CODE. + * + * Replacing the strategy does NOT affect operations already in flight. + * Each public op (`upsert`/`delete`/`search`/`reindex`) resolves its + * strategy at entry via {@link requireStrategy} and holds that reference + * for its lifetime; a concurrent `setAiStrategy(B)` mid-`upsert(A)` + * leaves the in-progress upsert running on strategy A and routes the + * next public op to strategy B. + * + * Pass `undefined` to detach the strategy — subsequent public-op calls + * throw with a setup hint instead of running. */ - async upsertDocument(document: Document): Promise { - const serialized = JSON.stringify(document.toJSON()); - - const insertEntity: InsertDocumentStorageEntity = { - doc_id: document.doc_id, - data: serialized, - }; - const entity = await this.tabularStorage.put(insertEntity); - - if (document.doc_id !== entity.doc_id) { - document.setDocId(entity.doc_id); - } - - if (this.onDocumentUpsert) { - await this.onDocumentUpsert(this, document); - } + setAiStrategy(strategy: IKbAiStrategy | undefined): void { + this.aiStrategy = strategy; + } - return document; + getAiStrategy(): IKbAiStrategy | undefined { + return this.aiStrategy; } /** - * Get a document by ID + * Snapshot the currently installed strategy or throw if none. + * + * Returns the field value as-is — callers should hold the returned + * reference for the duration of one public op so a concurrent + * `setAiStrategy(...)` doesn't swap the strategy mid-operation. See + * {@link setAiStrategy} for the full concurrency contract. */ - async getDocument(doc_id: string): Promise { - const entity = await this.tabularStorage.get({ doc_id }); - if (!entity) { - return undefined; + private requireStrategy(forOp: string): IKbAiStrategy { + if (!this.aiStrategy) { + throw new Error( + `KnowledgeBase.${forOp}() requires an AI strategy. ` + + `Install one via kb.setAiStrategy(strategy) — see createStandardKbStrategy from @workglow/ai.` + ); } - return Document.fromJSON(entity.data, entity.doc_id); + return this.aiStrategy; } + // =========================================================================== + // Public RAG API — strategy-driven + // =========================================================================== + /** - * Delete a document and all its chunks (cascading delete). + * Ingest a document end-to-end: chunk + embed + write. Delegates to the + * installed strategy. + * + * The strategy is snapshotted at entry: a concurrent `setAiStrategy(...)` + * during the upsert won't redirect the in-flight call to the new + * strategy. See {@link setAiStrategy}. */ - async deleteDocument(doc_id: string): Promise { - await this.deleteChunksForDocument(doc_id); - await this.tabularStorage.delete({ doc_id }); + async upsert(doc: Document, runConfig?: Partial): Promise { + const strategy = this.requireStrategy("upsert"); + return strategy.ingest(this, doc, runConfig); + } - if (this.onDocumentDelete) { - await this.onDocumentDelete(this, doc_id); - } + /** + * Remove a document and its chunks. Delegates to the installed strategy + * (snapshotted at entry — see {@link setAiStrategy}). + */ + async delete(doc_id: string, runConfig?: Partial): Promise { + const strategy = this.requireStrategy("delete"); + return strategy.delete(this, doc_id, runConfig); } /** - * List all document IDs + * Run a text query. Retrieval flavor (text / similarity / hybrid / + * rerank) is decided by the installed strategy — typically derived from + * this KB's `searchMode` field. + * + * The strategy is snapshotted at entry: a concurrent `setAiStrategy(...)` + * during the search won't redirect the in-flight call. See + * {@link setAiStrategy}. */ - async listDocuments(): Promise { - const entities = await this.tabularStorage.getAll(); - if (!entities) { - return []; - } - return entities.map((e: DocumentStorageEntity) => e.doc_id); + async search( + query: string, + options?: ISearchOptions, + runConfig?: Partial + ): Promise { + const strategy = this.requireStrategy("search"); + return strategy.search(this, query, options, runConfig); } // =========================================================================== - // Tree traversal + // Strategy-facing building blocks + // + // These methods are public so strategies (and subclasses like + // `ScopedKnowledgeBase`) can call them, but application code should go + // through `upsert` / `delete` / `search` above. The contract: every one + // of these goes through virtual dispatch, so a subclass can intercept + // any of them without the strategy knowing. // =========================================================================== /** - * Get a specific node by ID from a document + * Store a document JSON record. Does NOT chunk or embed; the strategy + * does that orchestration and then calls back into this method. + * @returns The document with the generated doc_id if it was auto-generated */ - async getNode(doc_id: string, nodeId: string): Promise { - const doc = await this.getDocument(doc_id); - if (!doc) { - return undefined; - } + async upsertDocument(document: Document): Promise { + const serialized = JSON.stringify(document.toJSON()); - const traverse = (node: DocumentNode): DocumentNode | undefined => { - if (node.nodeId === nodeId) { - return node; - } - if ("children" in node && Array.isArray(node.children)) { - for (const child of node.children) { - const found = traverse(child); - if (found) return found; - } - } - return undefined; + const insertEntity: InsertDocumentStorageEntity = { + doc_id: document.doc_id, + data: serialized, }; + const entity = await this.tabularStorage.put(insertEntity); - return traverse(doc.root); + if (document.doc_id !== entity.doc_id) { + document.setDocId(entity.doc_id); + } + + return document; } /** - * Get ancestors of a node (from root to target node) + * Cascading delete: chunks first, then the document row. Strategies call + * this directly when their `delete()` doesn't need extra logic. */ - async getAncestors(doc_id: string, nodeId: string): Promise { - const doc = await this.getDocument(doc_id); - if (!doc) { - return []; - } - - const path: string[] = []; - const findPath = (node: DocumentNode): boolean => { - path.push(node.nodeId); - if (node.nodeId === nodeId) { - return true; - } - if ("children" in node && Array.isArray(node.children)) { - for (const child of node.children) { - if (findPath(child)) { - return true; - } - } - } - path.pop(); - return false; - }; - - if (!findPath(doc.root)) { - return []; - } - - const ancestors: DocumentNode[] = []; - let currentNode: DocumentNode = doc.root; - ancestors.push(currentNode); + async deleteDocument(doc_id: string): Promise { + await this.deleteChunksForDocument(doc_id); + await this.tabularStorage.delete({ doc_id }); + } - for (let i = 1; i < path.length; i++) { - const targetId = path[i]; - if ("children" in currentNode && Array.isArray(currentNode.children)) { - const found = currentNode.children.find((child: DocumentNode) => child.nodeId === targetId); - if (found) { - currentNode = found; - ancestors.push(currentNode); - } else { - break; - } - } else { - break; - } - } + async getDocument(doc_id: string): Promise { + const entity = await this.tabularStorage.get({ doc_id }); + if (!entity) return undefined; + return Document.fromJSON(entity.data, entity.doc_id); + } - return ancestors; + async listDocuments(): Promise { + const entities = await this.tabularStorage.getAll(); + if (!entities) return []; + return entities.map((e: DocumentStorageEntity) => e.doc_id); } - // =========================================================================== - // Chunk CRUD - // =========================================================================== + // ----- chunks ----- - /** - * Upsert a single chunk vector entity - */ async upsertChunk(chunk: InsertChunkVectorEntity): Promise { const expected = this.getVectorDimensions(); if (expected > 0 && chunk.vector.length !== expected) { @@ -498,9 +517,6 @@ export class KnowledgeBase { return stored; } - /** - * Upsert multiple chunk vector entities - */ async upsertChunksBulk(chunks: InsertChunkVectorEntity[]): Promise { const expected = this.getVectorDimensions(); if (expected > 0) { @@ -517,9 +533,6 @@ export class KnowledgeBase { return stored; } - /** - * Delete all chunks for a specific document - */ async deleteChunksForDocument(doc_id: string): Promise { await this.chunkStorage.deleteSearch({ doc_id }); if (this.textIndex) { @@ -527,24 +540,13 @@ export class KnowledgeBase { } } - /** - * Get all chunks for a specific document - */ async getChunksForDocument(doc_id: string): Promise { const results = await this.chunkStorage.query({ doc_id }); return (results ?? []) as ChunkVectorEntity[]; } - // =========================================================================== - // Search - // =========================================================================== + // ----- vector retrieval ----- - /** - * Search for similar chunks using vector similarity. This is the canonical - * scope-aware entry point — subclasses (e.g. a scoped KB that isolates by - * tenant) override this to inject filter predicates before delegating to - * the underlying storage. - */ async similaritySearch( query: TypedArray, options?: VectorSearchOptions @@ -553,24 +555,6 @@ export class KnowledgeBase { return raw.map((r) => ({ ...r, scoreType: "cosine" }) as ChunkSearchResult); } - /** - * Hybrid search combining vector similarity and full-text BM25(F) ranking. - * The two rankers run in parallel, then their per-rank contributions are - * fused with Reciprocal Rank Fusion: - * - * ``` - * fused(d) = vectorWeight / (rrfK + rank_v(d)) - * + (1 - vectorWeight) / (rrfK + rank_t(d)) - * ``` - * - * RRF rewards items that appear in both rankings. Score values are *not* - * comparable to cosine scores — use `topK` to size the result, not a score - * threshold. - * - * Canonical scope-aware entry point; subclasses override for filter injection. - * - * @throws Error if no text index is installed (call {@link installTextIndex} first). - */ async hybridSearch( query: TypedArray, options: HybridSearchOptions @@ -592,9 +576,6 @@ export class KnowledgeBase { candidatePoolMultiplier = 5, } = options; - // Empty / whitespace-only textQuery has no signal for the BM25 ranker. - // Returning RRF-shaped scores in that case would surprise callers, so - // delegate to the cosine similarity path and return cosine scores. if (!textQuery || textQuery.trim().length === 0) { return this.similaritySearch(query, { topK, filter }); } @@ -621,8 +602,6 @@ export class KnowledgeBase { vectorResults.forEach((entity, rank) => { const contribution = vectorWeightClamped / (safeRrfK + rank + 1); - // Strip the cosine scoreType from the wrapped similarity result; the - // outer fused entity will carry "rrf" once we re-emit it. const { scoreType: _drop, ...rest } = entity; fused.set(entity.chunk_id, { score: contribution, @@ -646,9 +625,7 @@ export class KnowledgeBase { .map(([chunkId]) => chunkId); if (missing.length > 0) { - const hydrated = await this.chunkStorage.getBulk( - missing.map((chunk_id) => ({ chunk_id })) - ); + const hydrated = await this.chunkStorage.getBulk(missing.map((chunk_id) => ({ chunk_id }))); const byId = new Map(); for (const entity of hydrated as ChunkVectorEntity[]) { byId.set(entity.chunk_id, entity); @@ -677,13 +654,6 @@ export class KnowledgeBase { return ranked.slice(0, topK); } - /** - * Pure full-text search via the installed text index. Hydrates ranked - * chunkIds from chunk storage and applies optional metadata filtering - * post-hoc. - * - * @throws Error if no text index is installed. - */ async textSearch( query: string, options: TextOnlySearchOptions = {} @@ -702,9 +672,7 @@ export class KnowledgeBase { const hits = await index.search(query, { topK: poolSize }); if (hits.length === 0) return []; - const hydrated = await this.chunkStorage.getBulk( - hits.map((h) => ({ chunk_id: h.chunkId })) - ); + const hydrated = await this.chunkStorage.getBulk(hits.map((h) => ({ chunk_id: h.chunkId }))); const byId = new Map(); for (const entity of hydrated as ChunkVectorEntity[]) { byId.set(entity.chunk_id, entity); @@ -721,81 +689,121 @@ export class KnowledgeBase { return results; } - /** - * Whether {@link hybridSearch} is available — i.e. a text index has been - * installed. - */ supportsHybridSearch(): boolean { return this.textIndex !== undefined; } - /** - * High-level text search. Delegates to the `onSearch` callback, which is - * responsible for embedding the query and executing the appropriate search - * (similarity, hybrid, keyword, etc.). Install `onSearch` via - * `createKnowledgeBase({ onSearch })` or the KnowledgeBase constructor options. - * - * If `onSearch` calls back into `kb.similaritySearch()` / `kb.hybridSearch()`, - * those calls still go through virtual dispatch — so subclass filter injection - * (e.g. tenant scope) applies even when the entry point is `kb.search()`. - * - * @throws Error if `onSearch` is not configured. - */ - async search(query: string, options?: ISearchOptions): Promise { - if (!this.onSearch) { - throw new Error( - "KnowledgeBase.search() requires an `onSearch` callback. " + - "Pass one via createKnowledgeBase({ onSearch }) or the KnowledgeBase " + - "constructor options. For raw vector search, use " + - "`kb.similaritySearch()` or `kb.vectorStorage.similaritySearch()` directly." - ); + // =========================================================================== + // Tree traversal helpers (unchanged) + // =========================================================================== + + async getNode(doc_id: string, nodeId: string): Promise { + const doc = await this.getDocument(doc_id); + if (!doc) return undefined; + + const traverse = (node: DocumentNode): DocumentNode | undefined => { + if (node.nodeId === nodeId) return node; + if ("children" in node && Array.isArray(node.children)) { + for (const child of node.children) { + const found = traverse(child); + if (found) return found; + } + } + return undefined; + }; + + return traverse(doc.root); + } + + async getAncestors(doc_id: string, nodeId: string): Promise { + const doc = await this.getDocument(doc_id); + if (!doc) return []; + + const path: string[] = []; + const findPath = (node: DocumentNode): boolean => { + path.push(node.nodeId); + if (node.nodeId === nodeId) return true; + if ("children" in node && Array.isArray(node.children)) { + for (const child of node.children) { + if (findPath(child)) return true; + } + } + path.pop(); + return false; + }; + + if (!findPath(doc.root)) return []; + + const ancestors: DocumentNode[] = []; + let currentNode: DocumentNode = doc.root; + ancestors.push(currentNode); + + for (let i = 1; i < path.length; i++) { + const targetId = path[i]; + if ("children" in currentNode && Array.isArray(currentNode.children)) { + const found = currentNode.children.find((child: DocumentNode) => child.nodeId === targetId); + if (found) { + currentNode = found; + ancestors.push(currentNode); + } else { + break; + } + } else { + break; + } } - return this.onSearch(this, query, options); + + return ancestors; } // =========================================================================== - // Accessors for raw storage + // Lifecycle / accessors // =========================================================================== - /** - * The underlying chunk/vector storage. Use when you need raw, unscoped - * access to low-level vector operations — e.g. bulk maintenance, metrics, - * or behavior that explicitly should bypass any subclass scoping. For - * normal search, prefer `kb.similaritySearch()` / `kb.hybridSearch()`, - * which subclasses can override to inject scope. - */ + /** Underlying chunk store; for maintenance and inspection. */ get vectorStorage(): ChunkVectorStorage { return this.chunkStorage; } - // =========================================================================== - // Lifecycle - // =========================================================================== - /** - * Prepare a document for re-indexing: deletes all chunks but keeps the document. - * @returns The document if found, undefined otherwise + * Prepare a document for re-indexing: deletes all chunks but keeps the + * document. Used by re-index flows; routine callers should use + * `kb.upsert(doc)` to fully replace. */ async prepareReindex(doc_id: string): Promise { const doc = await this.getDocument(doc_id); - if (!doc) { - return undefined; - } + if (!doc) return undefined; await this.deleteChunksForDocument(doc_id); return doc; } /** - * Setup the underlying databases - */ + * Re-index every document by re-running ingest. Requires a strategy. + * + * The strategy is captured once at entry — every doc in the run uses + * the same strategy, even if `setAiStrategy(...)` is called concurrently + * partway through the loop. The next `reindex()` call would pick up + * the new strategy. See {@link setAiStrategy}. + */ + async reindex(runConfig?: Partial): Promise { + const strategy = this.requireStrategy("reindex"); + const docIds = await this.listDocuments(); + let count = 0; + for (const doc_id of docIds) { + if (runConfig?.signal?.aborted) throw runConfig.signal.reason; + const doc = await this.getDocument(doc_id); + if (!doc) continue; + await strategy.ingest(this, doc, runConfig); + count++; + } + return count; + } + async setupDatabase(): Promise { await this.tabularStorage.setupDatabase(); await this.chunkStorage.setupDatabase(); } - /** - * Destroy storage instances - */ destroy(): void { this.tabularStorage.destroy(); this.chunkStorage.destroy(); @@ -809,13 +817,6 @@ export class KnowledgeBase { this.destroy(); } - // =========================================================================== - // Accessors - // =========================================================================== - - /** - * Get a chunk by ID - */ async getChunk(chunk_id: string): Promise { return this.chunkStorage.get({ chunk_id }); } @@ -836,23 +837,14 @@ export class KnowledgeBase { return this.upsertChunksBulk(chunks); } - /** - * Get all chunks - */ async getAllChunks(): Promise { return this.chunkStorage.getAll() as Promise; } - /** - * Get chunk count - */ async chunkCount(): Promise { return this.chunkStorage.size(); } - /** - * Clear all chunks - */ async clearChunks(): Promise { await this.chunkStorage.deleteAll(); if (this.textIndex) { @@ -860,36 +852,19 @@ export class KnowledgeBase { } } - /** - * Get vector dimensions - */ getVectorDimensions(): number { return this.chunkStorage.getVectorDimensions(); } - // =========================================================================== - // Document chunk helpers - // =========================================================================== - - /** - * Get chunks from the document JSON (not from vector storage) - */ async getDocumentChunks(doc_id: string): Promise { const doc = await this.getDocument(doc_id); - if (!doc) { - return []; - } + if (!doc) return []; return doc.getChunks(); } - /** - * Find chunks in document JSON that contain a specific nodeId in their path - */ async findChunksByNodeId(doc_id: string, nodeId: string): Promise { const doc = await this.getDocument(doc_id); - if (!doc) { - return []; - } + if (!doc) return []; return doc.findChunksByNodeId(nodeId); } } diff --git a/packages/knowledge-base/src/knowledge-base/createKnowledgeBase.ts b/packages/knowledge-base/src/knowledge-base/createKnowledgeBase.ts index 902ed30e6..10a8a1182 100644 --- a/packages/knowledge-base/src/knowledge-base/createKnowledgeBase.ts +++ b/packages/knowledge-base/src/knowledge-base/createKnowledgeBase.ts @@ -11,11 +11,7 @@ import type { ChunkVectorStorage } from "../chunk/ChunkVectorStorageSchema"; import { ChunkVectorPrimaryKey, ChunkVectorStorageSchema } from "../chunk/ChunkVectorStorageSchema"; import type { DocumentTabularStorage } from "../document/DocumentStorageSchema"; import { DocumentStorageKey, DocumentStorageSchema } from "../document/DocumentStorageSchema"; -import type { - OnDocumentDeleteCallback, - OnDocumentUpsertCallback, - OnSearchCallback, -} from "./KnowledgeBase"; +import type { ChunkStrategy, IKbAiStrategy, SearchMode } from "./IKbAiStrategy"; import { KnowledgeBase } from "./KnowledgeBase"; import { registerKnowledgeBase } from "./KnowledgeBaseRegistry"; @@ -26,15 +22,18 @@ export interface CreateKnowledgeBaseOptions { readonly register?: boolean; readonly title?: string; readonly description?: string; - readonly onDocumentUpsert?: OnDocumentUpsertCallback; - readonly onDocumentDelete?: OnDocumentDeleteCallback; - readonly onSearch?: OnSearchCallback; /** * Optional full-text index. When provided, the KB enables * {@link KnowledgeBase.hybridSearch} and auto-writes chunks to the index * on upsert. */ readonly textIndex?: ITextIndex; + readonly docEmbeddingModel?: string; + readonly queryEmbeddingModel?: string; + readonly rerankerModel?: string; + readonly chunkStrategy?: ChunkStrategy; + readonly searchMode?: SearchMode; + readonly aiStrategy?: IKbAiStrategy; } /** @@ -45,7 +44,9 @@ export interface CreateKnowledgeBaseOptions { * const kb = await createKnowledgeBase({ * name: "my-kb", * vectorDimensions: 1024, + * docEmbeddingModel: "onnx:Xenova/bge-base-en-v1.5:q8", * }); + * kb.setAiStrategy(createStandardKbStrategy()); * ``` */ export async function createKnowledgeBase( @@ -58,10 +59,13 @@ export async function createKnowledgeBase( register: shouldRegister = true, title, description, - onDocumentUpsert, - onDocumentDelete, - onSearch, textIndex, + docEmbeddingModel, + queryEmbeddingModel, + rerankerModel, + chunkStrategy, + searchMode, + aiStrategy, } = options; const vectorCtor = vectorCtorOption ?? Float32Array; @@ -93,7 +97,17 @@ export async function createKnowledgeBase( name, tabularStorage as unknown as DocumentTabularStorage, vectorStorage as unknown as ChunkVectorStorage, - { title, description, onDocumentUpsert, onDocumentDelete, onSearch, textIndex } + { + title, + description, + textIndex, + docEmbeddingModel, + queryEmbeddingModel, + rerankerModel, + chunkStrategy, + searchMode, + aiStrategy, + } ); if (shouldRegister) { diff --git a/packages/knowledge-base/tsconfig.json b/packages/knowledge-base/tsconfig.json index 77696be01..1fef966a8 100644 --- a/packages/knowledge-base/tsconfig.json +++ b/packages/knowledge-base/tsconfig.json @@ -8,8 +8,5 @@ "outDir": "./dist", "rootDir": "./src" }, - "references": [ - { "path": "../util" }, - { "path": "../storage" } - ] + "references": [{ "path": "../util" }, { "path": "../storage" }, { "path": "../task-graph" }] } diff --git a/packages/mcp/CHANGELOG.md b/packages/mcp/CHANGELOG.md index 92be90787..1a024e91e 100644 --- a/packages/mcp/CHANGELOG.md +++ b/packages/mcp/CHANGELOG.md @@ -1,5 +1,11 @@ # Changelog +## 0.2.36 + +### Bug Fixes + +- mcp + ## 0.2.35 ### Chores diff --git a/packages/mcp/package.json b/packages/mcp/package.json index 4b669bf24..d176ae578 100644 --- a/packages/mcp/package.json +++ b/packages/mcp/package.json @@ -1,7 +1,7 @@ { "name": "@workglow/mcp", "type": "module", - "version": "0.2.35", + "version": "0.2.36", "repository": { "type": "git", "url": "https://github.com/workglow-dev/libs.git", @@ -76,4 +76,4 @@ "publishConfig": { "access": "public" } -} \ No newline at end of file +} diff --git a/packages/storage/CHANGELOG.md b/packages/storage/CHANGELOG.md index 1f5b38a75..2afe3a71c 100644 --- a/packages/storage/CHANGELOG.md +++ b/packages/storage/CHANGELOG.md @@ -1,5 +1,7 @@ # @workglow/storage +## 0.2.36 + ## 0.2.35 ### Features diff --git a/packages/storage/package.json b/packages/storage/package.json index bfcdc15fa..4c7d9d962 100644 --- a/packages/storage/package.json +++ b/packages/storage/package.json @@ -1,7 +1,7 @@ { "name": "@workglow/storage", "type": "module", - "version": "0.2.35", + "version": "0.2.36", "repository": { "type": "git", "url": "https://github.com/workglow-dev/libs.git", diff --git a/packages/task-graph/CHANGELOG.md b/packages/task-graph/CHANGELOG.md index bb4c5a5f9..7878bfbaf 100644 --- a/packages/task-graph/CHANGELOG.md +++ b/packages/task-graph/CHANGELOG.md @@ -1,5 +1,7 @@ # @workglow/task-graph +## 0.2.36 + ## 0.2.35 ### Features diff --git a/packages/task-graph/package.json b/packages/task-graph/package.json index 55d4dbab7..b10a4475f 100644 --- a/packages/task-graph/package.json +++ b/packages/task-graph/package.json @@ -1,7 +1,7 @@ { "name": "@workglow/task-graph", "type": "module", - "version": "0.2.35", + "version": "0.2.36", "repository": { "type": "git", "url": "https://github.com/workglow-dev/libs.git", diff --git a/packages/tasks/CHANGELOG.md b/packages/tasks/CHANGELOG.md index 022bf388f..c3c156afe 100644 --- a/packages/tasks/CHANGELOG.md +++ b/packages/tasks/CHANGELOG.md @@ -1,5 +1,15 @@ # @workglow/tasks +## 0.2.36 + +### Chores + +- update deps + +### Updated Dependencies + +- `undici`: ^8.3.0 + ## 0.2.35 ### Performance diff --git a/packages/tasks/package.json b/packages/tasks/package.json index 9ff3a26b1..33e79bc8b 100644 --- a/packages/tasks/package.json +++ b/packages/tasks/package.json @@ -1,7 +1,7 @@ { "name": "@workglow/tasks", "type": "module", - "version": "0.2.35", + "version": "0.2.36", "repository": { "type": "git", "url": "https://github.com/workglow-dev/libs.git", diff --git a/packages/test/CHANGELOG.md b/packages/test/CHANGELOG.md index d49636721..a2f3353c2 100644 --- a/packages/test/CHANGELOG.md +++ b/packages/test/CHANGELOG.md @@ -1,5 +1,38 @@ # @workglow/test +## 0.2.36 + +### Features + +#### kb + +- pluggable strategy with model config, IRunConfig threading, and document tasks + +#### hft + +- add text reranker capability with pipeline shape validation + +### Bug Fixes + +#### util/worker, ai/task + +- TTL-based pendingAborts eviction; clarify runWithIterable bond (#500) + +#### ai + +- avoid NaN reranker scores for empty queries +- escape regex metacharacters in RerankerTask.simpleRerank +- classify provider-error vs no-finish in AiTask.execute +- replace runWithIterable Proxy with shallow clone + +#### util + +- snapshot-then-delete eviction in WorkerServerBase Set caps + +### Tests + +- move tests so we we don't need a list of files per provider in the test script + ## 0.2.35 ### Features diff --git a/packages/test/package.json b/packages/test/package.json index 9968eac1a..69a94808e 100644 --- a/packages/test/package.json +++ b/packages/test/package.json @@ -2,7 +2,7 @@ "name": "@workglow/test", "type": "module", "private": true, - "version": "0.2.35", + "version": "0.2.36", "description": "Testing utilities and test suites for Workglow packages, providing comprehensive testing infrastructure.", "scripts": { "watch": "concurrently -c 'auto' 'bun:watch-*'", diff --git a/packages/test/src/samples/ONNXModelSamples.ts b/packages/test/src/samples/ONNXModelSamples.ts index f4364a370..56a873f98 100644 --- a/packages/test/src/samples/ONNXModelSamples.ts +++ b/packages/test/src/samples/ONNXModelSamples.ts @@ -1,6 +1,6 @@ import { getGlobalModelRepository } from "@workglow/ai"; -import { HF_TRANSFORMERS_ONNX } from "@workglow/huggingface-transformers/ai"; import type { HfTransformersOnnxModelRecord } from "@workglow/huggingface-transformers/ai"; +import { HF_TRANSFORMERS_ONNX } from "@workglow/huggingface-transformers/ai"; export async function registerHuggingfaceLocalModels(): Promise { const onnxModels: HfTransformersOnnxModelRecord[] = [ diff --git a/packages/test/src/test/rag/CreateStandardKbStrategyFirstStage.test.ts b/packages/test/src/test/rag/CreateStandardKbStrategyFirstStage.test.ts new file mode 100644 index 000000000..9ff0a8a72 --- /dev/null +++ b/packages/test/src/test/rag/CreateStandardKbStrategyFirstStage.test.ts @@ -0,0 +1,126 @@ +/** + * @license + * Copyright 2026 Steven Roussey + * SPDX-License-Identifier: Apache-2.0 + */ + +import type { AiProviderRunFn, ModelRecord, TextEmbeddingTaskInput } from "@workglow/ai"; +import { + createStandardKbStrategy, + getAiProviderRegistry, + getGlobalModelRepository, +} from "@workglow/ai"; +import { createKnowledgeBase } from "@workglow/knowledge-base"; +import { uuid4 } from "@workglow/util"; +import { beforeAll, describe, expect, it, vi } from "vitest"; + +/** + * Tests for the first-stage candidate-pool sizing in + * `createStandardKbStrategy` rerank mode. The pool size is + * `max(topK * firstStageMultiplier, firstStageMinimum)`; the minimum + * exists so that a very small `topK` (e.g. 1) does not collapse the + * reranker's input down to a useless handful of candidates. + */ +const TEST_PROVIDER = "test-firststage-provider"; +const TEST_EMBED_MODEL_ID = "test:firststage:embed"; + +describe("createStandardKbStrategy first-stage sizing (rerank mode)", () => { + beforeAll(async () => { + const registry = getAiProviderRegistry(); + const runFn: AiProviderRunFn = async (input, _model, _signal, emit) => { + const embeddingInput = input as TextEmbeddingTaskInput; + const texts = Array.isArray(embeddingInput.text) + ? embeddingInput.text + : [embeddingInput.text ?? ""]; + const vectors = texts.map(() => new Float32Array([1, 0, 0])); + emit({ + type: "finish", + data: { + vector: vectors.length === 1 ? vectors[0] : vectors, + }, + }); + }; + registry.registerRunFn(TEST_PROVIDER, { + serves: ["text.embedding"], + runFn, + }); + + const modelRepo = getGlobalModelRepository(); + const existing = await modelRepo.findByName(TEST_EMBED_MODEL_ID).catch(() => undefined); + if (!existing) { + await modelRepo.addModel({ + model_id: TEST_EMBED_MODEL_ID, + capabilities: ["text.embedding"], + title: "First-stage sizing test embed model", + description: "Stub embed model for first-stage sizing tests", + provider: TEST_PROVIDER, + provider_config: { native_dimensions: 3 }, + metadata: {}, + } as ModelRecord); + } + }); + + /** + * Set up a rerank-mode KB and spy on the first-stage retrieval call so + * we can assert exactly what `topK` the strategy hands down to it. + * Heuristic-reranker path (no `rerankerModel`) is used because we only + * care about the first-stage `topK`, not the reranker output. + */ + async function captureFirstStageTopK( + strategyOptions: Parameters[0], + searchTopK: number + ): Promise { + const kb = await createKnowledgeBase({ + name: `kb-first-stage-${uuid4()}`, + vectorDimensions: 3, + register: false, + docEmbeddingModel: TEST_EMBED_MODEL_ID, + searchMode: "rerank", + }); + kb.setAiStrategy(createStandardKbStrategy(strategyOptions)); + + // Spy on whichever first-stage method the strategy will pick. Both + // are stubbed to return [] so the rerank path short-circuits and we + // don't have to seed real chunks. + const hybridSpy = vi.spyOn(kb, "hybridSearch").mockResolvedValue([]); + const similaritySpy = vi.spyOn(kb, "similaritySearch").mockResolvedValue([]); + + await kb.search("hi", { topK: searchTopK }); + + // Exactly one of the spies should fire (mutually exclusive branches + // in the strategy: hybrid if supportsHybridSearch, else similarity). + const calls = [...hybridSpy.mock.calls, ...similaritySpy.mock.calls]; + expect(calls.length).toBe(1); + const opts = calls[0][1] as { topK: number }; + return opts.topK; + } + + it("topK=1, multiplier defaults to 5, minimum defaults to 20 → first-stage topK=20", async () => { + const firstStage = await captureFirstStageTopK(undefined, 1); + expect(firstStage).toBe(20); + }); + + it("topK=10, multiplier defaults to 5 → first-stage topK=50 (above minimum)", async () => { + const firstStage = await captureFirstStageTopK(undefined, 10); + expect(firstStage).toBe(50); + }); + + it("topK=2, multiplier=1, firstStageMinimum=20 → first-stage topK=20 (minimum wins)", async () => { + const firstStage = await captureFirstStageTopK( + { firstStageMultiplier: 1, firstStageMinimum: 20 }, + 2 + ); + expect(firstStage).toBe(20); + }); + + it("topK=10, multiplier=10, firstStageMinimum=20 → first-stage topK=100 (product wins over minimum)", async () => { + // Locks in precedence: when `topK * firstStageMultiplier` exceeds + // `firstStageMinimum`, the product wins. A regression that flipped the + // Math.max arguments (or replaced `max` with `min`) would surface here. + const firstStage = await captureFirstStageTopK( + { firstStageMultiplier: 10, firstStageMinimum: 20 }, + 10 + ); + expect(firstStage).toBe(100); + }); +}); diff --git a/packages/test/src/test/rag/DocumentRepository.test.ts b/packages/test/src/test/rag/DocumentRepository.test.ts index f2b92c04a..8089f08ac 100644 --- a/packages/test/src/test/rag/DocumentRepository.test.ts +++ b/packages/test/src/test/rag/DocumentRepository.test.ts @@ -4,6 +4,11 @@ * SPDX-License-Identifier: Apache-2.0 */ +import type { + ChunkVectorStorage, + DocumentTabularStorage, + SectionNode, +} from "@workglow/knowledge-base"; import { ChunkVectorPrimaryKey, ChunkVectorStorageSchema, @@ -15,17 +20,12 @@ import { StructuralParser, createKnowledgeBase, } from "@workglow/knowledge-base"; -import type { - ChunkVectorStorage, - DocumentTabularStorage, - SectionNode, -} from "@workglow/knowledge-base"; import { InMemoryTabularStorage, InMemoryVectorStorage } from "@workglow/storage"; import { setLogger, uuid4 } from "@workglow/util"; import { afterEach, beforeEach, describe, expect, it } from "vitest"; import { getTestingLogger } from "../../binding/TestingLogger"; -import { snap, report } from "../../binding/testTiming"; +import { report, snap } from "../../binding/testTiming"; let _snap = snap(); beforeEach(() => { @@ -442,99 +442,95 @@ Paragraph.`; }); }); - describe("callbacks", () => { - it("should invoke onDocumentUpsert when a document is upserted", async () => { - const calls: Array<{ kbName: string; docId: string | undefined }> = []; - const kbWithCb = await createKnowledgeBase({ - name: `test-kb-cb-${uuid4()}`, + describe("ai strategy", () => { + it("should throw a helpful error when kb.search() is called without a strategy", async () => { + const bareKb = await createKnowledgeBase({ + name: `test-kb-nostrategy-${uuid4()}`, vectorDimensions: 3, register: false, - onDocumentUpsert: async (instance, doc) => { - calls.push({ kbName: instance.name, docId: doc.doc_id }); - }, }); - const doc_id = uuid4(); - const root = await StructuralParser.parseMarkdown(doc_id, "# Test\n\nContent.", "Test"); - const doc = new Document(root, { title: "Test" }); - - await kbWithCb.upsertDocument(doc); - - expect(calls).toHaveLength(1); - expect(calls[0].kbName).toBe(kbWithCb.name); - expect(calls[0].docId).toBeDefined(); + await expect(bareKb.search("hello")).rejects.toThrow(/AI strategy/); }); - it("should invoke onDocumentDelete when a document is deleted", async () => { - const deletedIds: string[] = []; - const kbWithCb = await createKnowledgeBase({ - name: `test-kb-del-${uuid4()}`, + it("should delegate kb.search to the installed strategy", async () => { + const calls: Array<{ query: string; topK: number | undefined }> = []; + const kb = await createKnowledgeBase({ + name: `test-kb-search-${uuid4()}`, vectorDimensions: 3, register: false, - onDocumentDelete: async (_instance, doc_id) => { - deletedIds.push(doc_id); + }); + kb.setAiStrategy({ + ingest: async (_kb, doc) => doc, + delete: async () => {}, + search: async (_kb, query, options) => { + calls.push({ query, topK: options?.topK }); + return []; }, }); - const doc_id = uuid4(); - const root = await StructuralParser.parseMarkdown(doc_id, "# T\n\nx.", "T"); - const doc = new Document(root, { title: "T" }); - const inserted = await kbWithCb.upsertDocument(doc); - - await kbWithCb.deleteDocument(inserted.doc_id!); + await kb.search("hello", { topK: 4 }); - expect(deletedIds).toEqual([inserted.doc_id]); + expect(calls).toEqual([{ query: "hello", topK: 4 }]); }); - it("should reject upsertDocument when onDocumentUpsert throws, with storage already committed", async () => { - const kbWithCb = await createKnowledgeBase({ - name: `test-kb-throw-${uuid4()}`, + it("should delegate kb.upsert / kb.delete to the strategy", async () => { + const ingested: string[] = []; + const deleted: string[] = []; + const kb = await createKnowledgeBase({ + name: `test-kb-ingest-${uuid4()}`, vectorDimensions: 3, register: false, - onDocumentUpsert: async () => { - throw new Error("callback boom"); + }); + kb.setAiStrategy({ + ingest: async (target, doc) => { + await target.upsertDocument(doc); + ingested.push(doc.doc_id ?? ""); + return doc; + }, + delete: async (target, doc_id) => { + deleted.push(doc_id); + await target.deleteDocument(doc_id); }, + search: async () => [], }); - const doc_id = uuid4(); - const root = await StructuralParser.parseMarkdown(doc_id, "# T\n\nx.", "T"); + const root = await StructuralParser.parseMarkdown(uuid4(), "# T\n\nx.", "T"); const doc = new Document(root, { title: "T" }); + doc.setDocId("d1"); + await kb.upsert(doc); + expect(ingested).toEqual(["d1"]); - await expect(kbWithCb.upsertDocument(doc)).rejects.toThrow("callback boom"); - - // Contract: storage is committed before the callback runs, so the document - // must still be retrievable even though upsertDocument rejected. - const retrieved = await kbWithCb.getDocument(doc.doc_id!); - expect(retrieved).toBeDefined(); - expect(retrieved?.doc_id).toBe(doc.doc_id); + await kb.delete("d1"); + expect(deleted).toEqual(["d1"]); + expect(await kb.getDocument("d1")).toBeUndefined(); }); - it("should throw a helpful error when kb.search() is called without onSearch", async () => { - const bareKb = await createKnowledgeBase({ - name: `test-kb-nosearch-${uuid4()}`, + it("should expose model + chunk/search-mode config to the strategy", async () => { + let observed: { docModel?: string; mode?: string; chunk?: string } = {}; + const kb = await createKnowledgeBase({ + name: `test-kb-config-${uuid4()}`, vectorDimensions: 3, register: false, + docEmbeddingModel: "test:doc", + rerankerModel: "test:rerank", + chunkStrategy: "flat", + searchMode: "rerank", }); - - await expect(bareKb.search("hello")).rejects.toThrow(/onSearch/); - }); - - it("should invoke onSearch with the query and options when kb.search() is called", async () => { - const received: Array<{ query: string; topK: number | undefined }> = []; - const kbWithSearch = await createKnowledgeBase({ - name: `test-kb-search-${uuid4()}`, - vectorDimensions: 3, - register: false, - onSearch: async (_kb, query, options) => { - received.push({ query, topK: options?.topK }); + kb.setAiStrategy({ + ingest: async (_k, d) => d, + delete: async () => {}, + search: async (target) => { + observed = { + docModel: target.docEmbeddingModel, + mode: target.searchMode, + chunk: target.chunkStrategy, + }; return []; }, }); - - const results = await kbWithSearch.search("query text", { topK: 4 }); - - expect(received).toEqual([{ query: "query text", topK: 4 }]); - expect(results).toEqual([]); + await kb.search("q"); + expect(observed).toEqual({ docModel: "test:doc", mode: "rerank", chunk: "flat" }); }); }); }); @@ -762,4 +758,165 @@ Paragraph.`; ); }); }); + + describe("strategy contract", () => { + it("captures the strategy at op entry — mid-search setAiStrategy(B) doesn't redirect an in-flight search", async () => { + const kb = await createKnowledgeBase({ + name: `kb-strategy-snapshot-search-${uuid4()}`, + vectorDimensions: 3, + register: false, + }); + + let releaseA: () => void = () => {}; + const aPending = new Promise((resolve) => { + releaseA = resolve; + }); + const aCalls: string[] = []; + const bCalls: string[] = []; + + const strategyA = { + ingest: async (_kb: KnowledgeBase, d: Document) => d, + delete: async () => {}, + search: async () => { + aCalls.push("search"); + await aPending; + return []; + }, + }; + const strategyB = { + ingest: async (_kb: KnowledgeBase, d: Document) => d, + delete: async () => {}, + search: async () => { + bCalls.push("search"); + return []; + }, + }; + + kb.setAiStrategy(strategyA); + const inFlight = kb.search("q1"); + // Swap mid-flight; the in-flight call must still resolve via A. + kb.setAiStrategy(strategyB); + releaseA(); + await inFlight; + expect(aCalls).toEqual(["search"]); + expect(bCalls).toEqual([]); + + // Subsequent call routes to B as expected. + await kb.search("q2"); + expect(aCalls).toEqual(["search"]); + expect(bCalls).toEqual(["search"]); + }); + + it("captures the strategy at op entry — mid-upsert setAiStrategy(B) doesn't redirect an in-flight upsert", async () => { + const kb = await createKnowledgeBase({ + name: `kb-strategy-snapshot-upsert-${uuid4()}`, + vectorDimensions: 3, + register: false, + }); + + let releaseA: () => void = () => {}; + const aPending = new Promise((resolve) => { + releaseA = resolve; + }); + const aCalls: string[] = []; + const bCalls: string[] = []; + + const strategyA = { + ingest: async (_target: KnowledgeBase, d: Document) => { + aCalls.push("ingest"); + await aPending; + return d; + }, + delete: async () => {}, + search: async () => [], + }; + const strategyB = { + ingest: async (_target: KnowledgeBase, d: Document) => { + bCalls.push("ingest"); + return d; + }, + delete: async () => {}, + search: async () => [], + }; + + kb.setAiStrategy(strategyA); + const root = await StructuralParser.parseMarkdown(uuid4(), "# T\n\nx.", "T"); + const doc = new Document(root, { title: "T" }); + doc.setDocId("doc-snapshot-upsert"); + const inFlight = kb.upsert(doc); + kb.setAiStrategy(strategyB); + releaseA(); + await inFlight; + expect(aCalls).toEqual(["ingest"]); + expect(bCalls).toEqual([]); + }); + + it("chunkText helper throws with chunk_id when metadata.text is missing", async () => { + const { chunkText } = await import("@workglow/knowledge-base"); + expect(() => + chunkText({ + chunk_id: "c-no-text", + metadata: { custom: "x" } as unknown as Parameters[0]["metadata"], + }) + ).toThrow(/c-no-text/); + }); + + it("chunkText helper returns metadata.text when present", async () => { + const { chunkText } = await import("@workglow/knowledge-base"); + const text = chunkText({ + chunk_id: "c-has-text", + metadata: { text: "hello" } as unknown as Parameters[0]["metadata"], + }); + expect(text).toBe("hello"); + }); + + it("captures the strategy once at reindex() entry — mid-loop swap doesn't redirect remaining iterations", async () => { + const kb = await createKnowledgeBase({ + name: `kb-strategy-snapshot-reindex-${uuid4()}`, + vectorDimensions: 3, + register: false, + }); + + const aIngested: string[] = []; + const bIngested: string[] = []; + const strategyA = { + ingest: async (target: KnowledgeBase, d: Document) => { + aIngested.push(d.doc_id ?? ""); + // Swap to B partway through; the reindex loop should keep + // ingesting via A for the rest of this run. + target.setAiStrategy(strategyB); + return d; + }, + delete: async () => {}, + search: async () => [], + }; + const strategyB = { + ingest: async (_target: KnowledgeBase, d: Document) => { + bIngested.push(d.doc_id ?? ""); + return d; + }, + delete: async () => {}, + search: async () => [], + }; + + // Seed three documents through the storage layer directly so they're + // present for reindex to iterate. + for (let i = 0; i < 3; i++) { + const root = await StructuralParser.parseMarkdown(uuid4(), `# D${i}\n\nx.`, `D${i}`); + const doc = new Document(root, { title: `D${i}` }); + doc.setDocId(`doc-reindex-${i}`); + await kb.upsertDocument(doc); + } + + kb.setAiStrategy(strategyA); + const processed = await kb.reindex(); + expect(processed).toBe(3); + // All three iterations stayed on A even though A re-installed B + // after the first call. + expect(aIngested).toHaveLength(3); + expect(bIngested).toHaveLength(0); + // The KB's current strategy is now B (set during the loop). + expect(kb.getAiStrategy()).toBe(strategyB); + }); + }); }); diff --git a/packages/test/src/test/rag/HFT_TextReranker.shape.test.ts b/packages/test/src/test/rag/HFT_TextReranker.shape.test.ts new file mode 100644 index 000000000..adca70a50 --- /dev/null +++ b/packages/test/src/test/rag/HFT_TextReranker.shape.test.ts @@ -0,0 +1,70 @@ +/** + * @license + * Copyright 2026 Steven Roussey + * SPDX-License-Identifier: Apache-2.0 + */ + +import { KbRerankerOutputError } from "@workglow/ai"; +import { validateAndExtractRerankerScores } from "@workglow/huggingface-transformers/ai-runtime"; +import { describe, expect, it } from "vitest"; + +const MODEL_PATH = "Xenova/bge-reranker-base"; + +/** + * The transformers.js text-classification pipeline returns one entry per + * input pair. With `top_k > 1` each entry is an array of `{ label, score }` + * objects; with `top_k = 1` it is the bare object. The reranker run-fn + * accepts both forms. These tests pin down the shape contract: anything + * else must fail loudly so a misconfigured model produces an actionable + * error instead of silently returning zero scores. + */ +describe("validateAndExtractRerankerScores (HFT_TextReranker shape guard)", () => { + it("throws KbRerankerOutputError when an entry lacks a numeric `score`", () => { + const run = () => validateAndExtractRerankerScores([{ foo: "bar" }], MODEL_PATH); + expect(run).toThrow(KbRerankerOutputError); + expect(run).toThrow(/unexpected pipeline output shape/); + expect(run).toThrow(new RegExp(MODEL_PATH)); + }); + + it("accepts both `{ label, score }` and `[{ label, score }]` entries and returns scores in input order", () => { + // First entry is the bare object; second is a one-element array — the + // shape transformers.js emits when `top_k > 1`. Both must validate and + // produce per-document scores in the original input order. + const scores = validateAndExtractRerankerScores( + [{ label: "LABEL_0", score: 0.42 }, [{ label: "LABEL_0", score: 0.1 }]], + MODEL_PATH + ); + expect(scores).toEqual([0.42, 0.1]); + }); + + it("throws when one entry in a batch is valid and another is malformed", () => { + // A silent `?? 0` would have hidden this and returned a 0 score for the + // bad entry. The strict guard surfaces the misconfiguration immediately. + const run = () => + validateAndExtractRerankerScores( + [{ label: "LABEL_0", score: 0.9 }, { label: "LABEL_0" /* score missing */ }], + MODEL_PATH + ); + expect(run).toThrow(KbRerankerOutputError); + }); + + it("throws when the top-level value is not an array", () => { + // Defensive: if the pipeline returns a non-array (e.g. a tensor or null) + // we still error out rather than letting Array.prototype.map throw + // somewhere downstream with a less actionable message. + const run = () => validateAndExtractRerankerScores({ score: 0.5 }, MODEL_PATH); + expect(run).toThrow(KbRerankerOutputError); + }); + + it("includes a truncated shape snippet on the error for diagnostics", () => { + try { + validateAndExtractRerankerScores([{ foo: "bar" }], MODEL_PATH); + throw new Error("expected throw"); + } catch (err) { + expect(err).toBeInstanceOf(KbRerankerOutputError); + const e = err as KbRerankerOutputError; + expect(typeof e.actualShape).toBe("string"); + expect(e.actualShape as string).toContain("foo"); + } + }); +}); diff --git a/packages/test/src/test/rag/KbAddDocumentTask.test.ts b/packages/test/src/test/rag/KbAddDocumentTask.test.ts new file mode 100644 index 000000000..66d4efde8 --- /dev/null +++ b/packages/test/src/test/rag/KbAddDocumentTask.test.ts @@ -0,0 +1,57 @@ +/** + * @license + * Copyright 2026 Steven Roussey + * SPDX-License-Identifier: Apache-2.0 + */ + +import { kbAddDocument } from "@workglow/ai"; +import { Document, createKnowledgeBase } from "@workglow/knowledge-base"; +import { uuid4 } from "@workglow/util"; +import { describe, expect, it, vi } from "vitest"; + +describe("KbAddDocumentTask", () => { + function makeDoc(docId?: string): Document { + const doc = new Document({ type: "root", title: "T", children: [] } as never, { + title: "Test", + }); + if (docId) doc.setDocId(docId); + return doc; + } + + async function makeKbWithUpsertSpy() { + const kb = await createKnowledgeBase({ + name: `kb-add-doc-${uuid4()}`, + vectorDimensions: 3, + register: false, + }); + const doc = makeDoc("returned-id"); + const upsertSpy = vi.spyOn(kb, "upsert").mockResolvedValue(doc); + return { kb, doc, upsertSpy }; + } + + it("calls kb.upsert with the provided document", async () => { + const { kb, doc, upsertSpy } = await makeKbWithUpsertSpy(); + + await kbAddDocument({ knowledgeBase: kb, document: doc }); + + expect(upsertSpy).toHaveBeenCalledTimes(1); + expect(upsertSpy.mock.calls[0][0]).toBe(doc); + }); + + it("returns the doc_id from the upserted document", async () => { + const { kb, doc } = await makeKbWithUpsertSpy(); + + const result = await kbAddDocument({ knowledgeBase: kb, document: doc }); + + expect(result.doc_id).toBe("returned-id"); + }); + + it("threads run context (signal) to kb.upsert", async () => { + const { kb, upsertSpy } = await makeKbWithUpsertSpy(); + + await kbAddDocument({ knowledgeBase: kb, document: makeDoc("x") }); + + const forwardedRunConfig = upsertSpy.mock.calls[0][1]; + expect(forwardedRunConfig).toMatchObject({ signal: expect.any(AbortSignal) }); + }); +}); diff --git a/packages/test/src/test/rag/KbDeleteTask.test.ts b/packages/test/src/test/rag/KbDeleteTask.test.ts new file mode 100644 index 000000000..0cdb86b4f --- /dev/null +++ b/packages/test/src/test/rag/KbDeleteTask.test.ts @@ -0,0 +1,48 @@ +/** + * @license + * Copyright 2026 Steven Roussey + * SPDX-License-Identifier: Apache-2.0 + */ + +import { kbDelete } from "@workglow/ai"; +import { createKnowledgeBase } from "@workglow/knowledge-base"; +import { uuid4 } from "@workglow/util"; +import { describe, expect, it, vi } from "vitest"; + +describe("KbDeleteTask", () => { + async function makeKbWithDeleteSpy() { + const kb = await createKnowledgeBase({ + name: `kb-delete-${uuid4()}`, + vectorDimensions: 3, + register: false, + }); + const deleteSpy = vi.spyOn(kb, "delete").mockResolvedValue(undefined); + return { kb, deleteSpy }; + } + + it("calls kb.delete with the given doc_id", async () => { + const { kb, deleteSpy } = await makeKbWithDeleteSpy(); + + await kbDelete({ knowledgeBase: kb, doc_id: "my-doc" }); + + expect(deleteSpy).toHaveBeenCalledTimes(1); + expect(deleteSpy.mock.calls[0][0]).toBe("my-doc"); + }); + + it("echoes doc_id in the output", async () => { + const { kb } = await makeKbWithDeleteSpy(); + + const result = await kbDelete({ knowledgeBase: kb, doc_id: "my-doc" }); + + expect(result.doc_id).toBe("my-doc"); + }); + + it("threads run context (signal) to kb.delete", async () => { + const { kb, deleteSpy } = await makeKbWithDeleteSpy(); + + await kbDelete({ knowledgeBase: kb, doc_id: "my-doc" }); + + const forwardedRunConfig = deleteSpy.mock.calls[0][1]; + expect(forwardedRunConfig).toMatchObject({ signal: expect.any(AbortSignal) }); + }); +}); diff --git a/packages/test/src/test/rag/KbSearchTask.test.ts b/packages/test/src/test/rag/KbSearchTask.test.ts new file mode 100644 index 000000000..0ad7d738f --- /dev/null +++ b/packages/test/src/test/rag/KbSearchTask.test.ts @@ -0,0 +1,79 @@ +/** + * @license + * Copyright 2026 Steven Roussey + * SPDX-License-Identifier: Apache-2.0 + */ + +import { kbSearch } from "@workglow/ai"; +import { createKnowledgeBase } from "@workglow/knowledge-base"; +import { uuid4 } from "@workglow/util"; +import { describe, expect, it, vi } from "vitest"; + +/** + * Tests that `KbSearchTask` forwards `scoreThreshold` to `kb.search`. + * The task is the public surface that downstream callers wire into a + * workflow; if the schema accepts `scoreThreshold` but `execute` drops + * it, callers silently get unfiltered results. + */ +describe("KbSearchTask scoreThreshold forwarding", () => { + /** + * Build a KB whose `search` method is replaced with a spy that returns + * an empty result list. The spy lets us assert exactly which options + * the task hands down. + */ + async function makeKbWithSearchSpy() { + const kb = await createKnowledgeBase({ + name: `kb-search-task-${uuid4()}`, + vectorDimensions: 3, + register: false, + }); + const searchSpy = vi.spyOn(kb, "search").mockResolvedValue([]); + return { kb, searchSpy }; + } + + it("forwards `scoreThreshold` to kb.search when provided", async () => { + const { kb, searchSpy } = await makeKbWithSearchSpy(); + + await kbSearch({ + knowledgeBase: kb, + query: "hello", + topK: 3, + scoreThreshold: 0.42, + }); + + expect(searchSpy).toHaveBeenCalledTimes(1); + const [forwardedQuery, forwardedOpts] = searchSpy.mock.calls[0]; + expect(forwardedQuery).toBe("hello"); + expect(forwardedOpts).toMatchObject({ topK: 3, scoreThreshold: 0.42 }); + }); + + it("passes `scoreThreshold: undefined` to kb.search when omitted", async () => { + const { kb, searchSpy } = await makeKbWithSearchSpy(); + + await kbSearch({ + knowledgeBase: kb, + query: "hello", + topK: 3, + }); + + expect(searchSpy).toHaveBeenCalledTimes(1); + const [, forwardedOpts] = searchSpy.mock.calls[0]; + // Property may either be absent or explicitly undefined; both + // behave identically downstream. What we really care about is + // that it's NOT a stale value carried over from a previous call. + expect((forwardedOpts as { scoreThreshold?: number }).scoreThreshold).toBeUndefined(); + }); + + it("threads run context (signal) to kb.search", async () => { + const { kb, searchSpy } = await makeKbWithSearchSpy(); + + await kbSearch({ knowledgeBase: kb, query: "hello" }); + + expect(searchSpy).toHaveBeenCalledTimes(1); + const call = searchSpy.mock.calls[0]; + // Third argument is the runConfig extracted from IExecuteContext. + const forwardedRunConfig = call[2]; + expect(forwardedRunConfig).toBeDefined(); + expect(forwardedRunConfig).toMatchObject({ signal: expect.any(AbortSignal) }); + }); +}); diff --git a/packages/test/src/test/rag/KnowledgeBaseStandardStrategy.test.ts b/packages/test/src/test/rag/KnowledgeBaseStandardStrategy.test.ts new file mode 100644 index 000000000..ca8fab88c --- /dev/null +++ b/packages/test/src/test/rag/KnowledgeBaseStandardStrategy.test.ts @@ -0,0 +1,200 @@ +/** + * @license + * Copyright 2026 Steven Roussey + * SPDX-License-Identifier: Apache-2.0 + */ + +import type { AiProviderRunFn, ModelRecord, TextEmbeddingTaskInput } from "@workglow/ai"; +import { + createStandardKbStrategy, + getAiProviderRegistry, + getGlobalModelRepository, +} from "@workglow/ai"; +import type { + ChunkVectorStorage, + DocumentTabularStorage, + InsertChunkVectorEntity, +} from "@workglow/knowledge-base"; +import { + ChunkVectorPrimaryKey, + ChunkVectorStorageSchema, + Document, + DocumentStorageKey, + DocumentStorageSchema, + KnowledgeBase, + StructuralParser, + createKnowledgeBase, +} from "@workglow/knowledge-base"; +import { InMemoryTabularStorage, InMemoryVectorStorage } from "@workglow/storage"; +import { uuid4 } from "@workglow/util"; +import { beforeAll, describe, expect, it } from "vitest"; + +/** + * Tests exercising `createStandardKbStrategy` directly. Setup registers a + * tiny stub provider for `TextEmbeddingTask` so we don't need a real + * runtime (HuggingFace etc.) to assert order/tagging contracts. + */ +const TEST_PROVIDER = "test-strategy-provider"; +const TEST_EMBED_MODEL_ID = "test:strategy:embed"; + +describe("createStandardKbStrategy", () => { + beforeAll(async () => { + const registry = getAiProviderRegistry(); + const runFn: AiProviderRunFn = async (input, _model, _signal, emit) => { + // Deterministic 3-D unit vector keyed off the first text character; + // we don't need vector meaning here, only that embedTexts resolves. + const embeddingInput = input as TextEmbeddingTaskInput; + const texts = Array.isArray(embeddingInput.text) + ? embeddingInput.text + : [embeddingInput.text ?? ""]; + const vectors = texts.map(() => new Float32Array([1, 0, 0])); + emit({ + type: "finish", + data: { + vector: vectors.length === 1 ? vectors[0] : vectors, + }, + }); + }; + registry.registerRunFn(TEST_PROVIDER, { + serves: ["text.embedding"], + runFn, + }); + + const modelRepo = getGlobalModelRepository(); + const existing = await modelRepo.findByName(TEST_EMBED_MODEL_ID).catch(() => undefined); + if (!existing) { + await modelRepo.addModel({ + model_id: TEST_EMBED_MODEL_ID, + capabilities: ["text.embedding"], + title: "Strategy test embed model", + description: "Stub embed model used by createStandardKbStrategy tests", + provider: TEST_PROVIDER, + provider_config: { native_dimensions: 3 }, + metadata: {}, + } as ModelRecord); + } + }); + + /** + * Seed the KB with a single pre-existing chunk so a re-ingest has + * something to delete and the partial-failure test can verify it's gone. + */ + async function seedChunk(kb: KnowledgeBase, doc_id: string, chunk_id: string): Promise { + const insert: InsertChunkVectorEntity = { + chunk_id, + doc_id, + vector: new Float32Array([1, 0, 0]), + metadata: { + chunk_id, + doc_id, + text: "old chunk text", + nodePath: [], + depth: 0, + } as never, + }; + await kb.upsertChunksBulk([insert]); + } + + it("ingest deletes existing chunks BEFORE upsertDocument when doc_id is set; partial failure leaves no orphan chunks", async () => { + // KB subclass that rejects on upsertChunksBulk to simulate a failure + // partway through ingest (after delete, after document upsert, after + // chunker, after embed, but during the bulk insert). + class FailingKb extends KnowledgeBase { + failOnBulkInsert = false; + override async upsertChunksBulk(chunks: InsertChunkVectorEntity[]) { + if (this.failOnBulkInsert) { + throw new Error("simulated bulk-insert failure"); + } + return super.upsertChunksBulk(chunks); + } + } + + const tabular = new InMemoryTabularStorage(DocumentStorageSchema, DocumentStorageKey); + await tabular.setupDatabase(); + const vector = new InMemoryVectorStorage( + ChunkVectorStorageSchema, + ChunkVectorPrimaryKey, + [], + 3, + Float32Array + ); + await vector.setupDatabase(); + + const kb = new FailingKb( + `kb-ingest-order-${uuid4()}`, + tabular as unknown as DocumentTabularStorage, + vector as unknown as ChunkVectorStorage, + { docEmbeddingModel: TEST_EMBED_MODEL_ID } + ); + kb.setAiStrategy(createStandardKbStrategy()); + + const docId = "doc-ingest-order"; + // First, plant the document + a stale chunk that the next ingest + // should clear out. + const initialRoot = await StructuralParser.parseMarkdown( + uuid4(), + "# Initial\n\nold content.", + "Initial" + ); + const initialDoc = new Document(initialRoot, { title: "Initial" }); + initialDoc.setDocId(docId); + await kb.upsertDocument(initialDoc); + await seedChunk(kb, docId, "stale-chunk-1"); + expect((await kb.getChunksForDocument(docId)).length).toBe(1); + + // Now arm the failure and re-ingest the same doc_id. The strategy + // should: (1) delete the stale chunk, (2) upsert the new document + // version, (3) chunk + embed, (4) call upsertChunksBulk which + // throws. Post-failure: stale chunk still gone, document row + // reflects the new (re-upserted) content. + kb.failOnBulkInsert = true; + const newRoot = await StructuralParser.parseMarkdown( + uuid4(), + "# Updated\n\nnew content.", + "Updated" + ); + const updatedDoc = new Document(newRoot, { title: "Updated" }); + updatedDoc.setDocId(docId); + + await expect(kb.upsert(updatedDoc)).rejects.toThrow(/simulated bulk-insert failure/); + + // Chunks: empty (stale gone, new ones never inserted) — the data- + // integrity invariant of the new ordering. + expect(await kb.getChunksForDocument(docId)).toEqual([]); + // Document row: present (upserted before the failure), with the new + // title — the new content "won" even though the chunks didn't. + const storedDoc = await kb.getDocument(docId); + expect(storedDoc).toBeDefined(); + expect(storedDoc!.metadata.title).toBe("Updated"); + }); + + it("rerank mode tags results with scoreType: 'rerank' via the heuristic fallback (no rerankerModel)", async () => { + const kb = await createKnowledgeBase({ + name: `kb-rerank-tag-${uuid4()}`, + vectorDimensions: 3, + register: false, + docEmbeddingModel: TEST_EMBED_MODEL_ID, + searchMode: "rerank", + // Intentionally no rerankerModel → heuristic RerankerTask fallback. + }); + kb.setAiStrategy(createStandardKbStrategy()); + + // Plant a doc + chunk so the first stage retrieves something for the + // reranker to score. + const docId = "doc-rerank-tag"; + const root = await StructuralParser.parseMarkdown(uuid4(), "# T\n\nhi.", "T"); + const doc = new Document(root, { title: "T" }); + doc.setDocId(docId); + await kb.upsertDocument(doc); + await seedChunk(kb, docId, "chunk-rerank-tag"); + + const results = await kb.search("hi", { topK: 1 }); + + expect(results.length).toBeGreaterThan(0); + // The first-stage retrieval would have produced "cosine" scores; the + // rerank fallback MUST override them to "rerank". + for (const r of results) { + expect(r.scoreType).toBe("rerank"); + } + }); +}); diff --git a/packages/test/src/test/task/KbSearchTask.test.ts b/packages/test/src/test/task/KbSearchTask.test.ts index 44e888e16..74cf26c03 100644 --- a/packages/test/src/test/task/KbSearchTask.test.ts +++ b/packages/test/src/test/task/KbSearchTask.test.ts @@ -82,4 +82,23 @@ describe("KbSearchTask — execute()", () => { // @ts-expect-error — test helper expect(kb._calls[0].opts).toMatchObject({ topK: 5 }); }); + + it("throws (with the offending chunk_id) when a result is missing metadata.text", async () => { + // Custom-strategy KBs are free to return whatever shape they like, but + // chunks without `metadata.text` violate the documented contract on + // InsertChunkVectorEntity. `KbSearchTask` enforces that contract via + // `chunkText`; previously it silently fell back to + // JSON.stringify(metadata), which surfaced as garbage downstream. + const offending: ChunkSearchResult = { + chunk_id: "c-no-text", + doc_id: "d1", + vector: new Float32Array([1, 0, 0]), + // `text` intentionally absent — custom chunker forgot it. + metadata: { custom: "x" } as never, + score: 1, + }; + const kb = makeFakeKb([offending]); + const task = new KbSearchTask(); + await expect(task.run({ knowledgeBase: kb, query: "q" })).rejects.toThrow(/c-no-text/); + }); }); diff --git a/packages/util/CHANGELOG.md b/packages/util/CHANGELOG.md index dce05ca6e..a91dbe7c6 100644 --- a/packages/util/CHANGELOG.md +++ b/packages/util/CHANGELOG.md @@ -1,5 +1,17 @@ # @workglow/util +## 0.2.36 + +### Bug Fixes + +#### util/worker, ai/task + +- TTL-based pendingAborts eviction; clarify runWithIterable bond (#500) + +#### util + +- snapshot-then-delete eviction in WorkerServerBase Set caps + ## 0.2.35 ### Features diff --git a/packages/util/package.json b/packages/util/package.json index 443a572f4..94f387a63 100644 --- a/packages/util/package.json +++ b/packages/util/package.json @@ -1,7 +1,7 @@ { "name": "@workglow/util", "type": "module", - "version": "0.2.35", + "version": "0.2.36", "repository": { "type": "git", "url": "https://github.com/workglow-dev/libs.git", diff --git a/packages/workglow/CHANGELOG.md b/packages/workglow/CHANGELOG.md index 61497509f..19d1a9567 100644 --- a/packages/workglow/CHANGELOG.md +++ b/packages/workglow/CHANGELOG.md @@ -1,5 +1,7 @@ # workglow +## 0.2.36 + ## 0.2.35 ### Performance diff --git a/packages/workglow/package.json b/packages/workglow/package.json index c877d4868..28990e481 100644 --- a/packages/workglow/package.json +++ b/packages/workglow/package.json @@ -1,7 +1,7 @@ { "name": "workglow", "type": "module", - "version": "0.2.35", + "version": "0.2.36", "repository": { "type": "git", "url": "https://github.com/workglow-dev/libs.git", diff --git a/providers/anthropic/CHANGELOG.md b/providers/anthropic/CHANGELOG.md index 8891d36dc..969210069 100644 --- a/providers/anthropic/CHANGELOG.md +++ b/providers/anthropic/CHANGELOG.md @@ -1,5 +1,7 @@ # Changelog +## 0.2.36 + ## 0.2.35 ### Features diff --git a/providers/anthropic/package.json b/providers/anthropic/package.json index 825e1dbcb..508626e0a 100644 --- a/providers/anthropic/package.json +++ b/providers/anthropic/package.json @@ -2,7 +2,7 @@ "name": "@workglow/anthropic", "type": "module", "sideEffects": false, - "version": "0.2.35", + "version": "0.2.36", "repository": { "type": "git", "url": "https://github.com/workglow-dev/libs.git", diff --git a/providers/bun-webview/CHANGELOG.md b/providers/bun-webview/CHANGELOG.md index 6371a8c99..c0525838f 100644 --- a/providers/bun-webview/CHANGELOG.md +++ b/providers/bun-webview/CHANGELOG.md @@ -1,5 +1,7 @@ # Changelog +## 0.2.36 + ## 0.2.35 ### Chores diff --git a/providers/bun-webview/package.json b/providers/bun-webview/package.json index c8c3d8879..6f4e0ddbd 100644 --- a/providers/bun-webview/package.json +++ b/providers/bun-webview/package.json @@ -2,7 +2,7 @@ "name": "@workglow/bun-webview", "type": "module", "sideEffects": false, - "version": "0.2.35", + "version": "0.2.36", "repository": { "type": "git", "url": "https://github.com/workglow-dev/libs.git", diff --git a/providers/chrome-ai/CHANGELOG.md b/providers/chrome-ai/CHANGELOG.md index 0d57efd3c..f36c971a6 100644 --- a/providers/chrome-ai/CHANGELOG.md +++ b/providers/chrome-ai/CHANGELOG.md @@ -1,5 +1,7 @@ # Changelog +## 0.2.36 + ## 0.2.35 ### Features diff --git a/providers/chrome-ai/package.json b/providers/chrome-ai/package.json index 35b1f708e..a2a811848 100644 --- a/providers/chrome-ai/package.json +++ b/providers/chrome-ai/package.json @@ -2,7 +2,7 @@ "name": "@workglow/chrome-ai", "type": "module", "sideEffects": false, - "version": "0.2.35", + "version": "0.2.36", "repository": { "type": "git", "url": "https://github.com/workglow-dev/libs.git", diff --git a/providers/electron/CHANGELOG.md b/providers/electron/CHANGELOG.md index 6371a8c99..c0525838f 100644 --- a/providers/electron/CHANGELOG.md +++ b/providers/electron/CHANGELOG.md @@ -1,5 +1,7 @@ # Changelog +## 0.2.36 + ## 0.2.35 ### Chores diff --git a/providers/electron/package.json b/providers/electron/package.json index 5558fb4cc..72c2988c0 100644 --- a/providers/electron/package.json +++ b/providers/electron/package.json @@ -2,7 +2,7 @@ "name": "@workglow/electron", "type": "module", "sideEffects": false, - "version": "0.2.35", + "version": "0.2.36", "repository": { "type": "git", "url": "https://github.com/workglow-dev/libs.git", diff --git a/providers/google-gemini/CHANGELOG.md b/providers/google-gemini/CHANGELOG.md index 13c6fbb08..ae4c0bc56 100644 --- a/providers/google-gemini/CHANGELOG.md +++ b/providers/google-gemini/CHANGELOG.md @@ -1,5 +1,7 @@ # Changelog +## 0.2.36 + ## 0.2.35 ### Features diff --git a/providers/google-gemini/package.json b/providers/google-gemini/package.json index a44dc9896..a2a45ae5c 100644 --- a/providers/google-gemini/package.json +++ b/providers/google-gemini/package.json @@ -2,7 +2,7 @@ "name": "@workglow/google-gemini", "type": "module", "sideEffects": false, - "version": "0.2.35", + "version": "0.2.36", "repository": { "type": "git", "url": "https://github.com/workglow-dev/libs.git", diff --git a/providers/huggingface-inference/CHANGELOG.md b/providers/huggingface-inference/CHANGELOG.md index c05579cb4..059ded0f7 100644 --- a/providers/huggingface-inference/CHANGELOG.md +++ b/providers/huggingface-inference/CHANGELOG.md @@ -1,5 +1,7 @@ # Changelog +## 0.2.36 + ## 0.2.35 ### Features diff --git a/providers/huggingface-inference/package.json b/providers/huggingface-inference/package.json index fd47cc913..0b9ed1f1f 100644 --- a/providers/huggingface-inference/package.json +++ b/providers/huggingface-inference/package.json @@ -2,7 +2,7 @@ "name": "@workglow/huggingface-inference", "type": "module", "sideEffects": false, - "version": "0.2.35", + "version": "0.2.36", "repository": { "type": "git", "url": "https://github.com/workglow-dev/libs.git", diff --git a/providers/huggingface-transformers/CHANGELOG.md b/providers/huggingface-transformers/CHANGELOG.md index 80e6f7205..e25ce8ae1 100644 --- a/providers/huggingface-transformers/CHANGELOG.md +++ b/providers/huggingface-transformers/CHANGELOG.md @@ -1,5 +1,13 @@ # Changelog +## 0.2.36 + +### Features + +#### hft + +- add text reranker capability with pipeline shape validation + ## 0.2.35 ### Features diff --git a/providers/huggingface-transformers/package.json b/providers/huggingface-transformers/package.json index 63a12b48a..1bbaacf4e 100644 --- a/providers/huggingface-transformers/package.json +++ b/providers/huggingface-transformers/package.json @@ -2,7 +2,7 @@ "name": "@workglow/huggingface-transformers", "type": "module", "sideEffects": false, - "version": "0.2.35", + "version": "0.2.36", "repository": { "type": "git", "url": "https://github.com/workglow-dev/libs.git", diff --git a/providers/huggingface-transformers/src/ai/common/HFT_CapabilitySets.ts b/providers/huggingface-transformers/src/ai/common/HFT_CapabilitySets.ts index 7389c7215..207706785 100644 --- a/providers/huggingface-transformers/src/ai/common/HFT_CapabilitySets.ts +++ b/providers/huggingface-transformers/src/ai/common/HFT_CapabilitySets.ts @@ -28,6 +28,7 @@ export const HFT_TEXT_CLASSIFICATION = ["text.classification"] as const satisfie export const HFT_TEXT_LANGUAGE_DETECTION = [ "text.language-detection", ] as const satisfies Capability[]; +export const HFT_TEXT_RERANKING = ["text.reranking"] as const satisfies Capability[]; export const HFT_TEXT_FILL_MASK = ["text.fill-mask"] as const satisfies Capability[]; export const HFT_TEXT_NER = ["text.ner"] as const satisfies Capability[]; export const HFT_IMAGE_CLASSIFICATION = ["image.classification"] as const satisfies Capability[]; @@ -58,6 +59,7 @@ export const HFT_CAPABILITY_SETS = [ HFT_TEXT_EMBEDDING, HFT_TEXT_CLASSIFICATION, HFT_TEXT_LANGUAGE_DETECTION, + HFT_TEXT_RERANKING, HFT_TEXT_FILL_MASK, HFT_TEXT_NER, HFT_IMAGE_CLASSIFICATION, diff --git a/providers/huggingface-transformers/src/ai/common/HFT_JobRunFns.ts b/providers/huggingface-transformers/src/ai/common/HFT_JobRunFns.ts index 519d9d260..dcaeb5374 100644 --- a/providers/huggingface-transformers/src/ai/common/HFT_JobRunFns.ts +++ b/providers/huggingface-transformers/src/ai/common/HFT_JobRunFns.ts @@ -29,6 +29,7 @@ import { HFT_TEXT_LANGUAGE_DETECTION, HFT_TEXT_NER, HFT_TEXT_QUESTION_ANSWERING, + HFT_TEXT_RERANKING, HFT_TEXT_REWRITER, HFT_TEXT_SUMMARY, HFT_TEXT_TRANSLATION, @@ -56,6 +57,7 @@ import { HFT_TextGeneration } from "./HFT_TextGeneration"; import { HFT_TextLanguageDetection } from "./HFT_TextLanguageDetection"; import { HFT_TextNamedEntityRecognition } from "./HFT_TextNamedEntityRecognition"; import { HFT_TextQuestionAnswer } from "./HFT_TextQuestionAnswer"; +import { HFT_TextReranker } from "./HFT_TextReranker"; import { HFT_TextRewriter } from "./HFT_TextRewriter"; import { HFT_TextSummary } from "./HFT_TextSummary"; import { HFT_TextTranslation } from "./HFT_TextTranslation"; @@ -109,6 +111,7 @@ export const HFT_RUN_FNS: readonly AiProviderRunFnRegistration< { serves: HFT_TEXT_EMBEDDING, runFn: HFT_TextEmbedding }, { serves: HFT_TEXT_CLASSIFICATION, runFn: HFT_TextClassification }, { serves: HFT_TEXT_LANGUAGE_DETECTION, runFn: HFT_TextLanguageDetection }, + { serves: HFT_TEXT_RERANKING, runFn: HFT_TextReranker }, { serves: HFT_TEXT_FILL_MASK, runFn: HFT_TextFillMask }, { serves: HFT_TEXT_NER, runFn: HFT_TextNamedEntityRecognition }, { serves: HFT_IMAGE_CLASSIFICATION, runFn: HFT_ImageClassification }, diff --git a/providers/huggingface-transformers/src/ai/common/HFT_TextReranker.ts b/providers/huggingface-transformers/src/ai/common/HFT_TextReranker.ts new file mode 100644 index 000000000..3ca56e676 --- /dev/null +++ b/providers/huggingface-transformers/src/ai/common/HFT_TextReranker.ts @@ -0,0 +1,140 @@ +/** + * @license + * Copyright 2026 Steven Roussey + * SPDX-License-Identifier: Apache-2.0 + */ + +import type { TextClassificationPipeline } from "@huggingface/transformers"; +import type { AiProviderRunFn, TextRerankerTaskInput, TextRerankerTaskOutput } from "@workglow/ai"; +import { KbRerankerOutputError } from "@workglow/ai"; +import { getLogger } from "@workglow/util/worker"; +import type { HfTransformersOnnxModelConfig } from "./HFT_ModelSchema"; +import { getPipeline } from "./HFT_Pipeline"; + +/** + * Narrow guard: an entry from the text-classification pipeline must be an + * object with a numeric `score`. We deliberately do not require `label` + * to be present because some downstream models omit it; only `score` is + * load-bearing for reranking. + */ +function isScored(v: unknown): v is { label?: string; score: number } { + return ( + typeof v === "object" && v !== null && typeof (v as { score?: unknown }).score === "number" + ); +} + +/** + * Serialize an offending shape for the error payload. We cap the length to + * keep error messages bounded — a misconfigured model could otherwise dump + * unbounded data into logs. + */ +function truncateShape(value: unknown): string { + try { + const json = JSON.stringify(value); + if (typeof json !== "string") return String(value); + return json.length > 200 ? `${json.slice(0, 200)}…` : json; + } catch { + return String(value); + } +} + +/** + * Validate a transformers.js text-classification pipeline output and + * extract per-document scores. Each entry must be either a `{ score }` + * object or a non-empty array of such objects (transformers.js returns + * the array form when `top_k > 1`). + * + * Exported so the shape contract can be exercised in tests directly, + * without needing to spin up a real pipeline or mock the loader. Throws + * {@link KbRerankerOutputError} on mismatch — silently coercing to 0 + * would hide real model config bugs. + * + * @param rawResults Whatever the pipeline call returned. Typed as + * `unknown` so the caller is forced through this validation. + * @param modelPath Used in the error message to point operators at the + * misconfigured model. + */ +export function validateAndExtractRerankerScores( + rawResults: unknown, + modelPath: string | undefined +): number[] { + if (!Array.isArray(rawResults)) { + throw new KbRerankerOutputError( + `HFT_TextReranker: unexpected pipeline output shape for model ${modelPath}`, + truncateShape(rawResults) + ); + } + + const scores: number[] = new Array(rawResults.length); + for (let i = 0; i < rawResults.length; i++) { + const entry = rawResults[i]; + const candidate = Array.isArray(entry) ? entry[0] : entry; + if (!isScored(candidate)) { + throw new KbRerankerOutputError( + `HFT_TextReranker: unexpected pipeline output shape for model ${modelPath}`, + truncateShape(entry) + ); + } + scores[i] = candidate.score; + } + return scores; +} + +/** + * Cross-encoder reranker run-fn. Loads a `text-classification` pipeline + * (the way transformers.js exposes cross-encoder models like + * `Xenova/bge-reranker-base`) and scores each `[query, doc]` pair. + * + * Output `indices` is sorted best-first; `scores` is the per-document score + * in the original input order so callers can join back to their candidate + * list without re-sorting. + * + * Each pipeline result entry is validated at runtime via + * {@link validateAndExtractRerankerScores}: either a `{ score }` object or + * a non-empty array of such objects. On mismatch we throw + * {@link KbRerankerOutputError} with the model path and a truncated shape + * snippet — silently coercing missing scores to 0 would hide real model + * config bugs. + */ +export const HFT_TextReranker: AiProviderRunFn< + TextRerankerTaskInput, + TextRerankerTaskOutput, + HfTransformersOnnxModelConfig +> = async (input, model, signal, emit) => { + const logger = getLogger(); + const modelPath = model?.provider_config.model_path; + const timerLabel = `hft:TextReranker:${modelPath}`; + logger.time(timerLabel, { docs: input.documents.length }); + + const reranker: TextClassificationPipeline = await getPipeline(model!, emit, {}, signal); + + // Transformers.js' text-classification pipeline accepts an array of + // { text, text_pair } objects for sentence-pair tasks (which cross-encoder + // rerankers are). The pipeline returns one score per input pair (or an + // array of scored entries per pair when `top_k > 1`). + const pairs = input.documents.map((doc) => ({ text: input.query, text_pair: doc })); + + // We cast the *input parameter* shape (transformers.js types the pipeline + // very loosely) but keep the return as `Promise` so the guard at + // `validateAndExtractRerankerScores` is the only path to a typed score. + // The original bug was casting the return to `Promise>`, + // which erased the safety net and let a mismatched pipeline silently + // return garbage. + const callable = reranker as unknown as ( + inputs: ReadonlyArray<{ text: string; text_pair: string }>, + options?: Record + ) => Promise; + const rawResults: unknown = await callable(pairs, { top_k: 1 }); + + const scores = validateAndExtractRerankerScores(rawResults, modelPath); + + const indices = scores + .map((score, idx) => ({ score, idx })) + .sort((a, b) => b.score - a.score) + .map((p) => p.idx); + + const limited = typeof input.topK === "number" ? indices.slice(0, input.topK) : indices; + + logger.timeEnd(timerLabel, { docs: input.documents.length }); + emit({ type: "finish", data: { scores, indices: limited } }); +}; diff --git a/providers/huggingface-transformers/src/ai/runtime.ts b/providers/huggingface-transformers/src/ai/runtime.ts index cc51b4fd5..7ad044462 100644 --- a/providers/huggingface-transformers/src/ai/runtime.ts +++ b/providers/huggingface-transformers/src/ai/runtime.ts @@ -19,6 +19,7 @@ export * from "./common/HFT_Constants"; export * from "./common/HFT_ModelSchema"; export * from "./common/HFT_OnnxDtypes"; export * from "./common/HFT_Pipeline"; +export * from "./common/HFT_TextReranker"; export * from "./common/HFT_ToolMarkup"; export * from "./registerHuggingFaceTransformersInline"; export * from "./registerHuggingFaceTransformersWorker"; diff --git a/providers/node-llama-cpp/CHANGELOG.md b/providers/node-llama-cpp/CHANGELOG.md index a276737a1..c77c55017 100644 --- a/providers/node-llama-cpp/CHANGELOG.md +++ b/providers/node-llama-cpp/CHANGELOG.md @@ -1,5 +1,7 @@ # Changelog +## 0.2.36 + ## 0.2.35 ### Features diff --git a/providers/node-llama-cpp/package.json b/providers/node-llama-cpp/package.json index ff09b7acd..b0fc49490 100644 --- a/providers/node-llama-cpp/package.json +++ b/providers/node-llama-cpp/package.json @@ -2,7 +2,7 @@ "name": "@workglow/node-llama-cpp", "type": "module", "sideEffects": false, - "version": "0.2.35", + "version": "0.2.36", "repository": { "type": "git", "url": "https://github.com/workglow-dev/libs.git", diff --git a/providers/ollama/CHANGELOG.md b/providers/ollama/CHANGELOG.md index 9c8b2268e..0f0b0e17b 100644 --- a/providers/ollama/CHANGELOG.md +++ b/providers/ollama/CHANGELOG.md @@ -1,5 +1,7 @@ # Changelog +## 0.2.36 + ## 0.2.35 ### Features diff --git a/providers/ollama/package.json b/providers/ollama/package.json index 43a6713d9..b26c5f852 100644 --- a/providers/ollama/package.json +++ b/providers/ollama/package.json @@ -2,7 +2,7 @@ "name": "@workglow/ollama", "type": "module", "sideEffects": false, - "version": "0.2.35", + "version": "0.2.36", "repository": { "type": "git", "url": "https://github.com/workglow-dev/libs.git", diff --git a/providers/openai/CHANGELOG.md b/providers/openai/CHANGELOG.md index 714004cd6..97f44721e 100644 --- a/providers/openai/CHANGELOG.md +++ b/providers/openai/CHANGELOG.md @@ -1,5 +1,7 @@ # Changelog +## 0.2.36 + ## 0.2.35 ### Features diff --git a/providers/openai/package.json b/providers/openai/package.json index 3675d3036..6c85c66fa 100644 --- a/providers/openai/package.json +++ b/providers/openai/package.json @@ -2,7 +2,7 @@ "name": "@workglow/openai", "type": "module", "sideEffects": false, - "version": "0.2.35", + "version": "0.2.36", "repository": { "type": "git", "url": "https://github.com/workglow-dev/libs.git", diff --git a/providers/playwright/CHANGELOG.md b/providers/playwright/CHANGELOG.md index 6371a8c99..c0525838f 100644 --- a/providers/playwright/CHANGELOG.md +++ b/providers/playwright/CHANGELOG.md @@ -1,5 +1,7 @@ # Changelog +## 0.2.36 + ## 0.2.35 ### Chores diff --git a/providers/playwright/package.json b/providers/playwright/package.json index b3071023c..c40108fad 100644 --- a/providers/playwright/package.json +++ b/providers/playwright/package.json @@ -2,7 +2,7 @@ "name": "@workglow/playwright", "type": "module", "sideEffects": false, - "version": "0.2.35", + "version": "0.2.36", "repository": { "type": "git", "url": "https://github.com/workglow-dev/libs.git", diff --git a/providers/postgres/CHANGELOG.md b/providers/postgres/CHANGELOG.md index 068ad8034..ccce9f716 100644 --- a/providers/postgres/CHANGELOG.md +++ b/providers/postgres/CHANGELOG.md @@ -1,5 +1,7 @@ # Changelog +## 0.2.36 + ## 0.2.35 ### Features diff --git a/providers/postgres/package.json b/providers/postgres/package.json index 6d89fa21d..c86c5c75f 100644 --- a/providers/postgres/package.json +++ b/providers/postgres/package.json @@ -1,7 +1,7 @@ { "name": "@workglow/postgres", "type": "module", - "version": "0.2.35", + "version": "0.2.36", "repository": { "type": "git", "url": "https://github.com/workglow-dev/libs.git", diff --git a/providers/sqlite/CHANGELOG.md b/providers/sqlite/CHANGELOG.md index c582423ff..1fb25d1b0 100644 --- a/providers/sqlite/CHANGELOG.md +++ b/providers/sqlite/CHANGELOG.md @@ -1,5 +1,7 @@ # Changelog +## 0.2.36 + ## 0.2.35 ### Features diff --git a/providers/sqlite/package.json b/providers/sqlite/package.json index b56dc5764..92df30d30 100644 --- a/providers/sqlite/package.json +++ b/providers/sqlite/package.json @@ -1,7 +1,7 @@ { "name": "@workglow/sqlite", "type": "module", - "version": "0.2.35", + "version": "0.2.36", "repository": { "type": "git", "url": "https://github.com/workglow-dev/libs.git", diff --git a/providers/supabase/CHANGELOG.md b/providers/supabase/CHANGELOG.md index 161271fce..a44bea7e5 100644 --- a/providers/supabase/CHANGELOG.md +++ b/providers/supabase/CHANGELOG.md @@ -1,5 +1,7 @@ # Changelog +## 0.2.36 + ## 0.2.35 ### Documentation diff --git a/providers/supabase/package.json b/providers/supabase/package.json index bd9ece763..60152d7be 100644 --- a/providers/supabase/package.json +++ b/providers/supabase/package.json @@ -1,7 +1,7 @@ { "name": "@workglow/supabase", "type": "module", - "version": "0.2.35", + "version": "0.2.36", "repository": { "type": "git", "url": "https://github.com/workglow-dev/libs.git", diff --git a/providers/tf-mediapipe/CHANGELOG.md b/providers/tf-mediapipe/CHANGELOG.md index 379201bd6..2faac2e7c 100644 --- a/providers/tf-mediapipe/CHANGELOG.md +++ b/providers/tf-mediapipe/CHANGELOG.md @@ -1,5 +1,7 @@ # Changelog +## 0.2.36 + ## 0.2.35 ### Features diff --git a/providers/tf-mediapipe/package.json b/providers/tf-mediapipe/package.json index 885d0925b..dce0eb222 100644 --- a/providers/tf-mediapipe/package.json +++ b/providers/tf-mediapipe/package.json @@ -2,7 +2,7 @@ "name": "@workglow/tf-mediapipe", "type": "module", "sideEffects": false, - "version": "0.2.35", + "version": "0.2.36", "repository": { "type": "git", "url": "https://github.com/workglow-dev/libs.git",