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

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 call_api(prompt_text):
    url = "https://muryshev-mixtral-api-local.hf.space/completion"
    payload = {"prompt": prompt_text, "temperature": 0, "seed": 42, "repeat_penalty": 1, "top_p": 0.95, "stream": False, "n_predict": 2000}
    
    try:
        response = requests.post(url, json=payload)
        response.raise_for_status()  # Raise an exception for 4xx or 5xx status codes
        result = response.json().get('content', '') # Extract the text result from the JSON response
        print(result)
        return result
    except requests.exceptions.RequestException as e:
        print("Error:", e)
        return None
    
def generate(
    prompt, history=[], temperature=0, max_new_tokens=2000, top_p=0.95, repetition_penalty=1.0,
):
    return call_api(prompt)
    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]Исправь в тексте, сделанным роботом при транскрибации аудиозаписи совещания, плохо-читаемые предложения. Проверь каждое предложение на согласованность падежей, окончаний, на наличие и правильное использование предлогов, в реплике каждого спикера исправь пунктуацию, чтобы смысл и был четкий и понятный, также удали дискурсивные слова, не несущие смысла, сделай текст аккуратным и легко читабельным и пиши только на русском языке: {page}[/INST]'''
    #     response = generate(summary_prompt[:52000])
    #     result = result + '\n'+response
    
    summary_prompt = f'''<s>[INST]Исправь в тексте, сделанным роботом при транскрибации аудиозаписи совещания, плохо-читаемые предложения. Проверь каждое предложение на согласованность падежей, окончаний, на наличие и правильное использование предлогов, в реплике каждого спикера исправь пунктуацию, чтобы смысл и был четкий и понятный, также удали дискурсивные слова, не несущие смысла, сделай текст аккуратным и легко читабельным и пиши только на русском языке: {text}[/INST]'''
    result = generate(summary_prompt[:52000])
    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})

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


@app.route('/getprojectinfo', methods=['POST'])
def getprojectinfo_route():
    data = request.get_json()
    text = data.get('text', '')

    main_prompts = []
    prompt1 = (f'Ты всегда отвечаешь на РУССКОМ языке.+\
Представь, что ты лучший в мире смысловик и копирайтер. +\
Сделай глубокий вдох и, думая шаг за шагом, выполни задание:+\
Идентифицируй участников в этом скрипте. Кто говорил во время совещания? +\
Формат ответа: "Участники: *здесь ФИО говорящих из текста*". Если имя говорящего не указано, то ты тоже его НЕ указываешь. +\
Ты больше ничего не говоришь, не комментируешь, не выражаешь свои мысли, вообще ничего больше не говоришь. Отвечай на русском языке. +\
Скрипт: {text}')
    
    prompt2 = (f'Ты всегда отвечаешь на РУССКОМ языке.+\
Представь, что ты лучший в мире смысловик и копирайтер. +\
Сделай глубокий вдох и, думая шаг за шагом, выполни задание:+\
Каков ход совещания в этом тексте?+\
Формат ответа: "Повестка встречи: *здесь текст*".+\
Ты больше ничего не говоришь, не комментируешь, не выражаешь свои мысли, вообще ничего больше не говоришь. Отвечай на русском языке. +\
Скрипт: {text}')
    main_prompts.append(prompt1)
    main_prompts.append(prompt2)
    
    main_info =''
    for i in main_prompts:
        result = generate(i[:52000])
        main_info += result

    proj_prompt = []
    prompt = (f'Вы отвечаете только на РУССКОМ языке.+\
1. Найти все упоминания слов "проект" или "проекты" в тексте.+\
2. Прочитать контекст вокруг упоминаний слова "проект" или "проекты" и определить, какие проекты упоминаются.+\
3. Составить список названий проектов, которые были найдены в тексте.+\
Используй только такой формат ответа: "Проект №: *суть кратко*". Конец ответа.+\
Ты больше ничего не говоришь, не комментируешь, не добавляешь.+\
Текст: {text}')
    proj_prompt.append(prompt)
    
    list_of_projects =''
    for i in proj_prompt:
        result = generate(i[:52000])
        list_of_projects += result
    
    delimiter = 'Проект '
    proj  = [delimiter+x for x in list_of_projects.split(delimiter) if x]
    proj = proj[1:]
    
    proj_ = []
    for i in proj:
        a = i.replace("Проект №", "")
        a = a.replace("Конец ответа", "")
        a = a.replace("данный проект", "")  ###убираю слово "проект", чтобы модель не опиралась на него,
        a = a.replace("проект ", "")             # при ответе на вопрос, проект это или нет
        a = a.replace('\n', ' ')
        a = a.replace('</s>', ' ')
        a = a.strip()
        proj_.append(a)
        
    check_prompts = []

    checking = (f'Ты всегда отвечаешь на РУССКОМ языке. Сделай глубокий вдох и, думая шаг за шагом, выполни задание:+\
    У меня есть список проектов. Такие проекты имеются в тексте и являются основными? Список проектов: {proj_}. +\
    Ты обязательно используешь ТОЛЬКО такой формат ответа:+\
    "Да: Проект № *название проекта*." ИЛИ "Нет, такого проекта нет". +\
    Ты больше ничего не говоришь, не комментируешь, не добавляешь.+\
    Текст: {text}')
    check_prompts.append(checking)

    real_projects = ''
    for i in check_prompts:
        result = generate(i[:52000])
        real_projects += result
    
    real_projects_list = re.findall(r'Да:\s*(.*?)\s*(?:\n\n|$)', real_projects)
    prepared_prompts = []
    if real_projects_list:
        for i in real_projects_list:
            prompt_aim = (f'Ты всегда отвечаешь на РУССКОМ языке. Представь, что ты лучший в мире смысловик и копирайтер. +\
Сделай глубокий вдох и, думая шаг за шагом, выполни задание:+\
Опираясь на данный текст, определи цель ЭТОГО проекта: {i}.  Выбирай и анализируй информацию только об ЭТОМ проекте. +\
Формат ответа: "Проект: *название проекта*. Цель этого проекта: *здесь текст*".+\
Ты больше ничего не говоришь, не комментируешь, не выражаешь свои мысли, вообще ничего больше не говоришь. Отвечай на русском языке. +\
Текст: {text}')
            gk = (f'Ты всегда отвечаешь на РУССКОМ языке. Представь, что ты лучший в мире смысловик и копирайтер. +\
Сделай глубокий вдох и, думая шаг за шагом, ответь на вопрос: Какие преимущества (выгоду) получит группа компаний по итогам этого проекта: {i}.+\
Формат ответа: "По факту реализации проекта Группа Компаний Получит: *здесь текст*".+\
Ты больше ничего не говоришь, не комментируешь, не выражаешь свои мысли, вообще ничего больше не говоришь. Отвечай на русском языке. +\
Скрипт: {text}')
            budget = (f'Ты всегда отвечаешь на РУССКОМ языке. +\
Найди в тексте информацию о бюджете этого проекта: {i} . +\
Формат ответа: "Бюджет проекта: *здесь цифра*". Если бюджет проекта не указан, то отвечай так: "Такой информации нет."+\
Ты больше ничего не говоришь, не комментируешь, не выражаешь свои мысли, вообще ничего больше не говоришь. Отвечай на русском языке. +\
Скрипт: {text}')
            ec_ef = (f'Ты всегда отвечаешь на РУССКОМ языке. +\
Поразмышляй об этом тексте. Как ты думаешь, в чем заключается экономический эффект (экономические преимущества) для компании от этого проекта: {i}. Выбирай и анализируй информацию только об ЭТОМ проекте. +\
Формат ответа: "Экономический эффект от проекта: *здесь ответ на вопрос*".+\
Ты больше ничего не говоришь, не комментируешь, не выражаешь свои мысли, вообще ничего больше не говоришь. Отвечай на русском языке. +\
Скрипт: {text}')
            deadline = (f'Ты всегда отвечаешь на РУССКОМ языке. +\
Найди в тексте дату, когда должен быть реализован ЭТОТ проект: {i}. Выбирай и анализируй информацию только об ЭТОМ проекте: {i}.+\
Формат ответа: "Срок реализации: *здесь текст*".+\
Ты больше ничего не говоришь, не комментируешь, не выражаешь свои мысли, вообще ничего больше не говоришь. Отвечай на русском языке. +\
Скрипт: {text}')
            new_plan = (f'Ты всегда отвечаешь на РУССКОМ языке! Только на русском языке. Выбирай и анализируй информацию на русском языке только об ЭТОМ проекте: {i} .+\
 Представь, что ты лучший в мире смысловик и копирайтер. +\
Сделай глубокий вдох и, думая шаг за шагом, ответь на русском языке на вопрос: Какие действия участники решили предпринять, чтобы отлично выполнить проект?+\
Используй такой формат ответа: "Решения: *тест*". Ищи ответ в тексте: {text}. Отвечай на русском языке!!! ')
            conclusion = (f'Ты всегда отвечаешь на РУССКОМ языке. Поразмышляй о тексте.+\
Какой вывод заключили участники относительно ЭТОГО проекта: {i}.+\
Используй такой формат ответа: "Вывод: *Тут описывается принятое решение.*." Конец ответа.+\
Текст: {text}')
            p = [prompt_aim, gk, budget, ec_ef, deadline, new_plan, conclusion]
            prepared_prompts.extend(p)
    elif not real_projects_list or len(real_projects_list) == 1:
        prompt_aim = (f'Ты всегда отвечаешь на РУССКОМ языке. Поразмышляй об этом тексте.+\
Определи основную цель проекта из текста.+\
Формат ответа: "Цель проекта: *здесь текст*".+\
Ты больше ничего не говоришь, не комментируешь, не выражаешь свои мысли, вообще ничего больше не говоришь. Отвечай на русском языке. +\
Текст: {text}')

        gk = (f'Ты всегда отвечаешь на РУССКОМ языке.+\
Формат ответа: "По факту реализации проекта Группа Компаний Получит: *здесь текст*".+\
Ты больше ничего не говоришь, не комментируешь, не выражаешь свои мысли, вообще ничего больше не говоришь. Отвечай на русском языке. +\
Скрипт: {text}')

        budget = (f'Ты всегда отвечаешь на РУССКОМ языке. +\
Найди в тексте информацию о бюджете проекта.+\
Формат ответа: "Бюджет проекта: *здесь цифра*". Если бюджет проекта не указан, то отвечай так: "Такой информации нет."+\
Ты больше ничего не говоришь, не комментируешь, не выражаешь свои мысли, вообще ничего больше не говоришь. Отвечай на русском языке. +\
Скрипт: {text}')

        ec_ef = (f'Ты всегда отвечаешь на РУССКОМ языке. +\
Поразмышляй об этом тексте. Как ты думаешь, в чем заключается экономический эффект для компании от проектов, которые обсуждают в тексте?+\
Формат ответа: "Экономический эффект от проекта: *здесь ответ на вопрос*".+\
Ты больше ничего не говоришь, не комментируешь, не выражаешь свои мысли, вообще ничего больше не говоришь. Отвечай на русском языке. +\
Скрипт: {text}')

        deadline = (f'Ты всегда отвечаешь на РУССКОМ языке. +\
Найди в тексте дату, когда должен быть реализован проект. +\
Формат ответа: "Срок реализации: *здесь текст*".+\
Ты больше ничего не говоришь, не комментируешь, не выражаешь свои мысли, вообще ничего больше не говоришь. Отвечай на русском языке. +\
Скрипт: {text}')
        new_plan = (f'Ты всегда отвечаешь на РУССКОМ языке! Только на русском языке.+\
 Представь, что ты лучший в мире смысловик и копирайтер. +\
Сделай глубокий вдох и, думая шаг за шагом, ответь на русском языке на вопрос: Какие действия участники решили предпринять, чтобы отлично выполнить проект?+\
Используй такой формат ответа: "Решения: *тест*". Ищи ответ в тексте: {text}. Отвечай на русском языке!!! ')
        conclusion = (f'Ты всегда отвечаешь на РУССКОМ языке. Поразмышляй о тексте.+\
Какой вывод заключили участники относительно ЭТОГО проекта: {i}.+\
Используй такой формат ответа: "Вывод: *Тут описывается принятое решение.*." Конец ответа.+\
Текст: {text}')
        p = [prompt_aim, gk, budget, ec_ef, deadline, new_plan, conclusion]
        prepared_prompts.extend(p)
    
    final = ''
    for i in prepared_prompts:
        result = generate(i[:52000])
        final += result
    final = main_info + final
    final = final.replace("Конец ответа", "")
    final = final.replace('\n', ' ')
    final = final.replace('</s>', ' ')
    final = final.strip()
    return jsonify({'result': final})

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