# services/agents/ner_agent.py
"""
Named Entity Recognition Agent - Wraps utilities/ner.py
"""
from typing import Dict, Any
from services.agents.base_agent import BaseUtilityAgent
from utilities.ner import ner_remote


class NERAgent(BaseUtilityAgent):
    """
    Autonomous agent for named entity recognition.
    """
    
    def __init__(self):
        super().__init__(
            name="ner",
            role="Named Entity Recognition Specialist",
            goal="Identify and extract named entities (people, organizations, locations, dates) with high precision",
            backstory="""You are an expert in named entity recognition and information extraction.
            You can identify people names, organizations, locations, dates, and other entities
            in text with high accuracy. You understand context and can disambiguate entities.
            You validate NER results for completeness and accuracy.""",
            utility_function=ner_remote
        )
    
    def _prepare_task_description(self, input_data: Dict[str, Any]) -> str:
        """Prepare task description for the agent."""
        has_text = "text" in input_data
        filename = input_data.get("filename", "document")
        source = "provided text" if has_text else f"{filename}"
        
        return f"""Validate the named entity recognition results from {source}.
        
        Assess NER quality:
        - Completeness: Were all entities identified?
        - Accuracy: Are entity labels correct?
        - Precision: Are boundaries correctly identified?
        - Entity types: Are PERSON, ORG, LOC, DATE etc. properly classified?
        
        Provide confidence score (0.0-1.0)."""