templates: Switch neo4j generation template to LLMGraphTransformer (#19024)

templates/neo4j-generation/neo4j_generation/chain.py

@@ -1,19 +1,9 @@
 from typing import List, Optional
 
-from langchain.chains.openai_functions import (
-    create_structured_output_chain,
-)
-from langchain_community.chat_models import ChatOpenAI
 from langchain_community.graphs import Neo4jGraph
-from langchain_community.graphs.graph_document import GraphDocument
 from langchain_core.documents import Document
-from langchain_core.prompts import ChatPromptTemplate
-
-from neo4j_generation.utils import (
-    KnowledgeGraph,
-    map_to_base_node,
-    map_to_base_relationship,
-)
+from langchain_experimental.graph_transformers import LLMGraphTransformer
+from langchain_openai import ChatOpenAI
 
 graph = Neo4jGraph()
 
@@ -21,63 +11,6 @@ graph = Neo4jGraph()
 llm = ChatOpenAI(model="gpt-3.5-turbo-16k", temperature=0)
 
 
-def get_extraction_chain(
-    allowed_nodes: Optional[List[str]] = None, allowed_rels: Optional[List[str]] = None
-):
-    """
-    Constructs and returns an extraction chain for building a knowledge graph based on specified parameters.
-
-    The function generates a chat prompt template, outlining the instructions for an LLM to extract information
-    and construct a knowledge graph. It primarily focuses on consistency in labeling nodes, handling numerical data
-    and dates, coreference resolution, and strict compliance with the provided rules.
-
-    Parameters:
-    - allowed_nodes (Optional[List[str]]): A list of node labels that are allowed to be used in the knowledge graph.
-                                           If not provided, there won't be any specific restriction on node labels.
-    - allowed_rels (Optional[List[str]]): A list of relationship types that are allowed in the knowledge graph.
-                                         If not provided, there won't be any specific restriction on relationship types.
-    """  # noqa: E501
-    prompt = ChatPromptTemplate.from_messages(
-        [
-            (
-                "system",
-                f"""# Knowledge Graph Instructions for GPT-4
-## 1. Overview
-You are a top-tier algorithm designed for extracting information in structured formats to build a knowledge graph.
-- **Nodes** represent entities and concepts. They're akin to Wikipedia nodes.
-- The aim is to achieve simplicity and clarity in the knowledge graph, making it accessible for a vast audience.
-## 2. Labeling Nodes
-- **Consistency**: Ensure you use basic or elementary types for node labels.
-  - For example, when you identify an entity representing a person, always label it as **"person"**. Avoid using more specific terms like "mathematician" or "scientist".
-- **Node IDs**: Never utilize integers as node IDs. Node IDs should be names or human-readable identifiers found in the text.
-{'- **Allowed Node Labels:**' + ", ".join(allowed_nodes) if allowed_nodes else ""}
-{'- **Allowed Relationship Types**:' + ", ".join(allowed_rels) if allowed_rels else ""}
-## 3. Handling Numerical Data and Dates
-- Numerical data, like age or other related information, should be incorporated as attributes or properties of the respective nodes.
-- **No Separate Nodes for Dates/Numbers**: Do not create separate nodes for dates or numerical values. Always attach them as attributes or properties of nodes.
-- **Property Format**: Properties must be in a key-value format.
-- **Quotation Marks**: Never use escaped single or double quotes within property values.
-- **Naming Convention**: Use camelCase for property keys, e.g., `birthDate`.
-## 4. Coreference Resolution
-- **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 throughout the knowledge graph. In this example, use "John Doe" as the entity ID.
-Remember, the knowledge graph should be coherent and easily understandable, so maintaining consistency in entity references is crucial.
-## 5. Strict Compliance
-Adhere to the rules strictly. Non-compliance will result in termination.
-          """,  # noqa: E501
-            ),
-            (
-                "human",
-                "Use the given format to extract information from the "
-                "following input: {input}",
-            ),
-            ("human", "Tip: Make sure to answer in the correct format"),
-        ]
-    )
-    return create_structured_output_chain(KnowledgeGraph, llm, prompt, verbose=False)
-
-
 def chain(
     text: str,
     allowed_nodes: Optional[List[str]] = None,
@@ -97,15 +30,15 @@ def chain(
     Returns:
     str: A confirmation message indicating the completion of the graph construction.
     """  # noqa: E501
+    # Construct document based on text
+    documents = [Document(page_content=text)]
     # Extract graph data using OpenAI functions
-    extract_chain = get_extraction_chain(allowed_nodes, allowed_relationships)
-    data = extract_chain.run(text)
-    # Construct a graph document
-    graph_document = GraphDocument(
-        nodes=[map_to_base_node(node) for node in data.nodes],
-        relationships=[map_to_base_relationship(rel) for rel in data.rels],
-        source=Document(page_content=text),
+    llm_graph_transformer = LLMGraphTransformer(
+        llm=llm,
+        allowed_nodes=allowed_nodes,
+        allowed_relationships=allowed_relationships,
     )
+    graph_documents = llm_graph_transformer.convert_to_graph_documents(documents)
     # Store information into a graph
-    graph.add_graph_documents([graph_document])
+    graph.add_graph_documents(graph_documents)
     return "Graph construction finished"