src/mcp_ollama_python/main.py
Rama main · réplica de lectura del repositorio original.
"""
Ollama MCP Server - Main entry point (MCP Python SDK 2.x).
"""
from __future__ import annotations
import asyncio
import logging
import signal
import sys
from importlib.metadata import PackageNotFoundError, version
from typing import Any, Optional
# Configure logging
logging.basicConfig(
level=logging.INFO,
format="%(asctime)s - %(name)s - %(levelname)s - %(message)s",
)
logger = logging.getLogger(__name__)
# Use absolute imports so PyInstaller/standalone execution works even when __package__ is not set
try:
from mcp_ollama_python.server import OllamaMCPServer
from mcp_ollama_python.ollama_client import OllamaClient
except ImportError:
from .server import OllamaMCPServer
from .ollama_client import OllamaClient
try:
from mcp.server import Server, ServerRequestContext
from mcp.server.stdio import stdio_server
from mcp.types import (
CallToolRequestParams,
CallToolResult,
GetPromptRequestParams,
GetPromptResult,
ListPromptsResult,
ListResourcesResult,
ListToolsResult,
PaginatedRequestParams,
Prompt,
PromptArgument,
PromptMessage,
ReadResourceRequestParams,
ReadResourceResult,
Resource,
TextContent,
TextResourceContents,
Tool,
)
except ImportError as e:
logger.error("MCP package import failed: %s", e)
print(f"Error: mcp package import failed: {e}")
print("Please install it with: pip install 'mcp>=2.0.0,<3.0.0'")
sys.exit(1)
# Global server instance for signal handling
_server_instance: Optional[Server] = None
_shutdown_event: Optional[asyncio.Event] = None
def _package_version() -> str:
try:
return version("mcp-ollama-python")
except PackageNotFoundError:
return ""
def safe_print(*args, **kwargs) -> None:
"""
Safe print that handles closed stdout/stderr.
Args:
*args: Arguments to print
**kwargs: Keyword arguments for print function
"""
try:
print(*args, **kwargs)
except (ValueError, OSError):
# stdout/stderr is closed, ignore
pass
def signal_handler(signum: int, frame) -> None:
"""
Handle shutdown signals gracefully.
Args:
signum: Signal number received
frame: Current stack frame
"""
# Only perform signal-safe operations - set the shutdown event
if _shutdown_event:
_shutdown_event.set()
def _prompt_arguments(raw: Any) -> list[PromptArgument]:
"""Convert server prompt argument dicts to MCP PromptArgument models."""
if not isinstance(raw, list):
return []
out: list[PromptArgument] = []
for item in raw:
if not isinstance(item, dict):
continue
name = item.get("name")
if not name:
continue
out.append(
PromptArgument(
name=str(name),
description=item.get("description"),
required=bool(item.get("required", False)),
)
)
return out
def _prompt_messages(raw_messages: Any) -> list[PromptMessage]:
"""Convert get_prompt message dicts to MCP PromptMessage models."""
if not isinstance(raw_messages, list):
return []
messages: list[PromptMessage] = []
for item in raw_messages:
if not isinstance(item, dict):
continue
role = item.get("role") or "user"
content = item.get("content") or {}
text = ""
if isinstance(content, dict):
text = str(content.get("text") or "")
elif isinstance(content, str):
text = content
messages.append(
PromptMessage(role=role, content=TextContent(type="text", text=text))
)
return messages
async def main() -> None:
"""
Main function to start the MCP server.
Initializes the Ollama client, MCP server, and handles graceful shutdown.
Registers handlers for tools, resources, and prompts.
"""
global _server_instance, _shutdown_event
logger.info("Initializing Ollama MCP Server")
# Create shutdown event
_shutdown_event = asyncio.Event()
# Register signal handlers for graceful shutdown
signal.signal(signal.SIGINT, signal_handler)
signal.signal(signal.SIGTERM, signal_handler)
# Initialize Ollama client
try:
ollama_client = OllamaClient()
logger.info("Ollama client initialized successfully")
except Exception as e:
logger.error("Failed to initialize Ollama client: %s", e)
raise
# Initialize our MCP server
server = OllamaMCPServer(ollama_client)
async def handle_list_tools(
ctx: ServerRequestContext, params: PaginatedRequestParams | None
) -> ListToolsResult:
"""Handle list_tools request."""
del ctx, params # unused; required by MCP SDK 2 handler signature
try:
result = await server.handle_list_tools()
tools = [
Tool(
name=tool_data["name"],
description=tool_data["description"],
input_schema=tool_data["inputSchema"],
)
for tool_data in result["tools"]
]
logger.debug("Listed %d tools", len(tools))
return ListToolsResult(tools=tools)
except Exception as e:
logger.error("Error listing tools: %s", e)
raise
async def handle_call_tool(
ctx: ServerRequestContext, params: CallToolRequestParams
) -> CallToolResult:
"""Handle call_tool request."""
del ctx
try:
logger.debug("Calling tool: %s", params.name)
result = await server.handle_call_tool(params.name, params.arguments or {})
content = [
TextContent(type="text", text=item["text"])
for item in result.get("content") or []
if item.get("type") == "text" and "text" in item
]
return CallToolResult(
content=content,
structured_content=result.get("structuredContent"),
is_error=bool(result.get("isError", False)),
)
except Exception as e:
logger.error("Error calling tool %s: %s", params.name, e)
raise
async def handle_list_resources(
ctx: ServerRequestContext, params: PaginatedRequestParams | None
) -> ListResourcesResult:
"""Handle list_resources request."""
del ctx, params
try:
result = await server.handle_list_resources()
resources = [
Resource(
uri=resource_data["uri"],
name=resource_data["name"],
description=resource_data.get("description"),
mime_type=resource_data.get("mimeType", "text/plain"),
)
for resource_data in result["resources"]
]
logger.debug("Listed %d resources", len(resources))
return ListResourcesResult(resources=resources)
except Exception as e:
logger.error("Error listing resources: %s", e)
raise
async def handle_read_resource(
ctx: ServerRequestContext, params: ReadResourceRequestParams
) -> ReadResourceResult:
"""Handle read_resource request."""
del ctx
uri = str(params.uri)
try:
logger.debug("Reading resource: %s", uri)
result = await server.handle_read_resource(uri)
contents: list[TextResourceContents] = []
for item in result.get("contents") or []:
contents.append(
TextResourceContents(
uri=str(item.get("uri") or uri),
mime_type=item.get("mimeType", "text/plain"),
text=str(item.get("text") or ""),
)
)
return ReadResourceResult(contents=contents)
except Exception as e:
logger.error("Error reading resource %s: %s", uri, e)
raise
async def handle_list_prompts(
ctx: ServerRequestContext, params: PaginatedRequestParams | None
) -> ListPromptsResult:
"""Handle list_prompts request."""
del ctx, params
try:
result = await server.handle_list_prompts()
prompts = [
Prompt(
name=prompt_data["name"],
description=prompt_data.get("description"),
arguments=_prompt_arguments(prompt_data.get("arguments", [])),
)
for prompt_data in result["prompts"]
]
logger.debug("Listed %d prompts", len(prompts))
return ListPromptsResult(prompts=prompts)
except Exception as e:
logger.error("Error listing prompts: %s", e)
raise
async def handle_get_prompt(
ctx: ServerRequestContext, params: GetPromptRequestParams
) -> GetPromptResult:
"""Handle get_prompt request."""
del ctx
try:
logger.debug("Getting prompt: %s", params.name)
result = await server.handle_get_prompt(params.name, params.arguments)
return GetPromptResult(
description=result.get("description"),
messages=_prompt_messages(result.get("messages")),
)
except Exception as e:
logger.error("Error getting prompt %s: %s", params.name, e)
raise
# Create MCP server instance (SDK 2: handlers via on_* constructor kwargs)
mcp_server = Server(
"ollama-mcp",
version=_package_version(),
on_list_tools=handle_list_tools,
on_call_tool=handle_call_tool,
on_list_resources=handle_list_resources,
on_read_resource=handle_read_resource,
on_list_prompts=handle_list_prompts,
on_get_prompt=handle_get_prompt,
)
_server_instance = mcp_server
safe_print("Starting Ollama MCP Server...")
safe_print("Press Ctrl+C to stop the server")
# Run the server with graceful shutdown support
try:
async with stdio_server() as (read_stream, write_stream):
safe_print("Server started successfully!")
safe_print("Waiting for MCP client connections...")
# Create server task
server_task = asyncio.create_task(
mcp_server.run(
read_stream,
write_stream,
mcp_server.create_initialization_options(),
)
)
# Wait for either server completion or shutdown signal
shutdown_task = asyncio.create_task(_shutdown_event.wait())
done, pending = await asyncio.wait(
[server_task, shutdown_task], return_when=asyncio.FIRST_COMPLETED
)
# Cancel pending tasks
for task in pending:
task.cancel()
try:
await task
except asyncio.CancelledError:
pass
safe_print("Server stopped.")
except KeyboardInterrupt:
logger.info("Server interrupted by user")
safe_print("\nServer interrupted by user")
except Exception as e:
logger.error("Server error: %s", e, exc_info=True)
safe_print(f"Server error: {e}")
raise
finally:
# Cleanup
logger.info("Cleaning up resources")
await ollama_client.client.aclose()
safe_print("Cleanup completed.")
def run() -> None:
"""
Entry point for the mcp-ollama-python command.
Runs the async main function and handles shutdown gracefully.
"""
try:
logger.info("Starting mcp-ollama-python")
asyncio.run(main())
except KeyboardInterrupt:
logger.info("Shutdown complete")
safe_print("\nShutdown complete.")
sys.exit(0)
except Exception as e:
logger.error("Fatal error: %s", e, exc_info=True)
safe_print(f"Fatal error: {e}")
sys.exit(1)
def stop() -> None:
"""
Stop the running server (for programmatic control).
Sets the shutdown event to trigger graceful server shutdown.
"""
if _shutdown_event:
logger.info("Stop signal sent to server")
_shutdown_event.set()
safe_print("Stop signal sent to server")
else:
logger.warning("No server instance running")
safe_print("No server instance running")
if __name__ == "__main__":
run()