Update app.py

app.py

@@ -1,5 +1,4 @@
 import os
-import tempfile
 import streamlit as st
 import pandas as pd
 from pypdf import PdfReader
@@ -9,14 +8,15 @@ from knowledge_graph_maker import (
     GraphMaker, Ontology, Document, OpenAIClient
 )
 
+# ── Page setup ──────────────────────────────────────────────────────────────────
 st.set_page_config(page_title="Knowledge Graph (OpenRouter)", layout="wide")
 st.title("Knowledge Graph from Text/PDF — OpenRouter")
-st.caption("Uses knowledge-graph-maker via OpenRouter. Paste text or upload a PDF; choose an OpenRouter model.")
+st.caption("Builds a knowledge graph with knowledge-graph-maker via OpenRouter. Paste text or upload a PDF; choose a model.")
 
-# --- OpenRouter only ---
+# ── Secrets / env ───────────────────────────────────────────────────────────────
 OPENROUTER_API_KEY = os.getenv("OPENROUTER_API_KEY", "")
 
-# Model choices
+# Preset OpenRouter models (you can add more)
 OPENROUTER_MODELS = [
     "openai/gpt-oss-20b:free",
     "moonshotai/kimi-k2:free",
@@ -24,7 +24,7 @@ OPENROUTER_MODELS = [
     "google/gemma-3-27b-it:free",
 ]
 
-# Sidebar controls
+# ── Sidebar controls ───────────────────────────────────────────────────────────
 with st.sidebar:
     st.subheader("Model & Generation Settings")
     model_choice = st.selectbox("OpenRouter model", OPENROUTER_MODELS, index=0)
@@ -43,9 +43,10 @@ with st.sidebar:
         value="Relation between any pair of Entities",
     )
 
-def parse_labels(text):
+# ── Helpers ────────────────────────────────────────────────────────────────────
+def parse_labels(text: str):
     return [lbl.strip() for lbl in text.split(",") if lbl.strip()] or [
-        "Person","Object","Event","Place","Document","Organisation","Action","Miscellanous"
+        "Person", "Object", "Event", "Place", "Document", "Organisation", "Action", "Miscellanous"
     ]
 
 def pdf_to_text(file) -> str:
@@ -63,11 +64,20 @@ def chunk_text(text: str, chars: int = 3500) -> list[Document]:
     for i in range(0, len(text), chars):
         chunk = text[i:i+chars].strip()
         if chunk:
-            docs.append(Document(text=chunk, metadata={"chunk_id": i//chars}))
+            docs.append(Document(text=chunk, metadata={"chunk_id": i // chars}))
     return docs
 
 def edges_to_pyvis(edges):
-    net = Network(height="700px", width="100%", bgcolor="#ffffff", font_color="#222222", notebook=False, directed=False)
+    # IMPORTANT: cdn_resources="in_line" prevents PyVis from creating a ./lib folder
+    net = Network(
+        height="700px",
+        width="100%",
+        bgcolor="#ffffff",
+        font_color="#222222",
+        notebook=False,
+        directed=False,
+        cdn_resources="in_line",
+    )
     node_ids = {}
     for e in edges:
         n1 = f"{e.node_1.label}:{e.node_1.name}"
@@ -82,7 +92,7 @@ def edges_to_pyvis(edges):
     net.toggle_physics(True)
     return net
 
-# Input tabs
+# ── Input tabs ─────────────────────────────────────────────────────────────────
 tab_text, tab_pdf = st.tabs(["📝 Paste Text", "📄 Upload PDF"])
 input_text = ""
 with tab_text:
@@ -92,6 +102,7 @@ with tab_pdf:
     if pdf_file:
         input_text = pdf_to_text(pdf_file)
 
+# ── Action ─────────────────────────────────────────────────────────────────────
 if st.button("Generate Knowledge Graph", type="primary"):
     if not input_text.strip():
         st.warning("Please provide text or a PDF.")
@@ -100,11 +111,10 @@ if st.button("Generate Knowledge Graph", type="primary"):
         st.error("OPENROUTER_API_KEY is not set in Space Secrets.")
         st.stop()
 
-    # Point OpenAI SDK at OpenRouter
+    # Route OpenAI SDK traffic through OpenRouter (OpenAI-compatible)
     os.environ["OPENAI_API_KEY"] = OPENROUTER_API_KEY
     os.environ["OPENAI_BASE_URL"] = "https://openrouter.ai/api/v1"
-
-    # Optional: attribution headers for OpenRouter
+    # Optional attribution headers for OpenRouter analytics/ranking
     os.environ["OPENAI_DEFAULT_HEADERS"] = (
         '{"HTTP-Referer":"https://huggingface.co/spaces/blazingbunny/rahulnyk_knowledge_graph",'
         '"X-Title":"Knowledge Graph (OpenRouter)"}'
@@ -116,20 +126,21 @@ if st.button("Generate Knowledge Graph", type="primary"):
     ontology = Ontology(
         labels=parse_labels(labels_text),
         relationships=[r.strip() for r in relationships_text.split(",") if r.strip()] or
-                      ["Relation between any pair of Entities"]
+                      ["Relation between any pair of Entities"],
     )
 
     st.info("Chunking input and building graph…")
     docs = chunk_text(input_text)
 
-    # LLM client (OpenRouter through OpenAI-compatible client)
+    # LLM client (OpenRouter via OpenAI client)
     llm = OpenAIClient(model=selected_model, temperature=temperature, top_p=top_p)
 
     gm = GraphMaker(ontology=ontology, llm_client=llm, verbose=False)
-    edges = gm.from_documents(docs, delay_s_between=0)
+    edges = gm.from_documents(docs, delay_s_between=0)  # tweak delay for rate limits if needed
 
     st.success(f"Graph built with {len(edges)} edges.")
 
+    # Show edge table
     df = pd.DataFrame([{
         "node_1_label": e.node_1.label, "node_1": e.node_1.name,
         "node_2_label": e.node_2.label, "node_2": e.node_2.name,
@@ -137,13 +148,10 @@ if st.button("Generate Knowledge Graph", type="primary"):
     } for e in edges])
     st.dataframe(df, use_container_width=True)
 
+    # Render graph in-memory (no writes)
     net = edges_to_pyvis(edges)
-    with tempfile.TemporaryDirectory() as td:
-        html_path = os.path.join(td, "graph.html")
-        net.save_graph(html_path)
-        with open(html_path, "r", encoding="utf-8") as f:
-            html = f.read()
-        st.components.v1.html(html, height=750, scrolling=True)
+    html = net.generate_html()      # <- avoids creating ./lib
+    st.components.v1.html(html, height=750, scrolling=True)
 
 st.markdown("---")
 st.caption("Powered by knowledge-graph-maker via OpenRouter.")