From fa47c61061503fd3684793cc5774cbed00d2df14 Mon Sep 17 00:00:00 2001
From: blueprintparadise <rhiray03@gmail.com>
Date: Mon, 28 Apr 2025 00:53:38 +0100
Subject: [PATCH] added the download image button

Added an option to download the graph as a png for future reference,
---
 app.py               | 52 ++++++++++++++++++++++++++++++++++++++------
 requirements.txt     |  3 ++-
 templates/index.html | 13 ++++++++++-
 3 files changed, 59 insertions(+), 9 deletions(-)

diff --git a/app.py b/app.py
index 172aa74..d318963 100644
--- a/app.py
+++ b/app.py
@@ -1,4 +1,4 @@
-from flask import Flask, render_template, request, jsonify, Response, stream_with_context
+from flask import Flask, render_template, request, jsonify, Response
 import requests
 import json
 import time
@@ -7,12 +7,13 @@
 import numpy as np
 from sklearn.metrics.pairwise import cosine_similarity
 from annoy import AnnoyIndex
-import os
-import networkx as nx
-import heapq
 
+import heapq
+import matplotlib.pyplot as plt
+from io import BytesIO
+import networkx as nx
 app = Flask(__name__)
-
+latest_graph_data = {"nodes": [], "edges": []}
 # Function to get embeddings from the API
 def get_embedding(text):
     headers = {'Content-Type': 'application/json'}
@@ -40,6 +41,8 @@ def create_database():
     conn.commit()
     return conn
 
+
+
 def insert_data(conn, text, embedding, is_question):
     c = conn.cursor()
     c.execute("INSERT INTO embeddings (text, embedding, is_question) VALUES (?, ?, ?)",
@@ -314,6 +317,10 @@ def calculate_top_similarities(embeddings, current_step, top_k=2):
             graph_data['nodes'][-1]['value'] = 20  # Set a default size if no connections
 
         serialized_graph_data = serialize_graph_data(graph_data)
+
+        global latest_graph_data
+        latest_graph_data = serialized_graph_data
+
         strongest_path, path_weights, avg_similarity = calculate_strongest_path(serialized_graph_data, step_count)
         
         path_data = {
@@ -401,7 +408,7 @@ def calculate_top_similarities(embeddings, current_step, top_k=2):
         'label': f"Final Answer: {get_short_title(final_answer)}"
     })
     
-    # Calculate similarities with previous steps for the final answer
+
     top_similarities = calculate_top_similarities(embeddings + [final_embedding], step_count - 1, top_k=2)
     
     for prev_step, similarity in top_similarities:
@@ -501,6 +508,37 @@ def check_consistency(final_answer, evaluation):
     #print("Failed to get a valid consistency check after 5 attempts. Defaulting to inconsistent.")
     #return False
     return True
-
+@app.route("/export/image")
+def export_png():
+    global latest_graph_data
+    data = latest_graph_data
+
+    # Build a NetworkX graph
+    G = nx.DiGraph()
+    for n in data.get("nodes", []):
+        G.add_node(n["id"], label=n["label"])
+    for e in data.get("edges", []):
+        G.add_edge(e["from"], e["to"], weight=e["value"])
+
+
+    pos = nx.spring_layout(G, k=0.5, seed=42)
+
+    # Draw with Matplotlib
+    fig, ax = plt.subplots(figsize=(8, 6), dpi=150)
+    nx.draw(G, pos, ax=ax, with_labels=False, node_color="#3182bd",
+            edge_color="#bbb", width=1.2, node_size=800)
+    labels = {i: d["label"] for i, d in G.nodes(data=True)}
+    nx.draw_networkx_labels(G, pos, labels, font_size=8)
+    ax.axis("off")
+    buf = BytesIO()
+    fig.savefig(buf, format="png", bbox_inches="tight")
+    plt.close(fig)
+    buf.seek(0)
+
+    return Response(buf.read(),
+                    mimetype="image/png",
+                    headers={
+                        "Content-Disposition": "attachment; filename=graph.png"
+                    })
 if __name__ == '__main__':
     app.run(host='0.0.0.0', port=5100, debug=True)
\ No newline at end of file
diff --git a/requirements.txt b/requirements.txt
index 1404439..d38ef95 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -3,4 +3,5 @@ numpy
 scikit-learn
 annoy
 networkx
-requests
\ No newline at end of file
+requests
+matplotlib
\ No newline at end of file
diff --git a/templates/index.html b/templates/index.html
index 922b9a9..f2c600c 100644
--- a/templates/index.html
+++ b/templates/index.html
@@ -6,12 +6,14 @@
     <title>Local Llama Knowledge Graph</title>
     <script src="https://cdn.jsdelivr.net/npm/marked/marked.min.js"></script>
     <script src="https://unpkg.com/vis-network/standalone/umd/vis-network.min.js"></script>
+    <script src="https://cdnjs.cloudflare.com/ajax/libs/FileSaver.js/2.0.5/FileSaver.min.js"></script>
     <style>
         body { font-family: Arial, sans-serif; line-height: 1.6; padding: 20px; }
         #query { width: 100%; padding: 10px; margin-bottom: 10px; }
         #submit { padding: 10px 20px; }
         .step, .final-answer, .similar-item { margin-bottom: 20px; border: 1px solid #ddd; padding: 10px; }
         .step h3, .final-answer h3, .similar-item h3 { margin-top: 0; }
+
         #container { display: flex; }
         #left-panel { flex: 1; margin-right: 20px; }
         #graph { flex: 1; height: 600px; border: 1px solid #ddd; }
@@ -48,7 +50,7 @@ <h1>Local Llama Knowledge Graph</h1>
     
     <input type="text" id="query" placeholder="e.g., What is the capital of France?">
     <button id="submit">Submit</button>
-
+    <button id="download-img">Download PNG</button>
     <div id="container">
         <div id="left-panel">
             <div id="response"></div>
@@ -63,6 +65,8 @@ <h1>Local Llama Knowledge Graph</h1>
         const response = document.getElementById('response');
         const similar = document.getElementById('similar');
         const graphContainer = document.getElementById('graph');
+        const downloadBtn = document.getElementById('download-img');
+
 
         let network;
         let nodes = new vis.DataSet();
@@ -243,6 +247,13 @@ <h3>Final Answer</h3>
             pathHtml += `</ul><p>Overall Path Similarity (Weighted Average): ${avgSimilarity.toFixed(4)}</p>`;
             return pathHtml;
         }
+
+        downloadBtn.addEventListener('click', () => {
+  window.location = '/export/image';
+});
+
+
     </script>
+
 </body>
 </html>
\ No newline at end of file