parse_cmdline.go
1 // Copyright (c) 2024-2026 Tencent Zhuque Lab. All rights reserved. 2 // 3 // Licensed under the Apache License, Version 2.0 (the "License"); 4 // you may not use this file except in compliance with the License. 5 // You may obtain a copy of the License at 6 // 7 // http://www.apache.org/licenses/LICENSE-2.0 8 // 9 // Unless required by applicable law or agreed to in writing, software 10 // distributed under the License is distributed on an "AS IS" BASIS, 11 // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 // See the License for the specific language governing permissions and 13 // limitations under the License. 14 // 15 // Requirement: Any integration or derivative work must explicitly attribute 16 // Tencent Zhuque Lab (https://github.com/Tencent/AI-Infra-Guard) in its 17 // documentation or user interface, as detailed in the NOTICE file. 18 19 package agent 20 21 import ( 22 "encoding/json" 23 "fmt" 24 "path" 25 26 "github.com/Tencent/AI-Infra-Guard/common/utils" 27 "github.com/Tencent/AI-Infra-Guard/internal/gologger" 28 "github.com/google/uuid" 29 ) 30 31 type CmdNewPlanStep struct { 32 Title string `json:"title"` 33 StepId string `json:"stepId"` 34 } 35 36 type CmdStatusUpdate struct { 37 Brief string `json:"brief"` 38 Description string `json:"description"` 39 StepId string `json:"stepId"` 40 Status string `json:"status"` 41 } 42 43 type CmdToolUsed struct { 44 ToolId string `json:"tool_id"` 45 ToolName string `json:"tool_name"` 46 Brief string `json:"brief"` 47 Status string `json:"status"` 48 StepId string `json:"stepId"` 49 Params string `json:"params,omitempty"` 50 } 51 52 type CmdActionLog struct { 53 ToolId string `json:"tool_id"` 54 ToolName string `json:"tool_name"` 55 Log string `json:"log"` 56 StepId string `json:"stepId"` 57 } 58 59 type CmdContent struct { 60 Type string `json:"type"` 61 Content json.RawMessage `json:"content"` 62 } 63 64 type CmdConfig struct { 65 StatusId string 66 Status string 67 } 68 69 type PromptContent struct { 70 Results []PromptResults `json:"results"` 71 Total int `json:"total"` 72 Score float32 `json:"score"` 73 Attachment string `json:"attachment"` 74 Jailbreak int `json:"jailbreak"` 75 } 76 type PromptResults struct { 77 Status string `json:"status"` 78 ModelName string `json:"modelName"` 79 Vulnerability string `json:"vulnerability"` 80 AttackMethod string `json:"attackMethod"` 81 Input string `json:"input"` 82 Output string `json:"output"` 83 Reason string `json:"reason"` 84 } 85 86 func ParseStdoutLine(server, rootDir string, tasks []SubTask, line string, callbacks TaskCallbacks, config *CmdConfig, upload bool) { 87 var cmd CmdContent 88 if len(line) > 1 { 89 if line[0] == '{' { 90 if err := json.Unmarshal([]byte(line), &cmd); err != nil { 91 fmt.Println(line) 92 return 93 } 94 } else { 95 fmt.Println(line) 96 return 97 } 98 } else { 99 return 100 } 101 switch cmd.Type { 102 case AgentMsgTypeNewPlanStep: 103 var content CmdNewPlanStep 104 if err := json.Unmarshal(cmd.Content, &content); err != nil { 105 gologger.WithError(err).Errorln("Failed to AgentMsgTypeNewPlanStep unmarshal command", cmd.Content) 106 return 107 } 108 callbacks.NewPlanStepCallback(content.StepId, content.Title) 109 // 更新任务状态 110 for i, v := range tasks { 111 if v.Status == SubTaskStatusDone { 112 continue 113 } else { 114 if v.StepId == content.StepId { 115 tasks[i].Status = SubTaskStatusDoing 116 if i > 0 { 117 for j := 0; j < i; j++ { 118 tasks[j].Status = SubTaskStatusDone 119 } 120 } 121 } 122 } 123 } 124 callbacks.PlanUpdateCallback(tasks) 125 case AgentMsgTypeStatusUpdate: 126 var content CmdStatusUpdate 127 if err := json.Unmarshal(cmd.Content, &content); err != nil { 128 gologger.WithError(err).Errorln("Failed to AgentMsgTypeStatusUpdate unmarshal command", cmd.Content) 129 return 130 } 131 if content.Status == AgentStatusRunning { 132 config.StatusId = uuid.NewString() 133 config.Status = "running" 134 } else if content.Status == AgentStatusCompleted { 135 if config.Status == "completed" { 136 config.StatusId = uuid.NewString() 137 } 138 config.Status = "completed" 139 } 140 callbacks.StepStatusUpdateCallback(content.StepId, config.StatusId, content.Status, content.Brief, content.Description) 141 case AgentMsgTypeToolUsed: 142 var content CmdToolUsed 143 if err := json.Unmarshal(cmd.Content, &content); err != nil { 144 gologger.WithError(err).Errorln("Failed to AgentMsgTypeToolUsed unmarshal command", cmd.Content) 145 return 146 } 147 tool := CreateTool(content.ToolId, content.ToolId, statusString(content.Status), content.Brief, content.Brief, "", content.Params) 148 callbacks.ToolUsedCallback(content.StepId, config.StatusId, content.Brief, []Tool{tool}) 149 case AgentMsgTypeActionLog: 150 var content CmdActionLog 151 if err := json.Unmarshal(cmd.Content, &content); err != nil { 152 gologger.WithError(err).Errorln("Failed to AgentMsgTypeActionLog unmarshal command", cmd.Content) 153 return 154 } 155 callbacks.ToolUseLogCallback(content.ToolId, content.ToolName, content.StepId, content.Log) 156 case AgentMsgTypeResultUpdate: 157 for i, _ := range tasks { 158 tasks[i].Status = SubTaskStatusDone 159 } 160 callbacks.PlanUpdateCallback(tasks) 161 var content map[string]interface{} 162 if err := json.Unmarshal(cmd.Content, &content); err != nil { 163 gologger.WithError(err).Errorln("Failed to AgentMsgTypeResultUpdate unmarshal command", cmd.Content) 164 return 165 } 166 if upload { 167 var ret []map[string]interface{} 168 dd, err := json.Marshal(content["content"]) 169 if err != nil { 170 gologger.WithError(err).Errorln("Failed to parse result file json") 171 return 172 } 173 if err := json.Unmarshal(dd, &ret); err != nil { 174 gologger.WithError(err).Errorln("Failed to parse result file") 175 return 176 } 177 gologger.Infoln("开始上传文件") 178 for i, v := range ret { 179 attachment, ok := v["attachment"] 180 if !ok || attachment == "" { 181 continue 182 } 183 info, err := utils.UploadFile(server, path.Join(rootDir, attachment.(string))) 184 if err != nil { 185 gologger.WithError(err).Errorln("Failed to upload file") 186 return 187 } 188 gologger.Infoln("上传文件成功") 189 v["attachment"] = info.Data.FileUrl 190 ret[i] = v 191 } 192 dd, err = json.Marshal(ret) 193 if err != nil { 194 gologger.WithError(err).Errorln("Failed to parse result file json") 195 return 196 } 197 var content2 []map[string]interface{} 198 if err := json.Unmarshal(dd, &content2); err != nil { 199 gologger.WithError(err).Errorln("Failed to parse result file json") 200 return 201 } 202 content["content"] = content2 203 } 204 callbacks.ResultCallback(content) 205 case AgentMsgTypeError: 206 content := string(cmd.Content) 207 callbacks.ErrorCallback(content) 208 } 209 }