forked from React-Group/interstellar_ai
61 lines
2.1 KiB
Python
61 lines
2.1 KiB
Python
from flask import Flask, request, jsonify
|
|
from flask_cors import CORS
|
|
import ollama
|
|
import secrets
|
|
|
|
|
|
class AI:
|
|
@staticmethod
|
|
def process_local(model, message, system, return_class, access_token):
|
|
stream = ollama.chat(
|
|
model=model,
|
|
messages=[{'role': 'user', 'content': message}, {'role': 'system', 'content': system}],
|
|
stream=True,
|
|
)
|
|
|
|
for chunk in stream:
|
|
print(chunk['message']['content'])
|
|
return_class.ai_response[access_token] += chunk['message']['content']
|
|
|
|
|
|
class API:
|
|
def __init__(self):
|
|
self.app = Flask(__name__)
|
|
self.ai_response = {}
|
|
self.ai = AI()
|
|
CORS(self.app)
|
|
|
|
def run(self):
|
|
@self.app.route('/interstellar/api/ai_create', methods=['GET'])
|
|
def create_ai():
|
|
access_token = secrets.token_urlsafe(4096)
|
|
self.ai_response[access_token] = ""
|
|
return jsonify({'status': 200, 'access_token': access_token})
|
|
|
|
@self.app.route('/interstellar/api/ai_send', methods=['POST'])
|
|
def send_ai():
|
|
data = request.get_json()
|
|
message = data.get('message')
|
|
ai_model = data.get('ai_model')
|
|
system_prompt = data.get('system_prompt')
|
|
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, message, system_prompt, self, access_token)
|
|
return jsonify({'status': 200})
|
|
|
|
@self.app.route('/interstellar/api/ai_get', methods=['GET'])
|
|
def get_ai():
|
|
data = request.args.get('access_token')
|
|
if data not in self.ai_response:
|
|
return jsonify({'status': 401, 'error': 'Invalid access token'})
|
|
return jsonify({'status': 200, 'response': self.ai_response[data]})
|
|
|
|
ssl_context = ('cert.pem', 'key.pem')
|
|
self.app.run(debug=True, host='0.0.0.0', port=5000, ssl_context=ssl_context)
|
|
|
|
|
|
if __name__ == '__main__':
|
|
api = API()
|
|
api.run()
|
|
|