src/tools.py

"""
MCP tool definitions for sentiment analysis server.

This module defines the tools available through the Model Context Protocol,
including sentiment analysis, batch processing, and analyzer information tools.
"""

import logging
from typing import Dict, Any, List, Optional
from pydantic import BaseModel, Field
import asyncio

from .sentiment_analyzer import get_analyzer, SentimentAnalyzer


class SentimentAnalysisInput(BaseModel):
    """Input schema for sentiment analysis tool."""
    text: str = Field(..., description="Text to analyze for sentiment", min_length=1, max_length=10000)
    backend: Optional[str] = Field("auto", description="Analysis backend: 'textblob', 'transformers', or 'auto'")


class BatchSentimentAnalysisInput(BaseModel):
    """Input schema for batch sentiment analysis tool."""
    texts: List[str] = Field(..., description="List of texts to analyze for sentiment", min_items=1, max_items=100)
    backend: Optional[str] = Field("auto", description="Analysis backend: 'textblob', 'transformers', or 'auto'")


class AnalyzerInfoInput(BaseModel):
    """Input schema for analyzer info tool."""
    backend: Optional[str] = Field("auto", description="Backend to get info for")


class MCPTools:
    """
    MCP tool registry and handlers for sentiment analysis.
    
    This class manages the registration and execution of tools available
    through the Model Context Protocol interface.
    """
    
    def __init__(self):
        self.logger = logging.getLogger(__name__)
        self._tools = {}
        self._register_tools()
    
    def _register_tools(self) -> None:
        """Register all available MCP tools."""
        self._tools = {
            "analyze_sentiment": {
                "name": "analyze_sentiment",
                "description": "Analyze the sentiment of a given text",
                "inputSchema": {
                    "type": "object",
                    "properties": {
                        "text": {
                            "type": "string",
                            "description": "Text to analyze for sentiment",
                            "minLength": 1,
                            "maxLength": 10000
                        },
                        "backend": {
                            "type": "string",
                            "description": "Analysis backend: 'textblob', 'transformers', or 'auto'",
                            "enum": ["textblob", "transformers", "auto"],
                            "default": "auto"
                        }
                    },
                    "required": ["text"]
                },
                "handler": self._handle_analyze_sentiment
            },
            
            "analyze_sentiment_batch": {
                "name": "analyze_sentiment_batch",
                "description": "Analyze sentiment for multiple texts in batch",
                "inputSchema": {
                    "type": "object",
                    "properties": {
                        "texts": {
                            "type": "array",
                            "description": "List of texts to analyze for sentiment",
                            "items": {
                                "type": "string",
                                "minLength": 1,
                                "maxLength": 10000
                            },
                            "minItems": 1,
                            "maxItems": 100
                        },
                        "backend": {
                            "type": "string",
                            "description": "Analysis backend: 'textblob', 'transformers', or 'auto'",
                            "enum": ["textblob", "transformers", "auto"],
                            "default": "auto"
                        }
                    },
                    "required": ["texts"]
                },
                "handler": self._handle_analyze_sentiment_batch
            },
            
            "get_analyzer_info": {
                "name": "get_analyzer_info",
                "description": "Get information about the sentiment analyzer configuration",
                "inputSchema": {
                    "type": "object",
                    "properties": {
                        "backend": {
                            "type": "string",
                            "description": "Backend to get info for",
                            "enum": ["textblob", "transformers", "auto"],
                            "default": "auto"
                        }
                    },
                    "required": []
                },
                "handler": self._handle_get_analyzer_info
            },
            
            "health_check": {
                "name": "health_check",
                "description": "Check the health status of the sentiment analysis service",
                "inputSchema": {
                    "type": "object",
                    "properties": {},
                    "required": []
                },
                "handler": self._handle_health_check
            }
        }
        
        self.logger.info(f"Registered {len(self._tools)} MCP tools")
    
    def get_tools(self) -> List[Dict[str, Any]]:
        """
        Get list of available tools for MCP protocol.
        
        Returns:
            List of tool definitions
        """
        return [
            {
                "name": tool["name"],
                "description": tool["description"],
                "inputSchema": tool["inputSchema"]
            }
            for tool in self._tools.values()
        ]
    
    async def call_tool(self, name: str, arguments: Dict[str, Any]) -> Dict[str, Any]:
        """
        Call a registered tool with given arguments.
        
        Args:
            name: Tool name
            arguments: Tool arguments
            
        Returns:
            Tool execution result
            
        Raises:
            ValueError: If tool not found or arguments invalid
            RuntimeError: If tool execution fails
        """
        if name not in self._tools:
            raise ValueError(f"Tool '{name}' not found. Available tools: {list(self._tools.keys())}")
        
        tool = self._tools[name]
        handler = tool["handler"]
        
        try:
            self.logger.info(f"Calling tool '{name}' with arguments: {arguments}")
            result = await handler(arguments)
            self.logger.info(f"Tool '{name}' completed successfully")
            return result
            
        except Exception as e:
            self.logger.error(f"Tool '{name}' failed: {e}")
            raise RuntimeError(f"Tool execution failed: {e}")
    
    async def _handle_analyze_sentiment(self, arguments: Dict[str, Any]) -> Dict[str, Any]:
        """
        Handle sentiment analysis tool call.
        
        Args:
            arguments: Tool arguments containing text and optional backend
            
        Returns:
            Sentiment analysis result
        """
        try:
            # Validate input
            input_data = SentimentAnalysisInput(**arguments)
            
            # Get analyzer and perform analysis
            analyzer = await get_analyzer(input_data.backend)
            result = await analyzer.analyze(input_data.text)
            
            return {
                "success": True,
                "result": result.to_dict(),
                "metadata": {
                    "backend": analyzer.backend,
                    "text_length": len(input_data.text),
                    "model_info": analyzer.get_info()
                }
            }
            
        except Exception as e:
            return {
                "success": False,
                "error": str(e),
                "error_type": type(e).__name__
            }
    
    async def _handle_analyze_sentiment_batch(self, arguments: Dict[str, Any]) -> Dict[str, Any]:
        """
        Handle batch sentiment analysis tool call.
        
        Args:
            arguments: Tool arguments containing texts and optional backend
            
        Returns:
            Batch sentiment analysis results
        """
        try:
            # Validate input
            input_data = BatchSentimentAnalysisInput(**arguments)
            
            # Get analyzer and perform batch analysis
            analyzer = await get_analyzer(input_data.backend)
            results = await analyzer.analyze_batch(input_data.texts)
            
            # Convert results to dictionaries
            result_dicts = [result.to_dict() for result in results]
            
            # Calculate summary statistics
            labels = [result.label.value for result in results]
            label_counts = {
                "positive": labels.count("positive"),
                "negative": labels.count("negative"),
                "neutral": labels.count("neutral")
            }
            
            avg_confidence = sum(result.confidence for result in results) / len(results)
            
            return {
                "success": True,
                "results": result_dicts,
                "summary": {
                    "total_texts": len(input_data.texts),
                    "label_distribution": label_counts,
                    "average_confidence": round(avg_confidence, 4)
                },
                "metadata": {
                    "backend": analyzer.backend,
                    "model_info": analyzer.get_info()
                }
            }
            
        except Exception as e:
            return {
                "success": False,
                "error": str(e),
                "error_type": type(e).__name__
            }
    
    async def _handle_get_analyzer_info(self, arguments: Dict[str, Any]) -> Dict[str, Any]:
        """
        Handle analyzer info tool call.
        
        Args:
            arguments: Tool arguments containing optional backend
            
        Returns:
            Analyzer configuration information
        """
        try:
            # Validate input
            input_data = AnalyzerInfoInput(**arguments)
            
            # Get analyzer info
            analyzer = await get_analyzer(input_data.backend)
            info = analyzer.get_info()
            
            return {
                "success": True,
                "info": info,
                "available_backends": ["textblob", "transformers", "auto"],
                "recommended_backend": "transformers" if info.get("transformers_available") else "textblob"
            }
            
        except Exception as e:
            return {
                "success": False,
                "error": str(e),
                "error_type": type(e).__name__
            }
    
    async def _handle_health_check(self, arguments: Dict[str, Any]) -> Dict[str, Any]:
        """
        Handle health check tool call.
        
        Args:
            arguments: Tool arguments (empty for health check)
            
        Returns:
            Health status information
        """
        try:
            # Test basic functionality
            test_text = "This is a test message for health check."
            analyzer = await get_analyzer("auto")
            result = await analyzer.analyze(test_text)
            
            return {
                "success": True,
                "status": "healthy",
                "test_result": result.to_dict(),
                "analyzer_info": analyzer.get_info(),
                "timestamp": asyncio.get_event_loop().time()
            }
            
        except Exception as e:
            return {
                "success": False,
                "status": "unhealthy",
                "error": str(e),
                "error_type": type(e).__name__,
                "timestamp": asyncio.get_event_loop().time()
            }


# Global tools instance
_global_tools: Optional[MCPTools] = None


def get_tools() -> MCPTools:
    """
    Get or create global MCP tools instance.
    
    Returns:
        MCPTools instance
    """
    global _global_tools
    
    if _global_tools is None:
        _global_tools = MCPTools()
    
    return _global_tools


async def list_tools() -> List[Dict[str, Any]]:
    """
    Get list of available MCP tools.
    
    Returns:
        List of tool definitions
    """
    tools = get_tools()
    return tools.get_tools()


async def call_tool(name: str, arguments: Dict[str, Any]) -> Dict[str, Any]:
    """
    Call an MCP tool with given arguments.
    
    Args:
        name: Tool name
        arguments: Tool arguments
        
    Returns:
        Tool execution result
    """
    tools = get_tools()
    return await tools.call_tool(name, arguments)


if __name__ == "__main__":
    # Example usage
    async def main():
        tools = get_tools()
        
        # List available tools
        available_tools = tools.get_tools()
        print("Available tools:")
        for tool in available_tools:
            print(f"- {tool['name']}: {tool['description']}")
        
        print("\n" + "="*50 + "\n")
        
        # Test sentiment analysis tool
        result = await tools.call_tool("analyze_sentiment", {
            "text": "I love this new feature! It's absolutely amazing!",
            "backend": "textblob"
        })
        print("Sentiment analysis result:")
        print(result)
        
        print("\n" + "="*50 + "\n")
        
        # Test batch analysis
        batch_result = await tools.call_tool("analyze_sentiment_batch", {
            "texts": [
                "This is great!",
                "I hate this.",
                "It's okay, I guess."
            ],
            "backend": "textblob"
        })
        print("Batch analysis result:")
        print(batch_result)
        
        print("\n" + "="*50 + "\n")
        
        # Test health check
        health_result = await tools.call_tool("health_check", {})
        print("Health check result:")
        print(health_result)
    
    asyncio.run(main())