Update app.py

app.py

@@ -62,6 +62,7 @@ class AgentState(TypedDict):
     timings: Dict[str, float]
 
 def agent(template: str, state: AgentState, agent_key: str, timing_label: str):
+    st.write(f'🛠 Running agent: {agent_key}')
     start = time.time()
     model, tokenizer = load_agent_model(**AGENT_MODEL_CONFIG[agent_key])
 
@@ -71,7 +72,11 @@ def agent(template: str, state: AgentState, agent_key: str, timing_label: str):
         or state["messages"][-1]["content"]
     )
     prompt = template.format(user_input=latest_input, html=state.get("html", ""), final_prompt=state.get("final_prompt", ""))
+    st.write(f'📤 Prompt for {agent_key}:', prompt)
+
     response = call_model(prompt, model, tokenizer)
+    st.write(f'📥 Response from {agent_key}:', response[:500])
+
     state["messages"].append({"role": agent_key, "content": response})
     state["timings"][timing_label] = time.time() - start
     gc.collect()
@@ -150,6 +155,7 @@ def main():
         with st.spinner("Agents working..."):
             final = generate_ui(prompt, max_iter)
             st.success("✅ UI Generated")
+            st.write("🧠 Final state:", final)
             st.subheader("🔍 Output HTML")
             st.components.v1.html(final["html"], height=600, scrolling=True)
             st.subheader("🧠 Agent Messages")
@@ -157,4 +163,4 @@ def main():
                 st.markdown(f"**{msg['role'].title()}**:\n```\n{msg['content']}\n```")
 
 if __name__ == "__main__":
-    main()
+    main()