mirror of
https://github.com/Ed1s0nZ/CyberStrikeAI.git
synced 2026-08-15 07:30:53 +02:00
Add files via upload
This commit is contained in:
@@ -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
|
||||
}
|
||||
@@ -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()
|
||||
}
|
||||
@@ -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)
|
||||
}
|
||||
@@ -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
|
||||
}
|
||||
}
|
||||
@@ -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")
|
||||
}
|
||||
}
|
||||
}()
|
||||
}
|
||||
@@ -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
|
||||
}
|
||||
}
|
||||
@@ -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
|
||||
}
|
||||
@@ -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
|
||||
}
|
||||
}
|
||||
@@ -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)
|
||||
}
|
||||
@@ -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)
|
||||
}
|
||||
@@ -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")
|
||||
}
|
||||
}
|
||||
}()
|
||||
}
|
||||
@@ -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)
|
||||
}
|
||||
}
|
||||
@@ -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,
|
||||
})
|
||||
}
|
||||
@@ -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)
|
||||
}
|
||||
@@ -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)
|
||||
@@ -0,0 +1,24 @@
|
||||
package knowledge
|
||||
|
||||
import (
|
||||
"context"
|
||||
"fmt"
|
||||
|
||||
"github.com/cloudwego/eino/compose"
|
||||
"github.com/cloudwego/eino/schema"
|
||||
)
|
||||
|
||||
// BuildKnowledgeRetrieveChain 编译「查询字符串 → 文档列表」的 Eino Chain(MultiQuery → 向量 → 重排 → 后处理)。
|
||||
func BuildKnowledgeRetrieveChain(ctx context.Context, r *Retriever) (compose.Runnable[string, []*schema.Document], error) {
|
||||
if r == nil {
|
||||
return nil, fmt.Errorf("retriever is nil")
|
||||
}
|
||||
ch := compose.NewChain[string, []*schema.Document]()
|
||||
ch.AppendRetriever(r.AsEinoRetriever())
|
||||
return ch.Compile(ctx)
|
||||
}
|
||||
|
||||
// CompileRetrieveChain 等价于 [BuildKnowledgeRetrieveChain](ctx, r)。
|
||||
func (r *Retriever) CompileRetrieveChain(ctx context.Context) (compose.Runnable[string, []*schema.Document], error) {
|
||||
return BuildKnowledgeRetrieveChain(ctx, r)
|
||||
}
|
||||
@@ -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)
|
||||
@@ -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)
|
||||
@@ -0,0 +1,251 @@
|
||||
package knowledge
|
||||
|
||||
import (
|
||||
"context"
|
||||
"fmt"
|
||||
"net/http"
|
||||
"strings"
|
||||
"sync"
|
||||
"time"
|
||||
|
||||
"cyberstrike-ai/internal/config"
|
||||
|
||||
einoembedopenai "github.com/cloudwego/eino-ext/components/embedding/openai"
|
||||
"github.com/cloudwego/eino/components/embedding"
|
||||
"go.uber.org/zap"
|
||||
"golang.org/x/time/rate"
|
||||
)
|
||||
|
||||
// Embedder 使用 CloudWeGo Eino 的 OpenAI Embedding 组件,并保留速率限制与重试。
|
||||
type Embedder struct {
|
||||
eino embedding.Embedder
|
||||
config *config.KnowledgeConfig
|
||||
logger *zap.Logger
|
||||
|
||||
rateLimiter *rate.Limiter
|
||||
rateLimitDelay time.Duration
|
||||
maxRetries int
|
||||
retryDelay time.Duration
|
||||
mu sync.Mutex
|
||||
}
|
||||
|
||||
// NewEmbedder 基于 Eino eino-ext OpenAI Embedder;openAIConfig 用于在知识库未单独配置 key 时回退 API Key。
|
||||
func NewEmbedder(ctx context.Context, cfg *config.KnowledgeConfig, openAIConfig *config.OpenAIConfig, logger *zap.Logger) (*Embedder, error) {
|
||||
if cfg == nil {
|
||||
return nil, fmt.Errorf("knowledge config is nil")
|
||||
}
|
||||
|
||||
var rateLimiter *rate.Limiter
|
||||
var rateLimitDelay time.Duration
|
||||
if cfg.Indexing.MaxRPM > 0 {
|
||||
rpm := cfg.Indexing.MaxRPM
|
||||
rateLimiter = rate.NewLimiter(rate.Every(time.Minute/time.Duration(rpm)), rpm)
|
||||
if logger != nil {
|
||||
logger.Info("知识库索引速率限制已启用", zap.Int("maxRPM", rpm))
|
||||
}
|
||||
} else if cfg.Indexing.RateLimitDelayMs > 0 {
|
||||
rateLimitDelay = time.Duration(cfg.Indexing.RateLimitDelayMs) * time.Millisecond
|
||||
if logger != nil {
|
||||
logger.Info("知识库索引固定延迟已启用", zap.Duration("delay", rateLimitDelay))
|
||||
}
|
||||
}
|
||||
|
||||
maxRetries := 3
|
||||
retryDelay := 1000 * time.Millisecond
|
||||
if cfg.Indexing.MaxRetries > 0 {
|
||||
maxRetries = cfg.Indexing.MaxRetries
|
||||
}
|
||||
if cfg.Indexing.RetryDelayMs > 0 {
|
||||
retryDelay = time.Duration(cfg.Indexing.RetryDelayMs) * time.Millisecond
|
||||
}
|
||||
|
||||
model := strings.TrimSpace(cfg.Embedding.Model)
|
||||
if model == "" {
|
||||
model = "text-embedding-3-small"
|
||||
}
|
||||
|
||||
baseURL := strings.TrimSpace(cfg.Embedding.BaseURL)
|
||||
baseURL = strings.TrimSuffix(baseURL, "/")
|
||||
if baseURL == "" {
|
||||
baseURL = "https://api.openai.com/v1"
|
||||
}
|
||||
|
||||
apiKey := strings.TrimSpace(cfg.Embedding.APIKey)
|
||||
if apiKey == "" && openAIConfig != nil {
|
||||
apiKey = strings.TrimSpace(openAIConfig.APIKey)
|
||||
}
|
||||
if apiKey == "" {
|
||||
return nil, fmt.Errorf("embedding API key 未配置")
|
||||
}
|
||||
|
||||
timeout := 120 * time.Second
|
||||
if cfg.Indexing.RequestTimeoutSeconds > 0 {
|
||||
timeout = time.Duration(cfg.Indexing.RequestTimeoutSeconds) * time.Second
|
||||
}
|
||||
httpClient := &http.Client{Timeout: timeout}
|
||||
|
||||
inner, err := einoembedopenai.NewEmbedder(ctx, &einoembedopenai.EmbeddingConfig{
|
||||
APIKey: apiKey,
|
||||
BaseURL: baseURL,
|
||||
ByAzure: false,
|
||||
Model: model,
|
||||
HTTPClient: httpClient,
|
||||
})
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("eino OpenAI embedder: %w", err)
|
||||
}
|
||||
|
||||
return &Embedder{
|
||||
eino: inner,
|
||||
config: cfg,
|
||||
logger: logger,
|
||||
rateLimiter: rateLimiter,
|
||||
rateLimitDelay: rateLimitDelay,
|
||||
maxRetries: maxRetries,
|
||||
retryDelay: retryDelay,
|
||||
}, nil
|
||||
}
|
||||
|
||||
// EmbeddingModelName 返回配置的嵌入模型名(用于 tiktoken 分块与向量行元数据)。
|
||||
func (e *Embedder) EmbeddingModelName() string {
|
||||
if e == nil || e.config == nil {
|
||||
return ""
|
||||
}
|
||||
s := strings.TrimSpace(e.config.Embedding.Model)
|
||||
if s != "" {
|
||||
return s
|
||||
}
|
||||
return "text-embedding-3-small"
|
||||
}
|
||||
|
||||
func (e *Embedder) waitRateLimiter() {
|
||||
e.mu.Lock()
|
||||
defer e.mu.Unlock()
|
||||
|
||||
if e.rateLimiter != nil {
|
||||
ctx := context.Background()
|
||||
if err := e.rateLimiter.Wait(ctx); err != nil && e.logger != nil {
|
||||
e.logger.Warn("速率限制器等待失败", zap.Error(err))
|
||||
}
|
||||
}
|
||||
if e.rateLimitDelay > 0 {
|
||||
time.Sleep(e.rateLimitDelay)
|
||||
}
|
||||
}
|
||||
|
||||
// EmbedText 单条嵌入(float32,与历史存储格式一致)。
|
||||
func (e *Embedder) EmbedText(ctx context.Context, text string) ([]float32, error) {
|
||||
vecs, err := e.EmbedStrings(ctx, []string{text})
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
if len(vecs) != 1 {
|
||||
return nil, fmt.Errorf("unexpected embedding count: %d", len(vecs))
|
||||
}
|
||||
return vecs[0], nil
|
||||
}
|
||||
|
||||
// EmbedStrings 批量嵌入,带重试;实现 [embedding.Embedder],可供 Eino Indexer 使用。
|
||||
func (e *Embedder) EmbedStrings(ctx context.Context, texts []string, opts ...embedding.Option) ([][]float32, error) {
|
||||
if e == nil || e.eino == nil {
|
||||
return nil, fmt.Errorf("embedder not initialized")
|
||||
}
|
||||
if len(texts) == 0 {
|
||||
return nil, nil
|
||||
}
|
||||
|
||||
var lastErr error
|
||||
for attempt := 0; attempt < e.maxRetries; attempt++ {
|
||||
if attempt > 0 {
|
||||
wait := e.retryDelay * time.Duration(attempt)
|
||||
if e.logger != nil {
|
||||
e.logger.Debug("嵌入重试前等待", zap.Int("attempt", attempt+1), zap.Duration("wait", wait))
|
||||
}
|
||||
select {
|
||||
case <-ctx.Done():
|
||||
return nil, ctx.Err()
|
||||
case <-time.After(wait):
|
||||
}
|
||||
} else {
|
||||
e.waitRateLimiter()
|
||||
}
|
||||
|
||||
raw, err := e.eino.EmbedStrings(ctx, texts, opts...)
|
||||
if err == nil {
|
||||
out := make([][]float32, len(raw))
|
||||
for i, row := range raw {
|
||||
out[i] = make([]float32, len(row))
|
||||
for j, v := range row {
|
||||
out[i][j] = float32(v)
|
||||
}
|
||||
}
|
||||
return out, nil
|
||||
}
|
||||
lastErr = err
|
||||
if !e.isRetryableError(err) {
|
||||
return nil, err
|
||||
}
|
||||
if e.logger != nil {
|
||||
e.logger.Debug("嵌入失败,将重试", zap.Int("attempt", attempt+1), zap.Error(err))
|
||||
}
|
||||
}
|
||||
return nil, fmt.Errorf("达到最大重试次数 (%d): %v", e.maxRetries, lastErr)
|
||||
}
|
||||
|
||||
// EmbedTexts 批量 float32 嵌入(兼容旧调用;单次请求批量以减小延迟)。
|
||||
func (e *Embedder) EmbedTexts(ctx context.Context, texts []string) ([][]float32, error) {
|
||||
return e.EmbedStrings(ctx, texts)
|
||||
}
|
||||
|
||||
func (e *Embedder) isRetryableError(err error) bool {
|
||||
if err == nil {
|
||||
return false
|
||||
}
|
||||
errStr := err.Error()
|
||||
if strings.Contains(errStr, "429") || strings.Contains(errStr, "rate limit") {
|
||||
return true
|
||||
}
|
||||
if strings.Contains(errStr, "500") || strings.Contains(errStr, "502") ||
|
||||
strings.Contains(errStr, "503") || strings.Contains(errStr, "504") {
|
||||
return true
|
||||
}
|
||||
if strings.Contains(errStr, "timeout") || strings.Contains(errStr, "connection") ||
|
||||
strings.Contains(errStr, "network") || strings.Contains(errStr, "EOF") {
|
||||
return true
|
||||
}
|
||||
return false
|
||||
}
|
||||
|
||||
// einoFloatEmbedder adapts [][]float32 embedder to Eino's [][]float64 [embedding.Embedder] for Indexer.Store.
|
||||
type einoFloatEmbedder struct {
|
||||
inner *Embedder
|
||||
}
|
||||
|
||||
func (w *einoFloatEmbedder) EmbedStrings(ctx context.Context, texts []string, opts ...embedding.Option) ([][]float64, error) {
|
||||
vec32, err := w.inner.EmbedStrings(ctx, texts, opts...)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
out := make([][]float64, len(vec32))
|
||||
for i, row := range vec32 {
|
||||
out[i] = make([]float64, len(row))
|
||||
for j, v := range row {
|
||||
out[i][j] = float64(v)
|
||||
}
|
||||
}
|
||||
return out, nil
|
||||
}
|
||||
|
||||
func (w *einoFloatEmbedder) GetType() string {
|
||||
return "CyberStrikeKnowledgeEmbedder"
|
||||
}
|
||||
|
||||
func (w *einoFloatEmbedder) IsCallbacksEnabled() bool {
|
||||
return false
|
||||
}
|
||||
|
||||
// EinoEmbeddingComponent returns an [embedding.Embedder] that uses the same retry/rate-limit path
|
||||
// and produces float64 vectors expected by generic Eino indexer helpers.
|
||||
func (e *Embedder) EinoEmbeddingComponent() embedding.Embedder {
|
||||
return &einoFloatEmbedder{inner: e}
|
||||
}
|
||||
@@ -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
|
||||
})
|
||||
}
|
||||
@@ -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)
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -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()
|
||||
}
|
||||
@@ -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
|
||||
}
|
||||
@@ -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)
|
||||
@@ -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")
|
||||
}
|
||||
}
|
||||
@@ -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))
|
||||
}
|
||||
}
|
||||
@@ -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,
|
||||
}
|
||||
}
|
||||
@@ -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)
|
||||
}
|
||||
@@ -0,0 +1,323 @@
|
||||
package knowledge
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"sort"
|
||||
"strings"
|
||||
|
||||
"cyberstrike-ai/internal/mcp"
|
||||
"cyberstrike-ai/internal/mcp/builtin"
|
||||
|
||||
"go.uber.org/zap"
|
||||
)
|
||||
|
||||
// RegisterKnowledgeTool 注册知识检索工具到MCP服务器
|
||||
func RegisterKnowledgeTool(
|
||||
mcpServer *mcp.Server,
|
||||
retriever *Retriever,
|
||||
manager *Manager,
|
||||
logger *zap.Logger,
|
||||
) {
|
||||
// 注册第一个工具:获取所有可用的风险类型列表
|
||||
listRiskTypesTool := mcp.Tool{
|
||||
Name: builtin.ToolListKnowledgeRiskTypes,
|
||||
Description: "获取知识库中所有可用的风险类型(risk_type)列表。在搜索知识库之前,可以先调用此工具获取可用的风险类型,然后使用正确的风险类型进行精确搜索,这样可以大幅减少检索时间并提高检索准确性。",
|
||||
ShortDescription: "获取知识库中所有可用的风险类型列表",
|
||||
InputSchema: map[string]interface{}{
|
||||
"type": "object",
|
||||
"properties": map[string]interface{}{},
|
||||
"required": []string{},
|
||||
},
|
||||
}
|
||||
|
||||
listRiskTypesHandler := func(ctx context.Context, args map[string]interface{}) (*mcp.ToolResult, error) {
|
||||
categories, err := manager.GetCategories()
|
||||
if err != nil {
|
||||
logger.Error("获取风险类型列表失败", zap.Error(err))
|
||||
return &mcp.ToolResult{
|
||||
Content: []mcp.Content{
|
||||
{
|
||||
Type: "text",
|
||||
Text: fmt.Sprintf("获取风险类型列表失败: %v", err),
|
||||
},
|
||||
},
|
||||
IsError: true,
|
||||
}, nil
|
||||
}
|
||||
|
||||
if len(categories) == 0 {
|
||||
return &mcp.ToolResult{
|
||||
Content: []mcp.Content{
|
||||
{
|
||||
Type: "text",
|
||||
Text: "知识库中暂无风险类型。",
|
||||
},
|
||||
},
|
||||
}, nil
|
||||
}
|
||||
|
||||
var resultText strings.Builder
|
||||
resultText.WriteString(fmt.Sprintf("知识库中共有 %d 个风险类型:\n\n", len(categories)))
|
||||
for i, category := range categories {
|
||||
resultText.WriteString(fmt.Sprintf("%d. %s\n", i+1, category))
|
||||
}
|
||||
resultText.WriteString("\n提示:在调用 " + builtin.ToolSearchKnowledgeBase + " 工具时,可以使用上述风险类型之一作为 risk_type 参数,以缩小搜索范围并提高检索效率。")
|
||||
|
||||
return &mcp.ToolResult{
|
||||
Content: []mcp.Content{
|
||||
{
|
||||
Type: "text",
|
||||
Text: resultText.String(),
|
||||
},
|
||||
},
|
||||
}, nil
|
||||
}
|
||||
|
||||
mcpServer.RegisterTool(listRiskTypesTool, listRiskTypesHandler)
|
||||
logger.Debug("风险类型列表工具已注册", zap.String("toolName", listRiskTypesTool.Name))
|
||||
|
||||
// 注册第二个工具:搜索知识库(保持原有功能)
|
||||
searchTool := mcp.Tool{
|
||||
Name: builtin.ToolSearchKnowledgeBase,
|
||||
Description: "在知识库中搜索相关的安全知识。当你需要了解特定漏洞类型、攻击技术、检测方法等安全知识时,可以使用此工具进行检索。工具基于向量嵌入与余弦相似度检索(与 Eino retriever 语义一致)。建议:在搜索前可以先调用 " + builtin.ToolListKnowledgeRiskTypes + " 工具获取可用的风险类型,然后使用正确的 risk_type 参数进行精确搜索,这样可以大幅减少检索时间。",
|
||||
ShortDescription: "搜索知识库中的安全知识(向量语义检索)",
|
||||
InputSchema: map[string]interface{}{
|
||||
"type": "object",
|
||||
"properties": map[string]interface{}{
|
||||
"query": map[string]interface{}{
|
||||
"type": "string",
|
||||
"description": "搜索查询内容,描述你想要了解的安全知识主题",
|
||||
},
|
||||
"risk_type": map[string]interface{}{
|
||||
"type": "string",
|
||||
"description": "可选:指定风险类型(如:SQL注入、XSS、文件上传等)。建议先调用 " + builtin.ToolListKnowledgeRiskTypes + " 工具获取可用的风险类型列表,然后使用正确的风险类型进行精确搜索,这样可以大幅减少检索时间。如果不指定则搜索所有类型。",
|
||||
},
|
||||
},
|
||||
"required": []string{"query"},
|
||||
},
|
||||
}
|
||||
|
||||
searchHandler := func(ctx context.Context, args map[string]interface{}) (*mcp.ToolResult, error) {
|
||||
query, ok := args["query"].(string)
|
||||
if !ok || query == "" {
|
||||
return &mcp.ToolResult{
|
||||
Content: []mcp.Content{
|
||||
{
|
||||
Type: "text",
|
||||
Text: "错误: 查询参数不能为空",
|
||||
},
|
||||
},
|
||||
IsError: true,
|
||||
}, nil
|
||||
}
|
||||
|
||||
riskType := ""
|
||||
if rt, ok := args["risk_type"].(string); ok && rt != "" {
|
||||
riskType = rt
|
||||
}
|
||||
|
||||
logger.Info("执行知识库检索",
|
||||
zap.String("query", query),
|
||||
zap.String("riskType", riskType),
|
||||
)
|
||||
|
||||
// 检索统一走 Retriever.Search → VectorEinoRetriever(Eino retriever 语义)。
|
||||
searchReq := &SearchRequest{
|
||||
Query: query,
|
||||
RiskType: riskType,
|
||||
TopK: 5,
|
||||
}
|
||||
|
||||
results, err := retriever.Search(ctx, searchReq)
|
||||
if err != nil {
|
||||
logger.Error("知识库检索失败", zap.Error(err))
|
||||
return &mcp.ToolResult{
|
||||
Content: []mcp.Content{
|
||||
{
|
||||
Type: "text",
|
||||
Text: fmt.Sprintf("检索失败: %v", err),
|
||||
},
|
||||
},
|
||||
IsError: true,
|
||||
}, nil
|
||||
}
|
||||
|
||||
if len(results) == 0 {
|
||||
return &mcp.ToolResult{
|
||||
Content: []mcp.Content{
|
||||
{
|
||||
Type: "text",
|
||||
Text: fmt.Sprintf("未找到与查询 '%s' 相关的知识。建议:\n1. 尝试使用不同的关键词\n2. 检查风险类型是否正确\n3. 确认知识库中是否包含相关内容", query),
|
||||
},
|
||||
},
|
||||
}, nil
|
||||
}
|
||||
|
||||
// 格式化结果
|
||||
var resultText strings.Builder
|
||||
|
||||
// 按余弦相似度(Score)降序
|
||||
sort.Slice(results, func(i, j int) bool {
|
||||
return results[i].Score > results[j].Score
|
||||
})
|
||||
|
||||
// 按文档分组结果,以便更好地展示上下文
|
||||
type itemGroup struct {
|
||||
itemID string
|
||||
results []*RetrievalResult
|
||||
maxScore float64 // 该文档块的最高相似度
|
||||
}
|
||||
itemGroups := make([]*itemGroup, 0)
|
||||
itemMap := make(map[string]*itemGroup)
|
||||
|
||||
for _, result := range results {
|
||||
itemID := result.Item.ID
|
||||
group, exists := itemMap[itemID]
|
||||
if !exists {
|
||||
group = &itemGroup{
|
||||
itemID: itemID,
|
||||
results: make([]*RetrievalResult, 0),
|
||||
maxScore: result.Score,
|
||||
}
|
||||
itemMap[itemID] = group
|
||||
itemGroups = append(itemGroups, group)
|
||||
}
|
||||
group.results = append(group.results, result)
|
||||
if result.Score > group.maxScore {
|
||||
group.maxScore = result.Score
|
||||
}
|
||||
}
|
||||
|
||||
// 按文档内最高相似度排序
|
||||
sort.Slice(itemGroups, func(i, j int) bool {
|
||||
return itemGroups[i].maxScore > itemGroups[j].maxScore
|
||||
})
|
||||
|
||||
// 收集检索到的知识项ID(用于日志)
|
||||
retrievedItemIDs := make([]string, 0, len(itemGroups))
|
||||
|
||||
resultText.WriteString(fmt.Sprintf("找到 %d 条相关知识片段:\n\n", len(results)))
|
||||
|
||||
resultIndex := 1
|
||||
for _, group := range itemGroups {
|
||||
itemResults := group.results
|
||||
mainResult := itemResults[0]
|
||||
maxScore := mainResult.Score
|
||||
for _, result := range itemResults {
|
||||
if result.Score > maxScore {
|
||||
maxScore = result.Score
|
||||
mainResult = result
|
||||
}
|
||||
}
|
||||
|
||||
// 按chunk_index排序,保证阅读的逻辑顺序(文档的原始顺序)
|
||||
sort.Slice(itemResults, func(i, j int) bool {
|
||||
return itemResults[i].Chunk.ChunkIndex < itemResults[j].Chunk.ChunkIndex
|
||||
})
|
||||
|
||||
resultText.WriteString(fmt.Sprintf("--- 结果 %d (相似度: %.2f%%) ---\n",
|
||||
resultIndex, mainResult.Similarity*100))
|
||||
resultText.WriteString(fmt.Sprintf("来源: [%s] %s (ID: %s)\n", mainResult.Item.Category, mainResult.Item.Title, mainResult.Item.ID))
|
||||
|
||||
// 按逻辑顺序显示所有chunk(包括主结果和扩展的chunk)
|
||||
if len(itemResults) == 1 {
|
||||
// 只有一个chunk,直接显示
|
||||
resultText.WriteString(fmt.Sprintf("内容片段:\n%s\n", mainResult.Chunk.ChunkText))
|
||||
} else {
|
||||
// 多个chunk,按逻辑顺序显示
|
||||
resultText.WriteString("内容片段(按文档顺序):\n")
|
||||
for i, result := range itemResults {
|
||||
// 标记主结果
|
||||
marker := ""
|
||||
if result.Chunk.ID == mainResult.Chunk.ID {
|
||||
marker = " [主匹配]"
|
||||
}
|
||||
resultText.WriteString(fmt.Sprintf(" [片段 %d%s]\n%s\n", i+1, marker, result.Chunk.ChunkText))
|
||||
}
|
||||
}
|
||||
resultText.WriteString("\n")
|
||||
|
||||
if !contains(retrievedItemIDs, group.itemID) {
|
||||
retrievedItemIDs = append(retrievedItemIDs, group.itemID)
|
||||
}
|
||||
resultIndex++
|
||||
}
|
||||
|
||||
// 在结果末尾添加元数据(JSON格式,用于提取知识项ID)
|
||||
// 使用特殊标记,避免影响AI阅读结果
|
||||
if len(retrievedItemIDs) > 0 {
|
||||
metadataJSON, _ := json.Marshal(map[string]interface{}{
|
||||
"_metadata": map[string]interface{}{
|
||||
"retrievedItemIDs": retrievedItemIDs,
|
||||
},
|
||||
})
|
||||
resultText.WriteString(fmt.Sprintf("\n<!-- 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()
|
||||
}
|
||||
@@ -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
|
||||
}
|
||||
@@ -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
|
||||
}
|
||||
@@ -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() }
|
||||
@@ -0,0 +1,11 @@
|
||||
package projectprompt
|
||||
|
||||
// ShellExecExecuteGuidanceSection 供单代理/多代理系统提示追加:exec 与 execute 分工(尽量短)。
|
||||
func ShellExecExecuteGuidanceSection() string {
|
||||
return `Shell(exec/execute):有专用 MCP 工具时优先专用工具;系统命令(管道、workdir、后台 &)用 exec;skills/ 内脚本(配合 read_file、skill)用 execute;多步扫描分拆调用,禁止一条 shell 串多个扫描器。长脚本、请求体或 Payload 必须先用 write_file 写入会话工作目录,再用 exec/execute 执行短命令;禁止把长内容嵌入 command。下载/临时文件须写入系统提示中的「会话工作目录」,禁止用 /tmp。`
|
||||
}
|
||||
|
||||
// ShellExecExecuteGuidanceReconSuffix 侦察子代理可选追加(一行)。
|
||||
func ShellExecExecuteGuidanceReconSuffix() string {
|
||||
return `枚举优先 subfinder、amass 等专用 MCP,勿 exec/execute 拼长链。`
|
||||
}
|
||||
@@ -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))
|
||||
}
|
||||
@@ -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")
|
||||
}
|
||||
}
|
||||
@@ -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
|
||||
}
|
||||
@@ -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
|
||||
}
|
||||
@@ -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)
|
||||
}
|
||||
@@ -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 个 output;output/end 不能有出边。
|
||||
- 节点 type 只能从这些字符串中选择:start、tool、agent、condition、hitl、output、end。
|
||||
- 每个 agent、tool、output 节点都必须配置唯一的 output_key;output 节点默认使用 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_seconds;arguments 必须是合法 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())
|
||||
}
|
||||
@@ -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)
|
||||
}
|
||||
}
|
||||
@@ -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
|
||||
}
|
||||
@@ -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
|
||||
}
|
||||
@@ -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,
|
||||
})
|
||||
}
|
||||
@@ -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
|
||||
}
|
||||
@@ -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")
|
||||
}
|
||||
}
|
||||
@@ -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
|
||||
}
|
||||
@@ -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
|
||||
}
|
||||
@@ -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)
|
||||
}
|
||||
@@ -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
|
||||
}
|
||||
@@ -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()
|
||||
}
|
||||
@@ -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"
|
||||
}
|
||||
@@ -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
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -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
|
||||
}
|
||||
@@ -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
|
||||
}
|
||||
@@ -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)
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -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),
|
||||
})
|
||||
}
|
||||
}
|
||||
@@ -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", ""
|
||||
}
|
||||
@@ -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)
|
||||
}
|
||||
}
|
||||
@@ -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
|
||||
}
|
||||
@@ -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()
|
||||
}
|
||||
@@ -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)
|
||||
}
|
||||
@@ -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
|
||||
}
|
||||
@@ -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())
|
||||
}
|
||||
@@ -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)
|
||||
}
|
||||
}
|
||||
@@ -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)
|
||||
}
|
||||
@@ -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
|
||||
}
|
||||
Reference in New Issue
Block a user