Yogesh18018/knowledge-graph-engine
1
1"""
2GraphMind — Knowledge Graph Construction & Reasoning Engine
3============================================================
4Main Streamlit application.
5"""
6
7import streamlit as st
8import streamlit.components.v1 as components
9import pandas as pd
10
11from src.extractor import EntityExtractor
12from src.graph_builder import KnowledgeGraph, ENTITY_COLORS
13from src.visualizer import (
14 create_pyvis_graph,
15 graph_stats_chart,
16 centrality_chart,
17 community_chart,
18)
19from src.sample_texts import SAMPLE_TEXTS
20
21# ======================================================================
22# Page configuration
23# ======================================================================
24
25st.set_page_config(
26 page_title="GraphMind | Knowledge Graph",
27 page_icon="G",
28 layout="wide",
29 initial_sidebar_state="expanded",
30)
31
32# ======================================================================
33# Custom CSS — dark theme with accent colours
34# ======================================================================
35
36st.markdown(
37 """
38 <style>
39 /* ---- Global ---- */
40 .stApp {
41 background-color: #0a0a0a;
42 color: #e0e0e0;
43 }
44
45 /* ---- Sidebar ---- */
46 section[data-testid="stSidebar"] {
47 background-color: #111111;
48 border-right: 1px solid #1e1e1e;
49 }
50
51 /* ---- Headers ---- */
52 h1, h2, h3, h4 {
53 color: #ffffff !important;
54 }
55
56 /* ---- Metric cards ---- */
57 div[data-testid="stMetric"] {
58 background: linear-gradient(135deg, #111111 0%, #1a1a2e 100%);
59 border: 1px solid #1e1e1e;
60 border-radius: 12px;
61 padding: 16px 20px;
62 }
63 div[data-testid="stMetric"] label {
64 color: #888888 !important;
65 }
66 div[data-testid="stMetric"] div[data-testid="stMetricValue"] {
67 color: #00ff88 !important;
68 font-weight: 700;
69 }
70
71 /* ---- Buttons ---- */
72 .stButton > button {
73 background: linear-gradient(135deg, #00ff88 0%, #00d4ff 100%);
74 color: #0a0a0a;
75 border: none;
76 border-radius: 8px;
77 font-weight: 700;
78 padding: 0.5rem 1.5rem;
79 transition: all 0.3s ease;
80 }
81 .stButton > button:hover {
82 transform: translateY(-2px);
83 box-shadow: 0 4px 20px rgba(0,255,136,0.3);
84 }
85
86 /* ---- Tabs ---- */
87 .stTabs [data-baseweb="tab-list"] {
88 gap: 8px;
89 }
90 .stTabs [data-baseweb="tab"] {
91 background-color: #1a1a1a;
92 border-radius: 8px 8px 0 0;
93 color: #888888;
94 padding: 8px 20px;
95 }
96 .stTabs [aria-selected="true"] {
97 background-color: #1e1e2e;
98 color: #00ff88 !important;
99 }
100
101 /* ---- DataFrame ---- */
102 .stDataFrame {
103 border: 1px solid #1e1e1e;
104 border-radius: 8px;
105 }
106
107 /* ---- Expanders ---- */
108 .streamlit-expanderHeader {
109 background-color: #111111;
110 border-radius: 8px;
111 }
112
113 /* ---- Success / info banners ---- */
114 .stAlert {
115 background-color: #111111;
116 border: 1px solid #1e1e1e;
117 border-radius: 8px;
118 }
119
120 /* ---- Accent text helpers ---- */
121 .accent-green { color: #00ff88; font-weight: 700; }
122 .accent-blue { color: #00d4ff; font-weight: 700; }
123
124 /* ---- Legend colour pills ---- */
125 .legend-pill {
126 display: inline-block;
127 padding: 3px 12px;
128 border-radius: 20px;
129 margin: 2px 4px;
130 font-size: 0.82rem;
131 font-weight: 600;
132 color: #0a0a0a;
133 }
134
135 /* ---- Divider ---- */
136 hr {
137 border-color: #1e1e1e;
138 }
139 </style>
140 """,
141 unsafe_allow_html=True,
142)
143
144
145# ======================================================================
146# Sidebar
147# ======================================================================
148
149with st.sidebar:
150 st.markdown("## GraphMind")
151 st.markdown(
152 "<span class='accent-green'>Knowledge Graph</span> "
153 "<span class='accent-blue'>Construction & Reasoning</span>",
154 unsafe_allow_html=True,
155 )
156 st.markdown("---")
157
158 # --- Input source ---
159 st.markdown("### Text Source")
160 input_mode = st.radio(
161 "Choose input method",
162 ["Demo Texts", "Paste Your Own"],
163 label_visibility="collapsed",
164 )
165
166 text_to_process = ""
167
168 if input_mode == "Demo Texts":
169 selected_demo = st.selectbox(
170 "Select a demo text",
171 list(SAMPLE_TEXTS.keys()),
172 )
173 text_to_process = SAMPLE_TEXTS[selected_demo]
174 with st.expander("Preview text", expanded=False):
175 st.caption(text_to_process[:500] + "…")
176 else:
177 text_to_process = st.text_area(
178 "Paste your text below",
179 height=250,
180 placeholder="Enter text containing named entities…",
181 )
182
183 st.markdown("---")
184
185 # --- Extraction settings ---
186 st.markdown("### Extraction Settings")
187 entity_types = st.multiselect(
188 "Entity types to extract",
189 ["PERSON", "ORG", "LOCATION", "DATE", "TECHNOLOGY"],
190 default=["PERSON", "ORG", "LOCATION", "DATE", "TECHNOLOGY"],
191 )
192
193 min_mentions = st.slider(
194 "Minimum mentions for nodes",
195 min_value=1,
196 max_value=5,
197 value=1,
198 help="Only show entities mentioned at least this many times.",
199 )
200
201 st.markdown("---")
202
203 # --- Build button ---
204 build_clicked = st.button(" Build Knowledge Graph", use_container_width=True)
205
206 st.markdown("---")
207 st.markdown(
208 "<div style='text-align:center;color:#555;font-size:0.75rem;'>"
209 "Built by <b>Yogesh Kuchimanchi</b><br>MIT License</div>",
210 unsafe_allow_html=True,
211 )
212
213
214# ======================================================================
215# Main area — Header
216# ======================================================================
217
218st.markdown(
219 "<h1 style='text-align:center;'>"
220 " Graph<span class='accent-green'>Mind</span></h1>",
221 unsafe_allow_html=True,
222)
223st.markdown(
224 "<p style='text-align:center;color:#888;margin-top:-10px;'>"
225 "Construct knowledge graphs from unstructured text using rule-based NER "
226 "and graph reasoning.</p>",
227 unsafe_allow_html=True,
228)
229
230# Colour legend
231legend_html = " ".join(
232 f"<span class='legend-pill' style='background:{color};'>{label}</span>"
233 for label, color in ENTITY_COLORS.items()
234)
235st.markdown(
236 f"<div style='text-align:center;margin-bottom:20px;'>{legend_html}</div>",
237 unsafe_allow_html=True,
238)
239
240
241# ======================================================================
242# Processing pipeline
243# ======================================================================
244
245@st.cache_data(show_spinner=False)
246def run_pipeline(text: str, types: tuple):
247 """Run NER + graph construction and cache results."""
248 extractor = EntityExtractor()
249 entities = extractor.extract(text)
250
251 # Filter entity types
252 entities = [e for e in entities if e["label"] in types]
253
254 relationships = extractor.extract_relationships(text, entities)
255
256 kg = KnowledgeGraph()
257 kg.add_entities(entities)
258 kg.add_relationships(relationships)
259
260 stats = kg.get_stats()
261 graph_html = create_pyvis_graph(kg)
262
263 return entities, relationships, kg, stats, graph_html
264
265
266# ======================================================================
267# Run on button click OR first load with demo text
268# ======================================================================
269
270if "has_run" not in st.session_state:
271 st.session_state.has_run = False
272
273if build_clicked and text_to_process.strip():
274 st.session_state.has_run = True
275 st.session_state.text = text_to_process
276 st.session_state.types = tuple(entity_types)
277
278# Auto-run on first visit with demo text
279if not st.session_state.has_run and input_mode == "Demo Texts":
280 st.session_state.has_run = True
281 st.session_state.text = text_to_process
282 st.session_state.types = tuple(entity_types)
283
284if st.session_state.has_run:
285 with st.spinner("Extracting entities and building graph…"):
286 entities, relationships, kg, stats, graph_html = run_pipeline(
287 st.session_state.text, st.session_state.types
288 )
289
290 # ==================================================================
291 # Metrics row
292 # ==================================================================
293 m1, m2, m3, m4 = st.columns(4)
294 m1.metric("Total Nodes", stats["total_nodes"])
295 m2.metric("Total Edges", stats["total_edges"])
296 m3.metric("Communities", stats["num_communities"])
297 m4.metric("Entity Types", len(stats["entity_type_counts"]))
298
299 st.markdown("---")
300
301 # ==================================================================
302 # Tabs
303 # ==================================================================
304 tab_graph, tab_entities, tab_relations, tab_stats = st.tabs(
305 [" Interactive Graph", " Entities", " Relationships", " Statistics"]
306 )
307
308 # --- Interactive Graph ---
309 with tab_graph:
310 st.markdown("#### Interactive Knowledge Graph")
311 st.caption("Drag, zoom, and hover nodes for details.")
312 components.html(graph_html, height=680, scrolling=False)
313
314 # --- Entities table ---
315 with tab_entities:
316 st.markdown("#### Extracted Entities")
317 if entities:
318 df_ent = pd.DataFrame(entities)
319 df_ent = df_ent[["text", "label", "start", "end"]]
320 df_ent.columns = ["Entity", "Type", "Start", "End"]
321
322 # Colour-coded type column
323 st.dataframe(
324 df_ent.style.apply(
325 lambda row: [
326 "",
327 f"color: {ENTITY_COLORS.get(row['Type'], '#888')}",
328 "",
329 "",
330 ],
331 axis=1,
332 ),
333 use_container_width=True,
334 height=450,
335 )
336 st.caption(f"Total: **{len(entities)}** entities extracted.")
337 else:
338 st.info("No entities found. Try different text or settings.")
339
340 # --- Relationships table ---
341 with tab_relations:
342 st.markdown("#### Extracted Relationships")
343 if relationships:
344 df_rel = pd.DataFrame(relationships)
345 df_rel = df_rel[["source", "relation", "target", "source_label", "target_label"]]
346 df_rel.columns = ["Source", "Relation", "Target", "Src Type", "Tgt Type"]
347 st.dataframe(df_rel, use_container_width=True, height=450)
348 st.caption(f"Total: **{len(relationships)}** relationships inferred.")
349 else:
350 st.info("No relationships found.")
351
352 # --- Statistics ---
353 with tab_stats:
354 st.markdown("#### Graph Analytics")
355
356 col_left, col_right = st.columns(2)
357
358 with col_left:
359 fig_dist = graph_stats_chart(stats)
360 st.plotly_chart(fig_dist, use_container_width=True)
361
362 with col_right:
363 fig_community = community_chart(stats["communities"])
364 st.plotly_chart(fig_community, use_container_width=True)
365
366 st.markdown("---")
367 fig_central = centrality_chart(stats["top_central_nodes"])
368 st.plotly_chart(fig_central, use_container_width=True)
369
370 with st.expander("Community Details"):
371 for i, comm in enumerate(stats["communities"]):
372 st.markdown(
373 f"**Community {i+1}** ({len(comm)} members): "
374 + ", ".join(comm)
375 )
376
377 with st.expander("Raw Statistics"):
378 st.json(
379 {
380 "density": round(stats["density"], 6),
381 "total_nodes": stats["total_nodes"],
382 "total_edges": stats["total_edges"],
383 "entity_type_counts": stats["entity_type_counts"],
384 "relation_type_counts": stats["relation_type_counts"],
385 "num_communities": stats["num_communities"],
386 }
387 )
388
389else:
390 # Placeholder when nothing has been processed yet
391 st.markdown(
392 "<div style='text-align:center;padding:80px 20px;color:#555;'>"
393 "<h3>Paste text or select a demo, then click "
394 "<span class='accent-green'>Build Knowledge Graph</span></h3>"
395 "<p>The engine will extract entities, infer relationships, "
396 "and visualise an interactive knowledge graph.</p>"
397 "</div>",
398 unsafe_allow_html=True,
399 )
400 