Demosthene-OR/Knowledge_Graph_Generator

1 2from langchain_experimental.graph_transformers import LLMGraphTransformer3from langchain_core.documents import Document4from langchain_openai import ChatOpenAI, OpenAIEmbeddings5from langchain_community.vectorstores import FAISS6from langchain_core.prompts import PromptTemplate7from pyvis.network import Network8from dotenv import load_dotenv9import os10import asyncio11 12load_dotenv()13api_key = os.getenv("OPENAI_API_KEY")14 15llm = ChatOpenAI(temperature=0, model_name="gpt-4o")16graph_transformer = LLMGraphTransformer(llm=llm)17 18async def extract_graph_data(text):19    documents = [Document(page_content=text)]20    graph_documents = await graph_transformer.aconvert_to_graph_documents(documents)21    return graph_documents22 23def visualize_graph(graph_documents):24    net = Network(height="600px", width="100%", directed=True, notebook=False, bgcolor="#222222", font_color="white", filter_menu=True, cdn_resources='remote')25    26    # Collect all nodes and relationships across all graph documents in the list27    all_nodes = []28    all_relationships = []29    for doc in graph_documents:30        all_nodes.extend(doc.nodes)31        all_relationships.extend(doc.relationships)32 33    # Build node dictionary to deduplicate nodes by ID34    node_dict = {node.id: node for node in all_nodes}35    36    # Dynamically add any nodes referenced in relationships but missing from nodes list37    for rel in all_relationships:38        if rel.source.id not in node_dict:39            node_dict[rel.source.id] = rel.source40        if rel.target.id not in node_dict:41            node_dict[rel.target.id] = rel.target42 43    valid_edges = []44    valid_node_ids = set()45 46    for rel in all_relationships:47        if rel.source.id in node_dict and rel.target.id in node_dict:48            valid_edges.append(rel)49            valid_node_ids.update([rel.source.id, rel.target.id])50 51    for node_id in valid_node_ids:52        node = node_dict[node_id]53        try:54            net.add_node(node.id, label=node.id, title=node.type, group=node.type)55        except:56            continue57 58    # Deduplicate edges to avoid rendering duplicates59    edges_added = set()60    for rel in valid_edges:61        edge_key = (rel.source.id, rel.target.id, rel.type.lower())62        if edge_key not in edges_added:63            try:64                net.add_edge(rel.source.id, rel.target.id, label=rel.type.lower())65                edges_added.add(edge_key)66            except:67                continue68 69    net.set_options('{"physics": {"forceAtlas2Based": {"gravitationalConstant": -100, "centralGravity": 0.01, "springLength": 200, "springConstant": 0.08}, "minVelocity": 0.75, "solver": "forceAtlas2Based"}}')70    return net71 72 73def generate_knowledge_graph(text):74    graph_documents = asyncio.run(extract_graph_data(text))75    net = visualize_graph(graph_documents)76    return net, graph_documents77 78def answer_question_with_graph(question, graph_documents, k_relations=7):79    all_relationships = []80    for doc in graph_documents:81        all_relationships.extend(doc.relationships)82 83    if not all_relationships:84        return "No relationships found in the graph.", visualize_graph(graph_documents)85 86    rel_docs = []87    for i, rel in enumerate(all_relationships):88        text_rep = f"The entity '{rel.source.id}' has a '{rel.type}' relationship with the entity '{rel.target.id}'."89        rel_docs.append(Document(page_content=text_rep, metadata={"rel_index": i}))90 91    embeddings = OpenAIEmbeddings(model="text-embedding-3-small")92    vectorstore = FAISS.from_documents(rel_docs, embeddings)93    retrieved_docs = vectorstore.similarity_search(question, k=k_relations)94 95    used_relationships = [all_relationships[doc.metadata["rel_index"]] for doc in retrieved_docs]96    context = "\n".join([doc.page_content for doc in retrieved_docs])97 98    prompt = PromptTemplate(99        template='''Tu es un assistant expert qui répond aux questions en se basant UNIQUEMENT sur ce sous-ensemble de relations extraites d'un graphe de connaissances.\n\nContexte (Relations pertinentes trouvées) :\n{context}\n\nQuestion : {question}\n\nRéponds de manière claire et concise en anglais. Réponds directement à la question sans jamais faire référence au contexte, aux relations fournies, au document ou au graphe (ne commence pas par "Based on...", "According to the...", etc.). Mentionne le contexte uniquement si la réponse n'y figure pas (dans ce cas, réponds explicitement : "I cannot answer this question based on the provided context.").''',100        input_variables=["context", "question"]101    )102 103    chain = prompt | llm104    answer = chain.invoke({"context": context, "question": question}).content105 106    net = Network(height="450px", width="100%", directed=True, bgcolor="#222222", font_color="white")107 108    nodes_added = set()109    for rel in used_relationships:110        if rel.source.id not in nodes_added:111            net.add_node(rel.source.id, label=rel.source.id, title=rel.source.type, group=rel.source.type)112            nodes_added.add(rel.source.id)113        if rel.target.id not in nodes_added:114            net.add_node(rel.target.id, label=rel.target.id, title=rel.target.type, group=rel.target.type)115            nodes_added.add(rel.target.id)116        try:117            net.add_edge(rel.source.id, rel.target.id, label=rel.type)118        except:119            pass120 121    net.set_options('{"physics": {"forceAtlas2Based": {"gravitationalConstant": -50}}}')122    net.save_graph("filtered_graph.html")123 124    return answer, net125