support loading text prompts (#12)

Author: mr-brobot

langchain_mcp_adapters/client.py

@@ -1,12 +1,14 @@
 from contextlib import AsyncExitStack
 from types import TracebackType
-from typing import Literal, TypedDict, cast
+from typing import Any, Literal, Optional, TypedDict, cast
 
+from langchain_core.messages import AIMessage, HumanMessage
 from langchain_core.tools import BaseTool
 from mcp import ClientSession, StdioServerParameters
 from mcp.client.sse import sse_client
 from mcp.client.stdio import stdio_client
 
+from langchain_mcp_adapters.prompts import load_mcp_prompt
 from langchain_mcp_adapters.tools import load_mcp_tools
 
 DEFAULT_ENCODING = "utf-8"
@@ -211,6 +213,13 @@ def get_tools(self) -> list[BaseTool]:
             all_tools.extend(server_tools)
         return all_tools
 
+    async def get_prompt(
+        self, server_name: str, prompt_name: str, arguments: Optional[dict[str, Any]]
+    ) -> list[HumanMessage | AIMessage]:
+        """Get a prompt from a given MCP server."""
+        session = self.sessions[server_name]
+        return await load_mcp_prompt(session, prompt_name, arguments)
+
     async def __aenter__(self) -> "MultiServerMCPClient":
         try:
             connections = self.connections or {}

langchain_mcp_adapters/prompts.py

@@ -0,0 +1,37 @@
+from typing import Any, Optional
+
+from langchain_core.messages import AIMessage, HumanMessage
+from mcp import ClientSession
+from mcp.types import PromptMessage
+
+
+def convert_mcp_prompt_message_to_langchain_message(
+    message: PromptMessage,
+) -> HumanMessage | AIMessage:
+    """Convert an MCP prompt message to a LangChain message.
+
+    Args:
+        message: MCP prompt message to convert
+
+    Returns:
+        a LangChain message
+    """
+    if message.content.type == "text":
+        if message.role == "user":
+            return HumanMessage(content=message.content.text)
+        elif message.role == "assistant":
+            return AIMessage(content=message.content.text)
+        else:
+            raise ValueError(f"Unsupported prompt message role: {message.role}")
+
+    raise ValueError(f"Unsupported prompt message content type: {message.content.type}")
+
+
+async def load_mcp_prompt(
+    session: ClientSession, name: str, arguments: Optional[dict[str, Any]] = None
+) -> list[HumanMessage | AIMessage]:
+    """Load MCP prompt and convert to LangChain messages."""
+    response = await session.get_prompt(name, arguments)
+    return [
+        convert_mcp_prompt_message_to_langchain_message(message) for message in response.messages
+    ]

pyproject.toml

@@ -7,7 +7,7 @@ name = "langchain-mcp-adapters"
 version = "0.0.4"
 description = "Make Anthropic Model Context Protocol (MCP) tools compatible with LangChain and LangGraph agents."
 authors = [
-    {name = "Vadym Barda", email = "[email protected] "}
+    { name = "Vadym Barda", email = "[email protected]" },
 ]
 readme = "README.md"
 requires-python = ">=3.10"
@@ -22,15 +22,14 @@ test = [
     "ruff>=0.9.4",
     "mypy>=1.8.0",
     "pytest-socket>=0.7.0",
+    "pytest-asyncio>=0.25.0",
     "types-setuptools>=69.0.0",
 ]
 
 [tool.pytest.ini_options]
 minversion = "8.0"
 addopts = "-ra -q -v"
-testpaths = [
-    "tests",
-]
+testpaths = ["tests"]
 python_files = ["test_*.py"]
 python_functions = ["test_*"]
 
@@ -40,14 +39,14 @@ target-version = "py310"
 
 [tool.ruff.lint]
 select = [
-    "E",  # pycodestyle errors
-    "W",  # pycodestyle warnings
-    "F",  # pyflakes
-    "I",  # isort
-    "B",  # flake8-bugbear
+    "E", # pycodestyle errors
+    "W", # pycodestyle warnings
+    "F", # pyflakes
+    "I", # isort
+    "B", # flake8-bugbear
 ]
 ignore = [
-  "E501" # line-length
+    "E501", # line-length
 ]
 
 
@@ -57,4 +56,3 @@ warn_return_any = true
 warn_unused_configs = true
 disallow_untyped_defs = true
 check_untyped_defs = true
-

tests/test_import.py

@@ -1,3 +1,3 @@
 def test_import() -> None:
     """Test that the code can be imported"""
-    from langchain_mcp_adapters import client, tools  # noqa: F401
+    from langchain_mcp_adapters import client, prompts, tools  # noqa: F401

tests/test_prompts.py

@@ -0,0 +1,75 @@
+from unittest.mock import AsyncMock
+
+import pytest
+from langchain_core.messages import AIMessage, HumanMessage
+from mcp.types import (
+    EmbeddedResource,
+    ImageContent,
+    PromptMessage,
+    TextContent,
+    TextResourceContents,
+)
+
+from langchain_mcp_adapters.prompts import (
+    convert_mcp_prompt_message_to_langchain_message,
+    load_mcp_prompt,
+)
+
+
+@pytest.mark.parametrize(
+    "role,text,expected_cls",
+    [
+        ("assistant", "Hello", AIMessage),
+        ("user", "Hello", HumanMessage),
+    ],
+)
+def test_convert_mcp_prompt_message_to_langchain_message_with_text_content(
+    role: str, text: str, expected_cls: type
+):
+    message = PromptMessage(role=role, content=TextContent(type="text", text=text))
+    result = convert_mcp_prompt_message_to_langchain_message(message)
+    assert isinstance(result, expected_cls)
+    assert result.content == text
+
+
+@pytest.mark.parametrize("role", ["assistant", "user"])
+def test_convert_mcp_prompt_message_to_langchain_message_with_resource_content(role: str):
+    message = PromptMessage(
+        role=role,
+        content=EmbeddedResource(
+            type="resource",
+            resource=TextResourceContents(
+                uri="message://greeting", mimeType="text/plain", text="hi"
+            ),
+        ),
+    )
+    with pytest.raises(ValueError):
+        convert_mcp_prompt_message_to_langchain_message(message)
+
+
+@pytest.mark.parametrize("role", ["assistant", "user"])
+def test_convert_mcp_prompt_message_to_langchain_message_with_image_content(role: str):
+    message = PromptMessage(
+        role=role, content=ImageContent(type="image", mimeType="image/png", data="base64data")
+    )
+    with pytest.raises(ValueError):
+        convert_mcp_prompt_message_to_langchain_message(message)
+
+
+@pytest.mark.asyncio
+async def test_load_mcp_prompt():
+    session = AsyncMock()
+    session.get_prompt = AsyncMock(
+        return_value=AsyncMock(
+            messages=[
+                PromptMessage(role="user", content=TextContent(type="text", text="Hello")),
+                PromptMessage(role="assistant", content=TextContent(type="text", text="Hi")),
+            ]
+        )
+    )
+    result = await load_mcp_prompt(session, "test_prompt")
+    assert len(result) == 2
+    assert isinstance(result[0], HumanMessage)
+    assert result[0].content == "Hello"
+    assert isinstance(result[1], AIMessage)
+    assert result[1].content == "Hi"