From 5297aef9e8c6e92987bdddec8e8ee48a9a61cd14 Mon Sep 17 00:00:00 2001 From: Ashwin Bharambe Date: Tue, 21 Jan 2025 21:19:19 -0800 Subject: [PATCH] reuse some variables --- .../inline/agents/meta_reference/agent_instance.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/llama_stack/providers/inline/agents/meta_reference/agent_instance.py b/llama_stack/providers/inline/agents/meta_reference/agent_instance.py index d49c1f749..5b5175cee 100644 --- a/llama_stack/providers/inline/agents/meta_reference/agent_instance.py +++ b/llama_stack/providers/inline/agents/meta_reference/agent_instance.py @@ -396,7 +396,7 @@ class ChatAgent(ShieldRunnerMixin): session_id, documents, input_messages, tool_defs ) - if "builtin::memory" in toolgroups and len(input_messages) > 0: + if MEMORY_GROUP in toolgroups and len(input_messages) > 0: with tracing.span(MEMORY_QUERY_TOOL) as span: step_id = str(uuid.uuid4()) yield AgentTurnResponseStreamChunk( @@ -408,7 +408,7 @@ class ChatAgent(ShieldRunnerMixin): ) ) - args = toolgroup_args.get("builtin::memory", {}) + args = toolgroup_args.get(MEMORY_GROUP, {}) vector_db_ids = args.get("vector_db_ids", []) session_info = await self.storage.get_session_info(session_id) @@ -475,7 +475,7 @@ class ChatAgent(ShieldRunnerMixin): "input", [m.model_dump_json() for m in input_messages] ) span.set_attribute("output", retrieved_context) - span.set_attribute("tool_name", "builtin::memory") + span.set_attribute("tool_name", MEMORY_QUERY_TOOL) if retrieved_context: last_message = input_messages[-1] last_message.context = retrieved_context