from __future__ import annotations import pandas as pd import streamlit as st from rag_chunk_visualizer.core.config import settings from rag_chunk_visualizer.services.chunking_service import ( ChunkingParameterError, build_chunks_from_documents, ) from rag_chunk_visualizer.services.document_service import process_uploaded_files from rag_chunk_visualizer.services.embedding_service import ( EmbeddingError, embed_chunks, embed_query_text, summarize_embedding_matrix, ) from rag_chunk_visualizer.services.llm_providers import ( PROVIDERS, get_default_model, get_env_api_key, ) from rag_chunk_visualizer.services.rag_service import ( RAGGenerationError, build_grounded_prompt, generate_grounded_answer, ) from rag_chunk_visualizer.services.retrieval_service import ( RetrievalError, enrich_query_results, get_retrieved_chunk_ids, validate_query_text, ) from rag_chunk_visualizer.storage.lancedb_store import ( VectorStoreError, fetch_table_preview, search_similar_chunks, write_chunk_embeddings, ) from rag_chunk_visualizer.visualization.plotting import ( PlotBuildError, create_embedding_scatter_plot, ) from rag_chunk_visualizer.visualization.projection import ( ProjectionError, project_chunks, transform_query_vector, ) def toast_success(message: str) -> None: st.toast(message, icon="✅") def reset_downstream_state() -> None: st.session_state["embedding_errors"] = [] st.session_state["vector_store_errors"] = [] st.session_state["projection_errors"] = [] st.session_state["retrieval_errors"] = [] st.session_state["generation_errors"] = [] st.session_state["embedding_matrix"] = None st.session_state["embedding_summary"] = None st.session_state["query_embedding"] = None st.session_state["vector_store_summary"] = None st.session_state["vector_search_results"] = [] st.session_state["projection_df"] = None st.session_state["projection_summary"] = None st.session_state["projection_model"] = None st.session_state["query_projection"] = None st.session_state["grounded_prompt"] = None st.session_state["retrieval_results"] = [] st.session_state["generated_answer"] = None def handle_document_uploads(uploaded_files: list) -> None: if not uploaded_files: st.session_state["document_errors"] = ["No files selected."] st.session_state["chunk_errors"] = [] reset_downstream_state() st.session_state["app_status"] = "No files uploaded" return processed_documents, errors = process_uploaded_files(uploaded_files) st.session_state["raw_documents"] = processed_documents st.session_state["document_errors"] = errors st.session_state["chunk_errors"] = [] st.session_state["selected_document_id"] = ( processed_documents[0]["doc_id"] if processed_documents else None ) st.session_state["chunks"] = [] st.session_state["selected_chunk_id"] = None reset_downstream_state() if processed_documents: st.session_state["app_status"] = f"Loaded {len(processed_documents)} document(s)" toast_success(f"Loaded {len(processed_documents)} document(s)") else: st.session_state["app_status"] = "Upload failed" def handle_chunk_build(chunk_size: int, chunk_overlap: int) -> None: raw_documents = st.session_state["raw_documents"] if not raw_documents: st.session_state["chunk_errors"] = ["Upload at least one document before chunking."] st.session_state["app_status"] = "Chunking failed" return try: chunks = build_chunks_from_documents( raw_documents=raw_documents, chunk_size=chunk_size, chunk_overlap=chunk_overlap, ) except ChunkingParameterError as exc: st.session_state["chunk_errors"] = [str(exc)] st.session_state["chunks"] = [] st.session_state["selected_chunk_id"] = None reset_downstream_state() st.session_state["app_status"] = "Chunking failed" return st.session_state["chunk_errors"] = [] st.session_state["chunks"] = chunks st.session_state["selected_chunk_id"] = chunks[0]["chunk_id"] if chunks else None reset_downstream_state() if chunks: st.session_state["app_status"] = f"Built {len(chunks)} chunk(s)" toast_success(f"Built {len(chunks)} chunk(s)") else: st.session_state["app_status"] = "No chunks produced" def handle_embedding_build() -> None: chunks = st.session_state["chunks"] if not chunks: st.session_state["embedding_errors"] = ["Build chunks before generating embeddings."] st.session_state["app_status"] = "Embedding failed" return try: matrix = embed_chunks(chunks) summary = summarize_embedding_matrix(matrix) except EmbeddingError as exc: st.session_state["embedding_errors"] = [str(exc)] st.session_state["embedding_matrix"] = None st.session_state["embedding_summary"] = None st.session_state["query_embedding"] = None st.session_state["vector_store_summary"] = None st.session_state["vector_store_errors"] = [] st.session_state["vector_search_results"] = [] st.session_state["projection_df"] = None st.session_state["projection_summary"] = None st.session_state["projection_model"] = None st.session_state["query_projection"] = None st.session_state["projection_errors"] = [] st.session_state["retrieval_errors"] = [] st.session_state["generation_errors"] = [] st.session_state["grounded_prompt"] = None st.session_state["retrieval_results"] = [] st.session_state["generated_answer"] = None st.session_state["app_status"] = "Embedding failed" return st.session_state["embedding_errors"] = [] st.session_state["vector_store_errors"] = [] st.session_state["projection_errors"] = [] st.session_state["retrieval_errors"] = [] st.session_state["generation_errors"] = [] st.session_state["embedding_matrix"] = matrix st.session_state["embedding_summary"] = summary st.session_state["query_embedding"] = None st.session_state["vector_store_summary"] = None st.session_state["vector_search_results"] = [] st.session_state["projection_df"] = None st.session_state["projection_summary"] = None st.session_state["projection_model"] = None st.session_state["query_projection"] = None st.session_state["grounded_prompt"] = None st.session_state["retrieval_results"] = [] st.session_state["generated_answer"] = None st.session_state["app_status"] = f"Embedded {matrix.shape[0]} chunk(s)" toast_success(f"Embedded {matrix.shape[0]} chunk(s)") def handle_projection_build(method: str) -> None: chunks = st.session_state["chunks"] embedding_matrix = st.session_state["embedding_matrix"] if not chunks or embedding_matrix is None: st.session_state["projection_errors"] = [ "Generate embeddings before building a 2D projection." ] st.session_state["app_status"] = "Projection failed" return try: projection_df, projection_summary, projection_model = project_chunks( chunks=chunks, embedding_matrix=embedding_matrix, method=method, ) except ProjectionError as exc: st.session_state["projection_errors"] = [str(exc)] st.session_state["projection_df"] = None st.session_state["projection_summary"] = None st.session_state["projection_model"] = None st.session_state["query_projection"] = None st.session_state["retrieval_results"] = [] st.session_state["generated_answer"] = None st.session_state["grounded_prompt"] = None st.session_state["app_status"] = "Projection failed" return st.session_state["projection_errors"] = [] st.session_state["projection_df"] = projection_df st.session_state["projection_summary"] = projection_summary st.session_state["projection_model"] = projection_model st.session_state["query_projection"] = None st.session_state["grounded_prompt"] = None st.session_state["retrieval_results"] = [] st.session_state["generated_answer"] = None st.session_state["app_status"] = f"Projected {len(projection_df)} chunk(s) to 2D" toast_success(f"Projected {len(projection_df)} chunk(s) to 2D") def handle_store_in_lancedb() -> None: chunks = st.session_state["chunks"] embedding_matrix = st.session_state["embedding_matrix"] if not chunks or embedding_matrix is None: st.session_state["vector_store_errors"] = [ "Generate embeddings before storing vectors in LanceDB." ] st.session_state["app_status"] = "Vector DB sync failed" return try: summary = write_chunk_embeddings( chunks=chunks, embedding_matrix=embedding_matrix, ) except VectorStoreError as exc: st.session_state["vector_store_errors"] = [str(exc)] st.session_state["vector_store_summary"] = None st.session_state["vector_search_results"] = [] st.session_state["app_status"] = "Vector DB sync failed" return st.session_state["vector_store_errors"] = [] st.session_state["vector_store_summary"] = summary st.session_state["vector_search_results"] = [] st.session_state["app_status"] = f"Stored {summary['row_count']} row(s) in LanceDB" toast_success(f"Stored {summary['row_count']} row(s) in LanceDB") def get_selected_document() -> dict | None: raw_documents = st.session_state["raw_documents"] selected_document_id = st.session_state["selected_document_id"] for document in raw_documents: if document["doc_id"] == selected_document_id: return document return raw_documents[0] if raw_documents else None def get_selected_chunk() -> dict | None: chunks = st.session_state["chunks"] selected_chunk_id = st.session_state["selected_chunk_id"] for chunk in chunks: if chunk["chunk_id"] == selected_chunk_id: return chunk return chunks[0] if chunks else None def get_selected_chunk_vector(): selected_chunk = get_selected_chunk() embedding_matrix = st.session_state["embedding_matrix"] chunks = st.session_state["chunks"] if selected_chunk is None or embedding_matrix is None: return None, None for index, chunk in enumerate(chunks): if chunk["chunk_id"] == selected_chunk["chunk_id"]: return selected_chunk, embedding_matrix[index] return None, None def handle_lancedb_neighbor_search(top_k: int) -> None: vector_store_summary = st.session_state["vector_store_summary"] if vector_store_summary is None: st.session_state["vector_store_errors"] = [ "Store vectors in LanceDB before running nearest-neighbor preview." ] st.session_state["app_status"] = "Vector DB search failed" return selected_chunk, query_vector = get_selected_chunk_vector() if selected_chunk is None or query_vector is None: st.session_state["vector_store_errors"] = [ "Select a chunk and generate embeddings before searching LanceDB." ] st.session_state["app_status"] = "Vector DB search failed" return try: results_df = search_similar_chunks( query_vector=query_vector, top_k=top_k, exclude_chunk_id=selected_chunk["chunk_id"], ) except VectorStoreError as exc: st.session_state["vector_store_errors"] = [str(exc)] st.session_state["vector_search_results"] = [] st.session_state["app_status"] = "Vector DB search failed" return st.session_state["vector_store_errors"] = [] st.session_state["vector_search_results"] = results_df.to_dict(orient="records") st.session_state["app_status"] = f"Retrieved {len(results_df)} neighbor(s) from LanceDB" toast_success(f"Retrieved {len(results_df)} neighbor(s) from LanceDB") def handle_query_retrieval(query_text: str, top_k: int) -> None: try: validate_query_text(query_text) except RetrievalError as exc: st.session_state["retrieval_errors"] = [str(exc)] st.session_state["app_status"] = "Retrieval failed" return if st.session_state["vector_store_summary"] is None: st.session_state["retrieval_errors"] = [ "Store vectors in LanceDB before running query retrieval." ] st.session_state["app_status"] = "Retrieval failed" return if st.session_state["projection_summary"] is None: st.session_state["retrieval_errors"] = [ "Project embeddings to 2D before running query retrieval." ] st.session_state["app_status"] = "Retrieval failed" return try: query_vector = embed_query_text(query_text) results_df = search_similar_chunks( query_vector=query_vector, top_k=top_k, exclude_chunk_id=None, ) retrieval_results = enrich_query_results( results_df=results_df, chunks=st.session_state["chunks"], embedding_matrix=st.session_state["embedding_matrix"], query_vector=query_vector, ) query_projection = transform_query_vector( query_vector=query_vector, projection_model=st.session_state["projection_model"], projection_summary=st.session_state["projection_summary"], ) grounded_prompt = build_grounded_prompt(query_text, retrieval_results) except ( EmbeddingError, VectorStoreError, RetrievalError, ProjectionError, RAGGenerationError, ) as exc: st.session_state["retrieval_errors"] = [str(exc)] st.session_state["query_embedding"] = None st.session_state["query_projection"] = None st.session_state["grounded_prompt"] = None st.session_state["retrieval_results"] = [] st.session_state["generated_answer"] = None st.session_state["app_status"] = "Retrieval failed" return st.session_state["retrieval_errors"] = [] st.session_state["generation_errors"] = [] st.session_state["query_embedding"] = query_vector st.session_state["query_projection"] = { **query_projection, "query_text": query_text.strip(), } st.session_state["grounded_prompt"] = grounded_prompt st.session_state["retrieval_results"] = retrieval_results st.session_state["generated_answer"] = None if retrieval_results: top_result = retrieval_results[0] st.session_state["selected_chunk_id"] = top_result["chunk_id"] st.session_state["selected_document_id"] = top_result["doc_id"] st.session_state["app_status"] = f"Retrieved {len(retrieval_results)} chunk(s) for query" toast_success(f"Retrieved {len(retrieval_results)} chunk(s) for query") else: st.session_state["app_status"] = "No retrieval results found" def handle_answer_generation(query_text: str, generation_settings: dict) -> None: retrieval_results = st.session_state["retrieval_results"] if not retrieval_results: st.session_state["generation_errors"] = [ "Run retrieval before generating a grounded answer." ] st.session_state["app_status"] = "Answer generation failed" return try: generated_answer = generate_grounded_answer( query_text=query_text, retrieval_results=retrieval_results, provider=generation_settings["provider"], model_name=generation_settings["model"], api_key=generation_settings["api_key"], ) except RAGGenerationError as exc: st.session_state["generation_errors"] = [str(exc)] st.session_state["generated_answer"] = None st.session_state["app_status"] = "Answer generation failed" return st.session_state["generation_errors"] = [] st.session_state["grounded_prompt"] = generated_answer["prompt"] st.session_state["generated_answer"] = generated_answer st.session_state["app_status"] = f"Generated grounded answer with {generated_answer['model']}" toast_success(f"Generated grounded answer with {generated_answer['model']}") def handle_projection_selection(event) -> None: if event is None: return if hasattr(event, "selection"): selection = event.selection elif isinstance(event, dict): selection = event.get("selection") else: selection = None if not selection: return points = selection.get("points", []) if not points: return point = points[0] custom_data = point.get("customdata", []) if len(custom_data) < 2: return chunk_id = custom_data[0] doc_id = custom_data[1] if chunk_id == "__query__": return st.session_state["selected_chunk_id"] = chunk_id st.session_state["selected_document_id"] = doc_id def clear_documents() -> None: st.session_state["uploaded_documents"] = [] st.session_state["raw_documents"] = [] st.session_state["document_errors"] = [] st.session_state["chunk_errors"] = [] st.session_state["embedding_errors"] = [] st.session_state["vector_store_errors"] = [] st.session_state["projection_errors"] = [] st.session_state["retrieval_errors"] = [] st.session_state["generation_errors"] = [] st.session_state["selected_document_id"] = None st.session_state["chunks"] = [] st.session_state["selected_chunk_id"] = None st.session_state["embedding_matrix"] = None st.session_state["embedding_summary"] = None st.session_state["query_embedding"] = None st.session_state["vector_store_summary"] = None st.session_state["vector_search_results"] = [] st.session_state["projection_df"] = None st.session_state["projection_summary"] = None st.session_state["projection_model"] = None st.session_state["query_projection"] = None st.session_state["query_text"] = "" st.session_state["grounded_prompt"] = None st.session_state["retrieval_results"] = [] st.session_state["generated_answer"] = None st.session_state["app_status"] = "Ready" toast_success("Cleared current session data") def get_selected_chunk_embedding_preview(max_dims: int = 8) -> pd.DataFrame | None: selected_chunk = get_selected_chunk() embedding_matrix = st.session_state["embedding_matrix"] chunks = st.session_state["chunks"] if selected_chunk is None or embedding_matrix is None: return None chunk_position = None for index, chunk in enumerate(chunks): if chunk["chunk_id"] == selected_chunk["chunk_id"]: chunk_position = index break if chunk_position is None: return None vector = embedding_matrix[chunk_position] preview_rows = [] for dimension in range(min(max_dims, len(vector))): preview_rows.append( { "dimension": dimension, "value": round(float(vector[dimension]), 6), } ) return pd.DataFrame(preview_rows) def render_generation_settings() -> dict: provider_ids = list(PROVIDERS.keys()) default_index = ( provider_ids.index(settings.llm_provider) if settings.llm_provider in provider_ids else 0 ) provider_id = st.selectbox( "LLM provider", options=provider_ids, index=default_index, format_func=lambda pid: PROVIDERS[pid].label, help="Provider used to generate the grounded answer.", ) spec = PROVIDERS[provider_id] model = st.text_input( "Model", value=get_default_model(provider_id), key=f"model_input_{provider_id}", help="Model name exactly as the provider's API expects it.", ) api_key = "" if spec.requires_api_key: api_key = st.text_input( "API key", value="", type="password", key=f"api_key_input_{provider_id}", help=( "Kept only in this browser session — never written to disk or logs. " f"Leave blank to use {spec.api_key_env_var} from your environment." ), ) if not api_key: if get_env_api_key(provider_id): st.caption(f"Using `{spec.api_key_env_var}` from the environment.") else: st.caption( f"No API key configured. Paste one above or set " f"`{spec.api_key_env_var}`. Get a key at {spec.key_console_url}" ) else: st.caption(f"Runs against your local Ollama server at {settings.ollama_base_url}.") return { "provider": provider_id, "model": model, "api_key": api_key, } def render_sidebar_pipeline_health() -> None: checks = [ ("Documents", bool(st.session_state["raw_documents"])), ("Chunks", bool(st.session_state["chunks"])), ("Embeddings", st.session_state["embedding_summary"] is not None), ("Projection", st.session_state["projection_summary"] is not None), ("Vector DB", st.session_state["vector_store_summary"] is not None), ("Retrieval", bool(st.session_state["retrieval_results"])), ("Answer", st.session_state["generated_answer"] is not None), ] with st.container(border=True): st.subheader("Pipeline Health") for label, ok in checks: icon = "✅" if ok else "⚪" st.write(f"{icon} {label}") def render_sidebar() -> dict: with st.sidebar: st.title("RAG Chunk Visualizer") st.caption("Visual debugging and teaching tool for RAG systems") st.divider() st.subheader("Document Upload") uploader_key = f"document_uploader_{st.session_state['uploader_key_seq']}" uploaded_files = st.file_uploader( "Upload .txt or .md documents", type=["txt", "md"], accept_multiple_files=True, key=uploader_key, help="Drag and drop one or more UTF-8 files — they are processed automatically.", ) # Auto-process whenever the selected files change, so users never have to # hunt for a separate submit button after picking a file. signature = ( tuple(sorted((f.name, f.size) for f in uploaded_files)) if uploaded_files else () ) if signature != st.session_state["uploaded_signature"]: st.session_state["uploaded_signature"] = signature if uploaded_files: with st.spinner("Processing uploaded documents..."): handle_document_uploads(uploaded_files) if st.button("Clear Documents", width="stretch"): st.session_state["uploader_key_seq"] += 1 st.session_state["uploaded_signature"] = () clear_documents() st.rerun() st.divider() st.subheader("Pipeline Controls") chunk_size = st.slider( "Chunk size", min_value=100, max_value=1200, value=settings.chunk_size, step=50, help="Target character window for chunking.", ) chunk_overlap = st.slider( "Chunk overlap", min_value=0, max_value=300, value=settings.chunk_overlap, step=25, help="Character overlap between consecutive chunks.", ) top_k = st.slider( "Top-k retrieval", min_value=1, max_value=10, value=settings.top_k, step=1, help="Number of chunks to retrieve for the query.", ) projection_method = st.selectbox( "Projection method", options=["pca", "umap"], index=0 if settings.projection_method == "pca" else 1, help="2D projection method for embedding visualization.", ) build_chunks = st.button( "Build Chunks", type="primary", width="stretch", disabled=not st.session_state["raw_documents"], ) if build_chunks: handle_chunk_build(chunk_size=chunk_size, chunk_overlap=chunk_overlap) build_embeddings = st.button( "Generate Embeddings", type="primary", width="stretch", disabled=not st.session_state["chunks"], ) if build_embeddings: with st.spinner("Loading model and generating chunk embeddings..."): handle_embedding_build() build_projection = st.button( "Project to 2D", type="primary", width="stretch", disabled=st.session_state["embedding_matrix"] is None, ) if build_projection: with st.spinner(f"Projecting embeddings with {projection_method.upper()}..."): handle_projection_build(method=projection_method) store_vectors = st.button( "Store in LanceDB", type="primary", width="stretch", disabled=st.session_state["embedding_matrix"] is None, ) if store_vectors: with st.spinner("Writing chunk embeddings to LanceDB..."): handle_store_in_lancedb() preview_neighbors = st.button( "Preview Similar Chunks", type="primary", width="stretch", disabled=st.session_state["vector_store_summary"] is None, ) if preview_neighbors: with st.spinner("Searching nearest chunks in LanceDB..."): handle_lancedb_neighbor_search(top_k=top_k) st.divider() st.subheader("Query") query_text = st.text_input( "Enter a query", value=st.session_state.get("query_text", ""), placeholder="e.g. What does the document say about vector search?", ) run_query = st.button( "Run Retrieval", type="primary", width="stretch", disabled=( st.session_state["vector_store_summary"] is None or st.session_state["projection_summary"] is None or not query_text.strip() ), ) if run_query: with st.spinner("Embedding query and retrieving relevant chunks..."): handle_query_retrieval(query_text=query_text, top_k=top_k) st.divider() st.subheader("Answer Generation") enable_generation = st.toggle( "Enable answer generation", value=False, help="Generate a grounded answer from the retrieved chunks with the selected LLM.", ) generation_settings = render_generation_settings() generate_answer = st.button( "Generate Answer", type="primary", width="stretch", disabled=( not enable_generation or not st.session_state["retrieval_results"] or not query_text.strip() ), ) if generate_answer: provider_label = PROVIDERS[generation_settings["provider"]].label with st.spinner(f"Generating grounded answer with {provider_label}..."): handle_answer_generation( query_text=query_text, generation_settings=generation_settings, ) st.divider() render_sidebar_pipeline_health() st.divider() st.subheader("Embedding Settings") st.text_input( "Embedding model", value=settings.embedding_model, disabled=True, ) st.text_input( "Embedding device", value=settings.embedding_device, disabled=True, ) st.text_input( "LanceDB table", value=settings.lancedb_table_name, disabled=True, ) st.divider() return { "chunk_size": chunk_size, "chunk_overlap": chunk_overlap, "top_k": top_k, "projection_method": projection_method, "query_text": query_text, "enable_generation": enable_generation, "generation_settings": generation_settings, "run_query": run_query, } def render_header() -> None: st.title("RAG Chunk Visualizer") st.markdown( "Inspect how documents become chunks, how chunks become vectors, " "and how retrieval selects grounded context." ) def render_overview_banner() -> None: docs = len(st.session_state["raw_documents"]) chunks = len(st.session_state["chunks"]) retrieved = len(st.session_state["retrieval_results"]) answer_ready = st.session_state["generated_answer"] is not None with st.container(border=True): st.markdown( f"**Overview** — Docs: `{docs}` | " f"Chunks: `{chunks}` | " f"Retrieved: `{retrieved}` | " f"Answer ready: `{answer_ready}`" ) def render_status_metrics(sidebar_values: dict) -> None: embedding_matrix = st.session_state["embedding_matrix"] vector_store_summary = st.session_state["vector_store_summary"] projection_df = st.session_state["projection_df"] projection_summary = st.session_state["projection_summary"] generated_answer = st.session_state["generated_answer"] embedded_count = 0 if embedding_matrix is None else int(embedding_matrix.shape[0]) embedding_dim = "—" if embedding_matrix is None else int(embedding_matrix.shape[1]) stored_count = 0 if vector_store_summary is None else vector_store_summary["row_count"] projected_count = 0 if projection_df is None else int(len(projection_df)) retrieved_count = len(st.session_state["retrieval_results"]) answer_ready = "Yes" if generated_answer is not None else "No" method_label = "—" if projection_summary is not None: method_label = str(projection_summary.get("method", "—")).upper() metric_cols = st.columns(10) metric_cols[0].metric("Documents", len(st.session_state["raw_documents"])) metric_cols[1].metric("Chunks", len(st.session_state["chunks"])) metric_cols[2].metric("Embedded", embedded_count) metric_cols[3].metric("Stored", stored_count) metric_cols[4].metric("Projected", projected_count) metric_cols[5].metric("Retrieved", retrieved_count) metric_cols[6].metric("Answer", answer_ready) metric_cols[7].metric("Dim", embedding_dim) metric_cols[8].metric("Method", method_label) metric_cols[9].metric("Status", st.session_state["app_status"]) def render_document_summary() -> None: st.subheader("Uploaded Documents") raw_documents = st.session_state["raw_documents"] document_errors = st.session_state["document_errors"] if document_errors: for error in document_errors: st.error(error) if not raw_documents: st.info("No documents uploaded yet.") return summary_df = pd.DataFrame( [ { "doc_id": doc["doc_id"], "filename": doc["filename"], "extension": doc["extension"], "chars": doc["char_count"], "words": doc["word_count"], } for doc in raw_documents ] ) st.dataframe(summary_df, width="stretch", hide_index=True) def render_document_preview_panel() -> None: st.subheader("Document Preview") raw_documents = st.session_state["raw_documents"] if not raw_documents: with st.container(border=True): st.write("No document available for preview yet.") return doc_ids = [doc["doc_id"] for doc in raw_documents] if st.session_state["selected_document_id"] not in doc_ids: st.session_state["selected_document_id"] = doc_ids[0] st.selectbox( "Select document", options=doc_ids, key="selected_document_id", format_func=lambda doc_id: next( doc["filename"] for doc in raw_documents if doc["doc_id"] == doc_id ), ) selected_document = get_selected_document() if selected_document is None: return with st.container(border=True): st.markdown(f"**Document ID:** `{selected_document['doc_id']}`") st.markdown(f"**Filename:** {selected_document['filename']}") st.markdown(f"**Characters:** {selected_document['char_count']}") st.markdown(f"**Words:** {selected_document['word_count']}") st.markdown(f"**Saved to:** `{selected_document['source_path']}`") st.markdown("**Extracted text preview:**") st.code(selected_document["text"][:3000], language="text") def render_chunk_summary_panel() -> None: st.subheader("Chunk Summary") chunk_errors = st.session_state["chunk_errors"] chunks = st.session_state["chunks"] if chunk_errors: for error in chunk_errors: st.error(error) if not chunks: st.info("No chunks built yet. Upload documents, then click 'Build Chunks'.") return chunk_df = pd.DataFrame( [ { "chunk_id": chunk["chunk_id"], "filename": chunk["filename"], "chunk_index": chunk["chunk_index"], "char_start": chunk["char_start"], "char_end": chunk["char_end"], "chars": chunk["char_count"], "words": chunk["word_count"], "preview": chunk["preview"], } for chunk in chunks ] ) st.dataframe(chunk_df, width="stretch", hide_index=True) def render_chunk_detail_panel() -> None: st.subheader("Selected Chunk Details") chunks = st.session_state["chunks"] selected_document = get_selected_document() if not chunks: with st.container(border=True): st.code( "Chunking has not been run yet.\n\n" "Build chunks, generate embeddings, and project them in this stage.", language="text", ) return if selected_document is not None: visible_chunks = [c for c in chunks if c["doc_id"] == selected_document["doc_id"]] else: visible_chunks = chunks if not visible_chunks: with st.container(border=True): st.write("No chunks available for the selected document.") return chunk_lookup = {chunk["chunk_id"]: chunk for chunk in visible_chunks} chunk_ids = list(chunk_lookup.keys()) if st.session_state["selected_chunk_id"] not in chunk_ids: st.session_state["selected_chunk_id"] = chunk_ids[0] st.selectbox( "Select chunk", options=chunk_ids, key="selected_chunk_id", format_func=lambda chunk_id: ( f"{chunk_lookup[chunk_id]['filename']} | chunk {chunk_lookup[chunk_id]['chunk_index']}" ), ) selected_chunk = get_selected_chunk() if selected_chunk is None: return with st.container(border=True): st.markdown(f"**Chunk ID:** `{selected_chunk['chunk_id']}`") st.markdown(f"**Document:** {selected_chunk['filename']}") st.markdown(f"**Chunk Index:** {selected_chunk['chunk_index']}") st.markdown( f"**Character Range:** {selected_chunk['char_start']} → {selected_chunk['char_end']}" ) st.markdown(f"**Characters:** {selected_chunk['char_count']}") st.markdown(f"**Words:** {selected_chunk['word_count']}") st.markdown("**Chunk Text:**") st.code(selected_chunk["text"], language="text") embedding_preview_df = get_selected_chunk_embedding_preview() if embedding_preview_df is not None: st.markdown("**Embedding Preview (first dimensions):**") st.dataframe( embedding_preview_df, width="stretch", hide_index=True, ) def render_embedding_summary_panel() -> None: st.subheader("Embedding Summary") embedding_errors = st.session_state["embedding_errors"] embedding_summary = st.session_state["embedding_summary"] if embedding_errors: for error in embedding_errors: st.error(error) if embedding_summary is None: st.info("No embeddings generated yet. Build chunks, then click 'Generate Embeddings'.") return summary_df = pd.DataFrame([embedding_summary]) st.dataframe(summary_df, width="stretch", hide_index=True) def render_projection_summary_panel() -> None: st.subheader("Projection Summary") projection_errors = st.session_state["projection_errors"] projection_summary = st.session_state["projection_summary"] if projection_errors: for error in projection_errors: st.error(error) if projection_summary is None: st.info("No 2D projection yet. Generate embeddings, then click 'Project to 2D'.") return summary_df = pd.DataFrame([projection_summary]) st.dataframe(summary_df, width="stretch", hide_index=True) def render_vector_store_panel() -> None: st.subheader("LanceDB Summary") vector_store_errors = st.session_state["vector_store_errors"] vector_store_summary = st.session_state["vector_store_summary"] if vector_store_errors: for error in vector_store_errors: st.error(error) if vector_store_summary is None: st.info("No LanceDB table yet. Generate embeddings, then click 'Store in LanceDB'.") return summary_df = pd.DataFrame([vector_store_summary]) st.dataframe(summary_df, width="stretch", hide_index=True) try: preview_df = fetch_table_preview(limit=8) except VectorStoreError as exc: st.warning(str(exc)) return st.caption("Local LanceDB row preview") st.dataframe(preview_df, width="stretch", hide_index=True) def render_embedding_map_panel() -> None: st.subheader("Embedding Map") projection_df = st.session_state["projection_df"] projection_summary = st.session_state["projection_summary"] query_projection = st.session_state["query_projection"] retrieval_results = st.session_state["retrieval_results"] if projection_df is None or projection_df.empty: placeholder_df = pd.DataFrame( [ {"x": -1.2, "y": 0.8, "label": "chunk_001", "document": "demo_doc_a"}, {"x": -0.4, "y": 0.2, "label": "chunk_002", "document": "demo_doc_a"}, {"x": 0.3, "y": -0.1, "label": "chunk_003", "document": "demo_doc_b"}, {"x": 1.1, "y": -0.6, "label": "chunk_004", "document": "demo_doc_b"}, ] ) st.caption("Project embeddings to 2D to replace this placeholder with the real map.") st.dataframe(placeholder_df, width="stretch", hide_index=True) return try: figure = create_embedding_scatter_plot( projection_df=projection_df, projection_summary=projection_summary, selected_chunk_id=st.session_state["selected_chunk_id"], retrieved_chunk_ids=get_retrieved_chunk_ids(retrieval_results), query_projection=query_projection, query_text=st.session_state["query_text"], ) except PlotBuildError as exc: st.error(str(exc)) return st.caption( "Retrieved chunks are outlined, the selected chunk is highlighted, " "and the query appears as its own point." ) event = st.plotly_chart( figure, width="stretch", height=560, key="embedding_map_chart", on_select="rerun", selection_mode=("points", "box"), config={ "displaylogo": False, "scrollZoom": False, }, ) handle_projection_selection(event) def render_query_results_panel() -> None: st.subheader("Retrieved Sources") retrieval_errors = st.session_state["retrieval_errors"] retrieval_results = st.session_state["retrieval_results"] if retrieval_errors: for error in retrieval_errors: st.error(error) if not retrieval_results: with st.container(border=True): st.write("No query retrieval results yet.") st.caption("Store vectors in LanceDB, project to 2D, then run a query.") return results_df = pd.DataFrame( [ { "rank": result["rank"], "filename": result["filename"], "chunk_index": result["chunk_index"], "similarity_score": result["similarity_score"], "preview": result["preview"], } for result in retrieval_results ] ) st.dataframe(results_df, width="stretch", hide_index=True) for result in retrieval_results: expander_title = ( f"#{result['rank']} | {result['filename']} | " f"chunk {result['chunk_index']} | score {result['similarity_score']}" ) with st.expander(expander_title, expanded=result["rank"] == 1): st.markdown(f"**Chunk ID:** `{result['chunk_id']}`") st.markdown(f"**Document:** {result['filename']}") st.markdown(f"**Similarity Score:** {result['similarity_score']}") st.code(result["text"], language="text") def render_debug_neighbor_panel() -> None: st.subheader("LanceDB Neighbor Preview") results = st.session_state["vector_search_results"] if not results: with st.container(border=True): st.write("No chunk-to-chunk neighbor preview yet.") st.caption("Run 'Preview Similar Chunks' to inspect raw nearest-neighbor results.") return results_df = pd.DataFrame(results) preferred_columns = [ "chunk_id", "filename", "chunk_index", "_distance", "preview", ] available_columns = [column for column in preferred_columns if column in results_df.columns] with st.container(border=True): st.caption("Raw LanceDB nearest-neighbor debug view.") st.dataframe( results_df.loc[:, available_columns], width="stretch", hide_index=True, ) def render_prompt_and_answer_panel() -> None: left_col, right_col = st.columns([1, 1]) with left_col: st.subheader("Grounded Prompt") with st.container(border=True): grounded_prompt = st.session_state["grounded_prompt"] if grounded_prompt: st.code(grounded_prompt, language="text") else: st.code( "Run retrieval first. The grounded prompt will be built from the " "query and retrieved chunks.", language="text", ) with right_col: st.subheader("Answer Panel") generation_errors = st.session_state["generation_errors"] generated_answer = st.session_state["generated_answer"] if generation_errors: for error in generation_errors: st.error(error) with st.container(border=True): if generated_answer is None: st.write("No answer generated yet.") st.caption("Enable generation and click 'Generate Answer' after retrieval.") else: st.markdown("**Answer**") st.caption( f"{generated_answer.get('provider_label', 'LLM')} · " f"{generated_answer.get('model', 'unknown model')}" ) st.write(generated_answer["answer"]) st.markdown("**Generation Metadata**") metadata_df = pd.DataFrame([generated_answer["metadata"]]) st.dataframe(metadata_df, width="stretch", hide_index=True) st.markdown("**Sources Used**") for result in st.session_state["retrieval_results"]: st.markdown( f"- [Source {result['rank']}] " f"{result['filename']} | " f"chunk {result['chunk_index']}" ) def render_pipeline_tab() -> None: upper_left, upper_right = st.columns([1.25, 1]) with upper_left: render_document_summary() st.divider() render_document_preview_panel() with upper_right: render_chunk_detail_panel() st.divider() summary_col_1, summary_col_2, summary_col_3 = st.columns(3) with summary_col_1: render_embedding_summary_panel() with summary_col_2: render_projection_summary_panel() with summary_col_3: render_vector_store_panel() st.divider() render_chunk_summary_panel() st.divider() render_embedding_map_panel() def render_retrieval_tab() -> None: left_col, right_col = st.columns([1.25, 0.75]) with left_col: render_embedding_map_panel() st.divider() render_query_results_panel() with right_col: with st.container(border=True): st.subheader("Retrieval Notes") st.write("Base map: projected chunk embeddings") st.write("Green X: current query point") st.write("Orange outlines: retrieved chunks") st.write("White diamond: selected chunk") st.divider() with st.expander("Debug: Chunk-to-Chunk Neighbor Preview", expanded=False): render_debug_neighbor_panel() def render_generation_tab() -> None: with st.container(border=True): st.subheader("Generation Checklist") st.write("1. Run retrieval") st.write("2. Review retrieved chunks") st.write("3. Enable answer generation") st.write("4. Generate the grounded answer") st.divider() render_prompt_and_answer_panel() def render_footer_notes(sidebar_values: dict) -> None: st.divider() with st.expander("Current UI State", expanded=False): st.json( { "documents_loaded": len(st.session_state["raw_documents"]), "chunks_built": len(st.session_state["chunks"]), "embeddings_ready": st.session_state["embedding_summary"] is not None, "projection_ready": st.session_state["projection_summary"] is not None, "vector_store_ready": st.session_state["vector_store_summary"] is not None, "retrieval_count": len(st.session_state["retrieval_results"]), "answer_ready": st.session_state["generated_answer"] is not None, "chunk_size": sidebar_values["chunk_size"], "chunk_overlap": sidebar_values["chunk_overlap"], "top_k": sidebar_values["top_k"], "projection_method": sidebar_values["projection_method"], "query_text": sidebar_values["query_text"], "enable_generation": sidebar_values["enable_generation"], "llm_provider": sidebar_values["generation_settings"]["provider"], "llm_model": sidebar_values["generation_settings"]["model"], } ) def render_app() -> None: render_header() sidebar_values = render_sidebar() st.session_state["query_text"] = sidebar_values["query_text"] render_status_metrics(sidebar_values) render_overview_banner() pipeline_tab, retrieval_tab, generation_tab = st.tabs( ["🧩 Pipeline", "🔎 Retrieval", "🧠 Generation"], width="stretch", key="main_workspace_tabs", on_change="rerun", ) if pipeline_tab.open: with pipeline_tab: render_pipeline_tab() if retrieval_tab.open: with retrieval_tab: render_retrieval_tab() if generation_tab.open: with generation_tab: render_generation_tab() render_footer_notes(sidebar_values)