WeKnora deployment source with the reviewed wiki identity, summary and recovery fixes

internal/workbench/agent_snapshot.go

@@ -0,0 +1,85 @@
+package workbench
+
+import (
+	"context"
+	"time"
+
+	"github.com/Tencent/WeKnora/internal/types"
+)
+
+// Store generation inputs, never provider credentials or executable tool config.
+type AgentSnapshot struct {
+	Name            string              `json:"name"`
+	ModelID         string              `json:"model_id"`
+	SystemPrompt    string              `json:"system_prompt"`
+	Instructions    string              `json:"instructions"`
+	SkillIDs        []string            `json:"skill_ids"`
+	KBSelectionMode string              `json:"kb_selection_mode"`
+	KnowledgeBases  []string            `json:"knowledge_bases"`
+	Temperature     float64             `json:"temperature"`
+	Thinking        *bool               `json:"thinking"`
+	MemoryEnabled   *bool               `json:"memory_enabled"`
+	OutputLimit     int                 `json:"output_limit"`
+	ContextWindow   int                 `json:"context_window"`
+	CapturedAt      time.Time           `json:"captured_at"`
+	Fusions         []FusionSnapshot    `json:"fusions,omitempty"`
+	FusionMemories  []*types.MemoryItem `json:"fusion_memories,omitempty"`
+}
+
+type FusionSnapshot struct {
+	ID            string   `json:"id"`
+	Name          string   `json:"name"`
+	SystemPrompt  string   `json:"system_prompt"`
+	Instructions  string   `json:"instructions"`
+	SkillIDs      []string `json:"skill_ids"`
+	MemoryEnabled *bool    `json:"memory_enabled"`
+}
+
+type agentSnapshotsKey struct{}
+
+func snapshotFromContext(ctx context.Context, id string) (AgentSnapshot, bool) {
+	m, _ := ctx.Value(agentSnapshotsKey{}).(map[string]AgentSnapshot)
+	v, ok := m[id]
+	return v, ok
+}
+
+func applyAgentSnapshot(agent *types.CustomAgent, snapshot AgentSnapshot) *types.CustomAgent {
+	copy := *agent
+	copy.Name = snapshot.Name
+	copy.Config.ModelID = snapshot.ModelID
+	copy.Config.SystemPrompt = snapshot.SystemPrompt
+	copy.Config.InstructionSkillIDs = append([]string(nil), snapshot.SkillIDs...)
+	copy.Config.KBSelectionMode = snapshot.KBSelectionMode
+	copy.Config.KnowledgeBases = append([]string(nil), snapshot.KnowledgeBases...)
+	copy.Config.Temperature = snapshot.Temperature
+	copy.Config.Thinking = snapshot.Thinking
+	copy.Config.MemoryEnabled = snapshot.MemoryEnabled
+	// A saved generation configuration cannot undo a later privacy opt-out.
+	if agent.Config.MemoryEnabled != nil && !*agent.Config.MemoryEnabled {
+		copy.Config.MemoryEnabled = agent.Config.MemoryEnabled
+	}
+	copy.Config.MaxCompletionTokens = snapshot.OutputLimit
+	copy.Config.FusionAgentIDs = nil
+	for _, fusion := range snapshot.Fusions {
+		copy.Config.FusionAgentIDs = append(copy.Config.FusionAgentIDs, fusion.ID)
+	}
+	return &copy
+}
+
+func rememberAgentSnapshot(ctx context.Context, agent *types.CustomAgent, instructions string, contextWindow, outputLimit int, fusions []FusionSnapshot, memories []*types.MemoryItem) {
+	m, ok := ctx.Value(agentSnapshotsKey{}).(map[string]AgentSnapshot)
+	if !ok {
+		return
+	}
+	if _, exists := m[agent.ID]; exists {
+		return
+	}
+	m[agent.ID] = AgentSnapshot{
+		Name: agent.Name, ModelID: agent.Config.ModelID, SystemPrompt: agent.Config.SystemPrompt,
+		Instructions: instructions, SkillIDs: append([]string(nil), agent.Config.InstructionSkillIDs...),
+		KBSelectionMode: agent.Config.KBSelectionMode, KnowledgeBases: append([]string(nil), agent.Config.KnowledgeBases...),
+		Temperature: agent.Config.Temperature, Thinking: agent.Config.Thinking, MemoryEnabled: agent.Config.MemoryEnabled,
+		OutputLimit: outputLimit, ContextWindow: contextWindow, CapturedAt: time.Now().UTC(),
+		Fusions: fusions, FusionMemories: memories,
+	}
+}