"""
MCP Integration for ModuleWrapper
This module provides MCP middleware for the ModuleWrapper, allowing it to be used
within the FastMCP2 framework for module discovery and semantic search.
This middleware follows the FastMCP pattern with proper hooks for request handling,
tool execution, and resource access, similar to the auth middleware.
"""
from config.enhanced_logging import setup_logger
logger = setup_logger()
import asyncio
import importlib
import inspect
import json
from fastmcp.server.dependencies import get_context
# Import MCP-related components
from fastmcp.server.middleware import Middleware, MiddlewareContext
from typing_extensions import Any, Dict, List, Optional
from config.enhanced_logging import setup_logger
# Import type definitions
from .adapter_types import (
ModuleComponentInfo,
ModuleComponentsResponse,
WrappedModuleInfo,
WrappedModulesResponse,
)
# Import ModuleWrapper
from .module_wrapper import ModuleWrapper
logger = setup_logger()
class ModuleWrapperMiddleware(Middleware):
"""
Middleware for integrating ModuleWrapper with MCP using proper hooks.
This middleware manages ModuleWrapper instances and provides context
for module discovery and semantic search operations. It follows the
FastMCP middleware pattern with on_request, on_call_tool, and on_read_resource hooks.
"""
def __init__(
self,
qdrant_host: Optional[str] = None,
qdrant_port: Optional[int] = None,
qdrant_url: Optional[str] = None,
qdrant_api_key: Optional[str] = None,
collection_prefix: str = "mcp_module_",
embedding_model: str = "sentence-transformers/all-MiniLM-L6-v2",
auto_discovery: bool = True,
modules_to_wrap: Optional[List[str]] = None,
):
"""
Initialize the ModuleWrapper middleware using unified settings.
Args:
qdrant_host: Qdrant server hostname (overrides settings)
qdrant_port: Qdrant server port (overrides settings)
qdrant_url: Full Qdrant URL (overrides settings)
qdrant_api_key: API key for Qdrant authentication (overrides settings)
collection_prefix: Prefix for Qdrant collections
embedding_model: Model to use for generating embeddings
auto_discovery: Whether to auto-discover modules
modules_to_wrap: List of module names to wrap
"""
# Import unified settings - same source as QdrantUnifiedMiddleware
from config.settings import settings
logger.info(
"🔧 MODULEWRAPPER DEBUG: Initializing ModuleWrapper with unified settings"
)
logger.debug(
f"🔧 Settings values: host={settings.qdrant_host}, port={settings.qdrant_port}, url='{settings.qdrant_url}', api_key={'***' if settings.qdrant_api_key else 'None'}"
)
# Use provided parameters or fall back to unified settings (same as QdrantUnifiedMiddleware)
self.qdrant_host = (
qdrant_host if qdrant_host is not None else settings.qdrant_host
)
self.qdrant_port = (
qdrant_port if qdrant_port is not None else settings.qdrant_port
)
self.qdrant_url = qdrant_url if qdrant_url is not None else settings.qdrant_url
self.qdrant_api_key = (
qdrant_api_key if qdrant_api_key is not None else settings.qdrant_api_key
)
# Log final configuration for debugging
logger.debug(
f"🔧 MODULEWRAPPER DEBUG: Final config: host={self.qdrant_host}, port={self.qdrant_port}, url='{self.qdrant_url}', api_key={'***' if self.qdrant_api_key else 'None'}"
)
self.collection_prefix = collection_prefix
self.embedding_model = embedding_model
self.auto_discovery = auto_discovery
self.modules_to_wrap = modules_to_wrap or []
# Store module wrappers
self.wrappers = {}
# Initialize if modules are provided
if self.modules_to_wrap:
self._initialize_wrappers()
logger.info(
"✅ ModuleWrapper middleware initialized with unified settings - same Qdrant instance as main middleware"
)
def _initialize_wrappers(self):
"""Initialize wrappers for specified modules."""
for module_name in self.modules_to_wrap:
try:
# Import the module
module = importlib.import_module(module_name)
# Create wrapper
collection_name = (
f"{self.collection_prefix}{module_name.replace('.', '_')}"
)
wrapper = ModuleWrapper(
module_or_name=module,
qdrant_host=self.qdrant_host,
qdrant_port=self.qdrant_port,
qdrant_url=self.qdrant_url,
qdrant_api_key=self.qdrant_api_key,
collection_name=collection_name,
embedding_model=self.embedding_model,
auto_initialize=True,
)
# Store wrapper
self.wrappers[module_name] = wrapper
logger.info(f"✅ Initialized wrapper for module: {module_name}")
except ImportError:
logger.warning(f"⚠️ Could not import module: {module_name}")
except Exception as e:
logger.error(f"❌ Failed to initialize wrapper for {module_name}: {e}")
async def on_call_tool(self, context: MiddlewareContext, call_next):
"""
Handle tool execution for module wrapper operations.
This middleware provides module wrapping functionality by:
1. Checking if the tool is a module wrapper tool
2. Providing access to module wrapper functionality through context
3. Handling module-specific operations
"""
tool_name = getattr(context.message, "name", "unknown")
# Check if this is a module wrapper tool
module_wrapper_tools = {
"wrap_module",
"search_module",
"get_module_component",
"list_module_components",
"list_wrapped_modules",
}
if tool_name in module_wrapper_tools:
logger.debug(
f"ModuleWrapperMiddleware: Processing module wrapper tool: {tool_name}"
)
# Make module wrapper functionality available through FastMCP context
if context.fastmcp_context:
# Store a reference to this middleware's functionality
await context.fastmcp_context.set_state("module_wrapper_instance", self)
try:
result = await call_next(context)
if tool_name in module_wrapper_tools:
logger.debug(f"Module wrapper tool {tool_name} executed successfully")
return result
except Exception as e:
if tool_name in module_wrapper_tools:
logger.error(f"Error executing module wrapper tool {tool_name}: {e}")
raise
async def wrap_module(self, module_name: str) -> bool:
"""
Wrap a module and index its components.
Args:
module_name: Name of the module to wrap
Returns:
True if successful, False otherwise
"""
try:
# Check if already wrapped
if module_name in self.wrappers:
logger.info(f"Module {module_name} is already wrapped")
return True
# Import the module
module = importlib.import_module(module_name)
# Create wrapper
collection_name = f"{self.collection_prefix}{module_name.replace('.', '_')}"
# Create wrapper in a separate thread to avoid blocking
def create_wrapper():
return ModuleWrapper(
module_or_name=module,
qdrant_host=self.qdrant_host,
qdrant_port=self.qdrant_port,
qdrant_url=self.qdrant_url,
qdrant_api_key=self.qdrant_api_key,
collection_name=collection_name,
embedding_model=self.embedding_model,
auto_initialize=True,
)
wrapper = await asyncio.to_thread(create_wrapper)
# Store wrapper
self.wrappers[module_name] = wrapper
logger.info(f"✅ Wrapped module: {module_name}")
return True
except Exception as e:
logger.error(f"❌ Failed to wrap module {module_name}: {e}")
return False
async def search_module(
self, module_name: str, query: str, limit: int = 5, score_threshold: float = 0.3
) -> List[Dict[str, Any]]:
"""
Search for components in a module.
Args:
module_name: Name of the module to search
query: Search query
limit: Maximum number of results
score_threshold: Minimum similarity score
Returns:
List of matching components
"""
# Check if module is wrapped
if module_name not in self.wrappers:
# Try to wrap it
success = await self.wrap_module(module_name)
if not success:
return []
# Get wrapper
wrapper = self.wrappers[module_name]
# Search
try:
results = await wrapper.search_async(query, limit, score_threshold)
return results
except Exception as e:
logger.error(f"❌ Search failed: {e}")
return []
async def get_component(
self, module_name: str, path: str
) -> Optional[Dict[str, Any]]:
"""
Get a component by its path.
Args:
module_name: Name of the module
path: Path to the component
Returns:
Component information if found, None otherwise
"""
# Check if module is wrapped
if module_name not in self.wrappers:
# Try to wrap it
success = await self.wrap_module(module_name)
if not success:
return None
# Get wrapper
wrapper = self.wrappers[module_name]
# Get component info
try:
info = wrapper.get_component_info(path)
if info:
# Add the actual component
component = wrapper.get_component_by_path(path)
if component:
return {**info, "component": component}
return info
except Exception as e:
logger.error(f"❌ Failed to get component: {e}")
return None
async def list_modules(self) -> List[str]:
"""
List all wrapped modules.
Returns:
List of module names
"""
return list(self.wrappers.keys())
def setup_module_wrapper_tools(mcp):
"""Setup MCP tools for ModuleWrapper using FastMCP context pattern."""
async def get_module_wrapper() -> Optional[ModuleWrapperMiddleware]:
"""Get the ModuleWrapper middleware from FastMCP context."""
try:
ctx = get_context()
return await ctx.get_state("module_wrapper_instance")
except Exception as e:
logger.error(f"Failed to get module wrapper from context: {e}")
return None
@mcp.tool(
name="wrap_module",
description="Wrap a Python module and index its components for semantic search",
tags={"module", "wrapper", "index", "qdrant", "vector", "search"},
annotations={
"title": "Wrap Module",
"readOnlyHint": False,
"destructiveHint": False,
"idempotentHint": True,
"openWorldHint": True,
},
)
async def wrap_module(module_name: str, index_components: bool = True) -> str:
"""
Wrap a Python module and index its components for semantic search.
Args:
module_name: Name of the module to wrap (e.g., "json", "os.path")
index_components: Whether to index the module's components for semantic search
Returns:
Result message
"""
try:
# Get middleware from FastMCP context
middleware = await get_module_wrapper()
if not middleware:
return "❌ Error: ModuleWrapper middleware not available"
logger.info(f"Attempting to wrap module: {module_name}")
success = await middleware.wrap_module(module_name)
logger.info(f"Wrap module result: success={success}")
if success:
return f"✅ Successfully wrapped module: {module_name}"
else:
# Include the keywords that the test is looking for
return f"❌ Error: Module not found or import failed: {module_name}"
except Exception as e:
logger.error(f"Exception in wrap_module: {e}")
return f"❌ Error wrapping module: {str(e)}"
@mcp.tool(
name="search_module",
description="Search for components (classes, functions, etc.) in a module using natural language",
tags={"module", "search", "semantic", "qdrant", "vector"},
annotations={
"title": "Search Module",
"readOnlyHint": True,
"destructiveHint": False,
"idempotentHint": True,
"openWorldHint": True,
},
)
async def search_module(
module_name: str, query: str, limit: int = 5, include_source: bool = False
) -> str:
"""
Search for components in a module using natural language.
Args:
module_name: Name of the module to search (e.g., "json", "os.path")
query: Natural language search query (e.g., "parse json string", "read file")
limit: Maximum number of results to return
include_source: Whether to include source code in results
Returns:
JSON string with search results
"""
# Validate limit parameter
if limit <= 0:
return json.dumps(
{
"module": module_name,
"query": query,
"error": f"Invalid limit value: {limit}. Limit must be a positive integer.",
},
indent=2,
)
try:
# Get middleware from FastMCP context
middleware = await get_module_wrapper()
if not middleware:
return json.dumps(
{
"module": module_name,
"query": query,
"error": "ModuleWrapper middleware not available",
},
indent=2,
)
results = await middleware.search_module(module_name, query, limit)
# Format results
formatted_results = []
for result in results:
formatted_result = {
"name": result["name"],
"path": result["path"],
"type": result["type"],
"score": result["score"],
"docstring": result["docstring"],
}
# Add source if requested
if include_source and "component" in result and result["component"]:
try:
source = inspect.getsource(result["component"])
formatted_result["source"] = source
except (TypeError, OSError):
formatted_result["source"] = "Source code not available"
formatted_results.append(formatted_result)
return json.dumps(
{
"module": module_name,
"query": query,
"results": formatted_results,
"count": len(formatted_results),
},
indent=2,
)
except Exception as e:
return f"❌ Search failed: {str(e)}"
@mcp.tool(
name="get_module_component",
description="Get detailed information about a specific component in a module",
tags={"module", "component", "info", "qdrant"},
annotations={
"title": "Get Module Component",
"readOnlyHint": True,
"destructiveHint": False,
"idempotentHint": True,
"openWorldHint": True,
},
)
async def get_module_component(
module_name: str, component_path: str, include_source: bool = False
) -> str:
"""
Get detailed information about a specific component in a module.
Args:
module_name: Name of the module (e.g., "json", "os.path")
component_path: Path to the component (e.g., "json.loads", "os.path.join")
include_source: Whether to include source code
Returns:
JSON string with component information
"""
try:
# Get middleware from FastMCP context
middleware = await get_module_wrapper()
if not middleware:
return json.dumps(
{"error": "ModuleWrapper middleware not available"}, indent=2
)
info = await middleware.get_component(module_name, component_path)
if not info:
return f"❌ Component not found: {component_path}"
# Format result
formatted_info = {
"name": info["name"],
"path": info["full_path"],
"type": info["type"],
"module_path": info["module_path"],
"docstring": info["docstring"],
}
# Add source if requested
if include_source and "component" in info and info["component"]:
try:
source = inspect.getsource(info["component"])
formatted_info["source"] = source
except (TypeError, OSError):
formatted_info["source"] = "Source code not available"
return json.dumps(formatted_info, indent=2)
except Exception as e:
return f"❌ Error getting component: {str(e)}"
@mcp.tool(
name="list_module_components",
description="List all components in a wrapped module",
tags={"module", "list", "components", "qdrant"},
annotations={
"title": "List Module Components",
"readOnlyHint": True,
"destructiveHint": False,
"idempotentHint": True,
"openWorldHint": False,
},
)
async def list_module_components(module_name: str) -> ModuleComponentsResponse:
"""
List all components in a wrapped module.
Args:
module_name: Name of the module to list components for
Returns:
ModuleComponentsResponse: Structured list of module components with metadata
"""
try:
# Get middleware from FastMCP context
middleware = await get_module_wrapper()
if not middleware:
return ModuleComponentsResponse(
components=[],
count=0,
module=module_name,
error="ModuleWrapper middleware not available",
)
# Check if module is wrapped
if module_name not in middleware.wrappers:
# Try to wrap it
success = await middleware.wrap_module(module_name)
if not success:
return ModuleComponentsResponse(
components=[],
count=0,
module=module_name,
error=f"Module not wrapped: {module_name}",
)
# Get wrapper
wrapper = middleware.wrappers[module_name]
# Get components
raw_components = wrapper.list_components()
# Convert to structured format
components: List[ModuleComponentInfo] = []
for comp in raw_components:
# Handle both string and dict component formats
if isinstance(comp, str):
# Simple string format - parse the path
comp_info: ModuleComponentInfo = {
"name": comp.split(".")[-1] if "." in comp else comp,
"path": comp,
"type": "unknown",
"score": None,
"docstring": "",
"source": None,
}
elif isinstance(comp, dict):
# Dictionary format with full metadata
comp_info: ModuleComponentInfo = {
"name": comp.get("name", ""),
"path": comp.get("path", comp.get("full_path", "")),
"type": comp.get("type", "unknown"),
"score": comp.get("score"),
"docstring": comp.get("docstring", ""),
"source": comp.get("source"),
}
else:
# Skip invalid component formats
continue
components.append(comp_info)
return ModuleComponentsResponse(
components=components,
count=len(components),
module=module_name,
error=None,
)
except Exception as e:
logger.error(f"❌ Failed to list components: {e}")
return ModuleComponentsResponse(
components=[],
count=0,
module=module_name,
error=f"Error listing components: {str(e)}",
)
@mcp.tool(
name="list_wrapped_modules",
description="List all modules that have been wrapped for semantic search",
tags={"module", "list", "qdrant"},
annotations={
"title": "List Wrapped Modules",
"readOnlyHint": True,
"destructiveHint": False,
"idempotentHint": True,
"openWorldHint": False,
},
)
async def list_wrapped_modules() -> WrappedModulesResponse:
"""
List all modules that have been wrapped for semantic search.
Returns:
WrappedModulesResponse: Structured list of wrapped modules with metadata
"""
try:
# Get middleware from FastMCP context
middleware = await get_module_wrapper()
if not middleware:
return WrappedModulesResponse(
modules=[], count=0, error="ModuleWrapper middleware not available"
)
module_names = await middleware.list_modules()
# Convert to structured format
modules: List[WrappedModuleInfo] = []
for module_name in module_names:
# Get wrapper to get component count
wrapper = middleware.wrappers.get(module_name)
component_count = None
if wrapper:
try:
components = wrapper.list_components()
component_count = len(components)
except:
pass # If we can't get component count, leave as None
module_info: WrappedModuleInfo = {
"name": module_name,
"indexed": True, # If it's in the list, it's been indexed
"component_count": component_count,
}
modules.append(module_info)
return WrappedModulesResponse(
modules=modules, count=len(modules), error=None
)
except Exception as e:
logger.error(f"❌ Error listing modules: {str(e)}")
return WrappedModulesResponse(
modules=[], count=0, error=f"Error listing modules: {str(e)}"
)
def setup_module_wrapper_middleware(mcp, modules_to_wrap=None, tool_pushdown=False):
"""
Set up the ModuleWrapper middleware and tools using proper middleware hooks.
Args:
mcp: MCP server instance
modules_to_wrap: List of module names to wrap initially
Returns:
The middleware instance
"""
# Create middleware
middleware = ModuleWrapperMiddleware(modules_to_wrap=modules_to_wrap or [])
# Register middleware - this will enable the hooks (on_call_tool)
mcp.add_middleware(middleware)
# Set up tools - no middleware parameter needed, tools get it from context
if tool_pushdown:
setup_module_wrapper_tools(mcp)
return middleware
