feat(community): Provide fallback relationshipType in case it is not …

…present in graph_transformer (#7521) Co-authored-by: quantropi-minh <[email protected]> Co-authored-by: jacoblee93 <[email protected]>
langchain-ai · Jan 22, 2025 · 0397e93 · 0397e93
1 parent a2ebf76
commit 0397e93
Showing 1 changed file with 43 additions and 20 deletions.
diff --git 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 (