from vectordb import Memory
import gradio as gr
import json
import numpy as np

class CustomEncoder(json.JSONEncoder):
    def default(self, obj):
        if isinstance(obj, np.float32):
            return float(obj)  # Convert float32 to native float
        return super().default(obj)
        
def process_json(json_input):
    try:
        input = json.loads(json_input)
        
        memory = Memory()#embedding_model="TaylorAI/bge-micro-v2")
        memory.save(input['terms'], input['metadata'])

        results = memory.search(input['prompt'], top_n=input['topN'])
        
        return json.dumps(results, indent=4, cls=CustomEncoder)
    except json.JSONDecodeError:
        return "Invalid JSON input."

with gr.Blocks() as demo:
    gr.Markdown("## *VectorDB* based Paragraph Embedder")
    input_json = gr.Textbox(label="Input", lines=10, placeholder='{"topN": 5, "prompt": "yellow", "metadata": [], "terms": ["banana", "blueberry", "apple"]}')
    output_json = gr.Textbox(label="Output", lines=10, interactive=False)
    process_button = gr.Button("Process")
    
    process_button.click(process_json, inputs=input_json, outputs=output_json)

# Launch the app
demo.launch()