From 6703d795c56d860ad32b6417ff7eff419e4b95b4 Mon Sep 17 00:00:00 2001 From: Tomaz Bratanic Date: Mon, 26 Aug 2024 13:26:26 +0200 Subject: [PATCH] Handle Ollama tool raw schema in llmgraphtransformer (#25752) --- .../graph_transformers/llm.py | 19 +++++++++++++++---- 1 file changed, 15 insertions(+), 4 deletions(-) diff --git a/libs/experimental/langchain_experimental/graph_transformers/llm.py b/libs/experimental/langchain_experimental/graph_transformers/llm.py index 7f1eb9b8fb3..61997a8715d 100644 --- a/libs/experimental/langchain_experimental/graph_transformers/llm.py +++ b/libs/experimental/langchain_experimental/graph_transformers/llm.py @@ -479,7 +479,7 @@ def _parse_and_clean_json( nodes.append( Node( id=node["id"], - type=node.get("type"), + type=node.get("type", "Node"), properties=node_properties, ) ) @@ -584,9 +584,20 @@ def _convert_to_graph_document( ] ) except Exception: # Google type response - argument_json = json.loads( - raw_schema["raw"].additional_kwargs["function_call"]["arguments"] - ) + try: + argument_json = json.loads( + raw_schema["raw"].additional_kwargs["function_call"][ + "arguments" + ] + ) + except Exception: # Ollama type response + argument_json = raw_schema["raw"].tool_calls[0]["args"] + if isinstance(argument_json["nodes"], str): + argument_json["nodes"] = json.loads(argument_json["nodes"]) + if isinstance(argument_json["relationships"], str): + argument_json["relationships"] = json.loads( + argument_json["relationships"] + ) nodes, relationships = _parse_and_clean_json(argument_json) except Exception: # If we can't parse JSON