Skip to content

Commit 2b4e0e1

Browse files
committed
refactor: tools register
1 parent 80444c1 commit 2b4e0e1

File tree

7 files changed

+127
-117
lines changed

7 files changed

+127
-117
lines changed

src/elasticsearch_mcp_server/es_client.py

-40
This file was deleted.

src/elasticsearch_mcp_server/server.py

+51-20
Original file line numberDiff line numberDiff line change
@@ -1,39 +1,70 @@
11
#!/usr/bin/env python3
22
import logging
3+
import os
4+
from dotenv import load_dotenv
5+
from elasticsearch import Elasticsearch
6+
import warnings
37
from fastmcp import FastMCP
4-
from .tools.index import IndexTools
5-
from .tools.document import DocumentTools
6-
from .tools.cluster import ClusterTools
7-
from .tools.alias import AliasTools
8+
from elasticsearch_mcp_server.tools.index import IndexTools
9+
from elasticsearch_mcp_server.tools.document import DocumentTools
10+
from elasticsearch_mcp_server.tools.cluster import ClusterTools
11+
from elasticsearch_mcp_server.tools.alias import AliasTools
12+
from elasticsearch_mcp_server.tools.register import ToolsRegister
813

914
class ElasticsearchMCPServer:
1015
def __init__(self):
1116
self.name = "elasticsearch_mcp_server"
1217
self.mcp = FastMCP(self.name)
13-
14-
# Configure logging
15-
logging.basicConfig(
16-
level=logging.INFO,
17-
format='%(asctime)s - %(name)s - %(levelname)s - %(message)s'
18-
)
1918
self.logger = logging.getLogger(self.name)
19+
self.es_client = self._create_elasticsearch_client()
2020

2121
# Initialize tools
2222
self._register_tools()
23+
24+
def _get_es_config(self):
25+
"""Get Elasticsearch configuration from environment variables."""
26+
# Load environment variables from .env file
27+
load_dotenv()
28+
config = {
29+
"host": os.getenv("ELASTIC_HOST"),
30+
"username": os.getenv("ELASTIC_USERNAME"),
31+
"password": os.getenv("ELASTIC_PASSWORD")
32+
}
33+
34+
if not all([config["username"], config["password"]]):
35+
self.logger.error("Missing required Elasticsearch configuration. Please check environment variables:")
36+
self.logger.error("ELASTIC_USERNAME and ELASTIC_PASSWORD are required")
37+
raise ValueError("Missing required Elasticsearch configuration")
38+
39+
return config
40+
41+
def _create_elasticsearch_client(self) -> Elasticsearch:
42+
"""Create and return an Elasticsearch client using configuration from environment."""
43+
config = self._get_es_config()
44+
45+
# Disable SSL warnings
46+
warnings.filterwarnings("ignore", message=".*TLS with verify_certs=False is insecure.*",)
47+
48+
return Elasticsearch(
49+
config["host"],
50+
basic_auth=(config["username"], config["password"]),
51+
verify_certs=False
52+
)
2353

2454
def _register_tools(self):
2555
"""Register all MCP tools."""
26-
# Initialize tool classes
27-
index_tools = IndexTools(self.logger)
28-
document_tools = DocumentTools(self.logger)
29-
cluster_tools = ClusterTools(self.logger)
30-
alias_tools = AliasTools(self.logger)
56+
# Create a tools register
57+
register = ToolsRegister(self.logger, self.es_client, self.mcp)
3158

32-
# Register tools from each module
33-
index_tools.register_tools(self.mcp)
34-
document_tools.register_tools(self.mcp)
35-
cluster_tools.register_tools(self.mcp)
36-
alias_tools.register_tools(self.mcp)
59+
# Define all tool classes to register
60+
tool_classes = [
61+
IndexTools,
62+
DocumentTools,
63+
ClusterTools,
64+
AliasTools
65+
]
66+
# Register all tools
67+
register.register_all_tools(tool_classes)
3768

3869
def run(self):
3970
"""Run the MCP server."""
+7-15
Original file line numberDiff line numberDiff line change
@@ -1,35 +1,27 @@
1-
from typing import Any
2-
from ..es_client import ElasticsearchClient
1+
from fastmcp import FastMCP
32
from mcp.types import TextContent
43

5-
6-
class AliasTools(ElasticsearchClient):
7-
def register_tools(self, mcp: Any):
8-
"""Register alias-related tools."""
9-
10-
@mcp.tool(description="List all aliases in the Elasticsearch cluster")
4+
class AliasTools:
5+
def register_tools(self, mcp: FastMCP):
6+
@mcp.tool()
117
async def list_aliases() -> list[TextContent]:
128
"""List all aliases in the Elasticsearch cluster."""
13-
self.logger.info("Listing aliases...")
149
try:
15-
aliases = self.es_client.cat.aliases(format="json")
16-
return [TextContent(type="text", text=str(aliases))]
10+
return self.es_client.cat.aliases(format="json")
1711
except Exception as e:
1812
self.logger.error(f"Error listing aliases: {e}")
1913
return [TextContent(type="text", text=f"Error: {str(e)}")]
2014

21-
@mcp.tool(description="Get alias information for an index")
15+
@mcp.tool()
2216
async def get_alias(index: str) -> list[TextContent]:
2317
"""
2418
Get alias information for a specific index.
2519
2620
Args:
2721
index: Name of the index
2822
"""
29-
self.logger.info(f"Getting alias information for index: {index}")
3023
try:
31-
response = self.es_client.indices.get_alias(index=index)
32-
return [TextContent(type="text", text=str(response))]
24+
return self.es_client.indices.get_alias(index=index)
3325
except Exception as e:
3426
self.logger.error(f"Error getting alias information: {e}")
3527
return [TextContent(type="text", text=f"Error: {str(e)}")]
+7-13
Original file line numberDiff line numberDiff line change
@@ -1,38 +1,32 @@
1-
import logging
2-
from typing import Dict, Any
3-
from ..es_client import ElasticsearchClient
1+
from fastmcp import FastMCP
42
from mcp.types import TextContent
53

6-
class ClusterTools(ElasticsearchClient):
7-
def register_tools(self, mcp: Any):
4+
class ClusterTools:
5+
def register_tools(self, mcp: FastMCP):
86
"""Register cluster-related tools."""
97

10-
@mcp.tool(description="Get cluster health status")
8+
@mcp.tool()
119
async def get_cluster_health() -> list[TextContent]:
1210
"""
1311
Get health status of the Elasticsearch cluster.
1412
Returns information about the number of nodes, shards, etc.
1513
"""
16-
self.logger.info("Getting cluster health")
1714
try:
18-
response = self.es_client.cluster.health()
19-
return [TextContent(type="text", text=str(response))]
15+
return self.es_client.cluster.health()
2016
except Exception as e:
2117
self.logger.error(f"Error getting cluster health: {e}")
2218
return [TextContent(type="text", text=f"Error: {str(e)}")]
2319

24-
@mcp.tool(description="Get cluster statistics")
20+
@mcp.tool()
2521
async def get_cluster_stats() -> list[TextContent]:
2622
"""
2723
Get statistics from a cluster wide perspective.
2824
The API returns basic index metrics (shard numbers, store size, memory usage) and information
2925
about the current nodes that form the cluster (number, roles, os, jvm versions, memory usage, cpu and installed plugins).
3026
https://www.elastic.co/guide/en/elasticsearch/reference/8.17/cluster-stats.html
3127
"""
32-
self.logger.info("Getting cluster stats")
3328
try:
34-
response = self.es_client.cluster.stats()
35-
return [TextContent(type="text", text=str(response))]
29+
return self.es_client.cluster.stats()
3630
except Exception as e:
3731
self.logger.error(f"Error getting cluster stats: {e}")
3832
return [TextContent(type="text", text=f"Error: {str(e)}")]
Original file line numberDiff line numberDiff line change
@@ -1,13 +1,11 @@
1-
import logging
2-
from typing import Dict, Any
3-
from ..es_client import ElasticsearchClient
1+
from fastmcp import FastMCP
42
from mcp.types import TextContent
53

6-
class DocumentTools(ElasticsearchClient):
7-
def register_tools(self, mcp: Any):
4+
class DocumentTools:
5+
def register_tools(self, mcp: FastMCP):
86
"""Register document-related tools."""
97

10-
@mcp.tool(description="Search documents in an index with a custom query")
8+
@mcp.tool()
119
async def search_documents(index: str, body: dict) -> list[TextContent]:
1210
"""
1311
Search documents in a specified index using a custom query.
@@ -18,8 +16,7 @@ async def search_documents(index: str, body: dict) -> list[TextContent]:
1816
"""
1917
self.logger.info(f"Searching in index: {index} with query: {body}")
2018
try:
21-
response = self.es_client.search(index=index, body=body)
22-
return [TextContent(type="text", text=str(response))]
19+
return self.es_client.search(index=index, body=body)
2320
except Exception as e:
2421
self.logger.error(f"Error searching documents: {e}")
2522
return [TextContent(type="text", text=f"Error: {str(e)}")]
+12-21
Original file line numberDiff line numberDiff line change
@@ -1,51 +1,42 @@
1-
import logging
2-
from typing import Dict, Any
3-
from ..es_client import ElasticsearchClient
1+
from fastmcp import FastMCP
42
from mcp.types import TextContent
3+
from typing import Dict
54

6-
class IndexTools(ElasticsearchClient):
7-
def register_tools(self, mcp: Any):
8-
"""Register index-related tools."""
9-
10-
@mcp.tool(description="List all indices in the Elasticsearch cluster")
11-
async def list_indices() -> list[TextContent]:
12-
"""List all indices in the Elasticsearch cluster."""
13-
self.logger.info("Listing indices...")
5+
class IndexTools:
6+
def register_tools(self, mcp: FastMCP):
7+
@mcp.tool()
8+
async def list_indices() -> Dict:
9+
"""List all indices."""
1410
try:
15-
indices = self.es_client.cat.indices(format="json")
16-
return [TextContent(type="text", text=str(indices))]
11+
return self.es_client.cat.indices()
1712
except Exception as e:
1813
self.logger.error(f"Error listing indices: {e}")
1914
return [TextContent(type="text", text=f"Error: {str(e)}")]
2015

21-
@mcp.tool(description="Get index mapping")
16+
@mcp.tool()
2217
async def get_mapping(index: str) -> list[TextContent]:
2318
"""
2419
Get the mapping for an index.
2520
2621
Args:
2722
index: Name of the index
2823
"""
29-
self.logger.info(f"Getting mapping for index: {index}")
3024
try:
31-
response = self.es_client.indices.get_mapping(index=index)
32-
return [TextContent(type="text", text=str(response))]
25+
return self.es_client.indices.get_mapping(index=index)
3326
except Exception as e:
3427
self.logger.error(f"Error getting mapping: {e}")
3528
return [TextContent(type="text", text=f"Error: {str(e)}")]
3629

37-
@mcp.tool(description="Get index settings")
30+
@mcp.tool()
3831
async def get_settings(index: str) -> list[TextContent]:
3932
"""
4033
Get the settings for an index.
4134
4235
Args:
4336
index: Name of the index
4437
"""
45-
self.logger.info(f"Getting settings for index: {index}")
4638
try:
47-
response = self.es_client.indices.get_settings(index=index)
48-
return [TextContent(type="text", text=str(response))]
39+
return self.es_client.indices.get_settings(index=index)
4940
except Exception as e:
5041
self.logger.error(f"Error getting settings: {e}")
5142
return [TextContent(type="text", text=f"Error: {str(e)}")]
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,45 @@
1+
import logging
2+
3+
from fastmcp import FastMCP
4+
from typing import List, Type
5+
from elasticsearch import Elasticsearch
6+
7+
class ToolsRegister:
8+
"""
9+
A class responsible for creating and registering all tool classes.
10+
This centralizes the tool registration process and reduces duplication.
11+
"""
12+
13+
def __init__(self, logger: logging.Logger, es_client: Elasticsearch, mcp: FastMCP):
14+
"""
15+
Initialize the tools registrar.
16+
17+
Args:
18+
logger: Logger instance for logging
19+
es_client: Elasticsearch client instance for API calls
20+
mcp: FastMCP instance for tool registration
21+
"""
22+
self.logger = logger
23+
self.es_client = es_client
24+
self.mcp = mcp
25+
26+
def register_all_tools(self, tool_classes: List[Type]):
27+
"""
28+
Create instances of all tool classes and register their tools.
29+
30+
Args:
31+
tool_classes: List of tool classes to instantiate and register
32+
"""
33+
for tool_class in tool_classes:
34+
# Create an instance of the tool class without initialization parameters
35+
tool_instance = tool_class()
36+
37+
# Set logger and es_client attributes directly
38+
tool_instance.logger = self.logger
39+
# Pass the raw Elasticsearch client instead of the wrapper
40+
tool_instance.es_client = self.es_client
41+
42+
# Register the tools from this instance
43+
tool_instance.register_tools(self.mcp)
44+
45+
self.logger.info(f"Registered tools from {tool_class.__name__}")

0 commit comments

Comments
 (0)