Update app.py

app.py

@@ -44,7 +44,7 @@ memory = MemorySaver()
 # Define node
 def ai_assistance(state: State):
     result = llm.invoke(state["messages"])
-    return {"messages": [result]}
+    return {"messages": state['messages']+[result]}
 
 def generate_code(state: State):
     user_prompt = state["input"]
@@ -68,7 +68,7 @@ def generate_code(state: State):
     full_prompt = system_prompt + user_prompt
     code = code_generator.invoke(full_prompt)
     
-    return {"input": user_prompt, "code": code.content}
+    return {"code": code.content}
 
 def router(state: State):
     last_msg = state["messages"][-1].content.lower()