visual-rag-toolkit 0.1.1__py3-none-any.whl

demo/ui/playground.py

@@ -0,0 +1,339 @@
+"""Playground tab component."""
+
+import streamlit as st
+
+from demo.config import AVAILABLE_MODELS, RETRIEVAL_MODES, STAGE1_MODES
+from demo.qdrant_utils import (
+    get_qdrant_credentials,
+    get_collections,
+    sample_points_cached,
+    search_collection,
+)
+
+
+def render_playground_tab():
+    st.subheader("🎮 Playground")
+    
+    active_collection = st.session_state.get("active_collection")
+    url, api_key = get_qdrant_credentials()
+    
+    if not active_collection:
+        collections = get_collections(url, api_key)
+        if collections:
+            active_collection = collections[0]
+    
+    if not active_collection:
+        st.warning("No collection available. Upload documents or select a collection.")
+        return
+    
+    points_for_model = sample_points_cached(active_collection, 1, 0, url, api_key)
+    model_name = None
+    if points_for_model:
+        model_name = points_for_model[0].get("payload", {}).get("model_name")
+    if not model_name:
+        model_name = AVAILABLE_MODELS[1]
+    
+    model_short = model_name.split("/")[-1] if model_name else "unknown"
+    cache_key = f"{active_collection}_{model_name}"
+    
+    if st.session_state.get("loaded_model_key") != cache_key:
+        st.session_state["model_loaded"] = False
+    
+    col_info, col_model = st.columns([2, 1])
+    with col_info:
+        st.info(f"**Collection:** `{active_collection}`")
+    with col_model:
+        if not st.session_state.get("model_loaded"):
+            with st.spinner(f"Loading {model_short}..."):
+                try:
+                    from visual_rag.retrieval import MultiVectorRetriever
+                    _ = MultiVectorRetriever(collection_name=active_collection, model_name=model_name)
+                    st.session_state["model_loaded"] = True
+                    st.session_state["loaded_model_key"] = cache_key
+                    st.session_state["loaded_model_name"] = model_name
+                except Exception as e:
+                    st.warning(f"Failed: {model_short}")
+        
+        if st.session_state.get("model_loaded"):
+            st.markdown(f"✅ Found <span style='color:#e74c3c;font-weight:bold;'>{model_short}</span> model", unsafe_allow_html=True)
+    
+    with st.expander("📦 Sample Points Explorer", expanded=True):
+        render_sample_explorer(active_collection, url, api_key)
+    
+    st.divider()
+    
+    st.subheader("🔍 RAG Query")
+    render_rag_query_interface(active_collection, model_name)
+
+
+def render_document_details(pt: dict, p: dict, score: float = None, rel_pct: float = None):
+    def _is_missing(v) -> bool:
+        if v is None:
+            return True
+        if isinstance(v, (list, tuple, dict)) and len(v) == 0:
+            return True
+        if isinstance(v, str):
+            s = v.strip()
+            return s == "" or s.lower() in {"na", "n/a", "none", "null", "unknown", "?", "-"}
+        return False
+
+    doc_id = p.get("doc_id") or p.get("union_doc_id") or p.get("source_doc_id") or "?"
+    corpus_id = p.get("corpus-id") or p.get("source_doc_id") or "?"
+    dataset = p.get("dataset") or p.get("source") or None
+    model = (p.get("model_name") or p.get("model") or None)
+    model = model.split("/")[-1] if isinstance(model, str) else None
+    doc_name = p.get("doc-id") or p.get("filename") or "Unknown"
+    
+    num_tiles = p.get("num_tiles")
+    visual_tokens = p.get("index_recovery_num_visual_tokens") or p.get("num_visual_tokens")
+    patches_per_tile = p.get("patches_per_tile")
+    torch_dtype = p.get("torch_dtype")
+    
+    orig_w = p.get("original_width")
+    orig_h = p.get("original_height")
+    crop_w = p.get("cropped_width")
+    crop_h = p.get("cropped_height")
+    resize_w = p.get("resized_width")
+    resize_h = p.get("resized_height")
+    crop_pct = p.get("crop_empty_percentage_to_remove")
+    crop_enabled = bool(p.get("crop_empty_enabled", False))
+    
+    col_meta, col_img = st.columns([1, 2])
+    
+    with col_meta:
+        st.markdown("##### 📄 Document Info")
+        st.markdown(f"**📁 Doc:** {doc_name}")
+        if not _is_missing(dataset):
+            st.markdown(f"**🏛️ Dataset:** {dataset}")
+        if not _is_missing(doc_id) and str(doc_id) != "?":
+            st.markdown(f"**🔑 Doc ID:** `{str(doc_id)[:20]}...`")
+        if not _is_missing(corpus_id) and str(corpus_id) != "?":
+            st.markdown(f"**📋 Corpus ID:** {corpus_id}")
+        
+        if score is not None:
+            st.divider()
+            st.markdown("##### 🎯 Relevance")
+            if rel_pct is not None:
+                st.markdown(f"**Relative:** 🟢 {rel_pct:.1f}%")
+                st.progress(rel_pct / 100)
+            st.caption(f"Raw score: {score:.4f}")
+        
+        st.divider()
+        visual_rows = []
+        if not _is_missing(model):
+            visual_rows.append(("🤖 Model", f"`{model}`"))
+        if not _is_missing(num_tiles):
+            visual_rows.append(("🔲 Tiles", str(num_tiles)))
+        if not _is_missing(visual_tokens):
+            visual_rows.append(("🔢 Visual Tokens", str(visual_tokens)))
+        if not _is_missing(patches_per_tile):
+            visual_rows.append(("📦 Patches/Tile", str(patches_per_tile)))
+        if not _is_missing(torch_dtype):
+            visual_rows.append(("⚙️ Dtype", str(torch_dtype)))
+        if visual_rows:
+            st.markdown("##### 🎨 Visual Metadata")
+            for k, v in visual_rows:
+                st.markdown(f"**{k}:** {v}")
+        
+        st.divider()
+        dim_rows = []
+        if not _is_missing(orig_w) and not _is_missing(orig_h):
+            dim_rows.append(("Original", f"{orig_w}×{orig_h}"))
+        if not _is_missing(resize_w) and not _is_missing(resize_h):
+            dim_rows.append(("Resized", f"{resize_w}×{resize_h}"))
+        if crop_enabled and not _is_missing(crop_w) and not _is_missing(crop_h):
+            dim_rows.append(("Cropped", f"{crop_w}×{crop_h}"))
+        if dim_rows:
+            st.markdown("##### 📐 Dimensions")
+            for k, v in dim_rows:
+                st.markdown(f"**{k}:** {v}")
+        if crop_enabled and not _is_missing(crop_pct):
+            try:
+                st.markdown(f"**Crop %:** {int(float(crop_pct) * 100)}%")
+            except Exception:
+                pass
+    
+    with col_img:
+        st.markdown("##### 📷 Document Page")
+        tabs = st.tabs(["🖼️ Original", "📷 Resized", "✂️ Cropped"])
+        
+        url_o = p.get("original_url")
+        url_r = p.get("resized_url") or p.get("page")
+        url_c = p.get("cropped_url")
+        
+        with tabs[0]:
+            if url_o:
+                st.image(url_o, width=600)
+                st.caption(f"📐 **{orig_w}×{orig_h}**")
+            else:
+                st.info("No original image available")
+        
+        with tabs[1]:
+            if url_r:
+                st.image(url_r, width=600)
+                st.caption(f"📐 **{resize_w}×{resize_h}**")
+            else:
+                st.info("No resized image available")
+        
+        with tabs[2]:
+            if url_c:
+                # Display on a checkerboard background to make the crop boundary obvious.
+                w_caption = f"{crop_w}×{crop_h}" if (not _is_missing(crop_w) and not _is_missing(crop_h)) else None
+                pct_caption = None
+                if not _is_missing(crop_pct):
+                    try:
+                        pct_caption = f"{int(float(crop_pct) * 100)}%"
+                    except Exception:
+                        pct_caption = None
+                st.markdown(
+                    f"""
+                    <div style="
+                        width: 600px;
+                        padding: 14px;
+                        border-radius: 10px;
+                        background-image:
+                          linear-gradient(45deg, #e6e6e6 25%, transparent 25%),
+                          linear-gradient(-45deg, #e6e6e6 25%, transparent 25%),
+                          linear-gradient(45deg, transparent 75%, #e6e6e6 75%),
+                          linear-gradient(-45deg, transparent 75%, #e6e6e6 75%);
+                        background-size: 24px 24px;
+                        background-position: 0 0, 0 12px, 12px -12px, -12px 0px;
+                        box-shadow: 0 10px 30px rgba(0,0,0,0.18);
+                        display: inline-block;
+                    ">
+                        <img src="{url_c}" style="width: 100%; border-radius: 6px; display:block;" />
+                    </div>
+                    """,
+                    unsafe_allow_html=True,
+                )
+                cap = []
+                if w_caption:
+                    cap.append(f"📐 **{w_caption}**")
+                if pct_caption:
+                    cap.append(f"Crop: {pct_caption}")
+                if cap:
+                    st.caption(" | ".join(cap))
+            else:
+                st.info("No cropped image available")
+        
+        with st.expander("🔗 Image URLs"):
+            if url_o:
+                st.code(url_o, language=None)
+            if url_r and url_r != url_o:
+                st.code(url_r, language=None)
+            if url_c:
+                st.code(url_c, language=None)
+
+
+def render_sample_explorer(collection_name: str, url: str, api_key: str):
+    sample_for_filters = sample_points_cached(collection_name, 50, 0, url, api_key)
+    datasets = set()
+    doc_ids = set()
+    for pt in sample_for_filters:
+        p = pt.get("payload", {})
+        if ds := p.get("dataset"):
+            datasets.add(ds)
+        if did := (p.get("doc-id") or p.get("filename")):
+            doc_ids.add(did)
+    
+    c1, c2, c3, c4 = st.columns([1, 1, 2, 1])
+    with c1:
+        n_samples = st.slider("Samples", 1, 20, 3, key="pg_n")
+    with c2:
+        seed = st.number_input("Seed", 0, 9999, 42, key="pg_seed")
+    with c3:
+        filter_ds = st.selectbox("Dataset", ["All"] + sorted(datasets), key="pg_filter_ds")
+    with c4:
+        st.write("")
+        do_sample = st.button("🎲 Sample", type="primary", key="pg_sample_btn")
+    
+    if do_sample:
+        points = sample_points_cached(collection_name, n_samples * 5, seed, url, api_key)
+        if filter_ds != "All":
+            points = [p for p in points if p.get("payload", {}).get("dataset") == filter_ds]
+        points = points[:n_samples]
+        st.session_state["pg_points"] = points
+    
+    points = st.session_state.get("pg_points", [])
+    
+    if not points:
+        st.caption("Click 'Sample' to load documents")
+        return
+    
+    st.success(f"**{len(points)} points loaded**")
+    
+    for i, pt in enumerate(points):
+        p = pt.get("payload", {})
+        
+        filename = p.get("filename") or p.get("doc_id") or p.get("source_doc_id") or "Unknown"
+        page_num = p.get("page_number") or p.get("page") or "?"
+        
+        with st.expander(f"**{i+1}.** {str(filename)[:40]} - Page {page_num}", expanded=(i == 0)):
+            render_document_details(pt, p)
+
+
+def render_rag_query_interface(collection_name: str, model_name: str = None):
+    if not collection_name:
+        return
+    
+    url, api_key = get_qdrant_credentials()
+    
+    if not model_name:
+        points = sample_points_cached(collection_name, 1, 0, url, api_key)
+        if points:
+            model_name = points[0].get("payload", {}).get("model_name")
+        if not model_name:
+            model_name = AVAILABLE_MODELS[1]
+    
+    st.caption(f"Model: **{model_name.split('/')[-1] if model_name else 'auto'}**")
+    
+    c1, c2, c3 = st.columns([2, 1, 1])
+    with c2:
+        mode = st.selectbox("Mode", RETRIEVAL_MODES, index=0, key="q_mode")
+    with c3:
+        top_k = st.slider("Top K", 1, 30, 10, key="q_topk")
+    
+    prefetch_k, stage1_mode, stage1_k, stage2_k = 256, "tokens_vs_tiles", 1000, 300
+    
+    if mode == "two_stage":
+        cc1, cc2 = st.columns(2)
+        with cc1:
+            stage1_mode = st.selectbox("Stage1", STAGE1_MODES, key="q_s1mode")
+        with cc2:
+            prefetch_k = st.slider("Prefetch K", 50, 500, 256, key="q_pk")
+    elif mode == "three_stage":
+        cc1, cc2 = st.columns(2)
+        with cc1:
+            stage1_k = st.number_input("Stage1 K", 100, 5000, 1000, key="q_s1k")
+        with cc2:
+            stage2_k = st.number_input("Stage2 K", 50, 1000, 300, key="q_s2k")
+    
+    with c1:
+        query = st.text_input("Query", placeholder="Enter your search query...", key="q_text")
+    
+    if st.button("🔍 Search", type="primary", disabled=not query, key="q_search"):
+        with st.spinner("Searching..."):
+            results, err = search_collection(
+                collection_name, query, top_k, mode, prefetch_k, stage1_mode, stage1_k, stage2_k, model_name
+            )
+            if err:
+                st.error("Search failed")
+                st.code(err)
+            else:
+                st.session_state["q_results"] = results
+    
+    results = st.session_state.get("q_results", [])
+    if results:
+        st.success(f"**{len(results)} results**")
+        max_score = max(r.get("score_final", r.get("score_stage1", 0)) for r in results) or 1
+        
+        for i, r in enumerate(results):
+            p = r.get("payload", {})
+            score = r.get("score_final", r.get("score_stage1", 0))
+            rel = score / max_score * 100
+            
+            filename = p.get("filename") or p.get("doc_id") or p.get("source_doc_id") or "Unknown"
+            page_num = p.get("page_number") or p.get("page") or "?"
+            
+            with st.expander(f"**#{i+1}** {str(filename)[:35]} - Page {page_num} | 🎯 {rel:.0f}%", expanded=(i < 3)):
+                render_document_details(r, p, score=score, rel_pct=rel)

demo/ui/sidebar.py

@@ -0,0 +1,162 @@
+"""Sidebar component."""
+
+import os
+import streamlit as st
+
+from demo.qdrant_utils import (
+    get_qdrant_credentials,
+    init_qdrant_client_with_creds,
+    get_collections,
+    get_collection_stats,
+    sample_points_cached,
+    get_vector_sizes,
+)
+
+
+def render_sidebar():
+    with st.sidebar:
+        st.subheader("🔑 Qdrant Credentials")
+        
+        env_url = os.getenv("SIGIR_QDRANT_URL") or os.getenv("DEST_QDRANT_URL") or os.getenv("QDRANT_URL") or ""
+        env_key = os.getenv("SIGIR_QDRANT_KEY") or os.getenv("SIGIR_QDRANT_API_KEY") or os.getenv("DEST_QDRANT_API_KEY") or os.getenv("QDRANT_API_KEY") or ""
+        
+        if "qdrant_url_input" not in st.session_state:
+            st.session_state["qdrant_url_input"] = env_url
+        if "qdrant_key_input" not in st.session_state:
+            st.session_state["qdrant_key_input"] = env_key
+        
+        qdrant_url = st.text_input(
+            "Qdrant URL",
+            value=st.session_state["qdrant_url_input"],
+            key="qdrant_url_widget",
+            placeholder="https://xxx.cloud.qdrant.io:6333",
+        )
+        qdrant_key = st.text_input(
+            "API Key",
+            value=st.session_state["qdrant_key_input"],
+            key="qdrant_key_widget",
+            type="password",
+        )
+        
+        if qdrant_url != st.session_state["qdrant_url_input"] or qdrant_key != st.session_state["qdrant_key_input"]:
+            st.session_state["qdrant_url_input"] = qdrant_url
+            st.session_state["qdrant_key_input"] = qdrant_key
+            get_collections.clear()
+            get_collection_stats.clear()
+            sample_points_cached.clear()
+        
+        st.divider()
+        
+        st.subheader("📡 Status")
+        url, api_key = get_qdrant_credentials()
+        client, err = init_qdrant_client_with_creds(url, api_key)
+        
+        col_s1, col_s2 = st.columns(2)
+        with col_s1:
+            if client:
+                st.success("Qdrant ✓", icon="✅")
+            else:
+                st.error("Qdrant ✗", icon="❌")
+        with col_s2:
+            cloudinary_ok = all([os.getenv("CLOUDINARY_CLOUD_NAME"), os.getenv("CLOUDINARY_API_KEY")])
+            if cloudinary_ok:
+                st.success("Cloudinary ✓", icon="✅")
+            else:
+                st.warning("Cloudinary ✗", icon="⚠️")
+        
+        st.divider()
+        
+        with st.expander("📦 Collection", expanded=True):
+            collections = get_collections(url, api_key)
+            if collections:
+                prev_collection = st.session_state.get("active_collection")
+                selected = st.selectbox(
+                    "Select Collection",
+                    options=collections,
+                    key="sidebar_collection",
+                    label_visibility="collapsed",
+                )
+                if selected:
+                    if selected != prev_collection:
+                        st.session_state["model_loaded"] = False
+                        st.session_state["loaded_model_key"] = None
+                    st.session_state["active_collection"] = selected
+                    stats = get_collection_stats(selected)
+                    if "error" not in stats:
+                        col1, col2 = st.columns(2)
+                        col1.metric("Points", f"{stats.get('points_count', 0):,}")
+                        status_raw = stats.get("status", "unknown").replace("CollectionStatus.", "").lower()
+                        status_icon = "🟢" if status_raw == "green" else "🟡" if status_raw == "yellow" else "🔴"
+                        col2.metric("Status", status_icon)
+                        
+                        points = stats.get("points_count", 0)
+                        indexed = stats.get("indexed_vectors_count", 0) or 0
+                        is_indexed = indexed >= points and points > 0
+                        col3, col4 = st.columns(2)
+                        col3.metric("Indexed", f"{indexed:,}")
+                        col4.metric("HNSW", "✅" if is_indexed else "⏳")
+                        
+                        vector_info = stats.get("vector_info", {})
+                        if vector_info:
+                            st.markdown("---")
+                            st.markdown("**🔢 Vectors**")
+                            vec_sizes = get_vector_sizes(selected, url, api_key)
+                            rows = []
+                            sorted_names = sorted(vector_info.keys(), key=lambda x: len(x))
+                            for vname in sorted_names:
+                                vinfo = vector_info[vname]
+                                dim = vinfo.get("size", "?")
+                                num_vec = vec_sizes.get(vname, vinfo.get("num_vectors", 1))
+                                dtype = vinfo.get("datatype", "?").upper()
+                                on_disk = vinfo.get("on_disk", False)
+                                disk_icon = "💾" if on_disk else "🧠"
+                                dim_str = f"{num_vec}×{dim}"
+                                rows.append(f"<tr><td style='text-align:left;padding-right:12px;'><code>{vname}</code></td><td style='text-align:right;'>{dim_str}, {dtype}, {disk_icon}</td></tr>")
+                            table_html = f"<table style='width:100%;font-size:0.85em;'>{''.join(rows)}</table>"
+                            st.markdown(table_html, unsafe_allow_html=True)
+                    else:
+                        st.error("Error loading stats")
+            else:
+                st.info("No collections")
+        
+        with st.expander("⚙️ Admin", expanded=False):
+            active = st.session_state.get("active_collection")
+            if active and client:
+                stats = get_collection_stats(active)
+                vector_info = stats.get("vector_info", {})
+                if vector_info:
+                    st.markdown("**Change Storage**")
+                    vector_names = sorted(vector_info.keys())
+                    sel_vec = st.selectbox("Vector", vector_names, key="admin_vec")
+                    if sel_vec:
+                        current_on_disk = vector_info.get(sel_vec, {}).get("on_disk", False)
+                        current_in_ram = not current_on_disk
+                        st.caption(f"Current: {'🧠 RAM' if current_in_ram else '💾 Disk'}")
+                        target_in_ram = st.toggle("Move to RAM", value=current_in_ram, key=f"admin_ram_{sel_vec}")
+                        if target_in_ram != current_in_ram:
+                            if st.button("💾 Apply Change", key="admin_apply"):
+                                try:
+                                    from qdrant_client.models import VectorParamsDiff
+                                    client.update_collection(
+                                        collection_name=active,
+                                        vectors_config={sel_vec: VectorParamsDiff(on_disk=not target_in_ram)}
+                                    )
+                                    get_collection_stats.clear()
+                                    st.success(f"Updated {sel_vec}")
+                                    st.rerun()
+                                except Exception as e:
+                                    st.error(f"Failed: {e}")
+                        else:
+                            st.caption("Toggle to change storage location")
+                else:
+                    st.info("No vectors")
+            else:
+                st.info("Select a collection")
+        
+        st.divider()
+        
+        if st.button("🔄 Refresh", type="secondary", use_container_width=True):
+            get_collections.clear()
+            get_collection_stats.clear()
+            sample_points_cached.clear()
+            st.rerun()