Update generate_knowledge_graph.py

generate_knowledge_graph.py

@@ -15,14 +15,17 @@ api_key = os.getenv("OPENAI_API_KEY")
 llm = ChatOpenAI(temperature=0, model_name="gpt-4o")
 graph_transformer = LLMGraphTransformer(llm=llm)
 
+
 async def extract_graph_data(text):
     documents = [Document(page_content=text)]
     graph_documents = await graph_transformer.aconvert_to_graph_documents(documents)
     return graph_documents
 
+
 def visualize_graph(graph_documents):
-    net = Network(height="600px", width="100%", directed=True, notebook=False, bgcolor="#222222", font_color="white", filter_menu=True, cdn_resources='remote')
-    
+    net = Network(height="600px", width="100%", directed=True, notebook=False,
+                  bgcolor="#222222", font_color="white", filter_menu=True, cdn_resources='remote')
+
     # Collect all nodes and relationships across all graph documents in the list
     all_nodes = []
     all_relationships = []
@@ -32,7 +35,7 @@ def visualize_graph(graph_documents):
 
     # Build node dictionary to deduplicate nodes by ID
     node_dict = {node.id: node for node in all_nodes}
-    
+
     # Dynamically add any nodes referenced in relationships but missing from nodes list
     for rel in all_relationships:
         if rel.source.id not in node_dict:
@@ -51,7 +54,8 @@ def visualize_graph(graph_documents):
     for node_id in valid_node_ids:
         node = node_dict[node_id]
         try:
-            net.add_node(node.id, label=node.id, title=node.type, group=node.type)
+            net.add_node(node.id, label=node.id,
+                         title=node.type, group=node.type)
         except:
             continue
 
@@ -61,7 +65,8 @@ def visualize_graph(graph_documents):
         edge_key = (rel.source.id, rel.target.id, rel.type.lower())
         if edge_key not in edges_added:
             try:
-                net.add_edge(rel.source.id, rel.target.id, label=rel.type.lower())
+                net.add_edge(rel.source.id, rel.target.id,
+                             label=rel.type.lower())
                 edges_added.add(edge_key)
             except:
                 continue
@@ -75,6 +80,7 @@ def generate_knowledge_graph(text):
     net = visualize_graph(graph_documents)
     return net, graph_documents
 
+
 def answer_question_with_graph(question, graph_documents, k_relations=7):
     all_relationships = []
     for doc in graph_documents:
@@ -86,13 +92,15 @@ def answer_question_with_graph(question, graph_documents, k_relations=7):
     rel_docs = []
     for i, rel in enumerate(all_relationships):
         text_rep = f"The entity '{rel.source.id}' has a '{rel.type}' relationship with the entity '{rel.target.id}'."
-        rel_docs.append(Document(page_content=text_rep, metadata={"rel_index": i}))
+        rel_docs.append(Document(page_content=text_rep,
+                        metadata={"rel_index": i}))
 
     embeddings = OpenAIEmbeddings(model="text-embedding-3-small")
     vectorstore = FAISS.from_documents(rel_docs, embeddings)
     retrieved_docs = vectorstore.similarity_search(question, k=k_relations)
 
-    used_relationships = [all_relationships[doc.metadata["rel_index"]] for doc in retrieved_docs]
+    used_relationships = [
+        all_relationships[doc.metadata["rel_index"]] for doc in retrieved_docs]
     context = "\n".join([doc.page_content for doc in retrieved_docs])
 
     prompt = PromptTemplate(
@@ -103,22 +111,26 @@ def answer_question_with_graph(question, graph_documents, k_relations=7):
     chain = prompt | llm
     answer = chain.invoke({"context": context, "question": question}).content
 
-    net = Network(height="450px", width="100%", directed=True, bgcolor="#222222", font_color="white")
+    net = Network(height="450px", width="100%", directed=True,
+                  bgcolor="#222222", font_color="white")
 
     nodes_added = set()
     for rel in used_relationships:
         if rel.source.id not in nodes_added:
-            net.add_node(rel.source.id, label=rel.source.id, title=rel.source.type, group=rel.source.type)
+            net.add_node(rel.source.id, label=rel.source.id,
+                         title=f"Label : {rel.source.id}\nGroup : {rel.source.type}", group=rel.source.type)
             nodes_added.add(rel.source.id)
         if rel.target.id not in nodes_added:
-            net.add_node(rel.target.id, label=rel.target.id, title=rel.target.type, group=rel.target.type)
+            net.add_node(rel.target.id, label=rel.target.id,
+                         title=f"Label : {rel.target.id}\nGroup : {rel.target.type}", group=rel.target.type)
             nodes_added.add(rel.target.id)
         try:
             net.add_edge(rel.source.id, rel.target.id, label=rel.type)
         except:
             pass
 
-    net.set_options('{"physics": {"forceAtlas2Based": {"gravitationalConstant": -50}}}')
+    net.set_options(
+        '{"physics": {"forceAtlas2Based": {"gravitationalConstant": -50}}}')
     net.save_graph("filtered_graph.html")
 
     return answer, net