Coverage for /Users/antonigmitruk/golf/src/golf/utilities/sampling.py: 0%

1"""Sampling utilities for Golf MCP tools. 

2 

3This module provides simplified LLM sampling functions that Golf tool authors 

4can use without needing to manage FastMCP Context objects directly. 

5""" 

6 

7from typing import Any 

8from collections.abc import Callable 

9 

10from .context import get_current_context 

11 

12# Apply telemetry instrumentation if available 

13try: 

14 from golf.telemetry import instrument_sampling 

15 

16 _instrumentation_available = True 

17except ImportError: 

18 _instrumentation_available = False 

19 

20 def instrument_sampling(func: Callable, sampling_type: str = "sample") -> Callable: 

21 """No-op instrumentation when telemetry is not available.""" 

22 return func 

23 

24 

25async def sample( 

26 messages: str | list[str], 

27 system_prompt: str | None = None, 

28 temperature: float | None = None, 

29 max_tokens: int | None = None, 

30 model_preferences: str | list[str] | None = None, 

31) -> str: 

32 """Request an LLM completion from the MCP client. 

33 

34 This is a simplified wrapper around FastMCP's Context.sample() method 

35 that automatically handles context retrieval and response processing. 

36 

37 Args: 

38 messages: The message(s) to send to the LLM: 

39 - str: Single user message 

40 - list[str]: Multiple user messages 

41 system_prompt: Optional system prompt to guide the LLM 

42 temperature: Optional temperature for sampling (0.0 to 1.0) 

43 max_tokens: Optional maximum tokens to generate (default: 512) 

44 model_preferences: Optional model preferences: 

45 - str: Single model name hint 

46 - list[str]: Multiple model name hints in preference order 

47 

48 Returns: 

49 The LLM's response as a string 

50 

51 Raises: 

52 RuntimeError: If called outside MCP context or sampling fails 

53 ValueError: If parameters are invalid 

54 

55 Examples: 

56 ```python 

57 from golf.utilities import sample 

58 

59 async def analyze_data(data: str): 

60 # Simple completion 

61 analysis = await sample(f"Analyze this data: {data}") 

62 

63 # With system prompt and temperature 

64 creative_response = await sample( 

65 "Write a creative story about this data", 

66 system_prompt="You are a creative writer", 

67 temperature=0.8, 

68 max_tokens=1000 

69 ) 

70 

71 # With model preferences 

72 technical_analysis = await sample( 

73 f"Provide technical analysis: {data}", 

74 model_preferences=["gpt-4", "claude-3-sonnet"] 

75 ) 

76 

77 return { 

78 "analysis": analysis, 

79 "creative": creative_response, 

80 "technical": technical_analysis 

81 } 

82 ``` 

83 """ 

84 try: 

85 # Get the current FastMCP context 

86 ctx = get_current_context() 

87 

88 # Call the context's sample method 

89 result = await ctx.sample( 

90 messages=messages, 

91 system_prompt=system_prompt, 

92 temperature=temperature, 

93 max_tokens=max_tokens, 

94 model_preferences=model_preferences, 

95 ) 

96 

97 # Extract text content from the ContentBlock response 

98 if hasattr(result, "text"): 

99 return result.text 

100 elif hasattr(result, "content"): 

101 # Handle different content block types 

102 if isinstance(result.content, str): 

103 return result.content 

104 elif hasattr(result.content, "text"): 

105 return result.content.text 

106 else: 

107 return str(result.content) 

108 else: 

109 return str(result) 

110 

111 except Exception as e: 

112 raise RuntimeError(f"LLM sampling failed: {str(e)}") from e 

113 

114 

115async def sample_structured( 

116 messages: str | list[str], 

117 format_instructions: str, 

118 system_prompt: str | None = None, 

119 temperature: float = 0.1, 

120 max_tokens: int | None = None, 

121) -> str: 

122 """Request a structured LLM completion with specific formatting. 

123 

124 This is a convenience function for requesting structured responses 

125 like JSON, XML, or other formatted output. 

126 

127 Args: 

128 messages: The message(s) to send to the LLM 

129 format_instructions: Instructions for the desired output format 

130 system_prompt: Optional system prompt 

131 temperature: Temperature for sampling (default: 0.1 for consistency) 

132 max_tokens: Optional maximum tokens to generate 

133 

134 Returns: 

135 The structured LLM response as a string 

136 

137 Example: 

138 ```python 

139 from golf.utilities import sample_structured 

140 

141 async def extract_entities(text: str): 

142 entities = await sample_structured( 

143 f"Extract entities from: {text}", 

144 format_instructions="Return as JSON with keys: persons, " 

145 "organizations, locations", 

146 system_prompt="You are an expert at named entity recognition" 

147 ) 

148 return entities 

149 ``` 

150 """ 

151 # Combine the format instructions with the messages 

152 if isinstance(messages, str): 

153 formatted_message = f"{messages}\n\n{format_instructions}" 

154 else: 

155 formatted_message = messages + [format_instructions] 

156 

157 return await sample( 

158 messages=formatted_message, 

159 system_prompt=system_prompt, 

160 temperature=temperature, 

161 max_tokens=max_tokens, 

162 ) 

163 

164 

165async def sample_with_context( 

166 messages: str | list[str], 

167 context_data: dict[str, Any], 

168 system_prompt: str | None = None, 

169 **kwargs: Any, 

170) -> str: 

171 """Request an LLM completion with additional context data. 

172 

173 This convenience function formats context data and includes it 

174 in the sampling request. 

175 

176 Args: 

177 messages: The message(s) to send to the LLM 

178 context_data: Dictionary of context data to include 

179 system_prompt: Optional system prompt 

180 **kwargs: Additional arguments passed to sample() 

181 

182 Returns: 

183 The LLM response as a string 

184 

185 Example: 

186 ```python 

187 from golf.utilities import sample_with_context 

188 

189 async def generate_report(topic: str, user_data: dict): 

190 report = await sample_with_context( 

191 f"Generate a report about {topic}", 

192 context_data={ 

193 "user_preferences": user_data, 

194 "timestamp": "2024-01-01", 

195 "format": "markdown" 

196 }, 

197 system_prompt="You are a professional report writer" 

198 ) 

199 return report 

200 ``` 

201 """ 

202 # Format context data as a readable string 

203 context_str = "\n".join([f"{k}: {v}" for k, v in context_data.items()]) 

204 

205 # Add context to the message 

206 if isinstance(messages, str): 

207 contextual_message = f"{messages}\n\nContext:\n{context_str}" 

208 else: 

209 contextual_message = messages + [f"Context:\n{context_str}"] 

210 

211 return await sample( 

212 messages=contextual_message, 

213 system_prompt=system_prompt, 

214 **kwargs, 

215 ) 

216 

217 

218# Apply instrumentation to all sampling functions 

219sample = instrument_sampling(sample, "sample") 

220sample_structured = instrument_sampling(sample_structured, "structured") 

221sample_with_context = instrument_sampling(sample_with_context, "context")