From 33604f2c97334b0091efb30314948a06d5314f9c Mon Sep 17 00:00:00 2001
From: jeromechoo <hello@jeromechoo.com>
Date: Mon, 5 Aug 2024 17:13:26 -0500
Subject: [PATCH] DiffbotGraphTransformer doesn't need an LLMGraphTransformer

---
 backend/src/llm.py | 22 +++++++++++++---------
 1 file changed, 13 insertions(+), 9 deletions(-)

diff --git a/backend/src/llm.py b/backend/src/llm.py
index f155d8fca..0ee61b650 100644
--- a/backend/src/llm.py
+++ b/backend/src/llm.py
@@ -145,16 +145,20 @@ def get_graph_document_list(
 ):
     futures = []
     graph_document_list = []
-    if llm.get_name() == "ChatOllama":
-        node_properties = False
+
+    if "diffbot_api_key" in dir(llm):
+        llm_transformer = llm
     else:
-        node_properties = ["description"]
-    llm_transformer = LLMGraphTransformer(
-        llm=llm,
-        node_properties=node_properties,
-        allowed_nodes=allowedNodes,
-        allowed_relationships=allowedRelationship,
-    )
+        if "get_name" in dir(llm) and llm.get_name() == "ChatOllama":
+            node_properties = False
+        else:
+            node_properties = ["description"]
+        llm_transformer = LLMGraphTransformer(
+            llm=llm,
+            node_properties=node_properties,
+            allowed_nodes=allowedNodes,
+            allowed_relationships=allowedRelationship,
+        )
     with ThreadPoolExecutor(max_workers=10) as executor:
         for chunk in combined_chunk_document_list:
             chunk_doc = Document(