import gradio as gr
import requests

API_URL = "https://sendthat.cc"

def search_document(index_name, query, k):
    url = f"{API_URL}/search/{index_name}"
    payload = {"text": query, "k": k}
    headers = {"Content-Type": "application/json"}
    response = requests.post(url, json=payload, headers=headers)
    results = response.json()
    
    formatted_results = []
    for result in results.get('results', []):
        metadata = result.get('metadata', {})
        formatted_result = f"Source: {metadata.get('source', 'Unknown')}\n"
        formatted_result += f"Page: {metadata.get('page', 'Unknown')}\n"
        formatted_result += f"Content: {metadata.get('content', 'No content available')}\n"
        formatted_result += f"Distance: {result.get('distance', 'Unknown')}\n"
        formatted_results.append(formatted_result)
    
    return "\n\n".join(formatted_results)

def qa_document(index_name, question, k):
    url = f"{API_URL}/qa/{index_name}"
    payload = {"text": question, "k": k}
    headers = {"Content-Type": "application/json"}
    response = requests.post(url, json=payload, headers=headers)
    result = response.json()
    
    answer = result.get('answer', 'No answer available')
    sources = result.get('sources', [])
    
    formatted_sources = []
    for source in sources:
        formatted_source = f"Source: {source.get('source', 'Unknown')}\n"
        formatted_source += f"Relevance Score: {source.get('relevance_score', 'Unknown')}"
        formatted_sources.append(formatted_source)
    
    formatted_result = f"Answer: {answer}\n\nSources:\n" + "\n\n".join(formatted_sources)
    return formatted_result

with gr.Blocks() as demo:
    gr.Markdown("# Document Search and Question Answering System")
    
    index_name = gr.Textbox(label="Index Name", value="default")
    
    with gr.Tab("Search"):
        search_input = gr.Textbox(label="Search Query")
        search_k = gr.Slider(1, 10, 5, step=1, label="Number of Results")
        search_button = gr.Button("Search")
        search_output = gr.Textbox(label="Search Results", lines=10)
        
        search_button.click(search_document, inputs=[index_name, search_input, search_k], outputs=search_output)
    
    with gr.Tab("Question Answering"):
        qa_input = gr.Textbox(label="Question")
        qa_k = gr.Slider(1, 10, 5, step=1, label="Number of Contexts to Consider")
        qa_button = gr.Button("Ask Question")
        qa_output = gr.Textbox(label="Answer and Sources", lines=10)
        
        qa_button.click(qa_document, inputs=[index_name, qa_input, qa_k], outputs=qa_output)

demo.launch()