mirror of
https://github.com/langgenius/dify.git
synced 2026-09-02 22:36:15 +08:00
428 lines
13 KiB
TypeScript
428 lines
13 KiB
TypeScript
import { type KnowledgeNode, PublicationGenerationIdSchema } from "@knowledge/core";
|
|
|
|
import { mapWithConcurrency } from "./bounded-concurrency";
|
|
import { type ExtractedEntity, extractedEntitiesFromNodeMetadata } from "./entity-extraction-flow";
|
|
import { RELATION_EXTRACTION_TYPES, type RelationExtractionType } from "./extraction-types";
|
|
import { cloneJsonObject, isPlainObject } from "./json-utils";
|
|
import { type KnowledgeNodeRepository, cloneKnowledgeNode } from "./knowledge-node-repository";
|
|
|
|
export interface ExtractedRelation {
|
|
readonly confidence: number;
|
|
readonly metadata?: Readonly<Record<string, unknown>> | undefined;
|
|
readonly object: string;
|
|
readonly subject: string;
|
|
readonly type: RelationExtractionType;
|
|
}
|
|
|
|
export interface RelationExtractionProviderInput {
|
|
readonly entities: readonly ExtractedEntity[];
|
|
readonly maxRelations: number;
|
|
readonly model: string;
|
|
readonly node: KnowledgeNode;
|
|
readonly prompt: string;
|
|
readonly promptVersion: string;
|
|
readonly tenantId?: string | undefined;
|
|
}
|
|
|
|
export interface RelationExtractionProviderResult {
|
|
readonly metadata?: Readonly<Record<string, unknown>> | undefined;
|
|
readonly relations: readonly ExtractedRelation[];
|
|
}
|
|
|
|
export interface RelationExtractionProvider {
|
|
extract(input: RelationExtractionProviderInput): Promise<RelationExtractionProviderResult>;
|
|
extractBatch?(
|
|
inputs: readonly RelationExtractionProviderInput[],
|
|
): Promise<readonly RelationExtractionProviderResult[]>;
|
|
}
|
|
|
|
export class RelationExtractionBatchContractError extends Error {
|
|
override readonly name = "RelationExtractionBatchContractError";
|
|
}
|
|
|
|
export interface RelationExtractionFlowOptions {
|
|
readonly maxBatchSize: number;
|
|
readonly maxConcurrency?: number | undefined;
|
|
readonly maxRelationsPerNode?: number | undefined;
|
|
readonly model: string;
|
|
readonly nodes: KnowledgeNodeRepository;
|
|
readonly now?: () => string;
|
|
readonly promptVersion?: string | undefined;
|
|
readonly provider: RelationExtractionProvider;
|
|
readonly providerBatchSize?: number | undefined;
|
|
}
|
|
|
|
export interface ExtractKnowledgeNodeRelationsInput {
|
|
readonly knowledgeSpaceId: string;
|
|
readonly nodeIds: readonly string[];
|
|
readonly publicationGenerationId?: string | undefined;
|
|
readonly tenantId?: string | undefined;
|
|
readonly traceId?: string | undefined;
|
|
}
|
|
|
|
export interface RelationExtractionResult {
|
|
readonly extractedNodes: KnowledgeNode[];
|
|
readonly missingNodeIds: readonly string[];
|
|
}
|
|
|
|
export interface RelationExtractionFlow {
|
|
extract(input: ExtractKnowledgeNodeRelationsInput): Promise<RelationExtractionResult>;
|
|
}
|
|
|
|
export function createRelationExtractionFlow({
|
|
maxBatchSize,
|
|
maxConcurrency = 4,
|
|
maxRelationsPerNode = 100,
|
|
model,
|
|
nodes,
|
|
now = () => new Date().toISOString(),
|
|
promptVersion = "relation-extraction-v1",
|
|
provider,
|
|
providerBatchSize = 8,
|
|
}: RelationExtractionFlowOptions): RelationExtractionFlow {
|
|
if (!Number.isInteger(maxBatchSize) || maxBatchSize < 1) {
|
|
throw new Error("Relation extraction maxBatchSize must be at least 1");
|
|
}
|
|
|
|
if (!Number.isInteger(maxConcurrency) || maxConcurrency < 1) {
|
|
throw new Error("Relation extraction maxConcurrency must be at least 1");
|
|
}
|
|
|
|
if (!Number.isInteger(maxRelationsPerNode) || maxRelationsPerNode < 1) {
|
|
throw new Error("Relation extraction maxRelationsPerNode must be at least 1");
|
|
}
|
|
|
|
if (!Number.isInteger(providerBatchSize) || providerBatchSize < 1) {
|
|
throw new Error("Relation extraction providerBatchSize must be at least 1");
|
|
}
|
|
|
|
if (!model.trim()) {
|
|
throw new Error("Relation extraction model is required");
|
|
}
|
|
|
|
if (!promptVersion.trim()) {
|
|
throw new Error("Relation extraction promptVersion is required");
|
|
}
|
|
|
|
return {
|
|
extract: async ({ knowledgeSpaceId, nodeIds, publicationGenerationId, tenantId, traceId }) => {
|
|
validateRelationExtractionInput({
|
|
knowledgeSpaceId,
|
|
maxBatchSize,
|
|
nodeIds,
|
|
publicationGenerationId,
|
|
});
|
|
const uniqueNodeIds = uniqueStrings(nodeIds);
|
|
const loadedNodes = await nodes.getMany({
|
|
ids: uniqueNodeIds,
|
|
knowledgeSpaceId,
|
|
...(publicationGenerationId ? { publicationGenerationId } : {}),
|
|
});
|
|
const nodesById = new Map(loadedNodes.map((node) => [node.id, node]));
|
|
const orderedNodes = uniqueNodeIds.flatMap((id) => {
|
|
const node = nodesById.get(id);
|
|
|
|
return node ? [cloneKnowledgeNode(node)] : [];
|
|
});
|
|
const missingNodeIds = uniqueNodeIds.filter((id) => !nodesById.has(id));
|
|
|
|
if (orderedNodes.length === 0) {
|
|
return {
|
|
extractedNodes: [],
|
|
missingNodeIds,
|
|
};
|
|
}
|
|
|
|
const inputsWithNodes = orderedNodes.map((node) => {
|
|
const entities = extractedEntitiesFromNodeMetadata(node);
|
|
return {
|
|
input: {
|
|
entities,
|
|
maxRelations: maxRelationsPerNode,
|
|
model,
|
|
node: cloneKnowledgeNode(node),
|
|
prompt: relationExtractionPrompt(node, entities),
|
|
promptVersion,
|
|
...(tenantId ? { tenantId } : {}),
|
|
},
|
|
node,
|
|
};
|
|
});
|
|
const callable = inputsWithNodes.filter(({ input }) => input.entities.length > 0);
|
|
const extractBatch = provider.extractBatch;
|
|
const providerResults = extractBatch
|
|
? await extractRelationsInBatches({
|
|
extractBatch,
|
|
inputs: callable.map(({ input }) => input),
|
|
maxConcurrency,
|
|
provider,
|
|
providerBatchSize,
|
|
})
|
|
: await mapWithConcurrency(
|
|
callable.map(({ input }) => input),
|
|
maxConcurrency,
|
|
(providerInput) => provider.extract(providerInput),
|
|
);
|
|
const resultsByNodeId = new Map(
|
|
callable.map(({ node }, index) => {
|
|
const result = providerResults[index];
|
|
if (!result) {
|
|
throw new RelationExtractionBatchContractError(
|
|
"Relation extraction provider returned an incomplete result set",
|
|
);
|
|
}
|
|
return [node.id, result] as const;
|
|
}),
|
|
);
|
|
const generated = inputsWithNodes.map(({ node }) => {
|
|
const result = resultsByNodeId.get(node.id) ?? { relations: [] };
|
|
const relations = validateExtractedRelations(result.relations, maxRelationsPerNode);
|
|
|
|
return {
|
|
id: node.id,
|
|
metadata: relationExtractionMetadata({
|
|
metadata: result.metadata,
|
|
model,
|
|
node,
|
|
now,
|
|
promptVersion,
|
|
relations,
|
|
traceId,
|
|
}),
|
|
};
|
|
});
|
|
const extractedNodes = await nodes.updateMetadataMany({
|
|
knowledgeSpaceId,
|
|
patches: generated,
|
|
...(publicationGenerationId ? { publicationGenerationId } : {}),
|
|
});
|
|
|
|
return {
|
|
extractedNodes: extractedNodes.map(cloneKnowledgeNode),
|
|
missingNodeIds,
|
|
};
|
|
},
|
|
};
|
|
}
|
|
|
|
async function extractRelationsInBatches({
|
|
extractBatch,
|
|
inputs,
|
|
maxConcurrency,
|
|
provider,
|
|
providerBatchSize,
|
|
}: {
|
|
readonly extractBatch: NonNullable<RelationExtractionProvider["extractBatch"]>;
|
|
readonly inputs: readonly RelationExtractionProviderInput[];
|
|
readonly maxConcurrency: number;
|
|
readonly provider: RelationExtractionProvider;
|
|
readonly providerBatchSize: number;
|
|
}): Promise<RelationExtractionProviderResult[]> {
|
|
const batches = chunk(inputs, providerBatchSize);
|
|
const results = await mapWithConcurrency(batches, maxConcurrency, async (batch) => {
|
|
try {
|
|
const extracted = await extractBatch(batch);
|
|
if (extracted.length !== batch.length) {
|
|
throw new RelationExtractionBatchContractError(
|
|
`Relation extraction batch returned ${extracted.length} results for ${batch.length} inputs`,
|
|
);
|
|
}
|
|
return [...extracted];
|
|
} catch (error) {
|
|
if (!(error instanceof RelationExtractionBatchContractError)) throw error;
|
|
return mapWithConcurrency(batch, maxConcurrency, (input) => provider.extract(input));
|
|
}
|
|
});
|
|
return results.flat();
|
|
}
|
|
|
|
function chunk<T>(items: readonly T[], size: number): T[][] {
|
|
const chunks: T[][] = [];
|
|
for (let index = 0; index < items.length; index += size) {
|
|
chunks.push(items.slice(index, index + size));
|
|
}
|
|
return chunks;
|
|
}
|
|
|
|
function validateRelationExtractionInput({
|
|
knowledgeSpaceId,
|
|
maxBatchSize,
|
|
nodeIds,
|
|
publicationGenerationId,
|
|
}: {
|
|
readonly knowledgeSpaceId: string;
|
|
readonly maxBatchSize: number;
|
|
readonly nodeIds: readonly string[];
|
|
readonly publicationGenerationId?: string | undefined;
|
|
}) {
|
|
if (!knowledgeSpaceId.trim()) {
|
|
throw new Error("Relation extraction knowledgeSpaceId is required");
|
|
}
|
|
|
|
if (nodeIds.length < 1) {
|
|
throw new Error("Relation extraction nodeIds must contain at least 1 node id");
|
|
}
|
|
|
|
if (nodeIds.length > maxBatchSize) {
|
|
throw new Error(`Relation extraction nodeIds exceeds maxBatchSize=${maxBatchSize}`);
|
|
}
|
|
|
|
if (publicationGenerationId !== undefined) {
|
|
PublicationGenerationIdSchema.parse(publicationGenerationId);
|
|
}
|
|
|
|
for (const nodeId of nodeIds) {
|
|
if (!nodeId.trim()) {
|
|
throw new Error("Relation extraction nodeIds must be non-empty strings");
|
|
}
|
|
}
|
|
}
|
|
|
|
function validateExtractedRelations(
|
|
relations: readonly ExtractedRelation[],
|
|
maxRelationsPerNode: number,
|
|
): ExtractedRelation[] {
|
|
const validated = relations.map((relation, index) => {
|
|
if (!RELATION_EXTRACTION_TYPES.has(relation.type)) {
|
|
throw new Error("Relation extraction relation type is unsupported");
|
|
}
|
|
|
|
if (!relation.subject.trim()) {
|
|
throw new Error("Relation extraction relation subject is required");
|
|
}
|
|
|
|
if (!relation.object.trim()) {
|
|
throw new Error("Relation extraction relation object is required");
|
|
}
|
|
|
|
if (
|
|
!Number.isFinite(relation.confidence) ||
|
|
relation.confidence < 0 ||
|
|
relation.confidence > 1
|
|
) {
|
|
throw new Error("Relation extraction relation confidence must be between 0 and 1");
|
|
}
|
|
|
|
return {
|
|
index,
|
|
relation: {
|
|
confidence: relation.confidence,
|
|
...(relation.metadata ? { metadata: cloneJsonObject(relation.metadata) } : {}),
|
|
object: relation.object.trim(),
|
|
subject: relation.subject.trim(),
|
|
type: relation.type,
|
|
},
|
|
};
|
|
});
|
|
|
|
return validated
|
|
.sort(
|
|
(left, right) =>
|
|
right.relation.confidence - left.relation.confidence || left.index - right.index,
|
|
)
|
|
.slice(0, maxRelationsPerNode)
|
|
.map(({ relation }) => relation);
|
|
}
|
|
|
|
function relationExtractionMetadata({
|
|
metadata,
|
|
model,
|
|
node,
|
|
now,
|
|
promptVersion,
|
|
relations,
|
|
traceId,
|
|
}: {
|
|
readonly metadata?: Readonly<Record<string, unknown>> | undefined;
|
|
readonly model: string;
|
|
readonly node: KnowledgeNode;
|
|
readonly now: () => string;
|
|
readonly promptVersion: string;
|
|
readonly relations: readonly ExtractedRelation[];
|
|
readonly traceId?: string | undefined;
|
|
}): Record<string, unknown> {
|
|
return {
|
|
...cloneJsonObject(node.metadata),
|
|
extractedRelations: relations.map((relation) => ({
|
|
confidence: relation.confidence,
|
|
...(relation.metadata ? { metadata: cloneJsonObject(relation.metadata) } : {}),
|
|
object: relation.object,
|
|
subject: relation.subject,
|
|
type: relation.type,
|
|
})),
|
|
relationExtraction: {
|
|
...cloneJsonObject(metadata ?? {}),
|
|
extractedAt: now(),
|
|
model,
|
|
promptVersion,
|
|
relationCount: relations.length,
|
|
...(traceId ? { traceId } : {}),
|
|
},
|
|
};
|
|
}
|
|
|
|
export function extractedRelationsFromNodeMetadata(node: KnowledgeNode): ExtractedRelation[] {
|
|
const relations = node.metadata.extractedRelations;
|
|
|
|
if (!Array.isArray(relations)) {
|
|
return [];
|
|
}
|
|
|
|
return relations.flatMap((relation) => {
|
|
if (!isPlainObject(relation)) {
|
|
return [];
|
|
}
|
|
|
|
if (
|
|
typeof relation.subject !== "string" ||
|
|
!relation.subject.trim() ||
|
|
typeof relation.object !== "string" ||
|
|
!relation.object.trim() ||
|
|
typeof relation.type !== "string" ||
|
|
!RELATION_EXTRACTION_TYPES.has(relation.type as RelationExtractionType) ||
|
|
typeof relation.confidence !== "number" ||
|
|
!Number.isFinite(relation.confidence) ||
|
|
relation.confidence < 0 ||
|
|
relation.confidence > 1
|
|
) {
|
|
return [];
|
|
}
|
|
|
|
return [
|
|
{
|
|
confidence: relation.confidence,
|
|
...(isPlainObject(relation.metadata)
|
|
? { metadata: cloneJsonObject(relation.metadata) }
|
|
: {}),
|
|
object: relation.object.trim(),
|
|
subject: relation.subject.trim(),
|
|
type: relation.type as RelationExtractionType,
|
|
},
|
|
];
|
|
});
|
|
}
|
|
|
|
function relationExtractionPrompt(
|
|
node: KnowledgeNode,
|
|
entities: readonly ExtractedEntity[],
|
|
): string {
|
|
const sectionPath = node.sourceLocation.sectionPath.join(" > ") || "Unknown section";
|
|
const entityList =
|
|
entities.length === 0
|
|
? "No pre-extracted entities."
|
|
: entities.map((entity) => `${entity.type}:${entity.text}`).join(", ");
|
|
|
|
return [
|
|
"Extract typed relations: mentions, defines, references, depends_on, supersedes, and contradicts.",
|
|
"Use the existing entity context when possible and return confidence scores.",
|
|
`Kind: ${node.kind}`,
|
|
`Section: ${sectionPath}`,
|
|
`Entities: ${entityList}`,
|
|
`Text: ${node.text}`,
|
|
].join("\n");
|
|
}
|
|
|
|
function uniqueStrings(values: readonly string[]): string[] {
|
|
return [...new Set(values)];
|
|
}
|