server.py•4.16 kB
"""
UseKeen MCP Server
This module provides the MCP server implementation for the UseKeen package documentation search.
It handles setting up the server, registering tool handlers, and connecting to the transport.
"""
import json
import logging
import asyncio
from typing import Any, Dict, List, Optional
from mcp.server.stdio import stdio_server
from mcp.server import Server, NotificationOptions
from mcp.server.models import InitializationOptions
import mcp.types as types
from .client import UseKeenClient
from .tools import PACKAGE_DOC_SEARCH_TOOL
logger = logging.getLogger("usekeen_mcp.server")
# Create the server
server = Server("UseKeen Documentation MCP Server")
# Create the UseKeen client (will be initialized in main)
usekeen_client = None
@server.list_tools()
async def handle_list_tools() -> List[types.Tool]:
"""
List available tools.
"""
return [
types.Tool(
name="usekeen_package_doc_search",
description="Search documentation of packages and services to find implementation details, examples, and specifications",
inputSchema={
"type": "object",
"properties": {
"package_name": {
"type": "string",
"description": "Name of the package or service to search documentation for (e.g. 'react', 'aws-s3', 'docker')"
},
"query": {
"type": "string",
"description": "Search term to find specific information within the package/service documentation (e.g. 'file upload example', 'authentication methods')"
}
},
"required": ["package_name"]
}
)
]
@server.call_tool()
async def handle_call_tool(
name: str, arguments: Dict[str, Any] | None
) -> List[types.TextContent | types.ImageContent]:
"""
Handle tool execution requests.
"""
global usekeen_client
logger.info(f"Received call to tool: {name} with arguments: {arguments}")
if name == "usekeen_package_doc_search":
if not arguments or "package_name" not in arguments:
raise ValueError("Missing required argument: package_name")
package_name = arguments["package_name"]
query = arguments.get("query")
try:
# Call the UseKeen API
response = usekeen_client.search_package_documentation(package_name, query)
return [types.TextContent(type="text", text=json.dumps(response))]
except Exception as e:
logger.error(f"Error searching package documentation: {e}")
return [types.TextContent(type="text", text=json.dumps({"error": str(e)}))]
else:
raise ValueError(f"Unknown tool: {name}")
async def main(api_key: str) -> None:
"""
Main function to run the server.
Args:
api_key: The UseKeen API key
"""
global usekeen_client
# Initialize the UseKeen client
usekeen_client = UseKeenClient(api_key)
# Run the server using stdin/stdout streams
logger.info("Connecting server to transport...")
async with stdio_server() as (read_stream, write_stream):
await server.run(
read_stream,
write_stream,
InitializationOptions(
server_name="UseKeen Documentation MCP Server",
server_version="1.0.0",
capabilities=server.get_capabilities(
notification_options=NotificationOptions(),
experimental_capabilities={},
),
),
)
logger.info("UseKeen Documentation MCP Server running on stdio")
def start_server(api_key: str) -> None:
"""
Start the MCP server.
Args:
api_key: The UseKeen API key
Returns:
None
"""
try:
asyncio.run(main(api_key))
except KeyboardInterrupt:
logger.info("Server stopped by user")
except Exception as e:
logger.error(f"Fatal error in server: {e}")
raise