mcp_gen/generator.py

"""
MCP Generation Engine - The Innovation That Wins

Dynamically generates custom MCP servers based on user needs.
This is the KILLER FEATURE that has never been done before.
"""

import os
import json
import asyncio
from typing import Dict, Any, List, Optional
from datetime import datetime
import hashlib
from pathlib import Path

from core.model_router import router, TaskType


class MCPGenerator:
    """
    Generates custom MCP servers on-the-fly using AI.

    INNOVATION: Instead of pre-built tools, this creates new tools as needed.
    - User needs web scraping? Generate scraper MCP
    - User needs data analysis? Generate analyzer MCP
    - User needs API integration? Generate connector MCP
    """

    def __init__(self, output_dir: str = "./generated_mcps"):
        self.output_dir = Path(output_dir)
        self.output_dir.mkdir(parents=True, exist_ok=True)
        self.generated_servers = {}

    async def generate_mcp_server(
        self,
        task_description: str,
        tool_name: Optional[str] = None,
        context: Optional[Dict[str, Any]] = None
    ) -> Dict[str, Any]:
        """
        Generate a complete MCP server from a task description.

        Args:
            task_description: What the tool should do (e.g., "scrape product prices from Amazon")
            tool_name: Optional custom name for the tool
            context: Additional context (APIs to use, data schemas, etc.)

        Returns:
            Dict with server code, deployment info, and usage instructions
        """
        print(f"[GEN] Generating MCP server for: {task_description}")

        # Step 1: Analyze task and plan MCP architecture
        planning_prompt = f"""You are an expert MCP (Model Context Protocol) server architect.

Task: {task_description}
Context: {json.dumps(context or {}, indent=2)}

Analyze this task and design an MCP server architecture:

1. What tools/functions does this MCP need? (1-5 functions)
2. What are the function signatures? (name, parameters, return types)
3. What external APIs or libraries are needed?
4. What are the edge cases and error handling needs?
5. What's the best way to structure this MCP for reusability?

Respond with a JSON object:
{{
    "server_name": "descriptive_name",
    "description": "what this MCP does",
    "tools": [
        {{
            "name": "tool_function_name",
            "description": "what it does",
            "parameters": {{"param1": "type", "param2": "type"}},
            "returns": "return_type",
            "implementation_notes": "how to implement"
        }}
    ],
    "dependencies": ["package1", "package2"],
    "complexity": "simple|medium|complex"
}}
"""

        plan_result = await router.generate(
            planning_prompt,
            task_type=TaskType.PLANNING,
            temperature=0.3
        )

        # Parse planning response
        try:
            plan_json = self._extract_json(plan_result["response"])
        except Exception as e:
            print(f"❌ Failed to parse planning response: {e}")
            # Fallback to simple single-tool server
            plan_json = {
                "server_name": tool_name or "custom_tool",
                "description": task_description,
                "tools": [{
                    "name": "execute",
                    "description": task_description,
                    "parameters": {"input": "str"},
                    "returns": "dict"
                }],
                "dependencies": [],
                "complexity": "simple"
            }

        print(f"[PLAN] {plan_json['server_name']} with {len(plan_json['tools'])} tools")

        # Step 2: Generate MCP server code
        code_prompt = f"""You are an expert Python developer specializing in MCP servers.

Generate a COMPLETE, PRODUCTION-READY Gradio MCP server based on this specification:

{json.dumps(plan_json, indent=2)}

Requirements:
1. Use Gradio for the MCP server interface
2. Implement ALL tools from the specification
3. Include proper error handling and logging
4. Add docstrings and type hints
5. Make it deployable to Hugging Face Spaces
6. Include a simple Gradio UI for testing the tools
7. Follow MCP protocol standards

Generate the COMPLETE app.py file with:
- All imports
- Tool implementations
- Gradio interface
- MCP endpoint setup
- Error handling
- Main execution block

IMPORTANT: Return ONLY the Python code, no explanations.
"""

        code_result = await router.generate(
            code_prompt,
            task_type=TaskType.CODE_GEN,
            max_tokens=4000,
            temperature=0.2
        )

        server_code = self._extract_code(code_result["response"])

        # Step 3: Generate requirements.txt
        requirements = self._generate_requirements(plan_json["dependencies"])

        # Step 4: Generate README.md
        readme = self._generate_readme(plan_json, task_description)

        # Step 5: Save generated files
        server_id = self._generate_server_id(plan_json["server_name"])
        server_dir = self.output_dir / server_id
        server_dir.mkdir(parents=True, exist_ok=True)

        # Write files with UTF-8 encoding (Windows compatibility)
        (server_dir / "app.py").write_text(server_code, encoding='utf-8')
        (server_dir / "requirements.txt").write_text(requirements, encoding='utf-8')
        (server_dir / "README.md").write_text(readme, encoding='utf-8')

        # Store metadata
        metadata = {
            "server_id": server_id,
            "server_name": plan_json["server_name"],
            "description": plan_json["description"],
            "tools": plan_json["tools"],
            "task_description": task_description,
            "generated_at": datetime.now().isoformat(),
            "directory": str(server_dir),
            "files": {
                "app": str(server_dir / "app.py"),
                "requirements": str(server_dir / "requirements.txt"),
                "readme": str(server_dir / "README.md")
            },
            "deployment_status": "generated",
            "complexity": plan_json.get("complexity", "medium")
        }

        self.generated_servers[server_id] = metadata

        print(f"[OK] Generated MCP server: {server_id}")
        print(f"[LOC] Location: {server_dir}")
        print(f"[TOOLS] Tools: {[t['name'] for t in plan_json['tools']]}")

        return metadata

    def _extract_json(self, text: str) -> Dict[str, Any]:
        """Extract JSON from LLM response"""
        import re

        # Try to find JSON block
        json_match = re.search(r'\{[\s\S]*\}', text)
        if json_match:
            return json.loads(json_match.group())

        # Try parsing entire response
        return json.loads(text)

    def _extract_code(self, text: str) -> str:
        """Extract Python code from LLM response"""
        import re

        # Try to find code block
        code_match = re.search(r'```python\n([\s\S]*?)\n```', text)
        if code_match:
            return code_match.group(1)

        code_match = re.search(r'```\n([\s\S]*?)\n```', text)
        if code_match:
            return code_match.group(1)

        # Return as-is if no code block found
        return text

    def _generate_requirements(self, dependencies: List[str]) -> str:
        """Generate requirements.txt content"""
        base_requirements = [
            "gradio>=6.0.0",
            "httpx>=0.28.0",
            "pydantic>=2.0.0"
        ]

        all_requirements = base_requirements + dependencies
        return "\n".join(all_requirements)

    def _generate_readme(self, plan: Dict[str, Any], task_description: str) -> str:
        """Generate README.md for the MCP server"""
        tools_md = "\n".join([
            f"- **{tool['name']}**: {tool['description']}"
            for tool in plan["tools"]
        ])

        return f"""# {plan['server_name']}

{plan['description']}

## Original Request
{task_description}

## Available Tools

{tools_md}

## Installation

```bash
pip install -r requirements.txt
```

## Usage

### As MCP Server

```python
# Connect to this MCP server from your agent
mcp_url = "https://huggingface.co/spaces/YOUR_USERNAME/{plan['server_name']}/gradio_api/mcp/sse"
```

### Standalone Testing

```bash
python app.py
```

Then open http://localhost:7860 in your browser.

## Auto-Generated

This MCP server was automatically generated by OmniMind Orchestrator.

Generated: {datetime.now().strftime("%Y-%m-%d %H:%M:%S")}
"""

    def _generate_server_id(self, server_name: str) -> str:
        """Generate unique server ID"""
        timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
        name_hash = hashlib.md5(server_name.encode()).hexdigest()[:6]
        return f"{server_name.lower().replace(' ', '_')}_{name_hash}_{timestamp}"

    async def improve_mcp_server(
        self,
        server_id: str,
        feedback: str,
        error_log: Optional[str] = None
    ) -> Dict[str, Any]:
        """
        Improve an existing MCP server based on feedback or errors.

        This makes the system SELF-EVOLVING - it learns and improves tools.
        """
        if server_id not in self.generated_servers:
            raise ValueError(f"Server {server_id} not found")

        metadata = self.generated_servers[server_id]
        current_code = Path(metadata["files"]["app"]).read_text()

        improvement_prompt = f"""You are improving an existing MCP server.

Current Implementation:
```python
{current_code}
```

Feedback: {feedback}

{f"Error Log: {error_log}" if error_log else ""}

Analyze the issues and generate an IMPROVED version of the code.
Fix bugs, optimize performance, add missing features.

Return ONLY the complete improved Python code.
"""

        result = await router.generate(
            improvement_prompt,
            task_type=TaskType.CODE_GEN,
            max_tokens=4000,
            temperature=0.2
        )

        improved_code = self._extract_code(result["response"])

        # Save improved version with UTF-8 encoding (Windows compatibility)
        server_dir = Path(metadata["directory"])
        backup_path = server_dir / f"app_backup_{datetime.now().strftime('%Y%m%d_%H%M%S')}.py"
        Path(metadata["files"]["app"]).rename(backup_path)
        Path(metadata["files"]["app"]).write_text(improved_code, encoding='utf-8')

        metadata["improved_at"] = datetime.now().isoformat()
        metadata["improvement_count"] = metadata.get("improvement_count", 0) + 1

        print(f"[OK] Improved MCP server: {server_id}")

        return metadata

    def list_servers(self) -> List[Dict[str, Any]]:
        """List all generated MCP servers"""
        return list(self.generated_servers.values())

    def get_server(self, server_id: str) -> Optional[Dict[str, Any]]:
        """Get metadata for a specific server"""
        return self.generated_servers.get(server_id)

    async def test_mcp_server(self, server_id: str, test_input: Dict[str, Any]) -> Dict[str, Any]:
        """
        Test a generated MCP server locally before deployment.

        Returns test results and any errors.
        """
        if server_id not in self.generated_servers:
            raise ValueError(f"Server {server_id} not found")

        metadata = self.generated_servers[server_id]

        # In production, this would actually run the MCP server
        # For now, return simulation
        return {
            "server_id": server_id,
            "status": "success",
            "test_input": test_input,
            "message": "Server would be tested here in production"
        }


# Global generator instance
generator = MCPGenerator()