MCP Search Server

# /// script # requires-python = ">=3.10" # dependencies = [ # "rich", # "typer", # "pydantic>=2.10.4", # "pydantic-settings>=2.7.0" # ] # /// """ Generate JSON schema for MCP Agent configuration (mcp-agent.config.yaml). """ import json import re import sys from pathlib import Path from typing import Any, Dict, Tuple import typer from rich.console import Console from pydantic import BaseModel from pydantic_settings import BaseSettings app = typer.Typer() console = Console() def extract_model_info(content: str) -> Dict[str, Dict[str, str]]: """ Extract docstrings for all models and their fields. Returns a dict mapping model names to their field descriptions. """ models = {} current_model = None # Split content into lines for processing lines = content.splitlines() for i, line in enumerate(lines): # Look for class definition class_match = re.match(r"\s*class\s+(\w+)(?:\([^)]+\))?\s*:", line.strip()) if class_match: current_model = class_match.group(1) models[current_model] = {"__doc__": ""} # Look for class docstring for j in range(i + 1, min(i + 4, len(lines))): doc_match = re.match(r'\s*"""(.+?)"""', lines[j], re.DOTALL) if doc_match: models[current_model]["__doc__"] = doc_match.group(1).strip() break continue # If we're inside a model definition, look for field definitions if current_model: # Check if we've exited the class definition (unindented line that's not empty or comment) if line and not line.startswith(" ") and not line.startswith("#"): current_model = None continue # Look for field definitions with type annotations field_match = re.match(r"\s+(\w+)\s*:", line) if field_match: field_name = field_match.group(1) # Skip if this is model_config or other special attributes if field_name in ("model_config", "Config"): continue description = None # Look for Field description in the current line field_desc_match = re.search(r'Field\([^)]*description="([^"]+)"', line) if field_desc_match: description = field_desc_match.group(1).strip() else: # Look ahead for docstring until we hit another field definition or non-empty, non-docstring line for j in range(i + 1, min(i + 4, len(lines))): next_line = lines[j].strip() # If we hit a non-empty line that's not a docstring, stop looking if next_line and not next_line.startswith('"""'): break # Try to match docstring doc_match = re.match(r'\s*"""(.+?)"""', lines[j], re.DOTALL) if doc_match: description = doc_match.group(1).strip() break if description: models[current_model][field_name] = description # Debug output console.print("\nFound models and their field descriptions:") for model, fields in models.items(): console.print(f"\n[bold]{model}[/bold]: {fields.get('__doc__', '')}") for field, desc in fields.items(): if field != "__doc__": console.print(f" {field}: {desc}") return models class MockModule: """Mock module that returns itself for any attribute access.""" def __getattr__(self, _: str) -> Any: return self def __call__(self, *args: Any, **kwargs: Any) -> Any: return self def create_mock_modules() -> None: """Create mock modules for imports we want to ignore.""" mocked_modules = [ "opentelemetry", "opentelemetry.sdk", "opentelemetry.sdk.trace", "opentelemetry.sdk.resources", "opentelemetry.exporter.otlp.proto.http", "opentelemetry.trace", "mcp_agent.logging", "mcp_agent.logging.logger", "yaml", ] for module_name in mocked_modules: if module_name not in sys.modules: sys.modules[module_name] = MockModule() def load_settings_class( file_path: Path, ) -> Tuple[type[BaseSettings], Dict[str, Dict[str, str]]]: """Load Settings class from a Python file.""" # Add src directory to Python path src_dir = file_path.parent.parent.parent / "src" sys.path.insert(0, str(src_dir)) # Mock required modules create_mock_modules() # Create namespace with required classes namespace = { "BaseModel": BaseModel, "BaseSettings": BaseSettings, "Path": Path, "Dict": dict, "List": list, "Literal": str, # Simplified for schema } with open(file_path, mode="r", encoding="utf-8") as f: content = f.read() # Extract all model info before executing model_info = extract_model_info(content) # Execute the file exec(content, namespace) return namespace["Settings"], model_info def apply_descriptions_to_schema( schema: Dict[str, Any], model_info: Dict[str, Dict[str, str]] ) -> None: """Recursively apply descriptions to schema and all its nested models.""" if not isinstance(schema, dict): return # Handle $defs (nested model definitions) if "$defs" in schema: for model_name, model_schema in schema["$defs"].items(): if model_name in model_info: # Apply class docstring doc = model_info[model_name].get("__doc__", "").strip() if doc: model_schema["description"] = doc # Apply field descriptions if "properties" in model_schema: for field_name, field_schema in model_schema["properties"].items(): if field_name in model_info[model_name]: field_schema["description"] = model_info[model_name][ field_name ].strip() # Handle root properties if "properties" in schema: for field_name, field_schema in schema["properties"].items(): if "Settings" in model_info and field_name in model_info["Settings"]: field_schema["description"] = model_info["Settings"][field_name].strip() @app.command() def generate( config_py: Path = typer.Option( Path("src/mcp_agent/config.py"), "--config", "-c", help="Path to the config.py file", ), output: Path = typer.Option( Path("schema/mcp-agent.config.schema.json"), "--output", "-o", help="Output path for the schema file", ), ): """Generate JSON schema from Pydantic models in config.py""" if not config_py.exists(): console.print(f"[red]Error:[/] File not found: {config_py}") raise typer.Exit(1) try: Settings, model_info = load_settings_class(config_py) schema = Settings.model_json_schema() # Debug: Print raw schema structure before modifications console.print("\nSchema structure:") if "$defs" in schema: console.print("Found models in $defs:", list(schema["$defs"].keys())) # Add schema metadata schema.update( { "$schema": "http://json-schema.org/draft-07/schema#", "title": "MCP Agent Configuration Schema", "description": "Configuration schema for MCP Agent applications", } ) # Apply descriptions to all nested models recursively apply_descriptions_to_schema(schema, model_info) # Ensure output directory exists output.parent.mkdir(parents=True, exist_ok=True) # Make output path absolute if it isn't already output = output.absolute() # Write schema with open(output, "w") as f: json.dump(schema, f, indent=2) console.print(f"[green]✓[/] Schema written to: {output}") # Get path relative to cwd for VS Code settings try: rel_path = f"./{output.relative_to(Path.cwd())}" except ValueError: # If can't make relative, use absolute path rel_path = str(output) # Print VS Code settings suggestion vscode_settings = { "yaml.schemas": { rel_path: [ "mcp-agent.config.yaml", "mcp_agent.config.yaml", "mcp-agent.secrets.yaml", "mcp_agent.secrets.yaml", ] } } console.print("\n[yellow]VS Code Integration:[/]") console.print("Add this to .vscode/settings.json:") console.print(json.dumps(vscode_settings, indent=2)) except Exception as e: console.print(f"[red]Error generating schema:[/] {str(e)}") raise typer.Exit(1) if __name__ == "__main__": app()