diff --git a/llama_toolchain/agentic_system/event_logger.py b/llama_toolchain/agentic_system/event_logger.py index 22d961a10..6b8d034a6 100644 --- a/llama_toolchain/agentic_system/event_logger.py +++ b/llama_toolchain/agentic_system/event_logger.py @@ -6,7 +6,7 @@ from typing import Optional -from llama_models.llama3.api.datatypes import ToolResponseMessage +from llama_models.llama3.api.datatypes import * # noqa: F403 from llama_models.llama3.api.tool_utils import ToolUtils from termcolor import cprint @@ -162,5 +162,19 @@ class EventLogger: color="green", ) + if ( + step_type == StepType.memory_retrieval + and event_type == EventType.step_complete.value + ): + details = event.payload.step_details + content = interleaved_text_media_as_str(details.inserted_context) + content = content[:200] + "..." if len(content) > 200 else content + + yield event, LogEvent( + role=step_type, + content=f"Retrieved context from banks: {details.memory_bank_ids}.\n====\n{content}\n>", + color="cyan", + ) + preivous_event_type = event_type previous_step_type = step_type diff --git a/llama_toolchain/agentic_system/meta_reference/agent_instance.py b/llama_toolchain/agentic_system/meta_reference/agent_instance.py index c559470cc..04560ff7b 100644 --- a/llama_toolchain/agentic_system/meta_reference/agent_instance.py +++ b/llama_toolchain/agentic_system/meta_reference/agent_instance.py @@ -631,6 +631,8 @@ class ChatAgent(ShieldRunnerMixin): for a in attachments ] await self.memory_api.insert_documents(bank.bank_id, documents) + elif session.memory_bank: + bank_ids.append(session.memory_bank.bank_id) if not bank_ids: # this can happen if the per-session memory bank is not yet populated