app.py

import gradio as gr
import requests
import json
import os
import markdown
from markdown.extensions.codehilite import CodeHiliteExtension
import markdown.extensions.fenced_code

API_URL = "https://host.palple.polrambora.com/pmsq"

sessions = {}
PRIMARY_SYSTEM_INSTRUCTIONS = "You are P-MSQ (Messaging Service Query), a friendly AI Chatbot that can help in any situations"
ASSISTANT_PIC_PATH = "https://huggingface.co/spaces/PLRMB/P-MSQ-API-PREVIEW/resolve/main/API.png"
USER_PIC_PATH = "https://huggingface.co/spaces/PLRMB/P-MSQ-API-PREVIEW/resolve/main/usr.png"

def render_avatars(userid):
    try:
        response = requests.post(
            'https://host.palple.polrambora.com/userexistence',
            json={'userid': userid},
            timeout=10
        )
        if response.status_code == 200:
            response_json = response.json()
            return response_json["avatar"]["link"]
    except Exception:
        return None

def authorize(user, api_key, system_message):
    test_data = {"user": user, "key": api_key}
    try:
        response = requests.post(
            "https://host.palple.polrambora.com/check_key_impv",
            json=test_data,
        )
        if response.status_code == 200:
            avatar = render_avatars(user) or USER_PIC_PATH
            if api_key not in sessions:
                sessions[api_key] = {
                    "history": [],
                    "headers": {
                        "authorization": api_key,
                        "Content-Type": 'application/json'
                    },
                    "avatar": avatar,
                    "system_message": system_message
                }
            return True
        elif response.status_code == 403:
            return 403
        else:
            return False
    except Exception:
        return False

def respond(message, api_key, max_tokens, top_p, temperature):
    session = sessions.get(api_key, {})
    history = session.get("history", [])
    headers = session.get("headers", {})
    system_message = session.get("system_message", PRIMARY_SYSTEM_INSTRUCTIONS)
    messages = []
    for user_message, assistant_message, user_profile, assistant_profile, user_pic, assistant_pic in history:
        if user_message:
            messages.append({
                "role": "user",
                "content": user_message,
                "profile": user_profile,
                "picture": user_pic
            })
        if assistant_message:
            messages.append({
                "role": "assistant",
                "content": assistant_message,
                "profile": assistant_profile,
                "picture": assistant_pic
            })

    data = {
        "preferences": {
            "max_char": max_tokens,
            "temperature": temperature,
            "top_p": top_p,
            "system_message": system_message
        },
        "conversation_history": messages,
        "input": message
    }

    with requests.post(API_URL, headers=headers, data=json.dumps(data), stream=True) as response:
        if response.status_code == 200:
            response_json = response.json()
            assistant_reply = response_json["msq"]["message"][0]
            history.append((message, assistant_reply, "You", "P-ALPLE", sessions[api_key]["avatar"], ASSISTANT_PIC_PATH))
            sessions[api_key]["history"] = history
            return history, assistant_reply
        else:
            return history, "Error: " + response.json().get("error", "Unknown error occurred.")

def render_message(history):
    messages_html = """
    <div id="chatbox-container" class="chatbox" style="height: 400px; overflow: auto; border: 1px solid #262626; padding: 10px; background-color: #171717; display: flex; flex-direction: column-reverse;">
        <div id="messages" style="display: block; margin-bottom: 10px;">"""

    for user_message, assistant_message, user_profile, assistant_profile, user_pic, assistant_pic in history:
        if user_message:
            user_message_html = markdown.markdown(
                user_message, extensions=["fenced_code", CodeHiliteExtension()]
            )
            user_message_html = markdown.markdown(escape_html(user_message), extensions=["fenced_code", "codehilite"])
            messages_html += f"""
            <div style='display: flex; align-items: center; margin-bottom: 10px;'>
                <img src='{user_pic}' style='width: 40px; height: 40px; border-radius: 50%; margin-right: 10px;'>
                <span style='color: white;'>{user_message_html}</span>
            </div>"""

        if assistant_message:
            messages_html += f"""
            <div style='display: flex; align-items: center; margin-bottom: 10px;'>
                <img src='{assistant_pic}' style='width: 40px; height: 40px; border-radius: 50%; margin-right: 10px;'>
                <span style='color: white;'>{assistant_message}</span>
            </div>"""

    return messages_html

def user_interaction(message, history, api_key, max_tokens, top_p, temperature):
    loading_message = history + [(message, "Loading...", "You", "P-ALPLE", sessions[api_key]["avatar"], ASSISTANT_PIC_PATH)]
    yield render_message(loading_message), loading_message, ""
    history, assistant_reply = respond(message, api_key, max_tokens, top_p, temperature)
    yield render_message(history), history, ""

def escape_html(unsafe_text):
    return (
        unsafe_text.replace("&", "&amp;")
        .replace("<", "&lt;")
        .replace(">", "&gt;")
        .replace('"', "&quot;")
        .replace("'", "&#039;")
    )

css = """
.chatbox {height: 400px; overflow: auto; border: 1px solid #262626; padding: 10px; background-color: #171717; display: flex;  flex-direction: column-reverse;}
"""

with gr.Blocks(css=css) as demo:
    with gr.Column(visible=True) as auth_view:
        gr.Markdown("## P-MSQ Authorization")
        gr.Markdown("P-MSQ is in closed alpha test! The model, api and more are subject to change.")
        api_user_input = gr.Textbox(placeholder="snowflake", label="UserID", type='email')
        api_key_input = gr.Textbox(placeholder="Enter your API key", label="Token", type='password')
        auth_button = gr.Button("Authorize")
        auth_status = gr.Textbox(label="Authorization Status", interactive=False)

    with gr.Column(visible=False) as chat_view:
        gr.Markdown("## P-MSQ Chat Interface")
        chatbot_output = gr.HTML(elem_id="chatbox-container")
        msg_input = gr.Text(show_label=False, placeholder="Type your message and press Shift+Enter...", lines=2)
        send_btn = gr.Button("Send")
        regen_btn = gr.Button("Clear")
        system_instructions_input = gr.Textbox(placeholder="Enter custom instructions (optional)", label="Custom System Instructions", lines=2)
        save_instructions_btn = gr.Button("Save Instructions")
        max_tokens = gr.Slider(minimum=1, maximum=2048, value=1024, step=1, label="Max new tokens")
        top_p = gr.Slider(minimum=0, maximum=2, value=0.8, step=0.1, label="Top P")
        temperature = gr.Slider(minimum=0.1, maximum=1, value=0.7, step=0.1, label="Temperature")
        history_state = gr.State([])
        last_message_state = gr.State("")

        msg_input.submit(
            user_interaction,
            inputs=[msg_input, history_state, api_key_input, max_tokens, top_p, temperature],
            outputs=[chatbot_output, history_state, msg_input],
        )
        send_btn.click(
            user_interaction,
            inputs=[msg_input, history_state, api_key_input, max_tokens, top_p, temperature],
            outputs=[chatbot_output, history_state, msg_input],
        )
        regen_btn.click(lambda _: ("", []), outputs=[chatbot_output, history_state])

    auth_button.click(
        lambda user, key: (authorize(user, key, PRIMARY_SYSTEM_INSTRUCTIONS), ""),
        inputs=[api_user_input, api_key_input],
        outputs=[chat_view, auth_status],
    )
    save_instructions_btn.click(
        lambda key, instructions: ("", ""),
        inputs=[api_key_input, system_instructions_input],
        outputs=[auth_status],
    )
    demo.launch(show_api=False)

if __name__ == "__main__":
    demo.queue = False