MCPToolkit that accepts an MCP ClientSession

Author: rectalogic

.github/workflows/publish.yml

@@ -20,7 +20,7 @@ jobs:
       uses: astral-sh/setup-uv@v2
       with:
         enable-cache: true
-        #cache-dependency-glob: "uv.lock"
+        cache-dependency-glob: "uv.lock"
     - name: Set up Python 3.12
       uses: actions/setup-python@v5
       with:

.vscode/settings.json

@@ -0,0 +1,9 @@
+{
+  "[python]": {
+    "editor.formatOnSave": true,
+    "editor.codeActionsOnSave": {
+      "source.organizeImports": "explicit"
+    },
+    "editor.defaultFormatter": "charliermarsh.ruff"
+  }
+}

LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2024 Andrew Wason
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

README.md

@@ -1,3 +1,3 @@
 # langchain-mcp
 
-Placeholder
+[Model Context Protocol](https://modelcontextprotocol.io) tool calling support in LangChain.

pyproject.toml

@@ -7,8 +7,62 @@ authors = [
     { name = "Andrew Wason", email = "[email protected]" }
 ]
 requires-python = ">=3.10"
-dependencies = []
+dependencies = [
+    "langchain-core~=0.3.21",
+    "mcp~=1.0.0",
+    "pydantic>=2.10.2",
+]
+classifiers = [
+    "License :: OSI Approved :: MIT License",
+]
 
 [build-system]
 requires = ["hatchling"]
 build-backend = "hatchling.build"
+
+[dependency-groups]
+dev = [
+    "langchain-tests~=0.3.4",
+    "pytest~=8.3.3",
+    "pytest-asyncio~=0.24.0",
+    "pytest-socket~=0.7.0",
+    "ruff~=0.8.0",
+]
+
+[project.urls]
+Repository = "https://github.com/rectalogic/langchain-mcp"
+Issues = "https://github.com/rectalogic/langchain-mcp/issues"
+Related = "https://modelcontextprotocol.io/"
+
+[tool.ruff]
+target-version = "py310"
+line-length = 120
+
+[tool.ruff.format]
+docstring-code-format = true
+
+[tool.ruff.lint]
+select = [
+    # flake8-2020
+    "YTT",
+    # flake8-bandit
+    "S",
+    # flake8-bugbear
+    "B",
+    # flake8-builtins
+    "A",
+    # Pyflakes
+    "F",
+    # Pycodestyle
+    "E",
+    "W",
+    # isort
+    "I",
+    # flake8-no-pep420
+    "INP",
+    # pyupgrade
+    "UP",
+]
+
+[tool.ruff.lint.per-file-ignores]
+"tests/*" = ["S", "INP001"]

src/langchain_mcp/__init__.py

@@ -0,0 +1,4 @@
+# Copyright (C) 2024 Andrew Wason
+# SPDX-License-Identifier: MIT
+
+from .toolkit import MCPToolkit  # noqa: F401

src/langchain_mcp/toolkit.py

@@ -0,0 +1,82 @@
+# Copyright (C) 2024 Andrew Wason
+# SPDX-License-Identifier: MIT
+
+import typing as t
+from collections.abc import Callable
+
+import pydantic
+import pydantic_core
+from langchain_core.tools.base import BaseTool, BaseToolkit, InjectedToolArg, ToolException
+from mcp import ClientSession
+
+
+class MCPToolkit(BaseToolkit):
+    """
+    MCP server toolkit
+    """
+
+    session: ClientSession
+    """The MCP session used to obtain the tools"""
+
+    _initialized: bool = False
+
+    model_config = pydantic.ConfigDict(arbitrary_types_allowed=True)
+
+    async def get_tools(self) -> list[BaseTool]:
+        if not self._initialized:
+            await self.session.initialize()
+            self._initialized = True
+
+        return [
+            MCPTool(
+                session=self.session,
+                name=tool.name,
+                description=tool.description,
+                args_schema=create_schema_model(tool.inputSchema),
+            )
+            # list_tools returns a PaginatedResult, but I don't see a way to pass the cursor to retrieve more tools
+            for tool in (await self.session.list_tools()).tools
+        ]
+
+
+def create_schema_model(schema: dict[str, t.Any]) -> type[pydantic.BaseModel]:
+    # Create a new model class that returns our JSON schema.
+    # LangChain requires a BaseModel class.
+    class Schema(pydantic.BaseModel):
+        model_config = pydantic.ConfigDict(extra="allow", arbitrary_types_allowed=True)
+
+        @classmethod
+        def model_json_schema(
+            cls,
+            by_alias: bool = True,
+            ref_template: str = pydantic.json_schema.DEFAULT_REF_TEMPLATE,
+            schema_generator: type[pydantic.json_schema.GenerateJsonSchema] = pydantic.json_schema.GenerateJsonSchema,
+            mode: pydantic.json_schema.JsonSchemaMode = "validation",
+        ) -> dict[str, t.Any]:
+            return schema
+
+    return Schema
+
+
+class MCPTool(BaseTool):
+    """
+    MCP server tool
+    """
+
+    session: ClientSession
+
+    handle_tool_error: bool | str | Callable[[ToolException], str] | None = True
+
+    def _run(self, *args: t.Any, **kwargs: t.Any) -> t.Any:
+        raise NotImplementedError("Must invoke tool asynchronously")
+
+    async def _arun(self, *args: t.Any, **kwargs: t.Any) -> t.Any:
+        result = await self.session.call_tool(self.name, arguments=kwargs)
+        content = pydantic_core.to_json(result.content).decode()
+        if result.isError:
+            raise ToolException(content)
+        return content
+
+    @property
+    def tool_call_schema(self) -> type[pydantic.BaseModel]:
+        return self.args_schema