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

1"""Sampling utilities for Golf MCP tools.

2

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

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

5"""

6

7 from typing import Any

8 from collections.abc import Callable

9

10 from .context import get_current_context

11

12# Apply telemetry instrumentation if available

13 try:

14 from golf.telemetry import instrument_sampling

15

16 _instrumentation_available = True

17 except 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

25 async 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

115 async 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

165 async 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

219 sample = instrument_sampling(sample, "sample")

220 sample_structured = instrument_sampling(sample_structured, "structured")

221 sample_with_context = instrument_sampling(sample_with_context, "context")