Coverage for src/dataknobs_bots/bot/registry.py: 30%

1"""Multi-tenant bot registry with pluggable storage backends.

2

3 This module provides a registry for managing bot configurations and instances

4 across multiple tenants. It combines:

5- Pluggable storage backends (via RegistryBackend protocol)

6- Environment-aware configuration resolution

7- Portability validation for cross-environment deployments

8- Bot instance caching with TTL

9

10 Example:

11 ```python

12 from dataknobs_bots.bot import BotRegistry

13 from dataknobs_bots.registry import InMemoryBackend

14

15 # Create registry with in-memory storage

16 registry = BotRegistry(

17 backend=InMemoryBackend(),

18 environment="production",

19 )

20 await registry.initialize()

21

22 # Register a portable bot configuration

23 await registry.register("my-bot", {

24 "bot": {

25 "llm": {"$resource": "default", "type": "llm_providers"},

26 "conversation_storage": {"$resource": "db", "type": "databases"},

27 }

28 })

29

30 # Get bot instance (resolves $resource references)

31 bot = await registry.get_bot("my-bot")

32 response = await bot.chat(message, context)

33

34 # Cleanup

35 await registry.close()

36 ```

37"""

38

39 from __future__ import annotations

40

41 import asyncio

42 import logging

43 import time

44 from pathlib import Path

45 from typing import TYPE_CHECKING, Any

46

47 from ..registry import InMemoryBackend, RegistryBackend, validate_portability

48 from .base import DynaBot

49

50 if TYPE_CHECKING:

51 from dataknobs_config import EnvironmentConfig

52

53 from ..registry import Registration

54

55 logger = logging.getLogger(__name__)

56

57

58 class BotRegistry:

59 """Multi-tenant bot registry with caching and environment support.

60

61 The BotRegistry manages multiple bot instances for different clients/tenants.

62 It provides:

63 - Pluggable storage backends via RegistryBackend protocol

64 - Environment-aware configuration resolution

65 - Portability validation to ensure configs work across environments

66 - LRU-style caching with TTL for bot instances

67 - Thread-safe access

68

69 This enables:

70 - Multi-tenant SaaS platforms

71 - A/B testing with different bot configurations

72 - Horizontal scaling with stateless bot instances

73 - Cross-environment deployment with portable configs

74

75 Attributes:

76 backend: Storage backend for configurations

77 environment: Environment for $resource resolution

78 cache_ttl: Time-to-live for cached bots in seconds

79 max_cache_size: Maximum number of bots to cache

80

81 Example:

82 ```python

83 from dataknobs_bots.bot import BotRegistry

84 from dataknobs_bots.registry import InMemoryBackend

85

86 # Create registry

87 registry = BotRegistry(

88 backend=InMemoryBackend(),

89 environment="production",

90 cache_ttl=300,

91 )

92 await registry.initialize()

93

94 # Register portable configuration

95 await registry.register("client-123", {

96 "bot": {

97 "llm": {"$resource": "default", "type": "llm_providers"},

98 }

99 })

100

101 # Get bot for a client

102 bot = await registry.get_bot("client-123")

103

104 # Use the bot

105 response = await bot.chat(message, context)

106 ```

107 """

108

109 def __init__(

110 self,

111 backend: RegistryBackend | None = None,

112 environment: EnvironmentConfig | str | None = None,

113 env_dir: str | Path = "config/environments",

114 cache_ttl: int = 300,

115 max_cache_size: int = 1000,

116 validate_on_register: bool = True,

117 config_key: str = "bot",

118 ):

119 """Initialize bot registry.

120

121 Args:

122 backend: Storage backend for configurations.

123 If None, uses InMemoryBackend.

124 environment: Environment name or EnvironmentConfig for

125 $resource resolution. If None, configs are used as-is

126 without environment resolution.

127 env_dir: Directory containing environment config files.

128 Only used if environment is a string name.

129 cache_ttl: Cache time-to-live in seconds (default: 300)

130 max_cache_size: Maximum cached bots (default: 1000)

131 validate_on_register: If True, validate config portability

132 when registering (default: True)

133 config_key: Key within config containing bot configuration.

134 Defaults to "bot". Used during environment resolution.

135 """

136 self._backend = backend or InMemoryBackend()

137 self._env_dir = Path(env_dir)

138 self._cache_ttl = cache_ttl

139 self._max_cache_size = max_cache_size

140 self._validate_on_register = validate_on_register

141 self._config_key = config_key

142

143 # Bot instance cache: bot_id -> (DynaBot, cached_timestamp)

144 self._cache: dict[str, tuple[DynaBot, float]] = {}

145 self._lock = asyncio.Lock()

146 self._initialized = False

147

148 # Load environment config if specified

149 self._environment: EnvironmentConfig | None = None

150 if environment is not None:

151 try:

152 from dataknobs_config import EnvironmentConfig as EnvConfig

153

154 if isinstance(environment, str):

155 self._environment = EnvConfig.load(environment, env_dir)

156 else:

157 self._environment = environment

158 logger.info(f"BotRegistry using environment: {self._environment.name}")

159 except ImportError:

160 logger.warning(

161 "dataknobs_config not installed, environment-aware features disabled"

162 )

163

164 @property

165 def backend(self) -> RegistryBackend:

166 """Get the storage backend."""

167 return self._backend

168

169 @property

170 def environment(self) -> EnvironmentConfig | None:

171 """Get current environment config, or None if not environment-aware."""

172 return self._environment

173

174 @property

175 def environment_name(self) -> str | None:

176 """Get current environment name, or None if not environment-aware."""

177 return self._environment.name if self._environment else None

178

179 @property

180 def cache_ttl(self) -> int:

181 """Get cache TTL in seconds."""

182 return self._cache_ttl

183

184 @property

185 def max_cache_size(self) -> int:

186 """Get maximum cache size."""

187 return self._max_cache_size

188

189 async def initialize(self) -> None:

190 """Initialize the registry and backend.

191

192 Must be called before using the registry.

193 """

194 if not self._initialized:

195 await self._backend.initialize()

196 self._initialized = True

197 logger.info("BotRegistry initialized")

198

199 async def close(self) -> None:

200 """Close the registry and backend.

201

202 Clears the bot cache and closes the storage backend.

203 """

204 async with self._lock:

205 self._cache.clear()

206 await self._backend.close()

207 self._initialized = False

208 logger.info("BotRegistry closed")

209

210 async def register(

211 self,

212 bot_id: str,

213 config: dict[str, Any],

214 status: str = "active",

215 skip_validation: bool = False,

216 ) -> Registration:

217 """Register or update a bot configuration.

218

219 Stores a portable configuration in the backend. By default, validates

220 that the configuration is portable (no resolved local values).

221

222 Args:

223 bot_id: Unique bot identifier

224 config: Bot configuration dictionary (should be portable)

225 status: Registration status (default: active)

226 skip_validation: If True, skip portability validation

227

228 Returns:

229 Registration object with metadata

230

231 Raises:

232 PortabilityError: If config is not portable and validation is enabled

233

234 Example:

235 ```python

236 # Register with portable config

237 reg = await registry.register("support-bot", {

238 "bot": {

239 "llm": {"$resource": "default", "type": "llm_providers"},

240 }

241 })

242 print(f"Registered at: {reg.created_at}")

243

244 # Update existing registration

245 reg = await registry.register("support-bot", new_config)

246 print(f"Updated at: {reg.updated_at}")

247 ```

248 """

249 # Validate portability if enabled

250 if self._validate_on_register and not skip_validation:

251 validate_portability(config)

252

253 # Store in backend

254 registration = await self._backend.register(bot_id, config, status)

255

256 # Invalidate cache for this bot

257 async with self._lock:

258 if bot_id in self._cache:

259 del self._cache[bot_id]

260 logger.debug(f"Invalidated cache for bot: {bot_id}")

261

262 logger.info(f"Registered bot: {bot_id}")

263 return registration

264

265 async def get_bot(

266 self,

267 bot_id: str,

268 force_refresh: bool = False,

269 ) -> DynaBot:

270 """Get bot instance for a client.

271

272 Bots are cached for performance. If a cached bot exists and hasn't

273 expired, it's returned. Otherwise, a new bot is created from the

274 stored configuration with environment resolution applied.

275

276 Args:

277 bot_id: Bot identifier

278 force_refresh: If True, bypass cache and create fresh bot

279

280 Returns:

281 DynaBot instance for the client

282

283 Raises:

284 KeyError: If no registration exists for the bot_id

285 ValueError: If bot configuration is invalid

286

287 Example:

288 ```python

289 # Get cached bot

290 bot = await registry.get_bot("client-123")

291

292 # Force refresh (e.g., after config change)

293 bot = await registry.get_bot("client-123", force_refresh=True)

294 ```

295 """

296 async with self._lock:

297 # Check cache

298 if not force_refresh and bot_id in self._cache:

299 bot, cached_at = self._cache[bot_id]

300 if time.time() - cached_at < self._cache_ttl:

301 logger.debug(f"Returning cached bot: {bot_id}")

302 return bot

303

304 # Load configuration from backend

305 config = await self._backend.get_config(bot_id)

306 if config is None:

307 raise KeyError(f"No bot configuration found for: {bot_id}")

308

309 # Create bot with environment resolution if configured

310 if self._environment is not None:

311 logger.debug(f"Creating bot with environment resolution: {bot_id}")

312 bot = await DynaBot.from_environment_aware_config(

313 config,

314 environment=self._environment,

315 env_dir=self._env_dir,

316 config_key=self._config_key,

317 )

318 else:

319 # Traditional path - use config as-is

320 # Extract bot config if wrapped in config_key

321 bot_config = config.get(self._config_key, config)

322 logger.debug(f"Creating bot without environment resolution: {bot_id}")

323 bot = await DynaBot.from_config(bot_config)

324

325 # Cache the bot

326 self._cache[bot_id] = (bot, time.time())

327 logger.info(f"Created bot: {bot_id}")

328

329 # Evict old entries if cache is full

330 if len(self._cache) > self._max_cache_size:

331 self._evict_oldest()

332

333 return bot

334

335 async def get_config(self, bot_id: str) -> dict[str, Any] | None:

336 """Get stored configuration for a bot.

337

338 Returns the portable configuration as stored, without

339 environment resolution applied.

340

341 Args:

342 bot_id: Bot identifier

343

344 Returns:

345 Configuration dict if found, None otherwise

346 """

347 return await self._backend.get_config(bot_id)

348

349 async def get_registration(self, bot_id: str) -> Registration | None:

350 """Get full registration including metadata.

351

352 Args:

353 bot_id: Bot identifier

354

355 Returns:

356 Registration if found, None otherwise

357 """

358 return await self._backend.get(bot_id)

359

360 async def unregister(self, bot_id: str) -> bool:

361 """Remove a bot registration (hard delete).

362

363 Args:

364 bot_id: Bot identifier

365

366 Returns:

367 True if removed, False if not found

368 """

369 # Remove from cache

370 async with self._lock:

371 if bot_id in self._cache:

372 del self._cache[bot_id]

373

374 result = await self._backend.unregister(bot_id)

375 if result:

376 logger.info(f"Unregistered bot: {bot_id}")

377 return result

378

379 async def deactivate(self, bot_id: str) -> bool:

380 """Deactivate a bot registration (soft delete).

381

382 Args:

383 bot_id: Bot identifier

384

385 Returns:

386 True if deactivated, False if not found

387 """

388 # Remove from cache

389 async with self._lock:

390 if bot_id in self._cache:

391 del self._cache[bot_id]

392

393 result = await self._backend.deactivate(bot_id)

394 if result:

395 logger.info(f"Deactivated bot: {bot_id}")

396 return result

397

398 async def exists(self, bot_id: str) -> bool:

399 """Check if an active bot registration exists.

400

401 Args:

402 bot_id: Bot identifier

403

404 Returns:

405 True if registration exists and is active

406 """

407 return await self._backend.exists(bot_id)

408

409 async def list_bots(self) -> list[str]:

410 """List all active bot IDs.

411

412 Returns:

413 List of active bot identifiers

414 """

415 return await self._backend.list_ids()

416

417 async def count(self) -> int:

418 """Count active bot registrations.

419

420 Returns:

421 Number of active registrations

422 """

423 return await self._backend.count()

424

425 def get_cached_bots(self) -> list[str]:

426 """Get list of currently cached bot IDs.

427

428 Returns:

429 List of bot IDs with cached instances

430 """

431 return list(self._cache.keys())

432

433 def clear_cache(self) -> None:

434 """Clear all cached bot instances.

435

436 Does not affect stored registrations.

437 """

438 self._cache.clear()

439 logger.debug("Cleared bot cache")

440

441 def _evict_oldest(self) -> None:

442 """Evict oldest cache entries when cache is full.

443

444 Removes 10% of the oldest entries to make room for new ones.

445 """

446 # Sort by timestamp (oldest first)

447 sorted_items = sorted(self._cache.items(), key=lambda x: x[1][1])

448

449 # Remove oldest 10%

450 num_to_remove = max(1, len(sorted_items) // 10)

451 for bot_id, _ in sorted_items[:num_to_remove]:

452 del self._cache[bot_id]

453 logger.debug(f"Evicted {num_to_remove} bots from cache")

454

455 # Legacy compatibility methods

456

457 async def register_client(

458 self, client_id: str, bot_config: dict[str, Any]

459 ) -> None:

460 """Register or update a client's bot configuration.

461

462 .. deprecated::

463 Use :meth:`register` instead.

464

465 Args:

466 client_id: Client/tenant identifier

467 bot_config: Bot configuration dictionary

468 """

469 await self.register(client_id, bot_config)

470

471 async def remove_client(self, client_id: str) -> None:

472 """Remove a client from the registry.

473

474 .. deprecated::

475 Use :meth:`unregister` instead.

476

477 Args:

478 client_id: Client/tenant identifier

479 """

480 await self.unregister(client_id)

481

482 def get_cached_clients(self) -> list[str]:

483 """Get list of currently cached client IDs.

484

485 .. deprecated::

486 Use :meth:`get_cached_bots` instead.

487

488 Returns:

489 List of client IDs with cached bots

490 """

491 return self.get_cached_bots()

492

493 def __repr__(self) -> str:

494 """String representation."""

495 env = f", environment={self._environment.name!r}" if self._environment else ""

496 return (

497 f"BotRegistry(backend={self._backend!r}, "

498 f"cached={len(self._cache)}{env})"

499 )

500

501

502 class InMemoryBotRegistry(BotRegistry):

503 """BotRegistry with in-memory storage backend.

504

505 A convenience subclass that uses InMemoryBackend for storage,

506 suitable for testing, CLIs, and single-instance deployments.

507

508 Unlike the base BotRegistry which accepts a pluggable backend,

509 this class always uses in-memory storage and doesn't require

510 external dependencies like databases.

511

512 Example:

513 ```python

514 from dataknobs_bots.bot import InMemoryBotRegistry

515

516 # For testing - no environment resolution

517 registry = InMemoryBotRegistry(validate_on_register=False)

518 await registry.initialize()

519

520 await registry.register("test-bot", {"llm": {"provider": "echo"}})

521 bot = await registry.get_bot("test-bot")

522

523 # For development with environment

524 registry = InMemoryBotRegistry(environment="development")

525 await registry.initialize()

526 ```

527 """

528

529 def __init__(

530 self,

531 environment: EnvironmentConfig | str | None = None,

532 env_dir: str | Path = "config/environments",

533 cache_ttl: int = 300,

534 max_cache_size: int = 1000,

535 validate_on_register: bool = True,

536 config_key: str = "bot",

537 ):

538 """Initialize in-memory bot registry.

539

540 Args:

541 environment: Environment name or EnvironmentConfig for

542 $resource resolution. If None, configs are used as-is

543 without environment resolution.

544 env_dir: Directory containing environment config files.

545 Only used if environment is a string name.

546 cache_ttl: Cache time-to-live in seconds (default: 300)

547 max_cache_size: Maximum cached bots (default: 1000)

548 validate_on_register: If True, validate config portability

549 when registering (default: True)

550 config_key: Key within config containing bot configuration.

551 Defaults to "bot". Used during environment resolution.

552 """

553 super().__init__(

554 backend=InMemoryBackend(),

555 environment=environment,

556 env_dir=env_dir,

557 cache_ttl=cache_ttl,

558 max_cache_size=max_cache_size,

559 validate_on_register=validate_on_register,

560 config_key=config_key,

561 )

562

563 async def clear(self) -> None:

564 """Clear all registrations and cached bots.

565

566 Convenience method for test cleanup that clears both the

567 backend storage and the bot instance cache.

568

569 Example:

570 ```python

571 # In tests - reset between test cases

572 await registry.clear()

573 assert await registry.count() == 0

574 ```

575 """

576 await self._backend.clear()

577 self._cache.clear()

578 logger.debug("Cleared all registrations and cache")

579

580 def __repr__(self) -> str:

581 """String representation."""

582 env = f", environment={self._environment.name!r}" if self._environment else ""

583 return f"InMemoryBotRegistry(cached={len(self._cache)}{env})"

584

585

586 def create_memory_registry(

587 environment: EnvironmentConfig | str | None = None,

588 env_dir: str | Path = "config/environments",

589 cache_ttl: int = 300,

590 max_cache_size: int = 1000,

591 validate_on_register: bool = True,

592 config_key: str = "bot",

593) -> InMemoryBotRegistry:

594 """Create an InMemoryBotRegistry.

595

596 Convenience factory for creating in-memory registries suitable for

597 testing, CLIs, or single-instance deployments.

598

599 Args:

600 environment: Environment name or EnvironmentConfig for

601 $resource resolution. If None, configs are used as-is.

602 env_dir: Directory containing environment config files.

603 cache_ttl: Cache time-to-live in seconds (default: 300)

604 max_cache_size: Maximum cached bots (default: 1000)

605 validate_on_register: If True, validate config portability

606 config_key: Key within config containing bot configuration

607

608 Returns:

609 InMemoryBotRegistry instance

610

611 Example:

612 ```python

613 from dataknobs_bots.bot import create_memory_registry

614

615 registry = create_memory_registry(validate_on_register=False)

616 await registry.initialize()

617

618 await registry.register("test-bot", {"llm": {"provider": "echo"}})

619 bot = await registry.get_bot("test-bot")

620 ```

621 """

622 return InMemoryBotRegistry(

623 environment=environment,

624 env_dir=env_dir,

625 cache_ttl=cache_ttl,

626 max_cache_size=max_cache_size,

627 validate_on_register=validate_on_register,

628 config_key=config_key,

629 )