2025-12-11 19:09:10 +08:00
|
|
|
|
"""
|
|
|
|
|
|
Agent 基类
|
|
|
|
|
|
定义 Agent 的基本接口和通用功能
|
2025-12-11 21:14:32 +08:00
|
|
|
|
|
2025-12-11 23:29:04 +08:00
|
|
|
|
核心原则:
|
|
|
|
|
|
1. LLM 是 Agent 的大脑,全程参与决策
|
|
|
|
|
|
2. Agent 之间通过 TaskHandoff 传递结构化上下文
|
|
|
|
|
|
3. 事件分为流式事件(前端展示)和持久化事件(数据库记录)
|
feat(agent): implement comprehensive agent architecture with knowledge base and persistence layer
- Add database migrations for agent checkpoints and tree node tracking
- Implement core agent execution framework with executor, state management, and message handling
- Create knowledge base system with framework-specific modules (Django, FastAPI, Flask, Express, React, Supabase)
- Add vulnerability knowledge modules covering authentication, cryptography, injection, XSS, XXE, SSRF, path traversal, deserialization, and race conditions
- Introduce new agent tools: thinking tool, reporting tool, and agent-specific utilities
- Implement LLM memory compression and prompt caching for improved performance
- Add agent registry and persistence layer for checkpoint management
- Refactor agent implementations (analysis, recon, verification, orchestrator) with enhanced capabilities
- Remove legacy agent implementations (analysis_v2, react_agent)
- Update API endpoints for agent task creation and project management
- Add frontend components for agent task creation and enhanced audit UI
- Consolidate agent service architecture with improved separation of concerns
- This refactoring provides a scalable foundation for multi-agent collaboration with knowledge-driven decision making and state persistence
2025-12-12 15:27:12 +08:00
|
|
|
|
4. 支持动态Agent树和专业知识模块
|
|
|
|
|
|
5. 完整的状态管理和Agent间通信
|
2025-12-11 19:09:10 +08:00
|
|
|
|
"""
|
|
|
|
|
|
|
|
|
|
|
|
from abc import ABC, abstractmethod
|
2025-12-11 23:29:04 +08:00
|
|
|
|
from typing import List, Dict, Any, Optional, AsyncGenerator, Tuple
|
2025-12-11 19:09:10 +08:00
|
|
|
|
from dataclasses import dataclass, field
|
|
|
|
|
|
from enum import Enum
|
2025-12-11 23:29:04 +08:00
|
|
|
|
from datetime import datetime, timezone
|
|
|
|
|
|
import asyncio
|
2025-12-14 17:20:54 +08:00
|
|
|
|
import json
|
2025-12-11 19:09:10 +08:00
|
|
|
|
import logging
|
2025-12-11 23:29:04 +08:00
|
|
|
|
import uuid
|
2025-12-11 19:09:10 +08:00
|
|
|
|
|
feat(agent): implement comprehensive agent architecture with knowledge base and persistence layer
- Add database migrations for agent checkpoints and tree node tracking
- Implement core agent execution framework with executor, state management, and message handling
- Create knowledge base system with framework-specific modules (Django, FastAPI, Flask, Express, React, Supabase)
- Add vulnerability knowledge modules covering authentication, cryptography, injection, XSS, XXE, SSRF, path traversal, deserialization, and race conditions
- Introduce new agent tools: thinking tool, reporting tool, and agent-specific utilities
- Implement LLM memory compression and prompt caching for improved performance
- Add agent registry and persistence layer for checkpoint management
- Refactor agent implementations (analysis, recon, verification, orchestrator) with enhanced capabilities
- Remove legacy agent implementations (analysis_v2, react_agent)
- Update API endpoints for agent task creation and project management
- Add frontend components for agent task creation and enhanced audit UI
- Consolidate agent service architecture with improved separation of concerns
- This refactoring provides a scalable foundation for multi-agent collaboration with knowledge-driven decision making and state persistence
2025-12-12 15:27:12 +08:00
|
|
|
|
from ..core.state import AgentState, AgentStatus
|
|
|
|
|
|
from ..core.registry import agent_registry
|
|
|
|
|
|
from ..core.message import message_bus, MessageType, AgentMessage
|
|
|
|
|
|
|
2025-12-11 19:09:10 +08:00
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
class AgentType(Enum):
|
|
|
|
|
|
"""Agent 类型"""
|
|
|
|
|
|
ORCHESTRATOR = "orchestrator"
|
|
|
|
|
|
RECON = "recon"
|
|
|
|
|
|
ANALYSIS = "analysis"
|
|
|
|
|
|
VERIFICATION = "verification"
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
class AgentPattern(Enum):
|
|
|
|
|
|
"""Agent 运行模式"""
|
|
|
|
|
|
REACT = "react" # 反应式:思考-行动-观察循环
|
|
|
|
|
|
PLAN_AND_EXECUTE = "plan_execute" # 计划执行:先规划后执行
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
@dataclass
|
|
|
|
|
|
class AgentConfig:
|
|
|
|
|
|
"""Agent 配置"""
|
|
|
|
|
|
name: str
|
|
|
|
|
|
agent_type: AgentType
|
|
|
|
|
|
pattern: AgentPattern = AgentPattern.REACT
|
|
|
|
|
|
|
|
|
|
|
|
# LLM 配置
|
|
|
|
|
|
model: Optional[str] = None
|
|
|
|
|
|
temperature: float = 0.1
|
2025-12-17 11:02:42 +08:00
|
|
|
|
max_tokens: int = 8192
|
2025-12-11 19:09:10 +08:00
|
|
|
|
|
|
|
|
|
|
# 执行限制
|
|
|
|
|
|
max_iterations: int = 20
|
|
|
|
|
|
timeout_seconds: int = 600
|
|
|
|
|
|
|
|
|
|
|
|
# 工具配置
|
|
|
|
|
|
tools: List[str] = field(default_factory=list)
|
|
|
|
|
|
|
|
|
|
|
|
# 系统提示词
|
|
|
|
|
|
system_prompt: Optional[str] = None
|
|
|
|
|
|
|
|
|
|
|
|
# 元数据
|
|
|
|
|
|
metadata: Dict[str, Any] = field(default_factory=dict)
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
@dataclass
|
|
|
|
|
|
class AgentResult:
|
|
|
|
|
|
"""Agent 执行结果"""
|
|
|
|
|
|
success: bool
|
|
|
|
|
|
data: Any = None
|
|
|
|
|
|
error: Optional[str] = None
|
|
|
|
|
|
|
|
|
|
|
|
# 执行统计
|
|
|
|
|
|
iterations: int = 0
|
|
|
|
|
|
tool_calls: int = 0
|
|
|
|
|
|
tokens_used: int = 0
|
|
|
|
|
|
duration_ms: int = 0
|
|
|
|
|
|
|
|
|
|
|
|
# 中间结果
|
|
|
|
|
|
intermediate_steps: List[Dict[str, Any]] = field(default_factory=list)
|
|
|
|
|
|
|
|
|
|
|
|
# 元数据
|
|
|
|
|
|
metadata: Dict[str, Any] = field(default_factory=dict)
|
|
|
|
|
|
|
2025-12-11 23:29:04 +08:00
|
|
|
|
# 🔥 协作信息 - Agent 传递给下一个 Agent 的结构化信息
|
|
|
|
|
|
handoff: Optional["TaskHandoff"] = None
|
|
|
|
|
|
|
2025-12-11 19:09:10 +08:00
|
|
|
|
def to_dict(self) -> Dict[str, Any]:
|
|
|
|
|
|
return {
|
|
|
|
|
|
"success": self.success,
|
|
|
|
|
|
"data": self.data,
|
|
|
|
|
|
"error": self.error,
|
|
|
|
|
|
"iterations": self.iterations,
|
|
|
|
|
|
"tool_calls": self.tool_calls,
|
|
|
|
|
|
"tokens_used": self.tokens_used,
|
|
|
|
|
|
"duration_ms": self.duration_ms,
|
|
|
|
|
|
"metadata": self.metadata,
|
2025-12-11 23:29:04 +08:00
|
|
|
|
"handoff": self.handoff.to_dict() if self.handoff else None,
|
2025-12-11 19:09:10 +08:00
|
|
|
|
}
|
|
|
|
|
|
|
|
|
|
|
|
|
2025-12-11 23:29:04 +08:00
|
|
|
|
@dataclass
|
|
|
|
|
|
class TaskHandoff:
|
|
|
|
|
|
"""
|
|
|
|
|
|
任务交接协议 - Agent 之间传递的结构化信息
|
|
|
|
|
|
|
|
|
|
|
|
设计原则:
|
|
|
|
|
|
1. 包含足够的上下文让下一个 Agent 理解前序工作
|
|
|
|
|
|
2. 提供明确的建议和关注点
|
|
|
|
|
|
3. 可直接转换为 LLM 可理解的 prompt
|
|
|
|
|
|
"""
|
|
|
|
|
|
# 基本信息
|
|
|
|
|
|
from_agent: str
|
|
|
|
|
|
to_agent: str
|
|
|
|
|
|
|
|
|
|
|
|
# 工作摘要
|
|
|
|
|
|
summary: str
|
|
|
|
|
|
work_completed: List[str] = field(default_factory=list)
|
|
|
|
|
|
|
|
|
|
|
|
# 关键发现和洞察
|
|
|
|
|
|
key_findings: List[Dict[str, Any]] = field(default_factory=list)
|
|
|
|
|
|
insights: List[str] = field(default_factory=list)
|
|
|
|
|
|
|
|
|
|
|
|
# 建议和关注点
|
|
|
|
|
|
suggested_actions: List[Dict[str, Any]] = field(default_factory=list)
|
|
|
|
|
|
attention_points: List[str] = field(default_factory=list)
|
|
|
|
|
|
priority_areas: List[str] = field(default_factory=list)
|
|
|
|
|
|
|
|
|
|
|
|
# 上下文数据
|
|
|
|
|
|
context_data: Dict[str, Any] = field(default_factory=dict)
|
|
|
|
|
|
|
|
|
|
|
|
# 置信度
|
|
|
|
|
|
confidence: float = 0.8
|
|
|
|
|
|
|
|
|
|
|
|
# 时间戳
|
|
|
|
|
|
timestamp: datetime = field(default_factory=lambda: datetime.now(timezone.utc))
|
|
|
|
|
|
|
|
|
|
|
|
def to_dict(self) -> Dict[str, Any]:
|
|
|
|
|
|
return {
|
|
|
|
|
|
"from_agent": self.from_agent,
|
|
|
|
|
|
"to_agent": self.to_agent,
|
|
|
|
|
|
"summary": self.summary,
|
|
|
|
|
|
"work_completed": self.work_completed,
|
|
|
|
|
|
"key_findings": self.key_findings,
|
|
|
|
|
|
"insights": self.insights,
|
|
|
|
|
|
"suggested_actions": self.suggested_actions,
|
|
|
|
|
|
"attention_points": self.attention_points,
|
|
|
|
|
|
"priority_areas": self.priority_areas,
|
|
|
|
|
|
"context_data": self.context_data,
|
|
|
|
|
|
"confidence": self.confidence,
|
|
|
|
|
|
"timestamp": self.timestamp.isoformat(),
|
|
|
|
|
|
}
|
|
|
|
|
|
|
|
|
|
|
|
@classmethod
|
|
|
|
|
|
def from_dict(cls, data: Dict[str, Any]) -> "TaskHandoff":
|
|
|
|
|
|
return cls(
|
|
|
|
|
|
from_agent=data.get("from_agent", ""),
|
|
|
|
|
|
to_agent=data.get("to_agent", ""),
|
|
|
|
|
|
summary=data.get("summary", ""),
|
|
|
|
|
|
work_completed=data.get("work_completed", []),
|
|
|
|
|
|
key_findings=data.get("key_findings", []),
|
|
|
|
|
|
insights=data.get("insights", []),
|
|
|
|
|
|
suggested_actions=data.get("suggested_actions", []),
|
|
|
|
|
|
attention_points=data.get("attention_points", []),
|
|
|
|
|
|
priority_areas=data.get("priority_areas", []),
|
|
|
|
|
|
context_data=data.get("context_data", {}),
|
|
|
|
|
|
confidence=data.get("confidence", 0.8),
|
|
|
|
|
|
)
|
|
|
|
|
|
|
|
|
|
|
|
def to_prompt_context(self) -> str:
|
|
|
|
|
|
"""
|
|
|
|
|
|
转换为 LLM 可理解的上下文格式
|
|
|
|
|
|
这是关键!让 LLM 能够理解前序 Agent 的工作
|
|
|
|
|
|
"""
|
|
|
|
|
|
lines = [
|
|
|
|
|
|
f"## 来自 {self.from_agent} Agent 的任务交接",
|
|
|
|
|
|
"",
|
|
|
|
|
|
f"### 工作摘要",
|
|
|
|
|
|
self.summary,
|
|
|
|
|
|
"",
|
|
|
|
|
|
]
|
|
|
|
|
|
|
|
|
|
|
|
if self.work_completed:
|
|
|
|
|
|
lines.append("### 已完成的工作")
|
|
|
|
|
|
for work in self.work_completed:
|
|
|
|
|
|
lines.append(f"- {work}")
|
|
|
|
|
|
lines.append("")
|
|
|
|
|
|
|
|
|
|
|
|
if self.key_findings:
|
|
|
|
|
|
lines.append("### 关键发现")
|
|
|
|
|
|
for i, finding in enumerate(self.key_findings[:15], 1):
|
|
|
|
|
|
severity = finding.get("severity", "medium")
|
|
|
|
|
|
title = finding.get("title", "Unknown")
|
|
|
|
|
|
file_path = finding.get("file_path", "")
|
|
|
|
|
|
lines.append(f"{i}. [{severity.upper()}] {title}")
|
|
|
|
|
|
if file_path:
|
|
|
|
|
|
lines.append(f" 位置: {file_path}:{finding.get('line_start', '')}")
|
|
|
|
|
|
if finding.get("description"):
|
|
|
|
|
|
lines.append(f" 描述: {finding['description'][:100]}")
|
|
|
|
|
|
lines.append("")
|
|
|
|
|
|
|
|
|
|
|
|
if self.insights:
|
|
|
|
|
|
lines.append("### 洞察和分析")
|
|
|
|
|
|
for insight in self.insights:
|
|
|
|
|
|
lines.append(f"- {insight}")
|
|
|
|
|
|
lines.append("")
|
|
|
|
|
|
|
|
|
|
|
|
if self.suggested_actions:
|
|
|
|
|
|
lines.append("### 建议的下一步行动")
|
|
|
|
|
|
for action in self.suggested_actions:
|
|
|
|
|
|
action_type = action.get("type", "general")
|
|
|
|
|
|
description = action.get("description", "")
|
|
|
|
|
|
priority = action.get("priority", "medium")
|
|
|
|
|
|
lines.append(f"- [{priority.upper()}] {action_type}: {description}")
|
|
|
|
|
|
lines.append("")
|
|
|
|
|
|
|
|
|
|
|
|
if self.attention_points:
|
|
|
|
|
|
lines.append("### ⚠️ 需要特别关注")
|
|
|
|
|
|
for point in self.attention_points:
|
|
|
|
|
|
lines.append(f"- {point}")
|
|
|
|
|
|
lines.append("")
|
|
|
|
|
|
|
|
|
|
|
|
if self.priority_areas:
|
|
|
|
|
|
lines.append("### 优先分析区域")
|
|
|
|
|
|
for area in self.priority_areas:
|
|
|
|
|
|
lines.append(f"- {area}")
|
|
|
|
|
|
|
|
|
|
|
|
return "\n".join(lines)
|
|
|
|
|
|
|
|
|
|
|
|
|
2025-12-11 19:09:10 +08:00
|
|
|
|
class BaseAgent(ABC):
|
|
|
|
|
|
"""
|
|
|
|
|
|
Agent 基类
|
2025-12-11 21:14:32 +08:00
|
|
|
|
|
|
|
|
|
|
核心原则:
|
|
|
|
|
|
1. LLM 是 Agent 的大脑,全程参与决策
|
|
|
|
|
|
2. 所有日志应该反映 LLM 的思考过程
|
|
|
|
|
|
3. 工具调用是 LLM 的决策结果
|
2025-12-11 23:29:04 +08:00
|
|
|
|
|
|
|
|
|
|
协作原则:
|
|
|
|
|
|
1. 通过 TaskHandoff 接收前序 Agent 的上下文
|
|
|
|
|
|
2. 执行完成后生成 TaskHandoff 传递给下一个 Agent
|
|
|
|
|
|
3. 洞察和发现应该结构化记录
|
feat(agent): implement comprehensive agent architecture with knowledge base and persistence layer
- Add database migrations for agent checkpoints and tree node tracking
- Implement core agent execution framework with executor, state management, and message handling
- Create knowledge base system with framework-specific modules (Django, FastAPI, Flask, Express, React, Supabase)
- Add vulnerability knowledge modules covering authentication, cryptography, injection, XSS, XXE, SSRF, path traversal, deserialization, and race conditions
- Introduce new agent tools: thinking tool, reporting tool, and agent-specific utilities
- Implement LLM memory compression and prompt caching for improved performance
- Add agent registry and persistence layer for checkpoint management
- Refactor agent implementations (analysis, recon, verification, orchestrator) with enhanced capabilities
- Remove legacy agent implementations (analysis_v2, react_agent)
- Update API endpoints for agent task creation and project management
- Add frontend components for agent task creation and enhanced audit UI
- Consolidate agent service architecture with improved separation of concerns
- This refactoring provides a scalable foundation for multi-agent collaboration with knowledge-driven decision making and state persistence
2025-12-12 15:27:12 +08:00
|
|
|
|
|
|
|
|
|
|
动态Agent树:
|
|
|
|
|
|
1. 支持动态创建子Agent
|
|
|
|
|
|
2. Agent间通过消息总线通信
|
|
|
|
|
|
3. 完整的状态管理和生命周期
|
2025-12-11 19:09:10 +08:00
|
|
|
|
"""
|
|
|
|
|
|
|
|
|
|
|
|
def __init__(
|
|
|
|
|
|
self,
|
|
|
|
|
|
config: AgentConfig,
|
|
|
|
|
|
llm_service,
|
|
|
|
|
|
tools: Dict[str, Any],
|
|
|
|
|
|
event_emitter=None,
|
feat(agent): implement comprehensive agent architecture with knowledge base and persistence layer
- Add database migrations for agent checkpoints and tree node tracking
- Implement core agent execution framework with executor, state management, and message handling
- Create knowledge base system with framework-specific modules (Django, FastAPI, Flask, Express, React, Supabase)
- Add vulnerability knowledge modules covering authentication, cryptography, injection, XSS, XXE, SSRF, path traversal, deserialization, and race conditions
- Introduce new agent tools: thinking tool, reporting tool, and agent-specific utilities
- Implement LLM memory compression and prompt caching for improved performance
- Add agent registry and persistence layer for checkpoint management
- Refactor agent implementations (analysis, recon, verification, orchestrator) with enhanced capabilities
- Remove legacy agent implementations (analysis_v2, react_agent)
- Update API endpoints for agent task creation and project management
- Add frontend components for agent task creation and enhanced audit UI
- Consolidate agent service architecture with improved separation of concerns
- This refactoring provides a scalable foundation for multi-agent collaboration with knowledge-driven decision making and state persistence
2025-12-12 15:27:12 +08:00
|
|
|
|
parent_id: Optional[str] = None,
|
|
|
|
|
|
knowledge_modules: Optional[List[str]] = None,
|
2025-12-11 19:09:10 +08:00
|
|
|
|
):
|
|
|
|
|
|
"""
|
|
|
|
|
|
初始化 Agent
|
|
|
|
|
|
|
|
|
|
|
|
Args:
|
|
|
|
|
|
config: Agent 配置
|
|
|
|
|
|
llm_service: LLM 服务
|
|
|
|
|
|
tools: 可用工具字典
|
|
|
|
|
|
event_emitter: 事件发射器
|
feat(agent): implement comprehensive agent architecture with knowledge base and persistence layer
- Add database migrations for agent checkpoints and tree node tracking
- Implement core agent execution framework with executor, state management, and message handling
- Create knowledge base system with framework-specific modules (Django, FastAPI, Flask, Express, React, Supabase)
- Add vulnerability knowledge modules covering authentication, cryptography, injection, XSS, XXE, SSRF, path traversal, deserialization, and race conditions
- Introduce new agent tools: thinking tool, reporting tool, and agent-specific utilities
- Implement LLM memory compression and prompt caching for improved performance
- Add agent registry and persistence layer for checkpoint management
- Refactor agent implementations (analysis, recon, verification, orchestrator) with enhanced capabilities
- Remove legacy agent implementations (analysis_v2, react_agent)
- Update API endpoints for agent task creation and project management
- Add frontend components for agent task creation and enhanced audit UI
- Consolidate agent service architecture with improved separation of concerns
- This refactoring provides a scalable foundation for multi-agent collaboration with knowledge-driven decision making and state persistence
2025-12-12 15:27:12 +08:00
|
|
|
|
parent_id: 父Agent ID(用于动态Agent树)
|
|
|
|
|
|
knowledge_modules: 要加载的知识模块
|
2025-12-11 19:09:10 +08:00
|
|
|
|
"""
|
|
|
|
|
|
self.config = config
|
|
|
|
|
|
self.llm_service = llm_service
|
|
|
|
|
|
self.tools = tools
|
|
|
|
|
|
self.event_emitter = event_emitter
|
feat(agent): implement comprehensive agent architecture with knowledge base and persistence layer
- Add database migrations for agent checkpoints and tree node tracking
- Implement core agent execution framework with executor, state management, and message handling
- Create knowledge base system with framework-specific modules (Django, FastAPI, Flask, Express, React, Supabase)
- Add vulnerability knowledge modules covering authentication, cryptography, injection, XSS, XXE, SSRF, path traversal, deserialization, and race conditions
- Introduce new agent tools: thinking tool, reporting tool, and agent-specific utilities
- Implement LLM memory compression and prompt caching for improved performance
- Add agent registry and persistence layer for checkpoint management
- Refactor agent implementations (analysis, recon, verification, orchestrator) with enhanced capabilities
- Remove legacy agent implementations (analysis_v2, react_agent)
- Update API endpoints for agent task creation and project management
- Add frontend components for agent task creation and enhanced audit UI
- Consolidate agent service architecture with improved separation of concerns
- This refactoring provides a scalable foundation for multi-agent collaboration with knowledge-driven decision making and state persistence
2025-12-12 15:27:12 +08:00
|
|
|
|
self.parent_id = parent_id
|
|
|
|
|
|
self.knowledge_modules = knowledge_modules or []
|
|
|
|
|
|
|
|
|
|
|
|
# 🔥 生成唯一ID
|
|
|
|
|
|
self._agent_id = f"agent_{uuid.uuid4().hex[:8]}"
|
|
|
|
|
|
|
|
|
|
|
|
# 🔥 增强的状态管理
|
|
|
|
|
|
self._state = AgentState(
|
|
|
|
|
|
agent_id=self._agent_id,
|
|
|
|
|
|
agent_name=config.name,
|
|
|
|
|
|
agent_type=config.agent_type.value,
|
|
|
|
|
|
parent_id=parent_id,
|
|
|
|
|
|
max_iterations=config.max_iterations,
|
|
|
|
|
|
knowledge_modules=self.knowledge_modules,
|
|
|
|
|
|
)
|
2025-12-11 19:09:10 +08:00
|
|
|
|
|
feat(agent): implement comprehensive agent architecture with knowledge base and persistence layer
- Add database migrations for agent checkpoints and tree node tracking
- Implement core agent execution framework with executor, state management, and message handling
- Create knowledge base system with framework-specific modules (Django, FastAPI, Flask, Express, React, Supabase)
- Add vulnerability knowledge modules covering authentication, cryptography, injection, XSS, XXE, SSRF, path traversal, deserialization, and race conditions
- Introduce new agent tools: thinking tool, reporting tool, and agent-specific utilities
- Implement LLM memory compression and prompt caching for improved performance
- Add agent registry and persistence layer for checkpoint management
- Refactor agent implementations (analysis, recon, verification, orchestrator) with enhanced capabilities
- Remove legacy agent implementations (analysis_v2, react_agent)
- Update API endpoints for agent task creation and project management
- Add frontend components for agent task creation and enhanced audit UI
- Consolidate agent service architecture with improved separation of concerns
- This refactoring provides a scalable foundation for multi-agent collaboration with knowledge-driven decision making and state persistence
2025-12-12 15:27:12 +08:00
|
|
|
|
# 运行状态(保持向后兼容)
|
2025-12-11 19:09:10 +08:00
|
|
|
|
self._iteration = 0
|
|
|
|
|
|
self._total_tokens = 0
|
|
|
|
|
|
self._tool_calls = 0
|
|
|
|
|
|
self._cancelled = False
|
2025-12-11 23:29:04 +08:00
|
|
|
|
|
|
|
|
|
|
# 🔥 协作状态
|
|
|
|
|
|
self._incoming_handoff: Optional[TaskHandoff] = None
|
|
|
|
|
|
self._insights: List[str] = [] # 收集的洞察
|
|
|
|
|
|
self._work_completed: List[str] = [] # 完成的工作记录
|
feat(agent): implement comprehensive agent architecture with knowledge base and persistence layer
- Add database migrations for agent checkpoints and tree node tracking
- Implement core agent execution framework with executor, state management, and message handling
- Create knowledge base system with framework-specific modules (Django, FastAPI, Flask, Express, React, Supabase)
- Add vulnerability knowledge modules covering authentication, cryptography, injection, XSS, XXE, SSRF, path traversal, deserialization, and race conditions
- Introduce new agent tools: thinking tool, reporting tool, and agent-specific utilities
- Implement LLM memory compression and prompt caching for improved performance
- Add agent registry and persistence layer for checkpoint management
- Refactor agent implementations (analysis, recon, verification, orchestrator) with enhanced capabilities
- Remove legacy agent implementations (analysis_v2, react_agent)
- Update API endpoints for agent task creation and project management
- Add frontend components for agent task creation and enhanced audit UI
- Consolidate agent service architecture with improved separation of concerns
- This refactoring provides a scalable foundation for multi-agent collaboration with knowledge-driven decision making and state persistence
2025-12-12 15:27:12 +08:00
|
|
|
|
|
|
|
|
|
|
# 🔥 是否已注册到注册表
|
|
|
|
|
|
self._registered = False
|
|
|
|
|
|
|
|
|
|
|
|
# 🔥 加载知识模块到系统提示词
|
|
|
|
|
|
if self.knowledge_modules:
|
|
|
|
|
|
self._load_knowledge_modules()
|
|
|
|
|
|
|
|
|
|
|
|
def _register_to_registry(self, task: Optional[str] = None) -> None:
|
|
|
|
|
|
"""注册到Agent注册表(延迟注册,在run时调用)"""
|
2025-12-12 15:50:48 +08:00
|
|
|
|
logger.debug(f"[AgentTree] _register_to_registry 被调用: {self.config.name} (id={self._agent_id}, parent={self.parent_id}, _registered={self._registered})")
|
feat(agent): implement comprehensive agent architecture with knowledge base and persistence layer
- Add database migrations for agent checkpoints and tree node tracking
- Implement core agent execution framework with executor, state management, and message handling
- Create knowledge base system with framework-specific modules (Django, FastAPI, Flask, Express, React, Supabase)
- Add vulnerability knowledge modules covering authentication, cryptography, injection, XSS, XXE, SSRF, path traversal, deserialization, and race conditions
- Introduce new agent tools: thinking tool, reporting tool, and agent-specific utilities
- Implement LLM memory compression and prompt caching for improved performance
- Add agent registry and persistence layer for checkpoint management
- Refactor agent implementations (analysis, recon, verification, orchestrator) with enhanced capabilities
- Remove legacy agent implementations (analysis_v2, react_agent)
- Update API endpoints for agent task creation and project management
- Add frontend components for agent task creation and enhanced audit UI
- Consolidate agent service architecture with improved separation of concerns
- This refactoring provides a scalable foundation for multi-agent collaboration with knowledge-driven decision making and state persistence
2025-12-12 15:27:12 +08:00
|
|
|
|
|
|
|
|
|
|
if self._registered:
|
2025-12-12 15:50:48 +08:00
|
|
|
|
logger.debug(f"[AgentTree] {self.config.name} 已注册,跳过 (id={self._agent_id})")
|
feat(agent): implement comprehensive agent architecture with knowledge base and persistence layer
- Add database migrations for agent checkpoints and tree node tracking
- Implement core agent execution framework with executor, state management, and message handling
- Create knowledge base system with framework-specific modules (Django, FastAPI, Flask, Express, React, Supabase)
- Add vulnerability knowledge modules covering authentication, cryptography, injection, XSS, XXE, SSRF, path traversal, deserialization, and race conditions
- Introduce new agent tools: thinking tool, reporting tool, and agent-specific utilities
- Implement LLM memory compression and prompt caching for improved performance
- Add agent registry and persistence layer for checkpoint management
- Refactor agent implementations (analysis, recon, verification, orchestrator) with enhanced capabilities
- Remove legacy agent implementations (analysis_v2, react_agent)
- Update API endpoints for agent task creation and project management
- Add frontend components for agent task creation and enhanced audit UI
- Consolidate agent service architecture with improved separation of concerns
- This refactoring provides a scalable foundation for multi-agent collaboration with knowledge-driven decision making and state persistence
2025-12-12 15:27:12 +08:00
|
|
|
|
return
|
|
|
|
|
|
|
2025-12-12 15:50:48 +08:00
|
|
|
|
logger.debug(f"[AgentTree] 正在注册 Agent: {self.config.name} (id={self._agent_id}, parent={self.parent_id})")
|
feat(agent): implement comprehensive agent architecture with knowledge base and persistence layer
- Add database migrations for agent checkpoints and tree node tracking
- Implement core agent execution framework with executor, state management, and message handling
- Create knowledge base system with framework-specific modules (Django, FastAPI, Flask, Express, React, Supabase)
- Add vulnerability knowledge modules covering authentication, cryptography, injection, XSS, XXE, SSRF, path traversal, deserialization, and race conditions
- Introduce new agent tools: thinking tool, reporting tool, and agent-specific utilities
- Implement LLM memory compression and prompt caching for improved performance
- Add agent registry and persistence layer for checkpoint management
- Refactor agent implementations (analysis, recon, verification, orchestrator) with enhanced capabilities
- Remove legacy agent implementations (analysis_v2, react_agent)
- Update API endpoints for agent task creation and project management
- Add frontend components for agent task creation and enhanced audit UI
- Consolidate agent service architecture with improved separation of concerns
- This refactoring provides a scalable foundation for multi-agent collaboration with knowledge-driven decision making and state persistence
2025-12-12 15:27:12 +08:00
|
|
|
|
|
|
|
|
|
|
agent_registry.register_agent(
|
|
|
|
|
|
agent_id=self._agent_id,
|
|
|
|
|
|
agent_name=self.config.name,
|
|
|
|
|
|
agent_type=self.config.agent_type.value,
|
|
|
|
|
|
task=task or self._state.task or "Initializing",
|
|
|
|
|
|
parent_id=self.parent_id,
|
|
|
|
|
|
agent_instance=self,
|
|
|
|
|
|
state=self._state,
|
|
|
|
|
|
knowledge_modules=self.knowledge_modules,
|
|
|
|
|
|
)
|
|
|
|
|
|
|
|
|
|
|
|
# 创建消息队列
|
|
|
|
|
|
message_bus.create_queue(self._agent_id)
|
|
|
|
|
|
self._registered = True
|
|
|
|
|
|
|
|
|
|
|
|
tree = agent_registry.get_agent_tree()
|
2025-12-12 15:50:48 +08:00
|
|
|
|
logger.debug(f"[AgentTree] Agent 注册完成: {self.config.name}, 当前树节点数: {len(tree['nodes'])}")
|
feat(agent): implement comprehensive agent architecture with knowledge base and persistence layer
- Add database migrations for agent checkpoints and tree node tracking
- Implement core agent execution framework with executor, state management, and message handling
- Create knowledge base system with framework-specific modules (Django, FastAPI, Flask, Express, React, Supabase)
- Add vulnerability knowledge modules covering authentication, cryptography, injection, XSS, XXE, SSRF, path traversal, deserialization, and race conditions
- Introduce new agent tools: thinking tool, reporting tool, and agent-specific utilities
- Implement LLM memory compression and prompt caching for improved performance
- Add agent registry and persistence layer for checkpoint management
- Refactor agent implementations (analysis, recon, verification, orchestrator) with enhanced capabilities
- Remove legacy agent implementations (analysis_v2, react_agent)
- Update API endpoints for agent task creation and project management
- Add frontend components for agent task creation and enhanced audit UI
- Consolidate agent service architecture with improved separation of concerns
- This refactoring provides a scalable foundation for multi-agent collaboration with knowledge-driven decision making and state persistence
2025-12-12 15:27:12 +08:00
|
|
|
|
|
|
|
|
|
|
def set_parent_id(self, parent_id: str) -> None:
|
|
|
|
|
|
"""设置父Agent ID(在调度时调用)"""
|
|
|
|
|
|
self.parent_id = parent_id
|
|
|
|
|
|
self._state.parent_id = parent_id
|
|
|
|
|
|
|
|
|
|
|
|
def _load_knowledge_modules(self) -> None:
|
|
|
|
|
|
"""加载知识模块到系统提示词"""
|
|
|
|
|
|
if not self.knowledge_modules:
|
|
|
|
|
|
return
|
|
|
|
|
|
|
|
|
|
|
|
try:
|
|
|
|
|
|
from ..knowledge import knowledge_loader
|
|
|
|
|
|
|
|
|
|
|
|
enhanced_prompt = knowledge_loader.build_system_prompt_with_modules(
|
|
|
|
|
|
self.config.system_prompt or "",
|
|
|
|
|
|
self.knowledge_modules,
|
|
|
|
|
|
)
|
|
|
|
|
|
self.config.system_prompt = enhanced_prompt
|
|
|
|
|
|
|
|
|
|
|
|
logger.info(f"[{self.name}] Loaded knowledge modules: {self.knowledge_modules}")
|
|
|
|
|
|
except Exception as e:
|
|
|
|
|
|
logger.warning(f"Failed to load knowledge modules: {e}")
|
2025-12-11 19:09:10 +08:00
|
|
|
|
|
|
|
|
|
|
@property
|
|
|
|
|
|
def name(self) -> str:
|
|
|
|
|
|
return self.config.name
|
|
|
|
|
|
|
feat(agent): implement comprehensive agent architecture with knowledge base and persistence layer
- Add database migrations for agent checkpoints and tree node tracking
- Implement core agent execution framework with executor, state management, and message handling
- Create knowledge base system with framework-specific modules (Django, FastAPI, Flask, Express, React, Supabase)
- Add vulnerability knowledge modules covering authentication, cryptography, injection, XSS, XXE, SSRF, path traversal, deserialization, and race conditions
- Introduce new agent tools: thinking tool, reporting tool, and agent-specific utilities
- Implement LLM memory compression and prompt caching for improved performance
- Add agent registry and persistence layer for checkpoint management
- Refactor agent implementations (analysis, recon, verification, orchestrator) with enhanced capabilities
- Remove legacy agent implementations (analysis_v2, react_agent)
- Update API endpoints for agent task creation and project management
- Add frontend components for agent task creation and enhanced audit UI
- Consolidate agent service architecture with improved separation of concerns
- This refactoring provides a scalable foundation for multi-agent collaboration with knowledge-driven decision making and state persistence
2025-12-12 15:27:12 +08:00
|
|
|
|
@property
|
|
|
|
|
|
def agent_id(self) -> str:
|
|
|
|
|
|
return self._agent_id
|
|
|
|
|
|
|
|
|
|
|
|
@property
|
|
|
|
|
|
def state(self) -> AgentState:
|
|
|
|
|
|
return self._state
|
|
|
|
|
|
|
2025-12-11 19:09:10 +08:00
|
|
|
|
@property
|
|
|
|
|
|
def agent_type(self) -> AgentType:
|
|
|
|
|
|
return self.config.agent_type
|
|
|
|
|
|
|
feat(agent): implement comprehensive agent architecture with knowledge base and persistence layer
- Add database migrations for agent checkpoints and tree node tracking
- Implement core agent execution framework with executor, state management, and message handling
- Create knowledge base system with framework-specific modules (Django, FastAPI, Flask, Express, React, Supabase)
- Add vulnerability knowledge modules covering authentication, cryptography, injection, XSS, XXE, SSRF, path traversal, deserialization, and race conditions
- Introduce new agent tools: thinking tool, reporting tool, and agent-specific utilities
- Implement LLM memory compression and prompt caching for improved performance
- Add agent registry and persistence layer for checkpoint management
- Refactor agent implementations (analysis, recon, verification, orchestrator) with enhanced capabilities
- Remove legacy agent implementations (analysis_v2, react_agent)
- Update API endpoints for agent task creation and project management
- Add frontend components for agent task creation and enhanced audit UI
- Consolidate agent service architecture with improved separation of concerns
- This refactoring provides a scalable foundation for multi-agent collaboration with knowledge-driven decision making and state persistence
2025-12-12 15:27:12 +08:00
|
|
|
|
# ============ Agent间消息处理 ============
|
|
|
|
|
|
|
|
|
|
|
|
def check_messages(self) -> List[AgentMessage]:
|
|
|
|
|
|
"""
|
|
|
|
|
|
检查并处理收到的消息
|
|
|
|
|
|
|
|
|
|
|
|
Returns:
|
|
|
|
|
|
未读消息列表
|
|
|
|
|
|
"""
|
|
|
|
|
|
messages = message_bus.get_messages(
|
|
|
|
|
|
self._agent_id,
|
|
|
|
|
|
unread_only=True,
|
|
|
|
|
|
mark_as_read=True,
|
|
|
|
|
|
)
|
|
|
|
|
|
|
|
|
|
|
|
for msg in messages:
|
|
|
|
|
|
# 处理消息
|
|
|
|
|
|
if msg.from_agent == "user":
|
|
|
|
|
|
# 用户消息直接添加到对话历史
|
|
|
|
|
|
self._state.add_message("user", msg.content)
|
|
|
|
|
|
else:
|
|
|
|
|
|
# Agent间消息使用XML格式
|
|
|
|
|
|
self._state.add_message("user", msg.to_xml())
|
|
|
|
|
|
|
|
|
|
|
|
# 如果在等待状态,恢复执行
|
|
|
|
|
|
if self._state.is_waiting_for_input():
|
|
|
|
|
|
self._state.resume_from_waiting()
|
|
|
|
|
|
agent_registry.update_agent_status(self._agent_id, "running")
|
|
|
|
|
|
|
|
|
|
|
|
return messages
|
|
|
|
|
|
|
|
|
|
|
|
def has_pending_messages(self) -> bool:
|
|
|
|
|
|
"""检查是否有待处理的消息"""
|
|
|
|
|
|
return message_bus.has_unread_messages(self._agent_id)
|
|
|
|
|
|
|
|
|
|
|
|
def send_message_to_parent(
|
|
|
|
|
|
self,
|
|
|
|
|
|
content: str,
|
|
|
|
|
|
message_type: MessageType = MessageType.INFORMATION,
|
|
|
|
|
|
) -> None:
|
|
|
|
|
|
"""向父Agent发送消息"""
|
|
|
|
|
|
if self.parent_id:
|
|
|
|
|
|
message_bus.send_message(
|
|
|
|
|
|
from_agent=self._agent_id,
|
|
|
|
|
|
to_agent=self.parent_id,
|
|
|
|
|
|
content=content,
|
|
|
|
|
|
message_type=message_type,
|
|
|
|
|
|
)
|
|
|
|
|
|
|
|
|
|
|
|
def send_message_to_agent(
|
|
|
|
|
|
self,
|
|
|
|
|
|
target_id: str,
|
|
|
|
|
|
content: str,
|
|
|
|
|
|
message_type: MessageType = MessageType.INFORMATION,
|
|
|
|
|
|
) -> None:
|
|
|
|
|
|
"""向指定Agent发送消息"""
|
|
|
|
|
|
message_bus.send_message(
|
|
|
|
|
|
from_agent=self._agent_id,
|
|
|
|
|
|
to_agent=target_id,
|
|
|
|
|
|
content=content,
|
|
|
|
|
|
message_type=message_type,
|
|
|
|
|
|
)
|
|
|
|
|
|
|
|
|
|
|
|
# ============ 生命周期管理 ============
|
|
|
|
|
|
|
|
|
|
|
|
def on_start(self) -> None:
|
|
|
|
|
|
"""Agent开始执行时调用"""
|
|
|
|
|
|
self._state.start()
|
|
|
|
|
|
agent_registry.update_agent_status(self._agent_id, "running")
|
|
|
|
|
|
|
|
|
|
|
|
def on_complete(self, result: Dict[str, Any]) -> None:
|
|
|
|
|
|
"""Agent完成时调用"""
|
|
|
|
|
|
self._state.set_completed(result)
|
|
|
|
|
|
agent_registry.update_agent_status(self._agent_id, "completed", result)
|
|
|
|
|
|
|
|
|
|
|
|
# 向父Agent报告完成
|
|
|
|
|
|
if self.parent_id:
|
|
|
|
|
|
message_bus.send_completion_report(
|
|
|
|
|
|
from_agent=self._agent_id,
|
|
|
|
|
|
to_agent=self.parent_id,
|
|
|
|
|
|
summary=result.get("summary", "Task completed"),
|
|
|
|
|
|
findings=result.get("findings", []),
|
|
|
|
|
|
success=True,
|
|
|
|
|
|
)
|
|
|
|
|
|
|
|
|
|
|
|
def on_error(self, error: str) -> None:
|
|
|
|
|
|
"""Agent出错时调用"""
|
|
|
|
|
|
self._state.set_failed(error)
|
|
|
|
|
|
agent_registry.update_agent_status(self._agent_id, "failed", {"error": error})
|
|
|
|
|
|
|
2025-12-11 19:09:10 +08:00
|
|
|
|
@abstractmethod
|
|
|
|
|
|
async def run(self, input_data: Dict[str, Any]) -> AgentResult:
|
|
|
|
|
|
"""
|
|
|
|
|
|
执行 Agent 任务
|
|
|
|
|
|
|
|
|
|
|
|
Args:
|
|
|
|
|
|
input_data: 输入数据
|
|
|
|
|
|
|
|
|
|
|
|
Returns:
|
|
|
|
|
|
Agent 执行结果
|
|
|
|
|
|
"""
|
|
|
|
|
|
pass
|
|
|
|
|
|
|
|
|
|
|
|
def cancel(self):
|
|
|
|
|
|
"""取消执行"""
|
|
|
|
|
|
self._cancelled = True
|
feat(agent): implement comprehensive agent architecture with knowledge base and persistence layer
- Add database migrations for agent checkpoints and tree node tracking
- Implement core agent execution framework with executor, state management, and message handling
- Create knowledge base system with framework-specific modules (Django, FastAPI, Flask, Express, React, Supabase)
- Add vulnerability knowledge modules covering authentication, cryptography, injection, XSS, XXE, SSRF, path traversal, deserialization, and race conditions
- Introduce new agent tools: thinking tool, reporting tool, and agent-specific utilities
- Implement LLM memory compression and prompt caching for improved performance
- Add agent registry and persistence layer for checkpoint management
- Refactor agent implementations (analysis, recon, verification, orchestrator) with enhanced capabilities
- Remove legacy agent implementations (analysis_v2, react_agent)
- Update API endpoints for agent task creation and project management
- Add frontend components for agent task creation and enhanced audit UI
- Consolidate agent service architecture with improved separation of concerns
- This refactoring provides a scalable foundation for multi-agent collaboration with knowledge-driven decision making and state persistence
2025-12-12 15:27:12 +08:00
|
|
|
|
logger.info(f"[{self.name}] Cancel requested")
|
2025-12-11 19:09:10 +08:00
|
|
|
|
|
2025-12-17 11:02:42 +08:00
|
|
|
|
# 🔥 外部取消检查回调
|
|
|
|
|
|
self._cancel_callback = None
|
|
|
|
|
|
|
|
|
|
|
|
def set_cancel_callback(self, callback) -> None:
|
|
|
|
|
|
"""设置外部取消检查回调"""
|
|
|
|
|
|
self._cancel_callback = callback
|
|
|
|
|
|
|
2025-12-11 19:09:10 +08:00
|
|
|
|
@property
|
|
|
|
|
|
def is_cancelled(self) -> bool:
|
2025-12-17 11:02:42 +08:00
|
|
|
|
"""检查是否已取消(包含内部标志和外部回调)"""
|
|
|
|
|
|
if self._cancelled:
|
|
|
|
|
|
return True
|
|
|
|
|
|
# 检查外部回调
|
|
|
|
|
|
if self._cancel_callback and self._cancel_callback():
|
|
|
|
|
|
self._cancelled = True
|
|
|
|
|
|
logger.info(f"[{self.name}] Detected cancellation from callback")
|
|
|
|
|
|
return True
|
|
|
|
|
|
return False
|
2025-12-11 19:09:10 +08:00
|
|
|
|
|
2025-12-11 23:29:04 +08:00
|
|
|
|
# ============ 协作方法 ============
|
|
|
|
|
|
|
|
|
|
|
|
def receive_handoff(self, handoff: TaskHandoff):
|
|
|
|
|
|
"""
|
|
|
|
|
|
接收来自前序 Agent 的任务交接
|
|
|
|
|
|
|
|
|
|
|
|
Args:
|
|
|
|
|
|
handoff: 任务交接对象
|
|
|
|
|
|
"""
|
|
|
|
|
|
self._incoming_handoff = handoff
|
|
|
|
|
|
logger.info(
|
|
|
|
|
|
f"[{self.name}] Received handoff from {handoff.from_agent}: "
|
|
|
|
|
|
f"{handoff.summary[:50]}..."
|
|
|
|
|
|
)
|
|
|
|
|
|
|
|
|
|
|
|
def get_handoff_context(self) -> str:
|
|
|
|
|
|
"""
|
|
|
|
|
|
获取交接上下文(用于构建 LLM prompt)
|
|
|
|
|
|
|
|
|
|
|
|
Returns:
|
|
|
|
|
|
格式化的上下文字符串
|
|
|
|
|
|
"""
|
|
|
|
|
|
if not self._incoming_handoff:
|
|
|
|
|
|
return ""
|
|
|
|
|
|
return self._incoming_handoff.to_prompt_context()
|
|
|
|
|
|
|
|
|
|
|
|
def add_insight(self, insight: str):
|
|
|
|
|
|
"""记录洞察"""
|
|
|
|
|
|
self._insights.append(insight)
|
|
|
|
|
|
|
|
|
|
|
|
def record_work(self, work: str):
|
|
|
|
|
|
"""记录完成的工作"""
|
|
|
|
|
|
self._work_completed.append(work)
|
|
|
|
|
|
|
|
|
|
|
|
def create_handoff(
|
|
|
|
|
|
self,
|
|
|
|
|
|
to_agent: str,
|
|
|
|
|
|
summary: str,
|
|
|
|
|
|
key_findings: List[Dict[str, Any]] = None,
|
|
|
|
|
|
suggested_actions: List[Dict[str, Any]] = None,
|
|
|
|
|
|
attention_points: List[str] = None,
|
|
|
|
|
|
priority_areas: List[str] = None,
|
|
|
|
|
|
context_data: Dict[str, Any] = None,
|
|
|
|
|
|
) -> TaskHandoff:
|
|
|
|
|
|
"""
|
|
|
|
|
|
创建任务交接
|
|
|
|
|
|
|
|
|
|
|
|
Args:
|
|
|
|
|
|
to_agent: 目标 Agent
|
|
|
|
|
|
summary: 工作摘要
|
|
|
|
|
|
key_findings: 关键发现
|
|
|
|
|
|
suggested_actions: 建议的行动
|
|
|
|
|
|
attention_points: 需要关注的点
|
|
|
|
|
|
priority_areas: 优先分析区域
|
|
|
|
|
|
context_data: 上下文数据
|
|
|
|
|
|
|
|
|
|
|
|
Returns:
|
|
|
|
|
|
TaskHandoff 对象
|
|
|
|
|
|
"""
|
|
|
|
|
|
return TaskHandoff(
|
|
|
|
|
|
from_agent=self.name,
|
|
|
|
|
|
to_agent=to_agent,
|
|
|
|
|
|
summary=summary,
|
|
|
|
|
|
work_completed=self._work_completed.copy(),
|
|
|
|
|
|
key_findings=key_findings or [],
|
|
|
|
|
|
insights=self._insights.copy(),
|
|
|
|
|
|
suggested_actions=suggested_actions or [],
|
|
|
|
|
|
attention_points=attention_points or [],
|
|
|
|
|
|
priority_areas=priority_areas or [],
|
|
|
|
|
|
context_data=context_data or {},
|
|
|
|
|
|
)
|
|
|
|
|
|
|
|
|
|
|
|
def build_prompt_with_handoff(self, base_prompt: str) -> str:
|
|
|
|
|
|
"""
|
|
|
|
|
|
构建包含交接上下文的 prompt
|
|
|
|
|
|
|
|
|
|
|
|
Args:
|
|
|
|
|
|
base_prompt: 基础 prompt
|
|
|
|
|
|
|
|
|
|
|
|
Returns:
|
|
|
|
|
|
增强后的 prompt
|
|
|
|
|
|
"""
|
|
|
|
|
|
handoff_context = self.get_handoff_context()
|
|
|
|
|
|
if not handoff_context:
|
|
|
|
|
|
return base_prompt
|
|
|
|
|
|
|
|
|
|
|
|
return f"""{base_prompt}
|
|
|
|
|
|
|
|
|
|
|
|
---
|
|
|
|
|
|
## 前序 Agent 交接信息
|
|
|
|
|
|
|
|
|
|
|
|
{handoff_context}
|
|
|
|
|
|
|
|
|
|
|
|
---
|
|
|
|
|
|
请基于以上来自前序 Agent 的信息,结合你的专业能力开展工作。
|
|
|
|
|
|
"""
|
|
|
|
|
|
|
2025-12-11 21:14:32 +08:00
|
|
|
|
# ============ 核心事件发射方法 ============
|
|
|
|
|
|
|
2025-12-11 19:09:10 +08:00
|
|
|
|
async def emit_event(
|
|
|
|
|
|
self,
|
|
|
|
|
|
event_type: str,
|
|
|
|
|
|
message: str,
|
|
|
|
|
|
**kwargs
|
|
|
|
|
|
):
|
|
|
|
|
|
"""发射事件"""
|
|
|
|
|
|
if self.event_emitter:
|
|
|
|
|
|
from ..event_manager import AgentEventData
|
2025-12-12 10:39:32 +08:00
|
|
|
|
|
|
|
|
|
|
# 准备 metadata
|
|
|
|
|
|
metadata = kwargs.get("metadata", {}) or {}
|
|
|
|
|
|
if "agent_name" not in metadata:
|
|
|
|
|
|
metadata["agent_name"] = self.name
|
|
|
|
|
|
|
|
|
|
|
|
# 分离已知字段和未知字段
|
|
|
|
|
|
known_fields = {
|
|
|
|
|
|
"phase", "tool_name", "tool_input", "tool_output",
|
|
|
|
|
|
"tool_duration_ms", "finding_id", "tokens_used"
|
|
|
|
|
|
}
|
|
|
|
|
|
|
|
|
|
|
|
event_kwargs = {}
|
|
|
|
|
|
for k, v in kwargs.items():
|
|
|
|
|
|
if k in known_fields:
|
|
|
|
|
|
event_kwargs[k] = v
|
|
|
|
|
|
elif k != "metadata":
|
|
|
|
|
|
# 将未知字段放入 metadata
|
|
|
|
|
|
metadata[k] = v
|
|
|
|
|
|
|
2025-12-11 19:09:10 +08:00
|
|
|
|
await self.event_emitter.emit(AgentEventData(
|
|
|
|
|
|
event_type=event_type,
|
|
|
|
|
|
message=message,
|
2025-12-12 10:39:32 +08:00
|
|
|
|
metadata=metadata,
|
|
|
|
|
|
**event_kwargs
|
2025-12-11 19:09:10 +08:00
|
|
|
|
))
|
|
|
|
|
|
|
2025-12-11 21:14:32 +08:00
|
|
|
|
# ============ LLM 思考相关事件 ============
|
|
|
|
|
|
|
2025-12-11 19:09:10 +08:00
|
|
|
|
async def emit_thinking(self, message: str):
|
2025-12-11 21:14:32 +08:00
|
|
|
|
"""发射 LLM 思考事件"""
|
2025-12-12 10:39:32 +08:00
|
|
|
|
await self.emit_event("thinking", message)
|
2025-12-11 21:14:32 +08:00
|
|
|
|
|
|
|
|
|
|
async def emit_llm_start(self, iteration: int):
|
|
|
|
|
|
"""发射 LLM 开始思考事件"""
|
|
|
|
|
|
await self.emit_event(
|
|
|
|
|
|
"llm_start",
|
2025-12-11 23:29:04 +08:00
|
|
|
|
f"[{self.name}] 第 {iteration} 轮迭代开始",
|
2025-12-11 21:14:32 +08:00
|
|
|
|
metadata={"iteration": iteration}
|
|
|
|
|
|
)
|
|
|
|
|
|
|
|
|
|
|
|
async def emit_llm_thought(self, thought: str, iteration: int):
|
|
|
|
|
|
"""发射 LLM 思考内容事件 - 这是核心!展示 LLM 在想什么"""
|
|
|
|
|
|
# 截断过长的思考内容
|
|
|
|
|
|
display_thought = thought[:500] + "..." if len(thought) > 500 else thought
|
|
|
|
|
|
await self.emit_event(
|
|
|
|
|
|
"llm_thought",
|
2025-12-11 23:29:04 +08:00
|
|
|
|
f"[{self.name}] 思考: {display_thought}",
|
2025-12-11 21:14:32 +08:00
|
|
|
|
metadata={
|
|
|
|
|
|
"thought": thought,
|
|
|
|
|
|
"iteration": iteration,
|
|
|
|
|
|
}
|
|
|
|
|
|
)
|
|
|
|
|
|
|
2025-12-11 23:29:04 +08:00
|
|
|
|
async def emit_thinking_start(self):
|
|
|
|
|
|
"""发射开始思考事件(流式输出用)"""
|
2025-12-12 10:39:32 +08:00
|
|
|
|
await self.emit_event("thinking_start", "开始思考...")
|
2025-12-11 23:29:04 +08:00
|
|
|
|
|
|
|
|
|
|
async def emit_thinking_token(self, token: str, accumulated: str):
|
|
|
|
|
|
"""发射思考 token 事件(流式输出用)"""
|
|
|
|
|
|
await self.emit_event(
|
|
|
|
|
|
"thinking_token",
|
|
|
|
|
|
"", # 不需要 message,前端从 metadata 获取
|
|
|
|
|
|
metadata={
|
|
|
|
|
|
"token": token,
|
|
|
|
|
|
"accumulated": accumulated,
|
|
|
|
|
|
}
|
|
|
|
|
|
)
|
|
|
|
|
|
|
|
|
|
|
|
async def emit_thinking_end(self, full_response: str):
|
|
|
|
|
|
"""发射思考结束事件(流式输出用)"""
|
|
|
|
|
|
await self.emit_event(
|
|
|
|
|
|
"thinking_end",
|
2025-12-12 10:39:32 +08:00
|
|
|
|
"思考完成",
|
2025-12-11 23:29:04 +08:00
|
|
|
|
metadata={"accumulated": full_response}
|
|
|
|
|
|
)
|
|
|
|
|
|
|
2025-12-11 21:14:32 +08:00
|
|
|
|
async def emit_llm_decision(self, decision: str, reason: str = ""):
|
|
|
|
|
|
"""发射 LLM 决策事件 - 展示 LLM 做了什么决定"""
|
|
|
|
|
|
await self.emit_event(
|
|
|
|
|
|
"llm_decision",
|
2025-12-11 23:29:04 +08:00
|
|
|
|
f"[{self.name}] 决策: {decision}" + (f" ({reason})" if reason else ""),
|
2025-12-11 21:14:32 +08:00
|
|
|
|
metadata={
|
|
|
|
|
|
"decision": decision,
|
|
|
|
|
|
"reason": reason,
|
|
|
|
|
|
}
|
|
|
|
|
|
)
|
|
|
|
|
|
|
2025-12-11 23:29:04 +08:00
|
|
|
|
async def emit_llm_complete(self, result_summary: str, tokens_used: int):
|
|
|
|
|
|
"""发射 LLM 完成事件"""
|
|
|
|
|
|
await self.emit_event(
|
|
|
|
|
|
"llm_complete",
|
|
|
|
|
|
f"[{self.name}] 完成: {result_summary} (消耗 {tokens_used} tokens)",
|
|
|
|
|
|
metadata={
|
|
|
|
|
|
"tokens_used": tokens_used,
|
|
|
|
|
|
}
|
|
|
|
|
|
)
|
|
|
|
|
|
|
2025-12-11 21:14:32 +08:00
|
|
|
|
async def emit_llm_action(self, action: str, action_input: Dict):
|
2025-12-11 23:29:04 +08:00
|
|
|
|
"""发射 LLM 动作决策事件"""
|
2025-12-11 21:14:32 +08:00
|
|
|
|
await self.emit_event(
|
|
|
|
|
|
"llm_action",
|
2025-12-11 23:29:04 +08:00
|
|
|
|
f"[{self.name}] 执行动作: {action}",
|
2025-12-11 21:14:32 +08:00
|
|
|
|
metadata={
|
|
|
|
|
|
"action": action,
|
|
|
|
|
|
"action_input": action_input,
|
|
|
|
|
|
}
|
|
|
|
|
|
)
|
|
|
|
|
|
|
|
|
|
|
|
async def emit_llm_observation(self, observation: str):
|
2025-12-11 23:29:04 +08:00
|
|
|
|
"""发射 LLM 观察事件"""
|
|
|
|
|
|
# 截断过长的观察结果
|
2025-12-11 21:14:32 +08:00
|
|
|
|
display_obs = observation[:300] + "..." if len(observation) > 300 else observation
|
|
|
|
|
|
await self.emit_event(
|
|
|
|
|
|
"llm_observation",
|
2025-12-11 23:29:04 +08:00
|
|
|
|
f"[{self.name}] 观察结果: {display_obs}",
|
2025-12-11 21:14:32 +08:00
|
|
|
|
metadata={
|
2025-12-11 23:29:04 +08:00
|
|
|
|
"observation": observation[:2000], # 限制存储长度
|
2025-12-11 21:14:32 +08:00
|
|
|
|
}
|
|
|
|
|
|
)
|
|
|
|
|
|
|
|
|
|
|
|
# ============ 工具调用相关事件 ============
|
2025-12-11 19:09:10 +08:00
|
|
|
|
|
|
|
|
|
|
async def emit_tool_call(self, tool_name: str, tool_input: Dict):
|
2025-12-11 23:29:04 +08:00
|
|
|
|
"""发射工具调用事件"""
|
2025-12-11 19:09:10 +08:00
|
|
|
|
await self.emit_event(
|
|
|
|
|
|
"tool_call",
|
2025-12-11 23:29:04 +08:00
|
|
|
|
f"[{self.name}] 调用工具: {tool_name}",
|
2025-12-11 19:09:10 +08:00
|
|
|
|
tool_name=tool_name,
|
|
|
|
|
|
tool_input=tool_input,
|
|
|
|
|
|
)
|
|
|
|
|
|
|
|
|
|
|
|
async def emit_tool_result(self, tool_name: str, result: str, duration_ms: int):
|
|
|
|
|
|
"""发射工具结果事件"""
|
2025-12-15 10:24:58 +08:00
|
|
|
|
# 🔥 修复:确保 result 不为 None,避免显示 "None" 字符串
|
|
|
|
|
|
safe_result = result if result and result != "None" else ""
|
|
|
|
|
|
tool_output_dict = {"result": safe_result[:2000] if safe_result else ""} # 截断长输出
|
2025-12-11 19:09:10 +08:00
|
|
|
|
await self.emit_event(
|
|
|
|
|
|
"tool_result",
|
2025-12-11 23:29:04 +08:00
|
|
|
|
f"[{self.name}] 工具 {tool_name} 完成 ({duration_ms}ms)",
|
2025-12-11 19:09:10 +08:00
|
|
|
|
tool_name=tool_name,
|
2025-12-13 12:35:03 +08:00
|
|
|
|
tool_output=tool_output_dict,
|
2025-12-11 19:09:10 +08:00
|
|
|
|
tool_duration_ms=duration_ms,
|
|
|
|
|
|
)
|
|
|
|
|
|
|
2025-12-11 21:14:32 +08:00
|
|
|
|
# ============ 发现相关事件 ============
|
2025-12-13 18:45:05 +08:00
|
|
|
|
|
|
|
|
|
|
async def emit_finding(self, title: str, severity: str, vuln_type: str, file_path: str = "", is_verified: bool = False):
|
2025-12-11 21:14:32 +08:00
|
|
|
|
"""发射漏洞发现事件"""
|
2025-12-13 18:45:05 +08:00
|
|
|
|
import uuid
|
|
|
|
|
|
finding_id = str(uuid.uuid4())
|
|
|
|
|
|
|
|
|
|
|
|
# 🔥 使用 EventManager.emit_finding 发送正确的事件类型
|
|
|
|
|
|
if self.event_emitter and hasattr(self.event_emitter, 'emit_finding'):
|
|
|
|
|
|
await self.event_emitter.emit_finding(
|
|
|
|
|
|
finding_id=finding_id,
|
|
|
|
|
|
title=title,
|
|
|
|
|
|
severity=severity,
|
|
|
|
|
|
vulnerability_type=vuln_type,
|
|
|
|
|
|
is_verified=is_verified,
|
|
|
|
|
|
)
|
|
|
|
|
|
else:
|
|
|
|
|
|
# 回退:使用通用事件发射
|
|
|
|
|
|
severity_emoji = {
|
|
|
|
|
|
"critical": "🔴",
|
|
|
|
|
|
"high": "🟠",
|
|
|
|
|
|
"medium": "🟡",
|
|
|
|
|
|
"low": "🟢",
|
|
|
|
|
|
}.get(severity.lower(), "⚪")
|
|
|
|
|
|
|
|
|
|
|
|
event_type = "finding_verified" if is_verified else "finding_new"
|
|
|
|
|
|
await self.emit_event(
|
|
|
|
|
|
event_type,
|
|
|
|
|
|
f"{severity_emoji} [{self.name}] 发现漏洞: [{severity.upper()}] {title}\n 类型: {vuln_type}\n 位置: {file_path}",
|
|
|
|
|
|
metadata={
|
|
|
|
|
|
"id": finding_id,
|
|
|
|
|
|
"title": title,
|
|
|
|
|
|
"severity": severity,
|
|
|
|
|
|
"vulnerability_type": vuln_type,
|
|
|
|
|
|
"file_path": file_path,
|
|
|
|
|
|
"is_verified": is_verified,
|
|
|
|
|
|
}
|
|
|
|
|
|
)
|
2025-12-11 21:14:32 +08:00
|
|
|
|
|
|
|
|
|
|
# ============ 通用工具方法 ============
|
|
|
|
|
|
|
2025-12-11 19:09:10 +08:00
|
|
|
|
async def call_tool(self, tool_name: str, **kwargs) -> Any:
|
|
|
|
|
|
"""
|
|
|
|
|
|
调用工具
|
|
|
|
|
|
|
|
|
|
|
|
Args:
|
|
|
|
|
|
tool_name: 工具名称
|
|
|
|
|
|
**kwargs: 工具参数
|
|
|
|
|
|
|
|
|
|
|
|
Returns:
|
|
|
|
|
|
工具执行结果
|
|
|
|
|
|
"""
|
|
|
|
|
|
tool = self.tools.get(tool_name)
|
|
|
|
|
|
if not tool:
|
|
|
|
|
|
logger.warning(f"Tool not found: {tool_name}")
|
|
|
|
|
|
return None
|
|
|
|
|
|
|
|
|
|
|
|
self._tool_calls += 1
|
|
|
|
|
|
await self.emit_tool_call(tool_name, kwargs)
|
|
|
|
|
|
|
|
|
|
|
|
import time
|
|
|
|
|
|
start = time.time()
|
|
|
|
|
|
|
|
|
|
|
|
result = await tool.execute(**kwargs)
|
|
|
|
|
|
|
|
|
|
|
|
duration_ms = int((time.time() - start) * 1000)
|
2025-12-11 21:14:32 +08:00
|
|
|
|
await self.emit_tool_result(tool_name, str(result.data)[:500], duration_ms)
|
2025-12-11 19:09:10 +08:00
|
|
|
|
|
|
|
|
|
|
return result
|
|
|
|
|
|
|
|
|
|
|
|
async def call_llm(
|
|
|
|
|
|
self,
|
|
|
|
|
|
messages: List[Dict[str, str]],
|
|
|
|
|
|
tools: Optional[List[Dict]] = None,
|
|
|
|
|
|
) -> Dict[str, Any]:
|
|
|
|
|
|
"""
|
|
|
|
|
|
调用 LLM
|
|
|
|
|
|
|
|
|
|
|
|
Args:
|
|
|
|
|
|
messages: 消息列表
|
|
|
|
|
|
tools: 可用工具描述
|
|
|
|
|
|
|
|
|
|
|
|
Returns:
|
|
|
|
|
|
LLM 响应
|
|
|
|
|
|
"""
|
|
|
|
|
|
self._iteration += 1
|
|
|
|
|
|
|
|
|
|
|
|
try:
|
|
|
|
|
|
response = await self.llm_service.chat_completion(
|
|
|
|
|
|
messages=messages,
|
|
|
|
|
|
temperature=self.config.temperature,
|
|
|
|
|
|
max_tokens=self.config.max_tokens,
|
|
|
|
|
|
tools=tools,
|
|
|
|
|
|
)
|
|
|
|
|
|
|
|
|
|
|
|
if response.get("usage"):
|
|
|
|
|
|
self._total_tokens += response["usage"].get("total_tokens", 0)
|
|
|
|
|
|
|
|
|
|
|
|
return response
|
|
|
|
|
|
|
|
|
|
|
|
except Exception as e:
|
|
|
|
|
|
logger.error(f"LLM call failed: {e}")
|
|
|
|
|
|
raise
|
|
|
|
|
|
|
|
|
|
|
|
def get_tool_descriptions(self) -> List[Dict[str, Any]]:
|
|
|
|
|
|
"""获取工具描述(用于 LLM)"""
|
|
|
|
|
|
descriptions = []
|
|
|
|
|
|
|
|
|
|
|
|
for name, tool in self.tools.items():
|
|
|
|
|
|
if name.startswith("_"):
|
|
|
|
|
|
continue
|
|
|
|
|
|
|
|
|
|
|
|
desc = {
|
|
|
|
|
|
"type": "function",
|
|
|
|
|
|
"function": {
|
|
|
|
|
|
"name": name,
|
|
|
|
|
|
"description": tool.description,
|
|
|
|
|
|
}
|
|
|
|
|
|
}
|
|
|
|
|
|
|
|
|
|
|
|
# 添加参数 schema
|
|
|
|
|
|
if hasattr(tool, 'args_schema') and tool.args_schema:
|
|
|
|
|
|
desc["function"]["parameters"] = tool.args_schema.schema()
|
|
|
|
|
|
|
|
|
|
|
|
descriptions.append(desc)
|
|
|
|
|
|
|
|
|
|
|
|
return descriptions
|
|
|
|
|
|
|
|
|
|
|
|
def get_stats(self) -> Dict[str, Any]:
|
|
|
|
|
|
"""获取执行统计"""
|
|
|
|
|
|
return {
|
|
|
|
|
|
"agent": self.name,
|
|
|
|
|
|
"type": self.agent_type.value,
|
|
|
|
|
|
"iterations": self._iteration,
|
|
|
|
|
|
"tool_calls": self._tool_calls,
|
|
|
|
|
|
"tokens_used": self._total_tokens,
|
|
|
|
|
|
}
|
2025-12-11 23:29:04 +08:00
|
|
|
|
|
feat(agent): implement comprehensive agent architecture with knowledge base and persistence layer
- Add database migrations for agent checkpoints and tree node tracking
- Implement core agent execution framework with executor, state management, and message handling
- Create knowledge base system with framework-specific modules (Django, FastAPI, Flask, Express, React, Supabase)
- Add vulnerability knowledge modules covering authentication, cryptography, injection, XSS, XXE, SSRF, path traversal, deserialization, and race conditions
- Introduce new agent tools: thinking tool, reporting tool, and agent-specific utilities
- Implement LLM memory compression and prompt caching for improved performance
- Add agent registry and persistence layer for checkpoint management
- Refactor agent implementations (analysis, recon, verification, orchestrator) with enhanced capabilities
- Remove legacy agent implementations (analysis_v2, react_agent)
- Update API endpoints for agent task creation and project management
- Add frontend components for agent task creation and enhanced audit UI
- Consolidate agent service architecture with improved separation of concerns
- This refactoring provides a scalable foundation for multi-agent collaboration with knowledge-driven decision making and state persistence
2025-12-12 15:27:12 +08:00
|
|
|
|
# ============ Memory Compression ============
|
|
|
|
|
|
|
|
|
|
|
|
def compress_messages_if_needed(
|
|
|
|
|
|
self,
|
|
|
|
|
|
messages: List[Dict[str, str]],
|
|
|
|
|
|
max_tokens: int = 100000,
|
|
|
|
|
|
) -> List[Dict[str, str]]:
|
|
|
|
|
|
"""
|
|
|
|
|
|
如果消息历史过长,自动压缩
|
|
|
|
|
|
|
|
|
|
|
|
Args:
|
|
|
|
|
|
messages: 消息列表
|
|
|
|
|
|
max_tokens: 最大token数
|
|
|
|
|
|
|
|
|
|
|
|
Returns:
|
|
|
|
|
|
压缩后的消息列表
|
|
|
|
|
|
"""
|
|
|
|
|
|
from ...llm.memory_compressor import MemoryCompressor
|
|
|
|
|
|
|
|
|
|
|
|
compressor = MemoryCompressor(max_total_tokens=max_tokens)
|
|
|
|
|
|
|
|
|
|
|
|
if compressor.should_compress(messages):
|
|
|
|
|
|
logger.info(f"[{self.name}] Compressing conversation history...")
|
|
|
|
|
|
compressed = compressor.compress_history(messages)
|
|
|
|
|
|
logger.info(f"[{self.name}] Compressed {len(messages)} -> {len(compressed)} messages")
|
|
|
|
|
|
return compressed
|
|
|
|
|
|
|
|
|
|
|
|
return messages
|
|
|
|
|
|
|
2025-12-11 23:29:04 +08:00
|
|
|
|
# ============ 统一的流式 LLM 调用 ============
|
|
|
|
|
|
|
|
|
|
|
|
async def stream_llm_call(
|
|
|
|
|
|
self,
|
|
|
|
|
|
messages: List[Dict[str, str]],
|
|
|
|
|
|
temperature: float = 0.1,
|
|
|
|
|
|
max_tokens: int = 2048,
|
feat(agent): implement comprehensive agent architecture with knowledge base and persistence layer
- Add database migrations for agent checkpoints and tree node tracking
- Implement core agent execution framework with executor, state management, and message handling
- Create knowledge base system with framework-specific modules (Django, FastAPI, Flask, Express, React, Supabase)
- Add vulnerability knowledge modules covering authentication, cryptography, injection, XSS, XXE, SSRF, path traversal, deserialization, and race conditions
- Introduce new agent tools: thinking tool, reporting tool, and agent-specific utilities
- Implement LLM memory compression and prompt caching for improved performance
- Add agent registry and persistence layer for checkpoint management
- Refactor agent implementations (analysis, recon, verification, orchestrator) with enhanced capabilities
- Remove legacy agent implementations (analysis_v2, react_agent)
- Update API endpoints for agent task creation and project management
- Add frontend components for agent task creation and enhanced audit UI
- Consolidate agent service architecture with improved separation of concerns
- This refactoring provides a scalable foundation for multi-agent collaboration with knowledge-driven decision making and state persistence
2025-12-12 15:27:12 +08:00
|
|
|
|
auto_compress: bool = True,
|
2025-12-11 23:29:04 +08:00
|
|
|
|
) -> Tuple[str, int]:
|
|
|
|
|
|
"""
|
|
|
|
|
|
统一的流式 LLM 调用方法
|
|
|
|
|
|
|
|
|
|
|
|
所有 Agent 共用此方法,避免重复代码
|
|
|
|
|
|
|
|
|
|
|
|
Args:
|
|
|
|
|
|
messages: 消息列表
|
|
|
|
|
|
temperature: 温度
|
|
|
|
|
|
max_tokens: 最大 token 数
|
feat(agent): implement comprehensive agent architecture with knowledge base and persistence layer
- Add database migrations for agent checkpoints and tree node tracking
- Implement core agent execution framework with executor, state management, and message handling
- Create knowledge base system with framework-specific modules (Django, FastAPI, Flask, Express, React, Supabase)
- Add vulnerability knowledge modules covering authentication, cryptography, injection, XSS, XXE, SSRF, path traversal, deserialization, and race conditions
- Introduce new agent tools: thinking tool, reporting tool, and agent-specific utilities
- Implement LLM memory compression and prompt caching for improved performance
- Add agent registry and persistence layer for checkpoint management
- Refactor agent implementations (analysis, recon, verification, orchestrator) with enhanced capabilities
- Remove legacy agent implementations (analysis_v2, react_agent)
- Update API endpoints for agent task creation and project management
- Add frontend components for agent task creation and enhanced audit UI
- Consolidate agent service architecture with improved separation of concerns
- This refactoring provides a scalable foundation for multi-agent collaboration with knowledge-driven decision making and state persistence
2025-12-12 15:27:12 +08:00
|
|
|
|
auto_compress: 是否自动压缩过长的消息历史
|
2025-12-11 23:29:04 +08:00
|
|
|
|
|
|
|
|
|
|
Returns:
|
|
|
|
|
|
(完整响应内容, token数量)
|
|
|
|
|
|
"""
|
feat(agent): implement comprehensive agent architecture with knowledge base and persistence layer
- Add database migrations for agent checkpoints and tree node tracking
- Implement core agent execution framework with executor, state management, and message handling
- Create knowledge base system with framework-specific modules (Django, FastAPI, Flask, Express, React, Supabase)
- Add vulnerability knowledge modules covering authentication, cryptography, injection, XSS, XXE, SSRF, path traversal, deserialization, and race conditions
- Introduce new agent tools: thinking tool, reporting tool, and agent-specific utilities
- Implement LLM memory compression and prompt caching for improved performance
- Add agent registry and persistence layer for checkpoint management
- Refactor agent implementations (analysis, recon, verification, orchestrator) with enhanced capabilities
- Remove legacy agent implementations (analysis_v2, react_agent)
- Update API endpoints for agent task creation and project management
- Add frontend components for agent task creation and enhanced audit UI
- Consolidate agent service architecture with improved separation of concerns
- This refactoring provides a scalable foundation for multi-agent collaboration with knowledge-driven decision making and state persistence
2025-12-12 15:27:12 +08:00
|
|
|
|
# 🔥 自动压缩过长的消息历史
|
|
|
|
|
|
if auto_compress:
|
|
|
|
|
|
messages = self.compress_messages_if_needed(messages)
|
|
|
|
|
|
|
2025-12-11 23:29:04 +08:00
|
|
|
|
accumulated = ""
|
|
|
|
|
|
total_tokens = 0
|
|
|
|
|
|
|
feat(agent): implement comprehensive agent architecture with knowledge base and persistence layer
- Add database migrations for agent checkpoints and tree node tracking
- Implement core agent execution framework with executor, state management, and message handling
- Create knowledge base system with framework-specific modules (Django, FastAPI, Flask, Express, React, Supabase)
- Add vulnerability knowledge modules covering authentication, cryptography, injection, XSS, XXE, SSRF, path traversal, deserialization, and race conditions
- Introduce new agent tools: thinking tool, reporting tool, and agent-specific utilities
- Implement LLM memory compression and prompt caching for improved performance
- Add agent registry and persistence layer for checkpoint management
- Refactor agent implementations (analysis, recon, verification, orchestrator) with enhanced capabilities
- Remove legacy agent implementations (analysis_v2, react_agent)
- Update API endpoints for agent task creation and project management
- Add frontend components for agent task creation and enhanced audit UI
- Consolidate agent service architecture with improved separation of concerns
- This refactoring provides a scalable foundation for multi-agent collaboration with knowledge-driven decision making and state persistence
2025-12-12 15:27:12 +08:00
|
|
|
|
# 🔥 在开始 LLM 调用前检查取消
|
|
|
|
|
|
if self.is_cancelled:
|
|
|
|
|
|
logger.info(f"[{self.name}] Cancelled before LLM call")
|
|
|
|
|
|
return "", 0
|
|
|
|
|
|
|
2025-12-13 20:21:30 +08:00
|
|
|
|
logger.info(f"[{self.name}] 🚀 Starting stream_llm_call, emitting thinking_start...")
|
2025-12-11 23:29:04 +08:00
|
|
|
|
await self.emit_thinking_start()
|
2025-12-13 20:21:30 +08:00
|
|
|
|
logger.info(f"[{self.name}] ✅ thinking_start emitted, starting LLM stream...")
|
2025-12-11 23:29:04 +08:00
|
|
|
|
|
|
|
|
|
|
try:
|
2025-12-17 11:02:42 +08:00
|
|
|
|
# 获取流式迭代器
|
|
|
|
|
|
stream = self.llm_service.chat_completion_stream(
|
2025-12-11 23:29:04 +08:00
|
|
|
|
messages=messages,
|
|
|
|
|
|
temperature=temperature,
|
|
|
|
|
|
max_tokens=max_tokens,
|
2025-12-17 11:02:42 +08:00
|
|
|
|
)
|
|
|
|
|
|
# 兼容不同版本的 python async generator
|
|
|
|
|
|
iterator = stream.__aiter__()
|
|
|
|
|
|
|
|
|
|
|
|
import time
|
|
|
|
|
|
first_token_received = False
|
|
|
|
|
|
last_activity = time.time()
|
|
|
|
|
|
|
|
|
|
|
|
while True:
|
2025-12-11 23:29:04 +08:00
|
|
|
|
# 检查取消
|
|
|
|
|
|
if self.is_cancelled:
|
2025-12-17 11:02:42 +08:00
|
|
|
|
logger.info(f"[{self.name}] Cancelled during LLM streaming loop")
|
2025-12-11 23:29:04 +08:00
|
|
|
|
break
|
|
|
|
|
|
|
2025-12-17 11:02:42 +08:00
|
|
|
|
try:
|
|
|
|
|
|
# 🔥 第一個 token 30秒超时,后续 token 60秒超时
|
|
|
|
|
|
# 这是一个应用层的安全网,防止底层 LLM 客户端挂死
|
|
|
|
|
|
timeout = 30.0 if not first_token_received else 60.0
|
2025-12-11 23:29:04 +08:00
|
|
|
|
|
2025-12-17 11:02:42 +08:00
|
|
|
|
chunk = await asyncio.wait_for(iterator.__anext__(), timeout=timeout)
|
|
|
|
|
|
|
|
|
|
|
|
last_activity = time.time()
|
2025-12-11 23:29:04 +08:00
|
|
|
|
|
2025-12-17 11:02:42 +08:00
|
|
|
|
if chunk["type"] == "token":
|
|
|
|
|
|
first_token_received = True
|
|
|
|
|
|
token = chunk["content"]
|
|
|
|
|
|
# 🔥 累积 content,确保 accumulated 变量更新
|
|
|
|
|
|
# 注意:某些 adapter 返回的 chunk["accumulated"] 可能已经包含了累积值,
|
|
|
|
|
|
# 但为了安全起见,如果不一致,我们自己累积
|
|
|
|
|
|
if "accumulated" in chunk:
|
|
|
|
|
|
accumulated = chunk["accumulated"]
|
|
|
|
|
|
else:
|
|
|
|
|
|
# 如果 adapter 没返回 accumulated,我们自己拼
|
|
|
|
|
|
# 注意:如果是 token 类型,content 是增量
|
|
|
|
|
|
# 如果 accumulated 被覆盖了,需要小心。
|
|
|
|
|
|
# 实际上 service.py 中 chat_completion_stream 保证了 accumulated 存在
|
|
|
|
|
|
# 这里我们信任 service 层的 accumulated
|
|
|
|
|
|
pass
|
|
|
|
|
|
|
|
|
|
|
|
# Double check if accumulated is empty but we have token
|
|
|
|
|
|
if not accumulated and token:
|
|
|
|
|
|
accumulated += token # Fallback
|
|
|
|
|
|
|
|
|
|
|
|
await self.emit_thinking_token(token, accumulated)
|
|
|
|
|
|
# 🔥 CRITICAL: 让出控制权给事件循环,让 SSE 有机会发送事件
|
|
|
|
|
|
await asyncio.sleep(0)
|
|
|
|
|
|
|
|
|
|
|
|
elif chunk["type"] == "done":
|
|
|
|
|
|
accumulated = chunk["content"]
|
|
|
|
|
|
if chunk.get("usage"):
|
|
|
|
|
|
total_tokens = chunk["usage"].get("total_tokens", 0)
|
|
|
|
|
|
break
|
|
|
|
|
|
|
|
|
|
|
|
elif chunk["type"] == "error":
|
|
|
|
|
|
accumulated = chunk.get("accumulated", "")
|
|
|
|
|
|
error_msg = chunk.get("error", "Unknown error")
|
|
|
|
|
|
logger.error(f"[{self.name}] Stream error: {error_msg}")
|
|
|
|
|
|
if accumulated:
|
|
|
|
|
|
total_tokens = chunk.get("usage", {}).get("total_tokens", 0)
|
|
|
|
|
|
else:
|
|
|
|
|
|
accumulated = f"[系统错误: {error_msg}] 请重新思考并输出你的决策。"
|
|
|
|
|
|
break
|
|
|
|
|
|
|
|
|
|
|
|
except StopAsyncIteration:
|
|
|
|
|
|
break
|
|
|
|
|
|
except asyncio.TimeoutError:
|
|
|
|
|
|
timeout_type = "First Token" if not first_token_received else "Stream"
|
|
|
|
|
|
logger.error(f"[{self.name}] LLM {timeout_type} Timeout ({timeout}s)")
|
|
|
|
|
|
error_msg = f"LLM 响应超时 ({timeout_type}, {timeout}s)"
|
|
|
|
|
|
await self.emit_event("error", error_msg)
|
|
|
|
|
|
if not accumulated:
|
|
|
|
|
|
accumulated = f"[超时错误: {timeout}s 无响应] 请尝试简化请求或重试。"
|
2025-12-11 23:29:04 +08:00
|
|
|
|
break
|
|
|
|
|
|
|
|
|
|
|
|
except asyncio.CancelledError:
|
|
|
|
|
|
logger.info(f"[{self.name}] LLM call cancelled")
|
|
|
|
|
|
raise
|
2025-12-12 16:36:39 +08:00
|
|
|
|
except Exception as e:
|
|
|
|
|
|
# 🔥 增强异常处理,避免吞掉错误
|
|
|
|
|
|
logger.error(f"[{self.name}] Unexpected error in stream_llm_call: {e}", exc_info=True)
|
|
|
|
|
|
await self.emit_event("error", f"LLM 调用错误: {str(e)}")
|
|
|
|
|
|
accumulated = f"[LLM调用错误: {str(e)}] 请重试。"
|
2025-12-11 23:29:04 +08:00
|
|
|
|
finally:
|
|
|
|
|
|
await self.emit_thinking_end(accumulated)
|
|
|
|
|
|
|
2025-12-12 16:36:39 +08:00
|
|
|
|
# 🔥 记录空响应警告,帮助调试
|
|
|
|
|
|
if not accumulated or not accumulated.strip():
|
|
|
|
|
|
logger.warning(f"[{self.name}] Empty LLM response returned (total_tokens: {total_tokens})")
|
|
|
|
|
|
|
2025-12-11 23:29:04 +08:00
|
|
|
|
return accumulated, total_tokens
|
|
|
|
|
|
|
|
|
|
|
|
async def execute_tool(self, tool_name: str, tool_input: Dict) -> str:
|
|
|
|
|
|
"""
|
2025-12-16 17:31:29 +08:00
|
|
|
|
统一的工具执行方法 - 支持取消和超时
|
|
|
|
|
|
|
2025-12-11 23:29:04 +08:00
|
|
|
|
Args:
|
|
|
|
|
|
tool_name: 工具名称
|
|
|
|
|
|
tool_input: 工具参数
|
2025-12-16 17:31:29 +08:00
|
|
|
|
|
2025-12-11 23:29:04 +08:00
|
|
|
|
Returns:
|
|
|
|
|
|
工具执行结果字符串
|
|
|
|
|
|
"""
|
feat(agent): implement comprehensive agent architecture with knowledge base and persistence layer
- Add database migrations for agent checkpoints and tree node tracking
- Implement core agent execution framework with executor, state management, and message handling
- Create knowledge base system with framework-specific modules (Django, FastAPI, Flask, Express, React, Supabase)
- Add vulnerability knowledge modules covering authentication, cryptography, injection, XSS, XXE, SSRF, path traversal, deserialization, and race conditions
- Introduce new agent tools: thinking tool, reporting tool, and agent-specific utilities
- Implement LLM memory compression and prompt caching for improved performance
- Add agent registry and persistence layer for checkpoint management
- Refactor agent implementations (analysis, recon, verification, orchestrator) with enhanced capabilities
- Remove legacy agent implementations (analysis_v2, react_agent)
- Update API endpoints for agent task creation and project management
- Add frontend components for agent task creation and enhanced audit UI
- Consolidate agent service architecture with improved separation of concerns
- This refactoring provides a scalable foundation for multi-agent collaboration with knowledge-driven decision making and state persistence
2025-12-12 15:27:12 +08:00
|
|
|
|
# 🔥 在执行工具前检查取消
|
|
|
|
|
|
if self.is_cancelled:
|
2025-12-16 17:31:29 +08:00
|
|
|
|
return "⚠️ 任务已取消"
|
|
|
|
|
|
|
2025-12-11 23:29:04 +08:00
|
|
|
|
tool = self.tools.get(tool_name)
|
2025-12-16 17:31:29 +08:00
|
|
|
|
|
2025-12-11 23:29:04 +08:00
|
|
|
|
if not tool:
|
|
|
|
|
|
return f"错误: 工具 '{tool_name}' 不存在。可用工具: {list(self.tools.keys())}"
|
2025-12-16 17:31:29 +08:00
|
|
|
|
|
2025-12-11 23:29:04 +08:00
|
|
|
|
try:
|
|
|
|
|
|
self._tool_calls += 1
|
|
|
|
|
|
await self.emit_tool_call(tool_name, tool_input)
|
2025-12-16 17:31:29 +08:00
|
|
|
|
|
2025-12-11 23:29:04 +08:00
|
|
|
|
import time
|
|
|
|
|
|
start = time.time()
|
2025-12-16 17:31:29 +08:00
|
|
|
|
|
|
|
|
|
|
# 🔥 根据工具类型设置不同的超时时间
|
|
|
|
|
|
tool_timeouts = {
|
|
|
|
|
|
"semgrep_scan": 120, # 外部扫描工具需要更长时间
|
|
|
|
|
|
"bandit_scan": 90,
|
|
|
|
|
|
"gitleaks_scan": 60,
|
|
|
|
|
|
"npm_audit": 90,
|
|
|
|
|
|
"safety_scan": 60,
|
|
|
|
|
|
"kunlun_scan": 180,
|
|
|
|
|
|
"osv_scanner": 60,
|
|
|
|
|
|
"trufflehog_scan": 90,
|
|
|
|
|
|
"sandbox_exec": 60,
|
|
|
|
|
|
"php_test": 30,
|
|
|
|
|
|
"command_injection_test": 30,
|
|
|
|
|
|
"sql_injection_test": 30,
|
|
|
|
|
|
"xss_test": 30,
|
|
|
|
|
|
}
|
|
|
|
|
|
timeout = tool_timeouts.get(tool_name, 30) # 默认30秒
|
|
|
|
|
|
|
|
|
|
|
|
# 🔥 使用 asyncio.wait_for 添加超时控制,同时支持取消
|
|
|
|
|
|
async def execute_with_cancel_check():
|
|
|
|
|
|
"""包装工具执行,定期检查取消状态"""
|
|
|
|
|
|
# 创建工具执行任务
|
|
|
|
|
|
execute_task = asyncio.create_task(tool.execute(**tool_input))
|
|
|
|
|
|
|
|
|
|
|
|
try:
|
|
|
|
|
|
# 使用循环定期检查取消状态
|
|
|
|
|
|
while not execute_task.done():
|
|
|
|
|
|
if self.is_cancelled:
|
|
|
|
|
|
execute_task.cancel()
|
|
|
|
|
|
try:
|
|
|
|
|
|
await execute_task
|
|
|
|
|
|
except asyncio.CancelledError:
|
|
|
|
|
|
pass
|
|
|
|
|
|
raise asyncio.CancelledError("任务已取消")
|
|
|
|
|
|
|
|
|
|
|
|
# 等待任务完成或超时检查间隔
|
|
|
|
|
|
try:
|
|
|
|
|
|
return await asyncio.wait_for(
|
|
|
|
|
|
asyncio.shield(execute_task),
|
|
|
|
|
|
timeout=0.5 # 每0.5秒检查一次取消状态
|
|
|
|
|
|
)
|
|
|
|
|
|
except asyncio.TimeoutError:
|
|
|
|
|
|
continue # 继续循环检查
|
|
|
|
|
|
|
|
|
|
|
|
return await execute_task
|
|
|
|
|
|
except asyncio.CancelledError:
|
|
|
|
|
|
if not execute_task.done():
|
|
|
|
|
|
execute_task.cancel()
|
|
|
|
|
|
raise
|
|
|
|
|
|
|
|
|
|
|
|
try:
|
|
|
|
|
|
result = await asyncio.wait_for(
|
|
|
|
|
|
execute_with_cancel_check(),
|
|
|
|
|
|
timeout=timeout
|
|
|
|
|
|
)
|
|
|
|
|
|
except asyncio.TimeoutError:
|
|
|
|
|
|
duration_ms = int((time.time() - start) * 1000)
|
|
|
|
|
|
await self.emit_tool_result(tool_name, f"超时 ({timeout}s)", duration_ms)
|
|
|
|
|
|
return f"⚠️ 工具 '{tool_name}' 执行超时 ({timeout}秒),请尝试其他方法或减小操作范围。"
|
|
|
|
|
|
except asyncio.CancelledError:
|
|
|
|
|
|
duration_ms = int((time.time() - start) * 1000)
|
|
|
|
|
|
await self.emit_tool_result(tool_name, "已取消", duration_ms)
|
|
|
|
|
|
return "⚠️ 任务已取消"
|
|
|
|
|
|
|
2025-12-11 23:29:04 +08:00
|
|
|
|
duration_ms = int((time.time() - start) * 1000)
|
2025-12-15 10:24:58 +08:00
|
|
|
|
# 🔥 修复:确保传递有意义的结果字符串,避免 "None"
|
|
|
|
|
|
result_preview = str(result.data)[:200] if result.data is not None else (result.error[:200] if result.error else "")
|
|
|
|
|
|
await self.emit_tool_result(tool_name, result_preview, duration_ms)
|
2025-12-16 17:31:29 +08:00
|
|
|
|
|
|
|
|
|
|
# 🔥 工具执行后再次检查取消
|
|
|
|
|
|
if self.is_cancelled:
|
|
|
|
|
|
return "⚠️ 任务已取消"
|
|
|
|
|
|
|
2025-12-11 23:29:04 +08:00
|
|
|
|
if result.success:
|
|
|
|
|
|
output = str(result.data)
|
2025-12-14 17:20:54 +08:00
|
|
|
|
|
2025-12-11 23:29:04 +08:00
|
|
|
|
# 包含 metadata 中的额外信息
|
|
|
|
|
|
if result.metadata:
|
|
|
|
|
|
if "issues" in result.metadata:
|
|
|
|
|
|
output += f"\n\n发现的问题:\n{json.dumps(result.metadata['issues'], ensure_ascii=False, indent=2)}"
|
|
|
|
|
|
if "findings" in result.metadata:
|
|
|
|
|
|
output += f"\n\n发现:\n{json.dumps(result.metadata['findings'][:10], ensure_ascii=False, indent=2)}"
|
2025-12-14 17:20:54 +08:00
|
|
|
|
|
2025-12-11 23:29:04 +08:00
|
|
|
|
# 截断过长输出
|
|
|
|
|
|
if len(output) > 6000:
|
|
|
|
|
|
output = output[:6000] + f"\n\n... [输出已截断,共 {len(str(result.data))} 字符]"
|
|
|
|
|
|
return output
|
|
|
|
|
|
else:
|
2025-12-14 17:20:54 +08:00
|
|
|
|
# 🔥 输出详细的错误信息,包括原始错误
|
|
|
|
|
|
error_msg = f"""⚠️ 工具执行失败
|
|
|
|
|
|
|
|
|
|
|
|
**工具**: {tool_name}
|
|
|
|
|
|
**参数**: {json.dumps(tool_input, ensure_ascii=False, indent=2) if tool_input else '无'}
|
|
|
|
|
|
**错误**: {result.error}
|
|
|
|
|
|
|
|
|
|
|
|
请根据错误信息调整参数或尝试其他方法。"""
|
|
|
|
|
|
return error_msg
|
|
|
|
|
|
|
2025-12-16 17:31:29 +08:00
|
|
|
|
except asyncio.CancelledError:
|
|
|
|
|
|
logger.info(f"[{self.name}] Tool '{tool_name}' execution cancelled")
|
|
|
|
|
|
return "⚠️ 任务已取消"
|
2025-12-11 23:29:04 +08:00
|
|
|
|
except Exception as e:
|
2025-12-14 17:20:54 +08:00
|
|
|
|
import traceback
|
2025-12-11 23:29:04 +08:00
|
|
|
|
logger.error(f"Tool execution error: {e}")
|
2025-12-14 17:20:54 +08:00
|
|
|
|
# 🔥 输出完整的原始错误信息,包括堆栈跟踪
|
|
|
|
|
|
error_msg = f"""❌ 工具执行异常
|
|
|
|
|
|
|
|
|
|
|
|
**工具**: {tool_name}
|
|
|
|
|
|
**参数**: {json.dumps(tool_input, ensure_ascii=False, indent=2) if tool_input else '无'}
|
|
|
|
|
|
**错误类型**: {type(e).__name__}
|
|
|
|
|
|
**错误信息**: {str(e)}
|
|
|
|
|
|
**堆栈跟踪**:
|
|
|
|
|
|
```
|
|
|
|
|
|
{traceback.format_exc()}
|
|
|
|
|
|
```
|
|
|
|
|
|
|
|
|
|
|
|
请分析错误原因,可能需要:
|
|
|
|
|
|
1. 检查参数格式是否正确
|
|
|
|
|
|
2. 尝试使用其他工具
|
|
|
|
|
|
3. 如果是权限或资源问题,跳过该操作"""
|
|
|
|
|
|
return error_msg
|
2025-12-11 23:29:04 +08:00
|
|
|
|
|
|
|
|
|
|
def get_tools_description(self) -> str:
|
|
|
|
|
|
"""生成工具描述文本(用于 prompt)"""
|
|
|
|
|
|
tools_info = []
|
|
|
|
|
|
for name, tool in self.tools.items():
|
|
|
|
|
|
if name.startswith("_"):
|
|
|
|
|
|
continue
|
|
|
|
|
|
desc = f"- {name}: {getattr(tool, 'description', 'No description')}"
|
|
|
|
|
|
tools_info.append(desc)
|
|
|
|
|
|
return "\n".join(tools_info)
|