Skip to content

Commit 276302a

Browse files
jeromechoojeromechoo
authored andcommitted
DiffbotGraphTransformer doesn't need an LLMGraphTransformer (#659)
Co-authored-by: jeromechoo <[email protected]>
1 parent 01fcb08 commit 276302a

File tree

1 file changed

+13
-9
lines changed

1 file changed

+13
-9
lines changed

backend/src/llm.py

+13-9
Original file line numberDiff line numberDiff line change
@@ -145,16 +145,20 @@ def get_graph_document_list(
145145
):
146146
futures = []
147147
graph_document_list = []
148-
if llm.get_name() == "ChatOllama":
149-
node_properties = False
148+
149+
if "diffbot_api_key" in dir(llm):
150+
llm_transformer = llm
150151
else:
151-
node_properties = ["description"]
152-
llm_transformer = LLMGraphTransformer(
153-
llm=llm,
154-
node_properties=node_properties,
155-
allowed_nodes=allowedNodes,
156-
allowed_relationships=allowedRelationship,
157-
)
152+
if "get_name" in dir(llm) and llm.get_name() == "ChatOllama":
153+
node_properties = False
154+
else:
155+
node_properties = ["description"]
156+
llm_transformer = LLMGraphTransformer(
157+
llm=llm,
158+
node_properties=node_properties,
159+
allowed_nodes=allowedNodes,
160+
allowed_relationships=allowedRelationship,
161+
)
158162
with ThreadPoolExecutor(max_workers=10) as executor:
159163
for chunk in combined_chunk_document_list:
160164
chunk_doc = Document(

0 commit comments

Comments
 (0)