RAG-MCP example

A lightweight RAG-MCP prototype that embeds MCP tool descriptions with sentence-transformers and retrieves the most semantically relevant tools for each user query before generation.

Features

• Semantic Tool Retrieval – Uses sentence-transformers to embed and rank MCP tools by semantic similarity to the user query (retrieval-augmented layer).

• MCP Integration – Connects to any MCP-compliant server via JSON-RPC to fetch available tools (tools/list) and build an indexed registry.

• Context Injection Layer – Dynamically constructs a compact prompt containing only the top-K most relevant tools to reduce prompt bloat and improve reasoning accuracy.

• RAG-MCP Core Components – Implements the four essential RAG-MCP steps: Index, Rank, Inject, and Generate.

• Embedding Utilities – Built-in embedding and cosine-similarity utilities for ranking and experimentation with different sentence-transformer models.

Installation

# Install Python dependencies
pip install -r requirements.txt

Configuration

The MCP endpoint is configured via a global constant in src/main.py:

# Configuration
DEFAULT_MCP_ENDPOINT = 'http://localhost:3001'

To use a different port, simply change this value in the code.

Setup

1. Start the MCP server (if using local server):

export WRITER_API_KEY="your_actual_api_key" && npx -y writer-sdk-mcp@latest --transport=http --port=3001 &

Note: This command starts the server on port 3001. You can change the port by modifying --port=3001 to your desired port (e.g., --port=8080).

2. Stop the MCP server when done:

# Find and kill the server process
# 3001 by default, but change this to whatever port you chose
lsof -i :3001
kill <PID>

# Or use this one-liner to stop by port
lsof -ti:3001 | xargs kill

Usage

async def example_usage():
    print("RAG-MCP Example")
    print("==================\n")
    
    try:
        # Initialize RAG-MCP system (uses DEFAULT_MCP_ENDPOINT constant)
        mcp_endpoint = DEFAULT_MCP_ENDPOINT
        rag_mcp = RAGMCP(mcp_endpoint)
        await rag_mcp.initialize_tool_registry()
        
        # Example queries (you can change this to what you want to query)
        queries = [
            "Create a new Knowledge Graph for our Q4 product updates",
            "Add this PDF to the marketing Knowledge Graph",
            "Ask the Knowledge Graph what we said about Palmyra X5 in last quarter's docs",
            "Update the description of the finance Knowledge Graph",
            "List all Knowledge Graphs linked to my chat agent",
            "Remove an outdated file from the engineering Knowledge Graph",
        ]

        
        for query in queries:
            print(f"Query: {query}")
            
            # RAG-MCP Steps 1-2: Retrieve and rank tools
            tools = rag_mcp.retrieve_tools(query, 2)
            print(f"  Retrieved {len(tools)} tools:")
            for tool in tools:
                description = format_description(tool['description'])
                print(f"    - {tool['name']}: {description} (score: {tool['similarity_score']:.3f})")
            
            # RAG-MCP Complete Flow: Retrieve → Rank → Inject → Generate
            result = await rag_mcp.handle_query(query)
            print(f"  Recommendation: {result.get('recommendations', 'No recommendations')}")
            print(f"  Context: {result.get('context', 'No context')[:100]}...")
            print()
        
    except Exception as error:
        print(f"❌ Error running example: {error}")

if __name__ == "__main__":
    asyncio.run(example_usage())

Run the Complete Example

# Run the Python example
python src/main.py

RAG-MCP Architecture

Method	Description
initialize_tool_registry()	Step 1 – Index tools with embeddings
retrieve_tools(query, top_k)	Step 2 – Rank tools by semantic similarity
build_prompt_with_tools(tools, query)	Step 3 – Inject relevant tools into context
handle_query(query)	Complete RAG-MCP flow (Retrieve → Rank → Inject → Generate)
search_tools(search_term)	Search tools by name or description
get_available_tools()	Return all available tools