Fixes and Multithreading

This commit is contained in:
Patrick_Pluto 2024-09-23 11:01:39 +02:00
parent 1d1705ebfa
commit a8f9ba27bd
3 changed files with 44 additions and 13 deletions

View file

@ -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,12 +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):
@ -34,10 +36,18 @@ class API:
return jsonify({'status': 401, 'error': 'Invalid access token'})
if model_type == "local":
self.ai.process_local(ai_model, messages, self, access_token)
if model_type == "mistral":
self.ai.process_mistralai(ai_model, messages, self, access_token)
return jsonify({'status': 200})
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():
@ -95,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__':