from flask import jsonify, Flask, request
from flasgger import Swagger
from llmware.models import ModelCatalog

app = Flask(__name__)
swagger = Swagger(app, config={'specs_route': '/swagger-ui'}, merge=True)


@app.route('/', methods=['GET'])
def test_concurrency():
    import random
    random_numbers = [random.randint(0, 100000) for _ in range(100)]
    total_sum = sum(random_numbers)
    print(f'Sum: {total_sum}')
    return jsonify({'message': 'Welcome'}), 200


@app.route('/slim-ner', methods=['POST'])
def llmware_model():
    """
    # llmware/slim-ner model to identify named entities such as people, organization, and place
    ---
    description: llmware/slim-ner model to identify named entities such as people, organization, and places
    produces:
      - application/json
    parameters:
      - in: body
        name: body
        description: Input data
        schema:
          $ref: '#/definitions/model'
    definitions:
      model:
        type: object
        properties:
          params:
            type: array
            items:
              type: string
            collectionFormat: multi
          input:
            type: string
        example: {"params":["people","organization","place"],"input":"Yesterday, in Redmond, Satya Nadella announced that Microsoft would be launching a new AI strategy."}
    responses:
      '200':
        description: Success
        schema:
          type: object
          example: {"person": ["Satya Nadella"], "organization": ["Microsoft"], "place": ["Redmond"]}
    """
    data = request.get_json()
    text = data['input']
    slim_model = ModelCatalog().load_model('llmware/slim-ner')
    response = slim_model.function_call(text, function='classify', params=data['params'])
    print(f'Input: {text} \nResponse: {response}')
    return jsonify(response['llm_response']), 200


if __name__ == '__main__':
    app.run(debug=True, host='0.0.0.0', port=7860)