nvidia-nat-mcp 1.4.0a20251103__py3-none-any.whl → 1.5.0a20260118__py3-none-any.whl

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Files changed (40) hide show
  1. nat/meta/pypi.md +1 -1
  2. nat/plugins/mcp/__init__.py +1 -1
  3. nat/plugins/mcp/auth/__init__.py +1 -1
  4. nat/plugins/mcp/auth/auth_flow_handler.py +1 -1
  5. nat/plugins/mcp/auth/auth_provider.py +3 -2
  6. nat/plugins/mcp/auth/auth_provider_config.py +2 -1
  7. nat/plugins/mcp/auth/register.py +9 -1
  8. nat/plugins/mcp/auth/service_account/__init__.py +14 -0
  9. nat/plugins/mcp/auth/service_account/provider.py +136 -0
  10. nat/plugins/mcp/auth/service_account/provider_config.py +137 -0
  11. nat/plugins/mcp/auth/service_account/token_client.py +156 -0
  12. nat/plugins/mcp/auth/token_storage.py +2 -2
  13. nat/plugins/mcp/cli/__init__.py +15 -0
  14. nat/plugins/mcp/cli/commands.py +1094 -0
  15. nat/plugins/mcp/client/__init__.py +15 -0
  16. nat/plugins/mcp/{client_base.py → client/client_base.py} +18 -10
  17. nat/plugins/mcp/{client_config.py → client/client_config.py} +24 -9
  18. nat/plugins/mcp/{client_impl.py → client/client_impl.py} +253 -62
  19. nat/plugins/mcp/exception_handler.py +1 -1
  20. nat/plugins/mcp/exceptions.py +1 -1
  21. nat/plugins/mcp/register.py +5 -4
  22. nat/plugins/mcp/server/__init__.py +15 -0
  23. nat/plugins/mcp/server/front_end_config.py +109 -0
  24. nat/plugins/mcp/server/front_end_plugin.py +155 -0
  25. nat/plugins/mcp/server/front_end_plugin_worker.py +415 -0
  26. nat/plugins/mcp/server/introspection_token_verifier.py +72 -0
  27. nat/plugins/mcp/server/memory_profiler.py +320 -0
  28. nat/plugins/mcp/server/register_frontend.py +27 -0
  29. nat/plugins/mcp/server/tool_converter.py +290 -0
  30. nat/plugins/mcp/utils.py +1 -1
  31. {nvidia_nat_mcp-1.4.0a20251103.dist-info → nvidia_nat_mcp-1.5.0a20260118.dist-info}/METADATA +5 -5
  32. nvidia_nat_mcp-1.5.0a20260118.dist-info/RECORD +37 -0
  33. nvidia_nat_mcp-1.5.0a20260118.dist-info/entry_points.txt +9 -0
  34. nat/plugins/mcp/tool.py +0 -138
  35. nvidia_nat_mcp-1.4.0a20251103.dist-info/RECORD +0 -23
  36. nvidia_nat_mcp-1.4.0a20251103.dist-info/entry_points.txt +0 -3
  37. {nvidia_nat_mcp-1.4.0a20251103.dist-info → nvidia_nat_mcp-1.5.0a20260118.dist-info}/WHEEL +0 -0
  38. {nvidia_nat_mcp-1.4.0a20251103.dist-info → nvidia_nat_mcp-1.5.0a20260118.dist-info}/licenses/LICENSE-3rd-party.txt +0 -0
  39. {nvidia_nat_mcp-1.4.0a20251103.dist-info → nvidia_nat_mcp-1.5.0a20260118.dist-info}/licenses/LICENSE.md +0 -0
  40. {nvidia_nat_mcp-1.4.0a20251103.dist-info → nvidia_nat_mcp-1.5.0a20260118.dist-info}/top_level.txt +0 -0
@@ -0,0 +1,290 @@
1
+ # SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
2
+ # SPDX-License-Identifier: Apache-2.0
3
+ #
4
+ # Licensed under the Apache License, Version 2.0 (the "License");
5
+ # you may not use this file except in compliance with the License.
6
+ # You may obtain a copy of the License at
7
+ #
8
+ # http://www.apache.org/licenses/LICENSE-2.0
9
+ #
10
+ # Unless required by applicable law or agreed to in writing, software
11
+ # distributed under the License is distributed on an "AS IS" BASIS,
12
+ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13
+ # See the License for the specific language governing permissions and
14
+ # limitations under the License.
15
+
16
+ import json
17
+ import logging
18
+ from inspect import Parameter
19
+ from inspect import Signature
20
+ from typing import TYPE_CHECKING
21
+ from typing import Any
22
+
23
+ from pydantic import BaseModel
24
+ from pydantic.fields import FieldInfo
25
+ from pydantic_core import PydanticUndefined
26
+
27
+ from mcp.server.fastmcp import FastMCP
28
+ from nat.builder.function import Function
29
+ from nat.builder.function_base import FunctionBase
30
+
31
+ if TYPE_CHECKING:
32
+ from nat.plugins.mcp.server.memory_profiler import MemoryProfiler
33
+ from nat.runtime.session import SessionManager
34
+
35
+ logger = logging.getLogger(__name__)
36
+
37
+ # Sentinel: marks "optional; let Pydantic supply default/factory"
38
+ _USE_PYDANTIC_DEFAULT = object()
39
+
40
+
41
+ def is_field_optional(field: FieldInfo) -> tuple[bool, Any]:
42
+ """Determine if a Pydantic field is optional and extract its default value for MCP signatures.
43
+
44
+ For MCP tool signatures, we need to distinguish:
45
+ - Required fields: marked with Parameter.empty
46
+ - Optional with concrete default: use that default
47
+ - Optional with factory: use sentinel so Pydantic can apply the factory later
48
+
49
+ Args:
50
+ field: The Pydantic FieldInfo to check
51
+
52
+ Returns:
53
+ A tuple of (is_optional, default_value):
54
+ - (False, Parameter.empty) for required fields
55
+ - (True, actual_default) for optional fields with explicit defaults
56
+ - (True, _USE_PYDANTIC_DEFAULT) for optional fields with default_factory
57
+ """
58
+ if field.is_required():
59
+ return False, Parameter.empty
60
+
61
+ # Field is optional - has either default or factory
62
+ if field.default is not PydanticUndefined:
63
+ return True, field.default
64
+
65
+ # Factory case: mark optional in signature but don't fabricate a value
66
+ if field.default_factory is not None:
67
+ return True, _USE_PYDANTIC_DEFAULT
68
+
69
+ # Rare corner case: non-required yet no default surfaced
70
+ return True, _USE_PYDANTIC_DEFAULT
71
+
72
+
73
+ def create_function_wrapper(
74
+ function_name: str,
75
+ session_manager: 'SessionManager',
76
+ schema: type[BaseModel],
77
+ memory_profiler: 'MemoryProfiler | None' = None,
78
+ ):
79
+ """Create a wrapper function that exposes a NAT Function as an MCP tool using SessionManager.
80
+
81
+ Here SessionManager.run() which is used to create a Runner that
82
+ automatically handles observability (emits intermediate step events, starts exporters, etc).
83
+
84
+ Args:
85
+ function_name (str): The name of the function/tool
86
+ session_manager (SessionManager): SessionManager wrapping the function/workflow
87
+ schema (type[BaseModel]): The input schema of the function
88
+ memory_profiler: Optional memory profiler to track requests
89
+
90
+ Returns:
91
+ A wrapper function suitable for registration with MCP
92
+ """
93
+ # Check if we're dealing with ChatRequest - special case
94
+ is_chat_request = False
95
+
96
+ # Check if the schema name is ChatRequest
97
+ if schema.__name__ == "ChatRequest" or (hasattr(schema, "__qualname__") and "ChatRequest" in schema.__qualname__):
98
+ is_chat_request = True
99
+ logger.info("Function %s uses ChatRequest - creating simplified interface", function_name)
100
+
101
+ # For ChatRequest, we'll create a simple wrapper with just a query parameter
102
+ parameters = [Parameter(
103
+ name="query",
104
+ kind=Parameter.KEYWORD_ONLY,
105
+ default=Parameter.empty,
106
+ annotation=str,
107
+ )]
108
+ else:
109
+ # Regular case - extract parameter information from the input schema
110
+ # Extract parameter information from the input schema
111
+ param_fields = schema.model_fields
112
+
113
+ parameters = []
114
+ for name, field in param_fields.items():
115
+ # Get the field type and convert to appropriate Python type
116
+ field_type = field.annotation
117
+
118
+ # Check if field is optional and get its default value
119
+ _is_optional, param_default = is_field_optional(field)
120
+
121
+ # Add the parameter to our list
122
+ parameters.append(
123
+ Parameter(
124
+ name=name,
125
+ kind=Parameter.KEYWORD_ONLY,
126
+ default=param_default,
127
+ annotation=field_type,
128
+ ))
129
+
130
+ # Create the function signature WITHOUT the ctx parameter
131
+ # We'll handle this in the wrapper function internally
132
+ sig = Signature(parameters=parameters, return_annotation=str)
133
+
134
+ # Define the actual wrapper function that accepts ctx but doesn't expose it
135
+ def create_wrapper():
136
+
137
+ async def wrapper_with_ctx(**kwargs):
138
+ """Internal wrapper that will be called by MCP.
139
+
140
+ Uses SessionManager.run() which creates a Runner that automatically handles observability.
141
+ """
142
+ # MCP will add a ctx parameter, extract it
143
+ ctx = kwargs.get("ctx")
144
+
145
+ # Remove ctx if present
146
+ if "ctx" in kwargs:
147
+ del kwargs["ctx"]
148
+
149
+ # Process the function call
150
+ if ctx:
151
+ ctx.info("Calling function %s with args: %s", function_name, json.dumps(kwargs, default=str))
152
+ await ctx.report_progress(0, 100)
153
+
154
+ try:
155
+ # Prepare input payload
156
+ if is_chat_request:
157
+ from nat.data_models.api_server import ChatRequest
158
+ # Create a chat request from the query string
159
+ query = kwargs.get("query", "")
160
+ payload = ChatRequest.from_string(query)
161
+ else:
162
+ # Strip sentinel values so Pydantic can apply defaults/factories
163
+ cleaned_kwargs = {k: v for k, v in kwargs.items() if v is not _USE_PYDANTIC_DEFAULT}
164
+ # Always validate with the declared schema
165
+ payload = schema.model_validate(cleaned_kwargs)
166
+
167
+ # Use SessionManager.run() pattern - this automatically handles all observability
168
+ # The Runner created by session_manager.run() will:
169
+ # 1. Start the exporter manager
170
+ # 2. Emit WORKFLOW_START/FUNCTION_START events
171
+ # 3. Execute the function/workflow
172
+ # 4. Emit WORKFLOW_END/FUNCTION_END events
173
+ # 5. Stop the exporter manager
174
+ async with session_manager.run(payload) as runner:
175
+ result = await runner.result()
176
+
177
+ # Report completion
178
+ if ctx:
179
+ await ctx.report_progress(100, 100)
180
+
181
+ # Track request completion for memory profiling
182
+ if memory_profiler:
183
+ memory_profiler.on_request_complete()
184
+
185
+ # Handle different result types for proper formatting
186
+ if isinstance(result, str):
187
+ return result
188
+ if isinstance(result, dict | list):
189
+ return json.dumps(result, default=str)
190
+ return str(result)
191
+ except Exception as e:
192
+ if ctx:
193
+ ctx.error("Error calling function %s: %s", function_name, str(e))
194
+
195
+ # Track request completion even on error
196
+ if memory_profiler:
197
+ memory_profiler.on_request_complete()
198
+
199
+ raise
200
+
201
+ return wrapper_with_ctx
202
+
203
+ # Create the wrapper function
204
+ wrapper = create_wrapper()
205
+
206
+ # Set the signature on the wrapper function (WITHOUT ctx)
207
+ wrapper.__signature__ = sig # type: ignore
208
+ wrapper.__name__ = function_name
209
+
210
+ # Return the wrapper with proper signature
211
+ return wrapper
212
+
213
+
214
+ def get_function_description(function: FunctionBase) -> str:
215
+ """
216
+ Retrieve a human-readable description for a NAT function or workflow.
217
+
218
+ The description is determined using the following precedence:
219
+ 1. If the function is a Workflow and has a 'description' attribute, use it.
220
+ 2. If the Workflow's config has a 'description', use it.
221
+ 3. If the Workflow's config has a 'topic', use it.
222
+ 4. If the function is a regular Function, use its 'description' attribute.
223
+
224
+ Args:
225
+ function: The NAT FunctionBase instance (Function or Workflow).
226
+
227
+ Returns:
228
+ The best available description string for the function.
229
+ """
230
+ function_description = ""
231
+
232
+ # Import here to avoid circular imports
233
+ from nat.builder.workflow import Workflow
234
+
235
+ if isinstance(function, Workflow):
236
+ config = function.config
237
+
238
+ # Workflow doesn't have a description, but probably should
239
+ if hasattr(function, "description") and function.description:
240
+ function_description = function.description
241
+ # Try to get description from config
242
+ elif hasattr(config, "description") and config.description:
243
+ function_description = config.description
244
+ # Try to get anything that might be a description
245
+ elif hasattr(config, "topic") and config.topic:
246
+ function_description = config.topic
247
+ # Try to get description from the workflow config
248
+ elif hasattr(config, "workflow") and hasattr(config.workflow, "description") and config.workflow.description:
249
+ function_description = config.workflow.description
250
+
251
+ elif isinstance(function, Function):
252
+ function_description = function.description
253
+
254
+ return function_description
255
+
256
+
257
+ def register_function_with_mcp(mcp: FastMCP,
258
+ function_name: str,
259
+ session_manager: 'SessionManager',
260
+ memory_profiler: 'MemoryProfiler | None' = None,
261
+ function: FunctionBase | None = None) -> None:
262
+ """Register a NAT Function as an MCP tool using SessionManager.
263
+
264
+ Each function is wrapped in a SessionManager
265
+ so that all calls go through Runner that automatically handles observability.
266
+
267
+ Args:
268
+ mcp: The FastMCP instance
269
+ function_name: The name to register the function under
270
+ session_manager: SessionManager wrapping the function/workflow
271
+ memory_profiler: Optional memory profiler to track requests
272
+ """
273
+ logger.info("Registering function %s with MCP", function_name)
274
+
275
+ # Get the workflow from the session manager
276
+ workflow = session_manager.workflow
277
+
278
+ # Prefer the function's schema/description when available, fall back to workflow
279
+ target_function = function or workflow
280
+
281
+ # Get the input schema from the most specific object available
282
+ input_schema = getattr(target_function, "input_schema", workflow.input_schema)
283
+ logger.info("Function %s has input schema: %s", function_name, input_schema)
284
+
285
+ # Get function description
286
+ function_description = get_function_description(target_function)
287
+
288
+ # Create and register the wrapper function with MCP
289
+ wrapper_func = create_function_wrapper(function_name, session_manager, input_schema, memory_profiler)
290
+ mcp.tool(name=function_name, description=function_description)(wrapper_func)
nat/plugins/mcp/utils.py CHANGED
@@ -1,4 +1,4 @@
1
- # SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
1
+ # SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
2
2
  # SPDX-License-Identifier: Apache-2.0
3
3
  #
4
4
  # Licensed under the Apache License, Version 2.0 (the "License");
@@ -1,6 +1,6 @@
1
1
  Metadata-Version: 2.4
2
2
  Name: nvidia-nat-mcp
3
- Version: 1.4.0a20251103
3
+ Version: 1.5.0a20260118
4
4
  Summary: Subpackage for MCP client integration in NeMo Agent toolkit
5
5
  Author: NVIDIA Corporation
6
6
  Maintainer: NVIDIA Corporation
@@ -16,13 +16,13 @@ Requires-Python: <3.14,>=3.11
16
16
  Description-Content-Type: text/markdown
17
17
  License-File: LICENSE-3rd-party.txt
18
18
  License-File: LICENSE.md
19
- Requires-Dist: nvidia-nat==v1.4.0a20251103
20
- Requires-Dist: aiorwlock~=1.5
21
- Requires-Dist: mcp~=1.14
19
+ Requires-Dist: nvidia-nat==v1.5.0a20260118
20
+ Requires-Dist: aiorwlock~=1.5.0
21
+ Requires-Dist: mcp~=1.25
22
22
  Dynamic: license-file
23
23
 
24
24
  <!--
25
- SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
25
+ SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
26
26
  SPDX-License-Identifier: Apache-2.0
27
27
 
28
28
  Licensed under the Apache License, Version 2.0 (the "License");
@@ -0,0 +1,37 @@
1
+ nat/meta/pypi.md,sha256=-P0JFDuIEDuCRryhXejNke80c_b6AeAzMNKMGgS-NT4,1489
2
+ nat/plugins/mcp/__init__.py,sha256=hiEJ2ajB5zhZuZOUHAyHfTA9w9UNhygij7hp4yfA6QE,685
3
+ nat/plugins/mcp/exception_handler.py,sha256=iEVSZFgiKm6Mr2XB3NRYbYBgyK7GJ6LnnPYXs8MYwJI,7653
4
+ nat/plugins/mcp/exceptions.py,sha256=87fm20cBNijam8n0dYx7drwCZqH1zr9tKls2ayfkK0A,5985
5
+ nat/plugins/mcp/register.py,sha256=-_SHPsqmVuFNlo6fwnXpZA9whIlMKifk4vnBfNX33mA,858
6
+ nat/plugins/mcp/utils.py,sha256=5foCyTdCp0p72FvlP_VjkQe6wgWLc9__XCfGk_iAUbY,9116
7
+ nat/plugins/mcp/auth/__init__.py,sha256=hiEJ2ajB5zhZuZOUHAyHfTA9w9UNhygij7hp4yfA6QE,685
8
+ nat/plugins/mcp/auth/auth_flow_handler.py,sha256=BvrclzC1mmVBsUZ2lx6Cv5NNFtMmUDxsTwLAKRrjiWU,9130
9
+ nat/plugins/mcp/auth/auth_provider.py,sha256=c-Y_FIDVUjEUixvqzeh7Sl0r5aiCR9JDyfudnfC6um8,21326
10
+ nat/plugins/mcp/auth/auth_provider_config.py,sha256=l0jbQZ_lYEK0B1DlkjkoKUcWIvW5pZas_rbE_WWqvq8,4226
11
+ nat/plugins/mcp/auth/register.py,sha256=FEp4L_ff4HQgFI4QDH04NIEIKRvh8bacns957dQlk9g,1759
12
+ nat/plugins/mcp/auth/token_storage.py,sha256=Zuj0tAWqkv8pPt4baeVEL2eEmnCE4AjU4h6eLhyKyKI,9263
13
+ nat/plugins/mcp/auth/service_account/__init__.py,sha256=hiEJ2ajB5zhZuZOUHAyHfTA9w9UNhygij7hp4yfA6QE,685
14
+ nat/plugins/mcp/auth/service_account/provider.py,sha256=Gsb44cy-5Ap_1BcG3Ny9fc94VtfPg4lVrwM3LKeTBA8,5920
15
+ nat/plugins/mcp/auth/service_account/provider_config.py,sha256=rrzHAo1qmK1yI5CNaadNTVsmh5qtYF1UrYz2Q1nLF2E,5332
16
+ nat/plugins/mcp/auth/service_account/token_client.py,sha256=bGAjq_vG8oo8ZoQfRJjEgigGmF0hgF0umaUOWv0_ZOY,5986
17
+ nat/plugins/mcp/cli/__init__.py,sha256=GLpBQw8r1Nj38vIOluSHsTlLINI24tZ_JJbcWWksBJE,709
18
+ nat/plugins/mcp/cli/commands.py,sha256=CFqLqaG2ldKqdbOAfNfUc1wyYRxuq2psoGX0mE_-2Go,48869
19
+ nat/plugins/mcp/client/__init__.py,sha256=SFdgClH-7zCAAcujmdvtfiqSz3Xe1KVRX7fviW-45Zg,714
20
+ nat/plugins/mcp/client/client_base.py,sha256=jWNwb20qWIUuIJg7Ldp_WgiJlQmYin4YvVSJ9_ECP5c,26831
21
+ nat/plugins/mcp/client/client_config.py,sha256=qrBjXCKHpMoS2A5i9DoXkN7hdYnGy77FBcZDGZlyb4w,7179
22
+ nat/plugins/mcp/client/client_impl.py,sha256=dygcLnhnWnMh7u45aLvYbGEP1uv02uSb886YsOucxgs,36661
23
+ nat/plugins/mcp/server/__init__.py,sha256=Ioza5tZX7qlOicklSVbhcXO3r7EjRLENvV5AtBSusAQ,723
24
+ nat/plugins/mcp/server/front_end_config.py,sha256=yFmEMKcz_ANJnef4Zu4dwGtnfNZ0uzVdc_dYYCLmgCg,5749
25
+ nat/plugins/mcp/server/front_end_plugin.py,sha256=WAqgYFHAswUK1WKk_PIUo6Lyf0pvJLelT3esHhmIHFM,6896
26
+ nat/plugins/mcp/server/front_end_plugin_worker.py,sha256=2_PQ938IbYKFKDsmDw6rcgBWHZdbxZetUopU4Ik1FDQ,18187
27
+ nat/plugins/mcp/server/introspection_token_verifier.py,sha256=oyhZkCTo6u3w7m469DicF57BNmhTPLcMYtWuHvcCsRg,2840
28
+ nat/plugins/mcp/server/memory_profiler.py,sha256=KjJw4ARSmEFult2MorFySiucog0bD8ClI9ZbkJhIwns,12816
29
+ nat/plugins/mcp/server/register_frontend.py,sha256=gfHfQaQ4NhkZkDSrFur1bxbG93kbAGBtumm4eKYOzKE,1178
30
+ nat/plugins/mcp/server/tool_converter.py,sha256=v35IILT5S2jE7ZTYQzm8-4CMu5Xa21BACeRP1MS1vzA,11608
31
+ nvidia_nat_mcp-1.5.0a20260118.dist-info/licenses/LICENSE-3rd-party.txt,sha256=fOk5jMmCX9YoKWyYzTtfgl-SUy477audFC5hNY4oP7Q,284609
32
+ nvidia_nat_mcp-1.5.0a20260118.dist-info/licenses/LICENSE.md,sha256=QwcOLU5TJoTeUhuIXzhdCEEDDvorGiC6-3YTOl4TecE,11356
33
+ nvidia_nat_mcp-1.5.0a20260118.dist-info/METADATA,sha256=Iep40v3cqEEYDqdEQrg4XzCSz0oltniRMDvtAdSTLLo,2326
34
+ nvidia_nat_mcp-1.5.0a20260118.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
35
+ nvidia_nat_mcp-1.5.0a20260118.dist-info/entry_points.txt,sha256=t5HYRfXR-dyZf6BWRz1u1TqPprKpCWt-WLHKQXfiLLU,231
36
+ nvidia_nat_mcp-1.5.0a20260118.dist-info/top_level.txt,sha256=8-CJ2cP6-f0ZReXe5Hzqp-5pvzzHz-5Ds5H2bGqh1-U,4
37
+ nvidia_nat_mcp-1.5.0a20260118.dist-info/RECORD,,
@@ -0,0 +1,9 @@
1
+ [nat.cli]
2
+ mcp = nat.plugins.mcp.cli.commands:mcp_command
3
+
4
+ [nat.components]
5
+ nat_mcp = nat.plugins.mcp.register
6
+ nat_mcp_auth = nat.plugins.mcp.auth.register
7
+
8
+ [nat.front_ends]
9
+ nat_mcp_server = nat.plugins.mcp.server.register_frontend
nat/plugins/mcp/tool.py DELETED
@@ -1,138 +0,0 @@
1
- # SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
2
- # SPDX-License-Identifier: Apache-2.0
3
- #
4
- # Licensed under the Apache License, Version 2.0 (the "License");
5
- # you may not use this file except in compliance with the License.
6
- # You may obtain a copy of the License at
7
- #
8
- # http://www.apache.org/licenses/LICENSE-2.0
9
- #
10
- # Unless required by applicable law or agreed to in writing, software
11
- # distributed under the License is distributed on an "AS IS" BASIS,
12
- # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13
- # See the License for the specific language governing permissions and
14
- # limitations under the License.
15
-
16
- import logging
17
- from typing import Literal
18
-
19
- from pydantic import BaseModel
20
- from pydantic import Field
21
- from pydantic import HttpUrl
22
- from pydantic import model_validator
23
-
24
- from nat.builder.builder import Builder
25
- from nat.builder.function_info import FunctionInfo
26
- from nat.cli.register_workflow import register_function
27
- from nat.data_models.function import FunctionBaseConfig
28
- from nat.plugins.mcp.client_base import MCPToolClient
29
- from nat.utils.decorators import deprecated
30
-
31
- logger = logging.getLogger(__name__)
32
-
33
-
34
- class MCPToolConfig(FunctionBaseConfig, name="mcp_tool_wrapper"):
35
- """
36
- Function which connects to a Model Context Protocol (MCP) server and wraps the selected tool as a NeMo Agent toolkit
37
- function.
38
- """
39
- # Add your custom configuration parameters here
40
- url: HttpUrl | None = Field(default=None,
41
- description="The URL of the MCP server (for streamable-http or sse modes)")
42
- mcp_tool_name: str = Field(description="The name of the tool served by the MCP Server that you want to use")
43
- transport: Literal["sse", "stdio", "streamable-http"] = Field(
44
- default="streamable-http",
45
- description="The type of transport to use (default: streamable-http, backwards compatible with sse)")
46
- command: str | None = Field(default=None,
47
- description="The command to run for stdio mode (e.g. 'docker' or 'python')")
48
- args: list[str] | None = Field(default=None, description="Additional arguments for the stdio command")
49
- env: dict[str, str] | None = Field(default=None, description="Environment variables to set for the stdio process")
50
- description: str | None = Field(default=None,
51
- description="""
52
- Description for the tool that will override the description provided by the MCP server. Should only be used if
53
- the description provided by the server is poor or nonexistent
54
- """)
55
-
56
- @model_validator(mode="after")
57
- def validate_model(self):
58
- """Validate that stdio and SSE/Streamable HTTP properties are mutually exclusive."""
59
- if self.transport == 'stdio':
60
- if self.url is not None:
61
- raise ValueError("url should not be set when using stdio client type")
62
- if not self.command:
63
- raise ValueError("command is required when using stdio client type")
64
- elif self.transport in ['streamable-http', 'sse']:
65
- if self.command is not None or self.args is not None or self.env is not None:
66
- raise ValueError(
67
- "command, args, and env should not be set when using streamable-http or sse client type")
68
- if not self.url:
69
- raise ValueError("url is required when using streamable-http or sse client type")
70
- return self
71
-
72
-
73
- def mcp_tool_function(tool: MCPToolClient) -> FunctionInfo:
74
- """
75
- Create a NeMo Agent toolkit function from an MCP tool.
76
-
77
- Args:
78
- tool: The MCP tool to wrap
79
-
80
- Returns:
81
- The NeMo Agent toolkit function
82
- """
83
-
84
- def _convert_from_str(input_str: str) -> tool.input_schema:
85
- return tool.input_schema.model_validate_json(input_str)
86
-
87
- async def _response_fn(tool_input: BaseModel | None = None, **kwargs) -> str:
88
- # Run the tool, catching any errors and sending to agent for correction
89
- try:
90
- if tool_input:
91
- args = tool_input.model_dump()
92
- return await tool.acall(args)
93
-
94
- _ = tool.input_schema.model_validate(kwargs)
95
- return await tool.acall(kwargs)
96
- except Exception as e:
97
- logger.warning("Error calling tool %s", tool.name, exc_info=True)
98
- return str(e)
99
-
100
- return FunctionInfo.create(single_fn=_response_fn,
101
- description=tool.description,
102
- input_schema=tool.input_schema,
103
- converters=[_convert_from_str])
104
-
105
-
106
- @register_function(config_type=MCPToolConfig)
107
- @deprecated(
108
- reason=
109
- "This function is being replaced with the new mcp_client function group that supports additional MCP features",
110
- feature_name="mcp_tool_wrapper")
111
- async def mcp_tool(config: MCPToolConfig, builder: Builder):
112
- """
113
- Generate a NeMo Agent Toolkit Function that wraps a tool provided by the MCP server.
114
- """
115
-
116
- from nat.plugins.mcp.client_base import MCPSSEClient
117
- from nat.plugins.mcp.client_base import MCPStdioClient
118
- from nat.plugins.mcp.client_base import MCPStreamableHTTPClient
119
-
120
- # Initialize the client
121
- if config.transport == 'stdio':
122
- client = MCPStdioClient(command=config.command, args=config.args, env=config.env)
123
- elif config.transport == 'streamable-http':
124
- client = MCPStreamableHTTPClient(url=str(config.url))
125
- elif config.transport == 'sse':
126
- client = MCPSSEClient(url=str(config.url))
127
- else:
128
- raise ValueError(f"Invalid transport type: {config.transport}")
129
-
130
- async with client:
131
- # If the tool is found create a MCPToolClient object and set the description if provided
132
- tool: MCPToolClient = await client.get_tool(config.mcp_tool_name)
133
- if config.description:
134
- tool.set_description(description=config.description)
135
-
136
- logger.info("Configured to use tool: %s from MCP server at %s", tool.name, client.server_name)
137
-
138
- yield mcp_tool_function(tool)
@@ -1,23 +0,0 @@
1
- nat/meta/pypi.md,sha256=EYyJTCCEOWzuuz-uNaYJ_WBk55Jiig87wcUr9E4g0yw,1484
2
- nat/plugins/mcp/__init__.py,sha256=GUJrgGtpvyMUCjUBvR3faAdv-tZzbU9W-izgx9aMEQg,680
3
- nat/plugins/mcp/client_base.py,sha256=JIyO2ZJsVkQ1g5BOU2zKXGHg_0yxv16g7_YJAqdCXTA,26504
4
- nat/plugins/mcp/client_config.py,sha256=l9tVUHe8WdFPJ9rXDg8dZkQi1dvHGYwoqQ8Glqg2LGs,6783
5
- nat/plugins/mcp/client_impl.py,sha256=j7cKAUBKtZAY3mt5Mm8VqgqMhRZk7kzvUd1nwMU_h0o,27072
6
- nat/plugins/mcp/exception_handler.py,sha256=4JVdZDJL4LyumZEcMIEBK2LYC6djuSMzqUhQDZZ6dUo,7648
7
- nat/plugins/mcp/exceptions.py,sha256=EGVOnYlui8xufm8dhJyPL1SUqBLnCGOTvRoeyNcmcWE,5980
8
- nat/plugins/mcp/register.py,sha256=HOT2Wl2isGuyFc7BUTi58-BbjI5-EtZMZo7stsv5pN4,831
9
- nat/plugins/mcp/tool.py,sha256=xNfBIF__ugJKFEjkYEM417wWM1PpuTaCMGtSFmxHSuA,6089
10
- nat/plugins/mcp/utils.py,sha256=dUIig7jeKz0ctb4o38jFGbe2uvM3DMR3PSJjfN_Lr5M,9111
11
- nat/plugins/mcp/auth/__init__.py,sha256=GUJrgGtpvyMUCjUBvR3faAdv-tZzbU9W-izgx9aMEQg,680
12
- nat/plugins/mcp/auth/auth_flow_handler.py,sha256=v21IK3IKZ2TLEP6wO9r-sJQiilWPq7Ry40M96SAxQFA,9125
13
- nat/plugins/mcp/auth/auth_provider.py,sha256=BgH66DlZgzhLDLO4cBERpHvNAmli5fMo_SCy11W9aBU,21251
14
- nat/plugins/mcp/auth/auth_provider_config.py,sha256=ZdiUObYU_Oj8KDDZ-JqkS6kJgup5EBy2ZREUOBw_kkI,4143
15
- nat/plugins/mcp/auth/register.py,sha256=L2x69NjJPS4s6CCE5myzWVrWn3e_ttHyojmGXvBipMg,1228
16
- nat/plugins/mcp/auth/token_storage.py,sha256=aS13ZvEJXcYzkZ0GSbrSor4i5bpjD5BkXHQw1iywC9k,9240
17
- nvidia_nat_mcp-1.4.0a20251103.dist-info/licenses/LICENSE-3rd-party.txt,sha256=fOk5jMmCX9YoKWyYzTtfgl-SUy477audFC5hNY4oP7Q,284609
18
- nvidia_nat_mcp-1.4.0a20251103.dist-info/licenses/LICENSE.md,sha256=QwcOLU5TJoTeUhuIXzhdCEEDDvorGiC6-3YTOl4TecE,11356
19
- nvidia_nat_mcp-1.4.0a20251103.dist-info/METADATA,sha256=fCZDLO4JXpWtFMO7BLRxedZUch6avusurmsdBG_Zz94,2319
20
- nvidia_nat_mcp-1.4.0a20251103.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
21
- nvidia_nat_mcp-1.4.0a20251103.dist-info/entry_points.txt,sha256=rYvUp4i-klBr3bVNh7zYOPXret704vTjvCk1qd7FooI,97
22
- nvidia_nat_mcp-1.4.0a20251103.dist-info/top_level.txt,sha256=8-CJ2cP6-f0ZReXe5Hzqp-5pvzzHz-5Ds5H2bGqh1-U,4
23
- nvidia_nat_mcp-1.4.0a20251103.dist-info/RECORD,,
@@ -1,3 +0,0 @@
1
- [nat.components]
2
- nat_mcp = nat.plugins.mcp.register
3
- nat_mcp_auth = nat.plugins.mcp.auth.register