325 lines
11 KiB
Plaintext
325 lines
11 KiB
Plaintext
{
|
||
"cells": [
|
||
{
|
||
"cell_type": "markdown",
|
||
"metadata": {},
|
||
"source": [
|
||
"# Building a Local MCP Client with LlamaIndex\n",
|
||
"\n",
|
||
"This Jupyter notebook walks you through creating a **local MCP (Model Context Protocol) client** that can chat with a database through tools exposed by an MCP server—completely on your machine. Follow the cells in order for a smooth, self‑contained tutorial."
|
||
]
|
||
},
|
||
{
|
||
"cell_type": "code",
|
||
"execution_count": 1,
|
||
"metadata": {},
|
||
"outputs": [],
|
||
"source": [
|
||
"import nest_asyncio\n",
|
||
"nest_asyncio.apply()"
|
||
]
|
||
},
|
||
{
|
||
"cell_type": "markdown",
|
||
"metadata": {},
|
||
"source": [
|
||
"## 2 Setup a local LLM"
|
||
]
|
||
},
|
||
{
|
||
"cell_type": "code",
|
||
"execution_count": 3,
|
||
"metadata": {},
|
||
"outputs": [],
|
||
"source": [
|
||
"from llama_index.llms.ollama import Ollama\n",
|
||
"from llama_index.core import Settings\n",
|
||
"\n",
|
||
"llm = Ollama(model=\"llama3.2\", request_timeout=120.0)\n",
|
||
"Settings.llm = llm"
|
||
]
|
||
},
|
||
{
|
||
"cell_type": "markdown",
|
||
"metadata": {},
|
||
"source": [
|
||
"## 3 Initialize the MCP client and build the agent\n",
|
||
"Point the client at your local MCP server’s **SSE endpoint** (default shown below), and list the available tools."
|
||
]
|
||
},
|
||
{
|
||
"cell_type": "code",
|
||
"execution_count": 4,
|
||
"metadata": {},
|
||
"outputs": [],
|
||
"source": [
|
||
"from llama_index.tools.mcp import BasicMCPClient, McpToolSpec\n",
|
||
"\n",
|
||
"mcp_client = BasicMCPClient(\"http://127.0.0.1:8000/sse\")\n",
|
||
"mcp_tools = McpToolSpec(client=mcp_client) # you can also pass list of allowed tools"
|
||
]
|
||
},
|
||
{
|
||
"cell_type": "code",
|
||
"execution_count": 5,
|
||
"metadata": {},
|
||
"outputs": [
|
||
{
|
||
"name": "stdout",
|
||
"output_type": "stream",
|
||
"text": [
|
||
"add_data Add new data to the people table using a SQL INSERT query.\n",
|
||
"\n",
|
||
" Args:\n",
|
||
" query (str): SQL INSERT query following this format:\n",
|
||
" INSERT INTO people (name, age, profession)\n",
|
||
" VALUES ('John Doe', 30, 'Engineer')\n",
|
||
" \n",
|
||
" Schema:\n",
|
||
" - name: Text field (required)\n",
|
||
" - age: Integer field (required)\n",
|
||
" - profession: Text field (required)\n",
|
||
" Note: 'id' field is auto-generated\n",
|
||
" \n",
|
||
" Returns:\n",
|
||
" bool: True if data was added successfully, False otherwise\n",
|
||
" \n",
|
||
" Example:\n",
|
||
" >>> query = '''\n",
|
||
" ... INSERT INTO people (name, age, profession)\n",
|
||
" ... VALUES ('Alice Smith', 25, 'Developer')\n",
|
||
" ... '''\n",
|
||
" >>> add_data(query)\n",
|
||
" True\n",
|
||
" \n",
|
||
"read_data Read data from the people table using a SQL SELECT query.\n",
|
||
"\n",
|
||
" Args:\n",
|
||
" query (str, optional): SQL SELECT query. Defaults to \"SELECT * FROM people\".\n",
|
||
" Examples:\n",
|
||
" - \"SELECT * FROM people\"\n",
|
||
" - \"SELECT name, age FROM people WHERE age > 25\"\n",
|
||
" - \"SELECT * FROM people ORDER BY age DESC\"\n",
|
||
" \n",
|
||
" Returns:\n",
|
||
" list: List of tuples containing the query results.\n",
|
||
" For default query, tuple format is (id, name, age, profession)\n",
|
||
" \n",
|
||
" Example:\n",
|
||
" >>> # Read all records\n",
|
||
" >>> read_data()\n",
|
||
" [(1, 'John Doe', 30, 'Engineer'), (2, 'Alice Smith', 25, 'Developer')]\n",
|
||
" \n",
|
||
" >>> # Read with custom query\n",
|
||
" >>> read_data(\"SELECT name, profession FROM people WHERE age < 30\")\n",
|
||
" [('Alice Smith', 'Developer')]\n",
|
||
" \n"
|
||
]
|
||
}
|
||
],
|
||
"source": [
|
||
"tools = await mcp_tools.to_tool_list_async()\n",
|
||
"for tool in tools:\n",
|
||
" print(tool.metadata.name, tool.metadata.description)"
|
||
]
|
||
},
|
||
{
|
||
"cell_type": "markdown",
|
||
"metadata": {},
|
||
"source": [
|
||
"## 3 Define the system prompt\n",
|
||
"This prompt steers the LLM when it needs to decide how and when to call tools."
|
||
]
|
||
},
|
||
{
|
||
"cell_type": "code",
|
||
"execution_count": 6,
|
||
"metadata": {},
|
||
"outputs": [],
|
||
"source": [
|
||
"SYSTEM_PROMPT = \"\"\"\\\n",
|
||
"You are an AI assistant for Tool Calling.\n",
|
||
"\n",
|
||
"Before you help a user, you need to work with tools to interact with Our Database\n",
|
||
"\"\"\""
|
||
]
|
||
},
|
||
{
|
||
"cell_type": "markdown",
|
||
"metadata": {},
|
||
"source": [
|
||
"## 4 Helper function: `get_agent()`\n",
|
||
"Creates a `FunctionAgent` wired up with the MCP tool list and your chosen LLM."
|
||
]
|
||
},
|
||
{
|
||
"cell_type": "code",
|
||
"execution_count": 7,
|
||
"metadata": {},
|
||
"outputs": [],
|
||
"source": [
|
||
"from llama_index.tools.mcp import McpToolSpec\n",
|
||
"from llama_index.core.agent.workflow import FunctionAgent\n",
|
||
"\n",
|
||
"async def get_agent(tools: McpToolSpec):\n",
|
||
" tools = await tools.to_tool_list_async()\n",
|
||
" agent = FunctionAgent(\n",
|
||
" name=\"Agent\",\n",
|
||
" description=\"An agent that can work with Our Database software.\",\n",
|
||
" tools=tools,\n",
|
||
" llm=OpenAI(model=\"gpt-4\"),\n",
|
||
" system_prompt=SYSTEM_PROMPT,\n",
|
||
" )\n",
|
||
" return agent"
|
||
]
|
||
},
|
||
{
|
||
"cell_type": "markdown",
|
||
"metadata": {},
|
||
"source": [
|
||
"## 5 Helper function: `handle_user_message()`\n",
|
||
"Streams intermediate tool calls (for transparency) and returns the final response."
|
||
]
|
||
},
|
||
{
|
||
"cell_type": "code",
|
||
"execution_count": 8,
|
||
"metadata": {},
|
||
"outputs": [],
|
||
"source": [
|
||
"from llama_index.core.agent.workflow import (\n",
|
||
" FunctionAgent, \n",
|
||
" ToolCallResult, \n",
|
||
" ToolCall)\n",
|
||
"\n",
|
||
"from llama_index.core.workflow import Context\n",
|
||
"\n",
|
||
"async def handle_user_message(\n",
|
||
" message_content: str,\n",
|
||
" agent: FunctionAgent,\n",
|
||
" agent_context: Context,\n",
|
||
" verbose: bool = False,\n",
|
||
"):\n",
|
||
" handler = agent.run(message_content, ctx=agent_context)\n",
|
||
" async for event in handler.stream_events():\n",
|
||
" if verbose and type(event) == ToolCall:\n",
|
||
" print(f\"Calling tool {event.tool_name} with kwargs {event.tool_kwargs}\")\n",
|
||
" elif verbose and type(event) == ToolCallResult:\n",
|
||
" print(f\"Tool {event.tool_name} returned {event.tool_output}\")\n",
|
||
"\n",
|
||
" response = await handler\n",
|
||
" return str(response)"
|
||
]
|
||
},
|
||
{
|
||
"cell_type": "markdown",
|
||
"metadata": {},
|
||
"source": [
|
||
"## 6 Initialize the MCP client and build the agent\n",
|
||
"Point the client at your local MCP server’s **SSE endpoint** (default shown below), build the agent, and setup agent context."
|
||
]
|
||
},
|
||
{
|
||
"cell_type": "code",
|
||
"execution_count": 9,
|
||
"metadata": {},
|
||
"outputs": [],
|
||
"source": [
|
||
"from llama_index.tools.mcp import BasicMCPClient, McpToolSpec\n",
|
||
"\n",
|
||
"\n",
|
||
"mcp_client = BasicMCPClient(\"http://127.0.0.1:8000/sse\")\n",
|
||
"mcp_tool = McpToolSpec(client=mcp_client)\n",
|
||
"\n",
|
||
"# get the agent\n",
|
||
"agent = await get_agent(mcp_tool)\n",
|
||
"\n",
|
||
"# create the agent context\n",
|
||
"agent_context = Context(agent)"
|
||
]
|
||
},
|
||
{
|
||
"cell_type": "code",
|
||
"execution_count": null,
|
||
"metadata": {},
|
||
"outputs": [
|
||
{
|
||
"name": "stdout",
|
||
"output_type": "stream",
|
||
"text": [
|
||
"User: Add to the db: Rafael Nadal whose age is 39 and is a tennis player\n",
|
||
"Calling tool add_data with kwargs {'query': \"INSERT INTO people (name, age, profession) VALUES ('Rafael Nadal', 39, 'Tennis Player')\"}\n",
|
||
"Tool add_data returned meta=None content=[TextContent(type='text', text='true', annotations=None)] isError=False\n",
|
||
"Agent: The data has been added successfully.\n",
|
||
"User: fetch data\n",
|
||
"Calling tool read_data with kwargs {'query': 'SELECT * FROM people'}\n",
|
||
"Tool read_data returned meta=None content=[TextContent(type='text', text='1', annotations=None), TextContent(type='text', text='Rafael Nadal', annotations=None), TextContent(type='text', text='39', annotations=None), TextContent(type='text', text='Tennis Player', annotations=None)] isError=False\n",
|
||
"Agent: Here is the data from the database:\n",
|
||
"\n",
|
||
"1. ID: 1\n",
|
||
" Name: Rafael Nadal\n",
|
||
" Age: 39\n",
|
||
" Profession: Tennis Player\n",
|
||
"User: add to the db: Roger federer whose age is 42 and is a tennis player\n",
|
||
"Calling tool add_data with kwargs {'query': \"INSERT INTO people (name, age, profession) VALUES ('Roger Federer', 42, 'Tennis Player')\"}\n",
|
||
"Tool add_data returned meta=None content=[TextContent(type='text', text='true', annotations=None)] isError=False\n",
|
||
"Agent: The data has been added successfully.\n",
|
||
"User: fetch data\n",
|
||
"Calling tool read_data with kwargs {'query': 'SELECT * FROM people'}\n",
|
||
"Tool read_data returned meta=None content=[TextContent(type='text', text='1', annotations=None), TextContent(type='text', text='Rafael Nadal', annotations=None), TextContent(type='text', text='39', annotations=None), TextContent(type='text', text='Tennis Player', annotations=None), TextContent(type='text', text='2', annotations=None), TextContent(type='text', text='Roger Federer', annotations=None), TextContent(type='text', text='42', annotations=None), TextContent(type='text', text='Tennis Player', annotations=None)] isError=False\n",
|
||
"Agent: Here is the data from the database:\n",
|
||
"\n",
|
||
"1. ID: 1\n",
|
||
" Name: Rafael Nadal\n",
|
||
" Age: 39\n",
|
||
" Profession: Tennis Player\n",
|
||
"\n",
|
||
"2. ID: 2\n",
|
||
" Name: Roger Federer\n",
|
||
" Age: 42\n",
|
||
" Profession: Tennis Player\n"
|
||
]
|
||
}
|
||
],
|
||
"source": [
|
||
"# Run the agent!\n",
|
||
"while True:\n",
|
||
" user_input = input(\"Enter your message: \")\n",
|
||
" if user_input == \"exit\":\n",
|
||
" break\n",
|
||
" print(\"User: \", user_input)\n",
|
||
" response = await handle_user_message(user_input, agent, agent_context, verbose=True)\n",
|
||
" print(\"Agent: \", response)"
|
||
]
|
||
},
|
||
{
|
||
"cell_type": "code",
|
||
"execution_count": null,
|
||
"metadata": {},
|
||
"outputs": [],
|
||
"source": []
|
||
}
|
||
],
|
||
"metadata": {
|
||
"kernelspec": {
|
||
"display_name": ".venv",
|
||
"language": "python",
|
||
"name": "python3"
|
||
},
|
||
"language_info": {
|
||
"codemirror_mode": {
|
||
"name": "ipython",
|
||
"version": 3
|
||
},
|
||
"file_extension": ".py",
|
||
"mimetype": "text/x-python",
|
||
"name": "python",
|
||
"nbconvert_exporter": "python",
|
||
"pygments_lexer": "ipython3",
|
||
"version": "3.12.9"
|
||
}
|
||
},
|
||
"nbformat": 4,
|
||
"nbformat_minor": 2
|
||
}
|