Coverage for src/dataknobs_bots/bot/base.py: 11%

1"""Core DynaBot implementation."""

2

3 from __future__ import annotations

4

5 from collections.abc import AsyncGenerator

6 from pathlib import Path

7 from types import TracebackType

8 from typing import TYPE_CHECKING, Any

9

10 from typing_extensions import Self

11

12 from dataknobs_llm.conversations import ConversationManager, DataknobsConversationStorage

13 from dataknobs_llm.llm import AsyncLLMProvider

14 from dataknobs_llm.prompts import AsyncPromptBuilder

15 from dataknobs_llm.tools import ToolRegistry

16

17 from .context import BotContext

18 from ..memory.base import Memory

19

20 if TYPE_CHECKING:

21 from dataknobs_config import EnvironmentAwareConfig, EnvironmentConfig

22

23

24 class DynaBot:

25 """Configuration-driven chatbot leveraging the DataKnobs ecosystem.

26

27 DynaBot provides a flexible, configuration-driven bot that can be customized

28 for different use cases through YAML/JSON configuration files.

29

30 Attributes:

31 llm: LLM provider for generating responses

32 prompt_builder: Prompt builder for managing prompts

33 conversation_storage: Storage backend for conversations

34 tool_registry: Registry of available tools

35 memory: Optional memory implementation for context

36 knowledge_base: Optional knowledge base for RAG

37 reasoning_strategy: Optional reasoning strategy

38 middleware: List of middleware for request/response processing

39 system_prompt_name: Name of the system prompt template to use

40 system_prompt_content: Inline system prompt content (alternative to name)

41 system_prompt_rag_configs: RAG configurations for inline system prompts

42 default_temperature: Default temperature for LLM generation

43 default_max_tokens: Default max tokens for LLM generation

44 """

45

46 def __init__(

47 self,

48 llm: AsyncLLMProvider,

49 prompt_builder: AsyncPromptBuilder,

50 conversation_storage: DataknobsConversationStorage,

51 tool_registry: ToolRegistry | None = None,

52 memory: Memory | None = None,

53 knowledge_base: Any | None = None,

54 reasoning_strategy: Any | None = None,

55 middleware: list[Any] | None = None,

56 system_prompt_name: str | None = None,

57 system_prompt_content: str | None = None,

58 system_prompt_rag_configs: list[dict[str, Any]] | None = None,

59 default_temperature: float = 0.7,

60 default_max_tokens: int = 1000,

61 ):

62 """Initialize DynaBot.

63

64 Args:

65 llm: LLM provider instance

66 prompt_builder: Prompt builder instance

67 conversation_storage: Conversation storage backend

68 tool_registry: Optional tool registry

69 memory: Optional memory implementation

70 knowledge_base: Optional knowledge base

71 reasoning_strategy: Optional reasoning strategy

72 middleware: Optional middleware list

73 system_prompt_name: Name of system prompt template (mutually exclusive with content)

74 system_prompt_content: Inline system prompt content (mutually exclusive with name)

75 system_prompt_rag_configs: RAG configurations for inline system prompts

76 default_temperature: Default temperature (0-1)

77 default_max_tokens: Default max tokens to generate

78 """

79 self.llm = llm

80 self.prompt_builder = prompt_builder

81 self.conversation_storage = conversation_storage

82 self.tool_registry = tool_registry or ToolRegistry()

83 self.memory = memory

84 self.knowledge_base = knowledge_base

85 self.reasoning_strategy = reasoning_strategy

86 self.middleware = middleware or []

87 self.system_prompt_name = system_prompt_name

88 self.system_prompt_content = system_prompt_content

89 self.system_prompt_rag_configs = system_prompt_rag_configs

90 self.default_temperature = default_temperature

91 self.default_max_tokens = default_max_tokens

92 self._conversation_managers: dict[str, ConversationManager] = {}

93

94 @classmethod

95 async def from_config(cls, config: dict[str, Any]) -> DynaBot:

96 """Create DynaBot from configuration.

97

98 Args:

99 config: Configuration dictionary containing:

100 - llm: LLM configuration (provider, model, etc.)

101 - conversation_storage: Storage configuration

102 - tools: Optional list of tool configurations

103 - memory: Optional memory configuration

104 - knowledge_base: Optional knowledge base configuration

105 - reasoning: Optional reasoning strategy configuration

106 - middleware: Optional middleware configurations

107 - prompts: Optional prompts library (dict of name -> content)

108 - system_prompt: Optional system prompt configuration (see below)

109

110 Returns:

111 Configured DynaBot instance

112

113 System Prompt Formats:

114 The system_prompt can be specified in multiple ways:

115

116 - String: Smart detection - if the string exists as a template name

117 in the prompt library, it's used as a template reference; otherwise

118 it's treated as inline content.

119

120 - Dict with name: `{"name": "template_name"}` - explicit template reference

121 - Dict with name + strict: `{"name": "template_name", "strict": true}` -

122 raises error if template doesn't exist

123 - Dict with content: `{"content": "inline prompt text"}` - inline content

124 - Dict with content + rag_configs: inline content with RAG enhancement

125

126 Example:

127 ```python

128 # Smart detection: uses as template if it exists in prompts library

129 config = {

130 "llm": {"provider": "openai", "model": "gpt-4"},

131 "conversation_storage": {"backend": "memory"},

132 "prompts": {

133 "helpful_assistant": "You are a helpful AI assistant."

134 },

135 "system_prompt": "helpful_assistant" # Found in prompts, used as template

136 }

137

138 # Smart detection: treated as inline content (not in prompts library)

139 config = {

140 "llm": {"provider": "openai", "model": "gpt-4"},

141 "conversation_storage": {"backend": "memory"},

142 "system_prompt": "You are a helpful assistant." # Not a template name

143 }

144

145 # Explicit inline content with RAG enhancement

146 config = {

147 "llm": {"provider": "openai", "model": "gpt-4"},

148 "conversation_storage": {"backend": "memory"},

149 "system_prompt": {

150 "content": "You are a helpful assistant. Use this context: {{ CONTEXT }}",

151 "rag_configs": [{

152 "adapter_name": "docs",

153 "query": "assistant guidelines",

154 "placeholder": "CONTEXT",

155 "k": 3

156 }]

157 }

158 }

159

160 # Strict mode: error if template doesn't exist

161 config = {

162 "llm": {"provider": "openai", "model": "gpt-4"},

163 "conversation_storage": {"backend": "memory"},

164 "system_prompt": {

165 "name": "my_template",

166 "strict": true # Raises ValueError if my_template doesn't exist

167 }

168 }

169

170 bot = await DynaBot.from_config(config)

171 ```

172 """

173 from dataknobs_data.factory import AsyncDatabaseFactory

174 from dataknobs_llm.llm import LLMProviderFactory

175 from dataknobs_llm.prompts import AsyncPromptBuilder

176 from dataknobs_llm.prompts.implementations import CompositePromptLibrary

177 from ..memory import create_memory_from_config

178

179 # Create LLM provider

180 llm_config = config["llm"]

181 factory = LLMProviderFactory(is_async=True)

182 llm = factory.create(llm_config)

183 await llm.initialize()

184

185 # Create conversation storage

186 storage_config = config["conversation_storage"].copy()

187

188 # Create database backend using factory

189 db_factory = AsyncDatabaseFactory()

190 backend = db_factory.create(**storage_config)

191 await backend.connect()

192 conversation_storage = DataknobsConversationStorage(backend)

193

194 # Create prompt builder

195 # Support optional prompts configuration

196 prompt_libraries = []

197 if "prompts" in config:

198 from dataknobs_llm.prompts.implementations import ConfigPromptLibrary

199

200 prompts_config = config["prompts"]

201

202 # If prompts are provided as a dict, create a config-based library

203 if isinstance(prompts_config, dict):

204 # Convert simple string prompts to proper template structure

205 structured_config = {"system": {}, "user": {}}

206

207 for prompt_name, prompt_content in prompts_config.items():

208 if isinstance(prompt_content, dict):

209 # Already structured - use as-is

210 # Assume it's a system prompt unless specified

211 prompt_type = prompt_content.get("type", "system")

212 if prompt_type in structured_config:

213 structured_config[prompt_type][prompt_name] = prompt_content

214 else:

215 # Simple string - treat as system prompt template

216 structured_config["system"][prompt_name] = {

217 "template": prompt_content

218 }

219

220 library = ConfigPromptLibrary(structured_config)

221 prompt_libraries.append(library)

222

223 # Create composite library (empty if no prompts configured)

224 library = CompositePromptLibrary(libraries=prompt_libraries)

225 prompt_builder = AsyncPromptBuilder(library)

226

227 # Create tools

228 tool_registry = ToolRegistry()

229 if "tools" in config:

230 for tool_config in config["tools"]:

231 tool = cls._resolve_tool(tool_config, config)

232 if tool:

233 tool_registry.register_tool(tool)

234

235 # Create memory

236 memory = None

237 if "memory" in config:

238 memory = await create_memory_from_config(config["memory"])

239

240 # Create knowledge base

241 knowledge_base = None

242 kb_config = config.get("knowledge_base", {})

243 if kb_config.get("enabled"):

244 from ..knowledge import create_knowledge_base_from_config

245 import logging

246 logger = logging.getLogger(__name__)

247 logger.info(f"Initializing knowledge base with config: {kb_config.get('type', 'unknown')}")

248 knowledge_base = await create_knowledge_base_from_config(kb_config)

249 logger.info("Knowledge base initialized successfully")

250

251 # Create reasoning strategy

252 reasoning_strategy = None

253 if "reasoning" in config:

254 from ..reasoning import create_reasoning_from_config

255

256 reasoning_strategy = create_reasoning_from_config(config["reasoning"])

257

258 # Create middleware

259 middleware = []

260 if "middleware" in config:

261 for mw_config in config["middleware"]:

262 mw = cls._create_middleware(mw_config)

263 if mw:

264 middleware.append(mw)

265

266 # Extract system prompt (supports template name or inline content)

267 system_prompt_name = None

268 system_prompt_content = None

269 system_prompt_rag_configs = None

270 if "system_prompt" in config:

271 system_prompt_config = config["system_prompt"]

272 if isinstance(system_prompt_config, dict):

273 # Explicit dict format: {name: "template"} or {content: "inline..."}

274 system_prompt_name = system_prompt_config.get("name")

275 system_prompt_content = system_prompt_config.get("content")

276 system_prompt_rag_configs = system_prompt_config.get("rag_configs")

277

278 # If strict mode is enabled, require the template to exist

279 if system_prompt_name and system_prompt_config.get("strict"):

280 if library.get_system_prompt(system_prompt_name) is None:

281 raise ValueError(

282 f"System prompt template not found: {system_prompt_name} "

283 "(strict mode enabled)"

284 )

285 elif isinstance(system_prompt_config, str):

286 # String format: smart detection

287 # If it exists in the library, use as template name; otherwise treat as inline

288 if library.get_system_prompt(system_prompt_config) is not None:

289 system_prompt_name = system_prompt_config

290 else:

291 system_prompt_content = system_prompt_config

292

293 return cls(

294 llm=llm,

295 prompt_builder=prompt_builder,

296 conversation_storage=conversation_storage,

297 tool_registry=tool_registry,

298 memory=memory,

299 knowledge_base=knowledge_base,

300 reasoning_strategy=reasoning_strategy,

301 middleware=middleware,

302 system_prompt_name=system_prompt_name,

303 system_prompt_content=system_prompt_content,

304 system_prompt_rag_configs=system_prompt_rag_configs,

305 default_temperature=llm_config.get("temperature", 0.7),

306 default_max_tokens=llm_config.get("max_tokens", 1000),

307 )

308

309 @classmethod

310 async def from_environment_aware_config(

311 cls,

312 config: EnvironmentAwareConfig | dict[str, Any],

313 environment: EnvironmentConfig | str | None = None,

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

315 config_key: str = "bot",

316 ) -> DynaBot:

317 """Create DynaBot with environment-aware configuration.

318

319 This is the recommended entry point for environment-portable bots.

320 Resource references ($resource) are resolved against the environment

321 config, and environment variables are substituted at instantiation time

322 (late binding).

323

324 Args:

325 config: EnvironmentAwareConfig instance or dict with $resource references.

326 If dict, will be wrapped in EnvironmentAwareConfig.

327 environment: Environment name or EnvironmentConfig instance.

328 If None, auto-detects from DATAKNOBS_ENVIRONMENT env var.

329 Ignored if config is already an EnvironmentAwareConfig.

330 env_dir: Directory containing environment config files.

331 Only used if environment is a string name.

332 config_key: Key within config containing bot configuration.

333 Defaults to "bot". Set to None to use root config.

334

335 Returns:

336 Fully initialized DynaBot instance with resolved resources

337

338 Example:

339 ```python

340 # With portable config dict

341 config = {

342 "bot": {

343 "llm": {

344 "$resource": "default",

345 "type": "llm_providers",

346 "temperature": 0.7,

347 },

348 "conversation_storage": {

349 "$resource": "conversations",

350 "type": "databases",

351 },

352 }

353 }

354 bot = await DynaBot.from_environment_aware_config(config)

355

356 # With explicit environment

357 bot = await DynaBot.from_environment_aware_config(

358 config,

359 environment="production",

360 env_dir="configs/environments"

361 )

362

363 # With EnvironmentAwareConfig instance

364 from dataknobs_config import EnvironmentAwareConfig

365 env_config = EnvironmentAwareConfig.load_app("my-bot", ...)

366 bot = await DynaBot.from_environment_aware_config(env_config)

367 ```

368

369 Note:

370 The config should use $resource references for infrastructure:

371 ```yaml

372 bot:

373 llm:

374 $resource: default # Logical name

375 type: llm_providers # Resource type

376 temperature: 0.7 # Behavioral param (portable)

377 ```

378

379 The environment config provides concrete bindings:

380 ```yaml

381 resources:

382 llm_providers:

383 default:

384 provider: openai

385 model: gpt-4

386 api_key: ${OPENAI_API_KEY}

387 ```

388 """

389 from dataknobs_config import EnvironmentAwareConfig, EnvironmentConfig

390

391 # Wrap dict in EnvironmentAwareConfig if needed

392 if isinstance(config, dict):

393 # Load or use provided environment

394 if isinstance(environment, EnvironmentConfig):

395 env_config = environment

396 else:

397 env_config = EnvironmentConfig.load(environment, env_dir)

398

399 config = EnvironmentAwareConfig(

400 config=config,

401 environment=env_config,

402 )

403 elif environment is not None:

404 # Switch environment on existing EnvironmentAwareConfig

405 config = config.with_environment(environment, env_dir)

406

407 # Resolve resources and env vars (late binding happens here)

408 if config_key:

409 resolved = config.resolve_for_build(config_key)

410 else:

411 resolved = config.resolve_for_build()

412

413 # Delegate to existing from_config

414 return await cls.from_config(resolved)

415

416 @staticmethod

417 def get_portable_config(

418 config: EnvironmentAwareConfig | dict[str, Any],

419 ) -> dict[str, Any]:

420 """Extract portable configuration for storage.

421

422 Returns configuration with $resource references intact

423 and environment variables unresolved. This is the config

424 that should be stored in registries or databases for

425 cross-environment portability.

426

427 Args:

428 config: EnvironmentAwareConfig instance or portable dict

429

430 Returns:

431 Portable configuration dictionary

432

433 Example:

434 ```python

435 from dataknobs_config import EnvironmentAwareConfig

436

437 # From EnvironmentAwareConfig

438 env_config = EnvironmentAwareConfig.load_app("my-bot", ...)

439 portable = DynaBot.get_portable_config(env_config)

440

441 # Store portable config in registry

442 await registry.store(bot_id, portable)

443

444 # Dict passes through unchanged

445 portable = DynaBot.get_portable_config({"bot": {...}})

446 ```

447 """

448 # Import here to avoid circular dependency at module level

449 try:

450 from dataknobs_config import EnvironmentAwareConfig

451

452 if isinstance(config, EnvironmentAwareConfig):

453 return config.get_portable_config()

454 except ImportError:

455 pass

456

457 # Dict passes through (assumed already portable)

458 return config

459

460 async def chat(

461 self,

462 message: str,

463 context: BotContext,

464 temperature: float | None = None,

465 max_tokens: int | None = None,

466 stream: bool = False,

467 rag_query: str | None = None,

468 llm_config_overrides: dict[str, Any] | None = None,

469 **kwargs: Any,

470 ) -> str:

471 """Process a chat message.

472

473 Args:

474 message: User message to process

475 context: Bot execution context

476 temperature: Optional temperature override

477 max_tokens: Optional max tokens override

478 stream: Whether to stream the response

479 rag_query: Optional explicit query for knowledge base retrieval.

480 If provided, this is used instead of the message for RAG.

481 Useful when the message contains literal text to analyze

482 (e.g., "Analyze this prompt: [prompt text]") but you want

483 to search for analysis techniques instead.

484 llm_config_overrides: Optional dict to override LLM config fields

485 for this request only. Supported fields: model, temperature,

486 max_tokens, top_p, stop_sequences, seed, options.

487 **kwargs: Additional arguments

488

489 Returns:

490 Bot response as string

491

492 Example:

493 ```python

494 context = BotContext(

495 conversation_id="conv-123",

496 client_id="client-456",

497 user_id="user-789"

498 )

499 response = await bot.chat("Hello!", context)

500

501 # With explicit RAG query

502 response = await bot.chat(

503 "Analyze this: Write a poem about cats",

504 context,

505 rag_query="prompt analysis techniques evaluation"

506 )

507

508 # With LLM config overrides (switch model per-request)

509 response = await bot.chat(

510 "Explain quantum computing",

511 context,

512 llm_config_overrides={"model": "gpt-4-turbo", "temperature": 0.9}

513 )

514 ```

515 """

516 # Apply middleware (before)

517 for mw in self.middleware:

518 if hasattr(mw, "before_message"):

519 await mw.before_message(message, context)

520

521 # Build message with context from memory and knowledge

522 full_message = await self._build_message_with_context(message, rag_query=rag_query)

523

524 # Get or create conversation manager

525 manager = await self._get_or_create_conversation(context)

526

527 # Add user message

528 await manager.add_message(content=full_message, role="user")

529

530 # Update memory

531 if self.memory:

532 await self.memory.add_message(message, role="user")

533

534 # Generate response

535 if self.reasoning_strategy:

536 response = await self.reasoning_strategy.generate(

537 manager=manager,

538 llm=self.llm,

539 tools=list(self.tool_registry),

540 temperature=temperature or self.default_temperature,

541 max_tokens=max_tokens or self.default_max_tokens,

542 llm_config_overrides=llm_config_overrides,

543 )

544 else:

545 response = await manager.complete(

546 llm_config_overrides=llm_config_overrides,

547 temperature=temperature or self.default_temperature,

548 max_tokens=max_tokens or self.default_max_tokens,

549 )

550

551 # Extract response content

552 response_content = response.content if hasattr(response, "content") else str(response)

553

554 # Update memory

555 if self.memory:

556 await self.memory.add_message(response_content, role="assistant")

557

558 # Apply middleware (after)

559 for mw in self.middleware:

560 if hasattr(mw, "after_message"):

561 await mw.after_message(response, context)

562

563 return response_content

564

565 async def stream_chat(

566 self,

567 message: str,

568 context: BotContext,

569 temperature: float | None = None,

570 max_tokens: int | None = None,

571 rag_query: str | None = None,

572 llm_config_overrides: dict[str, Any] | None = None,

573 **kwargs: Any,

574 ) -> AsyncGenerator[str, None]:

575 """Stream chat response token by token.

576

577 Similar to chat() but yields response chunks as they are generated,

578 providing better UX for interactive applications.

579

580 Args:

581 message: User message to process

582 context: Bot execution context

583 temperature: Optional temperature override

584 max_tokens: Optional max tokens override

585 rag_query: Optional explicit query for knowledge base retrieval.

586 If provided, this is used instead of the message for RAG.

587 llm_config_overrides: Optional dict to override LLM config fields

588 for this request only. Supported fields: model, temperature,

589 max_tokens, top_p, stop_sequences, seed, options.

590 **kwargs: Additional arguments passed to LLM

591

592 Yields:

593 Response text chunks as strings

594

595 Example:

596 ```python

597 context = BotContext(

598 conversation_id="conv-123",

599 client_id="client-456",

600 user_id="user-789"

601 )

602

603 # Stream and display in real-time

604 async for chunk in bot.stream_chat("Explain quantum computing", context):

605 print(chunk, end="", flush=True)

606 print() # Newline after streaming

607

608 # Accumulate response

609 full_response = ""

610 async for chunk in bot.stream_chat("Hello!", context):

611 full_response += chunk

612

613 # With LLM config overrides

614 async for chunk in bot.stream_chat(

615 "Explain quantum computing",

616 context,

617 llm_config_overrides={"model": "gpt-4-turbo"}

618 ):

619 print(chunk, end="", flush=True)

620 ```

621

622 Note:

623 Conversation history is automatically updated after streaming completes.

624 The reasoning_strategy is not supported with streaming - use chat() instead.

625 """

626 # Apply middleware (before)

627 for mw in self.middleware:

628 if hasattr(mw, "before_message"):

629 await mw.before_message(message, context)

630

631 # Build message with context from memory and knowledge

632 full_message = await self._build_message_with_context(message, rag_query=rag_query)

633

634 # Get or create conversation manager

635 manager = await self._get_or_create_conversation(context)

636

637 # Add user message

638 await manager.add_message(content=full_message, role="user")

639

640 # Update memory

641 if self.memory:

642 await self.memory.add_message(message, role="user")

643

644 # Stream response (reasoning_strategy not supported for streaming)

645 full_response_chunks: list[str] = []

646 streaming_error: Exception | None = None

647

648 try:

649 async for chunk in manager.stream_complete(

650 llm_config_overrides=llm_config_overrides,

651 temperature=temperature or self.default_temperature,

652 max_tokens=max_tokens or self.default_max_tokens,

653 **kwargs,

654 ):

655 full_response_chunks.append(chunk.delta)

656 yield chunk.delta

657 except Exception as e:

658 streaming_error = e

659 # Call on_error middleware

660 for mw in self.middleware:

661 if hasattr(mw, "on_error"):

662 await mw.on_error(e, message, context)

663 # Re-raise to inform the caller

664 raise

665

666 # Only update memory and run post_stream middleware on success

667 if streaming_error is None:

668 complete_response = "".join(full_response_chunks)

669

670 # Update memory with complete response

671 if self.memory:

672 await self.memory.add_message(complete_response, role="assistant")

673

674 # Apply post_stream middleware hook (provides both message and response)

675 for mw in self.middleware:

676 if hasattr(mw, "post_stream"):

677 await mw.post_stream(message, complete_response, context)

678

679 async def get_conversation(self, conversation_id: str) -> Any:

680 """Retrieve conversation history.

681

682 This method fetches the complete conversation state including all messages,

683 metadata, and the message tree structure. Useful for displaying conversation

684 history, debugging, analytics, or exporting conversations.

685

686 Args:

687 conversation_id: Unique identifier of the conversation to retrieve

688

689 Returns:

690 ConversationState object containing the full conversation history,

691 or None if the conversation does not exist

692

693 Example:

694 ```python

695 # Retrieve a conversation

696 conv_state = await bot.get_conversation("conv-123")

697

698 # Access messages

699 messages = conv_state.message_tree

700

701 # Access metadata

702 print(conv_state.metadata)

703 ```

704

705 See Also:

706 - clear_conversation(): Clear/delete a conversation

707 - chat(): Add messages to a conversation

708 """

709 return await self.conversation_storage.load_conversation(conversation_id)

710

711 async def clear_conversation(self, conversation_id: str) -> bool:

712 """Clear a conversation's history.

713

714 This method removes the conversation from both persistent storage and the

715 internal cache. The next chat() call with this conversation_id will start

716 a fresh conversation. Useful for:

717

718 - Implementing "start over" functionality

719 - Privacy/data deletion requirements

720 - Testing and cleanup

721 - Resetting conversation context

722

723 Args:

724 conversation_id: Unique identifier of the conversation to clear

725

726 Returns:

727 True if the conversation was deleted, False if it didn't exist

728

729 Example:

730 ```python

731 # Clear a conversation

732 deleted = await bot.clear_conversation("conv-123")

733

734 if deleted:

735 print("Conversation deleted")

736 else:

737 print("Conversation not found")

738

739 # Next chat will start fresh

740 response = await bot.chat("Hello!", context)

741 ```

742

743 Note:

744 This operation is permanent and cannot be undone. The conversation

745 cannot be recovered after deletion.

746

747 See Also:

748 - get_conversation(): Retrieve conversation before clearing

749 - chat(): Will create new conversation after clearing

750 """

751 # Remove from cache if present

752 if conversation_id in self._conversation_managers:

753 del self._conversation_managers[conversation_id]

754

755 # Delete from storage

756 return await self.conversation_storage.delete_conversation(conversation_id)

757

758 async def close(self) -> None:

759 """Close the bot and clean up resources.

760

761 This method closes the LLM provider, conversation storage backend,

762 and releases associated resources like HTTP connections and database

763 connections. Should be called when the bot is no longer needed,

764 especially in testing or when creating temporary bot instances.

765

766 Example:

767 ```python

768 bot = await DynaBot.from_config(config)

769 try:

770 response = await bot.chat("Hello", context)

771 finally:

772 await bot.close()

773 ```

774

775 Note:

776 After calling close(), the bot should not be used for further operations.

777 Create a new bot instance if needed.

778 """

779 # Close LLM provider

780 if self.llm and hasattr(self.llm, 'close'):

781 await self.llm.close()

782

783 # Close conversation storage backend

784 if self.conversation_storage and hasattr(self.conversation_storage, 'backend'):

785 backend = self.conversation_storage.backend

786 if backend and hasattr(backend, 'close'):

787 await backend.close()

788

789 # Close knowledge base (releases embedding provider HTTP sessions)

790 if self.knowledge_base and hasattr(self.knowledge_base, 'close'):

791 await self.knowledge_base.close()

792

793 # Close memory store

794 if self.memory and hasattr(self.memory, 'close'):

795 await self.memory.close()

796

797 async def __aenter__(self) -> Self:

798 """Async context manager entry.

799

800 Returns:

801 Self for use in async with statement

802 """

803 return self

804

805 async def __aexit__(

806 self,

807 exc_type: type[BaseException] | None,

808 exc_val: BaseException | None,

809 exc_tb: TracebackType | None,

810 ) -> None:

811 """Async context manager exit - ensures cleanup.

812

813 Args:

814 exc_type: Exception type if an exception occurred

815 exc_val: Exception value if an exception occurred

816 exc_tb: Exception traceback if an exception occurred

817 """

818 await self.close()

819

820 async def _get_or_create_conversation(

821 self, context: BotContext

822 ) -> ConversationManager:

823 """Get or create conversation manager for context.

824

825 Args:

826 context: Bot execution context

827

828 Returns:

829 ConversationManager instance

830 """

831 conv_id = context.conversation_id

832

833 # Check cache

834 if conv_id in self._conversation_managers:

835 return self._conversation_managers[conv_id]

836

837 # Try to resume existing conversation

838 try:

839 manager = await ConversationManager.resume(

840 conversation_id=conv_id,

841 llm=self.llm,

842 prompt_builder=self.prompt_builder,

843 storage=self.conversation_storage,

844 )

845 except Exception:

846 # Create new conversation with specified conversation_id

847 from dataknobs_llm.conversations import ConversationNode, ConversationState

848 from dataknobs_llm.llm.base import LLMMessage

849 from dataknobs_structures.tree import Tree

850

851 metadata = {

852 "client_id": context.client_id,

853 "user_id": context.user_id,

854 **context.session_metadata,

855 }

856

857 # Create initial state with specified conversation_id

858 # Start with empty root node (will be replaced by system prompt if provided)

859 root_message = LLMMessage(role="system", content="")

860 root_node = ConversationNode(

861 message=root_message,

862 node_id="",

863 )

864 tree = Tree(root_node)

865 state = ConversationState(

866 conversation_id=conv_id, # Use the conversation_id from context

867 message_tree=tree,

868 current_node_id="",

869 metadata=metadata,

870 )

871

872 # Create manager with pre-initialized state

873 manager = ConversationManager(

874 llm=self.llm,

875 prompt_builder=self.prompt_builder,

876 storage=self.conversation_storage,

877 state=state,

878 metadata=metadata,

879 )

880

881 # Add system prompt if specified (either as template name or inline content)

882 if self.system_prompt_name:

883 # Use template name - will be rendered by prompt builder

884 await manager.add_message(

885 prompt_name=self.system_prompt_name,

886 role="system",

887 )

888 elif self.system_prompt_content:

889 # Use inline content - pass RAG configs if available

890 await manager.add_message(

891 content=self.system_prompt_content,

892 role="system",

893 rag_configs=self.system_prompt_rag_configs,

894 include_rag=bool(self.system_prompt_rag_configs),

895 )

896

897 # Cache manager

898 self._conversation_managers[conv_id] = manager

899 return manager

900

901 async def _build_message_with_context(

902 self,

903 message: str,

904 rag_query: str | None = None,

905 ) -> str:

906 """Build message with knowledge and memory context.

907

908 Args:

909 message: Original user message

910 rag_query: Optional explicit query for knowledge base retrieval.

911 If provided, this is used instead of the message for RAG.

912

913 Returns:

914 Message augmented with context

915 """

916 contexts = []

917

918 # Add knowledge context

919 if self.knowledge_base:

920 # Use explicit rag_query if provided, otherwise use message

921 search_query = rag_query if rag_query else message

922 kb_results = await self.knowledge_base.query(search_query, k=5)

923 if kb_results:

924 # Use format_context if available (new RAG utilities)

925 if hasattr(self.knowledge_base, "format_context"):

926 kb_context = self.knowledge_base.format_context(

927 kb_results, wrap_in_tags=True

928 )

929 contexts.append(kb_context)

930 else:

931 # Fallback to legacy formatting

932 formatted_chunks = []

933 for i, r in enumerate(kb_results, 1):

934 text = r["text"]

935 source = r.get("source", "")

936 heading = r.get("heading_path", "")

937

938 chunk_text = f"[{i}] {heading}\n{text}"

939 if source:

940 chunk_text += f"\n(Source: {source})"

941 formatted_chunks.append(chunk_text)

942

943 kb_context = "\n\n---\n\n".join(formatted_chunks)

944 contexts.append(f"<knowledge_base>\n{kb_context}\n</knowledge_base>")

945

946 # Add memory context

947 if self.memory:

948 mem_results = await self.memory.get_context(message)

949 if mem_results:

950 mem_context = "\n\n".join([r["content"] for r in mem_results])

951 contexts.append(f"<conversation_history>\n{mem_context}\n</conversation_history>")

952

953 # Build full message with clear separation

954 if contexts:

955 context_section = "\n\n".join(contexts)

956 return f"{context_section}\n\n<question>\n{message}\n</question>"

957 return message

958

959 @staticmethod

960 def _resolve_tool(tool_config: dict[str, Any] | str, config: dict[str, Any]) -> Any | None:

961 """Resolve tool from configuration.

962

963 Supports two patterns:

964 1. Direct class instantiation: {"class": "module.ToolClass", "params": {...}}

965 2. XRef resolution: "xref:tools[tool_name]" or {"xref": "tools[tool_name]"}

966

967 Args:

968 tool_config: Tool configuration (dict or string xref)

969 config: Full bot configuration for xref resolution

970

971 Returns:

972 Tool instance or None if resolution fails

973

974 Example:

975 # Direct instantiation

976 tool_config = {

977 "class": "my_tools.CalculatorTool",

978 "params": {"precision": 2}

979 }

980

981 # XRef to pre-defined tool

982 tool_config = "xref:tools[calculator]"

983 # Requires config to have:

984 # {

985 # "tool_definitions": {

986 # "calculator": {

987 # "class": "my_tools.CalculatorTool",

988 # "params": {}

989 # }

990 # }

991 # }

992 """

993 import importlib

994 import logging

995

996 logger = logging.getLogger(__name__)

997

998 try:

999 # Handle xref string format

1000 if isinstance(tool_config, str):

1001 if tool_config.startswith("xref:"):

1002 # Parse xref (e.g., "xref:tools[calculator]")

1003 # Extract the reference name

1004 import re

1005

1006 match = re.match(r"xref:tools\[([^\]]+)\]", tool_config)

1007 if not match:

1008 logger.error(f"Invalid xref format: {tool_config}")

1009 return None

1010

1011 tool_name = match.group(1)

1012

1013 # Look up in tool_definitions

1014 tool_definitions = config.get("tool_definitions", {})

1015 if tool_name not in tool_definitions:

1016 logger.error(

1017 f"Tool definition not found: {tool_name}. "

1018 f"Available: {list(tool_definitions.keys())}"

1019 )

1020 return None

1021

1022 # Recursively resolve the referenced config

1023 return DynaBot._resolve_tool(tool_definitions[tool_name], config)

1024 else:

1025 logger.error(f"String tool config must be xref format: {tool_config}")

1026 return None

1027

1028 # Handle dict with xref key

1029 if isinstance(tool_config, dict) and "xref" in tool_config:

1030 return DynaBot._resolve_tool(tool_config["xref"], config)

1031

1032 # Handle dict with class key (direct instantiation)

1033 if isinstance(tool_config, dict) and "class" in tool_config:

1034 class_path = tool_config["class"]

1035 params = tool_config.get("params", {})

1036

1037 # Import the tool class

1038 module_path, class_name = class_path.rsplit(".", 1)

1039 module = importlib.import_module(module_path)

1040 tool_class = getattr(module, class_name)

1041

1042 # Instantiate the tool

1043 tool = tool_class(**params)

1044

1045 # Validate it's a Tool instance

1046 from dataknobs_llm.tools import Tool

1047

1048 if not isinstance(tool, Tool):

1049 logger.error(

1050 f"Resolved class {class_path} is not a Tool instance: {type(tool)}"

1051 )

1052 return None

1053

1054 logger.info(f"Successfully loaded tool: {tool.name} ({class_path})")

1055 return tool

1056 else:

1057 logger.error(

1058 f"Invalid tool config format. Expected dict with 'class' or 'xref' key, "

1059 f"or xref string. Got: {type(tool_config)}"

1060 )

1061 return None

1062

1063 except ImportError as e:

1064 logger.error(f"Failed to import tool class: {e}")

1065 return None

1066 except AttributeError as e:

1067 logger.error(f"Failed to find tool class: {e}")

1068 return None

1069 except Exception as e:

1070 logger.error(f"Failed to instantiate tool: {e}")

1071 return None

1072

1073 @staticmethod

1074 def _create_middleware(config: dict[str, Any]) -> Any | None:

1075 """Create middleware from configuration.

1076

1077 Args:

1078 config: Middleware configuration

1079

1080 Returns:

1081 Middleware instance or None

1082 """

1083 try:

1084 import importlib

1085

1086 module_path, class_name = config["class"].rsplit(".", 1)

1087 module = importlib.import_module(module_path)

1088 middleware_class = getattr(module, class_name)

1089 return middleware_class(**config.get("params", {}))

1090 except Exception:

1091 return None