diff --git a/libs/langchain-community/src/experimental/graph_transformers/llm.ts b/libs/langchain-community/src/experimental/graph_transformers/llm.ts index a46033801016..42eb8784a37f 100644 --- a/libs/langchain-community/src/experimental/graph_transformers/llm.ts +++ b/libs/langchain-community/src/experimental/graph_transformers/llm.ts @@ -197,25 +197,33 @@ function mapToBaseNode(node: any): Node { } // eslint-disable-next-line @typescript-eslint/no-explicit-any -function mapToBaseRelationship(relationship: any): Relationship { - return new Relationship({ - source: new Node({ - id: relationship.sourceNodeId, - type: relationship.sourceNodeType - ? toTitleCase(relationship.sourceNodeType) - : "", - }), - target: new Node({ - id: relationship.targetNodeId, - type: relationship.targetNodeType - ? toTitleCase(relationship.targetNodeType) - : "", - }), - type: relationship.relationshipType.replace(" ", "_").toUpperCase(), - properties: relationship.properties - ? convertPropertiesToRecord(relationship.properties) - : {}, - }); +function mapToBaseRelationship({ + fallbackRelationshipType, +}: { + fallbackRelationshipType: string | null; +}) { + return function (relationship: any): Relationship { + return new Relationship({ + source: new Node({ + id: relationship.sourceNodeId, + type: relationship.sourceNodeType + ? toTitleCase(relationship.sourceNodeType) + : "", + }), + target: new Node({ + id: relationship.targetNodeId, + type: relationship.targetNodeType + ? toTitleCase(relationship.targetNodeType) + : "", + }), + type: (relationship.relationshipType || fallbackRelationshipType) + .replace(" ", "_") + .toUpperCase(), + properties: relationship.properties + ? convertPropertiesToRecord(relationship.properties) + : {}, + }); + }; } export interface LLMGraphTransformerProps { @@ -226,6 +234,13 @@ export interface LLMGraphTransformerProps { strictMode?: boolean; nodeProperties?: string[]; relationshipProperties?: string[]; + + /** + * @description + * The LLM may rarely create relationships without a type, causing extraction to fail. + * Use this to provide a fallback relationship type in such case. + */ + fallbackRelationshipType?: string | null; } export class LLMGraphTransformer { @@ -242,6 +257,8 @@ export class LLMGraphTransformer { relationshipProperties: string[]; + fallbackRelationshipType: string | null = null; + constructor({ llm, allowedNodes = [], @@ -250,6 +267,7 @@ export class LLMGraphTransformer { strictMode = true, nodeProperties = [], relationshipProperties = [], + fallbackRelationshipType = null, }: LLMGraphTransformerProps) { if (typeof llm.withStructuredOutput !== "function") { throw new Error( @@ -262,6 +280,7 @@ export class LLMGraphTransformer { this.strictMode = strictMode; this.nodeProperties = nodeProperties; this.relationshipProperties = relationshipProperties; + this.fallbackRelationshipType = fallbackRelationshipType; // Define chain const schema = createSchema( @@ -292,7 +311,11 @@ export class LLMGraphTransformer { let relationships: Relationship[] = []; if (rawSchema?.relationships) { - relationships = rawSchema.relationships.map(mapToBaseRelationship); + relationships = rawSchema.relationships.map( + mapToBaseRelationship({ + fallbackRelationshipType: this.fallbackRelationshipType, + }) + ); } if (