server.py•2.93 kB
import asyncio
import logging
import sys
from typing import Optional
import click
import mcp
import mcp.types as types
from mcp.server import Server, NotificationOptions
from mcp.server.models import InitializationOptions
from .llm import LLMConnector
logging.basicConfig(level=logging.DEBUG)
logger = logging.getLogger(__name__)
def serve(openai_api_key: str) -> Server:
server = Server("openai-server")
connector = LLMConnector(openai_api_key)
@server.list_tools()
async def handle_list_tools() -> list[types.Tool]:
return [
types.Tool(
name="ask-openai",
description="Ask my assistant models a direct question",
inputSchema={
"type": "object",
"properties": {
"query": {"type": "string", "description": "Ask assistant"},
"model": {"type": "string", "default": "o3-mini", "enum": ["o3-mini", "gpt-4o-mini"]}
},
"required": ["query"]
}
)
]
@server.call_tool()
async def handle_tool_call(name: str, arguments: dict | None) -> list[types.TextContent]:
try:
if not arguments:
raise ValueError("No arguments provided")
if name == "ask-openai":
response = await connector.ask_openai(
query=arguments["query"],
model=arguments.get("model", "o3-mini")
)
return [types.TextContent(type="text", text=response)]
raise ValueError(f"Unknown tool: {name}")
except Exception as e:
logger.error(f"Tool call failed: {str(e)}")
return [types.TextContent(type="text", text=f"Error: {str(e)}")]
return server
@click.command()
@click.option("--openai-api-key", envvar="OPENAI_API_KEY", required=True)
def main(openai_api_key: str):
try:
async def _run():
async with mcp.server.stdio.stdio_server() as (read_stream, write_stream):
server = serve(openai_api_key)
await server.run(
read_stream, write_stream,
InitializationOptions(
server_name="openai-server",
server_version="0.1.0",
capabilities=server.get_capabilities(
notification_options=NotificationOptions(),
experimental_capabilities={}
)
)
)
asyncio.run(_run())
except KeyboardInterrupt:
logger.info("Server stopped by user")
except Exception as e:
logger.exception("Server failed")
sys.exit(1)
if __name__ == "__main__":
main()