"""
Search tools for the MCP server.
This module implements MCP tools for searching and querying infrastructure entities.
Tools provide capabilities for:
- Full-text search across entities and observations
- Filtered queries based on entity types and metadata
- Relationship-based entity discovery
- Pattern matching and similarity search
Each tool follows standard patterns:
- Database integration through SQLAlchemy sessions
- Proper error handling with ValidationError and DatabaseError
- Consistent return structures with typed dictionaries
- Clean transaction management with commits and rollbacks
"""
from typing import Dict, Any, List, Optional
from sqlalchemy.orm import Session
from sqlalchemy import or_
from mcp.server.fastmcp import FastMCP
from ..db.connection import get_db
from ..db.models.entities import Entity
from ..db.models.observations import Observation
from ..utils.errors import DatabaseError, ValidationError
def register_tools(mcp: FastMCP) -> None:
"""Register search tools with the MCP server.
This function registers all search-related tools with the MCP server instance.
Tools include:
- search_entities: Full-text search across entities
- search_observations: Search entity observations
- find_related: Find entities via relationships
Each tool is registered with proper type hints, error handling, and database integration.
Tools follow MCP protocol patterns for consistent behavior and error reporting.
Search tools handle:
- Full-text entity search
- Metadata field matching
- Relationship traversal
- Pattern recognition
- Similarity scoring
Tool capabilities include:
- Query construction
- Result filtering
- Relevance scoring
- Pagination support
- Sort ordering
"""
@mcp.tool()
def search_entities(
query: str,
entity_type: Optional[str] = None,
limit: int = 100,
offset: int = 0,
db: Session = next(get_db()),
) -> List[Dict[str, Any]]:
"""Search for entities matching the query.
Args:
query: Search query string
entity_type: Optional entity type filter
limit: Maximum number of results (default: 100)
offset: Number of results to skip (default: 0)
db: Database session
Returns:
List of matching entities with relevance scores
Raises:
ValidationError: If query is empty
DatabaseError: If search fails
Example:
>>> results = search_entities(
... query="web server",
... entity_type="instance",
... limit=10
... )
>>> print(results[0])
{
"id": 1,
"name": "web-server-1",
"type": "instance",
"score": 0.85
}
"""
try:
if not query or not query.strip():
raise ValidationError("Search query cannot be empty")
# Build search query
search = db.query(Entity)
# Add type filter if specified
if entity_type:
search = search.filter(Entity.type == entity_type.lower())
# Add text search conditions
terms = query.strip().lower().split()
conditions = []
for term in terms:
conditions.append(Entity.name.ilike(f"%{term}%"))
conditions.append(
Entity.metadata["description"].astext.ilike(f"%{term}%")
)
search = search.filter(or_(*conditions))
# Add pagination
search = search.offset(offset).limit(limit)
# Execute search
results = search.all()
# Format results
return [
{
"id": entity.id,
"name": entity.name,
"type": entity.type,
"score": 1.0, # Placeholder for relevance score
}
for entity in results
]
except Exception as e:
raise DatabaseError(f"Search failed: {str(e)}")