240 lines
9.9 KiB
Python
240 lines
9.9 KiB
Python
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()
|