File size: 5,661 Bytes
29232b4
 
 
 
 
 
 
 
 
 
 
 
 
 
 
5456854
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
b58bc4a
5456854
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
3d1964a
5456854
cf13d53
5456854
 
 
 
 
 
 
 
 
 
3d1964a
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
#
# Copyright (C) Hadad <[email protected]>
# All rights reserved.
#
# This code is made available under the Creative Commons Attribution-NonCommercial-ShareAlike 4.0 International (CC BY-NC-SA 4.0) License.
# You are free to share and adapt the code for non-commercial purposes, as long as you provide appropriate credit,
# do not use it for commercial purposes, and distribute your contributions under the same license.
#
# Contributions can be made by directly submitting pull requests.
#
# For inquiries or permission requests, please contact [email protected].
#
# License: Creative Commons Attribution-NonCommercial-ShareAlike 4.0 International (CC BY-NC-SA 4.0)
#

import gradio as gr
import requests
import json
import os
import threading
import random
import time

LINUX_SERVER_HOST = os.getenv("LINUX_SERVER_HOST")
LINUX_SERVER_PROVIDER_KEY = [key for key in json.loads(os.getenv("LINUX_SERVER_PROVIDER_KEY", "[]")) if key]

AI_TYPES = {f"AI_TYPE_{i}": os.getenv(f"AI_TYPE_{i}") for i in range(1, 6)}
RESPONSES = {f"RESPONSE_{i}": os.getenv(f"RESPONSE_{i}") for i in range(1, 10)}

MODEL_MAPPING = json.loads(os.getenv("MODEL_MAPPING", "{}"))
MODEL_CONFIG = json.loads(os.getenv("MODEL_CONFIG", "{}"))
MODEL_CHOICES = list(MODEL_MAPPING.values())
DEFAULT_CONFIG = json.loads(os.getenv("DEFAULT_CONFIG", "{}"))

META_TAGS = os.getenv("META_TAGS")

stop_event = threading.Event()
session = requests.Session()

def get_model_key(display_name):
    return next((k for k, v in MODEL_MAPPING.items() if v == display_name), MODEL_CHOICES[0])

def simulate_streaming_response(text):
    for line in text.splitlines():
        if stop_event.is_set():
            return
        yield line + "\n"
        time.sleep(0.05)

def chat_with_model(history, user_input, selected_model_display):
    if stop_event.is_set():
        yield RESPONSES["RESPONSE_1"]
        return

    if not LINUX_SERVER_PROVIDER_KEY or not LINUX_SERVER_HOST:
        yield RESPONSES["RESPONSE_3"]
        return

    selected_model = get_model_key(selected_model_display)
    model_config = MODEL_CONFIG.get(selected_model, DEFAULT_CONFIG)

    messages = [{"role": "user", "content": user} for user, _ in history]
    messages += [{"role": "assistant", "content": assistant} for _, assistant in history if assistant]
    messages.append({"role": "user", "content": user_input})

    data = {"model": selected_model, "messages": messages, **model_config}
    random.shuffle(LINUX_SERVER_PROVIDER_KEY)

    for api_key in LINUX_SERVER_PROVIDER_KEY[:2]:
        if stop_event.is_set():
            yield RESPONSES["RESPONSE_1"]
            return
        try:
            response = session.post(LINUX_SERVER_HOST, json=data, headers={"Authorization": f"Bearer {api_key}"})
            if stop_event.is_set():
                yield RESPONSES["RESPONSE_1"]
                return
            if response.status_code < 400:
                ai_text = response.json().get("choices", [{}])[0].get("message", {}).get("content", RESPONSES["RESPONSE_2"])
                yield from simulate_streaming_response(ai_text)
                return
        except requests.exceptions.RequestException:
            continue

    yield RESPONSES["RESPONSE_3"]

def respond(user_input, history, selected_model_display):
    if not user_input.strip():
        yield history, gr.update(value=""), gr.update(visible=False, interactive=False), gr.update(visible=True)
        return

    stop_event.clear()
    history.append([user_input, RESPONSES["RESPONSE_8"]])

    yield history, gr.update(value=""), gr.update(visible=False), gr.update(visible=True)

    ai_response = ""
    for chunk in chat_with_model(history, user_input, selected_model_display):
        if stop_event.is_set():
            session.close()
            history[-1][1] = RESPONSES["RESPONSE_1"]
            yield history, gr.update(value=""), gr.update(visible=True), gr.update(visible=False)
            return
        ai_response += chunk
        history[-1][1] = ai_response
        yield history, gr.update(value=""), gr.update(visible=False), gr.update(visible=True)

    session.close()
    yield history, gr.update(value=""), gr.update(visible=True), gr.update(visible=False)

def stop_response():
    stop_event.set()
    session.close()

def change_model(new_model_display):
    return [], new_model_display

def check_send_button_enabled(msg):
    return gr.update(visible=bool(msg.strip()), interactive=bool(msg.strip()))

with gr.Blocks(fill_height=True, fill_width=True, title=AI_TYPES["AI_TYPE_4"], head=META_TAGS) as demo:
    user_history = gr.State([])
    selected_model = gr.State(MODEL_CHOICES[0])

    chatbot = gr.Chatbot(label=AI_TYPES["AI_TYPE_1"], show_copy_button=True, show_share_button=False, scale=1, elem_id=AI_TYPES["AI_TYPE_2"])
    model_dropdown = gr.Dropdown(label=AI_TYPES["AI_TYPE_3"], show_label=False, choices=MODEL_CHOICES, value=MODEL_CHOICES[0], interactive=True)
    msg = gr.Textbox(label=RESPONSES["RESPONSE_4"], show_label=False, scale=0, placeholder=RESPONSES["RESPONSE_5"])

    with gr.Row():
        send_btn = gr.Button(RESPONSES["RESPONSE_6"], visible=True, interactive=False)
        stop_btn = gr.Button(RESPONSES["RESPONSE_7"], variant=RESPONSES["RESPONSE_9"], visible=False)

    model_dropdown.change(fn=change_model, inputs=[model_dropdown], outputs=[user_history, selected_model])
    send_btn.click(respond, inputs=[msg, user_history, selected_model], outputs=[chatbot, msg, send_btn, stop_btn])
    msg.change(fn=check_send_button_enabled, inputs=[msg], outputs=[send_btn])
    stop_btn.click(fn=stop_response, outputs=[send_btn, stop_btn])

demo.launch(show_api=False)