71 lines
2.2 KiB
Python
71 lines
2.2 KiB
Python
from flask import Flask, request, jsonify
|
|
import os, ai.processor
|
|
from dotenv import load_dotenv
|
|
import io
|
|
from PIL import ImageGrab
|
|
|
|
load_dotenv()
|
|
|
|
class WebServerApp:
|
|
def __init__(self, aip):
|
|
self.app = Flask(__name__)
|
|
self._register_routes()
|
|
self.aip: ai.processor.AIProcessor = aip
|
|
|
|
def _register_routes(self):
|
|
@self.app.route('/api/test')
|
|
def test():
|
|
return jsonify({"message": "Hello, World!"})
|
|
|
|
@self.app.route('/api/request', methods=['POST'])
|
|
def handle_request():
|
|
# sent as form-data
|
|
data = request.form.to_dict()
|
|
if not data:
|
|
return jsonify({"error": "No data provided"}), 400
|
|
|
|
# Process the data as needed
|
|
prompt = data.get('prompt', '')
|
|
|
|
|
|
|
|
if not prompt:
|
|
return jsonify({"error": "No prompt provided"}), 400
|
|
img_data = None
|
|
if 'img' in request.files:
|
|
img_file = request.files['img']
|
|
if img_file:
|
|
img_data = img_file.read()
|
|
else:
|
|
img_data = None
|
|
else:
|
|
if 'host_screenshot' in data:
|
|
# take a screenshot right here
|
|
# capture the full screen
|
|
screenshot_img = ImageGrab.grab()
|
|
buf = io.BytesIO()
|
|
screenshot_img.save(buf, format='PNG')
|
|
img_data = buf.getvalue()
|
|
|
|
import base64
|
|
# Convert image data to base64 if provided
|
|
if img_data and isinstance(img_data, bytes):
|
|
img_data = base64.b64encode(img_data).decode('utf-8')
|
|
elif img_data and isinstance(img_data, str):
|
|
img_data = img_data.encode('utf-8')
|
|
|
|
response = self.aip.process(prompt, img_data)
|
|
return jsonify({"response": response}), 200
|
|
|
|
@self.app.route('/api/health')
|
|
def health():
|
|
return jsonify({"status": "healthy"})
|
|
|
|
def run(self, *args, **kwargs):
|
|
self.app.run(*args, **kwargs)
|
|
|
|
# Example usage:
|
|
# if __name__ == "__main__":
|
|
# server = WebServerApp()
|
|
# server.run()
|