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 delete_last_response(fernet_key, chatbot, user_name, sp_name): # conn, cursor = get_db_connection() # cursor.execute("SELECT messages FROM history WHERE user_name=? AND sp_name=?", (user_name, sp_name)) # conn.commit() # conn.close() 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: # Insert a new record if none exists 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 $") # delete_last_response_button = gr.Button("Delete Last Response") chat_input.submit( update_chat, [password_input, user_input, sp_path_input, chat_input, history_flag_input], [chatbot, tokens_display, cost_display, chat_input] ) # delete_last_response_button.click( # handle_delete, # [user_input, password_input, sp_path_input] # ) start_button.click( start_chat, [password_input, user_input, sp_path_input], [chat_interface, chatbot, user_input, password_input] ) demo.launch()