"""
MCP server implementation for local RAG.
"""
import logging
import os
from typing import List
from dotenv import load_dotenv
from pydantic import BaseModel, Field
from mcp.server import Server
from mcp import types
from indexer.retriever import create_retriever_from_env
load_dotenv()
logger = logging.getLogger(__name__)
NAME='mcp-ragnar'
VERSION='0.1.0'
server = Server(NAME)
retriever = create_retriever_from_env(os.getenv('INDEX_ROOT', '/tmp'))
DESCRIPTION = os.getenv('MCP_DESCRIPTION', "RAG retriever")
class Retriever(BaseModel):
question: str = Field(description="the question to use for querying the vector store")
@server.list_tools()
async def handle_list_tools() -> List[types.Tool]:
"""
List available tools.
Each tool specifies its arguments using JSON Schema validation.
"""
return [
types.Tool(
name="retrieve",
description=DESCRIPTION,
inputSchema=Retriever.model_json_schema(),
)
]
@server.call_tool()
async def handle_call_tool(
name: str,
arguments: dict | None
) -> List[types.TextContent | types.ImageContent | types.EmbeddedResource]:
"""
Handle tool execution requests.
"""
if not arguments:
raise ValueError("Missing arguments")
if name == "retrieve":
validated_data = Retriever.model_validate(arguments)
result = retriever.query_as_text(**validated_data.model_dump())
return [
types.TextContent(
type="text",
text=str(result)
)
]
else:
raise ValueError(f"Unknown tool: {name}")
