Add files via upload

This commit is contained in:
公明
2026-08-15 02:14:36 +08:00
committed by GitHub
parent 9f38fda15f
commit ee4676a591
72 changed files with 11131 additions and 0 deletions
+55
View File
@@ -0,0 +1,55 @@
package audit
import (
"strings"
"cyberstrike-ai/internal/database"
"cyberstrike-ai/internal/security"
"github.com/gin-gonic/gin"
)
// RegisterConversationCreateHook records platform audit rows for every new conversation.
func RegisterConversationCreateHook(s *Service) {
if s == nil {
return
}
database.SetConversationCreateHook(func(conv *database.Conversation, meta database.ConversationCreateMeta) {
detail := map[string]interface{}{
"title": conv.Title,
"source": meta.Source,
}
if meta.WebShellConnectionID != "" {
detail["webshell_connection_id"] = meta.WebShellConnectionID
}
s.Record(nil, Entry{
Category: "conversation",
Action: "create",
Result: "success",
Message: "创建对话",
ResourceType: "conversation",
ResourceID: conv.ID,
Detail: detail,
ClientIP: meta.ClientIP,
SessionHint: meta.SessionHint,
})
})
}
// ConversationCreateMeta builds audit metadata for conversation creation.
func ConversationCreateMeta(source string) database.ConversationCreateMeta {
return database.ConversationCreateMeta{Source: strings.TrimSpace(source)}
}
// ConversationCreateMetaFromGin includes client IP and session hint when available.
func ConversationCreateMetaFromGin(c *gin.Context, source string) database.ConversationCreateMeta {
m := ConversationCreateMeta(source)
if c == nil {
return m
}
m.ClientIP = c.ClientIP()
if token := c.GetString(security.ContextAuthTokenKey); token != "" {
m.SessionHint = sessionHint(token)
}
return m
}
+9
View File
@@ -0,0 +1,9 @@
package audit
// RetentionDays returns configured retention; 0 means keep forever.
func (s *Service) RetentionDays() int {
if s == nil || s.cfg == nil {
return 0
}
return s.cfg.Audit.RetentionDaysEffective()
}
+29
View File
@@ -0,0 +1,29 @@
package audit
import "github.com/gin-gonic/gin"
// RecordAction writes a platform audit row with common defaults.
func (s *Service) RecordAction(c *gin.Context, category, action, result, message, resourceType, resourceID string, detail map[string]interface{}) {
if s == nil {
return
}
s.Record(c, Entry{
Category: category,
Action: action,
Result: result,
Message: message,
ResourceType: resourceType,
ResourceID: resourceID,
Detail: detail,
})
}
// RecordOK is a shorthand for successful operations.
func (s *Service) RecordOK(c *gin.Context, category, action, message, resourceType, resourceID string, detail map[string]interface{}) {
s.RecordAction(c, category, action, "success", message, resourceType, resourceID, detail)
}
// RecordFail is a shorthand for failed operations.
func (s *Service) RecordFail(c *gin.Context, category, action, message string, detail map[string]interface{}) {
s.RecordAction(c, category, action, "failure", message, "", "", detail)
}
+86
View File
@@ -0,0 +1,86 @@
package audit
import (
"strings"
"cyberstrike-ai/internal/database"
)
var auditActionsResourceRemoved = map[string]bool{
"delete": true,
"item_delete": true,
"connection_delete": true,
"listener_delete": true,
"session_delete": true,
"task_delete": true,
"execution_delete": true,
"execution_delete_batch": true,
"delete_queue": true,
"delete_batch_task": true,
"markdown_delete": true,
}
// ApplyResourceAvailability sets log.ResourceAvailable when the linked resource can be checked.
func ApplyResourceAvailability(db *database.DB, log *database.AuditLog) {
if log == nil || strings.TrimSpace(log.ResourceID) == "" {
return
}
if auditActionsResourceRemoved[log.Action] {
f := false
log.ResourceAvailable = &f
return
}
if db == nil {
return
}
available, known := resourceStillExists(db, log.ResourceType, log.ResourceID)
if known {
log.ResourceAvailable = &available
}
}
func resourceStillExists(db *database.DB, resourceType, resourceID string) (bool, bool) {
resourceID = strings.TrimSpace(resourceID)
if resourceID == "" {
return false, false
}
t := strings.TrimSpace(resourceType)
if t == "" {
if len(resourceID) > 8 && !strings.HasPrefix(resourceID, "c2_") {
t = "conversation"
} else {
return false, false
}
}
switch t {
case "conversation":
ok, err := db.ConversationExists(resourceID)
return ok, err == nil
case "vulnerability":
_, err := db.GetVulnerability(resourceID)
if err != nil {
return false, strings.Contains(err.Error(), "不存在")
}
return true, true
case "batch_queue":
_, err := db.GetBatchQueue(resourceID)
return err == nil, true
case "c2_listener":
_, err := db.GetC2Listener(resourceID)
return err == nil, true
case "c2_session":
_, err := db.GetC2Session(resourceID)
return err == nil, true
case "c2_task":
_, err := db.GetC2Task(resourceID)
return err == nil, true
case "webshell_connection":
c, err := db.GetWebshellConnection(resourceID)
return err == nil && c != nil, true
case "tool_execution":
_, err := db.GetToolExecution(resourceID)
return err == nil, true
default:
return false, false
}
}
+27
View File
@@ -0,0 +1,27 @@
package audit
import (
"time"
"go.uber.org/zap"
)
// auditRetentionPurgeInterval is how often PurgeExpired runs while the process is up (startup also purges once).
const auditRetentionPurgeInterval = time.Hour
// StartRetentionLoop periodically purges expired audit rows.
func StartRetentionLoop(s *Service, logger *zap.Logger) {
if s == nil {
return
}
go func() {
ticker := time.NewTicker(auditRetentionPurgeInterval)
defer ticker.Stop()
for range ticker.C {
s.PurgeExpired()
if logger != nil {
logger.Debug("audit retention tick completed")
}
}
}()
}
+58
View File
@@ -0,0 +1,58 @@
package audit
import (
"encoding/json"
"strings"
)
var sensitiveKeySubstrings = []string{
"password", "api_key", "apikey", "secret", "token", "authorization",
"credential", "private_key", "access_key",
}
// SanitizeDetail redacts sensitive keys and truncates serialized size.
func SanitizeDetail(detail map[string]interface{}, maxBytes int) map[string]interface{} {
if detail == nil {
return nil
}
if maxBytes <= 0 {
maxBytes = 8192
}
out := sanitizeValue("", detail)
if m, ok := out.(map[string]interface{}); ok {
b, _ := json.Marshal(m)
if len(b) > maxBytes {
return map[string]interface{}{
"_truncated": true,
"_preview": string(b[:maxBytes]),
}
}
return m
}
return map[string]interface{}{"value": out}
}
func sanitizeValue(key string, v interface{}) interface{} {
kl := strings.ToLower(key)
for _, sub := range sensitiveKeySubstrings {
if strings.Contains(kl, sub) {
return "***"
}
}
switch t := v.(type) {
case map[string]interface{}:
m := make(map[string]interface{}, len(t))
for k, val := range t {
m[k] = sanitizeValue(k, val)
}
return m
case []interface{}:
arr := make([]interface{}, len(t))
for i, val := range t {
arr[i] = sanitizeValue(key, val)
}
return arr
default:
return v
}
}
+177
View File
@@ -0,0 +1,177 @@
package audit
import (
"crypto/sha256"
"encoding/hex"
"strings"
"time"
"cyberstrike-ai/internal/config"
"cyberstrike-ai/internal/database"
"cyberstrike-ai/internal/security"
"github.com/gin-gonic/gin"
"github.com/google/uuid"
"go.uber.org/zap"
)
// Service persists platform audit logs.
type Service struct {
db *database.DB
cfg *config.Config
logger *zap.Logger
failThrottle *failureThrottle
}
// NewService creates an audit service.
func NewService(db *database.DB, cfg *config.Config, logger *zap.Logger) *Service {
return &Service{
db: db,
cfg: cfg,
logger: logger,
failThrottle: newFailureThrottle(),
}
}
// Enabled reports whether audit persistence is on.
func (s *Service) Enabled() bool {
if s == nil || s.cfg == nil {
return false
}
return s.cfg.Audit.EnabledEffective()
}
// Record writes one audit row from a Gin request context.
func (s *Service) Record(c *gin.Context, e Entry) {
if s == nil || !s.Enabled() || s.db == nil {
return
}
if strings.TrimSpace(e.Category) == "" || strings.TrimSpace(e.Action) == "" {
return
}
if e.Result == "failure" && !s.allowFailureAudit(c, e) {
return
}
if strings.TrimSpace(e.Result) == "" {
e.Result = "success"
}
if strings.TrimSpace(e.Level) == "" {
if e.Result == "failure" {
e.Level = "warn"
} else {
e.Level = "info"
}
}
if strings.TrimSpace(e.Actor) == "" {
if c != nil {
e.Actor = strings.TrimSpace(c.GetString(security.ContextUsernameKey))
}
if e.Actor == "" {
e.Actor = "admin"
}
}
maxDetail := s.cfg.Audit.MaxDetailBytesEffective()
detail := SanitizeDetail(e.Detail, maxDetail)
sessionHintVal := e.SessionHint
if sessionHintVal == "" && c != nil {
if token := c.GetString(security.ContextAuthTokenKey); token != "" {
sessionHintVal = sessionHint(token)
}
}
clientIPVal := e.ClientIP
if clientIPVal == "" {
clientIPVal = clientIP(c)
}
row := &database.AuditLog{
ID: "audit_" + strings.ReplaceAll(uuid.New().String(), "-", ""),
CreatedAt: time.Now(),
Level: e.Level,
Category: e.Category,
Action: e.Action,
Result: e.Result,
Actor: e.Actor,
SessionHint: sessionHintVal,
ClientIP: clientIPVal,
UserAgent: userAgent(c),
ResourceType: e.ResourceType,
ResourceID: e.ResourceID,
Message: e.Message,
Detail: detail,
}
if err := s.db.AppendAuditLog(row); err != nil && s.logger != nil {
s.logger.Warn("写入审计日志失败",
zap.String("action", e.Action),
zap.Error(err),
)
}
}
// RecordSystem writes an audit row without HTTP context (e.g. retention cleanup).
func (s *Service) RecordSystem(e Entry) {
s.Record(nil, e)
}
// PurgeExpired deletes rows older than retention_days when configured.
func (s *Service) PurgeExpired() {
if s == nil || s.db == nil || s.cfg == nil {
return
}
days := s.cfg.Audit.RetentionDaysEffective()
if days <= 0 {
return
}
cutoff := time.Now().AddDate(0, 0, -days)
n, err := s.db.DeleteAuditLogsBefore(cutoff)
if err != nil {
if s.logger != nil {
s.logger.Warn("清理过期审计日志失败", zap.Error(err))
}
return
}
if n > 0 && s.logger != nil {
s.logger.Info("已清理过期审计日志", zap.Int64("deleted", n))
}
}
// HintFromToken returns a short stable hash prefix for a session token.
func HintFromToken(token string) string {
return sessionHint(token)
}
func sessionHint(token string) string {
token = strings.TrimSpace(token)
if token == "" {
return ""
}
sum := sha256.Sum256([]byte(token))
return hex.EncodeToString(sum[:4])
}
func (s *Service) allowFailureAudit(c *gin.Context, e Entry) bool {
if !isAuthFailureThrottled(e.Category, e.Action) {
return true
}
cooldown := time.Duration(s.cfg.Audit.AuthFailureCooldownEffective()) * time.Second
key := authFailureThrottleKey(e.Category, e.Action, clientIP(c))
return s.failThrottle.allow(key, cooldown)
}
func clientIP(c *gin.Context) string {
if c == nil {
return ""
}
return c.ClientIP()
}
func userAgent(c *gin.Context) string {
if c == nil {
return ""
}
ua := c.GetHeader("User-Agent")
if len(ua) > 512 {
return ua[:512]
}
return ua
}
+55
View File
@@ -0,0 +1,55 @@
package audit
import (
"sync"
"time"
)
// failureThrottle deduplicates high-frequency failure audit rows (e.g. wrong password).
type failureThrottle struct {
mu sync.Mutex
last map[string]time.Time
}
func newFailureThrottle() *failureThrottle {
return &failureThrottle{last: make(map[string]time.Time)}
}
// allow reports whether a row with the given key may be written now.
func (t *failureThrottle) allow(key string, cooldown time.Duration) bool {
if t == nil || cooldown <= 0 || key == "" {
return true
}
now := time.Now()
t.mu.Lock()
defer t.mu.Unlock()
if prev, ok := t.last[key]; ok && now.Sub(prev) < cooldown {
return false
}
t.last[key] = now
if len(t.last) > 4096 {
for k, ts := range t.last {
if now.Sub(ts) > cooldown*2 {
delete(t.last, k)
}
}
}
return true
}
// authFailureThrottleKey builds a per-IP key for auth failure deduplication.
func authFailureThrottleKey(category, action, clientIP string) string {
return category + ":" + action + ":" + clientIP
}
func isAuthFailureThrottled(category, action string) bool {
if category != "auth" {
return false
}
switch action {
case "login", "change_password":
return true
default:
return false
}
}
+16
View File
@@ -0,0 +1,16 @@
package audit
// Entry describes one platform audit record (not chat/tool execution bodies).
type Entry struct {
Level string
Category string
Action string
Result string // success | failure
Actor string
SessionHint string
ResourceType string
ResourceID string
Message string
Detail map[string]interface{}
ClientIP string // optional when c is nil (robot, batch, DB hook)
}
+72
View File
@@ -0,0 +1,72 @@
package authctx
import (
"context"
"strings"
)
// Principal is the immutable authorization identity propagated beyond the
// transport layer into Agent, MCP and background task contexts.
type Principal struct {
UserID string
Username string
Permissions map[string]bool
PermissionScopes map[string]string
Scope string
}
type principalContextKey struct{}
func NewPrincipal(userID, username, scope string, permissions map[string]bool) Principal {
return NewPrincipalWithScopes(userID, username, scope, permissions, nil)
}
func NewPrincipalWithScopes(userID, username, scope string, permissions map[string]bool, permissionScopes map[string]string) Principal {
permissionCopy := make(map[string]bool, len(permissions))
scopeCopy := make(map[string]string, len(permissionScopes))
for permission, allowed := range permissions {
if allowed {
permissionCopy[permission] = true
if permissionScope := strings.TrimSpace(permissionScopes[permission]); permissionScope != "" {
scopeCopy[permission] = permissionScope
}
}
}
return Principal{
UserID: strings.TrimSpace(userID), Username: strings.TrimSpace(username),
Scope: strings.TrimSpace(scope), Permissions: permissionCopy, PermissionScopes: scopeCopy,
}
}
func WithPrincipal(ctx context.Context, principal Principal) context.Context {
if ctx == nil {
ctx = context.Background()
}
if strings.TrimSpace(principal.UserID) == "" {
return ctx
}
return context.WithValue(ctx, principalContextKey{}, principal)
}
func PrincipalFromContext(ctx context.Context) (Principal, bool) {
if ctx == nil {
return Principal{}, false
}
principal, ok := ctx.Value(principalContextKey{}).(Principal)
return principal, ok && strings.TrimSpace(principal.UserID) != ""
}
func (p Principal) HasPermission(permission string) bool {
return p.Permissions[strings.TrimSpace(permission)]
}
// ScopeFor returns the scope attached to the permission that authorizes the
// current action. Falling back to Scope keeps explicit service principals and
// legacy callers compatible without reintroducing cross-role scope widening.
func (p Principal) ScopeFor(permission string) string {
permission = strings.TrimSpace(permission)
if scope := strings.TrimSpace(p.PermissionScopes[permission]); scope != "" {
return scope
}
return strings.TrimSpace(p.Scope)
}
+71
View File
@@ -0,0 +1,71 @@
package hitl
import (
"time"
"cyberstrike-ai/internal/config"
"cyberstrike-ai/internal/database"
"go.uber.org/zap"
)
const retentionPurgeInterval = time.Hour
// Service manages HITL audit log retention (decided hitl_interrupts rows).
type Service struct {
db *database.DB
cfg *config.Config
logger *zap.Logger
}
// NewService creates a HITL audit log retention service.
func NewService(db *database.DB, cfg *config.Config, logger *zap.Logger) *Service {
return &Service{db: db, cfg: cfg, logger: logger}
}
// RetentionDays returns configured retention; 0 means keep forever.
func (s *Service) RetentionDays() int {
if s == nil || s.cfg == nil {
return config.HitlConfig{}.RetentionDaysEffective()
}
return s.cfg.Hitl.RetentionDaysEffective()
}
// PurgeExpired deletes decided HITL log rows older than retention_days when configured.
func (s *Service) PurgeExpired() {
if s == nil || s.db == nil || s.cfg == nil {
return
}
days := s.cfg.Hitl.RetentionDaysEffective()
if days <= 0 {
return
}
cutoff := time.Now().AddDate(0, 0, -days)
n, err := s.db.PurgeHitlInterruptLogsBefore(cutoff)
if err != nil {
if s.logger != nil {
s.logger.Warn("清理过期人机协同审计日志失败", zap.Error(err))
}
return
}
if n > 0 && s.logger != nil {
s.logger.Info("已清理过期人机协同审计日志", zap.Int64("deleted", n), zap.Int("retention_days", days))
}
}
// StartRetentionLoop periodically purges expired HITL audit log rows.
func StartRetentionLoop(s *Service, logger *zap.Logger) {
if s == nil {
return
}
go func() {
ticker := time.NewTicker(retentionPurgeInterval)
defer ticker.Stop()
for range ticker.C {
s.PurgeExpired()
if logger != nil {
logger.Debug("hitl audit log retention tick completed")
}
}
}()
}
+50
View File
@@ -0,0 +1,50 @@
package hitl
import (
"path/filepath"
"testing"
"time"
appconfig "cyberstrike-ai/internal/config"
"cyberstrike-ai/internal/database"
"go.uber.org/zap"
)
func TestServicePurgeExpired_respectsZeroRetention(t *testing.T) {
dbPath := filepath.Join(t.TempDir(), "hitl.db")
db, err := database.NewDB(dbPath, zap.NewNop())
if err != nil {
t.Fatalf("NewDB: %v", err)
}
defer db.Close()
if _, err := db.Exec(`CREATE TABLE IF NOT EXISTS hitl_interrupts (
id TEXT PRIMARY KEY,
conversation_id TEXT NOT NULL,
mode TEXT NOT NULL,
tool_name TEXT NOT NULL,
status TEXT NOT NULL,
decision TEXT,
created_at DATETIME NOT NULL,
decided_at DATETIME
)`); err != nil {
t.Fatalf("create table: %v", err)
}
old := time.Now().AddDate(0, 0, -100).UTC().Format(time.RFC3339)
if _, err := db.Exec(`INSERT INTO hitl_interrupts
(id, conversation_id, mode, tool_name, status, decision, created_at, decided_at)
VALUES ('old-1', 'c1', 'approval', 'exec', 'decided', 'approve', ?, ?)`, old, old); err != nil {
t.Fatalf("insert: %v", err)
}
zero := 0
svc := NewService(db, &appconfig.Config{
Hitl: appconfig.HitlConfig{RetentionDays: &zero},
}, zap.NewNop())
svc.PurgeExpired()
if err := db.QueryRow(`SELECT id FROM hitl_interrupts WHERE id = 'old-1'`).Scan(new(string)); err != nil {
t.Fatalf("record should remain when retention_days=0: %v", err)
}
}
+67
View File
@@ -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,
})
}
+129
View File
@@ -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)
}
+14
View File
@@ -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")
}
}
@@ -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)
+24
View File
@@ -0,0 +1,24 @@
package knowledge
import (
"context"
"fmt"
"github.com/cloudwego/eino/compose"
"github.com/cloudwego/eino/schema"
)
// BuildKnowledgeRetrieveChain 编译「查询字符串 → 文档列表」的 Eino ChainMultiQuery → 向量 → 重排 → 后处理)。
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)
}
@@ -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")
}
}
@@ -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)
+142
View File
@@ -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)
+251
View File
@@ -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 EmbedderopenAIConfig 用于在知识库未单独配置 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}
}
+91
View File
@@ -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
})
}
+21
View File
@@ -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)
}
}
}
+435
View File
@@ -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
}
@@ -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()
}
+885
View File
@@ -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
}
+226
View File
@@ -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)
+97
View File
@@ -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")
}
}
+216
View File
@@ -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
}
@@ -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))
}
}
+334
View File
@@ -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,
}
}
+51
View File
@@ -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)
}
+323
View File
@@ -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 → VectorEinoRetrieverEino 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<!-- METADATA: %s -->", 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()
}
+123
View File
@@ -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
}
+76
View File
@@ -0,0 +1,76 @@
package knowledge
import (
"context"
"fmt"
"net/http"
"strings"
"time"
"cyberstrike-ai/internal/config"
"cyberstrike-ai/internal/openai"
einoopenai "github.com/cloudwego/eino-ext/components/model/openai"
"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
httpClient := openai.NewEinoHTTPClient(openAI, &http.Client{Timeout: 120 * time.Second})
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"
}
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
}
+132
View File
@@ -0,0 +1,132 @@
// Package projectprompt 提供项目黑板相关的系统提示文本(纯字符串,无 database 依赖)。
// 供 agent / multiagent 等包引用,避免 agent → project 导入环导致 gopls 元数据失败。
package projectprompt
import (
"strings"
"cyberstrike-ai/internal/mcp/builtin"
)
const (
factRhythmCore = "勿等会话结束或收尾再批量写入。每**确认**一条新认知(开放端口/服务版本、入口路径、认证态或凭据特征、可利用点或攻击面变化)后,**立即**调用 `upsert_project_fact`(同 fact_key 覆盖更新)。每**验证**出一条可复现漏洞(含 POC/影响)后,**立即**调用 `record_vulnerability`;与事实可各记一次。继续下一步工作前优先落库,避免上下文压缩后细节丢失。未绑项目时说明无法写黑板,仍在本轮保留证据摘要。"
factRhythmCoordinatorSuffix = "委派/子任务返回新认知或漏洞时,由协调者及时写入,勿假定子代理已记。"
factRhythmSubAgentSuffix = "若工具集中无上述工具,须在交付物末尾给出「待落库」结构化条目(fact_key 建议、summary、body/POC 要点),供协调者**立即**写入。"
)
// FactRecordingIncrementalRhythmMarkdown 返回边渗透边记录节奏(Markdown,供 agents/*.md 与文档对齐)。
func FactRecordingIncrementalRhythmMarkdown(coordinator, subAgent bool) string {
var b strings.Builder
b.WriteString("- **边渗透边记录(强制节奏)**:")
b.WriteString(factRhythmCore)
if coordinator {
b.WriteString(factRhythmCoordinatorSuffix)
}
if subAgent {
b.WriteString(factRhythmSubAgentSuffix)
}
return b.String()
}
func factRecordingIncrementalRhythmBuiltin(coordinator, subAgent bool) string {
var b strings.Builder
b.WriteString("- **边渗透边记录(强制节奏)**:勿等会话结束或收尾再批量写入。每**确认**一条新认知(开放端口/服务版本、入口路径、认证态或凭据特征、可利用点或攻击面变化)后,**立即**调用 ")
b.WriteString(builtin.ToolUpsertProjectFact)
b.WriteString("(同 fact_key 覆盖更新)。每**验证**出一条可复现漏洞(含 POC/影响)后,**立即**调用 ")
b.WriteString(builtin.ToolRecordVulnerability)
b.WriteString(";与事实可各记一次。继续下一步工作前优先落库,避免上下文压缩后细节丢失。未绑项目时说明无法写黑板,仍在本轮保留证据摘要。")
if coordinator {
b.WriteString(factRhythmCoordinatorSuffix)
}
if subAgent {
b.WriteString(factRhythmSubAgentSuffix)
}
return b.String()
}
func factEdgeRecordingGuidance() string {
return `### 事实关系边(links
- 写入 **finding / chain / exploit / poc** 时,**必须**在 ` + "`upsert_project_fact`" + ` 中提供 ` + "`links`" + `**推荐 ` + "`from`" + `**:来源 fact 指向当前 fact,即 ` + "`from`" + ` → 当前 ` + "`fact_key`" + `)。
- **最少要求**finding 类至少 1 条 from=target/* + type=discovered_on(即 target → finding);在 finding 上记录 exploit 用 from=exploit/* + type=exploits(即 exploit → finding)。
- **常用 type**` + "`discovered_on`" + `(发现在哪)、` + "`depends_on`" + `(复现前置)、` + "`leads_to`" + `(认知推进)、` + "`enables`" + `(扩大攻击面)、` + "`exploits`" + `(利用关系)、` + "`contains`" + `(资产包含)、` + "`part_of`" + `(属于链/组)、` + "`supports`" + `(证据支撑)。
- 更新时:**省略 links 保留已有边**;传入 links 则**替换**全部关系边(from → 当前 fact)。
- body 中「依赖事实」段落可与 links 并存(人读);结构化关系以 links 为准。`
}
func factRecordingGuidanceBlock() string {
return `### 事实写入规范(审计复现 / 知识沉淀)
- **summary**:索引用一行,须含「什么 + 在哪 + 如何触发/验证」要点,禁止只写结论(如仅写「存在 SQLi」)。
- **body**:完整可复现上下文,写入 ` + "`upsert_project_fact`" + ` 的 body 字段;索引不含 body,后续会话须靠 ` + "`get_project_fact`" + ` 取回。
- **category / fact_key 建议**
- 环境认知:` + "`target/`" + `` + "`auth/`" + `` + "`infra/`" + `` + "`business/`" + `body 用环境模板即可)
- 发现与利用:` + "`finding/`" + `` + "`chain/`" + `` + "`exploit/`" + `` + "`poc/`" + `(**必须**用攻击链模板填满 body:入口、逐步攻击链、原始请求/响应或命令、证据、关联漏洞 ID)
- **与漏洞记录分工**` + "`record_vulnerability`" + ` 记可交付 findings;事实记**复现所需的全部上下文**(含失败尝试、绕过、依赖会话),二者可各记一次。
- 更新同一发现时保持相同 ` + "`fact_key`" + ` 覆盖写入,勿散落多个 key 导致上下文丢失。`
}
// FactRecordingBlackboardSection 项目黑板与漏洞记录的完整系统提示块(单/多 Agent 主代理共用)。
func FactRecordingBlackboardSection(coordinatorDelegate bool) string {
var b strings.Builder
b.WriteString("## 项目黑板(事实)与漏洞记录(分离)\n\n")
b.WriteString("当前对话若已绑定项目,系统会自动注入「项目黑板索引」(仅 fact_key + 摘要)。**摘要不足时必须调用 ")
b.WriteString(builtin.ToolGetProjectFact)
b.WriteString("(fact_key) 获取 body,禁止凭摘要臆造细节。**\n\n")
b.WriteString(factRecordingIncrementalRhythmBuiltin(coordinatorDelegate, false))
b.WriteString("\n\n")
b.WriteString("- **环境/目标/认证等认知**(非正式漏洞条目):使用 ")
b.WriteString(builtin.ToolUpsertProjectFact)
b.WriteString("fact_key 建议 `category/slug`(如 target/primary_domain),同 key 覆盖更新;body 记端口/版本/凭据特征与证据来源。\n")
b.WriteString("- **发现与利用上下文**(审计复现):fact_key 建议 finding/、chain/、exploit/、poc/ 前缀;**body 必填**完整攻击链(入口 → 步骤 → 原始请求/响应或命令 → 现象 → 关联 related_vulnerability_id),**禁止仅写结论**summary 写「什么 + 在哪 + 如何验证」一行要点。\n")
b.WriteString("- **可交付漏洞**:使用 ")
b.WriteString(builtin.ToolRecordVulnerability)
b.WriteString(",含标题、严重程度、类型、目标、证明(POC)、影响、修复建议。记前可先 ")
b.WriteString(builtin.ToolListVulnerabilities)
b.WriteString(" 查重,详情用 ")
b.WriteString(builtin.ToolGetVulnerability)
b.WriteString("(id)(默认仅当前项目/会话)。\n")
b.WriteString("- 同一发现可能需**各记一次**(事实记**完整攻击链与 exploit 细节**供复现,漏洞记正式 findings)。误报用 ")
b.WriteString(builtin.ToolDeprecateProjectFact)
b.WriteString(" 或漏洞状态 false_positive。\n")
b.WriteString("- 事实多时用 ")
b.WriteString(builtin.ToolListProjectFacts)
b.WriteString(" / ")
b.WriteString(builtin.ToolSearchProjectFacts)
b.WriteString(" 检索。\n\n")
b.WriteString(factEdgeRecordingGuidance())
b.WriteString("\n\n")
b.WriteString(factRecordingGuidanceBlock())
b.WriteString("\n\n严重程度:critical / high / medium / low / info。证明须含足够证据(请求响应、截图、命令输出等)。")
return b.String()
}
// FactRecordingSubAgentSection 子代理边渗透边记录(无工具时输出待落库条目)。
func FactRecordingSubAgentSection() string {
return "## 边渗透边记录\n\n" + factRecordingIncrementalRhythmBuiltin(false, true) + "\n"
}
// FactRecordingBlackboardSectionMarkdown 与 FactRecordingBlackboardSection 等价的 Markdown(工具名为字面量,供 agents/*.md)。
func FactRecordingBlackboardSectionMarkdown(coordinatorDelegate bool) string {
var b strings.Builder
b.WriteString("## 项目黑板(事实)与漏洞记录(分离)\n\n")
b.WriteString("当前对话若已绑定项目,系统会自动注入「项目黑板索引」(仅 `fact_key` + 摘要)。**摘要不足时必须调用 `get_project_fact(fact_key)` 获取 body,禁止凭摘要臆造细节。**\n\n")
b.WriteString(FactRecordingIncrementalRhythmMarkdown(coordinatorDelegate, false))
b.WriteString("\n\n")
b.WriteString("- **环境/目标/认证等认知**(非正式漏洞):使用 **`upsert_project_fact`**`fact_key` 建议 `category/slug`(如 `target/primary_domain`),同 key 覆盖更新;body 记端口/版本/凭据特征与证据来源。\n")
b.WriteString("- **发现与利用上下文**(审计复现):`fact_key` 建议 `finding/`、`chain/`、`exploit/`、`poc/` 前缀;**body 必填**完整攻击链(入口 → 步骤 → 原始请求/响应或命令 → 现象 → 关联 `related_vulnerability_id`),**禁止仅写结论**summary 写「什么 + 在哪 + 如何验证」一行要点。\n")
b.WriteString("- **可交付漏洞**:使用 **`record_vulnerability`**(标题、描述、严重程度、类型、目标、证明 POC、影响、修复建议)。严重程度 critical / high / medium / low / info。\n")
b.WriteString("- 同一发现可能需**各记一次**(事实记可复现攻击链,漏洞记正式 findings)。误报用 **`deprecate_project_fact`** 或漏洞状态 false_positive。\n")
b.WriteString("- 事实多时用 **`list_project_facts`** / **`search_project_facts`** 检索。\n\n")
b.WriteString(factEdgeRecordingGuidance())
b.WriteString("\n\n")
b.WriteString(factRecordingGuidanceBlock())
b.WriteString("\n\n严重程度:critical / high / medium / low / info。证明须含足够证据(请求响应、截图、命令输出等)。")
return b.String()
}
// FactEdgeRecordingGuidance 写入边时的 Agent 规范(供 project 包复用)。
func FactEdgeRecordingGuidance() string { return factEdgeRecordingGuidance() }
// FactRecordingGuidanceBlock 事实写入规范块(供 project 包复用)。
func FactRecordingGuidanceBlock() string { return factRecordingGuidanceBlock() }
+11
View File
@@ -0,0 +1,11 @@
package projectprompt
// ShellExecExecuteGuidanceSection 供单代理/多代理系统提示追加:exec 与 execute 分工(尽量短)。
func ShellExecExecuteGuidanceSection() string {
return `Shellexec/execute):有专用 MCP 工具时优先专用工具;系统命令(管道、workdir、后台 &)用 execskills/ 内脚本(配合 read_file、skill)用 execute;多步扫描分拆调用,禁止一条 shell 串多个扫描器。长脚本、请求体或 Payload 必须先用 write_file 写入会话工作目录,再用 exec/execute 执行短命令;禁止把长内容嵌入 command。下载/临时文件须写入系统提示中的「会话工作目录」,禁止用 /tmp。`
}
// ShellExecExecuteGuidanceReconSuffix 侦察子代理可选追加(一行)。
func ShellExecExecuteGuidanceReconSuffix() string {
return `枚举优先 subfinder、amass 等专用 MCP,勿 exec/execute 拼长链。`
}
+428
View File
@@ -0,0 +1,428 @@
// Package reasoning maps user/config intent to CloudWeGo Eino OpenAI ChatModel fields
// (ReasoningEffort, ExtraFields such as thinking / reasoning_effort / output_config).
package reasoning
import (
"strings"
"cyberstrike-ai/internal/config"
einoopenai "github.com/cloudwego/eino-ext/components/model/openai"
)
// ClientIntent is optional per-request override from ChatRequest.reasoning.
type ClientIntent struct {
Mode string
Effort string
}
type wireProfile int
const (
wireNone wireProfile = iota
wireClaude
wireDeepseek
wireOpenAI
wireOutputConfig
)
// ApplyPlanExecutePlannerModelConfig configures the plan_execute planner/replanner
// ChatModel. Those Eino agents call WithToolChoice(Forced); several gateways reject
// thinking / reasoning fields on the same request (tool_choice required/object).
// Executor should keep the normal ApplyToEinoChatModelConfig path.
func ApplyPlanExecutePlannerModelConfig(cfg *einoopenai.ChatModelConfig, oa *config.OpenAIConfig) {
if cfg == nil || oa == nil {
return
}
mergeExtraRequestFields(cfg, oa.Reasoning.ExtraRequestFields)
clearReasoningFromChatModelConfig(cfg)
if resolveWireProfile(oa, &oa.Reasoning) == wireDeepseek || oa.IsDeepSeekEndpointOrModel() {
// DeepSeek enables thinking by default, so omission would not actually
// disable it for the planner's forced tool-choice requests.
applyThinkingDisabled(cfg)
}
}
func clearReasoningFromChatModelConfig(cfg *einoopenai.ChatModelConfig) {
if cfg == nil {
return
}
cfg.ReasoningEffort = ""
if cfg.ExtraFields != nil {
for _, key := range []string{"thinking", "reasoning_effort", "output_config", "reasoning"} {
delete(cfg.ExtraFields, key)
}
if len(cfg.ExtraFields) == 0 {
cfg.ExtraFields = nil
}
}
}
func mergeExtraRequestFields(cfg *einoopenai.ChatModelConfig, fields map[string]interface{}) {
if cfg == nil || len(fields) == 0 {
return
}
if cfg.ExtraFields == nil {
cfg.ExtraFields = make(map[string]any, len(fields))
}
for k, v := range fields {
cfg.ExtraFields[k] = v
}
}
// ApplyToEinoChatModelConfig merges reasoning-related options into cfg.
// Precondition: cfg already has APIKey, BaseURL, Model, HTTPClient set.
func ApplyToEinoChatModelConfig(cfg *einoopenai.ChatModelConfig, oa *config.OpenAIConfig, client *ClientIntent) {
if cfg == nil || oa == nil {
return
}
sr := &oa.Reasoning
allowClient := sr.AllowClientReasoningEffective()
mode := effectiveMode(sr, client, allowClient)
// Admin-defined root fields are independent of the selected reasoning wire
// profile. Merge them first so mode=off can remove only reasoning controls
// while preserving unrelated gateway options.
mergeExtraRequestFields(cfg, sr.ExtraRequestFields)
if mode == "off" {
clearReasoningFromChatModelConfig(cfg)
// Strict OpenAI endpoints reject unknown `thinking` fields, whereas the
// DeepSeek API enables thinking by default and requires an explicit
// thinking.type=disabled switch. Detect the actual DeepSeek target even
// when the configured reasoning profile was left as openai_compat.
if resolveWireProfile(oa, sr) == wireDeepseek || oa.IsDeepSeekEndpointOrModel() {
applyThinkingDisabled(cfg)
}
return
}
// Claude (Anthropic): merge admin extras first; optional extended thinking maps to top-level `thinking`
// (see internal/openai convertOpenAIToClaude). DeepSeek/OpenAI-style fields are not sent.
if strings.EqualFold(strings.TrimSpace(oa.Provider), "claude") ||
strings.EqualFold(strings.TrimSpace(oa.Provider), "anthropic") {
applyClaudeExtendedThinking(cfg, mode, effectiveEffort(sr, client, allowClient), oa.Model)
return
}
effort := effectiveEffort(sr, client, allowClient)
prof := resolveWireProfile(oa, sr)
switch prof {
case wireClaude, wireNone:
return
case wireDeepseek:
applyDeepseek(cfg, mode, effort)
case wireOutputConfig:
applyOutputConfigEffort(cfg, mode, effort)
default: // wireOpenAI
applyOpenAICompat(cfg, mode, effort)
}
}
// AgenticOpenAIExtraFields returns reasoning-related request fields for
// agenticopenai.ChatConfig. The agentic chat backend currently exposes provider
// extensions through ExtraFields instead of typed ReasoningEffort fields.
func AgenticOpenAIExtraFields(oa *config.OpenAIConfig, client *ClientIntent) map[string]any {
if oa == nil {
return nil
}
sr := &oa.Reasoning
allowClient := sr.AllowClientReasoningEffective()
mode := effectiveMode(sr, client, allowClient)
fields := cloneExtraRequestFields(sr.ExtraRequestFields)
if mode == "off" {
clearReasoningExtraFields(fields)
if resolveWireProfile(oa, sr) == wireDeepseek || oa.IsDeepSeekEndpointOrModel() {
if fields == nil {
fields = make(map[string]any)
}
fields["thinking"] = map[string]any{"type": "disabled"}
}
return fields
}
if strings.EqualFold(strings.TrimSpace(oa.Provider), "claude") ||
strings.EqualFold(strings.TrimSpace(oa.Provider), "anthropic") {
return fields
}
effort := effectiveEffort(sr, client, allowClient)
switch resolveWireProfile(oa, sr) {
case wireDeepseek:
if mode == "auto" || mode == "on" {
if fields == nil {
fields = make(map[string]any)
}
fields["thinking"] = map[string]any{"type": "enabled"}
}
if effort != "" {
if fields == nil {
fields = make(map[string]any)
}
fields["reasoning_effort"] = effortStringForAPI(effort)
}
case wireOutputConfig:
e := effort
if mode == "on" && e == "" {
e = "high"
}
if e != "" {
if fields == nil {
fields = make(map[string]any)
}
fields["output_config"] = map[string]any{"effort": effortStringForAPI(e)}
}
default:
e := effort
if mode == "on" && e == "" {
e = "medium"
}
if e != "" {
if fields == nil {
fields = make(map[string]any)
}
fields["reasoning_effort"] = effortStringForAPI(e)
}
}
return fields
}
// AgenticOpenAIPlannerExtraFields mirrors ApplyPlanExecutePlannerModelConfig for
// agenticopenai.ChatConfig: keep admin extras, strip reasoning controls, and
// explicitly disable DeepSeek thinking where omission would still think.
func AgenticOpenAIPlannerExtraFields(oa *config.OpenAIConfig) map[string]any {
if oa == nil {
return nil
}
fields := cloneExtraRequestFields(oa.Reasoning.ExtraRequestFields)
clearReasoningExtraFields(fields)
if resolveWireProfile(oa, &oa.Reasoning) == wireDeepseek || oa.IsDeepSeekEndpointOrModel() {
if fields == nil {
fields = make(map[string]any)
}
fields["thinking"] = map[string]any{"type": "disabled"}
}
return fields
}
func cloneExtraRequestFields(fields map[string]interface{}) map[string]any {
if len(fields) == 0 {
return nil
}
out := make(map[string]any, len(fields))
for k, v := range fields {
out[k] = v
}
return out
}
func clearReasoningExtraFields(fields map[string]any) {
for _, key := range []string{"thinking", "reasoning_effort", "output_config", "reasoning"} {
delete(fields, key)
}
}
// applyClaudeExtendedThinking sets Anthropic Messages API fields per official guidance:
// - Adaptive models (4.6+): thinking.type=adaptive; output_config.effort only when user sets effort (API default is high).
// - Sonnet 3.7: thinking.type=enabled + budget_tokens=10000 (doc example); effort is not mapped — use extra_request_fields for custom budget.
func applyClaudeExtendedThinking(cfg *einoopenai.ChatModelConfig, mode, effort, model string) {
if cfg == nil || mode == "off" {
return
}
if cfg.ExtraFields == nil {
cfg.ExtraFields = make(map[string]any)
}
m := strings.ToLower(strings.TrimSpace(model))
sonnet37 := isClaudeSonnet37(m)
if _, exists := cfg.ExtraFields["thinking"]; !exists {
cfg.ExtraFields["thinking"] = claudeThinkingForModel(m, sonnet37)
}
applyClaudeOutputConfigEffort(cfg, effort, sonnet37)
}
// claudeSonnet37DefaultBudgetTokens matches Anthropic extended-thinking documentation examples (budget_tokens with max_tokens 16000).
const claudeSonnet37DefaultBudgetTokens = 10000
func isClaudeSonnet37(m string) bool {
return strings.Contains(m, "claude-3-7-sonnet") ||
strings.Contains(m, "3-7-sonnet") ||
strings.Contains(m, "sonnet-3.7")
}
func claudeThinkingForModel(m string, sonnet37 bool) map[string]any {
if sonnet37 {
return map[string]any{
"type": "enabled",
"budget_tokens": claudeSonnet37DefaultBudgetTokens,
"display": "summarized",
}
}
// Opus 4.7+: manual enabled+budget rejected — adaptive only.
if strings.Contains(m, "opus-4-7") || strings.Contains(m, "opus-4.7") {
return map[string]any{
"type": "adaptive",
"display": "summarized",
}
}
return map[string]any{
"type": "adaptive",
"display": "summarized",
}
}
// applyClaudeOutputConfigEffort sets top-level output_config.effort only when effort is explicitly configured.
// Omitted effort uses the API default (high); do not inject effort on mode:on alone.
func applyClaudeOutputConfigEffort(cfg *einoopenai.ChatModelConfig, effort string, sonnet37 bool) {
if cfg == nil || sonnet37 {
return
}
if _, exists := cfg.ExtraFields["output_config"]; exists {
return
}
e := effortStringForAPI(effort)
if e == "" {
return
}
cfg.ExtraFields["output_config"] = map[string]any{"effort": e}
}
func effectiveMode(sr *config.OpenAIReasoningConfig, client *ClientIntent, allowClient bool) string {
server := strings.ToLower(strings.TrimSpace(sr.ModeEffective()))
if server == "" || server == "default" {
server = "auto"
}
if !allowClient || client == nil {
return server
}
cm := strings.ToLower(strings.TrimSpace(client.Mode))
if cm == "" || cm == "default" {
return server
}
return cm
}
func effectiveEffort(sr *config.OpenAIReasoningConfig, client *ClientIntent, allowClient bool) string {
se := normalizeEffort(sr.Effort)
if !allowClient || client == nil {
return se
}
ce := normalizeEffort(client.Effort)
if ce != "" {
return ce
}
return se
}
func normalizeEffort(s string) string {
e := strings.ToLower(strings.TrimSpace(s))
switch e {
case "low", "medium", "high", "max", "xhigh":
return e
default:
return ""
}
}
// usesExtraFieldsReasoningEffort 为 Eino 无枚举的最高档 effort,经 ExtraFields 原样下发(max / xhigh 由网关自行识别,不做互转)。
func usesExtraFieldsReasoningEffort(e string) bool {
return e == "max" || e == "xhigh"
}
func resolveWireProfile(oa *config.OpenAIConfig, sr *config.OpenAIReasoningConfig) wireProfile {
provider := strings.TrimSpace(oa.Provider)
if strings.EqualFold(provider, "claude") || strings.EqualFold(provider, "anthropic") {
return wireClaude
}
p := strings.ToLower(strings.TrimSpace(sr.ProfileEffective()))
switch p {
case "output_config", "output_config_effort":
return wireOutputConfig
case "openai", "openai_compat":
return wireOpenAI
case "deepseek", "deepseek_compat":
return wireDeepseek
case "auto", "":
if oa.IsDeepSeekEndpointOrModel() {
return wireDeepseek
}
return wireOpenAI
default:
return wireOpenAI
}
}
func applyThinkingDisabled(cfg *einoopenai.ChatModelConfig) {
if cfg == nil {
return
}
if cfg.ExtraFields == nil {
cfg.ExtraFields = make(map[string]any)
}
cfg.ExtraFields["thinking"] = map[string]any{"type": "disabled"}
}
func applyDeepseek(cfg *einoopenai.ChatModelConfig, mode, effort string) {
// auto: enable thinking for DeepSeek line; on: same; auto without effort still opens thinking.
if mode == "auto" || mode == "on" {
if cfg.ExtraFields == nil {
cfg.ExtraFields = make(map[string]any)
}
cfg.ExtraFields["thinking"] = map[string]any{"type": "enabled"}
}
if effort != "" {
if cfg.ExtraFields == nil {
cfg.ExtraFields = make(map[string]any)
}
cfg.ExtraFields["reasoning_effort"] = effortStringForAPI(effort)
}
}
func applyOpenAICompat(cfg *einoopenai.ChatModelConfig, mode, effort string) {
if mode == "auto" && effort == "" {
return
}
e := effort
if mode == "on" && e == "" {
e = "medium"
}
if e == "" {
return
}
if usesExtraFieldsReasoningEffort(e) {
if cfg.ExtraFields == nil {
cfg.ExtraFields = make(map[string]any)
}
cfg.ExtraFields["reasoning_effort"] = effortStringForAPI(e)
return
}
switch e {
case "low":
cfg.ReasoningEffort = einoopenai.ReasoningEffortLevelLow
case "medium":
cfg.ReasoningEffort = einoopenai.ReasoningEffortLevelMedium
case "high":
cfg.ReasoningEffort = einoopenai.ReasoningEffortLevelHigh
}
}
func applyOutputConfigEffort(cfg *einoopenai.ChatModelConfig, mode, effort string) {
if mode == "auto" && effort == "" {
return
}
e := effort
if mode == "on" && e == "" {
e = "high"
}
if e == "" {
return
}
if cfg.ExtraFields == nil {
cfg.ExtraFields = make(map[string]any)
}
cfg.ExtraFields["output_config"] = map[string]any{"effort": effortStringForAPI(e)}
}
func effortStringForAPI(e string) string {
// 原样透传:OpenAI 官方多为 xhigh,部分兼容网关为 max,由配置/对话 effort 选择。
return strings.ToLower(strings.TrimSpace(e))
}
+424
View File
@@ -0,0 +1,424 @@
package reasoning
import (
"context"
"encoding/json"
"io"
"net/http"
"net/http/httptest"
"testing"
"cyberstrike-ai/internal/config"
einoopenai "github.com/cloudwego/eino-ext/components/model/openai"
"github.com/cloudwego/eino/schema"
)
var reasoningPayloadKeysForTest = []string{"thinking", "reasoning_effort", "output_config", "reasoning"}
func assertNoReasoningFields(t *testing.T, cfg *einoopenai.ChatModelConfig) {
t.Helper()
if cfg.ReasoningEffort != "" {
t.Fatalf("expected ReasoningEffort omitted, got %q", cfg.ReasoningEffort)
}
for _, key := range reasoningPayloadKeysForTest {
if _, ok := cfg.ExtraFields[key]; ok {
t.Fatalf("expected %q omitted, got %#v", key, cfg.ExtraFields)
}
}
}
func TestEffortStringForAPI_passthrough(t *testing.T) {
cases := map[string]string{
"max": "max",
"xhigh": "xhigh",
"HIGH": "high",
"Medium": "medium",
}
for in, want := range cases {
if got := effortStringForAPI(in); got != want {
t.Fatalf("%q -> %q, want %q", in, got, want)
}
}
}
func TestNormalizeEffort_maxAndXhigh(t *testing.T) {
if normalizeEffort("xhigh") != "xhigh" {
t.Fatal("xhigh not accepted")
}
if normalizeEffort("max") != "max" {
t.Fatal("max not accepted")
}
}
func TestApplyOpenAICompat_xhighExtraField(t *testing.T) {
cfg := &einoopenai.ChatModelConfig{}
oa := &config.OpenAIConfig{
Reasoning: config.OpenAIReasoningConfig{
Profile: "openai_compat",
Mode: "on",
Effort: "xhigh",
},
}
ApplyToEinoChatModelConfig(cfg, oa, nil)
if cfg.ExtraFields == nil {
t.Fatal("expected ExtraFields")
}
if got, _ := cfg.ExtraFields["reasoning_effort"].(string); got != "xhigh" {
t.Fatalf("reasoning_effort=%q", got)
}
}
func TestAgenticOpenAIExtraFields_openAICompatReasoningEffort(t *testing.T) {
oa := &config.OpenAIConfig{
Reasoning: config.OpenAIReasoningConfig{
Profile: "openai_compat",
Mode: "on",
Effort: "high",
ExtraRequestFields: map[string]interface{}{
"vendor_option": true,
},
},
}
got := AgenticOpenAIExtraFields(oa, nil)
if got["reasoning_effort"] != "high" {
t.Fatalf("reasoning_effort=%#v, want high in %#v", got["reasoning_effort"], got)
}
if got["vendor_option"] != true {
t.Fatalf("vendor option not preserved: %#v", got)
}
}
func TestAgenticOpenAIExtraFields_reasoningOffPreservesUnrelatedFields(t *testing.T) {
oa := &config.OpenAIConfig{
Model: "gpt-4o-mini",
Reasoning: config.OpenAIReasoningConfig{
Profile: "openai_compat",
Mode: "off",
Effort: "high",
ExtraRequestFields: map[string]interface{}{
"reasoning_effort": "high",
"thinking": map[string]any{"type": "enabled"},
"vendor_option": true,
},
},
}
got := AgenticOpenAIExtraFields(oa, nil)
for _, key := range reasoningPayloadKeysForTest {
if _, ok := got[key]; ok {
t.Fatalf("agentic fields unexpectedly contain %q: %#v", key, got)
}
}
if got["vendor_option"] != true {
t.Fatalf("vendor option not preserved: %#v", got)
}
}
func TestAgenticOpenAIPlannerExtraFields_deepseekDisablesThinking(t *testing.T) {
oa := &config.OpenAIConfig{
BaseURL: "https://api.deepseek.com",
Model: "deepseek-chat",
Reasoning: config.OpenAIReasoningConfig{
Profile: "auto",
Mode: "on",
ExtraRequestFields: map[string]interface{}{
"reasoning_effort": "high",
"vendor_option": true,
},
},
}
got := AgenticOpenAIPlannerExtraFields(oa)
if got["reasoning_effort"] != nil {
t.Fatalf("planner should strip reasoning_effort: %#v", got)
}
thinking, ok := got["thinking"].(map[string]any)
if !ok || thinking["type"] != "disabled" {
t.Fatalf("expected deepseek thinking disabled, got %#v", got)
}
if got["vendor_option"] != true {
t.Fatalf("vendor option not preserved: %#v", got)
}
}
func TestAgenticOpenAIPlannerExtraFields_deepseekEndpointWinsOverOpenAIProfile(t *testing.T) {
oa := &config.OpenAIConfig{
BaseURL: "https://api.deepseek.com/v1",
Model: "deepseek-v4-flash",
Reasoning: config.OpenAIReasoningConfig{
Profile: "openai_compat",
Mode: "on",
Effort: "high",
ExtraRequestFields: map[string]interface{}{
"reasoning_effort": "high",
"vendor_option": true,
},
},
}
got := AgenticOpenAIPlannerExtraFields(oa)
if _, ok := got["reasoning_effort"]; ok {
t.Fatalf("planner should strip reasoning_effort: %#v", got)
}
thinking, ok := got["thinking"].(map[string]any)
if !ok || thinking["type"] != "disabled" {
t.Fatalf("expected deepseek thinking disabled despite openai_compat profile, got %#v", got)
}
if got["vendor_option"] != true {
t.Fatalf("vendor option not preserved: %#v", got)
}
}
func TestApplyPlanExecutePlannerModelConfig_stripsReasoningWhenGlobalOn(t *testing.T) {
cfg := &einoopenai.ChatModelConfig{ExtraFields: map[string]any{
"thinking": map[string]any{"type": "enabled"},
"reasoning_effort": "high",
"vendor_option": true,
}}
oa := &config.OpenAIConfig{
BaseURL: "https://antchat.example.com/v1",
Model: "minimax-m3",
Reasoning: config.OpenAIReasoningConfig{
Profile: "openai_compat",
Mode: "on",
Effort: "high",
},
}
ApplyPlanExecutePlannerModelConfig(cfg, oa)
assertNoReasoningFields(t, cfg)
if cfg.ExtraFields["vendor_option"] != true {
t.Fatalf("expected unrelated extra field preserved, got %#v", cfg.ExtraFields)
}
}
func TestApplyPlanExecutePlannerModelConfig_deepseekEndpointWinsOverOpenAIProfile(t *testing.T) {
cfg := &einoopenai.ChatModelConfig{ExtraFields: map[string]any{
"thinking": map[string]any{"type": "enabled"},
"reasoning_effort": "high",
"vendor_option": true,
}}
oa := &config.OpenAIConfig{
BaseURL: "https://api.deepseek.com/v1",
Model: "deepseek-v4-flash",
Reasoning: config.OpenAIReasoningConfig{
Profile: "openai_compat",
Mode: "on",
Effort: "high",
},
}
ApplyPlanExecutePlannerModelConfig(cfg, oa)
if cfg.ReasoningEffort != "" {
t.Fatalf("expected ReasoningEffort omitted, got %q", cfg.ReasoningEffort)
}
if _, ok := cfg.ExtraFields["reasoning_effort"]; ok {
t.Fatalf("expected reasoning_effort omitted, got %#v", cfg.ExtraFields)
}
thinking, ok := cfg.ExtraFields["thinking"].(map[string]any)
if !ok || thinking["type"] != "disabled" {
t.Fatalf("expected deepseek thinking disabled despite openai_compat profile, got %#v", cfg.ExtraFields)
}
if cfg.ExtraFields["vendor_option"] != true {
t.Fatalf("expected unrelated extra field preserved, got %#v", cfg.ExtraFields)
}
}
func TestApplyReasoningOff_omitsAllReasoningFields(t *testing.T) {
cfg := &einoopenai.ChatModelConfig{ExtraFields: map[string]any{
"thinking": map[string]any{"type": "enabled"},
"output_config": map[string]any{"effort": "high"},
}}
oa := &config.OpenAIConfig{
BaseURL: "https://api.openai.com/v1",
Model: "gpt-4o-mini",
Reasoning: config.OpenAIReasoningConfig{
Mode: "off",
Effort: "high",
Profile: "openai_compat",
ExtraRequestFields: map[string]interface{}{
"thinking": map[string]any{"type": "disabled"},
"reasoning": map[string]any{"effort": "high"},
"vendor_option": true,
},
},
}
ApplyToEinoChatModelConfig(cfg, oa, nil)
assertNoReasoningFields(t, cfg)
if cfg.ExtraFields["vendor_option"] != true {
t.Fatalf("expected unrelated extra field preserved, got %#v", cfg.ExtraFields)
}
}
func TestApplyReasoningOff_clientOverrideOmit(t *testing.T) {
cfg := &einoopenai.ChatModelConfig{}
oa := &config.OpenAIConfig{Reasoning: config.OpenAIReasoningConfig{
Mode: "on", Effort: "high", Profile: "openai_compat",
}}
ApplyToEinoChatModelConfig(cfg, oa, &ClientIntent{Mode: "off", Effort: "high"})
assertNoReasoningFields(t, cfg)
}
func TestApplyReasoningOff_deepseekExplicitlyDisablesDefaultThinking(t *testing.T) {
for _, profile := range []string{"deepseek_compat", "auto", "openai_compat"} {
t.Run(profile, func(t *testing.T) {
cfg := &einoopenai.ChatModelConfig{ExtraFields: map[string]any{
"reasoning_effort": "high",
"vendor_option": true,
}}
oa := &config.OpenAIConfig{
BaseURL: "https://api.deepseek.com",
Model: "deepseek-v4-pro",
Reasoning: config.OpenAIReasoningConfig{
Mode: "off", Effort: "high", Profile: profile,
},
}
ApplyToEinoChatModelConfig(cfg, oa, nil)
if cfg.ReasoningEffort != "" {
t.Fatalf("expected ReasoningEffort omitted, got %q", cfg.ReasoningEffort)
}
if _, ok := cfg.ExtraFields["reasoning_effort"]; ok {
t.Fatalf("expected reasoning_effort omitted, got %#v", cfg.ExtraFields)
}
thinking, ok := cfg.ExtraFields["thinking"].(map[string]any)
if !ok || thinking["type"] != "disabled" {
t.Fatalf("expected DeepSeek thinking disabled, got %#v", cfg.ExtraFields)
}
if cfg.ExtraFields["vendor_option"] != true {
t.Fatalf("expected unrelated extra field preserved, got %#v", cfg.ExtraFields)
}
})
}
}
func TestApplyReasoningOff_wirePayloadOmitsThinking(t *testing.T) {
var requestBody map[string]any
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
defer r.Body.Close()
body, err := io.ReadAll(r.Body)
if err != nil {
t.Errorf("read request body: %v", err)
}
if err := json.Unmarshal(body, &requestBody); err != nil {
t.Errorf("decode request body: %v; body=%s", err, body)
}
w.Header().Set("Content-Type", "application/json")
_, _ = io.WriteString(w, `{"id":"chatcmpl-test","object":"chat.completion","created":1,"model":"gpt-4o-mini","choices":[{"index":0,"message":{"role":"assistant","content":"ok"},"finish_reason":"stop"}],"usage":{"prompt_tokens":1,"completion_tokens":1,"total_tokens":2}}`)
}))
defer srv.Close()
cfg := &einoopenai.ChatModelConfig{
APIKey: "test-key",
BaseURL: srv.URL,
Model: "gpt-4o-mini",
}
oa := &config.OpenAIConfig{
BaseURL: "https://api.openai.com/v1",
Model: "gpt-4o-mini",
Reasoning: config.OpenAIReasoningConfig{
Mode: "off", Effort: "high", Profile: "openai_compat",
},
}
ApplyToEinoChatModelConfig(cfg, oa, nil)
model, err := einoopenai.NewChatModel(context.Background(), cfg)
if err != nil {
t.Fatalf("new chat model: %v", err)
}
if _, err := model.Generate(context.Background(), []*schema.Message{schema.UserMessage("hello")}); err != nil {
t.Fatalf("generate: %v", err)
}
for _, key := range reasoningPayloadKeysForTest {
if _, ok := requestBody[key]; ok {
t.Fatalf("wire payload unexpectedly contains %q: %#v", key, requestBody)
}
}
}
func TestApplyOpenAICompat_maxPassthrough(t *testing.T) {
cfg := &einoopenai.ChatModelConfig{}
oa := &config.OpenAIConfig{
Reasoning: config.OpenAIReasoningConfig{
Profile: "openai_compat",
Mode: "on",
Effort: "max",
},
}
ApplyToEinoChatModelConfig(cfg, oa, nil)
got, _ := cfg.ExtraFields["reasoning_effort"].(string)
if got != "max" {
t.Fatalf("max effort wire=%q, want max", got)
}
}
func TestApplyClaude_adaptiveOutputConfigEffort(t *testing.T) {
cfg := &einoopenai.ChatModelConfig{}
oa := &config.OpenAIConfig{
Provider: "claude",
Model: "claude-opus-4-8",
Reasoning: config.OpenAIReasoningConfig{
Mode: "on",
Effort: "xhigh",
},
}
ApplyToEinoChatModelConfig(cfg, oa, nil)
th, ok := cfg.ExtraFields["thinking"].(map[string]any)
if !ok || th["type"] != "adaptive" {
t.Fatalf("thinking=%#v", cfg.ExtraFields["thinking"])
}
oc, ok := cfg.ExtraFields["output_config"].(map[string]any)
if !ok {
t.Fatal("expected output_config")
}
if oc["effort"] != "xhigh" {
t.Fatalf("effort=%v", oc["effort"])
}
}
func TestApplyClaude_sonnet37OfficialBudget(t *testing.T) {
cfg := &einoopenai.ChatModelConfig{}
oa := &config.OpenAIConfig{
Provider: "claude",
Model: "claude-3-7-sonnet-latest",
Reasoning: config.OpenAIReasoningConfig{
Mode: "on",
Effort: "low", // 3.7 has no output_config.effort; effort is not mapped to budget_tokens
},
}
ApplyToEinoChatModelConfig(cfg, oa, nil)
th, ok := cfg.ExtraFields["thinking"].(map[string]any)
if !ok || th["type"] != "enabled" {
t.Fatalf("thinking=%#v", cfg.ExtraFields["thinking"])
}
if th["budget_tokens"] != claudeSonnet37DefaultBudgetTokens {
t.Fatalf("budget_tokens=%v, want official example %d", th["budget_tokens"], claudeSonnet37DefaultBudgetTokens)
}
if _, hasOC := cfg.ExtraFields["output_config"]; hasOC {
t.Fatal("sonnet 3.7 should not set output_config")
}
}
func TestApplyClaude_onWithoutEffortOmitsOutputConfig(t *testing.T) {
cfg := &einoopenai.ChatModelConfig{}
oa := &config.OpenAIConfig{
Provider: "claude",
Model: "claude-sonnet-4-6",
Reasoning: config.OpenAIReasoningConfig{
Mode: "on",
},
}
ApplyToEinoChatModelConfig(cfg, oa, nil)
if _, hasOC := cfg.ExtraFields["output_config"]; hasOC {
t.Fatal("on without explicit effort should omit output_config (API default high)")
}
}
func TestApplyClaude_autoWithoutEffortSkipsOutputConfig(t *testing.T) {
cfg := &einoopenai.ChatModelConfig{}
oa := &config.OpenAIConfig{
Provider: "claude",
Model: "claude-sonnet-4-6",
Reasoning: config.OpenAIReasoningConfig{
Mode: "auto",
},
}
ApplyToEinoChatModelConfig(cfg, oa, nil)
if _, hasOC := cfg.ExtraFields["output_config"]; hasOC {
t.Fatal("auto without effort should omit output_config")
}
}
+48
View File
@@ -0,0 +1,48 @@
package workflow
import (
"context"
"github.com/cloudwego/eino/compose"
)
// compileAgentSubgraph wraps an Agent canvas node as an Eino subgraph (AddGraphNode best practice).
func compileAgentSubgraph(_ context.Context, node graphNode) (compose.AnyGraph, error) {
n := node
prepareID := n.ID + "__agent_prepare"
executeID := n.ID + "__agent_execute"
finalizeID := n.ID + "__agent_finalize"
g := compose.NewGraph[WorkflowNodeOutput, WorkflowNodeOutput]()
_ = g.AddLambdaNode(prepareID, compose.InvokableLambda(func(_ context.Context, input WorkflowNodeOutput) (WorkflowNodeOutput, error) {
if input == nil {
input = WorkflowNodeOutput{}
}
input["agent_subgraph_stage"] = "prepare"
input["agent_node_id"] = n.ID
return input, nil
}))
_ = g.AddLambdaNode(executeID, compose.InvokableLambda(func(runCtx context.Context, _ WorkflowNodeOutput) (WorkflowNodeOutput, error) {
return runWorkflowNodeLambda(runCtx, n)
}))
_ = g.AddLambdaNode(finalizeID, compose.InvokableLambda(func(_ context.Context, output WorkflowNodeOutput) (WorkflowNodeOutput, error) {
if output == nil {
output = WorkflowNodeOutput{}
}
output["agent_subgraph_stage"] = "finalize"
output["agent_node_id"] = n.ID
return output, nil
}))
if err := g.AddEdge(compose.START, prepareID); err != nil {
return nil, err
}
if err := g.AddEdge(prepareID, executeID); err != nil {
return nil, err
}
if err := g.AddEdge(executeID, finalizeID); err != nil {
return nil, err
}
if err := g.AddEdge(finalizeID, compose.END); err != nil {
return nil, err
}
return g, nil
}
+153
View File
@@ -0,0 +1,153 @@
package workflow
import (
"encoding/json"
"fmt"
"strings"
)
// FieldBinding selects a value from workflow state (replaces {{...}} templates).
type FieldBinding struct {
From string `json:"from"` // inputs | previous | <nodeId>
Field string `json:"field"` // e.g. output, message
}
func parseFieldBinding(cfg map[string]any, keys ...string) (FieldBinding, bool) {
for _, key := range keys {
if cfg == nil {
continue
}
raw, ok := cfg[key]
if !ok || raw == nil {
continue
}
switch v := raw.(type) {
case map[string]any:
return FieldBinding{
From: strings.TrimSpace(fmt.Sprint(v["from"])),
Field: strings.TrimSpace(fmt.Sprint(v["field"])),
}, true
case string:
s := strings.TrimSpace(v)
if s == "" {
continue
}
var b FieldBinding
if err := json.Unmarshal([]byte(s), &b); err == nil && (b.From != "" || b.Field != "") {
return b, true
}
}
}
return FieldBinding{}, false
}
func defaultBinding(from, field string) FieldBinding {
return FieldBinding{From: from, Field: field}
}
func resolveBinding(b FieldBinding, state *WorkflowLocalState) any {
from := strings.TrimSpace(b.From)
field := strings.TrimSpace(b.Field)
if field == "" {
field = "output"
}
if from == "" || from == "previous" || from == "prev" {
if strings.HasPrefix(field, "$") || strings.HasPrefix(field, ".") {
return evalJSONPathValue(state.LastOutput, field)
}
if field == "output" && state.LastOutput != nil {
return state.LastOutput["output"]
}
return valueFromPath("previous."+field, state)
}
if from == "inputs" || from == "input" {
if strings.HasPrefix(field, "$") || strings.HasPrefix(field, ".") {
return evalJSONPathValue(state.Inputs, field)
}
if field == "" {
return state.Inputs
}
return valueFromPath("inputs."+field, state)
}
if from == "outputs" {
if strings.HasPrefix(field, "$") || strings.HasPrefix(field, ".") {
return evalJSONPathValue(state.Outputs, field)
}
return valueFromPath("outputs."+field, state)
}
if strings.HasPrefix(field, "$") || strings.HasPrefix(field, ".") {
return evalJSONPathValue(valueFromPath(from, state), field)
}
return valueFromPath(from+"."+field, state)
}
func resolveBindingString(b FieldBinding, state *WorkflowLocalState) string {
return strings.TrimSpace(fmt.Sprint(resolveBinding(b, state)))
}
func resolveNodeInputBinding(cfg map[string]any, state *WorkflowLocalState) string {
if b, ok := parseFieldBinding(cfg, "input_binding"); ok {
return resolveBindingString(b, state)
}
// legacy template field removed — default previous.output
return resolveBindingString(defaultBinding("previous", "output"), state)
}
func resolveOutputSourceBinding(cfg map[string]any, state *WorkflowLocalState) any {
if b, ok := parseFieldBinding(cfg, "source_binding"); ok {
return resolveBinding(b, state)
}
return resolveBinding(defaultBinding("previous", "output"), state)
}
func resolveHITLPromptBinding(cfg map[string]any, state *WorkflowLocalState) string {
if b, ok := parseFieldBinding(cfg, "prompt_binding"); ok {
return resolveBindingString(b, state)
}
if s := cfgString(cfg, "prompt"); s != "" {
return s
}
return resolveBindingString(defaultBinding("previous", "output"), state)
}
func toolArgumentBindings(cfg map[string]any) map[string]FieldBinding {
raw, ok := cfg["argument_bindings"].(map[string]any)
if !ok || len(raw) == 0 {
return nil
}
out := make(map[string]FieldBinding, len(raw))
for argName, v := range raw {
m, ok := v.(map[string]any)
if !ok {
continue
}
out[argName] = FieldBinding{
From: strings.TrimSpace(fmt.Sprint(m["from"])),
Field: strings.TrimSpace(fmt.Sprint(m["field"])),
}
}
return out
}
func resolveToolArguments(cfg map[string]any, state *WorkflowLocalState) (map[string]interface{}, error) {
bindings := toolArgumentBindings(cfg)
if len(bindings) > 0 {
args := make(map[string]interface{}, len(bindings))
for k, b := range bindings {
args[k] = resolveBinding(b, state)
}
return args, nil
}
raw := cfgString(cfg, "arguments")
if raw == "" {
return map[string]interface{}{}, nil
}
var args map[string]interface{}
if err := json.Unmarshal([]byte(raw), &args); err != nil {
return nil, err
}
if args == nil {
args = map[string]interface{}{}
}
return args, nil
}
+69
View File
@@ -0,0 +1,69 @@
package workflow
import (
"context"
"fmt"
"os"
"path/filepath"
"strings"
"sync"
)
// fileCheckPointStore persists Eino workflow checkpoints on disk (per run id).
type fileCheckPointStore struct {
dir string
mu sync.RWMutex
}
func newFileCheckPointStore(dir string) (*fileCheckPointStore, error) {
dir = strings.TrimSpace(dir)
if dir == "" {
dir = filepath.Join("data", "workflow-checkpoints")
}
if err := os.MkdirAll(dir, 0o755); err != nil {
return nil, fmt.Errorf("create workflow checkpoint dir: %w", err)
}
return &fileCheckPointStore{dir: dir}, nil
}
func (s *fileCheckPointStore) path(id string) (string, error) {
id = strings.TrimSpace(id)
if id == "" {
return "", fmt.Errorf("checkpoint id is empty")
}
if strings.Contains(id, "..") || strings.ContainsAny(id, `/\`) {
return "", fmt.Errorf("invalid checkpoint id")
}
return filepath.Join(s.dir, id+".ckpt"), nil
}
func (s *fileCheckPointStore) Get(_ context.Context, checkPointID string) ([]byte, bool, error) {
s.mu.RLock()
defer s.mu.RUnlock()
p, err := s.path(checkPointID)
if err != nil {
return nil, false, err
}
data, err := os.ReadFile(p)
if err != nil {
if os.IsNotExist(err) {
return nil, false, nil
}
return nil, false, err
}
return data, true, nil
}
func (s *fileCheckPointStore) Set(_ context.Context, checkPointID string, checkPoint []byte) error {
s.mu.Lock()
defer s.mu.Unlock()
p, err := s.path(checkPointID)
if err != nil {
return err
}
tmp := p + ".tmp"
if err := os.WriteFile(tmp, checkPoint, 0o600); err != nil {
return err
}
return os.Rename(tmp, p)
}
+782
View File
@@ -0,0 +1,782 @@
package workflow
import (
"context"
"encoding/json"
"fmt"
"hash/fnv"
"regexp"
"strings"
"time"
"unicode/utf8"
"cyberstrike-ai/internal/config"
"cyberstrike-ai/internal/openai"
"go.uber.org/zap"
)
type DraftTool struct {
Key string `json:"key"`
Name string `json:"name,omitempty"`
Enabled bool `json:"enabled"`
}
type DraftOptions struct {
IncludeObjective bool `json:"include_objective"`
AllowSchedule bool `json:"allow_schedule"`
AllowHighRisk bool `json:"allow_high_risk"`
}
type DraftRequest struct {
Prompt string `json:"prompt"`
Options DraftOptions `json:"options"`
AvailableTools []DraftTool `json:"available_tools,omitempty"`
}
type DraftMeta struct {
ID string `json:"id"`
Name string `json:"name"`
Description string `json:"description"`
Enabled bool `json:"enabled"`
}
type DraftCapability struct {
Label string `json:"label"`
ToolName string `json:"tool_name,omitempty"`
ToolCandidates []string `json:"tool_candidates,omitempty"`
}
type DraftAudit struct {
Savable bool `json:"savable"`
Validation []string `json:"validation,omitempty"`
MissingFields []string `json:"missing_fields,omitempty"`
RiskWarnings []string `json:"risk_warnings,omitempty"`
Assumptions []string `json:"assumptions,omitempty"`
HighRisk bool `json:"high_risk"`
NeedsHITL bool `json:"needs_hitl"`
}
type DraftResult struct {
Graph *graphDef `json:"graph"`
Meta DraftMeta `json:"meta"`
Generator string `json:"generator"`
Audit DraftAudit `json:"audit"`
Capabilities []DraftCapability `json:"capabilities,omitempty"`
Stats map[string]int `json:"stats"`
}
type llmDraftEnvelope struct {
Graph graphDef `json:"graph"`
Meta DraftMeta `json:"meta"`
Capabilities []DraftCapability `json:"capabilities,omitempty"`
Audit DraftAudit `json:"audit,omitempty"`
}
type draftToolHint struct {
Label string
Keywords []string
Tools []string
}
var draftToolHints = []draftToolHint{
{Label: "子域名发现", Keywords: []string{"子域名", "subdomain", "subfinder", "amass"}, Tools: []string{"subfinder", "amass"}},
{Label: "端口扫描", Keywords: []string{"端口", "port", "nmap", "rustscan", "masscan"}, Tools: []string{"nmap", "rustscan", "masscan"}},
{Label: "漏洞扫描", Keywords: []string{"漏洞", "vuln", "漏洞扫描", "nuclei", "nikto", "zap"}, Tools: []string{"nuclei", "nikto", "zap"}},
{Label: "暴露面探测", Keywords: []string{"目录", "路径", "暴露页面", "dir", "ffuf", "gobuster", "feroxbuster"}, Tools: []string{"ffuf", "gobuster", "feroxbuster", "dirsearch"}},
{Label: "证书与域名线索收集", Keywords: []string{"证书", "certificate", "crt"}, Tools: []string{"subfinder"}},
{Label: "云配置审计", Keywords: []string{"云", "cloud", "配置审计", "prowler", "scout"}, Tools: []string{"prowler", "scout-suite"}},
{Label: "容器安全检查", Keywords: []string{"容器", "镜像", "k8s", "kubernetes", "trivy", "kube"}, Tools: []string{"trivy", "kube-bench", "kube-hunter"}},
{Label: "威胁情报收集", Keywords: []string{"情报", "威胁情报", "threat", "ioc", "virustotal", "shodan", "fofa"}, Tools: []string{"virustotal_search", "shodan_search", "fofa_search"}},
}
var highRiskDraftRE = regexp.MustCompile(`(?i)(隔离|封禁|加固|修复|执行|命令|脚本|删除|清理|阻断|封锁|攻击|利用|getshell|shell|payload|exploit|isolate|block|execute|script|delete|exploit|payload)`)
func GenerateDraftFromNaturalLanguage(ctx context.Context, req DraftRequest) (*DraftResult, error) {
prompt := strings.TrimSpace(req.Prompt)
if prompt == "" {
return nil, fmt.Errorf("工作流需求不能为空")
}
capabilities := detectDraftCapabilities(prompt, req.AvailableTools)
wantsApproval := containsAnyFold(prompt, "审批", "确认", "审核", "负责人", "人工", "review", "approve", "approval", "human")
wantsReport := containsAnyFold(prompt, "报告", "汇总", "输出", "通知", "任务", "工单", "report", "summary", "notify", "ticket")
wantsCondition := containsAnyFold(prompt, "如果", "发现", "存在", "高危", "新增", "失败", "通过", "否则", "if", "when", "high", "critical", "new", "fail")
highRisk := highRiskDraftRE.MatchString(prompt)
builder := &draftGraphBuilder{x: 120, y: 150}
assumptions := make([]string, 0)
riskWarnings := make([]string, 0)
missingFields := make([]string, 0)
start := builder.add("start", "开始", map[string]any{"input_keys": "message, conversationId, projectId, target"}, 0)
previous := start
for _, capability := range capabilities {
hasTool := strings.TrimSpace(capability.ToolName) != ""
var id string
if hasTool {
id = builder.add("tool", capability.Label, map[string]any{
"tool_name": capability.ToolName,
"arguments": `{"target":"{{inputs.target}}","message":"{{inputs.message}}"}`,
"timeout_seconds": "120",
"join_strategy": "all_merge",
}, 0)
} else {
id = builder.add("agent", capability.Label, map[string]any{
"agent_mode": "eino_single",
"input_binding": map[string]any{"from": "previous", "field": "output"},
"instruction": capability.Label + "。根据用户需求执行安全流程步骤,并输出结构化结果:" + prompt,
"output_key": "agent_result",
"join_strategy": "all_merge",
"missing_tool_candidates": strings.Join(capability.ToolCandidates, ", "),
}, 0)
if len(capability.ToolCandidates) > 0 {
assumptions = append(assumptions, capability.Label+" 未匹配到已启用工具,已生成 Agent 草稿节点。")
missingFields = append(missingFields, capability.Label+": 选择或启用对应 MCP 工具")
}
}
builder.connect(previous, id, "", nil)
previous = id
}
openConditionID := ""
if wantsCondition {
expr := `{{previous.output}} != ""`
label := "是否满足触发条件"
if highRisk {
expr = `{{previous.output}} contains "高危"`
label = "是否需要高风险处置"
}
condition := builder.add("condition", label, map[string]any{"expression": expr, "join_strategy": "all_merge"}, 0)
builder.connect(previous, condition, "", nil)
openConditionID = condition
report := builder.add("output", draftOutputLabel(wantsReport), map[string]any{
"output_key": "result",
"source_binding": map[string]any{"from": "previous", "field": "output"},
"static_value": "",
"join_strategy": "all_merge",
}, 130)
builder.connect(condition, report, "否", map[string]any{"condition": `{{previous.matched}} == "false"`, "branch": "false"})
previous = condition
}
insertedHITL := false
if highRisk {
if !req.Options.AllowHighRisk || wantsApproval {
approval := builder.add("hitl", "人工审批", map[string]any{
"prompt": "请确认是否允许继续执行高风险处置:" + prompt,
"prompt_binding": map[string]any{"from": "previous", "field": "output"},
"reviewer": "human",
"join_strategy": "all_merge",
"risk_level": "high",
}, 0)
builder.connect(previous, approval, branchLabel(previous, openConditionID), branchConfig(previous, openConditionID, true))
if previous == openConditionID {
openConditionID = ""
}
previous = approval
insertedHITL = true
}
action := builder.add("agent", "执行受控处置", map[string]any{
"agent_mode": "eino_single",
"input_binding": map[string]any{"from": "previous", "field": "output"},
"instruction": "仅在授权范围内生成处置步骤草稿;实际执行前必须由人工确认。用户需求:" + prompt,
"output_key": "remediation_plan",
"join_strategy": "all_merge",
"risk_level": "high",
"requires_human_confirmation": "true",
}, 0)
builder.connect(previous, action, branchLabel(previous, openConditionID), branchConfig(previous, openConditionID, true))
if previous == openConditionID {
openConditionID = ""
}
previous = action
if insertedHITL {
riskWarnings = append(riskWarnings, "检测到高风险动作,已加入人工审批与 requires_human_confirmation 标记。")
} else {
riskWarnings = append(riskWarnings, "检测到高风险动作,已保留为草稿并添加 requires_human_confirmation 标记。")
}
} else if wantsApproval {
approval := builder.add("hitl", "人工审批", map[string]any{
"prompt": "请审核工作流阶段结果:" + prompt,
"prompt_binding": map[string]any{"from": "previous", "field": "output"},
"reviewer": "human",
"join_strategy": "all_merge",
}, 0)
builder.connect(previous, approval, "", nil)
previous = approval
insertedHITL = true
}
output := builder.add("output", draftOutputLabel(wantsReport), map[string]any{
"output_key": "result",
"source_binding": map[string]any{"from": "previous", "field": "output"},
"static_value": "",
"join_strategy": "all_merge",
}, 0)
builder.connect(previous, output, branchLabel(previous, openConditionID), branchConfig(previous, openConditionID, true))
graph := &graphDef{Nodes: builder.nodes, Edges: builder.edges, Config: map[string]any{
"schema_version": 1,
"generated_by": "natural_language",
"source_prompt": prompt,
}}
if req.Options.IncludeObjective {
graph.Config["objective"] = prompt
}
if req.Options.AllowSchedule && containsAnyFold(prompt, "每天", "每周", "定时", "周期", "持续", "daily", "weekly", "schedule", "monitor") {
if containsAnyFold(prompt, "每天", "daily") {
graph.Config["trigger_suggestion"] = "daily"
} else {
graph.Config["trigger_suggestion"] = "scheduled"
}
assumptions = append(assumptions, "已记录定时触发建议;保存后仍需在触发器或角色绑定处配置。")
}
raw, _ := json.Marshal(graph)
validation := make([]string, 0)
if err := ValidateGraphJSON(ctx, string(raw)); err != nil {
validation = append(validation, err.Error())
}
return &DraftResult{
Graph: graph,
Meta: DraftMeta{ID: draftSlug(prompt), Name: draftName(prompt), Description: prompt, Enabled: true},
Generator: "deterministic",
Audit: DraftAudit{
Savable: len(validation) == 0,
Validation: validation,
MissingFields: missingFields,
RiskWarnings: riskWarnings,
Assumptions: assumptions,
HighRisk: highRisk,
NeedsHITL: insertedHITL,
},
Capabilities: capabilities,
Stats: map[string]int{"nodes": len(graph.Nodes), "edges": len(graph.Edges)},
}, nil
}
func GenerateDraftFromLLM(ctx context.Context, req DraftRequest, oa config.OpenAIConfig, logger *zap.Logger) (*DraftResult, error) {
prompt := strings.TrimSpace(req.Prompt)
if prompt == "" {
return nil, fmt.Errorf("工作流需求不能为空")
}
if strings.TrimSpace(oa.APIKey) == "" || strings.TrimSpace(oa.Model) == "" {
return nil, fmt.Errorf("AI 通道未配置 api_key 或 model")
}
if logger == nil {
logger = zap.NewNop()
}
callCtx, cancel := context.WithTimeout(ctx, 90*time.Second)
defer cancel()
toolJSON, _ := json.Marshal(req.AvailableTools)
systemPrompt := `你是 CyberStrikeAI 的工作流编排助手。你必须把用户的一句话需求转换为可保存的工作流草稿 JSON。
只返回 JSON 对象,不要 Markdown,不要解释。JSON 必须符合:
{
"meta": {"id":"kebab-case-id","name":"短名称","description":"用户需求","enabled":true},
"graph": {
"nodes": [{"id":"start-1","type":"start","label":"显示名","position":{"x":120,"y":150},"config":{}}],
"edges": [{"id":"edge-1","source":"start-1","target":"node-2","label":"","config":{}}],
"config": {"schema_version":1,"generated_by":"llm","source_prompt":"用户原文"}
},
"capabilities": [{"label":"能力名","tool_name":"已匹配工具名","tool_candidates":["候选工具"]}],
"audit": {"assumptions":[],"missing_fields":[],"risk_warnings":[]}
}
硬性规则:
- 只能输出一个合法 JSON object;不要输出 JSON Schema、注释、解释文字、Markdown 代码块或多余前后缀。
- 不要在 JSON 字符串值中使用竖线枚举写法;type 字段一次只能填写一个节点类型字符串。
- 至少 1 个 start 和 1 个 outputoutput/end 不能有出边。
- 节点 type 只能从这些字符串中选择:start、tool、agent、condition、hitl、output、end。
- 每个 agent、tool、output 节点都必须配置唯一的 output_keyoutput 节点默认使用 result。
- agent 节点必须配置 instruction 或 input_binding;默认 input_binding 为 {"from":"previous","field":"output"}。
- output 节点必须配置 source_binding 或 static_value;默认 source_binding 为 {"from":"previous","field":"output"}。
- tool 节点必须配置 tool_name、arguments、timeout_secondsarguments 必须是合法 JSON 字符串。
- 所有非 start 且可能有多个上游的节点必须配置 join_strategy:"all_merge"。
- condition 最多 2 条出边,必须用 branch true/false,并用 label 是/否。
- tool 节点只有在 available_tools 中存在启用工具时才使用,否则用 agent 节点并在 audit.missing_fields 写明缺失工具。
- 高风险动作(执行脚本、隔离、封禁、删除、利用、payload、命令执行等)必须加入 hitl 审批,或在高风险节点 config 中标记 requires_human_confirmation:"true"、risk_level:"high"。
- 不要生成会真实执行攻击的参数;工具参数使用 {{inputs.target}}{{inputs.message}} 占位。
- 所有节点 config 加 generated_by:"llm" 和 needs_review:"true"。`
userPrompt := fmt.Sprintf("用户需求:%s\n\n选项:%+v\n\n可用工具 JSON%s", prompt, req.Options, string(toolJSON))
requestBody := map[string]interface{}{
"model": strings.TrimSpace(oa.Model),
"messages": []map[string]interface{}{
{"role": "system", "content": systemPrompt},
{"role": "user", "content": userPrompt},
},
"temperature": 0,
"max_completion_tokens": 4096,
"response_format": map[string]interface{}{"type": "json_object"},
"thinking": map[string]interface{}{"type": "disabled"},
}
var apiResponse struct {
Choices []struct {
Message struct {
Content string `json:"content"`
ReasoningContent string `json:"reasoning_content"`
} `json:"message"`
} `json:"choices"`
}
client := openai.NewClient(&oa, nil, logger)
if err := client.ChatCompletion(callCtx, requestBody, &apiResponse); err != nil {
return nil, fmt.Errorf("调用大模型失败: %w", err)
}
if len(apiResponse.Choices) == 0 {
return nil, fmt.Errorf("大模型未返回候选结果")
}
raw := strings.TrimSpace(apiResponse.Choices[0].Message.Content)
if raw == "" {
raw = strings.TrimSpace(apiResponse.Choices[0].Message.ReasoningContent)
}
env, err := parseLLMDraftEnvelope(raw)
if err != nil {
return nil, err
}
result := normalizeLLMDraft(prompt, req, env)
graphRaw, _ := json.Marshal(result.Graph)
validation := make([]string, 0)
if err := ValidateGraphJSON(ctx, string(graphRaw)); err != nil {
validation = append(validation, err.Error())
}
result.Audit.Validation = validation
result.Audit.Savable = len(validation) == 0
if !result.Audit.Savable {
return nil, fmt.Errorf("大模型生成的工作流未通过校验: %s", strings.Join(validation, ""))
}
return result, nil
}
type draftGraphBuilder struct {
nodes []graphNode
edges []graphEdge
x float64
y float64
nodeSeq int
edgeSeq int
}
func (b *draftGraphBuilder) add(nodeType, label string, config map[string]any, yOffset float64) string {
b.nodeSeq++
id := fmt.Sprintf("%s-%d", nodeType, b.nodeSeq)
if config == nil {
config = make(map[string]any)
}
config["generated_by"] = "natural_language"
config["needs_review"] = "true"
b.nodes = append(b.nodes, graphNode{
ID: id,
Type: nodeType,
Label: label,
Position: graphPosition{X: b.x, Y: b.y + yOffset},
Config: config,
})
b.x += 210
return id
}
func (b *draftGraphBuilder) connect(source, target, label string, config map[string]any) {
b.edgeSeq++
if config == nil {
config = make(map[string]any)
}
b.edges = append(b.edges, graphEdge{ID: fmt.Sprintf("edge-ai-%d", b.edgeSeq), Source: source, Target: target, Label: label, Config: config})
}
func parseLLMDraftEnvelope(raw string) (llmDraftEnvelope, error) {
var lastErr error
for _, candidate := range jsonObjectCandidates(raw) {
var env llmDraftEnvelope
if err := json.Unmarshal([]byte(candidate), &env); err == nil {
if len(env.Graph.Nodes) == 0 {
lastErr = fmt.Errorf("大模型 JSON 缺少 graph.nodes")
continue
}
return env, nil
} else {
lastErr = err
}
}
if lastErr == nil {
lastErr = fmt.Errorf("大模型响应为空")
}
return llmDraftEnvelope{}, fmt.Errorf("解析大模型工作流 JSON 失败: %w", lastErr)
}
func jsonObjectCandidates(raw string) []string {
s := strings.TrimSpace(raw)
s = strings.TrimPrefix(s, "```json")
s = strings.TrimPrefix(s, "```")
s = strings.TrimSuffix(s, "```")
s = strings.TrimSpace(s)
candidates := []string{s}
if start := strings.Index(s, "{"); start >= 0 {
if end := strings.LastIndex(s, "}"); end > start {
candidates = append(candidates, s[start:end+1])
}
}
return candidates
}
func normalizeLLMDraft(prompt string, req DraftRequest, env llmDraftEnvelope) *DraftResult {
g := env.Graph
if g.Config == nil {
g.Config = make(map[string]any)
}
g.Config["schema_version"] = 1
g.Config["generated_by"] = "llm"
g.Config["source_prompt"] = prompt
if req.Options.IncludeObjective {
g.Config["objective"] = prompt
}
enabledTools := enabledDraftToolNames(req.AvailableTools)
usedOutputKeys := make(map[string]bool)
nodeTypes := make(map[string]string, len(g.Nodes))
for i := range g.Nodes {
if strings.TrimSpace(g.Nodes[i].ID) == "" {
g.Nodes[i].ID = fmt.Sprintf("%s-%d", firstNonEmpty(g.Nodes[i].Type, "node"), i+1)
}
if strings.TrimSpace(g.Nodes[i].Type) == "" {
g.Nodes[i].Type = "agent"
}
if strings.TrimSpace(g.Nodes[i].Label) == "" {
g.Nodes[i].Label = displayNodeType(g.Nodes[i].Type)
}
if g.Nodes[i].Position.X == 0 && g.Nodes[i].Position.Y == 0 {
g.Nodes[i].Position = graphPosition{X: 120 + float64(i)*210, Y: 150}
}
if g.Nodes[i].Config == nil {
g.Nodes[i].Config = make(map[string]any)
}
g.Nodes[i].Config["generated_by"] = "llm"
g.Nodes[i].Config["needs_review"] = "true"
normalizeLLMNodeConfig(prompt, &g.Nodes[i], enabledTools, usedOutputKeys)
nodeTypes[g.Nodes[i].ID] = strings.ToLower(strings.TrimSpace(g.Nodes[i].Type))
}
conditionBranchCounts := make(map[string]int)
for i := range g.Edges {
if strings.TrimSpace(g.Edges[i].ID) == "" {
g.Edges[i].ID = fmt.Sprintf("edge-llm-%d", i+1)
}
if g.Edges[i].Config == nil {
g.Edges[i].Config = make(map[string]any)
}
normalizeLLMEdgeConfig(&g.Edges[i], nodeTypes, conditionBranchCounts)
}
audit := env.Audit
highRisk := highRiskDraftRE.MatchString(prompt) || graphHasHighRisk(g)
audit.HighRisk = highRisk
audit.NeedsHITL = graphHasNodeType(g, "hitl")
if highRisk && !audit.NeedsHITL && !graphHasConfirmation(g) {
audit.RiskWarnings = append(audit.RiskWarnings, "大模型生成包含高风险语义,请补充人工审批或确认标记后再运行。")
}
if len(audit.RiskWarnings) == 0 && highRisk {
audit.RiskWarnings = append(audit.RiskWarnings, "检测到高风险动作,已标记为需要重点审计。")
}
meta := env.Meta
if strings.TrimSpace(meta.Description) == "" {
meta.Description = prompt
}
if strings.TrimSpace(meta.Name) == "" {
meta.Name = draftName(prompt)
}
if strings.TrimSpace(meta.ID) == "" {
meta.ID = draftSlug(prompt)
}
meta.Enabled = true
return &DraftResult{
Graph: &g,
Meta: meta,
Generator: "llm",
Audit: audit,
Capabilities: env.Capabilities,
Stats: map[string]int{"nodes": len(g.Nodes), "edges": len(g.Edges)},
}
}
func normalizeLLMEdgeConfig(edge *graphEdge, nodeTypes map[string]string, conditionBranchCounts map[string]int) {
if nodeTypes[strings.TrimSpace(edge.Source)] != "condition" {
return
}
if conditionBranchHint(*edge) != "" {
return
}
conditionBranchCounts[edge.Source]++
branch := "true"
label := "是"
if conditionBranchCounts[edge.Source] > 1 {
branch = "false"
label = "否"
}
edge.Label = label
edge.Config["branch"] = branch
}
func normalizeLLMNodeConfig(prompt string, node *graphNode, enabledTools map[string]bool, usedOutputKeys map[string]bool) {
nodeType := strings.ToLower(strings.TrimSpace(node.Type))
switch nodeType {
case "start":
if cfgString(node.Config, "input_keys") == "" {
node.Config["input_keys"] = "message, conversationId, projectId, target"
}
case "tool":
toolName := cfgString(node.Config, "tool_name")
if toolName == "" || !enabledTools[strings.ToLower(toolName)] {
node.Type = "agent"
node.Config["missing_tool_name"] = toolName
normalizeAgentDraftConfig(prompt, node, usedOutputKeys)
return
}
if cfgString(node.Config, "arguments") == "" {
node.Config["arguments"] = `{"target":"{{inputs.target}}","message":"{{inputs.message}}"}`
}
if cfgString(node.Config, "timeout_seconds") == "" {
node.Config["timeout_seconds"] = "120"
}
ensureNodeOutputKey(node, usedOutputKeys, draftOutputKeyBase(node, "tool_result"))
ensureJoinStrategy(node)
case "agent":
normalizeAgentDraftConfig(prompt, node, usedOutputKeys)
case "condition":
if cfgString(node.Config, "expression") == "" {
node.Config["expression"] = `{{previous.output}} != ""`
}
ensureJoinStrategy(node)
case "hitl":
if cfgString(node.Config, "prompt") == "" {
node.Config["prompt"] = "请审核工作流阶段结果:" + prompt
}
if cfgString(node.Config, "reviewer") == "" {
node.Config["reviewer"] = "human"
}
ensureJoinStrategy(node)
case "output":
ensureNodeOutputKey(node, usedOutputKeys, "result")
if cfgString(node.Config, "static_value") == "" {
if _, ok := parseFieldBinding(node.Config, "source_binding"); !ok {
node.Config["source_binding"] = map[string]any{"from": "previous", "field": "output"}
}
}
ensureJoinStrategy(node)
case "end":
ensureJoinStrategy(node)
}
}
func normalizeAgentDraftConfig(prompt string, node *graphNode, usedOutputKeys map[string]bool) {
if cfgString(node.Config, "agent_mode") == "" {
node.Config["agent_mode"] = "eino_single"
}
if cfgString(node.Config, "instruction") == "" {
node.Config["instruction"] = node.Label + "。根据用户需求执行安全流程步骤,并输出结构化结果:" + prompt
}
if _, ok := parseFieldBinding(node.Config, "input_binding"); !ok {
node.Config["input_binding"] = map[string]any{"from": "previous", "field": "output"}
}
ensureNodeOutputKey(node, usedOutputKeys, draftOutputKeyBase(node, "agent_result"))
ensureJoinStrategy(node)
}
func ensureJoinStrategy(node *graphNode) {
if cfgString(node.Config, "join_strategy") == "" {
node.Config["join_strategy"] = "all_merge"
}
}
func ensureNodeOutputKey(node *graphNode, used map[string]bool, fallback string) {
current := sanitizeOutputKey(cfgString(node.Config, "output_key"))
if current == "" {
current = sanitizeOutputKey(fallback)
}
if current == "" {
current = "result"
}
base := current
for i := 2; used[current]; i++ {
current = fmt.Sprintf("%s_%d", base, i)
}
node.Config["output_key"] = current
used[current] = true
}
func draftOutputKeyBase(node *graphNode, fallback string) string {
if name := cfgString(node.Config, "tool_name"); name != "" {
return name + "_result"
}
if node.ID != "" {
return node.ID + "_result"
}
return fallback
}
func sanitizeOutputKey(value string) string {
value = strings.ToLower(strings.TrimSpace(value))
var b strings.Builder
lastUnderscore := false
for _, r := range value {
if (r >= 'a' && r <= 'z') || (r >= '0' && r <= '9') {
b.WriteRune(r)
lastUnderscore = false
continue
}
if b.Len() > 0 && !lastUnderscore {
b.WriteByte('_')
lastUnderscore = true
}
}
return strings.Trim(b.String(), "_")
}
func enabledDraftToolNames(tools []DraftTool) map[string]bool {
names := make(map[string]bool, len(tools)*2)
for _, tool := range tools {
if !tool.Enabled {
continue
}
if key := strings.ToLower(strings.TrimSpace(tool.Key)); key != "" {
names[key] = true
}
if name := strings.ToLower(strings.TrimSpace(tool.Name)); name != "" {
names[name] = true
}
}
return names
}
func graphHasNodeType(g graphDef, nodeType string) bool {
for _, node := range g.Nodes {
if strings.EqualFold(node.Type, nodeType) {
return true
}
}
return false
}
func graphHasConfirmation(g graphDef) bool {
for _, node := range g.Nodes {
if cfgString(node.Config, "requires_human_confirmation") == "true" {
return true
}
}
return false
}
func graphHasHighRisk(g graphDef) bool {
for _, node := range g.Nodes {
if cfgString(node.Config, "risk_level") == "high" || cfgString(node.Config, "requires_human_confirmation") == "true" {
return true
}
if highRiskDraftRE.MatchString(node.Label) || highRiskDraftRE.MatchString(cfgString(node.Config, "instruction")) {
return true
}
}
return false
}
func detectDraftCapabilities(prompt string, tools []DraftTool) []DraftCapability {
capabilities := make([]DraftCapability, 0)
for _, hint := range draftToolHints {
if containsAnyFold(prompt, hint.Keywords...) {
capabilities = append(capabilities, DraftCapability{
Label: hint.Label,
ToolName: matchDraftTool(hint.Tools, tools),
ToolCandidates: append([]string(nil), hint.Tools...),
})
}
}
if len(capabilities) == 0 {
capabilities = append(capabilities, DraftCapability{Label: "节点能力", ToolCandidates: nil})
}
return capabilities
}
func matchDraftTool(candidates []string, tools []DraftTool) string {
if len(candidates) == 0 || len(tools) == 0 {
return ""
}
for _, enabledOnly := range []bool{true, false} {
for _, candidate := range candidates {
candidate = strings.ToLower(strings.TrimSpace(candidate))
for _, tool := range tools {
if enabledOnly && !tool.Enabled {
continue
}
key := strings.ToLower(strings.TrimSpace(firstNonEmpty(tool.Key, tool.Name)))
if key != "" && strings.Contains(key, candidate) {
return firstNonEmpty(tool.Key, tool.Name)
}
}
}
}
return ""
}
func containsAnyFold(text string, needles ...string) bool {
lower := strings.ToLower(text)
for _, needle := range needles {
if strings.Contains(lower, strings.ToLower(needle)) {
return true
}
}
return false
}
func draftOutputLabel(wantsReport bool) string {
if wantsReport {
return "输出报告"
}
return "输出"
}
func branchLabel(source, conditionID string) string {
if source == conditionID && conditionID != "" {
return "是"
}
return ""
}
func branchConfig(source, conditionID string, yes bool) map[string]any {
if source != conditionID || conditionID == "" {
return nil
}
if yes {
return map[string]any{"condition": `{{previous.matched}} == "true"`, "branch": "true"}
}
return map[string]any{"condition": `{{previous.matched}} == "false"`, "branch": "false"}
}
func draftName(prompt string) string {
runes := []rune(strings.TrimSpace(prompt))
if len(runes) > 22 {
return string(runes[:22]) + "..."
}
return string(runes)
}
func draftSlug(prompt string) string {
lower := strings.ToLower(strings.TrimSpace(prompt))
var b strings.Builder
lastDash := false
for _, r := range lower {
if r >= 'a' && r <= 'z' || r >= '0' && r <= '9' {
b.WriteRune(r)
lastDash = false
continue
}
if !lastDash && b.Len() > 0 {
b.WriteByte('-')
lastDash = true
}
}
slug := strings.Trim(b.String(), "-")
if slug != "" {
if len(slug) > 48 {
return strings.Trim(slug[:48], "-")
}
return slug
}
h := fnv.New32a()
_, _ = h.Write([]byte(lower))
if !utf8.ValidString(lower) || lower == "" {
lower = "workflow"
}
return fmt.Sprintf("ai-workflow-%x", h.Sum32())
}
+213
View File
@@ -0,0 +1,213 @@
package workflow
import (
"context"
"encoding/json"
"net/http"
"net/http/httptest"
"strings"
"testing"
"cyberstrike-ai/internal/config"
)
func TestGenerateDraftFromNaturalLanguageHighRiskAddsHITLAndValidGraph(t *testing.T) {
result, err := GenerateDraftFromNaturalLanguage(context.Background(), DraftRequest{
Prompt: "对目标资产做端口扫描,如果发现高危端口就执行加固脚本,最后输出报告",
Options: DraftOptions{
IncludeObjective: true,
AllowSchedule: false,
AllowHighRisk: false,
},
AvailableTools: []DraftTool{{Key: "nmap", Name: "nmap", Enabled: true}},
})
if err != nil {
t.Fatalf("GenerateDraftFromNaturalLanguage: %v", err)
}
raw, _ := json.Marshal(result.Graph)
if err := ValidateGraphJSON(context.Background(), string(raw)); err != nil {
t.Fatalf("generated graph should validate: %v\n%s", err, raw)
}
if !result.Audit.HighRisk || !result.Audit.NeedsHITL || len(result.Audit.RiskWarnings) == 0 {
t.Fatalf("audit did not flag high-risk HITL path: %#v", result.Audit)
}
var hasTool, hasHITL, hasConfirmation bool
for _, node := range result.Graph.Nodes {
if node.Type == "tool" && cfgString(node.Config, "tool_name") == "nmap" {
hasTool = true
}
if node.Type == "hitl" {
hasHITL = true
}
if cfgString(node.Config, "requires_human_confirmation") == "true" {
hasConfirmation = true
}
}
if !hasTool || !hasHITL || !hasConfirmation {
t.Fatalf("expected nmap tool, HITL, and confirmation marker; tool=%v hitl=%v confirmation=%v", hasTool, hasHITL, hasConfirmation)
}
}
func TestGenerateDraftAllowHighRiskStillLabelsConditionBranch(t *testing.T) {
result, err := GenerateDraftFromNaturalLanguage(context.Background(), DraftRequest{
Prompt: "如果漏洞扫描发现高危漏洞,允许生成执行修复脚本的草稿并输出报告",
Options: DraftOptions{
AllowHighRisk: true,
},
AvailableTools: []DraftTool{{Key: "nuclei", Name: "nuclei", Enabled: true}},
})
if err != nil {
t.Fatalf("GenerateDraftFromNaturalLanguage: %v", err)
}
raw, _ := json.Marshal(result.Graph)
if err := ValidateGraphJSON(context.Background(), string(raw)); err != nil {
t.Fatalf("generated graph should validate: %v\n%s", err, raw)
}
branches := map[string]bool{}
for _, edge := range result.Graph.Edges {
if branch := cfgString(edge.Config, "branch"); branch != "" {
branches[branch] = true
}
}
if !branches["true"] || !branches["false"] {
t.Fatalf("condition branches = %#v, want true and false", branches)
}
}
func TestGenerateDraftFromLLMUsesOpenAICompatibleEndpoint(t *testing.T) {
called := false
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
called = true
if r.URL.Path != "/chat/completions" {
t.Fatalf("path = %s, want /chat/completions", r.URL.Path)
}
if got := r.Header.Get("Authorization"); got != "Bearer test-key" {
t.Fatalf("authorization = %q", got)
}
var payload struct {
Temperature float64 `json:"temperature"`
ResponseFormat struct {
Type string `json:"type"`
} `json:"response_format"`
Messages []struct {
Role string `json:"role"`
Content string `json:"content"`
} `json:"messages"`
}
if err := json.NewDecoder(r.Body).Decode(&payload); err != nil {
t.Fatalf("decode request: %v", err)
}
if payload.Temperature != 0 || payload.ResponseFormat.Type != "json_object" {
t.Fatalf("unexpected structured output controls: temperature=%v response_format=%#v", payload.Temperature, payload.ResponseFormat)
}
if len(payload.Messages) == 0 || strings.Contains(payload.Messages[0].Content, "start|tool|agent") {
t.Fatalf("system prompt still contains pipe enum: %q", payload.Messages[0].Content)
}
w.Header().Set("Content-Type", "application/json")
_, _ = w.Write([]byte(`{"choices":[{"message":{"content":"{\"meta\":{\"id\":\"llm-port-scan\",\"name\":\"端口扫描\",\"description\":\"端口扫描\",\"enabled\":true},\"graph\":{\"nodes\":[{\"id\":\"start-1\",\"type\":\"start\",\"label\":\"开始\",\"position\":{\"x\":120,\"y\":150},\"config\":{\"input_keys\":\"message, target\"}},{\"id\":\"tool-2\",\"type\":\"tool\",\"label\":\"端口扫描\",\"position\":{\"x\":330,\"y\":150},\"config\":{\"tool_name\":\"nmap\",\"arguments\":\"{\\\"target\\\":\\\"{{inputs.target}}\\\"}\",\"timeout_seconds\":\"120\",\"join_strategy\":\"all_merge\"}},{\"id\":\"output-3\",\"type\":\"output\",\"label\":\"输出报告\",\"position\":{\"x\":540,\"y\":150},\"config\":{\"source_binding\":{\"from\":\"previous\",\"field\":\"output\"},\"join_strategy\":\"all_merge\"}}],\"edges\":[{\"id\":\"edge-1\",\"source\":\"start-1\",\"target\":\"tool-2\"},{\"id\":\"edge-2\",\"source\":\"tool-2\",\"target\":\"output-3\"}],\"config\":{\"schema_version\":1}},\"capabilities\":[{\"label\":\"端口扫描\",\"tool_name\":\"nmap\",\"tool_candidates\":[\"nmap\"]}],\"audit\":{\"assumptions\":[]}}"}}]}`))
}))
defer srv.Close()
result, err := GenerateDraftFromLLM(context.Background(), DraftRequest{
Prompt: "对目标做端口扫描并输出报告",
AvailableTools: []DraftTool{{Key: "nmap", Name: "nmap", Enabled: true}},
}, config.OpenAIConfig{APIKey: "test-key", BaseURL: srv.URL, Model: "test-model"}, nil)
if err != nil {
t.Fatalf("GenerateDraftFromLLM: %v", err)
}
if !called {
t.Fatal("expected LLM endpoint to be called")
}
if result.Generator != "llm" || !result.Audit.Savable || result.Meta.ID != "llm-port-scan" {
t.Fatalf("unexpected result: %#v", result)
}
for _, node := range result.Graph.Nodes {
if node.Type == "output" && cfgString(node.Config, "output_key") != "result" {
t.Fatalf("output_key = %q, want result", cfgString(node.Config, "output_key"))
}
}
}
func TestGenerateDraftFromLLMReturnsErrorOnMalformedJSON(t *testing.T) {
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
w.Header().Set("Content-Type", "application/json")
_, _ = w.Write([]byte(`{"choices":[{"message":{"content":"{\"meta\":{\"id\":\"bad\"},|\"graph\":{\"nodes\":[]}}"}}]}`))
}))
defer srv.Close()
_, err := GenerateDraftFromLLM(context.Background(), DraftRequest{
Prompt: "随便生成一个工作流,要求所有节点都用到输出变量",
Options: DraftOptions{
IncludeObjective: true,
},
}, config.OpenAIConfig{APIKey: "test-key", BaseURL: srv.URL, Model: "test-model"}, nil)
if err == nil {
t.Fatal("expected malformed JSON error")
}
if !strings.Contains(err.Error(), "解析大模型工作流 JSON 失败") {
t.Fatalf("unexpected error: %v", err)
}
}
func TestNormalizeLLMDraftRepairsMissingRequiredConfig(t *testing.T) {
result := normalizeLLMDraft("随便生成一个工作流", DraftRequest{}, llmDraftEnvelope{
Graph: graphDef{
Nodes: []graphNode{
{ID: "start-1", Type: "start", Label: "开始", Config: map[string]any{}},
{ID: "agent-1", Type: "agent", Label: "分析", Config: map[string]any{}},
{ID: "out-1", Type: "output", Label: "输出结果", Config: map[string]any{}},
},
Edges: []graphEdge{
{ID: "e1", Source: "start-1", Target: "agent-1"},
{ID: "e2", Source: "agent-1", Target: "out-1"},
},
},
})
raw, _ := json.Marshal(result.Graph)
if err := ValidateGraphJSON(context.Background(), string(raw)); err != nil {
t.Fatalf("normalized graph should validate: %v\n%s", err, raw)
}
var agentKey, outputKey string
for _, node := range result.Graph.Nodes {
switch node.Type {
case "agent":
agentKey = cfgString(node.Config, "output_key")
case "output":
outputKey = cfgString(node.Config, "output_key")
}
}
if agentKey == "" || outputKey != "result" {
t.Fatalf("agentKey=%q outputKey=%q", agentKey, outputKey)
}
}
func TestNormalizeLLMDraftRepairsConditionBranches(t *testing.T) {
result := normalizeLLMDraft("如果发现异常则输出详情,否则输出正常", DraftRequest{}, llmDraftEnvelope{
Graph: graphDef{
Nodes: []graphNode{
{ID: "start-1", Type: "start", Label: "开始", Config: map[string]any{}},
{ID: "cond-1", Type: "condition", Label: "判断", Config: map[string]any{"expression": `{{inputs.message}} != ""`}},
{ID: "out-yes", Type: "output", Label: "异常", Config: map[string]any{}},
{ID: "out-no", Type: "output", Label: "正常", Config: map[string]any{}},
},
Edges: []graphEdge{
{ID: "e1", Source: "start-1", Target: "cond-1"},
{ID: "e2", Source: "cond-1", Target: "out-yes"},
{ID: "e3", Source: "cond-1", Target: "out-no"},
},
},
})
raw, _ := json.Marshal(result.Graph)
if err := ValidateGraphJSON(context.Background(), string(raw)); err != nil {
t.Fatalf("normalized graph should validate: %v\n%s", err, raw)
}
branches := map[string]bool{}
for _, edge := range result.Graph.Edges {
if edge.Source == "cond-1" {
branches[cfgString(edge.Config, "branch")] = true
}
}
if !branches["true"] || !branches["false"] {
t.Fatalf("branches = %#v, want true and false", branches)
}
}
+173
View File
@@ -0,0 +1,173 @@
package workflow
import (
"context"
"encoding/json"
"fmt"
"strings"
"time"
)
type DryRunResult struct {
Outputs map[string]any `json:"outputs"`
NodeOutputs map[string]map[string]any `json:"nodeOutputs"`
Executed []string `json:"executed"`
Skipped []string `json:"skipped"`
Trace []map[string]any `json:"trace"`
Metrics map[string]any `json:"metrics"`
ReplayScript []map[string]any `json:"replayScript"`
}
func DryRunGraphJSON(ctx context.Context, graphJSON string, inputs map[string]any) (*DryRunResult, error) {
g, err := parseGraph(graphJSON)
if err != nil {
return nil, err
}
idx := indexGraph(g)
if err := validateGraphDefinition(g, idx); err != nil {
return nil, err
}
in := make(map[string]interface{}, len(inputs))
for k, v := range inputs {
in[k] = v
}
if _, ok := in["message"]; !ok {
in["message"] = ""
}
state := newWorkflowLocalState(in, "dry-run")
rt := &workflowRuntime{runID: "dry-run", idx: idx, state: state}
trace := []map[string]any{}
executedIDs := map[string]bool{}
queue := findStartNodeIDs(idx)
for len(queue) > 0 {
select {
case <-ctx.Done():
return nil, ctx.Err()
default:
}
nodeID := queue[0]
queue = queue[1:]
if executedIDs[nodeID] {
continue
}
node := idx.nodes[nodeID]
if !dryRunPredecessorsReady(idx, nodeID, executedIDs) {
queue = append(queue, nodeID)
continue
}
if err := prepareNodeInputState(rt, node); err != nil {
return nil, err
}
started := time.Now()
out, proceed, status, errText := dryRunNode(node, state)
out["duration_ms"] = time.Since(started).Milliseconds()
out["status"] = status
state.NodeOutputs[node.ID] = out
state.LastOutput = out
executedIDs[nodeID] = true
if status == "skipped" {
state.Skipped = append(state.Skipped, firstNonEmpty(node.Label, node.ID))
} else {
state.Executed = append(state.Executed, firstNonEmpty(node.Label, node.ID))
}
trace = append(trace, map[string]any{
"nodeId": node.ID,
"label": firstNonEmpty(node.Label, node.ID),
"type": node.Type,
"status": status,
"error": errText,
"output": out,
"previous": state.LastOutput,
})
if !proceed {
continue
}
for edgeIdx, edge := range idx.outgoing[nodeID] {
if edgeAllowed(edge, node, edgeIdx, state) {
queue = append(queue, edge.Target)
}
}
}
for id, node := range idx.nodes {
if !executedIDs[id] {
state.Skipped = append(state.Skipped, firstNonEmpty(node.Label, id))
}
}
return &DryRunResult{
Outputs: state.Outputs,
NodeOutputs: state.NodeOutputs,
Executed: state.Executed,
Skipped: state.Skipped,
Trace: trace,
Metrics: state.Metrics,
ReplayScript: buildReplayScript(trace),
}, nil
}
func dryRunPredecessorsReady(idx *graphIndex, nodeID string, executed map[string]bool) bool {
for _, edge := range idx.incoming[nodeID] {
if !executed[edge.Source] {
return false
}
}
return true
}
func dryRunNode(node graphNode, state *WorkflowLocalState) (map[string]any, bool, string, string) {
switch strings.ToLower(strings.TrimSpace(node.Type)) {
case "start":
return startOutputMap(node, state.Inputs["message"], state.Inputs["conversationId"], state.Inputs["projectId"]), true, "completed", ""
case "condition":
expr := cfgString(node.Config, "expression")
matched := evalCondition(expr, state)
return conditionOutputMap(node, expr, matched), true, "completed", ""
case "output":
key := cfgString(node.Config, "output_key")
value := resolveOutputSourceBinding(node.Config, state)
if static := cfgString(node.Config, "static_value"); static != "" {
value = static
}
state.Outputs[key] = value
return outputNodeOutputMap(node, key, value), true, "completed", ""
case "end":
value := resolveOutputSourceBinding(node.Config, state)
if b, ok := parseFieldBinding(node.Config, "result_binding"); ok {
value = resolveBinding(b, state)
}
return endOutputMap(node, value), false, "completed", ""
case "tool":
args, err := resolveToolArguments(node.Config, state)
if err != nil {
errText := fmt.Sprintf("工具参数不是合法 JSON%v", err)
return outputMap(envelope("tool", node.ID, node.Type, "failed", ""), map[string]any{"error": errText}), false, "failed", errText
}
return toolOutputMap(node, "[dry-run] tool call skipped", cfgString(node.Config, "tool_name"), args, "dry-run", false), true, "simulated", ""
case "agent":
mode := firstNonEmpty(cfgString(node.Config, "agent_mode"), "eino_single")
response := "[dry-run] agent execution skipped"
if key := cfgString(node.Config, "output_key"); key != "" {
state.Outputs[key] = response
}
return agentOutputMap(node, response, mode, nil), true, "simulated", ""
case "hitl":
prompt := resolveHITLPromptBinding(node.Config, state)
return hitlOutputMap(node, "simulated", prompt, prompt, firstNonEmpty(cfgString(node.Config, "reviewer"), "human"), true), true, "simulated", ""
default:
return outputMap(envelope("unknown", node.ID, node.Type, "skipped", ""), map[string]any{"reason": "未知节点类型"}), true, "skipped", "未知节点类型"
}
}
func buildReplayScript(trace []map[string]any) []map[string]any {
out := make([]map[string]any, 0, len(trace))
for i, step := range trace {
raw, _ := json.Marshal(step["output"])
out = append(out, map[string]any{
"step": i + 1,
"nodeId": step["nodeId"],
"type": step["type"],
"status": step["status"],
"output": string(raw),
})
}
return out
}
+107
View File
@@ -0,0 +1,107 @@
package workflow
import (
"context"
"fmt"
"github.com/cloudwego/eino/compose"
)
func hasConditionalOutgoingEdges(idx *graphIndex, nodeID string) bool {
for _, edge := range idx.outgoing[nodeID] {
cond := firstNonEmpty(cfgString(edge.Config, "condition"), cfgString(edge.Config, "expression"))
if cond != "" {
return true
}
}
return false
}
func wireConditionBranch(
wf *compose.Workflow[WorkflowInput, WorkflowOutput],
nodeRefs map[string]*compose.WorkflowNode,
idx *graphIndex,
condID string,
condNode graphNode,
) error {
edges := idx.outgoing[condID]
if len(edges) == 0 {
return nil
}
branchID := branchNodeID(condID)
wf.AddPassthroughNode(branchID).AddInput(condID)
endNodes := map[string]bool{compose.END: true}
for _, edge := range edges {
endNodes[edge.Target] = true
}
sortedEdges := append([]graphEdge(nil), edges...)
sortEdgesByCanvas(sortedEdges, idx.nodes)
branch := compose.NewGraphBranch(func(runCtx context.Context, _ map[string]any) (string, error) {
rt := workflowRuntimeFrom(runCtx)
if rt == nil {
return compose.END, fmt.Errorf("workflow runtime missing in context")
}
emitConditionBranchProgress(rt.args, rt.runID, condNode, sortedEdges, idx.nodes, rt.state)
for edgeIdx, edge := range sortedEdges {
if conditionBranchAllowed(edge, edgeIdx, rt.state) {
return edge.Target, nil
}
}
return compose.END, nil
}, endNodes)
wf.AddBranch(branchID, branch)
for _, edge := range edges {
if target, ok := nodeRefs[edge.Target]; ok {
target.AddInput(branchID)
}
}
return nil
}
func wireEdgeConditionBranch(
wf *compose.Workflow[WorkflowInput, WorkflowOutput],
nodeRefs map[string]*compose.WorkflowNode,
idx *graphIndex,
sourceID string,
sourceNode graphNode,
) error {
edges := idx.outgoing[sourceID]
if len(edges) == 0 {
return nil
}
branchID := edgeBranchNodeID(sourceID)
wf.AddPassthroughNode(branchID).AddInput(sourceID)
endNodes := map[string]bool{compose.END: true}
for _, edge := range edges {
endNodes[edge.Target] = true
}
sortedEdges := append([]graphEdge(nil), edges...)
sortEdgesByCanvas(sortedEdges, idx.nodes)
branch := compose.NewGraphBranch(func(runCtx context.Context, _ map[string]any) (string, error) {
rt := workflowRuntimeFrom(runCtx)
if rt == nil {
return compose.END, fmt.Errorf("workflow runtime missing in context")
}
for edgeIdx, edge := range sortedEdges {
if edgeAllowed(edge, sourceNode, edgeIdx, rt.state) {
return edge.Target, nil
}
}
return compose.END, nil
}, endNodes)
wf.AddBranch(branchID, branch)
for _, edge := range edges {
if target, ok := nodeRefs[edge.Target]; ok {
target.AddInput(branchID)
}
}
return nil
}
+22
View File
@@ -0,0 +1,22 @@
package workflow
import (
"context"
"cyberstrike-ai/internal/config"
"cyberstrike-ai/internal/einoobserve"
)
func attachWorkflowCallbacks(ctx context.Context, cfg *config.Config, args RunArgs, workflowName string) context.Context {
if cfg == nil {
return ctx
}
cbCfg := &cfg.MultiAgent.EinoCallbacks
return einoobserve.AttachAgentRunCallbacks(ctx, cbCfg, einoobserve.Params{
Logger: args.Logger,
Progress: args.Progress,
ConversationID: args.ConversationID,
OrchMode: "workflow",
OrchestratorName: workflowName,
})
}
+243
View File
@@ -0,0 +1,243 @@
package workflow
import (
"context"
"encoding/json"
"fmt"
"strings"
"github.com/cloudwego/eino/compose"
)
func executeEinoGraph(ctx context.Context, args RunArgs, runID string, workflowID string, version int, g *graphDef, state *WorkflowLocalState) error {
_, err := invokeEinoGraph(ctx, args, runID, workflowID, version, g, state, false)
return err
}
func invokeEinoGraph(ctx context.Context, args RunArgs, runID string, workflowID string, version int, g *graphDef, state *WorkflowLocalState, resume bool) (bool, error) {
wfInput := workflowInputFromMap(state.Inputs)
if resume {
wfInput = WorkflowInput{}
}
rt := &workflowRuntime{
args: args,
runID: runID,
idx: indexGraph(g),
state: state,
}
art, err := defaultEngine.getOrCompile(ctx, workflowID, version, g)
if err != nil {
return false, fmt.Errorf("编译 Eino Workflow 失败: %w", err)
}
rt.idx = art.idx
runCtx := withWorkflowRuntime(ctx, rt)
runCtx = attachWorkflowCallbacks(runCtx, args.AppCfg, args, workflowID)
invokeOpts := []compose.Option{compose.WithCheckPointID(runID)}
for {
_, err = art.runnable.Invoke(runCtx, wfInput, invokeOpts...)
if err == nil {
return false, nil
}
if hitlErr := extractAwaitingHITL(err, art, runID, args, state); hitlErr != nil {
return true, hitlErr
}
return false, err
}
}
func extractAwaitingHITL(err error, art *compiledArtifact, runID string, args RunArgs, state *WorkflowLocalState) error {
info, ok := compose.ExtractInterruptInfo(err)
if !ok || len(art.hitlIDs) == 0 {
return nil
}
nodeID := nextHITLNodeID(info, art.hitlIDs)
node := art.idx.nodes[nodeID]
if nodeID == "" {
return nil
}
prompt := resolveHITLPromptBinding(node.Config, state)
label := firstNonEmpty(node.Label, nodeID)
if args.DB != nil {
pending := map[string]any{
"nodeId": nodeID,
"label": label,
"prompt": prompt,
"reviewer": cfgString(node.Config, "reviewer"),
"checkpointId": runID,
"interrupt": workflowInterruptMetadata(info),
"resumePayload": map[string]any{"approved": "bool", "comment": "string"},
}
pendingJSON, _ := json.Marshal(pending)
_ = args.DB.SetWorkflowRunAwaitingHITL(runID, nodeID, string(pendingJSON))
}
if args.Progress != nil {
args.Progress("workflow_hitl_waiting", fmt.Sprintf("等待人工确认:%s", label), map[string]any{
"workflowRunId": runID,
"nodeId": nodeID,
"label": label,
"prompt": prompt,
"reviewer": cfgString(node.Config, "reviewer"),
"mode": "interactive",
"resumeApi": fmt.Sprintf("/api/workflows/runs/%s/resume", runID),
})
}
return &AwaitingHITLError{
RunID: runID,
NodeID: nodeID,
NodeLabel: label,
Prompt: prompt,
Reviewer: cfgString(node.Config, "reviewer"),
}
}
func workflowInterruptMetadata(info *compose.InterruptInfo) map[string]any {
if info == nil {
return map[string]any{}
}
before := append([]string(nil), info.BeforeNodes...)
return map[string]any{
"beforeNodes": before,
"resumeTarget": firstString(before),
"address": map[string]any{
"kind": "compose_interrupt",
"beforeNodes": before,
"path": strings.Join(before, "/"),
},
"raw": fmt.Sprintf("%+v", info),
}
}
func firstString(values []string) string {
if len(values) == 0 {
return ""
}
return values[0]
}
func nextHITLNodeID(info *compose.InterruptInfo, hitlIDs []string) string {
if info != nil && len(info.BeforeNodes) > 0 {
for _, id := range info.BeforeNodes {
for _, hitl := range hitlIDs {
if id == hitl {
return id
}
}
}
return info.BeforeNodes[0]
}
if len(hitlIDs) == 0 {
return ""
}
return hitlIDs[0]
}
// ResumeWorkflowRun continues a run paused at HITL after human decision.
func ResumeWorkflowRun(ctx context.Context, args RunArgs, runID string, approved bool, comment string) (*RunResult, error) {
run, err := args.DB.GetWorkflowRun(runID)
if err != nil {
return nil, err
}
if run == nil {
return nil, fmt.Errorf("工作流运行不存在")
}
if run.Status != "awaiting_hitl" {
return nil, fmt.Errorf("工作流运行不在等待审批状态: %s", run.Status)
}
wf, err := args.DB.GetWorkflowDefinition(run.WorkflowID)
if err != nil || wf == nil {
return nil, fmt.Errorf("工作流定义不存在")
}
graph, err := parseGraph(wf.GraphJSON)
if err != nil {
return nil, err
}
var input map[string]interface{}
_ = json.Unmarshal([]byte(run.InputJSON), &input)
state := newWorkflowLocalState(input, runID)
if state.Inputs == nil {
state.Inputs = map[string]any{}
}
state.Inputs["_hitl_approved"] = approved
state.Inputs["_hitl_comment"] = strings.TrimSpace(comment)
state.Inputs["_hitl_node_id"] = run.PendingHITLNodeID
if !approved {
errText := strings.TrimSpace(comment)
if errText == "" {
errText = "人工审批拒绝"
}
_ = args.DB.FinishWorkflowRun(runID, "rejected", "", errText)
if args.Progress != nil {
args.Progress("workflow_hitl_rejected", fmt.Sprintf("工作流已在审批节点「%s」被拒绝。", run.PendingHITLNodeID), map[string]interface{}{
"workflowRunId": runID,
"nodeId": run.PendingHITLNodeID,
"comment": errText,
})
}
return &RunResult{
RunID: runID,
Response: fmt.Sprintf("工作流已在审批节点「%s」被拒绝。", run.PendingHITLNodeID),
Status: "rejected",
}, nil
}
if args.Progress != nil {
args.Progress("workflow_hitl_resumed", "人工审批已通过,继续执行", map[string]interface{}{
"workflowRunId": runID,
"nodeId": run.PendingHITLNodeID,
"comment": strings.TrimSpace(comment),
})
}
_ = args.DB.SetWorkflowRunStatus(runID, "running")
resumeArgs := args
if strings.TrimSpace(resumeArgs.ConversationID) == "" {
resumeArgs.ConversationID = run.ConversationID
}
awaiting, err := invokeEinoGraph(ctx, resumeArgs, runID, wf.ID, run.WorkflowVersion, graph, state, true)
if err != nil {
if IsAwaitingHITL(err) {
return &RunResult{
RunID: runID,
Status: "awaiting_hitl",
Response: fmt.Sprintf("工作流在节点「%s」等待下一次人工确认。", err.(*AwaitingHITLError).NodeID),
AwaitingHITL: true,
}, nil
}
_ = args.DB.FinishWorkflowRun(runID, "failed", "", err.Error())
return nil, err
}
_ = awaiting
output := map[string]interface{}{
"workflowId": wf.ID,
"workflowName": wf.Name,
"workflowVersion": wf.Version,
"workflowRunId": runID,
"status": "completed",
"outputs": state.Outputs,
"metrics": state.Metrics,
"executedNodes": state.Executed,
"skippedNodes": state.Skipped,
"engine": "eino_workflow",
}
outputJSON, _ := json.Marshal(output)
response := renderWorkflowResponse(args.Role.Name, wf.Name, wf.Version, runID, state)
_ = args.DB.FinishWorkflowRun(runID, "completed", string(outputJSON), "")
if args.Progress != nil {
args.Progress("workflow_done", fmt.Sprintf("流程「%s」运行完成", wf.Name), map[string]interface{}{
"workflowRunId": runID,
"workflowId": wf.ID,
"outputs": state.Outputs,
"metrics": state.Metrics,
"response": response,
"engine": "eino_workflow",
})
}
return &RunResult{Response: response, RunID: runID, Status: "completed"}, nil
}
+412
View File
@@ -0,0 +1,412 @@
package workflow
import (
"context"
"fmt"
"path/filepath"
"strings"
"testing"
"cyberstrike-ai/internal/config"
"cyberstrike-ai/internal/database"
"github.com/cloudwego/eino/compose"
"go.uber.org/zap"
)
func testWorkflowDB(t *testing.T) *database.DB {
t.Helper()
dir := t.TempDir()
db, err := database.NewDB(filepath.Join(dir, "workflow.db"), zap.NewNop())
if err != nil {
t.Fatalf("NewDB: %v", err)
}
t.Cleanup(func() { _ = db.Close() })
return db
}
func linearStartOutputGraph() string {
return `{
"nodes": [
{"id": "start-1", "type": "start", "label": "开始", "position": {"x": 0, "y": 0}, "config": {}},
{"id": "out-1", "type": "output", "label": "输出", "position": {"x": 0, "y": 120}, "config": {"output_key": "result", "source_binding": {"from": "inputs", "field": "message"}}}
],
"edges": [
{"id": "e1", "source": "start-1", "target": "out-1"}
],
"config": {"schema_version": 1}
}`
}
func conditionBranchGraph() string {
return `{
"nodes": [
{"id": "start-1", "type": "start", "label": "开始", "position": {"x": 0, "y": 0}, "config": {}},
{"id": "cond-1", "type": "condition", "label": "判断", "position": {"x": 0, "y": 80}, "config": {"expression": "{{inputs.message}} == yes"}},
{"id": "out-yes", "type": "output", "label": "是", "position": {"x": -80, "y": 160}, "config": {"output_key": "branch", "static_value": "yes"}},
{"id": "out-no", "type": "output", "label": "否", "position": {"x": 80, "y": 160}, "config": {"output_key": "branch", "static_value": "no"}}
],
"edges": [
{"id": "e1", "source": "start-1", "target": "cond-1"},
{"id": "e2", "source": "cond-1", "target": "out-yes", "label": "是"},
{"id": "e3", "source": "cond-1", "target": "out-no", "label": "否"}
],
"config": {"schema_version": 1}
}`
}
func TestValidateGraphJSON_linear(t *testing.T) {
if err := ValidateGraphJSON(context.Background(), linearStartOutputGraph()); err != nil {
t.Fatalf("validate: %v", err)
}
}
func TestValidateGraphJSON_rejectsInvalidGraphs(t *testing.T) {
tests := []struct {
name string
graph string
wantErr string
}{
{
name: "start with incoming edge",
graph: `{
"nodes": [
{"id": "start-1", "type": "start", "label": "开始", "position": {"x": 0, "y": 0}, "config": {}},
{"id": "agent-1", "type": "agent", "label": "Agent", "position": {"x": 0, "y": 80}, "config": {"instruction": "noop"}},
{"id": "out-1", "type": "output", "label": "输出", "position": {"x": 0, "y": 160}, "config": {"output_key": "result"}}
],
"edges": [
{"id": "e1", "source": "start-1", "target": "agent-1"},
{"id": "e2", "source": "agent-1", "target": "start-1"},
{"id": "e3", "source": "agent-1", "target": "out-1"}
]
}`,
wantErr: "开始节点",
},
{
name: "output with outgoing edge",
graph: `{
"nodes": [
{"id": "start-1", "type": "start", "label": "开始", "position": {"x": 0, "y": 0}, "config": {}},
{"id": "out-1", "type": "output", "label": "输出", "position": {"x": 0, "y": 80}, "config": {"output_key": "result"}},
{"id": "end-1", "type": "end", "label": "结束", "position": {"x": 0, "y": 160}, "config": {}}
],
"edges": [
{"id": "e1", "source": "start-1", "target": "out-1"},
{"id": "e2", "source": "out-1", "target": "end-1"}
]
}`,
wantErr: "不能有出边",
},
{
name: "tool without name",
graph: `{
"nodes": [
{"id": "start-1", "type": "start", "label": "开始", "position": {"x": 0, "y": 0}, "config": {}},
{"id": "tool-1", "type": "tool", "label": "工具", "position": {"x": 0, "y": 80}, "config": {}},
{"id": "out-1", "type": "output", "label": "输出", "position": {"x": 0, "y": 160}, "config": {"output_key": "result"}}
],
"edges": [
{"id": "e1", "source": "start-1", "target": "tool-1"},
{"id": "e2", "source": "tool-1", "target": "out-1"}
]
}`,
wantErr: "必须选择 MCP 工具",
},
{
name: "condition with too many branches",
graph: `{
"nodes": [
{"id": "start-1", "type": "start", "label": "开始", "position": {"x": 0, "y": 0}, "config": {}},
{"id": "cond-1", "type": "condition", "label": "判断", "position": {"x": 0, "y": 80}, "config": {"expression": "{{inputs.message}}"}},
{"id": "out-1", "type": "output", "label": "输出1", "position": {"x": -80, "y": 160}, "config": {"output_key": "a"}},
{"id": "out-2", "type": "output", "label": "输出2", "position": {"x": 0, "y": 160}, "config": {"output_key": "b"}},
{"id": "out-3", "type": "output", "label": "输出3", "position": {"x": 80, "y": 160}, "config": {"output_key": "c"}}
],
"edges": [
{"id": "e1", "source": "start-1", "target": "cond-1"},
{"id": "e2", "source": "cond-1", "target": "out-1"},
{"id": "e3", "source": "cond-1", "target": "out-2"},
{"id": "e4", "source": "cond-1", "target": "out-3"}
]
}`,
wantErr: "1 到 2 条出边",
},
{
name: "orphan node",
graph: `{
"nodes": [
{"id": "start-1", "type": "start", "label": "开始", "position": {"x": 0, "y": 0}, "config": {}},
{"id": "out-1", "type": "output", "label": "输出", "position": {"x": 0, "y": 80}, "config": {"output_key": "result"}},
{"id": "agent-1", "type": "agent", "label": "孤岛", "position": {"x": 200, "y": 80}, "config": {"instruction": "noop"}}
],
"edges": [
{"id": "e1", "source": "start-1", "target": "out-1"}
]
}`,
wantErr: "不可达",
},
{
name: "cycle",
graph: `{
"nodes": [
{"id": "start-1", "type": "start", "label": "开始", "position": {"x": 0, "y": 0}, "config": {}},
{"id": "agent-1", "type": "agent", "label": "Agent1", "position": {"x": 0, "y": 80}, "config": {"instruction": "noop", "output_key": "a1"}},
{"id": "agent-2", "type": "agent", "label": "Agent2", "position": {"x": 0, "y": 160}, "config": {"instruction": "noop", "output_key": "a2"}},
{"id": "out-1", "type": "output", "label": "输出", "position": {"x": 0, "y": 240}, "config": {"output_key": "result"}}
],
"edges": [
{"id": "e1", "source": "start-1", "target": "agent-1"},
{"id": "e2", "source": "agent-1", "target": "agent-2"},
{"id": "e3", "source": "agent-2", "target": "agent-1"},
{"id": "e4", "source": "agent-2", "target": "out-1"}
]
}`,
wantErr: "环路",
},
{
name: "output without key",
graph: `{
"nodes": [
{"id": "start-1", "type": "start", "label": "开始", "position": {"x": 0, "y": 0}, "config": {}},
{"id": "out-1", "type": "output", "label": "输出", "position": {"x": 0, "y": 80}, "config": {}}
],
"edges": [
{"id": "e1", "source": "start-1", "target": "out-1"}
]
}`,
wantErr: "输出变量名",
},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
err := ValidateGraphJSON(context.Background(), tt.graph)
if err == nil {
t.Fatal("expected validation error")
}
if !strings.Contains(err.Error(), tt.wantErr) {
t.Fatalf("error = %q, want substring %q", err.Error(), tt.wantErr)
}
})
}
}
func TestCompileEngine_linear(t *testing.T) {
ctx := context.Background()
SetCheckpointDir(t.TempDir())
g, err := parseGraph(linearStartOutputGraph())
if err != nil {
t.Fatal(err)
}
if _, err := defaultEngine.compile(ctx, g); err != nil {
t.Fatalf("compile: %v", err)
}
}
func createTestWorkflowRun(t *testing.T, db *database.DB, runID string) {
t.Helper()
if err := db.CreateWorkflowRun(&database.WorkflowRun{
ID: runID,
WorkflowID: "test-wf",
Status: "running",
}); err != nil {
t.Fatalf("CreateWorkflowRun: %v", err)
}
}
func TestExecuteEinoGraph_linearStartOutput(t *testing.T) {
ctx := context.Background()
SetCheckpointDir(t.TempDir())
db := testWorkflowDB(t)
createTestWorkflowRun(t, db, "run-linear")
g, err := parseGraph(linearStartOutputGraph())
if err != nil {
t.Fatal(err)
}
state := newWorkflowLocalState(map[string]interface{}{"message": "ping"}, "run-linear")
args := RunArgs{DB: db}
if err := executeEinoGraph(ctx, args, "run-linear", "test-wf", 1, g, state); err != nil {
t.Fatalf("execute: %v", err)
}
if got := state.Outputs["result"]; got != "ping" {
t.Fatalf("outputs[result] = %v, want ping", got)
}
if len(state.Executed) != 2 {
t.Fatalf("executed nodes = %d, want 2", len(state.Executed))
}
}
func TestExecuteEinoGraph_checkpointRestoresStartOutput(t *testing.T) {
ctx := context.Background()
checkpointStore, err := newFileCheckPointStore(t.TempDir())
if err != nil {
t.Fatalf("new checkpoint store: %v", err)
}
state := newWorkflowLocalState(map[string]interface{}{"message": "ping"}, "run-checkpoint")
node := graphNode{ID: "start-1", Type: "start"}
wf := compose.NewWorkflow[WorkflowInput, WorkflowOutput](
compose.WithGenLocalState(func(context.Context) *WorkflowLocalState { return state }),
)
start := wf.AddLambdaNode("start-1", compose.InvokableLambda(func(_ context.Context, input WorkflowInput) (WorkflowNodeOutput, error) {
result := startOutputMap(node, input.Message, input.ConversationID, input.ProjectID)
state.NodeOutputs[node.ID] = result
state.NodeOutputs["condition-1"] = conditionOutputMap(graphNode{ID: "condition-1", Type: "condition"}, "{{inputs.message}} == ping", true)
state.NodeOutputs["tool-1"] = toolOutputMap(graphNode{ID: "tool-1", Type: "tool"}, "tool result", "lookup", map[string]any{"id": "1"}, "exec-1", false)
state.NodeOutputs["agent-1"] = agentOutputMap(graphNode{ID: "agent-1", Type: "agent"}, "agent result", "chat", []string{"exec-1"})
state.NodeOutputs["hitl-1"] = hitlOutputMap(graphNode{ID: "hitl-1", Type: "hitl"}, "completed", "approved", "continue?", "reviewer", true)
state.NodeOutputs["output-1"] = outputNodeOutputMap(graphNode{ID: "output-1", Type: "output"}, "result", "ping")
state.NodeOutputs["end-1"] = endOutputMap(graphNode{ID: "end-1", Type: "end"}, "done")
state.LastOutput = result
state.Outputs["seed"] = "preserved"
return result, nil
}))
outputNode := wf.AddLambdaNode("out-1", compose.InvokableLambda(func(_ context.Context, input WorkflowNodeOutput) (WorkflowNodeOutput, error) {
return input, nil
}))
start.AddInput(compose.START)
outputNode.AddInput("start-1")
wf.End().AddInput("out-1", compose.ToField("out-1"))
runnable, err := wf.Compile(ctx,
compose.WithCheckPointStore(checkpointStore),
compose.WithInterruptAfterNodes([]string{"start-1"}),
)
if err != nil {
t.Fatalf("compile: %v", err)
}
_, err = runnable.Invoke(ctx, workflowInputFromMap(state.Inputs), compose.WithCheckPointID("run-checkpoint"))
info, ok := compose.ExtractInterruptInfo(err)
if !ok {
t.Fatalf("invoke error = %v, want checkpoint interrupt", err)
}
restored, ok := info.State.(*WorkflowLocalState)
if !ok {
t.Fatalf("checkpoint state = %T, want *WorkflowLocalState", info.State)
}
for nodeID, wantType := range map[string]string{
"start-1": "StartOutput",
"condition-1": "ConditionOutput",
"tool-1": "ToolOutput",
"agent-1": "AgentOutput",
"hitl-1": "HITLOutput",
"output-1": "OutputNodeOutput",
"end-1": "NodeOutputEnvelope",
} {
if got := fmt.Sprintf("%T", restored.NodeOutputs[nodeID]["typed"]); got != "workflow."+wantType {
t.Fatalf("restored %s typed output = %s, want workflow.%s", nodeID, got, wantType)
}
}
if got := valueFromPath("previous.message", restored); got != "ping" {
t.Fatalf("restored previous.message = %v, want ping", got)
}
if got := valueFromPath("inputs.message", restored); got != "ping" {
t.Fatalf("restored inputs.message = %v, want ping", got)
}
if got := valueFromPath("outputs.seed", restored); got != "preserved" {
t.Fatalf("restored outputs.seed = %v, want preserved", got)
}
result, err := runnable.Invoke(ctx, WorkflowInput{}, compose.WithCheckPointID("run-checkpoint"))
if err != nil {
t.Fatalf("resume checkpoint: %v", err)
}
output, ok := result["out-1"].(map[string]any)
if !ok {
t.Fatalf("resumed output type = %T, want map[string]any", result["out-1"])
}
if got := output["output"]; got != "ping" {
t.Fatalf("resumed output = %v, want ping", got)
}
}
func TestExecuteEinoGraph_conditionBranch(t *testing.T) {
ctx := context.Background()
SetCheckpointDir(t.TempDir())
db := testWorkflowDB(t)
createTestWorkflowRun(t, db, "run-yes")
createTestWorkflowRun(t, db, "run-no")
g, err := parseGraph(conditionBranchGraph())
if err != nil {
t.Fatal(err)
}
stateYes := newWorkflowLocalState(map[string]interface{}{"message": "yes"}, "run-yes")
if err := executeEinoGraph(ctx, RunArgs{DB: db}, "run-yes", "test-wf-branch", 1, g, stateYes); err != nil {
t.Fatalf("execute yes: %v", err)
}
if got := stateYes.Outputs["branch"]; got != "yes" {
t.Fatalf("yes branch output = %v", got)
}
stateNo := newWorkflowLocalState(map[string]interface{}{"message": "no"}, "run-no")
if err := executeEinoGraph(ctx, RunArgs{DB: db}, "run-no", "test-wf-branch", 1, g, stateNo); err != nil {
t.Fatalf("execute no: %v", err)
}
if got := stateNo.Outputs["branch"]; got != "no" {
t.Fatalf("no branch output = %v", got)
}
}
func TestRunRoleBoundWorkflow_integration(t *testing.T) {
ctx := context.Background()
SetCheckpointDir(t.TempDir())
db := testWorkflowDB(t)
graph := linearStartOutputGraph()
if err := db.UpsertWorkflowDefinition(&database.WorkflowDefinition{
ID: "wf-linear",
Name: "线性流程",
Version: 1,
GraphJSON: graph,
Enabled: true,
}); err != nil {
t.Fatal(err)
}
role := config.RoleConfig{
Name: "tester",
Enabled: true,
WorkflowID: "wf-linear",
WorkflowPolicy: "auto",
}
result, err := RunRoleBoundWorkflow(ctx, RunArgs{
DB: db,
Logger: zap.NewNop(),
Role: role,
UserMessage: "from-role",
})
if err != nil {
t.Fatalf("RunRoleBoundWorkflow: %v", err)
}
if result == nil || result.RunID == "" {
t.Fatal("expected run result")
}
}
func TestCompiledCache_reuse(t *testing.T) {
ctx := context.Background()
SetCheckpointDir(t.TempDir())
InvalidateCompiledCache("cache-wf")
g, err := parseGraph(linearStartOutputGraph())
if err != nil {
t.Fatal(err)
}
a1, err := defaultEngine.getOrCompile(ctx, "cache-wf", 1, g)
if err != nil {
t.Fatal(err)
}
a2, err := defaultEngine.getOrCompile(ctx, "cache-wf", 1, g)
if err != nil {
t.Fatal(err)
}
if a1 != a2 {
t.Fatal("expected cached artifact pointer reuse")
}
InvalidateCompiledCache("cache-wf")
a3, err := defaultEngine.getOrCompile(ctx, "cache-wf", 1, g)
if err != nil {
t.Fatal(err)
}
if a1 == a3 {
t.Fatal("expected new artifact after invalidation")
}
}
+64
View File
@@ -0,0 +1,64 @@
package workflow
import (
"context"
"cyberstrike-ai/internal/agent"
"cyberstrike-ai/internal/config"
"cyberstrike-ai/internal/database"
"go.uber.org/zap"
)
type workflowRuntimeCtxKey struct{}
// workflowRuntime carries per-run execution context into Eino Workflow local state.
type workflowRuntime struct {
args RunArgs
runID string
idx *graphIndex
state *WorkflowLocalState
}
func withWorkflowRuntime(ctx context.Context, rt *workflowRuntime) context.Context {
return context.WithValue(ctx, workflowRuntimeCtxKey{}, rt)
}
func workflowRuntimeFrom(ctx context.Context) *workflowRuntime {
rt, _ := ctx.Value(workflowRuntimeCtxKey{}).(*workflowRuntime)
return rt
}
func newWorkflowRuntime(args RunArgs, runID string, idx *graphIndex, inputs map[string]interface{}) *workflowRuntime {
return &workflowRuntime{
args: args,
runID: runID,
idx: idx,
state: newWorkflowLocalState(inputs, runID),
}
}
// RunArgs is the execution context for a role-bound workflow run.
type RunArgs struct {
DB *database.DB
Logger *zap.Logger
Role config.RoleConfig
AppCfg *config.Config
Agent *agent.Agent
ConversationID string
ProjectID string
UserMessage string
History []agent.ChatMessage
RoleTools []string
AgentsMarkdownDir string
SystemPromptExtra string
AssistantMessageID string
Progress agent.ProgressCallback
}
type RunResult struct {
Response string
RunID string
Status string
AwaitingHITL bool
}
+239
View File
@@ -0,0 +1,239 @@
package workflow
import (
"context"
"fmt"
"strings"
"sync"
"github.com/cloudwego/eino/compose"
)
type compiledArtifact struct {
runnable compose.Runnable[WorkflowInput, WorkflowOutput]
idx *graphIndex
hitlIDs []string
}
// Engine compiles and caches Eino Workflow artifacts.
type Engine struct {
mu sync.RWMutex
cache map[string]*compiledArtifact
cpStore compose.CheckPointStore
cpStoreMu sync.Once
cpStoreErr error
checkpointDir string
}
var defaultEngine = &Engine{
cache: make(map[string]*compiledArtifact),
checkpointDir: "data/workflow-checkpoints",
}
// SetCheckpointDir overrides the workflow checkpoint root (mainly for tests).
func SetCheckpointDir(dir string) {
defaultEngine.mu.Lock()
defer defaultEngine.mu.Unlock()
defaultEngine.checkpointDir = strings.TrimSpace(dir)
defaultEngine.cpStore = nil
defaultEngine.cpStoreErr = nil
defaultEngine.cpStoreMu = sync.Once{}
}
func (e *Engine) checkpointStore() (compose.CheckPointStore, error) {
e.cpStoreMu.Do(func() {
e.cpStore, e.cpStoreErr = newFileCheckPointStore(e.checkpointDir)
})
return e.cpStore, e.cpStoreErr
}
// InvalidateCompiledCache drops cached compilations for a workflow id.
func InvalidateCompiledCache(workflowID string) {
workflowID = strings.TrimSpace(workflowID)
if workflowID == "" {
return
}
defaultEngine.mu.Lock()
defer defaultEngine.mu.Unlock()
for key := range defaultEngine.cache {
if strings.HasPrefix(key, workflowID+":") {
delete(defaultEngine.cache, key)
}
}
}
// ValidateGraphJSON parses and trial-compiles a canvas graph (save-time gate).
func ValidateGraphJSON(ctx context.Context, graphJSON string) error {
g, err := parseGraph(graphJSON)
if err != nil {
return err
}
idx := indexGraph(g)
if err := validateGraphDefinition(g, idx); err != nil {
return err
}
_, err = defaultEngine.compile(ctx, g)
return err
}
func hasTerminalNode(idx *graphIndex) bool {
for id, node := range idx.nodes {
if len(idx.outgoing[id]) == 0 {
return true
}
if strings.EqualFold(node.Type, "end") || strings.EqualFold(node.Type, "output") {
return true
}
}
return false
}
func (e *Engine) getOrCompile(ctx context.Context, workflowID string, version int, g *graphDef) (*compiledArtifact, error) {
key := cacheKey(workflowID, version)
e.mu.RLock()
if art, ok := e.cache[key]; ok {
e.mu.RUnlock()
return art, nil
}
e.mu.RUnlock()
art, err := e.compile(ctx, g)
if err != nil {
return nil, err
}
e.mu.Lock()
if existing, ok := e.cache[key]; ok {
e.mu.Unlock()
return existing, nil
}
e.cache[key] = art
e.mu.Unlock()
return art, nil
}
func (e *Engine) compile(ctx context.Context, g *graphDef) (*compiledArtifact, error) {
cpStore, err := e.checkpointStore()
if err != nil {
return nil, err
}
idx := indexGraph(g)
if err := validateGraphDefinition(g, idx); err != nil {
return nil, err
}
hitlIDs := collectHITLNodeIDs(idx)
compileOpts := []compose.GraphCompileOption{
compose.WithGraphName("CyberStrikeWorkflow"),
compose.WithCheckPointStore(cpStore),
}
if len(hitlIDs) > 0 {
compileOpts = append(compileOpts, compose.WithInterruptBeforeNodes(hitlIDs))
}
wf := compose.NewWorkflow[WorkflowInput, WorkflowOutput](
compose.WithGenLocalState(func(runCtx context.Context) *WorkflowLocalState {
if rt := workflowRuntimeFrom(runCtx); rt != nil && rt.state != nil {
return rt.state
}
return &WorkflowLocalState{
Outputs: make(map[string]any),
NodeOutputs: make(map[string]map[string]any),
NodeProceed: make(map[string]bool),
}
}),
)
nodeRefs := make(map[string]*compose.WorkflowNode, len(idx.nodes))
for id, node := range idx.nodes {
n := node
if strings.EqualFold(n.Type, "agent") {
sub, err := compileAgentSubgraph(ctx, n)
if err != nil {
return nil, fmt.Errorf("编译 Agent 子图 %s 失败: %w", id, err)
}
nodeRefs[id] = wf.AddGraphNode(id, sub)
continue
}
if strings.EqualFold(n.Type, "start") {
nodeRefs[id] = wf.AddLambdaNode(id, compose.InvokableLambda(func(runCtx context.Context, _ WorkflowInput) (WorkflowNodeOutput, error) {
return runWorkflowNodeLambda(runCtx, n)
}))
continue
}
nodeRefs[id] = wf.AddLambdaNode(id, compose.InvokableLambda(func(runCtx context.Context, _ WorkflowNodeOutput) (WorkflowNodeOutput, error) {
return runWorkflowNodeLambda(runCtx, n)
}))
}
for id, node := range idx.nodes {
if strings.EqualFold(node.Type, "condition") {
if err := wireConditionBranch(wf, nodeRefs, idx, id, node); err != nil {
return nil, err
}
continue
}
if hasConditionalOutgoingEdges(idx, id) {
if err := wireEdgeConditionBranch(wf, nodeRefs, idx, id, node); err != nil {
return nil, err
}
continue
}
for _, edge := range idx.outgoing[id] {
if target, ok := nodeRefs[edge.Target]; ok {
target.AddInput(id)
}
}
}
for _, startID := range findStartNodeIDs(idx) {
if ref, ok := nodeRefs[startID]; ok {
ref.AddInput(compose.START)
}
}
endNode := wf.End()
for id, node := range idx.nodes {
if len(idx.outgoing[id]) == 0 || strings.EqualFold(node.Type, "end") {
endNode.AddInput(id, compose.ToField(id))
}
}
runnable, err := wf.Compile(ctx, compileOpts...)
if err != nil {
return nil, err
}
return &compiledArtifact{runnable: runnable, idx: idx, hitlIDs: hitlIDs}, nil
}
func collectHITLNodeIDs(idx *graphIndex) []string {
var ids []string
for id, node := range idx.nodes {
if strings.EqualFold(node.Type, "hitl") {
ids = append(ids, id)
}
}
return ids
}
func runWorkflowNodeLambda(runCtx context.Context, n graphNode) (WorkflowNodeOutput, error) {
localRT := workflowRuntimeFrom(runCtx)
if localRT == nil {
return nil, fmt.Errorf("workflow runtime missing in context")
}
if err := prepareNodeInputState(localRT, n); err != nil {
return nil, err
}
result, proceed, err := executeNode(runCtx, localRT.args, localRT.runID, n, localRT.state)
if err != nil {
return nil, err
}
localRT.state.NodeOutputs[n.ID] = result
localRT.state.LastOutput = result
if !proceed && !strings.EqualFold(n.Type, "end") {
label := firstNonEmpty(n.Label, n.ID)
if errText := cfgString(result, "error"); errText != "" {
return result, fmt.Errorf("节点「%s」失败: %s", label, errText)
}
return result, fmt.Errorf("节点「%s」未继续执行", label)
}
return result, nil
}
+24
View File
@@ -0,0 +1,24 @@
package workflow
import "errors"
// AwaitingHITLError indicates the workflow paused before a HITL node for human approval.
type AwaitingHITLError struct {
RunID string
NodeID string
NodeLabel string
Prompt string
Reviewer string
}
func (e *AwaitingHITLError) Error() string {
if e == nil {
return "workflow awaiting human approval"
}
return "workflow awaiting human approval at node " + e.NodeID
}
func IsAwaitingHITL(err error) bool {
var target *AwaitingHITLError
return errors.As(err, &target)
}
+186
View File
@@ -0,0 +1,186 @@
package workflow
import (
"fmt"
"regexp"
"strconv"
"strings"
)
var expressionOps = []string{">=", "<=", "==", "!=", " contains ", " matches ", ">", "<"}
var jsonFuncRe = regexp.MustCompile(`^(jsonpath|jq)\((.*),\s*(['"][^'"]+['"])\)$`)
var jsonFuncFindRe = regexp.MustCompile(`(jsonpath|jq)\([^)]*\)`)
var singleTemplateVarRe = regexp.MustCompile(`^\{\{\s*([a-zA-Z0-9_.-]+)\s*\}\}$`)
func validateConditionExpression(expr string) error {
expr = strings.TrimSpace(expr)
if expr == "" {
return fmt.Errorf("条件表达式不能为空")
}
for _, part := range splitBoolExpr(expr, "||") {
for _, atom := range splitBoolExpr(part, "&&") {
if err := validateConditionAtom(atom); err != nil {
return err
}
}
}
return nil
}
func validateConditionAtom(expr string) error {
expr = strings.TrimSpace(expr)
if expr == "" {
return fmt.Errorf("条件表达式存在空片段")
}
if strings.Count(expr, "{{") != strings.Count(expr, "}}") {
return fmt.Errorf("条件表达式模板括号不匹配: %s", expr)
}
if err := validateJSONFunctions(expr); err != nil {
return err
}
if left, right, ok := splitExpressionAtom(expr, " matches "); ok {
if strings.TrimSpace(left) == "" || strings.TrimSpace(right) == "" {
return fmt.Errorf("matches 表达式两侧不能为空: %s", expr)
}
pattern := cleanComparable(resolveStaticTemplate(right))
if _, err := regexp.Compile(pattern); err != nil {
return fmt.Errorf("matches 正则非法: %w", err)
}
return nil
}
for _, op := range expressionOps {
if op == " matches " {
continue
}
if left, right, ok := splitExpressionAtom(expr, op); ok {
if strings.TrimSpace(left) == "" || strings.TrimSpace(right) == "" {
return fmt.Errorf("表达式 %q 两侧不能为空: %s", strings.TrimSpace(op), expr)
}
return nil
}
}
return nil
}
func evalCondition(expr string, state *WorkflowLocalState) bool {
expr = strings.TrimSpace(expr)
if expr == "" {
return true
}
orParts := splitBoolExpr(expr, "||")
for _, orPart := range orParts {
andOK := true
for _, atom := range splitBoolExpr(orPart, "&&") {
if !evalConditionAtom(atom, state) {
andOK = false
break
}
}
if andOK {
return true
}
}
return false
}
func evalConditionAtom(expr string, state *WorkflowLocalState) bool {
expr = strings.TrimSpace(expr)
for _, op := range expressionOps {
if left, right, ok := splitExpressionAtom(expr, op); ok {
left = strings.TrimSpace(fmt.Sprint(resolveExpressionOperand(left, state)))
right = strings.TrimSpace(fmt.Sprint(resolveExpressionOperand(right, state)))
switch strings.TrimSpace(op) {
case "==":
return cleanComparable(left) == cleanComparable(right)
case "!=":
return cleanComparable(left) != cleanComparable(right)
case ">":
return compareNumeric(left, right, func(a, b float64) bool { return a > b })
case ">=":
return compareNumeric(left, right, func(a, b float64) bool { return a >= b })
case "<":
return compareNumeric(left, right, func(a, b float64) bool { return a < b })
case "<=":
return compareNumeric(left, right, func(a, b float64) bool { return a <= b })
case "contains":
return strings.Contains(cleanComparable(left), cleanComparable(right))
case "matches":
matched, _ := regexp.MatchString(cleanComparable(right), cleanComparable(left))
return matched
}
}
}
resolved := strings.TrimSpace(fmt.Sprint(resolveExpressionOperand(expr, state)))
v := strings.ToLower(cleanComparable(resolved))
return v != "" && v != "false" && v != "0" && v != "null"
}
func splitBoolExpr(expr, sep string) []string {
parts := strings.Split(expr, sep)
out := make([]string, 0, len(parts))
for _, part := range parts {
if s := strings.TrimSpace(part); s != "" {
out = append(out, s)
}
}
if len(out) == 0 {
return []string{strings.TrimSpace(expr)}
}
return out
}
func splitExpressionAtom(expr, op string) (string, string, bool) {
if strings.TrimSpace(op) == "contains" || strings.TrimSpace(op) == "matches" {
idx := strings.Index(expr, op)
if idx < 0 {
return "", "", false
}
return expr[:idx], expr[idx+len(op):], true
}
idx := strings.Index(expr, op)
if idx < 0 {
return "", "", false
}
return expr[:idx], expr[idx+len(op):], true
}
func compareNumeric(left, right string, cmp func(float64, float64) bool) bool {
a, errA := strconv.ParseFloat(cleanComparable(left), 64)
b, errB := strconv.ParseFloat(cleanComparable(right), 64)
if errA != nil || errB != nil {
return false
}
return cmp(a, b)
}
func resolveStaticTemplate(s string) string {
return templateVarRe.ReplaceAllString(s, "value")
}
func resolveExpressionOperand(raw string, state *WorkflowLocalState) any {
raw = strings.TrimSpace(raw)
if m := jsonFuncRe.FindStringSubmatch(raw); len(m) == 4 {
inputExpr := strings.TrimSpace(m[2])
path := strings.Trim(m[3], `"'`)
input := resolveExpressionOperand(inputExpr, state)
return evalJSONPathValue(input, path)
}
if m := singleTemplateVarRe.FindStringSubmatch(raw); len(m) == 2 {
return valueFromPath(m[1], state)
}
return resolveTemplate(raw, state)
}
func validateJSONFunctions(expr string) error {
for _, candidate := range jsonFuncFindRe.FindAllString(expr, -1) {
candidate = strings.TrimSpace(candidate)
m := jsonFuncRe.FindStringSubmatch(candidate)
if len(m) != 4 {
return fmt.Errorf("JSONPath/JQ 函数格式应为 jsonpath(value, \"$.path\") 或 jq(value, \".path\")")
}
if err := validateJSONPathSyntax(strings.Trim(m[3], `"'`)); err != nil {
return err
}
}
return nil
}
+107
View File
@@ -0,0 +1,107 @@
package workflow
import (
"context"
"testing"
)
func TestEvalCondition_extendedOperators(t *testing.T) {
state := newWorkflowLocalState(map[string]interface{}{"score": 9, "message": "status: ok"}, "run-expr")
state.LastOutput = map[string]any{"output": "asset-123.example.com"}
tests := []string{
"{{inputs.score}} >= 9",
"{{inputs.message}} contains ok",
"{{previous.output}} matches ^asset-[0-9]+\\.example\\.com$",
"{{inputs.score}} > 5 && {{inputs.message}} contains status",
}
for _, expr := range tests {
if err := validateConditionExpression(expr); err != nil {
t.Fatalf("validate %q: %v", expr, err)
}
if !evalCondition(expr, state) {
t.Fatalf("evalCondition(%q) = false, want true", expr)
}
}
}
func TestEvalCondition_jsonPathAndJQSafeSubset(t *testing.T) {
state := newWorkflowLocalState(map[string]interface{}{
"payload": map[string]any{
"risk": 9,
"items": []any{
map[string]any{"name": "first"},
},
},
}, "run-jsonpath")
state.LastOutput = map[string]any{"output": `{"status":"ok","score":7}`}
tests := []string{
`jsonpath({{inputs.payload}}, "$.risk") >= 8`,
`jq({{inputs.payload}}, ".items[0].name") == first`,
`jsonpath({{previous.output}}, "$.status") == ok`,
}
for _, expr := range tests {
if err := validateConditionExpression(expr); err != nil {
t.Fatalf("validate %q: %v", expr, err)
}
if !evalCondition(expr, state) {
t.Fatalf("evalCondition(%q) = false, want true", expr)
}
}
}
func TestMergeUpstreamOutputs_allMerge(t *testing.T) {
got := mergeUpstreamOutputs(JoinAllMerge, []map[string]any{
{"output": "a", "left": 1},
{"output": "b", "right": 2},
})
if got["kind"] != "join" || got["strategy"] != JoinAllMerge {
t.Fatalf("join metadata = %#v", got)
}
values, ok := got["output"].([]any)
if !ok || len(values) != 2 || values[0] != "a" || values[1] != "b" {
t.Fatalf("merged output = %#v", got["output"])
}
if got["left"] != 1 || got["right"] != 2 {
t.Fatalf("merged fields = %#v", got)
}
}
func TestMergeUpstreamOutputs_firstNonEmpty(t *testing.T) {
got := mergeUpstreamOutputs(JoinFirstNonEmpty, []map[string]any{
{"output": ""},
{"output": "winner"},
})
if got["output"] != "winner" {
t.Fatalf("output = %#v, want winner", got["output"])
}
}
func TestDryRunGraphJSON_simulatesUnsafeNodes(t *testing.T) {
graph := `{
"nodes": [
{"id": "start-1", "type": "start", "label": "开始", "position": {"x": 0, "y": 0}, "config": {}},
{"id": "agent-1", "type": "agent", "label": "Agent", "position": {"x": 0, "y": 80}, "config": {"instruction": "noop", "output_key": "agent_result"}},
{"id": "out-1", "type": "output", "label": "输出", "position": {"x": 0, "y": 160}, "config": {"output_key": "result", "source_binding": {"from": "outputs", "field": "agent_result"}}}
],
"edges": [
{"id": "e1", "source": "start-1", "target": "agent-1"},
{"id": "e2", "source": "agent-1", "target": "out-1"}
]
}`
result, err := DryRunGraphJSON(nilContext(), graph, map[string]any{"message": "hello"})
if err != nil {
t.Fatalf("DryRunGraphJSON: %v", err)
}
if got := result.Outputs["result"]; got != "[dry-run] agent execution skipped" {
t.Fatalf("result output = %#v", got)
}
if len(result.Trace) != 3 {
t.Fatalf("trace len = %d, want 3", len(result.Trace))
}
}
func nilContext() context.Context {
return context.Background()
}
+153
View File
@@ -0,0 +1,153 @@
package workflow
import (
"encoding/json"
"fmt"
"sort"
"strings"
)
type graphDef struct {
Nodes []graphNode `json:"nodes"`
Edges []graphEdge `json:"edges"`
Config map[string]any `json:"config"`
}
type graphNode struct {
ID string `json:"id"`
Type string `json:"type"`
Label string `json:"label"`
Position graphPosition `json:"position"`
Config map[string]any `json:"config"`
}
type graphEdge struct {
ID string `json:"id"`
Source string `json:"source"`
Target string `json:"target"`
Label string `json:"label"`
Config map[string]any `json:"config"`
}
type graphPosition struct {
X float64 `json:"x"`
Y float64 `json:"y"`
}
type graphIndex struct {
nodes map[string]graphNode
outgoing map[string][]graphEdge
incoming map[string][]graphEdge
}
func parseGraph(raw string) (*graphDef, error) {
var g graphDef
if err := json.Unmarshal([]byte(strings.TrimSpace(raw)), &g); err != nil {
return nil, fmt.Errorf("解析工作流图失败: %w", err)
}
if len(g.Nodes) == 0 {
return nil, fmt.Errorf("工作流没有节点")
}
if g.Config == nil {
g.Config = make(map[string]any)
}
return &g, nil
}
func indexGraph(g *graphDef) *graphIndex {
idx := &graphIndex{
nodes: make(map[string]graphNode, len(g.Nodes)),
outgoing: make(map[string][]graphEdge),
incoming: make(map[string][]graphEdge),
}
for _, node := range g.Nodes {
node.ID = strings.TrimSpace(node.ID)
if node.ID == "" {
continue
}
if strings.TrimSpace(node.Type) == "" {
node.Type = "tool"
}
if node.Config == nil {
node.Config = make(map[string]any)
}
idx.nodes[node.ID] = node
}
for _, edge := range g.Edges {
if _, ok := idx.nodes[edge.Source]; !ok {
continue
}
if _, ok := idx.nodes[edge.Target]; !ok {
continue
}
idx.outgoing[edge.Source] = append(idx.outgoing[edge.Source], edge)
idx.incoming[edge.Target] = append(idx.incoming[edge.Target], edge)
}
for source := range idx.outgoing {
sortEdgesByCanvas(idx.outgoing[source], idx.nodes)
}
return idx
}
func sortEdgesByCanvas(edges []graphEdge, nodes map[string]graphNode) {
sort.SliceStable(edges, func(i, j int) bool {
a := nodes[edges[i].Target]
b := nodes[edges[j].Target]
if a.Position.Y != b.Position.Y {
return a.Position.Y < b.Position.Y
}
if a.Position.X != b.Position.X {
return a.Position.X < b.Position.X
}
return edges[i].Target < edges[j].Target
})
}
func sortNodeIDsByCanvas(ids []string, nodes map[string]graphNode) {
sort.SliceStable(ids, func(i, j int) bool {
a := nodes[ids[i]]
b := nodes[ids[j]]
if a.Position.Y != b.Position.Y {
return a.Position.Y < b.Position.Y
}
if a.Position.X != b.Position.X {
return a.Position.X < b.Position.X
}
return ids[i] < ids[j]
})
}
func findStartNodeIDs(idx *graphIndex) []string {
var queue []string
for id, node := range idx.nodes {
if strings.EqualFold(node.Type, "start") {
queue = append(queue, id)
}
}
if len(queue) == 0 {
inDegree := make(map[string]int, len(idx.nodes))
for id := range idx.nodes {
inDegree[id] = 0
}
for _, edges := range idx.outgoing {
for _, edge := range edges {
inDegree[edge.Target]++
}
}
for id, deg := range inDegree {
if deg == 0 {
queue = append(queue, id)
}
}
}
sortNodeIDsByCanvas(queue, idx.nodes)
return queue
}
func branchNodeID(nodeID string) string {
return nodeID + "__eino_branch"
}
func edgeBranchNodeID(nodeID string) string {
return nodeID + "__eino_edge_branch"
}
+119
View File
@@ -0,0 +1,119 @@
package workflow
import (
"context"
"encoding/json"
"fmt"
"strings"
"sync"
"time"
"cyberstrike-ai/internal/database"
)
// HITLDecision is a human decision on a workflow approval node.
type HITLDecision struct {
Approved bool
Comment string
}
var hitlWaiters sync.Map // runID -> chan HITLDecision
func registerHITLWaiter(runID string) chan HITLDecision {
ch := make(chan HITLDecision, 1)
hitlWaiters.Store(runID, ch)
return ch
}
func unregisterHITLWaiter(runID string, ch chan HITLDecision) {
hitlWaiters.CompareAndDelete(runID, ch)
}
// NotifyHITLDecision wakes a streaming workflow run waiting at a HITL node.
// Returns true when an active waiter was signaled.
func NotifyHITLDecision(runID string, decision HITLDecision) bool {
v, ok := hitlWaiters.Load(runID)
if !ok {
return false
}
ch, ok := v.(chan HITLDecision)
if !ok {
return false
}
select {
case ch <- decision:
return true
default:
return true
}
}
func readHITLDecisionFromDB(db *database.DB, runID string) (HITLDecision, bool, error) {
if db == nil {
return HITLDecision{}, false, nil
}
run, err := db.GetWorkflowRun(runID)
if err != nil {
return HITLDecision{}, false, err
}
if run == nil || strings.TrimSpace(run.PendingHITLJSON) == "" {
return HITLDecision{}, false, nil
}
var pending map[string]interface{}
if err := json.Unmarshal([]byte(run.PendingHITLJSON), &pending); err != nil {
return HITLDecision{}, false, nil
}
raw, ok := pending["decision"]
if !ok {
return HITLDecision{}, false, nil
}
decision := strings.ToLower(strings.TrimSpace(fmt.Sprint(raw)))
switch decision {
case "approved", "approve":
comment := ""
if v, ok := pending["comment"]; ok {
comment = strings.TrimSpace(fmt.Sprint(v))
}
return HITLDecision{Approved: true, Comment: comment}, true, nil
case "rejected", "reject":
comment := ""
if v, ok := pending["comment"]; ok {
comment = strings.TrimSpace(fmt.Sprint(v))
}
return HITLDecision{Approved: false, Comment: comment}, true, nil
default:
return HITLDecision{}, false, nil
}
}
func waitWorkflowHITLDecision(ctx context.Context, db *database.DB, runID string) (HITLDecision, error) {
ch := registerHITLWaiter(runID)
defer unregisterHITLWaiter(runID, ch)
return waitWorkflowHITLDecisionWithChannel(ctx, db, runID, ch)
}
func waitWorkflowHITLDecisionWithChannel(ctx context.Context, db *database.DB, runID string, ch chan HITLDecision) (HITLDecision, error) {
if d, ok, err := readHITLDecisionFromDB(db, runID); err != nil {
return HITLDecision{}, err
} else if ok {
return d, nil
}
ticker := time.NewTicker(500 * time.Millisecond)
defer ticker.Stop()
for {
select {
case <-ctx.Done():
return HITLDecision{}, ctx.Err()
case d := <-ch:
return d, nil
case <-ticker.C:
if d, ok, err := readHITLDecisionFromDB(db, runID); err != nil {
return HITLDecision{}, err
} else if ok {
return d, nil
}
}
}
}
+117
View File
@@ -0,0 +1,117 @@
package workflow
import (
"fmt"
"strings"
)
const (
JoinAllMerge = "all_merge"
JoinLastByCanvas = "last_by_canvas"
JoinFirstNonEmpty = "first_non_empty"
JoinFailFast = "fail_fast"
)
var allowedJoinStrategies = map[string]bool{
JoinAllMerge: true,
JoinLastByCanvas: true,
JoinFirstNonEmpty: true,
JoinFailFast: true,
}
func joinStrategy(node graphNode) string {
strategy := strings.ToLower(cfgString(node.Config, "join_strategy"))
if strategy == "" {
return JoinAllMerge
}
return strategy
}
func prepareNodeInputState(rt *workflowRuntime, node graphNode) error {
if rt == nil || rt.idx == nil || rt.state == nil {
return nil
}
incoming := rt.idx.incoming[node.ID]
if len(incoming) <= 1 {
return nil
}
strategy := joinStrategy(node)
if !allowedJoinStrategies[strategy] {
return fmt.Errorf("节点「%s」使用了未知汇聚策略: %s", firstNonEmpty(node.Label, node.ID), strategy)
}
upstreams := make([]map[string]any, 0, len(incoming))
for _, edge := range incoming {
out := rt.state.NodeOutputs[edge.Source]
if out == nil {
continue
}
if isFailedNodeOutput(out) && strategy == JoinFailFast {
return fmt.Errorf("上游节点「%s」失败,汇聚策略 fail_fast 中止", edge.Source)
}
upstreams = append(upstreams, out)
}
if len(upstreams) == 0 {
return nil
}
rt.state.LastOutput = mergeUpstreamOutputs(strategy, upstreams)
return nil
}
func mergeUpstreamOutputs(strategy string, upstreams []map[string]any) map[string]any {
switch strategy {
case JoinLastByCanvas:
return cloneNodeOutput(upstreams[len(upstreams)-1])
case JoinFirstNonEmpty:
for _, out := range upstreams {
if !isEmptyOutputValue(out["output"]) {
return cloneNodeOutput(out)
}
}
return cloneNodeOutput(upstreams[0])
default:
merged := map[string]any{
"kind": "join",
"strategy": strategy,
"upstreams": upstreams,
}
values := make([]any, 0, len(upstreams))
for _, out := range upstreams {
values = append(values, out["output"])
for k, v := range out {
if _, exists := merged[k]; !exists {
merged[k] = v
}
}
}
merged["output"] = values
return merged
}
}
func cloneNodeOutput(in map[string]any) map[string]any {
out := make(map[string]any, len(in))
for k, v := range in {
out[k] = v
}
return out
}
func isEmptyOutputValue(v any) bool {
if v == nil {
return true
}
return strings.TrimSpace(fmt.Sprint(v)) == ""
}
func isFailedNodeOutput(out map[string]any) bool {
if out == nil {
return false
}
if v, ok := out["error"]; ok && strings.TrimSpace(fmt.Sprint(v)) != "" {
return true
}
if v, ok := out["is_error"]; ok {
return strings.EqualFold(fmt.Sprint(v), "true")
}
return false
}
+115
View File
@@ -0,0 +1,115 @@
package workflow
import (
"encoding/json"
"fmt"
"strconv"
"strings"
)
func evalJSONPathValue(input any, path string) any {
path = strings.TrimSpace(path)
if path == "" || path == "$" || path == "." {
return input
}
if strings.HasPrefix(path, "$.") {
path = strings.TrimPrefix(path, "$.")
} else if strings.HasPrefix(path, ".") {
path = strings.TrimPrefix(path, ".")
} else if strings.HasPrefix(path, "$") {
path = strings.TrimPrefix(path, "$")
}
cur := normalizeJSONInput(input)
for _, token := range parseJSONPathTokens(path) {
if token == "" {
continue
}
switch v := cur.(type) {
case map[string]any:
cur = v[token]
case []any:
idx, err := strconv.Atoi(token)
if err != nil || idx < 0 || idx >= len(v) {
return ""
}
cur = v[idx]
default:
return ""
}
}
if cur == nil {
return ""
}
return cur
}
func normalizeJSONInput(input any) any {
switch v := input.(type) {
case string:
var decoded any
if err := json.Unmarshal([]byte(v), &decoded); err == nil {
return decoded
}
return v
case []byte:
var decoded any
if err := json.Unmarshal(v, &decoded); err == nil {
return decoded
}
return string(v)
default:
return input
}
}
func parseJSONPathTokens(path string) []string {
var tokens []string
var buf strings.Builder
for i := 0; i < len(path); i++ {
ch := path[i]
switch ch {
case '.':
if buf.Len() > 0 {
tokens = append(tokens, buf.String())
buf.Reset()
}
case '[':
if buf.Len() > 0 {
tokens = append(tokens, buf.String())
buf.Reset()
}
j := i + 1
for j < len(path) && path[j] != ']' {
j++
}
if j <= len(path) {
token := strings.Trim(path[i+1:j], `"' `)
tokens = append(tokens, token)
i = j
}
default:
buf.WriteByte(ch)
}
}
if buf.Len() > 0 {
tokens = append(tokens, buf.String())
}
return tokens
}
func validateJSONPathSyntax(path string) error {
path = strings.TrimSpace(path)
if path == "" {
return fmt.Errorf("JSONPath 不能为空")
}
if !strings.HasPrefix(path, "$") && !strings.HasPrefix(path, ".") {
return fmt.Errorf("JSONPath/JQ 路径必须以 $ 或 . 开头")
}
if strings.Contains(path, "..") || strings.ContainsAny(path, "*?()|") {
return fmt.Errorf("仅支持安全路径子集,不支持通配符、递归或表达式")
}
if strings.Count(path, "[") != strings.Count(path, "]") {
return fmt.Errorf("JSONPath 方括号不匹配")
}
return nil
}
+57
View File
@@ -0,0 +1,57 @@
package workflow
import (
"fmt"
"strconv"
)
func accumulateWorkflowMetric(state *WorkflowLocalState, key string, delta any) {
if state == nil {
return
}
if state.Metrics == nil {
state.Metrics = make(map[string]any)
}
current := numericMetric(state.Metrics[key])
state.Metrics[key] = current + numericMetric(delta)
}
func numericMetric(v any) float64 {
switch n := v.(type) {
case int:
return float64(n)
case int32:
return float64(n)
case int64:
return float64(n)
case float32:
return float64(n)
case float64:
return n
case string:
f, _ := strconv.ParseFloat(n, 64)
return f
default:
f, _ := strconv.ParseFloat(fmt.Sprint(v), 64)
return f
}
}
func collectAgentMetrics(state *WorkflowLocalState, data interface{}) {
m, ok := data.(map[string]interface{})
if !ok || state == nil {
return
}
for _, key := range []string{"prompt_tokens", "completion_tokens", "total_tokens", "cost", "input_tokens", "output_tokens"} {
if v, ok := m[key]; ok {
accumulateWorkflowMetric(state, key, v)
}
}
if usage, ok := m["usage"].(map[string]interface{}); ok {
for _, key := range []string{"prompt_tokens", "completion_tokens", "total_tokens", "input_tokens", "output_tokens"} {
if v, ok := usage[key]; ok {
accumulateWorkflowMetric(state, key, v)
}
}
}
}
+153
View File
@@ -0,0 +1,153 @@
package workflow
import (
"context"
"encoding/json"
"fmt"
"strings"
"time"
"cyberstrike-ai/internal/database"
"github.com/google/uuid"
)
func executeNode(ctx context.Context, args RunArgs, runID string, node graphNode, state *WorkflowLocalState) (map[string]any, bool, error) {
label := node.Label
if strings.TrimSpace(label) == "" {
label = node.ID
}
nodeRunID := uuid.NewString()
startedAt := time.Now()
incomingCount := 0
if rt := workflowRuntimeFrom(ctx); rt != nil && rt.idx != nil {
incomingCount = len(rt.idx.incoming[node.ID])
}
input := map[string]any{
"nodeId": node.ID,
"nodeType": node.Type,
"label": label,
"inputs": state.Inputs,
"previous": state.LastOutput,
"join": map[string]any{
"strategy": joinStrategy(node),
"incoming": incomingCount,
},
}
inputJSON, _ := json.Marshal(input)
if err := args.DB.CreateWorkflowNodeRun(&database.WorkflowNodeRun{
ID: nodeRunID,
RunID: runID,
NodeID: node.ID,
Status: "running",
InputJSON: string(inputJSON),
StartedAt: startedAt,
}); err != nil {
return nil, false, err
}
if args.Progress != nil {
args.Progress("workflow_node_start", fmt.Sprintf("开始节点:%s", label), map[string]any{
"workflowRunId": runID,
"nodeRunId": nodeRunID,
"nodeId": node.ID,
"nodeType": node.Type,
"label": label,
})
}
result, proceed, status, errText := runBuiltinNode(ctx, args, node, state)
duration := time.Since(startedAt)
if result == nil {
result = map[string]any{}
}
result["duration_ms"] = duration.Milliseconds()
result["finished_at"] = time.Now().Format(time.RFC3339Nano)
result["status"] = status
accumulateWorkflowMetric(state, "node_count", 1)
accumulateWorkflowMetric(state, "duration_ms", duration.Milliseconds())
if strings.EqualFold(node.Type, "tool") {
accumulateWorkflowMetric(state, "tool_call_count", 1)
}
outputJSON, _ := json.Marshal(result)
if err := args.DB.FinishWorkflowNodeRun(nodeRunID, status, string(outputJSON), errText); err != nil {
return nil, false, err
}
if status == "skipped" {
state.Skipped = append(state.Skipped, label)
} else {
state.Executed = append(state.Executed, label)
}
if args.Progress != nil {
progressData := map[string]any{
"workflowRunId": runID,
"nodeRunId": nodeRunID,
"nodeId": node.ID,
"nodeType": node.Type,
"label": label,
"status": status,
"durationMs": duration.Milliseconds(),
"output": result,
}
progressMsg := fmt.Sprintf("节点完成:%s%s", label, status)
if strings.EqualFold(node.Type, "condition") {
matched := false
if v, ok := result["matched"].(bool); ok {
matched = v
}
expr := cfgString(node.Config, "expression")
if matched {
progressMsg = fmt.Sprintf("条件判断:%s → 是", label)
} else {
progressMsg = fmt.Sprintf("条件判断:%s → 否", label)
}
progressData["expression"] = expr
progressData["matched"] = matched
}
args.Progress("workflow_node_result", progressMsg, progressData)
}
state.NodeProceed[node.ID] = proceed
return result, proceed, nil
}
func emitConditionBranchProgress(args RunArgs, runID string, node graphNode, edges []graphEdge, nodes map[string]graphNode, state *WorkflowLocalState) {
if args.Progress == nil || len(edges) == 0 {
return
}
for edgeIdx, edge := range edges {
allowed := edgeAllowed(edge, node, edgeIdx, state)
target := nodes[edge.Target]
targetLabel := strings.TrimSpace(target.Label)
if targetLabel == "" {
targetLabel = edge.Target
}
branchLabel := strings.TrimSpace(edge.Label)
if branchLabel == "" {
switch edgeIdx {
case 0:
branchLabel = "是"
case 1:
branchLabel = "否"
default:
branchLabel = fmt.Sprintf("分支 %d", edgeIdx+1)
}
}
cond := firstNonEmpty(cfgString(edge.Config, "condition"), cfgString(edge.Config, "expression"))
eventType := "workflow_branch_skipped"
msg := fmt.Sprintf("跳过分支「%s」→ %s", branchLabel, targetLabel)
if allowed {
eventType = "workflow_branch_taken"
msg = fmt.Sprintf("执行分支「%s」→ %s", branchLabel, targetLabel)
}
args.Progress(eventType, msg, map[string]any{
"workflowRunId": runID,
"nodeId": node.ID,
"nodeType": node.Type,
"label": node.Label,
"branchLabel": branchLabel,
"targetId": edge.Target,
"targetLabel": targetLabel,
"edgeCondition": cond,
"matched": conditionMatched(state),
})
}
}
+333
View File
@@ -0,0 +1,333 @@
package workflow
import (
"context"
"fmt"
"strings"
"unicode/utf8"
"cyberstrike-ai/internal/agent"
"cyberstrike-ai/internal/config"
"cyberstrike-ai/internal/multiagent"
)
func runBuiltinNode(ctx context.Context, args RunArgs, node graphNode, state *WorkflowLocalState) (map[string]any, bool, string, string) {
cfg := node.Config
switch strings.ToLower(strings.TrimSpace(node.Type)) {
case "start":
return startOutputMap(node, state.Inputs["message"], state.Inputs["conversationId"], state.Inputs["projectId"]), true, "completed", ""
case "condition":
expr := cfgString(cfg, "expression")
ok := evalCondition(expr, state)
return conditionOutputMap(node, expr, ok), true, "completed", ""
case "output":
key := cfgString(cfg, "output_key")
if key == "" {
key = "result"
}
var value any
if v := cfgString(cfg, "static_value"); v != "" {
value = v
} else {
value = resolveOutputSourceBinding(cfg, state)
}
state.Outputs[key] = value
return outputNodeOutputMap(node, key, value), true, "completed", ""
case "end":
value := resolveOutputSourceBinding(cfg, state)
if b, ok := parseFieldBinding(cfg, "result_binding"); ok {
value = resolveBinding(b, state)
}
return endOutputMap(node, value), false, "completed", ""
case "tool":
return runToolNode(ctx, args, node, state)
case "agent":
return runAgentNode(ctx, args, node, state)
case "hitl":
return runHITLNode(args, node, state)
default:
reason := "未知节点类型"
out := outputMap(envelope("unknown", node.ID, node.Type, "skipped", ""), map[string]any{"skipped": true, "reason": reason})
return out, true, "skipped", reason
}
}
func runToolNode(ctx context.Context, args RunArgs, node graphNode, state *WorkflowLocalState) (map[string]any, bool, string, string) {
toolName := cfgString(node.Config, "tool_name")
if toolName == "" {
errText := "工具节点未选择 MCP 工具"
return outputMap(envelope("tool", node.ID, node.Type, "failed", ""), map[string]any{"error": errText}), false, "failed", errText
}
if args.Agent == nil {
errText := "工具节点执行失败:Agent 为空"
return outputMap(envelope("tool", node.ID, node.Type, "failed", ""), map[string]any{"tool_name": toolName, "error": errText}), false, "failed", errText
}
toolArgs, err := resolveToolArguments(node.Config, state)
if err != nil {
errText := fmt.Sprintf("工具参数不是合法 JSON%v", err)
return outputMap(envelope("tool", node.ID, node.Type, "failed", ""), map[string]any{"tool_name": toolName, "error": errText}), false, "failed", errText
}
if args.Progress != nil {
args.Progress("workflow_tool_start", fmt.Sprintf("调用工具:%s", toolName), map[string]any{
"nodeId": node.ID,
"tool": toolName,
"args": toolArgs,
})
}
result, err := args.Agent.ExecuteMCPToolForConversation(ctx, args.ConversationID, toolName, toolArgs)
if err != nil {
errText := err.Error()
return outputMap(envelope("tool", node.ID, node.Type, "failed", ""), map[string]any{"tool_name": toolName, "arguments": toolArgs, "error": errText}), false, "failed", errText
}
output := ""
executionID := ""
isError := false
if result != nil {
output = result.Result
executionID = result.ExecutionID
isError = result.IsError
}
maxToolOutputBytes := config.MultiAgentEinoMiddlewareConfig{}.ReductionMaxLengthForTruncEffective()
if args.AppCfg != nil {
maxToolOutputBytes = args.AppCfg.MultiAgent.EinoMiddleware.ReductionMaxLengthForTruncEffective()
}
output = truncateWorkflowToolOutput(output, maxToolOutputBytes, executionID)
out := toolOutputMap(node, output, toolName, toolArgs, executionID, isError)
if key := cfgString(node.Config, "output_key"); key != "" {
state.Outputs[key] = output
}
if isError {
errText := strings.TrimSpace(output)
if errText == "" {
errText = "工具返回错误"
}
return out, false, "failed", errText
}
return out, true, "completed", ""
}
func truncateWorkflowToolOutput(output string, maxBytes int, executionID string) string {
if maxBytes <= 0 || len(output) <= maxBytes {
return output
}
marker := fmt.Sprintf("\n\n...[workflow tool output truncated; full result is stored in execution %s]...\n\n", strings.TrimSpace(executionID))
if strings.TrimSpace(executionID) == "" {
marker = "\n\n...[workflow tool output truncated; full result remains in the tool execution record]...\n\n"
}
budget := maxBytes - len(marker)
if budget <= 0 {
return marker
}
head := budget / 2
tail := budget - head
for head > 0 && !utf8.RuneStart(output[head]) {
head--
}
tailStart := len(output) - tail
for tailStart < len(output) && !utf8.RuneStart(output[tailStart]) {
tailStart++
}
return output[:head] + marker + output[tailStart:]
}
func runAgentNode(ctx context.Context, args RunArgs, node graphNode, state *WorkflowLocalState) (map[string]any, bool, string, string) {
if args.AppCfg == nil || args.Agent == nil {
errText := "Agent 节点执行失败:应用配置或 Agent 为空"
return outputMap(envelope("agent", node.ID, node.Type, "failed", ""), map[string]any{"error": errText}), false, "failed", errText
}
mode := strings.ToLower(cfgString(node.Config, "agent_mode"))
if mode == "" {
mode = "eino_single"
}
inputSource := resolveNodeInputBinding(node.Config, state)
message := buildAgentNodeMessage(node, state, inputSource)
var result *multiagent.RunResult
var err error
state.SegmentMaxIteration = 0
agentProgress := workflowAgentProgress(args.Progress, state, node)
switch mode {
case "eino_single", "single", "chat":
result, err = multiagent.RunEinoSingleChatModelAgent(
ctx,
args.AppCfg,
&args.AppCfg.MultiAgent,
args.Agent,
args.DB,
args.Logger,
args.ConversationID,
args.ProjectID,
message,
args.History,
args.RoleTools,
agentProgress,
nil,
args.SystemPromptExtra,
)
default:
result, err = multiagent.RunDeepAgent(
ctx,
args.AppCfg,
&args.AppCfg.MultiAgent,
args.Agent,
args.DB,
args.Logger,
args.ConversationID,
args.ProjectID,
message,
args.History,
args.RoleTools,
agentProgress,
args.AgentsMarkdownDir,
mode,
nil,
args.SystemPromptExtra,
)
}
if err != nil {
errText := err.Error()
state.MainIterationOffset += state.SegmentMaxIteration
return outputMap(envelope("agent", node.ID, node.Type, "failed", ""), map[string]any{"mode": mode, "error": errText}), false, "failed", errText
}
state.MainIterationOffset += state.SegmentMaxIteration
response := ""
mcpIDs := []string{}
if result != nil {
response = result.Response
mcpIDs = result.MCPExecutionIDs
}
if args.Progress != nil {
args.Progress("workflow_agent_output", response, map[string]any{
"nodeId": node.ID,
"label": firstNonEmpty(node.Label, node.ID),
"mode": mode,
"inputSource": inputSource,
"inputPreview": truncateWorkflowPreview(inputSource, 500),
"mcpExecutionIds": mcpIDs,
})
}
if key := cfgString(node.Config, "output_key"); key != "" {
state.Outputs[key] = response
}
return agentOutputMap(node, response, mode, mcpIDs), true, "completed", ""
}
func buildAgentNodeMessage(node graphNode, state *WorkflowLocalState, upstreamInput string) string {
instruction := strings.TrimSpace(cfgString(node.Config, "instruction"))
upstreamInput = strings.TrimSpace(upstreamInput)
if instruction == "" {
if upstreamInput != "" {
return fmt.Sprintf("请基于上游节点输出继续处理:\n%s", upstreamInput)
}
return fmt.Sprintf("请基于上游节点输出继续处理:\n%v", state.LastOutput["output"])
}
if upstreamInput == "" {
return instruction
}
return strings.TrimSpace(fmt.Sprintf("上游输入:\n%s\n\n节点指令:\n%s", upstreamInput, instruction))
}
func workflowAgentProgress(progress agent.ProgressCallback, state *WorkflowLocalState, node graphNode) agent.ProgressCallback {
if progress == nil {
return nil
}
return func(eventType, message string, data interface{}) {
switch eventType {
case "response_start", "response_delta", "response", "done":
return
default:
enrichWorkflowAgentEventData(data, state, node)
collectAgentMetrics(state, data)
if eventType == "iteration" {
applyWorkflowMainIterationOffset(data, state)
}
progress(eventType, message, data)
}
}
}
func enrichWorkflowAgentEventData(data interface{}, state *WorkflowLocalState, node graphNode) {
m, ok := data.(map[string]interface{})
if !ok || m == nil {
return
}
if node.ID != "" {
m["workflowNodeId"] = node.ID
}
if state != nil && strings.TrimSpace(state.WorkflowRunID) != "" {
m["workflowRunId"] = state.WorkflowRunID
}
}
func applyWorkflowMainIterationOffset(data interface{}, state *WorkflowLocalState) {
if state == nil {
return
}
m, ok := data.(map[string]interface{})
if !ok || m == nil {
return
}
scope, _ := m["einoScope"].(string)
if strings.TrimSpace(scope) != "main" {
return
}
raw := iterationNumberFromProgressData(m)
if raw <= 0 {
return
}
if raw > state.SegmentMaxIteration {
state.SegmentMaxIteration = raw
}
m["iteration"] = raw + state.MainIterationOffset
}
func iterationNumberFromProgressData(m map[string]interface{}) int {
switch v := m["iteration"].(type) {
case int:
return v
case int32:
return int(v)
case int64:
return int(v)
case float64:
return int(v)
case float32:
return int(v)
default:
return 0
}
}
func runHITLNode(args RunArgs, node graphNode, state *WorkflowLocalState) (map[string]any, bool, string, string) {
prompt := resolveHITLPromptBinding(node.Config, state)
reviewer := cfgString(node.Config, "reviewer")
if reviewer == "" {
reviewer = "human"
}
approved := true
if state != nil && state.Inputs != nil {
if v, ok := state.Inputs["_hitl_approved"]; ok {
approved = fmt.Sprint(v) == "true"
}
}
if !approved {
reason := "人工审批已拒绝"
if state != nil && state.Inputs != nil {
if v, ok := state.Inputs["_hitl_comment"]; ok {
if s := strings.TrimSpace(fmt.Sprint(v)); s != "" {
reason = s
}
}
}
return hitlOutputMap(node, "failed", "", prompt, reviewer, false), false, "failed", reason
}
if args.Progress != nil {
args.Progress("workflow_hitl_checkpoint", "人工确认节点已通过", map[string]any{
"nodeId": node.ID,
"prompt": prompt,
"reviewer": reviewer,
"mode": "interactive",
"approved": true,
})
}
return hitlOutputMap(node, "completed", prompt, prompt, reviewer, true), true, "completed", ""
}
+98
View File
@@ -0,0 +1,98 @@
package workflowpackage
import (
"archive/zip"
"bytes"
"encoding/json"
"fmt"
"path"
"strings"
"time"
)
// Export builds a deterministic, human-readable single-workflow package.
func Export(source Document) ([]byte, ExportMetadata, error) {
source.ID = strings.TrimSpace(source.ID)
source.Name = strings.TrimSpace(source.Name)
if source.ID == "" || source.Name == "" || source.Version <= 0 || !safePackageWorkflowID(source.ID) {
return nil, ExportMetadata{}, fmt.Errorf("workflow id, name and version are required")
}
if strings.TrimSpace(source.GraphJSON) == "" {
return nil, ExportMetadata{}, fmt.Errorf("workflow graph_json is required")
}
contentHash, graphHash, payload, err := DocumentHashes(source)
if err != nil {
return nil, ExportMetadata{}, err
}
workflowPath := path.Join("workflows", source.ID+".json")
createdAt := source.UpdatedAt.UTC()
if createdAt.IsZero() {
createdAt = time.Unix(0, 0).UTC()
}
manifest := Manifest{
PackageFormat: PackageFormat,
FormatVersion: FormatVersion,
PackageID: "pkg_" + strings.TrimPrefix(contentHash, "sha256:")[:16],
CreatedAt: createdAt.Format(time.RFC3339),
Items: []ManifestItem{{
Type: "workflow",
Path: workflowPath,
SourceID: source.ID,
SourceRevision: source.Version,
ContentHash: contentHash,
GraphHash: graphHash,
}},
}
manifestBytes, err := json.Marshal(manifest)
if err != nil {
return nil, ExportMetadata{}, fmt.Errorf("marshal manifest: %w", err)
}
checksums := fmt.Sprintf("%s manifest.json\n%s %s\n", strings.TrimPrefix(sha256Prefixed(manifestBytes), "sha256:"), strings.TrimPrefix(contentHash, "sha256:"), workflowPath)
var out bytes.Buffer
zw := zip.NewWriter(&out)
for _, entry := range []struct {
name string
data []byte
}{
{name: "checksums.sha256", data: []byte(checksums)},
{name: "manifest.json", data: manifestBytes},
{name: workflowPath, data: payload},
} {
header := &zip.FileHeader{Name: entry.name, Method: zip.Store}
header.SetModTime(time.Unix(0, 0).UTC())
writer, err := zw.CreateHeader(header)
if err != nil {
return nil, ExportMetadata{}, fmt.Errorf("write %s: %w", entry.name, err)
}
if _, err := writer.Write(entry.data); err != nil {
return nil, ExportMetadata{}, fmt.Errorf("write %s: %w", entry.name, err)
}
}
if err := zw.Close(); err != nil {
return nil, ExportMetadata{}, fmt.Errorf("close package: %w", err)
}
pkg := out.Bytes()
return pkg, ExportMetadata{
PackageHash: sha256Prefixed(pkg),
ContentHash: contentHash,
GraphHash: graphHash,
SourceRevision: source.Version,
FileName: source.ID + ".csapkg.zip",
}, nil
}
// DocumentHashes returns the canonical package item and graph hashes together
// with the canonical item bytes used by export and inspection persistence.
func DocumentHashes(source Document) (string, string, []byte, error) {
graph, err := canonicalJSON([]byte(source.GraphJSON))
if err != nil {
return "", "", nil, fmt.Errorf("canonicalize graph_json: %w", err)
}
source.GraphJSON = string(graph)
payload, err := json.Marshal(source)
if err != nil {
return "", "", nil, fmt.Errorf("marshal workflow payload: %w", err)
}
return sha256Prefixed(payload), sha256Prefixed(graph), payload, nil
}
@@ -0,0 +1,58 @@
package workflowpackage
import (
"archive/zip"
"bytes"
"strings"
"testing"
"time"
)
func testDocument() Document {
return Document{
ID: "web-src-hunting",
Name: "Web SRC 猎洞",
Description: "面向 SRC Web 资产的侦察与漏洞候选流程",
Version: 18,
Enabled: true,
GraphJSON: `{"nodes":[{"id":"start-1","type":"start","label":"开始","position":{"x":0,"y":0},"config":{}},{"id":"out-1","type":"output","label":"输出","position":{"x":0,"y":120},"config":{"output_key":"result","source_binding":{"from":"inputs","field":"message"}}}],"edges":[{"id":"e1","source":"start-1","target":"out-1"}],"config":{"schema_version":1}}`,
UpdatedAt: time.Date(2026, 7, 13, 10, 0, 0, 0, time.UTC),
}
}
func TestExportIsDeterministicAndSelfDescribing(t *testing.T) {
first, firstMeta, err := Export(testDocument())
if err != nil {
t.Fatalf("first export: %v", err)
}
second, secondMeta, err := Export(testDocument())
if err != nil {
t.Fatalf("second export: %v", err)
}
if !bytes.Equal(first, second) {
t.Fatal("identical document must produce byte-identical package")
}
if firstMeta.PackageHash != secondMeta.PackageHash || !strings.HasPrefix(firstMeta.PackageHash, "sha256:") {
t.Fatalf("unexpected deterministic package hash: %#v / %#v", firstMeta, secondMeta)
}
zr, err := zip.NewReader(bytes.NewReader(first), int64(len(first)))
if err != nil {
t.Fatalf("open package: %v", err)
}
if len(zr.File) != 3 {
t.Fatalf("zip entry count = %d, want 3", len(zr.File))
}
wantNames := []string{"checksums.sha256", "manifest.json", "workflows/web-src-hunting.json"}
for i, f := range zr.File {
if f.Name != wantNames[i] {
t.Fatalf("entry %d = %q, want %q", i, f.Name, wantNames[i])
}
}
if firstMeta.SourceRevision != 18 {
t.Fatalf("source revision = %d, want 18", firstMeta.SourceRevision)
}
if !strings.HasPrefix(firstMeta.ContentHash, "sha256:") || !strings.HasPrefix(firstMeta.GraphHash, "sha256:") {
t.Fatalf("content/graph hashes must be sha256: %#v", firstMeta)
}
}
+225
View File
@@ -0,0 +1,225 @@
package workflowpackage
import (
"archive/zip"
"bytes"
"context"
"encoding/json"
"errors"
"fmt"
"io"
"os"
"path"
"strings"
"unicode"
)
const (
MaxArchiveBytes = 10 << 20
MaxExtractedBytes = 20 << 20
)
// PackageError contains only a contract error code and safe, client-facing fields.
type PackageError struct {
Code string
Message string
Details map[string]any
}
func (e *PackageError) Error() string { return e.Code + ": " + e.Message }
func packageError(code, message string) error {
return &PackageError{Code: code, Message: message}
}
// ErrorCode returns a package contract code without exposing internal errors.
func ErrorCode(err error) string {
var target *PackageError
if errors.As(err, &target) {
return target.Code
}
return ""
}
type InspectionResult struct {
PackageHash string
Manifest Manifest
Document Document
ContentHash string
GraphHash string
NodeCount int
EdgeCount int
}
// InspectArchive verifies an archive without executing any package content.
// validateGraph is injected by the application so this format package has no
// dependency on the workflow runtime or database driver.
func InspectArchive(ctx context.Context, archive []byte, validateGraph func(context.Context, string) error) (*InspectionResult, error) {
if len(archive) == 0 {
return nil, packageError("WFPKG_FILE_REQUIRED", "必须上传工作流包文件")
}
if len(archive) > MaxArchiveBytes {
return nil, packageError("WFPKG_FILE_TOO_LARGE", "工作流包文件超过大小限制")
}
zr, err := zip.NewReader(bytes.NewReader(archive), int64(len(archive)))
if err != nil {
return nil, packageError("WFPKG_INVALID_ARCHIVE", "工作流包不是有效 ZIP 文件")
}
entries := make(map[string][]byte, len(zr.File))
var extracted int64
for _, file := range zr.File {
if !safeArchivePath(file.Name) || file.FileInfo().IsDir() || file.FileInfo().Mode()&os.ModeSymlink != 0 {
return nil, packageError("WFPKG_INVALID_ARCHIVE", "工作流包包含不安全文件路径")
}
if _, exists := entries[file.Name]; exists {
return nil, packageError("WFPKG_INVALID_ARCHIVE", "工作流包包含重复文件")
}
if file.UncompressedSize64 > MaxExtractedBytes || extracted+int64(file.UncompressedSize64) > MaxExtractedBytes {
return nil, packageError("WFPKG_INVALID_ARCHIVE", "工作流包解压后超过大小限制")
}
reader, err := file.Open()
if err != nil {
return nil, packageError("WFPKG_INVALID_ARCHIVE", "无法读取工作流包文件")
}
data, readErr := io.ReadAll(io.LimitReader(reader, int64(MaxExtractedBytes)-extracted+1))
closeErr := reader.Close()
if readErr != nil || closeErr != nil || len(data) > MaxExtractedBytes-int(extracted) {
return nil, packageError("WFPKG_INVALID_ARCHIVE", "工作流包解压后超过大小限制")
}
extracted += int64(len(data))
entries[file.Name] = data
}
manifestRaw, hasManifest := entries["manifest.json"]
checksumsRaw, hasChecksums := entries["checksums.sha256"]
if !hasManifest || !hasChecksums {
return nil, packageError("WFPKG_UNSUPPORTED_FORMAT", "工作流包缺少必需文件")
}
manifest, err := parseManifest(manifestRaw)
if err != nil {
return nil, err
}
if len(manifest.Items) != 1 || manifest.Items[0].Type != "workflow" {
return nil, packageError("WFPKG_MULTIPLE_WORKFLOWS", "工作流包必须且只能包含一个工作流")
}
item := manifest.Items[0]
workflowRaw, exists := entries[item.Path]
if !exists || !safeWorkflowPath(item.Path) || len(entries) != 3 {
return nil, packageError("WFPKG_INVALID_ARCHIVE", "工作流包包含未声明文件")
}
checksums, err := parseChecksums(checksumsRaw)
if err != nil {
return nil, err
}
if len(checksums) != 2 || checksums["manifest.json"] != sha256Prefixed(manifestRaw) || checksums[item.Path] != sha256Prefixed(workflowRaw) {
return nil, packageError("WFPKG_CHECKSUM_MISMATCH", "工作流包校验和不匹配")
}
if item.ContentHash != sha256Prefixed(workflowRaw) || !validHash(item.ContentHash) || !validHash(item.GraphHash) {
return nil, packageError("WFPKG_CHECKSUM_MISMATCH", "工作流包内容校验和不匹配")
}
doc, err := parseDocument(workflowRaw)
if err != nil {
return nil, err
}
if !safePackageWorkflowID(doc.ID) || doc.ID != item.SourceID || doc.Version != item.SourceRevision {
return nil, packageError("WFPKG_INVALID_MANIFEST", "工作流包清单与工作流内容不一致")
}
graph, err := canonicalJSON([]byte(doc.GraphJSON))
if err != nil || item.GraphHash != sha256Prefixed(graph) {
return nil, packageError("WFPKG_CHECKSUM_MISMATCH", "工作流图校验和不匹配")
}
if validateGraph == nil || validateGraph(ctx, string(graph)) != nil {
return nil, packageError("WFPKG_WORKFLOW_INVALID", "工作流图校验失败")
}
var graphShape struct {
Nodes []json.RawMessage `json:"nodes"`
Edges []json.RawMessage `json:"edges"`
}
if err := json.Unmarshal(graph, &graphShape); err != nil {
return nil, packageError("WFPKG_WORKFLOW_INVALID", "工作流图不是有效 JSON")
}
return &InspectionResult{
PackageHash: sha256Prefixed(archive),
Manifest: manifest,
Document: doc,
ContentHash: item.ContentHash,
GraphHash: item.GraphHash,
NodeCount: len(graphShape.Nodes),
EdgeCount: len(graphShape.Edges),
}, nil
}
func safeArchivePath(name string) bool {
return name != "" && !strings.Contains(name, `\`) && !strings.HasPrefix(name, "/") && path.Clean(name) == name && !strings.HasPrefix(name, "../") && name != ".."
}
func safeWorkflowPath(name string) bool {
rest := strings.TrimPrefix(name, "workflows/")
return safeArchivePath(name) && strings.HasPrefix(name, "workflows/") && rest != "" && !strings.Contains(rest, "/") && strings.HasSuffix(rest, ".json")
}
func safePackageWorkflowID(id string) bool {
if id == "" || strings.ContainsAny(id, `/\`) {
return false
}
for _, r := range id {
if unicode.IsControl(r) {
return false
}
}
return true
}
func parseManifest(raw []byte) (Manifest, error) {
var manifest Manifest
dec := json.NewDecoder(bytes.NewReader(raw))
dec.DisallowUnknownFields()
if err := dec.Decode(&manifest); err != nil {
return Manifest{}, packageError("WFPKG_INVALID_MANIFEST", "工作流包清单格式无效")
}
if err := consumeJSONEnd(dec); err != nil || manifest.PackageFormat != PackageFormat || manifest.FormatVersion != FormatVersion || strings.TrimSpace(manifest.PackageID) == "" || len(manifest.Items) == 0 {
return Manifest{}, packageError("WFPKG_INVALID_MANIFEST", "工作流包清单格式不受支持")
}
return manifest, nil
}
func parseDocument(raw []byte) (Document, error) {
var doc Document
dec := json.NewDecoder(bytes.NewReader(raw))
dec.DisallowUnknownFields()
if err := dec.Decode(&doc); err != nil || consumeJSONEnd(dec) != nil {
return Document{}, packageError("WFPKG_WORKFLOW_INVALID", "工作流定义格式无效")
}
doc.ID = strings.TrimSpace(doc.ID)
doc.Name = strings.TrimSpace(doc.Name)
if doc.ID == "" || doc.Name == "" || doc.Version <= 0 || strings.TrimSpace(doc.GraphJSON) == "" {
return Document{}, packageError("WFPKG_WORKFLOW_INVALID", "工作流定义缺少必需字段")
}
return doc, nil
}
func consumeJSONEnd(dec *json.Decoder) error {
var extra any
if err := dec.Decode(&extra); err != io.EOF {
if err == nil {
return fmt.Errorf("multiple JSON values")
}
return err
}
return nil
}
func parseChecksums(raw []byte) (map[string]string, error) {
entries := make(map[string]string)
for _, line := range strings.Split(strings.TrimSpace(string(raw)), "\n") {
parts := strings.SplitN(strings.TrimSpace(line), " ", 2)
if len(parts) != 2 || !validHash("sha256:"+parts[0]) || !safeArchivePath(parts[1]) {
return nil, packageError("WFPKG_CHECKSUM_MISMATCH", "工作流包校验和格式无效")
}
if _, exists := entries[parts[1]]; exists {
return nil, packageError("WFPKG_CHECKSUM_MISMATCH", "工作流包校验和重复")
}
entries[parts[1]] = "sha256:" + parts[0]
}
return entries, nil
}
+154
View File
@@ -0,0 +1,154 @@
package workflowpackage
import (
"archive/zip"
"bytes"
"context"
"errors"
"fmt"
"io"
"os"
"testing"
)
func TestInspectArchiveAcceptsSingleVerifiedWorkflow(t *testing.T) {
pkg, meta, err := Export(testDocument())
if err != nil {
t.Fatal(err)
}
result, err := InspectArchive(context.Background(), pkg, func(context.Context, string) error { return nil })
if err != nil {
t.Fatalf("InspectArchive: %v", err)
}
if result.PackageHash != meta.PackageHash || result.Document.ID != "web-src-hunting" {
t.Fatalf("unexpected inspection: %#v", result)
}
if result.NodeCount != 2 || result.EdgeCount != 1 {
t.Fatalf("counts = %d/%d, want 2/1", result.NodeCount, result.EdgeCount)
}
}
func TestInspectArchiveRejectsUnsafeArchiveShapes(t *testing.T) {
valid, _, err := Export(testDocument())
if err != nil {
t.Fatal(err)
}
cases := []struct {
name string
archive []byte
}{
{name: "duplicate entry", archive: appendZipEntry(t, valid, "manifest.json", []byte(`{}`), 0)},
{name: "path traversal", archive: appendZipEntry(t, valid, "../payload.json", []byte(`{}`), 0)},
{name: "symlink", archive: appendZipEntry(t, valid, "workflows/link.json", []byte("target"), 0o120777)},
{name: "undeclared file", archive: appendZipEntry(t, valid, "notes.txt", []byte("not allowed"), 0)},
}
for _, tc := range cases {
t.Run(tc.name, func(t *testing.T) {
_, err := InspectArchive(context.Background(), tc.archive, func(context.Context, string) error { return nil })
if ErrorCode(err) != "WFPKG_INVALID_ARCHIVE" {
t.Fatalf("code = %q, err = %v", ErrorCode(err), err)
}
})
}
}
func TestInspectArchiveRejectsChecksumMismatchAndInvalidWorkflow(t *testing.T) {
pkg, _, err := Export(testDocument())
if err != nil {
t.Fatal(err)
}
badChecksum := replaceZipEntry(t, pkg, "checksums.sha256", []byte("00 manifest.json\n"), 0)
if _, err := InspectArchive(context.Background(), badChecksum, func(context.Context, string) error { return nil }); ErrorCode(err) != "WFPKG_CHECKSUM_MISMATCH" {
t.Fatalf("checksum code = %q, err = %v", ErrorCode(err), err)
}
if _, err := InspectArchive(context.Background(), pkg, func(context.Context, string) error { return errors.New("invalid graph") }); ErrorCode(err) != "WFPKG_WORKFLOW_INVALID" {
t.Fatalf("graph code = %q, err = %v", ErrorCode(err), err)
}
}
func appendZipEntry(t *testing.T, archive []byte, name string, data []byte, mode os.FileMode) []byte {
t.Helper()
return rewriteZip(t, archive, func(zw *zip.Writer) error {
h := &zip.FileHeader{Name: name, Method: zip.Store}
h.SetMode(mode)
w, err := zw.CreateHeader(h)
if err != nil {
return err
}
_, err = w.Write(data)
return err
})
}
func replaceZipEntry(t *testing.T, archive []byte, name string, data []byte, mode os.FileMode) []byte {
t.Helper()
zr, err := zip.NewReader(bytes.NewReader(archive), int64(len(archive)))
if err != nil {
t.Fatal(err)
}
var out bytes.Buffer
zw := zip.NewWriter(&out)
for _, f := range zr.File {
if f.Name == name {
h := &zip.FileHeader{Name: name, Method: zip.Store}
h.SetMode(mode)
w, err := zw.CreateHeader(h)
if err != nil {
t.Fatal(err)
}
if _, err := w.Write(data); err != nil {
t.Fatal(err)
}
continue
}
r, err := f.Open()
if err != nil {
t.Fatal(err)
}
h := &zip.FileHeader{Name: f.Name, Method: zip.Store}
w, err := zw.CreateHeader(h)
if err == nil {
_, err = io.Copy(w, r)
}
_ = r.Close()
if err != nil {
t.Fatal(err)
}
}
if err := zw.Close(); err != nil {
t.Fatal(err)
}
return out.Bytes()
}
func rewriteZip(t *testing.T, archive []byte, appendEntry func(*zip.Writer) error) []byte {
t.Helper()
zr, err := zip.NewReader(bytes.NewReader(archive), int64(len(archive)))
if err != nil {
t.Fatal(err)
}
var out bytes.Buffer
zw := zip.NewWriter(&out)
for _, f := range zr.File {
r, err := f.Open()
if err != nil {
t.Fatal(err)
}
h := &zip.FileHeader{Name: f.Name, Method: zip.Store}
w, err := zw.CreateHeader(h)
if err == nil {
_, err = io.Copy(w, r)
}
_ = r.Close()
if err != nil {
t.Fatal(err)
}
}
if err := appendEntry(zw); err != nil {
t.Fatal(fmt.Errorf("append entry: %w", err))
}
if err := zw.Close(); err != nil {
t.Fatal(err)
}
return out.Bytes()
}
+78
View File
@@ -0,0 +1,78 @@
package workflowpackage
import (
"crypto/sha256"
"encoding/hex"
"encoding/json"
"fmt"
"strings"
"time"
)
const (
PackageFormat = "cyberstrikeai.workflow-package"
FormatVersion = "1.0"
)
// Document is the single non-executable workflow definition carried by a package.
// Version is the source instance revision and is never applied as a target version.
type Document struct {
ID string `json:"id"`
Name string `json:"name"`
Description string `json:"description,omitempty"`
Version int `json:"version"`
GraphJSON string `json:"graph_json"`
Enabled bool `json:"enabled"`
UpdatedAt time.Time `json:"-"`
}
type Manifest struct {
PackageFormat string `json:"package_format"`
FormatVersion string `json:"format_version"`
PackageID string `json:"package_id"`
CreatedAt string `json:"created_at"`
Items []ManifestItem `json:"items"`
}
type ManifestItem struct {
Type string `json:"type"`
Path string `json:"path"`
SourceID string `json:"source_id"`
SourceRevision int `json:"source_revision"`
ContentHash string `json:"content_hash"`
GraphHash string `json:"graph_hash"`
}
type ExportMetadata struct {
PackageHash string
ContentHash string
GraphHash string
SourceRevision int
FileName string
}
func canonicalJSON(raw []byte) ([]byte, error) {
dec := json.NewDecoder(strings.NewReader(string(raw)))
dec.UseNumber()
var value any
if err := dec.Decode(&value); err != nil {
return nil, err
}
if dec.More() {
return nil, fmt.Errorf("extra JSON values")
}
return json.Marshal(value)
}
func sha256Prefixed(b []byte) string {
sum := sha256.Sum256(b)
return "sha256:" + hex.EncodeToString(sum[:])
}
func validHash(value string) bool {
if !strings.HasPrefix(value, "sha256:") || len(value) != len("sha256:")+64 {
return false
}
_, err := hex.DecodeString(strings.TrimPrefix(value, "sha256:"))
return err == nil && value == strings.ToLower(value)
}
+223
View File
@@ -0,0 +1,223 @@
package workflow
import (
"context"
"encoding/json"
"fmt"
"strings"
"time"
"cyberstrike-ai/internal/config"
"cyberstrike-ai/internal/database"
"github.com/google/uuid"
"go.uber.org/zap"
)
// ShouldAutoRunRoleWorkflow returns true when a role explicitly binds a workflow
// and does not turn it off. Empty policy defaults to auto to keep role UX simple.
func ShouldAutoRunRoleWorkflow(role config.RoleConfig) bool {
if strings.TrimSpace(role.WorkflowID) == "" {
return false
}
policy := strings.ToLower(strings.TrimSpace(role.WorkflowPolicy))
return policy == "" || policy == "auto"
}
// RunRoleBoundWorkflow executes the persisted role-bound workflow via cached Eino Workflow.
func RunRoleBoundWorkflow(ctx context.Context, args RunArgs) (*RunResult, error) {
if args.DB == nil {
return nil, fmt.Errorf("workflow db is nil")
}
workflowID := strings.TrimSpace(args.Role.WorkflowID)
if workflowID == "" {
return nil, fmt.Errorf("角色未绑定工作流")
}
wf, err := args.DB.GetWorkflowDefinition(workflowID)
if err != nil {
return nil, err
}
if wf == nil {
return nil, fmt.Errorf("角色绑定的工作流不存在: %s", workflowID)
}
if !wf.Enabled {
return nil, fmt.Errorf("角色绑定的工作流已禁用: %s", workflowID)
}
select {
case <-ctx.Done():
return nil, ctx.Err()
default:
}
runID := uuid.NewString()
input := map[string]interface{}{
"message": args.UserMessage,
"conversationId": args.ConversationID,
"projectId": args.ProjectID,
"role": args.Role.Name,
"workflowId": wf.ID,
"workflowVersion": wf.Version,
}
inputJSON, _ := json.Marshal(input)
run := &database.WorkflowRun{
ID: runID,
WorkflowID: wf.ID,
WorkflowVersion: wf.Version,
ConversationID: args.ConversationID,
ProjectID: args.ProjectID,
RoleID: args.Role.Name,
Status: "running",
InputJSON: string(inputJSON),
StartedAt: time.Now(),
}
if err := args.DB.CreateWorkflowRun(run); err != nil {
return nil, err
}
if args.Progress != nil {
args.Progress("workflow_start", fmt.Sprintf("开始运行流程「%s」", wf.Name), map[string]interface{}{
"workflowId": wf.ID,
"workflowName": wf.Name,
"workflowVersion": wf.Version,
"workflowRunId": runID,
"conversationId": args.ConversationID,
"engine": "eino_workflow",
})
}
graph, err := parseGraph(wf.GraphJSON)
if err != nil {
_ = args.DB.FinishWorkflowRun(runID, "failed", "", err.Error())
return nil, err
}
state := newWorkflowLocalState(input, runID)
streaming := args.Progress != nil
resuming := false
for {
_, err := invokeEinoGraph(ctx, args, runID, wf.ID, wf.Version, graph, state, resuming)
if err == nil {
break
}
if !IsAwaitingHITL(err) {
_ = args.DB.FinishWorkflowRun(runID, "failed", "", err.Error())
return nil, err
}
hitl := err.(*AwaitingHITLError)
partial := map[string]interface{}{
"workflowId": wf.ID,
"workflowName": wf.Name,
"workflowVersion": wf.Version,
"workflowRunId": runID,
"status": "awaiting_hitl",
"outputs": state.Outputs,
"executedNodes": state.Executed,
"skippedNodes": state.Skipped,
"pendingHitl": map[string]interface{}{
"nodeId": hitl.NodeID,
"label": hitl.NodeLabel,
"prompt": hitl.Prompt,
},
"engine": "eino_workflow",
}
partialJSON, _ := json.Marshal(partial)
_ = args.DB.SetWorkflowRunAwaitingHITL(runID, hitl.NodeID, string(partialJSON))
response := fmt.Sprintf("工作流「%s」已在节点「%s」暂停,等待人工审批。\n运行 ID:%s", wf.Name, firstNonEmpty(hitl.NodeLabel, hitl.NodeID), runID)
if args.Progress != nil {
args.Progress("workflow_paused", response, map[string]interface{}{
"workflowRunId": runID,
"status": "awaiting_hitl",
"nodeId": hitl.NodeID,
"resumeApi": fmt.Sprintf("/api/workflows/runs/%s/resume", runID),
})
}
if !streaming {
return &RunResult{
Response: response,
RunID: runID,
Status: "awaiting_hitl",
AwaitingHITL: true,
}, nil
}
ch := registerHITLWaiter(runID)
decision, waitErr := waitWorkflowHITLDecisionWithChannel(ctx, args.DB, runID, ch)
unregisterHITLWaiter(runID, ch)
if waitErr != nil {
_ = args.DB.FinishWorkflowRun(runID, "cancelled", "", waitErr.Error())
return nil, waitErr
}
if !decision.Approved {
errText := strings.TrimSpace(decision.Comment)
if errText == "" {
errText = "人工审批拒绝"
}
_ = args.DB.FinishWorkflowRun(runID, "rejected", "", errText)
rejectResponse := fmt.Sprintf("工作流已在审批节点「%s」被拒绝。", firstNonEmpty(hitl.NodeLabel, hitl.NodeID))
if args.Progress != nil {
args.Progress("workflow_hitl_rejected", rejectResponse, map[string]interface{}{
"workflowRunId": runID,
"nodeId": hitl.NodeID,
"comment": errText,
})
}
return &RunResult{
Response: rejectResponse,
RunID: runID,
Status: "rejected",
}, nil
}
if args.Progress != nil {
args.Progress("workflow_hitl_resumed", "人工审批已通过,继续执行", map[string]interface{}{
"workflowRunId": runID,
"nodeId": hitl.NodeID,
"comment": decision.Comment,
})
}
if state.Inputs == nil {
state.Inputs = map[string]any{}
}
state.Inputs["_hitl_approved"] = true
state.Inputs["_hitl_comment"] = decision.Comment
state.Inputs["_hitl_node_id"] = hitl.NodeID
_ = args.DB.SetWorkflowRunStatus(runID, "running")
resuming = true
}
output := map[string]interface{}{
"workflowId": wf.ID,
"workflowName": wf.Name,
"workflowVersion": wf.Version,
"workflowRunId": runID,
"status": "completed",
"outputs": state.Outputs,
"metrics": state.Metrics,
"executedNodes": state.Executed,
"skippedNodes": state.Skipped,
"engine": "eino_workflow",
}
outputJSON, _ := json.Marshal(output)
response := renderWorkflowResponse(args.Role.Name, wf.Name, wf.Version, runID, state)
if err := args.DB.FinishWorkflowRun(runID, "completed", string(outputJSON), ""); err != nil {
return nil, err
}
if args.Progress != nil {
args.Progress("workflow_done", fmt.Sprintf("流程「%s」运行完成", wf.Name), map[string]interface{}{
"workflowRunId": runID,
"workflowId": wf.ID,
"outputs": state.Outputs,
"metrics": state.Metrics,
"response": response,
"engine": "eino_workflow",
})
}
if args.Logger != nil {
args.Logger.Info("role-bound workflow completed",
zap.String("workflow_id", wf.ID),
zap.String("workflow_run_id", runID),
zap.String("conversation_id", args.ConversationID),
zap.String("role", args.Role.Name),
zap.String("engine", "eino_workflow"),
)
}
return &RunResult{Response: response, RunID: runID, Status: "completed"}, nil
}
+214
View File
@@ -0,0 +1,214 @@
package workflow
import (
"fmt"
"regexp"
"sort"
"strings"
"github.com/cloudwego/eino/schema"
)
func init() {
schema.RegisterName[*WorkflowLocalState]("_cyberstrike_workflow_local_state")
schema.RegisterName[NodeOutputEnvelope]("_cyberstrike_workflow_node_output_envelope")
schema.RegisterName[StartOutput]("_cyberstrike_workflow_start_output")
schema.RegisterName[ConditionOutput]("_cyberstrike_workflow_condition_output")
schema.RegisterName[ToolOutput]("_cyberstrike_workflow_tool_output")
schema.RegisterName[AgentOutput]("_cyberstrike_workflow_agent_output")
schema.RegisterName[HITLOutput]("_cyberstrike_workflow_hitl_output")
schema.RegisterName[OutputNodeOutput]("_cyberstrike_workflow_output_node_output")
}
// WorkflowLocalState is the Eino WithGenLocalState payload (checkpoint-serializable).
type WorkflowLocalState struct {
Inputs map[string]any `json:"inputs,omitempty"`
Outputs map[string]any `json:"outputs,omitempty"`
NodeOutputs map[string]map[string]any `json:"nodeOutputs,omitempty"`
NodeProceed map[string]bool `json:"nodeProceed,omitempty"`
LastOutput map[string]any `json:"lastOutput,omitempty"`
Metrics map[string]any `json:"metrics,omitempty"`
Executed []string `json:"executed,omitempty"`
Skipped []string `json:"skipped,omitempty"`
WorkflowRunID string `json:"workflowRunId,omitempty"`
MainIterationOffset int `json:"mainIterationOffset,omitempty"`
SegmentMaxIteration int `json:"segmentMaxIteration,omitempty"`
}
func newWorkflowLocalState(inputs map[string]interface{}, runID string) *WorkflowLocalState {
in := make(map[string]any, len(inputs))
for k, v := range inputs {
in[k] = v
}
return &WorkflowLocalState{
Inputs: in,
Outputs: make(map[string]any),
NodeOutputs: make(map[string]map[string]any),
NodeProceed: make(map[string]bool),
Metrics: make(map[string]any),
WorkflowRunID: runID,
}
}
var templateVarRe = regexp.MustCompile(`\{\{\s*([a-zA-Z0-9_.-]+)\s*\}\}`)
func resolveTemplate(s string, state *WorkflowLocalState) string {
if strings.TrimSpace(s) == "" {
return fmt.Sprint(valueFromPath("previous.output", state))
}
return templateVarRe.ReplaceAllStringFunc(s, func(match string) string {
m := templateVarRe.FindStringSubmatch(match)
if len(m) != 2 {
return match
}
return fmt.Sprint(valueFromPath(m[1], state))
})
}
func valueFromPath(path string, state *WorkflowLocalState) any {
parts := strings.Split(path, ".")
if len(parts) == 0 {
return ""
}
var cur any
switch parts[0] {
case "inputs", "input":
cur = state.Inputs
case "previous", "prev":
cur = state.LastOutput
case "outputs":
cur = state.Outputs
default:
if v, ok := state.Inputs[parts[0]]; ok {
cur = v
} else if v, ok := state.NodeOutputs[parts[0]]; ok {
cur = v
} else {
return ""
}
}
for _, p := range parts[1:] {
m, ok := cur.(map[string]any)
if !ok {
return ""
}
cur = m[p]
}
if cur == nil {
return ""
}
return cur
}
func cleanComparable(s string) string {
s = strings.TrimSpace(s)
s = strings.Trim(s, `"'`)
return s
}
func edgeAllowed(edge graphEdge, sourceNode graphNode, edgeIndex int, state *WorkflowLocalState) bool {
cond := firstNonEmpty(cfgString(edge.Config, "condition"), cfgString(edge.Config, "expression"))
if cond != "" {
return evalCondition(cond, state)
}
if strings.EqualFold(strings.TrimSpace(sourceNode.Type), "condition") {
return conditionBranchAllowed(edge, edgeIndex, state)
}
return true
}
func conditionBranchAllowed(edge graphEdge, edgeIndex int, state *WorkflowLocalState) bool {
matched := conditionMatched(state)
if branch := conditionBranchHint(edge); branch != "" {
return (branch == "true" && matched) || (branch == "false" && !matched)
}
switch edgeIndex {
case 0:
return matched
case 1:
return !matched
default:
return false
}
}
func conditionMatched(state *WorkflowLocalState) bool {
v := strings.ToLower(cleanComparable(fmt.Sprint(valueFromPath("previous.matched", state))))
return v == "true" || v == "1"
}
func conditionBranchHint(edge graphEdge) string {
if edge.Config != nil {
switch strings.ToLower(strings.TrimSpace(cfgString(edge.Config, "branch"))) {
case "true", "yes", "y", "是":
return "true"
case "false", "no", "n", "否":
return "false"
}
}
switch strings.ToLower(strings.TrimSpace(edge.Label)) {
case "true", "yes", "y", "是":
return "true"
case "false", "no", "n", "否":
return "false"
}
return ""
}
func cfgString(cfg map[string]any, key string) string {
if cfg == nil {
return ""
}
if v, ok := cfg[key]; ok {
return strings.TrimSpace(fmt.Sprint(v))
}
return ""
}
func firstNonEmpty(values ...string) string {
for _, value := range values {
if s := strings.TrimSpace(value); s != "" {
return s
}
}
return ""
}
func truncateWorkflowPreview(s string, limit int) string {
s = strings.TrimSpace(s)
if limit <= 0 || len([]rune(s)) <= limit {
return s
}
runes := []rune(s)
return string(runes[:limit]) + "..."
}
func renderWorkflowResponse(roleName, workflowName string, version int, runID string, state *WorkflowLocalState) string {
var sb strings.Builder
sb.WriteString(fmt.Sprintf("角色「%s」已完成工作流「%s」(版本 %d)。\n\n", roleName, workflowName, version))
sb.WriteString(fmt.Sprintf("运行 ID%s\n", runID))
sb.WriteString(fmt.Sprintf("已执行节点:%d", len(state.Executed)))
if len(state.Skipped) > 0 {
sb.WriteString(fmt.Sprintf(",跳过节点:%d", len(state.Skipped)))
}
sb.WriteString("\n\n")
if len(state.Outputs) > 0 {
sb.WriteString("输出:\n")
keys := make([]string, 0, len(state.Outputs))
for k := range state.Outputs {
keys = append(keys, k)
}
sort.Strings(keys)
for _, k := range keys {
sb.WriteString(fmt.Sprintf("- %s%v\n", k, state.Outputs[k]))
}
} else {
sb.WriteString("暂无输出。请检查是否配置了输出节点,或条件分支是否命中。\n")
}
if len(state.Skipped) > 0 {
sb.WriteString("\n未执行的节点类型仍会保留运行记录:")
sb.WriteString(strings.Join(state.Skipped, "、"))
sb.WriteString("。")
}
return strings.TrimSpace(sb.String())
}
+150
View File
@@ -0,0 +1,150 @@
package workflow
type NodeOutputEnvelope struct {
Kind string `json:"kind"`
NodeID string `json:"node_id"`
NodeType string `json:"node_type"`
Status string `json:"status"`
Output any `json:"output"`
}
type StartOutput struct {
NodeOutputEnvelope
Message any `json:"message"`
ConversationID any `json:"conversationId"`
ProjectID any `json:"projectId"`
}
type ConditionOutput struct {
NodeOutputEnvelope
Condition string `json:"condition"`
Matched bool `json:"matched"`
}
type ToolOutput struct {
NodeOutputEnvelope
ToolName string `json:"tool_name"`
Arguments map[string]any `json:"arguments"`
ExecutionID string `json:"execution_id"`
IsError bool `json:"is_error"`
}
type AgentOutput struct {
NodeOutputEnvelope
Mode string `json:"mode"`
MCPExecutionIDs []string `json:"mcp_execution_ids"`
}
type HITLOutput struct {
NodeOutputEnvelope
Prompt string `json:"prompt"`
Reviewer string `json:"reviewer"`
Approved bool `json:"approved"`
Mode string `json:"mode"`
}
type OutputNodeOutput struct {
NodeOutputEnvelope
OutputKey string `json:"output_key"`
Outputs map[string]any `json:"outputs"`
}
func envelope(kind, nodeID, nodeType, status string, output any) NodeOutputEnvelope {
return NodeOutputEnvelope{Kind: kind, NodeID: nodeID, NodeType: nodeType, Status: status, Output: output}
}
func outputMap(env NodeOutputEnvelope, extra map[string]any) map[string]any {
out := map[string]any{
"kind": env.Kind,
"node_id": env.NodeID,
"node_type": env.NodeType,
"status": env.Status,
"output": env.Output,
"typed": env,
}
for k, v := range extra {
out[k] = v
}
return out
}
func startOutputMap(node graphNode, message, conversationID, projectID any) map[string]any {
typed := StartOutput{
NodeOutputEnvelope: envelope("start", node.ID, node.Type, "completed", message),
Message: message,
ConversationID: conversationID,
ProjectID: projectID,
}
return outputMap(typed.NodeOutputEnvelope, map[string]any{
"message": typed.Message,
"conversationId": typed.ConversationID,
"projectId": typed.ProjectID,
"typed": typed,
})
}
func conditionOutputMap(node graphNode, expr string, matched bool) map[string]any {
typed := ConditionOutput{
NodeOutputEnvelope: envelope("condition", node.ID, node.Type, "completed", matched),
Condition: expr,
Matched: matched,
}
return outputMap(typed.NodeOutputEnvelope, map[string]any{"condition": expr, "matched": matched, "typed": typed})
}
func outputNodeOutputMap(node graphNode, key string, value any) map[string]any {
typed := OutputNodeOutput{
NodeOutputEnvelope: envelope("output", node.ID, node.Type, "completed", value),
OutputKey: key,
Outputs: map[string]any{key: value},
}
return outputMap(typed.NodeOutputEnvelope, map[string]any{"output_key": key, "outputs": typed.Outputs, "typed": typed})
}
func endOutputMap(node graphNode, value any) map[string]any {
typed := envelope("end", node.ID, node.Type, "completed", value)
return outputMap(typed, nil)
}
func toolOutputMap(node graphNode, output string, toolName string, args map[string]any, executionID string, isError bool) map[string]any {
typed := ToolOutput{
NodeOutputEnvelope: envelope("tool", node.ID, node.Type, "completed", output),
ToolName: toolName,
Arguments: args,
ExecutionID: executionID,
IsError: isError,
}
return outputMap(typed.NodeOutputEnvelope, map[string]any{
"tool_name": toolName,
"arguments": args,
"execution_id": executionID,
"is_error": isError,
"typed": typed,
})
}
func agentOutputMap(node graphNode, response, mode string, mcpIDs []string) map[string]any {
typed := AgentOutput{
NodeOutputEnvelope: envelope("agent", node.ID, node.Type, "completed", response),
Mode: mode,
MCPExecutionIDs: mcpIDs,
}
return outputMap(typed.NodeOutputEnvelope, map[string]any{"mode": mode, "mcp_execution_ids": mcpIDs, "typed": typed})
}
func hitlOutputMap(node graphNode, status string, output string, prompt string, reviewer string, approved bool) map[string]any {
typed := HITLOutput{
NodeOutputEnvelope: envelope("hitl", node.ID, node.Type, status, output),
Prompt: prompt,
Reviewer: reviewer,
Approved: approved,
Mode: "interactive",
}
return outputMap(typed.NodeOutputEnvelope, map[string]any{
"prompt": prompt,
"reviewer": reviewer,
"approved": approved,
"mode": "interactive",
"typed": typed,
})
}
@@ -0,0 +1,23 @@
package workflow
import (
"strings"
"testing"
)
func TestTruncateWorkflowToolOutputBoundsBytesAndKeepsExecutionReference(t *testing.T) {
out := truncateWorkflowToolOutput(strings.Repeat("响应正文", 1000), 256, "exec-123")
if len(out) > 256 {
t.Fatalf("workflow output bytes=%d, want <=256", len(out))
}
if !strings.Contains(out, "exec-123") || !strings.Contains(out, "truncated") {
t.Fatalf("missing truncation reference: %q", out)
}
}
func TestTruncateWorkflowToolOutputLeavesBoundedContentUntouched(t *testing.T) {
const want = "small-result"
if got := truncateWorkflowToolOutput(want, 256, "exec-123"); got != want {
t.Fatalf("got %q want %q", got, want)
}
}
+74
View File
@@ -0,0 +1,74 @@
package workflow
import (
"fmt"
"strconv"
)
// WorkflowInput is the typed entry for Eino compose.Workflow[I,O].
type WorkflowInput struct {
Message string `json:"message"`
ConversationID string `json:"conversationId"`
ProjectID string `json:"projectId"`
Role string `json:"role"`
WorkflowID string `json:"workflowId"`
WorkflowVersion int `json:"workflowVersion"`
}
// WorkflowOutput aggregates terminal node payloads keyed by canvas node id.
type WorkflowOutput map[string]any
// WorkflowNodeOutput is the per-node lambda payload (alias for Eino edge type alignment).
type WorkflowNodeOutput = map[string]interface{}
func workflowInputFromMap(m map[string]interface{}) WorkflowInput {
in := WorkflowInput{}
if m == nil {
return in
}
if v, ok := m["message"].(string); ok {
in.Message = v
} else if m["message"] != nil {
in.Message = fmt.Sprint(m["message"])
}
if v, ok := m["conversationId"].(string); ok {
in.ConversationID = v
}
if v, ok := m["projectId"].(string); ok {
in.ProjectID = v
}
if v, ok := m["role"].(string); ok {
in.Role = v
}
if v, ok := m["workflowId"].(string); ok {
in.WorkflowID = v
}
switch v := m["workflowVersion"].(type) {
case int:
in.WorkflowVersion = v
case int64:
in.WorkflowVersion = int(v)
case float64:
in.WorkflowVersion = int(v)
case string:
if n, err := strconv.Atoi(v); err == nil {
in.WorkflowVersion = n
}
}
return in
}
func (in WorkflowInput) toStateInputs() map[string]any {
return map[string]any{
"message": in.Message,
"conversationId": in.ConversationID,
"projectId": in.ProjectID,
"role": in.Role,
"workflowId": in.WorkflowID,
"workflowVersion": in.WorkflowVersion,
}
}
func cacheKey(workflowID string, version int) string {
return workflowID + ":" + strconv.Itoa(version)
}
+366
View File
@@ -0,0 +1,366 @@
package workflow
import (
"fmt"
"strconv"
"strings"
)
var allowedWorkflowNodeTypes = map[string]bool{
"start": true,
"tool": true,
"agent": true,
"condition": true,
"hitl": true,
"output": true,
"end": true,
}
func validateGraphDefinition(g *graphDef, idx *graphIndex) error {
if g == nil || idx == nil {
return fmt.Errorf("工作流图为空")
}
if err := validateNodeIDsAndTypes(g); err != nil {
return err
}
if err := validateEdges(g, idx); err != nil {
return err
}
if err := validateNodeTopology(idx); err != nil {
return err
}
if err := validateNodeConfigs(idx); err != nil {
return err
}
if err := validateDAG(idx); err != nil {
return err
}
if err := validateReachability(idx); err != nil {
return err
}
return nil
}
func validateNodeIDsAndTypes(g *graphDef) error {
seen := make(map[string]bool, len(g.Nodes))
for _, node := range g.Nodes {
id := strings.TrimSpace(node.ID)
if id == "" {
return fmt.Errorf("工作流存在空节点 ID")
}
if seen[id] {
return fmt.Errorf("工作流存在重复节点 ID: %s", id)
}
seen[id] = true
nodeType := strings.ToLower(strings.TrimSpace(node.Type))
if nodeType == "" {
return fmt.Errorf("节点「%s」缺少节点类型", id)
}
if !allowedWorkflowNodeTypes[nodeType] {
return fmt.Errorf("节点「%s」使用了未知节点类型: %s", id, node.Type)
}
}
return nil
}
func validateEdges(g *graphDef, idx *graphIndex) error {
seen := make(map[string]bool, len(g.Edges))
for _, edge := range g.Edges {
if id := strings.TrimSpace(edge.ID); id != "" {
if seen[id] {
return fmt.Errorf("工作流存在重复连线 ID: %s", id)
}
seen[id] = true
}
source := strings.TrimSpace(edge.Source)
target := strings.TrimSpace(edge.Target)
if source == "" || target == "" {
return fmt.Errorf("工作流存在源或目标为空的连线")
}
if source == target {
return fmt.Errorf("连线「%s」不能自环", firstNonEmpty(edge.ID, source))
}
if _, ok := idx.nodes[source]; !ok {
return fmt.Errorf("连线「%s」引用了不存在的源节点: %s", firstNonEmpty(edge.ID, source), source)
}
if _, ok := idx.nodes[target]; !ok {
return fmt.Errorf("连线「%s」引用了不存在的目标节点: %s", firstNonEmpty(edge.ID, target), target)
}
}
return nil
}
func validateNodeTopology(idx *graphIndex) error {
starts := explicitStartNodeIDs(idx)
if len(starts) == 0 {
return fmt.Errorf("工作流至少需要一个开始节点")
}
outputs := outputNodeIDs(idx)
if len(outputs) == 0 {
return fmt.Errorf("工作流至少需要一个输出节点")
}
for id, node := range idx.nodes {
inDegree := len(idx.incoming[id])
outDegree := len(idx.outgoing[id])
nodeType := strings.ToLower(strings.TrimSpace(node.Type))
switch nodeType {
case "start":
if inDegree > 0 {
return fmt.Errorf("开始节点「%s」不能有入边", firstNonEmpty(node.Label, id))
}
if outDegree == 0 {
return fmt.Errorf("开始节点「%s」至少需要一条出边", firstNonEmpty(node.Label, id))
}
case "output", "end":
if outDegree > 0 {
return fmt.Errorf("%s 节点「%s」不能有出边", displayNodeType(nodeType), firstNonEmpty(node.Label, id))
}
if inDegree == 0 {
return fmt.Errorf("%s 节点「%s」至少需要一条入边", displayNodeType(nodeType), firstNonEmpty(node.Label, id))
}
default:
if inDegree == 0 {
return fmt.Errorf("节点「%s」不可达:非开始节点必须有入边", firstNonEmpty(node.Label, id))
}
if outDegree == 0 {
return fmt.Errorf("节点「%s」没有出边;请连接到 output/end 节点", firstNonEmpty(node.Label, id))
}
}
}
return nil
}
func validateNodeConfigs(idx *graphIndex) error {
for id, node := range idx.nodes {
label := firstNonEmpty(node.Label, id)
switch strings.ToLower(strings.TrimSpace(node.Type)) {
case "tool":
if cfgString(node.Config, "tool_name") == "" {
return fmt.Errorf("工具节点「%s」必须选择 MCP 工具", label)
}
if err := validateToolConfig(node); err != nil {
return err
}
case "agent":
if cfgString(node.Config, "instruction") == "" {
if _, ok := parseFieldBinding(node.Config, "input_binding"); !ok {
return fmt.Errorf("Agent 节点「%s」必须填写节点指令或输入绑定", label)
}
}
if cfgString(node.Config, "output_key") == "" {
return fmt.Errorf("Agent 节点「%s」必须填写输出变量名", label)
}
case "condition":
if cfgString(node.Config, "expression") == "" {
return fmt.Errorf("条件节点「%s」必须填写表达式", label)
}
if err := validateConditionExpression(cfgString(node.Config, "expression")); err != nil {
return fmt.Errorf("条件节点「%s」表达式非法: %w", label, err)
}
if n := len(idx.outgoing[id]); n < 1 || n > 2 {
return fmt.Errorf("条件节点「%s」需要 1 到 2 条出边(是/否)", label)
}
if err := validateConditionBranchLabels(idx, id, node); err != nil {
return err
}
case "output":
if cfgString(node.Config, "output_key") == "" {
return fmt.Errorf("输出节点「%s」必须填写输出变量名", label)
}
}
if err := validateJoinConfig(idx, id, node); err != nil {
return err
}
if hasConditionalOutgoingEdges(idx, id) {
if err := validateConditionalOutgoingEdges(idx, id, node); err != nil {
return err
}
}
}
return nil
}
func validateConditionalOutgoingEdges(idx *graphIndex, nodeID string, node graphNode) error {
unconditional := 0
for _, edge := range idx.outgoing[nodeID] {
cond := firstNonEmpty(cfgString(edge.Config, "condition"), cfgString(edge.Config, "expression"))
if cond != "" {
if err := validateConditionExpression(cond); err != nil {
return fmt.Errorf("节点「%s」的连线条件非法: %w", firstNonEmpty(node.Label, nodeID), err)
}
}
if cond == "" {
unconditional++
}
}
if unconditional > 1 {
return fmt.Errorf("节点「%s」的条件出边最多只能有一条默认分支", firstNonEmpty(node.Label, nodeID))
}
return nil
}
func validateToolConfig(node graphNode) error {
rawArgs := cfgString(node.Config, "arguments")
if rawArgs != "" {
if _, err := resolveToolArguments(node.Config, &WorkflowLocalState{}); err != nil {
return fmt.Errorf("工具节点「%s」参数 JSON 非法: %w", firstNonEmpty(node.Label, node.ID), err)
}
}
if timeout := cfgString(node.Config, "timeout_seconds"); timeout != "" {
if _, err := parsePositiveInt(timeout); err != nil {
return fmt.Errorf("工具节点「%s」超时时间必须是正整数", firstNonEmpty(node.Label, node.ID))
}
}
return nil
}
func validateJoinConfig(idx *graphIndex, nodeID string, node graphNode) error {
strategy := joinStrategy(node)
if !allowedJoinStrategies[strategy] {
return fmt.Errorf("节点「%s」使用了未知汇聚策略: %s", firstNonEmpty(node.Label, nodeID), strategy)
}
if len(idx.incoming[nodeID]) > 1 && strategy == "" {
return fmt.Errorf("节点「%s」有多个上游时必须声明汇聚策略", firstNonEmpty(node.Label, nodeID))
}
return nil
}
func validateConditionBranchLabels(idx *graphIndex, nodeID string, node graphNode) error {
seen := map[string]bool{}
for _, edge := range idx.outgoing[nodeID] {
hint := conditionBranchHint(edge)
if hint == "" {
return fmt.Errorf("条件节点「%s」的出边必须标记为是/否或 true/false", firstNonEmpty(node.Label, nodeID))
}
if seen[hint] {
return fmt.Errorf("条件节点「%s」存在重复分支标签: %s", firstNonEmpty(node.Label, nodeID), hint)
}
seen[hint] = true
}
return nil
}
func validateDAG(idx *graphIndex) error {
color := make(map[string]int, len(idx.nodes))
var visit func(string) error
visit = func(id string) error {
switch color[id] {
case 1:
return fmt.Errorf("工作流存在环路,Workflow 编排必须是 DAG: %s", id)
case 2:
return nil
}
color[id] = 1
for _, edge := range idx.outgoing[id] {
if err := visit(edge.Target); err != nil {
return err
}
}
color[id] = 2
return nil
}
for id := range idx.nodes {
if err := visit(id); err != nil {
return err
}
}
return nil
}
func validateReachability(idx *graphIndex) error {
starts := explicitStartNodeIDs(idx)
reached := make(map[string]bool, len(idx.nodes))
queue := append([]string(nil), starts...)
for len(queue) > 0 {
id := queue[0]
queue = queue[1:]
if reached[id] {
continue
}
reached[id] = true
for _, edge := range idx.outgoing[id] {
queue = append(queue, edge.Target)
}
}
for id, node := range idx.nodes {
if !reached[id] {
return fmt.Errorf("节点「%s」不可达:没有从开始节点连通到该节点", firstNonEmpty(node.Label, id))
}
}
canReachTerminal := make(map[string]bool, len(idx.nodes))
visiting := make(map[string]bool, len(idx.nodes))
var reachesTerminal func(string) bool
reachesTerminal = func(id string) bool {
if canReachTerminal[id] {
return true
}
if visiting[id] {
return false
}
visiting[id] = true
node := idx.nodes[id]
nodeType := strings.ToLower(strings.TrimSpace(node.Type))
if nodeType == "output" || nodeType == "end" {
canReachTerminal[id] = true
visiting[id] = false
return true
}
for _, edge := range idx.outgoing[id] {
if reachesTerminal(edge.Target) {
canReachTerminal[id] = true
visiting[id] = false
return true
}
}
visiting[id] = false
return false
}
for id, node := range idx.nodes {
if !reachesTerminal(id) {
return fmt.Errorf("节点「%s」无法到达 output/end 终点", firstNonEmpty(node.Label, id))
}
}
return nil
}
func explicitStartNodeIDs(idx *graphIndex) []string {
var ids []string
for id, node := range idx.nodes {
if strings.EqualFold(node.Type, "start") {
ids = append(ids, id)
}
}
sortNodeIDsByCanvas(ids, idx.nodes)
return ids
}
func outputNodeIDs(idx *graphIndex) []string {
var ids []string
for id, node := range idx.nodes {
if strings.EqualFold(node.Type, "output") {
ids = append(ids, id)
}
}
sortNodeIDsByCanvas(ids, idx.nodes)
return ids
}
func displayNodeType(nodeType string) string {
switch strings.ToLower(strings.TrimSpace(nodeType)) {
case "output":
return "输出"
case "end":
return "结束"
default:
return nodeType
}
}
func parsePositiveInt(s string) (int, error) {
n, err := strconv.Atoi(strings.TrimSpace(s))
if err != nil || n <= 0 {
return 0, fmt.Errorf("not positive integer")
}
return n, nil
}