fix: query will be None if the query_prompt_template not exists (#11031)

Signed-off-by: -LAN- <laipz8200@outlook.com>
langgenius · Nov 24, 2024 · 60b5dac · 60b5dac
1 parent 8565c18
commit 60b5dac
Showing 1 changed file with 8 additions and 3 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,15 @@ 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
+            if query is None and (
+                query_variable := self.graph_runtime_state.variable_pool.get(
+                    (SYSTEM_VARIABLE_NODE_ID, SystemVariableKey.QUERY)
+                )
+            ):
+                query = query_variable.text
 
             prompt_messages, stop = self._fetch_prompt_messages(
                 user_query=query,