from smolagents import Tool
from datasets import load_dataset

class LawTool(Tool):
    name = "law_tool"
    description = """
    This is a tool that returns law content by input the category and number."""
    inputs = {
        "category": {
            "type": "string",
            "description": "the law category (such as 民法, 中華民國刑法, 民事訴訟法, 刑事訴訟法, 律師法 etc)",
        },
        "number": {
            "type": "integer",
            "description": "the law number (such as 23)"
        }
    }
    output_type = "string"
    law = None

    def __init__(self):
        dataset = load_dataset("robin0307/law", split='train')
        self.law = dataset.to_pandas()
        super().__init__()

    def forward(self, category: str, number: int):
        if category == "刑法":
            category = "中華民國刑法"
        
        data = self.law.loc[(self.law["category"]==category) & (self.law["number"]==number), "content"].values[0]
        return data