1  #!/usr/bin/env python3
  2  """
  3  MCP Sampling Example
  4  
  5  Demonstrates the Sampling API per MCP 2025-11-25 specification.
  6  Sampling allows servers to request LLM completions from clients.
  7  
  8  Features:
  9  - Tool calling support with toolChoice modes
 10  - Model preferences
 11  - System prompts
 12  
 13  Usage:
 14      python mcp_sampling_example.py
 15  """
 16  
 17  import asyncio
 18  from praisonai.mcp_server.sampling import (
 19      SamplingHandler,
 20      SamplingRequest,
 21      SamplingResponse,
 22      SamplingMessage,
 23      ToolChoice,
 24      ToolDefinition,
 25      ModelPreferences,
 26      create_sampling_request,
 27  )
 28  
 29  
 30  async def mock_llm_callback(request: SamplingRequest) -> SamplingResponse:
 31      """Mock LLM callback for demonstration."""
 32      # In real usage, this would call an actual LLM
 33      messages_text = " | ".join([f"{m.role}: {m.content}" for m in request.messages])
 34      
 35      if request.tools:
 36          # Simulate tool use response
 37          return SamplingResponse(
 38              role="assistant",
 39              content="",
 40              model="mock-model",
 41              stop_reason="toolUse",
 42              tool_calls=[{
 43                  "id": "call_123",
 44                  "name": request.tools[0].name,
 45                  "arguments": {"query": "test"},
 46              }],
 47          )
 48      
 49      return SamplingResponse(
 50          role="assistant",
 51          content=f"Mock response to: {messages_text}",
 52          model="mock-model",
 53          stop_reason="end_turn",
 54      )
 55  
 56  
 57  async def main():
 58      print("=" * 60)
 59      print("MCP Sampling Example (2025-11-25 Specification)")
 60      print("=" * 60)
 61      
 62      # Create handler with mock callback
 63      handler = SamplingHandler(default_model="gpt-4o-mini")
 64      handler.set_callback(mock_llm_callback)
 65      
 66      # 1. Basic sampling request
 67      print("\n1. Basic Sampling Request")
 68      print("-" * 40)
 69      
 70      basic_request = create_sampling_request(
 71          prompt="What is the capital of France?",
 72          system_prompt="You are a helpful geography assistant.",
 73          max_tokens=100,
 74      )
 75      
 76      print(f"   Messages: {[m.content for m in basic_request.messages]}")
 77      print(f"   System prompt: {basic_request.system_prompt}")
 78      print(f"   Max tokens: {basic_request.max_tokens}")
 79      
 80      response = await handler.create_message(basic_request)
 81      print(f"\n   Response: {response.content}")
 82      print(f"   Model: {response.model}")
 83      print(f"   Stop reason: {response.stop_reason}")
 84      
 85      # 2. Sampling with tools
 86      print("\n2. Sampling with Tools")
 87      print("-" * 40)
 88      
 89      tool_request = create_sampling_request(
 90          prompt="Search for the latest AI news",
 91          tools=[
 92              {
 93                  "name": "web_search",
 94                  "description": "Search the web for information",
 95                  "inputSchema": {
 96                      "type": "object",
 97                      "properties": {
 98                          "query": {"type": "string", "description": "Search query"},
 99                      },
100                      "required": ["query"],
101                  },
102              }
103          ],
104          tool_choice="auto",
105      )
106      
107      print(f"   Tools: {[t.name for t in tool_request.tools]}")
108      print(f"   Tool choice: {tool_request.tool_choice.to_dict()}")
109      
110      response = await handler.create_message(tool_request)
111      print(f"\n   Response stop reason: {response.stop_reason}")
112      if response.tool_calls:
113          print(f"   Tool calls: {response.tool_calls}")
114      
115      # 3. ToolChoice modes
116      print("\n3. ToolChoice Modes (MCP 2025-11-25)")
117      print("-" * 40)
118      
119      modes = [
120          ("auto", ToolChoice.auto()),
121          ("none", ToolChoice.none()),
122          ("any", ToolChoice.any()),
123          ("tool (specific)", ToolChoice.tool("web_search")),
124      ]
125      
126      for name, tc in modes:
127          print(f"   {name}: {tc.to_dict()}")
128      
129      # 4. Model preferences
130      print("\n4. Model Preferences")
131      print("-" * 40)
132      
133      prefs = ModelPreferences(
134          hints=[{"name": "claude-3-sonnet"}, {"name": "gpt-4"}],
135          cost_priority=0.3,
136          speed_priority=0.5,
137          intelligence_priority=0.8,
138      )
139      
140      print(f"   Preferences: {prefs.to_dict()}")
141      
142      # 5. Full request with all options
143      print("\n5. Full Request with All Options")
144      print("-" * 40)
145      
146      full_request = SamplingRequest(
147          messages=[
148              SamplingMessage(role="user", content="Hello!"),
149              SamplingMessage(role="assistant", content="Hi there!"),
150              SamplingMessage(role="user", content="What can you do?"),
151          ],
152          system_prompt="You are a helpful assistant.",
153          model_preferences=prefs,
154          max_tokens=500,
155          temperature=0.7,
156          tools=[
157              ToolDefinition(
158                  name="calculator",
159                  description="Perform calculations",
160                  input_schema={"type": "object", "properties": {"expression": {"type": "string"}}},
161              )
162          ],
163          tool_choice=ToolChoice.auto(),
164      )
165      
166      print(f"   Message count: {len(full_request.messages)}")
167      print(f"   Has tools: {bool(full_request.tools)}")
168      print(f"   Temperature: {full_request.temperature}")
169      
170      request_dict = full_request.to_dict()
171      print(f"\n   MCP Request Format (keys): {list(request_dict.keys())}")
172      
173      print("\n" + "=" * 60)
174      print("Sampling Example Complete!")
175      print("=" * 60)
176  
177  
178  if __name__ == "__main__":
179      asyncio.run(main())