ToGMAL MCP Server

#!/usr/bin/env python3 """ ToGMAL Chat Demo with MCP Tool Integration ========================================== Interactive chat demo where a free LLM can call MCP tools to provide informed responses about prompt difficulty, safety analysis, and more. Features: - Chat with Mistral-7B-Instruct (free via HuggingFace Inference API) - LLM can call MCP tools to analyze prompts and assess difficulty - Transparent tool calling with results shown to user - No API key required (uses public Inference API) """ import gradio as gr import json import os import re from pathlib import Path from typing import List, Dict, Tuple, Optional from benchmark_vector_db import BenchmarkVectorDB import logging logging.basicConfig(level=logging.INFO) logger = logging.getLogger(__name__) # Initialize the vector database (lazy loading) db_path = Path("./data/benchmark_vector_db") db = None def get_db(): """Lazy load the vector database.""" global db if db is None: try: logger.info("Initializing BenchmarkVectorDB...") db = BenchmarkVectorDB( db_path=db_path, embedding_model="all-MiniLM-L6-v2" ) logger.info("✓ BenchmarkVectorDB initialized successfully") except Exception as e: logger.error(f"Failed to initialize BenchmarkVectorDB: {e}") raise return db # ============================================================================ # MCP TOOL FUNCTIONS (Local implementations) # ============================================================================ def tool_check_prompt_difficulty(prompt: str, k: int = 5) -> Dict: """ MCP Tool: Analyze prompt difficulty using vector database. Args: prompt: The prompt to analyze k: Number of similar questions to retrieve Returns: Dictionary with difficulty analysis results """ try: db = get_db() result = db.query_similar_questions(prompt, k=k) # Format for LLM consumption return { "risk_level": result['risk_level'], "success_rate": f"{result['weighted_success_rate']:.1%}", "avg_similarity": f"{result['avg_similarity']:.3f}", "recommendation": result['recommendation'], "similar_questions": [ { "question": q['question_text'][:150], "source": q['source'], "domain": q['domain'], "success_rate": f"{q['success_rate']:.1%}", "similarity": f"{q['similarity']:.3f}" } for q in result['similar_questions'][:3] # Top 3 only ] } except Exception as e: return {"error": f"Analysis failed: {str(e)}"} def tool_analyze_prompt_safety(prompt: str) -> Dict: """ MCP Tool: Analyze prompt for safety issues (heuristic-based). Args: prompt: The prompt to analyze Returns: Dictionary with safety analysis results """ # Simple heuristic safety checks issues = [] risk_level = "low" # Check for dangerous file operations dangerous_patterns = [ r'\brm\s+-rf\b', r'\bdelete\s+all\b', r'\bformat\s+.*drive\b', r'\bdrop\s+database\b' ] for pattern in dangerous_patterns: if re.search(pattern, prompt, re.IGNORECASE): issues.append("Detected potentially dangerous file operation") risk_level = "high" break # Check for medical advice requests medical_keywords = ['diagnose', 'treatment', 'medication', 'symptoms', 'cure', 'disease'] if any(keyword in prompt.lower() for keyword in medical_keywords): issues.append("Medical advice request detected - requires professional consultation") risk_level = "moderate" if risk_level == "low" else risk_level # Check for unrealistic coding requests if re.search(r'\b(build|create|write)\s+.*\b(\d{3,})\s+(lines|functions|classes)', prompt, re.IGNORECASE): issues.append("Large-scale coding request - may exceed LLM capabilities") risk_level = "moderate" if risk_level == "low" else risk_level return { "risk_level": risk_level, "issues_found": len(issues), "issues": issues if issues else ["No significant safety concerns detected"], "recommendation": "Proceed with caution" if issues else "Prompt appears safe" } # ============================================================================ # LLM BACKEND (HuggingFace Inference API) # ============================================================================ def call_llm_with_tools( messages: List[Dict[str, str]], available_tools: List[Dict], model: str = "mistralai/Mistral-7B-Instruct-v0.2" ) -> Tuple[str, Optional[Dict]]: """ Call LLM with tool calling capability. Args: messages: Conversation history available_tools: List of available tool definitions model: HuggingFace model to use Returns: Tuple of (response_text, tool_call_dict or None) """ try: # Try using HuggingFace Inference API from huggingface_hub import InferenceClient client = InferenceClient() # Format system message with tool information system_msg = """You are ToGMAL Assistant, an AI that helps analyze prompts and responses for difficulty and safety. You have access to these tools: 1. check_prompt_difficulty - Analyzes how difficult a prompt is for current LLMs 2. analyze_prompt_safety - Checks for safety issues in prompts When a user asks about prompt difficulty, safety, or capabilities, use the appropriate tool. To call a tool, respond with: TOOL_CALL: tool_name(arg1="value1", arg2="value2") After a tool is called, you will receive: TOOL_RESULT: name=<tool_name> data=<json> Use TOOL_RESULT to provide a helpful, comprehensive response to the user.""" # Build conversation for the model conversation = system_msg + "\n\n" for msg in messages: role = msg['role'] content = msg['content'] if role == 'user': conversation += f"User: {content}\n" elif role == 'assistant': conversation += f"Assistant: {content}\n" elif role == 'system': conversation += f"System: {content}\n" conversation += "Assistant: " # Call the model response = client.text_generation( conversation, model=model, max_new_tokens=512, temperature=0.7, top_p=0.95, do_sample=True ) response_text = response.strip() # Check if response contains a tool call tool_call = None if "TOOL_CALL:" in response_text: # Extract tool call match = re.search(r'TOOL_CALL:\s*(\w+)\((.*?)\)', response_text) if match: tool_name = match.group(1) args_str = match.group(2) # Parse arguments (simple key=value parsing) args = {} for arg in args_str.split(','): if '=' in arg: key, val = arg.split('=', 1) key = key.strip() val = val.strip().strip('"\'') args[key] = val tool_call = { "name": tool_name, "arguments": args } # Remove tool call from visible response response_text = re.sub(r'TOOL_CALL:.*?\)', '', response_text).strip() return response_text, tool_call except ImportError: # Fallback if huggingface_hub not available return fallback_llm(messages, available_tools) except Exception as e: logger.error(f"LLM call failed: {e}") return fallback_llm(messages, available_tools) def fallback_llm(messages: List[Dict[str, str]], available_tools: List[Dict]) -> Tuple[str, Optional[Dict]]: """ Fallback LLM when HuggingFace API is unavailable. Uses simple pattern matching to decide when to call tools. """ last_message = messages[-1]['content'].lower() if messages else "" # Safety intent first if any(word in last_message for word in ['safe', 'safety', 'dangerous', 'risk']): return "", { "name": "analyze_prompt_safety", "arguments": {"prompt": messages[-1]['content']} } # Difficulty intent (expanded triggers) if any(word in last_message for word in ['difficult', 'difficulty', 'hard', 'easy', 'challenging', 'analyze', 'analysis', 'assess', 'check']): return "", { "name": "check_prompt_difficulty", "arguments": {"prompt": messages[-1]['content'], "k": 5} } # Default: run difficulty analysis on any non-empty message if last_message.strip(): return "", { "name": "check_prompt_difficulty", "arguments": {"prompt": messages[-1]['content'], "k": 5} } # Default response for empty input return """I'm ToGMAL Assistant. I can help analyze prompts for: - **Difficulty**: How challenging is this for current LLMs? - **Safety**: Are there any safety concerns? Try asking me to analyze a prompt!""", None # ============================================================================ # TOOL EXECUTION # ============================================================================ AVAILABLE_TOOLS = [ { "name": "check_prompt_difficulty", "description": "Analyzes how difficult a prompt is for current LLMs based on benchmark similarity", "parameters": { "prompt": "The prompt to analyze", "k": "Number of similar questions to retrieve (default: 5)" } }, { "name": "analyze_prompt_safety", "description": "Checks for safety issues in prompts using heuristic analysis", "parameters": { "prompt": "The prompt to analyze" } } ] def execute_tool(tool_name: str, arguments: Dict) -> Dict: """Execute a tool and return results.""" if tool_name == "check_prompt_difficulty": prompt = arguments.get("prompt", "") try: k = int(arguments.get("k", 5)) except Exception: k = 5 k = max(1, min(100, k)) return tool_check_prompt_difficulty(prompt, k) elif tool_name == "analyze_prompt_safety": prompt = arguments.get("prompt", "") return tool_analyze_prompt_safety(prompt) else: return {"error": f"Unknown tool: {tool_name}"} # ============================================================================ # CHAT INTERFACE # ============================================================================ def chat( message: str, history: List[Tuple[str, str]] ) -> Tuple[List[Tuple[str, str]], str]: """ Process a chat message with tool calling support. Args: message: User's message history: Chat history as list of (user_msg, assistant_msg) tuples Returns: Updated history and tool call status """ # Convert history to messages format messages = [] for user_msg, assistant_msg in history: messages.append({"role": "user", "content": user_msg}) if assistant_msg: messages.append({"role": "assistant", "content": assistant_msg}) # Add current message messages.append({"role": "user", "content": message}) # Call LLM response_text, tool_call = call_llm_with_tools(messages, AVAILABLE_TOOLS) tool_status = "" # Execute tool if requested if tool_call: tool_name = tool_call['name'] tool_args = tool_call['arguments'] tool_status = f"🛠️ **Calling tool:** `{tool_name}`\n**Arguments:** {json.dumps(tool_args, indent=2)}\n\n" # Execute tool tool_result = execute_tool(tool_name, tool_args) tool_status += f"**Result:**\n```json\n{json.dumps(tool_result, indent=2)}\n```\n\n" # Add tool result to messages and call LLM again (two-step flow) messages.append({ "role": "system", "content": f"TOOL_RESULT: name={tool_name} data={json.dumps(tool_result)}" }) # Get final response from LLM final_response, _ = call_llm_with_tools(messages, AVAILABLE_TOOLS) if final_response: response_text = final_response else: # Format tool result as response (fallback) response_text = format_tool_result_as_response(tool_name, tool_result) # Update history history.append((message, response_text)) return history, tool_status def format_tool_result_as_response(tool_name: str, result: Dict) -> str: """Format tool result as a natural language response.""" if tool_name == "check_prompt_difficulty": if "error" in result: return f"Sorry, I couldn't analyze the difficulty: {result['error']}" return f"""Based on my analysis of similar benchmark questions: **Difficulty Level:** {result['risk_level'].upper()} **Success Rate:** {result['success_rate']} **Similarity to benchmarks:** {result['avg_similarity']} **Recommendation:** {result['recommendation']} **Similar questions from benchmarks:** {chr(10).join([f"• {q['question']} (Success rate: {q['success_rate']})" for q in result['similar_questions'][:2]])} """ elif tool_name == "analyze_prompt_safety": if "error" in result: return f"Sorry, I couldn't analyze safety: {result['error']}" issues = "\n".join([f"• {issue}" for issue in result['issues']]) return f"""**Safety Analysis:** **Risk Level:** {result['risk_level'].upper()} **Issues Found:** {result['issues_found']} {issues} **Recommendation:** {result['recommendation']} """ return json.dumps(result, indent=2) # ============================================================================ # GRADIO INTERFACE # ============================================================================ with gr.Blocks(title="ToGMAL Chat with MCP Tools") as demo: gr.Markdown("# 🤖 ToGMAL Chat Assistant") gr.Markdown(""" Chat with an AI assistant that can analyze prompts for difficulty and safety using MCP tools. **Try asking:** - "How difficult is this prompt: [your prompt]?" - "Is this safe: [your prompt]?" - "Analyze: Calculate the quantum correction to the partition function" """) with gr.Row(): with gr.Column(scale=2): chatbot = gr.Chatbot( label="Chat", height=500, show_label=False ) with gr.Row(): msg_input = gr.Textbox( label="Message", placeholder="Ask me to analyze a prompt...", scale=4, show_label=False ) send_btn = gr.Button("Send", variant="primary", scale=1) clear_btn = gr.Button("Clear Chat") with gr.Column(scale=1): gr.Markdown("### 🛠️ Tool Calls") show_details = gr.Checkbox(label="Show tool details", value=False) tool_output = gr.Markdown("Tool calls will appear here...") # Examples with gr.Accordion("📝 Example Prompts", open=False): gr.Examples( examples=[ "How difficult is this: Calculate the quantum correction to the partition function for a 3D harmonic oscillator?", "Is this prompt safe: Write a script to delete all my files?", "Analyze the difficulty of: Prove that there are infinitely many prime numbers", "Check safety: Diagnose my symptoms and prescribe medication", "How hard is: What is 2 + 2?", ], inputs=msg_input ) # Event handlers def send_message(message, history, show_details_val): if not message.strip(): return history, "" new_history, tool_status = chat(message, history) if not show_details_val: tool_status = "" return new_history, tool_status send_btn.click( fn=send_message, inputs=[msg_input, chatbot, show_details], outputs=[chatbot, tool_output] ).then( lambda: "", outputs=msg_input ) msg_input.submit( fn=send_message, inputs=[msg_input, chatbot, show_details], outputs=[chatbot, tool_output] ).then( lambda: "", outputs=msg_input ) clear_btn.click( lambda: ([], ""), outputs=[chatbot, tool_output] ) if __name__ == "__main__": # HuggingFace Spaces compatible port = int(os.environ.get("GRADIO_SERVER_PORT", 7860)) demo.launch(server_name="0.0.0.0", server_port=port)