models/mcp_models.py

from pydantic import BaseModel, Field, ConfigDict
from typing import Any, Dict, List, Optional

class ToolCallRequest(BaseModel):
    """Request for calling an MCP tool"""
    tool_name: str
    arguments: Dict[str, Any] = {}

class ToolCallResponse(BaseModel):
    """Response from MCP tool call"""
    success: bool
    result: Optional[Dict[str, Any]] = None
    error: Optional[str] = None
    tool_name: str

# Response models for individual MCP tools
class DetectStanceResponse(BaseModel):
    """Response model for stance detection"""
    model_config = ConfigDict(
        json_schema_extra={
            "example": {
                "predicted_stance": "PRO",
                "confidence": 0.9598,
                "probability_con": 0.0402,
                "probability_pro": 0.9598
            }
        }
    )
    
    predicted_stance: str = Field(..., description="PRO or CON")
    confidence: float = Field(..., ge=0.0, le=1.0, description="Confidence score")
    probability_con: float = Field(..., ge=0.0, le=1.0)
    probability_pro: float = Field(..., ge=0.0, le=1.0)

class MatchKeypointResponse(BaseModel):
    """Response model for keypoint matching"""
    model_config = ConfigDict(
        json_schema_extra={
            "example": {
                "prediction": 1,
                "label": "apparie",
                "confidence": 0.8157,
                "probabilities": {
                    "non_apparie": 0.1843,
                    "apparie": 0.8157
                }
            }
        }
    )
    
    prediction: int = Field(..., description="1 = apparie, 0 = non_apparie")
    label: str = Field(..., description="apparie or non_apparie")
    confidence: float = Field(..., ge=0.0, le=1.0, description="Confidence score")
    probabilities: Dict[str, float] = Field(..., description="Dictionary of class probabilities")

class TranscribeAudioResponse(BaseModel):
    """Response model for audio transcription"""
    model_config = ConfigDict(
        json_schema_extra={
            "example": {
                "text": "Hello, this is the transcribed text from the audio file."
            }
        }
    )
    
    text: str = Field(..., description="Transcribed text from audio")

class GenerateSpeechResponse(BaseModel):
    """Response model for speech generation"""
    model_config = ConfigDict(
        json_schema_extra={
            "example": {
                "audio_path": "temp_audio/tts_e9b78164.wav"
            }
        }
    )
    
    audio_path: str = Field(..., description="Path to generated audio file")

class GenerateArgumentResponse(BaseModel):
    """Response model for argument generation"""
    model_config = ConfigDict(
        json_schema_extra={
            "example": {
                "argument": "Climate change is a pressing issue that requires immediate action..."
            }
        }
    )
    
    argument: str = Field(..., description="Generated debate argument")

class ResourceInfo(BaseModel):
    """Information about an MCP resource"""
    uri: str
    name: str
    description: Optional[str] = None
    mime_type: str

class ToolInfo(BaseModel):
    """Information about an MCP tool"""
    name: str
    description: str
    input_schema: Dict[str, Any]

class ResourceListResponse(BaseModel):
    """Response for listing resources"""
    resources: List[ResourceInfo]
    count: int

class ToolListResponse(BaseModel):
    """Response for listing tools"""
    tools: List[ToolInfo]
    count: int