This is an automated email from the ASF dual-hosted git repository. xtsong pushed a commit to branch main in repository https://gitbox.apache.org/repos/asf/flink-agents.git
commit 8a9ad0ffb5f3042d57772d2f409cdb0ce8158ec2 Author: WenjinXie <[email protected]> AuthorDate: Mon Jan 5 10:02:05 2026 +0800 [plan] Use sensory memory for tool call context. --- .../flink/agents/plan/actions/ChatModelAction.java | 30 +++++++++++----------- .../flink_agents/plan/actions/chat_model_action.py | 22 ++++++++-------- 2 files changed, 26 insertions(+), 26 deletions(-) diff --git a/plan/src/main/java/org/apache/flink/agents/plan/actions/ChatModelAction.java b/plan/src/main/java/org/apache/flink/agents/plan/actions/ChatModelAction.java index 5465ccf..2b153ae 100644 --- a/plan/src/main/java/org/apache/flink/agents/plan/actions/ChatModelAction.java +++ b/plan/src/main/java/org/apache/flink/agents/plan/actions/ChatModelAction.java @@ -67,12 +67,12 @@ public class ChatModelAction { (BaseChatModelSetup) ctx.getResource(model, ResourceType.CHAT_MODEL); ChatMessage response = chatModel.chat(messages, Map.of()); - MemoryObject stm = ctx.getShortTermMemory(); + MemoryObject sensoryMem = ctx.getSensoryMemory(); if (!response.getToolCalls().isEmpty()) { Map<UUID, Object> toolCallContext; - if (stm.isExist(TOOL_CALL_CONTEXT)) { - toolCallContext = (Map<UUID, Object>) stm.get(TOOL_CALL_CONTEXT).getValue(); + if (sensoryMem.isExist(TOOL_CALL_CONTEXT)) { + toolCallContext = (Map<UUID, Object>) sensoryMem.get(TOOL_CALL_CONTEXT).getValue(); } else { toolCallContext = new HashMap<>(); } @@ -84,32 +84,32 @@ public class ChatModelAction { messageContext.add(response); toolCallContext.put(initialRequestId, messageContext); - stm.set(TOOL_CALL_CONTEXT, toolCallContext); + sensoryMem.set(TOOL_CALL_CONTEXT, toolCallContext); ToolRequestEvent toolRequestEvent = new ToolRequestEvent(model, response.getToolCalls()); Map<UUID, Object> toolRequestEventContext; - if (stm.isExist(TOOL_REQUEST_EVENT_CONTEXT)) { + if (sensoryMem.isExist(TOOL_REQUEST_EVENT_CONTEXT)) { toolRequestEventContext = - (Map<UUID, Object>) stm.get(TOOL_REQUEST_EVENT_CONTEXT).getValue(); + (Map<UUID, Object>) sensoryMem.get(TOOL_REQUEST_EVENT_CONTEXT).getValue(); } else { toolRequestEventContext = new HashMap<>(); } toolRequestEventContext.put( toolRequestEvent.getId(), Map.of(INITIAL_REQUEST_ID, initialRequestId, MODEL, model)); - stm.set(TOOL_REQUEST_EVENT_CONTEXT, toolRequestEventContext); + sensoryMem.set(TOOL_REQUEST_EVENT_CONTEXT, toolRequestEventContext); ctx.sendEvent(toolRequestEvent); } else { // clean tool call context - if (stm.isExist(TOOL_CALL_CONTEXT)) { + if (sensoryMem.isExist(TOOL_CALL_CONTEXT)) { Map<UUID, Object> toolCallContext = - (Map<UUID, Object>) stm.get(TOOL_CALL_CONTEXT).getValue(); + (Map<UUID, Object>) sensoryMem.get(TOOL_CALL_CONTEXT).getValue(); if (toolCallContext.containsKey(initialRequestId)) { toolCallContext.remove(initialRequestId); - stm.set(TOOL_CALL_CONTEXT, toolCallContext); + sensoryMem.set(TOOL_CALL_CONTEXT, toolCallContext); } } @@ -131,7 +131,7 @@ public class ChatModelAction { @SuppressWarnings("unchecked") public static void processChatRequestOrToolResponse(Event event, RunnerContext ctx) throws Exception { - MemoryObject stm = ctx.getShortTermMemory(); + MemoryObject sensoryMem = ctx.getSensoryMemory(); if (event instanceof ChatRequestEvent) { ChatRequestEvent chatRequestEvent = (ChatRequestEvent) event; chat( @@ -144,19 +144,19 @@ public class ChatModelAction { UUID toolRequestId = toolResponseEvent.getRequestId(); // get tool request context from memory Map<UUID, Object> toolRequestEventContext = - (Map<UUID, Object>) stm.get(TOOL_REQUEST_EVENT_CONTEXT).getValue(); + (Map<UUID, Object>) sensoryMem.get(TOOL_REQUEST_EVENT_CONTEXT).getValue(); Map<String, Object> context = (Map<String, Object>) toolRequestEventContext.get(toolRequestId); UUID initialRequestId = (UUID) context.get(INITIAL_REQUEST_ID); String model = (String) context.get(MODEL); toolRequestEventContext.remove(toolRequestId); - stm.set(TOOL_REQUEST_EVENT_CONTEXT, toolRequestEventContext); + sensoryMem.set(TOOL_REQUEST_EVENT_CONTEXT, toolRequestEventContext); Map<String, ToolResponse> responses = toolResponseEvent.getResponses(); Map<String, Boolean> success = toolResponseEvent.getSuccess(); // get tool call context Map<UUID, Object> toolCallContext = - (Map<UUID, Object>) stm.get(TOOL_CALL_CONTEXT).getValue(); + (Map<UUID, Object>) sensoryMem.get(TOOL_CALL_CONTEXT).getValue(); // update tool call context List<ChatMessage> messages = new ArrayList<>((List<ChatMessage>) toolCallContext.get(initialRequestId)); @@ -185,7 +185,7 @@ public class ChatModelAction { } toolCallContext.put(initialRequestId, messages); // overwrite tool call context - stm.set(TOOL_CALL_CONTEXT, toolCallContext); + sensoryMem.set(TOOL_CALL_CONTEXT, toolCallContext); chat(initialRequestId, model, messages, ctx); } else { diff --git a/python/flink_agents/plan/actions/chat_model_action.py b/python/flink_agents/plan/actions/chat_model_action.py index bb74a5f..6db3908 100644 --- a/python/flink_agents/plan/actions/chat_model_action.py +++ b/python/flink_agents/plan/actions/chat_model_action.py @@ -53,7 +53,7 @@ def chat( # TODO: support async execution of chat. response = chat_model.chat(messages) - short_term_memory = ctx.short_term_memory + sensory_memory = ctx.sensory_memory # generate tool request event according tool calls in response if len(response.tool_calls) > 0: @@ -66,7 +66,7 @@ def chat( # to store and remove the specific tool context directly. # save tool call context - tool_call_context = short_term_memory.get(_TOOL_CALL_CONTEXT) + tool_call_context = sensory_memory.get(_TOOL_CALL_CONTEXT) if not tool_call_context: tool_call_context = {} if initial_request_id not in tool_call_context: @@ -74,7 +74,7 @@ def chat( # append response to tool call context tool_call_context[initial_request_id].append(response) # update tool call context - short_term_memory.set(_TOOL_CALL_CONTEXT, tool_call_context) + sensory_memory.set(_TOOL_CALL_CONTEXT, tool_call_context) tool_request_event = ToolRequestEvent( model=model, @@ -89,16 +89,16 @@ def chat( "initial_request_id": initial_request_id, "model": model, } - short_term_memory.set(_TOOL_REQUEST_EVENT_CONTEXT, tool_request_event_context) + sensory_memory.set(_TOOL_REQUEST_EVENT_CONTEXT, tool_request_event_context) ctx.send_event(tool_request_event) # if there is no tool call generated, return chat response directly else: # clear tool call context related to specific request id - tool_call_context = short_term_memory.get(_TOOL_CALL_CONTEXT) + tool_call_context = sensory_memory.get(_TOOL_CALL_CONTEXT) if tool_call_context and initial_request_id in tool_call_context: tool_call_context.pop(initial_request_id) - short_term_memory.set(_TOOL_CALL_CONTEXT, tool_call_context) + sensory_memory.set(_TOOL_CALL_CONTEXT, tool_call_context) ctx.send_event( ChatResponseEvent( request_id=initial_request_id, @@ -113,7 +113,7 @@ def process_chat_request_or_tool_response(event: Event, ctx: RunnerContext) -> N Internally, this action will use short term memory to save the tool call context, which is a dict mapping request id to chat messages. """ - short_term_memory = ctx.short_term_memory + sensory_memory = ctx.sensory_memory if isinstance(event, ChatRequestEvent): chat( initial_request_id=event.id, @@ -126,18 +126,18 @@ def process_chat_request_or_tool_response(event: Event, ctx: RunnerContext) -> N request_id = event.request_id # get correspond tool request event context - tool_request_event_context = short_term_memory.get(_TOOL_REQUEST_EVENT_CONTEXT) + tool_request_event_context = sensory_memory.get(_TOOL_REQUEST_EVENT_CONTEXT) initial_request_id = tool_request_event_context[request_id][ "initial_request_id" ] model = tool_request_event_context[request_id]["model"] # clear tool request event context tool_request_event_context.pop(request_id) - short_term_memory.set(_TOOL_REQUEST_EVENT_CONTEXT, tool_request_event_context) + sensory_memory.set(_TOOL_REQUEST_EVENT_CONTEXT, tool_request_event_context) responses = event.responses # update tool call context - tool_call_context = short_term_memory.get(_TOOL_CALL_CONTEXT) + tool_call_context = sensory_memory.get(_TOOL_CALL_CONTEXT) for id, response in responses.items(): tool_call_context[initial_request_id].append( ChatMessage( @@ -148,7 +148,7 @@ def process_chat_request_or_tool_response(event: Event, ctx: RunnerContext) -> N else {}, ) ) - short_term_memory.set(_TOOL_CALL_CONTEXT, tool_call_context) + sensory_memory.set(_TOOL_CALL_CONTEXT, tool_call_context) chat( initial_request_id=initial_request_id,
