utils.py

from pydantic import BaseModel
import re


class DiagramRequest(BaseModel):
    description: str


class DiagramResponse(BaseModel):
    plantuml_code: str


def preprocess_description(desc: str) -> str:
    """
    Улучшенный препроцессинг:
    - Убираем лишние пустые строки в начале/конце
    - Нормализуем множественные пробелы и табы
    - Сохраняем \n внутри — это важно для таблиц "Шаг | Роль" и списков
    - Убираем только trailing/leading whitespace
    """
    # Удаляем BOM, если есть
    desc = desc.strip().lstrip('\ufeff')

    # Заменяем множественные пробелы/табы на один, но сохраняем строки
    desc = re.sub(r'[ \t]+', ' ', desc)

    # Убираем пустые строки в начале и конце, но оставляем внутри
    lines = [line.rstrip() for line in desc.splitlines() if line.strip()]
    desc = '\n'.join(lines).strip()

    return desc


def get_prompt(description: str) -> list[dict]:
    """
    Возвращает список сообщений в формате ChatML для Qwen2.5-Instruct.
    Это сильно улучшает качество по сравнению с plain text.
    """
    system_prompt = """You are an expert PlantUML diagram generator. Your ONLY task is to convert the given process description into a valid, complete PlantUML activity diagram.

Rules you MUST follow strictly:
- Always output ONLY the PlantUML code: from @startuml to @enduml inclusive.
- NEVER add explanations, comments, markdown, extra text, apologies or anything else.
- Use the new activity diagram syntax (beta).
- Use swimlanes/partitions (|Actor|) for different roles or actors when mentioned (Шаг | Роль, Actor: action, etc.).
- Use start / stop / (*) / (*) for begin/end.
- Use :action; for steps.
- Use if (Condition?) then (yes) ... else (no) ... endif for decisions.
- Use -> [label]; for arrows with conditions if needed.
- Make diagrams clean, logical and complete — close all branches, ifs, etc.
- Detect type: simple sequence → linear; table Шаг/Роль → swimlanes; branches (Найдена?, Автоматически?) → if/else.

Examples follow. Follow their style and completeness exactly."""

    messages = [
        {"role": "system", "content": system_prompt},
    ]

    # Few-shot примеры (полные, без '...', валидные)
    examples = [
        # Пример 1: Таблица с ролями (тех. стек)
        {
            "user": """Шаг | Роль

1. Создание запроса | Инициатор
2. Внесение технологии в тех. стек в статусе "Consideration" | Координатор
3. Оценка по техническим критериям | Инициатор
4. Уточнение требований к технологии | Совет по технологиям, Координатор, Владелец продукта
5. Оценка ограничений и рисков принятия новой мажорной версии | Координатор
6. Презентация новой мажорной версии | Инициатор
7. Принятие решения | Главный корпоративный архитектор, Совет по технологиям
8. Внесение записи в тех. стек, оповещение | Координатор
9. Перевод прошлой мажорной версии в статус Deprecated | Координатор
10. Оповещение | Координатор
11. Планирование обновления продукта | Координатор
12. Планирование автоматизации ресурсов | Координатор
13. Изменение атрибутов технологии в стеке | Координатор""",
            "assistant": """@startuml
skinparam shadowing false
skinparam monochrome true

|Инициатор|
start
:Создание запроса;

|Координатор|
:Внесение технологии в тех. стек (Consideration);

|Инициатор|
:Оценка по техническим критериям;

|Совет по технологиям, Координатор, Владелец продукта|
:Уточнение требований к технологии;

|Координатор|
:Оценка ограничений и рисков;

|Инициатор|
:Презентация новой мажорной версии;

|Главный корпоративный архитектор, Совет по технологиям|
:Принятие решения;

|Координатор|
:Внесение записи в тех. стек;
:Оповещение;
:Перевод прошлой версии в Deprecated;
:Оповещение;
:Планирование обновления продукта;
:Планирование автоматизации ресурсов;

:Изменение атрибутов технологии;

stop
@enduml"""
        },

        # Пример 2: Ещё одна таблица с ролями
        {
            "user": """Шаг | Роль

1. Создание запроса | Инициатор
2. Подтверждение потребности в изменении условий использования | Координатор
3. Определение пользователей технологии | Совет по технологиям
4. Оценка трудозатрат | Совет по технологиям
5. Принятие решения | Главный корпоративный архитектор, Совет по технологиям
6. Изменение статуса технологии | Координатор
7. Оповещение клиентов | Команда разработки
8. Выполнение работ | ДЭиС""",
            "assistant": """@startuml
skinparam shadowing false

|Инициатор|
start
:Создание запроса;

|Координатор|
:Подтверждение потребности в изменении;

|Совет по технологиям|
:Определение пользователей технологии;
:Оценка трудозатрат;

|Главный корпоративный архитектор, Совет по технологиям|
:Принятие решения;

|Координатор|
:Изменение статуса технологии;

|Команда разработки|
:Оповещение клиентов;

|ДЭиС|
:Выполнение работ;

stop
@enduml"""
        },

        # Пример 3: Простая последовательность без ролей
        {
            "user": """Шаг

1. Создание запроса
2. Рассмотрение запроса
3. Анализ запроса
4. Принятие решения о реализации изменения
5. Реализация улучшений и исправление несоответствий
6. Оценка результатов""",
            "assistant": """@startuml
start
:Создание запроса;
:Рассмотрение запроса;
:Анализ запроса;
:Принятие решения о реализации;
:Реализация улучшений и исправлений;
:Оценка результатов;
stop
@enduml"""
        },

        # Пример 4: Чайник с событием и ветвлением
        {
            "user": """Шаг

1. Начало
2. Наполнить чайник
3. Включить чайник
4. Ждать закипания
5. Событие: Чайник закипел
6. Таймер: 30 секунд
7. Автовыключение
8. Ручное выключение
9. Чай готов""",
            "assistant": """@startuml
start
:Наполнить чайник;
:Включить чайник;
:Ждать закипания;
note right: Событие: Чайник закипел
:Таймер: 30 секунд;
if (Выключен автоматически?) then (да)
  :Автовыключение;
else (нет)
  :Ручное выключение;
endif
:Чай готов;
stop
@enduml"""
        },

        # Пример 5: Решение с ветвлением (работа/учёба)
        {
            "user": """Шаг

1. Зачисление
2. Обновить резюме
3. Искать гибкую работу
4. Решение: Найдена?
5. Договориться о графике
6. Совмещение работы и учебы
7. Фокус на учебе
8. Только учеба""",
            "assistant": """@startuml
start
:Зачисление;
:Обновить резюме;
:Искать гибкую работу;
if (Работа найдена?) then (да)
  :Договориться о графике;
  :Совмещение работы и учебы;
else (нет)
  :Фокус на учебе;
  :Только учеба;
endif
stop
@enduml"""
        },
    ]

    # Добавляем все few-shot
    for ex in examples:
        messages.append({"role": "user", "content": ex["user"]})
        messages.append({"role": "assistant", "content": ex["assistant"]})

    # Финальный запрос
    messages.append(
        {"role": "user", "content": f"Generate PlantUML activity diagram for this description:\n\n{description}"})

    return messages