import base64from typing import Union, List, Dict, Any
import pymupdfimport pymupdf4llmimport numpy as npfrom PIL import Imagefrom langchain_core.language_models import BaseLanguageModelfrom langchain_core.messages import SystemMessage, HumanMessage

def ocr_images_pytesseract(images: List[Union[np.ndarray | Image.Image]]) -> str:import pytesseract
all_text: str = ""for image in images:if isinstance(image, Image.Image):image = image.filter(ImageFilter.SMOOTH())all_text += '\n' +pytesseract.image_to_string(image, lang="eng", config='--psm 3 --dpi 300 --oem 1')return all_text.strip('\n \t')
# Use PyMuPDF to open the document and PyMuPDF4LLM to get the markdown outputdoc = pymupdf.Document(pdf_path)doc_markdown = pymupdf4llm.to_markdown(doc, page_chunks=True, write_images=False, force_text=True)

def extract_page_info(page_num: int, doc_metadata: dict, toc: list = None) -> Dict[str, Any]:"""Extracts text and metadata from a single page of a PDF document.
Args:page_num (int): The page numberdoc_metadata (dict): The whole document metadata to store along with each page metadatatoc (list | None): The list that represents the table-of-contents of the documentReturns:A dictionary with the following keys:1) text: Text content extracted from the page2) page_metadata: Metadata specific to the page only like page number, chapter number etc.3) doc_metadata: Metadata common to the whole document like filename, author etc."""page_info = {}
# Read the page of page = doc[page_num]
# doc_markdown stores the page-by-page markdown output of the documenttext_content: str = self.doc_markdown[page_num]['text']
# Get a list of all the images on this page - automatically, perform OCR on all these images and store the output as page text# There are 3 options available: each with different preprocessing steps and all of them implemented in the self._ocr_func methodimages_list: list[list] = page.get_images()if len(images_list) > 0:imgs = []for img in images_list:xref = img[0]pix = pymupdf.Pixmap(doc, xref)# using PyTesseract requires storing the image as PIL.Image though it could've been bytes or np.ndarray as wellcspace = pix.colorspaceif cspace is None:mode: str = "L"elif cspace.n == 1:mode = "L" if pix.alpha == 0 else "LA"elif cspace.n == 3:mode = "RGB" if pix.alpha == 0 else "RGBA"else:mode = "CMYK"img = Image.frombytes(mode, (pix.width, pix.height), pix.samples)if mode != "L":img = img.convert("L")imgs.append(img)text_content += '\n' + ocr_images_pytesseract(imgs)text_content = text_content.strip(' \n\t')page_info["text"] = text_contentpage_info["page_metadata"] = get_page_metadata(page_num=page_num+1, page_text=text_content, toc=toc)page_info["doc_metadata"] = doc_metadatareturn page_info

from pydantic import BaseModel, Field
class Property(BaseModel):"""A single property consisting of key and value"""key: str = Field(..., description="key")value: str = Field(..., description="value")
class Node(BaseModel):id: str = Field(..., description="The identifying property of the node in Title Case")type: str = Field(..., description="The entity type / label of the node in PascalCase.")properties: Optional[List[Property]] = Field(default=[], description="Detailed properties of the node")aliases: List[str] = Field(default=[], description="Alternative names or identifiers for the entity in Title Case")definition: Optional[str] = Field(None, description="A concise definition or description of the entity")
class Relationship(BaseModel):start_node_id: str = Field(..., description="The id of the first node in the relationship")end_node_id: str = Field(..., description="The id of the second node in the relationship")type: str = Field(..., description="TThe specific, descriptive label of the relationship in SCREAMING_SNAKE_CASE")properties: List[Property] = Field(default=[], description="Detailed properties of the relationship")context: Optional[str] = Field(None, description="Additional contextual information about the relationship")class KnowledgeGraph(BaseModel):"""Generate a knowledge graph with entities and relationships."""nodes: list[Node] = Field(..., description="List of nodes in the knowledge graph")rels: list[Relationship] = Field(..., description="List of relationships in the knowledge graph")

from langchain_core.prompts import ChatPromptTemplate
DATA_EXTRACTION_SYSTEM = """# Knowledge Graph Extraction for Rich Information Retrieval## 1. OverviewYou are an advanced algorithm designed to extract knowledge from various types of informational content. \Your task is to build a knowledge graph that provides rich, contextual information for downstream tasks.
## 2. Content Focus- Extract detailed information about concepts, entities, processes, and their relationships from the given text.- Prioritize information that provides rich context and is likely to be useful for answering a wide range of questions.- Include relevant attributes, properties, and descriptive information for each extracted entity.
## 3. Node Extraction- **Node IDs**: Use clear, unambiguous identifiers in Title Case. Avoid integers, abbreviations, and acronyms.- **Node Types**: Use PascalCase. Be as specific and descriptive as possible to aid in Wikidata matching.- Include all relevant attributes of the entity in the node properties.- Extract and include alternative names or aliases for entities when present in the text.
## 4. Relationship Extraction- Use SCREAMING_SNAKE_CASE for relationship types.- Create detailed, informative relationship types that clearly describe the nature of the connection.- Include directional relationships where applicable (e.g., PRECEDED_BY, FOLLOWED_BY instead of just RELATED_TO).
## 5. Contextual Information- For each node and relationship, strive to capture contextual information that might be useful for answering questions.- Include temporal information when available (e.g., dates, time periods, sequence of events).- Capture geographical or spatial information if relevant.
## 6. Handling Definitions and Descriptions- For key concepts, include concise definitions or descriptions as node properties.- Capture any notable characteristics, functions, or use cases of entities.
## 7. Coreference and Consistency- Maintain consistent entity references throughout the graph.- Resolve coreferences to their most complete form, including potential aliases or alternative names.
## 8. Granularity- Strike a balance between detailed extraction and maintaining a coherent graph structure.- Create separate nodes for distinct concepts, even if closely related."""

prompt = ChatPromptTemplate.from_messages([("system", DATA_EXTRACTION_SYSTEM),("human", "Use the given format to extract information from the following input which is a small sample from a much larger text belonging to the same subject matter: {input}"),("human", "Tip: Make sure to answer in the correct format"),])

nodes = []rels = []
for doc in docs:output: KnowledgeGraph = data_ext_chain.invoke({"input": doc.page_content})nodes.extend(format_nodes(output.nodes))rels.extend(format_rels(output.rels))

DATA_EXTRACTION_SYSTEM = """# Knowledge Graph Extraction for Rich Information Retrieval## 1. OverviewYou are an advanced algorithm designed to extract knowledge from various types of informational content. \Your task is to build a knowledge graph that provides rich, contextual information for downstream tasks.{subject}## 2. Content Focus- Extract detailed information about concepts, entities, processes, and their relationships from the given text.- Prioritize information that provides rich context and is likely to be useful for answering a wide range of questions.- Include relevant attributes, properties, and descriptive information for each extracted entity.
## 3. Node Extraction- **Node IDs**: Use clear, unambiguous identifiers in Title Case. Avoid integers, abbreviations, and acronyms.- **Node Types**: Use PascalCase. Be as specific and descriptive as possible to aid in Wikidata matching.- Include all relevant attributes of the entity in the node properties.- Extract and include alternative names or aliases for entities when present in the text.- Following are some existing node types that were extracted from previous samples of the same document:\n{node_types}
## 4. Relationship Extraction- Use SCREAMING_SNAKE_CASE for relationship types.- Create detailed, informative relationship types that clearly describe the nature of the connection.- Include directional relationships where applicable (e.g., PRECEDED_BY, FOLLOWED_BY instead of just RELATED_TO).
## 5. Contextual Information- For each node and relationship, strive to capture contextual information that might be useful for answering questions.- Include temporal information when available (e.g., dates, time periods, sequence of events).- Capture geographical or spatial information if relevant.
## 6. Handling Definitions and Descriptions- For key concepts, include concise definitions or descriptions as node properties.- Capture any notable characteristics, functions, or use cases of entities.
## 7. Coreference and Consistency- Maintain consistent entity references throughout the graph.- Resolve coreferences to their most complete form, including potential aliases or alternative names.
## 8. Granularity- Strike a balance between detailed extraction and maintaining a coherent graph structure.- Create separate nodes for distinct concepts, even if closely related."""from langchain_core.documents.base import Document

def merge_nodes(nodes: List[Node]) -> Node:'''Merges all the nodes with the same ID and type into a single node
The only difference in these nodes is their list of properties, aliases and definition so I extract and combine a list of unique properties and aliases from all these nodesand pick the longest definition to use as the definition of the merged node
Args: nodes (List[Node]):- List of nodes to merge into oneReturns: A single Node object created by merging all the nodes'''props: List[Property] = []definition = ""aliases = set([])max_def_len = -1for node in nodes:props.extend([p for p in node.propertiesif not p in props])if len(node.definition) >= max_def_len:# Pick the longest definition based on an inaccurate assumption: # longest description = most descriptive definitiondefinition = node.definitionmax_def_len = len(node.definition)aliases.update(set(node.aliases))return Node(id=nodes[0].id, type=nodes[0].type, properties=props, definition=definition, aliases=aliases)
def docs2nodes_and_rels(docs: List[Document], text_subject: str = '', use_existing_node_types: bool = False) -> Tuple[List[Document], List[Node], List[Relationship]]:'''Extract the list of nodes and relationships from the list of documents
Args:docs (List[Document]): List of documentstext_subject (str): The overall subject of the text that the given documents belong to. Default=''use_existing_node_types (bool): Whether to use the node types of the already existing nodes in theKG as a guide to the LLM. Default=False.
Returns:A tuple containing the following lists:docs (List[Document]): A list of the slightly modified input docsnodes (List[Node]): The list of unique extracted nodesrels (List[Relationship]): The list of extracted relationships'''# Only store unique nodes - Helps later when adding to KGnodes_dict: Dict[str, list] = dict({})rels: list = []
ex_node_types = set({})if use_existing_node_types:# get_existing_labels is a method that returns the set of labels of the already existing nodes in the KGex_node_types = get_existing_labels()
for i, doc in enumerate(docs):output: KnowledgeGraph = data_ext_chain.invoke({"subject": text_subject,"input": doc.page_content, "node_types": ','.join(list(ex_node_types))})
output_nodes: List[Node] = format_nodes(output.nodes)output_rels: List[Relationship] = format_rels(output.rels)
ntypes = set([n.type for n in output_nodes])
# Store nodes with same type and ID together so that they can be merged together laterdnodes_dict: Dict[str, List[str]] = {}for n in output_nodes:nk = f"{n.id}::{n.type}"if nk in dnodes_dict:if not n in dnodes_dict[nk]:dnodes_dict[nk].append(n)else:dnodes_dict[nk] = [n]
for r in output_rels:if not r in rels:rels.append(r)
# Store Node IDs+Node Type of nodes mentioned in this doc/chunk in doc metadatadoc.metadata['mentions'] = list(dnodes_dict.keys())nodes_dict = {**nodes_dict, **dnodes_dict}# Update existing node types for next iterationex_node_types: set[str] = ex_node_types.union(ntypes)
# Merge all duplicate nodes (nodes with the same ID and type) into a single node# This involves combining their properties, aliases and definitionnodes: List[Node] = [merge_nodes(nds) for _, nds in nodes_dict.items()]
return docs, nodes, list(rels)