fix(llm_node): Fetch query from system variable.

Signed-off-by: -LAN- <laipz8200@outlook.com>
langgenius · Nov 24, 2024 · ca1a1b5 · ca1a1b5
1 parent ae3a2cb
commit ca1a1b5
Showing 1 changed file with 7 additions and 4 deletions.
diff --git a/api/core/workflow/nodes/llm/node.py b/api/core/workflow/nodes/llm/node.py
@@ -38,6 +38,7 @@
     ObjectSegment,
     StringSegment,
 )
+from core.workflow.constants import SYSTEM_VARIABLE_NODE_ID
 from core.workflow.entities.node_entities import NodeRunMetadataKey, NodeRunResult
 from core.workflow.entities.variable_entities import VariableSelector
 from core.workflow.entities.variable_pool import VariablePool
@@ -133,11 +134,13 @@ def _run(self) -> NodeRunResult | Generator[NodeEvent | InNodeEvent, None, None]
             # fetch memory
             memory = self._fetch_memory(node_data_memory=self.node_data.memory, model_instance=model_instance)
 
-            # fetch prompt messages
+            query = None
             if self.node_data.memory:
-                query = self.node_data.memory.query_prompt_template
-            else:
-                query = None
+                query_variable = self.graph_runtime_state.variable_pool.get(
+                    (SYSTEM_VARIABLE_NODE_ID, SystemVariableKey.QUERY)
+                )
+                if query_variable:
+                    query = query_variable.text
 
             prompt_messages, stop = self._fetch_prompt_messages(
                 user_query=query,