"""
DocMind — Reusable Streamlit UI Components
Provides render functions for chat messages, source cards,
grounding bars, confidence badges, and pipeline progress.
"""
import streamlit as st
from typing import Dict, List, Optional
from pipeline.chunker import ChunkMetadata
from pipeline.grounding import ConfidenceLevel, GroundingResult, SentenceScore
from pipeline.retriever import RetrievalStats
# ── Chat Messages ───────────────────────────────────────────────────
def render_chat_message(
role: str,
content: str,
grounding_result: Optional[GroundingResult] = None,
sources: Optional[List[ChunkMetadata]] = None,
doc_index_map: Optional[Dict[str, int]] = None,
) -> None:
"""
Render a single chat message with optional grounding info and sources.
Args:
role: "user" or "bot"
content: The message text
grounding_result: Optional grounding gate result for bot messages
sources: Optional list of source chunks for bot messages
doc_index_map: Optional mapping of doc_id → color index
"""
if role == "user":
avatar = "👤"
bubble_class = "user"
else:
avatar = "🧠"
bubble_class = "bot"
st.markdown(f"""
""", unsafe_allow_html=True)
# For bot messages, show grounding info and sources
if role == "bot" and grounding_result and not grounding_result.is_refused:
render_grounding_bar(grounding_result)
render_confidence_badge(grounding_result.confidence)
# Show per-sentence scores in an expander
if grounding_result.sentence_scores:
with st.expander("📊 Per-sentence grounding scores"):
for ss in grounding_result.sentence_scores:
_render_sentence_score(ss)
if role == "bot" and sources:
with st.expander(f"📎 Source chunks ({len(sources)})"):
for chunk in sources:
doc_idx = 0
if doc_index_map and chunk.doc_id in doc_index_map:
doc_idx = doc_index_map[chunk.doc_id]
render_source_card(chunk, doc_idx)
# ── Source Cards ────────────────────────────────────────────────────
def render_source_card(chunk: ChunkMetadata, doc_color_index: int = 0) -> None:
"""Render a collapsible source chunk preview with document color tag."""
tag_class = f"doc-tag-{doc_color_index % 3}"
preview_text = chunk.text[:300] + ("..." if len(chunk.text) > 300 else "")
st.markdown(f"""
""", unsafe_allow_html=True)
# ── Grounding Bar ──────────────────────────────────────────────────
def render_grounding_bar(grounding_result: GroundingResult) -> None:
"""Render an animated grounding score progress bar."""
score = grounding_result.overall_score
pct = max(0, min(100, int(score * 100)))
if grounding_result.confidence == ConfidenceLevel.HIGH:
fill_class = "high"
elif grounding_result.confidence == ConfidenceLevel.MODERATE:
fill_class = "moderate"
else:
fill_class = "low"
st.markdown(f"""
Grounding Score
{score:.1%}
""", unsafe_allow_html=True)
# ── Confidence Badge ───────────────────────────────────────────────
def render_confidence_badge(level: ConfidenceLevel) -> None:
"""Render a confidence level badge."""
badges = {
ConfidenceLevel.HIGH: ("✅ High Confidence", "badge-high"),
ConfidenceLevel.MODERATE: ("⚠️ Moderate Confidence", "badge-moderate"),
ConfidenceLevel.LOW: ("❌ Insufficient Grounding", "badge-low"),
}
text, css_class = badges.get(level, ("❓ Unknown", "badge-low"))
st.markdown(f"""
{text}
""", unsafe_allow_html=True)
def _render_sentence_score(ss: SentenceScore) -> None:
"""Render a single sentence's grounding score."""
if ss.confidence == ConfidenceLevel.HIGH:
color = "#34D399"
icon = "✅"
elif ss.confidence == ConfidenceLevel.MODERATE:
color = "#FBBF24"
icon = "⚠️"
else:
color = "#F87171"
icon = "❌"
st.markdown(f"""
{icon}
{ss.entailment_score:.1%}
[{ss.chunk_id}]
{ss.sentence}
""", unsafe_allow_html=True)
# ── Retrieval Stats ────────────────────────────────────────────────
def render_retrieval_stats(stats: RetrievalStats) -> None:
"""Render retrieval debug information."""
st.markdown(f"""
Retrieval Stats
BM25: {stats.bm25_hits} hits |
Dense: {stats.dense_hits} hits |
After RRF: {stats.rrf_results} |
Latency: {stats.latency_ms:.0f}ms
""", unsafe_allow_html=True)
# ── Document Status ────────────────────────────────────────────────
def render_document_status(
doc_name: str,
chunk_count: int,
page_count: int,
doc_color_index: int = 0,
) -> None:
"""Render a sidebar document status card."""
tag_class = f"doc-tag-{doc_color_index % 3}"
st.markdown(f"""
📄 {doc_name}
{chunk_count} chunks • {page_count} pages
""", unsafe_allow_html=True)
# ── Pipeline Progress ──────────────────────────────────────────────
def render_pipeline_progress(stages: List[dict]) -> None:
"""
Render a multi-stage pipeline progress indicator.
Each stage is a dict: {"name": str, "status": "done"|"active"|"pending"}
"""
icons = {"done": "✅", "active": "⏳", "pending": "⏸️"}
html_parts = []
for stage in stages:
icon = icons.get(stage["status"], "⏸️")
css_class = stage["status"]
html_parts.append(
f'{icon} {stage["name"]}
'
)
st.markdown("\n".join(html_parts), unsafe_allow_html=True)
# ── Empty State / Hero ─────────────────────────────────────────────
def render_empty_state() -> None:
"""Render the advanced empty state (Hero Section) when no documents are uploaded."""
st.html("""
Welcome to DocMind
Upload your documents in the sidebar to securely chat, summarize, and compare content using Grounded RAG.
⚡
Hybrid Retrieval
Combines dense vector search (BGE-M3) with sparse keyword search (BM25) to ensure maximum context recall across your documents.
🛡️
NLI Grounding
Every generated sentence is rigorously verified against the source context using Natural Language Inference models to eliminate hallucinations.
🧠
Llama 3 Powered
Utilizes Groq's lightning-fast inference for deep reasoning, automated summarization, and side-by-side document comparison.
""")
# ── Comparison Table ───────────────────────────────────────────────
def render_comparison_table(comparison_text: str) -> None:
"""Render a document comparison result."""
st.markdown(comparison_text, unsafe_allow_html=True)
# ── Dashboard Metrics ──────────────────────────────────────────────
def render_dashboard_metrics(doc_count: int, chunk_count: int, memory_mb: float = 0.0) -> None:
"""Render a row of top-level metric cards for the dashboard."""
st.markdown(f"""
Documents
{doc_count}
Currently Indexed
Knowledge Chunks
{chunk_count}
Vector DB + BM25
System Status
Ready
Hybrid Retrieval Online
""", unsafe_allow_html=True)
# ── Key Point Card ─────────────────────────────────────────────────
def render_keypoint_card(text: str, page_ref: str = "") -> None:
"""Render a styled HTML card for a key point."""
page_html = f" {page_ref}" if page_ref else ""
st.markdown(f"""
""", unsafe_allow_html=True)