Update app.py

app.py

@@ -94,6 +94,12 @@ def generate_image(prompt):
         traceback.print_exc()
         return None
 
+def create_relations_table(graph_documents_filtered):
+    df = pd.DataFrame(columns=['node1', 'node2', 'relation'])
+    for edge in graph_documents_filtered[0].relationships:
+        df = pd.concat([df, pd.DataFrame({'node1': [edge.source.id], 'node2': [edge.target.id], 'relation': [edge.type]})], ignore_index=True)
+    return df
+
 def process_text(text, question):
     try:
         print("Creating graph...")
@@ -114,6 +120,10 @@ def process_text(text, question):
         graph_viz_path = visualize_graph(graph)
         print(f"Graph visualization saved to: {graph_viz_path}")
 
+        print("Creating relations table...")
+        relations_table = create_relations_table(graph_documents_filtered)
+        print("Relations table created")
+
         print("Generating summary...")
         summary_prompt = f"Summarize the following text in one sentence: {text}"
         summary = llm.invoke(summary_prompt).content
@@ -126,29 +136,33 @@ def process_text(text, question):
         else:
             print("Failed to generate or save image")
 
-        return answer, graph_viz_path, summary, image_path
+        return answer, graph_viz_path, relations_table, summary, image_path
     except Exception as e:
         print(f"An error occurred in process_text: {str(e)}")
         import traceback
         traceback.print_exc()
-        return str(e), None, str(e), None
+        return str(e), None, None, str(e), None
 
 def ui_function(text, question):
-    answer, graph_viz_path, summary, image_path = process_text(text, question)
+    answer, graph_viz_path, relations_table, summary, image_path = process_text(text, question)
     if isinstance(answer, str) and answer.startswith("An error occurred"):
-        return answer, None, answer, None
-    return answer, graph_viz_path, summary, image_path
+        return answer, None, None, answer, None
+    return answer, graph_viz_path, relations_table, summary, image_path
+
+# Example text
+example_text = """The Apollo 11 mission, launched by NASA in July 1969, was the first manned mission to land on the Moon. Commanded by Neil Armstrong and piloted by Buzz Aldrin and Michael Collins, it was the culmination of the Space Race between the United States and the Soviet Union. On July 20, 1969, Armstrong and Aldrin became the first humans to set foot on the lunar surface, while Collins orbited above in the command module."""
 
 # Create Gradio interface
 iface = gr.Interface(
     fn=ui_function,
     inputs=[
-        gr.Textbox(label="Input Text"),
+        gr.Textbox(label="Input Text", value=example_text),
         gr.Textbox(label="Question")
     ],
     outputs=[
         gr.Textbox(label="Answer"),
         gr.Image(label="Graph Visualization", type="filepath"),
+        gr.Dataframe(label="Relations Table"),
         gr.Textbox(label="Summary"),
         gr.Image(label="Generated Image", type="filepath")
     ],