diff --git a/.gitignore b/.gitignore index ea5d9c5..93a7e59 100644 --- a/.gitignore +++ b/.gitignore @@ -38,6 +38,7 @@ next-env.d.ts .idea/ venv/ +__pycache__/ key.pem cert.pem diff --git a/py/__pycache__/ai.cpython-312.pyc b/py/__pycache__/ai.cpython-312.pyc deleted file mode 100644 index a1c10a8..0000000 Binary files a/py/__pycache__/ai.cpython-312.pyc and /dev/null differ diff --git a/py/ai.py b/py/ai.py index c51f2d2..bfde713 100644 --- a/py/ai.py +++ b/py/ai.py @@ -1,4 +1,5 @@ from mistralai import Mistral +from openai import OpenAI import ollama @@ -12,19 +13,15 @@ class AI: options={"temperature": 0.5}, ) - for i in messages: - print(i) - - return_class.ai_response[access_token] = "" + with return_class.ai_response_lock: + return_class.ai_response[access_token] = "" for chunk in stream: - print(chunk['message']['content']) - return_class.ai_response[access_token] += chunk['message']['content'] + with return_class.ai_response_lock: + return_class.ai_response[access_token] += chunk['message']['content'] @staticmethod - def process_mistralai(model, messages, return_class, access_token): - with open("api_key.txt", 'r') as f: - api_key = f.read().strip() + def process_mistralai(model, messages, return_class, access_token, api_key): client = Mistral(api_key=api_key) @@ -33,7 +30,26 @@ class AI: messages=messages ) - return_class.ai_response[access_token] = "" + with return_class.ai_response_lock: + return_class.ai_response[access_token] = "" for chunk in stream_response: - return_class.ai_response[access_token] += chunk.data.choices[0].delta.content + with return_class.ai_response_lock: + return_class.ai_response[access_token] += chunk.data.choices[0].delta.content + + @staticmethod + def process_openai(model, messages, return_class, access_token, api_key): + + client = OpenAI(api_key=api_key) + + stream_response = client.chat.completions.create( + model=model, + messages=messages + ) + + with return_class.ai_response_lock: + return_class.ai_response[access_token] = "" + + for chunk in stream_response: + with return_class.ai_response_lock: + return_class.ai_response[access_token] += chunk.choices[0].delta.content \ No newline at end of file diff --git a/py/api.py b/py/api.py index 7b09e1d..386cd4f 100644 --- a/py/api.py +++ b/py/api.py @@ -1,6 +1,7 @@ from flask import Flask, request, jsonify from flask_cors import CORS import secrets +import threading from ai import AI from db import DB from OpenSSL import crypto @@ -8,11 +9,13 @@ from OpenSSL import crypto class API: def __init__(self): - self.crypt_size = 4096 + self.crypt_size = 1 self.app = Flask(__name__) self.ai_response = {} self.ai = AI() self.db = DB() + self.db.load_database() + self.ai_response_lock = threading.Lock() CORS(self.app) def run(self): @@ -26,12 +29,25 @@ class API: def send_ai(): data = request.get_json() messages = data.get('messages') + model_type = data.get('model_type') ai_model = data.get('ai_model') access_token = data.get('access_token') if access_token not in self.ai_response: return jsonify({'status': 401, 'error': 'Invalid access token'}) - self.ai.process_local(ai_model, messages, self, access_token) - return jsonify({'status': 200}) + + if model_type == "local": + thread = threading.Thread(target=self.ai.process_local, args=(ai_model, messages, self, access_token)) + thread.start() + thread.join() + return jsonify({'status': 200}) + elif model_type == "mistral": + api_key = data.get('api_key') + thread = threading.Thread(target=self.ai.process_mistralai, args=(ai_model, messages, self, access_token, api_key)) + thread.start() + thread.join() + return jsonify({'status': 200}) + + return jsonify({'status': 401, 'error': 'Invalid AI model type'}) @self.app.route('/interstellar/api/ai_get', methods=['GET']) def get_ai(): @@ -46,13 +62,13 @@ class API: data = request.args.get('data') if action == "create_account": self.db.add_user(data) - if action == "change_password": + elif action == "change_password": self.db.update_password(data) - if action == "get_data": + elif action == "get_data": self.db.get_data(data) - if action == "change_data": + elif action == "change_data": self.db.change_data(data) - if action == "check_credentials": + elif action == "check_credentials": self.db.check_credentials(data) email_address = "emailAddress" @@ -89,7 +105,7 @@ class API: f.write(crypto.dump_privatekey(crypto.FILETYPE_PEM, k).decode("utf-8")) ssl_context = ("cert.pem", "key.pem") - self.app.run(debug=True, host='0.0.0.0', port=5000, ssl_context=ssl_context) + self.app.run(debug=True, host='0.0.0.0', port=5000) if __name__ == '__main__': diff --git a/py/db.py b/py/db.py index ef82111..76470c9 100644 --- a/py/db.py +++ b/py/db.py @@ -1,5 +1,5 @@ -import json import hashlib +import json class DB: @@ -26,6 +26,7 @@ class DB: return False self.database[username]['data'] = data + self.save_database() return True def update_password(self, data): @@ -36,6 +37,7 @@ class DB: hashed_new_password = self.hash_password(new_password) self.database[username].update({"hashed_password": hashed_new_password}) + self.save_database() return True def check_credentials(self, data): @@ -55,3 +57,14 @@ class DB: send_back = self.database[username].get['data'] return send_back + + def save_database(self): + with open("database.json", 'w') as file: + json.dump(self.database, file) + + def load_database(self): + try: + with open("database.json", 'r') as file: + self.database = json.load(file) + except FileNotFoundError: + pass diff --git a/py/requirements.txt b/py/requirements.txt index 144c571..3c2be3b 100644 --- a/py/requirements.txt +++ b/py/requirements.txt @@ -2,4 +2,5 @@ flask flask-cors ollama mistralai +openai pyOpenSSL \ No newline at end of file