hanzo-mcp 0.6.10__py3-none-any.whl → 0.6.13__py3-none-any.whl

hanzo_mcp/tools/llm/llm_tool.py

@@ -1,9 +1,10 @@
-"""Universal LLM tool using LiteLLM."""
+"""Unified LLM tool with multiple actions including consensus mode."""
 
+from typing import Annotated, TypedDict, Unpack, final, override, Optional, List, Dict, Any
+import asyncio
 import os
 import json
-from typing import Annotated, Optional, TypedDict, Unpack, final, override, List, Dict, Any
-import asyncio
+from pathlib import Path
 
 from mcp.server.fastmcp import Context as MCPContext
 from pydantic import Field
@@ -11,27 +12,44 @@ from pydantic import Field
 from hanzo_mcp.tools.common.base import BaseTool
 from hanzo_mcp.tools.common.context import create_tool_context
 
+# Check if litellm is available
 try:
     import litellm
-    from litellm import completion, acompletion
     LITELLM_AVAILABLE = True
 except ImportError:
     LITELLM_AVAILABLE = False
 
 
-Model = Annotated[
+# Parameter types
+Action = Annotated[
     str,
     Field(
-        description="Model name in LiteLLM format (e.g., 'gpt-4', 'claude-3-opus-20240229', 'gemini/gemini-pro')",
-        min_length=1,
+        description="Action to perform: query, consensus, list, models, enable, disable, test",
+        default="query",
+    ),
+]
+
+Model = Annotated[
+    Optional[str],
+    Field(
+        description="Model name (e.g., gpt-4, claude-3-opus-20240229)",
+        default=None,
+    ),
+]
+
+Models = Annotated[
+    Optional[List[str]],
+    Field(
+        description="List of models for consensus mode",
+        default=None,
     ),
 ]
 
 Prompt = Annotated[
-    str,
+    Optional[str],
     Field(
-        description="The prompt or question to send to the model",
-        min_length=1,
+        description="The prompt to send to the LLM",
+        default=None,
     ),
 ]
 
@@ -46,7 +64,7 @@ SystemPrompt = Annotated[
 Temperature = Annotated[
     float,
     Field(
-        description="Temperature for response randomness (0.0-2.0)",
+        description="Temperature for response randomness (0-2)",
         default=0.7,
     ),
 ]
@@ -75,82 +93,140 @@ Stream = Annotated[
     ),
 ]
 
+Provider = Annotated[
+    Optional[str],
+    Field(
+        description="Provider name for list/enable/disable actions",
+        default=None,
+    ),
+]
 
-class LLMToolParams(TypedDict, total=False):
-    """Parameters for LLM tool."""
+IncludeRaw = Annotated[
+    bool,
+    Field(
+        description="Include raw responses in consensus mode",
+        default=False,
+    ),
+]
 
-    model: str
-    prompt: str
+JudgeModel = Annotated[
+    Optional[str],
+    Field(
+        description="Model to use as judge/aggregator in consensus",
+        default=None,
+    ),
+]
+
+DevilsAdvocate = Annotated[
+    bool,
+    Field(
+        description="Enable devil's advocate mode (10th model critiques others)",
+        default=False,
+    ),
+]
+
+ConsensusSize = Annotated[
+    Optional[int],
+    Field(
+        description="Number of models to use in consensus (default: 3)",
+        default=None,
+    ),
+]
+
+
+class LLMParams(TypedDict, total=False):
+    """Parameters for LLM tool."""
+    action: str
+    model: Optional[str]
+    models: Optional[List[str]]
+    prompt: Optional[str]
     system_prompt: Optional[str]
     temperature: float
     max_tokens: Optional[int]
     json_mode: bool
     stream: bool
+    provider: Optional[str]
+    include_raw: bool
+    judge_model: Optional[str]
+    devils_advocate: bool
+    consensus_size: Optional[int]
 
 
 @final
 class LLMTool(BaseTool):
-    """Universal LLM tool using LiteLLM."""
+    """Unified LLM tool with multiple actions."""
+    
+    # Config file for settings
+    CONFIG_FILE = Path.home() / ".hanzo" / "mcp" / "llm_config.json"
     
-    # Common environment variables for API keys
+    # Default consensus models in order of preference
+    DEFAULT_CONSENSUS_MODELS = [
+        "gpt-4o",  # OpenAI's latest
+        "claude-3-opus-20240229",  # Claude's most capable
+        "gemini/gemini-1.5-pro",  # Google's best
+        "groq/llama3-70b-8192",  # Fast Groq
+        "mistral/mistral-large-latest",  # Mistral's best
+        "perplexity/llama-3.1-sonar-large-128k-chat",  # Perplexity with search
+    ]
+    
+    # API key environment variables
     API_KEY_ENV_VARS = {
         "openai": ["OPENAI_API_KEY"],
         "anthropic": ["ANTHROPIC_API_KEY", "CLAUDE_API_KEY"],
-        "google": ["GOOGLE_API_KEY", "GEMINI_API_KEY", "GOOGLE_APPLICATION_CREDENTIALS"],
+        "google": ["GOOGLE_API_KEY", "GEMINI_API_KEY"],
         "groq": ["GROQ_API_KEY"],
+        "mistral": ["MISTRAL_API_KEY"],
+        "perplexity": ["PERPLEXITY_API_KEY", "PERPLEXITYAI_API_KEY"],
+        "together": ["TOGETHER_API_KEY", "TOGETHERAI_API_KEY"],
         "cohere": ["COHERE_API_KEY"],
         "replicate": ["REPLICATE_API_KEY"],
         "huggingface": ["HUGGINGFACE_API_KEY", "HF_TOKEN"],
-        "together": ["TOGETHER_API_KEY", "TOGETHERAI_API_KEY"],
-        "mistral": ["MISTRAL_API_KEY"],
-        "perplexity": ["PERPLEXITY_API_KEY"],
-        "anyscale": ["ANYSCALE_API_KEY"],
-        "deepinfra": ["DEEPINFRA_API_KEY"],
-        "ai21": ["AI21_API_KEY"],
-        "nvidia": ["NVIDIA_API_KEY"],
+        "bedrock": ["AWS_ACCESS_KEY_ID", "AWS_SECRET_ACCESS_KEY"],
+        "vertex": ["GOOGLE_APPLICATION_CREDENTIALS"],
+        "azure": ["AZURE_API_KEY"],
         "voyage": ["VOYAGE_API_KEY"],
-        "aws": ["AWS_ACCESS_KEY_ID", "AWS_SECRET_ACCESS_KEY"],  # For Bedrock
-        "azure": ["AZURE_API_KEY", "AZURE_OPENAI_API_KEY"],
+        "deepseek": ["DEEPSEEK_API_KEY"],
     }
     
-    # Model prefixes for each provider
-    PROVIDER_MODELS = {
-        "openai": ["gpt-4", "gpt-3.5", "o1", "davinci", "curie", "babbage", "ada"],
-        "anthropic": ["claude-3", "claude-2", "claude-instant"],
-        "google": ["gemini", "palm", "bison", "gecko"],
-        "groq": ["mixtral", "llama2", "llama3"],
-        "cohere": ["command", "command-light"],
-        "mistral": ["mistral-tiny", "mistral-small", "mistral-medium", "mistral-large"],
-        "perplexity": ["pplx", "sonar"],
-        "together": ["together"],
-        "bedrock": ["bedrock/"],
-        "azure": ["azure/"],
-    }
-
     def __init__(self):
-        """Initialize the LLM tool."""
+        """Initialize the unified LLM tool."""
         self.available_providers = self._detect_available_providers()
+        self.config = self._load_config()
         
-        # Configure LiteLLM settings
-        if LITELLM_AVAILABLE:
-            # Enable verbose logging for debugging
-            litellm.set_verbose = False
-            # Set default timeout
-            litellm.request_timeout = 120
-
     def _detect_available_providers(self) -> Dict[str, List[str]]:
-        """Detect which LLM providers have API keys configured."""
+        """Detect which providers have API keys configured."""
         available = {}
         
         for provider, env_vars in self.API_KEY_ENV_VARS.items():
             for var in env_vars:
                 if os.getenv(var):
-                    if provider not in available:
-                        available[provider] = []
-                    available[provider].append(var)
+                    available[provider] = env_vars
                     break
         
         return available
+    
+    def _load_config(self) -> Dict[str, Any]:
+        """Load configuration from file."""
+        if self.CONFIG_FILE.exists():
+            try:
+                with open(self.CONFIG_FILE, 'r') as f:
+                    return json.load(f)
+            except:
+                pass
+        
+        # Default config
+        return {
+            "disabled_providers": [],
+            "consensus_models": None,  # Use defaults if None
+            "default_judge_model": "gpt-4o",
+            "consensus_size": 3,
+        }
+    
+    def _save_config(self):
+        """Save configuration to file."""
+        self.CONFIG_FILE.parent.mkdir(parents=True, exist_ok=True)
+        with open(self.CONFIG_FILE, 'w') as f:
+            json.dump(self.config, f, indent=2)
 
     @property
     @override
@@ -162,149 +238,578 @@ class LLMTool(BaseTool):
     @override
     def description(self) -> str:
         """Get the tool description."""
-        providers_list = ", ".join(sorted(self.available_providers.keys())) if self.available_providers else "None"
+        available = list(self.available_providers.keys())
         
-        return f"""Query any LLM using LiteLLM's unified interface.
-
-Supports 100+ models from various providers through a single interface.
-Automatically uses API keys from environment variables.
+        return f"""Query LLMs. Default: single query. Actions: consensus, list, models, test.
 
-Detected providers: {providers_list}
+Usage:
+llm "What is the capital of France?"
+llm "Explain this code" --model gpt-4o
+llm --action consensus "Is this approach correct?" --devils-advocate
+llm --action models --provider openai
 
-Common models:
-- OpenAI: gpt-4o, gpt-4, gpt-3.5-turbo, o1-preview, o1-mini
-- Anthropic: claude-3-opus-20240229, claude-3-sonnet-20240229, claude-3-haiku-20240307
-- Google: gemini/gemini-pro, gemini/gemini-1.5-pro, gemini/gemini-1.5-flash
-- Groq: groq/mixtral-8x7b-32768, groq/llama3-70b-8192
-- Mistral: mistral/mistral-large-latest, mistral/mistral-medium
-- Perplexity: perplexity/sonar-medium-online
-- Together: together/mixtral-8x22b
-
-Examples:
-- llm --model "gpt-4" --prompt "Explain quantum computing"
-- llm --model "claude-3-opus-20240229" --prompt "Write a haiku about coding"
-- llm --model "gemini/gemini-pro" --prompt "What is the meaning of life?" --temperature 0.9
-- llm --model "groq/mixtral-8x7b-32768" --prompt "Generate a JSON schema" --json-mode
-
-For provider-specific tools, use: openai, anthropic, gemini, groq, etc.
-For consensus across models, use: consensus
-"""
+Available: {', '.join(available) if available else 'None'}"""
 
     @override
     async def call(
         self,
         ctx: MCPContext,
-        **params: Unpack[LLMToolParams],
+        **params: Unpack[LLMParams],
     ) -> str:
-        """Query an LLM.
-
-        Args:
-            ctx: MCP context
-            **params: Tool parameters
-
-        Returns:
-            LLM response
-        """
-        tool_ctx = create_tool_context(ctx)
-        await tool_ctx.set_tool_info(self.name)
+        """Execute LLM action."""
+        # Create tool context only if we have a proper MCP context
+        tool_ctx = None
+        try:
+            if hasattr(ctx, 'client') and ctx.client and hasattr(ctx.client, 'server'):
+                tool_ctx = create_tool_context(ctx)
+                if tool_ctx:
+                    await tool_ctx.set_tool_info(self.name)
+        except:
+            # Running in test mode without MCP context
+            pass
 
         if not LITELLM_AVAILABLE:
             return "Error: LiteLLM is not installed. Install it with: pip install litellm"
 
-        # Extract parameters
+        # Extract action
+        action = params.get("action", "query")
+        
+        # Route to appropriate handler
+        if action == "query":
+            return await self._handle_query(tool_ctx, params)
+        elif action == "consensus":
+            return await self._handle_consensus(tool_ctx, params)
+        elif action == "list":
+            return self._handle_list()
+        elif action == "models":
+            return self._handle_models(params.get("provider"))
+        elif action == "enable":
+            return self._handle_enable(params.get("provider"))
+        elif action == "disable":
+            return self._handle_disable(params.get("provider"))
+        elif action == "test":
+            return await self._handle_test(tool_ctx, params.get("model"), params.get("provider"))
+        else:
+            return f"Error: Unknown action '{action}'. Valid actions: query, consensus, list, models, enable, disable, test"
+
+    async def _handle_query(self, tool_ctx, params: Dict[str, Any]) -> str:
+        """Handle single model query."""
         model = params.get("model")
-        if not model:
-            return "Error: model is required"
-
         prompt = params.get("prompt")
+        
         if not prompt:
-            return "Error: prompt is required"
-
-        system_prompt = params.get("system_prompt")
-        temperature = params.get("temperature", 0.7)
-        max_tokens = params.get("max_tokens")
-        json_mode = params.get("json_mode", False)
-        stream = params.get("stream", False)
-
+            return "Error: prompt is required for query action"
+        
+        # Auto-select model if not specified
+        if not model:
+            if self.available_providers:
+                # Use first available model
+                if "openai" in self.available_providers:
+                    model = "gpt-4o-mini"
+                elif "anthropic" in self.available_providers:
+                    model = "claude-3-haiku-20240307"
+                elif "google" in self.available_providers:
+                    model = "gemini/gemini-1.5-flash"
+                else:
+                    # Use first provider's default
+                    provider = list(self.available_providers.keys())[0]
+                    model = f"{provider}/default"
+            else:
+                return "Error: No model specified and no API keys found"
+        
         # Check if we have API key for this model
         provider = self._get_provider_for_model(model)
         if provider and provider not in self.available_providers:
             env_vars = self.API_KEY_ENV_VARS.get(provider, [])
             return f"Error: No API key found for {provider}. Set one of: {', '.join(env_vars)}"
-
+        
         # Build messages
         messages = []
-        if system_prompt:
-            messages.append({"role": "system", "content": system_prompt})
+        if params.get("system_prompt"):
+            messages.append({"role": "system", "content": params["system_prompt"]})
         messages.append({"role": "user", "content": prompt})
-
+        
         # Build kwargs
         kwargs = {
             "model": model,
             "messages": messages,
-            "temperature": temperature,
+            "temperature": params.get("temperature", 0.7),
         }
         
-        if max_tokens:
-            kwargs["max_tokens"] = max_tokens
+        if params.get("max_tokens"):
+            kwargs["max_tokens"] = params["max_tokens"]
         
-        if json_mode:
+        if params.get("json_mode"):
             kwargs["response_format"] = {"type": "json_object"}
-
-        await tool_ctx.info(f"Querying {model}...")
-
+        
+        if params.get("stream"):
+            kwargs["stream"] = True
+        
         try:
-            if stream:
-                # Streaming response
+            if tool_ctx:
+                await tool_ctx.info(f"Querying {model}...")
+            
+            if kwargs.get("stream"):
+                # Handle streaming response
                 response_text = ""
-                async for chunk in await acompletion(**kwargs, stream=True):
+                async for chunk in await litellm.acompletion(**kwargs):
                     if chunk.choices[0].delta.content:
                         response_text += chunk.choices[0].delta.content
-                        # Could emit progress here if needed
-                
                 return response_text
             else:
-                # Non-streaming response
-                response = await acompletion(**kwargs)
+                # Regular response
+                response = await litellm.acompletion(**kwargs)
                 return response.choices[0].message.content
-
+                
         except Exception as e:
             error_msg = str(e)
-            
-            # Provide helpful error messages
-            if "api_key" in error_msg.lower():
-                provider = self._get_provider_for_model(model)
-                env_vars = self.API_KEY_ENV_VARS.get(provider, [])
-                return f"Error: API key issue for {provider}. Make sure one of these is set: {', '.join(env_vars)}\n\nOriginal error: {error_msg}"
-            elif "model" in error_msg.lower() and "not found" in error_msg.lower():
-                return f"Error: Model '{model}' not found or not accessible. Check the model name and your API permissions.\n\nOriginal error: {error_msg}"
+            if "model_not_found" in error_msg or "does not exist" in error_msg:
+                return f"Error: Model '{model}' not found. Use 'llm --action models' to see available models."
             else:
                 return f"Error calling LLM: {error_msg}"
 
+    async def _handle_consensus(self, tool_ctx, params: Dict[str, Any]) -> str:
+        """Handle consensus mode with multiple models."""
+        prompt = params.get("prompt")
+        if not prompt:
+            return "Error: prompt is required for consensus action"
+        
+        # Determine models to use
+        models = params.get("models")
+        if not models:
+            # Use configured or default models
+            consensus_size = params.get("consensus_size") or self.config.get("consensus_size", 3)
+            models = self._get_consensus_models(consensus_size)
+        
+        if not models:
+            return "Error: No models available for consensus. Set API keys for at least 2 providers."
+        
+        if len(models) < 2:
+            return "Error: Consensus requires at least 2 models"
+        
+        # Check for devil's advocate mode
+        devils_advocate = params.get("devils_advocate", False)
+        if devils_advocate and len(models) < 3:
+            return "Error: Devil's advocate mode requires at least 3 models"
+        
+        if tool_ctx:
+            await tool_ctx.info(f"Running consensus with {len(models)} models...")
+        
+        # Query models in parallel
+        system_prompt = params.get("system_prompt")
+        temperature = params.get("temperature", 0.7)
+        max_tokens = params.get("max_tokens")
+        
+        # Split models if using devil's advocate
+        if devils_advocate:
+            consensus_models = models[:-1]
+            devil_model = models[-1]
+        else:
+            consensus_models = models
+            devil_model = None
+        
+        # Query consensus models
+        responses = await self._query_models_parallel(
+            consensus_models, prompt, system_prompt, temperature, max_tokens, tool_ctx
+        )
+        
+        # Get devil's advocate response if enabled
+        devil_response = None
+        if devil_model:
+            # Create devil's advocate prompt
+            responses_text = "\n\n".join([
+                f"Model {i+1}: {resp['response']}" 
+                for i, resp in enumerate(responses) if resp['response']
+            ])
+            
+            devil_prompt = f"""You are a critical analyst. Review these responses to the question below and provide a devil's advocate perspective. Challenge assumptions, point out weaknesses, and suggest alternative viewpoints.
+
+Original Question: {prompt}
+
+Responses from other models:
+{responses_text}
+
+Provide your critical analysis:"""
+            
+            devil_result = await self._query_single_model(
+                devil_model, devil_prompt, system_prompt, temperature, max_tokens
+            )
+            
+            if devil_result['success']:
+                devil_response = {
+                    'model': devil_model,
+                    'response': devil_result['response'],
+                    'time_ms': devil_result['time_ms']
+                }
+        
+        # Aggregate responses
+        judge_model = params.get("judge_model") or self.config.get("default_judge_model", "gpt-4o")
+        include_raw = params.get("include_raw", False)
+        
+        return await self._aggregate_consensus(
+            responses, prompt, judge_model, include_raw, devil_response, tool_ctx
+        )
+
+    def _handle_list(self) -> str:
+        """List available providers."""
+        output = ["=== LLM Providers ==="]
+        
+        # Get all possible providers
+        all_providers = sorted(self.API_KEY_ENV_VARS.keys())
+        disabled = self.config.get("disabled_providers", [])
+        
+        output.append(f"Total providers: {len(all_providers)}")
+        output.append(f"Available: {len(self.available_providers)}")
+        output.append(f"Disabled: {len(disabled)}\n")
+        
+        for provider in all_providers:
+            status_parts = []
+            
+            # Check if API key exists
+            if provider in self.available_providers:
+                status_parts.append("✅ API key found")
+            else:
+                status_parts.append("❌ No API key")
+            
+            # Check if disabled
+            if provider in disabled:
+                status_parts.append("🚫 Disabled")
+            
+            # Show environment variables
+            env_vars = self.API_KEY_ENV_VARS.get(provider, [])
+            status = " | ".join(status_parts)
+            
+            output.append(f"{provider}: {status}")
+            output.append(f"  Environment variables: {', '.join(env_vars)}")
+        
+        output.append("\nUse 'llm --action enable/disable --provider <name>' to manage providers")
+        
+        return "\n".join(output)
+
+    def _handle_models(self, provider: Optional[str] = None) -> str:
+        """List available models."""
+        try:
+            all_models = self._get_all_models()
+            
+            if not all_models:
+                return "No models available or LiteLLM not properly initialized"
+            
+            output = ["=== Available LLM Models ==="]
+            
+            if provider:
+                # Show models for specific provider
+                provider_lower = provider.lower()
+                models = all_models.get(provider_lower, [])
+                
+                if not models:
+                    return f"No models found for provider '{provider}'"
+                
+                output.append(f"\n{provider.upper()} ({len(models)} models):")
+                output.append("-" * 40)
+                
+                # Show first 50 models
+                for model in models[:50]:
+                    output.append(f"  {model}")
+                
+                if len(models) > 50:
+                    output.append(f"  ... and {len(models) - 50} more")
+            else:
+                # Show summary of all providers
+                total_models = sum(len(models) for models in all_models.values())
+                output.append(f"Total models available: {total_models}")
+                output.append("")
+                
+                # Show providers with counts
+                for provider_name, models in sorted(all_models.items()):
+                    if models:
+                        available = "✅" if provider_name in self.available_providers else "❌"
+                        output.append(f"{available} {provider_name}: {len(models)} models")
+                
+                output.append("\nUse 'llm --action models --provider <name>' to see specific models")
+            
+            return "\n".join(output)
+            
+        except Exception as e:
+            return f"Error listing models: {str(e)}"
+
+    def _handle_enable(self, provider: Optional[str]) -> str:
+        """Enable a provider."""
+        if not provider:
+            return "Error: provider is required for enable action"
+        
+        provider = provider.lower()
+        disabled = self.config.get("disabled_providers", [])
+        
+        if provider in disabled:
+            disabled.remove(provider)
+            self.config["disabled_providers"] = disabled
+            self._save_config()
+            return f"Successfully enabled {provider}"
+        else:
+            return f"{provider} is already enabled"
+
+    def _handle_disable(self, provider: Optional[str]) -> str:
+        """Disable a provider."""
+        if not provider:
+            return "Error: provider is required for disable action"
+        
+        provider = provider.lower()
+        disabled = self.config.get("disabled_providers", [])
+        
+        if provider not in disabled:
+            disabled.append(provider)
+            self.config["disabled_providers"] = disabled
+            self._save_config()
+            return f"Successfully disabled {provider}"
+        else:
+            return f"{provider} is already disabled"
+
+    async def _handle_test(self, tool_ctx, model: Optional[str], provider: Optional[str]) -> str:
+        """Test a model or provider."""
+        if not model and not provider:
+            return "Error: Either model or provider is required for test action"
+        
+        # If provider specified, test its default model
+        if provider and not model:
+            provider = provider.lower()
+            if provider == "openai":
+                model = "gpt-3.5-turbo"
+            elif provider == "anthropic":
+                model = "claude-3-haiku-20240307"
+            elif provider == "google":
+                model = "gemini/gemini-1.5-flash"
+            elif provider == "groq":
+                model = "groq/llama3-8b-8192"
+            else:
+                model = f"{provider}/default"
+        
+        # Test the model
+        test_prompt = "Say 'Hello from Hanzo AI!' in exactly 5 words."
+        
+        try:
+            if tool_ctx:
+                await tool_ctx.info(f"Testing {model}...")
+            
+            response = await litellm.acompletion(
+                model=model,
+                messages=[{"role": "user", "content": test_prompt}],
+                temperature=0,
+                max_tokens=20
+            )
+            
+            result = response.choices[0].message.content
+            return f"✅ {model} is working!\nResponse: {result}"
+            
+        except Exception as e:
+            return f"❌ {model} failed: {str(e)}"
+
+    def _get_consensus_models(self, size: int) -> List[str]:
+        """Get models for consensus based on availability."""
+        # Use configured models if set
+        configured = self.config.get("consensus_models")
+        if configured:
+            return configured[:size]
+        
+        # Otherwise, build list from available providers
+        models = []
+        disabled = self.config.get("disabled_providers", [])
+        
+        # Try default models first
+        for model in self.DEFAULT_CONSENSUS_MODELS:
+            if len(models) >= size:
+                break
+                
+            provider = self._get_provider_for_model(model)
+            if provider and provider in self.available_providers and provider not in disabled:
+                models.append(model)
+        
+        # If still need more, add from available providers
+        if len(models) < size:
+            for provider in self.available_providers:
+                if provider in disabled:
+                    continue
+                    
+                if provider == "openai" and "gpt-4o" not in models:
+                    models.append("gpt-4o")
+                elif provider == "anthropic" and "claude-3-opus-20240229" not in models:
+                    models.append("claude-3-opus-20240229")
+                elif provider == "google" and "gemini/gemini-1.5-pro" not in models:
+                    models.append("gemini/gemini-1.5-pro")
+                
+                if len(models) >= size:
+                    break
+        
+        return models
+
+    async def _query_models_parallel(
+        self, models: List[str], prompt: str, system_prompt: Optional[str],
+        temperature: float, max_tokens: Optional[int], tool_ctx
+    ) -> List[Dict[str, Any]]:
+        """Query multiple models in parallel."""
+        async def query_with_info(model: str) -> Dict[str, Any]:
+            result = await self._query_single_model(model, prompt, system_prompt, temperature, max_tokens)
+            return {
+                'model': model,
+                'response': result.get('response'),
+                'success': result.get('success', False),
+                'error': result.get('error'),
+                'time_ms': result.get('time_ms', 0)
+            }
+        
+        # Run all queries in parallel
+        tasks = [query_with_info(model) for model in models]
+        results = await asyncio.gather(*tasks)
+        
+        # Report results
+        successful = sum(1 for r in results if r['success'])
+        if tool_ctx:
+            await tool_ctx.info(f"Completed {successful}/{len(models)} model queries")
+        
+        return results
+
+    async def _query_single_model(
+        self, model: str, prompt: str, system_prompt: Optional[str],
+        temperature: float, max_tokens: Optional[int]
+    ) -> Dict[str, Any]:
+        """Query a single model and return result with metadata."""
+        import time
+        start_time = time.time()
+        
+        try:
+            messages = []
+            if system_prompt:
+                messages.append({"role": "system", "content": system_prompt})
+            messages.append({"role": "user", "content": prompt})
+            
+            kwargs = {
+                "model": model,
+                "messages": messages,
+                "temperature": temperature,
+            }
+            if max_tokens:
+                kwargs["max_tokens"] = max_tokens
+            
+            response = await litellm.acompletion(**kwargs)
+            
+            return {
+                'success': True,
+                'response': response.choices[0].message.content,
+                'time_ms': int((time.time() - start_time) * 1000)
+            }
+            
+        except Exception as e:
+            return {
+                'success': False,
+                'error': str(e),
+                'time_ms': int((time.time() - start_time) * 1000)
+            }
+
+    async def _aggregate_consensus(
+        self, responses: List[Dict[str, Any]], original_prompt: str,
+        judge_model: str, include_raw: bool, devil_response: Optional[Dict[str, Any]],
+        tool_ctx
+    ) -> str:
+        """Aggregate consensus responses using a judge model."""
+        # Prepare response data
+        successful_responses = [r for r in responses if r['success']]
+        
+        if not successful_responses:
+            return "Error: All models failed to respond"
+        
+        # Format responses for aggregation
+        responses_text = "\n\n".join([
+            f"Model: {r['model']}\nResponse: {r['response']}"
+            for r in successful_responses
+        ])
+        
+        if devil_response:
+            responses_text += f"\n\nDevil's Advocate ({devil_response['model']}):\n{devil_response['response']}"
+        
+        # Create aggregation prompt
+        aggregation_prompt = f"""Analyze the following responses from multiple AI models to this question:
+
+<original_question>
+{original_prompt}
+</original_question>
+
+<model_responses>
+{responses_text}
+</model_responses>
+
+Please provide:
+1. A synthesis of the key points where models agree
+2. Notable differences or disagreements between responses
+3. A balanced conclusion incorporating the best insights
+{f"4. Evaluation of the devil's advocate critique" if devil_response else ""}
+
+Be concise and highlight the most important findings."""
+        
+        # Get aggregation
+        try:
+            if tool_ctx:
+                await tool_ctx.info(f"Aggregating responses with {judge_model}...")
+            
+            judge_result = await self._query_single_model(
+                judge_model, aggregation_prompt, None, 0.3, None
+            )
+            
+            if not judge_result['success']:
+                return f"Error: Judge model failed: {judge_result.get('error', 'Unknown error')}"
+            
+            # Format output
+            output = [f"=== Consensus Analysis ({len(successful_responses)} models) ===\n"]
+            output.append(judge_result['response'])
+            
+            # Add model list
+            output.append(f"\nModels consulted: {', '.join([r['model'] for r in successful_responses])}")
+            if devil_response:
+                output.append(f"Devil's Advocate: {devil_response['model']}")
+            
+            # Add timing info
+            avg_time = sum(r['time_ms'] for r in responses) / len(responses)
+            output.append(f"\nAverage response time: {avg_time:.0f}ms")
+            
+            # Include raw responses if requested
+            if include_raw:
+                output.append("\n\n=== Raw Responses ===")
+                for r in successful_responses:
+                    output.append(f"\n{r['model']}:")
+                    output.append("-" * 40)
+                    output.append(r['response'])
+                
+                if devil_response:
+                    output.append(f"\nDevil's Advocate ({devil_response['model']}):")
+                    output.append("-" * 40)
+                    output.append(devil_response['response'])
+            
+            return "\n".join(output)
+            
+        except Exception as e:
+            return f"Error during aggregation: {str(e)}"
+
     def _get_provider_for_model(self, model: str) -> Optional[str]:
         """Determine the provider for a given model."""
         model_lower = model.lower()
         
-        # Check explicit provider prefix (e.g., "groq/mixtral")
+        # Check explicit provider prefix
         if "/" in model:
-            provider = model.split("/")[0]
-            return provider
+            return model.split("/")[0]
         
         # Check model prefixes
-        for provider, prefixes in self.PROVIDER_MODELS.items():
-            for prefix in prefixes:
-                if model_lower.startswith(prefix.lower()):
-                    return provider
+        if model_lower.startswith("gpt"):
+            return "openai"
+        elif model_lower.startswith("claude"):
+            return "anthropic"
+        elif model_lower.startswith("gemini"):
+            return "google"
+        elif model_lower.startswith("command"):
+            return "cohere"
         
-        # Default to OpenAI for unknown models
+        # Default to OpenAI
         return "openai"
 
-    @classmethod
-    def get_all_models(cls) -> Dict[str, List[str]]:
-        """Get all available models from LiteLLM organized by provider."""
-        if not LITELLM_AVAILABLE:
-            return {}
-        
+    def _get_all_models(self) -> Dict[str, List[str]]:
+        """Get all available models from LiteLLM."""
         try:
             import litellm
             
@@ -343,4 +848,4 @@ For consensus across models, use: consensus
 
     def register(self, mcp_server) -> None:
         """Register this tool with the MCP server."""
-        pass
+        pass