llm_classification.py

from llm import llm
from langchain.prompts import PromptTemplate
from langchain_core.messages import SystemMessage, HumanMessage


def get_answer(q: str) -> int:
    system_prompt = (
        "You are an assistant classifying a sentence sentiment. "
        "Output strictly '1' for positive or '0' for negative. No extra text."
    )

    user_prompt = PromptTemplate(
        input_variables=["q"],
        template=(
            "Classify this sentence: {q}\n"
            "Output only 1 if positive or 0 if negative."
        ),
    )

    try:
        resp = llm.invoke([
            SystemMessage(content=system_prompt),
            HumanMessage(content=user_prompt.format(q=q))
        ])
        text = resp.content.strip()  # lấy nội dung model trả về
        return 1 if text == "1" else 0
    except Exception as e:
        print(f"Error: {e}")
        return 0  # hoặc None tùy bạn muốn