mirror of
https://github.com/Ed1s0nZ/CyberStrikeAI.git
synced 2026-07-06 20:47:57 +02:00
Add files via upload
This commit is contained in:
+43
-4
@@ -303,16 +303,17 @@ type ProgressCallback func(eventType, message string, data interface{})
|
||||
|
||||
// AgentLoop 执行Agent循环
|
||||
func (a *Agent) AgentLoop(ctx context.Context, userInput string, historyMessages []ChatMessage) (*AgentLoopResult, error) {
|
||||
return a.AgentLoopWithProgress(ctx, userInput, historyMessages, "", nil, nil)
|
||||
return a.AgentLoopWithProgress(ctx, userInput, historyMessages, "", nil, nil, nil)
|
||||
}
|
||||
|
||||
// AgentLoopWithConversationID 执行Agent循环(带对话ID)
|
||||
func (a *Agent) AgentLoopWithConversationID(ctx context.Context, userInput string, historyMessages []ChatMessage, conversationID string) (*AgentLoopResult, error) {
|
||||
return a.AgentLoopWithProgress(ctx, userInput, historyMessages, conversationID, nil, nil)
|
||||
return a.AgentLoopWithProgress(ctx, userInput, historyMessages, conversationID, nil, nil, nil)
|
||||
}
|
||||
|
||||
// AgentLoopWithProgress 执行Agent循环(带进度回调和对话ID)
|
||||
func (a *Agent) AgentLoopWithProgress(ctx context.Context, userInput string, historyMessages []ChatMessage, conversationID string, callback ProgressCallback, roleTools []string) (*AgentLoopResult, error) {
|
||||
// roleSkills: 角色配置的skills列表(用于在系统提示词中提示AI,但不硬编码内容)
|
||||
func (a *Agent) AgentLoopWithProgress(ctx context.Context, userInput string, historyMessages []ChatMessage, conversationID string, callback ProgressCallback, roleTools []string, roleSkills []string) (*AgentLoopResult, error) {
|
||||
// 设置当前对话ID
|
||||
a.mu.Lock()
|
||||
a.currentConversationID = conversationID
|
||||
@@ -411,7 +412,45 @@ func (a *Agent) AgentLoopWithProgress(ctx context.Context, userInput string, his
|
||||
* low(低):影响较小,难以利用或影响范围有限
|
||||
* info(信息):安全配置问题、信息泄露但不直接可利用等
|
||||
- 确保漏洞证明(proof)包含足够的证据,如请求/响应、截图、命令输出等
|
||||
- 在记录漏洞后,继续测试以发现更多问题`
|
||||
- 在记录漏洞后,继续测试以发现更多问题
|
||||
|
||||
技能库(Skills):
|
||||
- 系统提供了技能库(Skills),包含各种安全测试的专业技能和方法论文档
|
||||
- 技能库与知识库的区别:
|
||||
* 知识库(Knowledge Base):用于检索分散的知识片段,适合快速查找特定信息
|
||||
* 技能库(Skills):包含完整的专业技能文档,适合深入学习某个领域的测试方法、工具使用、绕过技巧等
|
||||
- 当你需要特定领域的专业技能时,可以使用以下工具按需获取:
|
||||
* ` + builtin.ToolListSkills + `: 获取所有可用的skills列表,查看有哪些专业技能可用
|
||||
* ` + builtin.ToolReadSkill + `: 读取指定skill的详细内容,获取该领域的专业技能文档
|
||||
- 建议在执行相关任务前,先使用 ` + builtin.ToolListSkills + ` 查看可用skills,然后根据任务需要调用 ` + builtin.ToolReadSkill + ` 获取相关专业技能
|
||||
- 例如:如果需要测试SQL注入,可以先调用 ` + builtin.ToolListSkills + ` 查看是否有sql-injection相关的skill,然后调用 ` + builtin.ToolReadSkill + ` 读取该skill的内容
|
||||
- Skills内容包含完整的测试方法、工具使用、绕过技巧、最佳实践等专业技能文档,可以帮助你更专业地执行任务`
|
||||
|
||||
// 如果角色配置了skills,在系统提示词中提示AI(但不硬编码内容)
|
||||
if len(roleSkills) > 0 {
|
||||
var skillsHint strings.Builder
|
||||
skillsHint.WriteString("\n\n本角色推荐使用的Skills:\n")
|
||||
for i, skillName := range roleSkills {
|
||||
if i > 0 {
|
||||
skillsHint.WriteString("、")
|
||||
}
|
||||
skillsHint.WriteString("`")
|
||||
skillsHint.WriteString(skillName)
|
||||
skillsHint.WriteString("`")
|
||||
}
|
||||
skillsHint.WriteString("\n- 这些skills包含了与本角色相关的专业技能文档,建议在执行相关任务时使用 `")
|
||||
skillsHint.WriteString(builtin.ToolReadSkill)
|
||||
skillsHint.WriteString("` 工具读取这些skills的内容")
|
||||
skillsHint.WriteString("\n- 例如:`")
|
||||
skillsHint.WriteString(builtin.ToolReadSkill)
|
||||
skillsHint.WriteString("(skill_name=\"")
|
||||
skillsHint.WriteString(roleSkills[0])
|
||||
skillsHint.WriteString("\")` 可以读取第一个推荐skill的内容")
|
||||
skillsHint.WriteString("\n- 注意:这些skills的内容不会自动注入,需要你根据任务需要主动调用 `")
|
||||
skillsHint.WriteString(builtin.ToolReadSkill)
|
||||
skillsHint.WriteString("` 工具获取")
|
||||
systemPrompt += skillsHint.String()
|
||||
}
|
||||
|
||||
messages := []ChatMessage{
|
||||
{
|
||||
|
||||
+88
-46
@@ -19,6 +19,7 @@ import (
|
||||
"cyberstrike-ai/internal/mcp/builtin"
|
||||
"cyberstrike-ai/internal/openai"
|
||||
"cyberstrike-ai/internal/security"
|
||||
"cyberstrike-ai/internal/skills"
|
||||
"cyberstrike-ai/internal/storage"
|
||||
|
||||
"github.com/gin-gonic/gin"
|
||||
@@ -215,53 +216,53 @@ func New(cfg *config.Config, log *logger.Logger) (*App, error) {
|
||||
return
|
||||
}
|
||||
|
||||
if hasIndex {
|
||||
// 如果已有索引,只索引新添加或更新的项
|
||||
if len(itemsToIndex) > 0 {
|
||||
log.Logger.Info("检测到已有知识库索引,开始增量索引", zap.Int("count", len(itemsToIndex)))
|
||||
ctx := context.Background()
|
||||
consecutiveFailures := 0
|
||||
var firstFailureItemID string
|
||||
var firstFailureError error
|
||||
failedCount := 0
|
||||
|
||||
for _, itemID := range itemsToIndex {
|
||||
if err := knowledgeIndexer.IndexItem(ctx, itemID); err != nil {
|
||||
failedCount++
|
||||
consecutiveFailures++
|
||||
|
||||
if consecutiveFailures == 1 {
|
||||
firstFailureItemID = itemID
|
||||
firstFailureError = err
|
||||
log.Logger.Warn("索引知识项失败", zap.String("itemId", itemID), zap.Error(err))
|
||||
}
|
||||
|
||||
// 如果连续失败2次,立即停止增量索引
|
||||
if consecutiveFailures >= 2 {
|
||||
log.Logger.Error("连续索引失败次数过多,立即停止增量索引",
|
||||
zap.Int("consecutiveFailures", consecutiveFailures),
|
||||
zap.Int("totalItems", len(itemsToIndex)),
|
||||
zap.String("firstFailureItemId", firstFailureItemID),
|
||||
zap.Error(firstFailureError),
|
||||
)
|
||||
break
|
||||
}
|
||||
continue
|
||||
if hasIndex {
|
||||
// 如果已有索引,只索引新添加或更新的项
|
||||
if len(itemsToIndex) > 0 {
|
||||
log.Logger.Info("检测到已有知识库索引,开始增量索引", zap.Int("count", len(itemsToIndex)))
|
||||
ctx := context.Background()
|
||||
consecutiveFailures := 0
|
||||
var firstFailureItemID string
|
||||
var firstFailureError error
|
||||
failedCount := 0
|
||||
|
||||
for _, itemID := range itemsToIndex {
|
||||
if err := knowledgeIndexer.IndexItem(ctx, itemID); err != nil {
|
||||
failedCount++
|
||||
consecutiveFailures++
|
||||
|
||||
if consecutiveFailures == 1 {
|
||||
firstFailureItemID = itemID
|
||||
firstFailureError = err
|
||||
log.Logger.Warn("索引知识项失败", zap.String("itemId", itemID), zap.Error(err))
|
||||
}
|
||||
|
||||
// 成功时重置连续失败计数
|
||||
if consecutiveFailures > 0 {
|
||||
consecutiveFailures = 0
|
||||
firstFailureItemID = ""
|
||||
firstFailureError = nil
|
||||
|
||||
// 如果连续失败2次,立即停止增量索引
|
||||
if consecutiveFailures >= 2 {
|
||||
log.Logger.Error("连续索引失败次数过多,立即停止增量索引",
|
||||
zap.Int("consecutiveFailures", consecutiveFailures),
|
||||
zap.Int("totalItems", len(itemsToIndex)),
|
||||
zap.String("firstFailureItemId", firstFailureItemID),
|
||||
zap.Error(firstFailureError),
|
||||
)
|
||||
break
|
||||
}
|
||||
continue
|
||||
}
|
||||
|
||||
// 成功时重置连续失败计数
|
||||
if consecutiveFailures > 0 {
|
||||
consecutiveFailures = 0
|
||||
firstFailureItemID = ""
|
||||
firstFailureError = nil
|
||||
}
|
||||
log.Logger.Info("增量索引完成", zap.Int("totalItems", len(itemsToIndex)), zap.Int("failedCount", failedCount))
|
||||
} else {
|
||||
log.Logger.Info("检测到已有知识库索引,没有需要索引的新项或更新项")
|
||||
}
|
||||
return
|
||||
log.Logger.Info("增量索引完成", zap.Int("totalItems", len(itemsToIndex)), zap.Int("failedCount", failedCount))
|
||||
} else {
|
||||
log.Logger.Info("检测到已有知识库索引,没有需要索引的新项或更新项")
|
||||
}
|
||||
return
|
||||
}
|
||||
|
||||
// 只有在没有索引时才自动重建
|
||||
log.Logger.Info("未检测到知识库索引,开始自动构建索引")
|
||||
@@ -278,8 +279,30 @@ func New(cfg *config.Config, log *logger.Logger) (*App, error) {
|
||||
configPath = os.Args[1]
|
||||
}
|
||||
|
||||
// 初始化Skills管理器
|
||||
skillsDir := cfg.SkillsDir
|
||||
if skillsDir == "" {
|
||||
skillsDir = "skills" // 默认目录
|
||||
}
|
||||
// 如果是相对路径,相对于配置文件所在目录
|
||||
configDir := filepath.Dir(configPath)
|
||||
if !filepath.IsAbs(skillsDir) {
|
||||
skillsDir = filepath.Join(configDir, skillsDir)
|
||||
}
|
||||
skillsManager := skills.NewManager(skillsDir, log.Logger)
|
||||
log.Logger.Info("Skills管理器已初始化", zap.String("skillsDir", skillsDir))
|
||||
|
||||
// 注册Skills工具到MCP服务器(让AI可以按需调用,带数据库存储支持统计)
|
||||
// 创建一个适配器,将database.DB适配为SkillStatsStorage接口
|
||||
var skillStatsStorage skills.SkillStatsStorage
|
||||
if db != nil {
|
||||
skillStatsStorage = &skillStatsDBAdapter{db: db}
|
||||
}
|
||||
skills.RegisterSkillsToolWithStorage(mcpServer, skillsManager, skillStatsStorage, log.Logger)
|
||||
|
||||
// 创建处理器
|
||||
agentHandler := handler.NewAgentHandler(agent, db, cfg, log.Logger)
|
||||
agentHandler.SetSkillsManager(skillsManager) // 设置Skills管理器
|
||||
// 如果知识库已启用,设置知识库管理器到AgentHandler以便记录检索日志
|
||||
if knowledgeManager != nil {
|
||||
agentHandler.SetKnowledgeManager(knowledgeManager)
|
||||
@@ -294,6 +317,11 @@ func New(cfg *config.Config, log *logger.Logger) (*App, error) {
|
||||
configHandler := handler.NewConfigHandler(configPath, cfg, mcpServer, executor, agent, attackChainHandler, externalMCPMgr, log.Logger)
|
||||
externalMCPHandler := handler.NewExternalMCPHandler(externalMCPMgr, cfg, configPath, log.Logger)
|
||||
roleHandler := handler.NewRoleHandler(cfg, configPath, log.Logger)
|
||||
roleHandler.SetSkillsManager(skillsManager) // 设置Skills管理器到RoleHandler
|
||||
skillsHandler := handler.NewSkillsHandler(skillsManager, cfg, configPath, log.Logger)
|
||||
if db != nil {
|
||||
skillsHandler.SetDB(db) // 设置数据库连接以便获取调用统计
|
||||
}
|
||||
|
||||
// 创建 App 实例(部分字段稍后填充)
|
||||
app := &App{
|
||||
@@ -371,6 +399,7 @@ func New(cfg *config.Config, log *logger.Logger) (*App, error) {
|
||||
app, // 传递 App 实例以便动态获取 knowledgeHandler
|
||||
vulnerabilityHandler,
|
||||
roleHandler,
|
||||
skillsHandler,
|
||||
mcpServer,
|
||||
authManager,
|
||||
)
|
||||
@@ -432,6 +461,7 @@ func setupRoutes(
|
||||
app *App, // 传递 App 实例以便动态获取 knowledgeHandler
|
||||
vulnerabilityHandler *handler.VulnerabilityHandler,
|
||||
roleHandler *handler.RoleHandler,
|
||||
skillsHandler *handler.SkillsHandler,
|
||||
mcpServer *mcp.Server,
|
||||
authManager *security.AuthManager,
|
||||
) {
|
||||
@@ -660,10 +690,22 @@ func setupRoutes(
|
||||
// 角色管理
|
||||
protected.GET("/roles", roleHandler.GetRoles)
|
||||
protected.GET("/roles/:name", roleHandler.GetRole)
|
||||
protected.GET("/roles/skills/list", roleHandler.GetSkills)
|
||||
protected.POST("/roles", roleHandler.CreateRole)
|
||||
protected.PUT("/roles/:name", roleHandler.UpdateRole)
|
||||
protected.DELETE("/roles/:name", roleHandler.DeleteRole)
|
||||
|
||||
// Skills管理
|
||||
protected.GET("/skills", skillsHandler.GetSkills)
|
||||
protected.GET("/skills/stats", skillsHandler.GetSkillStats)
|
||||
protected.DELETE("/skills/stats", skillsHandler.ClearSkillStats)
|
||||
protected.GET("/skills/:name", skillsHandler.GetSkill)
|
||||
protected.GET("/skills/:name/bound-roles", skillsHandler.GetSkillBoundRoles)
|
||||
protected.POST("/skills", skillsHandler.CreateSkill)
|
||||
protected.PUT("/skills/:name", skillsHandler.UpdateSkill)
|
||||
protected.DELETE("/skills/:name", skillsHandler.DeleteSkill)
|
||||
protected.DELETE("/skills/:name/stats", skillsHandler.ClearSkillStatsByName)
|
||||
|
||||
// MCP端点
|
||||
protected.POST("/mcp", func(c *gin.Context) {
|
||||
mcpServer.HandleHTTP(c.Writer, c.Request)
|
||||
@@ -979,18 +1021,18 @@ func initializeKnowledge(
|
||||
var firstFailureItemID string
|
||||
var firstFailureError error
|
||||
failedCount := 0
|
||||
|
||||
|
||||
for _, itemID := range itemsToIndex {
|
||||
if err := knowledgeIndexer.IndexItem(ctx, itemID); err != nil {
|
||||
failedCount++
|
||||
consecutiveFailures++
|
||||
|
||||
|
||||
if consecutiveFailures == 1 {
|
||||
firstFailureItemID = itemID
|
||||
firstFailureError = err
|
||||
logger.Warn("索引知识项失败", zap.String("itemId", itemID), zap.Error(err))
|
||||
}
|
||||
|
||||
|
||||
// 如果连续失败2次,立即停止增量索引
|
||||
if consecutiveFailures >= 2 {
|
||||
logger.Error("连续索引失败次数过多,立即停止增量索引",
|
||||
@@ -1003,7 +1045,7 @@ func initializeKnowledge(
|
||||
}
|
||||
continue
|
||||
}
|
||||
|
||||
|
||||
// 成功时重置连续失败计数
|
||||
if consecutiveFailures > 0 {
|
||||
consecutiveFailures = 0
|
||||
|
||||
@@ -0,0 +1,40 @@
|
||||
package app
|
||||
|
||||
import (
|
||||
"time"
|
||||
|
||||
"cyberstrike-ai/internal/database"
|
||||
"cyberstrike-ai/internal/skills"
|
||||
)
|
||||
|
||||
// skillStatsDBAdapter 将database.DB适配为skills.SkillStatsStorage接口
|
||||
type skillStatsDBAdapter struct {
|
||||
db *database.DB
|
||||
}
|
||||
|
||||
// UpdateSkillStats 更新Skills统计信息
|
||||
func (a *skillStatsDBAdapter) UpdateSkillStats(skillName string, totalCalls, successCalls, failedCalls int, lastCallTime *time.Time) error {
|
||||
return a.db.UpdateSkillStats(skillName, totalCalls, successCalls, failedCalls, lastCallTime)
|
||||
}
|
||||
|
||||
// LoadSkillStats 加载所有Skills统计信息
|
||||
func (a *skillStatsDBAdapter) LoadSkillStats() (map[string]*skills.SkillStats, error) {
|
||||
dbStats, err := a.db.LoadSkillStats()
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
// 转换为skills.SkillStats格式
|
||||
result := make(map[string]*skills.SkillStats)
|
||||
for name, stat := range dbStats {
|
||||
result[name] = &skills.SkillStats{
|
||||
SkillName: stat.SkillName,
|
||||
TotalCalls: stat.TotalCalls,
|
||||
SuccessCalls: stat.SuccessCalls,
|
||||
FailedCalls: stat.FailedCalls,
|
||||
LastCallTime: stat.LastCallTime,
|
||||
}
|
||||
}
|
||||
|
||||
return result, nil
|
||||
}
|
||||
@@ -25,6 +25,7 @@ type Config struct {
|
||||
Knowledge KnowledgeConfig `yaml:"knowledge,omitempty"`
|
||||
RolesDir string `yaml:"roles_dir,omitempty" json:"roles_dir,omitempty"` // 角色配置文件目录(新方式)
|
||||
Roles map[string]RoleConfig `yaml:"roles,omitempty" json:"roles,omitempty"` // 向后兼容:支持在主配置文件中定义角色
|
||||
SkillsDir string `yaml:"skills_dir,omitempty" json:"skills_dir,omitempty"` // Skills配置文件目录
|
||||
}
|
||||
|
||||
type ServerConfig struct {
|
||||
@@ -581,5 +582,6 @@ type RoleConfig struct {
|
||||
Icon string `yaml:"icon,omitempty" json:"icon,omitempty"` // 角色图标(可选)
|
||||
Tools []string `yaml:"tools,omitempty" json:"tools,omitempty"` // 关联的工具列表(toolKey格式,如 "toolName" 或 "mcpName::toolName")
|
||||
MCPs []string `yaml:"mcps,omitempty" json:"mcps,omitempty"` // 向后兼容:关联的MCP服务器列表(已废弃,使用tools替代)
|
||||
Skills []string `yaml:"skills,omitempty" json:"skills,omitempty"` // 关联的skills列表(skill名称列表,在执行任务前会读取这些skills的内容)
|
||||
Enabled bool `yaml:"enabled" json:"enabled"` // 是否启用
|
||||
}
|
||||
|
||||
@@ -104,6 +104,17 @@ func (db *DB) initTables() error {
|
||||
updated_at DATETIME NOT NULL DEFAULT CURRENT_TIMESTAMP
|
||||
);`
|
||||
|
||||
// 创建Skills统计表
|
||||
createSkillStatsTable := `
|
||||
CREATE TABLE IF NOT EXISTS skill_stats (
|
||||
skill_name TEXT PRIMARY KEY,
|
||||
total_calls INTEGER NOT NULL DEFAULT 0,
|
||||
success_calls INTEGER NOT NULL DEFAULT 0,
|
||||
failed_calls INTEGER NOT NULL DEFAULT 0,
|
||||
last_call_time DATETIME,
|
||||
updated_at DATETIME NOT NULL DEFAULT CURRENT_TIMESTAMP
|
||||
);`
|
||||
|
||||
// 创建攻击链节点表
|
||||
createAttackChainNodesTable := `
|
||||
CREATE TABLE IF NOT EXISTS attack_chain_nodes (
|
||||
@@ -264,6 +275,10 @@ func (db *DB) initTables() error {
|
||||
return fmt.Errorf("创建tool_stats表失败: %w", err)
|
||||
}
|
||||
|
||||
if _, err := db.Exec(createSkillStatsTable); err != nil {
|
||||
return fmt.Errorf("创建skill_stats表失败: %w", err)
|
||||
}
|
||||
|
||||
if _, err := db.Exec(createAttackChainNodesTable); err != nil {
|
||||
return fmt.Errorf("创建attack_chain_nodes表失败: %w", err)
|
||||
}
|
||||
|
||||
@@ -0,0 +1,142 @@
|
||||
package database
|
||||
|
||||
import (
|
||||
"database/sql"
|
||||
"time"
|
||||
|
||||
"go.uber.org/zap"
|
||||
)
|
||||
|
||||
// SkillStats Skills统计信息
|
||||
type SkillStats struct {
|
||||
SkillName string
|
||||
TotalCalls int
|
||||
SuccessCalls int
|
||||
FailedCalls int
|
||||
LastCallTime *time.Time
|
||||
}
|
||||
|
||||
// SaveSkillStats 保存Skills统计信息
|
||||
func (db *DB) SaveSkillStats(skillName string, stats *SkillStats) error {
|
||||
var lastCallTime sql.NullTime
|
||||
if stats.LastCallTime != nil {
|
||||
lastCallTime = sql.NullTime{Time: *stats.LastCallTime, Valid: true}
|
||||
}
|
||||
|
||||
query := `
|
||||
INSERT OR REPLACE INTO skill_stats
|
||||
(skill_name, total_calls, success_calls, failed_calls, last_call_time, updated_at)
|
||||
VALUES (?, ?, ?, ?, ?, ?)
|
||||
`
|
||||
|
||||
_, err := db.Exec(query,
|
||||
skillName,
|
||||
stats.TotalCalls,
|
||||
stats.SuccessCalls,
|
||||
stats.FailedCalls,
|
||||
lastCallTime,
|
||||
time.Now(),
|
||||
)
|
||||
|
||||
if err != nil {
|
||||
db.logger.Error("保存Skills统计信息失败", zap.Error(err), zap.String("skillName", skillName))
|
||||
return err
|
||||
}
|
||||
|
||||
return nil
|
||||
}
|
||||
|
||||
// LoadSkillStats 加载所有Skills统计信息
|
||||
func (db *DB) LoadSkillStats() (map[string]*SkillStats, error) {
|
||||
query := `
|
||||
SELECT skill_name, total_calls, success_calls, failed_calls, last_call_time
|
||||
FROM skill_stats
|
||||
`
|
||||
|
||||
rows, err := db.Query(query)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
defer rows.Close()
|
||||
|
||||
stats := make(map[string]*SkillStats)
|
||||
for rows.Next() {
|
||||
var stat SkillStats
|
||||
var lastCallTime sql.NullTime
|
||||
|
||||
err := rows.Scan(
|
||||
&stat.SkillName,
|
||||
&stat.TotalCalls,
|
||||
&stat.SuccessCalls,
|
||||
&stat.FailedCalls,
|
||||
&lastCallTime,
|
||||
)
|
||||
if err != nil {
|
||||
db.logger.Warn("加载Skills统计信息失败", zap.Error(err))
|
||||
continue
|
||||
}
|
||||
|
||||
if lastCallTime.Valid {
|
||||
stat.LastCallTime = &lastCallTime.Time
|
||||
}
|
||||
|
||||
stats[stat.SkillName] = &stat
|
||||
}
|
||||
|
||||
return stats, nil
|
||||
}
|
||||
|
||||
// UpdateSkillStats 更新Skills统计信息(累加模式)
|
||||
func (db *DB) UpdateSkillStats(skillName string, totalCalls, successCalls, failedCalls int, lastCallTime *time.Time) error {
|
||||
var lastCallTimeSQL sql.NullTime
|
||||
if lastCallTime != nil {
|
||||
lastCallTimeSQL = sql.NullTime{Time: *lastCallTime, Valid: true}
|
||||
}
|
||||
|
||||
query := `
|
||||
INSERT INTO skill_stats (skill_name, total_calls, success_calls, failed_calls, last_call_time, updated_at)
|
||||
VALUES (?, ?, ?, ?, ?, ?)
|
||||
ON CONFLICT(skill_name) DO UPDATE SET
|
||||
total_calls = total_calls + ?,
|
||||
success_calls = success_calls + ?,
|
||||
failed_calls = failed_calls + ?,
|
||||
last_call_time = COALESCE(?, last_call_time),
|
||||
updated_at = ?
|
||||
`
|
||||
|
||||
_, err := db.Exec(query,
|
||||
skillName, totalCalls, successCalls, failedCalls, lastCallTimeSQL, time.Now(),
|
||||
totalCalls, successCalls, failedCalls, lastCallTimeSQL, time.Now(),
|
||||
)
|
||||
|
||||
if err != nil {
|
||||
db.logger.Error("更新Skills统计信息失败", zap.Error(err), zap.String("skillName", skillName))
|
||||
return err
|
||||
}
|
||||
|
||||
return nil
|
||||
}
|
||||
|
||||
// ClearSkillStats 清空所有Skills统计信息
|
||||
func (db *DB) ClearSkillStats() error {
|
||||
query := `DELETE FROM skill_stats`
|
||||
_, err := db.Exec(query)
|
||||
if err != nil {
|
||||
db.logger.Error("清空Skills统计信息失败", zap.Error(err))
|
||||
return err
|
||||
}
|
||||
db.logger.Info("已清空所有Skills统计信息")
|
||||
return nil
|
||||
}
|
||||
|
||||
// ClearSkillStatsByName 清空指定skill的统计信息
|
||||
func (db *DB) ClearSkillStatsByName(skillName string) error {
|
||||
query := `DELETE FROM skill_stats WHERE skill_name = ?`
|
||||
_, err := db.Exec(query, skillName)
|
||||
if err != nil {
|
||||
db.logger.Error("清空指定skill统计信息失败", zap.Error(err), zap.String("skillName", skillName))
|
||||
return err
|
||||
}
|
||||
db.logger.Info("已清空指定skill统计信息", zap.String("skillName", skillName))
|
||||
return nil
|
||||
}
|
||||
@@ -15,6 +15,7 @@ import (
|
||||
"cyberstrike-ai/internal/config"
|
||||
"cyberstrike-ai/internal/database"
|
||||
"cyberstrike-ai/internal/mcp/builtin"
|
||||
"cyberstrike-ai/internal/skills"
|
||||
|
||||
"github.com/gin-gonic/gin"
|
||||
"go.uber.org/zap"
|
||||
@@ -72,6 +73,7 @@ type AgentHandler struct {
|
||||
knowledgeManager interface { // 知识库管理器接口
|
||||
LogRetrieval(conversationID, messageID, query, riskType string, retrievedItems []string) error
|
||||
}
|
||||
skillsManager *skills.Manager // Skills管理器
|
||||
}
|
||||
|
||||
// NewAgentHandler 创建新的Agent处理器
|
||||
@@ -101,6 +103,11 @@ func (h *AgentHandler) SetKnowledgeManager(manager interface {
|
||||
h.knowledgeManager = manager
|
||||
}
|
||||
|
||||
// SetSkillsManager 设置Skills管理器
|
||||
func (h *AgentHandler) SetSkillsManager(manager *skills.Manager) {
|
||||
h.skillsManager = manager
|
||||
}
|
||||
|
||||
// ChatRequest 聊天请求
|
||||
type ChatRequest struct {
|
||||
Message string `json:"message" binding:"required"`
|
||||
@@ -169,6 +176,7 @@ func (h *AgentHandler) AgentLoop(c *gin.Context) {
|
||||
// 应用角色用户提示词和工具配置
|
||||
finalMessage := req.Message
|
||||
var roleTools []string // 角色配置的工具列表
|
||||
var roleSkills []string // 角色配置的skills列表(用于提示AI,但不硬编码内容)
|
||||
if req.Role != "" && req.Role != "默认" {
|
||||
if h.config.Roles != nil {
|
||||
if role, exists := h.config.Roles[req.Role]; exists && role.Enabled {
|
||||
@@ -182,6 +190,11 @@ func (h *AgentHandler) AgentLoop(c *gin.Context) {
|
||||
roleTools = role.Tools
|
||||
h.logger.Info("使用角色配置的工具列表", zap.String("role", req.Role), zap.Int("toolCount", len(roleTools)))
|
||||
}
|
||||
// 获取角色配置的skills列表(用于在系统提示词中提示AI,但不硬编码内容)
|
||||
if len(role.Skills) > 0 {
|
||||
roleSkills = role.Skills
|
||||
h.logger.Info("角色配置了skills,将在系统提示词中提示AI", zap.String("role", req.Role), zap.Int("skillCount", len(roleSkills)), zap.Strings("skills", roleSkills))
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -193,7 +206,8 @@ func (h *AgentHandler) AgentLoop(c *gin.Context) {
|
||||
}
|
||||
|
||||
// 执行Agent Loop,传入历史消息和对话ID(使用包含角色提示词的finalMessage和角色工具列表)
|
||||
result, err := h.agent.AgentLoopWithProgress(c.Request.Context(), finalMessage, agentHistoryMessages, conversationID, nil, roleTools)
|
||||
// 注意:skills不会硬编码注入,但会在系统提示词中提示AI这个角色推荐使用哪些skills
|
||||
result, err := h.agent.AgentLoopWithProgress(c.Request.Context(), finalMessage, agentHistoryMessages, conversationID, nil, roleTools, roleSkills)
|
||||
if err != nil {
|
||||
h.logger.Error("Agent Loop执行失败", zap.Error(err))
|
||||
|
||||
@@ -515,6 +529,10 @@ func (h *AgentHandler) AgentLoopStream(c *gin.Context) {
|
||||
// 因为mcps是MCP服务器名称,不是工具列表
|
||||
h.logger.Info("角色配置使用旧的mcps字段,将使用所有工具", zap.String("role", req.Role))
|
||||
}
|
||||
// 注意:角色配置的skills不再硬编码注入,AI可以通过list_skills和read_skill工具按需调用
|
||||
if len(role.Skills) > 0 {
|
||||
h.logger.Info("角色配置了skills,AI可通过工具按需调用", zap.String("role", req.Role), zap.Int("skillCount", len(role.Skills)), zap.Strings("skills", role.Skills))
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -599,7 +617,18 @@ func (h *AgentHandler) AgentLoopStream(c *gin.Context) {
|
||||
|
||||
// 执行Agent Loop,传入独立的上下文,确保任务不会因客户端断开而中断(使用包含角色提示词的finalMessage和角色工具列表)
|
||||
sendEvent("progress", "正在分析您的请求...", nil)
|
||||
result, err := h.agent.AgentLoopWithProgress(taskCtx, finalMessage, agentHistoryMessages, conversationID, progressCallback, roleTools)
|
||||
// 注意:skills不会硬编码注入,但会在系统提示词中提示AI这个角色推荐使用哪些skills
|
||||
var roleSkills []string // 角色配置的skills列表(用于提示AI,但不硬编码内容)
|
||||
if req.Role != "" && req.Role != "默认" {
|
||||
if h.config.Roles != nil {
|
||||
if role, exists := h.config.Roles[req.Role]; exists && role.Enabled {
|
||||
if len(role.Skills) > 0 {
|
||||
roleSkills = role.Skills
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
result, err := h.agent.AgentLoopWithProgress(taskCtx, finalMessage, agentHistoryMessages, conversationID, progressCallback, roleTools, roleSkills)
|
||||
if err != nil {
|
||||
h.logger.Error("Agent Loop执行失败", zap.Error(err))
|
||||
cause := context.Cause(baseCtx)
|
||||
@@ -1099,6 +1128,7 @@ func (h *AgentHandler) executeBatchQueue(queueID string) {
|
||||
// 应用角色用户提示词和工具配置
|
||||
finalMessage := task.Message
|
||||
var roleTools []string // 角色配置的工具列表
|
||||
var roleSkills []string // 角色配置的skills列表(用于提示AI,但不硬编码内容)
|
||||
if queue.Role != "" && queue.Role != "默认" {
|
||||
if h.config.Roles != nil {
|
||||
if role, exists := h.config.Roles[queue.Role]; exists && role.Enabled {
|
||||
@@ -1112,6 +1142,11 @@ func (h *AgentHandler) executeBatchQueue(queueID string) {
|
||||
roleTools = role.Tools
|
||||
h.logger.Info("使用角色配置的工具列表", zap.String("queueId", queueID), zap.String("taskId", task.ID), zap.String("role", queue.Role), zap.Int("toolCount", len(roleTools)))
|
||||
}
|
||||
// 获取角色配置的skills列表(用于在系统提示词中提示AI,但不硬编码内容)
|
||||
if len(role.Skills) > 0 {
|
||||
roleSkills = role.Skills
|
||||
h.logger.Info("角色配置了skills,将在系统提示词中提示AI", zap.String("queueId", queueID), zap.String("taskId", task.ID), zap.String("role", queue.Role), zap.Int("skillCount", len(roleSkills)), zap.Strings("skills", roleSkills))
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -1144,7 +1179,8 @@ func (h *AgentHandler) executeBatchQueue(queueID string) {
|
||||
// 存储取消函数,以便在取消队列时能够取消当前任务
|
||||
h.batchTaskManager.SetTaskCancel(queueID, cancel)
|
||||
// 使用队列配置的角色工具列表(如果为空,表示使用所有工具)
|
||||
result, err := h.agent.AgentLoopWithProgress(ctx, finalMessage, []agent.ChatMessage{}, conversationID, progressCallback, roleTools)
|
||||
// 注意:skills不会硬编码注入,但会在系统提示词中提示AI这个角色推荐使用哪些skills
|
||||
result, err := h.agent.AgentLoopWithProgress(ctx, finalMessage, []agent.ChatMessage{}, conversationID, progressCallback, roleTools, roleSkills)
|
||||
// 任务执行完成,清理取消函数
|
||||
h.batchTaskManager.SetTaskCancel(queueID, nil)
|
||||
cancel()
|
||||
|
||||
@@ -18,9 +18,15 @@ import (
|
||||
|
||||
// RoleHandler 角色处理器
|
||||
type RoleHandler struct {
|
||||
config *config.Config
|
||||
configPath string
|
||||
logger *zap.Logger
|
||||
config *config.Config
|
||||
configPath string
|
||||
logger *zap.Logger
|
||||
skillsManager SkillsManager // Skills管理器接口(可选)
|
||||
}
|
||||
|
||||
// SkillsManager Skills管理器接口
|
||||
type SkillsManager interface {
|
||||
ListSkills() ([]string, error)
|
||||
}
|
||||
|
||||
// NewRoleHandler 创建新的角色处理器
|
||||
@@ -32,6 +38,34 @@ func NewRoleHandler(cfg *config.Config, configPath string, logger *zap.Logger) *
|
||||
}
|
||||
}
|
||||
|
||||
// SetSkillsManager 设置Skills管理器
|
||||
func (h *RoleHandler) SetSkillsManager(manager SkillsManager) {
|
||||
h.skillsManager = manager
|
||||
}
|
||||
|
||||
// GetSkills 获取所有可用的skills列表
|
||||
func (h *RoleHandler) GetSkills(c *gin.Context) {
|
||||
if h.skillsManager == nil {
|
||||
c.JSON(http.StatusOK, gin.H{
|
||||
"skills": []string{},
|
||||
})
|
||||
return
|
||||
}
|
||||
|
||||
skills, err := h.skillsManager.ListSkills()
|
||||
if err != nil {
|
||||
h.logger.Warn("获取skills列表失败", zap.Error(err))
|
||||
c.JSON(http.StatusOK, gin.H{
|
||||
"skills": []string{},
|
||||
})
|
||||
return
|
||||
}
|
||||
|
||||
c.JSON(http.StatusOK, gin.H{
|
||||
"skills": skills,
|
||||
})
|
||||
}
|
||||
|
||||
// GetRoles 获取所有角色
|
||||
func (h *RoleHandler) GetRoles(c *gin.Context) {
|
||||
if h.config.Roles == nil {
|
||||
|
||||
@@ -0,0 +1,778 @@
|
||||
package handler
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
"net/http"
|
||||
"os"
|
||||
"path/filepath"
|
||||
"regexp"
|
||||
"strings"
|
||||
|
||||
"cyberstrike-ai/internal/config"
|
||||
"cyberstrike-ai/internal/database"
|
||||
"cyberstrike-ai/internal/skills"
|
||||
|
||||
"github.com/gin-gonic/gin"
|
||||
"go.uber.org/zap"
|
||||
"gopkg.in/yaml.v3"
|
||||
)
|
||||
|
||||
// SkillsHandler Skills处理器
|
||||
type SkillsHandler struct {
|
||||
manager *skills.Manager
|
||||
config *config.Config
|
||||
configPath string
|
||||
logger *zap.Logger
|
||||
db *database.DB // 数据库连接(用于获取调用统计)
|
||||
}
|
||||
|
||||
// NewSkillsHandler 创建新的Skills处理器
|
||||
func NewSkillsHandler(manager *skills.Manager, cfg *config.Config, configPath string, logger *zap.Logger) *SkillsHandler {
|
||||
return &SkillsHandler{
|
||||
manager: manager,
|
||||
config: cfg,
|
||||
configPath: configPath,
|
||||
logger: logger,
|
||||
}
|
||||
}
|
||||
|
||||
// SetDB 设置数据库连接(用于获取调用统计)
|
||||
func (h *SkillsHandler) SetDB(db *database.DB) {
|
||||
h.db = db
|
||||
}
|
||||
|
||||
// GetSkills 获取所有skills列表(支持分页和搜索)
|
||||
func (h *SkillsHandler) GetSkills(c *gin.Context) {
|
||||
skillList, err := h.manager.ListSkills()
|
||||
if err != nil {
|
||||
h.logger.Error("获取skills列表失败", zap.Error(err))
|
||||
c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
|
||||
return
|
||||
}
|
||||
|
||||
// 搜索参数
|
||||
searchKeyword := strings.TrimSpace(c.Query("search"))
|
||||
|
||||
// 先加载所有skills的详细信息用于搜索过滤
|
||||
allSkillsInfo := make([]map[string]interface{}, 0, len(skillList))
|
||||
for _, skillName := range skillList {
|
||||
skill, err := h.manager.LoadSkill(skillName)
|
||||
if err != nil {
|
||||
h.logger.Warn("加载skill失败", zap.String("skill", skillName), zap.Error(err))
|
||||
continue
|
||||
}
|
||||
|
||||
// 获取文件信息
|
||||
skillPath := skill.Path
|
||||
skillFile := filepath.Join(skillPath, "SKILL.md")
|
||||
// 尝试其他可能的文件名
|
||||
if _, err := os.Stat(skillFile); os.IsNotExist(err) {
|
||||
alternatives := []string{
|
||||
filepath.Join(skillPath, "skill.md"),
|
||||
filepath.Join(skillPath, "README.md"),
|
||||
filepath.Join(skillPath, "readme.md"),
|
||||
}
|
||||
for _, alt := range alternatives {
|
||||
if _, err := os.Stat(alt); err == nil {
|
||||
skillFile = alt
|
||||
break
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
fileInfo, _ := os.Stat(skillFile)
|
||||
var fileSize int64
|
||||
var modTime string
|
||||
if fileInfo != nil {
|
||||
fileSize = fileInfo.Size()
|
||||
modTime = fileInfo.ModTime().Format("2006-01-02 15:04:05")
|
||||
}
|
||||
|
||||
skillInfo := map[string]interface{}{
|
||||
"name": skill.Name,
|
||||
"description": skill.Description,
|
||||
"path": skill.Path,
|
||||
"file_size": fileSize,
|
||||
"mod_time": modTime,
|
||||
}
|
||||
allSkillsInfo = append(allSkillsInfo, skillInfo)
|
||||
}
|
||||
|
||||
// 如果有搜索关键词,进行过滤
|
||||
filteredSkillsInfo := allSkillsInfo
|
||||
if searchKeyword != "" {
|
||||
keywordLower := strings.ToLower(searchKeyword)
|
||||
filteredSkillsInfo = make([]map[string]interface{}, 0)
|
||||
for _, skillInfo := range allSkillsInfo {
|
||||
name := strings.ToLower(fmt.Sprintf("%v", skillInfo["name"]))
|
||||
description := strings.ToLower(fmt.Sprintf("%v", skillInfo["description"]))
|
||||
path := strings.ToLower(fmt.Sprintf("%v", skillInfo["path"]))
|
||||
|
||||
if strings.Contains(name, keywordLower) ||
|
||||
strings.Contains(description, keywordLower) ||
|
||||
strings.Contains(path, keywordLower) {
|
||||
filteredSkillsInfo = append(filteredSkillsInfo, skillInfo)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// 分页参数
|
||||
limit := 20 // 默认每页20条
|
||||
offset := 0
|
||||
if limitStr := c.Query("limit"); limitStr != "" {
|
||||
if parsed, err := parseInt(limitStr); err == nil && parsed > 0 {
|
||||
// 允许更大的limit用于搜索场景,但设置一个合理的上限(10000)
|
||||
if parsed <= 10000 {
|
||||
limit = parsed
|
||||
} else {
|
||||
limit = 10000
|
||||
}
|
||||
}
|
||||
}
|
||||
if offsetStr := c.Query("offset"); offsetStr != "" {
|
||||
if parsed, err := parseInt(offsetStr); err == nil && parsed >= 0 {
|
||||
offset = parsed
|
||||
}
|
||||
}
|
||||
|
||||
// 计算分页范围
|
||||
total := len(filteredSkillsInfo)
|
||||
start := offset
|
||||
end := offset + limit
|
||||
if start > total {
|
||||
start = total
|
||||
}
|
||||
if end > total {
|
||||
end = total
|
||||
}
|
||||
|
||||
// 获取当前页的skill列表
|
||||
var paginatedSkillsInfo []map[string]interface{}
|
||||
if start < end {
|
||||
paginatedSkillsInfo = filteredSkillsInfo[start:end]
|
||||
} else {
|
||||
paginatedSkillsInfo = []map[string]interface{}{}
|
||||
}
|
||||
|
||||
c.JSON(http.StatusOK, gin.H{
|
||||
"skills": paginatedSkillsInfo,
|
||||
"total": total,
|
||||
"limit": limit,
|
||||
"offset": offset,
|
||||
})
|
||||
}
|
||||
|
||||
// GetSkill 获取单个skill的详细信息
|
||||
func (h *SkillsHandler) GetSkill(c *gin.Context) {
|
||||
skillName := c.Param("name")
|
||||
if skillName == "" {
|
||||
c.JSON(http.StatusBadRequest, gin.H{"error": "skill名称不能为空"})
|
||||
return
|
||||
}
|
||||
|
||||
skill, err := h.manager.LoadSkill(skillName)
|
||||
if err != nil {
|
||||
h.logger.Warn("加载skill失败", zap.String("skill", skillName), zap.Error(err))
|
||||
c.JSON(http.StatusNotFound, gin.H{"error": "skill不存在: " + err.Error()})
|
||||
return
|
||||
}
|
||||
|
||||
// 获取文件信息
|
||||
skillPath := skill.Path
|
||||
skillFile := filepath.Join(skillPath, "SKILL.md")
|
||||
if _, err := os.Stat(skillFile); os.IsNotExist(err) {
|
||||
alternatives := []string{
|
||||
filepath.Join(skillPath, "skill.md"),
|
||||
filepath.Join(skillPath, "README.md"),
|
||||
filepath.Join(skillPath, "readme.md"),
|
||||
}
|
||||
for _, alt := range alternatives {
|
||||
if _, err := os.Stat(alt); err == nil {
|
||||
skillFile = alt
|
||||
break
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
fileInfo, _ := os.Stat(skillFile)
|
||||
var fileSize int64
|
||||
var modTime string
|
||||
if fileInfo != nil {
|
||||
fileSize = fileInfo.Size()
|
||||
modTime = fileInfo.ModTime().Format("2006-01-02 15:04:05")
|
||||
}
|
||||
|
||||
c.JSON(http.StatusOK, gin.H{
|
||||
"skill": map[string]interface{}{
|
||||
"name": skill.Name,
|
||||
"description": skill.Description,
|
||||
"content": skill.Content,
|
||||
"path": skill.Path,
|
||||
"file_size": fileSize,
|
||||
"mod_time": modTime,
|
||||
},
|
||||
})
|
||||
}
|
||||
|
||||
// GetSkillBoundRoles 获取绑定指定skill的角色列表
|
||||
func (h *SkillsHandler) GetSkillBoundRoles(c *gin.Context) {
|
||||
skillName := c.Param("name")
|
||||
if skillName == "" {
|
||||
c.JSON(http.StatusBadRequest, gin.H{"error": "skill名称不能为空"})
|
||||
return
|
||||
}
|
||||
|
||||
boundRoles := h.getRolesBoundToSkill(skillName)
|
||||
c.JSON(http.StatusOK, gin.H{
|
||||
"skill": skillName,
|
||||
"bound_roles": boundRoles,
|
||||
"bound_count": len(boundRoles),
|
||||
})
|
||||
}
|
||||
|
||||
// getRolesBoundToSkill 获取绑定指定skill的角色列表(不修改配置)
|
||||
func (h *SkillsHandler) getRolesBoundToSkill(skillName string) []string {
|
||||
if h.config.Roles == nil {
|
||||
return []string{}
|
||||
}
|
||||
|
||||
boundRoles := make([]string, 0)
|
||||
for roleName, role := range h.config.Roles {
|
||||
// 确保角色名称正确设置
|
||||
if role.Name == "" {
|
||||
role.Name = roleName
|
||||
}
|
||||
|
||||
// 检查角色的Skills列表中是否包含该skill
|
||||
if len(role.Skills) > 0 {
|
||||
for _, skill := range role.Skills {
|
||||
if skill == skillName {
|
||||
boundRoles = append(boundRoles, roleName)
|
||||
break
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
return boundRoles
|
||||
}
|
||||
|
||||
// CreateSkill 创建新skill
|
||||
func (h *SkillsHandler) CreateSkill(c *gin.Context) {
|
||||
var req struct {
|
||||
Name string `json:"name" binding:"required"`
|
||||
Description string `json:"description"`
|
||||
Content string `json:"content" binding:"required"`
|
||||
}
|
||||
|
||||
if err := c.ShouldBindJSON(&req); err != nil {
|
||||
c.JSON(http.StatusBadRequest, gin.H{"error": "无效的请求参数: " + err.Error()})
|
||||
return
|
||||
}
|
||||
|
||||
// 验证skill名称(只允许字母、数字、连字符和下划线)
|
||||
if !isValidSkillName(req.Name) {
|
||||
c.JSON(http.StatusBadRequest, gin.H{"error": "skill名称只能包含字母、数字、连字符和下划线"})
|
||||
return
|
||||
}
|
||||
|
||||
// 获取skills目录
|
||||
skillsDir := h.config.SkillsDir
|
||||
if skillsDir == "" {
|
||||
skillsDir = "skills"
|
||||
}
|
||||
configDir := filepath.Dir(h.configPath)
|
||||
if !filepath.IsAbs(skillsDir) {
|
||||
skillsDir = filepath.Join(configDir, skillsDir)
|
||||
}
|
||||
|
||||
// 创建skill目录
|
||||
skillDir := filepath.Join(skillsDir, req.Name)
|
||||
if err := os.MkdirAll(skillDir, 0755); err != nil {
|
||||
h.logger.Error("创建skill目录失败", zap.String("skill", req.Name), zap.Error(err))
|
||||
c.JSON(http.StatusInternalServerError, gin.H{"error": "创建skill目录失败: " + err.Error()})
|
||||
return
|
||||
}
|
||||
|
||||
// 检查是否已存在
|
||||
skillFile := filepath.Join(skillDir, "SKILL.md")
|
||||
if _, err := os.Stat(skillFile); err == nil {
|
||||
c.JSON(http.StatusBadRequest, gin.H{"error": "skill已存在"})
|
||||
return
|
||||
}
|
||||
|
||||
// 构建SKILL.md内容
|
||||
var content strings.Builder
|
||||
content.WriteString("---\n")
|
||||
content.WriteString(fmt.Sprintf("name: %s\n", req.Name))
|
||||
if req.Description != "" {
|
||||
// 如果描述包含特殊字符,需要加引号
|
||||
desc := req.Description
|
||||
if strings.Contains(desc, ":") || strings.Contains(desc, "\n") {
|
||||
desc = fmt.Sprintf(`"%s"`, strings.ReplaceAll(desc, `"`, `\"`))
|
||||
}
|
||||
content.WriteString(fmt.Sprintf("description: %s\n", desc))
|
||||
}
|
||||
content.WriteString("version: 1.0.0\n")
|
||||
content.WriteString("---\n\n")
|
||||
content.WriteString(req.Content)
|
||||
|
||||
// 写入文件
|
||||
if err := os.WriteFile(skillFile, []byte(content.String()), 0644); err != nil {
|
||||
h.logger.Error("创建skill文件失败", zap.String("skill", req.Name), zap.Error(err))
|
||||
c.JSON(http.StatusInternalServerError, gin.H{"error": "创建skill文件失败: " + err.Error()})
|
||||
return
|
||||
}
|
||||
|
||||
h.logger.Info("创建skill成功", zap.String("skill", req.Name))
|
||||
c.JSON(http.StatusOK, gin.H{
|
||||
"message": "skill已创建",
|
||||
"skill": map[string]interface{}{
|
||||
"name": req.Name,
|
||||
"path": skillDir,
|
||||
},
|
||||
})
|
||||
}
|
||||
|
||||
// UpdateSkill 更新skill
|
||||
func (h *SkillsHandler) UpdateSkill(c *gin.Context) {
|
||||
skillName := c.Param("name")
|
||||
if skillName == "" {
|
||||
c.JSON(http.StatusBadRequest, gin.H{"error": "skill名称不能为空"})
|
||||
return
|
||||
}
|
||||
|
||||
var req struct {
|
||||
Description string `json:"description"`
|
||||
Content string `json:"content" binding:"required"`
|
||||
}
|
||||
|
||||
if err := c.ShouldBindJSON(&req); err != nil {
|
||||
c.JSON(http.StatusBadRequest, gin.H{"error": "无效的请求参数: " + err.Error()})
|
||||
return
|
||||
}
|
||||
|
||||
// 获取skills目录
|
||||
skillsDir := h.config.SkillsDir
|
||||
if skillsDir == "" {
|
||||
skillsDir = "skills"
|
||||
}
|
||||
configDir := filepath.Dir(h.configPath)
|
||||
if !filepath.IsAbs(skillsDir) {
|
||||
skillsDir = filepath.Join(configDir, skillsDir)
|
||||
}
|
||||
|
||||
// 查找skill文件
|
||||
skillDir := filepath.Join(skillsDir, skillName)
|
||||
skillFile := filepath.Join(skillDir, "SKILL.md")
|
||||
if _, err := os.Stat(skillFile); os.IsNotExist(err) {
|
||||
alternatives := []string{
|
||||
filepath.Join(skillDir, "skill.md"),
|
||||
filepath.Join(skillDir, "README.md"),
|
||||
filepath.Join(skillDir, "readme.md"),
|
||||
}
|
||||
found := false
|
||||
for _, alt := range alternatives {
|
||||
if _, err := os.Stat(alt); err == nil {
|
||||
skillFile = alt
|
||||
found = true
|
||||
break
|
||||
}
|
||||
}
|
||||
if !found {
|
||||
c.JSON(http.StatusNotFound, gin.H{"error": "skill不存在"})
|
||||
return
|
||||
}
|
||||
}
|
||||
|
||||
// 读取现有文件以保留front matter中的name
|
||||
existingContent, err := os.ReadFile(skillFile)
|
||||
if err != nil {
|
||||
h.logger.Error("读取skill文件失败", zap.String("skill", skillName), zap.Error(err))
|
||||
c.JSON(http.StatusInternalServerError, gin.H{"error": "读取skill文件失败: " + err.Error()})
|
||||
return
|
||||
}
|
||||
|
||||
// 解析现有内容,提取name
|
||||
existingName := skillName
|
||||
contentStr := string(existingContent)
|
||||
if strings.HasPrefix(contentStr, "---") {
|
||||
parts := strings.SplitN(contentStr, "---", 3)
|
||||
if len(parts) >= 2 {
|
||||
frontMatter := parts[1]
|
||||
lines := strings.Split(frontMatter, "\n")
|
||||
for _, line := range lines {
|
||||
line = strings.TrimSpace(line)
|
||||
if strings.HasPrefix(line, "name:") {
|
||||
name := strings.TrimSpace(strings.TrimPrefix(line, "name:"))
|
||||
name = strings.Trim(name, `"'`)
|
||||
if name != "" {
|
||||
existingName = name
|
||||
}
|
||||
break
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// 构建新的SKILL.md内容
|
||||
var newContent strings.Builder
|
||||
newContent.WriteString("---\n")
|
||||
newContent.WriteString(fmt.Sprintf("name: %s\n", existingName))
|
||||
if req.Description != "" {
|
||||
// 如果描述包含特殊字符,需要加引号
|
||||
desc := req.Description
|
||||
if strings.Contains(desc, ":") || strings.Contains(desc, "\n") {
|
||||
desc = fmt.Sprintf(`"%s"`, strings.ReplaceAll(desc, `"`, `\"`))
|
||||
}
|
||||
newContent.WriteString(fmt.Sprintf("description: %s\n", desc))
|
||||
}
|
||||
newContent.WriteString("version: 1.0.0\n")
|
||||
newContent.WriteString("---\n\n")
|
||||
newContent.WriteString(req.Content)
|
||||
|
||||
// 写入文件(统一使用SKILL.md)
|
||||
targetFile := filepath.Join(skillDir, "SKILL.md")
|
||||
if err := os.WriteFile(targetFile, []byte(newContent.String()), 0644); err != nil {
|
||||
h.logger.Error("更新skill文件失败", zap.String("skill", skillName), zap.Error(err))
|
||||
c.JSON(http.StatusInternalServerError, gin.H{"error": "更新skill文件失败: " + err.Error()})
|
||||
return
|
||||
}
|
||||
|
||||
// 如果原文件不是SKILL.md,删除旧文件
|
||||
if skillFile != targetFile {
|
||||
os.Remove(skillFile)
|
||||
}
|
||||
|
||||
h.logger.Info("更新skill成功", zap.String("skill", skillName))
|
||||
c.JSON(http.StatusOK, gin.H{
|
||||
"message": "skill已更新",
|
||||
})
|
||||
}
|
||||
|
||||
// DeleteSkill 删除skill
|
||||
func (h *SkillsHandler) DeleteSkill(c *gin.Context) {
|
||||
skillName := c.Param("name")
|
||||
if skillName == "" {
|
||||
c.JSON(http.StatusBadRequest, gin.H{"error": "skill名称不能为空"})
|
||||
return
|
||||
}
|
||||
|
||||
// 检查是否有角色绑定了该skill,如果有则自动移除绑定
|
||||
affectedRoles := h.removeSkillFromRoles(skillName)
|
||||
if len(affectedRoles) > 0 {
|
||||
h.logger.Info("从角色中移除skill绑定",
|
||||
zap.String("skill", skillName),
|
||||
zap.Strings("roles", affectedRoles))
|
||||
}
|
||||
|
||||
// 获取skills目录
|
||||
skillsDir := h.config.SkillsDir
|
||||
if skillsDir == "" {
|
||||
skillsDir = "skills"
|
||||
}
|
||||
configDir := filepath.Dir(h.configPath)
|
||||
if !filepath.IsAbs(skillsDir) {
|
||||
skillsDir = filepath.Join(configDir, skillsDir)
|
||||
}
|
||||
|
||||
// 删除skill目录
|
||||
skillDir := filepath.Join(skillsDir, skillName)
|
||||
if err := os.RemoveAll(skillDir); err != nil {
|
||||
h.logger.Error("删除skill失败", zap.String("skill", skillName), zap.Error(err))
|
||||
c.JSON(http.StatusInternalServerError, gin.H{"error": "删除skill失败: " + err.Error()})
|
||||
return
|
||||
}
|
||||
|
||||
responseMsg := "skill已删除"
|
||||
if len(affectedRoles) > 0 {
|
||||
responseMsg = fmt.Sprintf("skill已删除,已自动从 %d 个角色中移除绑定: %s",
|
||||
len(affectedRoles), strings.Join(affectedRoles, ", "))
|
||||
}
|
||||
|
||||
h.logger.Info("删除skill成功", zap.String("skill", skillName))
|
||||
c.JSON(http.StatusOK, gin.H{
|
||||
"message": responseMsg,
|
||||
"affected_roles": affectedRoles,
|
||||
})
|
||||
}
|
||||
|
||||
// GetSkillStats 获取skills调用统计信息
|
||||
func (h *SkillsHandler) GetSkillStats(c *gin.Context) {
|
||||
skillList, err := h.manager.ListSkills()
|
||||
if err != nil {
|
||||
h.logger.Error("获取skills列表失败", zap.Error(err))
|
||||
c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
|
||||
return
|
||||
}
|
||||
|
||||
// 获取skills目录
|
||||
skillsDir := h.config.SkillsDir
|
||||
if skillsDir == "" {
|
||||
skillsDir = "skills"
|
||||
}
|
||||
configDir := filepath.Dir(h.configPath)
|
||||
if !filepath.IsAbs(skillsDir) {
|
||||
skillsDir = filepath.Join(configDir, skillsDir)
|
||||
}
|
||||
|
||||
// 从数据库加载调用统计
|
||||
var skillStatsMap map[string]*database.SkillStats
|
||||
if h.db != nil {
|
||||
dbStats, err := h.db.LoadSkillStats()
|
||||
if err != nil {
|
||||
h.logger.Warn("从数据库加载Skills统计信息失败", zap.Error(err))
|
||||
skillStatsMap = make(map[string]*database.SkillStats)
|
||||
} else {
|
||||
skillStatsMap = dbStats
|
||||
}
|
||||
} else {
|
||||
skillStatsMap = make(map[string]*database.SkillStats)
|
||||
}
|
||||
|
||||
// 构建统计信息(包含所有skills,即使没有调用记录)
|
||||
statsList := make([]map[string]interface{}, 0, len(skillList))
|
||||
totalCalls := 0
|
||||
totalSuccess := 0
|
||||
totalFailed := 0
|
||||
|
||||
for _, skillName := range skillList {
|
||||
stat, exists := skillStatsMap[skillName]
|
||||
if !exists {
|
||||
stat = &database.SkillStats{
|
||||
SkillName: skillName,
|
||||
TotalCalls: 0,
|
||||
SuccessCalls: 0,
|
||||
FailedCalls: 0,
|
||||
}
|
||||
}
|
||||
|
||||
totalCalls += stat.TotalCalls
|
||||
totalSuccess += stat.SuccessCalls
|
||||
totalFailed += stat.FailedCalls
|
||||
|
||||
lastCallTimeStr := ""
|
||||
if stat.LastCallTime != nil {
|
||||
lastCallTimeStr = stat.LastCallTime.Format("2006-01-02 15:04:05")
|
||||
}
|
||||
|
||||
statsList = append(statsList, map[string]interface{}{
|
||||
"skill_name": stat.SkillName,
|
||||
"total_calls": stat.TotalCalls,
|
||||
"success_calls": stat.SuccessCalls,
|
||||
"failed_calls": stat.FailedCalls,
|
||||
"last_call_time": lastCallTimeStr,
|
||||
})
|
||||
}
|
||||
|
||||
c.JSON(http.StatusOK, gin.H{
|
||||
"total_skills": len(skillList),
|
||||
"total_calls": totalCalls,
|
||||
"total_success": totalSuccess,
|
||||
"total_failed": totalFailed,
|
||||
"skills_dir": skillsDir,
|
||||
"stats": statsList,
|
||||
})
|
||||
}
|
||||
|
||||
// ClearSkillStats 清空所有Skills统计信息
|
||||
func (h *SkillsHandler) ClearSkillStats(c *gin.Context) {
|
||||
if h.db == nil {
|
||||
c.JSON(http.StatusInternalServerError, gin.H{"error": "数据库连接未配置"})
|
||||
return
|
||||
}
|
||||
|
||||
if err := h.db.ClearSkillStats(); err != nil {
|
||||
h.logger.Error("清空Skills统计信息失败", zap.Error(err))
|
||||
c.JSON(http.StatusInternalServerError, gin.H{"error": "清空统计信息失败: " + err.Error()})
|
||||
return
|
||||
}
|
||||
|
||||
h.logger.Info("已清空所有Skills统计信息")
|
||||
c.JSON(http.StatusOK, gin.H{
|
||||
"message": "已清空所有Skills统计信息",
|
||||
})
|
||||
}
|
||||
|
||||
// ClearSkillStatsByName 清空指定skill的统计信息
|
||||
func (h *SkillsHandler) ClearSkillStatsByName(c *gin.Context) {
|
||||
skillName := c.Param("name")
|
||||
if skillName == "" {
|
||||
c.JSON(http.StatusBadRequest, gin.H{"error": "skill名称不能为空"})
|
||||
return
|
||||
}
|
||||
|
||||
if h.db == nil {
|
||||
c.JSON(http.StatusInternalServerError, gin.H{"error": "数据库连接未配置"})
|
||||
return
|
||||
}
|
||||
|
||||
if err := h.db.ClearSkillStatsByName(skillName); err != nil {
|
||||
h.logger.Error("清空指定skill统计信息失败", zap.String("skill", skillName), zap.Error(err))
|
||||
c.JSON(http.StatusInternalServerError, gin.H{"error": "清空统计信息失败: " + err.Error()})
|
||||
return
|
||||
}
|
||||
|
||||
h.logger.Info("已清空指定skill统计信息", zap.String("skill", skillName))
|
||||
c.JSON(http.StatusOK, gin.H{
|
||||
"message": fmt.Sprintf("已清空skill '%s' 的统计信息", skillName),
|
||||
})
|
||||
}
|
||||
|
||||
// removeSkillFromRoles 从所有角色中移除指定的skill绑定
|
||||
// 返回受影响角色名称列表
|
||||
func (h *SkillsHandler) removeSkillFromRoles(skillName string) []string {
|
||||
if h.config.Roles == nil {
|
||||
return []string{}
|
||||
}
|
||||
|
||||
affectedRoles := make([]string, 0)
|
||||
rolesToUpdate := make(map[string]config.RoleConfig)
|
||||
|
||||
// 遍历所有角色,查找并移除skill绑定
|
||||
for roleName, role := range h.config.Roles {
|
||||
// 确保角色名称正确设置
|
||||
if role.Name == "" {
|
||||
role.Name = roleName
|
||||
}
|
||||
|
||||
// 检查角色的Skills列表中是否包含要删除的skill
|
||||
if len(role.Skills) > 0 {
|
||||
updated := false
|
||||
newSkills := make([]string, 0, len(role.Skills))
|
||||
for _, skill := range role.Skills {
|
||||
if skill != skillName {
|
||||
newSkills = append(newSkills, skill)
|
||||
} else {
|
||||
updated = true
|
||||
}
|
||||
}
|
||||
if updated {
|
||||
role.Skills = newSkills
|
||||
rolesToUpdate[roleName] = role
|
||||
affectedRoles = append(affectedRoles, roleName)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// 如果有角色需要更新,保存到文件
|
||||
if len(rolesToUpdate) > 0 {
|
||||
// 更新内存中的配置
|
||||
for roleName, role := range rolesToUpdate {
|
||||
h.config.Roles[roleName] = role
|
||||
}
|
||||
// 保存更新后的角色配置到文件
|
||||
if err := h.saveRolesConfig(); err != nil {
|
||||
h.logger.Error("保存角色配置失败", zap.Error(err))
|
||||
}
|
||||
}
|
||||
|
||||
return affectedRoles
|
||||
}
|
||||
|
||||
// saveRolesConfig 保存角色配置到文件(从SkillsHandler调用)
|
||||
func (h *SkillsHandler) saveRolesConfig() error {
|
||||
configDir := filepath.Dir(h.configPath)
|
||||
rolesDir := h.config.RolesDir
|
||||
if rolesDir == "" {
|
||||
rolesDir = "roles" // 默认目录
|
||||
}
|
||||
|
||||
// 如果是相对路径,相对于配置文件所在目录
|
||||
if !filepath.IsAbs(rolesDir) {
|
||||
rolesDir = filepath.Join(configDir, rolesDir)
|
||||
}
|
||||
|
||||
// 确保目录存在
|
||||
if err := os.MkdirAll(rolesDir, 0755); err != nil {
|
||||
return fmt.Errorf("创建角色目录失败: %w", err)
|
||||
}
|
||||
|
||||
// 保存每个角色到独立的文件
|
||||
if h.config.Roles != nil {
|
||||
for roleName, role := range h.config.Roles {
|
||||
// 确保角色名称正确设置
|
||||
if role.Name == "" {
|
||||
role.Name = roleName
|
||||
}
|
||||
|
||||
// 使用角色名称作为文件名(安全化文件名,避免特殊字符)
|
||||
safeFileName := sanitizeRoleFileName(role.Name)
|
||||
roleFile := filepath.Join(rolesDir, safeFileName+".yaml")
|
||||
|
||||
// 将角色配置序列化为YAML
|
||||
roleData, err := yaml.Marshal(&role)
|
||||
if err != nil {
|
||||
h.logger.Error("序列化角色配置失败", zap.String("role", roleName), zap.Error(err))
|
||||
continue
|
||||
}
|
||||
|
||||
// 处理icon字段:确保包含\U的icon值被引号包围(YAML需要引号才能正确解析Unicode转义)
|
||||
roleDataStr := string(roleData)
|
||||
if role.Icon != "" && strings.HasPrefix(role.Icon, "\\U") {
|
||||
// 匹配 icon: \UXXXXXXXX 格式(没有引号),排除已经有引号的情况
|
||||
re := regexp.MustCompile(`(?m)^(icon:\s+)(\\U[0-9A-F]{8})(\s*)$`)
|
||||
roleDataStr = re.ReplaceAllString(roleDataStr, `${1}"${2}"${3}`)
|
||||
roleData = []byte(roleDataStr)
|
||||
}
|
||||
|
||||
// 写入文件
|
||||
if err := os.WriteFile(roleFile, roleData, 0644); err != nil {
|
||||
h.logger.Error("保存角色配置文件失败", zap.String("role", roleName), zap.String("file", roleFile), zap.Error(err))
|
||||
continue
|
||||
}
|
||||
|
||||
h.logger.Info("角色配置已保存到文件", zap.String("role", roleName), zap.String("file", roleFile))
|
||||
}
|
||||
}
|
||||
|
||||
return nil
|
||||
}
|
||||
|
||||
// sanitizeRoleFileName 将角色名称转换为安全的文件名
|
||||
func sanitizeRoleFileName(name string) string {
|
||||
// 替换可能不安全的字符
|
||||
replacer := map[rune]string{
|
||||
'/': "_",
|
||||
'\\': "_",
|
||||
':': "_",
|
||||
'*': "_",
|
||||
'?': "_",
|
||||
'"': "_",
|
||||
'<': "_",
|
||||
'>': "_",
|
||||
'|': "_",
|
||||
' ': "_",
|
||||
}
|
||||
|
||||
var result []rune
|
||||
for _, r := range name {
|
||||
if replacement, ok := replacer[r]; ok {
|
||||
result = append(result, []rune(replacement)...)
|
||||
} else {
|
||||
result = append(result, r)
|
||||
}
|
||||
}
|
||||
|
||||
fileName := string(result)
|
||||
// 如果文件名为空,使用默认名称
|
||||
if fileName == "" {
|
||||
fileName = "role"
|
||||
}
|
||||
|
||||
return fileName
|
||||
}
|
||||
|
||||
// isValidSkillName 验证skill名称是否有效
|
||||
func isValidSkillName(name string) bool {
|
||||
if name == "" || len(name) > 100 {
|
||||
return false
|
||||
}
|
||||
// 只允许字母、数字、连字符和下划线
|
||||
for _, r := range name {
|
||||
if !((r >= 'a' && r <= 'z') || (r >= 'A' && r <= 'Z') || (r >= '0' && r <= '9') || r == '-' || r == '_') {
|
||||
return false
|
||||
}
|
||||
}
|
||||
return true
|
||||
}
|
||||
@@ -9,6 +9,10 @@ const (
|
||||
// 知识库工具
|
||||
ToolListKnowledgeRiskTypes = "list_knowledge_risk_types"
|
||||
ToolSearchKnowledgeBase = "search_knowledge_base"
|
||||
|
||||
// Skills工具
|
||||
ToolListSkills = "list_skills"
|
||||
ToolReadSkill = "read_skill"
|
||||
)
|
||||
|
||||
// IsBuiltinTool 检查工具名称是否是内置工具
|
||||
@@ -16,7 +20,9 @@ func IsBuiltinTool(toolName string) bool {
|
||||
switch toolName {
|
||||
case ToolRecordVulnerability,
|
||||
ToolListKnowledgeRiskTypes,
|
||||
ToolSearchKnowledgeBase:
|
||||
ToolSearchKnowledgeBase,
|
||||
ToolListSkills,
|
||||
ToolReadSkill:
|
||||
return true
|
||||
default:
|
||||
return false
|
||||
@@ -29,5 +35,7 @@ func GetAllBuiltinTools() []string {
|
||||
ToolRecordVulnerability,
|
||||
ToolListKnowledgeRiskTypes,
|
||||
ToolSearchKnowledgeBase,
|
||||
ToolListSkills,
|
||||
ToolReadSkill,
|
||||
}
|
||||
}
|
||||
|
||||
@@ -0,0 +1,227 @@
|
||||
package skills
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
"os"
|
||||
"path/filepath"
|
||||
"strings"
|
||||
|
||||
"go.uber.org/zap"
|
||||
)
|
||||
|
||||
// Manager Skills管理器
|
||||
type Manager struct {
|
||||
skillsDir string
|
||||
logger *zap.Logger
|
||||
skills map[string]*Skill // 缓存已加载的skills
|
||||
}
|
||||
|
||||
// Skill Skill定义
|
||||
type Skill struct {
|
||||
Name string // Skill名称
|
||||
Description string // Skill描述
|
||||
Content string // Skill内容(从SKILL.md中提取)
|
||||
Path string // Skill路径
|
||||
}
|
||||
|
||||
// NewManager 创建新的Skills管理器
|
||||
func NewManager(skillsDir string, logger *zap.Logger) *Manager {
|
||||
return &Manager{
|
||||
skillsDir: skillsDir,
|
||||
logger: logger,
|
||||
skills: make(map[string]*Skill),
|
||||
}
|
||||
}
|
||||
|
||||
// LoadSkill 加载单个skill
|
||||
func (m *Manager) LoadSkill(skillName string) (*Skill, error) {
|
||||
// 检查缓存
|
||||
if skill, exists := m.skills[skillName]; exists {
|
||||
return skill, nil
|
||||
}
|
||||
|
||||
// 构建skill路径
|
||||
skillPath := filepath.Join(m.skillsDir, skillName)
|
||||
|
||||
// 检查目录是否存在
|
||||
if _, err := os.Stat(skillPath); os.IsNotExist(err) {
|
||||
return nil, fmt.Errorf("skill %s not found", skillName)
|
||||
}
|
||||
|
||||
// 查找SKILL.md文件
|
||||
skillFile := filepath.Join(skillPath, "SKILL.md")
|
||||
if _, err := os.Stat(skillFile); os.IsNotExist(err) {
|
||||
// 尝试其他可能的文件名
|
||||
alternatives := []string{
|
||||
filepath.Join(skillPath, "skill.md"),
|
||||
filepath.Join(skillPath, "README.md"),
|
||||
filepath.Join(skillPath, "readme.md"),
|
||||
}
|
||||
found := false
|
||||
for _, alt := range alternatives {
|
||||
if _, err := os.Stat(alt); err == nil {
|
||||
skillFile = alt
|
||||
found = true
|
||||
break
|
||||
}
|
||||
}
|
||||
if !found {
|
||||
return nil, fmt.Errorf("skill file not found for %s", skillName)
|
||||
}
|
||||
}
|
||||
|
||||
// 读取skill文件
|
||||
content, err := os.ReadFile(skillFile)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("failed to read skill file: %w", err)
|
||||
}
|
||||
|
||||
// 解析skill内容
|
||||
skill := m.parseSkillContent(string(content), skillName, skillPath)
|
||||
|
||||
// 缓存skill
|
||||
m.skills[skillName] = skill
|
||||
|
||||
return skill, nil
|
||||
}
|
||||
|
||||
// LoadSkills 批量加载skills
|
||||
func (m *Manager) LoadSkills(skillNames []string) ([]*Skill, error) {
|
||||
var skills []*Skill
|
||||
var errors []string
|
||||
|
||||
for _, name := range skillNames {
|
||||
skill, err := m.LoadSkill(name)
|
||||
if err != nil {
|
||||
errors = append(errors, fmt.Sprintf("failed to load skill %s: %v", name, err))
|
||||
m.logger.Warn("加载skill失败", zap.String("skill", name), zap.Error(err))
|
||||
continue
|
||||
}
|
||||
skills = append(skills, skill)
|
||||
}
|
||||
|
||||
if len(errors) > 0 && len(skills) == 0 {
|
||||
return nil, fmt.Errorf("failed to load any skills: %s", strings.Join(errors, "; "))
|
||||
}
|
||||
|
||||
return skills, nil
|
||||
}
|
||||
|
||||
// ListSkills 列出所有可用的skills
|
||||
func (m *Manager) ListSkills() ([]string, error) {
|
||||
if _, err := os.Stat(m.skillsDir); os.IsNotExist(err) {
|
||||
return []string{}, nil
|
||||
}
|
||||
|
||||
entries, err := os.ReadDir(m.skillsDir)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("failed to read skills directory: %w", err)
|
||||
}
|
||||
|
||||
var skills []string
|
||||
for _, entry := range entries {
|
||||
if !entry.IsDir() {
|
||||
continue
|
||||
}
|
||||
|
||||
skillName := entry.Name()
|
||||
// 检查是否有SKILL.md文件
|
||||
skillFile := filepath.Join(m.skillsDir, skillName, "SKILL.md")
|
||||
if _, err := os.Stat(skillFile); err == nil {
|
||||
skills = append(skills, skillName)
|
||||
continue
|
||||
}
|
||||
|
||||
// 尝试其他可能的文件名
|
||||
alternatives := []string{
|
||||
filepath.Join(m.skillsDir, skillName, "skill.md"),
|
||||
filepath.Join(m.skillsDir, skillName, "README.md"),
|
||||
filepath.Join(m.skillsDir, skillName, "readme.md"),
|
||||
}
|
||||
for _, alt := range alternatives {
|
||||
if _, err := os.Stat(alt); err == nil {
|
||||
skills = append(skills, skillName)
|
||||
break
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
return skills, nil
|
||||
}
|
||||
|
||||
// parseSkillContent 解析skill内容
|
||||
// 支持YAML front matter格式,类似goskills
|
||||
func (m *Manager) parseSkillContent(content, skillName, skillPath string) *Skill {
|
||||
skill := &Skill{
|
||||
Name: skillName,
|
||||
Path: skillPath,
|
||||
}
|
||||
|
||||
// 检查是否有YAML front matter
|
||||
if strings.HasPrefix(content, "---") {
|
||||
parts := strings.SplitN(content, "---", 3)
|
||||
if len(parts) >= 3 {
|
||||
// 解析front matter(简单实现,只提取name和description)
|
||||
frontMatter := parts[1]
|
||||
lines := strings.Split(frontMatter, "\n")
|
||||
for _, line := range lines {
|
||||
line = strings.TrimSpace(line)
|
||||
if strings.HasPrefix(line, "name:") {
|
||||
name := strings.TrimSpace(strings.TrimPrefix(line, "name:"))
|
||||
name = strings.Trim(name, `"'"`)
|
||||
if name != "" {
|
||||
skill.Name = name
|
||||
}
|
||||
} else if strings.HasPrefix(line, "description:") {
|
||||
desc := strings.TrimSpace(strings.TrimPrefix(line, "description:"))
|
||||
desc = strings.Trim(desc, `"'"`)
|
||||
skill.Description = desc
|
||||
}
|
||||
}
|
||||
// 剩余部分是内容
|
||||
if len(parts) == 3 {
|
||||
skill.Content = strings.TrimSpace(parts[2])
|
||||
}
|
||||
} else {
|
||||
// 没有front matter,整个内容就是skill内容
|
||||
skill.Content = content
|
||||
}
|
||||
} else {
|
||||
// 没有front matter,整个内容就是skill内容
|
||||
skill.Content = content
|
||||
}
|
||||
|
||||
// 如果内容为空,使用描述作为内容
|
||||
if skill.Content == "" {
|
||||
skill.Content = skill.Description
|
||||
}
|
||||
|
||||
return skill
|
||||
}
|
||||
|
||||
// GetSkillContent 获取skill的完整内容(用于注入到系统提示词)
|
||||
func (m *Manager) GetSkillContent(skillNames []string) (string, error) {
|
||||
skills, err := m.LoadSkills(skillNames)
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
|
||||
if len(skills) == 0 {
|
||||
return "", nil
|
||||
}
|
||||
|
||||
var builder strings.Builder
|
||||
builder.WriteString("## 可用Skills\n\n")
|
||||
builder.WriteString("在执行任务前,请仔细阅读以下skills内容,这些内容包含了相关的专业知识和方法:\n\n")
|
||||
|
||||
for _, skill := range skills {
|
||||
builder.WriteString(fmt.Sprintf("### Skill: %s\n", skill.Name))
|
||||
if skill.Description != "" {
|
||||
builder.WriteString(fmt.Sprintf("**描述**: %s\n\n", skill.Description))
|
||||
}
|
||||
builder.WriteString(skill.Content)
|
||||
builder.WriteString("\n\n---\n\n")
|
||||
}
|
||||
|
||||
return builder.String(), nil
|
||||
}
|
||||
@@ -0,0 +1,201 @@
|
||||
package skills
|
||||
|
||||
import (
|
||||
"context"
|
||||
"fmt"
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
"cyberstrike-ai/internal/mcp"
|
||||
"cyberstrike-ai/internal/mcp/builtin"
|
||||
|
||||
"go.uber.org/zap"
|
||||
)
|
||||
|
||||
// RegisterSkillsTool 注册Skills工具到MCP服务器
|
||||
func RegisterSkillsTool(
|
||||
mcpServer *mcp.Server,
|
||||
manager *Manager,
|
||||
logger *zap.Logger,
|
||||
) {
|
||||
RegisterSkillsToolWithStorage(mcpServer, manager, nil, logger)
|
||||
}
|
||||
|
||||
// RegisterSkillsToolWithStorage 注册Skills工具到MCP服务器(带存储支持)
|
||||
func RegisterSkillsToolWithStorage(
|
||||
mcpServer *mcp.Server,
|
||||
manager *Manager,
|
||||
storage SkillStatsStorage,
|
||||
logger *zap.Logger,
|
||||
) {
|
||||
// 注册第一个工具:获取所有可用的skills列表
|
||||
listSkillsTool := mcp.Tool{
|
||||
Name: builtin.ToolListSkills,
|
||||
Description: "获取所有可用的skills列表。Skills是专业知识文档,可以在执行任务前阅读以获取相关专业知识。使用此工具可以查看系统中所有可用的skills,然后使用read_skill工具读取特定skill的内容。",
|
||||
ShortDescription: "获取所有可用的skills列表",
|
||||
InputSchema: map[string]interface{}{
|
||||
"type": "object",
|
||||
"properties": map[string]interface{}{},
|
||||
"required": []string{},
|
||||
},
|
||||
}
|
||||
|
||||
listSkillsHandler := func(ctx context.Context, args map[string]interface{}) (*mcp.ToolResult, error) {
|
||||
skills, err := manager.ListSkills()
|
||||
if err != nil {
|
||||
logger.Error("获取skills列表失败", zap.Error(err))
|
||||
return &mcp.ToolResult{
|
||||
Content: []mcp.Content{
|
||||
{
|
||||
Type: "text",
|
||||
Text: fmt.Sprintf("获取skills列表失败: %v", err),
|
||||
},
|
||||
},
|
||||
IsError: true,
|
||||
}, nil
|
||||
}
|
||||
|
||||
if len(skills) == 0 {
|
||||
return &mcp.ToolResult{
|
||||
Content: []mcp.Content{
|
||||
{
|
||||
Type: "text",
|
||||
Text: "当前没有可用的skills。\n\nSkills是专业知识文档,可以在执行任务前阅读以获取相关专业知识。你可以在skills目录下创建新的skill。",
|
||||
},
|
||||
},
|
||||
IsError: false,
|
||||
}, nil
|
||||
}
|
||||
|
||||
var result strings.Builder
|
||||
result.WriteString(fmt.Sprintf("共有 %d 个可用的skills:\n\n", len(skills)))
|
||||
for i, skill := range skills {
|
||||
result.WriteString(fmt.Sprintf("%d. %s\n", i+1, skill))
|
||||
}
|
||||
result.WriteString("\n使用 read_skill 工具可以读取特定skill的详细内容。\n")
|
||||
result.WriteString("例如:read_skill(skill_name=\"sql-injection-testing\")")
|
||||
|
||||
return &mcp.ToolResult{
|
||||
Content: []mcp.Content{
|
||||
{
|
||||
Type: "text",
|
||||
Text: result.String(),
|
||||
},
|
||||
},
|
||||
IsError: false,
|
||||
}, nil
|
||||
}
|
||||
|
||||
mcpServer.RegisterTool(listSkillsTool, listSkillsHandler)
|
||||
logger.Info("注册skills列表工具成功")
|
||||
|
||||
// 注册第二个工具:读取特定skill的内容
|
||||
readSkillTool := mcp.Tool{
|
||||
Name: builtin.ToolReadSkill,
|
||||
Description: "读取指定skill的详细内容。Skills是专业知识文档,包含测试方法、工具使用、最佳实践等。在执行相关任务前,可以调用此工具读取相关skill的内容,以获取专业知识和指导。",
|
||||
ShortDescription: "读取指定skill的详细内容",
|
||||
InputSchema: map[string]interface{}{
|
||||
"type": "object",
|
||||
"properties": map[string]interface{}{
|
||||
"skill_name": map[string]interface{}{
|
||||
"type": "string",
|
||||
"description": "要读取的skill名称(必需)。可以使用list_skills工具获取所有可用的skill名称。",
|
||||
},
|
||||
},
|
||||
"required": []string{"skill_name"},
|
||||
},
|
||||
}
|
||||
|
||||
readSkillHandler := func(ctx context.Context, args map[string]interface{}) (*mcp.ToolResult, error) {
|
||||
skillName, ok := args["skill_name"].(string)
|
||||
if !ok || skillName == "" {
|
||||
return &mcp.ToolResult{
|
||||
Content: []mcp.Content{
|
||||
{
|
||||
Type: "text",
|
||||
Text: "错误: skill_name 参数必需且不能为空。请使用list_skills工具获取所有可用的skill名称。",
|
||||
},
|
||||
},
|
||||
IsError: true,
|
||||
}, nil
|
||||
}
|
||||
|
||||
skill, err := manager.LoadSkill(skillName)
|
||||
failed := err != nil
|
||||
now := time.Now()
|
||||
|
||||
// 记录调用统计
|
||||
if storage != nil {
|
||||
totalCalls := 1
|
||||
successCalls := 0
|
||||
failedCalls := 0
|
||||
if failed {
|
||||
failedCalls = 1
|
||||
} else {
|
||||
successCalls = 1
|
||||
}
|
||||
if err := storage.UpdateSkillStats(skillName, totalCalls, successCalls, failedCalls, &now); err != nil {
|
||||
logger.Warn("保存Skills统计信息失败", zap.String("skill", skillName), zap.Error(err))
|
||||
} else {
|
||||
logger.Info("Skills统计信息已更新",
|
||||
zap.String("skill", skillName),
|
||||
zap.Int("totalCalls", totalCalls),
|
||||
zap.Int("successCalls", successCalls),
|
||||
zap.Int("failedCalls", failedCalls))
|
||||
}
|
||||
} else {
|
||||
logger.Warn("Skills统计存储未配置,无法记录调用统计", zap.String("skill", skillName))
|
||||
}
|
||||
|
||||
if err != nil {
|
||||
logger.Warn("读取skill失败", zap.String("skill", skillName), zap.Error(err))
|
||||
return &mcp.ToolResult{
|
||||
Content: []mcp.Content{
|
||||
{
|
||||
Type: "text",
|
||||
Text: fmt.Sprintf("读取skill失败: %v\n\n请使用list_skills工具确认skill名称是否正确。", err),
|
||||
},
|
||||
},
|
||||
IsError: true,
|
||||
}, nil
|
||||
}
|
||||
|
||||
var result strings.Builder
|
||||
result.WriteString(fmt.Sprintf("## Skill: %s\n\n", skill.Name))
|
||||
if skill.Description != "" {
|
||||
result.WriteString(fmt.Sprintf("**描述**: %s\n\n", skill.Description))
|
||||
}
|
||||
result.WriteString("---\n\n")
|
||||
result.WriteString(skill.Content)
|
||||
result.WriteString("\n\n---\n\n")
|
||||
result.WriteString(fmt.Sprintf("*Skill路径: %s*", skill.Path))
|
||||
|
||||
return &mcp.ToolResult{
|
||||
Content: []mcp.Content{
|
||||
{
|
||||
Type: "text",
|
||||
Text: result.String(),
|
||||
},
|
||||
},
|
||||
IsError: false,
|
||||
}, nil
|
||||
}
|
||||
|
||||
mcpServer.RegisterTool(readSkillTool, readSkillHandler)
|
||||
logger.Info("注册skill读取工具成功")
|
||||
}
|
||||
|
||||
// SkillStatsStorage Skills统计存储接口
|
||||
type SkillStatsStorage interface {
|
||||
UpdateSkillStats(skillName string, totalCalls, successCalls, failedCalls int, lastCallTime *time.Time) error
|
||||
LoadSkillStats() (map[string]*SkillStats, error)
|
||||
}
|
||||
|
||||
// SkillStats Skills统计信息
|
||||
type SkillStats struct {
|
||||
SkillName string
|
||||
TotalCalls int
|
||||
SuccessCalls int
|
||||
FailedCalls int
|
||||
LastCallTime *time.Time
|
||||
}
|
||||
Reference in New Issue
Block a user