From 5a2e19916a89f4aefb3df53807fc4546fe7da4ce Mon Sep 17 00:00:00 2001 From: George Pickett Date: Tue, 6 Oct 2026 22:24:27 +0000 Subject: [PATCH] feat: add Streamable HTTP MCP and keyless Parallel search example --- README.md | 1 + docs/api/tools.mdx | 16 ++++- docs/topics/tools.mdx | 36 +++++++++- examples/parallel_search.py | 28 ++++++++ poetry.lock | 2 +- pyproject.toml | 2 +- rigging/tools/mcp.py | 66 +++++++++++++++-- tests/test_mcp.py | 138 ++++++++++++++++++++++++++++++++++++ 8 files changed, 278 insertions(+), 11 deletions(-) create mode 100644 examples/parallel_search.py create mode 100644 tests/test_mcp.py diff --git a/README.md b/README.md index 897c601c..83896891 100644 --- a/README.md +++ b/README.md @@ -139,6 +139,7 @@ Want more? - Jupyter code interpreter: [**jupyter.py**](examples/jupyter.py) - OverTheWire Bandit Agent: [**bandit.py**](examples/bandit.py) - Damn Vulnerable Restaurant Agent: [**dvra.py**](examples/dvra.py) +- Keyless web search with Parallel Search MCP: [**parallel_search.py**](examples/parallel_search.py) - RAG Pipeline: [**rag.py**](examples/rag.py) (from [kyleavery](https://github.com/kyleavery/)) - Integrating dreadnode-owned [robopages](https://github.com/dreadnode/robopages-cli) as a tool server (basic nmap scan example): [**rigging_example.py**](https://github.com/dreadnode/robopages-cli/blob/main/examples/rigging_example.py) diff --git a/docs/api/tools.mdx b/docs/api/tools.mdx index 4b33024e..ab78b58e 100644 --- a/docs/api/tools.mdx +++ b/docs/api/tools.mdx @@ -713,7 +713,9 @@ MCPClient ```python MCPClient( transport: Transport, - connection: StdioConnection | SSEConnection, + connection: StdioConnection + | SSEConnection + | StreamableHTTPConnection, ) ``` @@ -721,7 +723,11 @@ A client for communicating with MCP servers. ```python -def __init__(self, transport: Transport, connection: "StdioConnection | SSEConnection") -> None: +def __init__( + self, + transport: Transport, + connection: "StdioConnection | SSEConnection | StreamableHTTPConnection", +) -> None: self.transport = transport self.connection = connection self.tools = [] @@ -735,7 +741,11 @@ def __init__(self, transport: Transport, connection: "StdioConnection | SSEConne ### connection ```python -connection: StdioConnection | SSEConnection = connection +connection: ( + StdioConnection + | SSEConnection + | StreamableHTTPConnection +) = connection ``` Connection configuration diff --git a/docs/topics/tools.mdx b/docs/topics/tools.mdx index cf7319a0..778c4930 100644 --- a/docs/topics/tools.mdx +++ b/docs/topics/tools.mdx @@ -267,7 +267,7 @@ The [Model Context Protocol (MCP)](https://github.com/model-context-protocol/spe ### Using MCP Tools in Rigging -Use the `rigging.mcp` function, specifying the transport method (`stdio` or `sse`) and the connection parameters. It returns an *async context manager*. +Use the `rigging.mcp` function, specifying the transport method (`stdio`, `sse`, or `streamable-http`) and the connection parameters. It returns an *async context manager*. ```python Using "stdio" (Standard Input/Output) @@ -342,6 +342,40 @@ async with rg.mcp("stdio", command="claude", args=["mcp", "serve"]) as mcp_clien print(chat.conversation) ``` +### Keyless web search with Parallel + +[Parallel Search MCP](https://docs.parallel.ai/integrations/mcp/search-mcp) exposes +`web_search` and `web_fetch` over Streamable HTTP. The anonymous endpoint requires +no Parallel API key and provides free access at lower rate limits. Model inference +still uses your configured generator and its credentials. + +```python +import rigging as rg + +async with rg.mcp( + "streamable-http", + url="https://search.parallel.ai/mcp", + headers={"User-Agent": f"rigging/{rg.__version__}"}, +) as mcp_client: + chat = await ( + rg.get_generator("gpt-4o") + .chat("Search for Rigging's MCP documentation and summarize it with source URLs.") + .using(*mcp_client.tools) + .run() + ) + print(chat.conversation) +``` + +Keep tool calls inside the context manager so the connection stays open. The +`timeout` and `sse_read_timeout` options take seconds, defaulting to 5 and 300. + +For a search that does not require a model or any API keys, run the +[Parallel search example](https://github.com/dreadnode/rigging/blob/main/examples/parallel_search.py): + +```bash +python examples/parallel_search.py "Rigging MCP documentation" +``` + ### Serving Tools with MCP In addition to consuming tools from mcp servers, Rigging can **expose its own tools** as an MCP server. This allows you to write a singular tool implementation with the freedom to use it in other MCP-compliant clients, such as claude code. diff --git a/examples/parallel_search.py b/examples/parallel_search.py new file mode 100644 index 00000000..91fdcb6a --- /dev/null +++ b/examples/parallel_search.py @@ -0,0 +1,28 @@ +"""Search the web through Rigging's MCP tools without model or Parallel API keys. + +Run after installing Rigging from this checkout: + python examples/parallel_search.py "Rigging MCP documentation" +""" + +import argparse +import asyncio + +import rigging as rg + + +async def search(query: str) -> None: + async with rg.mcp( + "streamable-http", + url="https://search.parallel.ai/mcp", + headers={"User-Agent": f"rigging/{rg.__version__}"}, + ) as client: + tool = next(tool for tool in client.tools if tool.name == "web_search") + parts = await tool(objective=query, search_queries=[query]) + for part in parts: + print(part.text) + + +if __name__ == "__main__": + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("query", help="Web search query") + asyncio.run(search(parser.parse_args().query)) diff --git a/poetry.lock b/poetry.lock index aa1de253..77a4e35a 100644 --- a/poetry.lock +++ b/poetry.lock @@ -8539,4 +8539,4 @@ llm = ["accelerate", "transformers", "vllm"] [metadata] lock-version = "2.1" python-versions = ">=3.10,<3.14" -content-hash = "58dfb9bff0da998cf5bf028fed2b737c92a0ed4ce5ab5fad63b34c7d43df29e7" +content-hash = "931cdad90b925896134138f98f87b4425035ebe5fe28394fc708cd1ec81ef45f" diff --git a/pyproject.toml b/pyproject.toml index 688b915c..5d71852d 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -21,7 +21,7 @@ colorama = "^0.4.6" jsonpath-ng = "^1.7.0" ruamel-yaml = "^0.18.10" jsonref = "^1.1.0" -mcp = "^1.5.0" +mcp = "^1.10.0" dreadnode = ">=1.12.0" vllm = { version = "^0.5.0", optional = true } diff --git a/rigging/tools/mcp.py b/rigging/tools/mcp.py index 696e2b3f..261fa2dc 100644 --- a/rigging/tools/mcp.py +++ b/rigging/tools/mcp.py @@ -9,6 +9,7 @@ import typing as t from contextlib import AsyncExitStack from dataclasses import dataclass +from datetime import timedelta from pathlib import Path from types import TracebackType @@ -29,7 +30,7 @@ DEFAULT_HTTP_TIMEOUT = 5 DEFAULT_SSE_READ_TIMEOUT = 60 * 5 -Transport = t.Literal["stdio", "sse"] +Transport = t.Literal["stdio", "sse", "streamable-http"] P = t.ParamSpec("P") R = t.TypeVar("R") @@ -49,6 +50,10 @@ class SSEConnection(te.TypedDict): sse_read_timeout: float +class StreamableHTTPConnection(SSEConnection): + pass + + def _convert_mcp_result_to_message_parts(result: "CallToolResult") -> list[t.Any]: from mcp.types import TextResourceContents @@ -131,12 +136,16 @@ class MCPClient: transport: Transport """The transport to use""" - connection: StdioConnection | SSEConnection + connection: StdioConnection | SSEConnection | StreamableHTTPConnection """Connection configuration""" tools: list[Tool[..., t.Any]] """A list of tools available on the server""" - def __init__(self, transport: Transport, connection: "StdioConnection | SSEConnection") -> None: + def __init__( + self, + transport: Transport, + connection: "StdioConnection | SSEConnection | StreamableHTTPConnection", + ) -> None: self.transport = transport self.connection = connection self.tools = [] @@ -187,6 +196,24 @@ async def _connect_via_sse(self, connection: "SSEConnection") -> "ClientSession" read, write = await self._exit_stack.enter_async_context(client) return await self._exit_stack.enter_async_context(ClientSession(read, write)) + async def _connect_via_streamable_http( + self, + connection: "StreamableHTTPConnection", + ) -> "ClientSession": + from mcp import ClientSession + from mcp.client.streamable_http import streamablehttp_client + + client = streamablehttp_client( + url=connection["url"], + headers=connection.get("headers"), + timeout=timedelta(seconds=connection.get("timeout", DEFAULT_HTTP_TIMEOUT)), + sse_read_timeout=timedelta( + seconds=connection.get("sse_read_timeout", DEFAULT_SSE_READ_TIMEOUT), + ), + ) + read, write, _ = await self._exit_stack.enter_async_context(client) + return await self._exit_stack.enter_async_context(ClientSession(read, write)) + async def _shutdown(self) -> None: await self._exit_stack.aclose() self._session = None @@ -202,9 +229,13 @@ async def __aenter__(self) -> "MCPClient": self._session = await self._connect_via_sse( t.cast("SSEConnection", self.connection), ) + elif self.transport == "streamable-http": + self._session = await self._connect_via_streamable_http( + t.cast("StreamableHTTPConnection", self.connection), + ) else: raise TypeError( # noqa: TRY301 - f"Unsupported transport: {self.transport}. Must be 'stdio' or 'sse'", + f"Unsupported transport: {self.transport}. Must be 'stdio', 'sse' or 'streamable-http'", ) await asyncio.wait_for(self.session.initialize(), timeout=INITIALIZE_TIMEOUT) @@ -291,8 +322,33 @@ def mcp( """ +@t.overload +def mcp( + transport: t.Literal["streamable-http"], + *, + url: str, + headers: dict[str, str] | None = None, + timeout: float = DEFAULT_HTTP_TIMEOUT, + sse_read_timeout: float = DEFAULT_SSE_READ_TIMEOUT, +) -> MCPClient: + """ + Create an MCP client that communicates over Streamable HTTP. + + Args: + url: The URL of the MCP endpoint to connect to. + headers: HTTP headers sent on all requests, including tool calls. + timeout: HTTP timeout in seconds. + sse_read_timeout: Streaming read timeout in seconds. + + Returns: + An MCP client context manager. Discovered tools are available via tools. + """ + + def mcp(transport: Transport, **connection: t.Any) -> MCPClient: - return MCPClient(transport, t.cast("StdioConnection | SSEConnection", connection)) + return MCPClient( + transport, t.cast("StdioConnection | SSEConnection | StreamableHTTPConnection", connection) + ) def as_mcp( diff --git a/tests/test_mcp.py b/tests/test_mcp.py new file mode 100644 index 00000000..2972e606 --- /dev/null +++ b/tests/test_mcp.py @@ -0,0 +1,138 @@ +import importlib +from contextlib import asynccontextmanager +from datetime import timedelta +from types import SimpleNamespace +from unittest.mock import AsyncMock + +import pytest +from mcp.types import CallToolResult, ListToolsResult, TextContent, Tool + +import rigging as rg +from rigging.tools.base import FunctionCall, ToolCall +from rigging.tools.mcp import DEFAULT_HTTP_TIMEOUT, DEFAULT_SSE_READ_TIMEOUT + +# ruff: noqa: S101, SLF001, ARG001, PLR2004 + +mcp_module = importlib.import_module("rigging.tools.mcp") + + +@pytest.fixture +def http_server(monkeypatch: pytest.MonkeyPatch) -> SimpleNamespace: + state = SimpleNamespace(closed=False, options=None) + session = AsyncMock() + session.list_tools.return_value = ListToolsResult( + tools=[ + Tool( + name="web_search", + description="Search the web", + inputSchema={"type": "object", "properties": {"query": {"type": "string"}}}, + ), + ], + ) + session.call_tool.return_value = CallToolResult( + content=[TextContent(type="text", text="https://example.com: useful search excerpt")], + ) + state.session = session + + @asynccontextmanager + async def transport(**kwargs): + state.options = kwargs + try: + yield "read", "write", lambda: "session-id" + finally: + state.closed = True + + @asynccontextmanager + async def client_session(read, write): + assert (read, write) == ("read", "write") + yield session + + monkeypatch.setattr("mcp.client.streamable_http.streamablehttp_client", transport) + monkeypatch.setattr("mcp.ClientSession", client_session) + return state + + +async def test_streamable_http_tools(http_server: SimpleNamespace) -> None: + headers = {"User-Agent": "rigging/test"} + client = rg.mcp("streamable-http", url="https://example.com/mcp", headers=headers) + async with client: + assert http_server.options == { + "url": "https://example.com/mcp", + "headers": headers, + "timeout": timedelta(seconds=DEFAULT_HTTP_TIMEOUT), + "sse_read_timeout": timedelta(seconds=DEFAULT_SSE_READ_TIMEOUT), + } + http_server.session.initialize.assert_awaited_once() + tool = client.tools[0] + assert tool.name == "web_search" + assert tool.parameters_schema["properties"]["query"] == {"type": "string"} + message, stop = await tool.handle_tool_call( + ToolCall( + id="search", function=FunctionCall(name="web_search", arguments='{"query":"MCP"}') + ), + ) + assert "useful search excerpt" in message.content + assert message.tool_call_id == "search" + assert not stop + http_server.session.call_tool.assert_awaited_once_with("web_search", {"query": "MCP"}) + assert http_server.closed + assert client.tools == [] + with pytest.raises(RuntimeError, match="Session not initialized"): + _ = client.session + + +async def test_streamable_http_timeouts(http_server: SimpleNamespace) -> None: + async with rg.mcp( + "streamable-http", url="https://example.com/mcp", timeout=12, sse_read_timeout=45 + ): + assert http_server.options["timeout"] == timedelta(seconds=12) + assert http_server.options["sse_read_timeout"] == timedelta(seconds=45) + assert http_server.options["headers"] is None + + +@pytest.mark.parametrize("operation", ["initialize", "list_tools"]) +async def test_streamable_http_setup_failure(http_server: SimpleNamespace, operation: str) -> None: + getattr(http_server.session, operation).side_effect = RuntimeError("server failed") + client = rg.mcp("streamable-http", url="https://example.com/mcp") + with pytest.raises(RuntimeError, match="server failed"): + async with client: + pass + assert http_server.closed + assert client.tools == [] + assert client._session is None + + +@pytest.mark.parametrize("transport", ["stdio", "sse"]) +async def test_existing_transport_selection( + monkeypatch: pytest.MonkeyPatch, transport: str +) -> None: + session = AsyncMock() + session.list_tools.return_value = ListToolsResult(tools=[]) + connect = AsyncMock(return_value=session) + monkeypatch.setattr(mcp_module.MCPClient, f"_connect_via_{transport}", connect) + connection = ( + {"command": "server"} if transport == "stdio" else {"url": "https://example.com/sse"} + ) + client = mcp_module.MCPClient(transport, connection) + async with client: + connect.assert_awaited_once_with(connection) + session.initialize.assert_awaited_once() + + +async def test_unsupported_transport() -> None: + with pytest.raises(TypeError, match="streamable-http"): + async with mcp_module.MCPClient("unknown", {}): + pass + + +async def test_streamable_http_cancelled_tool(http_server: SimpleNamespace) -> None: + import asyncio + + http_server.session.call_tool.side_effect = asyncio.CancelledError + client = rg.mcp("streamable-http", url="https://example.com/mcp") + with pytest.raises(asyncio.CancelledError): + async with client: + await client.tools[0](query="MCP") + assert http_server.closed + assert client._session is None + assert client.tools == []