Update app.py

app.py

@@ -45,7 +45,7 @@ memory = MemorySaver()
 def ai_assistance(state: State):
     result = llm.invoke(state["messages"])
     return {"messages": state['messages']+[result]}
-
+import time
 def generate_code(state: State):
     user_prompt = state["input"]
 
@@ -66,9 +66,20 @@ def generate_code(state: State):
     User request:"""    
     
     full_prompt = system_prompt + user_prompt
-    code = code_generator.invoke(full_prompt)
+    # code = code_generator.invoke(full_prompt)
     
-    return {"code": str(code)}
+    # return {"code": str(code)}
+    for attempt in range(3):
+        try:
+            code_response = code_generator.invoke(full_prompt)
+            return {"code": str(code_response)}
+        except Exception as e:
+            if "503" in str(e):
+                print(f"[Retry {attempt+1}/3] Together API unavailable (503). Retrying...")
+                time.sleep(2)
+            else:
+                raise e
+    raise Exception("Together API failed after 3 retries.")
 
 def explain_code(state):
     code = state["code"]
@@ -210,9 +221,15 @@ with st.container():
             result = graph.invoke(state_input, config=config)
         
             # If it’s a code response, append it as-is
+            # if result.get("code"):
+            #     st.session_state.chat_history.append(AIMessage(content=result["code"]))
+            # else:
+            #     st.session_state.chat_history.append(result["messages"][-1])
             if result.get("code"):
                 st.session_state.chat_history.append(AIMessage(content=result["code"]))
-            else:
+            elif result.get("explanation"):
+                st.session_state.chat_history.append(AIMessage(content=result["explanation"]))
+            elif result.get("messages"):
                 st.session_state.chat_history.append(result["messages"][-1])
-        
+
             st.rerun()