import gradio as gr
import requests

# Base URL for LM Studio REST API
BASE_URL = "http://localhost:1234/api/v0"

# Chat completions function
def chat_with_lmstudio(messages, model, temperature=0.7, max_tokens=150):
    endpoint = f"{BASE_URL}/chat/completions"
    payload = {
        "model": model,
        "messages": messages,
        "temperature": temperature,
        "max_tokens": max_tokens,
        "stream": False
    }
    try:
        response = requests.post(endpoint, json=payload)
        response.raise_for_status()
        data = response.json()
        return data["choices"][0]["message"]["content"]
    except requests.RequestException as e:
        return f"Error: {str(e)}"

# Text completions function
def text_completion(prompt, model, temperature=0.7, max_tokens=150):
    endpoint = f"{BASE_URL}/completions"
    payload = {
        "model": model,
        "prompt": prompt,
        "temperature": temperature,
        "max_tokens": max_tokens,
        "stream": False
    }
    try:
        response = requests.post(endpoint, json=payload)
        response.raise_for_status()
        data = response.json()
        return data["choices"][0]["text"]
    except requests.RequestException as e:
        return f"Error: {str(e)}"

# Embeddings function
def text_embedding(text, model):
    endpoint = f"{BASE_URL}/embeddings"
    payload = {
        "model": model,
        "input": text
    }
    try:
        response = requests.post(endpoint, json=payload)
        response.raise_for_status()
        data = response.json()
        return data["data"][0]["embedding"]
    except requests.RequestException as e:
        return f"Error: {str(e)}"

# Gradio Interface
def chat_interface(user_message, history, model="granite-3.0-2b-instruct"):
    if history is None:
        history = []
    history.append({"role": "user", "content": user_message})
    assistant_response = chat_with_lmstudio(history, model=model)
    history.append({"role": "assistant", "content": assistant_response})
    conversation = [(h["content"], history[i+1]["content"]) for i, h in enumerate(history[:-1]) if h["role"] == "user"]
    return conversation, history

with gr.Blocks() as demo:
    gr.Markdown("# LM Studio API Interface")

    with gr.Tab("Chat with Model"):
        chat_history = gr.State([])
        chat_model = gr.Textbox(value="granite-3.0-2b-instruct", label="Model")
        chatbot = gr.Chatbot()
        msg = gr.Textbox(placeholder="Enter your message", label="User Input")
        submit_btn = gr.Button("Send")
        submit_btn.click(chat_interface, inputs=[msg, chat_history, chat_model], outputs=[chatbot, chat_history])

    with gr.Tab("Text Completion"):
        completion_prompt = gr.Textbox(placeholder="Enter a prompt for text completion", label="Prompt")
        completion_model = gr.Textbox(value="granite-3.0-2b-instruct", label="Model")
        completion_output = gr.Textbox(label="Completion")
        generate_btn = gr.Button("Generate")
        generate_btn.click(text_completion, inputs=[completion_prompt, completion_model], outputs=completion_output)

    with gr.Tab("Text Embeddings"):
        embedding_text = gr.Textbox(placeholder="Enter text for embeddings", label="Input Text")
        embedding_model = gr.Textbox(value="text-embedding-nomic-embed-text-v1.5", label="Model")
        embedding_output = gr.JSON(label="Embeddings")
        embed_btn = gr.Button("Get Embeddings")
        embed_btn.click(text_embedding, inputs=[embedding_text, embedding_model], outputs=embedding_output)

    demo.launch(share=True)