mypy

Author: rectalogic

.github/workflows/test.yml

@@ -28,4 +28,5 @@ jobs:
     - name: Lint
       run: |
         uv run --python ${{ matrix.python-version }} --python-preference only-system ruff format --check
-        uv run --python ${{ matrix.python-version }} --python-preference only-system ruff check 
+        uv run --python ${{ matrix.python-version }} --python-preference only-system ruff check
+        uv run --python ${{ matrix.python-version }} --python-preference only-system mypy 

pyproject.toml

@@ -27,6 +27,8 @@ dev = [
     "pytest-asyncio~=0.24.0",
     "pytest-socket~=0.7.0",
     "ruff~=0.8.0",
+    "mypy~=1.13.0",
+    "typing-extensions~=4.12.2",
 ]
 
 [project.urls]
@@ -70,3 +72,13 @@ select = [
 [tool.pytest.ini_options]
 asyncio_mode = "auto"
 asyncio_default_fixture_loop_scope = "class"
+
+[tool.mypy]
+disallow_untyped_defs = true
+warn_unused_configs = true
+warn_redundant_casts = true
+warn_unused_ignores = true
+strict_equality = true
+no_implicit_optional = true
+show_error_codes = true
+files = "src/**/*.py"

src/langchain_mcp/toolkit.py

@@ -2,12 +2,12 @@
 # SPDX-License-Identifier: MIT
 
 import asyncio
-import typing as t
 import warnings
 from collections.abc import Callable
 
 import pydantic
 import pydantic_core
+import typing_extensions as t
 from langchain_core.tools.base import BaseTool, BaseToolkit, ToolException
 from mcp import ClientSession
 
@@ -24,7 +24,8 @@ class MCPToolkit(BaseToolkit):
 
     model_config = pydantic.ConfigDict(arbitrary_types_allowed=True)
 
-    async def get_tools(self) -> list[BaseTool]:
+    @t.override
+    async def get_tools(self) -> list[BaseTool]:  # type: ignore[override]
         if not self._initialized:
             await self.session.initialize()
             self._initialized = True
@@ -33,7 +34,7 @@ async def get_tools(self) -> list[BaseTool]:
             MCPTool(
                 session=self.session,
                 name=tool.name,
-                description=tool.description,
+                description=tool.description or "",
                 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
@@ -47,6 +48,7 @@ def create_schema_model(schema: dict[str, t.Any]) -> type[pydantic.BaseModel]:
     class Schema(pydantic.BaseModel):
         model_config = pydantic.ConfigDict(extra="allow", arbitrary_types_allowed=True)
 
+        @t.override
         @classmethod
         def model_json_schema(
             cls,
@@ -66,22 +68,25 @@ class MCPTool(BaseTool):
     """
 
     session: ClientSession
-
+    args_schema: type[pydantic.BaseModel]
     handle_tool_error: bool | str | Callable[[ToolException], str] | None = True
 
+    @t.override
     def _run(self, *args: t.Any, **kwargs: t.Any) -> t.Any:
         warnings.warn(
             "Invoke this tool asynchronousely using `ainvoke`. This method exists only to satisfy tests.", stacklevel=1
         )
         return asyncio.run(self._arun(*args, **kwargs))
 
+    @t.override
     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
 
+    @t.override
     @property
     def tool_call_schema(self) -> type[pydantic.BaseModel]:
         return self.args_schema

tests/conftest.py

@@ -4,6 +4,7 @@
 from unittest import mock
 
 import pytest
+from langchain_tests.integration_tests import ToolsIntegrationTests
 from mcp import ClientSession, ListToolsResult, Tool
 from mcp.types import CallToolResult, TextContent
 
@@ -38,13 +39,13 @@ def mcptoolkit(request):
         isError=False,
     )
     toolkit = MCPToolkit(session=session_mock)
-    request.cls.toolkit = toolkit
     yield toolkit
-    # session_mock.call_tool.assert_called_with("read_file", arguments={"path": "LICENSE"})
+    if issubclass(request.cls, ToolsIntegrationTests):
+        session_mock.call_tool.assert_called_with("read_file", arguments={"path": "LICENSE"})
 
 
 @pytest.fixture(scope="class")
 async def mcptool(request, mcptoolkit):
-    tool = (await request.cls.toolkit.get_tools())[0]
+    tool = (await mcptoolkit.get_tools())[0]
     request.cls.tool = tool
     yield tool

tests/demo.py

@@ -1,8 +1,18 @@
 # Copyright (C) 2024 Andrew Wason
 # SPDX-License-Identifier: MIT
 
+# /// script
+# requires-python = ">=3.10"
+# dependencies = [
+#     "langchain-mcp",
+#     "langchain-groq",
+# ]
+# ///
+
+
 import asyncio
 import pathlib
+import sys
 
 from langchain_core.messages import HumanMessage
 from langchain_core.output_parsers import StrOutputParser
@@ -13,7 +23,7 @@
 from langchain_mcp import MCPToolkit
 
 
-async def main():
+async def main(prompt: str) -> None:
     model = ChatGroq(model="llama-3.1-8b-instant")  # requires GROQ_API_KEY
     server_params = StdioServerParameters(
         command="npx",
@@ -25,7 +35,7 @@ async def main():
             tools = await toolkit.get_tools()
             tools_map = {tool.name: tool for tool in tools}
             tools_model = model.bind_tools(tools)
-            messages = [HumanMessage("Read and summarize the file ./LICENSE")]
+            messages = [HumanMessage(prompt)]
             messages.append(await tools_model.ainvoke(messages))
             for tool_call in messages[-1].tool_calls:
                 selected_tool = tools_map[tool_call["name"].lower()]
@@ -36,4 +46,5 @@ async def main():
 
 
 if __name__ == "__main__":
-    asyncio.run(main())
+    prompt = sys.argv[1] if len(sys.argv) > 1 else "Read and summarize the file ./LICENSE"
+    asyncio.run(main(prompt))