Update app.py

app.py

@@ -1,5 +1,4 @@
 import os
-import io
 import tempfile
 import streamlit as st
 import pandas as pd
@@ -7,50 +6,49 @@ from pypdf import PdfReader
 from pyvis.network import Network
 
 from knowledge_graph_maker import (
-    GraphMaker, Ontology, Document,
-    OpenAIClient, GroqClient
+    GraphMaker, Ontology, Document, OpenAIClient
 )
 
-st.set_page_config(page_title="Knowledge Graph Maker", layout="wide")
+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.title("Knowledge Graph from Text/PDF (Docker Space)")
-st.caption("Uses knowledge-graph-maker with OpenAI or Groq. Paste text or upload a PDF; view the interactive graph below.")
+# --- OpenRouter only ---
+OPENROUTER_API_KEY = os.getenv("OPENROUTER_API_KEY", "")
 
-# Choose LLM client based on available env vars
-OPENAI_API_KEY = os.getenv("OPENAI_API_KEY", "")
-GROQ_API_KEY = os.getenv("GROQ_API_KEY", "")
+# Model choices
+OPENROUTER_MODELS = [
+    "openai/gpt-oss-20b:free",
+    "moonshotai/kimi-k2:free",
+    "google/gemini-2.0-flash-exp:free",
+    "google/gemma-3-27b-it:free",
+]
 
+# Sidebar controls
 with st.sidebar:
-    st.subheader("Model Settings")
+    st.subheader("Model & Generation Settings")
+    model_choice = st.selectbox("OpenRouter model", OPENROUTER_MODELS, index=0)
+    custom_model = st.text_input("Custom model id (optional)", placeholder="e.g. meta-llama/llama-3.1-8b-instruct")
     temperature = st.slider("Temperature", 0.0, 1.0, 0.1, 0.05)
     top_p = st.slider("Top-p", 0.0, 1.0, 0.5, 0.05)
-    provider = st.radio("Provider", ["OpenAI", "Groq"], index=0 if OPENAI_API_KEY else 1 if GROQ_API_KEY else 0)
-    if provider == "OpenAI":
-        oai_model = st.text_input("OpenAI model", value="gpt-3.5-turbo")
-    else:
-        groq_model = st.text_input("Groq model", value="mixtral-8x7b-32768")
 
     st.markdown("### Ontology (labels)")
-    default_labels = [
-        {"Person": "Person name without adjectives (may appear as name or pronoun)"},
-        {"Object": "Avoid the definite article 'the' in name"},
-        {"Event": "Events involving multiple people; no verbs like gives/leaves"},
-        "Place", "Document", "Organisation", "Action",
-        {"Miscellanous": "Important concept that fits none of the above"}
-    ]
-    labels_text = st.text_area("Labels (JSON or comma-separated)", value=", ".join(
-        [lbl if isinstance(lbl, str) else list(lbl.keys())[0] for lbl in default_labels]
-    ), height=80)
-    st.markdown("### Relationships focus")
-    relationships_text = st.text_input("Relationships (comma-separated)", value="Relation between any pair of Entities")
+    labels_text = st.text_area(
+        "Comma-separated labels",
+        value="Person, Object, Event, Place, Document, Organisation, Action, Miscellanous",
+        height=70,
+    )
+    relationships_text = st.text_input(
+        "Relationships (comma-separated)",
+        value="Relation between any pair of Entities",
+    )
 
 def parse_labels(text):
-    # Allow simple "A, B, C" input; fall back to defaults above if empty
-    if not text.strip():
-        return [ "Person","Object","Event","Place","Document","Organisation","Action","Miscellanous" ]
-    return [lbl.strip() for lbl in text.split(",") if lbl.strip()]
+    return [lbl.strip() for lbl in text.split(",") if lbl.strip()] or [
+        "Person","Object","Event","Place","Document","Organisation","Action","Miscellanous"
+    ]
 
-def split_pdf(file) -> str:
+def pdf_to_text(file) -> str:
     reader = PdfReader(file)
     parts = []
     for page in reader.pages:
@@ -60,42 +58,31 @@ def split_pdf(file) -> str:
             continue
     return "\n".join(parts)
 
-def build_graph_documents(text: str) -> list[Document]:
-    # Simple chunking: ~900-1000 tokens ≈ ~3000-4000 chars heuristic
-    # Adjust if needed.
-    CHARS = 3500
+def chunk_text(text: str, chars: int = 3500) -> list[Document]:
     docs = []
-    for i in range(0, len(text), CHARS):
-        chunk = text[i:i+CHARS].strip()
+    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)
-    # Simple map to keep unique node IDs
     node_ids = {}
-
-    def node_key(label, name): return f"{label}:{name}"
-
     for e in edges:
-        n1 = node_key(e.node_1.label, e.node_1.name)
-        n2 = node_key(e.node_2.label, e.node_2.name)
-
+        n1 = f"{e.node_1.label}:{e.node_1.name}"
+        n2 = f"{e.node_2.label}:{e.node_2.name}"
         if n1 not in node_ids:
             net.add_node(n1, label=e.node_1.name, title=e.node_1.label)
             node_ids[n1] = True
         if n2 not in node_ids:
             net.add_node(n2, label=e.node_2.name, title=e.node_2.label)
             node_ids[n2] = True
-
-        rel = e.relationship or ""
-        net.add_edge(n1, n2, title=rel, value=1)
-
+        net.add_edge(n1, n2, title=e.relationship or "", value=1)
     net.toggle_physics(True)
     return net
 
-# Input UI
+# Input tabs
 tab_text, tab_pdf = st.tabs(["📝 Paste Text", "📄 Upload PDF"])
 input_text = ""
 with tab_text:
@@ -103,40 +90,46 @@ with tab_text:
 with tab_pdf:
     pdf_file = st.file_uploader("Upload a PDF", type=["pdf"])
     if pdf_file:
-        input_text = split_pdf(pdf_file)
+        input_text = pdf_to_text(pdf_file)
 
 if st.button("Generate Knowledge Graph", type="primary"):
     if not input_text.strip():
         st.warning("Please provide text or a PDF.")
         st.stop()
+    if not OPENROUTER_API_KEY:
+        st.error("OPENROUTER_API_KEY is not set in Space Secrets.")
+        st.stop()
 
-    # Prepare LLM client
-    if provider == "OpenAI":
-        if not OPENAI_API_KEY:
-            st.error("OPENAI_API_KEY is not set in the Space Secrets.")
-            st.stop()
-        llm = OpenAIClient(model=oai_model, temperature=temperature, top_p=top_p)
-    else:
-        if not GROQ_API_KEY:
-            st.error("GROQ_API_KEY is not set in the Space Secrets.")
-            st.stop()
-        llm = GroqClient(model=groq_model, temperature=temperature, top_p=top_p)
+    # Point OpenAI SDK at OpenRouter
+    os.environ["OPENAI_API_KEY"] = OPENROUTER_API_KEY
+    os.environ["OPENAI_BASE_URL"] = "https://openrouter.ai/api/v1"
+
+    # Optional: attribution headers for OpenRouter
+    os.environ["OPENAI_DEFAULT_HEADERS"] = (
+        '{"HTTP-Referer":"https://huggingface.co/spaces/blazingbunny/rahulnyk_knowledge_graph",'
+        '"X-Title":"Knowledge Graph (OpenRouter)"}'
+    )
+
+    selected_model = custom_model.strip() if custom_model.strip() else model_choice
 
     # Ontology
     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"]
+        relationships=[r.strip() for r in relationships_text.split(",") if r.strip()] or
+                      ["Relation between any pair of Entities"]
     )
 
-    st.info("Chunking input and building graph… this may take a bit for longer texts.")
+    st.info("Chunking input and building graph…")
+    docs = chunk_text(input_text)
+
+    # LLM client (OpenRouter through OpenAI-compatible client)
+    llm = OpenAIClient(model=selected_model, temperature=temperature, top_p=top_p)
 
     gm = GraphMaker(ontology=ontology, llm_client=llm, verbose=False)
-    docs = build_graph_documents(input_text)
+    edges = gm.from_documents(docs, delay_s_between=0)
 
-    edges = gm.from_documents(docs, delay_s_between=0)  # tune delay for rate limits
     st.success(f"Graph built with {len(edges)} edges.")
 
-    # Show edges 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,
@@ -144,13 +137,13 @@ if st.button("Generate Knowledge Graph", type="primary"):
     } for e in edges])
     st.dataframe(df, use_container_width=True)
 
-    # Render with PyVis inside Streamlit
     net = edges_to_pyvis(edges)
     with tempfile.TemporaryDirectory() as td:
         html_path = os.path.join(td, "graph.html")
         net.save_graph(html_path)
-        html_content = open(html_path, "r", encoding="utf-8").read()
-        st.components.v1.html(html_content, height=750, scrolling=True)
+        with open(html_path, "r", encoding="utf-8") as f:
+            html = f.read()
+        st.components.v1.html(html, height=750, scrolling=True)
 
 st.markdown("---")
-st.caption("Built with [knowledge-graph-maker](https://github.com/rahulnyk/knowledge_graph_maker).")
+st.caption("Powered by knowledge-graph-maker via OpenRouter.")