main.py

import gradio as gr
from src.utils.dataset_loader import load_dataset
from src.models.model import load_model, run_inference

# Load dataset
df = load_dataset()

# List of available models (example)
models = ["Canstralian/text2shellcommands", "Canstralian/RabbitRedux", "Canstralian/CySec_Known_Exploit_Analyzer"]
prompt_ids = df.index.tolist() if df is not None else []

# Function to simulate conversation with model selection
def chat_interface(user_input, selected_model, prompt_id=None):
    if df is not None and prompt_id is not None:
        prompt = df.iloc[prompt_id]["prompt_text"]  # Replace with the actual column name
        # Run inference on the selected model
        response = run_inference(user_input, selected_model, prompt)
    else:
        response = f"[{selected_model}] says: You entered '{user_input}'. This is a simulated response."
    return response

# Gradio Interface
with gr.Blocks(css="./static/styles.css") as demo:
    with gr.Row():
        gr.Markdown("### Retro Hacker Chat with Debugging Prompts", elem_classes="retro-terminal")
    with gr.Row():
        user_input = gr.Textbox(
            label="Enter your message:",
            placeholder="Type your message here...",
            elem_classes="retro-terminal"
        )
        model_selector = gr.Dropdown(
            choices=models,
            label="Select Model",
            value=models[0],
            elem_classes="retro-terminal"
        )
        if prompt_ids:
            prompt_selector = gr.Dropdown(
                choices=prompt_ids,
                label="Select Debugging Prompt ID",
                value=prompt_ids[0],
                elem_classes="retro-terminal"
            )
        else:
            prompt_selector = None
    with gr.Row():
        response_box = gr.Textbox(
            label="Model Response:",
            placeholder="The model's response will appear here...",
            elem_classes="retro-terminal"
        )
    with gr.Row():
        send_button = gr.Button("Send", elem_classes="retro-terminal")

    # Link input and output
    if prompt_selector:
        send_button.click(chat_interface, inputs=[user_input, model_selector, prompt_selector], outputs=response_box)
    else:
        send_button.click(chat_interface, inputs=[user_input, model_selector], outputs=response_box)

# Launch the interface
demo.launch()