Update app.py

app.py

@@ -184,7 +184,7 @@ def discover_attachment(task_id: str, api_url: str) -> Optional[str]:
 class AgentState(TypedDict):
     question: str
     current_step: str
-    final_answer: str
+    answer: str
     history: List[Dict[str, str]]
     needs_search: bool
     search_query: str
@@ -192,7 +192,7 @@ class AgentState(TypedDict):
     logs: Dict[str, Any]
     file_url: str
     code_blocks: List[Dict[str, str]]
-    next_step: str  # Add this field for routing
+    next_step: str
 
 # -------------------------
 # BasicAgent implementation
@@ -249,7 +249,7 @@ class BasicAgent:
         state: AgentState = {
             "question": question,
             "current_step": "route",
-            "final_answer": "",
+            "answer": "",
             "history": [],
             "needs_search": False,
             "search_query": "",
@@ -266,11 +266,9 @@ class BasicAgent:
         
         try:
             final_state = self.workflow.invoke(state)
-            answer = final_state.get("final_answer", "")
-            if not answer:
-                print("Warning: No answer generated in final state")
-                return "No answer could be generated"
-            return answer
+            if "answer" not in final_state:
+                raise ValueError(f"No 'answer' in state: keys are {list(final_state.keys())}")
+            return final_state["answer"]
         except Exception as e:
             print(f"Error in workflow execution: {str(e)}")
             return f"Error processing question: {str(e)}"
@@ -332,16 +330,18 @@ class BasicAgent:
             
             print(f"Generated answer: {answer}")
             return {
-                "final_answer": answer,
+                "answer": answer,
                 "current_step": "done",
-                "next_step": "done"
+                "next_step": "done",
+                "history": state["history"] + [{"role": "assistant", "content": answer}]
             }
         except Exception as e:
             print(f"\nError processing image {state['file_url']}: {str(e)}")
             return {
-                "final_answer": f"Error processing image: {str(e)}",
+                "answer": f"Error processing image: {str(e)}",
                 "current_step": "done",
-                "next_step": "done"
+                "next_step": "done",
+                "history": state["history"] + [{"role": "assistant", "content": f"Error: {str(e)}"}]
             }
 
     def _process_video(self, state: AgentState) -> Dict[str, Any]:
@@ -358,16 +358,18 @@ class BasicAgent:
             
             print(f"Generated answer: {answer}")
             return {
-                "final_answer": answer,
+                "answer": answer,
                 "current_step": "done",
-                "next_step": "done"
+                "next_step": "done",
+                "history": state["history"] + [{"role": "assistant", "content": answer}]
             }
         except Exception as e:
             print(f"\nError processing video {state['file_url']}: {str(e)}")
             return {
-                "final_answer": f"Error processing video: {str(e)}",
+                "answer": f"Error processing video: {str(e)}",
                 "current_step": "done",
-                "next_step": "done"
+                "next_step": "done",
+                "history": state["history"] + [{"role": "assistant", "content": f"Error: {str(e)}"}]
             }
 
     def _process_spreadsheet(self, state: AgentState) -> Dict[str, Any]:
@@ -384,16 +386,18 @@ class BasicAgent:
             
             print(f"Generated answer: {answer}")
             return {
-                "final_answer": answer,
+                "answer": answer,
                 "current_step": "done",
-                "next_step": "done"
+                "next_step": "done",
+                "history": state["history"] + [{"role": "assistant", "content": answer}]
             }
         except Exception as e:
             print(f"\nError processing spreadsheet {state['file_url']}: {str(e)}")
             return {
-                "final_answer": f"Error processing spreadsheet: {str(e)}",
+                "answer": f"Error processing spreadsheet: {str(e)}",
                 "current_step": "done",
-                "next_step": "done"
+                "next_step": "done",
+                "history": state["history"] + [{"role": "assistant", "content": f"Error: {str(e)}"}]
             }
 
     def _process_python(self, state: AgentState) -> Dict[str, Any]:
@@ -410,16 +414,18 @@ class BasicAgent:
             
             print(f"Generated answer: {answer}")
             return {
-                "final_answer": answer,
+                "answer": answer,
                 "current_step": "done",
-                "next_step": "done"
+                "next_step": "done",
+                "history": state["history"] + [{"role": "assistant", "content": answer}]
             }
         except Exception as e:
             print(f"\nError processing Python file {state['file_url']}: {str(e)}")
             return {
-                "final_answer": f"Error processing Python file: {str(e)}",
+                "answer": f"Error processing Python file: {str(e)}",
                 "current_step": "done",
-                "next_step": "done"
+                "next_step": "done",
+                "history": state["history"] + [{"role": "assistant", "content": f"Error: {str(e)}"}]
             }
 
     def _process_text(self, state: AgentState) -> Dict[str, Any]:
@@ -438,16 +444,18 @@ QUESTION:
             answer = self._safe_parse(raw)
             print(f"Generated answer: {answer}")
             return {
-                "final_answer": answer,
+                "answer": answer,
                 "current_step": "done",
-                "next_step": "done"
+                "next_step": "done",
+                "history": state["history"] + [{"role": "assistant", "content": answer}]
             }
         except Exception as e:
             print(f"\nLLM Error in answer generation: {str(e)}")
             return {
-                "final_answer": "I encountered an error while generating the answer.",
+                "answer": "I encountered an error while generating the answer.",
                 "current_step": "done",
-                "next_step": "done"
+                "next_step": "done",
+                "history": state["history"] + [{"role": "assistant", "content": "Error generating answer"}]
             }
 
     def _build_workflow(self) -> Graph:
@@ -477,9 +485,9 @@ QUESTION:
             }
         )
         
-        # Set finish points for all tool nodes
+        # Add edges from each processing node to END
         for node in ["process_image", "process_video", "process_spreadsheet", "process_python", "process_text"]:
-            sg.set_finish_point(node)
+            sg.add_edge(node, "END")
             
         return sg.compile()