Add tests

Author: rectalogic

pyproject.toml

@@ -66,3 +66,7 @@ select = [
 
 [tool.ruff.lint.per-file-ignores]
 "tests/*" = ["S", "INP001"]
+
+[tool.pytest.ini_options]
+asyncio_mode = "auto"
+asyncio_default_fixture_loop_scope = "session"

tests/conftest.py

@@ -0,0 +1,50 @@
+# Copyright (C) 2024 Andrew Wason
+# SPDX-License-Identifier: MIT
+
+from unittest import mock
+
+import pytest
+from mcp import ClientSession, ListToolsResult, Tool
+from mcp.types import CallToolResult, TextContent
+
+from langchain_mcp import MCPToolkit
+
+
+@pytest.fixture(scope="class")
+def mcptoolkit(request):
+    session_mock = mock.AsyncMock(spec=ClientSession)
+    session_mock.list_tools.return_value = ListToolsResult(
+        tools=[
+            Tool(
+                name="read_file",
+                description=(
+                    "Read the complete contents of a file from the file system. Handles various text encodings "
+                    "and provides detailed error messages if the file cannot be read. "
+                    "Use this tool when you need to examine the contents of a single file. "
+                    "Only works within allowed directories."
+                ),
+                inputSchema={
+                    "type": "object",
+                    "properties": {"path": {"type": "string"}},
+                    "required": ["path"],
+                    "additionalProperties": False,
+                    "$schema": "http://json-schema.org/draft-07/schema#",
+                },
+            )
+        ]
+    )
+    session_mock.call_tool.return_value = CallToolResult(
+        content=[TextContent(type="text", text="MIT License\n\nCopyright (c) 2024 Andrew Wason\n")],
+        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"})
+
+
+@pytest.fixture(scope="class")
+async def mcptool(request, mcptoolkit):
+    tool = (await request.cls.toolkit.get_tools())[0]
+    request.cls.tool = tool
+    yield tool

tests/demo.py

@@ -0,0 +1,39 @@
+# Copyright (C) 2024 Andrew Wason
+# SPDX-License-Identifier: MIT
+
+import asyncio
+import pathlib
+
+from langchain_core.messages import HumanMessage
+from langchain_core.output_parsers import StrOutputParser
+from langchain_groq import ChatGroq
+from mcp import ClientSession, StdioServerParameters
+from mcp.client.stdio import stdio_client
+
+from langchain_mcp import MCPToolkit
+
+
+async def main():
+    model = ChatGroq(model="llama-3.1-8b-instant")  # requires GROQ_API_KEY
+    server_params = StdioServerParameters(
+        command="npx",
+        args=["-y", "@modelcontextprotocol/server-filesystem", str(pathlib.Path(__file__).parent.parent)],
+    )
+    async with stdio_client(server_params) as (read, write):
+        async with ClientSession(read, write) as session:
+            toolkit = MCPToolkit(session=session)
+            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.append(await tools_model.ainvoke(messages))
+            for tool_call in messages[-1].tool_calls:
+                selected_tool = tools_map[tool_call["name"].lower()]
+                tool_msg = await selected_tool.ainvoke(tool_call)
+                messages.append(tool_msg)
+            result = await (tools_model | StrOutputParser()).ainvoke(messages)
+            print(result)
+
+
+if __name__ == "__main__":
+    asyncio.run(main())

tests/integration_tests/__init__.py

@@ -0,0 +1,16 @@
+# Copyright (C) 2024 Andrew Wason
+# SPDX-License-Identifier: MIT
+
+import pytest
+from langchain_tests.unit_tests import ToolsUnitTests
+
+
+@pytest.mark.usefixtures("mcptool")
+class TestMCPToolIntegration(ToolsUnitTests):
+    @property
+    def tool_constructor(self):
+        return self.tool
+
+    @property
+    def tool_invoke_params_example(self) -> dict:
+        return {"path": "LICENSE"}

tests/integration_tests/test_tool.py

@@ -0,0 +1,16 @@
+# Copyright (C) 2024 Andrew Wason
+# SPDX-License-Identifier: MIT
+
+import pytest
+from langchain_tests.unit_tests import ToolsUnitTests
+
+
+@pytest.mark.usefixtures("mcptool")
+class TestMCPToolUnit(ToolsUnitTests):
+    @property
+    def tool_constructor(self):
+        return self.tool
+
+    @property
+    def tool_invoke_params_example(self) -> dict:
+        return {"path": "LICENSE"}