from smolagents import Tool
from typing import Any, Optional

class SimpleTool(Tool):
    name = "classify_topic"
    description = "Ph\u00e2n lo\u1ea1i xem v\u0103n b\u1ea3n ti\u1ebfng Vi\u1ec7t c\u00f3 li\u00ean quan \u0111\u1ebfn ch\u1ee7 \u0111\u1ec1 hay kh\u00f4ng."
    inputs = {'text': {'type': 'string', 'description': 'Nội dung văn bản tiếng Việt.'}, 'topic': {'type': 'string', 'description': 'Chủ đề cần kiểm tra liên quan.'}}
    output_type = "boolean"

    def forward(self, text: str, topic: str) -> bool:
        """
        Phân loại xem văn bản tiếng Việt có liên quan đến chủ đề hay không.

        Args:
            text (str): Nội dung văn bản tiếng Việt.
            topic (str): Chủ đề cần kiểm tra liên quan.

        Returns:
            bool: True nếu liên quan đến topic, False nếu không.
        """
        import torch
        device = 0 if torch.cuda.is_available() else -1  # -1 cho CPU
        from transformers import pipeline

        classifier = pipeline(
            "zero-shot-classification",
            model="vicgalle/xlm-roberta-large-xnli-anli",
            device=device
        )

        candidate_labels = [topic, f"không liên quan {topic}"]
        result = classifier(text, candidate_labels)
        predicted_label = result["labels"][0]

        return predicted_label == topic