You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@flink.apache.org by xt...@apache.org on 2025/09/17 09:41:17 UTC
(flink-agents) 04/04: [api][python] Add built-in mcp support.
This is an automated email from the ASF dual-hosted git repository.
xtsong pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/flink-agents.git
commit 62ff96030337f062f6804a89d5e79a8a4dadbeb8
Author: WenjinXie <we...@gmail.com>
AuthorDate: Fri Sep 12 14:58:44 2025 +0800
[api][python] Add built-in mcp support.
Co-authored-by: Hao Li <11...@users.noreply.github.com>
Co-authored-by: yanand0909 <ya...@confluent.io>
---
python/flink_agents/api/agent.py | 23 ++
python/flink_agents/api/decorators.py | 17 ++
python/flink_agents/api/execution_environment.py | 22 ++
python/flink_agents/api/resource.py | 2 +-
python/flink_agents/api/tests/mcp/__init__.py | 17 ++
python/flink_agents/api/tests/mcp/mcp_server.py | 47 ++++
python/flink_agents/api/tests/mcp/test_mcp.py | 56 ++++
python/flink_agents/api/tests/test_tool.py | 79 ++++++
python/flink_agents/api/tools/mcp.py | 309 +++++++++++++++++++++++
python/flink_agents/api/tools/utils.py | 51 +++-
python/flink_agents/plan/agent_plan.py | 44 +++-
python/pyproject.toml | 1 +
12 files changed, 663 insertions(+), 5 deletions(-)
diff --git a/python/flink_agents/api/agent.py b/python/flink_agents/api/agent.py
index cc9a3c6..6dc7cc9 100644
--- a/python/flink_agents/api/agent.py
+++ b/python/flink_agents/api/agent.py
@@ -29,6 +29,7 @@ from flink_agents.api.embedding_models.embedding_model import (
from flink_agents.api.events.event import Event
from flink_agents.api.prompts.prompt import Prompt
from flink_agents.api.resource import ResourceType
+from flink_agents.api.tools.mcp import MCPServer
class Agent(ABC):
@@ -77,6 +78,7 @@ class Agent(ABC):
_actions: Dict[str, Tuple[List[Type[Event]], Callable, Dict[str, Any]]]
_resources: Dict[ResourceType, Dict[str, Any]]
+ _mcp_servers: Dict[str, MCPServer]
def __init__(self) -> None:
"""Init method."""
@@ -280,3 +282,24 @@ class Agent(ABC):
kwargs["name"] = name
self._resources[ResourceType.EMBEDDING_MODEL][name] = (embedding_model, kwargs)
return self
+
+ def add_mcp_server(self, name: str, mcp_server: MCPServer) -> "Agent":
+ """Add an MCP server to the agent.
+
+ Parameters
+ ----------
+ name : str
+ The name of the MCP server, should be unique in the same Agent.
+ mcp_server : MCPServer
+ The MCP server resource instance.
+
+ Returns:
+ -------
+ Agent
+ The modified Agent instance.
+ """
+ if name in self._resources[ResourceType.MCP_SERVER]:
+ msg = f"MCP server {name} already defined"
+ raise ValueError(msg)
+ self._resources[ResourceType.MCP_SERVER][name] = mcp_server
+ return self
diff --git a/python/flink_agents/api/decorators.py b/python/flink_agents/api/decorators.py
index ef605a3..2889d0f 100644
--- a/python/flink_agents/api/decorators.py
+++ b/python/flink_agents/api/decorators.py
@@ -154,3 +154,20 @@ def prompt(func: Callable) -> Callable:
"""
func._is_prompt = True
return func
+
+
+def mcp_server(func: Callable) -> Callable:
+ """Decorator for marking a function declaring a MCP server.
+
+ Parameters
+ ----------
+ func : Callable
+ Function to be decorated.
+
+ Returns:
+ -------
+ Callable
+ Decorator function that marks the target function declare a MCP server.
+ """
+ func._is_mcp_server = True
+ return func
diff --git a/python/flink_agents/api/execution_environment.py b/python/flink_agents/api/execution_environment.py
index 80ee403..a46d04e 100644
--- a/python/flink_agents/api/execution_environment.py
+++ b/python/flink_agents/api/execution_environment.py
@@ -35,6 +35,7 @@ from flink_agents.api.embedding_models.embedding_model import (
)
from flink_agents.api.prompts.prompt import Prompt
from flink_agents.api.resource import ResourceType
+from flink_agents.api.tools.mcp import MCPServer
class AgentBuilder(ABC):
@@ -352,3 +353,24 @@ class AgentsExecutionEnvironment(ABC):
kwargs["name"] = name
self._resources[ResourceType.EMBEDDING_MODEL][name] = (embedding_model, kwargs)
return self
+
+ def add_mcp_server(self, name: str, mcp_server: MCPServer) -> "AgentsExecutionEnvironment":
+ """Add an MCP server to the agent execution environment.
+
+ Parameters
+ ----------
+ name : str
+ The name of the MCP server, should be unique in the same Agent.
+ mcp_server : MCPServer
+ The MCP server resource instance.
+
+ Returns:
+ -------
+ AgentsExecutionEnvironment
+ The environment contains registered embedding model setup.
+ """
+ if name in self._resources[ResourceType.MCP_SERVER]:
+ msg = f"MCP server {name} already defined"
+ raise ValueError(msg)
+ self._resources[ResourceType.MCP_SERVER][name] = mcp_server
+ return self
diff --git a/python/flink_agents/api/resource.py b/python/flink_agents/api/resource.py
index 0f8fe61..c599641 100644
--- a/python/flink_agents/api/resource.py
+++ b/python/flink_agents/api/resource.py
@@ -35,7 +35,7 @@ class ResourceType(Enum):
EMBEDDING_MODEL_CONNECTION = "embedding_model_connection"
PROMPT = "prompt"
# VECTOR_STORE = "vector_store"
- # MCP_SERVER = "mcp_server"
+ MCP_SERVER = "mcp_server"
class Resource(BaseModel, ABC):
diff --git a/python/flink_agents/api/tests/mcp/__init__.py b/python/flink_agents/api/tests/mcp/__init__.py
new file mode 100644
index 0000000..e154fad
--- /dev/null
+++ b/python/flink_agents/api/tests/mcp/__init__.py
@@ -0,0 +1,17 @@
+################################################################################
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#################################################################################
diff --git a/python/flink_agents/api/tests/mcp/mcp_server.py b/python/flink_agents/api/tests/mcp/mcp_server.py
new file mode 100644
index 0000000..4cba1ed
--- /dev/null
+++ b/python/flink_agents/api/tests/mcp/mcp_server.py
@@ -0,0 +1,47 @@
+################################################################################
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#################################################################################
+
+import dotenv
+from mcp.server.fastmcp import FastMCP
+
+dotenv.load_dotenv()
+
+# Create MCP server
+mcp = FastMCP("BasicServer")
+
+
+@mcp.prompt()
+def ask_sum(a: int, b: int) -> str:
+ """Prompt of add tool."""
+ return f"Can you please calculate the sum of {a} and {b}?"
+
+@mcp.tool()
+async def add(a: int, b: int) -> int:
+ """Get the detailed information of a specified IP address.
+
+ Args:
+ a: The first operand.
+ b: The second operand.
+
+ Returns:
+ int: The sum of a and b.
+ """
+ return a + b
+
+mcp.run("streamable-http")
+
diff --git a/python/flink_agents/api/tests/mcp/test_mcp.py b/python/flink_agents/api/tests/mcp/test_mcp.py
new file mode 100644
index 0000000..893e309
--- /dev/null
+++ b/python/flink_agents/api/tests/mcp/test_mcp.py
@@ -0,0 +1,56 @@
+################################################################################
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#################################################################################
+import multiprocessing
+import runpy
+import time
+from pathlib import Path
+
+from flink_agents.api.chat_message import ChatMessage, MessageRole
+from flink_agents.api.tools.mcp import MCPServer
+
+
+def run_server() -> None: # noqa : D103
+ runpy.run_path(f"{current_dir}/mcp_server.py")
+
+current_dir = Path(__file__).parent
+def test_mcp() -> None: # noqa : D103
+ process = multiprocessing.Process(target=run_server)
+ process.start()
+ time.sleep(5)
+
+ mcp_server = MCPServer(name="mcp_server", endpoint="http://127.0.0.1:8000/mcp")
+ prompts = mcp_server.list_prompts()
+ assert len(prompts) == 1
+ prompt = prompts[0]
+ assert prompt.name == "ask_sum"
+ message = prompt.format_messages(role=MessageRole.SYSTEM, a="1", b="2")
+ assert [ChatMessage(
+ role=MessageRole.USER,
+ content="Can you please calculate the sum of 1 and 2?",
+ )] == message
+ tools = mcp_server.list_tools()
+ assert len(tools) == 1
+ tool = tools[0]
+ assert tool.name == "add"
+
+ process.kill()
+
+
+
+
+
diff --git a/python/flink_agents/api/tests/test_tool.py b/python/flink_agents/api/tests/test_tool.py
index c8d36e7..b06effd 100644
--- a/python/flink_agents/api/tests/test_tool.py
+++ b/python/flink_agents/api/tests/test_tool.py
@@ -15,14 +15,26 @@
# See the License for the specific language governing permissions and
# limitations under the License.
#################################################################################
+from __future__ import annotations
+
import json
+from datetime import timedelta
from pathlib import Path
+from typing import TYPE_CHECKING
+from urllib.parse import parse_qs, urlparse
import pytest
+from mcp.client.auth import OAuthClientProvider, TokenStorage
+from mcp.shared.auth import OAuthClientMetadata
+from pydantic import AnyUrl
+from flink_agents.api.tools.mcp import MCPServer
from flink_agents.api.tools.tool import ToolMetadata
from flink_agents.api.tools.utils import create_schema_from_function
+if TYPE_CHECKING:
+ from mcp.shared.auth import OAuthClientInformationFull, OAuthToken
+
current_dir = Path(__file__).parent
@@ -67,3 +79,70 @@ def test_deserialize_tool_metadata(tool_metadata: ToolMetadata) -> None: # noqa
expected_json = f.read()
actual_tool_metadata = tool_metadata.model_validate_json(expected_json)
assert actual_tool_metadata == tool_metadata
+
+
+class InMemoryTokenStorage(TokenStorage):
+ """Demo In-memory token storage implementation."""
+
+ def __init__(self) -> None: # noqa:D107
+ self.tokens: OAuthToken | None = None
+ self.client_info: OAuthClientInformationFull | None = None
+
+ async def get_tokens(self) -> OAuthToken | None:
+ """Get stored tokens."""
+ return self.tokens
+
+ async def set_tokens(self, tokens: OAuthToken) -> None:
+ """Store tokens."""
+ self.tokens = tokens
+
+ async def get_client_info(self) -> OAuthClientInformationFull | None:
+ """Get stored client information."""
+ return self.client_info
+
+ async def set_client_info(self, client_info: OAuthClientInformationFull) -> None:
+ """Store client information."""
+ self.client_info = client_info
+
+
+async def handle_redirect(auth_url: str) -> None: # noqa:D103
+ print(f"Visit: {auth_url}")
+
+
+async def handle_callback() -> tuple[str, str | None]: # noqa:D103
+ callback_url = input("Paste callback URL: ")
+ params = parse_qs(urlparse(callback_url).query)
+ return params["code"][0], params.get("state", [None])[0]
+
+
+def test_serialize_mcp_server() -> None: # noqa:D103
+ oauth_auth = OAuthClientProvider(
+ server_url="http://localhost:8001",
+ client_metadata=OAuthClientMetadata(
+ client_name="Example MCP Client",
+ redirect_uris=[AnyUrl("http://localhost:3000/callback")],
+ grant_types=["authorization_code", "refresh_token"],
+ response_types=["code"],
+ scope="user",
+ ),
+ storage=InMemoryTokenStorage(),
+ redirect_handler=handle_redirect,
+ callback_handler=handle_callback,
+ )
+ mcp_server = MCPServer(
+ name="mcp",
+ endpoint="http://localhost:8080",
+ auth=oauth_auth,
+ timeout=timedelta(seconds=5),
+ )
+ data = mcp_server.model_dump_json(serialize_as_any=True)
+
+ deserialized = mcp_server.model_validate_json(data)
+ assert deserialized.name == mcp_server.name
+ assert deserialized.endpoint == mcp_server.endpoint
+ assert deserialized.timeout == mcp_server.timeout
+ assert deserialized.auth.context.server_url == mcp_server.auth.context.server_url
+ assert (
+ deserialized.auth.context.client_metadata
+ == mcp_server.auth.context.client_metadata
+ )
diff --git a/python/flink_agents/api/tools/mcp.py b/python/flink_agents/api/tools/mcp.py
new file mode 100644
index 0000000..47983fa
--- /dev/null
+++ b/python/flink_agents/api/tools/mcp.py
@@ -0,0 +1,309 @@
+################################################################################
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#################################################################################
+import asyncio
+import base64
+from abc import ABC
+from contextlib import AsyncExitStack, asynccontextmanager
+from datetime import timedelta
+from typing import Any, AsyncIterator, Dict, List
+from urllib.parse import urlparse
+
+import cloudpickle
+import httpx
+from mcp.client.session import ClientSession
+from mcp.client.streamable_http import streamablehttp_client
+from mcp.types import PromptArgument, TextContent
+from pydantic import (
+ ConfigDict,
+ Field,
+ field_serializer,
+ model_validator,
+)
+from typing_extensions import override
+
+from flink_agents.api.chat_message import ChatMessage, MessageRole
+from flink_agents.api.prompts.prompt import Prompt
+from flink_agents.api.resource import ResourceType, SerializableResource
+from flink_agents.api.tools.tool import BaseTool, ToolMetadata, ToolType
+from flink_agents.api.tools.utils import extract_mcp_content_item
+
+
+class MCPTool(BaseTool):
+ """MCP tool definition that can be called directly.
+
+ This represents a single tool from an MCP server.
+ """
+
+ mcp_server: "MCPServer" = Field(default=None, exclude=True)
+
+ @classmethod
+ @override
+ def tool_type(cls) -> ToolType:
+ return ToolType.MCP
+
+ @override
+ def call(self, *args: Any, **kwargs: Any) -> Any:
+ """Call the MCP tool with the given arguments."""
+ if self.mcp_server is None:
+ msg = "MCP tool call requires a reference to the MCP server"
+ raise ValueError(msg)
+
+ return asyncio.run(
+ self.mcp_server.call_tool_async(self.metadata.name, *args, **kwargs)
+ )
+
+
+class MCPPrompt(Prompt):
+ """MCP prompt definition that extends the base Prompt class.
+
+ This represents a prompt managed by an MCP server.
+ """
+
+ title: str | None
+ description: str | None = None
+ prompt_arguments: list[PromptArgument] = Field(default_factory=list)
+ mcp_server: "MCPServer" = Field(default=None, exclude=True)
+
+ def _check_arguments(self, **kwargs: str) -> Dict[str, str]:
+ if self.mcp_server is None:
+ msg = "MCP prompt requires a reference to the MCP server"
+ raise ValueError(msg)
+
+ arguments = {}
+ for argument in self.prompt_arguments:
+ if argument.required:
+ if argument.name in kwargs:
+ arguments[argument.name] = kwargs[argument.name]
+ else:
+ msg = f"Missing required argument {argument.name}"
+ raise RuntimeError(msg)
+ return arguments
+
+ @override
+ def format_string(self, **arguments: str) -> str:
+ """Get the prompt from the MCP server with optional arguments..
+
+ Returns a text representation of the prompt.
+ """
+ text = "\n".join(
+ message.content for message in self.format_messages(**arguments)
+ )
+ return text
+
+ @override
+ def format_messages(
+ self, role: MessageRole = MessageRole.SYSTEM, **arguments: str
+ ) -> List[ChatMessage]:
+ """Get the prompt from the MCP server with optional arguments.
+
+ Returns a list of messages formatted according to the MCP server.
+ """
+ arguments = self._check_arguments(**arguments)
+ return self.mcp_server.get_prompt(self.name, arguments)
+
+
+class MCPServer(SerializableResource, ABC):
+ """Resource representing an MCP server and exposing its tools/prompts.
+
+ This is a logical container for MCP tools and prompts; it is not directly invokable.
+ """
+
+ model_config = ConfigDict(arbitrary_types_allowed=True)
+
+ # HTTP connection parameters
+ endpoint: str
+ headers: Dict[str, Any] | None = None
+ timeout: timedelta = timedelta(seconds=30)
+ sse_read_timeout: timedelta = timedelta(seconds=60 * 5)
+ auth: httpx.Auth | None = None
+
+ session: ClientSession = Field(default=None, exclude=True)
+ connection_context: AsyncExitStack = Field(default=None, exclude=True)
+
+ @field_serializer("auth")
+ def __serialize_auth(self, auth: httpx.Auth | None) -> str | None:
+ if auth is None:
+ return auth
+ return base64.b64encode(cloudpickle.dumps(auth)).decode()
+
+ @model_validator(mode="before")
+ def __custom_deserialize(self) -> "MCPServer":
+ if "auth" in self:
+ auth = self["auth"]
+ if auth and isinstance(auth, str):
+ self["auth"] = cloudpickle.loads(base64.b64decode(auth))
+ return self
+
+ @classmethod
+ @override
+ def resource_type(cls) -> ResourceType:
+ return ResourceType.MCP_SERVER
+
+ def _is_valid_http_url(self) -> bool:
+ """Check if the endpoint is a valid HTTP URL."""
+ try:
+ parsed = urlparse(self.endpoint)
+ return parsed.scheme in ("http", "https") and bool(parsed.netloc)
+ except Exception:
+ return False
+
+ @asynccontextmanager
+ async def _get_session(self) -> AsyncIterator[ClientSession]:
+ """Get or create an MCP client session using the official SDK's streamable
+ HTTP transport.
+ """
+ if not self._is_valid_http_url():
+ msg = f"Invalid HTTP endpoint: {self.endpoint}"
+ raise ValueError(msg)
+
+ async with streamablehttp_client(
+ url=self.endpoint,
+ headers=self.headers,
+ timeout=self.timeout,
+ sse_read_timeout=self.sse_read_timeout,
+ auth=self.auth,
+ ) as (read, write, _), ClientSession(
+ read,
+ write
+ ) as session:
+ await session.initialize()
+ yield session
+
+
+ async def _cleanup_connection(self) -> None:
+ """Clean up connection resources."""
+ try:
+ if self.connection_context is not None:
+ await self.connection_context.aclose()
+ self.connection_context = None
+ self.session = None
+ except Exception:
+ pass
+
+ async def call_tool_async(self, tool_name: str, *args: Any, **kwargs: Any) -> Any:
+ """Call a tool on the MCP server asynchronously."""
+ async with self._get_session() as session:
+
+ arguments = kwargs if kwargs else (args[0] if args else {})
+
+ result = await session.call_tool(
+ tool_name, arguments, read_timeout_seconds=self.timeout
+ )
+
+ content = [extract_mcp_content_item(item) for item in result.content]
+
+ return content
+
+ def list_tools(self) -> List[MCPTool]:
+ """List available tool definitions from the MCP server."""
+ return asyncio.run(self._list_tools_async())
+
+ async def _list_tools_async(self) -> List[MCPTool]:
+ """Async implementation of list_tools."""
+ async with self._get_session() as session:
+ tools_response = await session.list_tools()
+
+ tools = []
+ for tool_data in tools_response.tools or []:
+ metadata = ToolMetadata(
+ name=tool_data.name,
+ description=tool_data.description or "",
+ args_schema=tool_data.inputSchema
+ or {"type": "object", "properties": {}},
+ )
+
+ tool = MCPTool(
+ name=tool_data.name,
+ metadata=metadata,
+ mcp_server=self,
+ )
+ tools.append(tool)
+
+ return tools
+
+ def get_tool(self, name: str) -> MCPTool:
+ """Get a single callable tool by name."""
+ return asyncio.run(self._get_tool_async(name))
+
+ async def _get_tool_async(self, name: str) -> MCPTool:
+ """Async implementation of get_tool."""
+ tools = await self._list_tools_async()
+ for tool in tools:
+ if tool.metadata.name == name:
+ return tool
+ msg = f"Tool '{name}' not found on MCP server at {self.endpoint}"
+ raise ValueError(msg)
+
+ def get_tool_metadata(self, name: str) -> ToolMetadata:
+ """Get a single tool definition metadata by name."""
+ tool = self.get_tool(name)
+ return tool.metadata
+
+ def list_prompts(self) -> List[MCPPrompt]:
+ """List available prompts from the MCP server."""
+ return asyncio.run(self._list_prompts_async())
+
+ async def _list_prompts_async(self) -> List[MCPPrompt]:
+ """Async implementation of list_prompts."""
+ async with self._get_session() as session:
+ prompts_response = await session.list_prompts()
+
+ prompts = []
+ for prompt_data in prompts_response.prompts or []:
+ arguments_dict = {}
+ if prompt_data.arguments:
+ for arg in prompt_data.arguments:
+ arguments_dict[arg.name] = {
+ "description": arg.description,
+ "required": getattr(arg, "required", False),
+ }
+
+ prompt = MCPPrompt(
+ name=prompt_data.name,
+ title=prompt_data.title,
+ description=prompt_data.description,
+ prompt_arguments=prompt_data.arguments,
+ mcp_server=self,
+ )
+ prompts.append(prompt)
+
+ return prompts
+
+ def get_prompt(
+ self, name: str, arguments: Dict[str, str] | None = None
+ ) -> List[ChatMessage]:
+ """Get a single prompt definition by name."""
+ return asyncio.run(self._get_prompt_async(name, arguments))
+
+ async def _get_prompt_async(
+ self, name: str, arguments: Dict[str, str] | None = None
+ ) -> List[ChatMessage]:
+ """Async implementation of get_prompt."""
+ async with self._get_session() as session:
+ prompt = await session.get_prompt(name, arguments)
+ chat_messages = []
+ for message in prompt.messages:
+ if isinstance(message.content, TextContent):
+ chat_messages.append(
+ ChatMessage(role=message.role, content=message.content.text)
+ )
+ else:
+ err_msg = f"Unsupported content type: {type(message.content)}"
+ raise TypeError(err_msg)
+
+ return chat_messages
diff --git a/python/flink_agents/api/tools/utils.py b/python/flink_agents/api/tools/utils.py
index 077728a..81784d4 100644
--- a/python/flink_agents/api/tools/utils.py
+++ b/python/flink_agents/api/tools/utils.py
@@ -17,9 +17,10 @@
#################################################################################
import typing
from inspect import signature
-from typing import Any, Callable, Optional, Type, Union
+from typing import Any, Callable, Dict, Optional, Type, Union
from docstring_parser import parse
+from mcp import types
from pydantic import BaseModel, create_model
from pydantic.fields import Field, FieldInfo
@@ -38,7 +39,11 @@ def create_schema_from_function(name: str, func: Callable) -> Type[BaseModel]:
for param_name in params:
param_type = params[param_name].annotation
param_default = params[param_name].default
- description = doc_params[param_name].description
+ description = doc_params.get(param_name)
+ if description is not None:
+ description = description.description
+ else:
+ description = f"Parameter: {param_name}"
if typing.get_origin(param_type) is typing.Annotated:
args = typing.get_args(param_type)
@@ -172,3 +177,45 @@ def create_model_from_schema(name: str, schema: dict) -> type[BaseModel]:
main_fields[field_name] = (field_type, Field(**field_params))
return create_model(name, **main_fields, __doc__=schema.get("description", ""))
+
+
+def extract_mcp_content_item(content_item: Any) -> Dict[str, Any] | str:
+ """Extract and normalize a single MCP content item.
+
+ Args:
+ content_item: A single MCP content item (TextContent, ImageContent, etc.)
+
+ Returns:
+ Dict representation of the content item
+
+ Raises:
+ ImportError: If MCP types are not available
+ """
+ if types is None:
+ err_msg = "MCP types not available. Please install the mcp package."
+ raise ImportError(err_msg)
+
+ if isinstance(content_item, types.TextContent):
+ return content_item.text
+ elif isinstance(content_item, types.ImageContent):
+ return {
+ "type": "image",
+ "data": content_item.data,
+ "mimeType": content_item.mimeType
+ }
+ elif isinstance(content_item, types.EmbeddedResource):
+ if isinstance(content_item.resource, types.TextResourceContents):
+ return {
+ "type": "resource",
+ "uri": content_item.resource.uri,
+ "text": content_item.resource.text
+ }
+ elif isinstance(content_item.resource, types.BlobResourceContents):
+ return {
+ "type": "resource",
+ "uri": content_item.resource.uri,
+ "blob": content_item.resource.blob
+ }
+ else:
+ # Handle unknown content types as generic dict
+ return content_item.model_dump() if hasattr(content_item, 'model_dump') else str(content_item)
diff --git a/python/flink_agents/plan/agent_plan.py b/python/flink_agents/plan/agent_plan.py
index 511c828..467819a 100644
--- a/python/flink_agents/plan/agent_plan.py
+++ b/python/flink_agents/plan/agent_plan.py
@@ -15,12 +15,13 @@
# See the License for the specific language governing permissions and
# limitations under the License.
#################################################################################
-from typing import Any, Dict, List
+from typing import Any, Dict, List, cast
from pydantic import BaseModel, field_serializer, model_validator
from flink_agents.api.agent import Agent
from flink_agents.api.resource import Resource, ResourceType
+from flink_agents.api.tools.mcp import MCPServer
from flink_agents.plan.actions.action import Action
from flink_agents.plan.actions.chat_model_action import CHAT_MODEL_ACTION
from flink_agents.plan.actions.tool_call_action import TOOL_CALL_ACTION
@@ -208,7 +209,9 @@ class AgentPlan(BaseModel):
self.__resources[type] = {}
if name not in self.__resources[type]:
resource_provider = self.resource_providers[type][name]
- resource = resource_provider.provide(get_resource=self.get_resource, config=self.config)
+ resource = resource_provider.provide(
+ get_resource=self.get_resource, config=self.config
+ )
self.__resources[type][name] = resource
return self.__resources[type][name]
@@ -345,6 +348,12 @@ def _get_resource_providers(agent: Agent) -> List[ResourceProvider]:
name=name, resource=prompt
)
)
+ elif hasattr(value, "_is_mcp_server"):
+ if isinstance(value, staticmethod):
+ value = value.__func__
+
+ mcp_server = value()
+ _add_mcp_server(name, resource_providers, mcp_server)
for name, prompt in agent.resources[ResourceType.PROMPT].items():
resource_providers.append(
@@ -357,6 +366,10 @@ def _get_resource_providers(agent: Agent) -> List[ResourceProvider]:
PythonSerializableResourceProvider.from_resource(name=name, resource=tool)
)
+ for name, mcp_server in agent.resources[ResourceType.MCP_SERVER].items():
+ mcp_server = cast("MCPServer", mcp_server)
+ _add_mcp_server(name, resource_providers, mcp_server)
+
for name, chat_model in agent.resources[ResourceType.CHAT_MODEL].items():
clazz, kwargs = chat_model
provider = PythonResourceProvider(
@@ -402,3 +415,30 @@ def _get_resource_providers(agent: Agent) -> List[ResourceProvider]:
resource_providers.append(provider)
return resource_providers
+
+
+def _add_mcp_server(
+ name: str, resource_providers: List[ResourceProvider], mcp_server: MCPServer
+) -> None:
+ resource_providers.append(
+ PythonSerializableResourceProvider.from_resource(name=name, resource=mcp_server)
+ )
+ resource_providers.extend(
+ [
+ PythonSerializableResourceProvider.from_resource(
+ name=prompt.name, resource=prompt
+ )
+ for prompt in mcp_server.list_prompts()
+ ]
+ )
+
+ resource_providers.extend(
+ [
+ PythonSerializableResourceProvider.from_resource(
+ name=tool.name, resource=tool
+ )
+ for tool in mcp_server.list_tools()
+ ]
+ )
+
+ mcp_server.close()
diff --git a/python/pyproject.toml b/python/pyproject.toml
index 38090c6..d45065b 100644
--- a/python/pyproject.toml
+++ b/python/pyproject.toml
@@ -46,6 +46,7 @@ dependencies = [
"pydantic==2.11.4",
"docstring-parser==0.16",
"pyyaml==6.0.2",
+ "mcp>=1.8.0",
#TODO: Seperate integration dependencies from project
"ollama==0.4.8",
"dashscope~=1.24.2",