Coverage for src / dataknobs_bots / bot / manager.py: 30%
100 statements
« prev ^ index » next coverage.py v7.13.0, created at 2025-12-16 10:13 -0700
1"""Bot manager for multi-tenant bot instances."""
3from __future__ import annotations
5import asyncio
6import inspect
7import logging
8from pathlib import Path
9from typing import TYPE_CHECKING, Any, Callable, Protocol, runtime_checkable
11from .base import DynaBot
13if TYPE_CHECKING:
14 from dataknobs_config import EnvironmentAwareConfig, EnvironmentConfig
16logger = logging.getLogger(__name__)
19@runtime_checkable
20class ConfigLoader(Protocol):
21 """Protocol for configuration loaders with a load method."""
23 def load(self, bot_id: str) -> dict[str, Any]:
24 """Load configuration for a bot."""
25 ...
28@runtime_checkable
29class AsyncConfigLoader(Protocol):
30 """Protocol for async configuration loaders."""
32 async def load(self, bot_id: str) -> dict[str, Any]:
33 """Load configuration for a bot asynchronously."""
34 ...
37ConfigLoaderType = (
38 ConfigLoader
39 | AsyncConfigLoader
40 | Callable[[str], dict[str, Any]]
41 | Callable[[str], Any] # For async callables
42)
45class BotManager:
46 """Manages multiple DynaBot instances for multi-tenancy.
48 BotManager handles:
49 - Bot instance creation and caching
50 - Client-level isolation
51 - Configuration loading and validation
52 - Bot lifecycle management
53 - Environment-aware resource resolution (optional)
55 Each client/tenant gets its own bot instance, which can serve multiple users.
56 The underlying DynaBot architecture ensures conversation isolation through
57 BotContext with different conversation_ids.
59 Attributes:
60 bots: Cache of bot_id -> DynaBot instances
61 config_loader: Optional configuration loader (sync or async)
62 environment_name: Current environment name (if environment-aware)
64 Example:
65 ```python
66 # Basic usage with inline configuration
67 manager = BotManager()
68 bot = await manager.get_or_create("my-bot", config={
69 "llm": {"provider": "openai", "model": "gpt-4o"},
70 "conversation_storage": {"backend": "memory"},
71 })
73 # With environment-aware configuration
74 manager = BotManager(environment="production")
75 bot = await manager.get_or_create("my-bot", config={
76 "bot": {
77 "llm": {"$resource": "default", "type": "llm_providers"},
78 "conversation_storage": {"$resource": "db", "type": "databases"},
79 }
80 })
82 # With config loader function
83 def load_config(bot_id: str) -> dict:
84 return load_yaml(f"configs/{bot_id}.yaml")
86 manager = BotManager(config_loader=load_config)
87 bot = await manager.get_or_create("my-bot")
89 # List active bots
90 active_bots = manager.list_bots()
91 ```
92 """
94 def __init__(
95 self,
96 config_loader: ConfigLoaderType | None = None,
97 environment: EnvironmentConfig | str | None = None,
98 env_dir: str | Path = "config/environments",
99 ):
100 """Initialize BotManager.
102 Args:
103 config_loader: Optional configuration loader.
104 Can be:
105 - An object with a `.load(bot_id)` method (sync or async)
106 - A callable function: bot_id -> config_dict (sync or async)
107 - None (configurations must be provided explicitly)
108 environment: Environment name or EnvironmentConfig for resource resolution.
109 If None, environment-aware features are disabled unless
110 an EnvironmentAwareConfig is passed to get_or_create().
111 If a string, loads environment config from env_dir.
112 env_dir: Directory containing environment config files.
113 Only used if environment is a string name.
114 """
115 self._bots: dict[str, DynaBot] = {}
116 self._config_loader = config_loader
117 self._env_dir = Path(env_dir)
119 # Load environment config if specified
120 self._environment: EnvironmentConfig | None = None
121 if environment is not None:
122 try:
123 from dataknobs_config import EnvironmentConfig
125 if isinstance(environment, str):
126 self._environment = EnvironmentConfig.load(environment, env_dir)
127 else:
128 self._environment = environment
129 logger.info(f"Initialized BotManager with environment: {self._environment.name}")
130 except ImportError:
131 logger.warning(
132 "dataknobs_config not installed, environment-aware features disabled"
133 )
134 else:
135 logger.info("Initialized BotManager")
137 @property
138 def environment_name(self) -> str | None:
139 """Get current environment name, or None if not environment-aware."""
140 return self._environment.name if self._environment else None
142 @property
143 def environment(self) -> EnvironmentConfig | None:
144 """Get current environment config, or None if not environment-aware."""
145 return self._environment
147 async def get_or_create(
148 self,
149 bot_id: str,
150 config: dict[str, Any] | EnvironmentAwareConfig | None = None,
151 use_environment: bool | None = None,
152 config_key: str = "bot",
153 ) -> DynaBot:
154 """Get existing bot or create new one.
156 Args:
157 bot_id: Bot identifier (e.g., "customer-support", "sales-assistant")
158 config: Optional bot configuration. Can be:
159 - dict with resolved values (traditional)
160 - dict with $resource references (requires environment)
161 - EnvironmentAwareConfig instance
162 If not provided and config_loader is set, will load configuration.
163 use_environment: Whether to use environment-aware resolution.
164 - True: Use environment for $resource resolution
165 - False: Use config as-is (no resolution)
166 - None (default): Auto-detect based on whether manager has
167 an environment configured or config is EnvironmentAwareConfig
168 config_key: Key within config containing bot configuration.
169 Defaults to "bot". Set to None to use root config.
170 Only used when use_environment is True.
172 Returns:
173 DynaBot instance
175 Raises:
176 ValueError: If config is None and no config_loader is set
178 Example:
179 ```python
180 # Traditional usage (no environment resolution)
181 manager = BotManager()
182 bot = await manager.get_or_create("support-bot", config={
183 "llm": {"provider": "openai", "model": "gpt-4"},
184 "conversation_storage": {"backend": "memory"},
185 })
187 # Environment-aware usage with $resource references
188 manager = BotManager(environment="production")
189 bot = await manager.get_or_create("support-bot", config={
190 "bot": {
191 "llm": {"$resource": "default", "type": "llm_providers"},
192 "conversation_storage": {"$resource": "db", "type": "databases"},
193 }
194 })
196 # Explicit environment resolution control
197 bot = await manager.get_or_create(
198 "support-bot",
199 config=my_config,
200 use_environment=True,
201 config_key="bot"
202 )
203 ```
204 """
205 # Return cached bot if exists
206 if bot_id in self._bots:
207 logger.debug(f"Returning cached bot: {bot_id}")
208 return self._bots[bot_id]
210 # Load configuration if not provided
211 if config is None:
212 if self._config_loader is None:
213 raise ValueError(
214 f"No configuration provided for bot '{bot_id}' "
215 "and no config_loader is set"
216 )
217 config = await self._load_config(bot_id)
219 # Determine whether to use environment resolution
220 is_env_aware_config = False
221 try:
222 from dataknobs_config import EnvironmentAwareConfig
224 is_env_aware_config = isinstance(config, EnvironmentAwareConfig)
225 except ImportError:
226 pass
228 should_use_environment = use_environment
229 if should_use_environment is None:
230 # Auto-detect: use environment if manager has one or config is EnvironmentAwareConfig
231 should_use_environment = self._environment is not None or is_env_aware_config
233 # Create new bot
234 logger.info(f"Creating new bot: {bot_id} (environment_aware={should_use_environment})")
236 if should_use_environment:
237 bot = await DynaBot.from_environment_aware_config(
238 config,
239 environment=self._environment,
240 env_dir=self._env_dir,
241 config_key=config_key,
242 )
243 else:
244 # Traditional path - use config as-is
245 bot = await DynaBot.from_config(config)
247 # Cache and return
248 self._bots[bot_id] = bot
249 return bot
251 async def get(self, bot_id: str) -> DynaBot | None:
252 """Get bot without creating if doesn't exist.
254 Args:
255 bot_id: Bot identifier
257 Returns:
258 DynaBot instance if exists, None otherwise
259 """
260 return self._bots.get(bot_id)
262 async def remove(self, bot_id: str) -> bool:
263 """Remove bot instance.
265 Args:
266 bot_id: Bot identifier
268 Returns:
269 True if bot was removed, False if didn't exist
270 """
271 if bot_id in self._bots:
272 logger.info(f"Removing bot: {bot_id}")
273 del self._bots[bot_id]
274 return True
275 return False
277 async def reload(self, bot_id: str) -> DynaBot:
278 """Reload bot instance with fresh configuration.
280 Args:
281 bot_id: Bot identifier
283 Returns:
284 New DynaBot instance
286 Raises:
287 ValueError: If no config_loader is set
288 """
289 if self._config_loader is None:
290 raise ValueError("Cannot reload without config_loader")
292 # Remove existing bot
293 await self.remove(bot_id)
295 # Create new one
296 return await self.get_or_create(bot_id)
298 def list_bots(self) -> list[str]:
299 """List all active bot IDs.
301 Returns:
302 List of bot identifiers
303 """
304 return list(self._bots.keys())
306 def get_bot_count(self) -> int:
307 """Get count of active bots.
309 Returns:
310 Number of active bot instances
311 """
312 return len(self._bots)
314 async def _load_config(self, bot_id: str) -> dict[str, Any]:
315 """Load configuration for bot using config_loader.
317 Supports both synchronous and asynchronous config loaders.
318 Handles both callable loaders and objects with a load() method.
320 Args:
321 bot_id: Bot identifier
323 Returns:
324 Bot configuration dictionary
325 """
326 logger.debug(f"Loading configuration for bot: {bot_id}")
328 if callable(self._config_loader):
329 # Handle callable config loader (function)
330 if inspect.iscoroutinefunction(self._config_loader):
331 # Async function
332 result = await self._config_loader(bot_id)
333 return dict(result) if isinstance(result, dict) else {}
334 else:
335 # Sync function - run in executor to avoid blocking
336 loop = asyncio.get_event_loop()
337 result = await loop.run_in_executor(None, self._config_loader, bot_id)
338 return dict(result) if isinstance(result, dict) else {}
339 else:
340 # Assume it's an object with a load method
341 load_method = self._config_loader.load # type: ignore
343 if inspect.iscoroutinefunction(load_method):
344 # Async method
345 result = await load_method(bot_id)
346 return dict(result) if isinstance(result, dict) else {}
347 else:
348 # Sync method - run in executor to avoid blocking
349 loop = asyncio.get_event_loop()
350 result = await loop.run_in_executor(None, load_method, bot_id)
351 return dict(result) if isinstance(result, dict) else {}
353 async def clear_all(self) -> None:
354 """Clear all bot instances.
356 Useful for testing or when restarting the service.
357 """
358 logger.info("Clearing all bot instances")
359 self._bots.clear()
361 def get_portable_config(
362 self,
363 config: dict[str, Any] | EnvironmentAwareConfig,
364 ) -> dict[str, Any]:
365 """Get portable configuration for storage.
367 Extracts portable config (with $resource references intact,
368 environment variables unresolved) suitable for storing in
369 registries or databases.
371 Args:
372 config: Configuration to make portable.
373 Can be dict or EnvironmentAwareConfig.
375 Returns:
376 Portable configuration dictionary
378 Example:
379 ```python
380 manager = BotManager(environment="production")
382 # Get portable config from EnvironmentAwareConfig
383 portable = manager.get_portable_config(env_aware_config)
385 # Store in registry (portable across environments)
386 await registry.store(bot_id, portable)
387 ```
388 """
389 return DynaBot.get_portable_config(config)
391 def __repr__(self) -> str:
392 """String representation."""
393 bots = ", ".join(self._bots.keys())
394 env = f", environment={self._environment.name!r}" if self._environment else ""
395 return f"BotManager(bots=[{bots}], count={len(self._bots)}{env})"