| import streamlit as st |
| import os |
| import tempfile |
| from uuid import uuid4 |
| from concurrent.futures import ThreadPoolExecutor, Future |
|
|
| from ingestion import ingest_pdf |
| from dotenv import load_dotenv, find_dotenv |
|
|
| load_dotenv(find_dotenv(), override=False) |
| st.set_page_config(page_title="PDF Ingestion", layout="wide") |
|
|
| if "session_id" not in st.session_state: |
| st.session_state.session_id = str(uuid4()) |
| if "ingestion_future" not in st.session_state: |
| st.session_state.ingestion_future = None |
| if "last_output" not in st.session_state: |
| st.session_state.last_output = None |
|
|
| st.title("Ingestion Layer") |
| st.caption("Upload a PDF to extract, normalize, chunk, and persist for later embedding.") |
|
|
| with st.sidebar: |
| st.header("Settings") |
| parser = st.selectbox("Parser", ["pypdf", "pdfplumber"], index=0) |
| chunk_size = st.number_input("Chunk size (chars)", min_value=500, max_value=10000, value=4000, step=100) |
| overlap = st.number_input("Overlap (chars)", min_value=0, max_value=2000, value=400, step=50) |
|
|
| uploaded = st.file_uploader("Upload PDF", type=["pdf"]) |
|
|
| def _run_ingestion(tmp_path: str, sid: str, parser: str, chunk_size: int, overlap: int): |
| return ingest_pdf(tmp_path, sid, parser=parser, chunk_size=chunk_size, overlap=overlap) |
|
|
| col1, col2 = st.columns([2, 1]) |
| with col1: |
| if uploaded is not None and st.button("Start Ingestion", type="primary"): |
| with tempfile.NamedTemporaryFile(delete=False, suffix=".pdf") as tmp: |
| tmp.write(uploaded.read()) |
| tmp_path = tmp.name |
|
|
| executor = ThreadPoolExecutor(max_workers=1) |
| future: Future = executor.submit(_run_ingestion, tmp_path, st.session_state.session_id, parser, int(chunk_size), int(overlap)) |
| st.session_state.ingestion_future = future |
| st.info("Ingestion started. Processing in background...") |
|
|
| future: Future = st.session_state.ingestion_future |
| if future is not None: |
| if future.done(): |
| try: |
| chunks, out_path = future.result() |
| st.session_state.last_output = (len(chunks), out_path) |
| st.success(f"Ingestion complete: {len(chunks)} chunks written to {out_path}") |
| if chunks: |
| st.subheader("Preview (first chunk)") |
| st.code(chunks[0].content[:2000]) |
| except Exception as e: |
| st.error(f"Ingestion failed: {e}") |
| finally: |
| st.session_state.ingestion_future = None |
| else: |
| st.status("Processing...", state="running") |
|
|
| with col2: |
| st.subheader("Session") |
| st.code(st.session_state.session_id) |
| if st.session_state.last_output: |
| cnt, path = st.session_state.last_output |
| st.metric("Chunks", cnt) |
| st.write("Output:") |
| st.code(path) |
| if uploaded is not None and st.button("Retry with fallback parser"): |
| fb = "pdfplumber" if parser == "pypdf" else "pypdf" |
| with tempfile.NamedTemporaryFile(delete=False, suffix=".pdf") as tmp: |
| tmp.write(uploaded.read()) |
| tmp_path = tmp.name |
| executor = ThreadPoolExecutor(max_workers=1) |
| future: Future = executor.submit(_run_ingestion, tmp_path, st.session_state.session_id, fb, int(chunk_size), int(overlap)) |
| st.session_state.ingestion_future = future |
| st.info(f"Retrying with {fb}...") |
|
|
