|
18 | 18 | from starlette.applications import Starlette
|
19 | 19 | from starlette.requests import Request
|
20 | 20 |
|
21 |
| -import mcp.types as types |
22 | 21 | from mcp.client.session import ClientSession
|
23 | 22 | from mcp.client.sse import sse_client
|
24 | 23 | from mcp.client.streamable_http import streamablehttp_client
|
25 | 24 | from mcp.server.fastmcp import FastMCP
|
26 | 25 | from mcp.server.fastmcp.resources import FunctionResource
|
27 | 26 | from mcp.shared.context import RequestContext
|
28 | 27 | from mcp.types import (
|
| 28 | + Completion, |
| 29 | + CompletionArgument, |
| 30 | + CompletionContext, |
29 | 31 | CreateMessageRequestParams,
|
30 | 32 | CreateMessageResult,
|
31 | 33 | GetPromptResult,
|
32 | 34 | InitializeResult,
|
| 35 | + LoggingMessageNotification, |
| 36 | + ProgressNotification, |
| 37 | + PromptReference, |
33 | 38 | ReadResourceResult,
|
| 39 | + ResourceListChangedNotification, |
| 40 | + ResourceTemplateReference, |
34 | 41 | SamplingMessage,
|
| 42 | + ServerNotification, |
35 | 43 | TextContent,
|
36 | 44 | TextResourceContents,
|
| 45 | + ToolListChangedNotification, |
37 | 46 | )
|
38 | 47 |
|
39 | 48 |
|
@@ -191,6 +200,40 @@ def complex_prompt(user_query: str, context: str = "general") -> str:
|
191 | 200 | # Since FastMCP doesn't support system messages in the same way
|
192 | 201 | return f"Context: {context}. Query: {user_query}"
|
193 | 202 |
|
| 203 | + # Resource template with completion support |
| 204 | + @mcp.resource("github://repos/{owner}/{repo}") |
| 205 | + def github_repo_resource(owner: str, repo: str) -> str: |
| 206 | + return f"Repository: {owner}/{repo}" |
| 207 | + |
| 208 | + # Add completion handler for the server |
| 209 | + @mcp.completion() |
| 210 | + async def handle_completion( |
| 211 | + ref: PromptReference | ResourceTemplateReference, |
| 212 | + argument: CompletionArgument, |
| 213 | + context: CompletionContext | None, |
| 214 | + ) -> Completion | None: |
| 215 | + # Handle GitHub repository completion |
| 216 | + if isinstance(ref, ResourceTemplateReference): |
| 217 | + if ref.uri == "github://repos/{owner}/{repo}" and argument.name == "repo": |
| 218 | + if context and context.arguments and context.arguments.get("owner") == "modelcontextprotocol": |
| 219 | + # Return repos for modelcontextprotocol org |
| 220 | + return Completion(values=["python-sdk", "typescript-sdk", "specification"], total=3, hasMore=False) |
| 221 | + elif context and context.arguments and context.arguments.get("owner") == "test-org": |
| 222 | + # Return repos for test-org |
| 223 | + return Completion(values=["test-repo1", "test-repo2"], total=2, hasMore=False) |
| 224 | + |
| 225 | + # Handle prompt completions |
| 226 | + if isinstance(ref, PromptReference): |
| 227 | + if ref.name == "complex_prompt" and argument.name == "context": |
| 228 | + # Complete context values |
| 229 | + contexts = ["general", "technical", "business", "academic"] |
| 230 | + return Completion( |
| 231 | + values=[c for c in contexts if c.startswith(argument.value)], total=None, hasMore=False |
| 232 | + ) |
| 233 | + |
| 234 | + # Default: no completion available |
| 235 | + return Completion(values=[], total=0, hasMore=False) |
| 236 | + |
194 | 237 | # Tool that echoes request headers from context
|
195 | 238 | @mcp.tool(description="Echo request headers from context")
|
196 | 239 | def echo_headers(ctx: Context[Any, Any, Request]) -> str:
|
@@ -597,15 +640,15 @@ async def handle_tool_list_changed(self, params) -> None:
|
597 | 640 |
|
598 | 641 | async def handle_generic_notification(self, message) -> None:
|
599 | 642 | # Check if this is a ServerNotification
|
600 |
| - if isinstance(message, types.ServerNotification): |
| 643 | + if isinstance(message, ServerNotification): |
601 | 644 | # Check the specific notification type
|
602 |
| - if isinstance(message.root, types.ProgressNotification): |
| 645 | + if isinstance(message.root, ProgressNotification): |
603 | 646 | await self.handle_progress(message.root.params)
|
604 |
| - elif isinstance(message.root, types.LoggingMessageNotification): |
| 647 | + elif isinstance(message.root, LoggingMessageNotification): |
605 | 648 | await self.handle_log(message.root.params)
|
606 |
| - elif isinstance(message.root, types.ResourceListChangedNotification): |
| 649 | + elif isinstance(message.root, ResourceListChangedNotification): |
607 | 650 | await self.handle_resource_list_changed(message.root.params)
|
608 |
| - elif isinstance(message.root, types.ToolListChangedNotification): |
| 651 | + elif isinstance(message.root, ToolListChangedNotification): |
609 | 652 | await self.handle_tool_list_changed(message.root.params)
|
610 | 653 |
|
611 | 654 |
|
@@ -781,6 +824,41 @@ async def progress_callback(progress: float, total: float | None, message: str |
|
781 | 824 | if context_data["method"]:
|
782 | 825 | assert context_data["method"] == "POST"
|
783 | 826 |
|
| 827 | + # Test completion functionality |
| 828 | + # 1. Test resource template completion with context |
| 829 | + repo_result = await session.complete( |
| 830 | + ref=ResourceTemplateReference(type="ref/resource", uri="github://repos/{owner}/{repo}"), |
| 831 | + argument={"name": "repo", "value": ""}, |
| 832 | + context_arguments={"owner": "modelcontextprotocol"}, |
| 833 | + ) |
| 834 | + assert repo_result.completion.values == ["python-sdk", "typescript-sdk", "specification"] |
| 835 | + assert repo_result.completion.total == 3 |
| 836 | + assert repo_result.completion.hasMore is False |
| 837 | + |
| 838 | + # 2. Test with different context |
| 839 | + repo_result2 = await session.complete( |
| 840 | + ref=ResourceTemplateReference(type="ref/resource", uri="github://repos/{owner}/{repo}"), |
| 841 | + argument={"name": "repo", "value": ""}, |
| 842 | + context_arguments={"owner": "test-org"}, |
| 843 | + ) |
| 844 | + assert repo_result2.completion.values == ["test-repo1", "test-repo2"] |
| 845 | + assert repo_result2.completion.total == 2 |
| 846 | + |
| 847 | + # 3. Test prompt argument completion |
| 848 | + context_result = await session.complete( |
| 849 | + ref=PromptReference(type="ref/prompt", name="complex_prompt"), |
| 850 | + argument={"name": "context", "value": "tech"}, |
| 851 | + ) |
| 852 | + assert "technical" in context_result.completion.values |
| 853 | + |
| 854 | + # 4. Test completion without context (should return empty) |
| 855 | + no_context_result = await session.complete( |
| 856 | + ref=ResourceTemplateReference(type="ref/resource", uri="github://repos/{owner}/{repo}"), |
| 857 | + argument={"name": "repo", "value": "test"}, |
| 858 | + ) |
| 859 | + assert no_context_result.completion.values == [] |
| 860 | + assert no_context_result.completion.total == 0 |
| 861 | + |
784 | 862 |
|
785 | 863 | async def sampling_callback(
|
786 | 864 | context: RequestContext[ClientSession, None],
|
|
0 commit comments