Merge pull request 'main' (#15) from React-Group/interstellar_ai:main into main

Reviewed-on: https://interstellardevelopment.org/code/code/YasinOnm08/interstellar_ai/pulls/15

.gitignore

@@ -38,6 +38,7 @@ next-env.d.ts
 
 .idea/
 venv/
+__pycache__/
 
 key.pem
 cert.pem

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

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__':

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

py/requirements.txt

@@ -2,4 +2,5 @@ flask
 flask-cors
 ollama
 mistralai
+openai
 pyOpenSSL