File size: 915 Bytes
f7ab812 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 |
from typing import List, Optional
from langchain.schema import Document
from langchain_community.graphs.graph_document import GraphDocument
from entities_extraction import get_extraction_chain
from utils import map_to_base_node, map_to_base_relationship
def extract_graph(
document: Document,
nodes: Optional[List[str]] = None,
rels: Optional[List[str]] = None,
) -> None:
# Extract graph data using OpenAI functions
extract_chain = get_extraction_chain(nodes, rels)
try:
data = extract_chain.invoke(document.page_content)["function"]
except Exception as e:
raise ValueError(f"Extraction failed: {e}")
# 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,
)
return graph_document
|