src/classifier.py - master - Vladimir_Py/cp_generator

src/classifier.py
68 строк4 KB
Test Bot
добавил фронтэнд и бэк
03 май 2026, 08:47
03 май 2026, 08:471be8190
import logging
from langchain_core.prompts import PromptTemplate
from langchain_openai import ChatOpenAI

from schemas.proj_classes import CategoriesSchema, ClassificationResult

logger = logging.getLogger(__name__)

class ProjectClassifier:
    """Классификация ИИ проектов по тексту ТЗ"""

    def __init__(self, llm: ChatOpenAI, schema: CategoriesSchema):
        self.llm = llm
        self.schema = schema

        self._categories_text = self._build_categories_text()
        self._exclusions_text = self._build_exclusions_text()

    def _build_categories_text(self) -> str:
        lines = []
        for key, info in self.schema.categories.items():
            name_part = f" ({info.name_ru})" if info.name_ru else ""
            lines.append(f"- '{key}'{name_part}: {info.description}")
        return "\n".join(lines)

    def _build_exclusions_text(self) -> str:
        lines = []
        for key, desc in self.schema.exclusion_rules.items():
            lines.append(f"- '{key}': {desc}")
        # Добавляем универсальную заглушку для всего остального
        lines.append("- 'skip_other': Проект не нарушает другие правила, но вообще не относится к целевым категориям.")
        return "\n".join(lines)

    def classify(self, raw_tz: str) -> ClassificationResult:
        logger.info("Начинаем классификацию ТЗ...")

        structured_llm = self.llm.with_structured_output(ClassificationResult)

        prompt = PromptTemplate.from_template(
            "Ты классификатор технических заданий.\n\n"
            "СТРАТЕГИЯ (МАКСИМАЛЬНЫЙ ОХВАТ):\n"
            "Твоя задача — не упустить НИ ОДНОГО потенциального проекта. Если в ТЗ есть упоминание, "
            "что бот должен понимать свободный человеческий язык, задавать\отвечать на вопросы, распознавать изображения"
            " и т.д., — ставь is_target=True. \n\n"
            "Генерация контента (статьи, видео, фото), \n"
            "АЛГОРИТМ ПРИНЯТИЯ РЕШЕНИЯ:\n"
            "1. ШАГ 1: Внимательно изучи ТЗ и примерь его к ЦЕЛЕВЫМ КАТЕГОРИЯМ. Если оно подходит хотя бы под одну из них — сразу возвращай is_target=True и ключ категории.\n"
            "2. ШАГ 2: если ТЗ ЯВНО НЕ ПОДХОДИТ ни в одну из целевых категорий, переходи к ПРАВИЛАМ ИСКЛЮЧЕНИЯ.\n"
            "3. ШАГ 3: Найди наиболее подходящую причину отбраковки, верни is_target=False и ключ исключения (если ничего не подошло, используй 'skip_other').\n\n"

            "ЦЕЛЕВЫЕ КАТЕГОРИИ (is_target = True):\n"
            "{categories_text}\n\n"
            "ПРАВИЛА ИСКЛЮЧЕНИЯ (is_target = False):\n"
            "{exclusions_text}\n\n"
            "СЫРОЕ ТЗ ОТ КЛИЕНТА:\n"
            "{raw_tz}"
        )

        chain = prompt | structured_llm
        # подстановка реальных данных в шаблон
        result: ClassificationResult = chain.invoke({
            "categories_text": self._categories_text,
            "exclusions_text": self._exclusions_text,
            "raw_tz": raw_tz
        })

        logger.debug(f"Результат классификации: {result.model_dump()}")
        return result