[fel] record LLM results to memory

Author: loveTsong

framework/fel/java/fel-flow/pom.xml

@@ -74,5 +74,10 @@
             <artifactId>assertj-core</artifactId>
             <scope>test</scope>
         </dependency>
+        <dependency>
+            <groupId>org.mockito</groupId>
+            <artifactId>mockito-core</artifactId>
+            <scope>test</scope>
+        </dependency>
     </dependencies>
 </project>

framework/fel/java/fel-flow/src/main/java/modelengine/fel/engine/operators/models/LlmEmitter.java

@@ -8,12 +8,15 @@
 
 import modelengine.fel.core.chat.ChatMessage;
 import modelengine.fel.core.chat.Prompt;
+import modelengine.fel.core.chat.support.HumanMessage;
+import modelengine.fel.core.memory.Memory;
 import modelengine.fel.engine.util.StateKey;
 import modelengine.fit.waterflow.bridge.fitflow.FitBoundedEmitter;
 import modelengine.fit.waterflow.domain.context.FlowSession;
 import modelengine.fitframework.flowable.Publisher;
 import modelengine.fitframework.inspection.Validation;
 import modelengine.fitframework.util.ObjectUtils;
+import modelengine.fitframework.util.StringUtils;
 
 /**
  * 流式模型发射器。
@@ -26,6 +29,8 @@ public class LlmEmitter<O extends ChatMessage> extends FitBoundedEmitter<O, Chat
 
     private final ChatChunk chunkAcc = new ChatChunk();
     private final StreamingConsumer<ChatMessage, ChatMessage> consumer;
+    private final Memory memory;
+    private final ChatMessage question;
 
     /**
      * 初始化 {@link LlmEmitter}。
@@ -38,6 +43,9 @@ public LlmEmitter(Publisher<O> publisher, Prompt prompt, FlowSession session) {
         super(publisher, data -> data);
         Validation.notNull(session, "The session cannot be null.");
         this.consumer = ObjectUtils.nullIf(session.getInnerState(StateKey.STREAMING_CONSUMER), EMPTY_CONSUMER);
+        this.memory = session.getInnerState(StateKey.HISTORY);
+        this.question =
+                ObjectUtils.getIfNull(session.getInnerState(StateKey.HISTORY_INPUT), () -> getDefaultQuestion(prompt));
     }
 
     @Override
@@ -46,4 +54,21 @@ public void emit(ChatMessage data, FlowSession trans) {
         this.chunkAcc.merge(data);
         this.consumer.accept(this.chunkAcc, data);
     }
+
+    @Override
+    public void complete() {
+        if (this.memory != null && this.chunkAcc.toolCalls().isEmpty()) {
+            this.memory.add(this.question);
+            this.memory.add(this.chunkAcc);
+        }
+        super.complete();
+    }
+
+    private static ChatMessage getDefaultQuestion(Prompt prompt) {
+        int size = prompt.messages().size();
+        if (size == 0) {
+            return new HumanMessage(StringUtils.EMPTY);
+        }
+        return prompt.messages().get(size - 1);
+    }
 }

framework/fel/java/fel-flow/src/test/java/modelengine/fel/engine/operators/models/LlmEmitterTest.java

@@ -0,0 +1,77 @@
+/*---------------------------------------------------------------------------------------------
+ *  Copyright (c) 2025 Huawei Technologies Co., Ltd. All rights reserved.
+ *  This file is a part of the ModelEngine Project.
+ *  Licensed under the MIT License. See License.txt in the project root for license information.
+ *--------------------------------------------------------------------------------------------*/
+
+package modelengine.fel.engine.operators.models;
+
+import modelengine.fel.core.chat.ChatMessage;
+import modelengine.fel.core.chat.Prompt;
+import modelengine.fel.core.chat.support.AiMessage;
+import modelengine.fel.core.chat.support.ChatMessages;
+import modelengine.fel.core.memory.Memory;
+import modelengine.fel.core.tool.ToolCall;
+import modelengine.fel.engine.util.StateKey;
+import modelengine.fit.waterflow.domain.context.FlowSession;
+import modelengine.fitframework.flowable.Choir;
+import modelengine.fitframework.util.StringUtils;
+
+import org.junit.jupiter.api.Test;
+import org.mockito.ArgumentCaptor;
+import org.mockito.Mockito;
+
+import java.util.Arrays;
+import java.util.Collections;
+import java.util.List;
+
+import static org.junit.jupiter.api.Assertions.*;
+
+/**
+ * 表示 {@link LlmEmitter} 的测试。
+ *
+ * @author 宋永坦
+ * @since 2025-07-05
+ */
+class LlmEmitterTest {
+    @Test
+    void shouldAddMemoryWhenCompleteGivenLlmOutput() {
+        String output = "data1";
+        Prompt prompt = ChatMessages.fromList(Collections.emptyList());
+        Choir<ChatMessage> dataSource = Choir.create(emitter -> {
+            emitter.emit(new AiMessage(output));
+            emitter.complete();
+        });
+        FlowSession flowSession = new FlowSession();
+        Memory mockMemory = Mockito.mock(Memory.class);
+        ArgumentCaptor<ChatMessage> captor = ArgumentCaptor.forClass(ChatMessage.class);
+        Mockito.doNothing().when(mockMemory).add(captor.capture());
+        flowSession.setInnerState(StateKey.HISTORY, mockMemory);
+
+        LlmEmitter<ChatMessage> llmEmitter = new LlmEmitter<>(dataSource, prompt, flowSession);
+        llmEmitter.start(flowSession);
+
+        List<ChatMessage> captured = captor.getAllValues();
+        assertEquals(2, captured.size());
+        assertEquals(StringUtils.EMPTY, captured.get(0).text());
+        assertEquals(output, captured.get(1).text());
+    }
+
+    @Test
+    void shouldNotAddMemoryWhenCompleteGivenLlmToolCallOutput() {
+        String output = "data1";
+        Prompt prompt = ChatMessages.fromList(Collections.emptyList());
+        Choir<ChatMessage> dataSource = Choir.create(emitter -> {
+            emitter.emit(new AiMessage(output, Arrays.asList(ToolCall.custom().id("id1").build())));
+            emitter.complete();
+        });
+        FlowSession flowSession = new FlowSession();
+        Memory mockMemory = Mockito.mock(Memory.class);
+        flowSession.setInnerState(StateKey.HISTORY, mockMemory);
+
+        LlmEmitter<ChatMessage> llmEmitter = new LlmEmitter<>(dataSource, prompt, flowSession);
+        llmEmitter.start(flowSession);
+
+        Mockito.verify(mockMemory, Mockito.times(0)).add(Mockito.any());
+    }
+}