import gradio as gr
import subprocess

def run_ollama_prompt(prompt, model):
    """Runs a prompt using Ollama CLI and returns the response."""
    command = ["ollama", "run", model, prompt]

    try:
        result = subprocess.run(command, capture_output=True, text=True, check=True, encoding="utf-8")
        response = result.stdout.strip()
        response = response.replace("<think>", "").replace("</think>", "").strip()
        return response
    except subprocess.CalledProcessError as e:
        return f"Error running Ollama: {e}"

# Gradio Interface
def generate_response(prompt, model):
    return run_ollama_prompt(prompt, model)

models = ["deepseek-r1:1.5b", "deepseek-r1:14b", "deepseek-r1:32b", "phi4:latest", "llama2:7b"]

with gr.Blocks() as app:
    gr.Markdown("# Ollama Chat App")
    
    with gr.Row():
        model_selector = gr.Dropdown(models, label="Select Model", value="deepseek-r1:1.5b")
    
    with gr.Row():
        prompt_input = gr.Textbox(label="Enter your prompt", placeholder="Type something...")
    
    with gr.Row():
        submit_button = gr.Button("Generate Response")
    
    response_output = gr.Textbox(label="Response", interactive=False)
    
    submit_button.click(fn=generate_response, inputs=[prompt_input, model_selector], outputs=response_output)
 


app.launch()