import gradio as gr
import requests
import json
import os

default_prompt = "The best thing about being a cat is"

def generate_completion(api_key, model, prompt, temperature, repetition_penalty, stop_phrase, max_tokens):
    if not api_key:
        return "Please enter your OpenRouter API key."

    try:
        # Prepare the request payload
        payload = {
            "model": model,
            "prompt": prompt,
            "temperature": temperature,
            "frequency_penalty": repetition_penalty,
            "max_tokens": max_tokens,
            "provider": {
                "require_parameters": True
            }
        }

        # Add stop phrase if provided
        if stop_phrase:
            payload["stop"] = [stop_phrase]

        # Prepare headers
        headers = {
            "Authorization": f"Bearer {api_key}",
            "Content-Type": "application/json"
        }

        # Make the API request
        response = requests.post(
            "https://openrouter.ai/api/v1/completions",
            headers=headers,
            json=payload,
            timeout=60
        )

        # Check if request was successful
        response.raise_for_status()

        # Parse the response
        result = response.json()

        # Extract the completion text
        if "choices" in result and len(result["choices"]) > 0:
            return result["choices"][0]["text"].strip()
        else:
            return "No completion generated."

    except requests.exceptions.RequestException as e:
        return f"Request error: {str(e)}"
    except json.JSONDecodeError as e:
        return f"JSON decode error: {str(e)}"
    except Exception as e:
        return f"An error occurred: {str(e)}"

def append_completion(prompt, completion):
    new_prompt = f"{prompt}{completion}".strip()
    return new_prompt, ""  # Return new prompt and empty completion

def clear_fields():
    return "", ""

with gr.Blocks(theme=gr.themes.Soft()) as iface:
    gr.Markdown("# Base Model Completion Interface via OpenRouter")

    # API Key input at the top
    api_key_input = gr.Textbox(
        label="OpenRouter API Key",
        type="password",
        placeholder="Enter your OpenRouter API key"
    )

    with gr.Row():
        with gr.Column(scale=2):
            prompt_input = gr.Textbox(label="Prompt", lines=6, value="The best thing about being a cat is")
        with gr.Column(scale=1):
            model_input = gr.Textbox(
                label="Model",
                value="meta-llama/llama-3.1-405b",
                placeholder="Enter model name"
            )
            temperature_slider = gr.Slider(minimum=0, maximum=1, value=0.7, step=0.1, label="Temperature")
            repetition_penalty_slider = gr.Slider(minimum=0, maximum=2, value=0.1, step=0.1, label="Repetition Penalty")
            max_tokens_slider = gr.Slider(minimum=1, maximum=4000, value=250, step=1, label="Max Tokens")
            stop_phrase_input = gr.Textbox(label="Stop Phrase", placeholder="Enter stop phrase (optional)")

    with gr.Row():
        generate_button = gr.Button("Generate Completion")
        append_button = gr.Button("Append Completion to Prompt")
        clear_button = gr.Button("Clear All Fields")

    output_text = gr.Textbox(label="Generated Completion", lines=10)

    generate_button.click(
        generate_completion,
        inputs=[api_key_input, model_input, prompt_input, temperature_slider, repetition_penalty_slider, stop_phrase_input, max_tokens_slider],
        outputs=output_text
    )

    append_button.click(
        append_completion,
        inputs=[prompt_input, output_text],
        outputs=[prompt_input, output_text]
    )

    clear_button.click(
        clear_fields,
        outputs=[prompt_input, output_text]
    )

iface.launch(share=True)