CalvinBackup/CyberStrikeAI
1
0
Fork 0
mirror of https://github.com/Ed1s0nZ/CyberStrikeAI.git synced 2026-08-09 04:16:10 +02:00
Code Issues Packages Projects Releases Wiki Activity

Add files via upload

Browse Source
This commit is contained in:
公明
2025-11-08 21:25:09 +08:00
committed by GitHub
parent 8b993b493c
commit 07363476ab
9 changed files with 1126 additions and 40 deletions
Download Patch File Download Diff File Expand all files Collapse all files
internal/agent/agent.go
+37 -5
View File
@@ -276,11 +276,17 @@ func (a *Agent) AgentLoopWithProgress(ctx context.Context, userInput string, his
// 获取可用工具
tools := a.getAvailableTools()
// 发送进度更新
// 发送迭代开始事件
if i == 0 {
sendProgress("progress", "正在分析请求并制定测试策略...", nil)
sendProgress("iteration", "开始分析请求并制定测试策略", map[string]interface{}{
"iteration": i + 1,
"total": maxIterations,
})
} else {
sendProgress("progress", fmt.Sprintf("正在继续分析（第 %d 轮迭代）...", i+1), nil)
sendProgress("iteration", fmt.Sprintf("第 %d 轮迭代", i+1), map[string]interface{}{
"iteration": i + 1,
"total": maxIterations,
})
}
// 记录每次调用OpenAI
@@ -333,6 +339,13 @@ func (a *Agent) AgentLoopWithProgress(ctx context.Context, userInput string, his
// 检查是否有工具调用
if len(choice.Message.ToolCalls) > 0 {
// 如果有思考内容，先发送思考事件
if choice.Message.Content != "" {
sendProgress("thinking", choice.Message.Content, map[string]interface{}{
"iteration": i + 1,
})
}
// 添加assistant消息（包含工具调用）
messages = append(messages, ChatMessage{
Role: "assistant",
@@ -341,7 +354,10 @@ func (a *Agent) AgentLoopWithProgress(ctx context.Context, userInput string, his
})
// 发送工具调用进度
sendProgress("progress", fmt.Sprintf("检测到 %d 个工具调用，开始执行...", len(choice.Message.ToolCalls)), nil)
sendProgress("tool_calls_detected", fmt.Sprintf("检测到 %d 个工具调用", len(choice.Message.ToolCalls)), map[string]interface{}{
"count": len(choice.Message.ToolCalls),
"iteration": i + 1,
})
// 执行所有工具调用
for idx, toolCall := range choice.Message.ToolCalls {
@@ -350,8 +366,11 @@ func (a *Agent) AgentLoopWithProgress(ctx context.Context, userInput string, his
sendProgress("tool_call", fmt.Sprintf("正在调用工具: %s", toolCall.Function.Name), map[string]interface{}{
"toolName": toolCall.Function.Name,
"arguments": string(toolArgsJSON),
"argumentsObj": toolCall.Function.Arguments,
"toolCallId": toolCall.ID,
"index": idx + 1,
"total": len(choice.Message.ToolCalls),
"iteration": i + 1,
})
// 执行工具
@@ -369,9 +388,12 @@ func (a *Agent) AgentLoopWithProgress(ctx context.Context, userInput string, his
sendProgress("tool_result", fmt.Sprintf("工具 %s 执行失败", toolCall.Function.Name), map[string]interface{}{
"toolName": toolCall.Function.Name,
"success": false,
"isError": true,
"error": err.Error(),
"toolCallId": toolCall.ID,
"index": idx + 1,
"total": len(choice.Message.ToolCalls),
"iteration": i + 1,
})
a.logger.Warn("工具执行失败，已返回详细错误信息",
@@ -399,10 +421,13 @@ func (a *Agent) AgentLoopWithProgress(ctx context.Context, userInput string, his
"toolName": toolCall.Function.Name,
"success": !execResult.IsError,
"isError": execResult.IsError,
"result": resultPreview,
"result": execResult.Result, // 完整结果
"resultPreview": resultPreview, // 预览结果
"executionId": execResult.ExecutionID,
"toolCallId": toolCall.ID,
"index": idx + 1,
"total": len(choice.Message.ToolCalls),
"iteration": i + 1,
})
// 如果工具返回了错误，记录日志但不中断流程
@@ -423,6 +448,13 @@ func (a *Agent) AgentLoopWithProgress(ctx context.Context, userInput string, his
Content: choice.Message.Content,
})
// 发送AI思考内容（如果没有工具调用）
if choice.Message.Content != "" {
sendProgress("thinking", choice.Message.Content, map[string]interface{}{
"iteration": i + 1,
})
}
// 如果完成，返回结果
if choice.FinishReason == "stop" {
sendProgress("progress", "正在生成最终回复...", nil)
internal/database/conversation.go
+148 -6
View File
@@ -21,12 +21,13 @@ type Conversation struct {
// Message 消息
type Message struct {
ID string `json:"id"`
ConversationID string `json:"conversationId"`
Role string `json:"role"`
Content string `json:"content"`
MCPExecutionIDs []string `json:"mcpExecutionIds,omitempty"`
CreatedAt time.Time `json:"createdAt"`
ID string `json:"id"`
ConversationID string `json:"conversationId"`
Role string `json:"role"`
Content string `json:"content"`
MCPExecutionIDs []string `json:"mcpExecutionIds,omitempty"`
ProcessDetails []map[string]interface{} `json:"processDetails,omitempty"`
CreatedAt time.Time `json:"createdAt"`
}
// CreateConversation 创建新对话
@@ -91,6 +92,39 @@ func (db *DB) GetConversation(id string) (*Conversation, error) {
}
conv.Messages = messages
// 加载过程详情（按消息ID分组）
processDetailsMap, err := db.GetProcessDetailsByConversation(id)
if err != nil {
db.logger.Warn("加载过程详情失败", zap.Error(err))
processDetailsMap = make(map[string][]ProcessDetail)
}
// 将过程详情附加到对应的消息上
for i := range conv.Messages {
if details, ok := processDetailsMap[conv.Messages[i].ID]; ok {
// 将ProcessDetail转换为JSON格式，以便前端使用
detailsJSON := make([]map[string]interface{}, len(details))
for j, detail := range details {
var data interface{}
if detail.Data != "" {
if err := json.Unmarshal([]byte(detail.Data), &data); err != nil {
db.logger.Warn("解析过程详情数据失败", zap.Error(err))
}
}
detailsJSON[j] = map[string]interface{}{
"id": detail.ID,
"messageId": detail.MessageID,
"conversationId": detail.ConversationID,
"eventType": detail.EventType,
"message": detail.Message,
"data": data,
"createdAt": detail.CreatedAt,
}
}
conv.Messages[i].ProcessDetails = detailsJSON
}
}
return &conv, nil
}
@@ -254,3 +288,111 @@ func (db *DB) GetMessages(conversationID string) ([]Message, error) {
return messages, nil
}
// ProcessDetail 过程详情事件
type ProcessDetail struct {
ID string `json:"id"`
MessageID string `json:"messageId"`
ConversationID string `json:"conversationId"`
EventType string `json:"eventType"` // iteration, thinking, tool_calls_detected, tool_call, tool_result, progress, error
Message string `json:"message"`
Data string `json:"data"` // JSON格式的数据
CreatedAt time.Time `json:"createdAt"`
}
// AddProcessDetail 添加过程详情事件
func (db *DB) AddProcessDetail(messageID, conversationID, eventType, message string, data interface{}) error {
id := uuid.New().String()
var dataJSON string
if data != nil {
jsonData, err := json.Marshal(data)
if err != nil {
db.logger.Warn("序列化过程详情数据失败", zap.Error(err))
} else {
dataJSON = string(jsonData)
}
}
_, err := db.Exec(
"INSERT INTO process_details (id, message_id, conversation_id, event_type, message, data, created_at) VALUES (?, ?, ?, ?, ?, ?, ?)",
id, messageID, conversationID, eventType, message, dataJSON, time.Now(),
)
if err != nil {
return fmt.Errorf("添加过程详情失败: %w", err)
}
return nil
}
// GetProcessDetails 获取消息的过程详情
func (db *DB) GetProcessDetails(messageID string) ([]ProcessDetail, error) {
rows, err := db.Query(
"SELECT id, message_id, conversation_id, event_type, message, data, created_at FROM process_details WHERE message_id = ? ORDER BY created_at ASC",
messageID,
)
if err != nil {
return nil, fmt.Errorf("查询过程详情失败: %w", err)
}
defer rows.Close()
var details []ProcessDetail
for rows.Next() {
var detail ProcessDetail
var createdAt string
if err := rows.Scan(&detail.ID, &detail.MessageID, &detail.ConversationID, &detail.EventType, &detail.Message, &detail.Data, &createdAt); err != nil {
return nil, fmt.Errorf("扫描过程详情失败: %w", err)
}
// 尝试多种时间格式解析
var err error
detail.CreatedAt, err = time.Parse("2006-01-02 15:04:05.999999999-07:00", createdAt)
if err != nil {
detail.CreatedAt, err = time.Parse("2006-01-02 15:04:05", createdAt)
}
if err != nil {
detail.CreatedAt, err = time.Parse(time.RFC3339, createdAt)
}
details = append(details, detail)
}
return details, nil
}
// GetProcessDetailsByConversation 获取对话的所有过程详情（按消息分组）
func (db *DB) GetProcessDetailsByConversation(conversationID string) (map[string][]ProcessDetail, error) {
rows, err := db.Query(
"SELECT id, message_id, conversation_id, event_type, message, data, created_at FROM process_details WHERE conversation_id = ? ORDER BY created_at ASC",
conversationID,
)
if err != nil {
return nil, fmt.Errorf("查询过程详情失败: %w", err)
}
defer rows.Close()
detailsMap := make(map[string][]ProcessDetail)
for rows.Next() {
var detail ProcessDetail
var createdAt string
if err := rows.Scan(&detail.ID, &detail.MessageID, &detail.ConversationID, &detail.EventType, &detail.Message, &detail.Data, &createdAt); err != nil {
return nil, fmt.Errorf("扫描过程详情失败: %w", err)
}
// 尝试多种时间格式解析
var err error
detail.CreatedAt, err = time.Parse("2006-01-02 15:04:05.999999999-07:00", createdAt)
if err != nil {
detail.CreatedAt, err = time.Parse("2006-01-02 15:04:05", createdAt)
}
if err != nil {
detail.CreatedAt, err = time.Parse(time.RFC3339, createdAt)
}
detailsMap[detail.MessageID] = append(detailsMap[detail.MessageID], detail)
}
return detailsMap, nil
}
internal/database/database.go
+20
View File
@@ -61,10 +61,26 @@ func (db *DB) initTables() error {
FOREIGN KEY (conversation_id) REFERENCES conversations(id) ON DELETE CASCADE
);`
// 创建过程详情表
createProcessDetailsTable := `
CREATE TABLE IF NOT EXISTS process_details (
id TEXT PRIMARY KEY,
message_id TEXT NOT NULL,
conversation_id TEXT NOT NULL,
event_type TEXT NOT NULL,
message TEXT,
data TEXT,
created_at DATETIME NOT NULL,
FOREIGN KEY (message_id) REFERENCES messages(id) ON DELETE CASCADE,
FOREIGN KEY (conversation_id) REFERENCES conversations(id) ON DELETE CASCADE
);`
// 创建索引
createIndexes := `
CREATE INDEX IF NOT EXISTS idx_messages_conversation_id ON messages(conversation_id);
CREATE INDEX IF NOT EXISTS idx_conversations_updated_at ON conversations(updated_at);
CREATE INDEX IF NOT EXISTS idx_process_details_message_id ON process_details(message_id);
CREATE INDEX IF NOT EXISTS idx_process_details_conversation_id ON process_details(conversation_id);
`
if _, err := db.Exec(createConversationsTable); err != nil {
@@ -75,6 +91,10 @@ func (db *DB) initTables() error {
return fmt.Errorf("创建messages表失败: %w", err)
}
if _, err := db.Exec(createProcessDetailsTable); err != nil {
return fmt.Errorf("创建process_details表失败: %w", err)
}
if _, err := db.Exec(createIndexes); err != nil {
return fmt.Errorf("创建索引失败: %w", err)
}
internal/handler/agent.go
+49 -5
View File
@@ -220,9 +220,29 @@ func (h *AgentHandler) AgentLoopStream(c *gin.Context) {
h.logger.Error("保存用户消息失败", zap.Error(err))
}
// 创建进度回调函数
// 预先创建助手消息，以便关联过程详情
assistantMsg, err := h.db.AddMessage(conversationID, "assistant", "处理中...", nil)
if err != nil {
h.logger.Error("创建助手消息失败", zap.Error(err))
// 如果创建失败，继续执行但不保存过程详情
assistantMsg = nil
}
// 创建进度回调函数，同时保存到数据库
var assistantMessageID string
if assistantMsg != nil {
assistantMessageID = assistantMsg.ID
}
progressCallback := func(eventType, message string, data interface{}) {
sendEvent(eventType, message, data)
// 保存过程详情到数据库（排除response和done事件，它们会在后面单独处理）
if assistantMessageID != "" && eventType != "response" && eventType != "done" {
if err := h.db.AddProcessDetail(assistantMessageID, conversationID, eventType, message, data); err != nil {
h.logger.Warn("保存过程详情失败", zap.Error(err), zap.String("eventType", eventType))
}
}
}
// 执行Agent Loop，传入进度回调
@@ -231,20 +251,44 @@ func (h *AgentHandler) AgentLoopStream(c *gin.Context) {
if err != nil {
h.logger.Error("Agent Loop执行失败", zap.Error(err))
sendEvent("error", "执行失败: "+err.Error(), nil)
// 保存错误事件
if assistantMessageID != "" {
h.db.AddProcessDetail(assistantMessageID, conversationID, "error", "执行失败: "+err.Error(), nil)
}
sendEvent("done", "", nil)
return
}
// 保存助手回复
_, err = h.db.AddMessage(conversationID, "assistant", result.Response, result.MCPExecutionIDs)
if err != nil {
h.logger.Error("保存助手消息失败", zap.Error(err))
// 更新助手消息内容
if assistantMsg != nil {
_, err = h.db.Exec(
"UPDATE messages SET content = ?, mcp_execution_ids = ? WHERE id = ?",
result.Response,
func() string {
if len(result.MCPExecutionIDs) > 0 {
jsonData, _ := json.Marshal(result.MCPExecutionIDs)
return string(jsonData)
}
return ""
}(),
assistantMessageID,
)
if err != nil {
h.logger.Error("更新助手消息失败", zap.Error(err))
}
} else {
// 如果之前创建失败，现在创建
_, err = h.db.AddMessage(conversationID, "assistant", result.Response, result.MCPExecutionIDs)
if err != nil {
h.logger.Error("保存助手消息失败", zap.Error(err))
}
}
// 发送最终响应
sendEvent("response", result.Response, map[string]interface{}{
"mcpExecutionIds": result.MCPExecutionIDs,
"conversationId": conversationID,
"messageId": assistantMessageID, // 包含消息ID，以便前端关联过程详情
})
sendEvent("done", "", map[string]interface{}{
"conversationId": conversationID,