578 lines
24 KiB
Python
578 lines
24 KiB
Python
"""Chat view — renders both the normal chat interface and the Coding Agent mode."""
|
|
|
|
import asyncio
|
|
import json
|
|
from pathlib import Path
|
|
|
|
import streamlit as st
|
|
from backend.managers.chat_manager import ChatManager
|
|
from backend.managers.system_prompter import SystemPrompter
|
|
from backend.managers.search_manager import SearchManager
|
|
from backend.agent.coding_agent import CodingAgent
|
|
from backend.managers.debug_logger import get_logger
|
|
|
|
logger = get_logger(__name__)
|
|
|
|
|
|
# ── Agent Mode helpers ────────────────────────────────────────────────────────
|
|
def _run_async(coro):
|
|
"""Execute an async coroutine from synchronous Streamlit code.
|
|
|
|
Streamlit runs in a synchronous context, but the CodingAgent uses async
|
|
methods (for MCP tool calls). This helper bridges the gap by reusing an
|
|
already-running event loop when one exists, or creating a new one otherwise.
|
|
|
|
Args:
|
|
coro: The coroutine to run.
|
|
|
|
Returns:
|
|
The return value of the coroutine.
|
|
"""
|
|
# REVIEW: asyncio.get_running_loop() always raises RuntimeError in a Streamlit context;
|
|
# the try branch is dead code. The except branch always runs.
|
|
try:
|
|
# Reuse the loop that is already running (e.g. inside pytest-asyncio).
|
|
loop = asyncio.get_running_loop()
|
|
except RuntimeError:
|
|
# No running loop in this thread — create a fresh one.
|
|
loop = asyncio.new_event_loop()
|
|
asyncio.set_event_loop(loop)
|
|
return loop.run_until_complete(coro)
|
|
|
|
def _start_agent(task: str):
|
|
"""Create a new CodingAgent, feed it the task, and propose the first action.
|
|
Stores the agent and its state in session_state so Streamlit can reference
|
|
them across reruns without losing progress.
|
|
"""
|
|
logger.info("Starting coding agent.")
|
|
agent = CodingAgent()
|
|
agent.start_task(task)
|
|
action = _run_async(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, log it, then immediately propose the next step."""
|
|
agent = st.session_state.coding_agent
|
|
pending = st.session_state.agent_pending_action
|
|
|
|
result = _run_async(agent.approve())
|
|
logger.info("Approve action and propose next step.")
|
|
|
|
# Append a record to the log so the user can review every completed step.
|
|
st.session_state.agent_log.append({
|
|
"thought": pending.get("thought", ""),
|
|
"tool": result["tool"],
|
|
"arguments": result.get("arguments", {}),
|
|
"result": result["result"],
|
|
})
|
|
|
|
if result["is_done"]:
|
|
# Agent called the "done" tool — task is fully complete.
|
|
st.session_state.agent_status = "done"
|
|
st.session_state.agent_pending_action = None
|
|
else:
|
|
next_action = _run_async(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 feedback so the agent replans.
|
|
The pending action is discarded; the agent receives the user's feedback and
|
|
proposes a different approach on the next call to propose_next_action().
|
|
"""
|
|
logger.info("Rejecting proposed action.")
|
|
agent = st.session_state.coding_agent
|
|
agent.reject(feedback or "Please try a different approach.")
|
|
next_action = _run_async(agent.propose_next_action())
|
|
st.session_state.agent_pending_action = next_action
|
|
st.session_state.agent_status = "waiting_approval"
|
|
|
|
|
|
def _followup_agent(question: str):
|
|
"""Continue a finished task by injecting a follow-up question and resuming the loop."""
|
|
logger.info("Asking follow up question")
|
|
agent = st.session_state.coding_agent
|
|
agent.follow_up(question)
|
|
action = _run_async(agent.propose_next_action())
|
|
st.session_state.agent_pending_action = action
|
|
st.session_state.agent_status = "waiting_approval"
|
|
|
|
|
|
def _reset_agent():
|
|
"""Clear all agent state and return to the idle (task input) screen."""
|
|
logger.info("Resetting Agent")
|
|
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_arguments(args: dict):
|
|
if not args:
|
|
return
|
|
|
|
with st.expander("📦 Arguments", expanded=False):
|
|
|
|
if args.get("path"):
|
|
st.markdown("##### 📁 Path")
|
|
st.code(args["path"])
|
|
|
|
if args.get("dir_path"):
|
|
st.markdown("##### 🌳 Directory")
|
|
st.code(args["dir_path"])
|
|
|
|
if args.get("query"):
|
|
st.markdown("##### 🔎 Query")
|
|
st.code(args["query"])
|
|
|
|
if args.get("url"):
|
|
st.markdown("##### 🌐 URL")
|
|
st.code(args["url"])
|
|
|
|
if args.get("content"):
|
|
st.markdown("##### 📝 Content")
|
|
st.code(args["content"])
|
|
|
|
if args.get("code"):
|
|
st.markdown("##### 🐍 Python Code")
|
|
st.code(args["code"], language="python")
|
|
|
|
if args.get("max_results") is not None:
|
|
st.markdown("##### 📊 Max Results")
|
|
st.code(str(args["max_results"]))
|
|
|
|
known_keys = {
|
|
"path",
|
|
"dir_path",
|
|
"query",
|
|
"content",
|
|
"url",
|
|
"code",
|
|
"max_results",
|
|
}
|
|
|
|
extra_args = {
|
|
k: v for k, v in args.items()
|
|
if k not in known_keys
|
|
}
|
|
|
|
if extra_args:
|
|
st.markdown("##### ⚙️ Other")
|
|
st.code(
|
|
json.dumps(extra_args, indent=2),
|
|
language="json"
|
|
)
|
|
|
|
def render_agent_mode():
|
|
"""Render the step-by-step agent UI.
|
|
|
|
Three distinct screens based on agent_status:
|
|
- "idle" → task description input + Start button
|
|
- "waiting_approval" → show proposed action, Approve / Reject / Abort
|
|
- "done" → success message, follow-up input, New Task button
|
|
"""
|
|
logger.info("Agent mode.")
|
|
# The toggle must always render 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 ────────────────────────────────────────────────────────────
|
|
# Collapsed by default so it doesn't clutter the UI during active tasks.
|
|
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"):
|
|
_render_arguments(step["arguments"])
|
|
result_text = step.get("result", "")
|
|
# Colour the result based on whether the tool succeeded or failed.
|
|
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:
|
|
_render_arguments(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 helpers ───────────────────────────────────────────────────────
|
|
|
|
def _detect_task_type(user_input: str) -> str:
|
|
"""Infer the task type from keywords in the user message."""
|
|
lower = user_input.lower()
|
|
if any(kw in lower for kw in ("error", "bug", "fix", "crash", "exception", "debug")):
|
|
return "debug"
|
|
if any(kw in lower for kw in ("explain", "what does", "how does", "why")):
|
|
return "explain"
|
|
if any(kw in lower for kw in ("optimize", "improve", "faster", "refactor", "clean")):
|
|
return "optimize"
|
|
return "default"
|
|
|
|
def _set_system_prompt(chat_manager: ChatManager, user_input: str) -> None:
|
|
"""Compute and inject the system prompt before every message.
|
|
|
|
Uses the custom prompt from Settings if set; otherwise generates one based
|
|
on the detected task type and active file context. Updates the existing
|
|
system message in-place so the history stays a single-system-message list.
|
|
"""
|
|
custom = st.session_state.get("custom_system_prompt", "").strip()
|
|
if custom:
|
|
prompt = custom
|
|
else:
|
|
prompt = SystemPrompter.generate_prompt(
|
|
user_message=user_input,
|
|
file_context=_build_file_context(),
|
|
task_type=_detect_task_type(user_input),
|
|
)
|
|
|
|
if chat_manager.chat_history and chat_manager.chat_history[0]["role"] == "system":
|
|
chat_manager.chat_history[0]["content"] = prompt
|
|
else:
|
|
chat_manager.chat_history.insert(0, {"role": "system", "content": prompt})
|
|
|
|
|
|
def _build_file_context() -> dict | None:
|
|
"""Return file context for the system prompt if a file is open and context is enabled.
|
|
|
|
Reads from files_content cache first; falls back to FileManager if the file
|
|
has not been loaded into the editor yet.
|
|
"""
|
|
if not st.session_state.get("include_file_context", True):
|
|
return None
|
|
active_file = st.session_state.get("active_file")
|
|
if not active_file:
|
|
return None
|
|
content = st.session_state.get("files_content", {}).get(active_file, "")
|
|
if not content:
|
|
try:
|
|
from backend.managers.file_manager import FileManager
|
|
fm = FileManager()
|
|
content = fm.read_file(Path(active_file)) or ""
|
|
except Exception:
|
|
return None
|
|
return {"name": Path(active_file).name, "content": content}
|
|
|
|
|
|
@st.dialog("Clear Chat")
|
|
def _clear_chat_dialog():
|
|
"""Confirmation dialog before wiping the full conversation history."""
|
|
st.warning("All messages will be deleted. This cannot be undone.")
|
|
col1, col2 = st.columns(2)
|
|
with col1:
|
|
if st.button("Clear", type="primary", use_container_width=True):
|
|
st.session_state.chat_manager.clear_history()
|
|
st.session_state.chat_history = []
|
|
st.rerun()
|
|
with col2:
|
|
if st.button("Cancel", use_container_width=True):
|
|
st.rerun()
|
|
|
|
|
|
# ── Normal Chat ───────────────────────────────────────────────────────────────
|
|
|
|
def _render_search_panel():
|
|
"""Render the collapsible web search panel above the chat history.
|
|
|
|
Stores results in session_state.search_results so they are automatically
|
|
injected as context into the next message the user sends.
|
|
"""
|
|
search_results = st.session_state.get("search_results", [])
|
|
label = f"🔍 Web Search ({len(search_results)} result{'s' if len(search_results) != 1 else ''} active)" if search_results else "🔍 Web Search"
|
|
|
|
with st.expander(label, expanded=False):
|
|
col_input, col_btn = st.columns([5, 1])
|
|
with col_input:
|
|
query = st.text_input(
|
|
"Search query",
|
|
key="search_query_input",
|
|
placeholder="e.g. Python asyncio best practices",
|
|
label_visibility="collapsed",
|
|
)
|
|
with col_btn:
|
|
search_clicked = st.button("Search", use_container_width=True)
|
|
|
|
if search_clicked and query.strip():
|
|
with st.spinner("Searching..."):
|
|
sm = SearchManager()
|
|
results = sm.perform_search(query.strip())
|
|
if results:
|
|
st.session_state.search_results = results
|
|
st.rerun()
|
|
else:
|
|
st.warning("No results found.")
|
|
|
|
# Display active results with a clear button.
|
|
if search_results:
|
|
st.caption("Results will be injected as context into your next message.")
|
|
for r in search_results:
|
|
st.markdown(f"**{r['title']}** \n{r['snippet']} \n[{r['url']}]({r['url']})")
|
|
st.divider()
|
|
if st.button("Clear search results", use_container_width=True):
|
|
st.session_state.search_results = []
|
|
st.rerun()
|
|
|
|
|
|
def render_normal_chat():
|
|
"""Render the standard multi-turn chat interface.
|
|
|
|
On the first message the system prompt is injected into the history,
|
|
including any active search results as context.
|
|
Each subsequent message appends to the same conversation so the AI retains
|
|
full context throughout the session. If search results are active when the
|
|
user sends a message, they are prepended to that message as a context block.
|
|
"""
|
|
logger.info("Chat mode")
|
|
chat_manager: ChatManager = st.session_state.chat_manager
|
|
|
|
# Apply model/token overrides from the Settings panel before any API call.
|
|
if st.session_state.get("selected_model"):
|
|
chat_manager.model = st.session_state.selected_model
|
|
if "chat_max_tokens" in st.session_state:
|
|
chat_manager.max_tokens = st.session_state.chat_max_tokens
|
|
|
|
# Consume a debug message forwarded from the editor's "Debug with AI" button.
|
|
pending_debug = st.session_state.pop("pending_debug_message", None)
|
|
if pending_debug:
|
|
_set_system_prompt(chat_manager, pending_debug)
|
|
with st.spinner("Sending debug info to AI..."):
|
|
try:
|
|
ai_response = chat_manager.send_message(pending_debug)
|
|
except Exception as e:
|
|
ai_response = f"Error: {e}"
|
|
st.session_state.chat_history.append({"role": "user", "content": pending_debug})
|
|
st.session_state.chat_history.append(
|
|
{"role": "assistant", "content": ai_response}
|
|
)
|
|
st.rerun()
|
|
return
|
|
|
|
# Replay the conversation history as chat bubbles (skip system messages).
|
|
for message in st.session_state.chat_history:
|
|
if message["role"] == "system":
|
|
continue
|
|
with st.chat_message(message["role"]):
|
|
st.markdown(message["content"])
|
|
|
|
# Chat input — Enter to send, no extra button needed.
|
|
# Supports /search <query> and /search clear as special commands.
|
|
user_input = st.chat_input("Type a message or /search <query>...")
|
|
if user_input:
|
|
stripped = user_input.strip()
|
|
|
|
# ── /search command ───────────────────────────────────────────────────
|
|
if stripped.lower().startswith("/search"):
|
|
arg = stripped[len("/search"):].strip()
|
|
|
|
with st.chat_message("user"):
|
|
st.markdown(stripped)
|
|
|
|
if arg.lower() == "clear" or arg == "":
|
|
# /search clear (or bare /search) — remove active results.
|
|
st.session_state.search_results = []
|
|
with st.chat_message("assistant"):
|
|
st.markdown("Search context cleared.")
|
|
st.session_state.chat_history.append({"role": "user", "content": stripped})
|
|
st.session_state.chat_history.append({"role": "assistant", "content": "Search context cleared."})
|
|
else:
|
|
# /search <query> — run search and store results in context.
|
|
with st.chat_message("assistant"):
|
|
with st.spinner(f'Searching for "{arg}"...'):
|
|
sm = SearchManager()
|
|
results = sm.perform_search(arg)
|
|
|
|
if results:
|
|
st.session_state.search_results = results
|
|
summary = f"Found {len(results)} result(s) for **{arg}**. They are now in context for this chat session.\n\n"
|
|
for i, r in enumerate(results, 1):
|
|
summary += f"**{i}. [{r['title']}]({r['url']})** \n{r['snippet']}\n\n"
|
|
st.markdown(summary)
|
|
response_text = summary
|
|
else:
|
|
msg = f'No results found for "{arg}".'
|
|
st.warning(msg)
|
|
response_text = msg
|
|
|
|
st.session_state.chat_history.append({"role": "user", "content": stripped})
|
|
st.session_state.chat_history.append({"role": "assistant", "content": response_text})
|
|
|
|
st.rerun()
|
|
return
|
|
|
|
# ── Normal chat message ───────────────────────────────────────────────
|
|
search_results = st.session_state.get("search_results", [])
|
|
|
|
# 5g — System-prompt logic: inject on first message, update on file change.
|
|
_set_system_prompt(chat_manager, user_input)
|
|
|
|
# If search results are active, prepend them as a context block so the
|
|
# AI can reference them regardless of where in the conversation we are.
|
|
if search_results:
|
|
context_block = "<search_context>\n"
|
|
for r in search_results:
|
|
context_block += (
|
|
f"Title: {r['title']}\n"
|
|
f"URL: {r['url']}\n"
|
|
f"Snippet: {r['snippet']}\n\n"
|
|
)
|
|
context_block += "</search_context>\n\n"
|
|
message_to_send = context_block + user_input
|
|
else:
|
|
message_to_send = user_input
|
|
|
|
# Show the original user text in the UI (not the context-enriched version).
|
|
with st.chat_message("user"):
|
|
st.markdown(user_input)
|
|
|
|
with st.chat_message("assistant"):
|
|
with st.spinner("Thinking..."):
|
|
try:
|
|
ai_response = chat_manager.send_message(message_to_send)
|
|
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()
|
|
|
|
# 5d — Clear Chat opens a confirmation dialog instead of deleting immediately.
|
|
if st.button("🗑️ Clear Chat"):
|
|
_clear_chat_dialog()
|
|
|
|
# Toggle to switch to Agent Mode (render_normal_chat and render_agent_mode are
|
|
# mutually exclusive, so the same key here causes no DuplicateWidgetID conflict).
|
|
st.toggle("Agent Mode", key="agent_mode")
|
|
|
|
# 5h — Settings expander: file context toggle, model, token limit, custom prompt.
|
|
with st.expander("⚙️ Settings", expanded=False):
|
|
st.toggle("Include current file as context", key="include_file_context", value=True)
|
|
|
|
st.divider()
|
|
|
|
default_model = chat_manager.model or ""
|
|
model_options = [default_model] if default_model else []
|
|
for m in ["claude-3-5-sonnet-20241022", "claude-3-haiku-20240307", "gpt-4o", "gpt-4o-mini"]:
|
|
if m not in model_options:
|
|
model_options.append(m)
|
|
st.selectbox("Model", model_options, key="selected_model")
|
|
|
|
st.slider(
|
|
"Max Response Tokens",
|
|
min_value=256, max_value=8000,
|
|
value=chat_manager.max_tokens,
|
|
step=256, key="chat_max_tokens",
|
|
)
|
|
|
|
st.divider()
|
|
|
|
st.text_area(
|
|
"Custom System Prompt (overrides default if set)",
|
|
key="custom_system_prompt",
|
|
height=120,
|
|
placeholder="Leave empty to use the default assistant prompt with optional file context.",
|
|
)
|
|
|
|
|
|
# ── Entry point ───────────────────────────────────────────────────────────────
|
|
|
|
def render_chat():
|
|
"""Top-level chat view — switches between Agent Mode and normal 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()
|