| 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 |