python/backend/core/nlp/intent_classifier.py - master - ksilisk/spbtechrun_hack

python/backend/core/nlp/intent_classifier.py

63 строки2 KB

Shaliko Salimov

all prompts moving to the prompts directory and README.me refactoring

02 дек 2025, 21:18

02 дек 2025, 21:180511ade

from __future__ import annotations

import json

from langchain_core.language_models import BaseChatModel
from langchain_core.prompts import ChatPromptTemplate

from core.nlp.prompts import build_answer_generation_messages
from core.prompt.prompt_store import PromptStore
from schemas.chat import Message
from schemas.intents import IntentClassificationResult, IntentType

_INTENT_CLASSIFICATION_PROMPT = "intent_classification.yaml"


def _build_prompt(system_prompt: str) -> ChatPromptTemplate:
    return ChatPromptTemplate.from_messages(
        [
            ("system", system_prompt),
            ("placeholder", "{chat_history}"),
            ("user", "Текущее сообщение пользователя для классификации:\n {current_user_message}"),
        ]
    )


def _extract_json(text: str) -> str:
    text = text.strip()
    if "```" in text:
        start = text.find("{")
        end = text.rfind("}")
        return text[start: end + 1]
    return text


class IntentClassifier:
    def __init__(self, llm: BaseChatModel, prompt_store: PromptStore) -> None:
        self._llm = llm
        self._prompt = _build_prompt(prompt_store.load(_INTENT_CLASSIFICATION_PROMPT))
        self._chain = self._prompt | self._llm
        self._prompt_store = prompt_store

    async def classify(self, messages: list[Message], current_message: str) -> IntentClassificationResult:
        resp = await self._chain.ainvoke({"chat_history": build_answer_generation_messages(messages),
                                          "current_user_message": current_message})
        raw = resp.content

        try:
            data = json.loads(_extract_json(raw))
        except json.JSONDecodeError:
            return IntentClassificationResult(
                intent=IntentType.UNKNOWN,
                confidence=0.0,
            )

        intent_name = data.get("intent", "UNKNOWN")
        confidence = float(data.get("confidence", 0.0))

        try:
            intent = IntentType(intent_name)
        except ValueError:
            intent = IntentType.UNKNOWN

        return IntentClassificationResult(intent=intent, confidence=confidence)