mcp/tools.py

from typing import Dict, Any, List
from mcp import Tool
import logging
from services import (
    kpa_model_manager,
    stance_model_manager,
    chat_service
)

logger = logging.getLogger(__name__)

async def predict_kpa_tool(arguments: Dict[str, Any]) -> Dict[str, Any]:
    """Tool for keypoint-argument matching prediction"""
    try:
        argument = arguments.get("argument", "")
        key_point = arguments.get("key_point", "")
        
        if not argument or not key_point:
            return {"error": "Both argument and key_point are required"}
        
        result = kpa_model_manager.predict(argument, key_point)
        return {
            "prediction": result["prediction"],
            "label": result["label"],
            "confidence": result["confidence"],
            "probabilities": result["probabilities"]
        }
    except Exception as e:
        logger.error(f"KPA tool error: {str(e)}")
        return {"error": str(e)}

async def predict_stance_tool(arguments: Dict[str, Any]) -> Dict[str, Any]:
    """Tool for stance detection prediction"""
    try:
        topic = arguments.get("topic", "")
        argument = arguments.get("argument", "")
        
        if not topic or not argument:
            return {"error": "Both topic and argument are required"}
        
        result = stance_model_manager.predict(topic, argument)
        return {
            "predicted_stance": result["predicted_stance"],
            "confidence": result["confidence"],
            "probability_con": result["probability_con"],
            "probability_pro": result["probability_pro"]
        }
    except Exception as e:
        logger.error(f"Stance tool error: {str(e)}")
        return {"error": str(e)}

async def batch_stance_tool(arguments: Dict[str, Any]) -> Dict[str, Any]:
    """Tool for batch stance detection"""
    try:
        items = arguments.get("items", [])
        
        if not items:
            return {"error": "Items list is required"}
        
        results = []
        for item in items:
            result = stance_model_manager.predict(item["topic"], item["argument"])
            results.append({
                "topic": item["topic"],
                "argument": item["argument"],
                **result
            })
        
        return {
            "results": results,
            "total_processed": len(results)
        }
    except Exception as e:
        logger.error(f"Batch stance tool error: {str(e)}")
        return {"error": str(e)}

async def generate_argument_tool(arguments: Dict[str, Any]) -> Dict[str, Any]:
    """Tool for argument generation (à compléter avec votre modèle)"""
    try:
        prompt = arguments.get("prompt", "")
        context = arguments.get("context", "")
        
        if not prompt:
            return {"error": "Prompt is required"}
        
        # TODO: Intégrer votre modèle d'argument generation ici
        # Pour l'instant, placeholder
        from services.chat_service import generate_chat_response
        
        response = generate_chat_response(
            user_input=f"Generate argument for: {prompt}. Context: {context}",
            system_prompt="You are an argument generation assistant. Generate persuasive arguments based on the given prompt and context."
        )
        
        return {
            "generated_argument": response,
            "prompt": prompt,
            "context": context
        }
    except Exception as e:
        logger.error(f"Argument generation tool error: {str(e)}")
        return {"error": str(e)}

async def voice_chat_tool(arguments: Dict[str, Any]) -> Dict[str, Any]:
    """Tool for voice chat interaction"""
    try:
        text = arguments.get("text", "")
        conversation_id = arguments.get("conversation_id", "")
        
        if not text:
            return {"error": "Text input is required"}
        
        # Utiliser le service de chat existant
        from services.chat_service import generate_chat_response
        
        response = generate_chat_response(
            user_input=text,
            conversation_id=conversation_id if conversation_id else None
        )
        
        # Optionnel: Ajouter TTS si nécessaire
        tts_required = arguments.get("tts", False)
        audio_url = None
        
        if tts_required:
            from services.tts_service import text_to_speech
            # TODO: Gérer le stockage et l'URL de l'audio
            
        return {
            "response": response,
            "conversation_id": conversation_id,
            "has_audio": tts_required,
            "audio_url": audio_url
        }
    except Exception as e:
        logger.error(f"Voice chat tool error: {str(e)}")
        return {"error": str(e)}

def get_tools() -> List[Tool]:
    """Retourne tous les outils disponibles"""
    return [
        Tool(
            name="predict_kpa",
            description="Predict keypoint-argument matching for a single pair",
            input_schema={
                "type": "object",
                "properties": {
                    "argument": {"type": "string", "description": "The argument text"},
                    "key_point": {"type": "string", "description": "The key point to evaluate"}
                },
                "required": ["argument", "key_point"]
            },
            execute=predict_kpa_tool
        ),
        Tool(
            name="predict_stance",
            description="Predict stance for a topic-argument pair",
            input_schema={
                "type": "object",
                "properties": {
                    "topic": {"type": "string", "description": "The debate topic"},
                    "argument": {"type": "string", "description": "The argument to classify"}
                },
                "required": ["topic", "argument"]
            },
            execute=predict_stance_tool
        ),
        Tool(
            name="batch_predict_stance",
            description="Predict stance for multiple topic-argument pairs",
            input_schema={
                "type": "object",
                "properties": {
                    "items": {
                        "type": "array",
                        "items": {
                            "type": "object",
                            "properties": {
                                "topic": {"type": "string"},
                                "argument": {"type": "string"}
                            },
                            "required": ["topic", "argument"]
                        },
                        "description": "List of topic-argument pairs"
                    }
                },
                "required": ["items"]
            },
            execute=batch_stance_tool
        ),
        Tool(
            name="generate_argument",
            description="Generate persuasive arguments based on prompt and context",
            input_schema={
                "type": "object",
                "properties": {
                    "prompt": {"type": "string", "description": "Main topic or question"},
                    "context": {"type": "string", "description": "Additional context"},
                    "stance": {
                        "type": "string", 
                        "enum": ["pro", "con", "neutral"],
                        "description": "Desired stance"
                    }
                },
                "required": ["prompt"]
            },
            execute=generate_argument_tool
        ),
        Tool(
            name="voice_chat",
            description="Chat with voice assistant capabilities",
            input_schema={
                "type": "object",
                "properties": {
                    "text": {"type": "string", "description": "Text input"},
                    "conversation_id": {"type": "string", "description": "Conversation ID for context"},
                    "tts": {"type": "boolean", "description": "Generate audio response"}
                },
                "required": ["text"]
            },
            execute=voice_chat_tool
        )
    ]