| 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 | |