from flask import Flask, request, Response, jsonify
from huggingface_hub import InferenceClient
from flask_cors import CORS
import json

client = InferenceClient("mistralai/Mixtral-8x7B-Instruct-v0.1")

summary_prompt = '<s>[INST]Ты ассистент. Отвечаешь на русском языке. Сформируй краткое изложение следующего текста: {}[/INST]'

def format_prompt(message, history):
  prompt = "<s>"
  for user_prompt, bot_response in history:
    prompt += f"[INST] {user_prompt} [/INST]"
    prompt += f" {bot_response}</s> "
  prompt += f"[INST] {message} [/INST]"
  return prompt

def split_text(text):
    max_chars = 3500
    sentences = text.split('.')
    lines = []
    for sentence in sentences:
        lines.extend(sentence.split('\n'))

    result = []
    current_chunk = ''
    for line in lines:
        if len(current_chunk) + len(line) < max_chars:
            current_chunk += line + '.'
        else:
            result.append(current_chunk.strip())
            current_chunk = line + '.'
    if current_chunk:
        result.append(current_chunk.strip())
    return result

def generate(
    prompt, history=[], temperature=0, max_new_tokens=2000, top_p=0.95, repetition_penalty=1.0,
):
    temperature = float(temperature)
    if temperature < 1e-2:
        temperature = 1e-2
    top_p = float(top_p)

    generate_kwargs = dict(
        temperature=temperature,
        max_new_tokens=max_new_tokens,
        top_p=top_p,
        repetition_penalty=repetition_penalty,
        do_sample=True,
        seed=42,
    )

    #formatted_prompt = format_prompt(prompt, history)

    #stream = client.text_generation(prompt, **generate_kwargs, stream=True, details=False, return_full_text=False)
    response = client.text_generation(prompt, **generate_kwargs, stream=False, details=False, return_full_text=False)
    print(response)
    return response
    #output = ""

    #for response in stream:
    #    yield response.token.text.encode('utf-8')

app = Flask(__name__)
CORS(app)

@app.route('/health', methods=['GET'])
def health():
    return jsonify({"status": "ok"})

@app.route('/completion', methods=['POST'])
def completion_route():
    data = request.get_json()
    prompt = data.get('prompt', '')
    #truncated_prompt = prompt[:32768]
    return Response(generate(prompt[:52768]), content_type='text/plain; charset=utf-8', status=200, direct_passthrough=True)

@app.route('/getsummary', methods=['POST'])
def getsummary_route():
    data = request.get_json()
    text = data.get('text', '')
    pages = split_text(text)
    result = ''
    for page in pages:
        summary_prompt = f'''<s>[INST]выпиши из текста в виде трех списков 1.какие вопросы, темы обсуждались, 2. какие проблемы были озвучены 3. какие предложения были сформулированы: {page}[/INST]'''
        response = generate(summary_prompt[:52000])
        result = result + '\n'+response
    
    return jsonify({'result': result})

@app.route('/cleantext', methods=['POST'])
def cleantext_route():
    data = request.get_json()
    text = data.get('text', '')
    
    summary_prompt = f'''<s>[INST]Ты ассистент. Отвечаешь на русском языке. Сформируй краткое изложение следующего текста: {text}[/INST]'''
    response = generate(summary_prompt[:52000])
    
    return jsonify({'result': response})

@app.route('/getfollowup', methods=['POST'])
def getfollowup_route():
    data = request.get_json()
    text = data.get('text', '')
    
    summary_prompt = f'''<s>[INST]Ты мой помощник. Ты отвечаешь только на РУССКОМ языке. Ты не отвечаешь на вопросы, не комментируешь,
не выражаешь эмоций, не выражаешь соображений по теме обращения. Ты выделяешь задачи из текста переписки и формируешь текст письма с поручениями сотрудникам для выполнения найденных задач.
Переписка: {text}[/INST]'''
    response = generate(summary_prompt[:52000])
    
    return jsonify({'result': response})

@app.route('/getagenda', methods=['POST'])
def getagenda_route():
    data = request.get_json()
    text = data.get('text', '')
    
    summary_prompt = f'''<s>[INST]Ты мой помощник. Ты отвечаешь только на РУССКОМ языке. Ты не отвечаешь на вопросы, не комментируешь,
не выражаешь эмоций, не выражаешь соображений по теме обращения. Ты выделяешь выпросы из текста переписки, которые надо дорешать и задачи, выполнение которых надо проконтролировать. Результат выводишь в виде нумерованного списка.
Переписка: {text}[/INST]'''
    response = generate(summary_prompt[:52000])
    
    return jsonify({'result': response})

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