| import os |
| import json |
| import time |
| import base64 |
| import hashlib |
| import sqlite3 |
|
|
| import tiktoken |
| import gradio as gr |
| from openai import OpenAI |
| from cryptography.fernet import Fernet |
|
|
| class FernetKey: |
| def __init__(self, password): |
| key = hashlib.pbkdf2_hmac('sha256', password.encode(), b'salt', 100000) |
| self.key = base64.urlsafe_b64encode(key) |
|
|
| def encrypt(self, content): |
| cipher_suite = Fernet(self.key) |
| encrypted_content = cipher_suite.encrypt(content.encode('utf-8')) |
| return encrypted_content |
|
|
| def decrypt(self, content): |
| cipher_suite = Fernet(self.key) |
| decrypted_content = cipher_suite.decrypt(content).decode('utf-8') |
| return decrypted_content |
|
|
| def get_db_connection(): |
| conn = sqlite3.connect('FKAI.db') |
| return conn, conn.cursor() |
|
|
| def list_prompts(): |
| conn, cursor = get_db_connection() |
| cursor.execute("SELECT name FROM sp") |
| prompt_names = [row[0] for row in cursor.fetchall()] |
| conn.close() |
| return prompt_names |
|
|
| def make_messages(sp_name, messages, chat, fernet_key): |
| conn, cursor = get_db_connection() |
| cursor.execute("SELECT content FROM sp WHERE name=?", (sp_name,)) |
| system_prompt = fernet_key.decrypt(cursor.fetchone()[0]) |
| conn.close() |
|
|
| with open('jb_head.txt', 'rb') as jb_file: |
| jailbreak_head = fernet_key.decrypt(jb_file.read()) |
|
|
| system_message = [{'role':'system', 'content': jailbreak_head+system_prompt}] |
| user_message = [{'role':'user', 'content': chat}] |
| messages = system_message + messages + user_message |
| return messages |
|
|
| def count_tokens(messages): |
| encoder = tiktoken.get_encoding('o200k_base') |
| num_tokens = sum(len(encoder.encode(message['content'])) for message in messages[:-1]) |
| num_tokens += 2 * len(messages) |
| cost = num_tokens / 1000 * 0.003 |
| completion = len(encoder.encode(messages[-1]['content'])) |
| num_tokens += completion + 1 |
| cost += completion / 1000 * 0.015 |
| return num_tokens, round(cost, 3) |
|
|
| def query_model(api_key, messages): |
| for message in messages: |
| message.pop("time", None) |
| client = OpenAI(api_key=api_key, base_url="https://api.deepseek.com") |
| response = client.chat.completions.create( |
| model='deepseek-reasoner', |
| temperature=1.0, |
| messages=messages |
| ) |
| return response.choices[0].message.content |
|
|
| |
| |
| |
| |
| |
|
|
| def ensure_user_sp_exists(user_name, sp_name, fernet_key): |
| conn, cursor = get_db_connection() |
| cursor.execute("SELECT COUNT(*) FROM history WHERE user_name=? AND sp_name=?", (user_name, sp_name)) |
| if cursor.fetchone()[0] == 0: |
| |
| cursor.execute("INSERT INTO history (user_name, sp_name, messages) VALUES (?, ?, ?)", |
| (user_name, sp_name, fernet_key.encrypt(json.dumps([],ensure_ascii=False)))) |
| conn.commit() |
| conn.close() |
|
|
| def handle_history_flag(history_flag, user_name, sp_name, fernet_key): |
| conn, cursor = get_db_connection() |
| ensure_user_sp_exists(user_name, sp_name, fernet_key) |
| history = [] |
| if history_flag: |
| cursor.execute("SELECT messages FROM history WHERE user_name=? AND sp_name=?", |
| (user_name, sp_name)) |
| row = cursor.fetchone() |
| if row: |
| history += json.loads(fernet_key.decrypt(row[0])) |
| conn.close() |
| return history |
|
|
| def update_chat(password, user_name, sp_name, chat, history_flag): |
| fernet_key = FernetKey(password) |
| conn, cursor = get_db_connection() |
| if chat.strip() == "": |
| return "", "", "", "" |
|
|
| api_key = os.getenv('API_KEY') |
| timestamp1 = time.strftime("%Y-%m-%d %H:%M:%S", time.localtime()) |
|
|
| messages = handle_history_flag(history_flag, user_name, sp_name, fernet_key) |
| content = query_model(api_key, make_messages(sp_name, messages, chat, fernet_key)) |
| timestamp2 = time.strftime("%Y-%m-%d %H:%M:%S", time.localtime()) |
|
|
| messages.append({'role': 'user', 'content': chat, "time": timestamp1}) |
| messages.append({'role': 'assistant', 'content': content, "time": timestamp2}) |
|
|
| tokens, cost = count_tokens(messages) |
|
|
| cursor.execute("UPDATE history SET messages = ? WHERE user_name = ? AND sp_name = ?", |
| (fernet_key.encrypt(json.dumps(messages, ensure_ascii=False)), user_name, sp_name)) |
| conn.commit() |
| conn.close() |
|
|
| return messages, tokens, cost, "" |
|
|
| def start_chat(password, user_name, sp_name): |
| fernet_key = FernetKey(password) |
| if not user_name.strip(): |
| return gr.update(), None, None, None |
|
|
| initial_history = handle_history_flag(True, user_name, sp_name, fernet_key) |
| return gr.update(visible=True), initial_history, gr.update(interactive=False), gr.update(interactive=False) |
|
|
| if __name__ == "__main__": |
| initial_prompts = list_prompts() |
|
|
| with gr.Blocks() as demo: |
| with gr.Column(): |
| user_input = gr.Textbox( |
| label="Your Name", |
| placeholder="Your Name" |
| ) |
| password_input = gr.Textbox( |
| label="Password", |
| placeholder="Password", |
| type="password" |
| ) |
| sp_path_input = gr.Dropdown( |
| label="System Prompt", |
| choices=initial_prompts |
| ) |
| start_button = gr.Button("Start Chat") |
|
|
| chat_interface = gr.Column(visible=False) |
| with chat_interface: |
| chatbot = gr.Chatbot( |
| type='messages' |
| ) |
| chat_input = gr.Textbox( |
| label="Your Chat Here", |
| placeholder="Shift Enter New Line" |
| ) |
| history_flag_input = gr.Checkbox( |
| label="Load History", |
| value=True |
| ) |
|
|
| with gr.Row(): |
| tokens_display = gr.Label(label="Tokens Used") |
| cost_display = gr.Label(label="Cost $") |
| |
|
|
| chat_input.submit( |
| update_chat, |
| [password_input, user_input, sp_path_input, chat_input, history_flag_input], |
| [chatbot, tokens_display, cost_display, chat_input] |
| ) |
|
|
| |
| |
| |
| |
|
|
| start_button.click( |
| start_chat, |
| [password_input, user_input, sp_path_input], |
| [chat_interface, chatbot, user_input, password_input] |
| ) |
|
|
| demo.launch() |
|
|
