From 85c00303cc71a839c2ea8ebdf7371f78b3e35de0 Mon Sep 17 00:00:00 2001 From: Mengqin Shen Date: Mon, 29 Dec 2025 16:34:13 -0800 Subject: [PATCH 1/4] feat(py):add mcp plugin and relevant tests --- py/packages/genkit/src/genkit/ai/_registry.py | 185 ++++++- .../genkit/src/genkit/blocks/prompt.py | 4 +- .../genkit/src/genkit/blocks/resource.py | 90 ++++ .../genkit/src/genkit/core/action/types.py | 1 + .../tests/genkit/blocks/resource_test.py | 259 ++++++++++ py/plugins/mcp/README.md | 3 + .../mcp/examples/client/simple_client.py | 53 ++ .../examples/server/prompts/port_code.prompt | 13 + .../mcp/examples/server/simple_server.py | 63 +++ py/plugins/mcp/pyproject.toml | 49 ++ .../mcp/src/genkit/plugins/mcp/__init__.py | 40 ++ .../src/genkit/plugins/mcp/client/__init__.py | 16 + .../src/genkit/plugins/mcp/client/client.py | 217 +++++++++ .../mcp/src/genkit/plugins/mcp/client/host.py | 64 +++ .../mcp/src/genkit/plugins/mcp/index.py | 40 ++ .../mcp/src/genkit/plugins/mcp/server.py | 455 ++++++++++++++++++ .../src/genkit/plugins/mcp/util/__init__.py | 58 +++ .../src/genkit/plugins/mcp/util/message.py | 169 +++++++ .../src/genkit/plugins/mcp/util/prompts.py | 137 ++++++ .../src/genkit/plugins/mcp/util/resource.py | 146 ++++++ .../mcp/src/genkit/plugins/mcp/util/tools.py | 146 ++++++ .../src/genkit/plugins/mcp/util/transport.py | 89 ++++ py/plugins/mcp/tests/fakes.py | 128 +++++ py/plugins/mcp/tests/test_mcp_conversion.py | 259 ++++++++++ py/plugins/mcp/tests/test_mcp_host.py | 64 +++ py/plugins/mcp/tests/test_mcp_integration.py | 311 ++++++++++++ py/plugins/mcp/tests/test_mcp_server.py | 341 +++++++++++++ .../mcp/tests/test_mcp_server_resources.py | 351 ++++++++++++++ py/pyproject.toml | 1 + py/uv.lock | 106 ++++ 30 files changed, 3835 insertions(+), 23 deletions(-) create mode 100644 py/packages/genkit/src/genkit/blocks/resource.py create mode 100644 py/packages/genkit/tests/genkit/blocks/resource_test.py create mode 100644 py/plugins/mcp/README.md create mode 100644 py/plugins/mcp/examples/client/simple_client.py create mode 100644 py/plugins/mcp/examples/server/prompts/port_code.prompt create mode 100644 py/plugins/mcp/examples/server/simple_server.py create mode 100644 py/plugins/mcp/pyproject.toml create mode 100644 py/plugins/mcp/src/genkit/plugins/mcp/__init__.py create mode 100644 py/plugins/mcp/src/genkit/plugins/mcp/client/__init__.py create mode 100644 py/plugins/mcp/src/genkit/plugins/mcp/client/client.py create mode 100644 py/plugins/mcp/src/genkit/plugins/mcp/client/host.py create mode 100644 py/plugins/mcp/src/genkit/plugins/mcp/index.py create mode 100644 py/plugins/mcp/src/genkit/plugins/mcp/server.py create mode 100644 py/plugins/mcp/src/genkit/plugins/mcp/util/__init__.py create mode 100644 py/plugins/mcp/src/genkit/plugins/mcp/util/message.py create mode 100644 py/plugins/mcp/src/genkit/plugins/mcp/util/prompts.py create mode 100644 py/plugins/mcp/src/genkit/plugins/mcp/util/resource.py create mode 100644 py/plugins/mcp/src/genkit/plugins/mcp/util/tools.py create mode 100644 py/plugins/mcp/src/genkit/plugins/mcp/util/transport.py create mode 100644 py/plugins/mcp/tests/fakes.py create mode 100644 py/plugins/mcp/tests/test_mcp_conversion.py create mode 100644 py/plugins/mcp/tests/test_mcp_host.py create mode 100644 py/plugins/mcp/tests/test_mcp_integration.py create mode 100644 py/plugins/mcp/tests/test_mcp_server.py create mode 100644 py/plugins/mcp/tests/test_mcp_server_resources.py diff --git a/py/packages/genkit/src/genkit/ai/_registry.py b/py/packages/genkit/src/genkit/ai/_registry.py index 8d62249981..f58256698a 100644 --- a/py/packages/genkit/src/genkit/ai/_registry.py +++ b/py/packages/genkit/src/genkit/ai/_registry.py @@ -30,6 +30,7 @@ | `'indexer'` | Indexer | | `'model'` | Model | | `'prompt'` | Prompt | +| `'resource'` | Resource | | `'retriever'` | Retriever | | `'text-llm'` | Text LLM | | `'tool'` | Tool | @@ -55,6 +56,12 @@ define_helper, define_prompt, lookup_prompt, + registry_definition_key, + to_generate_request, +) +from genkit.blocks.resource import ( + ResourceContent, + matches_uri_template, ) from genkit.blocks.retriever import IndexerFn, RetrieverFn from genkit.blocks.tools import ToolRunContext @@ -69,6 +76,8 @@ EvalRequest, EvalResponse, EvalStatusEnum, + GenerateActionOptions, + GenerateRequest, GenerationCommonConfig, Message, ModelInfo, @@ -573,6 +582,7 @@ def define_format(self, format: FormatDef) -> None: def define_prompt( self, + name: str | None = None, variant: str | None = None, model: str | None = None, config: GenerationCommonConfig | dict[str, Any] | None = None, @@ -598,31 +608,34 @@ def define_prompt( """Define a prompt. Args: - variant: Optional variant name for the prompt. - model: Optional model name to use for the prompt. - config: Optional configuration for the model. - description: Optional description for the prompt. - input_schema: Optional schema for the input to the prompt. - system: Optional system message for the prompt. - prompt: Optional prompt for the model. - messages: Optional messages for the model. - output_format: Optional output format for the prompt. - output_content_type: Optional output content type for the prompt. - output_instructions: Optional output instructions for the prompt. - output_schema: Optional schema for the output from the prompt. - output_constrained: Optional flag indicating whether the output - should be constrained. - max_turns: Optional maximum number of turns for the prompt. - return_tool_requests: Optional flag indicating whether tool requests - should be returned. - metadata: Optional metadata for the prompt. - tools: Optional list of tools to use for the prompt. - tool_choice: Optional tool choice for the prompt. - use: Optional list of model middlewares to use for the prompt. + name: The name of the prompt. + variant: The variant of the prompt. + model: The model to use for generation. + config: The generation configuration. + description: A description of the prompt. + input_schema: The input schema for the prompt. + system: The system message for the prompt. + prompt: The user prompt. + messages: A list of messages to include in the prompt. + output_format: The output format. + output_content_type: The output content type. + output_instructions: Instructions for formatting the output. + output_schema: The output schema. + output_constrained: Whether the output should be constrained to the output schema. + max_turns: The maximum number of turns in a conversation. + return_tool_requests: Whether to return tool requests. + metadata: Metadata to associate with the prompt. + tools: A list of tool names to use with the prompt. + tool_choice: The tool choice strategy. + use: A list of model middlewares to apply. + + Returns: + An ExecutablePrompt instance. """ - return define_prompt( + executable_prompt = define_prompt( self.registry, variant=variant, + _name=name, model=model, config=config, description=description, @@ -643,6 +656,50 @@ def define_prompt( use=use, ) + if name: + # Register actions for kind PROMPT and EXECUTABLE_PROMPT + # This allows discovery by MCP and Dev UI + + async def prompt_action_fn(input: Any = None) -> GenerateRequest: + """PROMPT action function - renders prompt and returns GenerateRequest.""" + options = await executable_prompt.render(input=input) + return await to_generate_request(self.registry, options) + + async def executable_prompt_action_fn(input: Any = None) -> GenerateActionOptions: + """EXECUTABLE_PROMPT action function - renders prompt and returns GenerateActionOptions.""" + return await executable_prompt.render(input=input) + + action_name = registry_definition_key(name, variant) + action_metadata = { + 'type': 'prompt', + 'lazy': False, + 'source': 'programmatic', + 'prompt': { + 'name': name, + 'variant': variant or '', + }, + } + + # Register the PROMPT action + prompt_action = self.registry.register_action( + kind=ActionKind.PROMPT, + name=action_name, + fn=prompt_action_fn, + metadata=action_metadata, + ) + executable_prompt._prompt_action = prompt_action + prompt_action._executable_prompt = executable_prompt + + # Register the EXECUTABLE_PROMPT action + self.registry.register_action( + kind=ActionKind.EXECUTABLE_PROMPT, + name=action_name, + fn=executable_prompt_action_fn, + metadata=action_metadata, + ) + + return executable_prompt + async def prompt( self, name: str, @@ -675,6 +732,90 @@ async def prompt( variant=variant, ) + def define_resource( + self, + name: str, + fn: Callable, + uri: str | None = None, + template: str | None = None, + description: str | None = None, + metadata: dict[str, Any] | None = None, + ) -> Action: + """Define a resource action. + + Resources provide content that can be accessed via URI. They can have: + - A fixed URI (e.g., "my://resource") + - A URI template with placeholders (e.g., "file://{path}") + + Args: + name: Name of the resource. + fn: Function implementing the resource behavior. Should accept a dict + with 'uri' key and return ResourceContent or dict with 'content' key. + uri: Optional fixed URI for the resource. + template: Optional URI template with {param} placeholders. + description: Optional description for the resource. + metadata: Optional metadata for the resource. + + Returns: + The registered Action for the resource. + + Raises: + ValueError: If neither uri nor template is provided. + + Examples: + # Fixed URI resource + ai.define_resource( + name="my_resource", + uri="my://resource", + fn=lambda req: {"content": [{"text": "resource content"}]} + ) + + # Template URI resource + ai.define_resource( + name="file", + template="file://{path}", + fn=lambda req: {"content": [{"text": f"contents of {req['uri']}"}]} + ) + """ + if not uri and not template: + raise ValueError("Either 'uri' or 'template' must be provided for a resource") + + resource_meta = metadata if metadata else {} + if 'resource' not in resource_meta: + resource_meta['resource'] = {} + + # Store URI or template in metadata + if uri: + resource_meta['resource']['uri'] = uri + if template: + resource_meta['resource']['template'] = template + + resource_description = get_func_description(fn, description) + + # Wrap the resource function to handle template matching and extraction + async def resource_wrapper(input_data: dict[str, Any]) -> ResourceContent: + req_uri = input_data.get('uri') + if template and req_uri: + # Extract parameters from URI based on template + params = matches_uri_template(template, req_uri) + if params: + # Merge extracted parameters into the request data + # This allows the resource function to access them as req['param_name'] + input_data = {**params, **input_data} + + result = fn(input_data) + if inspect.isawaitable(result): + return await result + return result + + return self.registry.register_action( + name=name, + kind=ActionKind.RESOURCE, + fn=resource_wrapper, + metadata=resource_meta, + description=resource_description, + ) + class FlowWrapper: """A wapper for flow functions to add `stream` method.""" diff --git a/py/packages/genkit/src/genkit/blocks/prompt.py b/py/packages/genkit/src/genkit/blocks/prompt.py index fcb3d65fa2..107f88e872 100644 --- a/py/packages/genkit/src/genkit/blocks/prompt.py +++ b/py/packages/genkit/src/genkit/blocks/prompt.py @@ -347,7 +347,7 @@ async def as_tool(self) -> Action: if self._name is None: raise GenkitError( status='FAILED_PRECONDITION', - message='Prompt name not available. This prompt was not created via define_prompt_async() or load_prompt().', + message='Prompt name not available. This prompt was not created via Genkit.define_prompt() or load_prompt().', ) lookup_key = registry_lookup_key(self._name, self._variant, self._ns) @@ -366,6 +366,7 @@ async def as_tool(self) -> Action: def define_prompt( registry: Registry, variant: str | None = None, + _name: str | None = None, model: str | None = None, config: GenerationCommonConfig | dict[str, Any] | None = None, description: str | None = None, @@ -435,6 +436,7 @@ def define_prompt( tools=tools, tool_choice=tool_choice, use=use, + _name=_name, ) diff --git a/py/packages/genkit/src/genkit/blocks/resource.py b/py/packages/genkit/src/genkit/blocks/resource.py new file mode 100644 index 0000000000..cd27bb77b7 --- /dev/null +++ b/py/packages/genkit/src/genkit/blocks/resource.py @@ -0,0 +1,90 @@ +# Copyright 2025 Google LLC +# +# Licensed 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. +# +# SPDX-License-Identifier: Apache-2.0 + +"""Resource types and functions for Genkit.""" + +import re +from collections.abc import Awaitable, Callable +from typing import Any + +from pydantic import BaseModel + +from genkit.core.typing import Part + + +class ResourceOptions(BaseModel): + """Options for defining a resource. + + Attributes: + name: The name of the resource. + uri: Optional fixed URI for the resource (e.g., "my://resource"). + template: Optional URI template with placeholders (e.g., "file://{path}"). + description: Optional description of the resource. + """ + + name: str + uri: str | None = None + template: str | None = None + description: str | None = None + + +class ResourceContent(BaseModel): + """Content returned by a resource. + + Attributes: + content: List of content parts (text, media, etc.). + """ + + content: list[Part] + + +# Type for resource function +ResourceFn = Callable[[dict[str, Any]], Awaitable[ResourceContent] | ResourceContent] + + +def matches_uri_template(template: str, uri: str) -> dict[str, str] | None: + """Check if a URI matches a template and extract parameters. + + Args: + template: URI template with {param} placeholders (e.g., "file://{path}"). + uri: The URI to match against the template. + + Returns: + Dictionary of extracted parameters if match, None otherwise. + + Examples: + >>> matches_uri_template('file://{path}', 'file:///home/user/doc.txt') + {'path': '/home/user/doc.txt'} + >>> matches_uri_template('user://{id}/profile', 'user://123/profile') + {'id': '123'} + """ + # Split template into parts: text and {param} placeholders + parts = re.split(r'(\{[\w]+\})', template) + pattern_parts = [] + for part in parts: + if part.startswith('{') and part.endswith('}'): + param_name = part[1:-1] + # Use .+? (non-greedy) to match parameters + pattern_parts.append(f'(?P<{param_name}>.+?)') + else: + pattern_parts.append(re.escape(part)) + + pattern = f'^{"".join(pattern_parts)}$' + + match = re.match(pattern, uri) + if match: + return match.groupdict() + return None diff --git a/py/packages/genkit/src/genkit/core/action/types.py b/py/packages/genkit/src/genkit/core/action/types.py index 9609285499..203000649e 100644 --- a/py/packages/genkit/src/genkit/core/action/types.py +++ b/py/packages/genkit/src/genkit/core/action/types.py @@ -57,6 +57,7 @@ class ActionKind(StrEnum): MODEL = 'model' PROMPT = 'prompt' RERANKER = 'reranker' + RESOURCE = 'resource' RETRIEVER = 'retriever' TOOL = 'tool' UTIL = 'util' diff --git a/py/packages/genkit/tests/genkit/blocks/resource_test.py b/py/packages/genkit/tests/genkit/blocks/resource_test.py new file mode 100644 index 0000000000..50800e039a --- /dev/null +++ b/py/packages/genkit/tests/genkit/blocks/resource_test.py @@ -0,0 +1,259 @@ +# Copyright 2025 Google LLC +# +# Licensed 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. +# +# SPDX-License-Identifier: Apache-2.0 + +"""Unit tests for resource registration and URI matching.""" + +import unittest + +import pytest + +from genkit.ai import Genkit +from genkit.blocks.resource import ResourceContent, ResourceOptions, matches_uri_template +from genkit.core.action.types import ActionKind + + +class TestResourceRegistration(unittest.TestCase): + """Tests for resource registration in Genkit registry.""" + + def setUp(self): + """Set up test fixtures.""" + self.ai = Genkit() + + def test_define_resource_with_fixed_uri(self): + """Test defining a resource with a fixed URI.""" + + def my_resource(req): + return {'content': [{'text': 'test content'}]} + + action = self.ai.define_resource(name='test_resource', uri='test://resource', fn=my_resource) + + # Verify action was registered + self.assertIsNotNone(action) + self.assertEqual(action.name, 'test_resource') + self.assertEqual(action.kind, ActionKind.RESOURCE) + + # Verify metadata + self.assertIn('resource', action.metadata) + self.assertEqual(action.metadata['resource']['uri'], 'test://resource') + + def test_define_resource_with_template(self): + """Test defining a resource with a URI template.""" + + def file_resource(req): + return {'content': [{'text': f'contents of {req["uri"]}'}]} + + action = self.ai.define_resource(name='file', template='file://{path}', fn=file_resource) + + # Verify action was registered + self.assertIsNotNone(action) + self.assertEqual(action.name, 'file') + + # Verify metadata + self.assertIn('resource', action.metadata) + self.assertEqual(action.metadata['resource']['template'], 'file://{path}') + + def test_define_resource_requires_uri_or_template(self): + """Test that defining a resource requires either uri or template.""" + + def my_resource(req): + return {'content': [{'text': 'test'}]} + + with self.assertRaises(ValueError) as context: + self.ai.define_resource(name='invalid_resource', fn=my_resource) + + self.assertIn('uri', str(context.exception).lower()) + self.assertIn('template', str(context.exception).lower()) + + def test_define_resource_with_description(self): + """Test defining a resource with a description.""" + + def my_resource(req): + return {'content': [{'text': 'test'}]} + + action = self.ai.define_resource( + name='described_resource', uri='test://resource', description='Test resource description', fn=my_resource + ) + + self.assertEqual(action.description, 'Test resource description') + + def test_define_resource_with_metadata(self): + """Test defining a resource with custom metadata.""" + + def my_resource(req): + return {'content': [{'text': 'test'}]} + + custom_metadata = {'custom_key': 'custom_value', 'mcp': {'_meta': {'version': '1.0'}}} + + action = self.ai.define_resource( + name='meta_resource', uri='test://resource', metadata=custom_metadata, fn=my_resource + ) + + self.assertIn('custom_key', action.metadata) + self.assertEqual(action.metadata['custom_key'], 'custom_value') + self.assertIn('mcp', action.metadata) + + +class TestURITemplateMatching(unittest.TestCase): + """Tests for URI template matching functionality.""" + + def test_exact_match(self): + """Test exact URI matching without parameters.""" + template = 'file:///exact/path' + uri = 'file:///exact/path' + + result = matches_uri_template(template, uri) + self.assertIsNotNone(result) + self.assertEqual(result, {}) + + def test_single_parameter_match(self): + """Test URI template with single parameter.""" + template = 'file://{path}' + uri = 'file:///home/user/document.txt' + + result = matches_uri_template(template, uri) + self.assertIsNotNone(result) + self.assertIn('path', result) + self.assertEqual(result['path'], '/home/user/document.txt') + + def test_multiple_parameters_match(self): + """Test URI template with multiple parameters.""" + template = 'user://{user_id}/profile/{section}' + uri = 'user://12345/profile/settings' + + result = matches_uri_template(template, uri) + self.assertIsNotNone(result) + self.assertEqual(result['user_id'], '12345') + self.assertEqual(result['section'], 'settings') + + def test_no_match(self): + """Test URI that doesn't match template.""" + template = 'file://{path}' + uri = 'http://example.com/file.txt' + + result = matches_uri_template(template, uri) + self.assertIsNone(result) + + def test_partial_match_fails(self): + """Test that partial matches fail.""" + template = 'file://{path}/document.txt' + uri = 'file:///home/user/other.txt' + + result = matches_uri_template(template, uri) + self.assertIsNone(result) + + def test_complex_template(self): + """Test complex URI template with multiple segments.""" + template = 'api://{version}/users/{user_id}/posts/{post_id}' + uri = 'api://v2/users/alice/posts/42' + + result = matches_uri_template(template, uri) + self.assertIsNotNone(result) + self.assertEqual(result['version'], 'v2') + self.assertEqual(result['user_id'], 'alice') + self.assertEqual(result['post_id'], '42') + + def test_special_characters_in_uri(self): + """Test URI with special characters.""" + template = 'file://{path}' + uri = 'file:///path/with-dashes_and_underscores.txt' + + result = matches_uri_template(template, uri) + self.assertIsNotNone(result) + # Note: The current implementation uses [^/]+ which may not capture all special chars + # This test documents current behavior + + def test_empty_parameter(self): + """Test template matching with empty parameter.""" + template = 'resource://{id}/data' + uri = 'resource:///data' + + result = matches_uri_template(template, uri) + # Should not match because {id} expects at least one character + self.assertIsNone(result) + + +@pytest.mark.asyncio +class TestResourceExecution(unittest.IsolatedAsyncioTestCase): + """Tests for executing resource actions.""" + + async def test_execute_fixed_uri_resource(self): + """Test executing a resource with fixed URI.""" + ai = Genkit() + + def my_resource(req): + return {'content': [{'text': 'Hello from resource!'}]} + + action = ai.define_resource(name='greeting', uri='app://greeting', fn=my_resource) + + # Execute the resource + result = await action.arun({'uri': 'app://greeting'}) + + self.assertIn('content', result.response) + self.assertEqual(len(result.response['content']), 1) + self.assertEqual(result.response['content'][0]['text'], 'Hello from resource!') + + async def test_execute_template_resource(self): + """Test executing a resource with URI template.""" + ai = Genkit() + + def user_profile(req): + user_id = req.get('user_id') + return {'content': [{'text': f'Profile for user {user_id}'}]} + + action = ai.define_resource(name='user_profile', template='user://{user_id}/profile', fn=user_profile) + + # Execute the resource + result = await action.arun({'uri': 'user://alice/profile'}) + + self.assertIn('content', result.response) + self.assertEqual(result.response['content'][0]['text'], 'Profile for user alice') + + async def test_resource_with_multiple_content_parts(self): + """Test resource returning multiple content parts.""" + ai = Genkit() + + def multi_part_resource(req): + return {'content': [{'text': 'Part 1'}, {'text': 'Part 2'}, {'text': 'Part 3'}]} + + action = ai.define_resource(name='multi', uri='test://multi', fn=multi_part_resource) + + result = await action.arun({'uri': 'test://multi'}) + + self.assertEqual(len(result.response['content']), 3) + self.assertEqual(result.response['content'][0]['text'], 'Part 1') + self.assertEqual(result.response['content'][1]['text'], 'Part 2') + self.assertEqual(result.response['content'][2]['text'], 'Part 3') + + async def test_async_resource_function(self): + """Test resource with async function.""" + ai = Genkit() + + async def async_resource(req): + # Simulate async operation + import asyncio + + await asyncio.sleep(0.01) + return {'content': [{'text': 'Async result'}]} + + action = ai.define_resource(name='async_res', uri='test://async', fn=async_resource) + + result = await action.arun({'uri': 'test://async'}) + + self.assertEqual(result.response['content'][0]['text'], 'Async result') + + +if __name__ == '__main__': + unittest.main() diff --git a/py/plugins/mcp/README.md b/py/plugins/mcp/README.md new file mode 100644 index 0000000000..1ad7262193 --- /dev/null +++ b/py/plugins/mcp/README.md @@ -0,0 +1,3 @@ +# Genkit MCP Plugin + +Integrate Model Context Protocol (MCP) with Genkit. diff --git a/py/plugins/mcp/examples/client/simple_client.py b/py/plugins/mcp/examples/client/simple_client.py new file mode 100644 index 0000000000..9512f12a5f --- /dev/null +++ b/py/plugins/mcp/examples/client/simple_client.py @@ -0,0 +1,53 @@ +# Copyright 2025 Google LLC +# +# Licensed 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. +# +# SPDX-License-Identifier: Apache-2.0 + +import asyncio + +from genkit.ai import Genkit +from genkit.plugins.mcp import McpServerConfig, create_mcp_client + +try: + from genkit.plugins.google_genai import GoogleAI +except ImportError: + GoogleAI = None + + +# Simple client example connecting to 'everything' server using npx +async def main(): + # Define the client plugin + everything_client = create_mcp_client( + name='everything', config=McpServerConfig(command='npx', args=['-y', '@modelcontextprotocol/server-everything']) + ) + + plugins = [everything_client] + if GoogleAI: + plugins.append(GoogleAI()) + + ai = Genkit(plugins=plugins) + + await everything_client.connect() + + print('Connected! Listing tools...') + + tools = await everything_client.list_tools() + for t in tools: + print(f'- {t.name}: {t.description}') + + await everything_client.close() + + +if __name__ == '__main__': + asyncio.run(main()) diff --git a/py/plugins/mcp/examples/server/prompts/port_code.prompt b/py/plugins/mcp/examples/server/prompts/port_code.prompt new file mode 100644 index 0000000000..77e8501b36 --- /dev/null +++ b/py/plugins/mcp/examples/server/prompts/port_code.prompt @@ -0,0 +1,13 @@ +--- +input: + schema: + code: string, the source code to port from one language to another + fromLang?: string, the original language of the source code (e.g. js, python) + toLang: string, the destination language of the source code (e.g. python, js) +--- + +You are assisting the user in translating code between two programming languages. Given the code below, translate it into {{toLang}}. + +```{{#if fromLang}}{{fromLang}}{{/if}} +{{code}} +``` diff --git a/py/plugins/mcp/examples/server/simple_server.py b/py/plugins/mcp/examples/server/simple_server.py new file mode 100644 index 0000000000..36e7bdbe2b --- /dev/null +++ b/py/plugins/mcp/examples/server/simple_server.py @@ -0,0 +1,63 @@ +# Copyright 2025 Google LLC +# +# Licensed 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. +# +# SPDX-License-Identifier: Apache-2.0 + +import asyncio + +from pydantic import BaseModel, Field + +from genkit.ai import Genkit +from genkit.plugins.mcp import McpServerOptions, create_mcp_server + + +# Define input model +class AddInput(BaseModel): + a: int = Field(..., description='First number') + b: int = Field(..., description='Second number') + + +import os + + +def main(): + # Load prompts from the 'prompts' directory relative to this script + script_dir = os.path.dirname(os.path.abspath(__file__)) + prompts_dir = os.path.join(script_dir, 'prompts') + + ai = Genkit(prompt_dir=prompts_dir) + + @ai.tool(name='add', description='add two numbers together') + def add(input: AddInput): + return input.a + input.b + + # Genkit Python prompt definition (simplified) + # Note: In Python, prompts are typically loaded from files via prompt_dir + # This inline definition is for demonstration purposes + happy_prompt = ai.define_prompt( + input_schema={'action': str}, + prompt="If you're happy and you know it, {{action}}.", + ) + + # Create and start MCP server + # Note: create_mcp_server returns McpServer instance. + # In JS example: .start() is called. + server = create_mcp_server(ai, McpServerOptions(name='example_server', version='0.0.1')) + + print('Starting MCP server on stdio...') + asyncio.run(server.start_stdio()) + + +if __name__ == '__main__': + main() diff --git a/py/plugins/mcp/pyproject.toml b/py/plugins/mcp/pyproject.toml new file mode 100644 index 0000000000..7f6936009f --- /dev/null +++ b/py/plugins/mcp/pyproject.toml @@ -0,0 +1,49 @@ +# Copyright 2025 Google LLC +# +# Licensed 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. +# +# SPDX-License-Identifier: Apache-2.0 + +[project] +authors = [{ name = "Google" }] +classifiers = [ + "Development Status :: 3 - Alpha", + "Environment :: Console", + "Environment :: Web Environment", + "Intended Audience :: Developers", + "Operating System :: OS Independent", + "License :: OSI Approved :: Apache Software License", + "Programming Language :: Python", + "Programming Language :: Python :: 3 :: Only", + "Programming Language :: Python :: 3.10", + "Programming Language :: Python :: 3.11", + "Programming Language :: Python :: 3.12", + "Programming Language :: Python :: 3.13", + "Programming Language :: Python :: 3.14", + "Topic :: Scientific/Engineering :: Artificial Intelligence", + "Topic :: Software Development :: Libraries", +] +dependencies = ["genkit", "mcp"] +description = "Genkit MCP Plugin" +license = { text = "Apache-2.0" } +name = "genkit-plugins-mcp" +readme = "README.md" +requires-python = ">=3.10" +version = "0.1.0" + +[build-system] +build-backend = "hatchling.build" +requires = ["hatchling"] + +[tool.hatch.build.targets.wheel] +packages = ["src/genkit", "src/genkit/plugins"] diff --git a/py/plugins/mcp/src/genkit/plugins/mcp/__init__.py b/py/plugins/mcp/src/genkit/plugins/mcp/__init__.py new file mode 100644 index 0000000000..8e5226f356 --- /dev/null +++ b/py/plugins/mcp/src/genkit/plugins/mcp/__init__.py @@ -0,0 +1,40 @@ +""" +Copyright 2025 Google LLC + +Licensed 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. +""" + +from .client.client import ( + McpClient, + McpServerConfig, + create_mcp_client, +) +from .client.host import McpHost, create_mcp_host +from .server import McpServer, McpServerOptions, create_mcp_server + + +def package_name() -> str: + return 'genkit.plugins.mcp' + + +__all__ = [ + 'McpClient', + 'McpHost', + 'McpServerConfig', + 'create_mcp_client', + 'create_mcp_host', + 'McpServer', + 'McpServerOptions', + 'create_mcp_server', + 'package_name', +] diff --git a/py/plugins/mcp/src/genkit/plugins/mcp/client/__init__.py b/py/plugins/mcp/src/genkit/plugins/mcp/client/__init__.py new file mode 100644 index 0000000000..19add86cb8 --- /dev/null +++ b/py/plugins/mcp/src/genkit/plugins/mcp/client/__init__.py @@ -0,0 +1,16 @@ +# Copyright 2025 Google LLC +# +# Licensed 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. +# +# SPDX-License-Identifier: Apache-2.0 + diff --git a/py/plugins/mcp/src/genkit/plugins/mcp/client/client.py b/py/plugins/mcp/src/genkit/plugins/mcp/client/client.py new file mode 100644 index 0000000000..218493c056 --- /dev/null +++ b/py/plugins/mcp/src/genkit/plugins/mcp/client/client.py @@ -0,0 +1,217 @@ +# Copyright 2025 Google LLC +# +# Licensed 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. +# +# SPDX-License-Identifier: Apache-2.0 + +import asyncio +import uuid +from typing import Any, Callable, Dict, List, Optional, Union + +import structlog +from pydantic import BaseModel + +from genkit.ai import Genkit +from genkit.ai._plugin import Plugin +from genkit.ai._registry import GenkitRegistry +from genkit.core.action.types import ActionKind +from mcp import ClientSession, StdioServerParameters +from mcp.client.sse import sse_client +from mcp.client.stdio import stdio_client +from mcp.types import CallToolResult, Prompt, Resource, Tool + +logger = structlog.get_logger(__name__) + + +class McpServerConfig(BaseModel): + command: Optional[str] = None + args: Optional[List[str]] = None + env: Optional[Dict[str, str]] = None + url: Optional[str] = None + disabled: bool = False + + +class McpClient(Plugin): + """Client for connecting to a single MCP server.""" + + def __init__(self, name: str, config: McpServerConfig, server_name: Optional[str] = None): + self.name = name + self.config = config + self.server_name = server_name or name + self.session: Optional[ClientSession] = None + self._exit_stack = None + self._session_context = None + self.ai: Optional[GenkitRegistry] = None + + def plugin_name(self) -> str: + return self.name + + def initialize(self, ai: GenkitRegistry) -> None: + self.ai = ai + + def resolve_action(self, ai: GenkitRegistry, kind: ActionKind, name: str) -> None: + # MCP tools are dynamic and currently registered upon connection/Discovery. + # This hook allows lazy resolution if we implement it. + pass + + async def connect(self): + """Connects to the MCP server.""" + if self.config.disabled: + logger.info(f'MCP server {self.server_name} is disabled.') + return + + try: + if self.config.command: + server_params = StdioServerParameters( + command=self.config.command, args=self.config.args or [], env=self.config.env + ) + # stdio_client returns (read, write) streams + stdio_context = stdio_client(server_params) + read, write = await stdio_context.__aenter__() + self._exit_stack = stdio_context + + # Create and initialize session + session_context = ClientSession(read, write) + self.session = await session_context.__aenter__() + self._session_context = session_context + + elif self.config.url: + # TODO: Verify SSE client usage in mcp python SDK + sse_context = sse_client(self.config.url) + read, write = await sse_context.__aenter__() + self._exit_stack = sse_context + + session_context = ClientSession(read, write) + self.session = await session_context.__aenter__() + self._session_context = session_context + + await self.session.initialize() + logger.info(f'Connected to MCP server: {self.server_name}') + + except Exception as e: + logger.error(f'Failed to connect to MCP server {self.server_name}: {e}') + self.config.disabled = True + # Clean up on error + if hasattr(self, '_session_context') and self._session_context: + try: + await self._session_context.__aexit__(None, None, None) + except: + pass + if self._exit_stack: + try: + await self._exit_stack.__aexit__(None, None, None) + except: + pass + raise e + + async def close(self): + """Closes the connection.""" + if hasattr(self, '_session_context') and self._session_context: + try: + await self._session_context.__aexit__(None, None, None) + except Exception as e: + logger.debug(f'Error closing session: {e}') + if self._exit_stack: + try: + await self._exit_stack.__aexit__(None, None, None) + except Exception as e: + logger.debug(f'Error closing transport: {e}') + + async def list_tools(self) -> List[Tool]: + if not self.session: + return [] + result = await self.session.list_tools() + return result.tools + + async def call_tool(self, tool_name: str, arguments: dict) -> Any: + if not self.session: + raise RuntimeError('MCP client is not connected') + result: CallToolResult = await self.session.call_tool(tool_name, arguments) + # Process result similarly to JS SDK + if result.isError: + raise RuntimeError(f'Tool execution failed: {result.content}') + + # Simple text extraction for now + texts = [c.text for c in result.content if c.type == 'text'] + return ''.join(texts) + + async def list_prompts(self) -> List[Prompt]: + if not self.session: + return [] + result = await self.session.list_prompts() + return result.prompts + + async def get_prompt(self, name: str, arguments: Optional[dict] = None) -> Any: + if not self.session: + raise RuntimeError('MCP client is not connected') + return await self.session.get_prompt(name, arguments) + + async def list_resources(self) -> List[Resource]: + if not self.session: + return [] + result = await self.session.list_resources() + return result.resources + + async def read_resource(self, uri: str) -> Any: + if not self.session: + raise RuntimeError('MCP client is not connected') + return await self.session.read_resource(uri) + + async def register_tools(self, ai: Optional[Genkit] = None): + """Registers all tools from connected client to Genkit.""" + registry = ai.registry if ai else (self.ai.registry if self.ai else None) + if not registry: + logger.warning('No Genkit registry available to register tools.') + return + + if not self.session: + return + + try: + tools = await self.list_tools() + for tool in tools: + # Create a wrapper function for the tool + # We need to capture tool and client in closure + async def tool_wrapper(args: Any = None, _tool_name=tool.name): + # args might be Pydantic model or dict. Genkit passes dict usually? + # TODO: Validate args against schema if needed + arguments = args + if hasattr(args, 'model_dump'): + arguments = args.model_dump() + return await self.call_tool(_tool_name, arguments or {}) + + # Use metadata to store MCP specific info + metadata = {'mcp': {'_meta': tool._meta}} if hasattr(tool, '_meta') else {} + + # Define the tool in Genkit registry + registry.register_action( + kind=ActionKind.TOOL, + name=f'{self.server_name}/{tool.name}', + fn=tool_wrapper, + description=tool.description, + metadata=metadata, + # TODO: json_schema conversion from tool.inputSchema + ) + logger.debug(f'Registered MCP tool: {self.server_name}/{tool.name}') + except Exception as e: + logger.error(f'Error registering tools for {self.server_name}: {e}') + + async def get_active_tools(self) -> List[Any]: + """Returns all active tools.""" + if not self.session: + return [] + return await self.list_tools() + + +def create_mcp_client(config: McpServerConfig, name: str = 'mcp-client') -> McpClient: + return McpClient(name, config) diff --git a/py/plugins/mcp/src/genkit/plugins/mcp/client/host.py b/py/plugins/mcp/src/genkit/plugins/mcp/client/host.py new file mode 100644 index 0000000000..365b6382d5 --- /dev/null +++ b/py/plugins/mcp/src/genkit/plugins/mcp/client/host.py @@ -0,0 +1,64 @@ +# Copyright 2025 Google LLC +# +# Licensed 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. +# +# SPDX-License-Identifier: Apache-2.0 + +from typing import Dict, List, Optional + +from genkit.ai import Genkit + +from .client import McpClient, McpServerConfig + + +class McpHost: + """Host for managing multiple MCP clients.""" + + def __init__(self, clients: Dict[str, McpServerConfig]): + self.clients_config = clients + self.clients: Dict[str, McpClient] = {name: McpClient(name, config) for name, config in clients.items()} + + async def start(self): + """Starts all enabled MCP clients.""" + for client in self.clients.values(): + if not client.config.disabled: + await client.connect() + + async def close(self): + """Closes all MCP clients.""" + for client in self.clients.values(): + await client.close() + + async def register_tools(self, ai: Genkit): + """Registers all tools from connected clients to Genkit.""" + for client in self.clients.values(): + if client.session: + await client.register_tools(ai) + + async def enable(self, name: str): + """Enables and connects an MCP client.""" + if name in self.clients: + client = self.clients[name] + client.config.disabled = False + await client.connect() + + async def disable(self, name: str): + """Disables and closes an MCP client.""" + if name in self.clients: + client = self.clients[name] + client.config.disabled = True + await client.close() + + +def create_mcp_host(configs: Dict[str, McpServerConfig]) -> McpHost: + return McpHost(configs) diff --git a/py/plugins/mcp/src/genkit/plugins/mcp/index.py b/py/plugins/mcp/src/genkit/plugins/mcp/index.py new file mode 100644 index 0000000000..38365925dc --- /dev/null +++ b/py/plugins/mcp/src/genkit/plugins/mcp/index.py @@ -0,0 +1,40 @@ +# Copyright 2025 Google LLC +# +# Licensed 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. +# +# SPDX-License-Identifier: Apache-2.0 + +""" +MCP Plugin Index + +This module serves as the main entry point for the MCP plugin, +similar to js/plugins/mcp/src/index.ts. + +In Python, the actual exports are handled by the parent __init__.py, +but this file exists for structural parity with the JS SDK. +""" + +from .client.client import McpClient, McpServerConfig, create_mcp_client +from .client.host import McpHost, create_mcp_host +from .server import McpServer, McpServerOptions, create_mcp_server + +__all__ = [ + 'McpClient', + 'McpHost', + 'McpServerConfig', + 'create_mcp_client', + 'create_mcp_host', + 'McpServer', + 'McpServerOptions', + 'create_mcp_server', +] diff --git a/py/plugins/mcp/src/genkit/plugins/mcp/server.py b/py/plugins/mcp/src/genkit/plugins/mcp/server.py new file mode 100644 index 0000000000..088a57e481 --- /dev/null +++ b/py/plugins/mcp/src/genkit/plugins/mcp/server.py @@ -0,0 +1,455 @@ +# Copyright 2025 Google LLC +# +# Licensed 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 +# distributed under the License. +# 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. +# +# SPDX-License-Identifier: Apache-2.0 + +"""MCP Server implementation for exposing Genkit actions via Model Context Protocol.""" + +import asyncio +from typing import Any, Optional + +import structlog +from pydantic import BaseModel + +from genkit.ai import Genkit +from genkit.blocks.resource import matches_uri_template +from genkit.core.action.types import ActionKind +from genkit.core.error import GenkitError +from genkit.core.schema import to_json_schema +from mcp.server import Server +from mcp.server.stdio import stdio_server +from mcp.types import ( + CallToolRequest, + CallToolResult, + GetPromptRequest, + GetPromptResult, + ListPromptsRequest, + ListPromptsResult, + ListResourcesRequest, + ListResourcesResult, + ListResourceTemplatesRequest, + ListResourceTemplatesResult, + ListToolsRequest, + ListToolsResult, + Prompt, + ReadResourceRequest, + ReadResourceResult, + Resource, + ResourceTemplate, + Tool, +) + +from .util import ( + to_mcp_prompt_arguments, + to_mcp_prompt_message, + to_mcp_resource_contents, + to_mcp_tool_result, +) + +logger = structlog.get_logger(__name__) + + +class McpServerOptions(BaseModel): + """Options for creating an MCP server. + + Attributes: + name: The name of the MCP server. + version: The version of the server (default: "1.0.0"). + """ + + name: str + version: str = '1.0.0' + + +class McpServer: + """Exposes Genkit tools, prompts, and resources as an MCP server. + + This class wraps a Genkit instance and makes its registered actions + (tools, prompts, resources) available to MCP clients via the Model Context Protocol. + """ + + def __init__(self, ai: Genkit, options: McpServerOptions): + """Initialize the MCP server. + + Args: + ai: The Genkit instance whose actions will be exposed. + options: Configuration options for the MCP server. + """ + self.ai = ai + self.options = options + self.server: Optional[Server] = None + self.actions_resolved = False + self.tool_actions: list[Any] = [] + self.prompt_actions: list[Any] = [] + self.resource_actions: list[Any] = [] + + async def setup(self) -> None: + """Initialize the MCP server and register request handlers. + + This method sets up the MCP Server instance, registers all request handlers, + and resolves all actions from the Genkit registry. It's idempotent and can + be called multiple times safely. + """ + if self.actions_resolved: + return + + # Create MCP Server instance + self.server = Server( + {'name': self.options.name, 'version': self.options.version}, + { + 'capabilities': { + 'prompts': {}, + 'tools': {}, + 'resources': {}, + } + }, + ) + + # Register request handlers + self.server.setRequestHandler(ListToolsRequestSchema, self.list_tools) + self.server.setRequestHandler(CallToolRequestSchema, self.call_tool) + self.server.setRequestHandler(ListPromptsRequestSchema, self.list_prompts) + self.server.setRequestHandler(GetPromptRequestSchema, self.get_prompt) + self.server.setRequestHandler(ListResourcesRequestSchema, self.list_resources) + self.server.setRequestHandler(ListResourceTemplatesRequestSchema, self.list_resource_templates) + self.server.setRequestHandler(ReadResourceRequestSchema, self.read_resource) + + # Resolve all actions from Genkit registry + # We need the actual Action objects, not just serializable dicts + self.tool_actions = [] + self.prompt_actions = [] + self.resource_actions = [] + + # Get all actions from the registry + # We use the internal _entries for local actions and plugins + with self.ai.registry._lock: + for kind, entries in self.ai.registry._entries.items(): + for name, action in entries.items(): + if kind == ActionKind.TOOL: + self.tool_actions.append(action) + elif kind == ActionKind.PROMPT: + self.prompt_actions.append(action) + elif kind == ActionKind.RESOURCE: + self.resource_actions.append(action) + + # Also get actions from plugins that might not be in _entries yet + # (though most plugins register them in _entries during initialization) + plugin_actions = self.ai.registry.list_actions() + for key in plugin_actions: + kind, name = parse_action_key(key) + action = self.ai.registry.lookup_action(kind, name) + if action: + if kind == ActionKind.TOOL and action not in self.tool_actions: + self.tool_actions.append(action) + elif kind == ActionKind.PROMPT and action not in self.prompt_actions: + self.prompt_actions.append(action) + elif kind == ActionKind.RESOURCE and action not in self.resource_actions: + self.resource_actions.append(action) + + self.actions_resolved = True + + logger.info( + f'MCP Server initialized', + tools=len(self.tool_actions), + prompts=len(self.prompt_actions), + resources=len(self.resource_actions), + ) + + async def list_tools(self, request: ListToolsRequest) -> ListToolsResult: + """Handle MCP requests to list available tools. + + Args: + request: The MCP ListToolsRequest. + + Returns: + ListToolsResult containing all registered Genkit tools. + """ + await self.setup() + + tools: list[Tool] = [] + for action in self.tool_actions: + # Get tool definition + input_schema = to_json_schema(action.input_schema) if action.input_schema else {'type': 'object'} + + tools.append( + Tool( + name=action.name, + description=action.description or '', + inputSchema=input_schema, + _meta=action.metadata.get('mcp', {}).get('_meta') if action.metadata else None, + ) + ) + + return ListToolsResult(tools=tools) + + async def call_tool(self, request: CallToolRequest) -> CallToolResult: + """Handle MCP requests to call a specific tool. + + Args: + request: The MCP CallToolRequest containing tool name and arguments. + + Returns: + CallToolResult with the tool execution result. + + Raises: + GenkitError: If the requested tool is not found. + """ + await self.setup() + + # Find the tool action + tool = next((t for t in self.tool_actions if t.name == request.params.name), None) + + if not tool: + raise GenkitError( + status='NOT_FOUND', message=f"Tried to call tool '{request.params.name}' but it could not be found." + ) + + # Execute the tool + result = await tool.arun(request.params.arguments) + result = result.response + + # Convert result to MCP format + return CallToolResult(content=to_mcp_tool_result(result)) + + async def list_prompts(self, request: ListPromptsRequest) -> ListPromptsResult: + """Handle MCP requests to list available prompts. + + Args: + request: The MCP ListPromptsRequest. + + Returns: + ListPromptsResult containing all registered Genkit prompts. + """ + await self.setup() + + prompts: list[Prompt] = [] + for action in self.prompt_actions: + # Convert input schema to MCP prompt arguments + input_schema = to_json_schema(action.input_schema) if action.input_schema else None + arguments = to_mcp_prompt_arguments(input_schema) if input_schema else None + + prompts.append( + Prompt( + name=action.name, + description=action.description or '', + arguments=arguments, + _meta=action.metadata.get('mcp', {}).get('_meta') if action.metadata else None, + ) + ) + + return ListPromptsResult(prompts=prompts) + + async def get_prompt(self, request: GetPromptRequest) -> GetPromptResult: + """Handle MCP requests to get (render) a specific prompt. + + Args: + request: The MCP GetPromptRequest containing prompt name and arguments. + + Returns: + GetPromptResult with the rendered prompt messages. + + Raises: + GenkitError: If the requested prompt is not found. + """ + await self.setup() + + # Find the prompt action + prompt = next((p for p in self.prompt_actions if p.name == request.params.name), None) + + if not prompt: + raise GenkitError( + status='NOT_FOUND', + message=f"[MCP Server] Tried to call prompt '{request.params.name}' but it could not be found.", + ) + + # Execute the prompt + result = await prompt.arun(request.params.arguments) + result = result.response + + # Convert messages to MCP format + messages = [to_mcp_prompt_message(msg) for msg in result.messages] + + return GetPromptResult(description=prompt.description, messages=messages) + + async def list_resources(self, request: ListResourcesRequest) -> ListResourcesResult: + """Handle MCP requests to list available resources with fixed URIs. + + Args: + request: The MCP ListResourcesRequest. + + Returns: + ListResourcesResult containing resources with fixed URIs. + """ + await self.setup() + + resources: list[Resource] = [] + for action in self.resource_actions: + metadata = action.metadata or {} + resource_meta = metadata.get('resource', {}) + + # Only include resources with fixed URIs (not templates) + if resource_meta.get('uri'): + resources.append( + Resource( + name=action.name, + description=action.description or '', + uri=resource_meta['uri'], + _meta=metadata.get('mcp', {}).get('_meta'), + ) + ) + + return ListResourcesResult(resources=resources) + + async def list_resource_templates(self, request: ListResourceTemplatesRequest) -> ListResourceTemplatesResult: + """Handle MCP requests to list available resource templates. + + Args: + request: The MCP ListResourceTemplatesRequest. + + Returns: + ListResourceTemplatesResult containing resources with URI templates. + """ + await self.setup() + + templates: list[ResourceTemplate] = [] + for action in self.resource_actions: + metadata = action.metadata or {} + resource_meta = metadata.get('resource', {}) + + # Only include resources with templates + if resource_meta.get('template'): + templates.append( + ResourceTemplate( + name=action.name, + description=action.description or '', + uriTemplate=resource_meta['template'], + _meta=metadata.get('mcp', {}).get('_meta'), + ) + ) + + return ListResourceTemplatesResult(resourceTemplates=templates) + + async def read_resource(self, request: ReadResourceRequest) -> ReadResourceResult: + """Handle MCP requests to read a specific resource. + + Args: + request: The MCP ReadResourceRequest containing the resource URI. + + Returns: + ReadResourceResult with the resource content. + + Raises: + GenkitError: If no matching resource is found. + """ + await self.setup() + + uri = request.params.uri + + # Find matching resource (either exact URI match or template match) + resource = None + for action in self.resource_actions: + metadata = action.metadata or {} + resource_meta = metadata.get('resource', {}) + + # Check for exact URI match + if resource_meta.get('uri') == uri: + resource = action + break + + # Check for template match + if resource_meta.get('template'): + if matches_uri_template(resource_meta['template'], uri): + resource = action + break + + if not resource: + raise GenkitError(status='NOT_FOUND', message=f"Tried to call resource '{uri}' but it could not be found.") + + # Execute the resource action + result = await resource.arun({'uri': uri}) + result = result.response + + # Convert content to MCP format + content = result.get('content', []) if isinstance(result, dict) else result.content + contents = to_mcp_resource_contents(uri, content) + + return ReadResourceResult(contents=contents) + + async def start(self, transport: Any = None) -> None: + """Start the MCP server with the specified transport. + + Args: + transport: Optional MCP transport instance. If not provided, + a StdioServerTransport will be created and used. + """ + if not transport: + transport = await stdio_server() + + await self.setup() + + # Connect the transport + async with transport as (read, write): + await self.server.run(read, write, self.server.create_initialization_options()) + + logger.debug(f"[MCP Server] MCP server '{self.options.name}' started successfully.") + + +# Schema imports (these would normally come from mcp.types) +# For now, we'll use the string names +ListToolsRequestSchema = 'ListToolsRequest' +CallToolRequestSchema = 'CallToolRequest' +ListPromptsRequestSchema = 'ListPromptsRequest' +GetPromptRequestSchema = 'GetPromptRequest' +ListResourcesRequestSchema = 'ListResourcesRequest' +ListResourceTemplatesRequestSchema = 'ListResourceTemplatesRequest' +ReadResourceRequestSchema = 'ReadResourceRequest' + + +def create_mcp_server(ai: Genkit, options: McpServerOptions) -> McpServer: + """Create an MCP server based on the supplied Genkit instance. + + All tools, prompts, and resources will be automatically converted to MCP compatibility. + + Args: + ai: Your Genkit instance with registered tools, prompts, and resources. + options: Configuration metadata for the server. + + Returns: + GenkitMcpServer instance. + + Example: + ```python + from genkit.ai import Genkit + from genkit.plugins.mcp import create_mcp_server, McpServerOptions + + ai = Genkit() + + + # Define some tools and resources + @ai.tool() + def add(a: int, b: int) -> int: + return a + b + + + ai.define_resource(name='my_resource', uri='my://resource', fn=lambda req: {'content': [{'text': 'resource content'}]}) + + # Create and start MCP server + server = create_mcp_server(ai, McpServerOptions(name='my-server')) + await server.start() + ``` + """ + return McpServer(ai, options) diff --git a/py/plugins/mcp/src/genkit/plugins/mcp/util/__init__.py b/py/plugins/mcp/src/genkit/plugins/mcp/util/__init__.py new file mode 100644 index 0000000000..9f438c5837 --- /dev/null +++ b/py/plugins/mcp/src/genkit/plugins/mcp/util/__init__.py @@ -0,0 +1,58 @@ +# Copyright 2025 Google LLC +# +# Licensed 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. +# +# SPDX-License-Identifier: Apache-2.0 + +""" +Utility functions for MCP plugin. + +This module contains helper functions for: +- Tool conversion and registration +- Prompt conversion and rendering +- Resource handling +- Message mapping between Genkit and MCP formats +- Transport utilities +""" + +from .message import from_mcp_part, from_mcp_prompt_message, to_mcp_prompt_message +from .prompts import convert_mcp_prompt_messages, convert_prompt_arguments_to_schema, to_mcp_prompt_arguments, to_schema +from .resource import ( + convert_resource_to_genkit_part, + from_mcp_resource_part, + process_resource_content, + to_mcp_resource_contents, +) +from .tools import convert_tool_schema, process_result, process_tool_result, to_mcp_tool_result, to_text +from .transport import create_stdio_params, transport_from + +__all__ = [ + 'process_tool_result', + 'process_result', + 'to_text', + 'convert_tool_schema', + 'convert_prompt_arguments_to_schema', + 'convert_mcp_prompt_messages', + 'to_schema', + 'from_mcp_prompt_message', + 'from_mcp_part', + 'process_resource_content', + 'convert_resource_to_genkit_part', + 'from_mcp_resource_part', + 'create_stdio_params', + 'transport_from', + 'to_mcp_prompt_message', + 'to_mcp_resource_contents', + 'to_mcp_tool_result', + 'to_mcp_prompt_arguments', +] diff --git a/py/plugins/mcp/src/genkit/plugins/mcp/util/message.py b/py/plugins/mcp/src/genkit/plugins/mcp/util/message.py new file mode 100644 index 0000000000..251bac968f --- /dev/null +++ b/py/plugins/mcp/src/genkit/plugins/mcp/util/message.py @@ -0,0 +1,169 @@ +# Copyright 2025 Google LLC +# +# Licensed 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. +# +# SPDX-License-Identifier: Apache-2.0 + +""" +Message utilities for MCP plugin. + +This module contains helper functions for converting between MCP message +formats and Genkit message formats. +""" + +from typing import Any, Dict + +import structlog + +from genkit.core.typing import Message +from mcp.types import ImageContent, PromptMessage, TextContent + +logger = structlog.get_logger(__name__) + +# Role mapping from MCP to Genkit +ROLE_MAP = { + 'user': 'user', + 'assistant': 'model', +} + + +def from_mcp_prompt_message(message: Dict[str, Any]) -> Dict[str, Any]: + """ + Convert MCP PromptMessage to Genkit MessageData format. + + This involves mapping MCP roles (user, assistant) to Genkit roles (user, model) + and transforming the MCP content part into a Genkit Part. + + Args: + message: MCP PromptMessage with 'role' and 'content' fields + + Returns: + Genkit MessageData object with 'role' and 'content' fields + """ + return { + 'role': ROLE_MAP.get(message.get('role', 'user'), 'user'), + 'content': [from_mcp_part(message.get('content', {}))], + } + + +def from_mcp_part(part: Dict[str, Any]) -> Dict[str, Any]: + """ + Convert MCP message content part to Genkit Part. + + Handles different content types: + - Text parts are directly mapped + - Image parts are converted to Genkit media parts with data URL + - Resource parts are mapped to Genkit resource format + + Args: + part: MCP PromptMessage content part + + Returns: + Genkit Part object + """ + part_type = part.get('type', '') + + if part_type == 'text': + return {'text': part.get('text', '')} + + elif part_type == 'image': + mime_type = part.get('mimeType', 'image/png') + data = part.get('data', '') + return { + 'media': { + 'contentType': mime_type, + 'url': f'data:{mime_type};base64,{data}', + } + } + + elif part_type == 'resource': + return { + 'resource': { + 'uri': str(part.get('uri', '')), + } + } + + # Default case for unknown types + return {} + + +def to_mcp_prompt_message(message: Message) -> PromptMessage: + """Convert a Genkit Message to an MCP PromptMessage. + + MCP only supports 'user' and 'assistant' roles. Genkit's 'model' role + is mapped to 'assistant'. + + Args: + message: The Genkit Message to convert. + + Returns: + An MCP PromptMessage. + + Raises: + ValueError: If the message role is not 'user' or 'model'. + ValueError: If media is not a base64 data URL. + """ + # Map Genkit roles to MCP roles + role_map = {'model': 'assistant', 'user': 'user'} + + if message.role not in role_map: + raise ValueError( + f"MCP prompt messages do not support role '{message.role}'. Only 'user' and 'model' messages are supported." + ) + + mcp_role = role_map[message.role] + + # Handle media content (images) + if message.content: + for part in message.content: + part_dict = part if isinstance(part, dict) else part.model_dump() + if part_dict.get('media'): + media = part_dict['media'] + url = media.get('url', '') + content_type = media.get('contentType', '') + + if not url.startswith('data:'): + raise ValueError('MCP prompt messages only support base64 data images.') + + # Extract MIME type and base64 data + mime_type = content_type or url[url.index(':') + 1 : url.index(';')] + data = url[url.index(',') + 1 :] + + return PromptMessage(role=mcp_role, content=ImageContent(type='image', data=data, mimeType=mime_type)) + elif part_dict.get('root') and isinstance(part_dict['root'], dict) and part_dict['root'].get('media'): + media = part_dict['root']['media'] + url = media.get('url', '') + content_type = media.get('contentType', '') + + if not url.startswith('data:'): + raise ValueError('MCP prompt messages only support base64 data images.') + + # Extract MIME type and base64 data + mime_type = content_type or url[url.index(':') + 1 : url.index(';')] + data = url[url.index(',') + 1 :] + + return PromptMessage(role=mcp_role, content=ImageContent(type='image', data=data, mimeType=mime_type)) + + # Handle text content + text = '' + if message.content: + for part in message.content: + part_dict = part if isinstance(part, dict) else part.model_dump() + if 'text' in part_dict and part_dict['text']: + text += part_dict['text'] + elif 'root' in part_dict and isinstance(part_dict['root'], dict) and 'text' in part_dict['root']: + text += part_dict['root']['text'] + elif isinstance(part, str): + text += part + + return PromptMessage(role=mcp_role, content=TextContent(type='text', text=text)) diff --git a/py/plugins/mcp/src/genkit/plugins/mcp/util/prompts.py b/py/plugins/mcp/src/genkit/plugins/mcp/util/prompts.py new file mode 100644 index 0000000000..1020e9fd20 --- /dev/null +++ b/py/plugins/mcp/src/genkit/plugins/mcp/util/prompts.py @@ -0,0 +1,137 @@ +# Copyright 2025 Google LLC +# +# Licensed 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. +# +# SPDX-License-Identifier: Apache-2.0 + +""" +Prompt utilities for MCP plugin. + +This module contains helper functions for converting between MCP prompts +and Genkit prompts, including schema and message conversion. +""" + +from typing import Any, Dict, List, Optional + +import structlog + +from mcp.types import GetPromptResult, Prompt + +logger = structlog.get_logger(__name__) + + +def to_schema(arguments: Optional[List[Dict[str, Any]]]) -> Dict[str, Any]: + """ + Convert MCP prompt arguments to JSON schema format. + + Args: + arguments: List of MCP prompt argument definitions with 'name', + 'description', and 'required' fields + + Returns: + JSON schema representing the prompt arguments + """ + if not arguments: + return {} + + schema: Dict[str, Any] = {'type': 'object', 'properties': {}, 'required': []} + + for arg in arguments: + arg_name = arg.get('name', '') + schema['properties'][arg_name] = { + 'type': 'string', + 'description': arg.get('description', ''), + } + if arg.get('required', False): + schema['required'].append(arg_name) + + return schema + + +def convert_prompt_arguments_to_schema(arguments: List[Any]) -> Dict[str, Any]: + """ + Convert MCP prompt arguments to JSON schema format. + + This is an alias for to_schema() for backwards compatibility. + + Args: + arguments: List of MCP prompt argument definitions + + Returns: + JSON schema representing the prompt arguments + """ + return to_schema(arguments) + + +def convert_mcp_prompt_messages(prompt_result: GetPromptResult) -> List[Dict[str, Any]]: + """ + Convert MCP prompt messages to Genkit message format. + + Args: + prompt_result: The GetPromptResult from MCP server containing messages + + Returns: + List of Genkit-formatted messages + """ + from .message import from_mcp_prompt_message + + if not hasattr(prompt_result, 'messages') or not prompt_result.messages: + return [] + + return [from_mcp_prompt_message(msg) for msg in prompt_result.messages] + + +def to_mcp_prompt_arguments(input_schema: dict[str, Any] | None) -> list[dict[str, Any]] | None: + """Convert Genkit input schema to MCP prompt arguments. + + MCP prompts only support string arguments. This function validates that + all properties in the schema are strings. + + Args: + input_schema: The Genkit input JSON schema. + + Returns: + List of MCP prompt argument definitions, or None if no schema. + + Raises: + ValueError: If the schema is not an object type. + ValueError: If any property is not a string type. + """ + if not input_schema: + return None + + if not input_schema.get('properties'): + raise ValueError('MCP prompts must take objects with properties as input schema.') + + args: list[dict[str, Any]] = [] + properties = input_schema['properties'] + required = input_schema.get('required', []) + + for name, prop in properties.items(): + prop_type = prop.get('type') + + # Check if type is string or includes string (for union types) + is_string = prop_type == 'string' or (isinstance(prop_type, list) and 'string' in prop_type) + + if not is_string: + raise ValueError( + f"MCP prompts may only take string arguments, but property '{name}' has type '{prop_type}'." + ) + + args.append({ + 'name': name, + 'description': prop.get('description'), + 'required': name in required, + }) + + return args diff --git a/py/plugins/mcp/src/genkit/plugins/mcp/util/resource.py b/py/plugins/mcp/src/genkit/plugins/mcp/util/resource.py new file mode 100644 index 0000000000..e0cdc87609 --- /dev/null +++ b/py/plugins/mcp/src/genkit/plugins/mcp/util/resource.py @@ -0,0 +1,146 @@ +# Copyright 2025 Google LLC +# +# Licensed 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. +# +# SPDX-License-Identifier: Apache-2.0 + +""" +Resource utilities for MCP plugin. + +This module contains helper functions for handling MCP resources, +including reading and converting resource content. +""" + +from typing import Any, Dict + +import structlog + +from genkit.core.typing import Part +from mcp.types import BlobResourceContents, ReadResourceResult, Resource, TextResourceContents + +logger = structlog.get_logger(__name__) + + +def from_mcp_resource_part(content: Dict[str, Any]) -> Dict[str, Any]: + """ + Convert MCP resource content to Genkit Part format. + + Handles different content types: + - Text content is mapped to text part + - Blob content is mapped to media part with base64 data + + Args: + content: MCP resource content part + + Returns: + Genkit Part representation + """ + content_type = content.get('type', '') + + if content_type == 'text': + return {'text': content.get('text', '')} + + elif content_type == 'blob': + mime_type = content.get('mimeType', 'application/octet-stream') + blob_data = content.get('blob', '') + return { + 'media': { + 'contentType': mime_type, + 'url': f'data:{mime_type};base64,{blob_data}', + } + } + + # Default case + return {'text': str(content)} + + +def process_resource_content(resource_result: ReadResourceResult) -> Any: + """ + Process MCP ReadResourceResult and extract content. + + Args: + resource_result: The ReadResourceResult from MCP server + + Returns: + Extracted resource content as Genkit Parts + """ + if not hasattr(resource_result, 'contents') or not resource_result.contents: + return [] + + return [from_mcp_resource_part(content) for content in resource_result.contents] + + +def convert_resource_to_genkit_part(resource: Resource) -> dict[str, Any]: + """ + Convert MCP resource to Genkit Part format. + + Args: + resource: MCP resource object + + Returns: + Genkit Part representation with resource URI + """ + return { + 'resource': { + 'uri': resource.uri, + 'name': resource.name, + 'description': resource.description if hasattr(resource, 'description') else None, + } + } + + +def to_mcp_resource_contents(uri: str, parts: list[Part]) -> list[TextResourceContents | BlobResourceContents]: + """Convert Genkit Parts to MCP resource contents. + + Args: + uri: The URI of the resource. + parts: List of Genkit Parts to convert. + + Returns: + List of MCP resource contents (text or blob). + + Raises: + ValueError: If media is not a base64 data URL. + ValueError: If part type is not supported. + """ + contents: list[TextResourceContents | BlobResourceContents] = [] + + for part in parts: + if isinstance(part, dict): + # Handle media/image content + if 'media' in part: + media = part['media'] + url = media.get('url', '') + content_type = media.get('contentType', '') + + if not url.startswith('data:'): + raise ValueError('MCP resource messages only support base64 data images.') + + # Extract MIME type and base64 data + mime_type = content_type or url[url.index(':') + 1 : url.index(';')] + blob_data = url[url.index(',') + 1 :] + + contents.append(BlobResourceContents(uri=uri, mimeType=mime_type, blob=blob_data)) + + # Handle text content + elif 'text' in part: + contents.append(TextResourceContents(uri=uri, text=part['text'])) + else: + raise ValueError( + f'MCP resource messages only support media and text parts. ' + f'Unsupported part type: {list(part.keys())}' + ) + elif isinstance(part, str): + contents.append(TextResourceContents(uri=uri, text=part)) + + return contents diff --git a/py/plugins/mcp/src/genkit/plugins/mcp/util/tools.py b/py/plugins/mcp/src/genkit/plugins/mcp/util/tools.py new file mode 100644 index 0000000000..b2ec40c792 --- /dev/null +++ b/py/plugins/mcp/src/genkit/plugins/mcp/util/tools.py @@ -0,0 +1,146 @@ +# Copyright 2025 Google LLC +# +# Licensed 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. +# +# SPDX-License-Identifier: Apache-2.0 + +""" +Tool utilities for MCP plugin. + +This module contains helper functions for converting between MCP tools +and Genkit actions, processing tool results, and registering tools. +""" + +import json +from typing import Any, Dict, List, Union + +import structlog + +from mcp.types import CallToolResult, ImageContent, TextContent, Tool + +logger = structlog.get_logger(__name__) + + +def to_text(content: List[Dict[str, Any]]) -> str: + """ + Extract text from MCP CallToolResult content. + + Args: + content: List of content parts from CallToolResult + + Returns: + Concatenated text from all text parts + """ + return ''.join(part.get('text', '') for part in content) + + +def process_result(result: CallToolResult) -> Any: + """ + Process MCP CallToolResult and extract/parse content. + + Handles different result types: + - Error results return error dict + - Text-only results attempt JSON parsing + - Single content results return the content directly + - Otherwise returns the full result + + Args: + result: The CallToolResult from MCP server + + Returns: + Processed result (parsed JSON, text, or raw content) + + Raises: + RuntimeError: If the tool execution failed (isError=True) + """ + if result.isError: + return {'error': to_text(result.content)} + + # Check if all content parts are text + if all(hasattr(c, 'text') and c.text for c in result.content): + text = to_text(result.content) + # Try to parse as JSON if it looks like JSON + text_stripped = text.strip() + if text_stripped.startswith('{') or text_stripped.startswith('['): + try: + return json.loads(text) + except (json.JSONDecodeError, ValueError): + return text + return text + + # Single content item + if len(result.content) == 1: + return result.content[0] + + # Return full result for complex cases + return result + + +def process_tool_result(result: CallToolResult) -> Any: + """ + Process MCP CallToolResult and extract content. + + This is an alias for process_result() for backwards compatibility. + + Args: + result: The CallToolResult from MCP server + + Returns: + Extracted text content from the result + + Raises: + RuntimeError: If the tool execution failed + """ + return process_result(result) + + +def convert_tool_schema(mcp_schema: Dict[str, Any]) -> Dict[str, Any]: + """ + Convert MCP tool input schema (JSONSchema7) to Genkit format. + + Args: + mcp_schema: MCP tool input schema + + Returns: + Genkit-compatible JSON schema + + Note: + Currently returns the schema as-is since both use JSON Schema. + Future enhancements may add validation or transformation. + """ + # MCP and Genkit both use JSON Schema, so minimal conversion needed + return mcp_schema + + +def to_mcp_tool_result(result: Any) -> list[TextContent | ImageContent]: + """Convert tool execution result to MCP CallToolResult content. + + Args: + result: The result from tool execution (can be string, dict, or other). + + Returns: + List of MCP content items (TextContent or ImageContent). + """ + if isinstance(result, str): + return [TextContent(type='text', text=result)] + elif isinstance(result, dict): + # If it's already in MCP format, return as-is + if 'type' in result and 'text' in result: + return [TextContent(type='text', text=result['text'])] + # Otherwise, serialize to JSON + import json + + return [TextContent(type='text', text=json.dumps(result))] + else: + # Convert to string for other types + return [TextContent(type='text', text=str(result))] diff --git a/py/plugins/mcp/src/genkit/plugins/mcp/util/transport.py b/py/plugins/mcp/src/genkit/plugins/mcp/util/transport.py new file mode 100644 index 0000000000..10c06601a7 --- /dev/null +++ b/py/plugins/mcp/src/genkit/plugins/mcp/util/transport.py @@ -0,0 +1,89 @@ +# Copyright 2025 Google LLC +# +# Licensed 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. +# +# SPDX-License-Identifier: Apache-2.0 + +""" +Transport utilities for MCP plugin. + +This module contains helper functions for creating and managing +MCP transport connections (stdio, SSE, custom). +""" + +from typing import Any, Dict, Optional, Tuple + +import structlog + +from mcp import StdioServerParameters + +logger = structlog.get_logger(__name__) + + +def create_stdio_params( + command: str, args: Optional[list] = None, env: Optional[Dict[str, str]] = None +) -> StdioServerParameters: + """ + Create StdioServerParameters for MCP connection. + + Args: + command: Command to execute + args: Command arguments + env: Environment variables + + Returns: + StdioServerParameters object + """ + return StdioServerParameters(command=command, args=args or [], env=env) + + +async def transport_from(config: Dict[str, Any], session_id: Optional[str] = None) -> Tuple[Any, str]: + """ + Create an MCP transport instance based on the provided server configuration. + + Supports creating SSE, Stdio, or using a pre-configured custom transport. + + Args: + config: Configuration for the MCP server + session_id: Optional session ID for HTTP transport + + Returns: + Tuple of (transport instance or None, transport type string) + + Note: + This function mirrors the JS SDK's transportFrom() function. + """ + # Handle pre-configured transport first + if 'transport' in config and config['transport']: + return (config['transport'], 'custom') + + # Handle SSE/HTTP config + if 'url' in config and config['url']: + try: + # Dynamic import to avoid hard dependency + from mcp.client.sse import sse_client + + # Note: Python MCP SDK may have different SSE client API + # This is a placeholder that matches the pattern + logger.info(f'Creating SSE transport for URL: {config["url"]}') + return (config['url'], 'http') # Simplified for now + except ImportError: + logger.warning('SSE client not available') + return (None, 'http') + + # Handle Stdio config + if 'command' in config and config['command']: + stdio_params = create_stdio_params(command=config['command'], args=config.get('args'), env=config.get('env')) + return (stdio_params, 'stdio') + + return (None, 'unknown') diff --git a/py/plugins/mcp/tests/fakes.py b/py/plugins/mcp/tests/fakes.py new file mode 100644 index 0000000000..a73cd0f0c8 --- /dev/null +++ b/py/plugins/mcp/tests/fakes.py @@ -0,0 +1,128 @@ +# Copyright 2025 Google LLC +# +# Licensed 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. +# +# SPDX-License-Identifier: Apache-2.0 + +import asyncio +import json +import sys +from typing import Any, Callable, Dict, List, Optional +from unittest.mock import MagicMock + +from genkit.ai import Genkit +from genkit.core.action.types import ActionKind + + +class MockSchema: + def __init__(self, **kwargs): + for k, v in kwargs.items(): + setattr(self, k, v) + + +def mock_mcp_modules(): + """Sets up comprehensive MCP mocks in sys.modules.""" + mock_mcp = MagicMock() + sys.modules['mcp'] = mock_mcp + sys.modules['mcp'].__path__ = [] + + types_mock = MagicMock() + sys.modules['mcp.types'] = types_mock + types_mock.ListToolsResult = MockSchema + types_mock.CallToolResult = MockSchema + types_mock.ListPromptsResult = MockSchema + types_mock.GetPromptResult = MockSchema + types_mock.ListResourcesResult = MockSchema + types_mock.ListResourceTemplatesResult = MockSchema + types_mock.ReadResourceResult = MockSchema + types_mock.Tool = MockSchema + types_mock.Prompt = MockSchema + types_mock.Resource = MockSchema + types_mock.ResourceTemplate = MockSchema + types_mock.TextContent = MockSchema + types_mock.PromptMessage = MockSchema + types_mock.TextResourceContents = MockSchema + types_mock.BlobResourceContents = MockSchema + types_mock.ImageContent = MockSchema + + sys.modules['mcp.server'] = MagicMock() + sys.modules['mcp.server.stdio'] = MagicMock() + sys.modules['mcp.client'] = MagicMock() + sys.modules['mcp.client'].__path__ = [] + sys.modules['mcp.client.stdio'] = MagicMock() + sys.modules['mcp.client.sse'] = MagicMock() + sys.modules['mcp.server.sse'] = MagicMock() + + return mock_mcp, types_mock + + +def define_echo_model(ai: Genkit): + """Defines a fake echo model for testing.""" + + @ai.tool(name='echoModel') + def echo_model(request: Any): + # This is a simplified mock of a model action + # Real model action would handle GenerateRequest and return GenerateResponse + + # logic to echo content + # For now, just a placeholder as we generally mock the model execution in tests + pass + + # In real usage, we would define a Model action properly. + # For unit tests here, we might not strictly need the full model implementation + # if we are mocking the generation or call. + # But matching JS behavior: + # JS defines 'echoModel' which returns "Echo: " + input. + + # We can use ai.define_model if available or just mock it. + pass + + +class FakeTransport: + """Fakes an MCP transport/server for testing.""" + + def __init__(self): + self.tools = [] + self.prompts = [] + self.resources = [] + self.resource_templates = [] + self.call_tool_result = None + self.get_prompt_result = None + self.read_resource_result = None + self.roots = [] + + # Callbacks that would simulate transport behavior + self.on_message = None + self.on_close = None + self.on_error = None + + async def start(self): + pass + + async def send(self, message: Dict[str, Any]): + """Handle incoming JSON-RPC message (simulating server).""" + request = message + # msg_id = request.get("id") + + # In a real transport we'd write back to the stream. + # Here we just store handling logic or print. + # Since we are mocking the ClientSession in our python tests, + # this logic might need to be hooked up to the mock session's methods. + pass + + # Helper methods to populate the fake state + def add_tool(self, name: str, description: str = '', schema: Dict = None): + self.tools.append({'name': name, 'description': description, 'inputSchema': schema or {'type': 'object'}}) + + def add_prompt(self, name: str, description: str = '', arguments: List = None): + self.prompts.append({'name': name, 'description': description, 'arguments': arguments or []}) diff --git a/py/plugins/mcp/tests/test_mcp_conversion.py b/py/plugins/mcp/tests/test_mcp_conversion.py new file mode 100644 index 0000000000..da2032ebd7 --- /dev/null +++ b/py/plugins/mcp/tests/test_mcp_conversion.py @@ -0,0 +1,259 @@ +# Copyright 2025 Google LLC +# +# Licensed 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. +# +# SPDX-License-Identifier: Apache-2.0 + +"""Tests for MCP conversion utilities.""" + +import os +import sys +import unittest + +sys.path.insert(0, os.path.dirname(__file__)) +sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), '../src'))) +from fakes import mock_mcp_modules + +mock_mcp_modules() + +from genkit.core.typing import Message +from genkit.plugins.mcp.util import ( + to_mcp_prompt_arguments, + to_mcp_prompt_message, + to_mcp_resource_contents, + to_mcp_tool_result, +) + + +class TestMessageConversion(unittest.TestCase): + """Tests for message conversion utilities.""" + + def test_convert_user_message(self): + """Test converting a user message.""" + message = Message(role='user', content=[{'text': 'Hello, world!'}]) + + result = to_mcp_prompt_message(message) + + self.assertEqual(result.role, 'user') + self.assertEqual(result.content.type, 'text') + self.assertEqual(result.content.text, 'Hello, world!') + + def test_convert_model_message(self): + """Test converting a model message (maps to assistant).""" + message = Message(role='model', content=[{'text': 'Hi there!'}]) + + result = to_mcp_prompt_message(message) + + self.assertEqual(result.role, 'assistant') + self.assertEqual(result.content.type, 'text') + self.assertEqual(result.content.text, 'Hi there!') + + def test_convert_message_with_multiple_text_parts(self): + """Test converting a message with multiple text parts.""" + message = Message(role='user', content=[{'text': 'Part 1 '}, {'text': 'Part 2 '}, {'text': 'Part 3'}]) + + result = to_mcp_prompt_message(message) + + self.assertEqual(result.content.text, 'Part 1 Part 2 Part 3') + + def test_convert_message_with_invalid_role(self): + """Test that converting a message with invalid role raises error.""" + message = Message(role='system', content=[{'text': 'System message'}]) + + with self.assertRaises(ValueError) as context: + to_mcp_prompt_message(message) + + self.assertIn('system', str(context.exception).lower()) + + def test_convert_message_with_image(self): + """Test converting a message with image content.""" + message = Message( + role='user', content=[{'media': {'url': 'data:image/png;base64,iVBORw0KG...', 'contentType': 'image/png'}}] + ) + + result = to_mcp_prompt_message(message) + + self.assertEqual(result.role, 'user') + self.assertEqual(result.content.type, 'image') + self.assertEqual(result.content.mimeType, 'image/png') + + def test_convert_message_with_non_data_url_fails(self): + """Test that non-data URLs raise an error.""" + message = Message(role='user', content=[{'media': {'url': 'http://example.com/image.png'}}]) + + with self.assertRaises(ValueError) as context: + to_mcp_prompt_message(message) + + self.assertIn('base64', str(context.exception).lower()) + + +class TestResourceConversion(unittest.TestCase): + """Tests for resource content conversion.""" + + def test_convert_text_resource(self): + """Test converting text resource content.""" + parts = [{'text': 'Resource content'}] + + result = to_mcp_resource_contents('test://resource', parts) + + self.assertEqual(len(result), 1) + self.assertEqual(result[0].uri, 'test://resource') + self.assertEqual(result[0].text, 'Resource content') + + def test_convert_multiple_text_parts(self): + """Test converting multiple text parts.""" + parts = [{'text': 'Part 1'}, {'text': 'Part 2'}, {'text': 'Part 3'}] + + result = to_mcp_resource_contents('test://resource', parts) + + self.assertEqual(len(result), 3) + for i, part in enumerate(result, 1): + self.assertEqual(part.text, f'Part {i}') + + def test_convert_string_parts(self): + """Test converting string parts.""" + parts = ['Text 1', 'Text 2'] + + result = to_mcp_resource_contents('test://resource', parts) + + self.assertEqual(len(result), 2) + self.assertEqual(result[0].text, 'Text 1') + self.assertEqual(result[1].text, 'Text 2') + + def test_convert_media_resource(self): + """Test converting media resource content.""" + parts = [{'media': {'url': 'data:image/png;base64,abc123', 'contentType': 'image/png'}}] + + result = to_mcp_resource_contents('test://image', parts) + + self.assertEqual(len(result), 1) + self.assertEqual(result[0].uri, 'test://image') + self.assertEqual(result[0].mimeType, 'image/png') + self.assertEqual(result[0].blob, 'abc123') + + def test_convert_mixed_content(self): + """Test converting mixed text and media content.""" + parts = [{'text': 'Description'}, {'media': {'url': 'data:image/png;base64,xyz', 'contentType': 'image/png'}}] + + result = to_mcp_resource_contents('test://mixed', parts) + + self.assertEqual(len(result), 2) + self.assertEqual(result[0].text, 'Description') + self.assertEqual(result[1].blob, 'xyz') + + +class TestToolResultConversion(unittest.TestCase): + """Tests for tool result conversion.""" + + def test_convert_string_result(self): + """Test converting string result.""" + result = to_mcp_tool_result('Hello, world!') + + self.assertEqual(len(result), 1) + self.assertEqual(result[0].type, 'text') + self.assertEqual(result[0].text, 'Hello, world!') + + def test_convert_dict_result(self): + """Test converting dict result.""" + result = to_mcp_tool_result({'key': 'value', 'number': 42}) + + self.assertEqual(len(result), 1) + self.assertEqual(result[0].type, 'text') + # Should be JSON serialized + import json + + parsed = json.loads(result[0].text) + self.assertEqual(parsed['key'], 'value') + self.assertEqual(parsed['number'], 42) + + def test_convert_number_result(self): + """Test converting number result.""" + result = to_mcp_tool_result(42) + + self.assertEqual(len(result), 1) + self.assertEqual(result[0].text, '42') + + def test_convert_boolean_result(self): + """Test converting boolean result.""" + result = to_mcp_tool_result(True) + + self.assertEqual(len(result), 1) + self.assertEqual(result[0].text, 'True') + + +class TestSchemaConversion(unittest.TestCase): + """Tests for schema conversion utilities.""" + + def test_convert_simple_schema(self): + """Test converting simple string schema.""" + schema = {'type': 'object', 'properties': {'name': {'type': 'string', 'description': 'User name'}}} + + result = to_mcp_prompt_arguments(schema) + + self.assertIsNotNone(result) + self.assertEqual(len(result), 1) + self.assertEqual(result[0]['name'], 'name') + self.assertEqual(result[0]['description'], 'User name') + + def test_convert_schema_with_required(self): + """Test converting schema with required fields.""" + schema = { + 'type': 'object', + 'properties': {'name': {'type': 'string'}, 'age': {'type': 'string'}}, + 'required': ['name'], + } + + result = to_mcp_prompt_arguments(schema) + + name_arg = next(arg for arg in result if arg['name'] == 'name') + age_arg = next(arg for arg in result if arg['name'] == 'age') + + self.assertTrue(name_arg['required']) + self.assertFalse(age_arg['required']) + + def test_convert_schema_with_non_string_fails(self): + """Test that non-string properties raise an error.""" + schema = {'type': 'object', 'properties': {'count': {'type': 'number'}}} + + with self.assertRaises(ValueError) as context: + to_mcp_prompt_arguments(schema) + + self.assertIn('string', str(context.exception).lower()) + + def test_convert_schema_with_union_type(self): + """Test converting schema with union type including string.""" + schema = {'type': 'object', 'properties': {'value': {'type': ['string', 'null']}}} + + result = to_mcp_prompt_arguments(schema) + + # Should succeed because string is in the union + self.assertEqual(len(result), 1) + + def test_convert_none_schema(self): + """Test converting None schema.""" + result = to_mcp_prompt_arguments(None) + + self.assertIsNone(result) + + def test_convert_schema_without_properties_fails(self): + """Test that schema without properties raises an error.""" + schema = {'type': 'object'} + + with self.assertRaises(ValueError) as context: + to_mcp_prompt_arguments(schema) + + self.assertIn('properties', str(context.exception).lower()) + + +if __name__ == '__main__': + unittest.main() diff --git a/py/plugins/mcp/tests/test_mcp_host.py b/py/plugins/mcp/tests/test_mcp_host.py new file mode 100644 index 0000000000..ce8832179d --- /dev/null +++ b/py/plugins/mcp/tests/test_mcp_host.py @@ -0,0 +1,64 @@ +# Copyright 2025 Google LLC +# +# Licensed 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. +# +# SPDX-License-Identifier: Apache-2.0 + +import os +import sys +from unittest.mock import AsyncMock, MagicMock + +sys.path.insert(0, os.path.dirname(__file__)) +sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), '../src'))) +from fakes import mock_mcp_modules + +mock_mcp_modules() + +import unittest +from unittest.mock import patch + +from genkit.ai import Genkit +from genkit.core.action.types import ActionKind + +# Now import plugin +from genkit.plugins.mcp import McpClient, McpHost, McpServerConfig, create_mcp_host + + +class TestMcpHost(unittest.IsolatedAsyncioTestCase): + async def test_connect_and_register(self): + # Setup configs + config1 = McpServerConfig(command='echo') + config2 = McpServerConfig(url='http://localhost:8000') + + host = create_mcp_host({'server1': config1, 'server2': config2}) + + # Mock clients within host + with patch('genkit.plugins.mcp.client.client.McpClient.connect', new_callable=AsyncMock) as mock_connect: + await host.start() + self.assertEqual(mock_connect.call_count, 2) + + # Mock session for registration + host.clients['server1'].session = AsyncMock() + mock_tool = MagicMock() + mock_tool.name = 'tool1' + host.clients['server1'].session.list_tools.return_value.tools = [mock_tool] + + ai = MagicMock(spec=Genkit) + ai.registry = MagicMock() + + await host.register_tools(ai) + + # Verify tool registration + ai.registry.register_action.assert_called() + call_args = ai.registry.register_action.call_args[1] + self.assertIn('server1/tool1', call_args['name']) diff --git a/py/plugins/mcp/tests/test_mcp_integration.py b/py/plugins/mcp/tests/test_mcp_integration.py new file mode 100644 index 0000000000..8a7c2ebce9 --- /dev/null +++ b/py/plugins/mcp/tests/test_mcp_integration.py @@ -0,0 +1,311 @@ +# Copyright 2025 Google LLC +# +# Licensed 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. +# +# SPDX-License-Identifier: Apache-2.0 + +"""Integration tests for MCP client-server communication.""" + +import asyncio +import os +import sys +import unittest +from unittest.mock import AsyncMock, MagicMock, patch + +sys.path.insert(0, os.path.dirname(__file__)) +sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), '../src'))) +from fakes import mock_mcp_modules + +mock_mcp_modules() + +import pytest + +from genkit.ai import Genkit +from genkit.plugins.mcp import McpClient, McpHost, McpServerConfig, create_mcp_host, create_mcp_server + + +@pytest.mark.asyncio +class TestClientServerIntegration(unittest.IsolatedAsyncioTestCase): + """Integration tests for MCP client-server communication.""" + + async def test_client_can_list_server_tools(self): + """Test that a client can list tools from a server.""" + # Create server with tools + server_ai = Genkit() + + @server_ai.tool() + def add(a: int, b: int) -> int: + return a + b + + # Create client + client = McpClient(name='test-client', config=McpServerConfig(command='echo', args=['test'])) + + # Mock the session to return tools + mock_session = AsyncMock() + mock_tool = MagicMock() + mock_tool.name = 'add' + mock_tool.description = 'Add two numbers' + mock_tool.inputSchema = {'type': 'object'} + + mock_session.list_tools.return_value.tools = [mock_tool] + client.session = mock_session + + # List tools + tools = await client.list_tools() + + # Verify + self.assertEqual(len(tools), 1) + self.assertEqual(tools[0].name, 'add') + + async def test_client_can_call_server_tool(self): + """Test that a client can call a tool on a server.""" + # Create client + client = McpClient(name='test-client', config=McpServerConfig(command='echo')) + + # Mock the session + mock_session = AsyncMock() + mock_result = MagicMock() + mock_result.isError = False + mock_content = MagicMock() + mock_content.type = 'text' + mock_content.text = '8' + mock_result.content = [mock_content] + + mock_session.call_tool.return_value = mock_result + client.session = mock_session + + # Call tool + result = await client.call_tool('add', {'a': 5, 'b': 3}) + + # Verify + self.assertEqual(result, '8') + mock_session.call_tool.assert_called_once_with('add', {'a': 5, 'b': 3}) + + async def test_client_can_list_server_resources(self): + """Test that a client can list resources from a server.""" + # Create client + client = McpClient(name='test-client', config=McpServerConfig(command='echo')) + + # Mock the session + mock_session = AsyncMock() + mock_resource = MagicMock() + mock_resource.name = 'config' + mock_resource.uri = 'app://config' + mock_resource.description = 'Configuration' + + mock_session.list_resources.return_value.resources = [mock_resource] + client.session = mock_session + + # List resources + resources = await client.list_resources() + + # Verify + self.assertEqual(len(resources), 1) + self.assertEqual(resources[0].name, 'config') + self.assertEqual(resources[0].uri, 'app://config') + + async def test_client_can_read_server_resource(self): + """Test that a client can read a resource from a server.""" + # Create client + client = McpClient(name='test-client', config=McpServerConfig(command='echo')) + + # Mock the session + mock_session = AsyncMock() + mock_result = MagicMock() + mock_result.contents = [MagicMock(text='Resource content')] + + mock_session.read_resource.return_value = mock_result + client.session = mock_session + + # Read resource + result = await client.read_resource('app://config') + + # Verify + self.assertIsNotNone(result) + mock_session.read_resource.assert_called_once_with('app://config') + + async def test_host_manages_multiple_clients(self): + """Test that a host can manage multiple clients.""" + # Create host with multiple servers + config1 = McpServerConfig(command='server1') + config2 = McpServerConfig(command='server2') + + host = create_mcp_host({'server1': config1, 'server2': config2}) + + # Verify clients were created + self.assertEqual(len(host.clients), 2) + self.assertIn('server1', host.clients) + self.assertIn('server2', host.clients) + + async def test_host_can_register_tools_from_multiple_servers(self): + """Test that a host can register tools from multiple servers.""" + # Create host + host = create_mcp_host({'server1': McpServerConfig(command='s1'), 'server2': McpServerConfig(command='s2')}) + + # Mock sessions for both clients + for client_name, client in host.clients.items(): + mock_session = AsyncMock() + mock_tool = MagicMock() + mock_tool.name = f'{client_name}_tool' + mock_tool.description = f'Tool from {client_name}' + mock_tool.inputSchema = {'type': 'object'} + + mock_session.list_tools.return_value.tools = [mock_tool] + client.session = mock_session + + # Register tools + ai = Genkit() + await host.register_tools(ai) + + # Verify tools were registered + # Each client should have registered one tool + # Tool names should be prefixed with server name + + async def test_client_handles_disabled_server(self): + """Test that a client handles disabled servers correctly.""" + # Create client with disabled config + config = McpServerConfig(command='echo', disabled=True) + client = McpClient(name='test-client', config=config) + + # Try to connect + await client.connect() + + # Should not have a session + self.assertIsNone(client.session) + + async def test_host_can_disable_and_enable_clients(self): + """Test that a host can disable and enable clients.""" + host = create_mcp_host({'test': McpServerConfig(command='echo')}) + + # Mock the client + client = host.clients['test'] + client.session = AsyncMock() + client.close = AsyncMock() + client.connect = AsyncMock() + + # Disable + await host.disable('test') + self.assertTrue(client.config.disabled) + + # Enable + await host.enable('test') + self.assertFalse(client.config.disabled) + + +@pytest.mark.asyncio +class TestResourceIntegration(unittest.IsolatedAsyncioTestCase): + """Integration tests specifically for resource handling.""" + + async def test_end_to_end_resource_flow(self): + """Test complete flow: define resource → expose via server → consume via client.""" + # This is a conceptual test showing the flow + # In practice, we'd need actual MCP transport for true end-to-end + + # 1. Server side: Define resource + server_ai = Genkit() + server_ai.define_resource( + name='config', uri='app://config', fn=lambda req: {'content': [{'text': 'config data'}]} + ) + + # 2. Create MCP server + from genkit.plugins.mcp import McpServerOptions + + server = create_mcp_server(server_ai, McpServerOptions(name='test-server')) + await server.setup() + + # 3. Verify server can list resources + resources_result = await server.list_resources({}) + self.assertEqual(len(resources_result.resources), 1) + self.assertEqual(resources_result.resources[0].uri, 'app://config') + + # 4. Verify server can read resource + request = MagicMock() + request.params.uri = 'app://config' + read_result = await server.read_resource(request) + self.assertEqual(read_result.contents[0].text, 'config data') + + async def test_template_resource_matching(self): + """Test that template resources match correctly.""" + server_ai = Genkit() + + def file_resource(req): + uri = req['uri'] + return {'content': [{'text': f'Contents of {uri}'}]} + + server_ai.define_resource(name='file', template='file://{path}', fn=file_resource) + + # Create server + from genkit.plugins.mcp import McpServerOptions + + server = create_mcp_server(server_ai, McpServerOptions(name='test-server')) + await server.setup() + + # List templates + templates_result = await server.list_resource_templates({}) + self.assertEqual(len(templates_result.resourceTemplates), 1) + self.assertEqual(templates_result.resourceTemplates[0].uriTemplate, 'file://{path}') + + # Read with different URIs + for test_uri in ['file:///path/to/file.txt', 'file:///another/file.md', 'file:///deep/nested/path/doc.pdf']: + request = MagicMock() + request.params.uri = test_uri + result = await server.read_resource(request) + self.assertIn(test_uri, result.contents[0].text) + + +@pytest.mark.asyncio +class TestErrorHandling(unittest.IsolatedAsyncioTestCase): + """Tests for error handling in client-server communication.""" + + async def test_server_handles_missing_tool(self): + """Test that server properly handles requests for non-existent tools.""" + server_ai = Genkit() + + @server_ai.tool() + def existing_tool(x: int) -> int: + return x + + from genkit.plugins.mcp import McpServerOptions + + server = create_mcp_server(server_ai, McpServerOptions(name='test-server')) + await server.setup() + + # Try to call non-existent tool + request = MagicMock() + request.params.name = 'nonexistent_tool' + request.params.arguments = {} + + from genkit.core.error import GenkitError + + with self.assertRaises(GenkitError) as context: + await server.call_tool(request) + + self.assertIn('NOT_FOUND', str(context.exception.status)) + + async def test_client_handles_connection_failure(self): + """Test that client handles connection failures gracefully.""" + client = McpClient(name='test-client', config=McpServerConfig(command='nonexistent_command')) + + # Mock the connection to fail + with patch('genkit.plugins.mcp.client.client.stdio_client') as mock_stdio: + mock_stdio.side_effect = Exception('Connection failed') + + with self.assertRaises(Exception): + await client.connect() + + # Client should mark server as disabled + self.assertTrue(client.config.disabled) + + +if __name__ == '__main__': + unittest.main() diff --git a/py/plugins/mcp/tests/test_mcp_server.py b/py/plugins/mcp/tests/test_mcp_server.py new file mode 100644 index 0000000000..c09ea9caf6 --- /dev/null +++ b/py/plugins/mcp/tests/test_mcp_server.py @@ -0,0 +1,341 @@ +# Copyright 2025 Google LLC +# +# Licensed 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. +# +# SPDX-License-Identifier: Apache-2.0 + +""" +MCP Server Tests + +Mirrors the functionality of js/plugins/mcp/tests/server_test.ts +Tests tools, prompts, and resources exposed via MCP server. +""" + +import os +import sys +import unittest +from unittest.mock import AsyncMock, MagicMock, patch + +sys.path.insert(0, os.path.dirname(__file__)) +sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), '../src'))) + +# Mock mcp module before importing +mock_mcp = MagicMock() +sys.modules['mcp'] = mock_mcp + + +class MockSchema: + def __init__(self, **kwargs): + for k, v in kwargs.items(): + setattr(self, k, v) + + +types_mock = MagicMock() +sys.modules['mcp.types'] = types_mock +types_mock.ListToolsResult = MockSchema +types_mock.CallToolResult = MockSchema +types_mock.ListPromptsResult = MockSchema +types_mock.GetPromptResult = MockSchema +types_mock.ListResourcesResult = MockSchema +types_mock.ListResourceTemplatesResult = MockSchema +types_mock.ReadResourceResult = MockSchema +types_mock.Tool = MockSchema +types_mock.Prompt = MockSchema +types_mock.Resource = MockSchema +types_mock.ResourceTemplate = MockSchema +types_mock.TextResourceContents = MockSchema +types_mock.BlobResourceContents = MockSchema +types_mock.ImageContent = MockSchema +types_mock.TextResourceContents = MockSchema +types_mock.BlobResourceContents = MockSchema +types_mock.ImageContent = MockSchema +types_mock.TextContent = MockSchema +types_mock.PromptMessage = MockSchema + +sys.modules['mcp.server'] = MagicMock() +sys.modules['mcp.server.stdio'] = MagicMock() +sys.modules['mcp.client'] = MagicMock() +sys.modules['mcp.client'].__path__ = [] +sys.modules['mcp.client.stdio'] = MagicMock() +sys.modules['mcp.client.sse'] = MagicMock() +sys.modules['mcp.server.sse'] = MagicMock() + +import pytest + +from genkit.ai import Genkit +from genkit.core.action.types import ActionKind +from genkit.plugins.mcp import McpServer, McpServerOptions, create_mcp_server + + +@pytest.mark.asyncio +class TestMcpServer(unittest.IsolatedAsyncioTestCase): + """Test MCP server functionality - mirrors JS server_test.ts""" + + def setUp(self): + """Set up test fixtures before each test.""" + self.ai = Genkit() + + # Define test tool + @self.ai.tool(description='test tool') + def test_tool(input: dict[str, str]) -> str: + foo = input.get('foo', '') + return f'yep {{"foo":"{foo}"}}' + + # Define test prompt + self.ai.define_prompt(name='testPrompt', model='test-model', prompt='prompt says: {{input}}') + + # Define test resource with fixed URI + self.ai.define_resource( + name='testResources', uri='my://resource', fn=lambda req: {'content': [{'text': 'my resource'}]} + ) + + # Define test resource with template + self.ai.define_resource( + name='testTmpl', + template='file://{path}', + fn=lambda req: {'content': [{'text': f'file contents for {req["uri"]}'}]}, + ) + + # Create MCP server + self.server = create_mcp_server(self.ai, McpServerOptions(name='test-server', version='0.0.1')) + + async def asyncSetUp(self): + """Async setup - initialize server.""" + await self.server.setup() + + # ===== TOOL TESTS ===== + + async def test_list_tools(self): + """Test listing tools - mirrors JS 'should list tools'.""" + result = await self.server.list_tools({}) + + # Verify we have the test tool + self.assertEqual(len(result.tools), 1) + tool = result.tools[0] + + self.assertEqual(tool.name, 'test_tool') + self.assertEqual(tool.description, 'test tool') + self.assertIsNotNone(tool.inputSchema) + + async def test_call_tool(self): + """Test calling a tool - mirrors JS 'should call the tool'.""" + # Create mock request + request = MagicMock() + request.params.name = 'test_tool' + request.params.arguments = {'foo': 'bar'} + + result = await self.server.call_tool(request) + + # Verify response + self.assertEqual(len(result.content), 1) + self.assertEqual(result.content[0].type, 'text') + self.assertEqual(result.content[0].text, 'yep {"foo":"bar"}') + + # ===== PROMPT TESTS ===== + + async def test_list_prompts(self): + """Test listing prompts - mirrors JS 'should list prompts'.""" + result = await self.server.list_prompts({}) + + # Verify we have the test prompt + prompt_names = [p.name for p in result.prompts] + self.assertIn('testPrompt', prompt_names) + + async def test_get_prompt(self): + """Test rendering a prompt - mirrors JS 'should render prompt'.""" + # Create mock request + request = MagicMock() + request.params.name = 'testPrompt' + request.params.arguments = {'input': 'hello'} + + result = await self.server.get_prompt(request) + + # Verify response + self.assertIsNotNone(result.messages) + self.assertGreater(len(result.messages), 0) + + # Check message content + message = result.messages[0] + self.assertEqual(message.role, 'user') + self.assertEqual(message.content.type, 'text') + self.assertIn('prompt says: hello', message.content.text) + + # ===== RESOURCE TESTS ===== + + async def test_list_resources(self): + """Test listing resources - mirrors JS 'should list resources'.""" + result = await self.server.list_resources({}) + + # Verify we have the fixed URI resource + self.assertEqual(len(result.resources), 1) + resource = result.resources[0] + + self.assertEqual(resource.name, 'testResources') + self.assertEqual(resource.uri, 'my://resource') + + async def test_list_resource_templates(self): + """Test listing resource templates - mirrors JS 'should list templates'.""" + result = await self.server.list_resource_templates({}) + + # Verify we have the template resource + self.assertEqual(len(result.resourceTemplates), 1) + template = result.resourceTemplates[0] + + self.assertEqual(template.name, 'testTmpl') + self.assertEqual(template.uriTemplate, 'file://{path}') + + async def test_read_resource(self): + """Test reading a resource - mirrors JS 'should read resource'.""" + # Create mock request + request = MagicMock() + request.params.uri = 'my://resource' + + result = await self.server.read_resource(request) + + # Verify response + self.assertEqual(len(result.contents), 1) + content = result.contents[0] + + self.assertEqual(content.uri, 'my://resource') + self.assertEqual(content.text, 'my resource') + + async def test_read_template_resource(self): + """Test reading a template resource.""" + # Create mock request + request = MagicMock() + request.params.uri = 'file:///path/to/file.txt' + + result = await self.server.read_resource(request) + + # Verify response + self.assertEqual(len(result.contents), 1) + content = result.contents[0] + + self.assertEqual(content.uri, 'file:///path/to/file.txt') + self.assertIn('file contents for file:///path/to/file.txt', content.text) + + # ===== ADDITIONAL TESTS ===== + + async def test_server_initialization(self): + """Test that server initializes correctly.""" + self.assertIsNotNone(self.server) + self.assertEqual(self.server.options.name, 'test-server') + self.assertEqual(self.server.options.version, '0.0.1') + self.assertTrue(self.server.actions_resolved) + + async def test_server_has_all_action_types(self): + """Test that server has tools, prompts, and resources.""" + self.assertGreater(len(self.server.tool_actions), 0) + self.assertGreater(len(self.server.prompt_actions), 0) + self.assertGreater(len(self.server.resource_actions), 0) + + async def test_tool_not_found(self): + """Test calling a non-existent tool.""" + from genkit.core.error import GenkitError + + request = MagicMock() + request.params.name = 'nonexistent_tool' + request.params.arguments = {} + + with self.assertRaises(GenkitError) as context: + await self.server.call_tool(request) + + self.assertEqual(context.exception.status, 'NOT_FOUND') + + async def test_prompt_not_found(self): + """Test getting a non-existent prompt.""" + from genkit.core.error import GenkitError + + request = MagicMock() + request.params.name = 'nonexistent_prompt' + request.params.arguments = {} + + with self.assertRaises(GenkitError) as context: + await self.server.get_prompt(request) + + self.assertEqual(context.exception.status, 'NOT_FOUND') + + async def test_resource_not_found(self): + """Test reading a non-existent resource.""" + from genkit.core.error import GenkitError + + request = MagicMock() + request.params.uri = 'nonexistent://resource' + + with self.assertRaises(GenkitError) as context: + await self.server.read_resource(request) + + self.assertEqual(context.exception.status, 'NOT_FOUND') + + +# Additional test class for resource-specific functionality +@pytest.mark.asyncio +class TestResourceFunctionality(unittest.IsolatedAsyncioTestCase): + """Test resource-specific functionality.""" + + async def test_resource_registration_with_fixed_uri(self): + """Test registering a resource with fixed URI.""" + ai = Genkit() + + action = ai.define_resource( + name='test_resource', uri='test://resource', fn=lambda req: {'content': [{'text': 'test'}]} + ) + + self.assertIsNotNone(action) + self.assertEqual(action.kind, ActionKind.RESOURCE) + self.assertEqual(action.metadata['resource']['uri'], 'test://resource') + + async def test_resource_registration_with_template(self): + """Test registering a resource with URI template.""" + ai = Genkit() + + action = ai.define_resource( + name='file', template='file://{path}', fn=lambda req: {'content': [{'text': 'file content'}]} + ) + + self.assertIsNotNone(action) + self.assertEqual(action.kind, ActionKind.RESOURCE) + self.assertEqual(action.metadata['resource']['template'], 'file://{path}') + + async def test_resource_requires_uri_or_template(self): + """Test that resource requires either uri or template.""" + ai = Genkit() + + with self.assertRaises(ValueError) as context: + ai.define_resource(name='invalid', fn=lambda req: {'content': []}) + + self.assertIn('uri', str(context.exception).lower()) + self.assertIn('template', str(context.exception).lower()) + + async def test_uri_template_matching(self): + """Test URI template matching.""" + from genkit.blocks.resource import matches_uri_template + + # Test exact match + result = matches_uri_template('file://{path}', 'file:///home/user/doc.txt') + self.assertIsNotNone(result) + self.assertIn('path', result) + + # Test no match + result = matches_uri_template('file://{path}', 'http://example.com') + self.assertIsNone(result) + + # Test multiple parameters + result = matches_uri_template('user://{id}/posts/{post_id}', 'user://123/posts/456') + self.assertIsNotNone(result) + self.assertEqual(result['id'], '123') + self.assertEqual(result['post_id'], '456') + + +if __name__ == '__main__': + unittest.main() diff --git a/py/plugins/mcp/tests/test_mcp_server_resources.py b/py/plugins/mcp/tests/test_mcp_server_resources.py new file mode 100644 index 0000000000..87ff45904b --- /dev/null +++ b/py/plugins/mcp/tests/test_mcp_server_resources.py @@ -0,0 +1,351 @@ +# Copyright 2025 Google LLC +# +# Licensed 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. +# +# SPDX-License-Identifier: Apache-2.0 + +"""Comprehensive tests for MCP server resource handling.""" + +import os +import sys +import unittest +from unittest.mock import AsyncMock, MagicMock, patch + +sys.path.insert(0, os.path.dirname(__file__)) +sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), '../src'))) +from fakes import mock_mcp_modules + +mock_mcp_modules() + +import pytest + +from genkit.ai import Genkit +from genkit.core.action.types import ActionKind +from genkit.plugins.mcp import McpServer, McpServerOptions, create_mcp_server + + +@pytest.mark.asyncio +class TestMcpServerResources(unittest.IsolatedAsyncioTestCase): + """Tests for MCP server resource handling.""" + + def setUp(self): + """Set up test fixtures.""" + self.ai = Genkit() + + async def test_list_resources_with_fixed_uri(self): + """Test listing resources with fixed URIs.""" + # Define resources + self.ai.define_resource(name='config', uri='app://config', fn=lambda req: {'content': [{'text': 'config'}]}) + + self.ai.define_resource(name='data', uri='app://data', fn=lambda req: {'content': [{'text': 'data'}]}) + + # Create server + server = create_mcp_server(self.ai, McpServerOptions(name='test-server')) + await server.setup() + + # List resources + result = await server.list_resources({}) + + # Verify + self.assertEqual(len(result.resources), 2) + resource_names = [r.name for r in result.resources] + self.assertIn('config', resource_names) + self.assertIn('data', resource_names) + + # Verify URIs + config_resource = next(r for r in result.resources if r.name == 'config') + self.assertEqual(config_resource.uri, 'app://config') + + async def test_list_resource_templates(self): + """Test listing resources with URI templates.""" + # Define template resources + self.ai.define_resource( + name='file', template='file://{path}', fn=lambda req: {'content': [{'text': 'file content'}]} + ) + + self.ai.define_resource( + name='user', template='user://{id}/profile', fn=lambda req: {'content': [{'text': 'user profile'}]} + ) + + # Create server + server = create_mcp_server(self.ai, McpServerOptions(name='test-server')) + await server.setup() + + # List resource templates + result = await server.list_resource_templates({}) + + # Verify + self.assertEqual(len(result.resourceTemplates), 2) + template_names = [t.name for t in result.resourceTemplates] + self.assertIn('file', template_names) + self.assertIn('user', template_names) + + # Verify templates + file_template = next(t for t in result.resourceTemplates if t.name == 'file') + self.assertEqual(file_template.uriTemplate, 'file://{path}') + + async def test_list_resources_excludes_templates(self): + """Test that list_resources excludes template resources.""" + # Define mixed resources + self.ai.define_resource(name='fixed', uri='app://fixed', fn=lambda req: {'content': [{'text': 'fixed'}]}) + + self.ai.define_resource( + name='template', template='app://{id}', fn=lambda req: {'content': [{'text': 'template'}]} + ) + + # Create server + server = create_mcp_server(self.ai, McpServerOptions(name='test-server')) + await server.setup() + + # List resources (should only include fixed URI) + result = await server.list_resources({}) + + self.assertEqual(len(result.resources), 1) + self.assertEqual(result.resources[0].name, 'fixed') + + async def test_list_resource_templates_excludes_fixed(self): + """Test that list_resource_templates excludes fixed URI resources.""" + # Define mixed resources + self.ai.define_resource(name='fixed', uri='app://fixed', fn=lambda req: {'content': [{'text': 'fixed'}]}) + + self.ai.define_resource( + name='template', template='app://{id}', fn=lambda req: {'content': [{'text': 'template'}]} + ) + + # Create server + server = create_mcp_server(self.ai, McpServerOptions(name='test-server')) + await server.setup() + + # List templates (should only include template) + result = await server.list_resource_templates({}) + + self.assertEqual(len(result.resourceTemplates), 1) + self.assertEqual(result.resourceTemplates[0].name, 'template') + + async def test_read_resource_with_fixed_uri(self): + """Test reading a resource with fixed URI.""" + + def config_resource(req): + return {'content': [{'text': 'Configuration data'}]} + + self.ai.define_resource(name='config', uri='app://config', fn=config_resource) + + # Create server + server = create_mcp_server(self.ai, McpServerOptions(name='test-server')) + await server.setup() + + # Read resource + from mcp.types import ReadResourceRequest + + request = MagicMock() + request.params.uri = 'app://config' + + result = await server.read_resource(request) + + # Verify + self.assertEqual(len(result.contents), 1) + self.assertEqual(result.contents[0].text, 'Configuration data') + + async def test_read_resource_with_template(self): + """Test reading a resource with URI template.""" + + def file_resource(req): + uri = req['uri'] + # Extract path from URI + path = uri.replace('file://', '') + return {'content': [{'text': f'Contents of {path}'}]} + + self.ai.define_resource(name='file', template='file://{path}', fn=file_resource) + + # Create server + server = create_mcp_server(self.ai, McpServerOptions(name='test-server')) + await server.setup() + + # Read resource + request = MagicMock() + request.params.uri = 'file:///home/user/document.txt' + + result = await server.read_resource(request) + + # Verify + self.assertEqual(len(result.contents), 1) + self.assertIn('/home/user/document.txt', result.contents[0].text) + + async def test_read_resource_not_found(self): + """Test reading a non-existent resource.""" + self.ai.define_resource(name='existing', uri='app://existing', fn=lambda req: {'content': [{'text': 'data'}]}) + + # Create server + server = create_mcp_server(self.ai, McpServerOptions(name='test-server')) + await server.setup() + + # Try to read non-existent resource + request = MagicMock() + request.params.uri = 'app://nonexistent' + + from genkit.core.error import GenkitError + + with self.assertRaises(GenkitError) as context: + await server.read_resource(request) + + self.assertIn('NOT_FOUND', str(context.exception.status)) + + async def test_read_resource_with_multiple_content_parts(self): + """Test reading a resource that returns multiple content parts.""" + + def multi_part_resource(req): + return {'content': [{'text': 'Part 1'}, {'text': 'Part 2'}, {'text': 'Part 3'}]} + + self.ai.define_resource(name='multi', uri='app://multi', fn=multi_part_resource) + + # Create server + server = create_mcp_server(self.ai, McpServerOptions(name='test-server')) + await server.setup() + + # Read resource + request = MagicMock() + request.params.uri = 'app://multi' + + result = await server.read_resource(request) + + # Verify + self.assertEqual(len(result.contents), 3) + self.assertEqual(result.contents[0].text, 'Part 1') + self.assertEqual(result.contents[1].text, 'Part 2') + self.assertEqual(result.contents[2].text, 'Part 3') + + +@pytest.mark.asyncio +class TestMcpServerToolsAndPrompts(unittest.IsolatedAsyncioTestCase): + """Tests for MCP server tool and prompt handling.""" + + def setUp(self): + """Set up test fixtures.""" + self.ai = Genkit() + + async def test_list_tools(self): + """Test listing tools.""" + + @self.ai.tool(description='Add two numbers') + def add(input: dict[str, int]) -> int: + return input['a'] + input['b'] + + @self.ai.tool(description='Multiply two numbers') + def multiply(input: dict[str, int]) -> int: + return input['a'] * input['b'] + + # Create server + server = create_mcp_server(self.ai, McpServerOptions(name='test-server')) + await server.setup() + + # List tools + result = await server.list_tools({}) + + # Verify + self.assertEqual(len(result.tools), 2) + tool_names = [t.name for t in result.tools] + self.assertIn('add', tool_names) + self.assertIn('multiply', tool_names) + + async def test_call_tool(self): + """Test calling a tool.""" + + @self.ai.tool() + def add(input: dict[str, int]) -> int: + return input['a'] + input['b'] + + # Create server + server = create_mcp_server(self.ai, McpServerOptions(name='test-server')) + await server.setup() + + # Call tool + request = MagicMock() + request.params.name = 'add' + request.params.arguments = {'a': 5, 'b': 3} + + result = await server.call_tool(request) + + # Verify + self.assertEqual(len(result.content), 1) + self.assertEqual(result.content[0].text, '8') + + async def test_list_prompts(self): + """Test listing prompts.""" + self.ai.define_prompt(name='greeting', prompt='Hello {{name}}!') + + self.ai.define_prompt(name='farewell', prompt='Goodbye {{name}}!') + + # Create server + server = create_mcp_server(self.ai, McpServerOptions(name='test-server')) + await server.setup() + + # List prompts + result = await server.list_prompts({}) + + # Verify + self.assertGreaterEqual(len(result.prompts), 2) + prompt_names = [p.name for p in result.prompts] + # Prompt names might have variant suffixes + + +@pytest.mark.asyncio +class TestMcpServerIntegration(unittest.IsolatedAsyncioTestCase): + """Integration tests for MCP server.""" + + async def test_server_exposes_all_action_types(self): + """Test that server exposes tools, prompts, and resources.""" + ai = Genkit() + + # Define tool + @ai.tool() + def test_tool(x: int) -> int: + return x * 2 + + # Define prompt + ai.define_prompt(name='test', prompt='Test prompt') + + # Define resource + ai.define_resource(name='test_resource', uri='test://resource', fn=lambda req: {'content': [{'text': 'test'}]}) + + # Create server + server = create_mcp_server(ai, McpServerOptions(name='integration-test')) + await server.setup() + + # Verify all action types are available + self.assertGreater(len(server.tool_actions), 0) + self.assertGreater(len(server.prompt_actions), 0) + self.assertGreater(len(server.resource_actions), 0) + + async def test_server_initialization_idempotent(self): + """Test that server setup is idempotent.""" + ai = Genkit() + + @ai.tool() + def test_tool(x: int) -> int: + return x + + server = create_mcp_server(ai, McpServerOptions(name='test')) + + # Setup multiple times + await server.setup() + count1 = len(server.tool_actions) + + await server.setup() + count2 = len(server.tool_actions) + + # Should be the same + self.assertEqual(count1, count2) + + +if __name__ == '__main__': + unittest.main() diff --git a/py/pyproject.toml b/py/pyproject.toml index 400fa73842..cac219b38a 100644 --- a/py/pyproject.toml +++ b/py/pyproject.toml @@ -114,6 +114,7 @@ genkit-plugin-google-genai = { workspace = true } genkit-plugin-ollama = { workspace = true } genkit-plugin-vertex-ai = { workspace = true } genkit-plugin-xai = { workspace = true } +genkit-plugins-mcp = { workspace = true } google-genai-hello = { workspace = true } google-genai-image = { workspace = true } prompt-demo = { workspace = true } diff --git a/py/uv.lock b/py/uv.lock index 560e88909a..c414fe7863 100644 --- a/py/uv.lock +++ b/py/uv.lock @@ -28,6 +28,7 @@ members = [ "genkit-plugin-ollama", "genkit-plugin-vertex-ai", "genkit-plugin-xai", + "genkit-plugins-mcp", "genkit-workspace", "google-genai-code-execution", "google-genai-context-caching", @@ -796,6 +797,7 @@ dependencies = [ ] sdist = { url = "https://files.pythonhosted.org/packages/13/1f/9fa001e74a1993a9cadd2333bb889e50c66327b8594ac538ab8a04f915b7/cryptography-45.0.3.tar.gz", hash = "sha256:ec21313dd335c51d7877baf2972569f40a4291b76a0ce51391523ae358d05899", size = 744738, upload-time = "2025-05-25T14:17:24.777Z" } wheels = [ + { url = "https://files.pythonhosted.org/packages/82/b2/2345dc595998caa6f68adf84e8f8b50d18e9fc4638d32b22ea8daedd4b7a/cryptography-45.0.3-cp311-abi3-macosx_10_9_universal2.whl", hash = "sha256:7573d9eebaeceeb55285205dbbb8753ac1e962af3d9640791d12b36864065e71", size = 7056239, upload-time = "2025-05-25T14:16:12.22Z" }, { url = "https://files.pythonhosted.org/packages/71/3d/ac361649a0bfffc105e2298b720d8b862330a767dab27c06adc2ddbef96a/cryptography-45.0.3-cp311-abi3-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:d377dde61c5d67eb4311eace661c3efda46c62113ff56bf05e2d679e02aebb5b", size = 4205541, upload-time = "2025-05-25T14:16:14.333Z" }, { url = "https://files.pythonhosted.org/packages/70/3e/c02a043750494d5c445f769e9c9f67e550d65060e0bfce52d91c1362693d/cryptography-45.0.3-cp311-abi3-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:fae1e637f527750811588e4582988932c222f8251f7b7ea93739acb624e1487f", size = 4433275, upload-time = "2025-05-25T14:16:16.421Z" }, { url = "https://files.pythonhosted.org/packages/40/7a/9af0bfd48784e80eef3eb6fd6fde96fe706b4fc156751ce1b2b965dada70/cryptography-45.0.3-cp311-abi3-manylinux_2_28_aarch64.whl", hash = "sha256:ca932e11218bcc9ef812aa497cdf669484870ecbcf2d99b765d6c27a86000942", size = 4209173, upload-time = "2025-05-25T14:16:18.163Z" }, @@ -805,6 +807,9 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/e7/53/8a130e22c1e432b3c14896ec5eb7ac01fb53c6737e1d705df7e0efb647c6/cryptography-45.0.3-cp311-abi3-manylinux_2_34_x86_64.whl", hash = "sha256:c824c9281cb628015bfc3c59335163d4ca0540d49de4582d6c2637312907e4b1", size = 4466300, upload-time = "2025-05-25T14:16:26.768Z" }, { url = "https://files.pythonhosted.org/packages/ba/75/6bb6579688ef805fd16a053005fce93944cdade465fc92ef32bbc5c40681/cryptography-45.0.3-cp311-abi3-musllinux_1_2_aarch64.whl", hash = "sha256:5833bb4355cb377ebd880457663a972cd044e7f49585aee39245c0d592904578", size = 4332483, upload-time = "2025-05-25T14:16:28.316Z" }, { url = "https://files.pythonhosted.org/packages/2f/11/2538f4e1ce05c6c4f81f43c1ef2bd6de7ae5e24ee284460ff6c77e42ca77/cryptography-45.0.3-cp311-abi3-musllinux_1_2_x86_64.whl", hash = "sha256:9bb5bf55dcb69f7067d80354d0a348368da907345a2c448b0babc4215ccd3497", size = 4573714, upload-time = "2025-05-25T14:16:30.474Z" }, + { url = "https://files.pythonhosted.org/packages/f5/bb/e86e9cf07f73a98d84a4084e8fd420b0e82330a901d9cac8149f994c3417/cryptography-45.0.3-cp311-abi3-win32.whl", hash = "sha256:3ad69eeb92a9de9421e1f6685e85a10fbcfb75c833b42cc9bc2ba9fb00da4710", size = 2934752, upload-time = "2025-05-25T14:16:32.204Z" }, + { url = "https://files.pythonhosted.org/packages/c7/75/063bc9ddc3d1c73e959054f1fc091b79572e716ef74d6caaa56e945b4af9/cryptography-45.0.3-cp311-abi3-win_amd64.whl", hash = "sha256:97787952246a77d77934d41b62fb1b6f3581d83f71b44796a4158d93b8f5c490", size = 3412465, upload-time = "2025-05-25T14:16:33.888Z" }, + { url = "https://files.pythonhosted.org/packages/71/9b/04ead6015229a9396890d7654ee35ef630860fb42dc9ff9ec27f72157952/cryptography-45.0.3-cp37-abi3-macosx_10_9_universal2.whl", hash = "sha256:c92519d242703b675ccefd0f0562eb45e74d438e001f8ab52d628e885751fb06", size = 7031892, upload-time = "2025-05-25T14:16:36.214Z" }, { url = "https://files.pythonhosted.org/packages/46/c7/c7d05d0e133a09fc677b8a87953815c522697bdf025e5cac13ba419e7240/cryptography-45.0.3-cp37-abi3-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:c5edcb90da1843df85292ef3a313513766a78fbbb83f584a5a58fb001a5a9d57", size = 4196181, upload-time = "2025-05-25T14:16:37.934Z" }, { url = "https://files.pythonhosted.org/packages/08/7a/6ad3aa796b18a683657cef930a986fac0045417e2dc428fd336cfc45ba52/cryptography-45.0.3-cp37-abi3-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:38deed72285c7ed699864f964a3f4cf11ab3fb38e8d39cfcd96710cd2b5bb716", size = 4423370, upload-time = "2025-05-25T14:16:39.502Z" }, { url = "https://files.pythonhosted.org/packages/4f/58/ec1461bfcb393525f597ac6a10a63938d18775b7803324072974b41a926b/cryptography-45.0.3-cp37-abi3-manylinux_2_28_aarch64.whl", hash = "sha256:5555365a50efe1f486eed6ac7062c33b97ccef409f5970a0b6f205a7cfab59c8", size = 4197839, upload-time = "2025-05-25T14:16:41.322Z" }, @@ -814,14 +819,20 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/71/7a/e002d5ce624ed46dfc32abe1deff32190f3ac47ede911789ee936f5a4255/cryptography-45.0.3-cp37-abi3-manylinux_2_34_x86_64.whl", hash = "sha256:57a6500d459e8035e813bd8b51b671977fb149a8c95ed814989da682314d0782", size = 4450308, upload-time = "2025-05-25T14:16:48.228Z" }, { url = "https://files.pythonhosted.org/packages/87/ad/3fbff9c28cf09b0a71e98af57d74f3662dea4a174b12acc493de00ea3f28/cryptography-45.0.3-cp37-abi3-musllinux_1_2_aarch64.whl", hash = "sha256:f22af3c78abfbc7cbcdf2c55d23c3e022e1a462ee2481011d518c7fb9c9f3d65", size = 4325125, upload-time = "2025-05-25T14:16:49.844Z" }, { url = "https://files.pythonhosted.org/packages/f5/b4/51417d0cc01802304c1984d76e9592f15e4801abd44ef7ba657060520bf0/cryptography-45.0.3-cp37-abi3-musllinux_1_2_x86_64.whl", hash = "sha256:232954730c362638544758a8160c4ee1b832dc011d2c41a306ad8f7cccc5bb0b", size = 4560038, upload-time = "2025-05-25T14:16:51.398Z" }, + { url = "https://files.pythonhosted.org/packages/80/38/d572f6482d45789a7202fb87d052deb7a7b136bf17473ebff33536727a2c/cryptography-45.0.3-cp37-abi3-win32.whl", hash = "sha256:cb6ab89421bc90e0422aca911c69044c2912fc3debb19bb3c1bfe28ee3dff6ab", size = 2924070, upload-time = "2025-05-25T14:16:53.472Z" }, + { url = "https://files.pythonhosted.org/packages/91/5a/61f39c0ff4443651cc64e626fa97ad3099249152039952be8f344d6b0c86/cryptography-45.0.3-cp37-abi3-win_amd64.whl", hash = "sha256:d54ae41e6bd70ea23707843021c778f151ca258081586f0cfa31d936ae43d1b2", size = 3395005, upload-time = "2025-05-25T14:16:55.134Z" }, + { url = "https://files.pythonhosted.org/packages/1b/63/ce30cb7204e8440df2f0b251dc0464a26c55916610d1ba4aa912f838bcc8/cryptography-45.0.3-pp310-pypy310_pp73-macosx_10_9_x86_64.whl", hash = "sha256:ed43d396f42028c1f47b5fec012e9e12631266e3825e95c00e3cf94d472dac49", size = 3578348, upload-time = "2025-05-25T14:16:56.792Z" }, { url = "https://files.pythonhosted.org/packages/45/0b/87556d3337f5e93c37fda0a0b5d3e7b4f23670777ce8820fce7962a7ed22/cryptography-45.0.3-pp310-pypy310_pp73-manylinux_2_28_aarch64.whl", hash = "sha256:fed5aaca1750e46db870874c9c273cd5182a9e9deb16f06f7bdffdb5c2bde4b9", size = 4142867, upload-time = "2025-05-25T14:16:58.459Z" }, { url = "https://files.pythonhosted.org/packages/72/ba/21356dd0bcb922b820211336e735989fe2cf0d8eaac206335a0906a5a38c/cryptography-45.0.3-pp310-pypy310_pp73-manylinux_2_28_x86_64.whl", hash = "sha256:00094838ecc7c6594171e8c8a9166124c1197b074cfca23645cee573910d76bc", size = 4385000, upload-time = "2025-05-25T14:17:00.656Z" }, { url = "https://files.pythonhosted.org/packages/2f/2b/71c78d18b804c317b66283be55e20329de5cd7e1aec28e4c5fbbe21fd046/cryptography-45.0.3-pp310-pypy310_pp73-manylinux_2_34_aarch64.whl", hash = "sha256:92d5f428c1a0439b2040435a1d6bc1b26ebf0af88b093c3628913dd464d13fa1", size = 4144195, upload-time = "2025-05-25T14:17:02.782Z" }, { url = "https://files.pythonhosted.org/packages/55/3e/9f9b468ea779b4dbfef6af224804abd93fbcb2c48605d7443b44aea77979/cryptography-45.0.3-pp310-pypy310_pp73-manylinux_2_34_x86_64.whl", hash = "sha256:ec64ee375b5aaa354b2b273c921144a660a511f9df8785e6d1c942967106438e", size = 4384540, upload-time = "2025-05-25T14:17:04.49Z" }, + { url = "https://files.pythonhosted.org/packages/97/f5/6e62d10cf29c50f8205c0dc9aec986dca40e8e3b41bf1a7878ea7b11e5ee/cryptography-45.0.3-pp310-pypy310_pp73-win_amd64.whl", hash = "sha256:71320fbefd05454ef2d457c481ba9a5b0e540f3753354fff6f780927c25d19b0", size = 3328796, upload-time = "2025-05-25T14:17:06.174Z" }, + { url = "https://files.pythonhosted.org/packages/e7/d4/58a246342093a66af8935d6aa59f790cbb4731adae3937b538d054bdc2f9/cryptography-45.0.3-pp311-pypy311_pp73-macosx_10_9_x86_64.whl", hash = "sha256:edd6d51869beb7f0d472e902ef231a9b7689508e83880ea16ca3311a00bf5ce7", size = 3589802, upload-time = "2025-05-25T14:17:07.792Z" }, { url = "https://files.pythonhosted.org/packages/96/61/751ebea58c87b5be533c429f01996050a72c7283b59eee250275746632ea/cryptography-45.0.3-pp311-pypy311_pp73-manylinux_2_28_aarch64.whl", hash = "sha256:555e5e2d3a53b4fabeca32835878b2818b3f23966a4efb0d566689777c5a12c8", size = 4146964, upload-time = "2025-05-25T14:17:09.538Z" }, { url = "https://files.pythonhosted.org/packages/8d/01/28c90601b199964de383da0b740b5156f5d71a1da25e7194fdf793d373ef/cryptography-45.0.3-pp311-pypy311_pp73-manylinux_2_28_x86_64.whl", hash = "sha256:25286aacb947286620a31f78f2ed1a32cded7be5d8b729ba3fb2c988457639e4", size = 4388103, upload-time = "2025-05-25T14:17:11.978Z" }, { url = "https://files.pythonhosted.org/packages/3d/ec/cd892180b9e42897446ef35c62442f5b8b039c3d63a05f618aa87ec9ebb5/cryptography-45.0.3-pp311-pypy311_pp73-manylinux_2_34_aarch64.whl", hash = "sha256:050ce5209d5072472971e6efbfc8ec5a8f9a841de5a4db0ebd9c2e392cb81972", size = 4150031, upload-time = "2025-05-25T14:17:14.131Z" }, { url = "https://files.pythonhosted.org/packages/db/d4/22628c2dedd99289960a682439c6d3aa248dff5215123ead94ac2d82f3f5/cryptography-45.0.3-pp311-pypy311_pp73-manylinux_2_34_x86_64.whl", hash = "sha256:dc10ec1e9f21f33420cc05214989544727e776286c1c16697178978327b95c9c", size = 4387389, upload-time = "2025-05-25T14:17:17.303Z" }, + { url = "https://files.pythonhosted.org/packages/39/ec/ba3961abbf8ecb79a3586a4ff0ee08c9d7a9938b4312fb2ae9b63f48a8ba/cryptography-45.0.3-pp311-pypy311_pp73-win_amd64.whl", hash = "sha256:9eda14f049d7f09c2e8fb411dda17dd6b16a3c76a1de5e249188a32aeb92de19", size = 3337432, upload-time = "2025-05-25T14:17:19.507Z" }, ] [[package]] @@ -1787,6 +1798,21 @@ requires-dist = [ { name = "xai-sdk", specifier = ">=0.0.1" }, ] +[[package]] +name = "genkit-plugins-mcp" +version = "0.1.0" +source = { editable = "plugins/mcp" } +dependencies = [ + { name = "genkit" }, + { name = "mcp" }, +] + +[package.metadata] +requires-dist = [ + { name = "genkit", editable = "packages/genkit" }, + { name = "mcp" }, +] + [[package]] name = "genkit-workspace" version = "0.1.0" @@ -2499,6 +2525,15 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/2a/39/e50c7c3a983047577ee07d2a9e53faf5a69493943ec3f6a384bdc792deb2/httpx-0.28.1-py3-none-any.whl", hash = "sha256:d909fcccc110f8c7faf814ca82a9a4d816bc5a6dbfea25d6591d6985b8ba59ad", size = 73517, upload-time = "2024-12-06T15:37:21.509Z" }, ] +[[package]] +name = "httpx-sse" +version = "0.4.3" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/0f/4c/751061ffa58615a32c31b2d82e8482be8dd4a89154f003147acee90f2be9/httpx_sse-0.4.3.tar.gz", hash = "sha256:9b1ed0127459a66014aec3c56bebd93da3c1bc8bb6618c8082039a44889a755d", size = 15943, upload-time = "2025-10-10T21:48:22.271Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/d2/fd/6668e5aec43ab844de6fc74927e155a3b37bf40d7c3790e49fc0406b6578/httpx_sse-0.4.3-py3-none-any.whl", hash = "sha256:0ac1c9fe3c0afad2e0ebb25a934a59f4c7823b60792691f779fad2c5568830fc", size = 8960, upload-time = "2025-10-10T21:48:21.158Z" }, +] + [[package]] name = "id" version = "1.5.0" @@ -3257,6 +3292,31 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/8f/8e/9ad090d3553c280a8060fbf6e24dc1c0c29704ee7d1c372f0c174aa59285/matplotlib_inline-0.1.7-py3-none-any.whl", hash = "sha256:df192d39a4ff8f21b1895d72e6a13f5fcc5099f00fa84384e0ea28c2cc0653ca", size = 9899, upload-time = "2024-04-15T13:44:43.265Z" }, ] +[[package]] +name = "mcp" +version = "1.25.0" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "anyio" }, + { name = "httpx" }, + { name = "httpx-sse" }, + { name = "jsonschema" }, + { name = "pydantic" }, + { name = "pydantic-settings" }, + { name = "pyjwt", extra = ["crypto"] }, + { name = "python-multipart" }, + { name = "pywin32", marker = "sys_platform == 'win32'" }, + { name = "sse-starlette" }, + { name = "starlette" }, + { name = "typing-extensions" }, + { name = "typing-inspection" }, + { name = "uvicorn", marker = "sys_platform != 'emscripten'" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/d5/2d/649d80a0ecf6a1f82632ca44bec21c0461a9d9fc8934d38cb5b319f2db5e/mcp-1.25.0.tar.gz", hash = "sha256:56310361ebf0364e2d438e5b45f7668cbb124e158bb358333cd06e49e83a6802", size = 605387, upload-time = "2025-12-19T10:19:56.985Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/e2/fc/6dc7659c2ae5ddf280477011f4213a74f806862856b796ef08f028e664bf/mcp-1.25.0-py3-none-any.whl", hash = "sha256:b37c38144a666add0862614cc79ec276e97d72aa8ca26d622818d4e278b9721a", size = 233076, upload-time = "2025-12-19T10:19:55.416Z" }, +] + [[package]] name = "mdurl" version = "0.1.2" @@ -4499,6 +4559,20 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/36/c7/cfc8e811f061c841d7990b0201912c3556bfeb99cdcb7ed24adc8d6f8704/pydantic_core-2.41.5-pp311-pypy311_pp73-win_amd64.whl", hash = "sha256:56121965f7a4dc965bff783d70b907ddf3d57f6eba29b6d2e5dabfaf07799c51", size = 2145302, upload-time = "2025-11-04T13:43:46.64Z" }, ] +[[package]] +name = "pydantic-settings" +version = "2.12.0" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "pydantic" }, + { name = "python-dotenv" }, + { name = "typing-inspection" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/43/4b/ac7e0aae12027748076d72a8764ff1c9d82ca75a7a52622e67ed3f765c54/pydantic_settings-2.12.0.tar.gz", hash = "sha256:005538ef951e3c2a68e1c08b292b5f2e71490def8589d4221b95dab00dafcfd0", size = 194184, upload-time = "2025-11-10T14:25:47.013Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/c1/60/5d4751ba3f4a40a6891f24eec885f51afd78d208498268c734e256fb13c4/pydantic_settings-2.12.0-py3-none-any.whl", hash = "sha256:fddb9fd99a5b18da837b29710391e945b1e30c135477f484084ee513adb93809", size = 51880, upload-time = "2025-11-10T14:25:45.546Z" }, +] + [[package]] name = "pygments" version = "2.19.1" @@ -4508,6 +4582,20 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/8a/0b/9fcc47d19c48b59121088dd6da2488a49d5f72dacf8262e2790a1d2c7d15/pygments-2.19.1-py3-none-any.whl", hash = "sha256:9ea1544ad55cecf4b8242fab6dd35a93bbce657034b0611ee383099054ab6d8c", size = 1225293, upload-time = "2025-01-06T17:26:25.553Z" }, ] +[[package]] +name = "pyjwt" +version = "2.10.1" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/e7/46/bd74733ff231675599650d3e47f361794b22ef3e3770998dda30d3b63726/pyjwt-2.10.1.tar.gz", hash = "sha256:3cc5772eb20009233caf06e9d8a0577824723b44e6648ee0a2aedb6cf9381953", size = 87785, upload-time = "2024-11-28T03:43:29.933Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/61/ad/689f02752eeec26aed679477e80e632ef1b682313be70793d798c1d5fc8f/PyJWT-2.10.1-py3-none-any.whl", hash = "sha256:dcdd193e30abefd5debf142f9adfcdd2b58004e644f25406ffaebd50bd98dacb", size = 22997, upload-time = "2024-11-28T03:43:27.893Z" }, +] + +[package.optional-dependencies] +crypto = [ + { name = "cryptography" }, +] + [[package]] name = "pypdf" version = "6.5.0" @@ -4613,6 +4701,15 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/ec/57/56b9bcc3c9c6a792fcbaf139543cee77261f3651ca9da0c93f5c1221264b/python_dateutil-2.9.0.post0-py2.py3-none-any.whl", hash = "sha256:a8b2bc7bffae282281c8140a97d3aa9c14da0b136dfe83f850eea9a5f7470427", size = 229892, upload-time = "2024-03-01T18:36:18.57Z" }, ] +[[package]] +name = "python-dotenv" +version = "1.2.1" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/f0/26/19cadc79a718c5edbec86fd4919a6b6d3f681039a2f6d66d14be94e75fb9/python_dotenv-1.2.1.tar.gz", hash = "sha256:42667e897e16ab0d66954af0e60a9caa94f0fd4ecf3aaf6d2d260eec1aa36ad6", size = 44221, upload-time = "2025-10-26T15:12:10.434Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/14/1b/a298b06749107c305e1fe0f814c6c74aea7b2f1e10989cb30f544a1b3253/python_dotenv-1.2.1-py3-none-any.whl", hash = "sha256:b81ee9561e9ca4004139c6cbba3a238c32b03e4894671e181b671e8cb8425d61", size = 21230, upload-time = "2025-10-26T15:12:09.109Z" }, +] + [[package]] name = "python-json-logger" version = "3.3.0" @@ -4622,6 +4719,15 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/08/20/0f2523b9e50a8052bc6a8b732dfc8568abbdc42010aef03a2d750bdab3b2/python_json_logger-3.3.0-py3-none-any.whl", hash = "sha256:dd980fae8cffb24c13caf6e158d3d61c0d6d22342f932cb6e9deedab3d35eec7", size = 15163, upload-time = "2025-03-07T07:08:25.627Z" }, ] +[[package]] +name = "python-multipart" +version = "0.0.21" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/78/96/804520d0850c7db98e5ccb70282e29208723f0964e88ffd9d0da2f52ea09/python_multipart-0.0.21.tar.gz", hash = "sha256:7137ebd4d3bbf70ea1622998f902b97a29434a9e8dc40eb203bbcf7c2a2cba92", size = 37196, upload-time = "2025-12-17T09:24:22.446Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/aa/76/03af049af4dcee5d27442f71b6924f01f3efb5d2bd34f23fcd563f2cc5f5/python_multipart-0.0.21-py3-none-any.whl", hash = "sha256:cf7a6713e01c87aa35387f4774e812c4361150938d20d232800f75ffcf266090", size = 24541, upload-time = "2025-12-17T09:24:21.153Z" }, +] + [[package]] name = "pywin32" version = "310" From 3a517c9137dcd9beb9c42e9a59b434be45b61fcf Mon Sep 17 00:00:00 2001 From: Mengqin Shen Date: Mon, 29 Dec 2025 20:08:46 -0800 Subject: [PATCH 2/4] feat(py):revise with gemini's comments --- .../mcp/examples/server/simple_server.py | 2 +- py/plugins/mcp/pyproject.toml | 2 +- .../src/genkit/plugins/mcp/client/client.py | 11 +-- .../mcp/src/genkit/plugins/mcp/server.py | 44 +++++++---- .../src/genkit/plugins/mcp/util/message.py | 78 +++++++++---------- 5 files changed, 71 insertions(+), 66 deletions(-) diff --git a/py/plugins/mcp/examples/server/simple_server.py b/py/plugins/mcp/examples/server/simple_server.py index 36e7bdbe2b..991abd3c96 100644 --- a/py/plugins/mcp/examples/server/simple_server.py +++ b/py/plugins/mcp/examples/server/simple_server.py @@ -56,7 +56,7 @@ def add(input: AddInput): server = create_mcp_server(ai, McpServerOptions(name='example_server', version='0.0.1')) print('Starting MCP server on stdio...') - asyncio.run(server.start_stdio()) + asyncio.run(server.start()) if __name__ == '__main__': diff --git a/py/plugins/mcp/pyproject.toml b/py/plugins/mcp/pyproject.toml index 7f6936009f..ee6ed05cc0 100644 --- a/py/plugins/mcp/pyproject.toml +++ b/py/plugins/mcp/pyproject.toml @@ -46,4 +46,4 @@ build-backend = "hatchling.build" requires = ["hatchling"] [tool.hatch.build.targets.wheel] -packages = ["src/genkit", "src/genkit/plugins"] +packages = ["src"] diff --git a/py/plugins/mcp/src/genkit/plugins/mcp/client/client.py b/py/plugins/mcp/src/genkit/plugins/mcp/client/client.py index 218493c056..15eba78727 100644 --- a/py/plugins/mcp/src/genkit/plugins/mcp/client/client.py +++ b/py/plugins/mcp/src/genkit/plugins/mcp/client/client.py @@ -102,16 +102,7 @@ async def connect(self): logger.error(f'Failed to connect to MCP server {self.server_name}: {e}') self.config.disabled = True # Clean up on error - if hasattr(self, '_session_context') and self._session_context: - try: - await self._session_context.__aexit__(None, None, None) - except: - pass - if self._exit_stack: - try: - await self._exit_stack.__aexit__(None, None, None) - except: - pass + await self.close() raise e async def close(self): diff --git a/py/plugins/mcp/src/genkit/plugins/mcp/server.py b/py/plugins/mcp/src/genkit/plugins/mcp/server.py index 088a57e481..972a4823a1 100644 --- a/py/plugins/mcp/src/genkit/plugins/mcp/server.py +++ b/py/plugins/mcp/src/genkit/plugins/mcp/server.py @@ -27,6 +27,7 @@ from genkit.ai import Genkit from genkit.blocks.resource import matches_uri_template +from genkit.core.action._key import parse_action_key from genkit.core.action.types import ActionKind from genkit.core.error import GenkitError from genkit.core.schema import to_json_schema @@ -96,6 +97,10 @@ def __init__(self, ai: Genkit, options: McpServerOptions): self.tool_actions: list[Any] = [] self.prompt_actions: list[Any] = [] self.resource_actions: list[Any] = [] + self.tool_actions_map: dict[str, Any] = {} + self.prompt_actions_map: dict[str, Any] = {} + self.resource_uri_map: dict[str, Any] = {} + self.resource_templates: list[tuple[str, Any]] = [] async def setup(self) -> None: """Initialize the MCP server and register request handlers. @@ -141,10 +146,18 @@ async def setup(self) -> None: for name, action in entries.items(): if kind == ActionKind.TOOL: self.tool_actions.append(action) + self.tool_actions_map[action.name] = action elif kind == ActionKind.PROMPT: self.prompt_actions.append(action) + self.prompt_actions_map[action.name] = action elif kind == ActionKind.RESOURCE: self.resource_actions.append(action) + metadata = action.metadata or {} + resource_meta = metadata.get('resource', {}) + if resource_meta.get('uri'): + self.resource_uri_map[resource_meta['uri']] = action + if resource_meta.get('template'): + self.resource_templates.append((resource_meta['template'], action)) # Also get actions from plugins that might not be in _entries yet # (though most plugins register them in _entries during initialization) @@ -155,10 +168,18 @@ async def setup(self) -> None: if action: if kind == ActionKind.TOOL and action not in self.tool_actions: self.tool_actions.append(action) + self.tool_actions_map[action.name] = action elif kind == ActionKind.PROMPT and action not in self.prompt_actions: self.prompt_actions.append(action) + self.prompt_actions_map[action.name] = action elif kind == ActionKind.RESOURCE and action not in self.resource_actions: self.resource_actions.append(action) + metadata = action.metadata or {} + resource_meta = metadata.get('resource', {}) + if resource_meta.get('uri'): + self.resource_uri_map[resource_meta['uri']] = action + if resource_meta.get('template'): + self.resource_templates.append((resource_meta['template'], action)) self.actions_resolved = True @@ -211,7 +232,7 @@ async def call_tool(self, request: CallToolRequest) -> CallToolResult: await self.setup() # Find the tool action - tool = next((t for t in self.tool_actions if t.name == request.params.name), None) + tool = self.tool_actions_map.get(request.params.name) if not tool: raise GenkitError( @@ -268,7 +289,7 @@ async def get_prompt(self, request: GetPromptRequest) -> GetPromptResult: await self.setup() # Find the prompt action - prompt = next((p for p in self.prompt_actions if p.name == request.params.name), None) + prompt = self.prompt_actions_map.get(request.params.name) if not prompt: raise GenkitError( @@ -359,20 +380,13 @@ async def read_resource(self, request: ReadResourceRequest) -> ReadResourceResul uri = request.params.uri - # Find matching resource (either exact URI match or template match) - resource = None - for action in self.resource_actions: - metadata = action.metadata or {} - resource_meta = metadata.get('resource', {}) + # Check for exact URI match + resource = self.resource_uri_map.get(uri) - # Check for exact URI match - if resource_meta.get('uri') == uri: - resource = action - break - - # Check for template match - if resource_meta.get('template'): - if matches_uri_template(resource_meta['template'], uri): + # Check for template match if not found by exact URI + if not resource: + for template, action in self.resource_templates: + if matches_uri_template(template, uri): resource = action break diff --git a/py/plugins/mcp/src/genkit/plugins/mcp/util/message.py b/py/plugins/mcp/src/genkit/plugins/mcp/util/message.py index 251bac968f..a1b0a13ebe 100644 --- a/py/plugins/mcp/src/genkit/plugins/mcp/util/message.py +++ b/py/plugins/mcp/src/genkit/plugins/mcp/util/message.py @@ -97,6 +97,34 @@ def from_mcp_part(part: Dict[str, Any]) -> Dict[str, Any]: return {} +def _get_part_data(part: Any) -> Dict[str, Any]: + """Extract data from a Part, handling potential 'root' nesting.""" + if isinstance(part, str): + return {'text': part} + part_dict = part if isinstance(part, dict) else part.model_dump() + if 'root' in part_dict and isinstance(part_dict['root'], dict): + return part_dict['root'] + return part_dict + + +def _parse_media_part(media: Dict[str, Any]) -> ImageContent: + """Extract MIME type and base64 data from a media part.""" + url = media.get('url', '') + content_type = media.get('contentType', '') + + if not url.startswith('data:'): + raise ValueError('MCP prompt messages only support base64 data images.') + + # Extract MIME type and base64 data + try: + mime_type = content_type or url[url.index(':') + 1 : url.index(';')] + data = url[url.index(',') + 1 :] + except ValueError as e: + raise ValueError(f'Invalid data URL format: {url}') from e + + return ImageContent(type='image', data=data, mimeType=mime_type) + + def to_mcp_prompt_message(message: Message) -> PromptMessage: """Convert a Genkit Message to an MCP PromptMessage. @@ -123,47 +151,19 @@ def to_mcp_prompt_message(message: Message) -> PromptMessage: mcp_role = role_map[message.role] - # Handle media content (images) + # First, look for any media content as MCP content is currently single-part if message.content: for part in message.content: - part_dict = part if isinstance(part, dict) else part.model_dump() - if part_dict.get('media'): - media = part_dict['media'] - url = media.get('url', '') - content_type = media.get('contentType', '') + data = _get_part_data(part) + if data.get('media'): + return PromptMessage(role=mcp_role, content=_parse_media_part(data['media'])) - if not url.startswith('data:'): - raise ValueError('MCP prompt messages only support base64 data images.') - - # Extract MIME type and base64 data - mime_type = content_type or url[url.index(':') + 1 : url.index(';')] - data = url[url.index(',') + 1 :] - - return PromptMessage(role=mcp_role, content=ImageContent(type='image', data=data, mimeType=mime_type)) - elif part_dict.get('root') and isinstance(part_dict['root'], dict) and part_dict['root'].get('media'): - media = part_dict['root']['media'] - url = media.get('url', '') - content_type = media.get('contentType', '') - - if not url.startswith('data:'): - raise ValueError('MCP prompt messages only support base64 data images.') - - # Extract MIME type and base64 data - mime_type = content_type or url[url.index(':') + 1 : url.index(';')] - data = url[url.index(',') + 1 :] - - return PromptMessage(role=mcp_role, content=ImageContent(type='image', data=data, mimeType=mime_type)) - - # Handle text content - text = '' + # If no media, aggregate all text content + text_content = [] if message.content: for part in message.content: - part_dict = part if isinstance(part, dict) else part.model_dump() - if 'text' in part_dict and part_dict['text']: - text += part_dict['text'] - elif 'root' in part_dict and isinstance(part_dict['root'], dict) and 'text' in part_dict['root']: - text += part_dict['root']['text'] - elif isinstance(part, str): - text += part - - return PromptMessage(role=mcp_role, content=TextContent(type='text', text=text)) + data = _get_part_data(part) + if data.get('text'): + text_content.append(data['text']) + + return PromptMessage(role=mcp_role, content=TextContent(type='text', text=''.join(text_content))) From 69544e8845ac9536e71bbf1a954feff1df69ede8 Mon Sep 17 00:00:00 2001 From: Mengqin Shen Date: Tue, 30 Dec 2025 10:24:55 -0800 Subject: [PATCH 3/4] feat(py):add sample test for mcp --- .../src/genkit/plugins/mcp/client/client.py | 15 +- .../mcp/src/genkit/plugins/mcp/client/host.py | 7 +- .../mcp/src/genkit/plugins/mcp/server.py | 174 ++++------ py/samples/mcp/README.md | 49 +++ py/samples/mcp/pyproject.toml | 24 ++ py/samples/mcp/src/http_server.py | 101 ++++++ py/samples/mcp/src/main.py | 318 ++++++++++++++++++ py/samples/mcp/src/server.py | 98 ++++++ py/samples/mcp/test-workspace/hello-world.txt | 4 + py/uv.lock | 22 +- 10 files changed, 695 insertions(+), 117 deletions(-) create mode 100644 py/samples/mcp/README.md create mode 100644 py/samples/mcp/pyproject.toml create mode 100644 py/samples/mcp/src/http_server.py create mode 100644 py/samples/mcp/src/main.py create mode 100644 py/samples/mcp/src/server.py create mode 100644 py/samples/mcp/test-workspace/hello-world.txt diff --git a/py/plugins/mcp/src/genkit/plugins/mcp/client/client.py b/py/plugins/mcp/src/genkit/plugins/mcp/client/client.py index 15eba78727..cd5466231a 100644 --- a/py/plugins/mcp/src/genkit/plugins/mcp/client/client.py +++ b/py/plugins/mcp/src/genkit/plugins/mcp/client/client.py @@ -158,16 +158,17 @@ async def read_resource(self, uri: str) -> Any: raise RuntimeError('MCP client is not connected') return await self.session.read_resource(uri) - async def register_tools(self, ai: Optional[Genkit] = None): + async def register_tools(self, ai: Optional[Genkit] = None) -> List[str]: """Registers all tools from connected client to Genkit.""" registry = ai.registry if ai else (self.ai.registry if self.ai else None) if not registry: logger.warning('No Genkit registry available to register tools.') - return + return [] if not self.session: - return + return [] + registered_tools = [] try: tools = await self.list_tools() for tool in tools: @@ -184,18 +185,22 @@ async def tool_wrapper(args: Any = None, _tool_name=tool.name): # Use metadata to store MCP specific info metadata = {'mcp': {'_meta': tool._meta}} if hasattr(tool, '_meta') else {} + tool_name = f'{self.server_name}/{tool.name}' # Define the tool in Genkit registry registry.register_action( kind=ActionKind.TOOL, - name=f'{self.server_name}/{tool.name}', + name=tool_name, fn=tool_wrapper, description=tool.description, metadata=metadata, # TODO: json_schema conversion from tool.inputSchema ) - logger.debug(f'Registered MCP tool: {self.server_name}/{tool.name}') + registered_tools.append(tool_name) + logger.debug(f'Registered MCP tool: {tool_name}') + return registered_tools except Exception as e: logger.error(f'Error registering tools for {self.server_name}: {e}') + return [] async def get_active_tools(self) -> List[Any]: """Returns all active tools.""" diff --git a/py/plugins/mcp/src/genkit/plugins/mcp/client/host.py b/py/plugins/mcp/src/genkit/plugins/mcp/client/host.py index 365b6382d5..b089cea98e 100644 --- a/py/plugins/mcp/src/genkit/plugins/mcp/client/host.py +++ b/py/plugins/mcp/src/genkit/plugins/mcp/client/host.py @@ -39,11 +39,14 @@ async def close(self): for client in self.clients.values(): await client.close() - async def register_tools(self, ai: Genkit): + async def register_tools(self, ai: Genkit) -> List[str]: """Registers all tools from connected clients to Genkit.""" + all_tools = [] for client in self.clients.values(): if client.session: - await client.register_tools(ai) + tools = await client.register_tools(ai) + all_tools.extend(tools) + return all_tools async def enable(self, name: str): """Enables and connects an MCP client.""" diff --git a/py/plugins/mcp/src/genkit/plugins/mcp/server.py b/py/plugins/mcp/src/genkit/plugins/mcp/server.py index 972a4823a1..bfe237d5dd 100644 --- a/py/plugins/mcp/src/genkit/plugins/mcp/server.py +++ b/py/plugins/mcp/src/genkit/plugins/mcp/server.py @@ -23,7 +23,8 @@ from typing import Any, Optional import structlog -from pydantic import BaseModel +from pydantic import BaseModel, AnyUrl +import mcp.types as types from genkit.ai import Genkit from genkit.blocks.resource import matches_uri_template @@ -113,25 +114,41 @@ async def setup(self) -> None: return # Create MCP Server instance - self.server = Server( - {'name': self.options.name, 'version': self.options.version}, - { - 'capabilities': { - 'prompts': {}, - 'tools': {}, - 'resources': {}, - } - }, - ) + self.server = Server(self.options.name) + + # Register request handlers using decorators + + @self.server.list_tools() + async def list_tools() -> list[types.Tool]: + return await self._list_tools() + + @self.server.call_tool() + async def call_tool(name: str, arguments: dict | None) -> list[types.TextContent | types.ImageContent | types.EmbeddedResource]: + return await self._call_tool(name, arguments) + + @self.server.list_prompts() + async def list_prompts() -> list[types.Prompt]: + return await self._list_prompts() + + @self.server.get_prompt() + async def get_prompt(name: str, arguments: dict[str, str] | None) -> types.GetPromptResult: + return await self._get_prompt(name, arguments) + + @self.server.list_resources() + async def list_resources() -> list[types.Resource]: + return await self._list_resources() + + @self.server.list_resource_templates() + async def list_resource_templates() -> list[types.ResourceTemplate]: + return await self._list_resource_templates() + + @self.server.read_resource() + async def read_resource(uri: AnyUrl) -> str | bytes: + # Note: The MCP SDK signature for read_resource expects returning content + # directly or a list of contents depending on version. + # Based on lowlevel/server.py it returns ReadResourceContents which is content list + return await self._read_resource(str(uri)) - # Register request handlers - self.server.setRequestHandler(ListToolsRequestSchema, self.list_tools) - self.server.setRequestHandler(CallToolRequestSchema, self.call_tool) - self.server.setRequestHandler(ListPromptsRequestSchema, self.list_prompts) - self.server.setRequestHandler(GetPromptRequestSchema, self.get_prompt) - self.server.setRequestHandler(ListResourcesRequestSchema, self.list_resources) - self.server.setRequestHandler(ListResourceTemplatesRequestSchema, self.list_resource_templates) - self.server.setRequestHandler(ReadResourceRequestSchema, self.read_resource) # Resolve all actions from Genkit registry # We need the actual Action objects, not just serializable dicts @@ -190,15 +207,8 @@ async def setup(self) -> None: resources=len(self.resource_actions), ) - async def list_tools(self, request: ListToolsRequest) -> ListToolsResult: - """Handle MCP requests to list available tools. - - Args: - request: The MCP ListToolsRequest. - - Returns: - ListToolsResult containing all registered Genkit tools. - """ + async def _list_tools(self) -> list[types.Tool]: + """Handle MCP requests to list available tools.""" await self.setup() tools: list[Tool] = [] @@ -215,46 +225,29 @@ async def list_tools(self, request: ListToolsRequest) -> ListToolsResult: ) ) - return ListToolsResult(tools=tools) - - async def call_tool(self, request: CallToolRequest) -> CallToolResult: - """Handle MCP requests to call a specific tool. - - Args: - request: The MCP CallToolRequest containing tool name and arguments. - - Returns: - CallToolResult with the tool execution result. + return tools - Raises: - GenkitError: If the requested tool is not found. - """ + async def _call_tool(self, name: str, arguments: dict | None) -> list[types.TextContent | types.ImageContent | types.EmbeddedResource]: + """Handle MCP requests to call a specific tool.""" await self.setup() # Find the tool action - tool = self.tool_actions_map.get(request.params.name) + tool = self.tool_actions_map.get(name) if not tool: raise GenkitError( - status='NOT_FOUND', message=f"Tried to call tool '{request.params.name}' but it could not be found." + status='NOT_FOUND', message=f"Tried to call tool '{name}' but it could not be found." ) # Execute the tool - result = await tool.arun(request.params.arguments) + result = await tool.arun(arguments) result = result.response - # Convert result to MCP format - return CallToolResult(content=to_mcp_tool_result(result)) - - async def list_prompts(self, request: ListPromptsRequest) -> ListPromptsResult: - """Handle MCP requests to list available prompts. - - Args: - request: The MCP ListPromptsRequest. + # Convert result to MCP format (list of contents) + return to_mcp_tool_result(result) - Returns: - ListPromptsResult containing all registered Genkit prompts. - """ + async def _list_prompts(self) -> list[types.Prompt]: + """Handle MCP requests to list available prompts.""" await self.setup() prompts: list[Prompt] = [] @@ -272,33 +265,23 @@ async def list_prompts(self, request: ListPromptsRequest) -> ListPromptsResult: ) ) - return ListPromptsResult(prompts=prompts) - - async def get_prompt(self, request: GetPromptRequest) -> GetPromptResult: - """Handle MCP requests to get (render) a specific prompt. - - Args: - request: The MCP GetPromptRequest containing prompt name and arguments. - - Returns: - GetPromptResult with the rendered prompt messages. + return prompts - Raises: - GenkitError: If the requested prompt is not found. - """ + async def _get_prompt(self, name: str, arguments: dict[str, str] | None) -> types.GetPromptResult: + """Handle MCP requests to get (render) a specific prompt.""" await self.setup() # Find the prompt action - prompt = self.prompt_actions_map.get(request.params.name) + prompt = self.prompt_actions_map.get(name) if not prompt: raise GenkitError( status='NOT_FOUND', - message=f"[MCP Server] Tried to call prompt '{request.params.name}' but it could not be found.", + message=f"[MCP Server] Tried to call prompt '{name}' but it could not be found.", ) # Execute the prompt - result = await prompt.arun(request.params.arguments) + result = await prompt.arun(arguments) result = result.response # Convert messages to MCP format @@ -306,15 +289,8 @@ async def get_prompt(self, request: GetPromptRequest) -> GetPromptResult: return GetPromptResult(description=prompt.description, messages=messages) - async def list_resources(self, request: ListResourcesRequest) -> ListResourcesResult: - """Handle MCP requests to list available resources with fixed URIs. - - Args: - request: The MCP ListResourcesRequest. - - Returns: - ListResourcesResult containing resources with fixed URIs. - """ + async def _list_resources(self) -> list[types.Resource]: + """Handle MCP requests to list available resources with fixed URIs.""" await self.setup() resources: list[Resource] = [] @@ -333,17 +309,10 @@ async def list_resources(self, request: ListResourcesRequest) -> ListResourcesRe ) ) - return ListResourcesResult(resources=resources) - - async def list_resource_templates(self, request: ListResourceTemplatesRequest) -> ListResourceTemplatesResult: - """Handle MCP requests to list available resource templates. - - Args: - request: The MCP ListResourceTemplatesRequest. + return resources - Returns: - ListResourceTemplatesResult containing resources with URI templates. - """ + async def _list_resource_templates(self) -> list[types.ResourceTemplate]: + """Handle MCP requests to list available resource templates.""" await self.setup() templates: list[ResourceTemplate] = [] @@ -362,24 +331,12 @@ async def list_resource_templates(self, request: ListResourceTemplatesRequest) - ) ) - return ListResourceTemplatesResult(resourceTemplates=templates) - - async def read_resource(self, request: ReadResourceRequest) -> ReadResourceResult: - """Handle MCP requests to read a specific resource. - - Args: - request: The MCP ReadResourceRequest containing the resource URI. - - Returns: - ReadResourceResult with the resource content. + return templates - Raises: - GenkitError: If no matching resource is found. - """ + async def _read_resource(self, uri: str) -> str | bytes | list[types.TextResourceContents | types.BlobResourceContents]: + """Handle MCP requests to read a specific resource.""" await self.setup() - uri = request.params.uri - # Check for exact URI match resource = self.resource_uri_map.get(uri) @@ -398,10 +355,9 @@ async def read_resource(self, request: ReadResourceRequest) -> ReadResourceResul result = result.response # Convert content to MCP format + # For SDK server usage, we return the contents list directly content = result.get('content', []) if isinstance(result, dict) else result.content - contents = to_mcp_resource_contents(uri, content) - - return ReadResourceResult(contents=contents) + return to_mcp_resource_contents(uri, content) async def start(self, transport: Any = None) -> None: """Start the MCP server with the specified transport. @@ -411,7 +367,7 @@ async def start(self, transport: Any = None) -> None: a StdioServerTransport will be created and used. """ if not transport: - transport = await stdio_server() + transport = stdio_server() await self.setup() diff --git a/py/samples/mcp/README.md b/py/samples/mcp/README.md new file mode 100644 index 0000000000..5c7e0c1e8b --- /dev/null +++ b/py/samples/mcp/README.md @@ -0,0 +1,49 @@ +# MCP Sample + +This sample demonstrates using the MCP (Model Context Protocol) plugin with Genkit Python SDK. + +## Setup environment + +Obtain an API key from [ai.dev](https://ai.dev). + +Export the API key as env variable `GEMINI\_API\_KEY` in your shell +configuration. + +### Run the MCP Client/Host +```bash +cd py/samples/mcp +genkit start -- uv run src/index.py +``` + +This will: +1. Connect to the configured MCP servers +2. Execute sample flows demonstrating tool usage +3. Clean up connections on exit + +### Run the MCP Server +```bash +cd py/samples/mcp +genkit start -- uv run src/server.py +``` + +This starts an MCP server on stdio that other MCP clients can connect to. + +## Requirements + +- Python 3.10+ +- `mcp` - Model Context Protocol Python SDK +- `genkit` - Genkit Python SDK +- `genkit-plugins-google-genai` - Google AI plugin for Genkit + +## MCP Servers Used + +The sample connects to these MCP servers (must be available): +- **mcp-server-git** - Install via `uvx mcp-server-git` +- **@modelcontextprotocol/server-filesystem** - Install via npm +- **@modelcontextprotocol/server-everything** - Install via npm + +## Learn More + +- [MCP Documentation](https://modelcontextprotocol.io/) +- [Genkit Python Documentation](https://firebase.google.com/docs/genkit) +- [MCP Plugin Source](../../plugins/mcp/) diff --git a/py/samples/mcp/pyproject.toml b/py/samples/mcp/pyproject.toml new file mode 100644 index 0000000000..01cb234ea1 --- /dev/null +++ b/py/samples/mcp/pyproject.toml @@ -0,0 +1,24 @@ +[project] +name = "mcp-sample" +version = "0.1.0" +description = "MCP sample application for Genkit Python SDK" +readme = "README.md" +requires-python = ">=3.10" +dependencies = [ + "genkit", + "genkit-plugin-google-genai", + "genkit-plugins-mcp", + "mcp", +] + +[tool.uv.sources] +genkit = { workspace = true } +genkit-plugin-google-genai = { workspace = true } +genkit-plugins-mcp = { workspace = true } + +[build-system] +requires = ["hatchling"] +build-backend = "hatchling.build" + +[tool.hatch.build.targets.wheel] +packages = ["src"] diff --git a/py/samples/mcp/src/http_server.py b/py/samples/mcp/src/http_server.py new file mode 100644 index 0000000000..69ae4b98e3 --- /dev/null +++ b/py/samples/mcp/src/http_server.py @@ -0,0 +1,101 @@ +# Copyright 2025 Google LLC +# +# Licensed 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. + +""" +HTTP MCP Server Example + +This demonstrates creating an HTTP-based MCP server using SSE transport +with Starlette and the official MCP Python SDK. +""" + +import asyncio +import logging +import uvicorn +from starlette.applications import Starlette +from starlette.routing import Route, Mount +from starlette.responses import Response + +from mcp.server import Server +from mcp.server.sse import SseServerTransport +import mcp.types as types + +# Configure logging +logging.basicConfig(level=logging.INFO) +logger = logging.getLogger(__name__) + +async def main(): + """Start the HTTP MCP server.""" + + # Create SSE transport logic + # The endpoint '/mcp/' is where clients will POST messages + sse = SseServerTransport("/mcp/") + + async def handle_sse(request): + """Handle incoming SSE connections.""" + async with sse.connect_sse( + request.scope, request.receive, request._send + ) as streams: + read_stream, write_stream = streams + + # Create a new server instance for this session + # This mirrors the JS logic of creating a new McpServer per connection + server = Server("example-server", version="1.0.0") + + @server.list_tools() + async def list_tools() -> list[types.Tool]: + return [ + types.Tool( + name="test_http", + description="Test HTTP transport", + inputSchema={"type": "object", "properties": {}}, + ) + ] + + @server.call_tool() + async def call_tool(name: str, arguments: dict) -> list[types.TextContent]: + if name == "test_http": + # In this SSE implementation, valid session ID is internal + # but we can return a confirmation. + return [types.TextContent(type="text", text="Session Active")] + raise ValueError(f"Unknown tool: {name}") + + # Run the server with the streams + await server.run( + read_stream, + write_stream, + server.create_initialization_options() + ) + + # Return empty response after connection closes + return Response() + + # Define routes + # GET /mcp -> Starts SSE stream + # POST /mcp/ -> Handles messages (via SseServerTransport) + routes = [ + Route("/mcp", endpoint=handle_sse, methods=["GET"]), + Mount("/mcp/", app=sse.handle_post_message), + ] + + app = Starlette(routes=routes) + + config = uvicorn.Config(app, host="0.0.0.0", port=3334, log_level="info") + server = uvicorn.Server(config) + + print("HTTP MCP server running on http://localhost:3334/mcp") + await server.serve() + + +if __name__ == '__main__': + asyncio.run(main()) diff --git a/py/samples/mcp/src/main.py b/py/samples/mcp/src/main.py new file mode 100644 index 0000000000..fe15079ed1 --- /dev/null +++ b/py/samples/mcp/src/main.py @@ -0,0 +1,318 @@ +# Copyright 2025 Google LLC +# +# Licensed 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. +# +# SPDX-License-Identifier: Apache-2.0 + + +import asyncio +import os +from pathlib import Path +from pydantic import BaseModel + +from genkit.ai import Genkit +from genkit.plugins.mcp import create_mcp_host, McpServerConfig +from genkit.plugins.google_genai import GoogleAI + +import structlog + +logger = structlog.get_logger(__name__) + +# Get the current directory +current_dir = Path(__file__).parent +workspace_dir = current_dir.parent / "test-workspace" +repo_root = current_dir.parent.parent.parent.parent + +# Initialize Genkit with GoogleAI +ai = Genkit( + plugins=[GoogleAI()], + model='googleai/gemini-2.5-flash' +) + +# Create MCP host with multiple servers +mcp_host = create_mcp_host({ + 'git-client': McpServerConfig( + command='uvx', + args=['mcp-server-git'] + ), + 'fs': McpServerConfig( + command='npx', + args=[ + '-y', + '@modelcontextprotocol/server-filesystem', + str(workspace_dir) + ] + ), + 'everything': McpServerConfig( + command='npx', + args=['-y', '@modelcontextprotocol/server-everything'] + ), +}) + + +@ai.flow(name='git_commits') +async def git_commits(query: str = ''): + """Summarize recent git commits using MCP git client.""" + await mcp_host.start() + tools = await mcp_host.register_tools(ai) + + result = await ai.generate( + prompt=f"summarize last 5 commits in '{repo_root}'", + tools=tools + ) + + await mcp_host.close() + return result.text + + +@ai.flow(name='dynamic_git_commits') +async def dynamic_git_commits(query: str = ''): + """Summarize recent git commits using 'all' tools matching pattern.""" + await mcp_host.start() + tools = await mcp_host.register_tools(ai) + + # Simulate wildcard matching "git-client:tool/*" by passing all tools + # (since registration prefixes with server name) + # JS: tools: ['test-mcp-manager:tool/*'] + + result = await ai.generate( + prompt=f"summarize last 5 commits in '{repo_root}'", + tools=tools + ) + + await mcp_host.close() + return result.text + +@ai.flow(name='get_file') +async def get_file(query: str = ''): + """Read and summarize a file using MCP filesystem client.""" + await mcp_host.start() + tools = await mcp_host.register_tools(ai) + + result = await ai.generate( + prompt=f"summarize contents of hello-world.txt (in '{workspace_dir}')", + tools=tools + ) + + await mcp_host.close() + return result.text + + +@ai.flow(name='dynamic_get_file') +async def dynamic_get_file(query: str = ''): + """Read file using specific tool selection.""" + await mcp_host.start() + tools = await mcp_host.register_tools(ai) + + # Filter for specific tool: 'fs/read_file' + # JS: tools: ['test-mcp-manager:tool/fs/read_file'] + import fnmatch + filtered_tools = [t for t in tools if fnmatch.fnmatch(t, '*/fs/read_file') or t.endswith('fs/read_file')] + + result = await ai.generate( + prompt=f"summarize contents of hello-world.txt (in '{workspace_dir}')", + tools=filtered_tools + ) + + await mcp_host.close() + return result.text + + +@ai.flow(name='dynamic_prefix_tool') +async def dynamic_prefix_tool(query: str = ''): + """Read file using prefix tool selection.""" + await mcp_host.start() + tools = await mcp_host.register_tools(ai) + + # Filter for prefix: 'fs/read_*' + # JS: tools: ['test-mcp-manager:tool/fs/read_*'] + import fnmatch + filtered_tools = [t for t in tools if fnmatch.fnmatch(t, '*/fs/read_*')] + + result = await ai.generate( + prompt=f"summarize contents of hello-world.txt (in '{workspace_dir}')", + tools=filtered_tools + ) + + await mcp_host.close() + return result.text + + +@ai.flow(name='dynamic_disable_enable') +async def dynamic_disable_enable(query: str = ''): + """Test disabling and re-enabling an MCP client.""" + await mcp_host.start() + tools = await mcp_host.register_tools(ai) + + import fnmatch + filtered_tools = [t for t in tools if fnmatch.fnmatch(t, '*/fs/read_file') or t.endswith('fs/read_file')] + + # 1. Run successfully + result1 = await ai.generate( + prompt=f"summarize contents of hello-world.txt (in '{workspace_dir}')", + tools=filtered_tools + ) + text1 = result1.text + + # 2. Disable 'fs' and try to run (should fail) + await mcp_host.disable('fs') + text2 = "" + try: + # Note: In Python, we might need to verify if tools list is updated + # or if the tool call fails. disable() closes connection. + # register_tools should ideally be called again or the tool invocation fails. + # Since we passed 'filtered_tools' (names), the model will try to call. + # The tool wrapper checks connection. + result = await ai.generate( + prompt=f"summarize contents of hello-world.txt (in '{workspace_dir}')", + tools=filtered_tools + ) + text2 = f"ERROR! This should have failed but succeeded: {result.text}" + except Exception as e: + text2 = str(e) + + # 3. Re-enable 'fs' and run + await mcp_host.enable('fs') + # Re-registering might be needed if registry was cleaned, but here we just re-connnect + # Implementation detail: Does register_tools need to be called again? + # Code shows wrappers capture client, client.session is updated on connect. + await mcp_host.clients['fs'].connect() + + result3 = await ai.generate( + prompt=f"summarize contents of hello-world.txt (in '{workspace_dir}')", + tools=filtered_tools + ) + text3 = result3.text + + await mcp_host.close() + + return f"Original:
{text1}
After Disable:
{text2}
After Enable:
{text3}" + + +@ai.flow(name='test_resource') +async def test_resource(query: str = ''): + """Test reading a resource (simulated).""" + await mcp_host.start() + + # Python SDK doesn't support 'resources' param in generate yet. + # We manually fetch the resource and add to prompt. + # JS: resources: await mcpHost.getActiveResources(ai) + + resource_content = "Resource not found" + uri = 'test://static/resource/1' + + # In a real implementation we would look up the resource provider. + # Here we search 'everything' client or similar. + found = False + for client in mcp_host.clients.values(): + if client.session and not client.config.disabled: + try: + # Try reading directly + res = await client.read_resource(uri) + if res and res.contents: + resource_content = res.contents[0].text + found = True + break + except Exception: + continue + + result = await ai.generate( + prompt=f"analyze this: {resource_content}", + ) + + await mcp_host.close() + return result.text + + +@ai.flow(name='dynamic_test_resources') +async def dynamic_test_resources(query: str = ''): + """Test reading resources with wildcard (simulated).""" + # Same simulation as test_resource + return await test_resource(query) + + +@ai.flow(name='dynamic_test_one_resource') +async def dynamic_test_one_resource(query: str = ''): + """Test reading one specific resource (simulated).""" + # Same simulation as test_resource + return await test_resource(query) + + + +@ai.flow(name='update_file') +async def update_file(query: str = ''): + """Update a file using MCP filesystem client.""" + await mcp_host.start() + tools = await mcp_host.register_tools(ai) + + result = await ai.generate( + prompt=f"Improve hello-world.txt (in '{workspace_dir}') by rewriting the text, making it longer, use your imagination.", + tools=tools + ) + + await mcp_host.close() + return result.text + + +class ControlMcpInput(BaseModel): + action: str # 'RECONNECT', 'ENABLE', 'DISABLE', 'DISCONNECT' + client_id: str = 'git-client' + + +@ai.flow(name='control_mcp') +async def control_mcp(input: ControlMcpInput): + """Control MCP client connections (enable/disable/reconnect).""" + client_id = input.client_id + action = input.action.upper() + + if action == 'DISABLE': + if client_id in mcp_host.clients: + mcp_host.clients[client_id].config.disabled = True + await mcp_host.clients[client_id].close() + elif action == 'DISCONNECT': + if client_id in mcp_host.clients: + await mcp_host.clients[client_id].close() + elif action == 'RECONNECT': + if client_id in mcp_host.clients: + await mcp_host.clients[client_id].connect() + elif action == 'ENABLE': + if client_id in mcp_host.clients: + mcp_host.clients[client_id].config.disabled = False + await mcp_host.clients[client_id].connect() + + return f"Action {action} completed for {client_id}" + + +async def main(): + """Run sample flows.""" + logger.info("Starting MCP sample application") + + # Test git commits flow + logger.info("Testing git_commits flow...") + try: + result = await git_commits() + logger.info("git_commits result", result=result[:200]) + except Exception as e: + logger.error("git_commits failed", error=str(e)) + + # Test get_file flow + logger.info("Testing get_file flow...") + try: + result = await get_file() + logger.info("get_file result", result=result[:200]) + except Exception as e: + logger.error("get_file failed", error=str(e)) + + +if __name__ == '__main__': + ai.run_main(main()) diff --git a/py/samples/mcp/src/server.py b/py/samples/mcp/src/server.py new file mode 100644 index 0000000000..a472143e2f --- /dev/null +++ b/py/samples/mcp/src/server.py @@ -0,0 +1,98 @@ +# Copyright 2025 Google LLC +# +# Licensed 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. +# +# SPDX-License-Identifier: Apache-2.0 + +""" +MCP Server Example + +This demonstrates creating an MCP server that exposes Genkit tools, prompts, +and resources through the Model Context Protocol. +""" + +import asyncio +from pydantic import BaseModel, Field + +from genkit.ai import Genkit +from genkit.plugins.mcp import create_mcp_server, McpServerOptions +from genkit.plugins.google_genai import GoogleAI + +# Initialize Genkit +ai = Genkit(plugins=[]) + + +# Define a tool +class AddInput(BaseModel): + a: int = Field(..., description="First number") + b: int = Field(..., description="Second number") + + +@ai.tool(name='add', description='add two numbers together') +def add(input: AddInput) -> int: + return input.a + input.b + + +# Define a prompt +happy_prompt = ai.define_prompt( + input_schema={'action': str}, + prompt="If you're happy and you know it, {{action}}.", +) + + +from genkit.core.action.types import ActionKind + +# Define resources (manually registering since define_resource is not yet in Genkit API) +def define_resource(name: str, uri: str, fn): + ai.registry.register_action( + kind=ActionKind.RESOURCE, + name=name, + fn=fn, + metadata={'resource': {'uri': uri}} + ) + +def define_resource_template(name: str, template: str, fn): + ai.registry.register_action( + kind=ActionKind.RESOURCE, + name=name, + fn=fn, + metadata={'resource': {'template': template}} + ) + +def my_resource_handler(inp): + return {'content': [{'text': 'my resource'}]} + +define_resource('my resources', 'test://static/resource/1', my_resource_handler) + + +def file_resource_handler(inp): + uri = inp.get('uri') + return {'content': [{'text': f'file contents for {uri}'}]} + +define_resource_template('file', 'file://{path}', file_resource_handler) + + +async def main(): + """Start the MCP server.""" + # Create MCP server + server = create_mcp_server( + ai, + McpServerOptions(name='example_server', version='0.0.1') + ) + + print("Starting MCP server on stdio...") + await server.start() + + +if __name__ == '__main__': + asyncio.run(main()) diff --git a/py/samples/mcp/test-workspace/hello-world.txt b/py/samples/mcp/test-workspace/hello-world.txt new file mode 100644 index 0000000000..723e9faf6d --- /dev/null +++ b/py/samples/mcp/test-workspace/hello-world.txt @@ -0,0 +1,4 @@ +Hello, World! + +This is a test file for the MCP filesystem sample. +It demonstrates reading and writing files through the MCP protocol. diff --git a/py/uv.lock b/py/uv.lock index c414fe7863..5de2fd654a 100644 --- a/py/uv.lock +++ b/py/uv.lock @@ -1,5 +1,5 @@ version = 1 -revision = 2 +revision = 3 requires-python = ">=3.10" resolution-markers = [ "python_full_version >= '3.14'", @@ -36,6 +36,7 @@ members = [ "google-genai-image", "google-genai-vertexai-hello", "google-genai-vertexai-image", + "mcp-sample", "menu", "model-garden-example", "multi-server", @@ -3317,6 +3318,25 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/e2/fc/6dc7659c2ae5ddf280477011f4213a74f806862856b796ef08f028e664bf/mcp-1.25.0-py3-none-any.whl", hash = "sha256:b37c38144a666add0862614cc79ec276e97d72aa8ca26d622818d4e278b9721a", size = 233076, upload-time = "2025-12-19T10:19:55.416Z" }, ] +[[package]] +name = "mcp-sample" +version = "0.1.0" +source = { editable = "samples/mcp" } +dependencies = [ + { name = "genkit" }, + { name = "genkit-plugin-google-genai" }, + { name = "genkit-plugins-mcp" }, + { name = "mcp" }, +] + +[package.metadata] +requires-dist = [ + { name = "genkit", editable = "packages/genkit" }, + { name = "genkit-plugin-google-genai", editable = "plugins/google-genai" }, + { name = "genkit-plugins-mcp", editable = "plugins/mcp" }, + { name = "mcp" }, +] + [[package]] name = "mdurl" version = "0.1.2" From 019939c995a32e14f5af038fa5245258d3aff0fc Mon Sep 17 00:00:00 2001 From: Mengqin Shen Date: Tue, 30 Dec 2025 10:32:07 -0800 Subject: [PATCH 4/4] feat(py):add sample test for mcp --- .../mcp/src/genkit/plugins/mcp/server.py | 27 ++-- py/samples/mcp/README.md | 13 +- py/samples/mcp/pyproject.toml | 38 +++-- py/samples/mcp/src/http_server.py | 55 ++++--- py/samples/mcp/src/main.py | 135 +++++++----------- py/samples/mcp/src/server.py | 31 ++-- 6 files changed, 146 insertions(+), 153 deletions(-) diff --git a/py/plugins/mcp/src/genkit/plugins/mcp/server.py b/py/plugins/mcp/src/genkit/plugins/mcp/server.py index bfe237d5dd..270fdc6662 100644 --- a/py/plugins/mcp/src/genkit/plugins/mcp/server.py +++ b/py/plugins/mcp/src/genkit/plugins/mcp/server.py @@ -23,9 +23,9 @@ from typing import Any, Optional import structlog -from pydantic import BaseModel, AnyUrl -import mcp.types as types +from pydantic import AnyUrl, BaseModel +import mcp.types as types from genkit.ai import Genkit from genkit.blocks.resource import matches_uri_template from genkit.core.action._key import parse_action_key @@ -117,13 +117,15 @@ async def setup(self) -> None: self.server = Server(self.options.name) # Register request handlers using decorators - + @self.server.list_tools() async def list_tools() -> list[types.Tool]: return await self._list_tools() @self.server.call_tool() - async def call_tool(name: str, arguments: dict | None) -> list[types.TextContent | types.ImageContent | types.EmbeddedResource]: + async def call_tool( + name: str, arguments: dict | None + ) -> list[types.TextContent | types.ImageContent | types.EmbeddedResource]: return await self._call_tool(name, arguments) @self.server.list_prompts() @@ -133,7 +135,7 @@ async def list_prompts() -> list[types.Prompt]: @self.server.get_prompt() async def get_prompt(name: str, arguments: dict[str, str] | None) -> types.GetPromptResult: return await self._get_prompt(name, arguments) - + @self.server.list_resources() async def list_resources() -> list[types.Resource]: return await self._list_resources() @@ -141,7 +143,7 @@ async def list_resources() -> list[types.Resource]: @self.server.list_resource_templates() async def list_resource_templates() -> list[types.ResourceTemplate]: return await self._list_resource_templates() - + @self.server.read_resource() async def read_resource(uri: AnyUrl) -> str | bytes: # Note: The MCP SDK signature for read_resource expects returning content @@ -149,7 +151,6 @@ async def read_resource(uri: AnyUrl) -> str | bytes: # Based on lowlevel/server.py it returns ReadResourceContents which is content list return await self._read_resource(str(uri)) - # Resolve all actions from Genkit registry # We need the actual Action objects, not just serializable dicts self.tool_actions = [] @@ -227,7 +228,9 @@ async def _list_tools(self) -> list[types.Tool]: return tools - async def _call_tool(self, name: str, arguments: dict | None) -> list[types.TextContent | types.ImageContent | types.EmbeddedResource]: + async def _call_tool( + self, name: str, arguments: dict | None + ) -> list[types.TextContent | types.ImageContent | types.EmbeddedResource]: """Handle MCP requests to call a specific tool.""" await self.setup() @@ -235,9 +238,7 @@ async def _call_tool(self, name: str, arguments: dict | None) -> list[types.Text tool = self.tool_actions_map.get(name) if not tool: - raise GenkitError( - status='NOT_FOUND', message=f"Tried to call tool '{name}' but it could not be found." - ) + raise GenkitError(status='NOT_FOUND', message=f"Tried to call tool '{name}' but it could not be found.") # Execute the tool result = await tool.arun(arguments) @@ -333,7 +334,9 @@ async def _list_resource_templates(self) -> list[types.ResourceTemplate]: return templates - async def _read_resource(self, uri: str) -> str | bytes | list[types.TextResourceContents | types.BlobResourceContents]: + async def _read_resource( + self, uri: str + ) -> str | bytes | list[types.TextResourceContents | types.BlobResourceContents]: """Handle MCP requests to read a specific resource.""" await self.setup() diff --git a/py/samples/mcp/README.md b/py/samples/mcp/README.md index 5c7e0c1e8b..650846a6fa 100644 --- a/py/samples/mcp/README.md +++ b/py/samples/mcp/README.md @@ -12,7 +12,18 @@ configuration. ### Run the MCP Client/Host ```bash cd py/samples/mcp -genkit start -- uv run src/index.py +genkit start -- uv run src/main.py +``` + +This will: +1. Connect to the configured MCP servers +2. Execute sample flows demonstrating tool usage +3. Clean up connections on exit + +### Run the MCP Client/Host +```bash +cd py/samples/mcp +genkit start -- uv run src/http_server.py ``` This will: diff --git a/py/samples/mcp/pyproject.toml b/py/samples/mcp/pyproject.toml index 01cb234ea1..4403a41fcb 100644 --- a/py/samples/mcp/pyproject.toml +++ b/py/samples/mcp/pyproject.toml @@ -1,24 +1,40 @@ +# Copyright 2025 Google LLC +# +# Licensed 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. +# +# SPDX-License-Identifier: Apache-2.0 + [project] -name = "mcp-sample" -version = "0.1.0" +dependencies = [ + "genkit", + "genkit-plugin-google-genai", + "genkit-plugins-mcp", + "mcp", +] description = "MCP sample application for Genkit Python SDK" +name = "mcp-sample" readme = "README.md" requires-python = ">=3.10" -dependencies = [ - "genkit", - "genkit-plugin-google-genai", - "genkit-plugins-mcp", - "mcp", -] +version = "0.1.0" [tool.uv.sources] -genkit = { workspace = true } +genkit = { workspace = true } genkit-plugin-google-genai = { workspace = true } -genkit-plugins-mcp = { workspace = true } +genkit-plugins-mcp = { workspace = true } [build-system] -requires = ["hatchling"] build-backend = "hatchling.build" +requires = ["hatchling"] [tool.hatch.build.targets.wheel] packages = ["src"] diff --git a/py/samples/mcp/src/http_server.py b/py/samples/mcp/src/http_server.py index 69ae4b98e3..d0040b5000 100644 --- a/py/samples/mcp/src/http_server.py +++ b/py/samples/mcp/src/http_server.py @@ -21,62 +21,57 @@ import asyncio import logging -import uvicorn -from starlette.applications import Starlette -from starlette.routing import Route, Mount -from starlette.responses import Response +import mcp.types as types +import uvicorn from mcp.server import Server from mcp.server.sse import SseServerTransport -import mcp.types as types +from starlette.applications import Starlette +from starlette.responses import Response +from starlette.routing import Mount, Route # Configure logging logging.basicConfig(level=logging.INFO) logger = logging.getLogger(__name__) + async def main(): """Start the HTTP MCP server.""" - + # Create SSE transport logic # The endpoint '/mcp/' is where clients will POST messages - sse = SseServerTransport("/mcp/") + sse = SseServerTransport('/mcp/') async def handle_sse(request): """Handle incoming SSE connections.""" - async with sse.connect_sse( - request.scope, request.receive, request._send - ) as streams: + async with sse.connect_sse(request.scope, request.receive, request._send) as streams: read_stream, write_stream = streams - + # Create a new server instance for this session # This mirrors the JS logic of creating a new McpServer per connection - server = Server("example-server", version="1.0.0") + server = Server('example-server', version='1.0.0') @server.list_tools() async def list_tools() -> list[types.Tool]: return [ types.Tool( - name="test_http", - description="Test HTTP transport", - inputSchema={"type": "object", "properties": {}}, + name='test_http', + description='Test HTTP transport', + inputSchema={'type': 'object', 'properties': {}}, ) ] @server.call_tool() async def call_tool(name: str, arguments: dict) -> list[types.TextContent]: - if name == "test_http": - # In this SSE implementation, valid session ID is internal + if name == 'test_http': + # In this SSE implementation, valid session ID is internal # but we can return a confirmation. - return [types.TextContent(type="text", text="Session Active")] - raise ValueError(f"Unknown tool: {name}") + return [types.TextContent(type='text', text='Session Active')] + raise ValueError(f'Unknown tool: {name}') # Run the server with the streams - await server.run( - read_stream, - write_stream, - server.create_initialization_options() - ) - + await server.run(read_stream, write_stream, server.create_initialization_options()) + # Return empty response after connection closes return Response() @@ -84,16 +79,16 @@ async def call_tool(name: str, arguments: dict) -> list[types.TextContent]: # GET /mcp -> Starts SSE stream # POST /mcp/ -> Handles messages (via SseServerTransport) routes = [ - Route("/mcp", endpoint=handle_sse, methods=["GET"]), - Mount("/mcp/", app=sse.handle_post_message), + Route('/mcp', endpoint=handle_sse, methods=['GET']), + Mount('/mcp/', app=sse.handle_post_message), ] app = Starlette(routes=routes) - config = uvicorn.Config(app, host="0.0.0.0", port=3334, log_level="info") + config = uvicorn.Config(app, host='0.0.0.0', port=3334, log_level='info') server = uvicorn.Server(config) - - print("HTTP MCP server running on http://localhost:3334/mcp") + + print('HTTP MCP server running on http://localhost:3334/mcp') await server.serve() diff --git a/py/samples/mcp/src/main.py b/py/samples/mcp/src/main.py index fe15079ed1..2e32301a95 100644 --- a/py/samples/mcp/src/main.py +++ b/py/samples/mcp/src/main.py @@ -18,45 +18,29 @@ import asyncio import os from pathlib import Path + +import structlog from pydantic import BaseModel from genkit.ai import Genkit -from genkit.plugins.mcp import create_mcp_host, McpServerConfig from genkit.plugins.google_genai import GoogleAI - -import structlog +from genkit.plugins.mcp import McpServerConfig, create_mcp_host logger = structlog.get_logger(__name__) # Get the current directory current_dir = Path(__file__).parent -workspace_dir = current_dir.parent / "test-workspace" +workspace_dir = current_dir.parent / 'test-workspace' repo_root = current_dir.parent.parent.parent.parent # Initialize Genkit with GoogleAI -ai = Genkit( - plugins=[GoogleAI()], - model='googleai/gemini-2.5-flash' -) +ai = Genkit(plugins=[GoogleAI()], model='googleai/gemini-2.5-flash') # Create MCP host with multiple servers mcp_host = create_mcp_host({ - 'git-client': McpServerConfig( - command='uvx', - args=['mcp-server-git'] - ), - 'fs': McpServerConfig( - command='npx', - args=[ - '-y', - '@modelcontextprotocol/server-filesystem', - str(workspace_dir) - ] - ), - 'everything': McpServerConfig( - command='npx', - args=['-y', '@modelcontextprotocol/server-everything'] - ), + 'git-client': McpServerConfig(command='uvx', args=['mcp-server-git']), + 'fs': McpServerConfig(command='npx', args=['-y', '@modelcontextprotocol/server-filesystem', str(workspace_dir)]), + 'everything': McpServerConfig(command='npx', args=['-y', '@modelcontextprotocol/server-everything']), }) @@ -66,10 +50,7 @@ async def git_commits(query: str = ''): await mcp_host.start() tools = await mcp_host.register_tools(ai) - result = await ai.generate( - prompt=f"summarize last 5 commits in '{repo_root}'", - tools=tools - ) + result = await ai.generate(prompt=f"summarize last 5 commits in '{repo_root}'", tools=tools) await mcp_host.close() return result.text @@ -80,29 +61,24 @@ async def dynamic_git_commits(query: str = ''): """Summarize recent git commits using 'all' tools matching pattern.""" await mcp_host.start() tools = await mcp_host.register_tools(ai) - - # Simulate wildcard matching "git-client:tool/*" by passing all tools + + # Simulate wildcard matching "git-client:tool/*" by passing all tools # (since registration prefixes with server name) # JS: tools: ['test-mcp-manager:tool/*'] - - result = await ai.generate( - prompt=f"summarize last 5 commits in '{repo_root}'", - tools=tools - ) - + + result = await ai.generate(prompt=f"summarize last 5 commits in '{repo_root}'", tools=tools) + await mcp_host.close() return result.text + @ai.flow(name='get_file') async def get_file(query: str = ''): """Read and summarize a file using MCP filesystem client.""" await mcp_host.start() tools = await mcp_host.register_tools(ai) - result = await ai.generate( - prompt=f"summarize contents of hello-world.txt (in '{workspace_dir}')", - tools=tools - ) + result = await ai.generate(prompt=f"summarize contents of hello-world.txt (in '{workspace_dir}')", tools=tools) await mcp_host.close() return result.text @@ -113,17 +89,17 @@ async def dynamic_get_file(query: str = ''): """Read file using specific tool selection.""" await mcp_host.start() tools = await mcp_host.register_tools(ai) - + # Filter for specific tool: 'fs/read_file' # JS: tools: ['test-mcp-manager:tool/fs/read_file'] import fnmatch + filtered_tools = [t for t in tools if fnmatch.fnmatch(t, '*/fs/read_file') or t.endswith('fs/read_file')] result = await ai.generate( - prompt=f"summarize contents of hello-world.txt (in '{workspace_dir}')", - tools=filtered_tools + prompt=f"summarize contents of hello-world.txt (in '{workspace_dir}')", tools=filtered_tools ) - + await mcp_host.close() return result.text @@ -133,17 +109,17 @@ async def dynamic_prefix_tool(query: str = ''): """Read file using prefix tool selection.""" await mcp_host.start() tools = await mcp_host.register_tools(ai) - + # Filter for prefix: 'fs/read_*' # JS: tools: ['test-mcp-manager:tool/fs/read_*'] import fnmatch + filtered_tools = [t for t in tools if fnmatch.fnmatch(t, '*/fs/read_*')] result = await ai.generate( - prompt=f"summarize contents of hello-world.txt (in '{workspace_dir}')", - tools=filtered_tools + prompt=f"summarize contents of hello-world.txt (in '{workspace_dir}')", tools=filtered_tools ) - + await mcp_host.close() return result.text @@ -153,64 +129,62 @@ async def dynamic_disable_enable(query: str = ''): """Test disabling and re-enabling an MCP client.""" await mcp_host.start() tools = await mcp_host.register_tools(ai) - + import fnmatch + filtered_tools = [t for t in tools if fnmatch.fnmatch(t, '*/fs/read_file') or t.endswith('fs/read_file')] # 1. Run successfully result1 = await ai.generate( - prompt=f"summarize contents of hello-world.txt (in '{workspace_dir}')", - tools=filtered_tools + prompt=f"summarize contents of hello-world.txt (in '{workspace_dir}')", tools=filtered_tools ) text1 = result1.text # 2. Disable 'fs' and try to run (should fail) await mcp_host.disable('fs') - text2 = "" + text2 = '' try: - # Note: In Python, we might need to verify if tools list is updated + # Note: In Python, we might need to verify if tools list is updated # or if the tool call fails. disable() closes connection. # register_tools should ideally be called again or the tool invocation fails. # Since we passed 'filtered_tools' (names), the model will try to call. # The tool wrapper checks connection. result = await ai.generate( - prompt=f"summarize contents of hello-world.txt (in '{workspace_dir}')", - tools=filtered_tools + prompt=f"summarize contents of hello-world.txt (in '{workspace_dir}')", tools=filtered_tools ) - text2 = f"ERROR! This should have failed but succeeded: {result.text}" + text2 = f'ERROR! This should have failed but succeeded: {result.text}' except Exception as e: text2 = str(e) # 3. Re-enable 'fs' and run await mcp_host.enable('fs') # Re-registering might be needed if registry was cleaned, but here we just re-connnect - # Implementation detail: Does register_tools need to be called again? + # Implementation detail: Does register_tools need to be called again? # Code shows wrappers capture client, client.session is updated on connect. await mcp_host.clients['fs'].connect() - + result3 = await ai.generate( - prompt=f"summarize contents of hello-world.txt (in '{workspace_dir}')", - tools=filtered_tools + prompt=f"summarize contents of hello-world.txt (in '{workspace_dir}')", tools=filtered_tools ) text3 = result3.text await mcp_host.close() - - return f"Original:
{text1}
After Disable:
{text2}
After Enable:
{text3}" + + return f'Original:
{text1}
After Disable:
{text2}
After Enable:
{text3}' @ai.flow(name='test_resource') async def test_resource(query: str = ''): """Test reading a resource (simulated).""" await mcp_host.start() - + # Python SDK doesn't support 'resources' param in generate yet. # We manually fetch the resource and add to prompt. # JS: resources: await mcpHost.getActiveResources(ai) - - resource_content = "Resource not found" + + resource_content = 'Resource not found' uri = 'test://static/resource/1' - + # In a real implementation we would look up the resource provider. # Here we search 'everything' client or similar. found = False @@ -220,16 +194,16 @@ async def test_resource(query: str = ''): # Try reading directly res = await client.read_resource(uri) if res and res.contents: - resource_content = res.contents[0].text - found = True - break + resource_content = res.contents[0].text + found = True + break except Exception: continue result = await ai.generate( - prompt=f"analyze this: {resource_content}", + prompt=f'analyze this: {resource_content}', ) - + await mcp_host.close() return result.text @@ -248,7 +222,6 @@ async def dynamic_test_one_resource(query: str = ''): return await test_resource(query) - @ai.flow(name='update_file') async def update_file(query: str = ''): """Update a file using MCP filesystem client.""" @@ -257,7 +230,7 @@ async def update_file(query: str = ''): result = await ai.generate( prompt=f"Improve hello-world.txt (in '{workspace_dir}') by rewriting the text, making it longer, use your imagination.", - tools=tools + tools=tools, ) await mcp_host.close() @@ -290,28 +263,28 @@ async def control_mcp(input: ControlMcpInput): mcp_host.clients[client_id].config.disabled = False await mcp_host.clients[client_id].connect() - return f"Action {action} completed for {client_id}" + return f'Action {action} completed for {client_id}' async def main(): """Run sample flows.""" - logger.info("Starting MCP sample application") + logger.info('Starting MCP sample application') # Test git commits flow - logger.info("Testing git_commits flow...") + logger.info('Testing git_commits flow...') try: result = await git_commits() - logger.info("git_commits result", result=result[:200]) + logger.info('git_commits result', result=result[:200]) except Exception as e: - logger.error("git_commits failed", error=str(e)) + logger.error('git_commits failed', error=str(e)) # Test get_file flow - logger.info("Testing get_file flow...") + logger.info('Testing get_file flow...') try: result = await get_file() - logger.info("get_file result", result=result[:200]) + logger.info('get_file result', result=result[:200]) except Exception as e: - logger.error("get_file failed", error=str(e)) + logger.error('get_file failed', error=str(e)) if __name__ == '__main__': diff --git a/py/samples/mcp/src/server.py b/py/samples/mcp/src/server.py index a472143e2f..3638dbfad5 100644 --- a/py/samples/mcp/src/server.py +++ b/py/samples/mcp/src/server.py @@ -22,11 +22,12 @@ """ import asyncio + from pydantic import BaseModel, Field from genkit.ai import Genkit -from genkit.plugins.mcp import create_mcp_server, McpServerOptions from genkit.plugins.google_genai import GoogleAI +from genkit.plugins.mcp import McpServerOptions, create_mcp_server # Initialize Genkit ai = Genkit(plugins=[]) @@ -34,8 +35,8 @@ # Define a tool class AddInput(BaseModel): - a: int = Field(..., description="First number") - b: int = Field(..., description="Second number") + a: int = Field(..., description='First number') + b: int = Field(..., description='Second number') @ai.tool(name='add', description='add two numbers together') @@ -52,26 +53,22 @@ def add(input: AddInput) -> int: from genkit.core.action.types import ActionKind + # Define resources (manually registering since define_resource is not yet in Genkit API) def define_resource(name: str, uri: str, fn): - ai.registry.register_action( - kind=ActionKind.RESOURCE, - name=name, - fn=fn, - metadata={'resource': {'uri': uri}} - ) + ai.registry.register_action(kind=ActionKind.RESOURCE, name=name, fn=fn, metadata={'resource': {'uri': uri}}) + def define_resource_template(name: str, template: str, fn): ai.registry.register_action( - kind=ActionKind.RESOURCE, - name=name, - fn=fn, - metadata={'resource': {'template': template}} + kind=ActionKind.RESOURCE, name=name, fn=fn, metadata={'resource': {'template': template}} ) + def my_resource_handler(inp): return {'content': [{'text': 'my resource'}]} + define_resource('my resources', 'test://static/resource/1', my_resource_handler) @@ -79,18 +76,16 @@ def file_resource_handler(inp): uri = inp.get('uri') return {'content': [{'text': f'file contents for {uri}'}]} + define_resource_template('file', 'file://{path}', file_resource_handler) async def main(): """Start the MCP server.""" # Create MCP server - server = create_mcp_server( - ai, - McpServerOptions(name='example_server', version='0.0.1') - ) + server = create_mcp_server(ai, McpServerOptions(name='example_server', version='0.0.1')) - print("Starting MCP server on stdio...") + print('Starting MCP server on stdio...') await server.start()