feat: enabling workflows without rewriting step

QuivrHQ · Jan 20, 2025 · 162e63c · 162e63c
1 parent 9681a9e
commit 162e63c
Showing 1 changed file with 9 additions and 2 deletions.
diff --git a/core/quivr_core/rag/quivr_rag_langgraph.py b/core/quivr_core/rag/quivr_rag_langgraph.py
@@ -615,8 +615,11 @@ async def retrieve(self, state: AgentState) -> AgentState:
         Returns:
             dict: The retrieved chunks
         """
+        if "tasks" in state:
+            tasks = state["tasks"]
+        else:
+            tasks = UserTasks([state["messages"][0].content])
 
-        tasks = state["tasks"]
         if not tasks.has_tasks():
             return {**state}
 
@@ -670,7 +673,11 @@ async def dynamic_retrieve(self, state: AgentState) -> AgentState:
 
         MAX_ITERATIONS = 3
 
-        tasks = state["tasks"]
+        if "tasks" in state:
+            tasks = state["tasks"]
+        else:
+            tasks = UserTasks([state["messages"][0].content])
+
         if not tasks or not tasks.has_tasks():
             return {**state}