diff --git a/README.md b/README.md
index 1140c14..7ea2a34 100644
--- a/README.md
+++ b/README.md
@@ -41,6 +41,11 @@ If you find AttackGen useful, please consider starring the repository on GitHub.
## Releases
+### v0.14 (unreleased)
+| What's new? | Why is it useful? |
+| ----------- | ----------------- |
+| Assistant Refines the Detection & Response Output | - Purple-Team Editing: The AttackGen Assistant can now refine the **Detection & Response** narrative, not just the scenario. When a scenario is generated with the π£ Purple-team narrative toggle (Threat Group / Custom pages), an **Editing:** selector appears on the Assistant page so the chat can target either the scenario or the defender's walkthrough β previously the Assistant only had the scenario text.
- Combined Editing Mode: A third **Scenario + Detection & Response** target refines both together, so a cross-cutting change β a different industry, threat actor, technique or timeline β is applied consistently across the attacker's scenario and the defender's narrative rather than being made twice and drifting apart.
- Deterministic Facts Untouched: Only the LLM narrative is a refinement target; the local STIX detection join is never LLM-edited. Each editing mode keeps its own chat history, and the combined pass is tagged `purple_team_narrative` in LangSmith. |
+
### v0.13.1
| What's new? | Why is it useful? |
| ----------- | ----------------- |
diff --git a/core/scenario_page.py b/core/scenario_page.py
index 5f6003c..9c7d47b 100644
--- a/core/scenario_page.py
+++ b/core/scenario_page.py
@@ -346,9 +346,15 @@ def _persist_and_render(
st.session_state[layer_key] = layer_payload
st.session_state[filename_key] = download_name
st.session_state[defense_key] = defense_state
- # Cross-page handoff for the AttackGen Assistant chat page.
+ # Cross-page handoff for the AttackGen Assistant chat page. The defense
+ # narrative rides along so the Assistant can refine it too; set it
+ # unconditionally (None when there's no narrative) so a stale one from an
+ # earlier generation can't linger after a plain-scenario regen.
st.session_state["last_scenario"] = True
st.session_state["last_scenario_text"] = cleaned
+ st.session_state["last_defense_narrative"] = (
+ defense_state.get("narrative_md") if defense_state else None
+ )
_render_result(
page_id=page_id,
diff --git "a/pages/4_\360\237\222\254_AttackGen_Assistant.py" "b/pages/4_\360\237\222\254_AttackGen_Assistant.py"
index 0f3b22b..dd4c53d 100644
--- "a/pages/4_\360\237\222\254_AttackGen_Assistant.py"
+++ "b/pages/4_\360\237\222\254_AttackGen_Assistant.py"
@@ -11,13 +11,34 @@
restore_from_query_params()
-SYSTEM_PROMPT = (
+SCENARIO_SYSTEM_PROMPT = (
"You are an AI assistant that helps users update and ask questions about their incident "
"response scenario. Only respond to questions or requests relating to the scenario, or "
"incident response testing in general. Format your responses using proper Markdown syntax "
"with headers, bullet points, and formatting for readability."
)
+DEFENSE_SYSTEM_PROMPT = (
+ "You are an AI assistant that helps users refine the purple-team Detection & Response "
+ "narrative that accompanies their incident response scenario. The narrative walks the "
+ "scenario from the defender's side β detection opportunities, log sources, and response "
+ "actions, stage by stage. Only respond to questions or requests relating to the detection "
+ "and response of this scenario, or purple-team testing in general. Keep your suggestions "
+ "grounded in the scenario provided for reference. Format your responses using proper "
+ "Markdown syntax with headers, bullet points, and formatting for readability."
+)
+
+BOTH_SYSTEM_PROMPT = (
+ "You are an AI assistant that helps users refine an incident response scenario and its "
+ "accompanying purple-team Detection & Response narrative together. When a requested change "
+ "affects both β a different threat actor, industry, technique, or timeline β apply it "
+ "consistently across the two so the attacker's scenario and the defender's walkthrough stay "
+ "aligned, and make clear which output each part of your response applies to. Only respond to "
+ "questions or requests relating to the scenario, its detection and response, or incident "
+ "response testing in general. Format your responses using proper Markdown syntax with "
+ "headers, bullet points, and formatting for readability."
+)
+
st.set_page_config(page_title="AttackGen Assistant", page_icon=":speech_balloon:")
inject_emoji_fonts()
@@ -25,87 +46,158 @@
st.markdown("# AttackGen Assistantπ¬", unsafe_allow_html=True)
-if 'last_scenario_text' in st.session_state and st.session_state.get('last_scenario'):
- input_scenario = st.session_state['last_scenario_text']
- with st.expander("Generated Scenario"):
- with st.container(height=400, border=True):
- st.markdown(input_scenario)
+scenario_text = (
+ st.session_state["last_scenario_text"]
+ if st.session_state.get("last_scenario") and "last_scenario_text" in st.session_state
+ else None
+)
+defense_narrative = st.session_state.get("last_defense_narrative")
- chat_container = st.empty()
+if not scenario_text:
+ st.info("No scenario found. Please generate a scenario first.")
+ st.stop()
+
+
+# Pick what to edit. The Detection & Response and combined options only appear
+# when a purple-team narrative was generated alongside the scenario (page 1/2
+# toggle). The combined option refines both together so a change made to one can
+# be carried consistently into the other.
+if defense_narrative:
+ choice = st.radio(
+ "Editing:",
+ ["Scenario", "Detection & Response", "Scenario + Detection & Response"],
+ horizontal=True,
+ key="assistant_target",
+ )
+else:
+ choice = "Scenario"
+target = {
+ "Detection & Response": "defense",
+ "Scenario + Detection & Response": "both",
+}.get(choice, "scenario")
+
+if target == "defense":
+ panels = [("Detection & Response Narrative", defense_narrative)]
+ system_prompt = DEFENSE_SYSTEM_PROMPT
+ greeting = "Hi, I can help you refine the Detection & Response narrative for your scenario."
+ trace_name = "AttackGen Assistant β Detection & Response"
+ trace_tags = ("assistant", "purple_team_narrative")
+elif target == "both":
+ panels = [
+ ("Generated Scenario", scenario_text),
+ ("Detection & Response Narrative", defense_narrative),
+ ]
+ system_prompt = BOTH_SYSTEM_PROMPT
+ greeting = (
+ "Hi, I can help you refine the scenario and its Detection & Response narrative "
+ "together, keeping changes consistent across both."
+ )
+ trace_name = "AttackGen Assistant β Scenario + Detection & Response"
+ trace_tags = ("assistant", "purple_team_narrative")
+else:
+ panels = [("Generated Scenario", scenario_text)]
+ system_prompt = SCENARIO_SYSTEM_PROMPT
+ greeting = "Hi, I can help you update and ask questions about your incident response scenario."
+ trace_name = "AttackGen Assistant"
+ trace_tags = ("assistant",)
- if 'messages' not in st.session_state:
- st.session_state.messages = [
- {"role": "assistant", "content": "Hi, I can help you update and ask questions about your incident response scenario."}
- ]
- with chat_container:
- for message in st.session_state.messages:
- with st.chat_message(message["role"]):
- st.markdown(message["content"])
-
- def generate_response(user_input, chat_history):
- messages = [
- {"role": "system", "content": SYSTEM_PROMPT},
- {
- "role": "user",
- "content": (
- f"Here is the scenario that the user previously generated:\n\n{input_scenario}\n\n"
- f"Chat history:\n{chat_history}\n\nUser: {user_input}"
- ),
- },
- ]
- config = LLMConfig.from_session_state(
- trace_name="AttackGen Assistant",
- trace_tags=("assistant",),
+for label, content in panels:
+ with st.expander(label):
+ with st.container(height=400, border=True):
+ st.markdown(content)
+
+chat_container = st.empty()
+
+# Keep a separate history per target so switching between the scenario and the
+# Detection & Response narrative doesn't feed one artifact's chat into the other.
+messages_key = f"assistant_messages_{target}"
+if messages_key not in st.session_state:
+ st.session_state[messages_key] = [{"role": "assistant", "content": greeting}]
+
+with chat_container:
+ for message in st.session_state[messages_key]:
+ with st.chat_message(message["role"]):
+ st.markdown(message["content"])
+
+
+def generate_response(user_input, chat_history):
+ if target == "both":
+ context = (
+ f"Here is the incident response scenario:\n\n{scenario_text}\n\n"
+ f"Here is the accompanying Detection & Response narrative:\n\n{defense_narrative}\n\n"
+ f"The user wants to refine both together. When a requested change affects both, "
+ f"apply it consistently across them and show the update to each.\n\n"
+ f"Chat history:\n{chat_history}\n\nUser: {user_input}"
+ )
+ elif target == "defense":
+ context = (
+ f"Here is the scenario, for reference:\n\n{scenario_text}\n\n"
+ f"Here is the current Detection & Response narrative the user wants to refine:"
+ f"\n\n{defense_narrative}\n\n"
+ f"Chat history:\n{chat_history}\n\nUser: {user_input}"
)
- raw_chunks: list[str] = []
-
- def _tee(chunks):
- for chunk in chunks:
- raw_chunks.append(chunk)
- yield chunk
-
- try:
- yield from stream_filter_thinking(_tee(call_llm_stream(config, messages)))
- except Exception as e:
- yield f"\n\nAn error occurred while calling the model: {e}"
- st.session_state["_last_assistant_cleaned"] = (
- f"An error occurred while calling the model: {e}"
- )
- return
-
- raw = "".join(raw_chunks)
- thinking, cleaned = clean_model_response(raw)
- if thinking:
- with st.expander("View Model's Reasoning"):
- st.markdown(thinking)
- st.session_state["_last_assistant_cleaned"] = cleaned
-
- if prompt := st.chat_input("Type your message here..."):
- st.session_state.messages.append({"role": "user", "content": prompt})
- with st.chat_message("user"):
- st.markdown(prompt)
+ else:
+ context = (
+ f"Here is the scenario that the user previously generated:\n\n{scenario_text}\n\n"
+ f"Chat history:\n{chat_history}\n\nUser: {user_input}"
+ )
+ messages = [
+ {"role": "system", "content": system_prompt},
+ {"role": "user", "content": context},
+ ]
+ config = LLMConfig.from_session_state(
+ trace_name=trace_name,
+ trace_tags=trace_tags,
+ )
+ raw_chunks: list[str] = []
+
+ def _tee(chunks):
+ for chunk in chunks:
+ raw_chunks.append(chunk)
+ yield chunk
+
+ try:
+ yield from stream_filter_thinking(_tee(call_llm_stream(config, messages)))
+ except Exception as e:
+ yield f"\n\nAn error occurred while calling the model: {e}"
+ st.session_state["_last_assistant_cleaned"] = (
+ f"An error occurred while calling the model: {e}"
+ )
+ return
- with st.chat_message("assistant"):
- history = "\n".join(f"{m['role']}: {m['content']}" for m in st.session_state.messages[:-1])
- st.write_stream(generate_response(prompt, history))
+ raw = "".join(raw_chunks)
+ thinking, cleaned = clean_model_response(raw)
+ if thinking:
+ with st.expander("View Model's Reasoning"):
+ st.markdown(thinking)
+ st.session_state["_last_assistant_cleaned"] = cleaned
- st.session_state.messages.append(
- {"role": "assistant", "content": st.session_state.pop("_last_assistant_cleaned", "")}
+
+if prompt := st.chat_input("Type your message here..."):
+ st.session_state[messages_key].append({"role": "user", "content": prompt})
+ with st.chat_message("user"):
+ st.markdown(prompt)
+
+ with st.chat_message("assistant"):
+ history = "\n".join(
+ f"{m['role']}: {m['content']}" for m in st.session_state[messages_key][:-1]
)
+ st.write_stream(generate_response(prompt, history))
- def clear_conversation():
- st.session_state.messages = [
- {"role": "assistant", "content": "Hi, I can help you update and ask questions about your incident response scenario."}
- ]
- chat_container.empty()
- with chat_container:
- with st.chat_message("assistant"):
- st.markdown(st.session_state.messages[0]["content"])
+ st.session_state[messages_key].append(
+ {"role": "assistant", "content": st.session_state.pop("_last_assistant_cleaned", "")}
+ )
- with st.container():
- if st.button("Clear Conversation", key='clear_button'):
- clear_conversation()
-else:
- st.info("No scenario found. Please generate a scenario first.")
+def clear_conversation():
+ st.session_state[messages_key] = [{"role": "assistant", "content": greeting}]
+ chat_container.empty()
+ with chat_container:
+ with st.chat_message("assistant"):
+ st.markdown(st.session_state[messages_key][0]["content"])
+
+
+with st.container():
+ if st.button("Clear Conversation", key='clear_button'):
+ clear_conversation()
diff --git a/tests/test_scenario_page.py b/tests/test_scenario_page.py
index afba353..816621c 100644
--- a/tests/test_scenario_page.py
+++ b/tests/test_scenario_page.py
@@ -203,6 +203,8 @@ def test_happy_path_calls_llm_cleans_response_and_persists(
# Cross-page handoff for the Assistant page.
assert fake_session_state["last_scenario"] is True
assert fake_session_state["last_scenario_text"] == cleaned
+ # No defense companion here, so nothing for the Assistant to refine there.
+ assert fake_session_state["last_defense_narrative"] is None
# The artifact flag is set.
assert fake_session_state["threat_group_scenario_generated"] is True
@@ -520,6 +522,8 @@ def test_defense_persisted_and_offered_for_download(
state = fake_session_state["threat_group_scenario_defense"]
assert state["narrative_md"] is None
assert "Command and Scripting Interpreter (T1059)" in state["deterministic_md"]
+ # Deterministic-only: no narrative for the Assistant to refine.
+ assert fake_session_state["last_defense_narrative"] is None
md_name = fake_session_state["threat_group_scenario_filename"]
detection_downloads = [
@@ -562,6 +566,8 @@ def test_defense_narrative_makes_second_llm_call_and_persists(
assert state["narrative_md"] and "Detection walkthrough" in state["narrative_md"]
# The combined download carries both the narrative and the reference section.
assert "Detection & Response Reference" in state["download_md"]
+ # The narrative is handed to the Assistant so it can be refined there too.
+ assert fake_session_state["last_defense_narrative"] == state["narrative_md"]
def test_no_defense_download_when_build_defense_returns_none(