Coverage for src/dataknobs_llm/prompts/versioning/ab_testing.py: 83%

1"""A/B testing management for prompt experiments.

2

3 This module provides:

4- Experiment creation and management

5- Random variant selection

6- User-sticky variant assignment

7- Traffic split management

8"""

9

10 import uuid

11 import hashlib

12 import random

13 from typing import Any, Dict, List, Optional

14 from datetime import datetime

15

16 from .types import (

17 PromptExperiment,

18 PromptVariant,

19 VersioningError,

20)

21

22

23 class ABTestManager:

24 """Manages A/B test experiments for prompts.

25

26 Supports multiple selection strategies:

27 - Random: Each request gets a random variant based on traffic split

28 - User-sticky: Same user always gets same variant (consistent experience)

29

30 Example:

31 ```python

32 manager = ABTestManager(storage_backend)

33

34 # Create experiment

35 experiment = await manager.create_experiment(

36 name="greeting",

37 prompt_type="system",

38 variants=[

39 PromptVariant("1.0.0", 0.5, "Control"),

40 PromptVariant("1.1.0", 0.5, "Treatment")

41 ]

42 )

43

44 # Get variant for user (sticky assignment)

45 variant_version = await manager.get_variant_for_user(

46 experiment.experiment_id,

47 user_id="user123"

48 )

49

50 # Get random variant

51 variant_version = await manager.get_random_variant(

52 experiment.experiment_id

53 )

54 ```

55 """

56

57 def __init__(self, storage: Optional[Any] = None):

58 """Initialize A/B test manager.

59

60 Args:

61 storage: Backend storage (dict for in-memory, database for persistence)

62 If None, uses in-memory dictionary

63 """

64 self.storage = storage if storage is not None else {}

65 self._experiments: Dict[str, PromptExperiment] = {} # experiment_id -> PromptExperiment

66 self._user_assignments: Dict[str, Dict[str, str]] = {} # experiment_id -> {user_id -> version}

67

68 async def create_experiment(

69 self,

70 name: str,

71 prompt_type: str,

72 variants: List[PromptVariant],

73 traffic_split: Optional[Dict[str, float]] = None,

74 metadata: Optional[Dict[str, Any]] = None,

75 ) -> PromptExperiment:

76 """Create a new A/B test experiment.

77

78 Args:

79 name: Prompt name

80 prompt_type: Prompt type

81 variants: List of variants to test

82 traffic_split: Optional custom traffic split (if None, derives from variant weights)

83 metadata: Additional metadata

84

85 Returns:

86 Created PromptExperiment

87

88 Raises:

89 VersioningError: If variants are invalid or traffic split doesn't sum to 1.0

90 """

91 if len(variants) < 2:

92 raise VersioningError("Experiment must have at least 2 variants")

93

94 # Generate experiment ID

95 experiment_id = str(uuid.uuid4())

96

97 # Derive traffic split from variant weights if not provided

98 if traffic_split is None:

99 # Normalize weights to ensure they sum to 1.0

100 total_weight = sum(v.weight for v in variants)

101 traffic_split = {

102 v.version: v.weight / total_weight

103 for v in variants

104 }

105

106 # Create experiment

107 experiment = PromptExperiment(

108 experiment_id=experiment_id,

109 name=name,

110 prompt_type=prompt_type,

111 variants=variants,

112 traffic_split=traffic_split,

113 start_date=datetime.utcnow(),

114 status="running",

115 metadata=metadata or {},

116 )

117

118 # Store experiment

119 self._experiments[experiment_id] = experiment

120

121 # Initialize user assignments

122 self._user_assignments[experiment_id] = {}

123

124 # Persist to backend if available

125 if hasattr(self.storage, "set"):

126 await self._persist_experiment(experiment)

127

128 return experiment

129

130 async def get_experiment(

131 self,

132 experiment_id: str,

133 ) -> Optional[PromptExperiment]:

134 """Retrieve an experiment by ID.

135

136 Args:

137 experiment_id: Experiment ID

138

139 Returns:

140 PromptExperiment if found, None otherwise

141 """

142 return self._experiments.get(experiment_id)

143

144 async def list_experiments(

145 self,

146 name: Optional[str] = None,

147 prompt_type: Optional[str] = None,

148 status: Optional[str] = None,

149 ) -> List[PromptExperiment]:

150 """List experiments with optional filters.

151

152 Args:

153 name: Filter by prompt name

154 prompt_type: Filter by prompt type

155 status: Filter by status ("running", "paused", "completed")

156

157 Returns:

158 List of matching experiments

159 """

160 experiments = list(self._experiments.values())

161

162 # Apply filters

163 if name:

164 experiments = [e for e in experiments if e.name == name]

165

166 if prompt_type:

167 experiments = [e for e in experiments if e.prompt_type == prompt_type]

168

169 if status:

170 experiments = [e for e in experiments if e.status == status]

171

172 return experiments

173

174 async def get_random_variant(

175 self,

176 experiment_id: str,

177 ) -> str:

178 """Get a random variant based on traffic split.

179

180 Each call returns a potentially different variant.

181

182 Args:

183 experiment_id: Experiment ID

184

185 Returns:

186 Version string of selected variant

187

188 Raises:

189 VersioningError: If experiment not found

190 """

191 experiment = self._experiments.get(experiment_id)

192 if not experiment:

193 raise VersioningError(f"Experiment not found: {experiment_id}")

194

195 if experiment.status != "running":

196 raise VersioningError(

197 f"Experiment {experiment_id} is not running (status: {experiment.status})"

198 )

199

200 # Weighted random selection

201 versions = list(experiment.traffic_split.keys())

202 weights = list(experiment.traffic_split.values())

203

204 return random.choices(versions, weights=weights)[0]

205

206 async def get_variant_for_user(

207 self,

208 experiment_id: str,

209 user_id: str,

210 ) -> str:

211 """Get variant for a specific user (sticky assignment).

212

213 The same user always gets the same variant for consistent experience.

214 Uses hash-based assignment to ensure deterministic selection.

215

216 Args:

217 experiment_id: Experiment ID

218 user_id: User identifier

219

220 Returns:

221 Version string of assigned variant

222

223 Raises:

224 VersioningError: If experiment not found

225 """

226 experiment = self._experiments.get(experiment_id)

227 if not experiment:

228 raise VersioningError(f"Experiment not found: {experiment_id}")

229

230 if experiment.status != "running":

231 raise VersioningError(

232 f"Experiment {experiment_id} is not running (status: {experiment.status})"

233 )

234

235 # Check if user already has assignment

236 if experiment_id in self._user_assignments:

237 existing = self._user_assignments[experiment_id].get(user_id)

238 if existing:

239 return existing

240

241 # Assign user to variant using hash-based selection

242 assigned_version = self._hash_based_assignment(

243 user_id,

244 experiment.traffic_split

245 )

246

247 # Store assignment

248 if experiment_id not in self._user_assignments:

249 self._user_assignments[experiment_id] = {}

250 self._user_assignments[experiment_id][user_id] = assigned_version

251

252 # Persist assignment if backend available

253 if hasattr(self.storage, "set"):

254 key = f"assignment:{experiment_id}:{user_id}"

255 await self.storage.set(key, assigned_version)

256

257 return assigned_version

258

259 async def update_experiment_status(

260 self,

261 experiment_id: str,

262 status: str,

263 end_date: Optional[datetime] = None,

264 ) -> PromptExperiment:

265 """Update experiment status.

266

267 Args:

268 experiment_id: Experiment ID

269 status: New status ("running", "paused", "completed")

270 end_date: Optional end date (auto-set to now if status is "completed")

271

272 Returns:

273 Updated experiment

274

275 Raises:

276 VersioningError: If experiment not found

277 """

278 experiment = self._experiments.get(experiment_id)

279 if not experiment:

280 raise VersioningError(f"Experiment not found: {experiment_id}")

281

282 experiment.status = status

283

284 if status == "completed" and end_date is None:

285 experiment.end_date = datetime.utcnow()

286 elif end_date:

287 experiment.end_date = end_date

288

289 # Persist if backend available

290 if hasattr(self.storage, "set"):

291 await self._persist_experiment(experiment)

292

293 return experiment

294

295 async def get_user_assignment(

296 self,

297 experiment_id: str,

298 user_id: str,

299 ) -> Optional[str]:

300 """Get existing user assignment without creating a new one.

301

302 Args:

303 experiment_id: Experiment ID

304 user_id: User ID

305

306 Returns:

307 Assigned version if exists, None otherwise

308 """

309 if experiment_id not in self._user_assignments:

310 return None

311 return self._user_assignments[experiment_id].get(user_id)

312

313 async def get_experiment_assignments(

314 self,

315 experiment_id: str,

316 ) -> Dict[str, str]:

317 """Get all user assignments for an experiment.

318

319 Args:

320 experiment_id: Experiment ID

321

322 Returns:

323 Dictionary mapping user_id to assigned version

324 """

325 return self._user_assignments.get(experiment_id, {})

326

327 async def delete_experiment(

328 self,

329 experiment_id: str,

330 ) -> bool:

331 """Delete an experiment.

332

333 Note: This also removes all user assignments.

334

335 Args:

336 experiment_id: Experiment ID

337

338 Returns:

339 True if deleted, False if not found

340 """

341 if experiment_id not in self._experiments:

342 return False

343

344 # Remove experiment

345 del self._experiments[experiment_id]

346

347 # Remove user assignments

348 if experiment_id in self._user_assignments:

349 del self._user_assignments[experiment_id]

350

351 # Persist deletion if backend available

352 if hasattr(self.storage, "delete"):

353 await self.storage.delete(f"experiment:{experiment_id}")

354

355 return True

356

357 # ===== Helper Methods =====

358

359 def _hash_based_assignment(

360 self,

361 user_id: str,

362 traffic_split: Dict[str, float],

363 ) -> str:

364 """Assign user to variant using consistent hash-based selection.

365

366 This ensures the same user always gets the same variant.

367

368 Args:

369 user_id: User identifier

370 traffic_split: Version to percentage mapping

371

372 Returns:

373 Selected version string

374 """

375 # Hash user_id to get a deterministic number

376 hash_hex = hashlib.md5(user_id.encode()).hexdigest()

377 hash_val = int(hash_hex, 16)

378

379 # Map to [0, 1) range

380 normalized = (hash_val % 1000) / 1000.0

381

382 # Select variant based on cumulative traffic split

383 cumulative = 0.0

384 versions = sorted(traffic_split.keys()) # Sort for consistency

385

386 for version in versions:

387 cumulative += traffic_split[version]

388 if normalized < cumulative:

389 return version

390

391 # Fallback to last version (handles floating point errors)

392 return versions[-1]

393

394 async def _persist_experiment(self, experiment: PromptExperiment):

395 """Persist experiment to backend storage."""

396 if hasattr(self.storage, "set"):

397 key = f"experiment:{experiment.experiment_id}"

398 await self.storage.set(key, experiment.to_dict())

399

400 async def get_variant_distribution(

401 self,

402 experiment_id: str,

403 ) -> Dict[str, int]:

404 """Get actual distribution of users across variants.

405

406 Args:

407 experiment_id: Experiment ID

408

409 Returns:

410 Dictionary mapping version to user count

411

412 Raises:

413 VersioningError: If experiment not found

414 """

415 experiment = self._experiments.get(experiment_id)

416 if not experiment:

417 raise VersioningError(f"Experiment not found: {experiment_id}")

418

419 assignments = self._user_assignments.get(experiment_id, {})

420 distribution: Dict[str, int] = {v.version: 0 for v in experiment.variants}

421

422 for version in assignments.values():

423 if version in distribution:

424 distribution[version] += 1

425

426 return distribution