diff --git a/internal/agents/markdown.go b/internal/agents/markdown.go new file mode 100644 index 00000000..b3aa8a0f --- /dev/null +++ b/internal/agents/markdown.go @@ -0,0 +1,526 @@ +// Package agents 从 agents/ 目录加载 Markdown 代理定义（子代理 + 可选主代理 orchestrator.md / kind: orchestrator）。 +package agents + +import ( + "fmt" + "os" + "path/filepath" + "sort" + "strings" + "unicode" + + "cyberstrike-ai/internal/config" + + "gopkg.in/yaml.v3" +) + +// OrchestratorMarkdownFilename 固定文件名：存在则视为 Deep 主代理定义，且不参与子代理列表。 +const OrchestratorMarkdownFilename = "orchestrator.md" + +// OrchestratorPlanExecuteMarkdownFilename plan_execute 模式主代理（规划侧）专用 Markdown 文件名。 +const OrchestratorPlanExecuteMarkdownFilename = "orchestrator-plan-execute.md" + +// OrchestratorSupervisorMarkdownFilename supervisor 模式主代理专用 Markdown 文件名。 +const OrchestratorSupervisorMarkdownFilename = "orchestrator-supervisor.md" + +// FrontMatter 对应 Markdown 文件头部字段（与文档示例一致）。 +type FrontMatter struct { + Name string `yaml:"name"` + ID string `yaml:"id"` + Description string `yaml:"description"` + Tools interface{} `yaml:"tools"` // 字符串 "A, B" 或 []string + MaxIterations int `yaml:"max_iterations"` + BindRole string `yaml:"bind_role,omitempty"` + Kind string `yaml:"kind,omitempty"` // orchestrator = 主代理（亦可仅用文件名 orchestrator.md） +} + +// OrchestratorMarkdown 从 agents 目录解析出的主代理（Deep 协调者）定义。 +type OrchestratorMarkdown struct { + Filename string + EinoName string // 写入 deep.Config.Name / 流式事件过滤 + DisplayName string + Description string + Instruction string +} + +// MarkdownDirLoad 一次扫描 agents 目录的结果（子代理不含主代理文件）。 +type MarkdownDirLoad struct { + SubAgents []config.MultiAgentSubConfig + Orchestrator *OrchestratorMarkdown // Deep 主代理 + OrchestratorPlanExecute *OrchestratorMarkdown // plan_execute 规划主代理 + OrchestratorSupervisor *OrchestratorMarkdown // supervisor 监督主代理 + FileEntries []FileAgent // 含主代理与所有子代理，供管理 API 列表 +} + +// OrchestratorMarkdownKind 按固定文件名返回主代理类型：deep、plan_execute、supervisor；否则返回空。 +func OrchestratorMarkdownKind(filename string) string { + base := filepath.Base(strings.TrimSpace(filename)) + switch { + case strings.EqualFold(base, OrchestratorPlanExecuteMarkdownFilename): + return "plan_execute" + case strings.EqualFold(base, OrchestratorSupervisorMarkdownFilename): + return "supervisor" + case strings.EqualFold(base, OrchestratorMarkdownFilename): + return "deep" + default: + return "" + } +} + +// IsOrchestratorMarkdown 判断该文件是否占用 **Deep** 主代理槽位：orchestrator.md、或 kind: orchestrator（不含 plan_execute / supervisor 专用文件名）。 +func IsOrchestratorMarkdown(filename string, fm FrontMatter) bool { + base := filepath.Base(strings.TrimSpace(filename)) + switch OrchestratorMarkdownKind(base) { + case "plan_execute", "supervisor": + return false + } + if strings.EqualFold(base, OrchestratorMarkdownFilename) { + return true + } + return strings.EqualFold(strings.TrimSpace(fm.Kind), "orchestrator") +} + +// IsOrchestratorLikeMarkdown 是否应在前端/API 中显示为「主代理类」文件。 +func IsOrchestratorLikeMarkdown(filename string, kind string) bool { + if OrchestratorMarkdownKind(filename) != "" { + return true + } + return IsOrchestratorMarkdown(filename, FrontMatter{Kind: kind}) +} + +// WantsMarkdownOrchestrator 保存前判断是否会把该文件作为主代理（用于唯一性校验）。 +func WantsMarkdownOrchestrator(filename string, kindField string, raw string) bool { + base := filepath.Base(strings.TrimSpace(filename)) + if OrchestratorMarkdownKind(base) != "" { + return true + } + if strings.EqualFold(strings.TrimSpace(kindField), "orchestrator") { + return true + } + if strings.EqualFold(base, OrchestratorMarkdownFilename) { + return true + } + if strings.TrimSpace(raw) == "" { + return false + } + sub, err := ParseMarkdownSubAgent(filename, raw) + if err != nil { + return false + } + return strings.EqualFold(strings.TrimSpace(sub.Kind), "orchestrator") +} + +// SplitFrontMatter 分离 YAML front matter 与正文（--- ... ---）。 +func SplitFrontMatter(content string) (frontYAML string, body string, err error) { + s := strings.TrimSpace(content) + if !strings.HasPrefix(s, "---") { + return "", s, nil + } + rest := strings.TrimPrefix(s, "---") + rest = strings.TrimLeft(rest, "\r\n") + end := strings.Index(rest, "\n---") + if end < 0 { + return "", "", fmt.Errorf("agents: 缺少结束的 --- 分隔符") + } + fm := strings.TrimSpace(rest[:end]) + body = strings.TrimSpace(rest[end+4:]) + body = strings.TrimLeft(body, "\r\n") + return fm, body, nil +} + +func parseToolsField(v interface{}) []string { + if v == nil { + return nil + } + switch t := v.(type) { + case string: + return splitToolList(t) + case []interface{}: + var out []string + for _, x := range t { + if s, ok := x.(string); ok && strings.TrimSpace(s) != "" { + out = append(out, strings.TrimSpace(s)) + } + } + return out + case []string: + var out []string + for _, s := range t { + if strings.TrimSpace(s) != "" { + out = append(out, strings.TrimSpace(s)) + } + } + return out + default: + return nil + } +} + +func splitToolList(s string) []string { + s = strings.TrimSpace(s) + if s == "" { + return nil + } + parts := strings.FieldsFunc(s, func(r rune) bool { + return r == ',' || r == ';' || r == '|' + }) + var out []string + for _, p := range parts { + p = strings.TrimSpace(p) + if p != "" { + out = append(out, p) + } + } + return out +} + +// SlugID 从 name 生成可用的代理 id（小写、连字符）。 +func SlugID(name string) string { + var b strings.Builder + name = strings.TrimSpace(strings.ToLower(name)) + lastDash := false + for _, r := range name { + switch { + case unicode.IsLetter(r) && r < unicode.MaxASCII, unicode.IsDigit(r): + b.WriteRune(r) + lastDash = false + case r == ' ' || r == '_' || r == '/' || r == '.': + if !lastDash && b.Len() > 0 { + b.WriteByte('-') + lastDash = true + } + } + } + s := strings.Trim(b.String(), "-") + if s == "" { + return "agent" + } + return s +} + +// sanitizeEinoAgentID 规范化 Deep 主代理在 Eino 中的 Name：小写 ASCII、数字、连字符，与默认 cyberstrike-deep 一致。 +func sanitizeEinoAgentID(s string) string { + s = strings.TrimSpace(strings.ToLower(s)) + var b strings.Builder + for _, r := range s { + switch { + case unicode.IsLetter(r) && r < unicode.MaxASCII, unicode.IsDigit(r): + b.WriteRune(r) + case r == '-': + b.WriteRune(r) + } + } + out := strings.Trim(b.String(), "-") + if out == "" { + return "cyberstrike-deep" + } + return out +} + +func parseMarkdownAgentRaw(filename string, content string) (FrontMatter, string, error) { + var fm FrontMatter + fmStr, body, err := SplitFrontMatter(content) + if err != nil { + return fm, "", err + } + if strings.TrimSpace(fmStr) == "" { + return fm, "", fmt.Errorf("agents: %s 无 YAML front matter", filename) + } + if err := yaml.Unmarshal([]byte(fmStr), &fm); err != nil { + return fm, "", fmt.Errorf("agents: 解析 front matter: %w", err) + } + return fm, body, nil +} + +func orchestratorFromParsed(filename string, fm FrontMatter, body string) (*OrchestratorMarkdown, error) { + display := strings.TrimSpace(fm.Name) + if display == "" { + display = "Orchestrator" + } + rawID := strings.TrimSpace(fm.ID) + if rawID == "" { + rawID = SlugID(display) + } + eino := sanitizeEinoAgentID(rawID) + return &OrchestratorMarkdown{ + Filename: filepath.Base(strings.TrimSpace(filename)), + EinoName: eino, + DisplayName: display, + Description: strings.TrimSpace(fm.Description), + Instruction: strings.TrimSpace(body), + }, nil +} + +func orchestratorConfigFromOrchestrator(o *OrchestratorMarkdown) config.MultiAgentSubConfig { + if o == nil { + return config.MultiAgentSubConfig{} + } + return config.MultiAgentSubConfig{ + ID: o.EinoName, + Name: o.DisplayName, + Description: o.Description, + Instruction: o.Instruction, + Kind: "orchestrator", + } +} + +func subAgentFromFrontMatter(filename string, fm FrontMatter, body string) (config.MultiAgentSubConfig, error) { + var out config.MultiAgentSubConfig + name := strings.TrimSpace(fm.Name) + if name == "" { + return out, fmt.Errorf("agents: %s 缺少 name 字段", filename) + } + id := strings.TrimSpace(fm.ID) + if id == "" { + id = SlugID(name) + } + out.ID = id + out.Name = name + out.Description = strings.TrimSpace(fm.Description) + out.Instruction = strings.TrimSpace(body) + out.RoleTools = parseToolsField(fm.Tools) + out.MaxIterations = fm.MaxIterations + out.BindRole = strings.TrimSpace(fm.BindRole) + out.Kind = strings.TrimSpace(fm.Kind) + return out, nil +} + +func collectMarkdownBasenames(dir string) ([]string, error) { + if strings.TrimSpace(dir) == "" { + return nil, nil + } + st, err := os.Stat(dir) + if err != nil { + if os.IsNotExist(err) { + return nil, nil + } + return nil, err + } + if !st.IsDir() { + return nil, fmt.Errorf("agents: 不是目录: %s", dir) + } + entries, err := os.ReadDir(dir) + if err != nil { + return nil, err + } + var names []string + for _, e := range entries { + if e.IsDir() { + continue + } + n := e.Name() + if strings.HasPrefix(n, ".") { + continue + } + if !strings.EqualFold(filepath.Ext(n), ".md") { + continue + } + if strings.EqualFold(n, "README.md") { + continue + } + names = append(names, n) + } + sort.Strings(names) + return names, nil +} + +// LoadMarkdownAgentsDir 扫描 agents 目录：拆出 Deep / plan_execute / supervisor 主代理各至多一个，及其余子代理。 +func LoadMarkdownAgentsDir(dir string) (*MarkdownDirLoad, error) { + out := &MarkdownDirLoad{} + names, err := collectMarkdownBasenames(dir) + if err != nil { + return nil, err + } + for _, n := range names { + p := filepath.Join(dir, n) + b, err := os.ReadFile(p) + if err != nil { + return nil, err + } + fm, body, err := parseMarkdownAgentRaw(n, string(b)) + if err != nil { + return nil, fmt.Errorf("%s: %w", n, err) + } + switch OrchestratorMarkdownKind(n) { + case "plan_execute": + if out.OrchestratorPlanExecute != nil { + return nil, fmt.Errorf("agents: 仅能定义一个 %s，已有 %s", OrchestratorPlanExecuteMarkdownFilename, out.OrchestratorPlanExecute.Filename) + } + orch, err := orchestratorFromParsed(n, fm, body) + if err != nil { + return nil, fmt.Errorf("%s: %w", n, err) + } + out.OrchestratorPlanExecute = orch + out.FileEntries = append(out.FileEntries, FileAgent{ + Filename: n, + Config: orchestratorConfigFromOrchestrator(orch), + IsOrchestrator: true, + }) + continue + case "supervisor": + if out.OrchestratorSupervisor != nil { + return nil, fmt.Errorf("agents: 仅能定义一个 %s，已有 %s", OrchestratorSupervisorMarkdownFilename, out.OrchestratorSupervisor.Filename) + } + orch, err := orchestratorFromParsed(n, fm, body) + if err != nil { + return nil, fmt.Errorf("%s: %w", n, err) + } + out.OrchestratorSupervisor = orch + out.FileEntries = append(out.FileEntries, FileAgent{ + Filename: n, + Config: orchestratorConfigFromOrchestrator(orch), + IsOrchestrator: true, + }) + continue + } + if IsOrchestratorMarkdown(n, fm) { + if out.Orchestrator != nil { + return nil, fmt.Errorf("agents: 仅能定义一个主代理（Deep 协调者），已有 %s，又与 %s 冲突", out.Orchestrator.Filename, n) + } + orch, err := orchestratorFromParsed(n, fm, body) + if err != nil { + return nil, fmt.Errorf("%s: %w", n, err) + } + out.Orchestrator = orch + out.FileEntries = append(out.FileEntries, FileAgent{ + Filename: n, + Config: orchestratorConfigFromOrchestrator(orch), + IsOrchestrator: true, + }) + continue + } + sub, err := subAgentFromFrontMatter(n, fm, body) + if err != nil { + return nil, fmt.Errorf("%s: %w", n, err) + } + out.SubAgents = append(out.SubAgents, sub) + out.FileEntries = append(out.FileEntries, FileAgent{Filename: n, Config: sub, IsOrchestrator: false}) + } + return out, nil +} + +// ParseMarkdownSubAgent 将单个 Markdown 文件解析为 MultiAgentSubConfig。 +func ParseMarkdownSubAgent(filename string, content string) (config.MultiAgentSubConfig, error) { + fm, body, err := parseMarkdownAgentRaw(filename, content) + if err != nil { + return config.MultiAgentSubConfig{}, err + } + if OrchestratorMarkdownKind(filename) != "" { + orch, err := orchestratorFromParsed(filename, fm, body) + if err != nil { + return config.MultiAgentSubConfig{}, err + } + return orchestratorConfigFromOrchestrator(orch), nil + } + if IsOrchestratorMarkdown(filename, fm) { + orch, err := orchestratorFromParsed(filename, fm, body) + if err != nil { + return config.MultiAgentSubConfig{}, err + } + return orchestratorConfigFromOrchestrator(orch), nil + } + return subAgentFromFrontMatter(filename, fm, body) +} + +// LoadMarkdownSubAgents 读取目录下所有子代理 .md（不含主代理 orchestrator.md / kind: orchestrator）。 +func LoadMarkdownSubAgents(dir string) ([]config.MultiAgentSubConfig, error) { + load, err := LoadMarkdownAgentsDir(dir) + if err != nil { + return nil, err + } + return load.SubAgents, nil +} + +// FileAgent 单个 Markdown 文件及其解析结果。 +type FileAgent struct { + Filename string + Config config.MultiAgentSubConfig + IsOrchestrator bool +} + +// LoadMarkdownAgentFiles 列出目录下全部 .md（含主代理），供管理 API 使用。 +func LoadMarkdownAgentFiles(dir string) ([]FileAgent, error) { + load, err := LoadMarkdownAgentsDir(dir) + if err != nil { + return nil, err + } + return load.FileEntries, nil +} + +// MergeYAMLAndMarkdown 合并 config.yaml 中的 sub_agents 与 Markdown 定义：同 id 时 Markdown 覆盖 YAML；仅存在于 Markdown 的条目追加在 YAML 顺序之后。 +func MergeYAMLAndMarkdown(yamlSubs []config.MultiAgentSubConfig, mdSubs []config.MultiAgentSubConfig) []config.MultiAgentSubConfig { + mdByID := make(map[string]config.MultiAgentSubConfig) + for _, m := range mdSubs { + id := strings.TrimSpace(m.ID) + if id == "" { + continue + } + mdByID[id] = m + } + yamlIDSet := make(map[string]bool) + for _, y := range yamlSubs { + yamlIDSet[strings.TrimSpace(y.ID)] = true + } + out := make([]config.MultiAgentSubConfig, 0, len(yamlSubs)+len(mdSubs)) + for _, y := range yamlSubs { + id := strings.TrimSpace(y.ID) + if id == "" { + continue + } + if m, ok := mdByID[id]; ok { + out = append(out, m) + } else { + out = append(out, y) + } + } + for _, m := range mdSubs { + id := strings.TrimSpace(m.ID) + if id == "" || yamlIDSet[id] { + continue + } + out = append(out, m) + } + return out +} + +// EffectiveSubAgents 供多代理运行时使用。 +func EffectiveSubAgents(yamlSubs []config.MultiAgentSubConfig, agentsDir string) ([]config.MultiAgentSubConfig, error) { + md, err := LoadMarkdownSubAgents(agentsDir) + if err != nil { + return nil, err + } + if len(md) == 0 { + return yamlSubs, nil + } + return MergeYAMLAndMarkdown(yamlSubs, md), nil +} + +// BuildMarkdownFile 根据配置序列化为可写回磁盘的 Markdown。 +func BuildMarkdownFile(sub config.MultiAgentSubConfig) ([]byte, error) { + fm := FrontMatter{ + Name: sub.Name, + ID: sub.ID, + Description: sub.Description, + MaxIterations: sub.MaxIterations, + BindRole: sub.BindRole, + } + if k := strings.TrimSpace(sub.Kind); k != "" { + fm.Kind = k + } + if len(sub.RoleTools) > 0 { + fm.Tools = sub.RoleTools + } + head, err := yaml.Marshal(fm) + if err != nil { + return nil, err + } + var b strings.Builder + b.WriteString("---\n") + b.Write(head) + b.WriteString("---\n\n") + b.WriteString(strings.TrimSpace(sub.Instruction)) + if !strings.HasSuffix(sub.Instruction, "\n") && sub.Instruction != "" { + b.WriteString("\n") + } + return []byte(b.String()), nil +} diff --git a/internal/agents/markdown_orchestrator_test.go b/internal/agents/markdown_orchestrator_test.go new file mode 100644 index 00000000..9ea7474d --- /dev/null +++ b/internal/agents/markdown_orchestrator_test.go @@ -0,0 +1,97 @@ +package agents + +import ( + "os" + "path/filepath" + "testing" +) + +func TestLoadMarkdownAgentsDir_OrchestratorExcludedFromSubs(t *testing.T) { + dir := t.TempDir() + orch := filepath.Join(dir, OrchestratorMarkdownFilename) + if err := os.WriteFile(orch, []byte(`--- +id: cyberstrike-deep +name: Main +description: Test desc +--- + +Hello orchestrator +`), 0644); err != nil { + t.Fatal(err) + } + subPath := filepath.Join(dir, "worker.md") + if err := os.WriteFile(subPath, []byte(`--- +id: worker +name: Worker +description: W +--- + +Do work +`), 0644); err != nil { + t.Fatal(err) + } + load, err := LoadMarkdownAgentsDir(dir) + if err != nil { + t.Fatal(err) + } + if load.Orchestrator == nil || load.Orchestrator.EinoName != "cyberstrike-deep" { + t.Fatalf("orchestrator: %+v", load.Orchestrator) + } + if len(load.SubAgents) != 1 || load.SubAgents[0].ID != "worker" { + t.Fatalf("subs: %+v", load.SubAgents) + } + if len(load.FileEntries) != 2 { + t.Fatalf("file entries: %d", len(load.FileEntries)) + } + var orchFile *FileAgent + for i := range load.FileEntries { + if load.FileEntries[i].IsOrchestrator { + orchFile = &load.FileEntries[i] + break + } + } + if orchFile == nil || orchFile.Filename != OrchestratorMarkdownFilename { + t.Fatal("missing orchestrator file entry") + } +} + +func TestLoadMarkdownAgentsDir_DuplicateOrchestrator(t *testing.T) { + dir := t.TempDir() + _ = os.WriteFile(filepath.Join(dir, OrchestratorMarkdownFilename), []byte("---\nname: A\n---\n\nx\n"), 0644) + _ = os.WriteFile(filepath.Join(dir, "b.md"), []byte("---\nname: B\nkind: orchestrator\n---\n\ny\n"), 0644) + _, err := LoadMarkdownAgentsDir(dir) + if err == nil { + t.Fatal("expected duplicate orchestrator error") + } +} + +func TestLoadMarkdownAgentsDir_ModeOrchestratorsCoexist(t *testing.T) { + dir := t.TempDir() + write := func(name, body string) { + t.Helper() + if err := os.WriteFile(filepath.Join(dir, name), []byte(body), 0644); err != nil { + t.Fatal(err) + } + } + write(OrchestratorMarkdownFilename, "---\nname: Deep\n---\n\ndeep\n") + write(OrchestratorPlanExecuteMarkdownFilename, "---\nname: PE\n---\n\npe\n") + write(OrchestratorSupervisorMarkdownFilename, "---\nname: SV\n---\n\nsv\n") + write("worker.md", "---\nid: worker\nname: Worker\n---\n\nw\n") + + load, err := LoadMarkdownAgentsDir(dir) + if err != nil { + t.Fatal(err) + } + if load.Orchestrator == nil || load.Orchestrator.Instruction != "deep" { + t.Fatalf("deep: %+v", load.Orchestrator) + } + if load.OrchestratorPlanExecute == nil || load.OrchestratorPlanExecute.Instruction != "pe" { + t.Fatalf("pe: %+v", load.OrchestratorPlanExecute) + } + if load.OrchestratorSupervisor == nil || load.OrchestratorSupervisor.Instruction != "sv" { + t.Fatalf("sv: %+v", load.OrchestratorSupervisor) + } + if len(load.SubAgents) != 1 || load.SubAgents[0].ID != "worker" { + t.Fatalf("subs: %+v", load.SubAgents) + } +} diff --git a/internal/config/config.go b/internal/config/config.go new file mode 100644 index 00000000..21efe233 --- /dev/null +++ b/internal/config/config.go @@ -0,0 +1,2108 @@ +package config + +import ( + "crypto/rand" + "encoding/hex" + "encoding/json" + "fmt" + "io/fs" + "os" + "path/filepath" + "strconv" + "strings" + + "cyberstrike-ai/internal/termout" + + "gopkg.in/yaml.v3" +) + +type Config struct { + Version string `yaml:"version,omitempty" json:"version,omitempty"` // 前端显示的版本号，如 v1.3.3 + Server ServerConfig `yaml:"server"` + Log LogConfig `yaml:"log"` + MCP MCPConfig `yaml:"mcp"` + AI AIConfig `yaml:"ai,omitempty" json:"ai,omitempty"` + OpenAI OpenAIConfig `yaml:"openai,omitempty" json:"openai,omitempty"` + FOFA FofaConfig `yaml:"fofa,omitempty" json:"fofa,omitempty"` + ZoomEye SpaceSearchConfig `yaml:"zoomeye,omitempty" json:"zoomeye,omitempty"` + Quake SpaceSearchConfig `yaml:"quake,omitempty" json:"quake,omitempty"` + Shodan SpaceSearchConfig `yaml:"shodan,omitempty" json:"shodan,omitempty"` + Agent AgentConfig `yaml:"agent"` + Hitl HitlConfig `yaml:"hitl,omitempty" json:"hitl,omitempty"` + Security SecurityConfig `yaml:"security"` + Database DatabaseConfig `yaml:"database"` + Auth AuthConfig `yaml:"auth"` + Audit AuditConfig `yaml:"audit,omitempty" json:"audit,omitempty"` + Monitor MonitorConfig `yaml:"monitor,omitempty" json:"monitor,omitempty"` + ExternalMCP ExternalMCPConfig `yaml:"external_mcp,omitempty"` + Knowledge KnowledgeConfig `yaml:"knowledge,omitempty"` + C2 C2Config `yaml:"c2,omitempty" json:"c2,omitempty"` // 内置 C2 总开关；未配置时默认启用 + Robots RobotsConfig `yaml:"robots,omitempty" json:"robots,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配置文件目录 + AgentsDir string `yaml:"agents_dir,omitempty" json:"agents_dir,omitempty"` // 多代理子 Agent Markdown 定义目录（*.md，YAML front matter） + MultiAgent MultiAgentConfig `yaml:"multi_agent,omitempty" json:"multi_agent,omitempty"` + Project ProjectConfig `yaml:"project,omitempty" json:"project,omitempty"` + Vision VisionConfig `yaml:"vision,omitempty" json:"vision,omitempty"` +} + +type EnsureLocalConfigResult struct { + Created bool + ExamplePath string +} + +const ( + DefaultMaxCompletionTokens = 16384 + DefaultSummarizationUserIntentLedgerMaxRunes = 96000 + DefaultSummarizationUserIntentLedgerEntryMaxRunes = 16000 + DefaultLatestUserMessageMaxRunes = 48000 + DefaultLatestUserMessageHeadRunes = 24000 + DefaultLatestUserMessageTailRunes = 24000 + DefaultSummarizationOutputReserveTokens = 8192 +) + +// ProjectConfig 项目黑板（跨对话共享事实）配置。 +type ProjectConfig struct { + Enabled bool `yaml:"enabled" json:"enabled"` + DefaultProjectID string `yaml:"default_project_id,omitempty" json:"default_project_id,omitempty"` // 机器人/批量等无显式项目时绑定的默认项目 + FactIndexMaxRunes int `yaml:"fact_index_max_runes,omitempty" json:"fact_index_max_runes,omitempty"` + FactIndexPathMaxRunes int `yaml:"fact_index_path_max_runes,omitempty" json:"fact_index_path_max_runes,omitempty"` + FactSummaryMaxRunes int `yaml:"fact_summary_max_runes,omitempty" json:"fact_summary_max_runes,omitempty"` + DefaultInjectDeprecated bool `yaml:"default_inject_deprecated,omitempty" json:"default_inject_deprecated,omitempty"` +} + +// FactIndexMaxRunesEffective 自动注入黑板索引的最大 rune 数。 +func (c ProjectConfig) FactIndexMaxRunesEffective() int { + if c.FactIndexMaxRunes <= 0 { + return 3500 + } + return c.FactIndexMaxRunes +} + +// FactIndexPathMaxRunesEffective 攻击路径速览段的最大 rune 数（从 fact_index_max_runes 预算中预留）。 +func (c ProjectConfig) FactIndexPathMaxRunesEffective() int { + if c.FactIndexPathMaxRunes <= 0 { + return 1000 + } + return c.FactIndexPathMaxRunes +} + +// FactSummaryMaxRunesEffective upsert 时 summary 最大 rune 数（索引一行，宜含验证要点）。 +func (c ProjectConfig) FactSummaryMaxRunesEffective() int { + if c.FactSummaryMaxRunes <= 0 { + return 200 + } + return c.FactSummaryMaxRunes +} + +// MultiAgentConfig 基于 CloudWeGo Eino adk/prebuilt 的多代理编排（deep | plan_execute | supervisor）。 +type MultiAgentConfig struct { + Enabled bool `yaml:"enabled" json:"enabled"` + RobotDefaultAgentMode string `yaml:"robot_default_agent_mode,omitempty" json:"robot_default_agent_mode,omitempty"` // eino_single | deep | plan_execute | supervisor + BatchUseMultiAgent bool `yaml:"batch_use_multi_agent" json:"batch_use_multi_agent"` // 为 true 时批量任务队列中每子任务走 Eino 多代理 + // Orchestration 已弃用：保留仅兼容旧版 config.yaml；编排由聊天/WebShell 请求体 orchestration 决定，未传时按 deep。 + Orchestration string `yaml:"orchestration,omitempty" json:"orchestration,omitempty"` + // MaxIteration 已废弃：统一使用 agent.max_iterations（YAML 中保留字段仅为兼容旧配置，运行时不读取）。 + MaxIteration int `yaml:"max_iteration,omitempty" json:"max_iteration,omitempty"` + // PlanExecuteLoopMaxIterations plan_execute 模式下 execute↔replan 外层循环上限；0 表示用 Eino 默认 10。 + PlanExecuteLoopMaxIterations int `yaml:"plan_execute_loop_max_iterations,omitempty" json:"plan_execute_loop_max_iterations,omitempty"` + // SubAgentMaxIterations 已废弃：子代理与主代理均使用 agent.max_iterations（Markdown max_iterations>0 可覆盖）。 + SubAgentMaxIterations int `yaml:"sub_agent_max_iterations,omitempty" json:"sub_agent_max_iterations,omitempty"` + WithoutGeneralSubAgent bool `yaml:"without_general_sub_agent" json:"without_general_sub_agent"` + WithoutWriteTodos bool `yaml:"without_write_todos" json:"without_write_todos"` + OrchestratorInstruction string `yaml:"orchestrator_instruction" json:"orchestrator_instruction"` + // OrchestratorInstructionPlanExecute plan_execute 主代理（规划侧）系统提示；非空且 agents/orchestrator-plan-execute.md 正文为空或未存在时生效。不与 Deep 的 orchestrator_instruction 混用。 + OrchestratorInstructionPlanExecute string `yaml:"orchestrator_instruction_plan_execute,omitempty" json:"orchestrator_instruction_plan_execute,omitempty"` + // OrchestratorInstructionSupervisor supervisor 主代理系统提示（transfer/exit 说明仍由运行追加）；非空且 agents/orchestrator-supervisor.md 正文为空或未存在时生效。 + OrchestratorInstructionSupervisor string `yaml:"orchestrator_instruction_supervisor,omitempty" json:"orchestrator_instruction_supervisor,omitempty"` + SubAgents []MultiAgentSubConfig `yaml:"sub_agents" json:"sub_agents"` + // SubAgentUserContextMaxRunes caps user-context supplement for sub-agent task descriptions. + // 0 (default) preserves all user turns verbatim; >0 caps total runes; negative disables injection. + SubAgentUserContextMaxRunes int `yaml:"sub_agent_user_context_max_runes,omitempty" json:"sub_agent_user_context_max_runes,omitempty"` + // EinoSkills configures CloudWeGo Eino ADK skill middleware + optional local filesystem/execute on DeepAgent. + EinoSkills MultiAgentEinoSkillsConfig `yaml:"eino_skills,omitempty" json:"eino_skills,omitempty"` + // EinoMiddleware wires optional ADK middleware (patchtoolcalls, toolsearch, plantask, reduction) and Deep extras. + EinoMiddleware MultiAgentEinoMiddlewareConfig `yaml:"eino_middleware,omitempty" json:"eino_middleware,omitempty"` + // EinoCallbacks attaches CloudWeGo eino callbacks.InitCallbacks on ADK Runner context (structured logs + optional SSE trace). + EinoCallbacks MultiAgentEinoCallbacksConfig `yaml:"eino_callbacks,omitempty" json:"eino_callbacks,omitempty"` +} + +// SubAgentUserContextMaxRunesEffective returns max runes for sub-agent task supplement; 0 = unlimited; negative = disabled. +func (c MultiAgentConfig) SubAgentUserContextMaxRunesEffective() int { + return c.SubAgentUserContextMaxRunes +} + +// MultiAgentEinoCallbacksConfig enables Eino unified callbacks on each ADK agent run (deep / plan_execute / supervisor / eino_single). +// Modes: log_only (zap + optional OTel; no SSE to browser), sse (adds client SSE eino_trace_* when sse_trace_to_client), full (sse rules + stream callback copies closed). +type MultiAgentEinoCallbacksConfig struct { + Enabled bool `yaml:"enabled" json:"enabled"` + Mode string `yaml:"mode,omitempty" json:"mode,omitempty"` // log_only | sse | full; empty with enabled=true defaults to log_only + // SseTraceToClient when true emits eino_trace_* SSE for UI (use only for admin/debug; nil/false recommended in production). + SseTraceToClient *bool `yaml:"sse_trace_to_client,omitempty" json:"sse_trace_to_client,omitempty"` + // Otel configures OpenTelemetry trace export (independent of mode; exporter none disables export even if enabled). + Otel MultiAgentEinoCallbacksOtelConfig `yaml:"otel,omitempty" json:"otel,omitempty"` + // MaxInputSummaryRunes / MaxOutputSummaryRunes cap text placed in SSE payloads and debug logs (not full payloads). + MaxInputSummaryRunes int `yaml:"max_input_summary_runes,omitempty" json:"max_input_summary_runes,omitempty"` + MaxOutputSummaryRunes int `yaml:"max_output_summary_runes,omitempty" json:"max_output_summary_runes,omitempty"` + // ZapVerbose when true logs input/output summaries at zap.Debug on start/end; false uses Info with short fields only. + ZapVerbose bool `yaml:"zap_verbose,omitempty" json:"zap_verbose,omitempty"` +} + +// MultiAgentEinoCallbacksOtelConfig OpenTelemetry for Eino callback spans (W3C trace in collector / stdout). +type MultiAgentEinoCallbacksOtelConfig struct { + Enabled bool `yaml:"enabled" json:"enabled"` + ServiceName string `yaml:"service_name,omitempty" json:"service_name,omitempty"` + Exporter string `yaml:"exporter,omitempty" json:"exporter,omitempty"` // none | stdout | otlphttp + OTLPEndpoint string `yaml:"otlp_endpoint,omitempty" json:"otlp_endpoint,omitempty"` // host:port, e.g. localhost:4318 (path /v1/traces) + SampleRatio float64 `yaml:"sample_ratio,omitempty" json:"sample_ratio,omitempty"` // 0–1, default 1.0 +} + +// EinoCallbacksModeEffective returns off | log_only | sse | full. +func (c MultiAgentEinoCallbacksConfig) EinoCallbacksModeEffective() string { + if !c.Enabled { + return "off" + } + m := strings.TrimSpace(strings.ToLower(c.Mode)) + switch m { + case "log_only": + return "log_only" + case "sse": + return "sse" + case "full": + return "full" + case "": + return "log_only" + default: + return "log_only" + } +} + +// SseTraceToClientEffective is false unless explicitly set true (best practice: do not expose framework traces to end users by default). +func (c MultiAgentEinoCallbacksConfig) SseTraceToClientEffective() bool { + if c.SseTraceToClient == nil { + return false + } + return *c.SseTraceToClient +} + +// ShouldEmitEinoTraceSSE is true when client-visible trace events should be sent over progress/SSE. +func (c MultiAgentEinoCallbacksConfig) ShouldEmitEinoTraceSSE(mode string) bool { + if !c.SseTraceToClientEffective() { + return false + } + return mode == "sse" || mode == "full" +} + +// OtelExporterEffective returns none | stdout | otlphttp. +func (c MultiAgentEinoCallbacksOtelConfig) OtelExporterEffective() string { + e := strings.TrimSpace(strings.ToLower(c.Exporter)) + switch e { + case "none", "stdout", "otlphttp": + return e + case "": + if c.Enabled { + return "stdout" + } + return "none" + default: + return "none" + } +} + +// OtelTracingActive is true when spans should be started (enabled + non-none exporter). +func (c MultiAgentEinoCallbacksConfig) OtelTracingActive() bool { + if !c.Otel.Enabled { + return false + } + return c.Otel.OtelExporterEffective() != "none" +} + +func (c MultiAgentEinoCallbacksOtelConfig) ServiceNameEffective() string { + s := strings.TrimSpace(c.ServiceName) + if s != "" { + return s + } + return "cyberstrike-ai" +} + +func (c MultiAgentEinoCallbacksOtelConfig) SampleRatioEffective() float64 { + r := c.SampleRatio + if r <= 0 { + return 1.0 + } + if r > 1 { + return 1.0 + } + return r +} + +func (c MultiAgentEinoCallbacksConfig) EinoCallbacksMaxInputSummaryRunes() int { + if c.MaxInputSummaryRunes > 0 { + return c.MaxInputSummaryRunes + } + return 400 +} + +func (c MultiAgentEinoCallbacksConfig) EinoCallbacksMaxOutputSummaryRunes() int { + if c.MaxOutputSummaryRunes > 0 { + return c.MaxOutputSummaryRunes + } + return 400 +} + +// MultiAgentEinoMiddlewareConfig optional Eino ADK middleware and Deep / supervisor tuning. +type MultiAgentEinoMiddlewareConfig struct { + // PatchToolCalls inserts placeholder tool results for dangling assistant tool_calls (nil = enabled). + PatchToolCalls *bool `yaml:"patch_tool_calls,omitempty" json:"patch_tool_calls,omitempty"` + // ToolSearch enables dynamictool/toolsearch: hide tail tools until model calls tool_search (reduces prompt tools). + ToolSearchEnable bool `yaml:"tool_search_enable,omitempty" json:"tool_search_enable,omitempty"` + ToolSearchMinTools int `yaml:"tool_search_min_tools,omitempty" json:"tool_search_min_tools,omitempty"` // default 20; applies when len(tools) >= this + ToolSearchAlwaysVisible int `yaml:"tool_search_always_visible,omitempty" json:"tool_search_always_visible,omitempty"` // default 12; first N tools stay always visible + // ToolSearchAlwaysVisibleTools keeps specified tool names always visible (never hidden by tool_search). + ToolSearchAlwaysVisibleTools []string `yaml:"tool_search_always_visible_tools,omitempty" json:"tool_search_always_visible_tools,omitempty"` + // Plantask adds TaskCreate/Get/Update/List (file-backed under skills dir); requires eino_skills + local backend. + PlantaskEnable bool `yaml:"plantask_enable,omitempty" json:"plantask_enable,omitempty"` + // PlantaskRelDir relative to skills_dir for per-conversation task boards (default .eino/plantask). + PlantaskRelDir string `yaml:"plantask_rel_dir,omitempty" json:"plantask_rel_dir,omitempty"` + // Reduction truncates/offloads large tool outputs (requires eino local backend for Write). + ReductionEnable bool `yaml:"reduction_enable,omitempty" json:"reduction_enable,omitempty"` + ReductionRootDir string `yaml:"reduction_root_dir,omitempty" json:"reduction_root_dir,omitempty"` // 非空：落盘根目录（默认 tmp/reduction）；其下按 projects/{id} 或 conversations/{id} 隔离 + ReductionMaxLengthForTrunc int `yaml:"reduction_max_length_for_trunc,omitempty" json:"reduction_max_length_for_trunc,omitempty"` // default 12000 + ReductionMaxTokensForClear int `yaml:"reduction_max_tokens_for_clear,omitempty" json:"reduction_max_tokens_for_clear,omitempty"` // default 50000 + ReductionClearExclude []string `yaml:"reduction_clear_exclude,omitempty" json:"reduction_clear_exclude,omitempty"` + ReductionSubAgents bool `yaml:"reduction_sub_agents,omitempty" json:"reduction_sub_agents,omitempty"` // also attach to sub-agents + // SummarizationTriggerRatio controls summarization trigger threshold as max_total_tokens * ratio (default 0.8). + SummarizationTriggerRatio float64 `yaml:"summarization_trigger_ratio,omitempty" json:"summarization_trigger_ratio,omitempty"` + // SummarizationOutputReserveTokens reserves completion headroom for the summarization model call (default 8192). + SummarizationOutputReserveTokens int `yaml:"summarization_output_reserve_tokens,omitempty" json:"summarization_output_reserve_tokens,omitempty"` + // SummarizationEmitInternalEvents controls middleware internal event emission (default true). + SummarizationEmitInternalEvents *bool `yaml:"summarization_emit_internal_events,omitempty" json:"summarization_emit_internal_events,omitempty"` + // SummarizationUserIntentLedgerMaxRunes caps the DB-backed immutable user input ledger injected into model context. + SummarizationUserIntentLedgerMaxRunes int `yaml:"summarization_user_intent_ledger_max_runes,omitempty" json:"summarization_user_intent_ledger_max_runes,omitempty"` + // SummarizationUserIntentLedgerEntryMaxRunes caps each user message entry inside the immutable user input ledger. + SummarizationUserIntentLedgerEntryMaxRunes int `yaml:"summarization_user_intent_ledger_entry_max_runes,omitempty" json:"summarization_user_intent_ledger_entry_max_runes,omitempty"` + // LatestUserMessageMaxRunes caps the current user turn inserted into model context; full text is persisted as an artifact when capped. + LatestUserMessageMaxRunes int `yaml:"latest_user_message_max_runes,omitempty" json:"latest_user_message_max_runes,omitempty"` + // LatestUserMessageHeadRunes keeps the head preview for an oversized current user turn. + LatestUserMessageHeadRunes int `yaml:"latest_user_message_head_runes,omitempty" json:"latest_user_message_head_runes,omitempty"` + // LatestUserMessageTailRunes keeps the tail preview for an oversized current user turn. + LatestUserMessageTailRunes int `yaml:"latest_user_message_tail_runes,omitempty" json:"latest_user_message_tail_runes,omitempty"` + // SummarizationRetryMaxAttempts 已废弃：summarization 与 Eino 原生 ModelRetry 共用 model_retry_max_retries 及 isEinoTransientRunError。 + SummarizationRetryMaxAttempts int `yaml:"summarization_retry_max_attempts,omitempty" json:"summarization_retry_max_attempts,omitempty"` + // PlanExecuteUserInputBudgetRatio caps planner/replanner/executor userInput prompt budget ratio (default 0.35). + PlanExecuteUserInputBudgetRatio float64 `yaml:"plan_execute_user_input_budget_ratio,omitempty" json:"plan_execute_user_input_budget_ratio,omitempty"` + // PlanExecuteExecutedStepsBudgetRatio caps executed_steps prompt budget ratio (default 0.2). + PlanExecuteExecutedStepsBudgetRatio float64 `yaml:"plan_execute_executed_steps_budget_ratio,omitempty" json:"plan_execute_executed_steps_budget_ratio,omitempty"` + // PlanExecuteMaxStepResultRunes caps each executed step result length for prompt view (default 4000). + PlanExecuteMaxStepResultRunes int `yaml:"plan_execute_max_step_result_runes,omitempty" json:"plan_execute_max_step_result_runes,omitempty"` + // PlanExecuteKeepLastSteps keeps only the tail steps in prompt view (default 8). + PlanExecuteKeepLastSteps int `yaml:"plan_execute_keep_last_steps,omitempty" json:"plan_execute_keep_last_steps,omitempty"` + // CheckpointDir when non-empty enables adk.Runner CheckPointStore (file-backed) for interrupt/resume persistence. + CheckpointDir string `yaml:"checkpoint_dir,omitempty" json:"checkpoint_dir,omitempty"` + // DeepOutputKey passed to deep.Config OutputKey (session final text); empty = off. + DeepOutputKey string `yaml:"deep_output_key,omitempty" json:"deep_output_key,omitempty"` + // DeepModelRetryMaxRetries 已废弃：请用 model_retry_max_retries；保留字段仅为兼容旧配置。 + DeepModelRetryMaxRetries int `yaml:"deep_model_retry_max_retries,omitempty" json:"deep_model_retry_max_retries,omitempty"` + // ModelRetryMaxRetries configures Eino ADK native ChatModel retry attempts; 0=default 4. + ModelRetryMaxRetries int `yaml:"model_retry_max_retries,omitempty" json:"model_retry_max_retries,omitempty"` + // ModelRetryMaxBackoffSec caps native model retry backoff seconds; 0=default 30. + ModelRetryMaxBackoffSec int `yaml:"model_retry_max_backoff_sec,omitempty" json:"model_retry_max_backoff_sec,omitempty"` + // ModelFailoverChannels lists ai.channels IDs to try after native model retry is exhausted. + ModelFailoverChannels []string `yaml:"model_failover_channels,omitempty" json:"model_failover_channels,omitempty"` + // ModelFailoverMaxRetries caps distinct failover channel attempts; 0=all configured failover channels. + ModelFailoverMaxRetries int `yaml:"model_failover_max_retries,omitempty" json:"model_failover_max_retries,omitempty"` + // RunRetryMaxAttempts 已废弃：模型临时错误由 Eino 原生 ModelRetry 处理；仅保留给非模型层 run loop 兜底与 summarization 旧字段。 + RunRetryMaxAttempts int `yaml:"run_retry_max_attempts,omitempty" json:"run_retry_max_attempts,omitempty"` + // RunRetryMaxBackoffSec 已废弃：请用 model_retry_max_backoff_sec；仅保留给非模型层 run loop 兜底与 summarization 旧字段。 + RunRetryMaxBackoffSec int `yaml:"run_retry_max_backoff_sec,omitempty" json:"run_retry_max_backoff_sec,omitempty"` + // EmptyResponseContinueMaxAttempts Run 成功但未捕获助手正文时 Handler 层退避续跑次数；0=默认 5。 + EmptyResponseContinueMaxAttempts int `yaml:"empty_response_continue_max_attempts,omitempty" json:"empty_response_continue_max_attempts,omitempty"` + // TaskToolDescriptionPrefix when non-empty sets deep.Config TaskToolDescriptionGenerator (sub-agent names appended). + TaskToolDescriptionPrefix string `yaml:"task_tool_description_prefix,omitempty" json:"task_tool_description_prefix,omitempty"` +} + +func (c MultiAgentEinoMiddlewareConfig) SummarizationTriggerRatioEffective() float64 { + v := c.SummarizationTriggerRatio + if v <= 0 { + return 0.8 + } + if v < 0.5 { + return 0.5 + } + if v > 0.95 { + return 0.95 + } + return v +} + +func (c MultiAgentEinoMiddlewareConfig) SummarizationOutputReserveTokensEffective() int { + if c.SummarizationOutputReserveTokens > 0 { + return c.SummarizationOutputReserveTokens + } + return DefaultSummarizationOutputReserveTokens +} + +func (c MultiAgentEinoMiddlewareConfig) SummarizationEmitInternalEventsEffective() bool { + if c.SummarizationEmitInternalEvents != nil { + return *c.SummarizationEmitInternalEvents + } + return true +} + +func (c MultiAgentEinoMiddlewareConfig) SummarizationUserIntentLedgerMaxRunesEffective() int { + if c.SummarizationUserIntentLedgerMaxRunes > 0 { + return c.SummarizationUserIntentLedgerMaxRunes + } + return DefaultSummarizationUserIntentLedgerMaxRunes +} + +func (c MultiAgentEinoMiddlewareConfig) SummarizationUserIntentLedgerEntryMaxRunesEffective() int { + if c.SummarizationUserIntentLedgerEntryMaxRunes > 0 { + return c.SummarizationUserIntentLedgerEntryMaxRunes + } + return DefaultSummarizationUserIntentLedgerEntryMaxRunes +} + +func (c MultiAgentEinoMiddlewareConfig) LatestUserMessageMaxRunesEffective() int { + if c.LatestUserMessageMaxRunes > 0 { + return c.LatestUserMessageMaxRunes + } + return DefaultLatestUserMessageMaxRunes +} + +func (c MultiAgentEinoMiddlewareConfig) LatestUserMessageHeadRunesEffective() int { + if c.LatestUserMessageHeadRunes > 0 { + return c.LatestUserMessageHeadRunes + } + return DefaultLatestUserMessageHeadRunes +} + +func (c MultiAgentEinoMiddlewareConfig) LatestUserMessageTailRunesEffective() int { + if c.LatestUserMessageTailRunes > 0 { + return c.LatestUserMessageTailRunes + } + return DefaultLatestUserMessageTailRunes +} + +func (c MultiAgentEinoMiddlewareConfig) PlanExecuteUserInputBudgetRatioEffective() float64 { + v := c.PlanExecuteUserInputBudgetRatio + if v <= 0 { + return 0.35 + } + if v < 0.1 { + return 0.1 + } + if v > 0.6 { + return 0.6 + } + return v +} + +func (c MultiAgentEinoMiddlewareConfig) PlanExecuteExecutedStepsBudgetRatioEffective() float64 { + v := c.PlanExecuteExecutedStepsBudgetRatio + if v <= 0 { + return 0.2 + } + if v < 0.08 { + return 0.08 + } + if v > 0.5 { + return 0.5 + } + return v +} + +func (c MultiAgentEinoMiddlewareConfig) PlanExecuteMaxStepResultRunesEffective() int { + if c.PlanExecuteMaxStepResultRunes > 0 { + return c.PlanExecuteMaxStepResultRunes + } + return 4000 +} + +func (c MultiAgentEinoMiddlewareConfig) PlanExecuteKeepLastStepsEffective() int { + if c.PlanExecuteKeepLastSteps > 0 { + return c.PlanExecuteKeepLastSteps + } + return 8 +} + +func (c MultiAgentEinoMiddlewareConfig) ReductionMaxLengthForTruncEffective() int { + if c.ReductionMaxLengthForTrunc > 0 { + return c.ReductionMaxLengthForTrunc + } + return 12000 +} + +func (c MultiAgentEinoMiddlewareConfig) ReductionMaxTokensForClearEffective() int { + if c.ReductionMaxTokensForClear > 0 { + return c.ReductionMaxTokensForClear + } + return 50000 +} + +// MultiAgentEinoSkillsConfig toggles Eino official skill progressive disclosure and host filesystem tools. +type MultiAgentEinoSkillsConfig struct { + // Disable skips skill middleware (and does not attach local FS tools for Deep). + Disable bool `yaml:"disable" json:"disable"` + // FilesystemTools registers read_file/glob/grep/write/edit/execute (eino-ext local backend). Nil/omitted = true. + FilesystemTools *bool `yaml:"filesystem_tools,omitempty" json:"filesystem_tools,omitempty"` + // SkillToolName overrides the default Eino tool name "skill". + SkillToolName string `yaml:"skill_tool_name,omitempty" json:"skill_tool_name,omitempty"` +} + +// EinoSkillFilesystemToolsEffective returns whether Deep/sub-agents should attach local filesystem + streaming shell. +func (c MultiAgentEinoSkillsConfig) EinoSkillFilesystemToolsEffective() bool { + if c.FilesystemTools != nil { + return *c.FilesystemTools + } + return true +} + +// PatchToolCallsEffective returns whether patchtoolcalls middleware should run (default true). +func (c MultiAgentEinoMiddlewareConfig) PatchToolCallsEffective() bool { + if c.PatchToolCalls != nil { + return *c.PatchToolCalls + } + return true +} + +// MultiAgentSubConfig 子代理（Eino ChatModelAgent）：deep 下由 task 调度；supervisor 下由 transfer 委派；plan_execute 不使用子代理列表。 +type MultiAgentSubConfig struct { + ID string `yaml:"id" json:"id"` + Name string `yaml:"name" json:"name"` + Description string `yaml:"description" json:"description"` + Instruction string `yaml:"instruction" json:"instruction"` + BindRole string `yaml:"bind_role,omitempty" json:"bind_role,omitempty"` // 可选：关联主配置 roles 中的角色名；未配 role_tools 时沿用该角色的 tools + RoleTools []string `yaml:"role_tools" json:"role_tools"` // 与单 Agent 角色工具相同 key；空表示全部工具（bind_role 可补全 tools） + MaxIterations int `yaml:"max_iterations" json:"max_iterations"` + Kind string `yaml:"kind,omitempty" json:"kind,omitempty"` // 仅 Markdown：kind=orchestrator 表示 Deep 主代理（与 orchestrator.md 二选一约定） +} + +// MultiAgentPublic 返回给前端的精简信息（不含子代理指令全文）。 +type MultiAgentPublic struct { + Enabled bool `json:"enabled"` + RobotDefaultAgentMode string `json:"robot_default_agent_mode,omitempty"` + BatchUseMultiAgent bool `json:"batch_use_multi_agent"` + SubAgentCount int `json:"sub_agent_count"` + Orchestration string `json:"orchestration,omitempty"` + PlanExecuteLoopMaxIterations int `json:"plan_execute_loop_max_iterations"` + SummarizationUserIntentLedgerMaxRunes int `json:"summarization_user_intent_ledger_max_runes"` + SummarizationUserIntentLedgerEntryMaxRunes int `json:"summarization_user_intent_ledger_entry_max_runes"` + LatestUserMessageMaxRunes int `json:"latest_user_message_max_runes"` + LatestUserMessageHeadRunes int `json:"latest_user_message_head_runes"` + LatestUserMessageTailRunes int `json:"latest_user_message_tail_runes"` + ModelRetryMaxRetries int `json:"model_retry_max_retries"` + ModelRetryMaxBackoffSec int `json:"model_retry_max_backoff_sec"` + ModelFailoverChannels []string `json:"model_failover_channels,omitempty"` + ModelFailoverMaxRetries int `json:"model_failover_max_retries"` + ToolSearchAlwaysVisibleTools []string `json:"tool_search_always_visible_tools,omitempty"` + ToolSearchAlwaysVisibleEffectiveTools []string `json:"tool_search_always_visible_effective_tools,omitempty"` +} + +// NormalizeAgentMode 解析代理模式（eino_single | deep | plan_execute | supervisor）；空值默认 eino_single。 +func NormalizeAgentMode(mode string) string { + s := strings.TrimSpace(strings.ToLower(mode)) + switch s { + case "", "eino_single": + return "eino_single" + case "deep": + return "deep" + case "plan_execute", "plan-execute", "planexecute", "pe": + return "plan_execute" + case "supervisor", "super", "sv": + return "supervisor" + default: + return "eino_single" + } +} + +// NormalizeRobotAgentMode 解析机器人默认对话模式。 +func NormalizeRobotAgentMode(ma MultiAgentConfig) string { + return NormalizeAgentMode(ma.RobotDefaultAgentMode) +} + +// NormalizeMultiAgentOrchestration 返回 deep、plan_execute 或 supervisor。 +func NormalizeMultiAgentOrchestration(s string) string { + v := strings.TrimSpace(strings.ToLower(s)) + switch v { + case "plan_execute", "plan-execute", "planexecute", "pe": + return "plan_execute" + case "supervisor", "super", "sv": + return "supervisor" + default: + return "deep" + } +} + +// MultiAgentAPIUpdate 设置页/API 仅更新多代理标量字段；写入 YAML 时不覆盖 sub_agents 等块。 +type MultiAgentAPIUpdate struct { + Enabled bool `json:"enabled"` + RobotDefaultAgentMode string `json:"robot_default_agent_mode,omitempty"` + BatchUseMultiAgent bool `json:"batch_use_multi_agent"` + PlanExecuteLoopMaxIterations *int `json:"plan_execute_loop_max_iterations,omitempty"` + SummarizationUserIntentLedgerMaxRunes *int `json:"summarization_user_intent_ledger_max_runes,omitempty"` + SummarizationUserIntentLedgerEntryMaxRunes *int `json:"summarization_user_intent_ledger_entry_max_runes,omitempty"` + LatestUserMessageMaxRunes *int `json:"latest_user_message_max_runes,omitempty"` + LatestUserMessageHeadRunes *int `json:"latest_user_message_head_runes,omitempty"` + LatestUserMessageTailRunes *int `json:"latest_user_message_tail_runes,omitempty"` + ModelRetryMaxRetries *int `json:"model_retry_max_retries,omitempty"` + ModelRetryMaxBackoffSec *int `json:"model_retry_max_backoff_sec,omitempty"` + ModelFailoverChannels *[]string `json:"model_failover_channels,omitempty"` + ModelFailoverMaxRetries *int `json:"model_failover_max_retries,omitempty"` + // 指针区分「JSON 未传该字段」与「传空数组要清空」；省略时不应覆盖 YAML 中的常驻工具白名单。 + ToolSearchAlwaysVisibleTools *[]string `json:"tool_search_always_visible_tools,omitempty"` +} + +// RobotsConfig 机器人配置（企业微信、钉钉、飞书、微信 iLink、Telegram、Slack、Discord、QQ 等） +type RobotsConfig struct { + Session RobotSessionConfig `yaml:"session,omitempty" json:"session,omitempty"` // 机器人会话隔离策略 + Wechat RobotWechatConfig `yaml:"wechat,omitempty" json:"wechat,omitempty"` // 微信（iLink 扫码绑定） + Wecom RobotWecomConfig `yaml:"wecom,omitempty" json:"wecom,omitempty"` // 企业微信 + Dingtalk RobotDingtalkConfig `yaml:"dingtalk,omitempty" json:"dingtalk,omitempty"` // 钉钉 + Lark RobotLarkConfig `yaml:"lark,omitempty" json:"lark,omitempty"` // 飞书 + Telegram RobotTelegramConfig `yaml:"telegram,omitempty" json:"telegram,omitempty"` // Telegram + Slack RobotSlackConfig `yaml:"slack,omitempty" json:"slack,omitempty"` // Slack + Discord RobotDiscordConfig `yaml:"discord,omitempty" json:"discord,omitempty"` // Discord + QQ RobotQQConfig `yaml:"qq,omitempty" json:"qq,omitempty"` // QQ 机器人 +} + +// RobotWechatConfig 微信 iLink 机器人配置（个人微信 ClawBot / iLink 协议） +type RobotWechatConfig struct { + Enabled bool `yaml:"enabled" json:"enabled"` + BotToken string `yaml:"bot_token,omitempty" json:"bot_token,omitempty"` + ILinkBotID string `yaml:"ilink_bot_id,omitempty" json:"ilink_bot_id,omitempty"` + ILinkUserID string `yaml:"ilink_user_id,omitempty" json:"ilink_user_id,omitempty"` + BaseURL string `yaml:"base_url,omitempty" json:"base_url,omitempty"` // 默认 https://ilinkai.weixin.qq.com + BotType string `yaml:"bot_type,omitempty" json:"bot_type,omitempty"` // get_bot_qrcode 参数，默认 3 + BotAgent string `yaml:"bot_agent,omitempty" json:"bot_agent,omitempty"` // base_info.bot_agent + GetUpdatesBuf string `yaml:"get_updates_buf,omitempty" json:"get_updates_buf,omitempty"` // 长轮询游标（运行时） + Auth RobotAuthorizationConfig `yaml:"auth,omitempty" json:"auth,omitempty"` +} + +const ( + RobotAuthModeUserBinding = "user_binding" + RobotAuthModeServiceAccount = "service_account" +) + +// RobotAuthorizationConfig controls how a verified platform sender becomes +// an RBAC principal. service_account is intentionally fail-closed unless an +// explicit non-admin service user and sender allowlist are both configured. +type RobotAuthorizationConfig struct { + Mode string `yaml:"mode,omitempty" json:"mode,omitempty"` + ServiceUserID string `yaml:"service_user_id,omitempty" json:"service_user_id,omitempty"` + AllowedExternalUsers []string `yaml:"allowed_external_users,omitempty" json:"allowed_external_users,omitempty"` +} + +func (c RobotAuthorizationConfig) EffectiveMode() string { + mode := strings.ToLower(strings.TrimSpace(c.Mode)) + if mode == "" { + return RobotAuthModeUserBinding + } + return mode +} + +func (c RobotAuthorizationConfig) ExternalUserAllowed(externalUserID string) bool { + externalUserID = strings.TrimSpace(externalUserID) + if externalUserID == "" { + return false + } + for _, allowed := range c.AllowedExternalUsers { + if strings.TrimSpace(allowed) == externalUserID { + return true + } + } + return false +} + +// RobotSessionConfig 机器人会话隔离策略 +type RobotSessionConfig struct { + StrictUserIdentity *bool `yaml:"strict_user_identity,omitempty" json:"strict_user_identity,omitempty"` // true 时只允许真实用户标识，不允许会话/群 ID 兜底 +} + +// StrictUserIdentityEnabled 返回是否启用严格用户身份模式；未配置时默认 true。 +func (c RobotSessionConfig) StrictUserIdentityEnabled() bool { + if c.StrictUserIdentity == nil { + return true + } + return *c.StrictUserIdentity +} + +// RobotWecomConfig 企业微信机器人配置 +type RobotWecomConfig struct { + Enabled bool `yaml:"enabled" json:"enabled"` + Token string `yaml:"token" json:"token"` // 回调 URL 校验 Token + EncodingAESKey string `yaml:"encoding_aes_key" json:"encoding_aes_key"` // EncodingAESKey + CorpID string `yaml:"corp_id" json:"corp_id"` // 企业 ID + Secret string `yaml:"secret" json:"secret"` // 应用 Secret + AgentID int64 `yaml:"agent_id" json:"agent_id"` // 应用 AgentId + Auth RobotAuthorizationConfig `yaml:"auth,omitempty" json:"auth,omitempty"` +} + +// ValidateWecomConfig 校验企业微信机器人配置；启用时必须配置 token，否则回调无法防伪造。 +func ValidateWecomConfig(w RobotWecomConfig) error { + if !w.Enabled { + return nil + } + if strings.TrimSpace(w.Token) == "" { + return fmt.Errorf("robots.wecom.enabled 为 true 时必须配置 robots.wecom.token") + } + return nil +} + +// RobotDingtalkConfig 钉钉机器人配置 +type RobotDingtalkConfig struct { + Enabled bool `yaml:"enabled" json:"enabled"` + ClientID string `yaml:"client_id" json:"client_id"` // 应用 Key (AppKey) + ClientSecret string `yaml:"client_secret" json:"client_secret"` // 应用 Secret + AllowConversationIDFallback bool `yaml:"allow_conversation_id_fallback" json:"allow_conversation_id_fallback"` // sender_id 缺失时是否允许回退到会话 ID + Auth RobotAuthorizationConfig `yaml:"auth,omitempty" json:"auth,omitempty"` +} + +// RobotLarkConfig 飞书机器人配置 +type RobotLarkConfig struct { + Enabled bool `yaml:"enabled" json:"enabled"` + AppID string `yaml:"app_id" json:"app_id"` // 应用 App ID + AppSecret string `yaml:"app_secret" json:"app_secret"` // 应用 App Secret + VerifyToken string `yaml:"verify_token" json:"verify_token"` // 事件订阅 Verification Token（可选） + AllowChatIDFallback bool `yaml:"allow_chat_id_fallback" json:"allow_chat_id_fallback"` // 用户 ID 缺失时是否允许回退到 chat_id + Auth RobotAuthorizationConfig `yaml:"auth,omitempty" json:"auth,omitempty"` +} + +// RobotTelegramConfig Telegram 机器人配置（Bot API 长轮询） +type RobotTelegramConfig struct { + Enabled bool `yaml:"enabled" json:"enabled"` + BotToken string `yaml:"bot_token" json:"bot_token"` + BotUsername string `yaml:"bot_username,omitempty" json:"bot_username,omitempty"` // 可选，用于群聊 @ 识别；留空则启动时 getMe + AllowGroupMessages bool `yaml:"allow_group_messages" json:"allow_group_messages"` // 群聊中仅响应 @ 机器人 + UpdateOffset int64 `yaml:"update_offset,omitempty" json:"update_offset,omitempty"` + Auth RobotAuthorizationConfig `yaml:"auth,omitempty" json:"auth,omitempty"` +} + +// RobotSlackConfig Slack 机器人配置（Socket Mode，无需公网回调） +type RobotSlackConfig struct { + Enabled bool `yaml:"enabled" json:"enabled"` + BotToken string `yaml:"bot_token" json:"bot_token"` // xoxb- + AppToken string `yaml:"app_token" json:"app_token"` // xapp-（connections:write） + Auth RobotAuthorizationConfig `yaml:"auth,omitempty" json:"auth,omitempty"` +} + +// RobotDiscordConfig Discord 机器人配置（Gateway WebSocket） +type RobotDiscordConfig struct { + Enabled bool `yaml:"enabled" json:"enabled"` + BotToken string `yaml:"bot_token" json:"bot_token"` + AllowGuildMessages bool `yaml:"allow_guild_messages" json:"allow_guild_messages"` // 服务器频道中仅响应 @ 机器人 + Auth RobotAuthorizationConfig `yaml:"auth,omitempty" json:"auth,omitempty"` +} + +// RobotQQConfig QQ 机器人配置（QQ 开放平台 WebSocket） +type RobotQQConfig struct { + Enabled bool `yaml:"enabled" json:"enabled"` + AppID string `yaml:"app_id" json:"app_id"` + ClientSecret string `yaml:"client_secret" json:"client_secret"` + Sandbox bool `yaml:"sandbox" json:"sandbox"` // 沙箱环境（上线前测试） + Auth RobotAuthorizationConfig `yaml:"auth,omitempty" json:"auth,omitempty"` +} + +func (c RobotsConfig) AuthorizationFor(platform string) RobotAuthorizationConfig { + switch strings.ToLower(strings.TrimSpace(platform)) { + case "wechat": + return c.Wechat.Auth + case "wecom": + return c.Wecom.Auth + case "dingtalk": + return c.Dingtalk.Auth + case "lark": + return c.Lark.Auth + case "telegram": + return c.Telegram.Auth + case "slack": + return c.Slack.Auth + case "discord": + return c.Discord.Auth + case "qq": + return c.QQ.Auth + default: + return RobotAuthorizationConfig{} + } +} + +func ValidateRobotAuthorization(c RobotAuthorizationConfig, path string) error { + switch c.EffectiveMode() { + case RobotAuthModeUserBinding: + return nil + case RobotAuthModeServiceAccount: + serviceUserID := strings.TrimSpace(c.ServiceUserID) + if serviceUserID == "" { + return fmt.Errorf("%s.auth.service_user_id 不能为空", path) + } + if len(c.AllowedExternalUsers) == 0 { + return fmt.Errorf("%s.auth.allowed_external_users 至少配置一个真实发送者", path) + } + seen := map[string]bool{} + for _, userID := range c.AllowedExternalUsers { + userID = strings.TrimSpace(userID) + if userID == "" || userID == "*" { + return fmt.Errorf("%s.auth.allowed_external_users 不允许空值或通配符", path) + } + if seen[userID] { + return fmt.Errorf("%s.auth.allowed_external_users 包含重复用户", path) + } + seen[userID] = true + } + return nil + default: + return fmt.Errorf("%s.auth.mode 仅支持 user_binding 或 service_account", path) + } +} + +func ValidateRobotsAuthorization(c RobotsConfig) error { + items := []struct { + path string + auth RobotAuthorizationConfig + }{ + {"robots.wechat", c.Wechat.Auth}, {"robots.wecom", c.Wecom.Auth}, + {"robots.dingtalk", c.Dingtalk.Auth}, {"robots.lark", c.Lark.Auth}, + {"robots.telegram", c.Telegram.Auth}, {"robots.slack", c.Slack.Auth}, + {"robots.discord", c.Discord.Auth}, {"robots.qq", c.QQ.Auth}, + } + for _, item := range items { + if err := ValidateRobotAuthorization(item.auth, item.path); err != nil { + return err + } + } + return nil +} + +func (c RobotsConfig) ServiceAccountUserIDs() map[string]string { + out := map[string]string{} + for _, platform := range []string{"wechat", "wecom", "dingtalk", "lark", "telegram", "slack", "discord", "qq"} { + auth := c.AuthorizationFor(platform) + if auth.EffectiveMode() == RobotAuthModeServiceAccount { + out[platform] = strings.TrimSpace(auth.ServiceUserID) + } + } + return out +} + +type ServerConfig struct { + Host string `yaml:"host" json:"host"` + Port int `yaml:"port" json:"port"` + // CORSAllowedOrigins contains additional, exact origins that may call the API. + // Same-origin browser requests are always allowed. Wildcards are intentionally unsupported. + CORSAllowedOrigins []string `yaml:"cors_allowed_origins,omitempty" json:"cors_allowed_origins,omitempty"` + // TLSEnabled 为 true 时主 Web UI 使用 HTTPS；现代浏览器在同源下会协商 HTTP/2，缓解 HTTP/1.1 每源并发连接数限制。 + TLSEnabled bool `yaml:"tls_enabled,omitempty" json:"tls_enabled,omitempty"` + // TLSCertPath / TLSKeyPath 非空时从 PEM 文件加载证书（生产环境推荐）。 + TLSCertPath string `yaml:"tls_cert_path,omitempty" json:"tls_cert_path,omitempty"` + TLSKeyPath string `yaml:"tls_key_path,omitempty" json:"tls_key_path,omitempty"` + // TLSAutoSelfSign 为 true 且未配置有效证书路径时，启动时生成内存自签证书（仅本地/测试；浏览器会提示不受信任）。 + TLSAutoSelfSign bool `yaml:"tls_auto_self_sign,omitempty" json:"tls_auto_self_sign,omitempty"` + // TLSHTTPRedirect 为 false 时禁用 HTTP→HTTPS 跳转；省略或为 true 且已启用 HTTPS 时，明文 HTTP 访问将 308 跳转到 HTTPS（同端口嗅探分流）。 + TLSHTTPRedirect *bool `yaml:"tls_http_redirect,omitempty" json:"tls_http_redirect,omitempty"` +} + +type LogConfig struct { + Level string `yaml:"level"` + Output string `yaml:"output"` +} + +type MCPConfig struct { + Enabled bool `yaml:"enabled"` + Host string `yaml:"host"` + Port int `yaml:"port"` + AuthHeader string `yaml:"auth_header,omitempty"` // 可选的全局服务凭证 header；普通调用优先使用用户 Bearer Token + AuthHeaderValue string `yaml:"auth_header_value,omitempty"` // 全局服务凭证，仅 allow_global_access=true 时接受 + AllowGlobalAccess bool `yaml:"allow_global_access,omitempty"` // 静态服务密钥是否映射为全局服务身份（默认关闭） +} + +type OpenAIConfig struct { + Provider string `yaml:"provider,omitempty" json:"provider,omitempty"` // API 提供商: "openai"(默认) 或 "claude"，claude 使用 Eino 原生 Anthropic Messages API + APIKey string `yaml:"api_key" json:"api_key"` + BaseURL string `yaml:"base_url" json:"base_url"` + Model string `yaml:"model" json:"model"` + MaxTotalTokens int `yaml:"max_total_tokens,omitempty" json:"max_total_tokens,omitempty"` + MaxCompletionTokens int `yaml:"max_completion_tokens,omitempty" json:"max_completion_tokens,omitempty"` + // Reasoning 控制 Eino ChatModel 的 thinking / reasoning_effort / output_config 等（Eino 单/多代理路径生效）。 + Reasoning OpenAIReasoningConfig `yaml:"reasoning,omitempty" json:"reasoning,omitempty"` +} + +// AIConfig stores first-class model channels. Runtime callers resolve a channel +// into OpenAIConfig at the edge instead of moving API credentials through chat requests. +type AIConfig struct { + DefaultChannel string `yaml:"default_channel,omitempty" json:"default_channel,omitempty"` + Channels map[string]AIChannelConfig `yaml:"channels,omitempty" json:"channels,omitempty"` +} + +type AIChannelConfig struct { + Name string `yaml:"name,omitempty" json:"name,omitempty"` + Provider string `yaml:"provider,omitempty" json:"provider,omitempty"` + APIKey string `yaml:"api_key" json:"api_key"` + BaseURL string `yaml:"base_url" json:"base_url"` + Model string `yaml:"model" json:"model"` + MaxTotalTokens int `yaml:"max_total_tokens,omitempty" json:"max_total_tokens,omitempty"` + MaxCompletionTokens int `yaml:"max_completion_tokens,omitempty" json:"max_completion_tokens,omitempty"` + Reasoning OpenAIReasoningConfig `yaml:"reasoning,omitempty" json:"reasoning,omitempty"` +} + +func (c AIChannelConfig) ToOpenAIConfig() OpenAIConfig { + provider := strings.TrimSpace(c.Provider) + if provider == "" || provider == "openai_compatible" { + provider = "openai" + } + return OpenAIConfig{ + Provider: provider, + APIKey: c.APIKey, + BaseURL: c.BaseURL, + Model: c.Model, + MaxTotalTokens: c.MaxTotalTokens, + MaxCompletionTokens: c.MaxCompletionTokens, + Reasoning: c.Reasoning, + } +} + +func AIChannelFromOpenAI(id, name string, oa OpenAIConfig) AIChannelConfig { + if strings.TrimSpace(name) == "" { + name = id + } + return AIChannelConfig{ + Name: name, + Provider: oa.Provider, + APIKey: oa.APIKey, + BaseURL: oa.BaseURL, + Model: oa.Model, + MaxTotalTokens: oa.MaxTotalTokens, + MaxCompletionTokens: oa.MaxCompletionTokens, + Reasoning: oa.Reasoning, + } +} + +func NormalizeAIChannelID(s string) string { + id := strings.ToLower(strings.TrimSpace(s)) + id = strings.ReplaceAll(id, "_", "-") + var b strings.Builder + lastDash := false + for _, r := range id { + ok := (r >= 'a' && r <= 'z') || (r >= '0' && r <= '9') + if ok { + b.WriteRune(r) + lastDash = false + continue + } + if !lastDash { + b.WriteByte('-') + lastDash = true + } + } + out := strings.Trim(b.String(), "-") + if out == "" { + return "default" + } + return out +} + +func (c *AIConfig) EnsureDefaultFromOpenAI(openAI OpenAIConfig) { + if c.Channels == nil { + c.Channels = make(map[string]AIChannelConfig) + } + def := NormalizeAIChannelID(c.DefaultChannel) + if def == "default" && strings.TrimSpace(c.DefaultChannel) == "" { + def = "default" + } + c.DefaultChannel = def + if _, ok := c.Channels[def]; !ok { + c.Channels[def] = AIChannelFromOpenAI(def, "Default", openAI) + } +} + +func (c AIConfig) ResolveChannel(channelID string) (OpenAIConfig, string, bool) { + id := NormalizeAIChannelID(channelID) + if strings.TrimSpace(channelID) == "" { + id = NormalizeAIChannelID(c.DefaultChannel) + } + if id == "" { + id = "default" + } + if c.Channels != nil { + if ch, ok := c.Channels[id]; ok { + return ch.ToOpenAIConfig(), id, true + } + } + return OpenAIConfig{}, id, false +} + +func (c *Config) ResolveAIChannel(channelID string) (OpenAIConfig, string, bool) { + if c == nil { + return OpenAIConfig{}, "", false + } + if oa, id, ok := c.AI.ResolveChannel(channelID); ok { + return oa, id, true + } + return c.OpenAI, NormalizeAIChannelID(channelID), strings.TrimSpace(c.OpenAI.Model) != "" || strings.TrimSpace(c.OpenAI.BaseURL) != "" +} + +func (c *Config) ApplyDefaultAIChannel() { + if c == nil { + return + } + c.AI.EnsureDefaultFromOpenAI(c.OpenAI) + if oa, _, ok := c.AI.ResolveChannel(c.AI.DefaultChannel); ok { + c.OpenAI = oa + } +} + +func (c OpenAIConfig) MaxCompletionTokensEffective() int { + if c.MaxCompletionTokens > 0 { + return c.MaxCompletionTokens + } + return DefaultMaxCompletionTokens +} + +// IsDeepSeekEndpointOrModel reports whether the channel targets DeepSeek's +// official-compatible API or a DeepSeek model family. This is separate from the +// reasoning profile: profile controls field mapping, while DeepSeek has provider +// constraints such as default thinking mode and no tool_choice in thinking mode. +func (c OpenAIConfig) IsDeepSeekEndpointOrModel() bool { + baseURL := strings.ToLower(strings.TrimSpace(c.BaseURL)) + model := strings.ToLower(strings.TrimSpace(c.Model)) + return strings.Contains(baseURL, "deepseek") || strings.Contains(model, "deepseek") +} + +// OpenAIReasoningConfig 全局默认与网关 profile（对话页可通过 ChatRequest.reasoning 覆盖，受 AllowClientReasoning 约束）。 +type OpenAIReasoningConfig struct { + // Mode: auto（默认）| on | off | default（与 auto 相同）。 + // off 在 OpenAI/Claude profile 下省略推理字段；DeepSeek profile 下发送 thinking.type=disabled（其默认开启思考）。 + Mode string `yaml:"mode,omitempty" json:"mode,omitempty"` + // Effort: low | medium | high | max | xhigh；max/xhigh 为不同网关最高档命名，原样下发、不互转。空表示不单独指定强度。 + Effort string `yaml:"effort,omitempty" json:"effort,omitempty"` + // AllowClientReasoning 为 false 时忽略请求体 reasoning；nil 或未设置等同于 true。 + AllowClientReasoning *bool `yaml:"allow_client_reasoning,omitempty" json:"allow_client_reasoning,omitempty"` + // Profile: auto | deepseek_compat | openai_compat | output_config_effort + Profile string `yaml:"profile,omitempty" json:"profile,omitempty"` + // ExtraRequestFields 合并进 Chat Completions 根 JSON（管理员用；与自动字段同名时后者覆盖）。 + // Mode=off 时会移除其中的推理控制字段，但保留其他扩展字段；DeepSeek profile 随后补充显式关闭开关。 + ExtraRequestFields map[string]interface{} `yaml:"extra_request_fields,omitempty" json:"extra_request_fields,omitempty"` +} + +// ModeEffective returns auto when empty or default. +func (c OpenAIReasoningConfig) ModeEffective() string { + m := strings.ToLower(strings.TrimSpace(c.Mode)) + if m == "" || m == "default" { + return "auto" + } + return m +} + +// ProfileEffective returns auto when empty. +func (c OpenAIReasoningConfig) ProfileEffective() string { + p := strings.ToLower(strings.TrimSpace(c.Profile)) + if p == "" { + return "auto" + } + return p +} + +// AllowClientReasoningEffective true when client may send ChatRequest.reasoning. +func (c OpenAIReasoningConfig) AllowClientReasoningEffective() bool { + if c.AllowClientReasoning == nil { + return true + } + return *c.AllowClientReasoning +} + +type FofaConfig struct { + // APIKey 为 FOFA API Key（建议使用只读权限的 Key） + APIKey string `yaml:"api_key,omitempty" json:"api_key,omitempty"` + BaseURL string `yaml:"base_url,omitempty" json:"base_url,omitempty"` // 默认 https://fofa.info/api/v1/search/all +} + +type SpaceSearchConfig struct { + APIKey string `yaml:"api_key,omitempty" json:"api_key,omitempty"` + BaseURL string `yaml:"base_url,omitempty" json:"base_url,omitempty"` +} + +type SecurityConfig struct { + Tools []ToolConfig `yaml:"tools,omitempty"` // 向后兼容：支持在主配置文件中定义工具 + ToolsDir string `yaml:"tools_dir,omitempty"` // 工具配置文件目录（新方式） + ToolDescriptionMode string `yaml:"tool_description_mode,omitempty"` // 工具描述模式: "short" | "full"，默认 short +} + +type DatabaseConfig struct { + Path string `yaml:"path"` // 会话数据库路径 + KnowledgeDBPath string `yaml:"knowledge_db_path,omitempty"` // 知识库数据库路径（可选，为空则使用会话数据库） +} + +type AgentConfig struct { + MaxIterations int `yaml:"max_iterations" json:"max_iterations"` + ToolTimeoutMinutes int `yaml:"tool_timeout_minutes" json:"tool_timeout_minutes"` // 单次工具执行最大时长（分钟），超时自动终止，防止长时间挂起；0 表示不限制（不推荐） + ToolWaitTimeoutSeconds int `yaml:"tool_wait_timeout_seconds" json:"tool_wait_timeout_seconds"` // 工具本轮等待秒数；到时返回 execution_id，worker 继续后台执行；0 表示等到完成 + ExternalMCPMaxConcurrentPerServer int `yaml:"external_mcp_max_concurrent_per_server" json:"external_mcp_max_concurrent_per_server"` // 单个外部 MCP server 同时运行的工具数；0 表示默认 2 + ExternalMCPMaxConcurrentTotal int `yaml:"external_mcp_max_concurrent_total" json:"external_mcp_max_concurrent_total"` // 所有外部 MCP 工具全局并发；0 表示默认 16 + ExternalMCPCircuitFailureThreshold int `yaml:"external_mcp_circuit_failure_threshold" json:"external_mcp_circuit_failure_threshold"` // 单个 MCP server 连续失败多少次后打开熔断；0 表示默认 3；负数关闭 + ExternalMCPCircuitCooldownSeconds int `yaml:"external_mcp_circuit_cooldown_seconds" json:"external_mcp_circuit_cooldown_seconds"` // 熔断后冷却秒数；0 表示默认 60 + // ShellNoOutputTimeoutSeconds execute/exec 无任何 stdout/stderr 时的空闲终止秒数（通用防挂死，不维护命令黑名单）；0=默认 300（5 分钟）；-1=关闭。 + ShellNoOutputTimeoutSeconds int `yaml:"shell_no_output_timeout_seconds" json:"shell_no_output_timeout_seconds"` + // WorkspaceRootDir 会话工作目录根路径（curl/wget 下载、read_file/glob/grep 本地分析）；空=tmp/workspace，其下按 projects/{id} 或 conversations/{id} 隔离。 + WorkspaceRootDir string `yaml:"workspace_root_dir,omitempty" json:"workspace_root_dir,omitempty"` + // SystemPromptPath 单代理系统提示 Markdown/文本文件路径（相对 config.yaml 所在目录，或可写绝对路径）。非空且可读时替换内置单代理提示；留空用内置。 + SystemPromptPath string `yaml:"system_prompt_path,omitempty" json:"system_prompt_path,omitempty"` +} + +// HitlConfig 人机协同全局选项；与会话侧栏/API 中的白名单合并为并集后参与判定。 +// tool_whitelist 可在侧栏「应用」时合并写入 config.yaml 并立即生效。 +// audit_agent_prompt / audit_agent_prompt_review_edit 可在人机协同页编辑并立即生效；空则使用内置默认。 +type HitlConfig struct { + // AuditModel 审计 Agent 专用模型；字段留空时继承 OpenAI 主配置，便于用小模型做审批。 + AuditModel OpenAIConfig `yaml:"audit_model,omitempty" json:"audit_model,omitempty"` + // ToolWhitelist 全局免审批工具名（与白名单内工具不触发 HITL 审批）。 + ToolWhitelist []string `yaml:"tool_whitelist,omitempty" json:"tool_whitelist,omitempty"` + // AuditAgentPrompt 审批模式（approval）下审计 Agent 系统提示词。 + AuditAgentPrompt string `yaml:"audit_agent_prompt,omitempty" json:"audit_agent_prompt,omitempty"` + // AuditAgentPromptReviewEdit 审查编辑模式（review_edit）下审计 Agent 系统提示词。 + AuditAgentPromptReviewEdit string `yaml:"audit_agent_prompt_review_edit,omitempty" json:"audit_agent_prompt_review_edit,omitempty"` + // RetentionDays 已决策审计日志（hitl_interrupts 非 pending）保留天数；省略时默认 90；0 表示不自动清理。 + RetentionDays *int `yaml:"retention_days,omitempty" json:"retention_days,omitempty"` + // DefaultReviewer 全局默认审批方（human | audit_agent）；未选会话时切换会写入 config.yaml；新建会话无独立配置时沿用。 + DefaultReviewer string `yaml:"default_reviewer,omitempty" json:"default_reviewer,omitempty"` +} + +// EffectiveDefaultReviewer returns human or audit_agent; omitted or unknown values default to human. +func (h HitlConfig) EffectiveDefaultReviewer() string { + switch strings.ToLower(strings.TrimSpace(h.DefaultReviewer)) { + case "audit_agent", "agent", "ai": + return "audit_agent" + default: + return "human" + } +} + +// RetentionDaysEffective returns retention; 0 means keep forever; omitted defaults to 90. +func (h HitlConfig) RetentionDaysEffective() int { + if h.RetentionDays == nil { + return 90 + } + if *h.RetentionDays < 0 { + return 0 + } + return *h.RetentionDays +} + +// AuditModelEffective returns the audit-agent model config with empty fields inherited from the main model config. +func (h HitlConfig) AuditModelEffective(main OpenAIConfig) OpenAIConfig { + out := main + am := h.AuditModel + if strings.TrimSpace(am.Provider) != "" { + out.Provider = strings.TrimSpace(am.Provider) + } + if strings.TrimSpace(am.BaseURL) != "" { + out.BaseURL = strings.TrimSpace(am.BaseURL) + } + if strings.TrimSpace(am.APIKey) != "" { + out.APIKey = strings.TrimSpace(am.APIKey) + } + if strings.TrimSpace(am.Model) != "" { + out.Model = strings.TrimSpace(am.Model) + } + if am.MaxTotalTokens > 0 { + out.MaxTotalTokens = am.MaxTotalTokens + } + if am.MaxCompletionTokens > 0 { + out.MaxCompletionTokens = am.MaxCompletionTokens + } + return out +} + +const hitlAuditAgentPromptBase = `你是 CyberStrikeAI 人机协同审计 Agent。审查 Agent 即将执行的工具调用是否会对系统造成实质性损害。 + +你会收到 JSON，包含 hitlMode、toolName、arguments/argumentsObj、userMessage、thinking、reasoningChain、planning 等字段。 + +规则优先级： +- 必须 reject 的规则 +- 必须 approve 的规则 +- 默认放行规则 + +规则匹配要求： +- 如果同时命中 reject 和 approve，必须 reject。 +- 每次裁决必须给出一个规则编号，不得输出“命中规则：无”。 +- 未命中明确 reject 或 approve 规则，但未发现明确破坏行为时，命中 D1：默认放行。 +- 不得仅凭“本地执行”“设置超时”“文件名包含 test”判断脚本无害；必须检查实际 command、代码或请求参数。 +- 如果脚本具体内容不可见，应表述为“当前参数未显示明确高危操作”，不能断言“无持久化影响”。 + +裁决基调（默认放行）： +- 绝大多数渗透测试操作（包括各类攻击 payload）都应放行 +- 在裁决之前需要判断出：1、此请求的是干什么（理解请求内容）？2、此请求造成什么后果？ +- 仅在「明确会直接破坏目标业务系统正常运行，或造成真实用户数据丢失/不可逆篡改」时 → reject +- 信息不足、参数略宽泛、仅为探测意图时 → approve +- 拿不准时 → approve（宁可放过可恢复的测试，不要误拦合法攻击） + +必须 reject 的高危情形（示例，非穷举）： +- 修改/重置任意用户或管理员密码、强制下线真实用户 +- 修改服务器或服务配置（系统配置、Web/数据库参数、防火墙、启动项、计划任务） +- 修改/创建/删除用户、角色、权限（账户体系篡改） +- 删除或大规模破坏数据/文件（DROP/TRUNCATE/DELETE 全表、rm 重要文件、清空数据库） +- 停止、禁用、重启业务服务（systemctl stop、kill 关键进程、shutdown、reboot） +- 大规模压测或高并发造成 DoS + +以下情形必须 approve（即使看起来“攻击性强”或包含写操作）： +- 各类漏洞探测与利用 payload（SQL 注入含读写型、命令注入、反序列化、文件包含、路径遍历、SSRF、SSTI、XXE 等） +- 文件上传测试，包括上传 webshell、写入 webshell 文件 +- 命令执行探测、反弹 Shell、dnslog / 数据外带探测 +- 数据库命令执行、写文件、写计划任务、内存马注入等 getshell 利用 +- 常规速率的口令爆破（非 DoS 级） +- 信息收集、端口/目录扫描、枚举、抓取 +- 读取、查询、探测类操作（含敏感文件） +- 攻陷后无害侦察命令（whoami/id/uname/ls/cat 等） +- 参数宽泛但无明确破坏意图的请求 + +判定原则： +- “写”本身不是拦截理由——渗透测试大量依赖写操作（注入、上传、爆破、getshell） +- 关键看是否破坏「业务可用性 / 关键数据完整性」 +- 只有明确命中上方 reject 清单才 reject` + +const hitlAuditAgentPromptApprovalOutput = ` +仅输出一行 JSON，不要 markdown 代码块： +{"decision":"approve"|"reject","comment":"实际操作：...；成功后的后果：...；命中规则：..."}` + +const hitlAuditAgentPromptReviewEditOutput = ` +仅输出一行 JSON，不要 markdown 代码块： +{"decision":"approve"|"reject","comment":"实际操作：...；成功后的后果：...；命中规则：...","editedArguments":{...}} + +editedArguments 规则（仅 approve 且需要改参时填写，否则省略该字段）： +- 提供完整替换后的工具参数对象，键名与 argumentsObj 一致 +- 只做最小必要修改以收窄范围、消除风险（如限制 path、去掉危险 flag） +- 禁止扩大攻击面：不得扩大目标范围、提升权限或引入破坏性参数 +- 无法安全改参时应 reject，不要勉强 approve` + +// DefaultHitlAuditAgentPrompt 内置审批模式审计 Agent 提示词。 +func DefaultHitlAuditAgentPrompt() string { + return hitlAuditAgentPromptBase + hitlAuditAgentPromptApprovalOutput +} + +// DefaultHitlAuditAgentPromptReviewEdit 内置审查编辑模式审计 Agent 提示词。 +func DefaultHitlAuditAgentPromptReviewEdit() string { + return hitlAuditAgentPromptBase + hitlAuditAgentPromptReviewEditOutput +} + +// EffectiveAuditAgentPrompt 返回审批模式生效的审计 Agent 提示词。 +func (c HitlConfig) EffectiveAuditAgentPrompt() string { + return c.EffectiveAuditAgentPromptForMode("approval") +} + +// EffectiveAuditAgentPromptForMode 按 HITL 模式返回生效的审计 Agent 提示词。 +func (c HitlConfig) EffectiveAuditAgentPromptForMode(mode string) string { + if normalizeHitlModeForPrompt(mode) == "review_edit" { + if s := strings.TrimSpace(c.AuditAgentPromptReviewEdit); s != "" { + return s + } + return DefaultHitlAuditAgentPromptReviewEdit() + } + if s := strings.TrimSpace(c.AuditAgentPrompt); s != "" { + return s + } + return DefaultHitlAuditAgentPrompt() +} + +func normalizeHitlModeForPrompt(mode string) string { + switch strings.ToLower(strings.TrimSpace(mode)) { + case "review_edit": + return "review_edit" + default: + return "approval" + } +} + +type AuthConfig struct { + SessionDurationHours int `yaml:"session_duration_hours" json:"session_duration_hours"` +} + +// MonitorConfig MCP 状态监控（tool_executions）保留策略。 +type MonitorConfig struct { + // RetentionDays 执行记录保留天数；省略时默认 90；0 表示不自动清理。 + RetentionDays *int `yaml:"retention_days,omitempty" json:"retention_days,omitempty"` +} + +// RetentionDaysEffective returns retention; 0 means keep forever; omitted defaults to 90. +func (m MonitorConfig) RetentionDaysEffective() int { + if m.RetentionDays == nil { + return 90 + } + if *m.RetentionDays < 0 { + return 0 + } + return *m.RetentionDays +} + +// AuditConfig platform operation audit log settings (not chat/tool execution bodies). +type AuditConfig struct { + // Enabled nil or true enables persistence; explicit false disables. + Enabled *bool `yaml:"enabled,omitempty" json:"enabled,omitempty"` + RetentionDays int `yaml:"retention_days,omitempty" json:"retention_days,omitempty"` + MaxDetailBytes int `yaml:"max_detail_bytes,omitempty" json:"max_detail_bytes,omitempty"` + // AuthFailureCooldownSeconds: per-IP cooldown for auth login/change_password failure audit rows; -1 disables; 0 uses default 60. + AuthFailureCooldownSeconds int `yaml:"auth_failure_cooldown_seconds,omitempty" json:"auth_failure_cooldown_seconds,omitempty"` +} + +// EnabledEffective returns true unless audit.enabled is explicitly false. +func (a AuditConfig) EnabledEffective() bool { + if a.Enabled == nil { + return true + } + return *a.Enabled +} + +// RetentionDaysEffective returns retention; 0 means keep forever. +func (a AuditConfig) RetentionDaysEffective() int { + if a.RetentionDays < 0 { + return 0 + } + return a.RetentionDays +} + +// MaxDetailBytesEffective caps serialized detail JSON size. +func (a AuditConfig) MaxDetailBytesEffective() int { + if a.MaxDetailBytes <= 0 { + return 8192 + } + return a.MaxDetailBytes +} + +// AuthFailureCooldownEffective returns seconds between duplicate auth-failure audit rows per IP (default 60; -1 disables). +func (a AuditConfig) AuthFailureCooldownEffective() int { + if a.AuthFailureCooldownSeconds < 0 { + return 0 + } + if a.AuthFailureCooldownSeconds == 0 { + return 60 + } + return a.AuthFailureCooldownSeconds +} + +// ExternalMCPConfig 外部MCP配置 +type ExternalMCPConfig struct { + Servers map[string]ExternalMCPServerConfig `yaml:"servers,omitempty" json:"servers,omitempty"` +} + +// ExternalMCPServerConfig 外部MCP服务器配置（遵循官方 MCP 配置格式，兼容 Claude Desktop / Cursor / VS Code）。 +// 所有字符串字段均支持 ${VAR} 和 ${VAR:-default} 环境变量展开语法。 +type ExternalMCPServerConfig struct { + // 传输类型: "stdio" | "sse" | "http"（Streamable HTTP）。 + // stdio 模式可省略，有 command 字段时自动推断。 + Type string `yaml:"type,omitempty" json:"type,omitempty"` + + // stdio 模式配置 + Command string `yaml:"command,omitempty" json:"command,omitempty"` + Args []string `yaml:"args,omitempty" json:"args,omitempty"` + Env map[string]string `yaml:"env,omitempty" json:"env,omitempty"` + + // HTTP/SSE 模式配置 + URL string `yaml:"url,omitempty" json:"url,omitempty"` + Headers map[string]string `yaml:"headers,omitempty" json:"headers,omitempty"` + + // 官方标准字段 + Disabled bool `yaml:"disabled,omitempty" json:"disabled,omitempty"` // 禁用服务器（官方字段） + AutoApprove []string `yaml:"autoApprove,omitempty" json:"autoApprove,omitempty"` // 自动批准的工具列表（官方字段） + + // SDK 高级配置（对应 MCP Go SDK 传输层参数） + MaxRetries int `yaml:"max_retries,omitempty" json:"max_retries,omitempty"` // Streamable HTTP 断线重连次数（默认 5） + TerminateDuration int `yaml:"terminate_duration,omitempty" json:"terminate_duration,omitempty"` // stdio 进程优雅关闭等待秒数（默认 5） + KeepAlive int `yaml:"keep_alive,omitempty" json:"keep_alive,omitempty"` // 客户端心跳间隔秒数（0 = 禁用） + + // 通用配置 + Description string `yaml:"description,omitempty" json:"description,omitempty"` + Timeout int `yaml:"timeout,omitempty" json:"timeout,omitempty"` // 连接超时（秒） + ExternalMCPEnable bool `yaml:"external_mcp_enable,omitempty" json:"external_mcp_enable,omitempty"` // 是否启用 + ToolEnabled map[string]bool `yaml:"tool_enabled,omitempty" json:"tool_enabled,omitempty"` // 每个工具的启用状态 +} + +// GetTransportType 返回实际传输类型。优先读 Type，否则根据 Command/URL 自动推断。 +func (c ExternalMCPServerConfig) GetTransportType() string { + if c.Type != "" { + return c.Type + } + if c.Command != "" { + return "stdio" + } + if c.URL != "" { + return "http" + } + return "" +} + +type ToolConfig struct { + Name string `yaml:"name"` + Command string `yaml:"command"` + Args []string `yaml:"args,omitempty"` // 固定参数（可选） + ShortDescription string `yaml:"short_description,omitempty"` // 简短描述（用于工具列表，减少token消耗） + Description string `yaml:"description"` // 详细描述（用于工具文档） + Enabled bool `yaml:"enabled"` + Parameters []ParameterConfig `yaml:"parameters,omitempty"` // 参数定义（可选） + ArgMapping string `yaml:"arg_mapping,omitempty"` // 参数映射方式: "auto", "manual", "template"（可选） + AllowedExitCodes []int `yaml:"allowed_exit_codes,omitempty"` // 允许的退出码列表（某些工具在成功时也返回非零退出码） +} + +// ParameterConfig 参数配置 +type ParameterConfig struct { + Name string `yaml:"name"` // 参数名称 + Type string `yaml:"type"` // 参数类型: string, int, bool, array + Description string `yaml:"description"` // 参数描述 + Required bool `yaml:"required,omitempty"` // 是否必需 + Default interface{} `yaml:"default,omitempty"` // 默认值 + ItemType string `yaml:"item_type,omitempty"` // 当 type 为 array 时，数组元素类型，如 string, number, object + Flag string `yaml:"flag,omitempty"` // 命令行标志，如 "-u", "--url", "-p" + Position *int `yaml:"position,omitempty"` // 位置参数的位置（从0开始） + Format string `yaml:"format,omitempty"` // 参数格式: "flag", "positional", "combined" (flag=value), "template" + Template string `yaml:"template,omitempty"` // 模板字符串，如 "{flag} {value}" 或 "{value}" + Options []string `yaml:"options,omitempty"` // 可选值列表（用于枚举） +} + +func Load(path string) (*Config, error) { + data, err := os.ReadFile(path) + if err != nil { + return nil, fmt.Errorf("读取配置文件失败: %w", err) + } + + var cfg Config + if err := yaml.Unmarshal(data, &cfg); err != nil { + return nil, fmt.Errorf("解析配置文件失败: %w", err) + } + + if cfg.Auth.SessionDurationHours <= 0 { + cfg.Auth.SessionDurationHours = 12 + } + if cfg.Audit.MaxDetailBytes <= 0 { + cfg.Audit.MaxDetailBytes = 8192 + } + cfg.ApplyDefaultAIChannel() + if err := validateOpenAIOutputLimits(cfg.OpenAI); err != nil { + return nil, err + } + // 如果配置了工具目录，从目录加载工具配置 + if cfg.Security.ToolsDir != "" { + inlineTools := append([]ToolConfig(nil), cfg.Security.Tools...) + toolsDir := ResolveToolsDir(cfg.Security.ToolsDir, path) + merged, err := MergeToolsFromDir(toolsDir, inlineTools) + if err != nil { + return nil, fmt.Errorf("从工具目录加载工具配置失败: %w", err) + } + cfg.Security.Tools = merged + } + + // 外部 MCP：迁移 + 环境变量展开 + if cfg.ExternalMCP.Servers != nil { + for name, serverCfg := range cfg.ExternalMCP.Servers { + // 官方 disabled 字段 → ExternalMCPEnable + if serverCfg.Disabled { + serverCfg.ExternalMCPEnable = false + } else if !serverCfg.ExternalMCPEnable { + // 默认启用 + serverCfg.ExternalMCPEnable = true + } + + // 展开所有 ${VAR} / ${VAR:-default} 环境变量引用 + ExpandConfigEnv(&serverCfg) + + cfg.ExternalMCP.Servers[name] = serverCfg + } + } + + // 从角色目录加载角色配置 + if cfg.RolesDir != "" { + configDir := filepath.Dir(path) + rolesDir := cfg.RolesDir + + // 如果是相对路径，相对于配置文件所在目录 + if !filepath.IsAbs(rolesDir) { + rolesDir = filepath.Join(configDir, rolesDir) + } + + roles, err := LoadRolesFromDir(rolesDir) + if err != nil { + return nil, fmt.Errorf("从角色目录加载角色配置失败: %w", err) + } + + cfg.Roles = roles + } else { + // 如果未配置 roles_dir，初始化为空 map + if cfg.Roles == nil { + cfg.Roles = make(map[string]RoleConfig) + } + } + + if err := ValidateWecomConfig(cfg.Robots.Wecom); err != nil { + return nil, err + } + if err := ValidateRobotsAuthorization(cfg.Robots); err != nil { + return nil, err + } + + return &cfg, nil +} + +func validateOpenAIOutputLimits(openAI OpenAIConfig) error { + if openAI.MaxCompletionTokens < 0 { + return fmt.Errorf("openai.max_completion_tokens 必须为正数") + } + return nil +} + +func EnsureLocalConfig(path string) (EnsureLocalConfigResult, error) { + path = strings.TrimSpace(path) + if path == "" { + path = "config.yaml" + } + + if _, err := os.Stat(path); err == nil { + return EnsureLocalConfigResult{}, nil + } else if !os.IsNotExist(err) { + return EnsureLocalConfigResult{}, fmt.Errorf("检查配置文件失败: %w", err) + } + + examplePath := filepath.Join(filepath.Dir(path), "config.example.yaml") + if _, err := os.Stat(examplePath); err != nil { + if os.IsNotExist(err) { + if alt := "config.example.yaml"; examplePath != alt { + if _, altErr := os.Stat(alt); altErr == nil { + examplePath = alt + } else { + return EnsureLocalConfigResult{}, fmt.Errorf("配置文件 %s 不存在，且未找到模板 %s", path, examplePath) + } + } else { + return EnsureLocalConfigResult{}, fmt.Errorf("配置文件 %s 不存在，且未找到模板 %s", path, examplePath) + } + } else { + return EnsureLocalConfigResult{}, fmt.Errorf("检查配置模板失败: %w", err) + } + } + + data, err := os.ReadFile(examplePath) + if err != nil { + return EnsureLocalConfigResult{}, fmt.Errorf("读取配置模板失败: %w", err) + } + + if dir := filepath.Dir(path); dir != "." && dir != "" { + if err := os.MkdirAll(dir, 0700); err != nil { + return EnsureLocalConfigResult{}, fmt.Errorf("创建配置目录失败: %w", err) + } + } + if err := os.WriteFile(path, data, fs.FileMode(0600)); err != nil { + return EnsureLocalConfigResult{}, fmt.Errorf("创建配置文件失败: %w", err) + } + + return EnsureLocalConfigResult{ + Created: true, + ExamplePath: examplePath, + }, nil +} + +func PrintBootstrapAdminPassword(password string) { + termout.PrintBootstrapAdminCredentials(password) +} + +// generateRandomToken 生成用于 MCP 鉴权的随机字符串（64 位十六进制） +func generateRandomToken() (string, error) { + b := make([]byte, 32) + if _, err := rand.Read(b); err != nil { + return "", err + } + return hex.EncodeToString(b), nil +} + +// persistMCPAuth 将 MCP 的 auth_header / auth_header_value 写回配置文件 +func persistMCPAuth(path string, mcp *MCPConfig) error { + data, err := os.ReadFile(path) + if err != nil { + return err + } + lines := strings.Split(string(data), "\n") + inMcpBlock := false + mcpIndent := -1 + + for i, line := range lines { + trimmed := strings.TrimSpace(line) + if !inMcpBlock { + if strings.HasPrefix(trimmed, "mcp:") { + inMcpBlock = true + mcpIndent = len(line) - len(strings.TrimLeft(line, " ")) + } + continue + } + if trimmed == "" || strings.HasPrefix(trimmed, "#") { + continue + } + leadingSpaces := len(line) - len(strings.TrimLeft(line, " ")) + if leadingSpaces <= mcpIndent { + inMcpBlock = false + mcpIndent = -1 + if strings.HasPrefix(trimmed, "mcp:") { + inMcpBlock = true + mcpIndent = leadingSpaces + } + continue + } + + prefix := line[:leadingSpaces] + rest := strings.TrimSpace(line[leadingSpaces:]) + comment := "" + if idx := strings.Index(line, "#"); idx >= 0 { + comment = strings.TrimRight(line[idx:], " ") + } + withComment := "" + if comment != "" { + if !strings.HasPrefix(comment, " ") { + withComment = " " + } + withComment += comment + } + + if strings.HasPrefix(rest, "auth_header_value:") { + lines[i] = fmt.Sprintf("%sauth_header_value: %q%s", prefix, mcp.AuthHeaderValue, withComment) + } else if strings.HasPrefix(rest, "auth_header:") { + lines[i] = fmt.Sprintf("%sauth_header: %q%s", prefix, mcp.AuthHeader, withComment) + } + } + + return os.WriteFile(path, []byte(strings.Join(lines, "\n")), 0644) +} + +// EnsureMCPAuth only provisions the privileged static service credential when +// global service access was explicitly enabled. +func EnsureMCPAuth(path string, cfg *Config) error { + if !cfg.MCP.Enabled || !cfg.MCP.AllowGlobalAccess || strings.TrimSpace(cfg.MCP.AuthHeaderValue) != "" { + return nil + } + token, err := generateRandomToken() + if err != nil { + return fmt.Errorf("生成 MCP 鉴权密钥失败: %w", err) + } + cfg.MCP.AuthHeaderValue = token + if strings.TrimSpace(cfg.MCP.AuthHeader) == "" { + cfg.MCP.AuthHeader = "X-MCP-Token" + } + return persistMCPAuth(path, &cfg.MCP) +} + +// PrintMCPConfigJSON 向终端输出 MCP 配置的 JSON，可直接复制到 Cursor / Claude Code 的 mcp 配置中使用 +func PrintMCPConfigJSON(mcp MCPConfig) { + if !mcp.Enabled { + return + } + hostForURL := strings.TrimSpace(mcp.Host) + if hostForURL == "" || hostForURL == "0.0.0.0" { + hostForURL = "localhost" + } + url := fmt.Sprintf("http://%s:%d/mcp", hostForURL, mcp.Port) + headers := map[string]string{"Authorization": "Bearer "} + if mcp.AllowGlobalAccess && mcp.AuthHeader != "" { + delete(headers, "Authorization") + headers[mcp.AuthHeader] = mcp.AuthHeaderValue + } + serverEntry := map[string]interface{}{ + "url": url, + } + if len(headers) > 0 { + serverEntry["headers"] = headers + } + // Claude Code 需要 type: "http" + serverEntry["type"] = "http" + out := map[string]interface{}{ + "mcpServers": map[string]interface{}{ + "cyberstrike-ai": serverEntry, + }, + } + b, _ := json.MarshalIndent(out, "", " ") + fmt.Println("[CyberStrikeAI] MCP 配置（可复制到 Cursor / Claude Code 使用）：") + fmt.Println(" Cursor: 放入 ~/.cursor/mcp.json 的 mcpServers，或项目 .cursor/mcp.json") + fmt.Println(" Claude Code: 放入 .mcp.json 或 ~/.claude.json 的 mcpServers") + fmt.Println("----------------------------------------------------------------") + fmt.Println(string(b)) + fmt.Println("----------------------------------------------------------------") +} + +// ResolveToolsDir 将 tools_dir 解析为绝对路径（相对路径相对于 configPath 所在目录）。 +func ResolveToolsDir(toolsDir, configPath string) string { + toolsDir = strings.TrimSpace(toolsDir) + if toolsDir == "" { + return "" + } + if filepath.IsAbs(toolsDir) { + return toolsDir + } + return filepath.Join(filepath.Dir(configPath), toolsDir) +} + +// MergeToolsFromDir 从目录加载工具并与 inline 列表合并：目录中的工具优先，主配置中的工具作为补充。 +func MergeToolsFromDir(toolsDir string, inlineTools []ToolConfig) ([]ToolConfig, error) { + dirTools, err := LoadToolsFromDir(toolsDir) + if err != nil { + return nil, err + } + existing := make(map[string]bool, len(dirTools)) + for _, tool := range dirTools { + existing[tool.Name] = true + } + merged := append([]ToolConfig(nil), dirTools...) + for _, tool := range inlineTools { + if !existing[tool.Name] { + merged = append(merged, tool) + } + } + return merged, nil +} + +// loadInlineSecurityToolsFromYAML 读取 config.yaml 中 security.tools（不含 tools_dir 扫描结果）。 +func loadInlineSecurityToolsFromYAML(configPath string) ([]ToolConfig, error) { + data, err := os.ReadFile(configPath) + if err != nil { + return nil, fmt.Errorf("读取配置文件失败: %w", err) + } + var partial struct { + Security struct { + Tools []ToolConfig `yaml:"tools"` + } `yaml:"security"` + } + if err := yaml.Unmarshal(data, &partial); err != nil { + return nil, fmt.Errorf("解析配置文件失败: %w", err) + } + if partial.Security.Tools == nil { + return []ToolConfig{}, nil + } + return partial.Security.Tools, nil +} + +// ReloadSecurityToolsFromDir 从 tools_dir 重新加载工具并更新 cfg.Security.Tools（ApplyConfig 热重载用）。 +func ReloadSecurityToolsFromDir(cfg *Config, configPath string) error { + if cfg == nil || strings.TrimSpace(cfg.Security.ToolsDir) == "" { + return nil + } + inlineTools, err := loadInlineSecurityToolsFromYAML(configPath) + if err != nil { + return err + } + toolsDir := ResolveToolsDir(cfg.Security.ToolsDir, configPath) + merged, err := MergeToolsFromDir(toolsDir, inlineTools) + if err != nil { + return fmt.Errorf("从工具目录加载工具配置失败: %w", err) + } + cfg.Security.Tools = merged + return nil +} + +// LoadToolsFromDir 从目录加载所有工具配置文件 +func LoadToolsFromDir(dir string) ([]ToolConfig, error) { + var tools []ToolConfig + + // 检查目录是否存在 + if _, err := os.Stat(dir); os.IsNotExist(err) { + return tools, nil // 目录不存在时返回空列表，不报错 + } + + // 读取目录中的所有 .yaml 和 .yml 文件 + entries, err := os.ReadDir(dir) + if err != nil { + return nil, fmt.Errorf("读取工具目录失败: %w", err) + } + + for _, entry := range entries { + if entry.IsDir() { + continue + } + + name := entry.Name() + if !strings.HasSuffix(name, ".yaml") && !strings.HasSuffix(name, ".yml") { + continue + } + + filePath := filepath.Join(dir, name) + tool, err := LoadToolFromFile(filePath) + if err != nil { + // 记录错误但继续加载其他文件 + fmt.Printf("警告: 加载工具配置文件 %s 失败: %v\n", filePath, err) + continue + } + + tools = append(tools, *tool) + } + + return tools, nil +} + +// LoadToolFromFile 从单个文件加载工具配置 +func LoadToolFromFile(path string) (*ToolConfig, error) { + data, err := os.ReadFile(path) + if err != nil { + return nil, fmt.Errorf("读取文件失败: %w", err) + } + + var tool ToolConfig + if err := yaml.Unmarshal(data, &tool); err != nil { + return nil, fmt.Errorf("解析工具配置失败: %w", err) + } + + // 验证必需字段 + if tool.Name == "" { + return nil, fmt.Errorf("工具名称不能为空") + } + if tool.Command == "" { + return nil, fmt.Errorf("工具命令不能为空") + } + + return &tool, nil +} + +// LoadRolesFromDir 从目录加载所有角色配置文件 +func LoadRolesFromDir(dir string) (map[string]RoleConfig, error) { + roles := make(map[string]RoleConfig) + + // 检查目录是否存在 + if _, err := os.Stat(dir); os.IsNotExist(err) { + return roles, nil // 目录不存在时返回空map，不报错 + } + + // 读取目录中的所有 .yaml 和 .yml 文件 + entries, err := os.ReadDir(dir) + if err != nil { + return nil, fmt.Errorf("读取角色目录失败: %w", err) + } + + for _, entry := range entries { + if entry.IsDir() { + continue + } + + name := entry.Name() + if !strings.HasSuffix(name, ".yaml") && !strings.HasSuffix(name, ".yml") { + continue + } + + filePath := filepath.Join(dir, name) + role, err := LoadRoleFromFile(filePath) + if err != nil { + // 记录错误但继续加载其他文件 + fmt.Printf("警告: 加载角色配置文件 %s 失败: %v\n", filePath, err) + continue + } + + // 使用角色名称作为key + roleName := role.Name + if roleName == "" { + // 如果角色名称为空，使用文件名（去掉扩展名）作为名称 + roleName = strings.TrimSuffix(strings.TrimSuffix(name, ".yaml"), ".yml") + role.Name = roleName + } + + roles[roleName] = *role + } + + return roles, nil +} + +// LoadRoleFromFile 从单个文件加载角色配置 +func LoadRoleFromFile(path string) (*RoleConfig, error) { + data, err := os.ReadFile(path) + if err != nil { + return nil, fmt.Errorf("读取文件失败: %w", err) + } + + var role RoleConfig + if err := yaml.Unmarshal(data, &role); err != nil { + return nil, fmt.Errorf("解析角色配置失败: %w", err) + } + + // 处理 icon 字段：如果包含 Unicode 转义格式（\U0001F3C6），转换为实际的 Unicode 字符 + // Go 的 yaml 库可能不会自动解析 \U 转义序列，需要手动转换 + if role.Icon != "" { + icon := role.Icon + // 去除可能的引号 + icon = strings.Trim(icon, `"`) + + // 检查是否是 Unicode 转义格式 \U0001F3C6（8位十六进制）或 \uXXXX（4位十六进制） + if len(icon) >= 3 && icon[0] == '\\' { + if icon[1] == 'U' && len(icon) >= 10 { + // \U0001F3C6 格式（8位十六进制） + if codePoint, err := strconv.ParseInt(icon[2:10], 16, 32); err == nil { + role.Icon = string(rune(codePoint)) + } + } else if icon[1] == 'u' && len(icon) >= 6 { + // \uXXXX 格式（4位十六进制） + if codePoint, err := strconv.ParseInt(icon[2:6], 16, 32); err == nil { + role.Icon = string(rune(codePoint)) + } + } + } + } + + // 验证必需字段 + if role.Name == "" { + // 如果名称为空，尝试从文件名获取 + baseName := filepath.Base(path) + role.Name = strings.TrimSuffix(strings.TrimSuffix(baseName, ".yaml"), ".yml") + } + + return &role, nil +} + +func Default() *Config { + strictRobotIdentity := true + return &Config{ + Server: ServerConfig{ + Host: "0.0.0.0", + Port: 8080, + }, + Log: LogConfig{ + Level: "info", + Output: "stdout", + }, + MCP: MCPConfig{ + Enabled: false, + Host: "127.0.0.1", + Port: 8081, + }, + AI: AIConfig{ + DefaultChannel: "default", + Channels: map[string]AIChannelConfig{ + "default": { + Name: "Default", + Provider: "openai_compatible", + BaseURL: "https://api.openai.com/v1", + Model: "gpt-4", + MaxTotalTokens: 120000, + MaxCompletionTokens: DefaultMaxCompletionTokens, + }, + }, + }, + OpenAI: OpenAIConfig{}, + Agent: AgentConfig{ + MaxIterations: 30, // 默认最大迭代次数 + ToolTimeoutMinutes: 10, // 单次工具执行默认最多 10 分钟，避免异常长时间占用 + ToolWaitTimeoutSeconds: 60, // 外部 MCP 工具单轮最多等待 60 秒，超时后返回 execution_id 可继续等待 + ExternalMCPMaxConcurrentPerServer: 2, // 单个外部 MCP server 默认最多 2 个工具同时执行 + ExternalMCPMaxConcurrentTotal: 16, // 外部 MCP 工具全局默认最多 16 个同时执行 + ExternalMCPCircuitFailureThreshold: 3, // 单个 server 连续 3 次失败后临时熔断 + ExternalMCPCircuitCooldownSeconds: 60, // 熔断默认冷却 60 秒 + ShellNoOutputTimeoutSeconds: 300, // execute/exec 无新输出空闲终止（秒）；-1 关闭 + }, + Security: SecurityConfig{ + Tools: []ToolConfig{}, // 工具配置应该从 config.yaml 或 tools/ 目录加载 + ToolsDir: "tools", // 默认工具目录 + }, + Database: DatabaseConfig{ + Path: "data/conversations.db", + KnowledgeDBPath: "data/knowledge.db", // 默认知识库数据库路径 + }, + Auth: AuthConfig{ + SessionDurationHours: 12, + }, + Audit: func() AuditConfig { + on := true + return AuditConfig{ + RetentionDays: 90, + MaxDetailBytes: 8192, + Enabled: &on, + } + }(), + Monitor: func() MonitorConfig { + days := 90 + return MonitorConfig{RetentionDays: &days} + }(), + Robots: RobotsConfig{ + Session: RobotSessionConfig{ + StrictUserIdentity: &strictRobotIdentity, + }, + }, + Knowledge: KnowledgeConfig{ + Enabled: true, + BasePath: "knowledge_base", + Embedding: EmbeddingConfig{ + Provider: "openai", + Model: "text-embedding-3-small", + BaseURL: "https://api.openai.com/v1", + }, + Retrieval: RetrievalConfig{ + TopK: 5, + SimilarityThreshold: 0.65, + MultiQuery: MultiQueryConfig{MaxQueries: 4}, + Rerank: RerankConfig{}, + PostRetrieve: PostRetrieveConfig{ + PrefetchTopK: 20, + }, + }, + Indexing: IndexingConfig{ + ChunkStrategy: "markdown_then_recursive", + RequestTimeoutSeconds: 120, + ChunkSize: 768, // 增加到 768，更好的上下文保持 + ChunkOverlap: 50, + MaxChunksPerItem: 20, // 限制单个知识项最多 20 个块，避免消耗过多配额 + BatchSize: 64, + PreferSourceFile: false, + MaxRPM: 100, // 默认 100 RPM，避免 429 错误 + RateLimitDelayMs: 600, // 600ms 间隔，对应 100 RPM + MaxRetries: 3, + RetryDelayMs: 1000, + SubIndexes: nil, + }, + }, + } +} + +// C2Config 内置 C2 模块开关（与知识库 enabled 语义一致：关闭后不初始化监听器、不注册 C2 MCP 工具）。 +type C2Config struct { + // Enabled 为 nil 表示未写配置，按 true 处理（兼容旧 config.yaml） + Enabled *bool `yaml:"enabled,omitempty" json:"enabled,omitempty"` +} + +// EnabledEffective 返回是否启用 C2；未显式配置时默认启用。 +func (c C2Config) EnabledEffective() bool { + if c.Enabled == nil { + return true + } + return *c.Enabled +} + +// C2Public 返回给前端的 C2 状态（仅标量）。 +type C2Public struct { + Enabled bool `json:"enabled"` +} + +// Public 将内部配置转为 API 响应。 +func (c C2Config) Public() C2Public { + return C2Public{Enabled: c.EnabledEffective()} +} + +// C2APIUpdate 设置页/API 更新 C2 开关。 +type C2APIUpdate struct { + Enabled bool `json:"enabled"` +} + +// KnowledgeConfig 知识库配置 +type KnowledgeConfig struct { + Enabled bool `yaml:"enabled" json:"enabled"` // 是否启用知识检索 + BasePath string `yaml:"base_path" json:"base_path"` // 知识库路径 + Embedding EmbeddingConfig `yaml:"embedding" json:"embedding"` + Retrieval RetrievalConfig `yaml:"retrieval" json:"retrieval"` + Indexing IndexingConfig `yaml:"indexing,omitempty" json:"indexing,omitempty"` // 索引构建配置 +} + +// IndexingConfig 索引构建配置（用于控制知识库索引构建时的行为） +type IndexingConfig struct { + // ChunkStrategy: "markdown_then_recursive"（默认，Eino Markdown 标题切分后再递归切）或 "recursive"（仅递归切分） + ChunkStrategy string `yaml:"chunk_strategy,omitempty" json:"chunk_strategy,omitempty"` + // RequestTimeoutSeconds 嵌入 HTTP 客户端超时（秒），0 表示使用默认 120 + RequestTimeoutSeconds int `yaml:"request_timeout_seconds,omitempty" json:"request_timeout_seconds,omitempty"` + // 分块配置 + ChunkSize int `yaml:"chunk_size,omitempty" json:"chunk_size,omitempty"` // 每个块的最大 token 数（估算），默认 512 + ChunkOverlap int `yaml:"chunk_overlap,omitempty" json:"chunk_overlap,omitempty"` // 块之间的重叠 token 数，默认 50 + MaxChunksPerItem int `yaml:"max_chunks_per_item,omitempty" json:"max_chunks_per_item,omitempty"` // 单个知识项的最大块数量，0 表示不限制 + + // PreferSourceFile 为 true 时优先用 Eino FileLoader 从 file_path 读原文再索引（与库内 content 不一致时以磁盘为准） + PreferSourceFile bool `yaml:"prefer_source_file,omitempty" json:"prefer_source_file,omitempty"` + + // 速率限制配置（用于避免 API 速率限制） + RateLimitDelayMs int `yaml:"rate_limit_delay_ms,omitempty" json:"rate_limit_delay_ms,omitempty"` // 请求间隔时间（毫秒），0 表示不使用固定延迟 + MaxRPM int `yaml:"max_rpm,omitempty" json:"max_rpm,omitempty"` // 每分钟最大请求数，0 表示不限制 + + // 重试配置（用于处理临时错误） + MaxRetries int `yaml:"max_retries,omitempty" json:"max_retries,omitempty"` // 最大重试次数，默认 3 + RetryDelayMs int `yaml:"retry_delay_ms,omitempty" json:"retry_delay_ms,omitempty"` // 重试间隔（毫秒），默认 1000 + + // BatchSize 嵌入批大小（SQLite 索引写入），0 表示默认 64 + BatchSize int `yaml:"batch_size,omitempty" json:"batch_size,omitempty"` + // SubIndexes 传入 Eino indexer.WithSubIndexes（逻辑分区标记，随 Document 元数据传递） + SubIndexes []string `yaml:"sub_indexes,omitempty" json:"sub_indexes,omitempty"` +} + +// EmbeddingConfig 嵌入配置 +type EmbeddingConfig struct { + Provider string `yaml:"provider" json:"provider"` // 嵌入模型提供商 + Model string `yaml:"model" json:"model"` // 模型名称 + BaseURL string `yaml:"base_url" json:"base_url"` // API Base URL + APIKey string `yaml:"api_key" json:"api_key"` // API Key（从OpenAI配置继承） +} + +// PostRetrieveConfig 检索后处理：固定对正文做规范化去重（最佳实践）、上下文预算截断；PrefetchTopK 用于多取候选再收敛到 top_k。 +type PostRetrieveConfig struct { + // PrefetchTopK 向量检索阶段每条 MultiQuery 变体最多保留的候选数；0 表示使用内置默认 max(top_k*4, 20)。 + PrefetchTopK int `yaml:"prefetch_top_k,omitempty" json:"prefetch_top_k,omitempty"` + // MaxContextChars 返回文档内容总 Unicode 字符数上限（整段 chunk，不截断半段）；0 表示不限制。 + MaxContextChars int `yaml:"max_context_chars,omitempty" json:"max_context_chars,omitempty"` + // MaxContextTokens 返回文档内容总 token 上限（tiktoken，按嵌入模型名映射，失败则 cl100k_base）；0 表示不限制。 + MaxContextTokens int `yaml:"max_context_tokens,omitempty" json:"max_context_tokens,omitempty"` +} + +// MultiQueryConfig Eino MultiQuery 查询改写（始终启用，无关闭开关）。 +type MultiQueryConfig struct { + // MaxQueries LLM 生成的检索变体上限（含原问语义覆盖）；0 表示默认 4。 + MaxQueries int `yaml:"max_queries,omitempty" json:"max_queries,omitempty"` +} + +func (c MultiQueryConfig) MaxQueriesEffective() int { + if c.MaxQueries <= 0 { + return 4 + } + if c.MaxQueries > 8 { + return 8 + } + return c.MaxQueries +} + +// RerankConfig 检索精排（始终启用）；支持 dashscope 与 Cohere 兼容 HTTP API。 +type RerankConfig struct { + // Provider: dashscope | cohere；空则按 base_url 自动推断。 + Provider string `yaml:"provider,omitempty" json:"provider,omitempty"` + Model string `yaml:"model,omitempty" json:"model,omitempty"` + BaseURL string `yaml:"base_url,omitempty" json:"base_url,omitempty"` + APIKey string `yaml:"api_key,omitempty" json:"api_key,omitempty"` +} + +func (c RerankConfig) ProviderEffective(baseURL string) string { + p := strings.TrimSpace(strings.ToLower(c.Provider)) + if p != "" { + return p + } + u := strings.ToLower(baseURL) + if strings.Contains(u, "dashscope") { + return "dashscope" + } + return "cohere" +} + +func (c RerankConfig) ModelEffective(provider string) string { + if m := strings.TrimSpace(c.Model); m != "" { + return m + } + if provider == "dashscope" { + return "gte-rerank" + } + return "rerank-multilingual-v3.0" +} + +// RetrievalConfig 检索配置 +type RetrievalConfig struct { + TopK int `yaml:"top_k" json:"top_k"` // 检索Top-K + SimilarityThreshold float64 `yaml:"similarity_threshold" json:"similarity_threshold"` // 余弦相似度阈值 + // SubIndexFilter 非空时仅保留 sub_indexes 含该标签（逗号分隔之一）的行；sub_indexes 为空的旧行仍返回。 + SubIndexFilter string `yaml:"sub_index_filter,omitempty" json:"sub_index_filter,omitempty"` + MultiQuery MultiQueryConfig `yaml:"multi_query" json:"multi_query"` + Rerank RerankConfig `yaml:"rerank" json:"rerank"` + // PostRetrieve 检索后处理（去重、预算截断）；精排在 MultiQuery 融合后执行。 + PostRetrieve PostRetrieveConfig `yaml:"post_retrieve,omitempty" json:"post_retrieve,omitempty"` +} + +// RolesConfig 角色配置（已废弃，使用 map[string]RoleConfig 替代） +// 保留此类型以兼容旧代码，但建议直接使用 map[string]RoleConfig +type RolesConfig struct { + Roles map[string]RoleConfig `yaml:"roles,omitempty" json:"roles,omitempty"` +} + +// RoleConfig 单个角色配置 +type RoleConfig struct { + Name string `yaml:"name" json:"name"` // 角色名称 + Description string `yaml:"description" json:"description"` // 角色描述 + UserPrompt string `yaml:"user_prompt" json:"user_prompt"` // 用户提示词(追加到用户消息前) + 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替代） + WorkflowID string `yaml:"workflow_id,omitempty" json:"workflow_id,omitempty"` // 可选：绑定工作流 ID + WorkflowVersion string `yaml:"workflow_version,omitempty" json:"workflow_version,omitempty"` // latest 或具体版本号；空等同 latest + WorkflowPolicy string `yaml:"workflow_policy,omitempty" json:"workflow_policy,omitempty"` // auto | off；空且 workflow_id 非空时按 auto + Enabled bool `yaml:"enabled" json:"enabled"` // 是否启用 +} diff --git a/internal/config/config_test.go b/internal/config/config_test.go new file mode 100644 index 00000000..a774f8c6 --- /dev/null +++ b/internal/config/config_test.go @@ -0,0 +1,207 @@ +package config + +import ( + "os" + "path/filepath" + "strings" + "testing" +) + +func TestEnsureLocalConfigCreatesFromExample(t *testing.T) { + dir := t.TempDir() + examplePath := filepath.Join(dir, "config.example.yaml") + configPath := filepath.Join(dir, "config.yaml") + + example := []byte(`auth: + session_duration_hours: 12 +server: + host: 127.0.0.1 + port: 8080 +`) + if err := os.WriteFile(examplePath, example, 0644); err != nil { + t.Fatalf("write example: %v", err) + } + + result, err := EnsureLocalConfig(configPath) + if err != nil { + t.Fatalf("EnsureLocalConfig: %v", err) + } + if !result.Created { + t.Fatal("Created = false, want true") + } + if result.ExamplePath != examplePath { + t.Fatalf("ExamplePath = %q, want %q", result.ExamplePath, examplePath) + } + + cfg, err := Load(configPath) + if err != nil { + t.Fatalf("Load generated config: %v", err) + } + if cfg.Auth.SessionDurationHours != 12 { + t.Fatalf("SessionDurationHours = %d, want 12", cfg.Auth.SessionDurationHours) + } + + second, err := EnsureLocalConfig(configPath) + if err != nil { + t.Fatalf("EnsureLocalConfig existing: %v", err) + } + if second.Created { + t.Fatal("Created = true for existing config, want false") + } +} + +func TestLoadIgnoresLegacyAuthPasswordField(t *testing.T) { + dir := t.TempDir() + path := filepath.Join(dir, "config.yaml") + initial := strings.Join([]string{ + "auth:", + ` password: "legacy-password"`, + " session_duration_hours: 12", + "server:", + " host: 127.0.0.1", + " port: 8080", + "", + }, "\n") + if err := os.WriteFile(path, []byte(initial), 0644); err != nil { + t.Fatalf("write config: %v", err) + } + + cfg, err := Load(path) + if err != nil { + t.Fatalf("Load: %v", err) + } + if cfg.Auth.SessionDurationHours != 12 { + t.Fatalf("SessionDurationHours = %d, want 12", cfg.Auth.SessionDurationHours) + } +} + +func TestHitlAuditModelEffectiveFallsBackToMainConfig(t *testing.T) { + main := OpenAIConfig{ + Provider: "openai", + BaseURL: "https://api.example.com/v1", + APIKey: "main-key", + Model: "large-model", + } + + got := (HitlConfig{ + AuditModel: OpenAIConfig{Model: "small-reviewer"}, + }).AuditModelEffective(main) + + if got.Provider != main.Provider || got.BaseURL != main.BaseURL || got.APIKey != main.APIKey { + t.Fatalf("expected provider/base_url/api_key to inherit main config, got %+v", got) + } + if got.Model != "small-reviewer" { + t.Fatalf("expected audit model override, got %q", got.Model) + } +} + +func TestLoadUsesAIDefaultChannelAsRuntimeOpenAI(t *testing.T) { + dir := t.TempDir() + path := filepath.Join(dir, "config.yaml") + initial := strings.Join([]string{ + "ai:", + " default_channel: deepseek", + " channels:", + " qwen:", + " name: Qwen", + " provider: openai_compatible", + " base_url: https://dashscope.example/v1", + " api_key: qwen-key", + " model: qwen-max", + " deepseek:", + " name: DeepSeek", + " provider: openai_compatible", + " base_url: https://deepseek.example/v1", + " api_key: deepseek-key", + " model: deepseek-chat", + " max_total_tokens: 64000", + "server:", + " host: 127.0.0.1", + " port: 8080", + "", + }, "\n") + if err := os.WriteFile(path, []byte(initial), 0644); err != nil { + t.Fatalf("write config: %v", err) + } + + cfg, err := Load(path) + if err != nil { + t.Fatalf("Load: %v", err) + } + if cfg.OpenAI.Model != "deepseek-chat" || cfg.OpenAI.APIKey != "deepseek-key" || cfg.OpenAI.MaxTotalTokens != 64000 { + t.Fatalf("runtime OpenAI config did not follow ai.default_channel: %+v", cfg.OpenAI) + } + oa, id, ok := cfg.ResolveAIChannel("qwen") + if !ok || id != "qwen" || oa.Model != "qwen-max" || oa.APIKey != "qwen-key" { + t.Fatalf("ResolveAIChannel(qwen) = (%+v, %q, %v)", oa, id, ok) + } +} + +func TestSummarizationUserIntentLedgerRunesEffective(t *testing.T) { + var zero MultiAgentEinoMiddlewareConfig + if got := zero.SummarizationUserIntentLedgerMaxRunesEffective(); got != DefaultSummarizationUserIntentLedgerMaxRunes { + t.Fatalf("default ledger max runes = %d, want %d", got, DefaultSummarizationUserIntentLedgerMaxRunes) + } + if got := zero.SummarizationUserIntentLedgerEntryMaxRunesEffective(); got != DefaultSummarizationUserIntentLedgerEntryMaxRunes { + t.Fatalf("default ledger entry max runes = %d, want %d", got, DefaultSummarizationUserIntentLedgerEntryMaxRunes) + } + + custom := MultiAgentEinoMiddlewareConfig{ + SummarizationUserIntentLedgerMaxRunes: 12345, + SummarizationUserIntentLedgerEntryMaxRunes: 2345, + } + if got := custom.SummarizationUserIntentLedgerMaxRunesEffective(); got != 12345 { + t.Fatalf("custom ledger max runes = %d", got) + } + if got := custom.SummarizationUserIntentLedgerEntryMaxRunesEffective(); got != 2345 { + t.Fatalf("custom ledger entry max runes = %d", got) + } +} + +func TestSummarizationOutputReserveTokensEffective(t *testing.T) { + var zero MultiAgentEinoMiddlewareConfig + if got := zero.SummarizationOutputReserveTokensEffective(); got != DefaultSummarizationOutputReserveTokens { + t.Fatalf("default output reserve = %d, want %d", got, DefaultSummarizationOutputReserveTokens) + } + custom := MultiAgentEinoMiddlewareConfig{SummarizationOutputReserveTokens: 4096} + if got := custom.SummarizationOutputReserveTokensEffective(); got != 4096 { + t.Fatalf("custom output reserve = %d", got) + } +} + +func TestOpenAIOutputLimitValidation(t *testing.T) { + if got := (OpenAIConfig{}).MaxCompletionTokensEffective(); got != DefaultMaxCompletionTokens { + t.Fatalf("max completion default=%d", got) + } + if err := validateOpenAIOutputLimits(OpenAIConfig{MaxCompletionTokens: -1}); err == nil { + t.Fatal("negative completion limit must fail") + } +} + +func TestLatestUserMessageRunesEffective(t *testing.T) { + var zero MultiAgentEinoMiddlewareConfig + if got := zero.LatestUserMessageMaxRunesEffective(); got != DefaultLatestUserMessageMaxRunes { + t.Fatalf("default latest user max runes = %d, want %d", got, DefaultLatestUserMessageMaxRunes) + } + if got := zero.LatestUserMessageHeadRunesEffective(); got != DefaultLatestUserMessageHeadRunes { + t.Fatalf("default latest user head runes = %d, want %d", got, DefaultLatestUserMessageHeadRunes) + } + if got := zero.LatestUserMessageTailRunesEffective(); got != DefaultLatestUserMessageTailRunes { + t.Fatalf("default latest user tail runes = %d, want %d", got, DefaultLatestUserMessageTailRunes) + } + + custom := MultiAgentEinoMiddlewareConfig{ + LatestUserMessageMaxRunes: 100, + LatestUserMessageHeadRunes: 40, + LatestUserMessageTailRunes: 60, + } + if got := custom.LatestUserMessageMaxRunesEffective(); got != 100 { + t.Fatalf("custom latest user max runes = %d", got) + } + if got := custom.LatestUserMessageHeadRunesEffective(); got != 40 { + t.Fatalf("custom latest user head runes = %d", got) + } + if got := custom.LatestUserMessageTailRunesEffective(); got != 60 { + t.Fatalf("custom latest user tail runes = %d", got) + } +} diff --git a/internal/config/envexpand.go b/internal/config/envexpand.go new file mode 100644 index 00000000..0ffc1784 --- /dev/null +++ b/internal/config/envexpand.go @@ -0,0 +1,66 @@ +package config + +import ( + "os" + "strings" +) + +// expandEnvVar 展开字符串中的 ${VAR} 和 ${VAR:-default} 环境变量引用。 +// 与官方 MCP 配置格式一致（Claude Desktop / Cursor / VS Code 均支持此语法）。 +func expandEnvVar(s string) string { + var b strings.Builder + i := 0 + for i < len(s) { + // 查找 ${ + idx := strings.Index(s[i:], "${") + if idx < 0 { + b.WriteString(s[i:]) + break + } + b.WriteString(s[i : i+idx]) + i += idx + 2 // skip ${ + + // 查找对应的 } + end := strings.IndexByte(s[i:], '}') + if end < 0 { + // 没有 }，原样保留 + b.WriteString("${") + continue + } + expr := s[i : i+end] + i += end + 1 // skip } + + // 解析 VAR:-default + varName := expr + defaultVal := "" + hasDefault := false + if colonIdx := strings.Index(expr, ":-"); colonIdx >= 0 { + varName = expr[:colonIdx] + defaultVal = expr[colonIdx+2:] + hasDefault = true + } + + val := os.Getenv(varName) + if val == "" && hasDefault { + val = defaultVal + } + b.WriteString(val) + } + return b.String() +} + +// ExpandConfigEnv 展开 ExternalMCPServerConfig 中所有支持环境变量的字段。 +// 展开范围：Command、Args、Env values、URL、Headers values。 +func ExpandConfigEnv(cfg *ExternalMCPServerConfig) { + cfg.Command = expandEnvVar(cfg.Command) + for i, arg := range cfg.Args { + cfg.Args[i] = expandEnvVar(arg) + } + for k, v := range cfg.Env { + cfg.Env[k] = expandEnvVar(v) + } + cfg.URL = expandEnvVar(cfg.URL) + for k, v := range cfg.Headers { + cfg.Headers[k] = expandEnvVar(v) + } +} diff --git a/internal/config/envexpand_test.go b/internal/config/envexpand_test.go new file mode 100644 index 00000000..a17c4514 --- /dev/null +++ b/internal/config/envexpand_test.go @@ -0,0 +1,81 @@ +package config + +import ( + "os" + "testing" +) + +func TestExpandEnvVar(t *testing.T) { + os.Setenv("TEST_MCP_VAR", "hello") + os.Setenv("TEST_MCP_PATH", "/usr/local/bin") + defer os.Unsetenv("TEST_MCP_VAR") + defer os.Unsetenv("TEST_MCP_PATH") + + tests := []struct { + name string + input string + expect string + }{ + {"plain string", "no vars here", "no vars here"}, + {"empty string", "", ""}, + {"simple var", "${TEST_MCP_VAR}", "hello"}, + {"var in middle", "prefix-${TEST_MCP_VAR}-suffix", "prefix-hello-suffix"}, + {"multiple vars", "${TEST_MCP_PATH}/${TEST_MCP_VAR}", "/usr/local/bin/hello"}, + {"missing var empty", "${NONEXISTENT_MCP_VAR_XYZ}", ""}, + {"default value used", "${NONEXISTENT_MCP_VAR_XYZ:-fallback}", "fallback"}, + {"default not used", "${TEST_MCP_VAR:-unused}", "hello"}, + {"default with path", "${NONEXISTENT_MCP_VAR_XYZ:-/tmp/default}", "/tmp/default"}, + {"unclosed brace", "${UNCLOSED", "${UNCLOSED"}, + {"dollar without brace", "$PLAIN", "$PLAIN"}, + {"empty var name", "${}", ""}, + {"default empty var", "${:-default}", "default"}, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + got := expandEnvVar(tt.input) + if got != tt.expect { + t.Errorf("expandEnvVar(%q) = %q, want %q", tt.input, got, tt.expect) + } + }) + } +} + +func TestExpandConfigEnv(t *testing.T) { + os.Setenv("TEST_MCP_CMD", "python3") + os.Setenv("TEST_MCP_TOKEN", "secret123") + defer os.Unsetenv("TEST_MCP_CMD") + defer os.Unsetenv("TEST_MCP_TOKEN") + + cfg := &ExternalMCPServerConfig{ + Command: "${TEST_MCP_CMD}", + Args: []string{"--token", "${TEST_MCP_TOKEN}", "${MISSING:-default_arg}"}, + Env: map[string]string{"API_KEY": "${TEST_MCP_TOKEN}", "LEVEL": "${MISSING:-INFO}"}, + URL: "https://${MISSING:-example.com}/mcp", + Headers: map[string]string{"Authorization": "Bearer ${TEST_MCP_TOKEN}"}, + } + + ExpandConfigEnv(cfg) + + if cfg.Command != "python3" { + t.Errorf("Command = %q, want %q", cfg.Command, "python3") + } + if cfg.Args[1] != "secret123" { + t.Errorf("Args[1] = %q, want %q", cfg.Args[1], "secret123") + } + if cfg.Args[2] != "default_arg" { + t.Errorf("Args[2] = %q, want %q", cfg.Args[2], "default_arg") + } + if cfg.Env["API_KEY"] != "secret123" { + t.Errorf("Env[API_KEY] = %q, want %q", cfg.Env["API_KEY"], "secret123") + } + if cfg.Env["LEVEL"] != "INFO" { + t.Errorf("Env[LEVEL] = %q, want %q", cfg.Env["LEVEL"], "INFO") + } + if cfg.URL != "https://example.com/mcp" { + t.Errorf("URL = %q, want %q", cfg.URL, "https://example.com/mcp") + } + if cfg.Headers["Authorization"] != "Bearer secret123" { + t.Errorf("Headers[Authorization] = %q, want %q", cfg.Headers["Authorization"], "Bearer secret123") + } +} diff --git a/internal/config/hitl_prompt_test.go b/internal/config/hitl_prompt_test.go new file mode 100644 index 00000000..d2fbdcb8 --- /dev/null +++ b/internal/config/hitl_prompt_test.go @@ -0,0 +1,31 @@ +package config + +import ( + "strings" + "testing" +) + +func TestDefaultHitlAuditAgentPromptIncludesPrioritizedRules(t *testing.T) { + prompt := DefaultHitlAuditAgentPrompt() + for _, want := range []string{ + "如果同时命中 reject 和 approve，必须 reject", + "修改/重置任意用户或管理员密码", + "修改/创建/删除用户、角色、权限", + "停止、禁用、重启业务服务", + "命中规则：...", + } { + if !strings.Contains(prompt, want) { + t.Fatalf("default approval prompt missing %q", want) + } + } +} + +func TestDefaultHitlAuditAgentPromptReviewEditKeepsEditedArguments(t *testing.T) { + prompt := DefaultHitlAuditAgentPromptReviewEdit() + if !strings.Contains(prompt, `"editedArguments":{...}`) { + t.Fatal("review-edit prompt must preserve editedArguments output") + } + if !strings.Contains(prompt, "命中规则：...") { + t.Fatal("review-edit prompt must require a matched rule") + } +} diff --git a/internal/config/robots_validate_test.go b/internal/config/robots_validate_test.go new file mode 100644 index 00000000..f012ce54 --- /dev/null +++ b/internal/config/robots_validate_test.go @@ -0,0 +1,69 @@ +package config + +import "testing" + +func TestValidateWecomConfig(t *testing.T) { + t.Parallel() + + tests := []struct { + name string + cfg RobotWecomConfig + wantErr bool + }{ + { + name: "disabled without token", + cfg: RobotWecomConfig{Enabled: false, Token: ""}, + wantErr: false, + }, + { + name: "enabled with token", + cfg: RobotWecomConfig{Enabled: true, Token: "secret"}, + wantErr: false, + }, + { + name: "enabled without token", + cfg: RobotWecomConfig{Enabled: true, Token: ""}, + wantErr: true, + }, + { + name: "enabled with whitespace token", + cfg: RobotWecomConfig{Enabled: true, Token: " "}, + wantErr: true, + }, + } + + for _, tt := range tests { + tt := tt + t.Run(tt.name, func(t *testing.T) { + t.Parallel() + err := ValidateWecomConfig(tt.cfg) + if (err != nil) != tt.wantErr { + t.Fatalf("ValidateWecomConfig() error = %v, wantErr %v", err, tt.wantErr) + } + }) + } +} + +func TestValidateRobotAuthorization(t *testing.T) { + tests := []struct { + name string + cfg RobotAuthorizationConfig + wantErr bool + }{ + {name: "default user binding", cfg: RobotAuthorizationConfig{}, wantErr: false}, + {name: "explicit user binding", cfg: RobotAuthorizationConfig{Mode: RobotAuthModeUserBinding}, wantErr: false}, + {name: "service account", cfg: RobotAuthorizationConfig{Mode: RobotAuthModeServiceAccount, ServiceUserID: "svc-1", AllowedExternalUsers: []string{"t:x|u:y"}}, wantErr: false}, + {name: "missing service user", cfg: RobotAuthorizationConfig{Mode: RobotAuthModeServiceAccount, AllowedExternalUsers: []string{"t:x|u:y"}}, wantErr: true}, + {name: "admin allowed with exact sender", cfg: RobotAuthorizationConfig{Mode: RobotAuthModeServiceAccount, ServiceUserID: "admin", AllowedExternalUsers: []string{"t:x|u:y"}}, wantErr: false}, + {name: "allowlist required", cfg: RobotAuthorizationConfig{Mode: RobotAuthModeServiceAccount, ServiceUserID: "svc-1"}, wantErr: true}, + {name: "wildcard forbidden", cfg: RobotAuthorizationConfig{Mode: RobotAuthModeServiceAccount, ServiceUserID: "svc-1", AllowedExternalUsers: []string{"*"}}, wantErr: true}, + {name: "unknown mode", cfg: RobotAuthorizationConfig{Mode: "open"}, wantErr: true}, + } + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + if err := ValidateRobotAuthorization(tt.cfg, "robots.lark"); (err != nil) != tt.wantErr { + t.Fatalf("ValidateRobotAuthorization() error=%v wantErr=%v", err, tt.wantErr) + } + }) + } +} diff --git a/internal/config/server_https_bootstrap.go b/internal/config/server_https_bootstrap.go new file mode 100644 index 00000000..6f67a601 --- /dev/null +++ b/internal/config/server_https_bootstrap.go @@ -0,0 +1,60 @@ +package config + +import "strings" + +// MainWebUIUsesHTTPS 判断主 Web UI 是否以 HTTPS 监听（与 internal/app.prepareMainServerTLS 前置条件一致）。 +func MainWebUIUsesHTTPS(s *ServerConfig) bool { + if s == nil { + return false + } + if s.TLSEnabled { + return true + } + if s.TLSAutoSelfSign { + return true + } + cert := strings.TrimSpace(s.TLSCertPath) + key := strings.TrimSpace(s.TLSKeyPath) + return cert != "" && key != "" +} + +// ServerHTTPRedirectEnabled 是否在主站启用 HTTPS 时把明文 HTTP 请求重定向到 HTTPS（默认开启）。 +func ServerHTTPRedirectEnabled(s *ServerConfig) bool { + if s == nil || !MainWebUIUsesHTTPS(s) { + return false + } + if s.TLSHTTPRedirect == nil { + return true + } + return *s.TLSHTTPRedirect +} + +// ApplyDevHTTPSBootstrap 供 --https / 一键脚本使用：强制开启主站 TLS。 +// 若已配置 tls_cert_path 与 tls_key_path 则仅用 PEM，不开启自签；否则启用 tls_auto_self_sign（内存证书，仅本地测试）。 +func ApplyDevHTTPSBootstrap(cfg *Config) { + if cfg == nil { + return + } + cfg.Server.TLSEnabled = true + cert := strings.TrimSpace(cfg.Server.TLSCertPath) + key := strings.TrimSpace(cfg.Server.TLSKeyPath) + if cert != "" && key != "" { + cfg.Server.TLSAutoSelfSign = false + return + } + cfg.Server.TLSAutoSelfSign = true +} + +// ApplyPlainHTTPBootstrap 供 --http / 一键脚本使用：强制主站使用明文 HTTP。 +// 它会覆盖配置文件中的 TLS 开关、自签证书以及证书路径，避免 --http 仍被配置中的 HTTPS 选项重新启用。 +func ApplyPlainHTTPBootstrap(cfg *Config) { + if cfg == nil { + return + } + cfg.Server.TLSEnabled = false + cfg.Server.TLSAutoSelfSign = false + cfg.Server.TLSCertPath = "" + cfg.Server.TLSKeyPath = "" + disabled := false + cfg.Server.TLSHTTPRedirect = &disabled +} diff --git a/internal/config/server_https_bootstrap_test.go b/internal/config/server_https_bootstrap_test.go new file mode 100644 index 00000000..0a7836ca --- /dev/null +++ b/internal/config/server_https_bootstrap_test.go @@ -0,0 +1,31 @@ +package config + +import "testing" + +func TestApplyPlainHTTPBootstrapDisablesConfiguredTLS(t *testing.T) { + enabled := true + cfg := &Config{ + Server: ServerConfig{ + TLSEnabled: true, + TLSAutoSelfSign: true, + TLSCertPath: "/tmp/server.crt", + TLSKeyPath: "/tmp/server.key", + TLSHTTPRedirect: &enabled, + }, + } + + ApplyPlainHTTPBootstrap(cfg) + + if MainWebUIUsesHTTPS(&cfg.Server) { + t.Fatal("expected --http bootstrap to disable main web UI HTTPS") + } + if ServerHTTPRedirectEnabled(&cfg.Server) { + t.Fatal("expected --http bootstrap to disable HTTP to HTTPS redirect") + } + if cfg.Server.TLSCertPath != "" || cfg.Server.TLSKeyPath != "" { + t.Fatalf("expected TLS cert paths to be cleared, got cert=%q key=%q", cfg.Server.TLSCertPath, cfg.Server.TLSKeyPath) + } + if cfg.Server.TLSHTTPRedirect == nil || *cfg.Server.TLSHTTPRedirect { + t.Fatal("expected TLSHTTPRedirect to be explicitly disabled") + } +} diff --git a/internal/config/tools_reload_test.go b/internal/config/tools_reload_test.go new file mode 100644 index 00000000..0fbbd074 --- /dev/null +++ b/internal/config/tools_reload_test.go @@ -0,0 +1,111 @@ +package config + +import ( + "os" + "path/filepath" + "testing" +) + +func TestReloadSecurityToolsFromDir(t *testing.T) { + root := t.TempDir() + toolsDir := filepath.Join(root, "tools") + if err := os.MkdirAll(toolsDir, 0755); err != nil { + t.Fatal(err) + } + + configPath := filepath.Join(root, "config.yaml") + if err := os.WriteFile(configPath, []byte(`security: + tools_dir: tools + tools: + - name: inline-only + command: inline-cmd + enabled: true + description: inline tool +`), 0644); err != nil { + t.Fatal(err) + } + + writeTool := func(name, command string) { + t.Helper() + content := "name: " + name + "\ncommand: " + command + "\nenabled: true\ndescription: test\n" + if err := os.WriteFile(filepath.Join(toolsDir, name+".yaml"), []byte(content), 0644); err != nil { + t.Fatal(err) + } + } + + writeTool("alpha", "alpha-cmd") + + cfg := &Config{ + Security: SecurityConfig{ + ToolsDir: "tools", + Tools: []ToolConfig{ + {Name: "stale", Command: "stale-cmd", Enabled: true, Description: "should be removed"}, + }, + }, + } + + if err := ReloadSecurityToolsFromDir(cfg, configPath); err != nil { + t.Fatalf("reload: %v", err) + } + if len(cfg.Security.Tools) != 2 { + t.Fatalf("expected 2 tools, got %d", len(cfg.Security.Tools)) + } + + names := map[string]string{} + for _, tool := range cfg.Security.Tools { + names[tool.Name] = tool.Command + } + if names["alpha"] != "alpha-cmd" { + t.Fatalf("alpha tool missing or wrong command: %#v", names) + } + if names["inline-only"] != "inline-cmd" { + t.Fatalf("inline-only tool missing: %#v", names) + } + if _, ok := names["stale"]; ok { + t.Fatal("stale in-memory tool should not survive reload") + } + + writeTool("beta", "beta-cmd") + if err := ReloadSecurityToolsFromDir(cfg, configPath); err != nil { + t.Fatalf("second reload: %v", err) + } + if len(cfg.Security.Tools) != 3 { + t.Fatalf("expected 3 tools after add, got %d", len(cfg.Security.Tools)) + } + foundBeta := false + for _, tool := range cfg.Security.Tools { + if tool.Name == "beta" { + foundBeta = true + break + } + } + if !foundBeta { + t.Fatal("beta tool not found after second reload") + } +} + +func TestMergeToolsFromDir_DirOverridesInline(t *testing.T) { + root := t.TempDir() + toolsDir := filepath.Join(root, "tools") + if err := os.MkdirAll(toolsDir, 0755); err != nil { + t.Fatal(err) + } + content := "name: shared\ncommand: dir-cmd\nenabled: true\ndescription: from dir\n" + if err := os.WriteFile(filepath.Join(toolsDir, "shared.yaml"), []byte(content), 0644); err != nil { + t.Fatal(err) + } + + inline := []ToolConfig{ + {Name: "shared", Command: "inline-cmd", Enabled: true, Description: "from inline"}, + } + merged, err := MergeToolsFromDir(toolsDir, inline) + if err != nil { + t.Fatal(err) + } + if len(merged) != 1 { + t.Fatalf("expected 1 tool, got %d", len(merged)) + } + if merged[0].Command != "dir-cmd" { + t.Fatalf("dir tool should win, got command %q", merged[0].Command) + } +} diff --git a/internal/config/vision.go b/internal/config/vision.go new file mode 100644 index 00000000..1052d3b9 --- /dev/null +++ b/internal/config/vision.go @@ -0,0 +1,97 @@ +package config + +import "strings" + +// VisionConfig 独立视觉模型与 analyze_image 工具参数；enabled 时注册 MCP 工具 analyze_image。 +type VisionConfig struct { + Enabled bool `yaml:"enabled" json:"enabled"` + APIKey string `yaml:"api_key,omitempty" json:"api_key,omitempty"` + BaseURL string `yaml:"base_url,omitempty" json:"base_url,omitempty"` + Model string `yaml:"model,omitempty" json:"model,omitempty"` + Provider string `yaml:"provider,omitempty" json:"provider,omitempty"` + TimeoutSeconds int `yaml:"timeout_seconds,omitempty" json:"timeout_seconds,omitempty"` + MaxImageBytes int64 `yaml:"max_image_bytes,omitempty" json:"max_image_bytes,omitempty"` + MaxDimension int `yaml:"max_dimension,omitempty" json:"max_dimension,omitempty"` + JPEGQuality int `yaml:"jpeg_quality,omitempty" json:"jpeg_quality,omitempty"` + MaxPayloadBytes int64 `yaml:"max_payload_bytes,omitempty" json:"max_payload_bytes,omitempty"` + SkipPreprocessBelowBytes int64 `yaml:"skip_preprocess_below_bytes,omitempty" json:"skip_preprocess_below_bytes,omitempty"` // 0=始终压缩；默认 2MB 且长边已<=max_dimension 时原图直传 + Detail string `yaml:"detail,omitempty" json:"detail,omitempty"` // low | high | auto +} + +func (v VisionConfig) TimeoutSecondsEffective() int { + if v.TimeoutSeconds <= 0 { + return 60 + } + return v.TimeoutSeconds +} + +func (v VisionConfig) MaxImageBytesEffective() int64 { + if v.MaxImageBytes <= 0 { + return 5 * 1024 * 1024 + } + return v.MaxImageBytes +} + +func (v VisionConfig) MaxDimensionEffective() int { + if v.MaxDimension <= 0 { + return 2048 + } + return v.MaxDimension +} + +func (v VisionConfig) JPEGQualityEffective() int { + if v.JPEGQuality <= 0 || v.JPEGQuality > 100 { + return 82 + } + return v.JPEGQuality +} + +func (v VisionConfig) MaxPayloadBytesEffective() int64 { + if v.MaxPayloadBytes <= 0 { + return 512 * 1024 + } + return v.MaxPayloadBytes +} + +// SkipPreprocessBelowBytesEffective 低于该字节数且长边<=max_dimension、且<=max_payload 时可原图直传；0 表示始终压缩。 +func (v VisionConfig) SkipPreprocessBelowBytesEffective() int64 { + if v.SkipPreprocessBelowBytes < 0 { + return 0 + } + return v.SkipPreprocessBelowBytes +} + +func (v VisionConfig) DetailEffective() string { + d := strings.ToLower(strings.TrimSpace(v.Detail)) + switch d { + case "high", "low", "auto": + return d + default: + return "low" + } +} + +// OpenAICfgEffective 合并主 openai 配置与 vision 覆盖项，供 VL ChatModel 使用。 +// vision.api_key / base_url / provider 留空或省略时，沿用 main（openai）对应字段；vision.model 必填（由 Ready 校验）。 +func (v VisionConfig) OpenAICfgEffective(main OpenAIConfig) OpenAIConfig { + out := main + if k := strings.TrimSpace(v.APIKey); k != "" { + out.APIKey = k + } + if u := strings.TrimSpace(v.BaseURL); u != "" { + out.BaseURL = u + } + if m := strings.TrimSpace(v.Model); m != "" { + out.Model = m + } + if p := strings.TrimSpace(v.Provider); p != "" { + out.Provider = p + } + out.Reasoning.Mode = "off" + return out +} + +// Ready 表示已启用且模型名非空。 +func (v VisionConfig) Ready() bool { + return v.Enabled && strings.TrimSpace(v.Model) != "" +} diff --git a/internal/config/vision_test.go b/internal/config/vision_test.go new file mode 100644 index 00000000..0620a181 --- /dev/null +++ b/internal/config/vision_test.go @@ -0,0 +1,55 @@ +package config + +import "testing" + +func TestVisionConfig_OpenAICfgEffective_fallbackToMain(t *testing.T) { + main := OpenAIConfig{ + APIKey: "main-key", + BaseURL: "https://main.example/v1", + Model: "main-model", + Provider: "openai", + } + v := VisionConfig{Model: "qwen-vl-max"} + out := v.OpenAICfgEffective(main) + if out.APIKey != main.APIKey || out.BaseURL != main.BaseURL || out.Provider != main.Provider { + t.Fatalf("expected openai fallback, got key=%q url=%q provider=%q", out.APIKey, out.BaseURL, out.Provider) + } + if out.Model != "qwen-vl-max" { + t.Fatalf("model: %s", out.Model) + } +} + +func TestVisionConfig_OpenAICfgEffective(t *testing.T) { + main := OpenAIConfig{ + APIKey: "main-key", + BaseURL: "https://main.example/v1", + Model: "main-model", + Provider: "openai", + Reasoning: OpenAIReasoningConfig{Mode: "on"}, + } + v := VisionConfig{ + Model: "vl-model", + APIKey: "vl-key", + BaseURL: "https://vl.example/v1", + Provider: "claude", + } + out := v.OpenAICfgEffective(main) + if out.APIKey != "vl-key" || out.BaseURL != "https://vl.example/v1" || out.Model != "vl-model" { + t.Fatalf("unexpected merge: %+v", out) + } + if out.Provider != "claude" { + t.Fatalf("provider: %s", out.Provider) + } + if out.Reasoning.Mode != "off" { + t.Fatalf("reasoning should be off for vision, got %s", out.Reasoning.Mode) + } +} + +func TestVisionConfig_Ready(t *testing.T) { + if (VisionConfig{Enabled: true, Model: "x"}).Ready() != true { + t.Fatal("expected ready") + } + if (VisionConfig{Enabled: true}).Ready() != false { + t.Fatal("expected not ready without model") + } +} diff --git a/internal/knowledge/chunk_eino.go b/internal/knowledge/chunk_eino.go new file mode 100644 index 00000000..6592f350 --- /dev/null +++ b/internal/knowledge/chunk_eino.go @@ -0,0 +1,67 @@ +package knowledge + +import ( + "context" + "fmt" + "strings" + + "github.com/cloudwego/eino-ext/components/document/transformer/splitter/markdown" + "github.com/cloudwego/eino-ext/components/document/transformer/splitter/recursive" + "github.com/cloudwego/eino/components/document" + "github.com/pkoukk/tiktoken-go" +) + +func tokenizerLenFunc(embeddingModel string) func(string) int { + fallback := func(s string) int { + r := []rune(s) + if len(r) == 0 { + return 0 + } + return (len(r) + 3) / 4 + } + m := strings.TrimSpace(embeddingModel) + if m == "" { + return fallback + } + tok, err := tiktoken.EncodingForModel(m) + if err != nil { + return fallback + } + return func(s string) int { + return len(tok.Encode(s, nil, nil)) + } +} + +// newKnowledgeSplitter builds an Eino recursive text splitter. LenFunc uses tiktoken for +// embeddingModel when available, else rune/4 approximation. +func newKnowledgeSplitter(chunkSize, overlap int, embeddingModel string) (document.Transformer, error) { + if chunkSize <= 0 { + return nil, fmt.Errorf("chunk size must be positive") + } + if overlap < 0 { + overlap = 0 + } + return recursive.NewSplitter(context.Background(), &recursive.Config{ + ChunkSize: chunkSize, + OverlapSize: overlap, + LenFunc: tokenizerLenFunc(embeddingModel), + Separators: []string{ + "\n\n", "\n## ", "\n### ", "\n#### ", "\n", + "。", "！", "？", ". ", "? ", "! ", + " ", + }, + }) +} + +// newMarkdownHeaderSplitter Eino-ext Markdown 按标题切分（#～####），适合技术/Markdown 知识库。 +func newMarkdownHeaderSplitter(ctx context.Context) (document.Transformer, error) { + return markdown.NewHeaderSplitter(ctx, &markdown.HeaderConfig{ + Headers: map[string]string{ + "#": "h1", + "##": "h2", + "###": "h3", + "####": "h4", + }, + TrimHeaders: false, + }) +} diff --git a/internal/knowledge/eino_meta.go b/internal/knowledge/eino_meta.go new file mode 100644 index 00000000..0ee7c41b --- /dev/null +++ b/internal/knowledge/eino_meta.go @@ -0,0 +1,129 @@ +package knowledge + +import ( + "fmt" + "strings" +) + +// Document metadata keys for Eino schema.Document flowing through the RAG pipeline. +const ( + metaKBCategory = "kb_category" + metaKBTitle = "kb_title" + metaKBItemID = "kb_item_id" + metaKBChunkIndex = "kb_chunk_index" + metaSimilarity = "similarity" +) + +// DSL keys for [VectorEinoRetriever.Retrieve] via [retriever.WithDSLInfo]. +const ( + DSLRiskType = "risk_type" + DSLSimilarityThreshold = "similarity_threshold" + DSLSubIndexFilter = "sub_index_filter" +) + +// FormatEmbeddingInput matches the historical indexing format so existing embeddings +// stay comparable if users skip reindex; new indexes use the same string shape. +func FormatEmbeddingInput(category, title, chunkText string) string { + return fmt.Sprintf("[风险类型：%s] [标题：%s]\n%s", category, title, chunkText) +} + +// FormatQueryEmbeddingText builds the string embedded at query time so it matches +// [FormatEmbeddingInput] for the same risk category (title left empty for queries). +func FormatQueryEmbeddingText(riskType, query string) string { + q := strings.TrimSpace(query) + rt := strings.TrimSpace(riskType) + if rt != "" { + return FormatEmbeddingInput(rt, "", q) + } + return q +} + +// MetaLookupString returns metadata string value or "" if absent. +func MetaLookupString(md map[string]any, key string) string { + if md == nil { + return "" + } + v, ok := md[key] + if !ok || v == nil { + return "" + } + switch t := v.(type) { + case string: + return t + default: + return strings.TrimSpace(fmt.Sprint(t)) + } +} + +// MetaStringOK returns trimmed non-empty string and true if present and non-empty. +func MetaStringOK(md map[string]any, key string) (string, bool) { + s := strings.TrimSpace(MetaLookupString(md, key)) + if s == "" { + return "", false + } + return s, true +} + +// RequireMetaString requires a non-empty string metadata field. +func RequireMetaString(md map[string]any, key string) (string, error) { + s, ok := MetaStringOK(md, key) + if !ok { + return "", fmt.Errorf("missing or empty metadata %q", key) + } + return s, nil +} + +// RequireMetaInt requires an integer metadata field. +func RequireMetaInt(md map[string]any, key string) (int, error) { + if md == nil { + return 0, fmt.Errorf("missing metadata key %q", key) + } + v, ok := md[key] + if !ok { + return 0, fmt.Errorf("missing metadata key %q", key) + } + switch t := v.(type) { + case int: + return t, nil + case int32: + return int(t), nil + case int64: + return int(t), nil + case float64: + return int(t), nil + default: + return 0, fmt.Errorf("metadata %q: unsupported type %T", key, v) + } +} + +// DSLNumeric coerces DSL map values (e.g. from JSON) to float64. +func DSLNumeric(v any) (float64, bool) { + switch t := v.(type) { + case float64: + return t, true + case float32: + return float64(t), true + case int: + return float64(t), true + case int64: + return float64(t), true + case uint32: + return float64(t), true + case uint64: + return float64(t), true + default: + return 0, false + } +} + +// MetaFloat64OK reads a float metadata value. +func MetaFloat64OK(md map[string]any, key string) (float64, bool) { + if md == nil { + return 0, false + } + v, ok := md[key] + if !ok { + return 0, false + } + return DSLNumeric(v) +} diff --git a/internal/knowledge/eino_meta_test.go b/internal/knowledge/eino_meta_test.go new file mode 100644 index 00000000..ba3f60da --- /dev/null +++ b/internal/knowledge/eino_meta_test.go @@ -0,0 +1,14 @@ +package knowledge + +import "testing" + +func TestFormatQueryEmbeddingText_AlignsWithIndexPrefix(t *testing.T) { + q := FormatQueryEmbeddingText("XSS", "payload") + want := FormatEmbeddingInput("XSS", "", "payload") + if q != want { + t.Fatalf("query embed text mismatch:\n got: %q\nwant: %q", q, want) + } + if FormatQueryEmbeddingText("", "hello") != "hello" { + t.Fatalf("expected bare query without risk type") + } +} diff --git a/internal/knowledge/eino_pipeline_retriever.go b/internal/knowledge/eino_pipeline_retriever.go new file mode 100644 index 00000000..487b439b --- /dev/null +++ b/internal/knowledge/eino_pipeline_retriever.go @@ -0,0 +1,96 @@ +package knowledge + +import ( + "context" + "fmt" + "strings" + + "cyberstrike-ai/internal/config" + + "github.com/cloudwego/eino/callbacks" + "github.com/cloudwego/eino/components" + "github.com/cloudwego/eino/components/retriever" + "github.com/cloudwego/eino/schema" + "go.uber.org/zap" +) + +// knowledgePipelineRetriever: MultiQuery → vector candidates → rerank → post-process. +type knowledgePipelineRetriever struct { + inner retriever.Retriever + base *Retriever +} + +func newKnowledgePipelineRetriever(inner retriever.Retriever, base *Retriever) *knowledgePipelineRetriever { + if inner == nil || base == nil { + return nil + } + return &knowledgePipelineRetriever{inner: inner, base: base} +} + +func (p *knowledgePipelineRetriever) GetType() string { + return "KnowledgeRAGPipeline" +} + +func (p *knowledgePipelineRetriever) Retrieve(ctx context.Context, query string, opts ...retriever.Option) (out []*schema.Document, err error) { + if p == nil || p.inner == nil || p.base == nil { + return nil, fmt.Errorf("knowledge pipeline retriever: nil") + } + q := strings.TrimSpace(query) + if q == "" { + return nil, fmt.Errorf("查询不能为空") + } + + ro := retriever.GetCommonOptions(nil, opts...) + finalTopK := p.base.config.TopK + if finalTopK <= 0 { + finalTopK = 5 + } + if ro.TopK != nil && *ro.TopK > 0 { + finalTopK = *ro.TopK + } + + ctx = callbacks.EnsureRunInfo(ctx, p.GetType(), components.ComponentOfRetriever) + ctx = callbacks.OnStart(ctx, &retriever.CallbackInput{Query: q, TopK: finalTopK, Extra: ro.DSLInfo}) + defer func() { + if err != nil { + _ = callbacks.OnError(ctx, err) + return + } + _ = callbacks.OnEnd(ctx, &retriever.CallbackOutput{Docs: out}) + }() + + out, err = p.inner.Retrieve(ctx, q, opts...) + if err != nil { + return nil, err + } + if len(out) == 0 { + return out, nil + } + + if rr := p.base.documentReranker(); rr != nil && len(out) > 1 { + reranked, rerr := rr.Rerank(ctx, q, out) + if rerr != nil { + if p.base.logger != nil { + p.base.logger.Warn("知识检索重排失败，已使用融合序", zap.Error(rerr)) + } + } else if len(reranked) > 0 { + out = reranked + } + } + + tokenModel := "" + if p.base.embedder != nil { + tokenModel = p.base.embedder.EmbeddingModelName() + } + var postPO *config.PostRetrieveConfig + if p.base.config != nil { + postPO = &p.base.config.PostRetrieve + } + out, err = ApplyPostRetrieve(out, postPO, tokenModel, finalTopK) + if err != nil { + return nil, err + } + return out, nil +} + +var _ retriever.Retriever = (*knowledgePipelineRetriever)(nil) diff --git a/internal/knowledge/eino_retrieve_chain.go b/internal/knowledge/eino_retrieve_chain.go new file mode 100644 index 00000000..81fa4159 --- /dev/null +++ b/internal/knowledge/eino_retrieve_chain.go @@ -0,0 +1,24 @@ +package knowledge + +import ( + "context" + "fmt" + + "github.com/cloudwego/eino/compose" + "github.com/cloudwego/eino/schema" +) + +// BuildKnowledgeRetrieveChain 编译「查询字符串 → 文档列表」的 Eino Chain（MultiQuery → 向量 → 重排 → 后处理）。 +func BuildKnowledgeRetrieveChain(ctx context.Context, r *Retriever) (compose.Runnable[string, []*schema.Document], error) { + if r == nil { + return nil, fmt.Errorf("retriever is nil") + } + ch := compose.NewChain[string, []*schema.Document]() + ch.AppendRetriever(r.AsEinoRetriever()) + return ch.Compile(ctx) +} + +// CompileRetrieveChain 等价于 [BuildKnowledgeRetrieveChain](ctx, r)。 +func (r *Retriever) CompileRetrieveChain(ctx context.Context) (compose.Runnable[string, []*schema.Document], error) { + return BuildKnowledgeRetrieveChain(ctx, r) +} diff --git a/internal/knowledge/eino_retrieve_chain_test.go b/internal/knowledge/eino_retrieve_chain_test.go new file mode 100644 index 00000000..c74a6900 --- /dev/null +++ b/internal/knowledge/eino_retrieve_chain_test.go @@ -0,0 +1,23 @@ +package knowledge + +import ( + "context" + "testing" + + "go.uber.org/zap" +) + +func TestBuildKnowledgeRetrieveChain_Compile(t *testing.T) { + r := NewRetriever(nil, nil, &RetrievalConfig{TopK: 3, SimilarityThreshold: 0.5}, zap.NewNop()) + _, err := BuildKnowledgeRetrieveChain(context.Background(), r) + if err != nil { + t.Fatal(err) + } +} + +func TestBuildKnowledgeRetrieveChain_NilRetriever(t *testing.T) { + _, err := BuildKnowledgeRetrieveChain(context.Background(), nil) + if err == nil { + t.Fatal("expected error for nil retriever") + } +} diff --git a/internal/knowledge/eino_retriever_adapter.go b/internal/knowledge/eino_retriever_adapter.go new file mode 100644 index 00000000..712b4734 --- /dev/null +++ b/internal/knowledge/eino_retriever_adapter.go @@ -0,0 +1,173 @@ +package knowledge + +import ( + "context" + "fmt" + "strings" + + "cyberstrike-ai/internal/config" + + "github.com/cloudwego/eino/callbacks" + "github.com/cloudwego/eino/components" + "github.com/cloudwego/eino/components/retriever" + "github.com/cloudwego/eino/schema" +) + +// VectorEinoRetriever implements [retriever.Retriever] on top of SQLite-stored embeddings + cosine similarity. +// It returns prefetch-sized vector candidates only; rerank and post-process run in [knowledgePipelineRetriever]. +type VectorEinoRetriever struct { + inner *Retriever +} + +// NewVectorEinoRetriever wraps r for Eino compose / tooling. +func NewVectorEinoRetriever(r *Retriever) *VectorEinoRetriever { + if r == nil { + return nil + } + return &VectorEinoRetriever{inner: r} +} + +// GetType identifies this retriever for Eino callbacks. +func (h *VectorEinoRetriever) GetType() string { + return "SQLiteVectorKnowledgeRetriever" +} + +// Retrieve runs vector search and returns [schema.Document] rows. +func (h *VectorEinoRetriever) Retrieve(ctx context.Context, query string, opts ...retriever.Option) (out []*schema.Document, err error) { + if h == nil || h.inner == nil { + return nil, fmt.Errorf("VectorEinoRetriever: nil retriever") + } + q := strings.TrimSpace(query) + if q == "" { + return nil, fmt.Errorf("查询不能为空") + } + + ro := retriever.GetCommonOptions(nil, opts...) + cfg := h.inner.config + + req := &SearchRequest{Query: q} + + if ro.TopK != nil && *ro.TopK > 0 { + req.TopK = *ro.TopK + } else if cfg != nil && cfg.TopK > 0 { + req.TopK = cfg.TopK + } else { + req.TopK = 5 + } + + req.Threshold = 0 + if ro.DSLInfo != nil { + if rt, ok := ro.DSLInfo[DSLRiskType].(string); ok { + req.RiskType = strings.TrimSpace(rt) + } + if v, ok := ro.DSLInfo[DSLSimilarityThreshold]; ok { + if f, ok2 := DSLNumeric(v); ok2 && f > 0 { + req.Threshold = f + } + } + if sf, ok := ro.DSLInfo[DSLSubIndexFilter].(string); ok { + req.SubIndexFilter = strings.TrimSpace(sf) + } + } + if req.SubIndexFilter == "" && cfg != nil && strings.TrimSpace(cfg.SubIndexFilter) != "" { + req.SubIndexFilter = strings.TrimSpace(cfg.SubIndexFilter) + } + if req.Threshold <= 0 && cfg != nil && cfg.SimilarityThreshold > 0 { + req.Threshold = cfg.SimilarityThreshold + } + if req.Threshold <= 0 { + req.Threshold = 0.7 + } + + finalTopK := req.TopK + var postPO *config.PostRetrieveConfig + if cfg != nil { + postPO = &cfg.PostRetrieve + } + fetchK := EffectivePrefetchTopK(finalTopK, postPO) + searchReq := *req + searchReq.TopK = fetchK + + ctx = callbacks.EnsureRunInfo(ctx, h.GetType(), components.ComponentOfRetriever) + th := req.Threshold + st := &th + ctx = callbacks.OnStart(ctx, &retriever.CallbackInput{ + Query: q, + TopK: finalTopK, + ScoreThreshold: st, + Extra: ro.DSLInfo, + }) + defer func() { + if err != nil { + _ = callbacks.OnError(ctx, err) + return + } + _ = callbacks.OnEnd(ctx, &retriever.CallbackOutput{Docs: out}) + }() + + results, err := h.inner.vectorSearch(ctx, &searchReq) + if err != nil { + return nil, err + } + out = retrievalResultsToDocuments(results) + return out, nil +} + +func retrievalResultsToDocuments(results []*RetrievalResult) []*schema.Document { + out := make([]*schema.Document, 0, len(results)) + for _, res := range results { + if res == nil || res.Chunk == nil || res.Item == nil { + continue + } + d := &schema.Document{ + ID: res.Chunk.ID, + Content: res.Chunk.ChunkText, + MetaData: map[string]any{ + metaKBItemID: res.Item.ID, + metaKBCategory: res.Item.Category, + metaKBTitle: res.Item.Title, + metaKBChunkIndex: res.Chunk.ChunkIndex, + metaSimilarity: res.Similarity, + }, + } + d.WithScore(res.Score) + out = append(out, d) + } + return out +} + +func documentsToRetrievalResults(docs []*schema.Document) ([]*RetrievalResult, error) { + out := make([]*RetrievalResult, 0, len(docs)) + for i, d := range docs { + if d == nil { + continue + } + itemID, err := RequireMetaString(d.MetaData, metaKBItemID) + if err != nil { + return nil, fmt.Errorf("document %d: %w", i, err) + } + cat := MetaLookupString(d.MetaData, metaKBCategory) + title := MetaLookupString(d.MetaData, metaKBTitle) + chunkIdx, err := RequireMetaInt(d.MetaData, metaKBChunkIndex) + if err != nil { + return nil, fmt.Errorf("document %d: %w", i, err) + } + sim, _ := MetaFloat64OK(d.MetaData, metaSimilarity) + item := &KnowledgeItem{ID: itemID, Category: cat, Title: title} + chunk := &KnowledgeChunk{ + ID: d.ID, + ItemID: itemID, + ChunkIndex: chunkIdx, + ChunkText: d.Content, + } + out = append(out, &RetrievalResult{ + Chunk: chunk, + Item: item, + Similarity: sim, + Score: d.Score(), + }) + } + return out, nil +} + +var _ retriever.Retriever = (*VectorEinoRetriever)(nil) diff --git a/internal/knowledge/eino_sqlite_indexer.go b/internal/knowledge/eino_sqlite_indexer.go new file mode 100644 index 00000000..a0bbdcdc --- /dev/null +++ b/internal/knowledge/eino_sqlite_indexer.go @@ -0,0 +1,142 @@ +package knowledge + +import ( + "context" + "database/sql" + "encoding/json" + "fmt" + "strings" + + "github.com/cloudwego/eino/callbacks" + "github.com/cloudwego/eino/components" + "github.com/cloudwego/eino/components/indexer" + "github.com/cloudwego/eino/schema" + "github.com/google/uuid" +) + +// SQLiteIndexer implements [indexer.Indexer] against knowledge_embeddings + existing schema. +type SQLiteIndexer struct { + db *sql.DB + batchSize int + embeddingModel string +} + +// NewSQLiteIndexer returns an indexer that writes chunk rows for one knowledge item per Store call. +// batchSize is the embedding batch size; if <= 0, default 64 is used. +// embeddingModel is persisted per row for retrieval-time consistency checks (may be empty). +func NewSQLiteIndexer(db *sql.DB, batchSize int, embeddingModel string) *SQLiteIndexer { + return &SQLiteIndexer{db: db, batchSize: batchSize, embeddingModel: strings.TrimSpace(embeddingModel)} +} + +// GetType implements eino callback run info. +func (s *SQLiteIndexer) GetType() string { + return "SQLiteKnowledgeIndexer" +} + +// Store embeds documents and inserts rows. Each doc must carry MetaData: +// kb_item_id, kb_category, kb_title, kb_chunk_index (int). Content is chunk text only. +func (s *SQLiteIndexer) Store(ctx context.Context, docs []*schema.Document, opts ...indexer.Option) (ids []string, err error) { + options := indexer.GetCommonOptions(nil, opts...) + if options.Embedding == nil { + return nil, fmt.Errorf("sqlite indexer: embedding is required") + } + if len(docs) == 0 { + return nil, nil + } + + ctx = callbacks.EnsureRunInfo(ctx, s.GetType(), components.ComponentOfIndexer) + ctx = callbacks.OnStart(ctx, &indexer.CallbackInput{Docs: docs}) + defer func() { + if err != nil { + _ = callbacks.OnError(ctx, err) + return + } + _ = callbacks.OnEnd(ctx, &indexer.CallbackOutput{IDs: ids}) + }() + + subIdxStr := strings.Join(options.SubIndexes, ",") + + texts := make([]string, len(docs)) + for i, d := range docs { + if d == nil { + return nil, fmt.Errorf("sqlite indexer: nil document at %d", i) + } + cat := MetaLookupString(d.MetaData, metaKBCategory) + title := MetaLookupString(d.MetaData, metaKBTitle) + texts[i] = FormatEmbeddingInput(cat, title, d.Content) + } + + bs := s.batchSize + if bs <= 0 { + bs = 64 + } + + var allVecs [][]float64 + for start := 0; start < len(texts); start += bs { + end := start + bs + if end > len(texts) { + end = len(texts) + } + batch := texts[start:end] + vecs, embedErr := options.Embedding.EmbedStrings(ctx, batch) + if embedErr != nil { + return nil, fmt.Errorf("sqlite indexer: embed batch %d-%d: %w", start, end, embedErr) + } + if len(vecs) != len(batch) { + return nil, fmt.Errorf("sqlite indexer: embed count mismatch: got %d want %d", len(vecs), len(batch)) + } + allVecs = append(allVecs, vecs...) + } + + embedDim := 0 + if len(allVecs) > 0 { + embedDim = len(allVecs[0]) + } + + tx, err := s.db.BeginTx(ctx, nil) + if err != nil { + return nil, fmt.Errorf("sqlite indexer: begin tx: %w", err) + } + defer tx.Rollback() + + ids = make([]string, 0, len(docs)) + for i, d := range docs { + chunkID := uuid.New().String() + itemID, metaErr := RequireMetaString(d.MetaData, metaKBItemID) + if metaErr != nil { + return nil, fmt.Errorf("sqlite indexer: doc %d: %w", i, metaErr) + } + chunkIdx, metaErr := RequireMetaInt(d.MetaData, metaKBChunkIndex) + if metaErr != nil { + return nil, fmt.Errorf("sqlite indexer: doc %d: %w", i, metaErr) + } + vec := allVecs[i] + if embedDim > 0 && len(vec) != embedDim { + return nil, fmt.Errorf("sqlite indexer: inconsistent embedding dim at doc %d: got %d want %d", i, len(vec), embedDim) + } + vec32 := make([]float32, len(vec)) + for j, v := range vec { + vec32[j] = float32(v) + } + embeddingJSON, jsonErr := json.Marshal(vec32) + if jsonErr != nil { + return nil, fmt.Errorf("sqlite indexer: marshal embedding: %w", jsonErr) + } + _, err = tx.ExecContext(ctx, + `INSERT INTO knowledge_embeddings (id, item_id, chunk_index, chunk_text, embedding, sub_indexes, embedding_model, embedding_dim, created_at) + VALUES (?, ?, ?, ?, ?, ?, ?, ?, datetime('now'))`, + chunkID, itemID, chunkIdx, d.Content, string(embeddingJSON), subIdxStr, s.embeddingModel, embedDim, + ) + if err != nil { + return nil, fmt.Errorf("sqlite indexer: insert chunk %d: %w", i, err) + } + ids = append(ids, chunkID) + } + + if err := tx.Commit(); err != nil { + return nil, fmt.Errorf("sqlite indexer: commit: %w", err) + } + return ids, nil +} + +var _ indexer.Indexer = (*SQLiteIndexer)(nil) diff --git a/internal/knowledge/embedder.go b/internal/knowledge/embedder.go new file mode 100644 index 00000000..d9ce8afa --- /dev/null +++ b/internal/knowledge/embedder.go @@ -0,0 +1,251 @@ +package knowledge + +import ( + "context" + "fmt" + "net/http" + "strings" + "sync" + "time" + + "cyberstrike-ai/internal/config" + + einoembedopenai "github.com/cloudwego/eino-ext/components/embedding/openai" + "github.com/cloudwego/eino/components/embedding" + "go.uber.org/zap" + "golang.org/x/time/rate" +) + +// Embedder 使用 CloudWeGo Eino 的 OpenAI Embedding 组件，并保留速率限制与重试。 +type Embedder struct { + eino embedding.Embedder + config *config.KnowledgeConfig + logger *zap.Logger + + rateLimiter *rate.Limiter + rateLimitDelay time.Duration + maxRetries int + retryDelay time.Duration + mu sync.Mutex +} + +// NewEmbedder 基于 Eino eino-ext OpenAI Embedder；openAIConfig 用于在知识库未单独配置 key 时回退 API Key。 +func NewEmbedder(ctx context.Context, cfg *config.KnowledgeConfig, openAIConfig *config.OpenAIConfig, logger *zap.Logger) (*Embedder, error) { + if cfg == nil { + return nil, fmt.Errorf("knowledge config is nil") + } + + var rateLimiter *rate.Limiter + var rateLimitDelay time.Duration + if cfg.Indexing.MaxRPM > 0 { + rpm := cfg.Indexing.MaxRPM + rateLimiter = rate.NewLimiter(rate.Every(time.Minute/time.Duration(rpm)), rpm) + if logger != nil { + logger.Info("知识库索引速率限制已启用", zap.Int("maxRPM", rpm)) + } + } else if cfg.Indexing.RateLimitDelayMs > 0 { + rateLimitDelay = time.Duration(cfg.Indexing.RateLimitDelayMs) * time.Millisecond + if logger != nil { + logger.Info("知识库索引固定延迟已启用", zap.Duration("delay", rateLimitDelay)) + } + } + + maxRetries := 3 + retryDelay := 1000 * time.Millisecond + if cfg.Indexing.MaxRetries > 0 { + maxRetries = cfg.Indexing.MaxRetries + } + if cfg.Indexing.RetryDelayMs > 0 { + retryDelay = time.Duration(cfg.Indexing.RetryDelayMs) * time.Millisecond + } + + model := strings.TrimSpace(cfg.Embedding.Model) + if model == "" { + model = "text-embedding-3-small" + } + + baseURL := strings.TrimSpace(cfg.Embedding.BaseURL) + baseURL = strings.TrimSuffix(baseURL, "/") + if baseURL == "" { + baseURL = "https://api.openai.com/v1" + } + + apiKey := strings.TrimSpace(cfg.Embedding.APIKey) + if apiKey == "" && openAIConfig != nil { + apiKey = strings.TrimSpace(openAIConfig.APIKey) + } + if apiKey == "" { + return nil, fmt.Errorf("embedding API key 未配置") + } + + timeout := 120 * time.Second + if cfg.Indexing.RequestTimeoutSeconds > 0 { + timeout = time.Duration(cfg.Indexing.RequestTimeoutSeconds) * time.Second + } + httpClient := &http.Client{Timeout: timeout} + + inner, err := einoembedopenai.NewEmbedder(ctx, &einoembedopenai.EmbeddingConfig{ + APIKey: apiKey, + BaseURL: baseURL, + ByAzure: false, + Model: model, + HTTPClient: httpClient, + }) + if err != nil { + return nil, fmt.Errorf("eino OpenAI embedder: %w", err) + } + + return &Embedder{ + eino: inner, + config: cfg, + logger: logger, + rateLimiter: rateLimiter, + rateLimitDelay: rateLimitDelay, + maxRetries: maxRetries, + retryDelay: retryDelay, + }, nil +} + +// EmbeddingModelName 返回配置的嵌入模型名（用于 tiktoken 分块与向量行元数据）。 +func (e *Embedder) EmbeddingModelName() string { + if e == nil || e.config == nil { + return "" + } + s := strings.TrimSpace(e.config.Embedding.Model) + if s != "" { + return s + } + return "text-embedding-3-small" +} + +func (e *Embedder) waitRateLimiter() { + e.mu.Lock() + defer e.mu.Unlock() + + if e.rateLimiter != nil { + ctx := context.Background() + if err := e.rateLimiter.Wait(ctx); err != nil && e.logger != nil { + e.logger.Warn("速率限制器等待失败", zap.Error(err)) + } + } + if e.rateLimitDelay > 0 { + time.Sleep(e.rateLimitDelay) + } +} + +// EmbedText 单条嵌入（float32，与历史存储格式一致）。 +func (e *Embedder) EmbedText(ctx context.Context, text string) ([]float32, error) { + vecs, err := e.EmbedStrings(ctx, []string{text}) + if err != nil { + return nil, err + } + if len(vecs) != 1 { + return nil, fmt.Errorf("unexpected embedding count: %d", len(vecs)) + } + return vecs[0], nil +} + +// EmbedStrings 批量嵌入，带重试；实现 [embedding.Embedder]，可供 Eino Indexer 使用。 +func (e *Embedder) EmbedStrings(ctx context.Context, texts []string, opts ...embedding.Option) ([][]float32, error) { + if e == nil || e.eino == nil { + return nil, fmt.Errorf("embedder not initialized") + } + if len(texts) == 0 { + return nil, nil + } + + var lastErr error + for attempt := 0; attempt < e.maxRetries; attempt++ { + if attempt > 0 { + wait := e.retryDelay * time.Duration(attempt) + if e.logger != nil { + e.logger.Debug("嵌入重试前等待", zap.Int("attempt", attempt+1), zap.Duration("wait", wait)) + } + select { + case <-ctx.Done(): + return nil, ctx.Err() + case <-time.After(wait): + } + } else { + e.waitRateLimiter() + } + + raw, err := e.eino.EmbedStrings(ctx, texts, opts...) + if err == nil { + out := make([][]float32, len(raw)) + for i, row := range raw { + out[i] = make([]float32, len(row)) + for j, v := range row { + out[i][j] = float32(v) + } + } + return out, nil + } + lastErr = err + if !e.isRetryableError(err) { + return nil, err + } + if e.logger != nil { + e.logger.Debug("嵌入失败，将重试", zap.Int("attempt", attempt+1), zap.Error(err)) + } + } + return nil, fmt.Errorf("达到最大重试次数 (%d): %v", e.maxRetries, lastErr) +} + +// EmbedTexts 批量 float32 嵌入（兼容旧调用；单次请求批量以减小延迟）。 +func (e *Embedder) EmbedTexts(ctx context.Context, texts []string) ([][]float32, error) { + return e.EmbedStrings(ctx, texts) +} + +func (e *Embedder) isRetryableError(err error) bool { + if err == nil { + return false + } + errStr := err.Error() + if strings.Contains(errStr, "429") || strings.Contains(errStr, "rate limit") { + return true + } + if strings.Contains(errStr, "500") || strings.Contains(errStr, "502") || + strings.Contains(errStr, "503") || strings.Contains(errStr, "504") { + return true + } + if strings.Contains(errStr, "timeout") || strings.Contains(errStr, "connection") || + strings.Contains(errStr, "network") || strings.Contains(errStr, "EOF") { + return true + } + return false +} + +// einoFloatEmbedder adapts [][]float32 embedder to Eino's [][]float64 [embedding.Embedder] for Indexer.Store. +type einoFloatEmbedder struct { + inner *Embedder +} + +func (w *einoFloatEmbedder) EmbedStrings(ctx context.Context, texts []string, opts ...embedding.Option) ([][]float64, error) { + vec32, err := w.inner.EmbedStrings(ctx, texts, opts...) + if err != nil { + return nil, err + } + out := make([][]float64, len(vec32)) + for i, row := range vec32 { + out[i] = make([]float64, len(row)) + for j, v := range row { + out[i][j] = float64(v) + } + } + return out, nil +} + +func (w *einoFloatEmbedder) GetType() string { + return "CyberStrikeKnowledgeEmbedder" +} + +func (w *einoFloatEmbedder) IsCallbacksEnabled() bool { + return false +} + +// EinoEmbeddingComponent returns an [embedding.Embedder] that uses the same retry/rate-limit path +// and produces float64 vectors expected by generic Eino indexer helpers. +func (e *Embedder) EinoEmbeddingComponent() embedding.Embedder { + return &einoFloatEmbedder{inner: e} +} diff --git a/internal/knowledge/index_pipeline.go b/internal/knowledge/index_pipeline.go new file mode 100644 index 00000000..a9b9a4c4 --- /dev/null +++ b/internal/knowledge/index_pipeline.go @@ -0,0 +1,91 @@ +package knowledge + +import ( + "context" + "database/sql" + "fmt" + "strings" + + "cyberstrike-ai/internal/config" + + "github.com/cloudwego/eino/components/document" + "github.com/cloudwego/eino/compose" + "github.com/cloudwego/eino/schema" +) + +// normalizeChunkStrategy returns "recursive" or "markdown_then_recursive". +func normalizeChunkStrategy(s string) string { + v := strings.TrimSpace(strings.ToLower(s)) + switch v { + case "recursive": + return "recursive" + case "markdown_then_recursive", "markdown_recursive", "markdown": + return "markdown_then_recursive" + case "": + return "markdown_then_recursive" + default: + return "markdown_then_recursive" + } +} + +func buildKnowledgeIndexChain( + ctx context.Context, + indexingCfg *config.IndexingConfig, + db *sql.DB, + recursive document.Transformer, + embeddingModel string, +) (compose.Runnable[[]*schema.Document, []string], error) { + if recursive == nil { + return nil, fmt.Errorf("recursive transformer is nil") + } + if db == nil { + return nil, fmt.Errorf("db is nil") + } + strategy := normalizeChunkStrategy("markdown_then_recursive") + batch := 64 + maxChunks := 0 + if indexingCfg != nil { + strategy = normalizeChunkStrategy(indexingCfg.ChunkStrategy) + if indexingCfg.BatchSize > 0 { + batch = indexingCfg.BatchSize + } + maxChunks = indexingCfg.MaxChunksPerItem + } + + si := NewSQLiteIndexer(db, batch, embeddingModel) + ch := compose.NewChain[[]*schema.Document, []string]() + if strategy != "recursive" { + md, err := newMarkdownHeaderSplitter(ctx) + if err != nil { + return nil, fmt.Errorf("markdown splitter: %w", err) + } + ch.AppendDocumentTransformer(md) + } + ch.AppendDocumentTransformer(recursive) + ch.AppendLambda(newChunkEnrichLambda(maxChunks)) + ch.AppendIndexer(si) + return ch.Compile(ctx) +} + +func newChunkEnrichLambda(maxChunks int) *compose.Lambda { + return compose.InvokableLambda(func(ctx context.Context, docs []*schema.Document) ([]*schema.Document, error) { + _ = ctx + out := make([]*schema.Document, 0, len(docs)) + for _, d := range docs { + if d == nil || strings.TrimSpace(d.Content) == "" { + continue + } + out = append(out, d) + } + if maxChunks > 0 && len(out) > maxChunks { + out = out[:maxChunks] + } + for i, d := range out { + if d.MetaData == nil { + d.MetaData = make(map[string]any) + } + d.MetaData[metaKBChunkIndex] = i + } + return out, nil + }) +} diff --git a/internal/knowledge/index_pipeline_test.go b/internal/knowledge/index_pipeline_test.go new file mode 100644 index 00000000..9e4b03fa --- /dev/null +++ b/internal/knowledge/index_pipeline_test.go @@ -0,0 +1,21 @@ +package knowledge + +import "testing" + +func TestNormalizeChunkStrategy(t *testing.T) { + cases := []struct { + in, want string + }{ + {"", "markdown_then_recursive"}, + {"recursive", "recursive"}, + {"RECURSIVE", "recursive"}, + {"markdown_then_recursive", "markdown_then_recursive"}, + {"markdown", "markdown_then_recursive"}, + {"unknown", "markdown_then_recursive"}, + } + for _, tc := range cases { + if got := normalizeChunkStrategy(tc.in); got != tc.want { + t.Errorf("normalizeChunkStrategy(%q) = %q, want %q", tc.in, got, tc.want) + } + } +} diff --git a/internal/knowledge/indexer.go b/internal/knowledge/indexer.go new file mode 100644 index 00000000..1c1e7cbe --- /dev/null +++ b/internal/knowledge/indexer.go @@ -0,0 +1,435 @@ +package knowledge + +import ( + "context" + "database/sql" + "fmt" + "strings" + "sync" + "time" + + "cyberstrike-ai/internal/config" + + fileloader "github.com/cloudwego/eino-ext/components/document/loader/file" + "github.com/cloudwego/eino/components/document" + "github.com/cloudwego/eino/components/indexer" + "github.com/cloudwego/eino/compose" + "github.com/cloudwego/eino/schema" + "go.uber.org/zap" +) + +// Indexer 使用 Eino Compose 索引链（Markdown/递归分块、Lambda enrich、SQLite 索引）与嵌入写入。 +type Indexer struct { + db *sql.DB + embedder *Embedder + logger *zap.Logger + chunkSize int + overlap int + indexingCfg *config.IndexingConfig + + indexChain compose.Runnable[[]*schema.Document, []string] + fileLoader *fileloader.FileLoader + + mu sync.RWMutex + lastError string + lastErrorTime time.Time + errorCount int + + rebuildMu sync.RWMutex + isRebuilding bool + rebuildTotalItems int + rebuildCurrent int + rebuildFailed int + rebuildStartTime time.Time + rebuildLastItemID string + rebuildLastChunks int +} + +// NewIndexer 创建索引器并编译 Eino 索引链；kcfg 为完整知识库配置（含 indexing 与路径相关行为）。 +func NewIndexer(ctx context.Context, db *sql.DB, embedder *Embedder, logger *zap.Logger, kcfg *config.KnowledgeConfig) (*Indexer, error) { + if db == nil { + return nil, fmt.Errorf("db is nil") + } + if embedder == nil { + return nil, fmt.Errorf("embedder is nil") + } + if err := EnsureKnowledgeEmbeddingsSchema(db); err != nil { + return nil, fmt.Errorf("knowledge_embeddings 结构迁移: %w", err) + } + if kcfg == nil { + kcfg = &config.KnowledgeConfig{} + } + indexingCfg := &kcfg.Indexing + + chunkSize := 512 + overlap := 50 + if indexingCfg.ChunkSize > 0 { + chunkSize = indexingCfg.ChunkSize + } + if indexingCfg.ChunkOverlap >= 0 { + overlap = indexingCfg.ChunkOverlap + } + + embedModel := embedder.EmbeddingModelName() + splitter, err := newKnowledgeSplitter(chunkSize, overlap, embedModel) + if err != nil { + return nil, fmt.Errorf("eino recursive splitter: %w", err) + } + + chain, err := buildKnowledgeIndexChain(ctx, indexingCfg, db, splitter, embedModel) + if err != nil { + return nil, fmt.Errorf("knowledge index chain: %w", err) + } + + var fl *fileloader.FileLoader + fl, err = fileloader.NewFileLoader(ctx, nil) + if err != nil { + if logger != nil { + logger.Warn("Eino FileLoader 初始化失败，prefer_source_file 将回退数据库正文", zap.Error(err)) + } + fl = nil + err = nil + } + + return &Indexer{ + db: db, + embedder: embedder, + logger: logger, + chunkSize: chunkSize, + overlap: overlap, + indexingCfg: indexingCfg, + indexChain: chain, + fileLoader: fl, + }, nil +} + +// RecompileIndexChain 在配置或嵌入模型变更后重建 Eino 索引链（无需重启进程）。 +func (idx *Indexer) RecompileIndexChain(ctx context.Context) error { + if idx == nil || idx.db == nil || idx.embedder == nil { + return fmt.Errorf("indexer 未初始化") + } + if err := EnsureKnowledgeEmbeddingsSchema(idx.db); err != nil { + return err + } + embedModel := idx.embedder.EmbeddingModelName() + splitter, err := newKnowledgeSplitter(idx.chunkSize, idx.overlap, embedModel) + if err != nil { + return fmt.Errorf("eino recursive splitter: %w", err) + } + chain, err := buildKnowledgeIndexChain(ctx, idx.indexingCfg, idx.db, splitter, embedModel) + if err != nil { + return fmt.Errorf("knowledge index chain: %w", err) + } + idx.indexChain = chain + return nil +} + +// IndexItem 索引单个知识项：先清空旧向量，再走 Compose 链（分块、嵌入、写入）。 +func (idx *Indexer) IndexItem(ctx context.Context, itemID string) error { + if idx.indexChain == nil { + return fmt.Errorf("索引链未初始化") + } + if idx.embedder == nil { + return fmt.Errorf("嵌入器未初始化") + } + + var content, category, title, filePath string + err := idx.db.QueryRow("SELECT content, category, title, file_path FROM knowledge_base_items WHERE id = ?", itemID).Scan(&content, &category, &title, &filePath) + if err != nil { + return fmt.Errorf("获取知识项失败：%w", err) + } + + if _, err := idx.db.Exec("DELETE FROM knowledge_embeddings WHERE item_id = ?", itemID); err != nil { + return fmt.Errorf("删除旧向量失败：%w", err) + } + + body := strings.TrimSpace(content) + if idx.indexingCfg != nil && idx.indexingCfg.PreferSourceFile && strings.TrimSpace(filePath) != "" && idx.fileLoader != nil { + docs, lerr := idx.fileLoader.Load(ctx, document.Source{URI: strings.TrimSpace(filePath)}) + if lerr == nil && len(docs) > 0 { + var b strings.Builder + for i, d := range docs { + if d == nil { + continue + } + if i > 0 { + b.WriteString("\n\n") + } + b.WriteString(d.Content) + } + if s := strings.TrimSpace(b.String()); s != "" { + body = s + } + } else if idx.logger != nil { + idx.logger.Warn("优先源文件读取失败，使用数据库正文", + zap.String("itemId", itemID), + zap.String("path", filePath), + zap.Error(lerr)) + } + } + + root := &schema.Document{ + ID: itemID, + Content: body, + MetaData: map[string]any{ + metaKBCategory: category, + metaKBTitle: title, + metaKBItemID: itemID, + }, + } + + idxOpts := []indexer.Option{indexer.WithEmbedding(idx.embedder.EinoEmbeddingComponent())} + if idx.indexingCfg != nil && len(idx.indexingCfg.SubIndexes) > 0 { + idxOpts = append(idxOpts, indexer.WithSubIndexes(idx.indexingCfg.SubIndexes)) + } + + ids, err := idx.indexChain.Invoke(ctx, []*schema.Document{root}, compose.WithIndexerOption(idxOpts...)) + if err != nil { + msg := fmt.Sprintf("索引写入失败 (知识项：%s): %v", itemID, err) + idx.mu.Lock() + idx.lastError = msg + idx.lastErrorTime = time.Now() + idx.mu.Unlock() + return err + } + + if idx.logger != nil { + idx.logger.Info("知识项索引完成", zap.String("itemId", itemID), zap.Int("chunks", len(ids))) + } + idx.rebuildMu.Lock() + idx.rebuildLastItemID = itemID + idx.rebuildLastChunks = len(ids) + idx.rebuildMu.Unlock() + return nil +} + +// HasIndex 检查是否存在索引 +func (idx *Indexer) HasIndex() (bool, error) { + var count int + err := idx.db.QueryRow("SELECT COUNT(*) FROM knowledge_embeddings").Scan(&count) + if err != nil { + return false, fmt.Errorf("检查索引失败：%w", err) + } + return count > 0, nil +} + +func (idx *Indexer) beginIndexRun() error { + idx.rebuildMu.Lock() + defer idx.rebuildMu.Unlock() + + if idx.isRebuilding { + return fmt.Errorf("索引任务已在进行中") + } + idx.isRebuilding = true + idx.rebuildTotalItems = 0 + idx.rebuildCurrent = 0 + idx.rebuildFailed = 0 + idx.rebuildStartTime = time.Now() + idx.rebuildLastItemID = "" + idx.rebuildLastChunks = 0 + return nil +} + +// TryBeginIndexRun 同步占用索引任务槽位；调用方必须在后台任务结束时调用 FinishIndexRun。 +func (idx *Indexer) TryBeginIndexRun() error { + return idx.beginIndexRun() +} + +func (idx *Indexer) FinishIndexRun() { + idx.rebuildMu.Lock() + idx.isRebuilding = false + idx.rebuildMu.Unlock() +} + +func (idx *Indexer) resetLastError() { + idx.mu.Lock() + idx.lastError = "" + idx.lastErrorTime = time.Time{} + idx.errorCount = 0 + idx.mu.Unlock() +} + +func (idx *Indexer) setIndexRunTotal(total int) { + idx.rebuildMu.Lock() + idx.rebuildTotalItems = total + idx.rebuildMu.Unlock() +} + +// IndexMissing 为尚无向量的知识项构建索引（默认推荐路径，适合冷启动与中断续跑）。 +func (idx *Indexer) IndexMissing(ctx context.Context) error { + if err := idx.beginIndexRun(); err != nil { + return err + } + defer idx.FinishIndexRun() + return idx.runIndexMissing(ctx) +} + +// RebuildIndex 全量重建所有知识项索引（显式 opt-in，成本更高）。 +func (idx *Indexer) RebuildIndex(ctx context.Context) error { + if err := idx.beginIndexRun(); err != nil { + return err + } + defer idx.FinishIndexRun() + return idx.runRebuildIndex(ctx) +} + +// RunRebuildIndex 在已占用索引任务槽位后执行全量重建（供 HTTP handler 后台任务使用）。 +func (idx *Indexer) RunRebuildIndex(ctx context.Context) error { + return idx.runRebuildIndex(ctx) +} + +// RunIndexMissing 在已占用索引任务槽位后执行缺失索引补齐（供 HTTP handler 后台任务使用）。 +func (idx *Indexer) RunIndexMissing(ctx context.Context) error { + return idx.runIndexMissing(ctx) +} + +func (idx *Indexer) runRebuildIndex(ctx context.Context) error { + idx.resetLastError() + + rows, err := idx.db.QueryContext(ctx, "SELECT id FROM knowledge_base_items ORDER BY updated_at ASC, id ASC") + if err != nil { + return fmt.Errorf("查询知识项失败：%w", err) + } + defer rows.Close() + + itemIDs, err := scanKnowledgeItemIDs(rows) + if err != nil { + return err + } + + idx.setIndexRunTotal(len(itemIDs)) + idx.logger.Info("开始重建索引", zap.Int("totalItems", len(itemIDs))) + + return idx.indexItemIDs(ctx, itemIDs, "索引重建完成") +} + +func (idx *Indexer) runIndexMissing(ctx context.Context) error { + idx.resetLastError() + + rows, err := idx.db.QueryContext(ctx, ` + SELECT i.id + FROM knowledge_base_items i + LEFT JOIN knowledge_embeddings e ON e.item_id = i.id + WHERE e.item_id IS NULL + ORDER BY i.updated_at ASC, i.id ASC + `) + if err != nil { + return fmt.Errorf("查询未索引知识项失败：%w", err) + } + defer rows.Close() + + itemIDs, err := scanKnowledgeItemIDs(rows) + if err != nil { + return fmt.Errorf("扫描未索引知识项 ID 失败：%w", err) + } + + idx.setIndexRunTotal(len(itemIDs)) + idx.logger.Info("开始补齐缺失索引", zap.Int("totalItems", len(itemIDs))) + + return idx.indexItemIDs(ctx, itemIDs, "索引构建完成") +} + +func scanKnowledgeItemIDs(rows *sql.Rows) ([]string, error) { + var itemIDs []string + for rows.Next() { + var id string + if err := rows.Scan(&id); err != nil { + return nil, fmt.Errorf("扫描知识项 ID 失败：%w", err) + } + itemIDs = append(itemIDs, id) + } + if err := rows.Err(); err != nil { + return nil, fmt.Errorf("扫描知识项 ID 失败：%w", err) + } + return itemIDs, nil +} + +func (idx *Indexer) indexItemIDs(ctx context.Context, itemIDs []string, doneMessage string) error { + failedCount := 0 + consecutiveFailures := 0 + maxConsecutiveFailures := 5 + firstFailureItemID := "" + var firstFailureError error + + for i, itemID := range itemIDs { + if err := idx.IndexItem(ctx, itemID); err != nil { + failedCount++ + consecutiveFailures++ + + if consecutiveFailures == 1 { + firstFailureItemID = itemID + firstFailureError = err + idx.logger.Warn("索引知识项失败", + zap.String("itemId", itemID), + zap.Int("totalItems", len(itemIDs)), + zap.Error(err), + ) + } + + if consecutiveFailures >= maxConsecutiveFailures { + errorMsg := fmt.Sprintf("连续 %d 个知识项索引失败，可能存在配置问题（如嵌入模型配置错误、API 密钥无效、余额不足等）。第一个失败项：%s, 错误：%v", consecutiveFailures, firstFailureItemID, firstFailureError) + idx.mu.Lock() + idx.lastError = errorMsg + idx.lastErrorTime = time.Now() + idx.mu.Unlock() + + idx.logger.Error("连续索引失败次数过多，立即停止索引", + zap.Int("consecutiveFailures", consecutiveFailures), + zap.Int("totalItems", len(itemIDs)), + zap.Int("processedItems", i+1), + zap.String("firstFailureItemId", firstFailureItemID), + zap.Error(firstFailureError), + ) + return fmt.Errorf("连续索引失败次数过多：%v", firstFailureError) + } + + if failedCount > len(itemIDs)*3/10 && failedCount == len(itemIDs)*3/10+1 { + errorMsg := fmt.Sprintf("索引失败的知识项过多 (%d/%d)，可能存在配置问题。第一个失败项：%s, 错误：%v", failedCount, len(itemIDs), firstFailureItemID, firstFailureError) + idx.mu.Lock() + idx.lastError = errorMsg + idx.lastErrorTime = time.Now() + idx.mu.Unlock() + + idx.logger.Error("索引失败的知识项过多，可能存在配置问题", + zap.Int("failedCount", failedCount), + zap.Int("totalItems", len(itemIDs)), + zap.String("firstFailureItemId", firstFailureItemID), + zap.Error(firstFailureError), + ) + } + continue + } + + if consecutiveFailures > 0 { + consecutiveFailures = 0 + firstFailureItemID = "" + firstFailureError = nil + } + + idx.rebuildMu.Lock() + idx.rebuildCurrent = i + 1 + idx.rebuildFailed = failedCount + idx.rebuildMu.Unlock() + + if (i+1)%10 == 0 || (len(itemIDs) > 0 && (i+1)*100/len(itemIDs)%10 == 0 && (i+1)*100/len(itemIDs) > 0) { + idx.logger.Info("索引进度", zap.Int("current", i+1), zap.Int("total", len(itemIDs)), zap.Int("failed", failedCount)) + } + } + + idx.logger.Info(doneMessage, zap.Int("totalItems", len(itemIDs)), zap.Int("failedCount", failedCount)) + return nil +} + +// GetLastError 获取最近一次错误信息 +func (idx *Indexer) GetLastError() (string, time.Time) { + idx.mu.RLock() + defer idx.mu.RUnlock() + return idx.lastError, idx.lastErrorTime +} + +// GetRebuildStatus 获取重建索引状态 +func (idx *Indexer) GetRebuildStatus() (isRebuilding bool, totalItems int, current int, failed int, lastItemID string, lastChunks int, startTime time.Time) { + idx.rebuildMu.RLock() + defer idx.rebuildMu.RUnlock() + return idx.isRebuilding, idx.rebuildTotalItems, idx.rebuildCurrent, idx.rebuildFailed, idx.rebuildLastItemID, idx.rebuildLastChunks, idx.rebuildStartTime +} diff --git a/internal/knowledge/indexer_rebuild_state_test.go b/internal/knowledge/indexer_rebuild_state_test.go new file mode 100644 index 00000000..1e04480e --- /dev/null +++ b/internal/knowledge/indexer_rebuild_state_test.go @@ -0,0 +1,20 @@ +package knowledge + +import "testing" + +func TestIndexerRejectsConcurrentIndexRuns(t *testing.T) { + idx := &Indexer{} + + if err := idx.beginIndexRun(); err != nil { + t.Fatalf("first index run should start: %v", err) + } + if err := idx.beginIndexRun(); err == nil { + t.Fatal("second index run should be rejected while one is active") + } + + idx.FinishIndexRun() + if err := idx.beginIndexRun(); err != nil { + t.Fatalf("index run should start again after finish: %v", err) + } + idx.FinishIndexRun() +} diff --git a/internal/knowledge/manager.go b/internal/knowledge/manager.go new file mode 100644 index 00000000..7309cc2a --- /dev/null +++ b/internal/knowledge/manager.go @@ -0,0 +1,885 @@ +package knowledge + +import ( + "database/sql" + "encoding/json" + "fmt" + "io/fs" + "os" + "path/filepath" + "strings" + "time" + + "github.com/google/uuid" + "go.uber.org/zap" +) + +// Manager 知识库管理器 +type Manager struct { + db *sql.DB + basePath string + logger *zap.Logger +} + +// NewManager 创建新的知识库管理器 +func NewManager(db *sql.DB, basePath string, logger *zap.Logger) *Manager { + return &Manager{ + db: db, + basePath: basePath, + logger: logger, + } +} + +// ScanKnowledgeBase 扫描知识库目录，更新数据库 +// 返回需要索引的知识项ID列表（新添加的或更新的） +func (m *Manager) ScanKnowledgeBase() ([]string, error) { + if m.basePath == "" { + return nil, fmt.Errorf("知识库路径未配置") + } + + // 确保目录存在 + if err := os.MkdirAll(m.basePath, 0755); err != nil { + return nil, fmt.Errorf("创建知识库目录失败: %w", err) + } + + var itemsToIndex []string + + // 遍历知识库目录 + err := filepath.WalkDir(m.basePath, func(path string, d fs.DirEntry, err error) error { + if err != nil { + return err + } + + // 跳过目录和非markdown文件 + if d.IsDir() || !strings.HasSuffix(strings.ToLower(path), ".md") { + return nil + } + + // 计算相对路径和分类 + relPath, err := filepath.Rel(m.basePath, path) + if err != nil { + return err + } + + // 第一个目录名作为分类（风险类型） + parts := strings.Split(relPath, string(filepath.Separator)) + category := "未分类" + if len(parts) > 1 { + category = parts[0] + } + + // 文件名为标题 + title := strings.TrimSuffix(filepath.Base(path), ".md") + + // 读取文件内容 + content, err := os.ReadFile(path) + if err != nil { + m.logger.Warn("读取知识库文件失败", zap.String("path", path), zap.Error(err)) + return nil // 继续处理其他文件 + } + + // 检查是否已存在 + var existingID string + var existingContent string + var existingUpdatedAt time.Time + err = m.db.QueryRow( + "SELECT id, content, updated_at FROM knowledge_base_items WHERE file_path = ?", + path, + ).Scan(&existingID, &existingContent, &existingUpdatedAt) + + if err == sql.ErrNoRows { + // 创建新项 + id := uuid.New().String() + now := time.Now() + _, err = m.db.Exec( + "INSERT INTO knowledge_base_items (id, category, title, file_path, content, created_at, updated_at) VALUES (?, ?, ?, ?, ?, ?, ?)", + id, category, title, path, string(content), now, now, + ) + if err != nil { + return fmt.Errorf("插入知识项失败: %w", err) + } + m.logger.Info("添加知识项", zap.String("id", id), zap.String("title", title), zap.String("category", category)) + // 新添加的项需要索引 + itemsToIndex = append(itemsToIndex, id) + } else if err == nil { + // 检查内容是否有变化 + contentChanged := existingContent != string(content) + if contentChanged { + // 更新现有项 + _, err = m.db.Exec( + "UPDATE knowledge_base_items SET category = ?, title = ?, content = ?, updated_at = ? WHERE id = ?", + category, title, string(content), time.Now(), existingID, + ) + if err != nil { + return fmt.Errorf("更新知识项失败: %w", err) + } + m.logger.Info("更新知识项", zap.String("id", existingID), zap.String("title", title)) + // 内容已更新的项需要重新索引 + itemsToIndex = append(itemsToIndex, existingID) + } else { + m.logger.Debug("知识项未变化，跳过", zap.String("id", existingID), zap.String("title", title)) + } + } else { + return fmt.Errorf("查询知识项失败: %w", err) + } + + return nil + }) + + if err != nil { + return nil, err + } + + return itemsToIndex, nil +} + +// GetCategories 获取所有分类（风险类型） +func (m *Manager) GetCategories() ([]string, error) { + rows, err := m.db.Query("SELECT DISTINCT category FROM knowledge_base_items ORDER BY category") + if err != nil { + return nil, fmt.Errorf("查询分类失败: %w", err) + } + defer rows.Close() + + var categories []string + for rows.Next() { + var category string + if err := rows.Scan(&category); err != nil { + return nil, fmt.Errorf("扫描分类失败: %w", err) + } + categories = append(categories, category) + } + + return categories, nil +} + +// GetStats 获取知识库统计信息 +func (m *Manager) GetStats() (int, int, error) { + // 获取分类总数 + categories, err := m.GetCategories() + if err != nil { + return 0, 0, fmt.Errorf("获取分类失败: %w", err) + } + totalCategories := len(categories) + + // 获取知识项总数 + var totalItems int + err = m.db.QueryRow("SELECT COUNT(*) FROM knowledge_base_items").Scan(&totalItems) + if err != nil { + return totalCategories, 0, fmt.Errorf("获取知识项总数失败: %w", err) + } + + return totalCategories, totalItems, nil +} + +// GetCategoriesWithItems 按分类分页获取知识项（每个分类包含其下的所有知识项） +// limit: 每页分类数量（0表示不限制） +// offset: 偏移量（按分类偏移） +func (m *Manager) GetCategoriesWithItems(limit, offset int) ([]*CategoryWithItems, int, error) { + // 首先获取所有分类（带数量统计） + rows, err := m.db.Query(` + SELECT category, COUNT(*) as item_count + FROM knowledge_base_items + GROUP BY category + ORDER BY category + `) + if err != nil { + return nil, 0, fmt.Errorf("查询分类失败: %w", err) + } + defer rows.Close() + + // 收集所有分类信息 + type categoryInfo struct { + name string + itemCount int + } + var allCategories []categoryInfo + for rows.Next() { + var info categoryInfo + if err := rows.Scan(&info.name, &info.itemCount); err != nil { + return nil, 0, fmt.Errorf("扫描分类失败: %w", err) + } + allCategories = append(allCategories, info) + } + + totalCategories := len(allCategories) + + // 应用分页（按分类分页） + var paginatedCategories []categoryInfo + if limit > 0 { + start := offset + end := offset + limit + if start >= totalCategories { + paginatedCategories = []categoryInfo{} + } else { + if end > totalCategories { + end = totalCategories + } + paginatedCategories = allCategories[start:end] + } + } else { + paginatedCategories = allCategories + } + + // 为每个分类获取其下的知识项（只返回摘要，不包含完整内容） + result := make([]*CategoryWithItems, 0, len(paginatedCategories)) + for _, catInfo := range paginatedCategories { + // 获取该分类下的所有知识项 + items, _, err := m.GetItemsSummary(catInfo.name, 0, 0) + if err != nil { + return nil, 0, fmt.Errorf("获取分类 %s 的知识项失败: %w", catInfo.name, err) + } + + result = append(result, &CategoryWithItems{ + Category: catInfo.name, + ItemCount: catInfo.itemCount, + Items: items, + }) + } + + return result, totalCategories, nil +} + +// GetItems 获取知识项列表（完整内容，用于向后兼容） +func (m *Manager) GetItems(category string) ([]*KnowledgeItem, error) { + return m.GetItemsWithOptions(category, 0, 0, true) +} + +// GetItemsWithOptions 获取知识项列表（支持分页和可选内容） +// category: 分类筛选（空字符串表示所有分类） +// limit: 每页数量（0表示不限制） +// offset: 偏移量 +// includeContent: 是否包含完整内容（false时只返回摘要） +func (m *Manager) GetItemsWithOptions(category string, limit, offset int, includeContent bool) ([]*KnowledgeItem, error) { + var rows *sql.Rows + var err error + + // 构建SQL查询 + var query string + var args []interface{} + + if includeContent { + query = "SELECT id, category, title, file_path, content, created_at, updated_at FROM knowledge_base_items" + } else { + query = "SELECT id, category, title, file_path, created_at, updated_at FROM knowledge_base_items" + } + + if category != "" { + query += " WHERE category = ?" + args = append(args, category) + } + + query += " ORDER BY category, title" + + if limit > 0 { + query += " LIMIT ?" + args = append(args, limit) + if offset > 0 { + query += " OFFSET ?" + args = append(args, offset) + } + } + + rows, err = m.db.Query(query, args...) + if err != nil { + return nil, fmt.Errorf("查询知识项失败: %w", err) + } + defer rows.Close() + + var items []*KnowledgeItem + for rows.Next() { + item := &KnowledgeItem{} + var createdAt, updatedAt string + + if includeContent { + if err := rows.Scan(&item.ID, &item.Category, &item.Title, &item.FilePath, &item.Content, &createdAt, &updatedAt); err != nil { + return nil, fmt.Errorf("扫描知识项失败: %w", err) + } + } else { + if err := rows.Scan(&item.ID, &item.Category, &item.Title, &item.FilePath, &createdAt, &updatedAt); err != nil { + return nil, fmt.Errorf("扫描知识项失败: %w", err) + } + // 不包含内容时，Content为空字符串 + item.Content = "" + } + + // 解析时间 - 支持多种格式 + timeFormats := []string{ + "2006-01-02 15:04:05.999999999-07:00", + "2006-01-02 15:04:05.999999999", + "2006-01-02T15:04:05.999999999Z07:00", + "2006-01-02T15:04:05Z", + "2006-01-02 15:04:05", + time.RFC3339, + time.RFC3339Nano, + } + + // 解析创建时间 + if createdAt != "" { + for _, format := range timeFormats { + parsed, err := time.Parse(format, createdAt) + if err == nil && !parsed.IsZero() { + item.CreatedAt = parsed + break + } + } + } + + // 解析更新时间 + if updatedAt != "" { + for _, format := range timeFormats { + parsed, err := time.Parse(format, updatedAt) + if err == nil && !parsed.IsZero() { + item.UpdatedAt = parsed + break + } + } + } + + // 如果更新时间为空，使用创建时间 + if item.UpdatedAt.IsZero() && !item.CreatedAt.IsZero() { + item.UpdatedAt = item.CreatedAt + } + + items = append(items, item) + } + + return items, nil +} + +// GetItemsCount 获取知识项总数 +func (m *Manager) GetItemsCount(category string) (int, error) { + var count int + var err error + + if category != "" { + err = m.db.QueryRow("SELECT COUNT(*) FROM knowledge_base_items WHERE category = ?", category).Scan(&count) + } else { + err = m.db.QueryRow("SELECT COUNT(*) FROM knowledge_base_items").Scan(&count) + } + + if err != nil { + return 0, fmt.Errorf("查询知识项总数失败: %w", err) + } + + return count, nil +} + +// SearchItemsByKeyword 按关键字搜索知识项（在所有数据中搜索，支持标题、分类、路径、内容匹配） +func (m *Manager) SearchItemsByKeyword(keyword string, category string) ([]*KnowledgeItemSummary, error) { + if keyword == "" { + return nil, fmt.Errorf("搜索关键字不能为空") + } + + // 构建SQL查询，使用LIKE进行关键字匹配（不区分大小写） + var query string + var args []interface{} + + // SQLite的LIKE不区分大小写，使用COLLATE NOCASE或LOWER()函数 + // 使用%keyword%进行模糊匹配 + searchPattern := "%" + keyword + "%" + + query = ` + SELECT id, category, title, file_path, created_at, updated_at + FROM knowledge_base_items + WHERE (LOWER(title) LIKE LOWER(?) OR LOWER(category) LIKE LOWER(?) OR LOWER(file_path) LIKE LOWER(?) OR LOWER(content) LIKE LOWER(?)) + ` + args = append(args, searchPattern, searchPattern, searchPattern, searchPattern) + + // 如果指定了分类，添加分类过滤 + if category != "" { + query += " AND category = ?" + args = append(args, category) + } + + query += " ORDER BY category, title" + + rows, err := m.db.Query(query, args...) + if err != nil { + return nil, fmt.Errorf("搜索知识项失败: %w", err) + } + defer rows.Close() + + var items []*KnowledgeItemSummary + for rows.Next() { + item := &KnowledgeItemSummary{} + var createdAt, updatedAt string + + if err := rows.Scan(&item.ID, &item.Category, &item.Title, &item.FilePath, &createdAt, &updatedAt); err != nil { + return nil, fmt.Errorf("扫描知识项失败: %w", err) + } + + // 解析时间 + timeFormats := []string{ + "2006-01-02 15:04:05.999999999-07:00", + "2006-01-02 15:04:05.999999999", + "2006-01-02T15:04:05.999999999Z07:00", + "2006-01-02T15:04:05Z", + "2006-01-02 15:04:05", + time.RFC3339, + time.RFC3339Nano, + } + + if createdAt != "" { + for _, format := range timeFormats { + parsed, err := time.Parse(format, createdAt) + if err == nil && !parsed.IsZero() { + item.CreatedAt = parsed + break + } + } + } + + if updatedAt != "" { + for _, format := range timeFormats { + parsed, err := time.Parse(format, updatedAt) + if err == nil && !parsed.IsZero() { + item.UpdatedAt = parsed + break + } + } + } + + if item.UpdatedAt.IsZero() && !item.CreatedAt.IsZero() { + item.UpdatedAt = item.CreatedAt + } + + items = append(items, item) + } + + return items, nil +} + +// GetItemsSummary 获取知识项摘要列表（不包含完整内容，支持分页） +func (m *Manager) GetItemsSummary(category string, limit, offset int) ([]*KnowledgeItemSummary, int, error) { + // 获取总数 + total, err := m.GetItemsCount(category) + if err != nil { + return nil, 0, err + } + + // 获取列表数据（不包含内容） + var rows *sql.Rows + var query string + var args []interface{} + + query = "SELECT id, category, title, file_path, created_at, updated_at FROM knowledge_base_items" + + if category != "" { + query += " WHERE category = ?" + args = append(args, category) + } + + query += " ORDER BY category, title" + + if limit > 0 { + query += " LIMIT ?" + args = append(args, limit) + if offset > 0 { + query += " OFFSET ?" + args = append(args, offset) + } + } + + rows, err = m.db.Query(query, args...) + if err != nil { + return nil, 0, fmt.Errorf("查询知识项失败: %w", err) + } + defer rows.Close() + + var items []*KnowledgeItemSummary + for rows.Next() { + item := &KnowledgeItemSummary{} + var createdAt, updatedAt string + + if err := rows.Scan(&item.ID, &item.Category, &item.Title, &item.FilePath, &createdAt, &updatedAt); err != nil { + return nil, 0, fmt.Errorf("扫描知识项失败: %w", err) + } + + // 解析时间 + timeFormats := []string{ + "2006-01-02 15:04:05.999999999-07:00", + "2006-01-02 15:04:05.999999999", + "2006-01-02T15:04:05.999999999Z07:00", + "2006-01-02T15:04:05Z", + "2006-01-02 15:04:05", + time.RFC3339, + time.RFC3339Nano, + } + + if createdAt != "" { + for _, format := range timeFormats { + parsed, err := time.Parse(format, createdAt) + if err == nil && !parsed.IsZero() { + item.CreatedAt = parsed + break + } + } + } + + if updatedAt != "" { + for _, format := range timeFormats { + parsed, err := time.Parse(format, updatedAt) + if err == nil && !parsed.IsZero() { + item.UpdatedAt = parsed + break + } + } + } + + if item.UpdatedAt.IsZero() && !item.CreatedAt.IsZero() { + item.UpdatedAt = item.CreatedAt + } + + items = append(items, item) + } + + return items, total, nil +} + +// GetItem 获取单个知识项 +func (m *Manager) GetItem(id string) (*KnowledgeItem, error) { + item := &KnowledgeItem{} + var createdAt, updatedAt string + err := m.db.QueryRow( + "SELECT id, category, title, file_path, content, created_at, updated_at FROM knowledge_base_items WHERE id = ?", + id, + ).Scan(&item.ID, &item.Category, &item.Title, &item.FilePath, &item.Content, &createdAt, &updatedAt) + + if err == sql.ErrNoRows { + return nil, fmt.Errorf("知识项不存在") + } + if err != nil { + return nil, fmt.Errorf("查询知识项失败: %w", err) + } + + // 解析时间 - 支持多种格式 + timeFormats := []string{ + "2006-01-02 15:04:05.999999999-07:00", + "2006-01-02 15:04:05.999999999", + "2006-01-02T15:04:05.999999999Z07:00", + "2006-01-02T15:04:05Z", + "2006-01-02 15:04:05", + time.RFC3339, + time.RFC3339Nano, + } + + // 解析创建时间 + if createdAt != "" { + for _, format := range timeFormats { + parsed, err := time.Parse(format, createdAt) + if err == nil && !parsed.IsZero() { + item.CreatedAt = parsed + break + } + } + } + + // 解析更新时间 + if updatedAt != "" { + for _, format := range timeFormats { + parsed, err := time.Parse(format, updatedAt) + if err == nil && !parsed.IsZero() { + item.UpdatedAt = parsed + break + } + } + } + + // 如果更新时间为空，使用创建时间 + if item.UpdatedAt.IsZero() && !item.CreatedAt.IsZero() { + item.UpdatedAt = item.CreatedAt + } + + return item, nil +} + +// CreateItem 创建知识项 +func (m *Manager) CreateItem(category, title, content string) (*KnowledgeItem, error) { + id := uuid.New().String() + now := time.Now() + + // 构建文件路径 + filePath := filepath.Join(m.basePath, category, title+".md") + + // 确保目录存在 + if err := os.MkdirAll(filepath.Dir(filePath), 0755); err != nil { + return nil, fmt.Errorf("创建目录失败: %w", err) + } + + // 写入文件 + if err := os.WriteFile(filePath, []byte(content), 0644); err != nil { + return nil, fmt.Errorf("写入文件失败: %w", err) + } + + // 插入数据库 + _, err := m.db.Exec( + "INSERT INTO knowledge_base_items (id, category, title, file_path, content, created_at, updated_at) VALUES (?, ?, ?, ?, ?, ?, ?)", + id, category, title, filePath, content, now, now, + ) + if err != nil { + return nil, fmt.Errorf("插入知识项失败: %w", err) + } + + return &KnowledgeItem{ + ID: id, + Category: category, + Title: title, + FilePath: filePath, + Content: content, + CreatedAt: now, + UpdatedAt: now, + }, nil +} + +// UpdateItem 更新知识项 +func (m *Manager) UpdateItem(id, category, title, content string) (*KnowledgeItem, error) { + // 获取现有项 + item, err := m.GetItem(id) + if err != nil { + return nil, err + } + + // 构建新文件路径 + newFilePath := filepath.Join(m.basePath, category, title+".md") + + // 如果路径改变，需要移动文件 + if item.FilePath != newFilePath { + // 确保新目录存在 + if err := os.MkdirAll(filepath.Dir(newFilePath), 0755); err != nil { + return nil, fmt.Errorf("创建目录失败: %w", err) + } + + // 移动文件 + if err := os.Rename(item.FilePath, newFilePath); err != nil { + return nil, fmt.Errorf("移动文件失败: %w", err) + } + + // 删除旧目录（如果为空） + oldDir := filepath.Dir(item.FilePath) + if isEmpty, _ := isEmptyDir(oldDir); isEmpty { + // 只有当目录不是知识库根目录时才删除（避免删除根目录） + if oldDir != m.basePath { + if err := os.Remove(oldDir); err != nil { + m.logger.Warn("删除空目录失败", zap.String("dir", oldDir), zap.Error(err)) + } + } + } + } + + // 写入文件 + if err := os.WriteFile(newFilePath, []byte(content), 0644); err != nil { + return nil, fmt.Errorf("写入文件失败: %w", err) + } + + // 更新数据库 + _, err = m.db.Exec( + "UPDATE knowledge_base_items SET category = ?, title = ?, file_path = ?, content = ?, updated_at = ? WHERE id = ?", + category, title, newFilePath, content, time.Now(), id, + ) + if err != nil { + return nil, fmt.Errorf("更新知识项失败: %w", err) + } + + // 删除旧的向量嵌入（需要重新索引） + _, err = m.db.Exec("DELETE FROM knowledge_embeddings WHERE item_id = ?", id) + if err != nil { + m.logger.Warn("删除旧向量嵌入失败", zap.Error(err)) + } + + return m.GetItem(id) +} + +// DeleteItem 删除知识项 +func (m *Manager) DeleteItem(id string) error { + // 获取文件路径 + var filePath string + err := m.db.QueryRow("SELECT file_path FROM knowledge_base_items WHERE id = ?", id).Scan(&filePath) + if err != nil { + return fmt.Errorf("查询知识项失败: %w", err) + } + + // 删除文件 + if err := os.Remove(filePath); err != nil && !os.IsNotExist(err) { + m.logger.Warn("删除文件失败", zap.String("path", filePath), zap.Error(err)) + } + + // 删除数据库记录（级联删除向量） + _, err = m.db.Exec("DELETE FROM knowledge_base_items WHERE id = ?", id) + if err != nil { + return fmt.Errorf("删除知识项失败: %w", err) + } + + // 删除空目录（如果为空） + dir := filepath.Dir(filePath) + if isEmpty, _ := isEmptyDir(dir); isEmpty { + // 只有当目录不是知识库根目录时才删除（避免删除根目录） + if dir != m.basePath { + if err := os.Remove(dir); err != nil { + m.logger.Warn("删除空目录失败", zap.String("dir", dir), zap.Error(err)) + } + } + } + + return nil +} + +// isEmptyDir 检查目录是否为空（忽略隐藏文件和 . 开头的文件） +func isEmptyDir(dir string) (bool, error) { + entries, err := os.ReadDir(dir) + if err != nil { + return false, err + } + for _, entry := range entries { + // 忽略隐藏文件（以 . 开头） + if !strings.HasPrefix(entry.Name(), ".") { + return false, nil + } + } + return true, nil +} + +// LogRetrieval 记录检索日志 +func (m *Manager) LogRetrieval(conversationID, messageID, query, riskType string, retrievedItems []string) error { + id := uuid.New().String() + itemsJSON, _ := json.Marshal(retrievedItems) + + _, err := m.db.Exec( + "INSERT INTO knowledge_retrieval_logs (id, conversation_id, message_id, query, risk_type, retrieved_items, created_at) VALUES (?, ?, ?, ?, ?, ?, ?)", + id, conversationID, messageID, query, riskType, string(itemsJSON), time.Now(), + ) + return err +} + +// GetIndexStatus 获取索引状态 +func (m *Manager) GetIndexStatus() (map[string]interface{}, error) { + // 获取总知识项数 + var totalItems int + err := m.db.QueryRow("SELECT COUNT(*) FROM knowledge_base_items").Scan(&totalItems) + if err != nil { + return nil, fmt.Errorf("查询总知识项数失败: %w", err) + } + + // 获取已索引的知识项数（有向量嵌入的） + var indexedItems int + err = m.db.QueryRow(` + SELECT COUNT(DISTINCT item_id) + FROM knowledge_embeddings + `).Scan(&indexedItems) + if err != nil { + return nil, fmt.Errorf("查询已索引项数失败: %w", err) + } + + // 计算进度百分比 + var progressPercent float64 + if totalItems > 0 { + progressPercent = float64(indexedItems) / float64(totalItems) * 100 + } else { + progressPercent = 100.0 + } + + // 判断是否完成 + isComplete := indexedItems >= totalItems && totalItems > 0 + + return map[string]interface{}{ + "total_items": totalItems, + "indexed_items": indexedItems, + "progress_percent": progressPercent, + "is_complete": isComplete, + }, nil +} + +// GetRetrievalLogs 获取检索日志 +func (m *Manager) GetRetrievalLogs(conversationID, messageID string, limit int) ([]*RetrievalLog, error) { + var rows *sql.Rows + var err error + + if messageID != "" { + rows, err = m.db.Query( + "SELECT id, conversation_id, message_id, query, risk_type, retrieved_items, created_at FROM knowledge_retrieval_logs WHERE message_id = ? ORDER BY created_at DESC LIMIT ?", + messageID, limit, + ) + } else if conversationID != "" { + rows, err = m.db.Query( + "SELECT id, conversation_id, message_id, query, risk_type, retrieved_items, created_at FROM knowledge_retrieval_logs WHERE conversation_id = ? ORDER BY created_at DESC LIMIT ?", + conversationID, limit, + ) + } else { + rows, err = m.db.Query( + "SELECT id, conversation_id, message_id, query, risk_type, retrieved_items, created_at FROM knowledge_retrieval_logs ORDER BY created_at DESC LIMIT ?", + limit, + ) + } + + if err != nil { + return nil, fmt.Errorf("查询检索日志失败: %w", err) + } + defer rows.Close() + + var logs []*RetrievalLog + for rows.Next() { + log := &RetrievalLog{} + var createdAt string + var itemsJSON sql.NullString + if err := rows.Scan(&log.ID, &log.ConversationID, &log.MessageID, &log.Query, &log.RiskType, &itemsJSON, &createdAt); err != nil { + return nil, fmt.Errorf("扫描检索日志失败: %w", err) + } + + // 解析时间 - 支持多种格式 + var err error + timeFormats := []string{ + "2006-01-02 15:04:05.999999999-07:00", + "2006-01-02 15:04:05.999999999", + "2006-01-02T15:04:05.999999999Z07:00", + "2006-01-02T15:04:05Z", + "2006-01-02 15:04:05", + time.RFC3339, + time.RFC3339Nano, + } + + for _, format := range timeFormats { + log.CreatedAt, err = time.Parse(format, createdAt) + if err == nil && !log.CreatedAt.IsZero() { + break + } + } + + // 如果所有格式都失败，记录警告但继续处理 + if log.CreatedAt.IsZero() { + m.logger.Warn("解析检索日志时间失败", + zap.String("timeStr", createdAt), + zap.Error(err), + ) + // 使用当前时间作为fallback + log.CreatedAt = time.Now() + } + + // 解析检索项 + if itemsJSON.Valid { + json.Unmarshal([]byte(itemsJSON.String), &log.RetrievedItems) + } + + logs = append(logs, log) + } + + return logs, nil +} + +// DeleteRetrievalLog 删除检索日志 +func (m *Manager) DeleteRetrievalLog(id string) error { + result, err := m.db.Exec("DELETE FROM knowledge_retrieval_logs WHERE id = ?", id) + if err != nil { + return fmt.Errorf("删除检索日志失败: %w", err) + } + + rowsAffected, err := result.RowsAffected() + if err != nil { + return fmt.Errorf("获取删除行数失败: %w", err) + } + + if rowsAffected == 0 { + return fmt.Errorf("检索日志不存在") + } + + return nil +} diff --git a/internal/knowledge/rerank_http.go b/internal/knowledge/rerank_http.go new file mode 100644 index 00000000..61e173ed --- /dev/null +++ b/internal/knowledge/rerank_http.go @@ -0,0 +1,226 @@ +package knowledge + +import ( + "bytes" + "context" + "encoding/json" + "fmt" + "io" + "net/http" + "strings" + "time" + + "cyberstrike-ai/internal/config" + + "github.com/cloudwego/eino/schema" + "go.uber.org/zap" +) + +// HTTPReranker calls a hosted rerank API (DashScope or Cohere-compatible). +type HTTPReranker struct { + provider string + model string + baseURL string + apiKey string + client *http.Client + logger *zap.Logger +} + +// NewHTTPReranker builds a rerank client from knowledge retrieval config; openAI supplies fallback credentials. +func NewHTTPReranker(rc *config.RerankConfig, openAI *config.OpenAIConfig, logger *zap.Logger) (*HTTPReranker, error) { + if rc == nil { + return nil, fmt.Errorf("rerank config is nil") + } + baseURL := strings.TrimSpace(rc.BaseURL) + apiKey := strings.TrimSpace(rc.APIKey) + if openAI != nil { + if baseURL == "" { + baseURL = strings.TrimSpace(openAI.BaseURL) + } + if apiKey == "" { + apiKey = strings.TrimSpace(openAI.APIKey) + } + } + if apiKey == "" { + return nil, fmt.Errorf("rerank api_key is required") + } + provider := rc.ProviderEffective(baseURL) + model := rc.ModelEffective(provider) + return &HTTPReranker{ + provider: provider, + model: model, + baseURL: strings.TrimSuffix(baseURL, "/"), + apiKey: apiKey, + client: &http.Client{Timeout: 60 * time.Second}, + logger: logger, + }, nil +} + +func (r *HTTPReranker) Rerank(ctx context.Context, query string, docs []*schema.Document) ([]*schema.Document, error) { + if r == nil { + return docs, nil + } + q := strings.TrimSpace(query) + if q == "" || len(docs) == 0 { + return docs, nil + } + if len(docs) == 1 { + return docs, nil + } + texts := make([]string, 0, len(docs)) + for _, d := range docs { + if d == nil { + texts = append(texts, "") + continue + } + texts = append(texts, d.Content) + } + var order []int + var err error + switch r.provider { + case "dashscope": + order, err = r.rerankDashScope(ctx, q, texts, len(docs)) + default: + order, err = r.rerankCohere(ctx, q, texts, len(docs)) + } + if err != nil { + return nil, err + } + out := make([]*schema.Document, 0, len(order)) + for _, idx := range order { + if idx < 0 || idx >= len(docs) || docs[idx] == nil { + continue + } + out = append(out, docs[idx]) + } + if len(out) == 0 { + return docs, nil + } + return out, nil +} + +func (r *HTTPReranker) rerankCohere(ctx context.Context, query string, documents []string, topN int) ([]int, error) { + url := r.cohereRerankURL() + body := map[string]any{ + "model": r.model, + "query": query, + "documents": documents, + "top_n": topN, + } + raw, err := json.Marshal(body) + if err != nil { + return nil, err + } + req, err := http.NewRequestWithContext(ctx, http.MethodPost, url, bytes.NewReader(raw)) + if err != nil { + return nil, err + } + req.Header.Set("Content-Type", "application/json") + req.Header.Set("Authorization", "Bearer "+r.apiKey) + resp, err := r.client.Do(req) + if err != nil { + return nil, fmt.Errorf("rerank request: %w", err) + } + defer resp.Body.Close() + respBody, _ := io.ReadAll(resp.Body) + if resp.StatusCode < 200 || resp.StatusCode >= 300 { + return nil, fmt.Errorf("rerank http %d: %s", resp.StatusCode, truncateForRerankLog(string(respBody))) + } + var parsed struct { + Results []struct { + Index int `json:"index"` + } `json:"results"` + } + if err := json.Unmarshal(respBody, &parsed); err != nil { + return nil, fmt.Errorf("rerank decode: %w", err) + } + order := make([]int, 0, len(parsed.Results)) + for _, row := range parsed.Results { + order = append(order, row.Index) + } + return order, nil +} + +func (r *HTTPReranker) rerankDashScope(ctx context.Context, query string, documents []string, topN int) ([]int, error) { + url := r.dashscopeRerankURL() + body := map[string]any{ + "model": r.model, + "input": map[string]any{ + "query": query, + "documents": documents, + }, + "parameters": map[string]any{ + "return_documents": false, + "top_n": topN, + }, + } + raw, err := json.Marshal(body) + if err != nil { + return nil, err + } + req, err := http.NewRequestWithContext(ctx, http.MethodPost, url, bytes.NewReader(raw)) + if err != nil { + return nil, err + } + req.Header.Set("Content-Type", "application/json") + req.Header.Set("Authorization", "Bearer "+r.apiKey) + resp, err := r.client.Do(req) + if err != nil { + return nil, fmt.Errorf("dashscope rerank: %w", err) + } + defer resp.Body.Close() + respBody, _ := io.ReadAll(resp.Body) + if resp.StatusCode < 200 || resp.StatusCode >= 300 { + return nil, fmt.Errorf("dashscope rerank http %d: %s", resp.StatusCode, truncateForRerankLog(string(respBody))) + } + var parsed struct { + Output struct { + Results []struct { + Index int `json:"index"` + } `json:"results"` + } `json:"output"` + } + if err := json.Unmarshal(respBody, &parsed); err != nil { + return nil, fmt.Errorf("dashscope rerank decode: %w", err) + } + order := make([]int, 0, len(parsed.Output.Results)) + for _, row := range parsed.Output.Results { + order = append(order, row.Index) + } + return order, nil +} + +func (r *HTTPReranker) cohereRerankURL() string { + base := r.baseURL + if base == "" { + base = "https://api.cohere.com" + } + if strings.HasSuffix(base, "/v1") { + return base + "/rerank" + } + return base + "/v1/rerank" +} + +func (r *HTTPReranker) dashscopeRerankURL() string { + base := strings.TrimSpace(r.baseURL) + if base == "" { + return "https://dashscope.aliyuncs.com/api/v1/services/rerank/text-rerank/text-rerank" + } + if strings.Contains(base, "/api/v1/services/rerank") { + return base + } + if strings.Contains(base, "dashscope.aliyuncs.com") || strings.Contains(base, "compatible-mode") { + return "https://dashscope.aliyuncs.com/api/v1/services/rerank/text-rerank/text-rerank" + } + return strings.TrimSuffix(base, "/") +} + +func truncateForRerankLog(s string) string { + s = strings.TrimSpace(s) + if len(s) > 512 { + return s[:512] + "..." + } + return s +} + +var _ DocumentReranker = (*HTTPReranker)(nil) diff --git a/internal/knowledge/rerank_http_test.go b/internal/knowledge/rerank_http_test.go new file mode 100644 index 00000000..013ad7d4 --- /dev/null +++ b/internal/knowledge/rerank_http_test.go @@ -0,0 +1,97 @@ +package knowledge + +import ( + "context" + "encoding/json" + "net/http" + "net/http/httptest" + "testing" + + "cyberstrike-ai/internal/config" + + "github.com/cloudwego/eino/schema" +) + +func TestHTTPReranker_CohereOrder(t *testing.T) { + t.Parallel() + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if r.URL.Path != "/v1/rerank" { + t.Fatalf("path %s", r.URL.Path) + } + _ = json.NewEncoder(w).Encode(map[string]any{ + "results": []map[string]any{ + {"index": 2, "relevance_score": 0.9}, + {"index": 0, "relevance_score": 0.5}, + }, + }) + })) + defer srv.Close() + + rr, err := NewHTTPReranker(&config.RerankConfig{ + Provider: "cohere", + Model: "rerank-multilingual-v3.0", + BaseURL: srv.URL, + APIKey: "test-key", + }, nil, nil) + if err != nil { + t.Fatal(err) + } + docs := []*schema.Document{ + {ID: "a", Content: "alpha"}, + {ID: "b", Content: "beta"}, + {ID: "c", Content: "gamma"}, + } + out, err := rr.Rerank(context.Background(), "query", docs) + if err != nil { + t.Fatal(err) + } + if len(out) != 2 || out[0].ID != "c" || out[1].ID != "a" { + t.Fatalf("order wrong: %#v", out) + } +} + +func TestHTTPReranker_DashScopeOrder(t *testing.T) { + t.Parallel() + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + _ = json.NewEncoder(w).Encode(map[string]any{ + "output": map[string]any{ + "results": []map[string]any{ + {"index": 1, "relevance_score": 0.88}, + }, + }, + }) + })) + defer srv.Close() + + rr, err := NewHTTPReranker(&config.RerankConfig{ + Provider: "dashscope", + Model: "gte-rerank", + BaseURL: srv.URL, + APIKey: "test-key", + }, nil, nil) + if err != nil { + t.Fatal(err) + } + docs := []*schema.Document{{ID: "a", Content: "a"}, {ID: "b", Content: "b"}} + out, err := rr.Rerank(context.Background(), "q", docs) + if err != nil { + t.Fatal(err) + } + if len(out) != 1 || out[0].ID != "b" { + t.Fatalf("got %#v", out) + } +} + +func TestRerankConfigDefaults(t *testing.T) { + t.Parallel() + rc := config.RerankConfig{} + if rc.ProviderEffective("https://dashscope.aliyuncs.com/x") != "dashscope" { + t.Fatal("dashscope detect") + } + if rc.ModelEffective("dashscope") != "gte-rerank" { + t.Fatal("dashscope model") + } + if rc.ModelEffective("cohere") != "rerank-multilingual-v3.0" { + t.Fatal("cohere model") + } +} diff --git a/internal/knowledge/retrieval_postprocess.go b/internal/knowledge/retrieval_postprocess.go new file mode 100644 index 00000000..20e07110 --- /dev/null +++ b/internal/knowledge/retrieval_postprocess.go @@ -0,0 +1,216 @@ +package knowledge + +import ( + "context" + "crypto/sha256" + "encoding/hex" + "fmt" + "strings" + "sync" + "unicode" + "unicode/utf8" + + "cyberstrike-ai/internal/config" + + "github.com/cloudwego/eino/schema" + "github.com/pkoukk/tiktoken-go" +) + +// postRetrieveMaxPrefetchCap 限制单次向量候选上限，避免误配置导致全表扫压力过大。 +const postRetrieveMaxPrefetchCap = 200 + +// DocumentReranker 精排（HTTP dashscope / Cohere 兼容 API），由 [WireRetrieverPipeline] 注入。 +type DocumentReranker interface { + Rerank(ctx context.Context, query string, docs []*schema.Document) ([]*schema.Document, error) +} + +// NopDocumentReranker 占位实现，便于测试或未启用重排时显式注入。 +type NopDocumentReranker struct{} + +// Rerank implements [DocumentReranker] as no-op. +func (NopDocumentReranker) Rerank(_ context.Context, _ string, docs []*schema.Document) ([]*schema.Document, error) { + return docs, nil +} + +var tiktokenEncMu sync.Mutex +var tiktokenEncCache = map[string]*tiktoken.Tiktoken{} + +func encodingForTokenizerModel(model string) (*tiktoken.Tiktoken, error) { + m := strings.TrimSpace(model) + if m == "" { + m = "gpt-4" + } + tiktokenEncMu.Lock() + defer tiktokenEncMu.Unlock() + if enc, ok := tiktokenEncCache[m]; ok { + return enc, nil + } + enc, err := tiktoken.EncodingForModel(m) + if err != nil { + enc, err = tiktoken.GetEncoding("cl100k_base") + if err != nil { + return nil, err + } + } + tiktokenEncCache[m] = enc + return enc, nil +} + +func countDocTokens(text, model string) (int, error) { + enc, err := encodingForTokenizerModel(model) + if err != nil { + return 0, err + } + toks := enc.Encode(text, nil, nil) + return len(toks), nil +} + +// normalizeContentFingerprintKey 去重键：trim + 空白折叠（不改动大小写，避免合并仅大小写不同的代码片段）。 +func normalizeContentFingerprintKey(s string) string { + s = strings.TrimSpace(s) + var b strings.Builder + b.Grow(len(s)) + prevSpace := false + for _, r := range s { + if unicode.IsSpace(r) { + if !prevSpace { + b.WriteByte(' ') + prevSpace = true + } + continue + } + prevSpace = false + b.WriteRune(r) + } + return b.String() +} + +func contentNormKey(d *schema.Document) string { + if d == nil { + return "" + } + n := normalizeContentFingerprintKey(d.Content) + if n == "" { + return "" + } + sum := sha256.Sum256([]byte(n)) + return hex.EncodeToString(sum[:]) +} + +// dedupeByNormalizedContent 按规范化正文去重，保留向量检索顺序中首次出现的文档（同正文仅保留一条）。 +func dedupeByNormalizedContent(docs []*schema.Document) []*schema.Document { + if len(docs) < 2 { + return docs + } + seen := make(map[string]struct{}, len(docs)) + out := make([]*schema.Document, 0, len(docs)) + for _, d := range docs { + if d == nil { + continue + } + k := contentNormKey(d) + if k == "" { + out = append(out, d) + continue + } + if _, ok := seen[k]; ok { + continue + } + seen[k] = struct{}{} + out = append(out, d) + } + return out +} + +// truncateDocumentsByBudget 按检索顺序整段保留文档，直至字符数或 token 数（任一启用）超限则停止。 +func truncateDocumentsByBudget(docs []*schema.Document, maxRunes, maxTokens int, tokenModel string) ([]*schema.Document, error) { + if len(docs) == 0 { + return docs, nil + } + unlimitedChars := maxRunes <= 0 + unlimitedTok := maxTokens <= 0 + if unlimitedChars && unlimitedTok { + return docs, nil + } + + remRunes := maxRunes + remTok := maxTokens + out := make([]*schema.Document, 0, len(docs)) + + for _, d := range docs { + if d == nil || strings.TrimSpace(d.Content) == "" { + continue + } + runes := utf8.RuneCountInString(d.Content) + if !unlimitedChars && runes > remRunes { + break + } + var tok int + var err error + if !unlimitedTok { + tok, err = countDocTokens(d.Content, tokenModel) + if err != nil { + return nil, fmt.Errorf("token count: %w", err) + } + if tok > remTok { + break + } + } + out = append(out, d) + if !unlimitedChars { + remRunes -= runes + } + if !unlimitedTok { + remTok -= tok + } + } + return out, nil +} + +// EffectivePrefetchTopK 计算每条 MultiQuery 变体在向量阶段的候选条数（供融合 / 重排 / 后处理）。 +func EffectivePrefetchTopK(topK int, po *config.PostRetrieveConfig) int { + if topK < 1 { + topK = 5 + } + fetch := topK * 4 + if fetch < 20 { + fetch = 20 + } + if po != nil && po.PrefetchTopK > 0 { + fetch = po.PrefetchTopK + } + if fetch > postRetrieveMaxPrefetchCap { + fetch = postRetrieveMaxPrefetchCap + } + return fetch +} + +// ApplyPostRetrieve 检索后处理：规范化正文去重 → 预算截断 → 最终 TopK（精排已在流水线中完成）。 +func ApplyPostRetrieve(docs []*schema.Document, po *config.PostRetrieveConfig, tokenModel string, finalTopK int) ([]*schema.Document, error) { + if finalTopK < 1 { + finalTopK = 5 + } + if len(docs) == 0 { + return docs, nil + } + + maxChars := 0 + maxTok := 0 + if po != nil { + maxChars = po.MaxContextChars + maxTok = po.MaxContextTokens + } + + out := dedupeByNormalizedContent(docs) + + var err error + out, err = truncateDocumentsByBudget(out, maxChars, maxTok, tokenModel) + if err != nil { + return nil, err + } + + if len(out) > finalTopK { + out = out[:finalTopK] + } + return out, nil +} diff --git a/internal/knowledge/retrieval_postprocess_test.go b/internal/knowledge/retrieval_postprocess_test.go new file mode 100644 index 00000000..889d5c62 --- /dev/null +++ b/internal/knowledge/retrieval_postprocess_test.go @@ -0,0 +1,62 @@ +package knowledge + +import ( + "testing" + + "cyberstrike-ai/internal/config" + + "github.com/cloudwego/eino/schema" +) + +func doc(id, content string, score float64) *schema.Document { + d := &schema.Document{ID: id, Content: content, MetaData: map[string]any{metaKBItemID: "it1"}} + d.WithScore(score) + return d +} + +func TestDedupeByNormalizedContent(t *testing.T) { + a := doc("1", "hello world", 0.9) + b := doc("2", "hello world", 0.8) + c := doc("3", "other", 0.7) + out := dedupeByNormalizedContent([]*schema.Document{a, b, c}) + if len(out) != 2 { + t.Fatalf("len=%d want 2", len(out)) + } + if out[0].ID != "1" || out[1].ID != "3" { + t.Fatalf("order/ids wrong: %#v", out) + } +} + +func TestEffectivePrefetchTopK(t *testing.T) { + if g := EffectivePrefetchTopK(5, nil); g != 20 { + t.Fatalf("default prefetch got %d want 20", g) + } + if g := EffectivePrefetchTopK(5, &config.PostRetrieveConfig{PrefetchTopK: 50}); g != 50 { + t.Fatalf("got %d", g) + } + if g := EffectivePrefetchTopK(5, &config.PostRetrieveConfig{PrefetchTopK: 9999}); g != postRetrieveMaxPrefetchCap { + t.Fatalf("cap: got %d", g) + } +} + +func TestApplyPostRetrieveTruncateAndTopK(t *testing.T) { + d1 := doc("1", "ab", 0.9) + d2 := doc("2", "cd", 0.8) + d3 := doc("3", "ef", 0.7) + po := &config.PostRetrieveConfig{MaxContextChars: 3} + out, err := ApplyPostRetrieve([]*schema.Document{d1, d2, d3}, po, "gpt-4", 5) + if err != nil { + t.Fatal(err) + } + if len(out) != 1 || out[0].ID != "1" { + t.Fatalf("got %#v", out) + } + + out2, err := ApplyPostRetrieve([]*schema.Document{d1, d2, d3}, nil, "gpt-4", 2) + if err != nil { + t.Fatal(err) + } + if len(out2) != 2 { + t.Fatalf("topk: len=%d", len(out2)) + } +} diff --git a/internal/knowledge/retriever.go b/internal/knowledge/retriever.go new file mode 100644 index 00000000..c75e8a13 --- /dev/null +++ b/internal/knowledge/retriever.go @@ -0,0 +1,334 @@ +package knowledge + +import ( + "context" + "database/sql" + "encoding/json" + "fmt" + "math" + "sort" + "strings" + "sync" + + "cyberstrike-ai/internal/config" + + "github.com/cloudwego/eino/components/retriever" + "github.com/cloudwego/eino/schema" + "go.uber.org/zap" +) + +// Retriever 检索器：SQLite 存向量 + Eino 嵌入，**纯向量检索**（余弦相似度、TopK、阈值）， +// 实现语义与 [retriever.Retriever] 适配层 [VectorEinoRetriever] 一致。 +type Retriever struct { + db *sql.DB + embedder *Embedder + config *RetrievalConfig + logger *zap.Logger + + rerankMu sync.RWMutex + reranker DocumentReranker + + pipeline retriever.Retriever + wireOpenAI *config.OpenAIConfig +} + +// RetrievalConfig 检索配置 +type RetrievalConfig struct { + TopK int + SimilarityThreshold float64 + SubIndexFilter string + MultiQuery config.MultiQueryConfig + Rerank config.RerankConfig + PostRetrieve config.PostRetrieveConfig +} + +// NewRetriever 创建新的检索器 +func NewRetriever(db *sql.DB, embedder *Embedder, config *RetrievalConfig, logger *zap.Logger) *Retriever { + return &Retriever{ + db: db, + embedder: embedder, + config: config, + logger: logger, + } +} + +// UpdateConfig 更新检索配置并重建 Eino MultiQuery + 重排流水线。 +func (r *Retriever) UpdateConfig(cfg *RetrievalConfig) { + if cfg != nil { + r.config = cfg + if r.logger != nil { + r.logger.Info("检索器配置已更新", + zap.Int("top_k", cfg.TopK), + zap.Float64("similarity_threshold", cfg.SimilarityThreshold), + zap.String("sub_index_filter", cfg.SubIndexFilter), + zap.Int("multi_query_max", cfg.MultiQuery.MaxQueriesEffective()), + zap.Int("post_retrieve_prefetch_top_k", cfg.PostRetrieve.PrefetchTopK), + zap.Int("post_retrieve_max_context_chars", cfg.PostRetrieve.MaxContextChars), + zap.Int("post_retrieve_max_context_tokens", cfg.PostRetrieve.MaxContextTokens), + ) + } + } + if r.wireOpenAI != nil { + if err := WireRetrieverPipeline(context.Background(), r, r.wireOpenAI); err != nil && r.logger != nil { + r.logger.Warn("检索流水线重建失败", zap.Error(err)) + } + } +} + +// SetDocumentReranker 注入可选重排器（并发安全）；nil 表示禁用。 +func (r *Retriever) SetDocumentReranker(rr DocumentReranker) { + if r == nil { + return + } + r.rerankMu.Lock() + defer r.rerankMu.Unlock() + r.reranker = rr +} + +func (r *Retriever) documentReranker() DocumentReranker { + if r == nil { + return nil + } + r.rerankMu.RLock() + defer r.rerankMu.RUnlock() + return r.reranker +} + +func cosineSimilarity(a, b []float32) float64 { + if len(a) != len(b) { + return 0.0 + } + + var dotProduct, normA, normB float64 + for i := range a { + dotProduct += float64(a[i] * b[i]) + normA += float64(a[i] * a[i]) + normB += float64(b[i] * b[i]) + } + + if normA == 0 || normB == 0 { + return 0.0 + } + + return dotProduct / (math.Sqrt(normA) * math.Sqrt(normB)) +} + +// Search 搜索知识库（Eino MultiQuery → 向量检索 → 重排 → 后处理）。 +func (r *Retriever) Search(ctx context.Context, req *SearchRequest) ([]*RetrievalResult, error) { + if req == nil { + return nil, fmt.Errorf("请求不能为空") + } + q := strings.TrimSpace(req.Query) + if q == "" { + return nil, fmt.Errorf("查询不能为空") + } + opts := r.einoRetrieverOptions(req) + docs, err := r.activeEinoRetriever().Retrieve(ctx, q, opts...) + if err != nil { + return nil, err + } + return documentsToRetrievalResults(docs) +} + +func (r *Retriever) einoRetrieverOptions(req *SearchRequest) []retriever.Option { + var opts []retriever.Option + if req.TopK > 0 { + opts = append(opts, retriever.WithTopK(req.TopK)) + } + dsl := map[string]any{} + if strings.TrimSpace(req.RiskType) != "" { + dsl[DSLRiskType] = strings.TrimSpace(req.RiskType) + } + if req.Threshold > 0 { + dsl[DSLSimilarityThreshold] = req.Threshold + } + if strings.TrimSpace(req.SubIndexFilter) != "" { + dsl[DSLSubIndexFilter] = strings.TrimSpace(req.SubIndexFilter) + } + if len(dsl) > 0 { + opts = append(opts, retriever.WithDSLInfo(dsl)) + } + return opts +} + +// EinoRetrieve 直接返回 [schema.Document]，供 Eino Graph / Chain 使用。 +func (r *Retriever) EinoRetrieve(ctx context.Context, query string, opts ...retriever.Option) ([]*schema.Document, error) { + return r.activeEinoRetriever().Retrieve(ctx, query, opts...) +} + +func (r *Retriever) activeEinoRetriever() retriever.Retriever { + if r != nil && r.pipeline != nil { + return r.pipeline + } + return NewVectorEinoRetriever(r) +} + +// AsEinoRetriever 将知识库检索流水线暴露为 Eino [retriever.Retriever]。 +func (r *Retriever) AsEinoRetriever() retriever.Retriever { + return r.activeEinoRetriever() +} + +func (r *Retriever) knowledgeEmbeddingSelectSQL(riskType, subIndexFilter string) (string, []interface{}) { + q := `SELECT e.id, e.item_id, e.chunk_index, e.chunk_text, e.embedding, e.embedding_model, e.embedding_dim, i.category, i.title +FROM knowledge_embeddings e +JOIN knowledge_base_items i ON e.item_id = i.id +WHERE 1=1` + var args []interface{} + if strings.TrimSpace(riskType) != "" { + q += ` AND TRIM(i.category) = TRIM(?) COLLATE NOCASE` + args = append(args, riskType) + } + if tag := strings.TrimSpace(subIndexFilter); tag != "" { + tag = strings.ToLower(strings.ReplaceAll(tag, " ", "")) + q += ` AND (TRIM(COALESCE(e.sub_indexes,'')) = '' OR INSTR(',' || LOWER(REPLACE(e.sub_indexes,' ','')) || ',', ',' || ? || ',') > 0)` + args = append(args, tag) + } + return q, args +} + +// vectorSearch 纯向量检索：余弦相似度排序，按相似度阈值与 TopK 截断（无 BM25、无混合分、无邻块扩展）。 +func (r *Retriever) vectorSearch(ctx context.Context, req *SearchRequest) ([]*RetrievalResult, error) { + if req.Query == "" { + return nil, fmt.Errorf("查询不能为空") + } + + topK := req.TopK + if topK <= 0 && r.config != nil { + topK = r.config.TopK + } + if topK <= 0 { + topK = 5 + } + + threshold := req.Threshold + if threshold <= 0 && r.config != nil { + threshold = r.config.SimilarityThreshold + } + if threshold <= 0 { + threshold = 0.7 + } + + subIdxFilter := strings.TrimSpace(req.SubIndexFilter) + if subIdxFilter == "" && r.config != nil { + subIdxFilter = strings.TrimSpace(r.config.SubIndexFilter) + } + + queryText := FormatQueryEmbeddingText(req.RiskType, req.Query) + queryEmbedding, err := r.embedder.EmbedText(ctx, queryText) + if err != nil { + return nil, fmt.Errorf("向量化查询失败: %w", err) + } + queryDim := len(queryEmbedding) + expectedModel := "" + if r.embedder != nil { + expectedModel = r.embedder.EmbeddingModelName() + } + + sqlStr, sqlArgs := r.knowledgeEmbeddingSelectSQL(strings.TrimSpace(req.RiskType), subIdxFilter) + rows, err := r.db.QueryContext(ctx, sqlStr, sqlArgs...) + if err != nil { + return nil, fmt.Errorf("查询向量失败: %w", err) + } + defer rows.Close() + + type candidate struct { + chunk *KnowledgeChunk + item *KnowledgeItem + similarity float64 + } + + candidates := make([]candidate, 0) + rowNum := 0 + for rows.Next() { + rowNum++ + if rowNum%48 == 0 { + select { + case <-ctx.Done(): + return nil, ctx.Err() + default: + } + } + + var chunkID, itemID, chunkText, embeddingJSON, category, title, rowModel string + var chunkIndex, rowDim int + + if err := rows.Scan(&chunkID, &itemID, &chunkIndex, &chunkText, &embeddingJSON, &rowModel, &rowDim, &category, &title); err != nil { + r.logger.Warn("扫描向量失败", zap.Error(err)) + continue + } + + var embedding []float32 + if err := json.Unmarshal([]byte(embeddingJSON), &embedding); err != nil { + r.logger.Warn("解析向量失败", zap.Error(err)) + continue + } + + if rowDim > 0 && len(embedding) != rowDim { + r.logger.Debug("跳过维度不一致的向量行", zap.String("chunkId", chunkID), zap.Int("rowDim", rowDim), zap.Int("got", len(embedding))) + continue + } + if queryDim > 0 && len(embedding) != queryDim { + r.logger.Debug("跳过与查询维度不一致的向量", zap.String("chunkId", chunkID), zap.Int("queryDim", queryDim), zap.Int("got", len(embedding))) + continue + } + if expectedModel != "" && strings.TrimSpace(rowModel) != "" && strings.TrimSpace(rowModel) != expectedModel { + r.logger.Debug("跳过嵌入模型不一致的行", zap.String("chunkId", chunkID), zap.String("rowModel", rowModel), zap.String("expected", expectedModel)) + continue + } + + similarity := cosineSimilarity(queryEmbedding, embedding) + candidates = append(candidates, candidate{ + chunk: &KnowledgeChunk{ + ID: chunkID, + ItemID: itemID, + ChunkIndex: chunkIndex, + ChunkText: chunkText, + Embedding: embedding, + }, + item: &KnowledgeItem{ + ID: itemID, + Category: category, + Title: title, + }, + similarity: similarity, + }) + } + + sort.Slice(candidates, func(i, j int) bool { + return candidates[i].similarity > candidates[j].similarity + }) + + filtered := make([]candidate, 0, len(candidates)) + for _, c := range candidates { + if c.similarity >= threshold { + filtered = append(filtered, c) + } + } + + if len(filtered) > topK { + filtered = filtered[:topK] + } + + results := make([]*RetrievalResult, len(filtered)) + for i, c := range filtered { + results[i] = &RetrievalResult{ + Chunk: c.chunk, + Item: c.item, + Similarity: c.similarity, + Score: c.similarity, + } + } + return results, nil +} + +// RetrievalConfigFromYAML maps API/YAML retrieval settings into the knowledge package. +func RetrievalConfigFromYAML(r config.RetrievalConfig) *RetrievalConfig { + return &RetrievalConfig{ + TopK: r.TopK, + SimilarityThreshold: r.SimilarityThreshold, + SubIndexFilter: r.SubIndexFilter, + MultiQuery: r.MultiQuery, + Rerank: r.Rerank, + PostRetrieve: r.PostRetrieve, + } +} diff --git a/internal/knowledge/schema_migrate.go b/internal/knowledge/schema_migrate.go new file mode 100644 index 00000000..85fd26e2 --- /dev/null +++ b/internal/knowledge/schema_migrate.go @@ -0,0 +1,51 @@ +package knowledge + +import ( + "database/sql" + "fmt" +) + +// EnsureKnowledgeEmbeddingsSchema migrates knowledge_embeddings for sub_indexes + embedding metadata. +func EnsureKnowledgeEmbeddingsSchema(db *sql.DB) error { + if db == nil { + return fmt.Errorf("db is nil") + } + var n int + if err := db.QueryRow(`SELECT COUNT(*) FROM sqlite_master WHERE type='table' AND name='knowledge_embeddings'`).Scan(&n); err != nil { + return err + } + if n == 0 { + return nil + } + if err := addKnowledgeEmbeddingsColumnIfMissing(db, "sub_indexes", + `ALTER TABLE knowledge_embeddings ADD COLUMN sub_indexes TEXT NOT NULL DEFAULT ''`); err != nil { + return err + } + if err := addKnowledgeEmbeddingsColumnIfMissing(db, "embedding_model", + `ALTER TABLE knowledge_embeddings ADD COLUMN embedding_model TEXT NOT NULL DEFAULT ''`); err != nil { + return err + } + if err := addKnowledgeEmbeddingsColumnIfMissing(db, "embedding_dim", + `ALTER TABLE knowledge_embeddings ADD COLUMN embedding_dim INTEGER NOT NULL DEFAULT 0`); err != nil { + return err + } + return nil +} + +func addKnowledgeEmbeddingsColumnIfMissing(db *sql.DB, column, alterSQL string) error { + var colCount int + q := `SELECT COUNT(*) FROM pragma_table_info('knowledge_embeddings') WHERE name = ?` + if err := db.QueryRow(q, column).Scan(&colCount); err != nil { + return err + } + if colCount > 0 { + return nil + } + _, err := db.Exec(alterSQL) + return err +} + +// ensureKnowledgeEmbeddingsSubIndexesColumn 向后兼容；请使用 [EnsureKnowledgeEmbeddingsSchema]。 +func ensureKnowledgeEmbeddingsSubIndexesColumn(db *sql.DB) error { + return EnsureKnowledgeEmbeddingsSchema(db) +} diff --git a/internal/knowledge/tool.go b/internal/knowledge/tool.go new file mode 100644 index 00000000..ffceb7e2 --- /dev/null +++ b/internal/knowledge/tool.go @@ -0,0 +1,323 @@ +package knowledge + +import ( + "context" + "encoding/json" + "fmt" + "sort" + "strings" + + "cyberstrike-ai/internal/mcp" + "cyberstrike-ai/internal/mcp/builtin" + + "go.uber.org/zap" +) + +// RegisterKnowledgeTool 注册知识检索工具到MCP服务器 +func RegisterKnowledgeTool( + mcpServer *mcp.Server, + retriever *Retriever, + manager *Manager, + logger *zap.Logger, +) { + // 注册第一个工具：获取所有可用的风险类型列表 + listRiskTypesTool := mcp.Tool{ + Name: builtin.ToolListKnowledgeRiskTypes, + Description: "获取知识库中所有可用的风险类型（risk_type）列表。在搜索知识库之前，可以先调用此工具获取可用的风险类型，然后使用正确的风险类型进行精确搜索，这样可以大幅减少检索时间并提高检索准确性。", + ShortDescription: "获取知识库中所有可用的风险类型列表", + InputSchema: map[string]interface{}{ + "type": "object", + "properties": map[string]interface{}{}, + "required": []string{}, + }, + } + + listRiskTypesHandler := func(ctx context.Context, args map[string]interface{}) (*mcp.ToolResult, error) { + categories, err := manager.GetCategories() + if err != nil { + logger.Error("获取风险类型列表失败", zap.Error(err)) + return &mcp.ToolResult{ + Content: []mcp.Content{ + { + Type: "text", + Text: fmt.Sprintf("获取风险类型列表失败: %v", err), + }, + }, + IsError: true, + }, nil + } + + if len(categories) == 0 { + return &mcp.ToolResult{ + Content: []mcp.Content{ + { + Type: "text", + Text: "知识库中暂无风险类型。", + }, + }, + }, nil + } + + var resultText strings.Builder + resultText.WriteString(fmt.Sprintf("知识库中共有 %d 个风险类型：\n\n", len(categories))) + for i, category := range categories { + resultText.WriteString(fmt.Sprintf("%d. %s\n", i+1, category)) + } + resultText.WriteString("\n提示：在调用 " + builtin.ToolSearchKnowledgeBase + " 工具时，可以使用上述风险类型之一作为 risk_type 参数，以缩小搜索范围并提高检索效率。") + + return &mcp.ToolResult{ + Content: []mcp.Content{ + { + Type: "text", + Text: resultText.String(), + }, + }, + }, nil + } + + mcpServer.RegisterTool(listRiskTypesTool, listRiskTypesHandler) + logger.Debug("风险类型列表工具已注册", zap.String("toolName", listRiskTypesTool.Name)) + + // 注册第二个工具：搜索知识库（保持原有功能） + searchTool := mcp.Tool{ + Name: builtin.ToolSearchKnowledgeBase, + Description: "在知识库中搜索相关的安全知识。当你需要了解特定漏洞类型、攻击技术、检测方法等安全知识时，可以使用此工具进行检索。工具基于向量嵌入与余弦相似度检索（与 Eino retriever 语义一致）。建议：在搜索前可以先调用 " + builtin.ToolListKnowledgeRiskTypes + " 工具获取可用的风险类型，然后使用正确的 risk_type 参数进行精确搜索，这样可以大幅减少检索时间。", + ShortDescription: "搜索知识库中的安全知识（向量语义检索）", + InputSchema: map[string]interface{}{ + "type": "object", + "properties": map[string]interface{}{ + "query": map[string]interface{}{ + "type": "string", + "description": "搜索查询内容，描述你想要了解的安全知识主题", + }, + "risk_type": map[string]interface{}{ + "type": "string", + "description": "可选：指定风险类型（如：SQL注入、XSS、文件上传等）。建议先调用 " + builtin.ToolListKnowledgeRiskTypes + " 工具获取可用的风险类型列表，然后使用正确的风险类型进行精确搜索，这样可以大幅减少检索时间。如果不指定则搜索所有类型。", + }, + }, + "required": []string{"query"}, + }, + } + + searchHandler := func(ctx context.Context, args map[string]interface{}) (*mcp.ToolResult, error) { + query, ok := args["query"].(string) + if !ok || query == "" { + return &mcp.ToolResult{ + Content: []mcp.Content{ + { + Type: "text", + Text: "错误: 查询参数不能为空", + }, + }, + IsError: true, + }, nil + } + + riskType := "" + if rt, ok := args["risk_type"].(string); ok && rt != "" { + riskType = rt + } + + logger.Info("执行知识库检索", + zap.String("query", query), + zap.String("riskType", riskType), + ) + + // 检索统一走 Retriever.Search → VectorEinoRetriever（Eino retriever 语义）。 + searchReq := &SearchRequest{ + Query: query, + RiskType: riskType, + TopK: 5, + } + + results, err := retriever.Search(ctx, searchReq) + if err != nil { + logger.Error("知识库检索失败", zap.Error(err)) + return &mcp.ToolResult{ + Content: []mcp.Content{ + { + Type: "text", + Text: fmt.Sprintf("检索失败: %v", err), + }, + }, + IsError: true, + }, nil + } + + if len(results) == 0 { + return &mcp.ToolResult{ + Content: []mcp.Content{ + { + Type: "text", + Text: fmt.Sprintf("未找到与查询 '%s' 相关的知识。建议：\n1. 尝试使用不同的关键词\n2. 检查风险类型是否正确\n3. 确认知识库中是否包含相关内容", query), + }, + }, + }, nil + } + + // 格式化结果 + var resultText strings.Builder + + // 按余弦相似度（Score）降序 + sort.Slice(results, func(i, j int) bool { + return results[i].Score > results[j].Score + }) + + // 按文档分组结果，以便更好地展示上下文 + type itemGroup struct { + itemID string + results []*RetrievalResult + maxScore float64 // 该文档块的最高相似度 + } + itemGroups := make([]*itemGroup, 0) + itemMap := make(map[string]*itemGroup) + + for _, result := range results { + itemID := result.Item.ID + group, exists := itemMap[itemID] + if !exists { + group = &itemGroup{ + itemID: itemID, + results: make([]*RetrievalResult, 0), + maxScore: result.Score, + } + itemMap[itemID] = group + itemGroups = append(itemGroups, group) + } + group.results = append(group.results, result) + if result.Score > group.maxScore { + group.maxScore = result.Score + } + } + + // 按文档内最高相似度排序 + sort.Slice(itemGroups, func(i, j int) bool { + return itemGroups[i].maxScore > itemGroups[j].maxScore + }) + + // 收集检索到的知识项ID（用于日志） + retrievedItemIDs := make([]string, 0, len(itemGroups)) + + resultText.WriteString(fmt.Sprintf("找到 %d 条相关知识片段：\n\n", len(results))) + + resultIndex := 1 + for _, group := range itemGroups { + itemResults := group.results + mainResult := itemResults[0] + maxScore := mainResult.Score + for _, result := range itemResults { + if result.Score > maxScore { + maxScore = result.Score + mainResult = result + } + } + + // 按chunk_index排序，保证阅读的逻辑顺序（文档的原始顺序） + sort.Slice(itemResults, func(i, j int) bool { + return itemResults[i].Chunk.ChunkIndex < itemResults[j].Chunk.ChunkIndex + }) + + resultText.WriteString(fmt.Sprintf("--- 结果 %d (相似度: %.2f%%) ---\n", + resultIndex, mainResult.Similarity*100)) + resultText.WriteString(fmt.Sprintf("来源: [%s] %s (ID: %s)\n", mainResult.Item.Category, mainResult.Item.Title, mainResult.Item.ID)) + + // 按逻辑顺序显示所有chunk（包括主结果和扩展的chunk） + if len(itemResults) == 1 { + // 只有一个chunk，直接显示 + resultText.WriteString(fmt.Sprintf("内容片段:\n%s\n", mainResult.Chunk.ChunkText)) + } else { + // 多个chunk，按逻辑顺序显示 + resultText.WriteString("内容片段（按文档顺序）:\n") + for i, result := range itemResults { + // 标记主结果 + marker := "" + if result.Chunk.ID == mainResult.Chunk.ID { + marker = " [主匹配]" + } + resultText.WriteString(fmt.Sprintf(" [片段 %d%s]\n%s\n", i+1, marker, result.Chunk.ChunkText)) + } + } + resultText.WriteString("\n") + + if !contains(retrievedItemIDs, group.itemID) { + retrievedItemIDs = append(retrievedItemIDs, group.itemID) + } + resultIndex++ + } + + // 在结果末尾添加元数据（JSON格式，用于提取知识项ID） + // 使用特殊标记，避免影响AI阅读结果 + if len(retrievedItemIDs) > 0 { + metadataJSON, _ := json.Marshal(map[string]interface{}{ + "_metadata": map[string]interface{}{ + "retrievedItemIDs": retrievedItemIDs, + }, + }) + resultText.WriteString(fmt.Sprintf("\n", string(metadataJSON))) + } + + // 记录检索日志（异步，不阻塞） + // 注意：这里没有conversationID和messageID，需要在Agent层面记录 + // 实际的日志记录应该在Agent的progressCallback中完成 + + return &mcp.ToolResult{ + Content: []mcp.Content{ + { + Type: "text", + Text: resultText.String(), + }, + }, + }, nil + } + + mcpServer.RegisterTool(searchTool, searchHandler) + logger.Debug("知识检索工具已注册", zap.String("toolName", searchTool.Name)) +} + +// contains 检查切片是否包含元素 +func contains(slice []string, item string) bool { + for _, s := range slice { + if s == item { + return true + } + } + return false +} + +// GetRetrievalMetadata 从工具调用中提取检索元数据（用于日志记录） +func GetRetrievalMetadata(args map[string]interface{}) (query string, riskType string) { + if q, ok := args["query"].(string); ok { + query = q + } + if rt, ok := args["risk_type"].(string); ok { + riskType = rt + } + return +} + +// FormatRetrievalResults 格式化检索结果为字符串（用于日志） +func FormatRetrievalResults(results []*RetrievalResult) string { + if len(results) == 0 { + return "未找到相关结果" + } + + var builder strings.Builder + builder.WriteString(fmt.Sprintf("检索到 %d 条结果:\n", len(results))) + + itemIDs := make(map[string]bool) + for i, result := range results { + builder.WriteString(fmt.Sprintf("%d. [%s] %s (相似度: %.2f%%)\n", + i+1, result.Item.Category, result.Item.Title, result.Similarity*100)) + itemIDs[result.Item.ID] = true + } + + // 返回知识项ID列表（JSON格式） + ids := make([]string, 0, len(itemIDs)) + for id := range itemIDs { + ids = append(ids, id) + } + idsJSON, _ := json.Marshal(ids) + builder.WriteString(fmt.Sprintf("\n检索到的知识项ID: %s", string(idsJSON))) + + return builder.String() +} diff --git a/internal/knowledge/types.go b/internal/knowledge/types.go new file mode 100644 index 00000000..42e35e76 --- /dev/null +++ b/internal/knowledge/types.go @@ -0,0 +1,123 @@ +package knowledge + +import ( + "encoding/json" + "time" +) + +// formatTime 格式化时间为 RFC3339 格式，零时间返回空字符串 +func formatTime(t time.Time) string { + if t.IsZero() { + return "" + } + return t.Format(time.RFC3339) +} + +// KnowledgeItem 知识库项 +type KnowledgeItem struct { + ID string `json:"id"` + Category string `json:"category"` // 风险类型（文件夹名） + Title string `json:"title"` // 标题（文件名） + FilePath string `json:"filePath"` // 文件路径 + Content string `json:"content"` // 文件内容 + CreatedAt time.Time `json:"createdAt"` + UpdatedAt time.Time `json:"updatedAt"` +} + +// KnowledgeItemSummary 知识库项摘要（用于列表，不包含完整内容） +type KnowledgeItemSummary struct { + ID string `json:"id"` + Category string `json:"category"` + Title string `json:"title"` + FilePath string `json:"filePath"` + Content string `json:"content,omitempty"` // 可选：内容预览（如果提供，通常只包含前 150 字符） + CreatedAt time.Time `json:"createdAt"` + UpdatedAt time.Time `json:"updatedAt"` +} + +// MarshalJSON 自定义 JSON 序列化，确保时间格式正确 +func (k *KnowledgeItemSummary) MarshalJSON() ([]byte, error) { + type Alias KnowledgeItemSummary + aux := &struct { + *Alias + CreatedAt string `json:"createdAt"` + UpdatedAt string `json:"updatedAt"` + }{ + Alias: (*Alias)(k), + } + aux.CreatedAt = formatTime(k.CreatedAt) + aux.UpdatedAt = formatTime(k.UpdatedAt) + return json.Marshal(aux) +} + +// MarshalJSON 自定义 JSON 序列化，确保时间格式正确 +func (k *KnowledgeItem) MarshalJSON() ([]byte, error) { + type Alias KnowledgeItem + aux := &struct { + *Alias + CreatedAt string `json:"createdAt"` + UpdatedAt string `json:"updatedAt"` + }{ + Alias: (*Alias)(k), + } + aux.CreatedAt = formatTime(k.CreatedAt) + aux.UpdatedAt = formatTime(k.UpdatedAt) + return json.Marshal(aux) +} + +// KnowledgeChunk 知识块（用于向量化） +type KnowledgeChunk struct { + ID string `json:"id"` + ItemID string `json:"itemId"` + ChunkIndex int `json:"chunkIndex"` + ChunkText string `json:"chunkText"` + Embedding []float32 `json:"-"` // 向量嵌入，不序列化到 JSON + CreatedAt time.Time `json:"createdAt"` +} + +// RetrievalResult 检索结果 +type RetrievalResult struct { + Chunk *KnowledgeChunk `json:"chunk"` + Item *KnowledgeItem `json:"item"` + Similarity float64 `json:"similarity"` // 相似度分数 + Score float64 `json:"score"` // 与 Similarity 相同：余弦相似度 +} + +// RetrievalLog 检索日志 +type RetrievalLog struct { + ID string `json:"id"` + ConversationID string `json:"conversationId,omitempty"` + MessageID string `json:"messageId,omitempty"` + Query string `json:"query"` + RiskType string `json:"riskType,omitempty"` + RetrievedItems []string `json:"retrievedItems"` // 检索到的知识项 ID 列表 + CreatedAt time.Time `json:"createdAt"` +} + +// MarshalJSON 自定义 JSON 序列化，确保时间格式正确 +func (r *RetrievalLog) MarshalJSON() ([]byte, error) { + type Alias RetrievalLog + return json.Marshal(&struct { + *Alias + CreatedAt string `json:"createdAt"` + }{ + Alias: (*Alias)(r), + CreatedAt: formatTime(r.CreatedAt), + }) +} + +// CategoryWithItems 分类及其下的知识项（用于按分类分页） +type CategoryWithItems struct { + Category string `json:"category"` // 分类名称 + ItemCount int `json:"itemCount"` // 该分类下的知识项总数 + Items []*KnowledgeItemSummary `json:"items"` // 该分类下的知识项列表 +} + +// SearchRequest 搜索请求 +type SearchRequest struct { + Query string `json:"query"` + RiskType string `json:"riskType,omitempty"` // 可选：指定风险类型 + SubIndexFilter string `json:"subIndexFilter,omitempty"` // 可选：仅保留 sub_indexes 含该标签的行（含未打标旧数据） + TopK int `json:"topK,omitempty"` // 返回 Top-K 结果，默认 5 + Threshold float64 `json:"threshold,omitempty"` // 相似度阈值，默认 0.7 +} diff --git a/internal/knowledge/wire_retriever.go b/internal/knowledge/wire_retriever.go new file mode 100644 index 00000000..5cca7258 --- /dev/null +++ b/internal/knowledge/wire_retriever.go @@ -0,0 +1,95 @@ +package knowledge + +import ( + "context" + "fmt" + "net/http" + "strings" + "time" + + "cyberstrike-ai/internal/config" + "cyberstrike-ai/internal/llm" + "cyberstrike-ai/internal/openai" + + einoopenai "github.com/cloudwego/eino-ext/components/model/openai" + "github.com/cloudwego/eino/components/model" + "github.com/cloudwego/eino/flow/retriever/multiquery" + "go.uber.org/zap" +) + +// WireRetrieverPipeline builds Eino MultiQuery + HTTP rerank + post-process pipeline on r. +// Call once after NewRetriever; UpdateConfig re-invokes when wireOpenAI is set. +func WireRetrieverPipeline(ctx context.Context, r *Retriever, openAI *config.OpenAIConfig) error { + if r == nil { + return fmt.Errorf("retriever is nil") + } + if openAI == nil { + return fmt.Errorf("openai config is nil") + } + if r.config == nil { + return fmt.Errorf("retrieval config is nil") + } + r.wireOpenAI = openAI + + baseHTTPClient := &http.Client{Timeout: 120 * time.Second} + var rewriteLLM model.ChatModel + if llm.IsClaudeProvider(openAI.Provider) { + nativeModel, err := llm.NewClaudeAgenticModel( + ctx, + *openAI, + baseHTTPClient, + openAI.MaxCompletionTokensEffective(), + nil, + ) + if err != nil { + return fmt.Errorf("multi_query native Claude rewrite model: %w", err) + } + rewriteLLM = llm.NewAgenticChatModelAdapter(nativeModel) + } else { + httpClient := openai.NewEinoHTTPClient(openAI, baseHTTPClient) + maxCompletionTokens := openAI.MaxCompletionTokensEffective() + chatCfg := &einoopenai.ChatModelConfig{ + APIKey: strings.TrimSpace(openAI.APIKey), + BaseURL: strings.TrimSuffix(strings.TrimSpace(openAI.BaseURL), "/"), + Model: strings.TrimSpace(openAI.Model), + HTTPClient: httpClient, + MaxCompletionTokens: &maxCompletionTokens, + } + if chatCfg.Model == "" { + chatCfg.Model = "gpt-4o" + } + var err error + rewriteLLM, err = einoopenai.NewChatModel(ctx, chatCfg) + if err != nil { + return fmt.Errorf("multi_query rewrite model: %w", err) + } + } + + reranker, err := NewHTTPReranker(&r.config.Rerank, openAI, r.logger) + if err != nil { + return fmt.Errorf("reranker: %w", err) + } + r.SetDocumentReranker(reranker) + + vec := NewVectorEinoRetriever(r) + mq, err := multiquery.NewRetriever(ctx, &multiquery.Config{ + RewriteLLM: rewriteLLM, + MaxQueriesNum: r.config.MultiQuery.MaxQueriesEffective(), + OrigRetriever: vec, + }) + if err != nil { + return fmt.Errorf("multi_query: %w", err) + } + + r.pipeline = newKnowledgePipelineRetriever(mq, r) + if r.logger != nil { + provider := r.config.Rerank.ProviderEffective(strings.TrimSpace(openAI.BaseURL)) + r.logger.Info("知识库检索流水线已启用", + zap.String("pipeline", "MultiQuery→Vector→Rerank→PostRetrieve"), + zap.Int("multi_query_max", r.config.MultiQuery.MaxQueriesEffective()), + zap.String("rerank_provider", provider), + zap.String("rerank_model", r.config.Rerank.ModelEffective(provider)), + ) + } + return nil +} diff --git a/internal/openai/claude_native.go b/internal/openai/claude_native.go new file mode 100644 index 00000000..d084139f --- /dev/null +++ b/internal/openai/claude_native.go @@ -0,0 +1,245 @@ +package openai + +import ( + "context" + "encoding/json" + "errors" + "fmt" + "io" + "strings" + + "cyberstrike-ai/internal/llm" + + "github.com/cloudwego/eino/components/model" + "github.com/cloudwego/eino/schema" +) + +type claudeNativePayload struct { + Model string `json:"model"` + Messages []claudeNativeMessage `json:"messages"` + Temperature *float32 `json:"temperature,omitempty"` + TopP *float32 `json:"top_p,omitempty"` + MaxCompletionTokens int `json:"max_completion_tokens,omitempty"` + MaxTokens int `json:"max_tokens,omitempty"` + Thinking any `json:"thinking,omitempty"` + OutputConfig any `json:"output_config,omitempty"` +} + +type claudeNativeMessage struct { + Role string `json:"role"` + Content json.RawMessage `json:"content"` +} + +func (c *Client) isClaude() bool { + return c != nil && c.config != nil && llm.IsClaudeProvider(c.config.Provider) +} + +func (c *Client) claudeNativeChatCompletion(ctx context.Context, payload, out any) error { + req, err := c.parseClaudeNativePayload(payload) + if err != nil { + return err + } + nativeModel, err := llm.NewClaudeAgenticModel( + ctx, + *c.config, + c.httpClient, + req.maxTokens(c.config.MaxCompletionTokensEffective()), + req.extraFields(), + ) + if err != nil { + return fmt.Errorf("create native Claude model: %w", err) + } + resp, err := nativeModel.Generate(ctx, req.agenticMessages(), req.options()...) + if err != nil { + return fmt.Errorf("native Claude generate: %w", err) + } + return marshalClaudeNativeResponse(resp, req.model(c.config.Model), out) +} + +func (c *Client) claudeNativeChatCompletionStream( + ctx context.Context, + payload any, + onDelta func(delta string) error, +) (string, error) { + req, err := c.parseClaudeNativePayload(payload) + if err != nil { + return "", err + } + nativeModel, err := llm.NewClaudeAgenticModel( + ctx, + *c.config, + c.httpClient, + req.maxTokens(c.config.MaxCompletionTokensEffective()), + req.extraFields(), + ) + if err != nil { + return "", fmt.Errorf("create native Claude model: %w", err) + } + stream, err := nativeModel.Stream(ctx, req.agenticMessages(), req.options()...) + if err != nil { + return "", fmt.Errorf("native Claude stream: %w", err) + } + defer stream.Close() + + var full strings.Builder + for { + chunk, recvErr := stream.Recv() + if errors.Is(recvErr, io.EOF) { + return full.String(), nil + } + if recvErr != nil { + return full.String(), fmt.Errorf("native Claude stream receive: %w", recvErr) + } + content, _ := llm.AgenticText(chunk) + if content == "" { + continue + } + full.WriteString(content) + if onDelta != nil { + if err := onDelta(content); err != nil { + return full.String(), err + } + } + } +} + +func (c *Client) parseClaudeNativePayload(payload any) (*claudeNativePayload, error) { + raw, err := json.Marshal(payload) + if err != nil { + return nil, fmt.Errorf("marshal Claude payload: %w", err) + } + var req claudeNativePayload + if err := json.Unmarshal(raw, &req); err != nil { + return nil, fmt.Errorf("unmarshal Claude payload: %w", err) + } + if strings.TrimSpace(req.model(c.config.Model)) == "" { + return nil, fmt.Errorf("native Claude model is empty") + } + if len(req.Messages) == 0 { + return nil, fmt.Errorf("native Claude messages are empty") + } + return &req, nil +} + +func (p *claudeNativePayload) model(fallback string) string { + if modelName := strings.TrimSpace(p.Model); modelName != "" { + return modelName + } + return strings.TrimSpace(fallback) +} + +func (p *claudeNativePayload) maxTokens(fallback int) int { + if p.MaxCompletionTokens > 0 { + return p.MaxCompletionTokens + } + if p.MaxTokens > 0 { + return p.MaxTokens + } + return fallback +} + +func (p *claudeNativePayload) extraFields() map[string]any { + fields := make(map[string]any, 2) + if p.Thinking != nil { + fields["thinking"] = p.Thinking + } + if p.OutputConfig != nil { + fields["output_config"] = p.OutputConfig + } + if len(fields) == 0 { + return nil + } + return fields +} + +func (p *claudeNativePayload) options() []model.Option { + opts := make([]model.Option, 0, 3) + if p.Temperature != nil { + opts = append(opts, model.WithTemperature(*p.Temperature)) + } + if p.TopP != nil { + opts = append(opts, model.WithTopP(*p.TopP)) + } + if p.MaxCompletionTokens > 0 || p.MaxTokens > 0 { + opts = append(opts, model.WithMaxTokens(p.maxTokens(0))) + } + return opts +} + +func (p *claudeNativePayload) agenticMessages() []*schema.AgenticMessage { + out := make([]*schema.AgenticMessage, 0, len(p.Messages)) + for _, msg := range p.Messages { + text := claudeNativeTextContent(msg.Content) + role := schema.AgenticRoleTypeUser + var block *schema.ContentBlock + switch strings.ToLower(strings.TrimSpace(msg.Role)) { + case "system": + role = schema.AgenticRoleTypeSystem + block = schema.NewContentBlock(&schema.UserInputText{Text: text}) + case "assistant": + role = schema.AgenticRoleTypeAssistant + block = schema.NewContentBlock(&schema.AssistantGenText{Text: text}) + default: + block = schema.NewContentBlock(&schema.UserInputText{Text: text}) + } + out = append(out, &schema.AgenticMessage{ + Role: role, + ContentBlocks: []*schema.ContentBlock{block}, + }) + } + return out +} + +func claudeNativeTextContent(raw json.RawMessage) string { + var text string + if err := json.Unmarshal(raw, &text); err == nil { + return text + } + var parts []struct { + Type string `json:"type"` + Text string `json:"text"` + } + if err := json.Unmarshal(raw, &parts); err == nil { + var out strings.Builder + for _, part := range parts { + if part.Type == "" || part.Type == "text" { + out.WriteString(part.Text) + } + } + return out.String() + } + return strings.TrimSpace(string(raw)) +} + +func marshalClaudeNativeResponse(resp *schema.AgenticMessage, modelName string, out any) error { + if out == nil { + return nil + } + content, reasoning := llm.AgenticText(resp) + id := "" + if resp != nil && resp.ResponseMeta != nil && resp.ResponseMeta.ClaudeExtension != nil { + id = resp.ResponseMeta.ClaudeExtension.ID + } + wire := map[string]any{ + "id": id, + "object": "chat.completion", + "model": modelName, + "choices": []any{map[string]any{ + "index": 0, + "message": map[string]any{ + "role": "assistant", + "content": content, + "reasoning_content": reasoning, + }, + "finish_reason": "stop", + }}, + } + raw, err := json.Marshal(wire) + if err != nil { + return fmt.Errorf("marshal native Claude response: %w", err) + } + if err := json.Unmarshal(raw, out); err != nil { + return fmt.Errorf("unmarshal native Claude response: %w", err) + } + return nil +} diff --git a/internal/openai/claude_native_test.go b/internal/openai/claude_native_test.go new file mode 100644 index 00000000..c4420bd5 --- /dev/null +++ b/internal/openai/claude_native_test.go @@ -0,0 +1,70 @@ +package openai + +import ( + "context" + "io" + "net/http" + "strings" + "testing" + + "cyberstrike-ai/internal/config" +) + +type roundTripFunc func(*http.Request) (*http.Response, error) + +func (f roundTripFunc) RoundTrip(req *http.Request) (*http.Response, error) { + return f(req) +} + +func TestChatCompletionUsesNativeClaudeMessagesAPI(t *testing.T) { + t.Parallel() + httpClient := &http.Client{Transport: roundTripFunc(func(req *http.Request) (*http.Response, error) { + if req.URL.Path != "/v1/messages" { + t.Fatalf("request path = %q, want /v1/messages", req.URL.Path) + } + if req.Header.Get("x-api-key") != "test-key" { + t.Fatalf("x-api-key header = %q", req.Header.Get("x-api-key")) + } + return &http.Response{ + StatusCode: http.StatusOK, + Header: http.Header{"Content-Type": []string{"application/json"}}, + Body: io.NopCloser(strings.NewReader(`{ + "id":"msg_1", + "type":"message", + "role":"assistant", + "model":"claude-test", + "content":[{"type":"text","text":"native ok"}], + "stop_reason":"end_turn", + "usage":{"input_tokens":1,"output_tokens":2} + }`)), + Request: req, + }, nil + })} + client := NewClient(&config.OpenAIConfig{ + Provider: "claude", + BaseURL: "https://example.test", + APIKey: "test-key", + Model: "claude-test", + }, httpClient, nil) + + var out struct { + Choices []struct { + Message struct { + Content string `json:"content"` + } `json:"message"` + } `json:"choices"` + } + err := client.ChatCompletion(context.Background(), map[string]any{ + "model": "claude-test", + "messages": []map[string]string{ + {"role": "user", "content": "hello"}, + }, + "max_completion_tokens": 16, + }, &out) + if err != nil { + t.Fatalf("ChatCompletion: %v", err) + } + if len(out.Choices) != 1 || out.Choices[0].Message.Content != "native ok" { + t.Fatalf("response = %#v", out) + } +} diff --git a/internal/openai/claude_reasoning_legacy.go b/internal/openai/claude_reasoning_legacy.go new file mode 100644 index 00000000..5dd32ccf --- /dev/null +++ b/internal/openai/claude_reasoning_legacy.go @@ -0,0 +1,21 @@ +package openai + +import "strings" + +// claudeReasoningRoundTripSep is retained only to render historical traces +// written by the removed OpenAI-to-Claude HTTP bridge. +const claudeReasoningRoundTripSep = "\n---CSAI_CLAUDE_THINKING_BLOCKS---\n" + +// DisplayReasoningContent strips the obsolete bridge metadata suffix from +// historical records. Native AgenticMessage reasoning does not add this suffix. +func DisplayReasoningContent(s string) string { + s = strings.TrimSpace(s) + if s == "" { + return "" + } + i := strings.LastIndex(s, claudeReasoningRoundTripSep) + if i < 0 { + return s + } + return strings.TrimSpace(s[:i]) +} diff --git a/internal/openai/claude_reasoning_legacy_test.go b/internal/openai/claude_reasoning_legacy_test.go new file mode 100644 index 00000000..6b24d6a0 --- /dev/null +++ b/internal/openai/claude_reasoning_legacy_test.go @@ -0,0 +1,10 @@ +package openai + +import "testing" + +func TestDisplayReasoningContentStripsLegacyClaudeSuffix(t *testing.T) { + raw := "hello" + claudeReasoningRoundTripSep + `[{"type":"thinking"}]` + if got := DisplayReasoningContent(raw); got != "hello" { + t.Fatalf("DisplayReasoningContent() = %q, want hello", got) + } +} diff --git a/internal/openai/eino_http.go b/internal/openai/eino_http.go new file mode 100644 index 00000000..7c2a8bc5 --- /dev/null +++ b/internal/openai/eino_http.go @@ -0,0 +1,24 @@ +package openai + +import ( + "net/http" + + "cyberstrike-ai/internal/config" +) + +// NewEinoHTTPClient adds OpenAI-compatible request fixes and SSE sanitation. +// Claude channels use Eino's native agenticclaude model and never enter here. +func NewEinoHTTPClient(cfg *config.OpenAIConfig, base *http.Client) *http.Client { + if base == nil { + base = http.DefaultClient + } + cloned := *base + transport := base.Transport + if transport == nil { + transport = http.DefaultTransport + } + transport = &reasoningToolChoiceCompatRoundTripper{base: transport, cfg: cfg} + transport = &einoSSESanitizingRoundTripper{base: transport} + cloned.Transport = transport + return &cloned +} diff --git a/internal/openai/eino_sse_sanitizer.go b/internal/openai/eino_sse_sanitizer.go new file mode 100644 index 00000000..43e07d5b --- /dev/null +++ b/internal/openai/eino_sse_sanitizer.go @@ -0,0 +1,149 @@ +package openai + +// eino_sse_sanitizer.go 解决 Eino 走 meguminnnnnnnnn/go-openai SDK 时， +// 中转站心跳/SSE 控制行累计 > 300 行触发 ErrTooManyEmptyStreamMessages +// （报错文案: "stream has sent too many empty messages"）的问题。 +// +// 触发链路: +// einoopenai.NewChatModel +// → eino-ext/libs/acl/openai → meguminnnnnnnnn/go-openai +// → streamReader.processLines() 对所有非 "data:" 行计数, > 300 即抛错。 +// +// 中转站常见的非 data: 行（合法 SSE 但 SDK 不接受）: +// ":" / ": keepalive" / ": ping" / "event: ping" / "retry: 3000" +// 以及思考型模型 prefill 期间穿插的大量心跳。 +// +// 兜底策略: 在 HTTP transport 层把响应 Body 包一层 reader, 只放行 "data:" +// 开头的行, 把心跳/注释/事件类型行就地吞掉。下游 SDK 永远见不到非 data: 行, +// 计数器始终为 0, 该错误不可能再发生。 +// +// 该层对调用方完全透明: +// - 仅当响应 Content-Type 是 text/event-stream 时介入；普通 JSON 响应原样透传 +// - data: payload (含 [DONE] 与 {"error":...}) 一字节不改 +// - 上游真断流 (EOF / connection reset / context cancel) 原样透传 + +import ( + "bufio" + "bytes" + "io" + "net/http" + "strings" +) + +const ( + // einoSSEReaderBufSize 给 bufio 一个较大的初始缓冲, 避免单行大 JSON chunk + // (含工具调用 arguments / reasoning_content) 频繁触发缓冲区扩容。 + einoSSEReaderBufSize = 64 * 1024 +) + +// einoSSESanitizingRoundTripper 包装下游 RoundTripper, 对 SSE 响应做行级清洗。 +type einoSSESanitizingRoundTripper struct { + base http.RoundTripper +} + +func (rt *einoSSESanitizingRoundTripper) RoundTrip(req *http.Request) (*http.Response, error) { + resp, err := rt.base.RoundTrip(req) + if err != nil || resp == nil { + return resp, err + } + if !isSSEResponse(resp) { + return resp, nil + } + resp.Body = newEinoSSESanitizingBody(resp.Body) + return resp, nil +} + +// isSSEResponse 仅对 200 + text/event-stream 的响应做清洗; +// 错误响应 (4xx/5xx 通常是 application/json) 不动, 由 SDK 走原错误路径。 +func isSSEResponse(resp *http.Response) bool { + if resp.StatusCode != http.StatusOK { + return false + } + ct := resp.Header.Get("Content-Type") + if ct == "" { + return false + } + ct = strings.ToLower(strings.TrimSpace(ct)) + // 兼容 "text/event-stream", "text/event-stream; charset=utf-8" 等。 + return strings.HasPrefix(ct, "text/event-stream") +} + +// einoSSESanitizingBody 是包装后的响应体: 只放行 data: 行, 其它行吞掉。 +type einoSSESanitizingBody struct { + upstream io.ReadCloser + reader *bufio.Reader + pending []byte // 已清洗、待返回给下游的字节 (永远以 \n 结尾的完整 data: 行) + err error // upstream 终态错误 (io.EOF 或网络错误) +} + +func newEinoSSESanitizingBody(body io.ReadCloser) *einoSSESanitizingBody { + return &einoSSESanitizingBody{ + upstream: body, + reader: bufio.NewReaderSize(body, einoSSEReaderBufSize), + } +} + +func (b *einoSSESanitizingBody) Read(p []byte) (int, error) { + if len(p) == 0 { + return 0, nil + } + if len(b.pending) > 0 { + n := copy(p, b.pending) + b.pending = b.pending[n:] + return n, nil + } + + // 从上游读, 直到攒出一行 data: 或拿到终态。 + // 单次循环可能丢弃任意多行心跳, 但只放行至多一行 data: 后退出, + // 避免一次 Read 阻塞过久 / pending 缓冲过大。 + for b.err == nil { + line, err := b.reader.ReadBytes('\n') + if len(line) > 0 { + if isPassThroughSSELine(line) { + if line[len(line)-1] != '\n' { + line = append(line, '\n') + } + b.pending = line + if err != nil { + b.err = err + } + break + } + // 非 data: 行 (空行 / ":" 注释 / event: / retry: / id: / 任何裸文本) + // 全部吞掉, 不向下游透出, 继续循环读下一行。 + } + if err != nil { + b.err = err + break + } + } + + if len(b.pending) > 0 { + n := copy(p, b.pending) + b.pending = b.pending[n:] + return n, nil + } + return 0, b.err +} + +func (b *einoSSESanitizingBody) Close() error { + return b.upstream.Close() +} + +// isPassThroughSSELine 判定该行是否需要原样放行给下游 SDK。 +// 仅 "data:" (大小写不敏感, 可有任意前导空白) 开头的行需要保留。 +// 注意: 不能用 TrimSpace 去尾部换行后再判, 否则 " data: x" 会被误判; +// 我们只 trim 前导空白, 与 SDK 内部 TrimSpace 后再正则 ^data:\s* 的语义一致。 +func isPassThroughSSELine(line []byte) bool { + trimmed := bytes.TrimLeft(line, " \t") + if len(trimmed) < 5 { + return false + } + // 大小写不敏感比较前 5 字节是否为 "data:"。SSE 规范要求字段名小写, + // 但宽松匹配可以兼容个别中转站的非规范实现。 + return (trimmed[0] == 'd' || trimmed[0] == 'D') && + (trimmed[1] == 'a' || trimmed[1] == 'A') && + (trimmed[2] == 't' || trimmed[2] == 'T') && + (trimmed[3] == 'a' || trimmed[3] == 'A') && + trimmed[4] == ':' +} diff --git a/internal/openai/eino_sse_sanitizer_test.go b/internal/openai/eino_sse_sanitizer_test.go new file mode 100644 index 00000000..ef52db39 --- /dev/null +++ b/internal/openai/eino_sse_sanitizer_test.go @@ -0,0 +1,303 @@ +package openai + +import ( + "bufio" + "bytes" + "errors" + "io" + "net/http" + "net/http/httptest" + "regexp" + "strings" + "testing" +) + +// 复现 meguminnnnnnnnn/go-openai 的 SSE 行计数算法 (默认 limit=300): +// - 逐行读 +// - 非 "data:" 行 (空行 / ":" 注释 / event: / retry:) 累计 emptyMessagesCount +// - > 300 抛 ErrTooManyEmptyStreamMessages +// - 遇到 data: 行 reset, 返回 payload +// +// 这一算法与上游 SDK 的 stream_reader.go processLines() 严格一致 (验证依据见 +// /Users/temp/go/pkg/mod/github.com/meguminnnnnnnnn/go-openai@v0.1.2/stream_reader.go)。 +// 测试中只复刻 "限制触发" 这一行为, 用来回归验证 sanitizer 的根因修复。 +var errTooManyEmptyStreamMessages = errors.New("stream has sent too many empty messages") + +func sdkLikeRecvAll(body io.Reader, limit uint) ([]string, error) { + headerData := regexp.MustCompile(`^data:\s*`) + r := bufio.NewReader(body) + var payloads []string + for { + var emptyMessagesCount uint + var payload []byte + for { + line, err := r.ReadBytes('\n') + if err != nil { + if err == io.EOF { + return payloads, nil + } + return payloads, err + } + noSpace := bytes.TrimSpace(line) + if !headerData.Match(noSpace) { + emptyMessagesCount++ + if emptyMessagesCount > limit { + return payloads, errTooManyEmptyStreamMessages + } + continue + } + payload = headerData.ReplaceAll(noSpace, nil) + break + } + if string(payload) == "[DONE]" { + return payloads, nil + } + payloads = append(payloads, string(payload)) + } +} + +func newSSEServer(t *testing.T, body string, contentType string, status int) *httptest.Server { + t.Helper() + return httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) { + if contentType != "" { + w.Header().Set("Content-Type", contentType) + } + w.WriteHeader(status) + _, _ = io.WriteString(w, body) + })) +} + +func sanitizingClient(base *http.Client) *http.Client { + if base == nil { + base = &http.Client{} + } + cloned := *base + transport := base.Transport + if transport == nil { + transport = http.DefaultTransport + } + cloned.Transport = &einoSSESanitizingRoundTripper{base: transport} + return &cloned +} + +func readAll(t *testing.T, body io.ReadCloser) string { + t.Helper() + defer body.Close() + out, err := io.ReadAll(body) + if err != nil { + t.Fatalf("read body: %v", err) + } + return string(out) +} + +// 1) 仅 data: 行 → 一字节不改地透传。 +func TestSSESanitizer_PassesDataLinesUnchanged(t *testing.T) { + body := "data: {\"a\":1}\ndata: {\"b\":2}\ndata: [DONE]\n" + srv := newSSEServer(t, body, "text/event-stream", 200) + defer srv.Close() + + resp, err := sanitizingClient(nil).Get(srv.URL) + if err != nil { + t.Fatalf("get: %v", err) + } + got := readAll(t, resp.Body) + if got != body { + t.Fatalf("body mismatch:\nwant %q\ngot %q", body, got) + } +} + +// 2) 心跳/注释/事件类型行被吞掉, 仅保留 data: 行。 +func TestSSESanitizer_DropsHeartbeatsAndControlLines(t *testing.T) { + body := strings.Join([]string{ + ": keepalive", + "", + "event: ping", + "retry: 3000", + "id: 42", + "data: {\"x\":1}", + ": ping", + "", + "data: {\"x\":2}", + "data: [DONE]", + "", + }, "\n") + srv := newSSEServer(t, body, "text/event-stream", 200) + defer srv.Close() + + resp, err := sanitizingClient(nil).Get(srv.URL) + if err != nil { + t.Fatalf("get: %v", err) + } + got := readAll(t, resp.Body) + want := "data: {\"x\":1}\ndata: {\"x\":2}\ndata: [DONE]\n" + if got != want { + t.Fatalf("sanitized body mismatch:\nwant %q\ngot %q", want, got) + } +} + +// 3) 根因回归: 上游堆 500 行心跳后才发 data:, 原始 SDK 算法会抛 +// ErrTooManyEmptyStreamMessages, sanitize 之后必须能正常拿到所有 data:。 +func TestSSESanitizer_ProtectsAgainstTooManyEmptyMessages(t *testing.T) { + const heartbeats = 500 + var buf bytes.Buffer + for i := 0; i < heartbeats; i++ { + buf.WriteString(": keepalive\n") + } + buf.WriteString("data: {\"chunk\":1}\n") + buf.WriteString("data: {\"chunk\":2}\n") + buf.WriteString("data: [DONE]\n") + + t.Run("baseline_without_sanitizer_must_fail", func(t *testing.T) { + _, err := sdkLikeRecvAll(bytes.NewReader(buf.Bytes()), 300) + if !errors.Is(err, errTooManyEmptyStreamMessages) { + t.Fatalf("expected ErrTooManyEmptyStreamMessages, got %v", err) + } + }) + + t.Run("with_sanitizer_must_succeed", func(t *testing.T) { + srv := newSSEServer(t, buf.String(), "text/event-stream", 200) + defer srv.Close() + + resp, err := sanitizingClient(nil).Get(srv.URL) + if err != nil { + t.Fatalf("get: %v", err) + } + defer resp.Body.Close() + + payloads, err := sdkLikeRecvAll(resp.Body, 300) + if err != nil { + t.Fatalf("sdk-like recv after sanitize: %v", err) + } + want := []string{`{"chunk":1}`, `{"chunk":2}`} + if len(payloads) != len(want) { + t.Fatalf("payload count mismatch: want %d got %d (%v)", len(want), len(payloads), payloads) + } + for i, w := range want { + if payloads[i] != w { + t.Fatalf("payload[%d] mismatch: want %q got %q", i, w, payloads[i]) + } + } + }) +} + +// 4) 心跳穿插在 data: 之间也能正确清洗 (思考型模型 prefill 期间常见)。 +func TestSSESanitizer_HeartbeatsInterleavedWithData(t *testing.T) { + var buf bytes.Buffer + buf.WriteString("data: {\"chunk\":1}\n") + for i := 0; i < 400; i++ { + buf.WriteString(": keepalive\n") + } + buf.WriteString("data: {\"chunk\":2}\n") + buf.WriteString("data: [DONE]\n") + + srv := newSSEServer(t, buf.String(), "text/event-stream", 200) + defer srv.Close() + + resp, err := sanitizingClient(nil).Get(srv.URL) + if err != nil { + t.Fatalf("get: %v", err) + } + defer resp.Body.Close() + + payloads, err := sdkLikeRecvAll(resp.Body, 300) + if err != nil { + t.Fatalf("sdk-like recv: %v", err) + } + if got, want := len(payloads), 2; got != want { + t.Fatalf("payload count: want %d got %d", want, got) + } +} + +// 5) 非 SSE 响应 (例如非流式 JSON) 不应被 sanitizer 介入。 +func TestSSESanitizer_PassesNonSSEResponseUntouched(t *testing.T) { + body := `{"id":"x","object":"chat.completion","choices":[]}` + srv := newSSEServer(t, body, "application/json", 200) + defer srv.Close() + + resp, err := sanitizingClient(nil).Get(srv.URL) + if err != nil { + t.Fatalf("get: %v", err) + } + got := readAll(t, resp.Body) + if got != body { + t.Fatalf("non-SSE body must be untouched:\nwant %q\ngot %q", body, got) + } +} + +// 6) 错误响应 (4xx/5xx) 不应被 sanitize, 即使 Content-Type 是 SSE 也不动, +// 避免吞掉类似 "data: " 之外的错误正文。 +func TestSSESanitizer_PassesNon200Untouched(t *testing.T) { + body := `{"error":{"message":"rate limit"}}` + srv := newSSEServer(t, body, "text/event-stream", 429) + defer srv.Close() + + resp, err := sanitizingClient(nil).Get(srv.URL) + if err != nil { + t.Fatalf("get: %v", err) + } + got := readAll(t, resp.Body) + if got != body { + t.Fatalf("error body must be untouched:\nwant %q\ngot %q", body, got) + } +} + +// 7) data: 行末尾若缺 \n (异常上游) sanitizer 也补齐, 保证下游按行解析。 +func TestSSESanitizer_AppendsTrailingNewlineIfMissing(t *testing.T) { + body := "data: {\"a\":1}" + srv := newSSEServer(t, body, "text/event-stream", 200) + defer srv.Close() + + resp, err := sanitizingClient(nil).Get(srv.URL) + if err != nil { + t.Fatalf("get: %v", err) + } + got := readAll(t, resp.Body) + want := "data: {\"a\":1}\n" + if got != want { + t.Fatalf("trailing newline:\nwant %q\ngot %q", want, got) + } +} + +// 8) 大 chunk (一行数十 KB) 也能完整透传, 不被切断。 +func TestSSESanitizer_LargeDataLinePassesIntact(t *testing.T) { + huge := strings.Repeat("x", 80*1024) + body := "data: {\"big\":\"" + huge + "\"}\ndata: [DONE]\n" + srv := newSSEServer(t, body, "text/event-stream", 200) + defer srv.Close() + + resp, err := sanitizingClient(nil).Get(srv.URL) + if err != nil { + t.Fatalf("get: %v", err) + } + got := readAll(t, resp.Body) + if got != body { + t.Fatalf("large body length mismatch: want %d got %d", len(body), len(got)) + } +} + +// 9) isPassThroughSSELine 单元覆盖。 +func TestIsPassThroughSSELine(t *testing.T) { + cases := []struct { + line string + want bool + }{ + {"data: {\"a\":1}\n", true}, + {"DATA: x\n", true}, + {" data: x\n", true}, + {"data:\n", true}, + {"\n", false}, + {"\r\n", false}, + {": keepalive\n", false}, + {":\n", false}, + {"event: ping\n", false}, + {"retry: 3000\n", false}, + {"id: 42\n", false}, + {"datax: y\n", false}, + {"da", false}, + } + for _, c := range cases { + if got := isPassThroughSSELine([]byte(c.line)); got != c.want { + t.Errorf("isPassThroughSSELine(%q) = %v, want %v", c.line, got, c.want) + } + } +} diff --git a/internal/openai/normalize_streaming_delta_test.go b/internal/openai/normalize_streaming_delta_test.go new file mode 100644 index 00000000..6959b590 --- /dev/null +++ b/internal/openai/normalize_streaming_delta_test.go @@ -0,0 +1,56 @@ +package openai + +import "testing" + +func TestNormalizeStreamingDelta_RepeatedCharBoundary(t *testing.T) { + // 流式在重复数字边界分片：不得把 "43" 的首字符与 "194" 尾字符误合并。 + cur, d := normalizeStreamingDelta("https://x:194", "43") + if want := "https://x:19443"; cur != want { + t.Fatalf("next: want %q got %q", want, cur) + } + if d != "43" { + t.Fatalf("delta: want %q got %q", "43", d) + } +} + +func TestNormalizeStreamingDelta_CumulativePrefix(t *testing.T) { + cur, d := normalizeStreamingDelta("今天", "今天天气") + if cur != "今天天气" || d != "天气" { + t.Fatalf("got cur=%q d=%q", cur, d) + } +} + +func TestNormalizeStreamingDelta_FullRetransmit(t *testing.T) { + cur, d := normalizeStreamingDelta("今天", "今天") + if d != "" || cur != "今天" { + t.Fatalf("got cur=%q d=%q", cur, d) + } +} + +func TestNormalizeStreamingDelta_SingleRuneRepeated(t *testing.T) { + cur, d := normalizeStreamingDelta("呀", "呀") + if want := "呀呀"; cur != want { + t.Fatalf("next: want %q got %q", want, cur) + } + if d != "呀" { + t.Fatalf("delta: want %q got %q", "呀", d) + } + cur, d = normalizeStreamingDelta("4", "4") + if want := "44"; cur != want { + t.Fatalf("next: want %q got %q", want, cur) + } + if d != "4" { + t.Fatalf("delta: want %q got %q", "4", d) + } +} + +func TestNormalizeStreamingDelta_CumulativeExtendsNumber(t *testing.T) { + // 已缓冲 "194" 后收到累计串 "19443"（注意 "1943" 并非 "19443" 的前缀，不能靠误写的中间态测 HasPrefix）。 + cur, d := normalizeStreamingDelta("194", "19443") + if want := "19443"; cur != want { + t.Fatalf("next: want %q got %q", want, cur) + } + if d != "43" { + t.Fatalf("delta: want %q got %q", "43", d) + } +} diff --git a/internal/openai/openai.go b/internal/openai/openai.go new file mode 100644 index 00000000..318813e4 --- /dev/null +++ b/internal/openai/openai.go @@ -0,0 +1,616 @@ +package openai + +import ( + "bufio" + "bytes" + "context" + "encoding/json" + "fmt" + "io" + "net/http" + "sort" + "strings" + "time" + "unicode/utf8" + + "cyberstrike-ai/internal/config" + + "go.uber.org/zap" +) + +// Client 统一封装与OpenAI兼容模型交互的HTTP客户端。 +type Client struct { + httpClient *http.Client + config *config.OpenAIConfig + logger *zap.Logger +} + +// APIError 表示OpenAI接口返回的非200错误。 +type APIError struct { + StatusCode int + Body string +} + +func (e *APIError) Error() string { + return fmt.Sprintf("openai api error: status=%d body=%s", e.StatusCode, e.Body) +} + +// normalizeStreamingDelta 将可能是“累计片段/重发片段”的内容归一化为“纯增量”。 +// 部分兼容网关会返回累计 content；若直接 append 会出现重复文本。 +// +// 注意： +// - 不做「任意后缀与前缀重叠」合并；流式可能在重复字符边界分片（"194"+"43"→"19443"）。 +// - HasPrefix 仅在 incoming 严格长于 current 时视为累计全文，否则会把分片产生的第二个相同 +// 单字/单码点（叠字、44、22 等）误判为「整段重复」而吞字。 +// - incoming==current 仅当 current 长度 >1 个码点时才视为整包重发；单码点重复必须走拼接。 +// - 不再使用「current 以 incoming 结尾则丢弃」：否则 "1943"+"43" 会误吞增量（19443 显示成 1943）。 +// 若网关重复发送尾部片段，应重复送完整累计串，由 HasPrefix 分支去重。 +func normalizeStreamingDelta(current, incoming string) (next, delta string) { + if incoming == "" { + return current, "" + } + if current == "" { + return incoming, incoming + } + if strings.HasPrefix(incoming, current) && len(incoming) > len(current) { + return incoming, incoming[len(current):] + } + if incoming == current && utf8.RuneCountInString(current) > 1 { + return current, "" + } + return current + incoming, incoming +} + +// NewClient 创建一个新的OpenAI客户端。 +func NewClient(cfg *config.OpenAIConfig, httpClient *http.Client, logger *zap.Logger) *Client { + if httpClient == nil { + httpClient = http.DefaultClient + } + if logger == nil { + logger = zap.NewNop() + } + return &Client{ + httpClient: httpClient, + config: cfg, + logger: logger, + } +} + +// UpdateConfig 动态更新OpenAI配置。 +func (c *Client) UpdateConfig(cfg *config.OpenAIConfig) { + c.config = cfg +} + +// ChatCompletion 调用 /chat/completions 接口。 +func (c *Client) ChatCompletion(ctx context.Context, payload interface{}, out interface{}) error { + if c == nil { + return fmt.Errorf("openai client is not initialized") + } + if c.config == nil { + return fmt.Errorf("openai config is nil") + } + if strings.TrimSpace(c.config.APIKey) == "" { + return fmt.Errorf("openai api key is empty") + } + if c.isClaude() { + return c.claudeNativeChatCompletion(ctx, payload, out) + } + + baseURL := strings.TrimSuffix(c.config.BaseURL, "/") + if baseURL == "" { + baseURL = "https://api.openai.com/v1" + } + + body, err := json.Marshal(payload) + if err != nil { + return fmt.Errorf("marshal openai payload: %w", err) + } + + c.logger.Debug("sending OpenAI chat completion request", + zap.Int("payloadSizeKB", len(body)/1024)) + + req, err := http.NewRequestWithContext(ctx, http.MethodPost, baseURL+"/chat/completions", bytes.NewReader(body)) + if err != nil { + return fmt.Errorf("build openai request: %w", err) + } + req.Header.Set("Content-Type", "application/json") + req.Header.Set("Authorization", "Bearer "+c.config.APIKey) + + requestStart := time.Now() + resp, err := c.httpClient.Do(req) + if err != nil { + return fmt.Errorf("call openai api: %w", err) + } + defer resp.Body.Close() + + bodyChan := make(chan []byte, 1) + errChan := make(chan error, 1) + go func() { + responseBody, err := io.ReadAll(resp.Body) + if err != nil { + errChan <- err + return + } + bodyChan <- responseBody + }() + + var respBody []byte + select { + case respBody = <-bodyChan: + case err := <-errChan: + return fmt.Errorf("read openai response: %w", err) + case <-ctx.Done(): + return fmt.Errorf("read openai response timeout: %w", ctx.Err()) + case <-time.After(25 * time.Minute): + return fmt.Errorf("read openai response timeout (25m)") + } + + c.logger.Debug("received OpenAI response", + zap.Int("status", resp.StatusCode), + zap.Duration("duration", time.Since(requestStart)), + zap.Int("responseSizeKB", len(respBody)/1024), + ) + + if resp.StatusCode != http.StatusOK { + c.logger.Warn("OpenAI chat completion returned non-200", + zap.Int("status", resp.StatusCode), + zap.String("body", string(respBody)), + ) + return &APIError{ + StatusCode: resp.StatusCode, + Body: string(respBody), + } + } + + if out != nil { + if err := json.Unmarshal(respBody, out); err != nil { + c.logger.Error("failed to unmarshal OpenAI response", + zap.Error(err), + zap.String("body", string(respBody)), + ) + return fmt.Errorf("unmarshal openai response: %w", err) + } + } + + return nil +} + +// ChatCompletionStream 调用 /chat/completions 的流式模式（stream=true），并在每个 delta 到达时回调 onDelta。 +// 返回最终拼接的 content（只拼 content delta；工具调用 delta 未做处理）。 +func (c *Client) ChatCompletionStream(ctx context.Context, payload interface{}, onDelta func(delta string) error) (string, error) { + if c == nil { + return "", fmt.Errorf("openai client is not initialized") + } + if c.config == nil { + return "", fmt.Errorf("openai config is nil") + } + if strings.TrimSpace(c.config.APIKey) == "" { + return "", fmt.Errorf("openai api key is empty") + } + if c.isClaude() { + return c.claudeNativeChatCompletionStream(ctx, payload, onDelta) + } + + baseURL := strings.TrimSuffix(c.config.BaseURL, "/") + if baseURL == "" { + baseURL = "https://api.openai.com/v1" + } + + body, err := json.Marshal(payload) + if err != nil { + return "", fmt.Errorf("marshal openai payload: %w", err) + } + + req, err := http.NewRequestWithContext(ctx, http.MethodPost, baseURL+"/chat/completions", bytes.NewReader(body)) + if err != nil { + return "", fmt.Errorf("build openai request: %w", err) + } + req.Header.Set("Content-Type", "application/json") + req.Header.Set("Authorization", "Bearer "+c.config.APIKey) + + requestStart := time.Now() + resp, err := c.httpClient.Do(req) + if err != nil { + return "", fmt.Errorf("call openai api: %w", err) + } + defer resp.Body.Close() + + // 非200：读完 body 返回 + if resp.StatusCode != http.StatusOK { + respBody, readErr := io.ReadAll(resp.Body) + if readErr != nil { + c.logger.Warn("failed to read OpenAI error response body", zap.Error(readErr)) + } + return "", &APIError{ + StatusCode: resp.StatusCode, + Body: string(respBody), + } + } + + type streamDelta struct { + // OpenAI 兼容流式通常使用 content；但部分兼容实现可能用 text。 + Content string `json:"content,omitempty"` + Text string `json:"text,omitempty"` + } + type streamChoice struct { + Delta streamDelta `json:"delta"` + FinishReason *string `json:"finish_reason,omitempty"` + } + type streamResponse struct { + ID string `json:"id,omitempty"` + Choices []streamChoice `json:"choices"` + Error *struct { + Message string `json:"message"` + Type string `json:"type"` + } `json:"error,omitempty"` + } + + reader := bufio.NewReader(resp.Body) + var full strings.Builder + fullText := "" + + // 典型 SSE 结构： + // data: {...}\n\n + // data: [DONE]\n\n + for { + line, readErr := reader.ReadString('\n') + if readErr != nil { + if readErr == io.EOF { + break + } + return full.String(), fmt.Errorf("read openai stream: %w", readErr) + } + trimmed := strings.TrimSpace(line) + if trimmed == "" { + continue + } + if !strings.HasPrefix(trimmed, "data:") { + continue + } + dataStr := strings.TrimSpace(strings.TrimPrefix(trimmed, "data:")) + if dataStr == "[DONE]" { + break + } + + var chunk streamResponse + if err := json.Unmarshal([]byte(dataStr), &chunk); err != nil { + // 解析失败跳过（兼容各种兼容层的差异） + continue + } + if chunk.Error != nil && strings.TrimSpace(chunk.Error.Message) != "" { + return full.String(), fmt.Errorf("openai stream error: %s", chunk.Error.Message) + } + if len(chunk.Choices) == 0 { + continue + } + + delta := chunk.Choices[0].Delta.Content + if delta == "" { + delta = chunk.Choices[0].Delta.Text + } + if delta == "" { + continue + } + + var deltaOut string + fullText, deltaOut = normalizeStreamingDelta(fullText, delta) + if deltaOut == "" { + continue + } + full.WriteString(deltaOut) + if onDelta != nil { + if err := onDelta(deltaOut); err != nil { + return full.String(), err + } + } + } + + c.logger.Debug("received OpenAI stream completion", + zap.Duration("duration", time.Since(requestStart)), + zap.Int("contentLen", full.Len()), + ) + + return full.String(), nil +} + +// StreamToolCall 流式工具调用的累积结果（arguments 以字符串形式拼接，留给上层再解析为 JSON）。 +type StreamToolCall struct { + Index int + ID string + Type string + FunctionName string + FunctionArgsStr string +} + +// ChatCompletionStreamWithToolCalls 流式模式：同时把 content delta 实时回调，并在结束后返回 tool_calls 和 finish_reason。 +func (c *Client) ChatCompletionStreamWithToolCalls( + ctx context.Context, + payload interface{}, + onContentDelta func(delta string) error, +) (string, []StreamToolCall, string, error) { + if c == nil { + return "", nil, "", fmt.Errorf("openai client is not initialized") + } + if c.config == nil { + return "", nil, "", fmt.Errorf("openai config is nil") + } + if strings.TrimSpace(c.config.APIKey) == "" { + return "", nil, "", fmt.Errorf("openai api key is empty") + } + if c.isClaude() { + return "", nil, "", fmt.Errorf("native Claude tool-call streaming requires Eino AgenticModel") + } + + baseURL := strings.TrimSuffix(c.config.BaseURL, "/") + if baseURL == "" { + baseURL = "https://api.openai.com/v1" + } + + body, err := json.Marshal(payload) + if err != nil { + return "", nil, "", fmt.Errorf("marshal openai payload: %w", err) + } + + req, err := http.NewRequestWithContext(ctx, http.MethodPost, baseURL+"/chat/completions", bytes.NewReader(body)) + if err != nil { + return "", nil, "", fmt.Errorf("build openai request: %w", err) + } + req.Header.Set("Content-Type", "application/json") + req.Header.Set("Authorization", "Bearer "+c.config.APIKey) + + requestStart := time.Now() + resp, err := c.httpClient.Do(req) + if err != nil { + return "", nil, "", fmt.Errorf("call openai api: %w", err) + } + defer resp.Body.Close() + + if resp.StatusCode != http.StatusOK { + respBody, readErr := io.ReadAll(resp.Body) + if readErr != nil { + c.logger.Warn("failed to read OpenAI error response body", zap.Error(readErr)) + } + return "", nil, "", &APIError{ + StatusCode: resp.StatusCode, + Body: string(respBody), + } + } + + // delta tool_calls 的增量结构 + type toolCallFunctionDelta struct { + Name string `json:"name,omitempty"` + Arguments string `json:"arguments,omitempty"` + } + type toolCallDelta struct { + Index int `json:"index,omitempty"` + ID string `json:"id,omitempty"` + Type string `json:"type,omitempty"` + Function toolCallFunctionDelta `json:"function,omitempty"` + } + type streamDelta2 struct { + Content string `json:"content,omitempty"` + Text string `json:"text,omitempty"` + ToolCalls []toolCallDelta `json:"tool_calls,omitempty"` + } + type streamChoice2 struct { + Delta streamDelta2 `json:"delta"` + FinishReason *string `json:"finish_reason,omitempty"` + } + type streamResponse2 struct { + Choices []streamChoice2 `json:"choices"` + Error *struct { + Message string `json:"message"` + Type string `json:"type"` + } `json:"error,omitempty"` + } + + type toolCallAccum struct { + id string + typ string + name string + args strings.Builder + } + toolCallAccums := make(map[int]*toolCallAccum) + + reader := bufio.NewReader(resp.Body) + var full strings.Builder + fullText := "" + finishReason := "" + + for { + line, readErr := reader.ReadString('\n') + if readErr != nil { + if readErr == io.EOF { + break + } + return full.String(), nil, finishReason, fmt.Errorf("read openai stream: %w", readErr) + } + trimmed := strings.TrimSpace(line) + if trimmed == "" { + continue + } + if !strings.HasPrefix(trimmed, "data:") { + continue + } + dataStr := strings.TrimSpace(strings.TrimPrefix(trimmed, "data:")) + if dataStr == "[DONE]" { + break + } + + var chunk streamResponse2 + if err := json.Unmarshal([]byte(dataStr), &chunk); err != nil { + // 兼容：解析失败跳过 + continue + } + if chunk.Error != nil && strings.TrimSpace(chunk.Error.Message) != "" { + return full.String(), nil, finishReason, fmt.Errorf("openai stream error: %s", chunk.Error.Message) + } + if len(chunk.Choices) == 0 { + continue + } + + choice := chunk.Choices[0] + if choice.FinishReason != nil && strings.TrimSpace(*choice.FinishReason) != "" { + finishReason = strings.TrimSpace(*choice.FinishReason) + } + + delta := choice.Delta + + content := delta.Content + if content == "" { + content = delta.Text + } + if content != "" { + var contentOut string + fullText, contentOut = normalizeStreamingDelta(fullText, content) + if contentOut != "" { + full.WriteString(contentOut) + if onContentDelta != nil { + if err := onContentDelta(contentOut); err != nil { + return full.String(), nil, finishReason, err + } + } + } + } + + if len(delta.ToolCalls) > 0 { + for _, tc := range delta.ToolCalls { + acc, ok := toolCallAccums[tc.Index] + if !ok { + acc = &toolCallAccum{} + toolCallAccums[tc.Index] = acc + } + if tc.ID != "" { + acc.id = tc.ID + } + if tc.Type != "" { + acc.typ = tc.Type + } + if tc.Function.Name != "" { + acc.name = tc.Function.Name + } + if tc.Function.Arguments != "" { + acc.args.WriteString(tc.Function.Arguments) + } + } + } + } + + // 组装 tool calls + indices := make([]int, 0, len(toolCallAccums)) + for idx := range toolCallAccums { + indices = append(indices, idx) + } + // 手写简单排序（避免额外 import） + for i := 0; i < len(indices); i++ { + for j := i + 1; j < len(indices); j++ { + if indices[j] < indices[i] { + indices[i], indices[j] = indices[j], indices[i] + } + } + } + + toolCalls := make([]StreamToolCall, 0, len(indices)) + for _, idx := range indices { + acc := toolCallAccums[idx] + tc := StreamToolCall{ + Index: idx, + ID: acc.id, + Type: acc.typ, + FunctionName: acc.name, + FunctionArgsStr: acc.args.String(), + } + toolCalls = append(toolCalls, tc) + } + + c.logger.Debug("received OpenAI stream completion (tool_calls)", + zap.Duration("duration", time.Since(requestStart)), + zap.Int("contentLen", full.Len()), + zap.Int("toolCalls", len(toolCalls)), + zap.String("finishReason", finishReason), + ) + + if strings.TrimSpace(finishReason) == "" { + finishReason = "stop" + } + + return full.String(), toolCalls, finishReason, nil +} + +// ModelsListResponse 表示 OpenAI 兼容 GET /models 响应。 +type ModelsListResponse struct { + Object string `json:"object"` + Data []struct { + ID string `json:"id"` + Object string `json:"object,omitempty"` + OwnedBy string `json:"owned_by,omitempty"` + } `json:"data"` +} + +// ListModels 调用 GET {baseURL}/models 获取可用模型 id 列表（按字典序）。 +func (c *Client) ListModels(ctx context.Context) ([]string, error) { + if c == nil { + return nil, fmt.Errorf("openai client is not initialized") + } + if c.config == nil { + return nil, fmt.Errorf("openai config is nil") + } + if strings.TrimSpace(c.config.APIKey) == "" { + return nil, fmt.Errorf("openai api key is empty") + } + if c.isClaude() { + return nil, fmt.Errorf("claude provider does not support models list API") + } + + baseURL := strings.TrimSuffix(c.config.BaseURL, "/") + if baseURL == "" { + baseURL = "https://api.openai.com/v1" + } + + req, err := http.NewRequestWithContext(ctx, http.MethodGet, baseURL+"/models", nil) + if err != nil { + return nil, fmt.Errorf("build openai models request: %w", err) + } + req.Header.Set("Authorization", "Bearer "+c.config.APIKey) + + resp, err := c.httpClient.Do(req) + if err != nil { + return nil, fmt.Errorf("call openai models api: %w", err) + } + defer resp.Body.Close() + + respBody, err := io.ReadAll(resp.Body) + if err != nil { + return nil, fmt.Errorf("read openai models response: %w", err) + } + if resp.StatusCode != http.StatusOK { + return nil, &APIError{ + StatusCode: resp.StatusCode, + Body: string(respBody), + } + } + + var list ModelsListResponse + if err := json.Unmarshal(respBody, &list); err != nil { + return nil, fmt.Errorf("decode openai models response: %w", err) + } + + seen := make(map[string]struct{}, len(list.Data)) + models := make([]string, 0, len(list.Data)) + for _, item := range list.Data { + id := strings.TrimSpace(item.ID) + if id == "" { + continue + } + if _, ok := seen[id]; ok { + continue + } + seen[id] = struct{}{} + models = append(models, id) + } + sort.Strings(models) + if len(models) == 0 { + return nil, fmt.Errorf("models list is empty") + } + return models, nil +} diff --git a/internal/openai/reasoning_payload.go b/internal/openai/reasoning_payload.go new file mode 100644 index 00000000..2110c90d --- /dev/null +++ b/internal/openai/reasoning_payload.go @@ -0,0 +1,117 @@ +package openai + +import ( + "strings" + + "github.com/bytedance/sonic" +) + +// reasoningPayloadKeys are OpenAI-compatible root fields that enable "thinking" / +// extended-reasoning modes on gateways such as DashScope/Qwen and MiniMax. +var reasoningPayloadKeys = []string{ + "thinking", + "reasoning_effort", + "output_config", + "reasoning", +} + +// StripReasoningFromChatCompletionBody removes thinking / reasoning fields from a +// chat-completions JSON body. +func StripReasoningFromChatCompletionBody(rawBody []byte) ([]byte, error) { + var payload map[string]any + if err := sonic.Unmarshal(rawBody, &payload); err != nil { + return rawBody, nil + } + if !stripReasoningFields(payload) { + return rawBody, nil + } + out, err := sonic.Marshal(payload) + if err != nil { + return rawBody, err + } + return out, nil +} + +// StripReasoningIfForcedToolChoice removes thinking / reasoning fields when the +// request sets tool_choice to "required" or an object. Several providers reject +// that combination (e.g. DashScope: "tool_choice does not support being set to +// required or object in thinking mode"). +func StripReasoningIfForcedToolChoice(rawBody []byte) ([]byte, error) { + var payload map[string]any + if err := sonic.Unmarshal(rawBody, &payload); err != nil { + return rawBody, nil + } + if !forcedToolChoiceIncompatibleWithThinking(payload) { + return rawBody, nil + } + if !stripReasoningFields(payload) { + return rawBody, nil + } + out, err := sonic.Marshal(payload) + if err != nil { + return rawBody, err + } + return out, nil +} + +// StripToolChoiceForThinkingMode removes tool_choice while preserving tools and +// thinking fields. DeepSeek thinking mode can use tools, but rejects the +// tool_choice parameter itself on some agent requests. +func StripToolChoiceForThinkingMode(rawBody []byte) ([]byte, error) { + var payload map[string]any + if err := sonic.Unmarshal(rawBody, &payload); err != nil { + return rawBody, nil + } + if !thinkingModeEnabledByPayload(payload) { + return rawBody, nil + } + if _, ok := payload["tool_choice"]; !ok { + return rawBody, nil + } + delete(payload, "tool_choice") + out, err := sonic.Marshal(payload) + if err != nil { + return rawBody, err + } + return out, nil +} + +func stripReasoningFields(payload map[string]any) bool { + changed := false + for _, key := range reasoningPayloadKeys { + if _, ok := payload[key]; ok { + delete(payload, key) + changed = true + } + } + return changed +} + +func forcedToolChoiceIncompatibleWithThinking(payload map[string]any) bool { + tc, ok := payload["tool_choice"] + if !ok || tc == nil { + return false + } + switch v := tc.(type) { + case string: + return v == "required" + case map[string]any: + return true + default: + return false + } +} + +func thinkingModeEnabledByPayload(payload map[string]any) bool { + thinking, ok := payload["thinking"] + if !ok || thinking == nil { + // DeepSeek enables thinking by default unless explicitly disabled. + return true + } + if m, ok := thinking.(map[string]any); ok { + if typ, ok := m["type"].(string); ok && strings.EqualFold(strings.TrimSpace(typ), "disabled") { + return false + } + } + return true +} diff --git a/internal/openai/reasoning_payload_test.go b/internal/openai/reasoning_payload_test.go new file mode 100644 index 00000000..9f7dccfe --- /dev/null +++ b/internal/openai/reasoning_payload_test.go @@ -0,0 +1,250 @@ +package openai + +import ( + "io" + "net/http" + "strings" + "testing" + + "cyberstrike-ai/internal/config" +) + +func TestStripReasoningFromChatCompletionBody(t *testing.T) { + in := []byte(`{"model":"deepseek-chat","messages":[],"thinking":{"type":"enabled"},"reasoning_effort":"high"}`) + out, err := StripReasoningFromChatCompletionBody(in) + if err != nil { + t.Fatal(err) + } + s := string(out) + if strings.Contains(s, "thinking") || strings.Contains(s, "reasoning_effort") { + t.Fatalf("expected reasoning fields stripped, got %s", s) + } + if !strings.Contains(s, `"model":"deepseek-chat"`) { + t.Fatalf("expected model preserved, got %s", s) + } + + plain := []byte(`{"model":"gpt-4o","messages":[]}`) + out2, err := StripReasoningFromChatCompletionBody(plain) + if err != nil { + t.Fatal(err) + } + if string(out2) != string(plain) { + t.Fatalf("expected unchanged payload, got %s", out2) + } +} + +func TestStripReasoningIfForcedToolChoice(t *testing.T) { + cases := []struct { + name string + in string + strip bool + contain string + }{ + { + name: "required strips thinking", + in: `{"model":"minimax","messages":[],"thinking":{"type":"enabled"},"tool_choice":"required","tools":[]}`, + strip: true, + }, + { + name: "object tool_choice strips thinking", + in: `{"model":"qwen","messages":[],"thinking":{"type":"enabled"},"tool_choice":{"type":"function","function":{"name":"respond"}}}`, + strip: true, + }, + { + name: "auto keeps thinking", + in: `{"model":"qwen","messages":[],"thinking":{"type":"enabled"},"tool_choice":"auto"}`, + strip: false, + contain: "thinking", + }, + { + name: "no tool_choice keeps thinking", + in: `{"model":"qwen","messages":[],"thinking":{"type":"enabled"}}`, + strip: false, + contain: "thinking", + }, + } + for _, tc := range cases { + t.Run(tc.name, func(t *testing.T) { + out, err := StripReasoningIfForcedToolChoice([]byte(tc.in)) + if err != nil { + t.Fatal(err) + } + s := string(out) + hasThinking := strings.Contains(s, "thinking") + if tc.strip && hasThinking { + t.Fatalf("expected thinking stripped, got %s", s) + } + if !tc.strip && tc.contain != "" && !strings.Contains(s, tc.contain) { + t.Fatalf("expected %q in %s", tc.contain, s) + } + if !tc.strip && string(out) != tc.in { + t.Fatalf("expected unchanged payload, got %s", s) + } + }) + } +} + +func TestStripToolChoiceForThinkingMode(t *testing.T) { + cases := []struct { + name string + in string + wantToolChoice bool + wantThinking bool + }{ + { + name: "enabled thinking removes tool_choice", + in: `{"model":"deepseek-v4","messages":[],"thinking":{"type":"enabled"},"tool_choice":"required","tools":[{"type":"function","function":{"name":"scan"}}]}`, + wantToolChoice: false, + wantThinking: true, + }, + { + name: "default thinking removes tool_choice", + in: `{"model":"deepseek-v4","messages":[],"tool_choice":"auto","tools":[]}`, + wantToolChoice: false, + wantThinking: false, + }, + { + name: "disabled thinking keeps tool_choice", + in: `{"model":"deepseek-v4","messages":[],"thinking":{"type":"disabled"},"tool_choice":"required","tools":[]}`, + wantToolChoice: true, + wantThinking: true, + }, + { + name: "no tool_choice unchanged", + in: `{"model":"deepseek-v4","messages":[],"thinking":{"type":"enabled"},"tools":[]}`, + wantToolChoice: false, + wantThinking: true, + }, + } + for _, tc := range cases { + t.Run(tc.name, func(t *testing.T) { + out, err := StripToolChoiceForThinkingMode([]byte(tc.in)) + if err != nil { + t.Fatal(err) + } + s := string(out) + if strings.Contains(s, "tool_choice") != tc.wantToolChoice { + t.Fatalf("tool_choice presence mismatch, got %s", s) + } + if strings.Contains(s, "thinking") != tc.wantThinking { + t.Fatalf("thinking presence mismatch, got %s", s) + } + if !strings.Contains(s, "tools") { + t.Fatalf("expected tools preserved, got %s", s) + } + }) + } +} + +func TestReasoningToolChoiceCompatRoundTripper(t *testing.T) { + var gotBody string + rt := &reasoningToolChoiceCompatRoundTripper{ + base: roundTripperFunc(func(req *http.Request) (*http.Response, error) { + b, _ := io.ReadAll(req.Body) + gotBody = string(b) + return &http.Response{ + StatusCode: 200, + Body: io.NopCloser(strings.NewReader(`{"choices":[{"message":{"content":"ok"}}]}`)), + Header: http.Header{"Content-Type": []string{"application/json"}}, + }, nil + }), + } + req, err := http.NewRequest(http.MethodPost, "https://example.com/v1/chat/completions", strings.NewReader( + `{"model":"m","thinking":{"type":"enabled"},"tool_choice":"required","messages":[]}`, + )) + if err != nil { + t.Fatal(err) + } + _, err = rt.RoundTrip(req) + if err != nil { + t.Fatal(err) + } + if strings.Contains(gotBody, "thinking") { + t.Fatalf("expected thinking stripped in transit, got %s", gotBody) + } + if !strings.Contains(gotBody, `"tool_choice":"required"`) { + t.Fatalf("expected tool_choice preserved, got %s", gotBody) + } +} + +func TestReasoningToolChoiceCompatRoundTripperDeepSeek(t *testing.T) { + var gotBody string + rt := &reasoningToolChoiceCompatRoundTripper{ + cfg: &config.OpenAIConfig{ + BaseURL: "https://api.deepseek.com/v1", + Model: "deepseek-v4", + }, + base: roundTripperFunc(func(req *http.Request) (*http.Response, error) { + b, _ := io.ReadAll(req.Body) + gotBody = string(b) + return &http.Response{ + StatusCode: 200, + Body: io.NopCloser(strings.NewReader(`{"choices":[{"message":{"content":"ok"}}]}`)), + Header: http.Header{"Content-Type": []string{"application/json"}}, + }, nil + }), + } + req, err := http.NewRequest(http.MethodPost, "https://api.deepseek.com/v1/chat/completions", strings.NewReader( + `{"model":"deepseek-v4","thinking":{"type":"enabled"},"tool_choice":"required","tools":[],"messages":[]}`, + )) + if err != nil { + t.Fatal(err) + } + _, err = rt.RoundTrip(req) + if err != nil { + t.Fatal(err) + } + if strings.Contains(gotBody, "tool_choice") { + t.Fatalf("expected DeepSeek tool_choice stripped in transit, got %s", gotBody) + } + if !strings.Contains(gotBody, "thinking") { + t.Fatalf("expected thinking preserved for DeepSeek, got %s", gotBody) + } + if !strings.Contains(gotBody, "tools") { + t.Fatalf("expected tools preserved for DeepSeek, got %s", gotBody) + } +} + +func TestReasoningToolChoiceCompatRoundTripperDeepSeekEndpointWinsOverProfile(t *testing.T) { + var gotBody string + rt := &reasoningToolChoiceCompatRoundTripper{ + cfg: &config.OpenAIConfig{ + BaseURL: "https://api.deepseek.com/v1", + Model: "deepseek-v4-flash", + Reasoning: config.OpenAIReasoningConfig{ + Profile: "openai_compat", + }, + }, + base: roundTripperFunc(func(req *http.Request) (*http.Response, error) { + b, _ := io.ReadAll(req.Body) + gotBody = string(b) + return &http.Response{ + StatusCode: 200, + Body: io.NopCloser(strings.NewReader(`{"choices":[{"message":{"content":"ok"}}]}`)), + Header: http.Header{"Content-Type": []string{"application/json"}}, + }, nil + }), + } + req, err := http.NewRequest(http.MethodPost, "https://api.deepseek.com/v1/chat/completions", strings.NewReader( + `{"model":"deepseek-v4-flash","tool_choice":"required","tools":[],"messages":[]}`, + )) + if err != nil { + t.Fatal(err) + } + _, err = rt.RoundTrip(req) + if err != nil { + t.Fatal(err) + } + if strings.Contains(gotBody, "tool_choice") { + t.Fatalf("expected DeepSeek tool_choice stripped despite openai_compat profile, got %s", gotBody) + } + if !strings.Contains(gotBody, "tools") { + t.Fatalf("expected tools preserved for DeepSeek, got %s", gotBody) + } +} + +type roundTripperFunc func(*http.Request) (*http.Response, error) + +func (f roundTripperFunc) RoundTrip(req *http.Request) (*http.Response, error) { + return f(req) +} diff --git a/internal/openai/reasoning_tool_choice_compat.go b/internal/openai/reasoning_tool_choice_compat.go new file mode 100644 index 00000000..fd703222 --- /dev/null +++ b/internal/openai/reasoning_tool_choice_compat.go @@ -0,0 +1,69 @@ +package openai + +import ( + "bytes" + "io" + "net/http" + "strconv" + "strings" + + "cyberstrike-ai/internal/config" +) + +// reasoningToolChoiceCompatRoundTripper strips thinking/reasoning fields from +// chat/completions requests that force tool_choice, which some gateways reject +// when thinking mode is enabled on the same request. +type reasoningToolChoiceCompatRoundTripper struct { + base http.RoundTripper + cfg *config.OpenAIConfig +} + +func (rt *reasoningToolChoiceCompatRoundTripper) RoundTrip(req *http.Request) (*http.Response, error) { + if rt == nil || rt.base == nil || req == nil || req.Body == nil { + if rt != nil && rt.base != nil { + return rt.base.RoundTrip(req) + } + return http.DefaultTransport.RoundTrip(req) + } + if req.Method != http.MethodPost || !strings.HasSuffix(req.URL.Path, "/chat/completions") { + return rt.base.RoundTrip(req) + } + + body, err := io.ReadAll(req.Body) + _ = req.Body.Close() + if err != nil { + return nil, err + } + + patched := body + var perr error + if isDeepSeekToolChoiceCompatProfile(rt.cfg) { + patched, perr = StripToolChoiceForThinkingMode(body) + } else { + patched, perr = StripReasoningIfForcedToolChoice(body) + } + if perr != nil { + patched = body + } + req.Body = io.NopCloser(bytes.NewReader(patched)) + req.ContentLength = int64(len(patched)) + req.Header.Set("Content-Length", strconv.Itoa(len(patched))) + return rt.base.RoundTrip(req) +} + +func isDeepSeekToolChoiceCompatProfile(cfg *config.OpenAIConfig) bool { + if cfg == nil { + return false + } + if cfg.IsDeepSeekEndpointOrModel() { + return true + } + profile := strings.ToLower(strings.TrimSpace(cfg.Reasoning.ProfileEffective())) + if profile == "deepseek" || profile == "deepseek_compat" { + return true + } + if profile != "" && profile != "auto" { + return false + } + return false +} diff --git a/internal/openai/sse_stream.go b/internal/openai/sse_stream.go new file mode 100644 index 00000000..a86d6306 --- /dev/null +++ b/internal/openai/sse_stream.go @@ -0,0 +1,20 @@ +package openai + +// SSEAccumulatedKey 为 SSE progress 事件 data 中的服务端权威流式全文快照字段。 +// 前端应优先用该字段更新 buffer，避免对 delta 二次 normalize 导致叠字。 +const SSEAccumulatedKey = "accumulated" + +// WithSSEAccumulated 在 progress data 中附带当前流式累计全文（权威快照）。 +func WithSSEAccumulated(data map[string]interface{}, accumulated string) map[string]interface{} { + if data == nil { + data = make(map[string]interface{}, 1) + } + data[SSEAccumulatedKey] = accumulated + return data +} + +// NormalizeStreamingDelta 将可能是“累计片段/重发片段”的内容归一化为“纯增量”。 +// 与 unexported normalizeStreamingDelta 相同，供 agent / multiagent 等包在发 SSE 前累计正文。 +func NormalizeStreamingDelta(current, incoming string) (next, delta string) { + return normalizeStreamingDelta(current, incoming) +} diff --git a/internal/openai/summarization_diag.go b/internal/openai/summarization_diag.go new file mode 100644 index 00000000..44465145 --- /dev/null +++ b/internal/openai/summarization_diag.go @@ -0,0 +1,108 @@ +package openai + +import ( + "bytes" + "io" + "net/http" + "strings" + + "github.com/bytedance/sonic" + "go.uber.org/zap" +) + +// SummarizationRequestHeader marks chat/completion requests issued by Eino summarization +// middleware (via model.WithExtraHeader). The diagnostic transport logs empty-choices bodies +// only for these requests so main-agent traffic stays quiet. +const SummarizationRequestHeader = "X-CyberStrike-Summarization" + +const summarizationDiagBodyMaxBytes = 8192 + +// AttachSummarizationDiagTransport wraps client.Transport to log raw API bodies when +// summarization receives HTTP 200 with an empty choices array. +func AttachSummarizationDiagTransport(client *http.Client, logger *zap.Logger) { + if client == nil || logger == nil { + return + } + base := client.Transport + if base == nil { + base = http.DefaultTransport + } + client.Transport = &summarizationDiagRoundTripper{base: base, logger: logger} +} + +type summarizationDiagRoundTripper struct { + base http.RoundTripper + logger *zap.Logger +} + +func (rt *summarizationDiagRoundTripper) RoundTrip(req *http.Request) (*http.Response, error) { + resp, err := rt.base.RoundTrip(req) + if err != nil || resp == nil || resp.Body == nil { + return resp, err + } + if !isSummarizationRequest(req) || !strings.Contains(strings.ToLower(resp.Header.Get("Content-Type")), "json") { + return resp, err + } + + body, readErr := io.ReadAll(resp.Body) + _ = resp.Body.Close() + if readErr != nil { + resp.Body = io.NopCloser(bytes.NewReader(nil)) + return resp, err + } + resp.Body = io.NopCloser(bytes.NewReader(body)) + resp.ContentLength = int64(len(body)) + + if rt.logger != nil && resp.StatusCode >= http.StatusBadRequest { + rt.logger.Warn("eino summarization: API request rejected", + zap.Int("status", resp.StatusCode), + zap.String("request_id", responseRequestID(resp)), + zap.Int("response_bytes", len(body)), + zap.String("raw_body", truncateForLog(string(body), summarizationDiagBodyMaxBytes)), + ) + } else if rt.logger != nil && summarizationResponseEmptyChoices(body) { + rt.logger.Warn("eino summarization: API returned empty choices", + zap.Int("status", resp.StatusCode), + zap.String("request_id", responseRequestID(resp)), + zap.Int("response_bytes", len(body)), + zap.String("raw_body", truncateForLog(string(body), summarizationDiagBodyMaxBytes)), + ) + } + return resp, err +} + +func responseRequestID(resp *http.Response) string { + if resp == nil { + return "" + } + for _, key := range []string{"x-request-id", "request-id", "x-trace-id"} { + if value := strings.TrimSpace(resp.Header.Get(key)); value != "" { + return value + } + } + return "" +} + +func isSummarizationRequest(req *http.Request) bool { + if req == nil { + return false + } + return strings.TrimSpace(req.Header.Get(SummarizationRequestHeader)) == "1" +} + +func summarizationResponseEmptyChoices(body []byte) bool { + var parsed struct { + Choices []any `json:"choices"` + } + if err := sonic.Unmarshal(body, &parsed); err != nil { + return false + } + return len(parsed.Choices) == 0 +} + +func truncateForLog(s string, maxBytes int) string { + if maxBytes <= 0 || len(s) <= maxBytes { + return s + } + return s[:maxBytes] + "…(truncated)" +} diff --git a/internal/openai/summarization_diag_test.go b/internal/openai/summarization_diag_test.go new file mode 100644 index 00000000..753a61ae --- /dev/null +++ b/internal/openai/summarization_diag_test.go @@ -0,0 +1,47 @@ +package openai + +import ( + "io" + "net/http" + "strings" + "testing" + + "go.uber.org/zap" +) + +type staticRoundTripper struct { + status int + body string +} + +func (s *staticRoundTripper) RoundTrip(req *http.Request) (*http.Response, error) { + return &http.Response{ + StatusCode: s.status, + Header: http.Header{"Content-Type": []string{"application/json"}}, + Body: io.NopCloser(strings.NewReader(s.body)), + }, nil +} + +func TestSummarizationResponseEmptyChoices(t *testing.T) { + if !summarizationResponseEmptyChoices([]byte(`{"choices":[]}`)) { + t.Fatal("expected empty choices") + } + if summarizationResponseEmptyChoices([]byte(`{"choices":[{"index":0}]}`)) { + t.Fatal("expected non-empty choices") + } +} + +func TestSummarizationDiagRoundTripper_SkipsWithoutHeader(t *testing.T) { + client := &http.Client{ + Transport: &summarizationDiagRoundTripper{ + base: &staticRoundTripper{status: 200, body: `{"choices":[]}`}, + logger: zap.NewNop(), + }, + } + req, _ := http.NewRequest(http.MethodPost, "https://example.com/v1/chat/completions", nil) + resp, err := client.Do(req) + if err != nil { + t.Fatal(err) + } + _ = resp.Body.Close() +} diff --git a/internal/skillpackage/content.go b/internal/skillpackage/content.go new file mode 100644 index 00000000..91a02310 --- /dev/null +++ b/internal/skillpackage/content.go @@ -0,0 +1,164 @@ +package skillpackage + +import ( + "fmt" + "regexp" + "strings" +) + +var reH2 = regexp.MustCompile(`(?m)^##\s+(.+)$`) + +const summaryContentRunes = 6000 + +type markdownSection struct { + Heading string + Title string + Content string +} + +func splitMarkdownSections(body string) []markdownSection { + body = strings.TrimSpace(body) + if body == "" { + return nil + } + idxs := reH2.FindAllStringIndex(body, -1) + titles := reH2.FindAllStringSubmatch(body, -1) + if len(idxs) == 0 { + return []markdownSection{{ + Heading: "", + Title: "_body", + Content: body, + }} + } + var out []markdownSection + for i := range idxs { + title := strings.TrimSpace(titles[i][1]) + start := idxs[i][0] + end := len(body) + if i+1 < len(idxs) { + end = idxs[i+1][0] + } + chunk := strings.TrimSpace(body[start:end]) + out = append(out, markdownSection{ + Heading: "## " + title, + Title: title, + Content: chunk, + }) + } + return out +} + +func deriveSections(body string) []SkillSection { + md := splitMarkdownSections(body) + out := make([]SkillSection, 0, len(md)) + for _, ms := range md { + if ms.Title == "_body" { + continue + } + out = append(out, SkillSection{ + ID: slugifySectionID(ms.Title), + Title: ms.Title, + Heading: ms.Heading, + Level: 2, + }) + } + return out +} + +func slugifySectionID(title string) string { + title = strings.TrimSpace(strings.ToLower(title)) + if title == "" { + return "section" + } + var b strings.Builder + for _, r := range title { + switch { + case r >= 'a' && r <= 'z', r >= '0' && r <= '9': + b.WriteRune(r) + case r == ' ', r == '-', r == '_': + b.WriteRune('-') + } + } + s := strings.Trim(b.String(), "-") + if s == "" { + return "section" + } + return s +} + +func findSectionContent(sections []markdownSection, sec string) string { + sec = strings.TrimSpace(sec) + if sec == "" { + return "" + } + want := strings.ToLower(sec) + for _, s := range sections { + if strings.EqualFold(slugifySectionID(s.Title), want) || strings.EqualFold(s.Title, sec) { + return s.Content + } + if strings.EqualFold(strings.ReplaceAll(s.Title, " ", "-"), want) { + return s.Content + } + } + return "" +} + +func buildSummaryMarkdown(name, description string, tags []string, scripts []SkillScriptInfo, sections []SkillSection, body string) string { + var b strings.Builder + if description != "" { + b.WriteString(description) + b.WriteString("\n\n") + } + if len(tags) > 0 { + b.WriteString("**Tags**: ") + b.WriteString(strings.Join(tags, ", ")) + b.WriteString("\n\n") + } + if len(scripts) > 0 { + b.WriteString("### Bundled scripts\n\n") + for _, sc := range scripts { + line := "- `" + sc.RelPath + "`" + if sc.Description != "" { + line += " — " + sc.Description + } + b.WriteString(line) + b.WriteString("\n") + } + b.WriteString("\n") + } + if len(sections) > 0 { + b.WriteString("### Sections\n\n") + for _, sec := range sections { + line := "- **" + sec.ID + "**" + if sec.Title != "" && sec.Title != sec.ID { + line += ": " + sec.Title + } + b.WriteString(line) + b.WriteString("\n") + } + b.WriteString("\n") + } + mdSecs := splitMarkdownSections(body) + preview := body + if len(mdSecs) > 0 && mdSecs[0].Title != "_body" { + preview = mdSecs[0].Content + } + b.WriteString("### Preview (SKILL.md)\n\n") + b.WriteString(truncateRunes(strings.TrimSpace(preview), summaryContentRunes)) + b.WriteString("\n\n---\n\n_(Summary for admin UI. Agents use Eino `skill` tool for full SKILL.md progressive loading.)_") + if name != "" { + b.WriteString(fmt.Sprintf("\n\n_Skill name: %s_", name)) + } + return b.String() +} + +func truncateRunes(s string, max int) string { + if max <= 0 || s == "" { + return s + } + r := []rune(s) + if len(r) <= max { + return s + } + return string(r[:max]) + "…" +} diff --git a/internal/skillpackage/frontmatter.go b/internal/skillpackage/frontmatter.go new file mode 100644 index 00000000..905156b1 --- /dev/null +++ b/internal/skillpackage/frontmatter.go @@ -0,0 +1,114 @@ +package skillpackage + +import ( + "fmt" + "strings" + + "gopkg.in/yaml.v3" +) + +// ExtractSkillMDFrontMatterYAML returns the YAML source inside the first --- ... --- block and the markdown body. +func ExtractSkillMDFrontMatterYAML(raw []byte) (fmYAML string, body string, err error) { + text := strings.TrimPrefix(string(raw), "\ufeff") + if strings.TrimSpace(text) == "" { + return "", "", fmt.Errorf("SKILL.md is empty") + } + lines := strings.Split(text, "\n") + if len(lines) < 2 || strings.TrimSpace(lines[0]) != "---" { + return "", "", fmt.Errorf("SKILL.md must start with YAML front matter (---) per Agent Skills standard") + } + var fmLines []string + i := 1 + for i < len(lines) { + if strings.TrimSpace(lines[i]) == "---" { + break + } + fmLines = append(fmLines, lines[i]) + i++ + } + if i >= len(lines) { + return "", "", fmt.Errorf("SKILL.md: front matter must end with a line containing only ---") + } + body = strings.Join(lines[i+1:], "\n") + body = strings.TrimSpace(body) + fmYAML = strings.Join(fmLines, "\n") + return fmYAML, body, nil +} + +// ParseSkillMD parses SKILL.md YAML head + body. +func ParseSkillMD(raw []byte) (*SkillManifest, string, error) { + fmYAML, body, err := ExtractSkillMDFrontMatterYAML(raw) + if err != nil { + return nil, "", err + } + var m SkillManifest + if err := yaml.Unmarshal([]byte(fmYAML), &m); err != nil { + return nil, "", fmt.Errorf("SKILL.md front matter: %w", err) + } + return &m, body, nil +} + +type skillFrontMatterExport struct { + Name string `yaml:"name"` + Description string `yaml:"description"` + License string `yaml:"license,omitempty"` + Compatibility string `yaml:"compatibility,omitempty"` + Metadata map[string]any `yaml:"metadata,omitempty"` + AllowedTools string `yaml:"allowed-tools,omitempty"` +} + +// BuildSkillMD serializes SKILL.md per agentskills.io. +func BuildSkillMD(m *SkillManifest, body string) ([]byte, error) { + if m == nil { + return nil, fmt.Errorf("nil manifest") + } + fm := skillFrontMatterExport{ + Name: strings.TrimSpace(m.Name), + Description: strings.TrimSpace(m.Description), + License: strings.TrimSpace(m.License), + Compatibility: strings.TrimSpace(m.Compatibility), + AllowedTools: strings.TrimSpace(m.AllowedTools), + } + if len(m.Metadata) > 0 { + fm.Metadata = m.Metadata + } + head, err := yaml.Marshal(&fm) + if err != nil { + return nil, err + } + s := strings.TrimSpace(string(head)) + out := "---\n" + s + "\n---\n\n" + strings.TrimSpace(body) + "\n" + return []byte(out), nil +} + +func manifestTags(m *SkillManifest) []string { + if m == nil || m.Metadata == nil { + return nil + } + var out []string + if raw, ok := m.Metadata["tags"]; ok { + switch v := raw.(type) { + case []any: + for _, x := range v { + if s, ok := x.(string); ok && s != "" { + out = append(out, s) + } + } + case []string: + out = append(out, v...) + } + } + return out +} + +func versionFromMetadata(m *SkillManifest) string { + if m == nil || m.Metadata == nil { + return "" + } + if v, ok := m.Metadata["version"]; ok { + if s, ok := v.(string); ok { + return strings.TrimSpace(s) + } + } + return "" +} diff --git a/internal/skillpackage/io.go b/internal/skillpackage/io.go new file mode 100644 index 00000000..8a2b7222 --- /dev/null +++ b/internal/skillpackage/io.go @@ -0,0 +1,200 @@ +package skillpackage + +import ( + "fmt" + "io/fs" + "os" + "path/filepath" + "strings" +) + +const ( + maxPackageFiles = 4000 + maxPackageDepth = 24 + maxScriptsDepth = 24 + defaultMaxRead = 10 << 20 +) + +// SafeRelPath resolves rel inside root (no ..). +func SafeRelPath(root, rel string) (string, error) { + rel = strings.TrimSpace(rel) + rel = filepath.ToSlash(rel) + rel = strings.TrimPrefix(rel, "/") + if rel == "" || rel == "." { + return "", fmt.Errorf("empty resource path") + } + if strings.Contains(rel, "..") { + return "", fmt.Errorf("invalid path %q", rel) + } + abs := filepath.Join(root, filepath.FromSlash(rel)) + cleanRoot := filepath.Clean(root) + cleanAbs := filepath.Clean(abs) + relOut, err := filepath.Rel(cleanRoot, cleanAbs) + if err != nil || relOut == ".." || strings.HasPrefix(relOut, ".."+string(filepath.Separator)) { + return "", fmt.Errorf("path escapes skill directory: %q", rel) + } + return cleanAbs, nil +} + +// ListPackageFiles lists files under a skill directory. +func ListPackageFiles(skillsRoot, skillID string) ([]PackageFileInfo, error) { + root := SkillDir(skillsRoot, skillID) + if _, err := ResolveSKILLPath(root); err != nil { + return nil, fmt.Errorf("skill %q: %w", skillID, err) + } + var out []PackageFileInfo + err := filepath.WalkDir(root, func(path string, d fs.DirEntry, err error) error { + if err != nil { + return err + } + rel, e := filepath.Rel(root, path) + if e != nil { + return e + } + if rel == "." { + return nil + } + depth := strings.Count(rel, string(os.PathSeparator)) + if depth > maxPackageDepth { + if d.IsDir() { + return filepath.SkipDir + } + return nil + } + if strings.HasPrefix(d.Name(), ".") { + if d.IsDir() { + return filepath.SkipDir + } + return nil + } + if len(out) >= maxPackageFiles { + return fmt.Errorf("skill package exceeds %d files", maxPackageFiles) + } + fi, err := d.Info() + if err != nil { + return err + } + out = append(out, PackageFileInfo{ + Path: filepath.ToSlash(rel), + Size: fi.Size(), + IsDir: d.IsDir(), + }) + return nil + }) + return out, err +} + +// ReadPackageFile reads a file relative to the skill package. +func ReadPackageFile(skillsRoot, skillID, relPath string, maxBytes int64) ([]byte, error) { + if maxBytes <= 0 { + maxBytes = defaultMaxRead + } + root := SkillDir(skillsRoot, skillID) + abs, err := SafeRelPath(root, relPath) + if err != nil { + return nil, err + } + fi, err := os.Stat(abs) + if err != nil { + return nil, err + } + if fi.IsDir() { + return nil, fmt.Errorf("path is a directory") + } + if fi.Size() > maxBytes { + return readFileHead(abs, maxBytes) + } + return os.ReadFile(abs) +} + +// WritePackageFile writes a file inside the skill package. +func WritePackageFile(skillsRoot, skillID, relPath string, content []byte) error { + root := SkillDir(skillsRoot, skillID) + if _, err := ResolveSKILLPath(root); err != nil { + return fmt.Errorf("skill %q: %w", skillID, err) + } + abs, err := SafeRelPath(root, relPath) + if err != nil { + return err + } + if err := os.MkdirAll(filepath.Dir(abs), 0755); err != nil { + return err + } + return os.WriteFile(abs, content, 0644) +} + +func readFileHead(path string, max int64) ([]byte, error) { + f, err := os.Open(path) + if err != nil { + return nil, err + } + defer f.Close() + buf := make([]byte, max) + n, err := f.Read(buf) + if err != nil && n == 0 { + return nil, err + } + return buf[:n], nil +} + +func listScripts(skillsRoot, skillID string) ([]SkillScriptInfo, error) { + root := filepath.Join(SkillDir(skillsRoot, skillID), "scripts") + st, err := os.Stat(root) + if err != nil { + if os.IsNotExist(err) { + return nil, nil + } + return nil, err + } + if !st.IsDir() { + return nil, nil + } + var out []SkillScriptInfo + err = filepath.WalkDir(root, func(path string, d os.DirEntry, err error) error { + if err != nil { + return err + } + rel, e := filepath.Rel(root, path) + if e != nil { + return e + } + if rel == "." { + return nil + } + if d.IsDir() { + if strings.HasPrefix(d.Name(), ".") { + return filepath.SkipDir + } + if strings.Count(rel, string(os.PathSeparator)) >= maxScriptsDepth { + return filepath.SkipDir + } + return nil + } + if strings.HasPrefix(d.Name(), ".") { + return nil + } + relSkill := filepath.Join("scripts", rel) + full := filepath.Join(root, rel) + fi, err := os.Stat(full) + if err != nil || fi.IsDir() { + return nil + } + out = append(out, SkillScriptInfo{ + Name: filepath.Base(rel), + RelPath: filepath.ToSlash(relSkill), + Size: fi.Size(), + }) + return nil + }) + return out, err +} + +func countNonDirFiles(files []PackageFileInfo) int { + n := 0 + for _, f := range files { + if !f.IsDir && f.Path != "SKILL.md" { + n++ + } + } + return n +} diff --git a/internal/skillpackage/layout.go b/internal/skillpackage/layout.go new file mode 100644 index 00000000..275e1924 --- /dev/null +++ b/internal/skillpackage/layout.go @@ -0,0 +1,66 @@ +package skillpackage + +import ( + "fmt" + "os" + "path/filepath" + "strings" +) + +// SkillDir returns the absolute path to a skill package directory. +func SkillDir(skillsRoot, skillID string) string { + return filepath.Join(skillsRoot, skillID) +} + +// ResolveSKILLPath returns SKILL.md path or error if missing. +func ResolveSKILLPath(skillPath string) (string, error) { + md := filepath.Join(skillPath, "SKILL.md") + if st, err := os.Stat(md); err != nil || st.IsDir() { + return "", fmt.Errorf("missing SKILL.md in %q (Agent Skills standard)", filepath.Base(skillPath)) + } + return md, nil +} + +// SkillsRootFromConfig resolves cfg.SkillsDir relative to the config file directory. +func SkillsRootFromConfig(skillsDir string, configPath string) string { + if skillsDir == "" { + skillsDir = "skills" + } + configDir := filepath.Dir(configPath) + if !filepath.IsAbs(skillsDir) { + skillsDir = filepath.Join(configDir, skillsDir) + } + return skillsDir +} + +// DirLister lists skill package directory names under SkillsRoot. +type DirLister struct { + SkillsRoot string +} + +// ListSkills returns skill package directory names that contain SKILL.md. +func (d DirLister) ListSkills() ([]string, error) { + return ListSkillDirNames(d.SkillsRoot) +} + +// ListSkillDirNames returns subdirectory names under skillsRoot that contain SKILL.md. +func ListSkillDirNames(skillsRoot string) ([]string, error) { + if _, err := os.Stat(skillsRoot); os.IsNotExist(err) { + return nil, nil + } + entries, err := os.ReadDir(skillsRoot) + if err != nil { + return nil, fmt.Errorf("read skills directory: %w", err) + } + var names []string + for _, entry := range entries { + if !entry.IsDir() || strings.HasPrefix(entry.Name(), ".") { + continue + } + skillPath := filepath.Join(skillsRoot, entry.Name()) + if _, err := ResolveSKILLPath(skillPath); err == nil { + names = append(names, entry.Name()) + } + } + return names, nil +} diff --git a/internal/skillpackage/service.go b/internal/skillpackage/service.go new file mode 100644 index 00000000..52dbe90a --- /dev/null +++ b/internal/skillpackage/service.go @@ -0,0 +1,155 @@ +package skillpackage + +import ( + "fmt" + "os" + "sort" + "strings" +) + +// ListSkillSummaries scans skillsRoot and returns index rows for the admin API. +func ListSkillSummaries(skillsRoot string) ([]SkillSummary, error) { + names, err := ListSkillDirNames(skillsRoot) + if err != nil { + return nil, err + } + sort.Strings(names) + out := make([]SkillSummary, 0, len(names)) + for _, dirName := range names { + su, err := loadSummary(skillsRoot, dirName) + if err != nil { + continue + } + out = append(out, su) + } + return out, nil +} + +func loadSummary(skillsRoot, dirName string) (SkillSummary, error) { + skillPath := SkillDir(skillsRoot, dirName) + mdPath, err := ResolveSKILLPath(skillPath) + if err != nil { + return SkillSummary{}, err + } + raw, err := os.ReadFile(mdPath) + if err != nil { + return SkillSummary{}, err + } + man, _, err := ParseSkillMD(raw) + if err != nil { + return SkillSummary{}, err + } + if err := ValidateAgentSkillManifestInPackage(man, dirName); err != nil { + return SkillSummary{}, err + } + fi, err := os.Stat(mdPath) + if err != nil { + return SkillSummary{}, err + } + pfiles, err := ListPackageFiles(skillsRoot, dirName) + if err != nil { + return SkillSummary{}, err + } + nFiles := 0 + for _, p := range pfiles { + if !p.IsDir { + nFiles++ + } + } + scripts, err := listScripts(skillsRoot, dirName) + if err != nil { + return SkillSummary{}, err + } + ver := versionFromMetadata(man) + return SkillSummary{ + ID: dirName, + DirName: dirName, + Name: man.Name, + Description: man.Description, + Version: ver, + Path: skillPath, + Tags: manifestTags(man), + ScriptCount: len(scripts), + FileCount: nFiles, + FileSize: fi.Size(), + ModTime: fi.ModTime().Format("2006-01-02 15:04:05"), + Progressive: true, + }, nil +} + +// LoadOptions mirrors legacy API query params for the web admin. +type LoadOptions struct { + Depth string // summary | full + Section string +} + +// LoadSkill returns manifest + body + package listing for admin. +func LoadSkill(skillsRoot, skillID string, opt LoadOptions) (*SkillView, error) { + skillPath := SkillDir(skillsRoot, skillID) + mdPath, err := ResolveSKILLPath(skillPath) + if err != nil { + return nil, err + } + raw, err := os.ReadFile(mdPath) + if err != nil { + return nil, err + } + man, body, err := ParseSkillMD(raw) + if err != nil { + return nil, err + } + if err := ValidateAgentSkillManifestInPackage(man, skillID); err != nil { + return nil, err + } + pfiles, err := ListPackageFiles(skillsRoot, skillID) + if err != nil { + return nil, err + } + scripts, err := listScripts(skillsRoot, skillID) + if err != nil { + return nil, err + } + sort.Slice(scripts, func(i, j int) bool { return scripts[i].RelPath < scripts[j].RelPath }) + sections := deriveSections(body) + ver := versionFromMetadata(man) + v := &SkillView{ + DirName: skillID, + Name: man.Name, + Description: man.Description, + Content: body, + Path: skillPath, + Version: ver, + Tags: manifestTags(man), + Scripts: scripts, + Sections: sections, + PackageFiles: pfiles, + } + depth := strings.ToLower(strings.TrimSpace(opt.Depth)) + if depth == "" { + depth = "full" + } + sec := strings.TrimSpace(opt.Section) + if sec != "" { + mds := splitMarkdownSections(body) + chunk := findSectionContent(mds, sec) + if chunk == "" { + v.Content = fmt.Sprintf("_(section %q not found in SKILL.md for skill %s)_", sec, skillID) + } else { + v.Content = chunk + } + return v, nil + } + if depth == "summary" { + v.Content = buildSummaryMarkdown(man.Name, man.Description, v.Tags, scripts, sections, body) + } + return v, nil +} + +// ReadScriptText returns file content as string (for HTTP resource_path). +func ReadScriptText(skillsRoot, skillID, relPath string, maxBytes int64) (string, error) { + b, err := ReadPackageFile(skillsRoot, skillID, relPath, maxBytes) + if err != nil { + return "", err + } + return string(b), nil +} diff --git a/internal/skillpackage/types.go b/internal/skillpackage/types.go new file mode 100644 index 00000000..bf313425 --- /dev/null +++ b/internal/skillpackage/types.go @@ -0,0 +1,67 @@ +// Package skillpackage provides filesystem-backed Agent Skills layout (SKILL.md + package files) +// for HTTP admin APIs. Runtime discovery and progressive loading for agents use Eino ADK skill middleware. +package skillpackage + +// SkillManifest is parsed from SKILL.md front matter (https://agentskills.io/specification.md). +type SkillManifest struct { + Name string `yaml:"name"` + Description string `yaml:"description"` + License string `yaml:"license,omitempty"` + Compatibility string `yaml:"compatibility,omitempty"` + Metadata map[string]any `yaml:"metadata,omitempty"` + AllowedTools string `yaml:"allowed-tools,omitempty"` +} + +// SkillSummary is API metadata for one skill directory. +type SkillSummary struct { + ID string `json:"id"` + DirName string `json:"dir_name"` + Name string `json:"name"` + Description string `json:"description"` + Version string `json:"version"` + Path string `json:"path"` + Tags []string `json:"tags"` + Triggers []string `json:"triggers,omitempty"` + ScriptCount int `json:"script_count"` + FileCount int `json:"file_count"` + FileSize int64 `json:"file_size"` + ModTime string `json:"mod_time"` + Progressive bool `json:"progressive"` +} + +// SkillScriptInfo describes a file under scripts/. +type SkillScriptInfo struct { + Name string `json:"name"` + RelPath string `json:"rel_path"` + Description string `json:"description,omitempty"` + Size int64 `json:"size"` +} + +// SkillSection is derived from ## headings in SKILL.md. +type SkillSection struct { + ID string `json:"id"` + Title string `json:"title"` + Heading string `json:"heading"` + Level int `json:"level"` +} + +// PackageFileInfo describes one file inside a package. +type PackageFileInfo struct { + Path string `json:"path"` + Size int64 `json:"size"` + IsDir bool `json:"is_dir,omitempty"` +} + +// SkillView is a loaded package for admin / API. +type SkillView struct { + DirName string `json:"dir_name"` + Name string `json:"name"` + Description string `json:"description"` + Content string `json:"content"` + Path string `json:"path"` + Version string `json:"version"` + Tags []string `json:"tags"` + Scripts []SkillScriptInfo `json:"scripts,omitempty"` + Sections []SkillSection `json:"sections,omitempty"` + PackageFiles []PackageFileInfo `json:"package_files,omitempty"` +} diff --git a/internal/skillpackage/validate.go b/internal/skillpackage/validate.go new file mode 100644 index 00000000..79d8255c --- /dev/null +++ b/internal/skillpackage/validate.go @@ -0,0 +1,102 @@ +package skillpackage + +import ( + "fmt" + "strings" + "unicode/utf8" + + "gopkg.in/yaml.v3" +) + +var agentSkillsSpecFrontMatterKeys = map[string]struct{}{ + "name": {}, "description": {}, "license": {}, "compatibility": {}, + "metadata": {}, "allowed-tools": {}, +} + +// ValidateAgentSkillManifest enforces Agent Skills rules for name and description. +func ValidateAgentSkillManifest(m *SkillManifest) error { + if m == nil { + return fmt.Errorf("skill manifest is nil") + } + if strings.TrimSpace(m.Name) == "" { + return fmt.Errorf("SKILL.md front matter: name is required") + } + if strings.TrimSpace(m.Description) == "" { + return fmt.Errorf("SKILL.md front matter: description is required") + } + if utf8.RuneCountInString(m.Name) > 64 { + return fmt.Errorf("name exceeds 64 characters (Agent Skills limit)") + } + if utf8.RuneCountInString(m.Description) > 1024 { + return fmt.Errorf("description exceeds 1024 characters (Agent Skills limit)") + } + if m.Name != strings.ToLower(m.Name) { + return fmt.Errorf("name must be lowercase (Agent Skills)") + } + for _, r := range m.Name { + if !((r >= 'a' && r <= 'z') || (r >= '0' && r <= '9') || r == '-') { + return fmt.Errorf("name must contain only lowercase letters, numbers, hyphens (Agent Skills)") + } + } + if strings.HasPrefix(m.Name, "-") || strings.HasSuffix(m.Name, "-") { + return fmt.Errorf("name must not start or end with a hyphen (Agent Skills spec)") + } + if strings.Contains(m.Name, "--") { + return fmt.Errorf("name must not contain consecutive hyphens (Agent Skills spec)") + } + lname := strings.ToLower(m.Name) + if strings.Contains(lname, "anthropic") || strings.Contains(lname, "claude") { + return fmt.Errorf("name must not contain reserved words anthropic or claude") + } + return nil +} + +// ValidateAgentSkillManifestInPackage checks manifest and that name matches package directory. +func ValidateAgentSkillManifestInPackage(m *SkillManifest, packageDirName string) error { + if err := ValidateAgentSkillManifest(m); err != nil { + return err + } + if strings.TrimSpace(packageDirName) == "" { + return nil + } + if m.Name != packageDirName { + return fmt.Errorf("SKILL.md name %q must match directory name %q (Agent Skills spec)", m.Name, packageDirName) + } + return nil +} + +// ValidateOfficialFrontMatterTopLevelKeys rejects keys not in the open spec. +func ValidateOfficialFrontMatterTopLevelKeys(fmYAML string) error { + var top map[string]interface{} + if err := yaml.Unmarshal([]byte(fmYAML), &top); err != nil { + return fmt.Errorf("SKILL.md front matter: %w", err) + } + for k := range top { + if _, ok := agentSkillsSpecFrontMatterKeys[k]; !ok { + return fmt.Errorf("SKILL.md front matter: unsupported key %q (allowed: name, description, license, compatibility, metadata, allowed-tools — see https://agentskills.io/specification.md)", k) + } + } + return nil +} + +// ValidateSkillMDPackage validates SKILL.md bytes for writes. +func ValidateSkillMDPackage(raw []byte, packageDirName string) error { + fmYAML, body, err := ExtractSkillMDFrontMatterYAML(raw) + if err != nil { + return err + } + if err := ValidateOfficialFrontMatterTopLevelKeys(fmYAML); err != nil { + return err + } + if strings.TrimSpace(body) == "" { + return fmt.Errorf("SKILL.md: markdown body after front matter must not be empty") + } + var fm SkillManifest + if err := yaml.Unmarshal([]byte(fmYAML), &fm); err != nil { + return fmt.Errorf("SKILL.md front matter: %w", err) + } + if c := strings.TrimSpace(fm.Compatibility); c != "" && utf8.RuneCountInString(c) > 500 { + return fmt.Errorf("compatibility exceeds 500 characters (Agent Skills spec)") + } + return ValidateAgentSkillManifestInPackage(&fm, packageDirName) +}