erp-ai-latest/src/api/routes/chatbot.py

import os
from flask import Blueprint, request, jsonify, current_app
from werkzeug.utils import secure_filename
from src.services.chatbot import Chatbot
from src.utils.utils import delete_all_files_in_directory
from src.utils.document_loader import load_document  
from src.services.chatbot import Chatbot


# Initialize the Blueprint
bot = Blueprint('chatbot', __name__)

# Allowed file extensions
ALLOWED_EXTENSIONS = {'pdf', 'doc', 'docx'}

def allowed_file(filename):
    """Check if the file has an allowed extension."""
    return '.' in filename and filename.rsplit('.', 1)[1].lower() in ALLOWED_EXTENSIONS

@bot.route('/validate_worker_document', methods=['POST'])
def validate_worker_document():
    try:
        # Retrieve form data
        question = request.form.get('question')
        file = request.files.get('document')

        if not question or not file:
            return jsonify({"error": "Missing data", "message": "Question or document not provided."}), 400
        
        if file.filename == '':
            return jsonify({"error": "No selected file", "message": "No file was selected for upload."}), 400

        if file and allowed_file(file.filename):
            # Secure the file name and save it
            filename = secure_filename(file.filename)
            upload_folder = current_app.config['UPLOAD_FOLDER']
            file_path = os.path.join(upload_folder, filename)
            file.save(file_path)

            # Load the document for processing
            docs = load_document(file_path)

            # Instantiate the chatbot service
            chatbot = Chatbot()

            # Validate the worker's response using the provided document
            validation_result = chatbot.validate_worker(question, docs)

            if not validation_result:
                return jsonify({"error": "Validation failed", "message": "Validation process failed."}), 400

            # Clean up uploaded file (optional)
            os.remove(file_path)

            # Return the validation result
            return jsonify({"result": validation_result["result"]}), 200
        
        else:
            return jsonify({"error": "Invalid file", "message": "File format not supported."}), 400

    except Exception as e:
        return jsonify({"error": "Internal Server Error", "message": str(e)}), 500


@bot.route('/predict_next_n_assessments', methods=['POST'])
def predict_next_n_assessments():
    try:
        # Retrieve JSON data from the request
        data = request.get_json()
        company_info = data.get('company_info')
        companyid = data.get('companyid')
        N = data.get('N')

        if not company_info or not companyid or N is None:
            return jsonify({"error": "Missing data", "message": "Company info, company ID, or N value not provided."}), 400

        # Instantiate the chatbot service
        chatbot = Chatbot()

        # Call the prediction method
        response = chatbot.predict_next_n_assessment(
            company_info=company_info,
            companyid=companyid,
            N=N
        )
        if not response:
            return jsonify({"error": "No predictions available", "message": "Prediction process failed."}), 400
        
        return jsonify({"predictions": response}), 200

    except Exception as e:
        return jsonify({"error": "Internal Server Error", "message": str(e)}), 500



@bot.route('/use_bot_predict_assessments', methods=['POST'])
def use_bot_predict_assessments():
    try:
        # Retrieve JSON data from the request
        data = request.get_json()
        company_info = data.get('company_info')
        companyid = data.get('companyid')
        query = data.get('query')

        if not company_info or not companyid or query is None:
            return jsonify({"error": "Missing data", "message": "Company info, company ID, or query value not provided."}), 400

        # Instantiate the chatbot service
        chatbot = Chatbot()

        # Call the prediction method
        response = chatbot.predict_based_on_past_assessment(
            company_info=company_info,
            companyid=companyid,
            query=query
        )
        if not response:
            return jsonify({"error": "No predictions available", "message": "Prediction process failed."}), 400
        
        return jsonify({"predictions": response}), 200

    except Exception as e:
        return jsonify({"error": "Internal Server Error", "message": str(e)}), 500