Langchain/CreatingGraphDocumentsfromText/Neo4j: KeyError: tail_type when using llm_transformer.convert_to_graph()

[ghinwa-mj]

Checked other resources

• I added a very descriptive title to this question.
• I searched the LangChain documentation with the integrated search.
• I used the GitHub search to find a similar question and didn't find it.

Commit to Help

• I commit to help with one of those options 👆

Example Code

Extract text from PDF
pdf_text = extract_text_from_pdf(pdf_path)

Create a Document object
document_pdf = [Document(page_content=pdf_text)]

#Create a GraphDocument object
llm_transformer = LLMGraphTransformer(llm=llm)
graph_documents = llm_transformer.convert_to_graph_documents(document_pdf)

Description

I am using the function llm_transformer.convert_to_graph() to create graph documents from pdf_texts and upload them to my neo4j database.

Sometimes I get a KeyError: “KeyError: 'tail_type’” or “KeyError: 'tail’”.

This doesn’t happen always (typically with the longer pdfs). Any ideas where this could be coming from and what I can do to fix it?

System Info

Using GoogleColab

[dosubot[bot]]

Hey there, @ghinwa-mj! I'm here to help you with your LangChain issue. I'm a bot that can assist you with bug fixes, answer your questions, and guide you to become a contributor. Let's tackle this together!

The KeyError: 'tail_type' or KeyError: 'tail' occurs because the code expects every relationship in the parsed JSON to have the keys head, head_type, relation, tail, and tail_type. If any of these keys are missing, the error will be raised.

To fix this issue, you can add a check to ensure that all required keys are present before attempting to access them. Here is the modified process_response method with the necessary checks:

def process_response(self, document: Document) -> GraphDocument:
    """
    Processes a single document, transforming it into a graph document using
    an LLM based on the model's schema and constraints.
    """
    text = document.page_content
    raw_schema = self.chain.invoke({"input": text})
    if self._function_call:
        raw_schema = cast(Dict[Any, Any], raw_schema)
        nodes, relationships = _convert_to_graph_document(raw_schema)
    else:
        nodes_set = set()
        relationships = []
        if not isinstance(raw_schema, str):
            raw_schema = raw_schema.content
        parsed_json = self.json_repair.loads(raw_schema)
        for rel in parsed_json:
            # Check if all required keys are present
            if all(key in rel for key in ["head", "head_type", "relation", "tail", "tail_type"]):
                # Nodes need to be deduplicated using a set
                nodes_set.add((rel["head"], rel["head_type"]))
                nodes_set.add((rel["tail"], rel["tail_type"]))

                source_node = Node(id=rel["head"], type=rel["head_type"])
                target_node = Node(id=rel["tail"], type=rel["tail_type"])
                relationships.append(
                    Relationship(
                        source=source_node, target=target_node, type=rel["relation"]
                    )
                )
            else:
                # Handle the case where keys are missing
                print(f"Skipping relationship due to missing keys: {rel}")

        # Create nodes list
        nodes = [Node(id=el[0], type=el[1]) for el in list(nodes_set)]

    # Strict mode filtering
    if self.strict_mode and (self.allowed_nodes or self.allowed_relationships):
        if self.allowed_nodes:
            lower_allowed_nodes = [el.lower() for el in self.allowed_nodes]
            nodes = [
                node for node in nodes if node.type.lower() in lower_allowed_nodes
            ]
            relationships = [
                rel
                for rel in relationships
                if rel.source.type.lower() in lower_allowed_nodes
                and rel.target.type.lower() in lower_allowed_nodes
            ]
        if self.allowed_relationships:
            relationships = [
                rel
                for rel in relationships
                if rel.type.lower()
                in [el.lower() for el in self.allowed_relationships]
            ]

    return GraphDocument(nodes=nodes, relationships=relationships, source=document)

This modification ensures that the code only processes relationships that contain all the required keys, thereby preventing the KeyError ^[1].

To continue talking to Dosu, mention @dosu.

[ghinwa-mj]

Thank you for your reply, however the issue I need to answer is 1. Why these relationships are empty 2. Why "tail" and "tail_end" are empty 3. Accordingly, what to do to stop this from happening. This is especially when I am automating a script.

[dosubot[bot]]

The relationships, "tail," and "tail_type" are empty when using llm_transformer.convert_to_graph() because the raw_schema returned by the LLM does not contain the expected structure or data. This can happen if the LLM's output is not correctly parsed or if the prompt used to generate the output is not properly formatted.

To prevent this issue, especially when automating a script, ensure that the prompt used for the LLM is correctly defined and that the LLM's output is properly parsed. Here is a step-by-step approach to address this:

1. Ensure Proper Prompt Definition: Make sure the prompt used to generate the LLM's output is correctly defined to extract the required entities and relationships.

2. Parse the LLM Output Correctly: Ensure that the output from the LLM is correctly parsed into the expected JSON format.

3. Check for Strict Mode Filtering: If strict mode is enabled, ensure that the allowed nodes and relationships are correctly specified.

Here is an example of how to define the prompt and parse the LLM output correctly:

def create_unstructured_prompt(node_labels: Optional[List[str]] = None, rel_types: Optional[List[str]] = None) -> ChatPromptTemplate:
    node_labels_str = str(node_labels) if node_labels else ""
    rel_types_str = str(rel_types) if rel_types else ""
    base_string_parts = [
        "You are a top-tier algorithm designed for extracting information in "
        "structured formats to build a knowledge graph. Your task is to identify "
        "the entities and relations requested with the user prompt from a given "
        "text. You must generate the output in a JSON format containing a list "
        'with JSON objects. Each object should have the keys: "head", '
        '"head_type", "relation", "tail", and "tail_type". The "head" '
        "key must contain the text of the extracted entity with one of the types "
        "from the provided list in the user prompt.",
        f'The "head_type" key must contain the type of the extracted head entity, '
        f"which must be one of the types from {node_labels_str}."
        if node_labels
        else "",
        f'The "relation" key must contain the type of relation between the "head" '
        f'and the "tail", which must be one of the relations from {rel_types_str}.'
        if rel_types
        else "",
        f'The "tail" key must represent the text of an extracted entity which is '
        f'the tail of the relation, and the "tail_type" key must contain the type '
        f"of the tail entity from {node_labels_str}."
        if node_labels
        else "",
        "Attempt to extract as many entities and relations as you can. Maintain "
        "Entity Consistency: When extracting entities, it's vital to ensure "
        'consistency. If an entity, such as "John Doe", is mentioned multiple '
        "times in the text but is referred to by different names or pronouns "
        '(e.g., "Joe", "he"), always use the most complete identifier for '
        "that entity. The knowledge graph should be coherent and easily "
        "understandable, so maintaining consistency in entity references is "
        "crucial.",
        "IMPORTANT NOTES:\n- Don't add any explanation and text.",
    ]
    system_prompt = "\n".join(filter(None, base_string_parts))
    system_message = SystemMessage(content=system_prompt)
    parser = JsonOutputParser(pydantic_object=UnstructuredRelation)

    human_prompt = PromptTemplate(
        template="""Based on the following example, extract entities and 
relations from the provided text.\n\n
Use the following entity types, don't use other entity that is not defined below:
# ENTITY TYPES:
{node_labels}

Use the following relation types, don't use other relation that is not defined below:
# RELATION TYPES:
{rel_types}

Below are a number of examples of text and their extracted entities and relationships.
{examples}

For the following text, extract entities and relations as in the provided example.
{format_instructions}\nText: {input}""",
        input_variables=["input"],
        partial_variables={
            "format_instructions": parser.get_format_instructions(),
            "node_labels": node_labels,
            "rel_types": rel_types,
            "examples": examples,
        },
    )

    human_message_prompt = HumanMessagePromptTemplate(prompt=human_prompt)

    chat_prompt = ChatPromptTemplate.from_messages(
        [system_message, human_message_prompt]
    )
    return chat_prompt

By ensuring the prompt is correctly defined and the LLM output is properly parsed, you can prevent the issue of empty "tail" and "tail_type" fields when converting PDF texts to graph documents ^[1].

To continue talking to Dosu, mention @dosu.