utils/kg/construct_kg.py

from langchain_community.graphs import Neo4jGraph
from langchain_experimental.graph_transformers import LLMGraphTransformer
from langchain_openai import ChatOpenAI , OpenAIEmbeddings
from langchain_core.documents import Document
from itext2kg import iText2KG


def get_graph(text,allowed_nodes=None,prompt=None,multiple_docs=False):
    
    llm = ChatOpenAI(temperature=0, model_name="gpt-4o-2024-08-06")

    if allowed_nodes:
        llm_transformer =   LLMGraphTransformer(llm=llm,allowed_nodes=allowed_nodes)
    else:
        llm_transformer = LLMGraphTransformer(llm=llm)
    documents = [Document(page_content=text)]

    graph_documents = llm_transformer.convert_to_graph_documents(documents)
    
    return graph_documents

def get_advanced_graph(semantic_chunk,knowledge_graph = None):
    openai_llm_model = ChatOpenAI(
                                model="gpt-4o",
                                temperature=0,
                                max_tokens=None,
                                timeout=None,
                                max_retries=2,
                            )
    
    openai_embeddings_model = OpenAIEmbeddings(
                                    model="text-embedding-3-large",
                                )
    itext2kg = iText2KG(llm_model = openai_llm_model, embeddings_model = openai_embeddings_model)
    if knowledge_graph:
        current_kg = itext2kg.build_graph(sections=[semantic_chunk],existing_knowledge_graph=knowledge_graph, ent_threshold=0.6, rel_threshold=0.6)
    else:
        current_kg = itext2kg.build_graph(sections=[semantic_chunk], ent_threshold=0.6, rel_threshold=0.6)

    return current_kg