import streamlit as st from backend.managers.chat_manager import ChatManager from backend.managers.system_prompter import SystemPrompter # ── Agent Mode helpers ──────────────────────────────────────────────────────── def _start_agent(task: str): """Initialise a fresh CodingAgent, start the task, and propose the first action.""" from backend.agent.coding_agent import CodingAgent agent = CodingAgent() agent.start_task(task) action = agent.propose_next_action() st.session_state.coding_agent = agent st.session_state.agent_pending_action = action st.session_state.agent_status = "waiting_approval" st.session_state.agent_log = [] def _approve_action(): """Execute the pending action, append it to the log, then propose the next step.""" agent = st.session_state.coding_agent pending = st.session_state.agent_pending_action result = agent.approve() st.session_state.agent_log.append({ "thought": pending.get("thought", ""), "tool": result["tool"], "arguments": result.get("arguments", {}), "result": result["result"], }) if result["is_done"]: st.session_state.agent_status = "done" st.session_state.agent_pending_action = None else: next_action = agent.propose_next_action() st.session_state.agent_pending_action = next_action st.session_state.agent_status = "waiting_approval" def _reject_action(feedback: str): """Reject the pending action with optional feedback, then replan.""" agent = st.session_state.coding_agent agent.reject(feedback or "Please try a different approach.") next_action = agent.propose_next_action() st.session_state.agent_pending_action = next_action st.session_state.agent_status = "waiting_approval" def _followup_agent(question: str): """Inject a follow-up question into the finished agent and resume the loop.""" agent = st.session_state.coding_agent agent.follow_up(question) action = agent.propose_next_action() st.session_state.agent_pending_action = action st.session_state.agent_status = "waiting_approval" def _reset_agent(): """Reset all agent state back to idle.""" st.session_state.coding_agent = None st.session_state.agent_status = "idle" st.session_state.agent_log = [] st.session_state.agent_pending_action = None # ── Agent Mode UI ───────────────────────────────────────────────────────────── def render_agent_mode(): # Toggle must always be rendered so Streamlit keeps agent_mode=True in session_state st.toggle("Agent Mode", key="agent_mode") agent_status = st.session_state.get("agent_status", "idle") agent_log = st.session_state.get("agent_log", []) # ── Agent Log ──────────────────────────────────────────────────────────── if agent_log: with st.expander(f"Agent Log — {len(agent_log)} step(s) completed", expanded=False): for i, step in enumerate(agent_log): with st.chat_message("assistant"): st.markdown(f"**Step {i + 1} — `{step['tool']}`**") st.caption(f"Thought: {step['thought']}") if step.get("arguments"): st.json(step["arguments"]) result_text = step.get("result", "") if result_text.startswith("ERROR") or result_text.startswith("SYNTAX ERROR"): st.error(result_text) elif result_text.startswith("OK") or result_text.startswith("DONE"): st.success(result_text) else: st.code(result_text, language=None) # ── Idle: task input ────────────────────────────────────────────────────── if agent_status == "idle": task = st.text_area( "Describe what the agent should do:", key="agent_task_input", height=120, placeholder="e.g. Write a function that sorts a list and saves it to sorted.py", ) if st.button("Start Agent", type="primary", use_container_width=True): if task.strip(): with st.spinner("Agent is thinking..."): _start_agent(task.strip()) st.rerun() else: st.warning("Please describe a task first.") # ── Waiting: show proposed action + Approve / Reject ───────────────────── elif agent_status == "waiting_approval": pending = st.session_state.get("agent_pending_action", {}) with st.status("Agent proposes the following step:", expanded=True): st.markdown(f"**Thought:** {pending.get('thought', '—')}") st.markdown(f"**Tool:** `{pending.get('tool', '—')}`") args = pending.get("arguments", {}) if args: if "content" in args: display_args = {k: v for k, v in args.items() if k != "content"} if display_args: st.json(display_args) st.code(args["content"], language="python") else: st.json(args) feedback = st.text_input( "Rejection feedback (optional):", key="agent_reject_feedback", placeholder="e.g. Use a different approach...", ) col1, col2, col3 = st.columns([3, 2, 2]) with col1: if st.button("Approve", type="primary", use_container_width=True): with st.spinner("Executing and planning next step..."): _approve_action() st.rerun() with col2: if st.button("Reject", use_container_width=True): with st.spinner("Agent is replanning..."): _reject_action(feedback) st.rerun() with col3: if st.button("Abort Task", use_container_width=True): _reset_agent() st.rerun() # ── Done ───────────────────────────────────────────────────────────────── elif agent_status == "done": last_result = agent_log[-1]["result"] if agent_log else "" st.success(f"Task completed! {last_result}") st.divider() followup = st.text_area( "Follow-up question or correction:", key="agent_followup_input", height=80, placeholder="e.g. The output is wrong — it should sort descending. Can you fix that?", ) col1, col2 = st.columns(2) with col1: if st.button("Ask Follow-up", type="primary", use_container_width=True): if followup.strip(): with st.spinner("Agent is thinking..."): _followup_agent(followup.strip()) st.rerun() else: st.warning("Please enter a follow-up question first.") with col2: if st.button("New Task", use_container_width=True): _reset_agent() st.rerun() # ── Normal Chat ─────────────────────────────────────────────────────────────── def render_normal_chat(): # Chat history as bubbles for message in st.session_state.chat_history: role = message["role"] if role == "system": continue with st.chat_message(role): st.markdown(message["content"]) # Chat input — Enter to send, no extra button needed user_input = st.chat_input("Type your message here...") if user_input: chat_manager = st.session_state.chat_manager # Inject system prompt on the first message if not chat_manager.get_history(): system_prompt = SystemPrompter.generate_prompt() chat_manager.add_message("system", system_prompt) # Show user message immediately without waiting for response with st.chat_message("user"): st.markdown(user_input) # Show response with spinner while API is called with st.chat_message("assistant"): with st.spinner("Thinking..."): try: ai_response = chat_manager.send_message(user_input) except Exception as e: ai_response = f"Error: {e}" st.markdown(ai_response) st.session_state.chat_history.append({"role": "user", "content": user_input}) st.session_state.chat_history.append({"role": "assistant", "content": ai_response}) st.rerun() # Rendered in the normal flow; JS above clones them to fixed positions # and hides these originals. st.toggle("Agent Mode", key="agent_mode") with st.expander("Settings", expanded=False): st.toggle("Use debug system prompt", key="use_system_prompt", value=True) # ── Entry point ─────────────────────────────────────────────────────────────── def render_chat(): if st.session_state.get("agent_mode", False): st.subheader("Coding Agent") render_agent_mode() else: st.subheader("Chat with AI Assistant") render_normal_chat() if __name__ == "__main__": render_chat()