app.py

import gradio as gr
import requests
import json
import os
import datetime
from requests.exceptions import RequestException

# gorillion totally mine OR keys go here :3
api_keys_env = os.environ.get('API_KEYS')
if api_keys_env:
    API_KEYS = [key.strip() for key in api_keys_env.strip().split('\n') if key.strip()]
else:
    raise ValueError("all keez ded go kys") # this will never happen due to the superior OR key acquisition method, but juuuust in case (if it actually happens it is joever for joEver.)

API_URL = os.environ.get('API_URL')

DEFAULT_PARAMS = {
    "temperature": 1.0,
    "top_p": 1,
    "top_k": 0,
    "min_p": 0,
    "top_a": 0.1,
    "frequency_penalty": 0,
    "presence_penalty": 0,
    "repetition_penalty": 1.1,
    "max_tokens": 512
}

# debug switches
USER_LOGGING_ENABLED = False
RESPONSE_LOGGING_ENABLED = True

def get_timestamp():
    return datetime.datetime.now().strftime("%H:%M:%S")

def predict(message, history, system_prompt, temperature, top_p, top_k, min_p, top_a,
            frequency_penalty, presence_penalty, repetition_penalty, max_tokens, stop_flag):
    history_format = [{"role": "system", "content": system_prompt}] if system_prompt else []
    for human, assistant in history:
        history_format.append({"role": "user", "content": human})
        if assistant:
            history_format.append({"role": "assistant", "content": assistant})
    history_format.append({"role": "user", "content": message})

    if USER_LOGGING_ENABLED and not message.startswith(('*', '"')):
        print(f"<|system|> {system_prompt}")
        print(f"{get_timestamp()} <|user|> {message}")

    current_params = {
        "temperature": temperature,
        "top_p": top_p,
        "top_k": top_k,
        "min_p": min_p,
        "top_a": top_a,
        "frequency_penalty": frequency_penalty,
        "presence_penalty": presence_penalty,
        "repetition_penalty": repetition_penalty,
        "max_tokens": max_tokens
    }

    non_default_params = {k: v for k, v in current_params.items() if v != DEFAULT_PARAMS[k]}
    
    if USER_LOGGING_ENABLED and non_default_params and not message.startswith(('*', '"')):
        for param, value in non_default_params.items():
            print(f"{param}={value}")

    data = {
        "model": "meta-llama/llama-3.1-405b-instruct:free",
        "messages": history_format,
        "stream": True,
        "temperature": temperature,
        "top_p": top_p,
        "top_k": top_k,
        "min_p": min_p,
        "top_a": top_a,
        "frequency_penalty": frequency_penalty,
        "presence_penalty": presence_penalty,
        "repetition_penalty": repetition_penalty,
        "max_tokens": max_tokens
    }

    api_key_index = 0
    retries = 0
    max_retries = len(API_KEYS)

    while retries < max_retries:
        current_api_key = API_KEYS[api_key_index]

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

        try:
            with requests.post(API_URL, headers=headers, data=json.dumps(data), stream=True) as response:
                if response.status_code == 200:
                    partial_message = ""
                    for line in response.iter_lines():
                        if stop_flag[0]:
                            response.close()
                            break
                        if line:
                            line = line.decode('utf-8')
                            if RESPONSE_LOGGING_ENABLED:
                                print(f"API Response: {line}")
                            if line.startswith("data: "):
                                if line.strip() == "data: [DONE]":
                                    break
                                try:
                                    json_data = json.loads(line[6:])
                                    if 'choices' in json_data and json_data['choices']:
                                        delta = json_data['choices'][0]['delta']
                                        content = delta.get('content', '')
                                        if content:
                                            partial_message += content
                                            yield partial_message
                                except json.JSONDecodeError:
                                    continue  
                    if partial_message:
                        yield partial_message
                    # successful response, break out of retry loop
                    break
                elif response.status_code == 429:
                    print("rate limit hit, cycling keys...")
                    retries += 1
                    api_key_index = (api_key_index + 1) % len(API_KEYS)
                    continue
                else:
                    # theothershits
                    error_message = f"Error: Received status code {response.status_code} - {response.text}"
                    print(error_message)
                    yield f"An error occurred: {error_message}"
                    break
        except RequestException as e:
            print(f"Request error: {e}")
            yield f"An error occurred: {str(e)}"
            break

def import_chat(custom_format_string):
    try:
        sections = custom_format_string.split('<|')

        imported_history = []
        system_prompt = ""

        for section in sections:
            if section.startswith('system|>'):
                system_prompt = section.replace('system|>', '').strip()
            elif section.startswith('user|>'):
                user_message = section.replace('user|>', '').strip()
                imported_history.append([user_message, None])
            elif section.startswith('assistant|>'):
                assistant_message = section.replace('assistant|>', '').strip()
                if imported_history:
                    imported_history[-1][1] = assistant_message
                else:
                    imported_history.append(["", assistant_message])

        return imported_history, system_prompt
    except Exception as e:
        print(f"Error importing chat: {e}")
        return None, None

def export_chat(history, system_prompt):
    export_data = f"<|system|> {system_prompt}\n\n"
    if history is not None:
        for user_msg, assistant_msg in history:
            export_data += f"<|user|> {user_msg}\n\n"
            if assistant_msg:
                export_data += f"<|assistant|> {assistant_msg}\n\n"
    return export_data

def stop_generation_func(stop_flag):
    stop_flag[0] = True
    return stop_flag

with gr.Blocks(theme='gradio/monochrome') as demo:
    stop_flag = gr.State([False])
    
    with gr.Row():
        with gr.Column(scale=2):
            chatbot = gr.Chatbot(value=[])
            msg = gr.Textbox(label="Message")
            with gr.Row():
                clear = gr.Button("Clear")
                regenerate = gr.Button("Regenerate")
                stop_btn = gr.Button("Stop")
            with gr.Row():
                with gr.Column(scale=4):
                    import_textbox = gr.Textbox(label="Import textbox", lines=5)
                with gr.Column(scale=1):
                    export_button = gr.Button("Export Chat")
                    import_button = gr.Button("Import Chat")

        with gr.Column(scale=1):
            system_prompt = gr.Textbox("", label="System Prompt", lines=5)
            temperature = gr.Slider(0, 2, value=DEFAULT_PARAMS["temperature"], step=0.01, label="Temperature")
            top_p = gr.Slider(0, 1, value=DEFAULT_PARAMS["top_p"], step=0.01, label="Top P")
            top_k = gr.Slider(0, 500, value=DEFAULT_PARAMS["top_k"], step=1, label="Top K")
            min_p = gr.Slider(0, 1, value=DEFAULT_PARAMS["min_p"], step=0.01, label="Min P")
            top_a = gr.Slider(0, 1, value=DEFAULT_PARAMS["top_a"], step=0.01, label="Top A")
            frequency_penalty = gr.Slider(-2, 2, value=DEFAULT_PARAMS["frequency_penalty"], step=0.1, label="Frequency Penalty")
            presence_penalty = gr.Slider(-2, 2, value=DEFAULT_PARAMS["presence_penalty"], step=0.1, label="Presence Penalty")
            repetition_penalty = gr.Slider(0.01, 5, value=DEFAULT_PARAMS["repetition_penalty"], step=0.01, label="Repetition Penalty")
            max_tokens = gr.Slider(1, 4096, value=DEFAULT_PARAMS["max_tokens"], step=1, label="Max Output (max_tokens)")

    def user(user_message, history):
        history = history or []
        return "", history + [[user_message, None]]

    def bot(history, system_prompt, temperature, top_p, top_k, min_p, top_a,
            frequency_penalty, presence_penalty, repetition_penalty, max_tokens, stop_flag):
        stop_flag[0] = False
        history = history or []
        if not history:
            return history
        user_message = history[-1][0]
        bot_message = predict(user_message, history[:-1], system_prompt, temperature, top_p, top_k, min_p, top_a,
                              frequency_penalty, presence_penalty, repetition_penalty, max_tokens, stop_flag)
        history[-1][1] = ""
        for chunk in bot_message:
            if stop_flag[0]:
                history[-1][1] += " [Generation stopped]"
                break
            history[-1][1] = chunk
            yield history

    def regenerate_response(history, system_prompt, temperature, top_p, top_k, min_p, top_a,
                            frequency_penalty, presence_penalty, repetition_penalty, max_tokens, stop_flag):
        if history and len(history) > 0:
            last_user_message = history[-1][0]
            history[-1][1] = None  
            for new_history in bot(history, system_prompt, temperature, top_p, top_k, min_p, top_a,
                                   frequency_penalty, presence_penalty, repetition_penalty, max_tokens, stop_flag):
                yield new_history
        else:
            yield []

    def import_chat_wrapper(custom_format_string):
        imported_history, imported_system_prompt = import_chat(custom_format_string)
        return imported_history, imported_system_prompt

    msg.submit(user, [msg, chatbot], [msg, chatbot], queue=False).then(
        bot, [chatbot, system_prompt, temperature, top_p, top_k, min_p, top_a,
              frequency_penalty, presence_penalty, repetition_penalty, max_tokens, stop_flag], chatbot
    )

    clear.click(lambda: None, None, chatbot, queue=False)

    regenerate.click(
        regenerate_response,
        [chatbot, system_prompt, temperature, top_p, top_k, min_p, top_a,
         frequency_penalty, presence_penalty, repetition_penalty, max_tokens, stop_flag],
        chatbot
    )

    import_button.click(import_chat_wrapper, inputs=[import_textbox], outputs=[chatbot, system_prompt])

    export_button.click(
        export_chat,
        inputs=[chatbot, system_prompt],
        outputs=[import_textbox]
    )

    stop_btn.click(stop_generation_func, inputs=[stop_flag], outputs=[stop_flag])

if __name__ == "__main__":
    demo.queue(max_size=3, default_concurrency_limit=3).launch(debug=True)