tools/classifier_tool.py

from tools.base_tool import BaseTool
from mistral_hf_wrapper import MistralInference

class ClassifierTool(BaseTool):
    name = 'open_classifier'
    description = "Classifies given items into categories from a specific knowledge area."

    inputs = {
        'knowledge_area': {
            'type': 'string',
            'description': 'Domain or field for classification (e.g., biology, technology).',
        },
        'environment': {
            'type': 'string',
            'description': 'Brief description of the context in which items are evaluated.',
        },
        'categories': {
            'type': 'string',
            'description': 'Comma-separated list of categories to classify into.',
        },
        'items': {
            'type': 'string',
            'description': 'Comma-separated list of items to classify.',
        },
    }
    output_type = 'string'

    def __init__(self, model=None):
        self.model = model or MistralInference()
        super().__init__()

    def forward(self, knowledge_area: str, environment: str, categories: str, items: str) -> str:
        prompt = self._build_prompt(knowledge_area, environment, categories, items)
        response = self.model(prompt)
        return response.strip()

    def _build_prompt(self, knowledge_area, environment, categories, items) -> str:
        return f"""You are an expert classifier in the field of {knowledge_area}.
Context: {environment}
You must assign each item below to one of the following categories:
{categories}

Items to classify:
{items}

Rules:
- Use your {knowledge_area} knowledge only.
- Use the environment only if it's essential for disambiguation.
- Use "Other" only if an item clearly doesn't belong to any provided category.
- Format:
Category 1: item1, item2
Category 2: item3
Other: item4
"""