community[patch]: Fix strict mode comparison and formatting for llm graph transformer (#4988)

* Fix strict mode comparison and formatting for llm graph transformer

* Map only once, extend serializable.

* formatting

* lc namespace

* formattting

---------

Co-authored-by: Anej Gorkič <[email protected]>

Author: tomasonjo

libs/langchain-community/src/experimental/graph_transformers/llm.int.test.ts

@@ -45,13 +45,52 @@ test("convertToGraphDocuments with allowed", async () => {
   expect(result).toEqual([
     new GraphDocument({
       nodes: [
-        new Node({ id: "Elon Musk", type: "PERSON" }),
-        new Node({ id: "OpenAI", type: "ORGANIZATION" }),
+        new Node({ id: "Elon Musk", type: "Person" }),
+        new Node({ id: "OpenAI", type: "Organization" }),
       ],
       relationships: [
         new Relationship({
-          source: new Node({ id: "Elon Musk", type: "PERSON" }),
-          target: new Node({ id: "OpenAI", type: "ORGANIZATION" }),
+          source: new Node({ id: "Elon Musk", type: "Person" }),
+          target: new Node({ id: "OpenAI", type: "Organization" }),
+          type: "SUES",
+        }),
+      ],
+      source: new Document({
+        pageContent: "Elon Musk is suing OpenAI",
+        metadata: {},
+      }),
+    }),
+  ]);
+});
+
+test("convertToGraphDocuments with allowed lowercased", async () => {
+  const model = new ChatOpenAI({
+    temperature: 0,
+    modelName: "gpt-4-turbo-preview",
+  });
+
+  const llmGraphTransformer = new LLMGraphTransformer({
+    llm: model,
+    allowedNodes: ["Person", "Organization"],
+    allowedRelationships: ["SUES"],
+  });
+
+  const result = await llmGraphTransformer.convertToGraphDocuments([
+    new Document({ pageContent: "Elon Musk is suing OpenAI" }),
+  ]);
+
+  console.log(JSON.stringify(result));
+
+  expect(result).toEqual([
+    new GraphDocument({
+      nodes: [
+        new Node({ id: "Elon Musk", type: "Person" }),
+        new Node({ id: "OpenAI", type: "Organization" }),
+      ],
+      relationships: [
+        new Relationship({
+          source: new Node({ id: "Elon Musk", type: "Person" }),
+          target: new Node({ id: "OpenAI", type: "Organization" }),
           type: "SUES",
         }),
       ],

libs/langchain-community/src/experimental/graph_transformers/llm.ts

@@ -47,6 +47,13 @@ interface OptionalEnumFieldProps {
   fieldKwargs?: object;
 }
 
+function toTitleCase(str: string): string {
+  return str
+    .split(" ")
+    .map((w) => w[0].toUpperCase() + w.substring(1).toLowerCase())
+    .join("");
+}
+
 function createOptionalEnumType({
   enumValues = undefined,
   description = "",
@@ -122,7 +129,7 @@ function createSchema(allowedNodes: string[], allowedRelationships: string[]) {
 function mapToBaseNode(node: any): Node {
   return new Node({
     id: node.id,
-    type: node.type.replace(" ", "_").toUpperCase(),
+    type: toTitleCase(node.type),
   });
 }
 
@@ -131,11 +138,11 @@ function mapToBaseRelationship(relationship: any): Relationship {
   return new Relationship({
     source: new Node({
       id: relationship.sourceNodeId,
-      type: relationship.sourceNodeType.replace(" ", "_").toUpperCase(),
+      type: toTitleCase(relationship.sourceNodeType),
     }),
     target: new Node({
       id: relationship.targetNodeId,
-      type: relationship.targetNodeType.replace(" ", "_").toUpperCase(),
+      type: toTitleCase(relationship.targetNodeType),
     }),
     type: relationship.relationshipType.replace(" ", "_").toUpperCase(),
   });
@@ -208,16 +215,29 @@ export class LLMGraphTransformer {
       (this.allowedNodes.length > 0 || this.allowedRelationships.length > 0)
     ) {
       if (this.allowedNodes.length > 0) {
-        nodes = nodes.filter((node) => this.allowedNodes.includes(node.type));
+        const allowedNodesLowerCase = this.allowedNodes.map((node) =>
+          node.toLowerCase()
+        );
+
+        // For nodes, compare lowercased types
+        nodes = nodes.filter((node) =>
+          allowedNodesLowerCase.includes(node.type.toLowerCase())
+        );
+
+        // For relationships, compare lowercased types for both source and target nodes
         relationships = relationships.filter(
           (rel) =>
-            this.allowedNodes.includes(rel.source.type) &&
-            this.allowedNodes.includes(rel.target.type)
+            allowedNodesLowerCase.includes(rel.source.type.toLowerCase()) &&
+            allowedNodesLowerCase.includes(rel.target.type.toLowerCase())
         );
       }
+
       if (this.allowedRelationships.length > 0) {
+        // For relationships, compare lowercased types
         relationships = relationships.filter((rel) =>
-          this.allowedRelationships.includes(rel.type)
+          this.allowedRelationships
+            .map((rel) => rel.toLowerCase())
+            .includes(rel.type.toLowerCase())
         );
       }
     }

libs/langchain-community/src/graphs/graph_document.ts

@@ -60,14 +60,14 @@ export class Relationship extends Serializable {
   }
 }
 
-export class GraphDocument extends Document {
+export class GraphDocument extends Serializable {
   nodes: Node[];
 
   relationships: Relationship[];
 
   source: Document;
 
-  lc_namespace = ["langchain", "graph", "document_node"];
+  lc_namespace = ["langchain", "graph", "graph_document"];
 
   constructor({
     nodes,
@@ -78,7 +78,11 @@ export class GraphDocument extends Document {
     relationships: Relationship[];
     source: Document;
   }) {
-    super(source);
+    super({
+      nodes,
+      relationships,
+      source,
+    });
     this.nodes = nodes;
     this.relationships = relationships;
     this.source = source;