From a8f9ba27bdd60d77be356edbec7eb5bf74e8ef5d Mon Sep 17 00:00:00 2001
From: Patrick_Pluto <patrick_pluto@noreply.codeberg.org>
Date: Mon, 23 Sep 2024 11:01:39 +0200
Subject: [PATCH] Fixes and Multithreading

---
 py/ai.py            | 34 +++++++++++++++++++++++++++-------
 py/api.py           | 22 ++++++++++++++++------
 py/requirements.txt |  1 +
 3 files changed, 44 insertions(+), 13 deletions(-)

diff --git a/py/ai.py b/py/ai.py
index 0a57b57..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,15 +13,15 @@ class AI:
             options={"temperature": 0.5},
         )
 
-        return_class.ai_response[access_token] = ""
+        with return_class.ai_response_lock:
+            return_class.ai_response[access_token] = ""
 
         for chunk in stream:
-            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)
 
@@ -29,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 b52870f..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,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__':
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
-- 
2.39.5