| from pathlib import Path | |
| import streamlit as st | |
| from src.pipeline import ingest, query | |
| def render_ui(): | |
| st.title("Neo4j GraphRAG Engine") | |
| st.markdown( | |
| "Ingest documents, extract entities and relationships into Neo4j, " | |
| "embed into Pinecone, and query via hybrid retrieval (vector + graph)." | |
| ) | |
| tab1, tab2 = st.tabs(["Query", "Ingest Documents"]) | |
| with tab1: | |
| q = st.text_input("Ask a question about your documents:", key="query_input") | |
| if st.button("Search", type="primary", disabled=not q): | |
| with st.spinner("Running hybrid retrieval..."): | |
| try: | |
| result = query(q) | |
| st.session_state["query_result"] = result | |
| except Exception as e: | |
| st.error(f"Query failed: {e}") | |
| if "query_result" in st.session_state: | |
| r = st.session_state["query_result"] | |
| st.subheader("Answer") | |
| st.write(r.answer) | |
| with st.expander("Vector Context (top chunks)"): | |
| for i, ctx in enumerate(r.vector_context[:5]): | |
| st.caption(f"Chunk {i + 1}") | |
| st.text(ctx[:500]) | |
| with st.expander("Graph Context (connected entities)"): | |
| for ctx in r.graph_context[:20]: | |
| st.write(f"- {ctx}") | |
| with tab2: | |
| docs_path = st.text_input("Documents directory", value="data/sample_docs") | |
| if st.button("Ingest", type="secondary"): | |
| _run_ingestion(docs_path) | |
| _render_sample_docs() | |
| def _render_sample_docs(): | |
| docs_dir = Path("data/sample_docs") | |
| if not docs_dir.exists(): | |
| return | |
| md_files = sorted(docs_dir.glob("*.md")) | |
| txt_files = sorted(docs_dir.glob("*.txt")) | |
| files = md_files + txt_files | |
| if not files: | |
| return | |
| st.divider() | |
| st.subheader("Source Documents") | |
| for f in files: | |
| content = f.read_text(encoding="utf-8") | |
| escaped = content.replace("&", "&").replace("<", "<").replace(">", ">") | |
| html = ( | |
| f"<div style='border:1px solid #e2e8f0;border-radius:8px;" | |
| f"padding:12px 16px;margin-bottom:8px;max-height:200px;" | |
| f"overflow-y:auto;font-size:13px;'>" | |
| f"<b>{f.name}</b><br>{escaped}</div>" | |
| ) | |
| st.markdown(html, unsafe_allow_html=True) | |
| def _run_ingestion(docs_path: str): | |
| with st.status("Preparing ingestion...", expanded=True) as status: | |
| def on_msg(msg: str): | |
| if msg.startswith("extracting:"): | |
| status.update(label=msg, state="running") | |
| else: | |
| status.update(label=msg, state="running") | |
| try: | |
| ingest(docs_path, status_callback=on_msg) | |
| status.update(label="Ingestion complete", state="complete") | |
| st.success("All data ingested successfully.") | |
| except Exception as e: | |
| status.update(label=f"Ingestion failed: {e}", state="error") | |
| st.error(f"Ingestion failed: {e}") | |