import gradio as gr
from vectordb import Memory

# Initialize Memory
memory = Memory()

# Define a function to save new text and metadata
def save_data(texts, metadata):
    try:
        # Split texts and metadata by lines for simplicity
        text_list = texts.strip().split("\n")
        metadata_list = [eval(meta.strip()) for meta in metadata.strip().split("\n")]
        memory.save(text_list, metadata_list)
        return "Data saved successfully!"
    except Exception as e:
        return f"Error saving data: {e}"

# Define a function for querying
def search_query(query, top_n):
    try:
        results = memory.search(query, top_n=int(top_n))  # Search for top_n results
        return results
    except Exception as e:
        return f"Error during search: {e}"

# Create Gradio interface
with gr.Blocks() as demo:
    gr.Markdown("### VectorDB Search App")
    
    # Save Data Section
    gr.Markdown("#### Save Data")
    with gr.Row():
        input_texts = gr.Textbox(
            label="Enter text (one per line)", 
            lines=5, 
            placeholder="Example:\napples are green\noranges are orange"
        )
        input_metadata = gr.Textbox(
            label="Enter metadata (one per line, matching the texts)", 
            lines=5, 
            placeholder='Example:\n{"url": "https://apples.com"}\n{"url": "https://oranges.com"}'
        )
    save_button = gr.Button("Save Data")
    save_status = gr.Textbox(label="Status", interactive=False)
    save_button.click(save_data, inputs=[input_texts, input_metadata], outputs=save_status)
    
    # Search Section
    gr.Markdown("#### Search")
    with gr.Row():
        input_query = gr.Textbox(label="Enter your query")
        input_top_n = gr.Number(label="Top N results", value=1)
        output_result = gr.Textbox(label="Search Results", interactive=False)

    search_button = gr.Button("Search")
    search_button.click(search_query, inputs=[input_query, input_top_n], outputs=output_result)

# Run the Gradio app
demo.launch()