Coverage for /home/admin/Documents/AI/applications/lexigram-dev/lexigram/experimental/ai/lexigram-ai-mcp/src/lexigram/ai/mcp/server/handlers/sampling.py: 56%

1"""MCP sampling handler — server-side LLM inference capability. 

2 

3When an MCP client sends a ``sampling/createMessage`` request, the server 

4uses its configured ``LLMClientProtocol`` to generate a completion and returns 

5it in the MCP sampling response format. 

6 

7This implements the MCP Sampling specification: 

8https://spec.modelcontextprotocol.io/specification/client/sampling/ 

9""" 

10 

11from __future__ import annotations 

12 

13from dataclasses import dataclass, field 

14from typing import TYPE_CHECKING, Any 

15 

16from lexigram.contracts.mcp.exceptions import MCPError, MCPToolCallError 

17from lexigram.logging import ( 

18 get_logger, 

19) 

20from lexigram.result import Err, Ok, Result 

21 

22if TYPE_CHECKING: 

23 from lexigram.contracts.ai.llm import LLMClientProtocol 

24 

25logger = get_logger(__name__) 

26 

27 

28@dataclass 

29class SamplingRequest: 

30 """MCP sampling/createMessage request payload. 

31 

32 Maps directly to the MCP sampling request schema. 

33 """ 

34 

35 messages: list[dict[str, Any]] 

36 """Conversation messages in MCP content format.""" 

37 

38 model_preferences: dict[str, Any] = field(default_factory=dict) 

39 """Client model hints (hints, cost_priority, speed_priority, intelligence_priority).""" 

40 

41 system_prompt: str | None = None 

42 """Optional system prompt to prepend.""" 

43 

44 max_tokens: int = 1024 

45 """Maximum tokens for the completion.""" 

46 

47 temperature: float | None = None 

48 """Sampling temperature (0.0-2.0). None uses the model default.""" 

49 

50 stop_sequences: list[str] = field(default_factory=list) 

51 """Optional stop sequences.""" 

52 

53 metadata: dict[str, Any] = field(default_factory=dict) 

54 """Pass-through metadata (ignored by the server, forwarded as-is).""" 

55 

56 

57@dataclass 

58class SamplingResponse: 

59 """MCP sampling/createMessage response payload.""" 

60 

61 role: str 

62 """Always 'assistant' for server-generated completions.""" 

63 

64 content: dict[str, Any] 

65 """MCP content item — typically ``{"type": "text", "text": "..."}``.""" 

66 

67 model: str 

68 """The model identifier used to generate the response.""" 

69 

70 stop_reason: str | None = None 

71 """Why generation stopped: 'end_turn', 'max_tokens', 'stop_sequence', etc.""" 

72 

73 def to_dict(self) -> dict[str, Any]: 

74 """Serialise to MCP sampling response format.""" 

75 data: dict[str, Any] = { 

76 "role": self.role, 

77 "content": self.content, 

78 "model": self.model, 

79 } 

80 if self.stop_reason is not None: 

81 data["stopReason"] = self.stop_reason 

82 return data 

83 

84 

85class SamplingHandler: 

86 """MCP sampling capability — server-side LLM inference. 

87 

88 When an MCP client sends a ``sampling/createMessage`` request, this 

89 handler converts it to an ``LLMClientProtocol.complete()`` call and 

90 returns the result. 

91 

92 The handler is optional — registered only when an LLM client is 

93 available in the container (see ``MCPProvider._boot_handlers``). 

94 

95 Example:: 

96 

97 handler = SamplingHandler(llm=openai_client) 

98 result = await handler.handle(request) 

99 if result.is_ok(): 

100 response = result.unwrap() 

101 """ 

102 

103 def __init__(self, llm: LLMClientProtocol) -> None: 

104 """Initialize the sampling handler. 

105 

106 Args: 

107 llm: LLM client used to generate completions. 

108 """ 

109 self._llm = llm 

110 

111 async def create_message( 

112 self, 

113 messages: list[dict[str, Any]] | None = None, 

114 maxTokens: int = 1024, 

115 modelPreferences: dict[str, Any] | None = None, 

116 systemPrompt: str | None = None, 

117 temperature: float | None = None, 

118 stopSequences: list[str] | None = None, 

119 metadata: dict[str, Any] | None = None, 

120 **_kwargs: Any, 

121 ) -> Result[dict[str, Any], MCPError]: 

122 """Handle the ``sampling/createMessage`` MCP method. 

123 

124 Accepts MCP-spec camelCase kwargs directly so the MCPServer dispatcher 

125 can call ``handler(**params)`` without transformation. 

126 

127 Args: 

128 messages: MCP conversation messages. 

129 maxTokens: Maximum tokens to generate. 

130 modelPreferences: Client model hints (ignored, best-effort). 

131 systemPrompt: Optional system prompt to prepend. 

132 temperature: Sampling temperature. 

133 stopSequences: Stop sequences. 

134 metadata: Pass-through metadata. 

135 **_kwargs: Ignored extra params for forward-compatibility. 

136 

137 Returns: 

138 ``Result`` containing the MCP sampling response payload. 

139 """ 

140 request = SamplingRequest( 

141 messages=messages or [], 

142 model_preferences=modelPreferences or {}, 

143 system_prompt=systemPrompt, 

144 max_tokens=maxTokens, 

145 temperature=temperature, 

146 stop_sequences=stopSequences or [], 

147 metadata=metadata or {}, 

148 ) 

149 result = await self._handle(request) 

150 if result.is_ok(): 

151 return Ok(result.unwrap().to_dict()) 

152 return Err(result.unwrap_err()) 

153 

154 async def _handle( 

155 self, request: SamplingRequest 

156 ) -> Result[SamplingResponse, MCPError]: 

157 """Convert MCP sampling request to LLM call and return a response. 

158 

159 Args: 

160 request: Parsed sampling request. 

161 

162 Returns: 

163 ``Ok(SamplingResponse)`` or ``Err(MCPError)``. 

164 """ 

165 from lexigram.contracts.ai.llm import ChatMessage 

166 

167 chat_messages: list[Any] = [] 

168 

169 # Convert MCP message format to LLM chat messages 

170 for msg in request.messages: 

171 role = msg.get("role", "user") 

172 content = msg.get("content", {}) 

173 if isinstance(content, dict): 

174 text = content.get("text", "") 

175 else: 

176 text = str(content) 

177 chat_messages.append(ChatMessage(role=role, content=text)) 

178 

179 kwargs: dict[str, Any] = {} 

180 if request.temperature is not None: 

181 kwargs["temperature"] = request.temperature 

182 

183 completion_result = await self._llm.complete( 

184 chat_messages, 

185 max_tokens=request.max_tokens, 

186 **kwargs, 

187 ) 

188 

189 if completion_result.is_err(): 

190 llm_error = completion_result.unwrap_err() 

191 return Err( 

192 MCPToolCallError( 

193 message=f"LLM sampling failed: {llm_error}", 

194 ) 

195 ) 

196 

197 completion = completion_result.unwrap() 

198 text = getattr(completion, "text", "") or "" 

199 model = getattr(completion, "model", "unknown") or "unknown" 

200 stop_reason = getattr(completion, "finish_reason", None) 

201 

202 return Ok( 

203 SamplingResponse( 

204 role="assistant", 

205 content={"type": "text", "text": text}, 

206 model=model, 

207 stop_reason=stop_reason, 

208 ) 

209 ) 

210 

211 

212__all__ = ["SamplingHandler", "SamplingRequest", "SamplingResponse"]