312 lines
11 KiB
Python
312 lines
11 KiB
Python
# -*- coding: utf-8 -*-
|
|
"""
|
|
Streamlit 实时 Agent 协作平台
|
|
功能：
|
|
1. 实时展示每个 Agent 的状态和动作
|
|
2. 自动保存生成的文件到 workspace/
|
|
3. 简单稳定的代码结构
|
|
"""
|
|
import streamlit as st
|
|
import os
|
|
from pathlib import Path
|
|
from datetime import datetime
|
|
import time
|
|
|
|
try:
|
|
from autogen import AssistantAgent, UserProxyAgent, GroupChat, GroupChatManager
|
|
AUTOGEN_AVAILABLE = True
|
|
except ImportError:
|
|
AUTOGEN_AVAILABLE = False
|
|
|
|
# 添加项目根目录到路径
|
|
import sys
|
|
sys.path.insert(0, str(Path(__file__).parent.parent))
|
|
|
|
from config.llm_config import get_llm_config, PM_PROMPT, QA_PROMPT, DEV_PROMPT, ORCH_PROMPT
|
|
|
|
# 页面配置
|
|
st.set_page_config(page_title="多 Agent 协作平台", page_icon="🤖", layout="wide")
|
|
|
|
# Agent 配置
|
|
AGENTS = {
|
|
"PM_Agent": {"name": "产品经理", "avatar": "📋", "color": "blue"},
|
|
"QA_Agent": {"name": "测试工程师", "avatar": "✅", "color": "green"},
|
|
"Dev_Agent": {"name": "开发工程师", "avatar": "💻", "color": "orange"},
|
|
"Orchestrator": {"name": "协调器", "avatar": "🎯", "color": "purple"},
|
|
"User_Proxy": {"name": "用户代理", "avatar": "👤", "color": "gray"}
|
|
}
|
|
|
|
def init_state():
|
|
"""初始化 session state"""
|
|
if "messages" not in st.session_state:
|
|
st.session_state.messages = []
|
|
if "running" not in st.session_state:
|
|
st.session_state.running = False
|
|
if "current_agent" not in st.session_state:
|
|
st.session_state.current_agent = None
|
|
if "agent_counts" not in st.session_state:
|
|
st.session_state.agent_counts = {k: 0 for k in AGENTS}
|
|
|
|
def add_message(agent, content, task=""):
|
|
"""添加消息"""
|
|
msg = {
|
|
"agent": agent,
|
|
"content": content,
|
|
"task": task,
|
|
"time": datetime.now().strftime("%H:%M:%S")
|
|
}
|
|
st.session_state.messages.append(msg)
|
|
st.session_state.agent_counts[agent] = st.session_state.agent_counts.get(agent, 0) + 1
|
|
st.session_state.current_agent = agent
|
|
|
|
def show_agent_status():
|
|
"""显示 Agent 状态"""
|
|
cols = st.columns(len(AGENTS))
|
|
for i, (agent_key, info) in enumerate(AGENTS.items()):
|
|
with cols[i]:
|
|
is_active = st.session_state.current_agent == agent_key
|
|
count = st.session_state.agent_counts.get(agent_key, 0)
|
|
|
|
border_color = info["color"]
|
|
bg_color = "#e8f5e9" if is_active else "white"
|
|
|
|
st.markdown(f"""
|
|
<div style='
|
|
padding: 15px;
|
|
border-radius: 10px;
|
|
border: 3px {"solid" if is_active else "dashed"} {border_color};
|
|
background: {bg_color};
|
|
text-align: center;
|
|
'>
|
|
<div style='font-size: 2.5rem;'>{info["avatar"]}</div>
|
|
<div style='font-weight: bold; margin: 5px 0;'>{info["name"]}</div>
|
|
<div style='font-size: 0.8rem; color: #666;'>
|
|
{"🟢 发言中" if is_active else "⚪ 等待中"}
|
|
</div>
|
|
<div style='font-size: 0.8rem; color: #999; margin-top: 5px;'>
|
|
💬 {count} 条消息
|
|
</div>
|
|
</div>
|
|
""", unsafe_allow_html=True)
|
|
|
|
def show_chat():
|
|
"""显示对话流"""
|
|
st.subheader("💬 Agent 对话流")
|
|
|
|
if not st.session_state.messages:
|
|
st.info("👈 暂无对话，请在下方输入需求并启动")
|
|
return
|
|
|
|
for msg in st.session_state.messages:
|
|
agent = msg["agent"]
|
|
info = AGENTS.get(agent, {"name": "未知", "avatar": "🤖", "color": "gray"})
|
|
|
|
with st.chat_message(agent.lower(), avatar=info["avatar"]):
|
|
st.markdown(f"**{info['name']}** *{msg['time']}* - {msg['task']}")
|
|
st.markdown(msg["content"][:800] + ("..." if len(msg["content"]) > 800 else ""))
|
|
|
|
def extract_code(content):
|
|
"""从 Markdown 代码块中提取纯代码"""
|
|
# 检查是否有 ```python 标记
|
|
if "```python" in content:
|
|
# 提取 ```python 和 ``` 之间的内容
|
|
parts = content.split("```python")
|
|
if len(parts) > 1:
|
|
code = parts[1].split("```")[0].strip()
|
|
return code
|
|
elif "```" in content:
|
|
# 通用的 ``` 标记
|
|
parts = content.split("```")
|
|
if len(parts) > 1:
|
|
code = parts[1].strip()
|
|
return code
|
|
# 没有标记，返回原内容
|
|
return content
|
|
|
|
def save_files():
|
|
"""保存生成的文件到 workspace/"""
|
|
workspace = Path("workspace")
|
|
workspace.mkdir(exist_ok=True)
|
|
|
|
files = []
|
|
|
|
# 遍历消息，提取并保存文件
|
|
for msg in st.session_state.messages:
|
|
agent = msg["agent"]
|
|
content = msg["content"]
|
|
|
|
# PM Agent 生成 SRS
|
|
if agent == "PM_Agent" and ("需求" in content or "SRS" in content):
|
|
file = workspace / "SRS.md"
|
|
with open(file, "w", encoding="utf-8") as f:
|
|
f.write(f"# 软件需求规格说明书\n\n生成时间：{datetime.now().strftime('%Y-%m-%d %H:%M:%S')}\n\n")
|
|
f.write(content)
|
|
files.append(str(file))
|
|
|
|
# QA Agent 生成测试
|
|
if agent == "QA_Agent" and ("test" in content.lower() or "测试" in content or "def test_" in content):
|
|
file = workspace / "test_sample.py"
|
|
# 提取纯代码
|
|
code = extract_code(content)
|
|
with open(file, "w", encoding="utf-8") as f:
|
|
f.write(f"# 测试用例\n# 生成时间：{datetime.now().strftime('%Y-%m-%d %H:%M:%S')}\n\n")
|
|
f.write(code)
|
|
files.append(str(file))
|
|
|
|
# Dev Agent 生成代码
|
|
if agent == "Dev_Agent" and ("def " in content or "class " in content):
|
|
file = workspace / "src_sample.py"
|
|
# 提取纯代码
|
|
code = extract_code(content)
|
|
with open(file, "w", encoding="utf-8") as f:
|
|
f.write(f"# 源代码\n# 生成时间：{datetime.now().strftime('%Y-%m-%d %H:%M:%S')}\n\n")
|
|
f.write(code)
|
|
files.append(str(file))
|
|
|
|
return files
|
|
|
|
def main():
|
|
st.title("🤖 多 Agent 协作平台")
|
|
st.markdown("**实时展示 Agent 状态 · 自动生成文件**")
|
|
|
|
init_state()
|
|
|
|
# 侧边栏
|
|
with st.sidebar:
|
|
st.title("⚙️ 配置")
|
|
|
|
api_key = st.text_input("API Key", type="password", value=os.getenv("DASHSCOPE_API_KEY", ""))
|
|
model = st.selectbox("模型", ["qwen3.5-flash", "qwen-max", "qwen-plus"], index=0)
|
|
max_round = st.slider("最大轮数", 5, 300, 15)
|
|
|
|
st.divider()
|
|
|
|
if st.button("▶️ 启动工作流", type="primary", use_container_width=True):
|
|
if not api_key:
|
|
st.error("请先设置 API Key")
|
|
elif not AUTOGEN_AVAILABLE:
|
|
st.error("请先安装 AutoGen")
|
|
else:
|
|
run_workflow(api_key, model, max_round)
|
|
|
|
st.divider()
|
|
|
|
if st.button("🗑️ 清空对话", use_container_width=True):
|
|
st.session_state.messages = []
|
|
st.session_state.current_agent = None
|
|
st.session_state.agent_counts = {k: 0 for k in AGENTS}
|
|
st.rerun()
|
|
|
|
# 显示生成的文件
|
|
st.divider()
|
|
st.subheader("📁 生成的文件")
|
|
workspace = Path("workspace")
|
|
if workspace.exists():
|
|
for file in workspace.glob("*"):
|
|
if file.is_file():
|
|
st.caption(f"📄 {file.name}")
|
|
|
|
# 主界面
|
|
show_agent_status()
|
|
st.divider()
|
|
show_chat()
|
|
|
|
# 输入框
|
|
st.divider()
|
|
if user_input := st.chat_input("输入需求..."):
|
|
add_message("User_Proxy", user_input, "提出需求")
|
|
st.rerun()
|
|
|
|
def run_workflow(api_key, model, max_round):
|
|
"""运行工作流"""
|
|
# 保存现有消息（包括用户需求）
|
|
existing_messages = list(st.session_state.messages)
|
|
|
|
# 只清空计数，不清空消息
|
|
st.session_state.running = True
|
|
st.session_state.agent_counts = {k: 0 for k in AGENTS}
|
|
|
|
progress = st.empty()
|
|
progress.info("🚀 启动工作流...")
|
|
|
|
try:
|
|
# 获取需求 - 在清空前已经保存了
|
|
user_msgs = [m for m in existing_messages if m["agent"] == "User_Proxy"]
|
|
requirement = user_msgs[-1]["content"] if user_msgs else "开发一个电池健康预测 API"
|
|
|
|
# 显示需求
|
|
st.info(f"📋 用户需求：{requirement}")
|
|
|
|
# 创建 Agent
|
|
llm_config = get_llm_config(model=model, api_key=api_key)
|
|
|
|
pm = AssistantAgent("PM_Agent", system_message=PM_PROMPT, llm_config=llm_config, human_input_mode="NEVER")
|
|
qa = AssistantAgent("QA_Agent", system_message=QA_PROMPT, llm_config=llm_config, human_input_mode="NEVER")
|
|
dev = AssistantAgent("Dev_Agent", system_message=DEV_PROMPT, llm_config=llm_config, human_input_mode="NEVER")
|
|
orch = AssistantAgent("Orchestrator", system_message=ORCH_PROMPT, llm_config=llm_config, human_input_mode="NEVER")
|
|
user = UserProxyAgent("User_Proxy", human_input_mode="NEVER", max_consecutive_auto_reply=0,
|
|
code_execution_config={"work_dir": "workspace", "use_docker": False})
|
|
|
|
# 创建 GroupChat
|
|
groupchat = GroupChat(
|
|
agents=[pm, qa, dev, orch, user],
|
|
messages=[],
|
|
max_round=max_round,
|
|
speaker_selection_method="round_robin"
|
|
)
|
|
manager = GroupChatManager(groupchat=groupchat, llm_config=llm_config)
|
|
|
|
# 初始消息
|
|
initial_msg = f"""请启动完整的 SDLC 流程：
|
|
|
|
【用户需求】{requirement}
|
|
|
|
【流程】
|
|
1. PM_Agent → SRS 文档
|
|
2. QA_Agent → 测试用例
|
|
3. Dev_Agent → 编写代码
|
|
4. User_Proxy → 执行测试
|
|
5. Orchestrator → 汇总
|
|
|
|
开始协作！"""
|
|
|
|
# 执行对话
|
|
with st.spinner("💬 Agent 们正在协作中..."):
|
|
chat_result = user.initiate_chat(manager, message=initial_msg, max_turns=max_round)
|
|
|
|
# 记录所有对话
|
|
task_map = {
|
|
"PM_Agent": "需求分析",
|
|
"QA_Agent": "测试设计",
|
|
"Dev_Agent": "代码实现",
|
|
"Orchestrator": "流程协调",
|
|
"User_Proxy": "测试执行"
|
|
}
|
|
|
|
for msg in groupchat.messages:
|
|
agent = msg.get("name", "Unknown")
|
|
content = msg.get("content", "")
|
|
task = task_map.get(agent, "工作中")
|
|
add_message(agent, content, task)
|
|
|
|
# 保存文件
|
|
progress.info("💾 正在保存文件...")
|
|
files = save_files()
|
|
|
|
if files:
|
|
progress.success(f"✅ 完成！已保存 {len(files)} 个文件到 workspace/")
|
|
else:
|
|
progress.success("✅ 工作流完成！")
|
|
|
|
st.session_state.running = False
|
|
st.rerun()
|
|
|
|
except Exception as e:
|
|
st.session_state.running = False
|
|
progress.error(f"❌ 错误：{e}")
|
|
import traceback
|
|
st.error(traceback.format_exc())
|
|
|
|
if __name__ == "__main__":
|
|
main()
|