dify/knowledge-fs/packages/api/src/model-capability-preflight.ts

664 lines
22 KiB
TypeScript

import { createHash } from "node:crypto";
import {
type KnowledgeSpaceModelSelection,
KnowledgeSpaceModelSelectionSchema,
stableJson,
} from "@knowledge/core";
import type { EmbeddingProvider, RerankerProvider } from "@knowledge/embeddings";
import { z } from "zod";
import { resolveVectorIndexCapability } from "./vector-index-capability";
export const ModelCapabilityKindSchema = z.enum(["embedding", "reasoning", "rerank"]);
export type ModelCapabilityKind = z.infer<typeof ModelCapabilityKindSchema>;
export const ModelCatalogEntrySchema = z
.object({
capabilities: z.record(z.unknown()).default({}),
kinds: z.array(ModelCapabilityKindSchema).min(1),
model: z.string().trim().min(1).max(256),
pluginId: z.string().trim().min(1).max(256),
pluginUniqueIdentifier: z.string().trim().min(1).max(1024),
pluginVersion: z.string().trim().min(1).max(256).optional(),
provider: z.string().trim().min(1).max(256),
schemaFingerprint: z.string().regex(/^sha256:[a-f0-9]{64}$/),
})
.strict();
export type ModelCatalogEntry = z.infer<typeof ModelCatalogEntrySchema>;
export interface ResolveModelCatalogEntryInput {
readonly kind: ModelCapabilityKind;
readonly selection: KnowledgeSpaceModelSelection;
readonly signal?: AbortSignal | undefined;
readonly tenantId: string;
}
export interface ListModelCatalogEntriesInput {
readonly cursor?: string | undefined;
readonly kind?: ModelCapabilityKind | undefined;
readonly limit: number;
readonly signal?: AbortSignal | undefined;
readonly tenantId: string;
}
export interface ListModelCatalogEntriesResult {
readonly items: readonly ModelCatalogEntry[];
readonly nextCursor?: string | undefined;
}
/** Tenant-scoped view of models that Dify reports as active and invokable. */
export interface ModelCapabilityCatalog {
list(input: ListModelCatalogEntriesInput): Promise<ListModelCatalogEntriesResult>;
resolve(input: ResolveModelCatalogEntryInput): Promise<ModelCatalogEntry | null>;
/** Optional runtime validation before the active invocation probe. */
validate?(input: ResolveModelCatalogEntryInput): Promise<boolean>;
}
export const ModelCapabilitySnapshotSchema = z
.object({
capabilityDigest: z.string().regex(/^sha256:[a-f0-9]{64}$/),
checkedAt: z.string().datetime({ offset: true }),
dimension: z.number().int().positive().optional(),
distanceMetric: z.enum(["cosine", "dot", "l2"]).optional(),
kind: ModelCapabilityKindSchema,
pluginUniqueIdentifier: z.string().trim().min(1).max(1024),
pluginVersion: z.string().trim().min(1).max(256).optional(),
schemaFingerprint: z.string().regex(/^sha256:[a-f0-9]{64}$/),
selection: KnowledgeSpaceModelSelectionSchema,
})
.strict()
.superRefine((snapshot, context) => {
if (snapshot.kind === "embedding" && snapshot.dimension === undefined) {
context.addIssue({
code: z.ZodIssueCode.custom,
message: "Embedding capability snapshots require an observed dimension",
path: ["dimension"],
});
}
if (snapshot.kind !== "embedding" && snapshot.dimension !== undefined) {
context.addIssue({
code: z.ZodIssueCode.custom,
message: "Only embedding capability snapshots may contain a dimension",
path: ["dimension"],
});
}
});
export type ModelCapabilitySnapshot = z.infer<typeof ModelCapabilitySnapshotSchema>;
export type ModelCapabilityPreflightErrorCode =
| "EMBEDDING_DIMENSION_INVALID"
| "EMBEDDING_DIMENSION_UNSUPPORTED"
| "MODEL_CAPABILITY_MISMATCH"
| "MODEL_CREDENTIAL_INVALID"
| "MODEL_CREDENTIAL_VALIDATION_UNAVAILABLE"
| "MODEL_IDENTITY_MISMATCH"
| "MODEL_PREFLIGHT_CANCELED"
| "MODEL_PREFLIGHT_FAILED"
| "MODEL_PREFLIGHT_TIMEOUT"
| "MODEL_PREFLIGHT_UNAVAILABLE"
| "MODEL_SELECTION_NOT_FOUND";
export class ModelCapabilityPreflightError extends Error {
readonly code: ModelCapabilityPreflightErrorCode;
readonly retryable: boolean;
constructor(
code: ModelCapabilityPreflightErrorCode,
message: string,
options: { readonly cause?: unknown; readonly retryable?: boolean } = {},
) {
super(message, options.cause === undefined ? undefined : { cause: options.cause });
this.name = "ModelCapabilityPreflightError";
this.code = code;
this.retryable = options.retryable ?? false;
}
}
export interface ModelCapabilityPreflightInput extends ResolveModelCatalogEntryInput {
readonly signal?: AbortSignal | undefined;
}
export interface ModelCapabilityPreflight {
/**
* Captures a tenant-active catalog declaration without invoking the model. This is used for
* configured non-embedding models whose credentials were already checked by Dify.
*/
resolveConfigured?(input: ModelCapabilityPreflightInput): Promise<ModelCapabilitySnapshot>;
verify(input: ModelCapabilityPreflightInput): Promise<ModelCapabilitySnapshot>;
}
export interface ModelCapabilityPreflightOptions {
readonly catalog: ModelCapabilityCatalog;
readonly embeddingProviderFactory: (selection: KnowledgeSpaceModelSelection) => EmbeddingProvider;
readonly now?: (() => string) | undefined;
readonly reasoningProviderFactory: (
selection: KnowledgeSpaceModelSelection,
) => ReasoningModelPreflightProvider;
readonly rerankerProviderFactory: (selection: KnowledgeSpaceModelSelection) => RerankerProvider;
readonly timeoutMs?: number | undefined;
/**
* Production vector storage dialect. Embedding models are probed dynamically, then rejected
* only when their observed dimension cannot be stored by this backend. Dimensions that merely
* exceed an ANN index limit remain valid and use the exact-search fallback.
*/
readonly vectorStorageDialect?: "postgres" | "tidb" | undefined;
}
/** Structural subset implemented by the Dify-managed LLM provider without coupling API to it. */
export interface ReasoningModelPreflightProvider {
generate(input: {
readonly maxOutputTokens: number;
readonly messages: readonly { readonly content: string; readonly role: "user" }[];
readonly model: string;
readonly signal: AbortSignal;
readonly temperature: number;
readonly tenantId: string;
}): Promise<{
readonly metadata: { readonly model: string };
readonly model: string;
readonly text: string;
}>;
}
const DEFAULT_PREFLIGHT_TIMEOUT_MS = 15_000;
const PREFLIGHT_EMBEDDING_SENTINEL = "knowledge-fs model capability preflight";
/**
* Verifies that a catalog declaration is actually invokable before a profile revision can be
* persisted. Provider errors are deliberately collapsed to a stable, non-secret response.
*/
export function createModelCapabilityPreflight({
catalog,
embeddingProviderFactory,
now = () => new Date().toISOString(),
reasoningProviderFactory,
rerankerProviderFactory,
timeoutMs = DEFAULT_PREFLIGHT_TIMEOUT_MS,
vectorStorageDialect,
}: ModelCapabilityPreflightOptions): ModelCapabilityPreflight {
if (!Number.isSafeInteger(timeoutMs) || timeoutMs < 1) {
throw new Error("Model capability preflight timeoutMs must be a positive integer");
}
return {
resolveConfigured: async (input) => {
const tenantId = input.tenantId.trim();
if (!tenantId) {
throw new ModelCapabilityPreflightError(
"MODEL_CAPABILITY_MISMATCH",
"Model capability resolution requires a tenant",
);
}
const kind = ModelCapabilityKindSchema.parse(input.kind);
if (kind === "embedding") {
throw new ModelCapabilityPreflightError(
"EMBEDDING_DIMENSION_INVALID",
"Embedding capabilities require an observed vector dimension",
);
}
const selection = KnowledgeSpaceModelSelectionSchema.parse(input.selection);
const scoped = createPreflightAbortScope(input.signal, timeoutMs);
try {
return await scoped.race(
(async () => {
const catalogEntry = await resolveCatalogDeclaration({
catalog,
kind,
selection,
signal: scoped.signal,
tenantId,
});
assertPreflightActive(scoped.signal);
return capabilitySnapshot({
catalogEntry,
checkedAt: z.string().datetime({ offset: true }).parse(now()),
kind,
selection,
});
})(),
);
} catch (error) {
if (error instanceof ModelCapabilityPreflightError) throw error;
throw normalizePreflightProviderError(error);
} finally {
scoped.dispose();
}
},
verify: async (input) => {
const tenantId = input.tenantId.trim();
if (!tenantId) {
throw new ModelCapabilityPreflightError(
"MODEL_CAPABILITY_MISMATCH",
"Model capability preflight requires a tenant",
);
}
const kind = ModelCapabilityKindSchema.parse(input.kind);
const selection = KnowledgeSpaceModelSelectionSchema.parse(input.selection);
const scoped = createPreflightAbortScope(input.signal, timeoutMs);
try {
return await scoped.race(
(async () => {
assertPreflightActive(scoped.signal);
const catalogEntry = await resolveCatalogDeclaration({
catalog,
kind,
selection,
signal: scoped.signal,
tenantId,
});
if (catalog.validate) {
let valid: boolean;
try {
valid = await catalog.validate({
kind,
selection,
signal: scoped.signal,
tenantId,
});
} catch (cause) {
throw new ModelCapabilityPreflightError(
"MODEL_CREDENTIAL_VALIDATION_UNAVAILABLE",
"The selected model's credentials could not be validated",
{ cause, retryable: true },
);
}
assertPreflightActive(scoped.signal);
if (!valid) {
throw new ModelCapabilityPreflightError(
"MODEL_CREDENTIAL_INVALID",
"The selected model's credentials are not valid",
);
}
}
const observed = await invokePreflight({
embeddingProviderFactory,
kind,
reasoningProviderFactory,
rerankerProviderFactory,
selection,
signal: scoped.signal,
tenantId,
});
if (
kind === "embedding" &&
vectorStorageDialect &&
observed.dimension !== undefined &&
observed.distanceMetric !== undefined
) {
const storage = resolveVectorIndexCapability({
dialect: vectorStorageDialect,
dimension: observed.dimension,
metric: observed.distanceMetric,
});
if (storage.status === "unsupported") {
throw new ModelCapabilityPreflightError(
"EMBEDDING_DIMENSION_UNSUPPORTED",
`The embedding model dimension=${observed.dimension} exceeds ${vectorStorageDialect} vector storage capacity`,
);
}
}
assertPreflightActive(scoped.signal);
return capabilitySnapshot({
catalogEntry,
checkedAt: z.string().datetime({ offset: true }).parse(now()),
kind,
observed,
selection,
});
})(),
);
} catch (error) {
if (error instanceof ModelCapabilityPreflightError) {
throw error;
}
throw normalizePreflightProviderError(error);
} finally {
scoped.dispose();
}
},
};
}
async function resolveCatalogDeclaration({
catalog,
kind,
selection,
signal,
tenantId,
}: {
readonly catalog: ModelCapabilityCatalog;
readonly kind: ModelCapabilityKind;
readonly selection: KnowledgeSpaceModelSelection;
readonly signal: AbortSignal;
readonly tenantId: string;
}): Promise<ModelCatalogEntry> {
assertPreflightActive(signal);
let entry: ModelCatalogEntry | null;
try {
entry = await catalog.resolve({ kind, selection, signal, tenantId });
} catch (cause) {
throw new ModelCapabilityPreflightError(
"MODEL_PREFLIGHT_UNAVAILABLE",
"Model capability catalog is temporarily unavailable",
{ cause, retryable: true },
);
}
assertPreflightActive(signal);
if (!entry) {
throw new ModelCapabilityPreflightError(
"MODEL_SELECTION_NOT_FOUND",
"The selected model is not installed for this tenant",
);
}
const catalogEntry = ModelCatalogEntrySchema.parse(entry);
assertCatalogIdentity({ catalogEntry, kind, selection });
return catalogEntry;
}
function capabilitySnapshot({
catalogEntry,
checkedAt,
kind,
observed = {},
selection,
}: {
readonly catalogEntry: ModelCatalogEntry;
readonly checkedAt: string;
readonly kind: ModelCapabilityKind;
readonly observed?: {
readonly dimension?: number | undefined;
readonly distanceMetric?: "cosine" | "dot" | "l2" | undefined;
};
readonly selection: KnowledgeSpaceModelSelection;
}): ModelCapabilitySnapshot {
const capabilityMaterial = {
...(observed.dimension === undefined ? {} : { dimension: observed.dimension }),
...(observed.distanceMetric === undefined ? {} : { distanceMetric: observed.distanceMetric }),
kind,
pluginUniqueIdentifier: catalogEntry.pluginUniqueIdentifier,
...(catalogEntry.pluginVersion ? { pluginVersion: catalogEntry.pluginVersion } : {}),
schemaFingerprint: catalogEntry.schemaFingerprint,
selection,
};
return ModelCapabilitySnapshotSchema.parse({
...capabilityMaterial,
capabilityDigest: `sha256:${createHash("sha256")
.update(stableJson({ ...capabilityMaterial, capabilities: catalogEntry.capabilities }))
.digest("hex")}`,
checkedAt,
});
}
function assertCatalogIdentity({
catalogEntry,
kind,
selection,
}: {
readonly catalogEntry: ModelCatalogEntry;
readonly kind: ModelCapabilityKind;
readonly selection: KnowledgeSpaceModelSelection;
}): void {
if (
catalogEntry.pluginId !== selection.pluginId ||
catalogEntry.provider !== selection.provider ||
catalogEntry.model !== selection.model
) {
throw new ModelCapabilityPreflightError(
"MODEL_IDENTITY_MISMATCH",
"The model catalog returned a different model identity",
);
}
if (!catalogEntry.kinds.includes(kind)) {
throw new ModelCapabilityPreflightError(
"MODEL_CAPABILITY_MISMATCH",
"The selected model does not support the requested capability",
);
}
}
async function invokePreflight({
embeddingProviderFactory,
kind,
reasoningProviderFactory,
rerankerProviderFactory,
selection,
signal,
tenantId,
}: {
readonly embeddingProviderFactory: ModelCapabilityPreflightOptions["embeddingProviderFactory"];
readonly kind: ModelCapabilityKind;
readonly reasoningProviderFactory: ModelCapabilityPreflightOptions["reasoningProviderFactory"];
readonly rerankerProviderFactory: ModelCapabilityPreflightOptions["rerankerProviderFactory"];
readonly selection: KnowledgeSpaceModelSelection;
readonly signal: AbortSignal;
readonly tenantId: string;
}): Promise<{ readonly dimension?: number; readonly distanceMetric?: "cosine" | "dot" | "l2" }> {
if (kind === "embedding") {
const provider = embeddingProviderFactory(selection);
const result = await provider.embed({
inputType: "search_query",
model: selection.model,
signal,
tenantId,
texts: [PREFLIGHT_EMBEDDING_SENTINEL],
});
assertPreflightActive(signal);
assertObservedIdentity(selection.model, result.model);
const vector = result.dense[0];
if (
result.dense.length !== 1 ||
!vector ||
vector.length < 1 ||
!vector.every(Number.isFinite) ||
(result.metadata.dimension !== undefined && result.metadata.dimension !== vector.length)
) {
throw new ModelCapabilityPreflightError(
"EMBEDDING_DIMENSION_INVALID",
"The embedding model returned an invalid vector dimension",
);
}
const modelInfo = (await provider.models()).find((model) => model.id === selection.model);
assertPreflightActive(signal);
if (modelInfo?.dimension !== undefined && modelInfo.dimension !== vector.length) {
throw new ModelCapabilityPreflightError(
"EMBEDDING_DIMENSION_INVALID",
"The embedding model returned a dimension that conflicts with its capability declaration",
);
}
return { dimension: vector.length, distanceMetric: modelInfo?.distanceMetric ?? "cosine" };
}
if (kind === "rerank") {
const documents = [
{ id: "preflight-relevant", text: "knowledge retrieval" },
{ id: "preflight-control", text: "unrelated control" },
];
const topN = 2;
const result = await rerankerProviderFactory(selection).rerank({
documents,
model: selection.model,
query: "knowledge retrieval",
signal,
tenantId,
topN,
});
assertPreflightActive(signal);
assertObservedIdentity(selection.model, result.model);
assertObservedIdentity(selection.model, result.metadata?.model);
const seenDocumentIds = new Set<string>();
const seenIndices = new Set<number>();
const items = Array.isArray(result.items) ? result.items : [];
const invalidItems =
items.length < 1 ||
items.length > topN ||
items.some((item) => {
const returnedDocument = item.document;
const original = documents[item.index];
const invalid =
!Number.isInteger(item.index) ||
!original ||
!returnedDocument ||
seenIndices.has(item.index) ||
seenDocumentIds.has(returnedDocument.id) ||
returnedDocument.id !== original.id ||
returnedDocument.text !== original.text ||
!Number.isFinite(item.score) ||
item.score < 0 ||
item.score > 1;
seenIndices.add(item.index);
if (returnedDocument) {
seenDocumentIds.add(returnedDocument.id);
}
return invalid;
});
if (invalidItems) {
throw new ModelCapabilityPreflightError(
"MODEL_CAPABILITY_MISMATCH",
"The rerank model returned an invalid capability response",
);
}
return {};
}
const result = await reasoningProviderFactory(selection).generate({
maxOutputTokens: 512,
messages: [{ content: "Reply OK.", role: "user" }],
model: selection.model,
signal,
temperature: 0,
tenantId,
});
assertPreflightActive(signal);
if (typeof result.text !== "string" || !result.text.trim()) {
throw new ModelCapabilityPreflightError(
"MODEL_CAPABILITY_MISMATCH",
"The reasoning model returned an invalid capability response",
);
}
assertObservedIdentity(selection.model, result.model);
assertObservedIdentity(selection.model, result.metadata?.model);
return {};
}
function assertObservedIdentity(requested: string, observed: unknown): void {
if (typeof observed !== "string" || !observed.trim() || observed.trim() !== requested) {
throw new ModelCapabilityPreflightError(
"MODEL_IDENTITY_MISMATCH",
"The model response identity did not match the selected model",
);
}
}
function assertPreflightActive(signal: AbortSignal): void {
if (!signal.aborted) {
return;
}
throw signal.reason instanceof Error
? signal.reason
: new Error("Model capability preflight was aborted");
}
function createPreflightAbortScope(
parentSignal: AbortSignal | undefined,
timeoutMs: number,
): {
readonly dispose: () => void;
readonly race: <T>(operation: Promise<T>) => Promise<T>;
readonly signal: AbortSignal;
} {
const controller = new AbortController();
let rejectBoundary: ((reason: unknown) => void) | undefined;
let settled = false;
const boundary = new Promise<never>((_resolve, reject) => {
rejectBoundary = reject;
});
const abort = (error: ModelCapabilityPreflightError) => {
if (settled) {
return;
}
settled = true;
controller.abort(error);
rejectBoundary?.(error);
};
const abortFromParent = () =>
abort(
new ModelCapabilityPreflightError(
"MODEL_PREFLIGHT_CANCELED",
"The selected model capability preflight was canceled",
{ cause: parentSignal?.reason, retryable: true },
),
);
if (parentSignal?.aborted) {
abortFromParent();
} else {
parentSignal?.addEventListener("abort", abortFromParent, { once: true });
}
const timeout = setTimeout(() => {
abort(
new ModelCapabilityPreflightError(
"MODEL_PREFLIGHT_TIMEOUT",
"The selected model capability preflight timed out",
{ retryable: true },
),
);
}, timeoutMs);
return {
dispose: () => {
settled = true;
clearTimeout(timeout);
parentSignal?.removeEventListener("abort", abortFromParent);
},
race: <T>(operation: Promise<T>) => Promise.race([operation, boundary]),
signal: controller.signal,
};
}
function normalizePreflightProviderError(error: unknown): ModelCapabilityPreflightError {
const providerCode =
error && typeof error === "object" && "code" in error
? (error as { readonly code?: unknown }).code
: undefined;
const normalizedProviderCode =
typeof providerCode === "string" ? providerCode.trim().toLowerCase() : "";
const retryable =
error && typeof error === "object" && "retryable" in error
? (error as { readonly retryable?: unknown }).retryable === true
: true;
if (normalizedProviderCode.includes("timeout")) {
return new ModelCapabilityPreflightError(
"MODEL_PREFLIGHT_TIMEOUT",
"The selected model capability preflight timed out",
{ cause: error, retryable: true },
);
}
if (
normalizedProviderCode.includes("abort") ||
normalizedProviderCode.includes("cancel") ||
(error instanceof Error && error.name === "AbortError")
) {
return new ModelCapabilityPreflightError(
"MODEL_PREFLIGHT_CANCELED",
"The selected model capability preflight was canceled",
{ cause: error, retryable: true },
);
}
if (
normalizedProviderCode.includes("request_failed") ||
normalizedProviderCode.includes("unavailable")
) {
return new ModelCapabilityPreflightError(
"MODEL_PREFLIGHT_UNAVAILABLE",
"The selected model service is temporarily unavailable",
{ cause: error, retryable },
);
}
return new ModelCapabilityPreflightError(
"MODEL_PREFLIGHT_FAILED",
"The selected model failed its capability preflight",
{ cause: error, retryable },
);
}