feat: add gradio web app

pyproject.toml

@@ -20,6 +20,9 @@ langchain-openai = "^0.2.3"
 matplotlib = "^3.9.2"
 pandas = "^2.2.3"
 chromadb = "^0.5.20"
+gradio = "^5.6.0"
+gradio-pdf = "^0.0.19"
+tabulate = "^0.9.0"
 
 
 [build-system]

src/config/output_formatting.py

@@ -341,7 +341,7 @@ def display_search_result_page(
     print("\nChunk distances:")
     print("-" * 80)
     for chunk_type, distance in result.chunk_distances.items():
-        status = f"distance: {distance:.3f}" if distance is not None else "not matched"
+        status = f"{distance:.3f}" if distance is not None else "not matched"
         print(f"{chunk_type}: {status}")
 
     # Display all chunks content

src/rag/storage.py

@@ -57,6 +57,10 @@ class ScoredChunk(BaseModel):
         """Get chunk type from metadata"""
         return self.document.metadata["chunk_type"]
 
+    @property
+    def page_num(self) -> int:
+        return int(self.document.metadata["page_num"])
+
     model_config = ConfigDict(arbitrary_types_allowed=True)
 
 
@@ -97,10 +101,14 @@ class SearchResultPage(BaseModel):
     def best_score(self):
         return self.matched_chunk.score
 
+    @property
+    def page_num(self):
+        return self.matched_chunk.page_num
+
     model_config = ConfigDict(arbitrary_types_allowed=True)
 
 
-class PresentationSearchResult(BaseModel):
+class SearchResultPresentation(BaseModel):
     """Container for presentation-level search results
 
     Represents all matching slides from a single presentation
@@ -166,7 +174,7 @@ class SlideIndexer:
         metadata = dict(
             # Basic slide info
             pdf_path=str(slide.pdf_path),
-            page_num=str(slide.page_num),
+            page_num=str(slide.page_num), # BUG: why str?
             # Chunk specific
             chunk_type=chunk_type,
             slide_id=f"{slide.pdf_path.stem}__{slide.page_num}",
@@ -453,7 +461,7 @@ class ChromaSlideStore:
         search_results = self.search_query(
             query=query,
             chunk_types=chunk_types,
-            n_results=n_results * 5,  # Get more to handle duplicates
+            n_results=n_results * 3,  # Get more to handle duplicates
             max_score=max_distance,
             metadata_filter=metadata_filter,
         )
@@ -500,20 +508,20 @@ class ChromaSlideStore:
             )
             page_results.append(result)
 
-            if len(page_results) == n_results:
-                break
+            # if len(page_results) == n_results:
+            #     break
 
-        return page_results[:n_results]
+        return page_results # [:n_results]
 
     def search_query_presentations(
         self,
         query: str,
         chunk_types: Optional[List[str]] = None,
         n_results: int = 3,
-        n_slides_per_presentation: int = 4,
+        n_slides_per_presentation: int = 3,
         max_distance: float = 2.0,
         metadata_filter: Optional[Dict] = None,
-    ) -> List[PresentationSearchResult]:
+    ) -> List[SearchResultPresentation]:
         """Search presentations based on query and return grouped results
 
         Args:
@@ -532,7 +540,7 @@ class ChromaSlideStore:
         search_results = self.search_query_pages(
             query=query,
             chunk_types=chunk_types,
-            n_results=n_results * n_slides_per_presentation * 2,
+            n_results=n_results * n_slides_per_presentation,
             max_distance=max_distance,
             metadata_filter=metadata_filter,
         )
@@ -553,12 +561,12 @@ class ChromaSlideStore:
             if len(presentations_map[pres_name]) < n_slides_per_presentation:
                 presentations_map[pres_name].append(result)
 
-        # Convert to PresentationSearchResult objects
+        # Convert to SearchResultPresentation objects
         presentation_results = []
 
         for pres_name, slides in presentations_map.items():
             # Create presentation result
-            pres_result = PresentationSearchResult(
+            pres_result = SearchResultPresentation(
                 slides=slides,
                 # NOTE: This is only for testing. Can be removed
                 metadata=dict(
@@ -570,13 +578,13 @@ class ChromaSlideStore:
             )
             presentation_results.append(pres_result)
 
-            if len(presentation_results) == n_results:
-                break
+            # if len(presentation_results) == n_results:
+            #     break
 
         # TODO: Gotta check different ways to sort
         presentation_results.sort(key=lambda x: x.mean_score)
 
-        return presentation_results[:n_results]
+        return presentation_results # [:n_results]
 
     def get_by_metadata(
         self, metadata_filter: Dict, n_results: Optional[int] = None

src/webapp/app.py

@@ -0,0 +1,286 @@
+import argparse
+import logging
+from pathlib import Path
+from textwrap import dedent
+from typing import Dict, List, Optional, Tuple
+
+import gradio as gr
+import pandas as pd
+from gradio_pdf import PDF
+from pymupdf.mupdf import ll_pdf_annot_modification_date
+
+from src.config import Config, Navigator
+from src.rag.storage import ChromaSlideStore, SearchResultPage, SearchResultPresentation
+
+logger = logging.getLogger(__name__)
+
+
+def format_page_results(result_page: SearchResultPage) -> str:
+    """Format individual slide results as markdown text"""
+    chunks = result_page.slide_chunks
+
+    text = dedent(
+        f"""\
+         ### Page: {result_page.page_num+1}
+         **Best matching chunk:** `{result_page.matched_chunk.chunk_type}`\\
+         **Chunk distances:** {result_page.matched_chunk.score:.4f}
+         """
+    )
+
+    # chunk_distances_str = ""
+    # for chunk_type, distance in result_page.chunk_distances.items():
+    #     distance_str = f"{distance:.4f}" if distance else "`not matched`"
+    #     chunk_distances_str += f"{chunk_type}: {distance_str}\\\n"
+
+    chunk_df = (
+        pd.DataFrame(result_page.chunk_distances, index=["distance"])
+        .T.assign(
+            distance=lambda df_: df_["distance"].apply(
+                lambda x: f"{x:.4f}" if x is not None else "not matched"
+            )
+        )
+        .reset_index(names="chunk type")
+        .sort_values("distance")
+    )
+    chunk_distances_str = chunk_df.to_markdown(index=False)
+    text += f"\n{chunk_distances_str}\n"
+
+    # Add matched chunks info
+    text += "#### Content:\n"
+    for i, (chunk_type, distance) in chunk_df.iterrows():
+        if distance != "not matched":
+            text += f"`{chunk_type}` d={distance}\n"
+
+            # Create an embed for text
+            chunk_text = chunks[chunk_type].page_content.replace("\n", "\n>\n> ")
+            chunk_text = "> " + chunk_text + "\n\n"  # Include first line into embed
+            text += chunk_text
+
+    return text
+
+
+def format_presentation_results(
+    pres_result: SearchResultPresentation,
+) -> Tuple[str, Path, int]:
+    """Format single presentation results"""
+    # Get best matching page
+    best_slide = pres_result.best_slide
+    pdf_path = Path(best_slide.pdf_path)
+    page_num = int(best_slide.page_num)
+
+    page_nums = [s.page_num + 1 for s in pres_result.slides]
+    page_scores = [s.best_score for s in pres_result.slides]
+    df = pd.DataFrame(
+        dict(
+            page_nums=page_nums,
+            page_scores=[f"{x:.4f}" for x in page_scores],
+        )
+    )
+
+    df_string = df.to_markdown(index=False)
+
+    # Format header
+    text = f"## {pdf_path.stem}\n"
+    text += f"\n{df_string}\n\n"
+    text += f"**Mean Score:** {pres_result.mean_score:.4f}\n"
+
+    # Format individual slides
+    for slide in pres_result.slides:
+        text += format_page_results(slide)
+        text += "\n\n---\n\n"
+
+    return text, pdf_path, page_num
+
+
+class RagInterface:
+    """Gradio interface for RAG application"""
+
+    def __init__(self, store: ChromaSlideStore, config: Optional[Config] = None):
+        """Initialize interface
+
+        Args:
+            store: Configured vector store
+            config: Optional application config
+        """
+        self.store = store
+        self.config = config or Config()
+        self.nav = self.config.navigator
+
+        # Create interface
+        self.interface = gr.Blocks()
+        self._build_interface()
+
+    def _build_interface(self):
+        """Build Gradio interface layout"""
+        with self.interface:
+            gr.Markdown("# Presentation Search")
+
+            with gr.Row():
+                # Input components
+                with gr.Column(scale=2):
+                    query = gr.Textbox(
+                        label="Search Query",
+                        placeholder="Enter your search query...",
+                        lines=3,
+                    )
+                    with gr.Row():
+                        n_results = gr.Number(
+                            label="Number of Presentations",
+                            scale=1,
+                            minimum=1,
+                            maximum=10,
+                            value=3,
+                            step=1,
+                        )
+                        n_pages_per_pres = gr.Number(
+                            label="Number of pages per presentation",
+                            scale=1,
+                            minimum=1,
+                            maximum=5,
+                            value=2,
+                            step=1,
+                        )
+                        max_distance = gr.Number(
+                            label="Maximum Distance",
+                            scale=1,
+                            minimum=0.1,
+                            maximum=2.0,
+                            value=2.0,
+                            step=0.1,
+                        )
+
+                        search_btn = gr.Button("Search", size="lg", scale=3)
+
+            # Results container
+            with gr.Column(scale=3):
+                with gr.Tabs() as results_tabs:
+                    # Create 3 identical result tabs
+                    result_components = []
+                    for i in range(3):
+                        with gr.Tab(f"Result {i+1}"):
+                            with gr.Column():
+                                # PDF viewer
+                                pdf = PDF(
+                                    label="Presentation",
+                                    height=500,
+                                    interactive=False,
+                                    visible=False,
+                                )
+                                # Results text
+                                results = gr.Markdown(
+                                    label="Search Results", visible=False
+                                )
+                            result_components.append((pdf, results))
+
+            # Wire up the search function
+            search_btn.click(
+                fn=self._search,
+                inputs=[query, n_results, n_pages_per_pres, max_distance],
+                outputs=[item for pair in result_components for item in pair],
+            )
+
+    def _search(
+        self,
+        query: str,
+        n_results: int,
+        n_pages: int,
+        max_distance: float,
+    ) -> List[gr.components.Component]:
+        """Search presentations and format results
+
+        Args:
+            query: Search query text
+            n_results: Number of presentations to return
+            max_distance: Maximum cosine distance threshold
+
+        Returns:
+            List of components to update in UI
+        """
+        try:
+            # Search presentations
+            results = self.store.search_query_presentations(
+                query=query,
+                n_results=n_results,
+                max_distance=max_distance,
+                n_slides_per_presentation=n_pages,
+            )
+
+            # Prepare outputs for all possible tabs
+            outputs = []
+            for i in range(3):
+                if i < len(results):
+                    # Format this result
+                    text, pdf_path, page = format_presentation_results(results[i])
+
+                    # Add components: PDF viewer and results text
+                    outputs.extend(
+                        [
+                            # PDF component
+                            PDF(
+                                value=str(pdf_path),
+                                starting_page=page
+                                + 1,  # Pages are 0-based in store but 1-based in PDF
+                                visible=True,
+                            ),
+                            # Results text
+                            gr.Markdown(value=text, visible=True),
+                        ]
+                    )
+                else:
+                    # Hide unused tabs
+                    outputs.extend(
+                        [
+                            PDF(visible=False),
+                            gr.Markdown(visible=False),
+                        ]
+                    )
+
+            return outputs
+
+        except Exception as e:
+            logger.exception("Search failed")
+            # Return empty results on error
+            return [PDF(visible=False), gr.Markdown(visible=False)] * 3
+
+    def launch(self, **kwargs):
+        """Launch the Gradio interface"""
+        self.interface.launch(**kwargs)
+
+
+def run_app(store: ChromaSlideStore, **kwargs):
+    """Run Gradio application
+
+    Args:
+        store: Configured ChromaSlideStore instance
+        **kwargs: Additional arguments for Gradio launch
+    """
+    viewer = RagInterface(store)
+    viewer.launch(**kwargs)
+
+
+def main():
+    """Run presentation search web application"""
+    # Load environment
+    from dotenv import load_dotenv
+
+    load_dotenv()
+
+    # Parse arguments
+    parser = argparse.ArgumentParser()
+    parser.add_argument(
+        "--collection", default="pres0", help="ChromaDB collection name"
+    )
+    parser.add_argument("--host", default="0.0.0.0", help="Host to run on")
+    parser.add_argument("--port", type=int, default=7860, help="Port to run on")
+    parser.add_argument("--share", action="store_true", help="Create public link")
+    args = parser.parse_args()
+
+    # Initialize store
+    store = ChromaSlideStore(collection_name=args.collection)
+
+    # Run app
+    run_app(store, server_name=args.host, server_port=args.port, share=args.share)
+
+
+if __name__ == "__main__":
+    main()