a0c8464e58
Build Package / build (ubuntu-latest) (push) Failing after 1s
CodeQL / Analyze (python) (push) Failing after 1s
Core Typecheck / core-typecheck (push) Failing after 1s
Linting / lint (push) Failing after 1s
llama-dev tests / test-llama-dev (push) Failing after 1s
Publish Sub-Package to PyPI if Needed / publish_subpackage_if_needed (push) Has been skipped
Sync Docs to Developer Hub / sync-docs (push) Failing after 0s
Build Package / build (windows-latest) (push) Has been cancelled
331 lines
10 KiB
Plaintext
331 lines
10 KiB
Plaintext
{
|
|
"cells": [
|
|
{
|
|
"cell_type": "markdown",
|
|
"metadata": {},
|
|
"source": [
|
|
"# LlamaIndex + MCP Usage\n",
|
|
"\n",
|
|
"The `llama-index-tools-mcp` package provides several tools for using MCP with LlamaIndex."
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"metadata": {},
|
|
"outputs": [],
|
|
"source": [
|
|
"%pip install llama-index-tools-mcp"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "markdown",
|
|
"metadata": {},
|
|
"source": [
|
|
"## Using Tools from an MCP Server\n",
|
|
"\n",
|
|
"Using the `get_tools_from_mcp_url` or `aget_tools_from_mcp_url` function, you can get a list of `FunctionTool`s from an MCP server."
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"metadata": {},
|
|
"outputs": [],
|
|
"source": [
|
|
"from llama_index.tools.mcp import (\n",
|
|
" get_tools_from_mcp_url,\n",
|
|
" aget_tools_from_mcp_url,\n",
|
|
")\n",
|
|
"\n",
|
|
"# async\n",
|
|
"tools = await aget_tools_from_mcp_url(\"http://127.0.0.1:8000/mcp\")"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "markdown",
|
|
"metadata": {},
|
|
"source": [
|
|
"By default, this will use our `BasicMCPClient`, which will run a command or connect to the URL and return the tools.\n",
|
|
"\n",
|
|
"You can also pass in a custom `ClientSession` to use a different client.\n",
|
|
"\n",
|
|
"You can also specify a list of allowed tools to filter the tools that are returned."
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"metadata": {},
|
|
"outputs": [],
|
|
"source": [
|
|
"from llama_index.tools.mcp import BasicMCPClient\n",
|
|
"\n",
|
|
"client = BasicMCPClient(\"http://127.0.0.1:8000/mcp\")\n",
|
|
"\n",
|
|
"tools = await aget_tools_from_mcp_url(\n",
|
|
" \"http://127.0.0.1:8000/mcp\",\n",
|
|
" client=client,\n",
|
|
" allowed_tools=[\"tool1\", \"tool2\"],\n",
|
|
")"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "markdown",
|
|
"metadata": {},
|
|
"source": [
|
|
"## Converting a Workflow to an MCP App\n",
|
|
"\n",
|
|
"If you have a custom `Workflow`, you can convert it to an MCP app using the `workflow_as_mcp` function.\n",
|
|
"\n",
|
|
"For example, let's use the following workflow that will make a string loud:"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"metadata": {},
|
|
"outputs": [],
|
|
"source": [
|
|
"from llama_index.core.workflow import (\n",
|
|
" Context,\n",
|
|
" Workflow,\n",
|
|
" Event,\n",
|
|
" StartEvent,\n",
|
|
" StopEvent,\n",
|
|
" step,\n",
|
|
")\n",
|
|
"from llama_index.tools.mcp.utils import workflow_as_mcp\n",
|
|
"\n",
|
|
"\n",
|
|
"class RunEvent(StartEvent):\n",
|
|
" msg: str\n",
|
|
"\n",
|
|
"\n",
|
|
"class InfoEvent(Event):\n",
|
|
" msg: str\n",
|
|
"\n",
|
|
"\n",
|
|
"class LoudWorkflow(Workflow):\n",
|
|
" \"\"\"Useful for converting strings to uppercase and making them louder.\"\"\"\n",
|
|
"\n",
|
|
" @step\n",
|
|
" def step_one(self, ctx: Context, ev: RunEvent) -> StopEvent:\n",
|
|
" ctx.write_event_to_stream(InfoEvent(msg=\"Hello, world!\"))\n",
|
|
"\n",
|
|
" return StopEvent(result=ev.msg.upper() + \"!\")\n",
|
|
"\n",
|
|
"\n",
|
|
"workflow = LoudWorkflow()\n",
|
|
"\n",
|
|
"mcp = workflow_as_mcp(workflow)"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "markdown",
|
|
"metadata": {},
|
|
"source": [
|
|
"This code will automatically generate a `FastMCP` server that will\n",
|
|
"- Use the workflow class name as the tool name\n",
|
|
"- Use our custom `RunEvent` as the typed inputs to the tool\n",
|
|
"- Automatically use the SSE stream for streaming json dumps of the workflow event stream"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "markdown",
|
|
"metadata": {},
|
|
"source": [
|
|
"If this code was in a script called `script.py`, you could launch the MCP server with:\n",
|
|
"\n",
|
|
"```bash\n",
|
|
"mcp dev script.py\n",
|
|
"```\n",
|
|
"\n",
|
|
"Or the other commands documented in the [MCP CLI README](https://github.com/modelcontextprotocol/python-sdk?tab=readme-ov-file#quickstart).\n",
|
|
"\n",
|
|
"Note that to launch from the CLI, you may need to install the MCP CLI:\n",
|
|
"\n",
|
|
"```bash\n",
|
|
"pip install \"mcp[cli]\"\n",
|
|
"```\n",
|
|
"\n"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "markdown",
|
|
"metadata": {},
|
|
"source": [
|
|
"You can further customize the `FastMCP` server by passing in additional arguments to the `workflow_as_mcp` function:\n",
|
|
"- `workflow_name`: The name of the workflow. Defaults to the class name.\n",
|
|
"- `workflow_description`: The description of the workflow. Defaults to the class docstring.\n",
|
|
"- `start_event_model`: The event model to use for the start event. You can either use a custom `StartEvent` class in your workflow or pass in your own pydantic model here to define the inputs to the workflow.\n",
|
|
"- `**fastmcp_init_kwargs`: Any extra arguments to pass to the `FastMCP()` server constructor."
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "markdown",
|
|
"metadata": {},
|
|
"source": [
|
|
"## MCP Client Usage\n",
|
|
"\n",
|
|
"The `BasicMCPClient` provides comprehensive access to MCP server capabilities beyond just tools.\n",
|
|
"\n",
|
|
"### Basic Client Operations"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"metadata": {},
|
|
"outputs": [],
|
|
"source": [
|
|
"from llama_index.tools.mcp import BasicMCPClient\n",
|
|
"\n",
|
|
"# Connect to an MCP server using different transports\n",
|
|
"http_client = BasicMCPClient(\"https://example.com/mcp\") # Streamable HTTP\n",
|
|
"sse_client = BasicMCPClient(\"https://example.com/sse\") # Server-Sent Events\n",
|
|
"local_client = BasicMCPClient(\"python\", args=[\"server.py\"]) # stdio\n",
|
|
"\n",
|
|
"# List available tools\n",
|
|
"tools = await http_client.list_tools()\n",
|
|
"\n",
|
|
"# Call a tool\n",
|
|
"result = await http_client.call_tool(\"calculate\", {\"x\": 5, \"y\": 10})\n",
|
|
"\n",
|
|
"# List available resources\n",
|
|
"resources = await http_client.list_resources()\n",
|
|
"\n",
|
|
"# Read a resource\n",
|
|
"content, mime_type = await http_client.read_resource(\"config://app\")\n",
|
|
"\n",
|
|
"# List available prompts\n",
|
|
"prompts = await http_client.list_prompts()\n",
|
|
"\n",
|
|
"# Get a prompt\n",
|
|
"prompt_result = await http_client.get_prompt(\"greet\", {\"name\": \"World\"})"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "markdown",
|
|
"metadata": {},
|
|
"source": [
|
|
"### OAuth Authentication\n",
|
|
"\n",
|
|
"The client supports OAuth 2.0 authentication for connecting to protected MCP servers.\n",
|
|
"\n",
|
|
"You can see the [MCP docs](https://github.com/modelcontextprotocol/python-sdk/blob/main/README.md) for full details on configuring the various aspects of OAuth for both [clients](https://github.com/modelcontextprotocol/python-sdk?tab=readme-ov-file#oauth-authentication-for-clients) and [servers](https://github.com/modelcontextprotocol/python-sdk?tab=readme-ov-file#authentication)."
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"metadata": {},
|
|
"outputs": [],
|
|
"source": [
|
|
"from llama_index.tools.mcp import BasicMCPClient\n",
|
|
"\n",
|
|
"# Simple authentication with in-memory token storage\n",
|
|
"client = BasicMCPClient.with_oauth(\n",
|
|
" \"https://api.example.com/mcp\",\n",
|
|
" client_name=\"My App\",\n",
|
|
" redirect_uris=[\"http://localhost:3000/callback\"],\n",
|
|
" # Function to handle the redirect URL (e.g., open a browser)\n",
|
|
" redirect_handler=lambda url: print(f\"Please visit: {url}\"),\n",
|
|
" # Function to get the authorization code from the user\n",
|
|
" callback_handler=lambda: (input(\"Enter the code: \"), None),\n",
|
|
")\n",
|
|
"\n",
|
|
"# Use the authenticated client\n",
|
|
"tools = await client.list_tools()"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "markdown",
|
|
"metadata": {},
|
|
"source": [
|
|
"By default, the client will use an in-memory token storage if no `token_storage` is provided. You can pass in a custom `TokenStorage` instance to use a different storage.\n",
|
|
"\n",
|
|
"Below is an example showing the default in-memory token storage implementation."
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"metadata": {},
|
|
"outputs": [],
|
|
"source": [
|
|
"from llama_index.tools.mcp import BasicMCPClient\n",
|
|
"from mcp.client.auth import TokenStorage\n",
|
|
"from mcp.shared.auth import OAuthToken, OAuthClientInformationFull\n",
|
|
"from typing import Optional\n",
|
|
"\n",
|
|
"\n",
|
|
"class DefaultInMemoryTokenStorage(TokenStorage):\n",
|
|
" \"\"\"\n",
|
|
" Simple in-memory token storage implementation for OAuth authentication.\n",
|
|
"\n",
|
|
" This is the default storage used when none is provided to with_oauth().\n",
|
|
" Not suitable for production use across restarts as tokens are only stored\n",
|
|
" in memory.\n",
|
|
" \"\"\"\n",
|
|
"\n",
|
|
" def __init__(self):\n",
|
|
" self._tokens: Optional[OAuthToken] = None\n",
|
|
" self._client_info: Optional[OAuthClientInformationFull] = None\n",
|
|
"\n",
|
|
" async def get_tokens(self) -> Optional[OAuthToken]:\n",
|
|
" \"\"\"Get the stored OAuth tokens.\"\"\"\n",
|
|
" return self._tokens\n",
|
|
"\n",
|
|
" async def set_tokens(self, tokens: OAuthToken) -> None:\n",
|
|
" \"\"\"Store OAuth tokens.\"\"\"\n",
|
|
" self._tokens = tokens\n",
|
|
"\n",
|
|
" async def get_client_info(self) -> Optional[OAuthClientInformationFull]:\n",
|
|
" \"\"\"Get the stored client information.\"\"\"\n",
|
|
" return self._client_info\n",
|
|
"\n",
|
|
" async def set_client_info(\n",
|
|
" self, client_info: OAuthClientInformationFull\n",
|
|
" ) -> None:\n",
|
|
" \"\"\"Store client information.\"\"\"\n",
|
|
" self._client_info = client_info\n",
|
|
"\n",
|
|
"\n",
|
|
"# Use custom storage\n",
|
|
"client = BasicMCPClient.with_oauth(\n",
|
|
" \"https://api.example.com/mcp\",\n",
|
|
" client_name=\"My App\",\n",
|
|
" redirect_uris=[\"http://localhost:3000/callback\"],\n",
|
|
" redirect_handler=lambda url: print(f\"Please visit: {url}\"),\n",
|
|
" callback_handler=lambda: (input(\"Enter the code: \"), None),\n",
|
|
" token_storage=DefaultInMemoryTokenStorage(),\n",
|
|
")"
|
|
]
|
|
}
|
|
],
|
|
"metadata": {
|
|
"colab": {
|
|
"provenance": []
|
|
},
|
|
"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"
|
|
}
|
|
},
|
|
"nbformat": 4,
|
|
"nbformat_minor": 0
|
|
}
|