From b46d94d87db5622f8e2dfa743a8dfed45b231f8a Mon Sep 17 00:00:00 2001 From: Dinesh Yeduguru Date: Wed, 8 Jan 2025 18:53:32 -0800 Subject: [PATCH] do not pass memory tools to inference --- .../inline/agents/meta_reference/agent_instance.py | 13 +++---------- .../providers/inline/tool_runtime/memory/memory.py | 2 +- 2 files changed, 4 insertions(+), 11 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 2cd86bcaa..24448a28f 100644 --- a/llama_stack/providers/inline/agents/meta_reference/agent_instance.py +++ b/llama_stack/providers/inline/agents/meta_reference/agent_instance.py @@ -487,20 +487,13 @@ class ChatAgent(ShieldRunnerMixin): stop_reason = None with tracing.span("inference") as span: - - def is_memory_group(tool): - memory_tool_group = tool_to_group.get(MEMORY_QUERY_TOOL, None) - has_memory_tool = MEMORY_QUERY_TOOL in tool_defs - return ( - has_memory_tool - and tool_to_group.get(tool.tool_name, None) != memory_tool_group - ) - async for chunk in await self.inference_api.chat_completion( self.agent_config.model, input_messages, tools=[ - tool for tool in tool_defs.values() if not is_memory_group(tool) + tool + for tool in tool_defs.values() + if tool_to_group.get(tool.tool_name, None) != MEMORY_GROUP ], tool_prompt_format=self.agent_config.tool_prompt_format, stream=True, diff --git a/llama_stack/providers/inline/tool_runtime/memory/memory.py b/llama_stack/providers/inline/tool_runtime/memory/memory.py index f46b37510..fe6325abb 100644 --- a/llama_stack/providers/inline/tool_runtime/memory/memory.py +++ b/llama_stack/providers/inline/tool_runtime/memory/memory.py @@ -60,7 +60,7 @@ class MemoryToolRuntimeImpl(ToolsProtocolPrivate, ToolRuntime): description="Retrieve context from memory", parameters=[ ToolParameter( - name="input_messages", + name="messages", description="The input messages to search for", parameter_type="array", ),