dify/knowledge-fs/apps/api/src/retriever-options.ts

701 lines
26 KiB
TypeScript

import {
type BasicHybridRetriever,
type DocumentOutlineRepository,
type GraphIndexRepository,
type HybridRetrievalItem,
type HybridRetrievalRepository,
type ImageBytesVisualEmbeddingProvider,
type PageIndexFindabilityRepository,
type PageIndexLayeredTreeSearch,
type PageIndexSemanticTreeSearch,
type PageIndexWholeTreeSelector,
type ProjectionSetPublicationMemberRepository,
type PublishedGraphIndexRepository,
type PublishedPageIndexRepository,
QUERY_IMAGE_VISUAL_LEG_UNAVAILABLE,
type ResearchEvidenceReasoning,
type ResearchQueryVectorizer,
type RetrievalCandidate,
type RetrievalOperationalMetrics,
type RetrievalPlanner,
type TidbFtsPostingReadinessGate,
createBasicHybridRetriever,
createDocumentOutlineRetrievalPath,
createFinalRerankRetrieval,
createGraphExpandedRetrievalPath,
createImageOcrRetrievalPath,
createPublishedPageIndexRetrievalPath,
createRequiredDeepGraphCapabilityGuard,
createResearchEvidenceRetrieval,
createResearchOutlineEvidenceRetrieval,
createTableSpecificRetrievalPath,
filterRetrievalCandidatesByMetadata,
filterRetrievalCandidatesByPermission,
filterRetrievalCandidatesByProjectionSet,
fuseRetrievalCandidates,
normalizeRetrievalMetadataFilters,
normalizeRetrievalPermissionScope,
recordRetrievalOperationalMetric,
} from "@knowledge/api";
import type { EmbeddingProvider } from "@knowledge/embeddings";
import {
type ApiGraphExpansionOptions,
DEFAULT_GRAPH_EXPANSION_OPTIONS,
} from "./graph-expansion-options";
import type { ApiRerankerOptions } from "./reranker-options";
export interface ApiRetrieverOptions {
/** Whether a dense embedding provider is configured; gates the dense leg. */
readonly embeddingEnabled: boolean;
/** Fail-closed latch for TiDB lexical postings used by every online retrieval mode. */
readonly ftsReadiness?: TidbFtsPostingReadinessGate | undefined;
/**
* Knowledge graph. When provided, wraps the stack with graph-expanded
* retrieval (GraphRAG-style: seed entities from base hits -> traverse ->
* re-retrieve boosted), which `shouldRunModeExtension` activates for `deep`
* mode only. Omit to disable graph expansion entirely.
*/
readonly graph?: GraphIndexRepository | undefined;
/**
* Graph-expansion tuning knobs (env-derived via `createApiGraphExpansionOptions`).
* Falls back to `DEFAULT_GRAPH_EXPANSION_OPTIONS` when omitted. Ignored without `graph`.
*/
readonly graphExpansion?: ApiGraphExpansionOptions | undefined;
/** Aggregation-only retrieval result telemetry. */
readonly metrics?: RetrievalOperationalMetrics | undefined;
/** Compatibility outline path used when the strict published PageIndex repository is absent. */
readonly outlines?: DocumentOutlineRepository | undefined;
/** Strict publication-member scoped PageIndex capability used by production Research. */
readonly pageIndex?: PublishedPageIndexRepository | undefined;
readonly pageIndexFindability?: Pick<PageIndexFindabilityRepository, "getManyRoutes"> | undefined;
/** V2-only profile-scoped LLM scorer retained for replaying old Research checkpoints. */
readonly pageIndexSemanticTreeSearch?: PageIndexSemanticTreeSearch | undefined;
/** V2-only sibling-level traversal retained for replaying old Research checkpoints. */
readonly pageIndexLayeredTreeSearch?: PageIndexLayeredTreeSearch | undefined;
/** Compatibility selector used only if a lower-level caller omits layered navigation. */
readonly pageIndexWholeTreeSelector?: PageIndexWholeTreeSelector | undefined;
/**
* Mode-aware planner. Optional for compatibility with the underlying
* `createBasicHybridRetriever`; when omitted the basic retriever falls back to
* `defaultRetrievalPlan` ("fast" for every request). Pass one so fast/deep/
* research resolve distinct fan-out — `index.ts` always does.
*/
readonly planner?: RetrievalPlanner | undefined;
/** Bounded defense-in-depth membership check for non-database repositories. */
readonly publishedProjectionMembership?:
| Pick<ProjectionSetPublicationMemberRepository, "filterComponentKeys">
| undefined;
/** Strict graph view bound to the immutable query-start publication snapshot. */
readonly publishedGraph?: PublishedGraphIndexRepository | undefined;
readonly repository: HybridRetrievalRepository;
/** Online Research V3. Omission retains the V2 path for lower-level compatibility tests. */
readonly researchEvidence?:
| {
readonly queryVectorizer: ResearchQueryVectorizer;
readonly reasoning: ResearchEvidenceReasoning;
}
| undefined;
readonly rerankerOptions?: ApiRerankerOptions | undefined;
/** Require a fixed published projection snapshot before any retrieval leg runs. */
readonly strictPublishedReads?: boolean | undefined;
/** Text-to-visual query embedding for the separate visual_vector search leg. */
readonly visualQuery?:
| {
readonly model: string;
readonly mode: "fallback" | "primary";
readonly provider: EmbeddingProvider;
}
| undefined;
/** Query-image embeddings in the same visual vector space; separately feature-gated. */
readonly imageQuery?:
| {
readonly model: string;
readonly mode: "fallback" | "primary";
readonly provider: ImageBytesVisualEmbeddingProvider;
}
| undefined;
}
/**
* Stable fail-closed signal when the text embedding capability is unavailable.
* Fast/Deep use dense hybrid recall, while Research uses the same immutable
* dense projections as its semantic Value Search leg.
*/
export class HybridEmbeddingCapabilityUnavailableError extends Error {
constructor() {
super("Fast, Deep, and Research retrieval require the configured text embedding capability");
this.name = "HybridEmbeddingCapabilityUnavailableError";
}
}
/**
* Builds the wired retrieval stack: final-rerank -> Research Evidence V3 ->
* graph-expansion -> deterministic published-outline expansion -> image-ocr ->
* table -> visual-dense + text-hybrid. Research V3 combines dense, FTS, outline,
* and optionally Graph candidates across the knowledge space, applies weighted
* RRF, then uses the profile reranker for comparable final scores. The reasoning
* model is called only for bounded query planning and one evidence-set judgement;
* retained V2 checkpoints alone use PageIndex LLM tree traversal.
*
* The `planner` is threaded into the basic hybrid retriever so the requested
* mode actually changes recall depth / fusion width / rerank gating. Without it
* the basic retriever falls back to `defaultRetrievalPlan` ("fast" for every
* request), collapsing fast/deep/research into one behaviour.
*/
export function createApiRetriever({
embeddingEnabled,
ftsReadiness,
graph,
graphExpansion,
metrics,
imageQuery,
outlines,
pageIndex,
pageIndexFindability,
pageIndexSemanticTreeSearch,
pageIndexLayeredTreeSearch,
pageIndexWholeTreeSelector,
planner,
publishedGraph,
publishedProjectionMembership,
repository,
researchEvidence,
rerankerOptions,
strictPublishedReads = false,
visualQuery,
}: ApiRetrieverOptions): BasicHybridRetriever {
if (strictPublishedReads && !pageIndex) {
throw new Error(
"Strict published retrieval requires the independent published PageIndex repository",
);
}
const basicRetriever = createBasicHybridRetriever({
planner,
...(publishedProjectionMembership ? { publishedProjectionMembership } : {}),
repository: {
...(repository.publishedMembershipEnforced ? { publishedMembershipEnforced: true } : {}),
searchDense: embeddingEnabled ? (input) => repository.searchDense(input) : async () => [],
searchFts: (input) => repository.searchFts(input),
},
strictPublishedReads,
});
const searchVisualDense = repository.searchVisualDense;
const visualAwareRetriever =
(visualQuery || imageQuery) && searchVisualDense
? createVisualDenseRetrievalPath({
...(imageQuery ? { imageQuery } : {}),
planner,
publishedMembershipEnforced: repository.publishedMembershipEnforced === true,
...(publishedProjectionMembership ? { publishedProjectionMembership } : {}),
retriever: basicRetriever,
searchVisualDense,
strictPublishedReads,
...(visualQuery ? { visualQuery } : {}),
})
: createUnavailableQueryImageVisualRetrievalPath(basicRetriever);
const multimodalStack = createImageOcrRetrievalPath({
imageBoost: 0.2,
maxImageCandidates: 5,
maxImageTopK: 10,
retriever: createTableSpecificRetrievalPath({
maxTableCandidates: 5,
maxTableTopK: 10,
retriever: visualAwareRetriever,
tableBoost: 0.25,
}),
});
let stack = multimodalStack;
let legacyResearchStack: BasicHybridRetriever | undefined;
if (pageIndex) {
if (!pageIndexSemanticTreeSearch) {
throw new Error("Published PageIndex retrieval requires semantic LLM tree search");
}
if (!pageIndexWholeTreeSelector) {
throw new Error("Published PageIndex retrieval requires its compatibility tree selector");
}
const pageIndexPlanner = planner;
if (!pageIndexPlanner) {
throw new Error("Published PageIndex retrieval requires a mode-aware planner");
}
// Keep V2 assembled behind a compatibility boundary so retained durable checkpoints remain
// replayable. Fresh V3 requests never enter this per-document LLM traversal.
legacyResearchStack = createPublishedPageIndexRetrievalPath({
...(pageIndexFindability ? { findability: pageIndexFindability } : {}),
...(pageIndexLayeredTreeSearch ? { layeredTreeSearch: pageIndexLayeredTreeSearch } : {}),
// Research's planner already caps semantic recall at RETRIEVAL_MAX_TOP_K.
maxSemanticCandidates: 100,
maxSemanticCandidatesPerCall: 5,
pageIndex,
planner: pageIndexPlanner,
retriever: multimodalStack,
semanticTreeSearch: pageIndexSemanticTreeSearch,
valueSearch: repository,
wholeTreeSelector: pageIndexWholeTreeSelector,
});
stack = researchEvidence
? createResearchOutlineEvidenceRetrieval({
pageIndex,
retriever: multimodalStack,
})
: legacyResearchStack;
} else if (outlines) {
stack = createDocumentOutlineRetrievalPath({
// This bounds outline I/O, not the final Top K. Research first keeps a
// wider PageIndex candidate pool, then truncates after outline scoring.
maxOutlinesPerQuery: 50,
outlines,
planner,
retriever: multimodalStack,
});
}
const publishedGraphAvailable = graph !== undefined && publishedGraph !== undefined;
const extendedStack =
graph && (!strictPublishedReads || publishedGraph)
? createGraphExpandedRetrievalPath({
...(graphExpansion ?? DEFAULT_GRAPH_EXPANSION_OPTIONS),
graph,
...(publishedGraph ? { publishedGraph } : {}),
retriever: stack,
strictPublishedReads,
})
: createRequiredDeepGraphCapabilityGuard({
available: !strictPublishedReads || publishedGraphAvailable,
retriever: stack,
});
const legacyDefaultConfigured =
rerankerOptions !== undefined && rerankerOptions.legacyDefaultConfigured !== false;
// Always keep the final-rerank gate in the stack. An omitted/disabled
// deployment capability must be observable as an error when a Fast/Deep
// knowledge-space profile explicitly enables reranking; otherwise the query
// would silently degrade to an un-reranked result. Legacy requests still use
// the deployment default when one exists, and still skip reranking when it
// does not.
const researchRetriever = researchEvidence
? createResearchEvidenceRetrieval({
...(legacyResearchStack ? { legacyResearchRetriever: legacyResearchStack } : {}),
planner,
queryVectorizer: researchEvidence.queryVectorizer,
reasoning: researchEvidence.reasoning,
rerankerFactory: (selection) => {
if (!rerankerOptions?.providerFactory) {
throw new Error("Research retrieval requires the reranker provider factory");
}
return rerankerOptions.providerFactory(selection);
},
retriever: extendedStack,
})
: extendedStack;
const finalRetriever = createFinalRerankRetrieval({
planner,
...(rerankerOptions?.providerFactory
? { rerankerFactory: rerankerOptions.providerFactory }
: {}),
...(rerankerOptions && legacyDefaultConfigured
? {
reranker: rerankerOptions.provider,
rerankerModel: rerankerOptions.model,
}
: {}),
retriever: researchRetriever,
});
return {
retrieve: async (input) => {
const resolvedMode =
planner?.plan({
hasQueryImages: (input.queryImages?.length ?? 0) > 0,
mode: input.mode,
query: input.query,
topK: input.topK,
traceId: input.traceId,
}).resolvedMode ?? (input.mode === "research" ? "research" : "fast");
if (!embeddingEnabled && input.query.trim()) {
throw new HybridEmbeddingCapabilityUnavailableError();
}
if (ftsReadiness && input.query.trim()) {
if (!input.tenantId) {
throw new Error("TiDB FTS readiness requires a tenant-scoped retrieval input");
}
await ftsReadiness.assertReady({
knowledgeSpaceId: input.knowledgeSpaceId,
tenantId: input.tenantId,
});
}
const result = await finalRetriever.retrieve(input);
recordRetrievalOperationalMetric(metrics, input, result);
return result;
},
};
}
function createUnavailableQueryImageVisualRetrievalPath(
retriever: BasicHybridRetriever,
): BasicHybridRetriever {
return {
retrieve: async (input) => {
const result = await retriever.retrieve(input);
if ((input.queryImages?.length ?? 0) === 0) return result;
return {
...result,
metrics: result.metrics
? {
...result.metrics,
degradationFlags: [
...(result.metrics.degradationFlags ?? []),
QUERY_IMAGE_VISUAL_LEG_UNAVAILABLE,
],
}
: undefined,
};
},
};
}
function createVisualDenseRetrievalPath({
imageQuery,
planner,
publishedMembershipEnforced,
publishedProjectionMembership,
retriever,
searchVisualDense,
strictPublishedReads,
visualQuery,
}: {
readonly imageQuery?: NonNullable<ApiRetrieverOptions["imageQuery"]> | undefined;
readonly planner?: RetrievalPlanner | undefined;
readonly publishedMembershipEnforced: boolean;
readonly publishedProjectionMembership?:
| Pick<ProjectionSetPublicationMemberRepository, "filterComponentKeys">
| undefined;
readonly retriever: BasicHybridRetriever;
readonly searchVisualDense: NonNullable<HybridRetrievalRepository["searchVisualDense"]>;
readonly strictPublishedReads: boolean;
readonly visualQuery?: NonNullable<ApiRetrieverOptions["visualQuery"]> | undefined;
}): BasicHybridRetriever {
return {
retrieve: async (input) => {
const snapshot = input.projectionSnapshot;
if (strictPublishedReads && !snapshot) {
throw new Error("Hybrid retrieval requires a published projection snapshot");
}
if (snapshot && snapshot.knowledgeSpaceId !== input.knowledgeSpaceId) {
throw new Error(
"Published projection snapshot knowledgeSpaceId does not match retrieval input",
);
}
if (snapshot && input.tenantId !== undefined && snapshot.tenantId !== input.tenantId) {
throw new Error("Published projection snapshot tenantId does not match retrieval input");
}
if (snapshot && !publishedMembershipEnforced && !publishedProjectionMembership) {
throw new Error(
"Visual retrieval published snapshot requires authoritative repository filtering or a membership checker",
);
}
const plan = planner?.plan({
hasQueryImages: (input.queryImages?.length ?? 0) > 0,
mode: input.mode,
query: input.query,
topK: input.topK,
traceId: input.traceId,
});
const searchVector = async (
queryVector: readonly number[],
resolvedModel: string,
): Promise<RetrievalCandidate[]> => {
if (queryVector.length === 0 || !queryVector.every((value) => Number.isFinite(value))) {
throw new Error("Visual query embedding provider returned an invalid query vector");
}
if (!resolvedModel.trim()) {
throw new Error("Visual query embedding provider returned an empty model");
}
const candidates = await searchVisualDense({
denseProjectionModel: resolvedModel,
filters: input.filters,
knowledgeSpaceId: input.knowledgeSpaceId,
permissionScope: input.permissionScope,
projectionSetCandidateFingerprint: input.projectionSetCandidateFingerprint,
projectionSetFingerprint: input.projectionSetFingerprint,
...(snapshot ? { projectionSetPublicationId: snapshot.publicationId } : {}),
projectionSetReadMode: input.projectionSetReadMode,
queryVector,
...(snapshot
? { tenantId: snapshot.tenantId }
: input.tenantId
? { tenantId: input.tenantId }
: {}),
topK: plan?.denseTopK ?? input.topK,
});
const metadataFiltered = filterRetrievalCandidatesByMetadata(
candidates,
normalizeRetrievalMetadataFilters(input.filters),
);
const permissionFiltered = filterRetrievalCandidatesByPermission(
metadataFiltered,
normalizeRetrievalPermissionScope(input.permissionScope),
);
const projectionFiltered = snapshot
? permissionFiltered
: filterRetrievalCandidatesByProjectionSet(permissionFiltered, {
candidateFingerprint: input.projectionSetCandidateFingerprint,
mode: input.projectionSetReadMode,
publishedFingerprint: input.projectionSetFingerprint,
});
if (!snapshot || !publishedProjectionMembership) {
return projectionFiltered;
}
const allowed = new Set(
await publishedProjectionMembership.filterComponentKeys({
componentKeys: [
...new Set(projectionFiltered.map((candidate) => candidate.projectionId)),
],
componentType: "index-projection",
knowledgeSpaceId: snapshot.knowledgeSpaceId,
publicationId: snapshot.publicationId,
tenantId: snapshot.tenantId,
}),
);
return projectionFiltered.filter((candidate) => allowed.has(candidate.projectionId));
};
const retrieveVisual = async () => {
try {
if ((input.queryImages?.length ?? 0) > 0) {
if (!imageQuery) {
return {
candidateLists: [] as RetrievalCandidate[][],
degradationFlag: QUERY_IMAGE_VISUAL_LEG_UNAVAILABLE,
ok: false as const,
};
}
const images = input.queryImages ?? [];
const embedding = await imageQuery.provider.embedImages({
images: images.map((image) => ({
assetRef: { uploadFileId: image.uploadFileId },
body: image.body,
contentType: image.mimeType,
documentAssetId: image.uploadFileId,
metadata: { queryImage: true, sha256: image.sha256 },
modality: "image",
nodeId: image.uploadFileId,
objectKey: image.uploadFileId,
sourceText: "",
})),
inputType: "query",
model: imageQuery.model,
...(snapshot
? { tenantId: snapshot.tenantId }
: input.tenantId
? { tenantId: input.tenantId }
: {}),
});
if (embedding.dense.length !== images.length) {
throw new Error(
`Visual query embedding provider returned ${embedding.dense.length} vectors for ${images.length} images`,
);
}
const resolvedModel = embedding.model.trim();
const candidateLists = await Promise.all(
embedding.dense.map((vector) => searchVector(vector, resolvedModel)),
);
return { candidateLists, ok: true as const };
}
if (!visualQuery || !input.query.trim()) {
return { candidateLists: [] as RetrievalCandidate[][], ok: true as const };
}
const embedding = await visualQuery.provider.embed({
inputType: "search_query",
model: visualQuery.model,
texts: [input.query],
...(snapshot
? { tenantId: snapshot.tenantId }
: input.tenantId
? { tenantId: input.tenantId }
: {}),
});
if (embedding.dense.length !== 1) {
throw new Error(
`Visual query embedding provider returned ${embedding.dense.length} vectors for 1 query`,
);
}
const queryVector = embedding.dense[0];
if (!queryVector) {
throw new Error("Visual query embedding provider returned no query vector");
}
if (
embedding.metadata.dimension !== undefined &&
embedding.metadata.dimension !== queryVector.length
) {
throw new Error(
`Visual query embedding provider reported dimension=${embedding.metadata.dimension}; query vector has dimension=${queryVector.length}`,
);
}
return {
candidateLists: [await searchVector(queryVector, embedding.model)],
ok: true as const,
};
} catch {
return {
candidateLists: [] as RetrievalCandidate[][],
degradationFlag:
(input.queryImages?.length ?? 0) > 0
? QUERY_IMAGE_VISUAL_LEG_UNAVAILABLE
: "visual-dense-failed:skipped",
ok: false as const,
};
}
};
const basePromise = retriever.retrieve(input);
const visualMode =
(input.queryImages?.length ?? 0) > 0 ? imageQuery?.mode : visualQuery?.mode;
const [baseResult, visualResult] =
visualMode === "fallback" && (input.queryImages?.length ?? 0) === 0
? await (async () => {
const base = await basePromise;
return base.items.length > 0
? [
base,
{
candidateLists: [] as RetrievalCandidate[][],
ok: true as const,
},
]
: [base, await retrieveVisual()];
})()
: await Promise.all([basePromise, retrieveVisual()]);
if (
visualMode === "fallback" &&
(input.queryImages?.length ?? 0) === 0 &&
baseResult.items.length > 0
) {
return baseResult;
}
if (!visualResult.ok) {
return baseResult.metrics
? {
...baseResult,
metrics: {
...baseResult.metrics,
degradationFlags: [
...(baseResult.metrics.degradationFlags ?? []),
visualResult.degradationFlag,
],
},
}
: baseResult;
}
return {
...baseResult,
items: mergeVisualDenseItems({
baseItems: baseResult.items,
limit: input.limit,
visualCandidateLists: visualResult.candidateLists,
visualWeight: visualMode === "primary" ? 1 : 0.5,
}),
metrics: baseResult.metrics
? {
...baseResult.metrics,
visualEmbeddingCandidates: visualResult.candidateLists.reduce(
(total, candidates) => total + candidates.length,
0,
),
}
: undefined,
};
},
};
}
function mergeVisualDenseItems({
baseItems,
limit,
visualCandidateLists,
visualWeight,
}: {
readonly baseItems: readonly HybridRetrievalItem[];
readonly limit: number;
readonly visualCandidateLists: readonly (readonly RetrievalCandidate[])[];
readonly visualWeight: number;
}): HybridRetrievalItem[] {
const byNodeId = new Map<string, HybridRetrievalItem>();
for (const item of baseItems) {
byNodeId.set(item.nodeId, cloneHybridItem(item));
}
const nonEmptyLists = visualCandidateLists.filter((candidates) => candidates.length > 0);
const perImageWeight =
visualCandidateLists.length > 0 ? visualWeight / visualCandidateLists.length : 0;
for (const visualCandidates of nonEmptyLists) {
const normalizedVisualItems = fuseRetrievalCandidates({
dense: visualCandidates,
fts: [],
limit: visualCandidates.length,
});
for (const visualItem of normalizedVisualItems) {
const contribution = visualItem.score * perImageWeight;
const existing = byNodeId.get(visualItem.nodeId);
if (existing) {
byNodeId.set(visualItem.nodeId, {
...existing,
metadata: { ...visualItem.metadata, ...existing.metadata },
projectionIds: uniqueStrings([...existing.projectionIds, ...visualItem.projectionIds]),
score: existing.score + contribution,
sources: uniqueStrings([...existing.sources, "visual"]) as HybridRetrievalItem["sources"],
});
continue;
}
byNodeId.set(visualItem.nodeId, {
citation: {
...visualItem.citation,
sectionPath: [...visualItem.citation.sectionPath],
},
metadata: { ...visualItem.metadata },
nodeId: visualItem.nodeId,
permissionScope: visualItem.permissionScope ? [...visualItem.permissionScope] : undefined,
projectionIds: [...visualItem.projectionIds],
score: contribution,
sources: ["visual"],
});
}
}
return [...byNodeId.values()]
.sort(
(first, second) => second.score - first.score || first.nodeId.localeCompare(second.nodeId),
)
.slice(0, limit);
}
function cloneHybridItem(item: HybridRetrievalItem): HybridRetrievalItem {
return {
...item,
citation: { ...item.citation, sectionPath: [...item.citation.sectionPath] },
metadata: { ...item.metadata },
permissionScope: item.permissionScope ? [...item.permissionScope] : undefined,
projectionIds: [...item.projectionIds],
sources: [...item.sources],
};
}
function uniqueStrings(values: readonly string[]): string[] {
return [...new Set(values)];
}