fe/gradio_app.py

import sys
import os
import gradio as gr
import requests
from config import BACKEND_HOST, BACKEND_PORT, FRONTEND_PORT, APP_NAME


# Add the project root directory to the Python path
sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
API_URL = f"http://{BACKEND_HOST}:{BACKEND_PORT}/search"


def search_prompt(query: str, n: int) -> list:
    """
    Query the backend search API and format the results.
    Args:
        query: The search query.
        n: The number of results to return.
    Returns:
        list: A list of formatted results, or an error message
    """
    try:
        response = requests.get(API_URL, params={"query": query, "n": n})
        response.raise_for_status()
        data = response.json()

        # Extract and sort results
        results = data["results"]
        sorted_results = sorted(results, key=lambda x: x[0], reverse=True)

        # Format results in a table
        tabular_results = [[result[1], f"{result[0]:.4f}"] for result in sorted_results]
        return tabular_results

    except Exception as e:
        return [["Error", str(e)]]


with gr.Blocks() as demo:
    gr.Markdown(f"# {APP_NAME}")

    query = gr.Textbox(label="Enter your query:")
    n = gr.Number(label="Number of top results", value=5, precision=0)
    output = gr.Dataframe(
        headers=["Prompt", "Similarity"],
        datatype=["str", "str"],
        interactive=False,
        label=f"Top results"
    )

    search_button = gr.Button("Search")
    search_button.click(search_prompt, inputs=[query, n], outputs=[output])


demo.launch(server_name="0.0.0.0", server_port=FRONTEND_PORT)