xinmi/internal/workbench/agent_snapshot.go

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,
	}
}