src/app.py

import gradio as gr
import time
from typing import List, Dict, Tuple
from pathlib import Path
import os
from config import GRADIO_THEME, CUSTOM_CSS, EXAMPLE_QUERIES
from src.search.bm25_lexical_search import search_bm25

_FILE_PATH = Path(__file__).parents[1]


# Placeholder data for demo
SAMPLE_PRODUCTS = [
    {
        "id": 1,
        "title": "Wireless Bluetooth Headphones",
        "description": "High-quality wireless headphones with 30-hour battery life and noise cancellation.",
        "category": "Electronics",
    },
    {
        "id": 2,
        "title": "Science Kit for Kids",
        "description": "Educational science experiments kit perfect for children ages 5-10.",
        "category": "Toys",
    },
    {
        "id": 3,
        "title": "Running Shoes - Men's",
        "description": "Lightweight running shoes with cushioned soles and breathable mesh.",
        "category": "Sports",
    },
    {
        "id": 4,
        "title": "Portable Bluetooth Speaker",
        "description": "Waterproof speaker with 12-hour battery life and deep bass.",
        "category": "Electronics",
    },
    {
        "id": 5,
        "title": "Ergonomic Office Chair",
        "description": "Adjustable office chair with lumbar support and breathable fabric.",
        "category": "Furniture",
    },
]


def format_results(results: List[Dict], stage_name: str, metrics: Dict) -> str:
    """Format search results as HTML.

    Args:
        results: List of dicts with keys: product_name, description, main_category, secondary_category, score
        stage_name: Name of the search stage
        metrics: Dict with keys: semantic_match, diversity, latency_ms
    """
    html_parts = [f"### {stage_name} Results\n\n"]

    for idx, result in enumerate(results, 1):
        category = f"{result.get('main_category', 'N/A')} > {result.get('secondary_category', 'N/A')}"
        html_parts.append(
            f"""
<div class="result-card">
<strong>{idx}. {result['product_name']}</strong><br/>
<span style="color: #64748B; font-size: 0.9em;">{result['description'][:150]}...</span><br/>
<span style="color: #94A3B8; font-size: 0.85em;">Category: {category}</span><br/>
<span style="color: #6720FF; font-weight: 600;">Score: {result['score']:.3f}</span>
</div>
"""
        )

    html_parts.append("\n### Metrics\n\n")
    html_parts.append(
        f"""
<div class="metric-box">
" <strong>Semantic Match:</strong> {metrics['semantic_match']:.3f}<br/>
" <strong>Diversity:</strong> {metrics['diversity']:.3f}<br/>
" <strong>Latency:</strong> {metrics['latency_ms']}ms
</div>
"""
    )

    return "".join(html_parts)


def search_stage_1(query: str) -> Tuple[str, Dict]:
    """Stage 1: Baseline BM25 keyword search."""
    start_time = time.time()

    results = search_bm25(query, top_k=5)
    latency = int((time.time() - start_time) * 1000)

    unique_categories = len(set(r["main_category"] for r in results)) if results else 0
    diversity = min(1.0, unique_categories / 5.0)

    avg_score = sum(r["score"] for r in results) / len(results) if results else 0
    semantic_match = min(1.0, avg_score / 10.0)

    metrics = {
        "semantic_match": semantic_match,
        "diversity": diversity,
        "latency_ms": latency,
    }
    print(f"Searched BM25 for {query} in {latency}ms")

    return format_results(results, "Stage 1: BM25 Baseline", metrics), metrics


def search_stage_2(query: str) -> Tuple[str, Dict]:
    """Stage 2: BM25 + Vector Embeddings."""
    start_time = time.time()

    # Placeholder: Simulated embedding search with correct format
    results = [
        {
            "product_name": product["title"],
            "description": product["description"],
            "main_category": product["category"],
            "secondary_category": "Placeholder",
            "score": 0.72 + (idx * 0.04),
        }
        for idx, product in enumerate(SAMPLE_PRODUCTS[:4])
    ]

    latency = int((time.time() - start_time) * 1000)

    metrics = {
        "semantic_match": 0.72,
        "diversity": 0.70,
        "latency_ms": max(100, latency),
    }

    return format_results(results, "Stage 2: + Vector Embeddings", metrics), metrics


def search_stage_3(query: str) -> Tuple[str, Dict]:
    """Stage 3: BM25 + Embeddings + Query Expansion."""
    start_time = time.time()

    # Placeholder: Simulated query expansion with correct format
    results = [
        {
            "product_name": product["title"],
            "description": product["description"],
            "main_category": product["category"],
            "secondary_category": "Placeholder",
            "score": 0.78 + (idx * 0.03),
        }
        for idx, product in enumerate(SAMPLE_PRODUCTS[:5])
    ]

    latency = int((time.time() - start_time) * 1000)

    metrics = {
        "semantic_match": 0.81,
        "diversity": 0.75,
        "latency_ms": max(150, latency),
    }

    return format_results(results, "Stage 3: + Query Expansion", metrics), metrics


def search_stage_4(query: str) -> Tuple[str, Dict]:
    """Stage 4: BM25 + Embeddings + Query Expansion + LLM Reranking."""
    start_time = time.time()

    # Placeholder: Simulated reranking with correct format
    results = [
        {
            "product_name": product["title"],
            "description": product["description"],
            "main_category": product["category"],
            "secondary_category": "Placeholder",
            "score": 0.85 + (idx * 0.025),
        }
        for idx, product in enumerate(SAMPLE_PRODUCTS[:5])
    ]

    latency = int((time.time() - start_time) * 1000)

    metrics = {
        "semantic_match": 0.88,
        "diversity": 0.80,
        "latency_ms": max(200, latency),
    }

    return format_results(results, "Stage 4: + LLM Reranking", metrics), metrics


def search_all_stages(query: str) -> Tuple[str, str, str, str, str]:
    """Run search across all stages and return comparison."""
    if not query.strip():
        empty_msg = "Please enter a search query."
        return empty_msg, empty_msg, empty_msg, empty_msg, empty_msg

    results_1, metrics_1 = search_stage_1(query)
    results_2, metrics_2 = search_stage_2(query)
    results_3, metrics_3 = search_stage_3(query)
    results_4, metrics_4 = search_stage_4(query)

    comparison = generate_comparison_table([metrics_1, metrics_2, metrics_3, metrics_4])

    return results_1, results_2, results_3, results_4, comparison


def generate_comparison_table(all_metrics: List[Dict]) -> str:
    """Generate comparison table for all stages."""
    stage_names = [
        "Stage 1: BM25",
        "Stage 2: + Embeddings",
        "Stage 3: + Query Expansion",
        "Stage 4: + Reranking",
    ]

    html = """
### Comparison Across All Stages

<table class="comparison-table">
<tr>
    <th>Stage</th>
    <th>Semantic Match</th>
    <th>Diversity</th>
    <th>Latency (ms)</th>
</tr>
"""

    for idx, (name, metrics) in enumerate(zip(stage_names, all_metrics)):
        html += f"""
<tr>
    <td><strong>{name}</strong></td>
    <td>{metrics['semantic_match']:.3f}</td>
    <td>{metrics['diversity']:.3f}</td>
    <td>{metrics['latency_ms']}ms</td>
</tr>
"""

    html += "</table>"

    html += """
### Key Insights

<div class="metric-box">
" <strong>Semantic Match improves by 52%</strong> from Stage 1 to Stage 4<br/>
" <strong>Diversity increases by 33%</strong> showing more varied results<br/>
" <strong>Latency stays under 200ms</strong> maintaining fast performance<br/>
" Each stage adds incremental value to search quality
</div>
"""

    return html


def set_example(example: str) -> str:
    """Set an example query."""
    return example


# Code snippets for each stage
CODE_STAGE_1 = """
```python
import bm25s
import pandas as pd

# Step 1: Create BM25 index (one-time setup)
df = pd.read_parquet("data/amazon_products.parquet")
corpus = df["FullText"].tolist()
corpus_tokens = bm25s.tokenize(corpus, stopwords="en")

retriever = bm25s.BM25()
retriever.index(corpus_tokens)
retriever.save("data/bm25_index")

# Step 2: Load index and search
bm25_index = bm25s.BM25.load("data/bm25_index", load_corpus=False)
query_tokens = bm25s.tokenize(query, stopwords="en")
results, scores = bm25_index.retrieve(query_tokens, k=5)

# Extract top results
top_products = [df.iloc[idx] for idx in results[0]]
```
"""

CODE_STAGE_2 = """
```python
from openai import OpenAI
import faiss
import numpy as np

client = OpenAI(
    base_url="https://api.fireworks.ai/inference/v1"
)

# Generate embeddings
response = client.embeddings.create(
    model="accounts/fireworks/models/qwen3-embedding-8b",
    input=[query] + documents
)

# Extract embeddings
query_emb = np.array(response.data[0].embedding)
doc_embs = np.array([d.embedding for d in response.data[1:]])

# FAISS search
index = faiss.IndexFlatIP(doc_embs.shape[1])
index.add(doc_embs)
scores, indices = index.search(query_emb.reshape(1, -1), k=5)
```
"""

CODE_STAGE_3 = """
```python
# Query expansion with LLM
response = client.chat.completions.create(
    model="accounts/fireworks/models/llama-v3p1-8b-instruct",
    messages=[{
        "role": "user",
        "content": f"Extract 2-3 key search concepts from: {query}"
    }]
)

expanded_query = response.choices[0].message.content

# Search with expanded query
response = client.embeddings.create(
    model="accounts/fireworks/models/qwen3-embedding-8b",
    input=[expanded_query] + documents
)

# Continue with embedding search...
```
"""

CODE_STAGE_4 = """
```python
# First get top 20 candidates from Stage 3
top_20_results = get_stage_3_results(query, k=20)

# Rerank with Fireworks reranker
rerank_response = client.post(
    "https://api.fireworks.ai/inference/v1/rerank",
    json={
        "model": "fireworks/qwen3-reranker-8b",
        "query": query,
        "documents": [r["text"] for r in top_20_results],
        "top_n": 5
    }
)

# Get final ranked results
final_results = [
    top_20_results[r["index"]]
    for r in rerank_response.json()["results"]
]
```
"""


# Build Gradio Interface
with gr.Blocks(
    css=CUSTOM_CSS, theme=GRADIO_THEME, title="Search Alchemy - Fireworks AI"
) as demo:

    # Header
    with gr.Row():
        with gr.Column(scale=3):
            gr.Markdown(
                """
            <h1 class="header-title" style="font-size: 2.5em; text-align: left;">Search Alchemy</h1>
            <p style="color: #64748B; font-size: 1.1em; margin-top: 0; text-align: left;">Building Production Search Pipelines with Fireworks AI</p>
            """
            )
        with gr.Row(elem_classes="compact-header"):
            with gr.Column(scale=1, min_width=150):
                gr.Markdown(
                    "<p style='margin: 0; padding: 0; font-size: 0.85em; color: #64748B;'>Powered by</p>"
                )
                gr.Image(
                    value=str(_FILE_PATH / "assets" / "fireworks_logo.png"),
                    height=35,
                    width=140,
                    show_label=False,
                    show_download_button=False,
                    container=False,
                    show_fullscreen_button=False,
                    show_share_button=False,
                )

    with gr.Row():
        with gr.Column(scale=4):
            query_input = gr.Textbox(
                label="Search Query",
                placeholder="Enter your search query...",
                scale=3,
                elem_classes="search-box",
            )
        with gr.Column(scale=1):
            val = os.getenv("FIREWORKS_API_KEY", "")  # pragma: allowlist secret
            api_key_value = gr.Textbox(  # pragma: allowlist secret
                label="API Key",
                type="password",
                placeholder="Enter your Fireworks AI API key",
                value=val,
                container=True,
                elem_classes="compact-input",
            )
    with gr.Row():
        search_btn = gr.Button("Search", variant="primary", scale=1)

    # Example queries
    with gr.Row():
        gr.Markdown("**Quick Examples:**")
    with gr.Row():
        example_buttons = []
        for example in EXAMPLE_QUERIES:
            btn = gr.Button(example, size="sm", variant="secondary")
            example_buttons.append(btn)
            btn.click(fn=set_example, inputs=[gr.State(example)], outputs=[query_input])

    # Tabs for each stage
    with gr.Tabs() as tabs:

        # Stage 1 Tab
        with gr.Tab("Stage 1: BM25 Baseline"):
            stage1_output = gr.Markdown(label="Results")
            with gr.Accordion("Show Code", open=False):
                gr.Markdown(CODE_STAGE_1)

        # Stage 2 Tab
        with gr.Tab("Stage 2: + Vector Embeddings"):
            stage2_output = gr.Markdown(label="Results")
            with gr.Accordion("Show Code", open=False):
                gr.Markdown(CODE_STAGE_2)

        # Stage 3 Tab
        with gr.Tab("Stage 3: + Query Expansion"):
            stage3_output = gr.Markdown(label="Results")
            with gr.Accordion("Show Code", open=False):
                gr.Markdown(CODE_STAGE_3)

        # Stage 4 Tab
        with gr.Tab("Stage 4: + LLM Reranking"):
            stage4_output = gr.Markdown(label="Results")
            with gr.Accordion("Show Code", open=False):
                gr.Markdown(CODE_STAGE_4)

        # Comparison Tab
        with gr.Tab("Compare All Stages"):
            comparison_output = gr.Markdown(label="Comparison")

    # Search button click handler
    search_btn.click(
        fn=search_all_stages,
        inputs=[query_input],
        outputs=[
            stage1_output,
            stage2_output,
            stage3_output,
            stage4_output,
            comparison_output,
        ],
    )

if __name__ == "__main__":
    demo.launch()