Update tools/classifier_tool.py

tools/classifier_tool.py

@@ -1,67 +1,54 @@
-from smolagents import Tool
-import requests
-import os
+from tools.base import Tool
+from mistral_hf_wrapper import MistralInference
 
-class ClassifyItemsTool(Tool):
-    name = 'hf_classifier'
-    description = """Categorizes a list of items using domain-specific knowledge. Intended for structured classification tasks."""
+class ClassifierTool(Tool):
+    name = 'open_classifier'
+    description = "Classifies given items into categories from a specific knowledge area."
 
     inputs = {
-        'domain': {
+        'knowledge_area': {
             'type': 'string',
-            'description': 'The knowledge domain to use for classification (e.g., medicine, education, etc.).'
+            'description': 'Domain or field for classification (e.g., biology, technology).',
         },
-        'context': {
+        'environment': {
             'type': 'string',
-            'description': 'A brief description of the environment where items appear (helps disambiguate meanings).'
+            'description': 'Brief description of the context in which items are evaluated.',
         },
         'categories': {
             'type': 'string',
-            'description': 'Comma-separated category list to classify items into.'
+            'description': 'Comma-separated list of categories to classify into.',
         },
         'items': {
             'type': 'string',
-            'description': 'Comma-separated list of items to classify.'
+            'description': 'Comma-separated list of items to classify.',
         },
     }
     output_type = 'string'
 
-    def __init__(self, hf_model_url: str | None = None, **kwargs):
-        self.api_url = hf_model_url or os.getenv("HF_ENDPOINT")
-        if not self.api_url:
-            raise ValueError("HF_ENDPOINT must be set as environment variable or passed to constructor.")
-        super().__init__(**kwargs)
+    def __init__(self, model=None):
+        self.model = model or MistralInference()
+        super().__init__()
 
-    def forward(self, domain: str, context: str, categories: str, items: str) -> str:
-        prompt = self._build_prompt(domain, context, categories, items)
+    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()
 
-        response = requests.post(
-            self.api_url,
-            headers={"Content-Type": "application/json"},
-            json={"inputs": prompt}
-        )
-        response.raise_for_status()
-        result = response.json()
-
-        if isinstance(result, list) and 'generated_text' in result[0]:
-            return result[0]['generated_text'].strip()
-        elif isinstance(result, dict) and 'generated_text' in result:
-            return result['generated_text'].strip()
-        return str(result)
-
-    def _build_prompt(self, domain: str, context: str, categories: str, items: str) -> str:
-        return f"""
-You are a {domain} expert working within a {context} context.
-Classify the following items into the specified categories using your domain expertise.
-
-Categories:
+    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}
 
-Return explanation followed by a classification in this format:
-Category 1: item list
-Category 2: item list
-Other (if uncertain): item list
+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
 """