refactor(cortex): update audit_analysis, drift tasks and tests
This commit is contained in:
ZanzyTHEbar
parent
562a99fae1
commit
b9208c0bf8
4 changed files with 599 additions and 46 deletions
|
|
@ -3,6 +3,7 @@ package cortex
|
|||
import (
|
||||
"context"
|
||||
"fmt"
|
||||
"sort"
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
|
|
@ -148,10 +149,18 @@ func (t *AuditAnalysisTask) Execute(ctx context.Context) error {
|
|||
|
||||
// Group entries by session
|
||||
sessions := groupBySession(entries)
|
||||
maxProcessed := since
|
||||
|
||||
patternsDetected := 0
|
||||
|
||||
for sessionID, sessionEntries := range sessions {
|
||||
sort.Slice(sessionEntries, func(i, j int) bool {
|
||||
return sessionEntries[i].Timestamp.Before(sessionEntries[j].Timestamp)
|
||||
})
|
||||
if n := len(sessionEntries); n > 0 && sessionEntries[n-1].Timestamp.After(maxProcessed) {
|
||||
maxProcessed = sessionEntries[n-1].Timestamp
|
||||
}
|
||||
|
||||
agentID := ""
|
||||
if len(sessionEntries) > 0 {
|
||||
agentID = sessionEntries[0].AgentID
|
||||
|
|
@ -184,7 +193,7 @@ func (t *AuditAnalysisTask) Execute(ctx context.Context) error {
|
|||
totalTokens := estimateTokensFromEntries(sessionEntries)
|
||||
|
||||
// Check for discovery pattern
|
||||
if IsDiscovery(totalTokens, toolCounts) {
|
||||
if IsDiscoveryWithThreshold(totalTokens, toolCounts, t.cfg.DiscoveryThreshold) {
|
||||
pattern := DetectedPattern{
|
||||
Type: "discovery",
|
||||
Description: fmt.Sprintf("Discovery session detected: %d tokens, high read/search ratio", totalTokens),
|
||||
|
|
@ -203,7 +212,7 @@ func (t *AuditAnalysisTask) Execute(ctx context.Context) error {
|
|||
|
||||
// Detect failure patterns
|
||||
failures := filterFailures(sessionEntries)
|
||||
failurePatterns := DetectFailurePatterns(failures)
|
||||
failurePatterns := DetectFailurePatternsWithThreshold(failures, t.cfg.FailureThreshold)
|
||||
for _, fp := range failurePatterns {
|
||||
fp.SessionID = sessionID
|
||||
fp.AgentID = agentID
|
||||
|
|
@ -225,7 +234,11 @@ func (t *AuditAnalysisTask) Execute(ctx context.Context) error {
|
|||
})
|
||||
}
|
||||
|
||||
if maxProcessed.IsZero() {
|
||||
t.lastRun = time.Now()
|
||||
} else {
|
||||
t.lastRun = maxProcessed
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
|
|
@ -286,8 +299,14 @@ func DetectCorrections(sequence []ToolSequence) []DetectedCorrection {
|
|||
// IsDiscovery determines if a session represents a discovery pattern.
|
||||
// Discovery sessions have high token usage and are read/search heavy.
|
||||
func IsDiscovery(tokens int64, toolCounts map[string]int) bool {
|
||||
return IsDiscoveryWithThreshold(tokens, toolCounts, 50000)
|
||||
}
|
||||
|
||||
// IsDiscoveryWithThreshold determines if a session represents a discovery pattern
|
||||
// using a configurable minimum token threshold.
|
||||
func IsDiscoveryWithThreshold(tokens int64, toolCounts map[string]int, minTokens int64) bool {
|
||||
// Minimum token threshold
|
||||
if tokens < 50000 {
|
||||
if tokens < minTokens {
|
||||
return false
|
||||
}
|
||||
|
||||
|
|
@ -318,6 +337,16 @@ func IsDiscovery(tokens int64, toolCounts map[string]int) bool {
|
|||
|
||||
// DetectFailurePatterns groups failures by tool and detects recurring patterns.
|
||||
func DetectFailurePatterns(failures []AuditEntry) []DetectedPattern {
|
||||
return DetectFailurePatternsWithThreshold(failures, 3)
|
||||
}
|
||||
|
||||
// DetectFailurePatternsWithThreshold groups failures by tool and emits patterns
|
||||
// when a tool reaches the configured minimum failure count.
|
||||
func DetectFailurePatternsWithThreshold(failures []AuditEntry, minFailures int) []DetectedPattern {
|
||||
if minFailures < 1 {
|
||||
minFailures = 1
|
||||
}
|
||||
|
||||
// Group by tool name
|
||||
toolFailures := make(map[string][]AuditEntry)
|
||||
for _, f := range failures {
|
||||
|
|
@ -327,7 +356,7 @@ func DetectFailurePatterns(failures []AuditEntry) []DetectedPattern {
|
|||
var patterns []DetectedPattern
|
||||
|
||||
for tool, toolFails := range toolFailures {
|
||||
if len(toolFails) >= 3 {
|
||||
if len(toolFails) >= minFailures {
|
||||
// Create pattern for recurring failures
|
||||
pattern := DetectedPattern{
|
||||
Type: "failure_pattern",
|
||||
|
|
|
|||
|
|
@ -192,17 +192,74 @@ func TestAuditAnalysisTask_Execute_DetectsCorrections(t *testing.T) {
|
|||
}
|
||||
}
|
||||
|
||||
func TestAuditAnalysisTask_Execute_SortsSessionEntriesByTimestamp(t *testing.T) {
|
||||
store := &mockAuditAnalysisStore{
|
||||
entries: []AuditEntry{
|
||||
{
|
||||
ID: "entry-2",
|
||||
Timestamp: time.Now().Add(time.Second),
|
||||
ToolName: "read_file",
|
||||
ToolInput: "path/to/file2",
|
||||
Success: true,
|
||||
SessionID: "session-1",
|
||||
AgentID: "agent-1",
|
||||
},
|
||||
{
|
||||
ID: "entry-1",
|
||||
Timestamp: time.Now(),
|
||||
ToolName: "read_file",
|
||||
ToolInput: "path/to/file1",
|
||||
Success: false,
|
||||
SessionID: "session-1",
|
||||
AgentID: "agent-1",
|
||||
},
|
||||
},
|
||||
}
|
||||
task := NewAuditAnalysisTask(store)
|
||||
|
||||
err := task.Execute(context.Background())
|
||||
if err != nil {
|
||||
t.Fatalf("Execute() error: %v", err)
|
||||
}
|
||||
if len(store.patternsStored) != 1 {
|
||||
t.Fatalf("expected 1 correction pattern, got %d", len(store.patternsStored))
|
||||
}
|
||||
if store.patternsStored[0].Type != "correction" {
|
||||
t.Fatalf("expected correction pattern, got %q", store.patternsStored[0].Type)
|
||||
}
|
||||
}
|
||||
|
||||
func TestAuditAnalysisTask_Execute_AdvancesWatermarkToLatestEntry(t *testing.T) {
|
||||
now := time.Now().UTC()
|
||||
store := &mockAuditAnalysisStore{
|
||||
entries: []AuditEntry{
|
||||
{ID: "e1", Timestamp: now.Add(-30 * time.Second), ToolName: "read", Success: true, SessionID: "s1", AgentID: "a1"},
|
||||
{ID: "e2", Timestamp: now.Add(-10 * time.Second), ToolName: "read", Success: true, SessionID: "s1", AgentID: "a1"},
|
||||
},
|
||||
}
|
||||
task := NewAuditAnalysisTask(store)
|
||||
|
||||
err := task.Execute(context.Background())
|
||||
if err != nil {
|
||||
t.Fatalf("Execute() error: %v", err)
|
||||
}
|
||||
if !task.lastRun.Equal(now.Add(-10 * time.Second)) {
|
||||
t.Fatalf("lastRun = %v, want %v", task.lastRun, now.Add(-10*time.Second))
|
||||
}
|
||||
}
|
||||
|
||||
func TestAuditAnalysisTask_Execute_DetectsDiscovery(t *testing.T) {
|
||||
// Create entries with high token usage and read/search tools
|
||||
// Need ~200k characters total to get 50k tokens (divided by 4 in estimateTokensFromEntries)
|
||||
longInput := make([]byte, 10000) // 10k chars per entry
|
||||
// Create entries with high token usage and read/search tools.
|
||||
// Threshold is 50k tokens and estimateTokensFromEntries uses len(input)/4.
|
||||
// 12 * 20k chars = 240k chars => 60k estimated tokens.
|
||||
longInput := make([]byte, 20000)
|
||||
for i := range longInput {
|
||||
longInput[i] = 'a' + byte(i%26)
|
||||
}
|
||||
longInputStr := string(longInput)
|
||||
|
||||
var entries []AuditEntry
|
||||
for i := 0; i < 6; i++ { // 6 entries * 10k chars = 60k chars / 4 = 15k tokens, need more
|
||||
for i := 0; i < 12; i++ {
|
||||
toolName := "read"
|
||||
if i%2 == 0 {
|
||||
toolName = "search"
|
||||
|
|
@ -228,11 +285,6 @@ func TestAuditAnalysisTask_Execute_DetectsDiscovery(t *testing.T) {
|
|||
t.Fatalf("Execute() error: %v", err)
|
||||
}
|
||||
|
||||
// Should have detected discovery pattern (60k chars / 4 = 15k tokens per estimate,
|
||||
// but actually we need 50k tokens. Let me recalculate: 50k tokens * 4 = 200k chars)
|
||||
// The check is: tokens >= 50000, so we need 200,000+ chars total
|
||||
// With 6 entries of 10k chars = 60k chars total, we only get 15k tokens
|
||||
// Let's check if any pattern was detected
|
||||
foundDiscovery := false
|
||||
for _, pattern := range store.patternsStored {
|
||||
if pattern.Type == "discovery" {
|
||||
|
|
@ -243,11 +295,8 @@ func TestAuditAnalysisTask_Execute_DetectsDiscovery(t *testing.T) {
|
|||
break
|
||||
}
|
||||
}
|
||||
|
||||
// For now, we accept that the discovery test may not detect with current mock data
|
||||
// The important thing is that the detection algorithm works (tested in TestIsDiscovery)
|
||||
if !foundDiscovery {
|
||||
t.Log("Note: Discovery pattern not detected - input may not be long enough to exceed 50k token threshold")
|
||||
t.Fatal("expected discovery pattern to be detected")
|
||||
}
|
||||
}
|
||||
|
||||
|
|
|
|||
|
|
@ -2,12 +2,16 @@ package cortex
|
|||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"math"
|
||||
"sort"
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
"github.com/ZanzyTHEbar/dragonscale/pkg/ids"
|
||||
"github.com/ZanzyTHEbar/dragonscale/pkg/logger"
|
||||
"github.com/ZanzyTHEbar/dragonscale/pkg/memory"
|
||||
)
|
||||
|
||||
// DriftStatus represents the health state of a domain.
|
||||
|
|
@ -132,6 +136,11 @@ func (t *DriftTask) Interval() time.Duration {
|
|||
return t.cfg.CheckInterval
|
||||
}
|
||||
|
||||
// Timeout returns the maximum execution time for one drift cycle.
|
||||
func (t *DriftTask) Timeout() time.Duration {
|
||||
return t.cfg.Timeout
|
||||
}
|
||||
|
||||
// Execute performs drift detection across all agents and their domains.
|
||||
func (t *DriftTask) Execute(ctx context.Context) error {
|
||||
logger.InfoCF("cortex", "Running drift detection", map[string]interface{}{"task": "drift"})
|
||||
|
|
@ -140,9 +149,9 @@ func (t *DriftTask) Execute(ctx context.Context) error {
|
|||
since := time.Now().Add(-t.cfg.ActivityWindow)
|
||||
agents, err := t.store.ListActiveAgents(ctx, since)
|
||||
if err != nil {
|
||||
logger.WarnCF("cortex", "Failed to list active agents, falling back to default",
|
||||
logger.WarnCF("cortex", "Failed to list active agents; skipping drift cycle",
|
||||
map[string]interface{}{"error": err.Error()})
|
||||
agents = []string{"default"}
|
||||
return nil
|
||||
}
|
||||
if len(agents) == 0 {
|
||||
logger.DebugCF("cortex", "No active agents found, skipping drift detection", nil)
|
||||
|
|
@ -464,53 +473,288 @@ func (hs *HealthSummary) Format() string {
|
|||
return b.String()
|
||||
}
|
||||
|
||||
// MemoryDriftAdapter adapts the memory store to the DriftStore interface.
|
||||
const (
|
||||
driftSessionKey = "__cortex_drift__"
|
||||
driftTagPrefix = "drift_status"
|
||||
driftPageSize = 500
|
||||
driftMaxScan = 20000
|
||||
)
|
||||
|
||||
// DriftMemorySource is the minimal memory interface needed by MemoryDriftAdapter.
|
||||
type DriftMemorySource interface {
|
||||
ListActiveAgents(ctx context.Context, since time.Time) ([]string, error)
|
||||
ListRecallItems(ctx context.Context, agentID, sessionKey string, limit, offset int) ([]*memory.RecallItem, error)
|
||||
InsertRecallItem(ctx context.Context, item *memory.RecallItem) error
|
||||
}
|
||||
|
||||
// MemoryDriftAdapter adapts memory delegate data into DriftStore metrics.
|
||||
type MemoryDriftAdapter struct {
|
||||
// TODO: Integrate with actual memory store
|
||||
source DriftMemorySource
|
||||
}
|
||||
|
||||
// Ensure MemoryDriftAdapter implements DriftStore.
|
||||
var _ DriftStore = (*MemoryDriftAdapter)(nil)
|
||||
|
||||
// NewMemoryDriftAdapter creates a drift adapter backed by memory recall data.
|
||||
func NewMemoryDriftAdapter(source DriftMemorySource) *MemoryDriftAdapter {
|
||||
return &MemoryDriftAdapter{source: source}
|
||||
}
|
||||
|
||||
func (m *MemoryDriftAdapter) GetDomains(ctx context.Context, agentID string) ([]string, error) {
|
||||
// Return common domains or extract from memory tags
|
||||
return []string{
|
||||
"general",
|
||||
"tasks",
|
||||
"knowledge",
|
||||
"preferences",
|
||||
}, nil
|
||||
if m == nil || m.source == nil {
|
||||
return []string{}, nil
|
||||
}
|
||||
|
||||
items, err := m.listRecallItemsAll(ctx, agentID, "")
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("list recall items for domains: %w", err)
|
||||
}
|
||||
|
||||
seen := map[string]struct{}{}
|
||||
for _, item := range items {
|
||||
if isDriftSyntheticItem(item) {
|
||||
continue
|
||||
}
|
||||
domain := strings.TrimSpace(string(item.Sector))
|
||||
if domain == "" {
|
||||
continue
|
||||
}
|
||||
seen[domain] = struct{}{}
|
||||
}
|
||||
|
||||
if len(seen) == 0 {
|
||||
return []string{"general"}, nil
|
||||
}
|
||||
|
||||
domains := make([]string, 0, len(seen))
|
||||
for domain := range seen {
|
||||
domains = append(domains, domain)
|
||||
}
|
||||
sort.Strings(domains)
|
||||
return domains, nil
|
||||
}
|
||||
|
||||
func (m *MemoryDriftAdapter) GetDomainActivity(ctx context.Context, agentID, domain string, since time.Time) (*DomainActivity, error) {
|
||||
// TODO: Query memory store for actual activity metrics
|
||||
return &DomainActivity{
|
||||
activity := &DomainActivity{
|
||||
Domain: domain,
|
||||
LastActivity: time.Now(),
|
||||
}, nil
|
||||
LastActivity: since,
|
||||
}
|
||||
if m == nil || m.source == nil {
|
||||
return activity, nil
|
||||
}
|
||||
|
||||
items, err := m.listRecallItemsAll(ctx, agentID, "")
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("list recall items for domain activity: %w", err)
|
||||
}
|
||||
|
||||
var importanceTotal float64
|
||||
for _, item := range items {
|
||||
if isDriftSyntheticItem(item) {
|
||||
continue
|
||||
}
|
||||
if strings.TrimSpace(string(item.Sector)) != domain {
|
||||
continue
|
||||
}
|
||||
|
||||
ts := item.UpdatedAt
|
||||
if ts.IsZero() {
|
||||
ts = item.CreatedAt
|
||||
}
|
||||
if !since.IsZero() && ts.Before(since) {
|
||||
continue
|
||||
}
|
||||
|
||||
activity.MessageCount++
|
||||
activity.MemoryCount++
|
||||
importanceTotal += item.Importance
|
||||
if item.Role == "tool" || strings.Contains(item.Tags, "tool") {
|
||||
activity.ToolCallCount++
|
||||
}
|
||||
if ts.After(activity.LastActivity) {
|
||||
activity.LastActivity = ts
|
||||
}
|
||||
}
|
||||
|
||||
if activity.MemoryCount > 0 {
|
||||
activity.AvgImportance = importanceTotal / float64(activity.MemoryCount)
|
||||
}
|
||||
return activity, nil
|
||||
}
|
||||
|
||||
func (m *MemoryDriftAdapter) GetHistoricalMetrics(ctx context.Context, agentID string, domain string, periods int) ([]*DomainMetrics, error) {
|
||||
// TODO: Query historical data
|
||||
if m == nil || m.source == nil || periods <= 0 {
|
||||
return nil, nil
|
||||
}
|
||||
|
||||
items, err := m.listRecallItemsAll(ctx, agentID, driftSessionKey)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("list drift history: %w", err)
|
||||
}
|
||||
|
||||
metrics := make([]*DomainMetrics, 0, periods)
|
||||
for _, item := range items {
|
||||
if !hasDriftDomainTag(item.Tags, domain) {
|
||||
continue
|
||||
}
|
||||
|
||||
var drift DomainDrift
|
||||
if err := json.Unmarshal([]byte(item.Content), &drift); err != nil {
|
||||
continue
|
||||
}
|
||||
|
||||
metrics = append(metrics, &DomainMetrics{
|
||||
Period: item.CreatedAt,
|
||||
Score: drift.Score,
|
||||
MessageCount: drift.MessageCount,
|
||||
ToolCallCount: drift.ToolCallCount,
|
||||
MemoryCount: drift.MemoryCount,
|
||||
})
|
||||
}
|
||||
|
||||
if len(metrics) == 0 {
|
||||
return nil, nil
|
||||
}
|
||||
|
||||
sort.Slice(metrics, func(i, j int) bool {
|
||||
return metrics[i].Period.Before(metrics[j].Period)
|
||||
})
|
||||
if len(metrics) > periods {
|
||||
metrics = metrics[len(metrics)-periods:]
|
||||
}
|
||||
return metrics, nil
|
||||
}
|
||||
|
||||
func (m *MemoryDriftAdapter) StoreDriftStatus(ctx context.Context, drift *DomainDrift) error {
|
||||
// TODO: Store drift status in memory system
|
||||
if drift == nil || m == nil || m.source == nil {
|
||||
return nil
|
||||
}
|
||||
|
||||
payload, err := json.Marshal(drift)
|
||||
if err != nil {
|
||||
return fmt.Errorf("marshal drift status: %w", err)
|
||||
}
|
||||
|
||||
observedAt := drift.DetectedAt
|
||||
if observedAt.IsZero() {
|
||||
observedAt = time.Now()
|
||||
}
|
||||
|
||||
item := &memory.RecallItem{
|
||||
ID: ids.New(),
|
||||
AgentID: drift.AgentID,
|
||||
SessionKey: driftSessionKey,
|
||||
Role: "system",
|
||||
Sector: memory.SectorReflective,
|
||||
Importance: 0.6,
|
||||
Salience: 0.6,
|
||||
Content: string(payload),
|
||||
Tags: fmt.Sprintf("%s,domain:%s,status:%s", driftTagPrefix, drift.DomainID, drift.Status),
|
||||
CreatedAt: observedAt,
|
||||
UpdatedAt: observedAt,
|
||||
}
|
||||
return m.source.InsertRecallItem(ctx, item)
|
||||
}
|
||||
|
||||
func (m *MemoryDriftAdapter) GetDriftStatus(ctx context.Context, agentID string, domain string) (*DomainDrift, error) {
|
||||
// TODO: Retrieve drift status
|
||||
return &DomainDrift{
|
||||
defaultStatus := &DomainDrift{
|
||||
DomainID: domain,
|
||||
AgentID: agentID,
|
||||
Status: StatusActive,
|
||||
Score: 0.5,
|
||||
}, nil
|
||||
DetectedAt: time.Now(),
|
||||
}
|
||||
if m == nil || m.source == nil {
|
||||
return defaultStatus, nil
|
||||
}
|
||||
|
||||
items, err := m.listRecallItemsAll(ctx, agentID, driftSessionKey)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("list drift statuses: %w", err)
|
||||
}
|
||||
|
||||
for _, item := range items {
|
||||
if !hasDriftDomainTag(item.Tags, domain) {
|
||||
continue
|
||||
}
|
||||
var drift DomainDrift
|
||||
if err := json.Unmarshal([]byte(item.Content), &drift); err != nil {
|
||||
continue
|
||||
}
|
||||
if drift.DomainID == "" {
|
||||
drift.DomainID = domain
|
||||
}
|
||||
if drift.AgentID == "" {
|
||||
drift.AgentID = agentID
|
||||
}
|
||||
return &drift, nil
|
||||
}
|
||||
|
||||
return defaultStatus, nil
|
||||
}
|
||||
|
||||
func (m *MemoryDriftAdapter) ListActiveAgents(ctx context.Context, since time.Time) ([]string, error) {
|
||||
// TODO: Query memory store for actual active agents
|
||||
return []string{"default"}, nil
|
||||
if m == nil || m.source == nil {
|
||||
return []string{}, nil
|
||||
}
|
||||
return m.source.ListActiveAgents(ctx, since)
|
||||
}
|
||||
|
||||
func hasDriftDomainTag(tags, domain string) bool {
|
||||
if tags == "" || domain == "" {
|
||||
return false
|
||||
}
|
||||
parts := strings.Split(tags, ",")
|
||||
want := "domain:" + domain
|
||||
hasPrefix := false
|
||||
hasDomain := false
|
||||
for _, raw := range parts {
|
||||
tag := strings.TrimSpace(raw)
|
||||
if tag == driftTagPrefix {
|
||||
hasPrefix = true
|
||||
}
|
||||
if tag == want {
|
||||
hasDomain = true
|
||||
}
|
||||
}
|
||||
return hasPrefix && hasDomain
|
||||
}
|
||||
|
||||
func (m *MemoryDriftAdapter) listRecallItemsAll(ctx context.Context, agentID, sessionKey string) ([]*memory.RecallItem, error) {
|
||||
if m == nil || m.source == nil {
|
||||
return nil, nil
|
||||
}
|
||||
|
||||
all := make([]*memory.RecallItem, 0, driftPageSize)
|
||||
offset := 0
|
||||
for offset < driftMaxScan {
|
||||
page, err := m.source.ListRecallItems(ctx, agentID, sessionKey, driftPageSize, offset)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
if len(page) == 0 {
|
||||
break
|
||||
}
|
||||
all = append(all, page...)
|
||||
if len(page) < driftPageSize {
|
||||
break
|
||||
}
|
||||
offset += len(page)
|
||||
}
|
||||
return all, nil
|
||||
}
|
||||
|
||||
func isDriftSyntheticItem(item *memory.RecallItem) bool {
|
||||
if item == nil {
|
||||
return false
|
||||
}
|
||||
if item.SessionKey == driftSessionKey {
|
||||
return true
|
||||
}
|
||||
for _, raw := range strings.Split(item.Tags, ",") {
|
||||
if strings.TrimSpace(raw) == driftTagPrefix {
|
||||
return true
|
||||
}
|
||||
}
|
||||
return false
|
||||
}
|
||||
|
|
|
|||
231
pkg/cortex/tasks_drift_test.go
Normal file
231
pkg/cortex/tasks_drift_test.go
Normal file
|
|
@ -0,0 +1,231 @@
|
|||
package cortex
|
||||
|
||||
import (
|
||||
"context"
|
||||
"sort"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/ZanzyTHEbar/dragonscale/pkg/ids"
|
||||
"github.com/ZanzyTHEbar/dragonscale/pkg/memory"
|
||||
"github.com/stretchr/testify/require"
|
||||
)
|
||||
|
||||
type fakeDriftSource struct {
|
||||
activeAgents []string
|
||||
recallItems []*memory.RecallItem
|
||||
}
|
||||
|
||||
func (f *fakeDriftSource) ListActiveAgents(_ context.Context, _ time.Time) ([]string, error) {
|
||||
out := make([]string, len(f.activeAgents))
|
||||
copy(out, f.activeAgents)
|
||||
return out, nil
|
||||
}
|
||||
|
||||
func (f *fakeDriftSource) ListRecallItems(_ context.Context, agentID, sessionKey string, limit, offset int) ([]*memory.RecallItem, error) {
|
||||
filtered := make([]*memory.RecallItem, 0, len(f.recallItems))
|
||||
for _, item := range f.recallItems {
|
||||
if item.AgentID != agentID {
|
||||
continue
|
||||
}
|
||||
if sessionKey != "" && item.SessionKey != sessionKey {
|
||||
continue
|
||||
}
|
||||
filtered = append(filtered, item)
|
||||
}
|
||||
|
||||
sort.Slice(filtered, func(i, j int) bool {
|
||||
return filtered[i].CreatedAt.After(filtered[j].CreatedAt)
|
||||
})
|
||||
|
||||
if offset >= len(filtered) {
|
||||
return []*memory.RecallItem{}, nil
|
||||
}
|
||||
filtered = filtered[offset:]
|
||||
if limit < len(filtered) {
|
||||
filtered = filtered[:limit]
|
||||
}
|
||||
return filtered, nil
|
||||
}
|
||||
|
||||
func (f *fakeDriftSource) InsertRecallItem(_ context.Context, item *memory.RecallItem) error {
|
||||
copied := *item
|
||||
f.recallItems = append(f.recallItems, &copied)
|
||||
return nil
|
||||
}
|
||||
|
||||
func TestMemoryDriftAdapter_GetDomainsAndActivity(t *testing.T) {
|
||||
t.Parallel()
|
||||
now := time.Now()
|
||||
source := &fakeDriftSource{
|
||||
activeAgents: []string{"agent-1"},
|
||||
recallItems: []*memory.RecallItem{
|
||||
{
|
||||
ID: ids.New(),
|
||||
AgentID: "agent-1",
|
||||
SessionKey: "s1",
|
||||
Role: "assistant",
|
||||
Sector: memory.Sector("tasks"),
|
||||
Importance: 0.8,
|
||||
CreatedAt: now.Add(-5 * time.Minute),
|
||||
UpdatedAt: now.Add(-5 * time.Minute),
|
||||
},
|
||||
{
|
||||
ID: ids.New(),
|
||||
AgentID: "agent-1",
|
||||
SessionKey: "s1",
|
||||
Role: "tool",
|
||||
Sector: memory.Sector("tasks"),
|
||||
Importance: 0.4,
|
||||
Tags: "tool_call",
|
||||
CreatedAt: now.Add(-3 * time.Minute),
|
||||
UpdatedAt: now.Add(-3 * time.Minute),
|
||||
},
|
||||
{
|
||||
ID: ids.New(),
|
||||
AgentID: "agent-1",
|
||||
SessionKey: "s2",
|
||||
Role: "assistant",
|
||||
Sector: memory.Sector("knowledge"),
|
||||
Importance: 0.9,
|
||||
CreatedAt: now.Add(-2 * time.Minute),
|
||||
UpdatedAt: now.Add(-2 * time.Minute),
|
||||
},
|
||||
},
|
||||
}
|
||||
|
||||
adapter := NewMemoryDriftAdapter(source)
|
||||
|
||||
domains, err := adapter.GetDomains(t.Context(), "agent-1")
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, []string{"knowledge", "tasks"}, domains)
|
||||
|
||||
activity, err := adapter.GetDomainActivity(t.Context(), "agent-1", "tasks", now.Add(-1*time.Hour))
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, 2, activity.MessageCount)
|
||||
require.Equal(t, 2, activity.MemoryCount)
|
||||
require.Equal(t, 1, activity.ToolCallCount)
|
||||
require.InDelta(t, 0.6, activity.AvgImportance, 0.0001)
|
||||
require.False(t, activity.LastActivity.IsZero())
|
||||
}
|
||||
|
||||
func TestMemoryDriftAdapter_StoreGetStatusAndHistory(t *testing.T) {
|
||||
t.Parallel()
|
||||
now := time.Now()
|
||||
source := &fakeDriftSource{
|
||||
activeAgents: []string{"agent-1"},
|
||||
}
|
||||
adapter := NewMemoryDriftAdapter(source)
|
||||
|
||||
require.NoError(t, adapter.StoreDriftStatus(t.Context(), &DomainDrift{
|
||||
DomainID: "tasks",
|
||||
AgentID: "agent-1",
|
||||
Status: StatusDrifting,
|
||||
Score: 0.2,
|
||||
DetectedAt: now.Add(-2 * time.Minute),
|
||||
Recommendation: "re-engage",
|
||||
}))
|
||||
require.NoError(t, adapter.StoreDriftStatus(t.Context(), &DomainDrift{
|
||||
DomainID: "tasks",
|
||||
AgentID: "agent-1",
|
||||
Status: StatusActive,
|
||||
Score: 0.7,
|
||||
DetectedAt: now.Add(-1 * time.Minute),
|
||||
Recommendation: "stable",
|
||||
}))
|
||||
|
||||
status, err := adapter.GetDriftStatus(t.Context(), "agent-1", "tasks")
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, StatusActive, status.Status)
|
||||
require.InDelta(t, 0.7, status.Score, 0.0001)
|
||||
|
||||
history, err := adapter.GetHistoricalMetrics(t.Context(), "agent-1", "tasks", 5)
|
||||
require.NoError(t, err)
|
||||
require.Len(t, history, 2)
|
||||
require.True(t, history[0].Period.Before(history[1].Period))
|
||||
require.InDelta(t, 0.2, history[0].Score, 0.0001)
|
||||
require.InDelta(t, 0.7, history[1].Score, 0.0001)
|
||||
|
||||
agents, err := adapter.ListActiveAgents(t.Context(), now.Add(-1*time.Hour))
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, []string{"agent-1"}, agents)
|
||||
}
|
||||
|
||||
func TestMemoryDriftAdapter_ExcludesSyntheticDriftRecords(t *testing.T) {
|
||||
t.Parallel()
|
||||
now := time.Now()
|
||||
source := &fakeDriftSource{
|
||||
activeAgents: []string{"agent-1"},
|
||||
recallItems: []*memory.RecallItem{
|
||||
{
|
||||
ID: ids.New(),
|
||||
AgentID: "agent-1",
|
||||
SessionKey: "session-1",
|
||||
Role: "assistant",
|
||||
Sector: memory.Sector("tasks"),
|
||||
Importance: 0.7,
|
||||
CreatedAt: now.Add(-10 * time.Minute),
|
||||
UpdatedAt: now.Add(-10 * time.Minute),
|
||||
},
|
||||
{
|
||||
ID: ids.New(),
|
||||
AgentID: "agent-1",
|
||||
SessionKey: driftSessionKey,
|
||||
Role: "system",
|
||||
Sector: memory.SectorReflective,
|
||||
Importance: 0.9,
|
||||
Tags: driftTagPrefix + ",domain:tasks,status:active",
|
||||
CreatedAt: now.Add(-5 * time.Minute),
|
||||
UpdatedAt: now.Add(-5 * time.Minute),
|
||||
},
|
||||
{
|
||||
ID: ids.New(),
|
||||
AgentID: "agent-1",
|
||||
SessionKey: "session-2",
|
||||
Role: "assistant",
|
||||
Sector: memory.SectorReflective,
|
||||
Importance: 0.4,
|
||||
Tags: driftTagPrefix + ",domain:tasks,status:drifting",
|
||||
CreatedAt: now.Add(-4 * time.Minute),
|
||||
UpdatedAt: now.Add(-4 * time.Minute),
|
||||
},
|
||||
},
|
||||
}
|
||||
adapter := NewMemoryDriftAdapter(source)
|
||||
|
||||
domains, err := adapter.GetDomains(t.Context(), "agent-1")
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, []string{"tasks"}, domains)
|
||||
|
||||
activity, err := adapter.GetDomainActivity(t.Context(), "agent-1", "tasks", now.Add(-1*time.Hour))
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, 1, activity.MessageCount)
|
||||
require.Equal(t, 1, activity.MemoryCount)
|
||||
}
|
||||
|
||||
func TestMemoryDriftAdapter_PaginatesRecallItems(t *testing.T) {
|
||||
t.Parallel()
|
||||
now := time.Now()
|
||||
items := make([]*memory.RecallItem, 0, driftPageSize+50)
|
||||
for i := 0; i < driftPageSize+50; i++ {
|
||||
items = append(items, &memory.RecallItem{
|
||||
ID: ids.New(),
|
||||
AgentID: "agent-1",
|
||||
SessionKey: "session-1",
|
||||
Role: "assistant",
|
||||
Sector: memory.Sector("tasks"),
|
||||
Importance: 0.5,
|
||||
CreatedAt: now.Add(-time.Duration(i) * time.Second),
|
||||
UpdatedAt: now.Add(-time.Duration(i) * time.Second),
|
||||
})
|
||||
}
|
||||
source := &fakeDriftSource{
|
||||
activeAgents: []string{"agent-1"},
|
||||
recallItems: items,
|
||||
}
|
||||
adapter := NewMemoryDriftAdapter(source)
|
||||
|
||||
activity, err := adapter.GetDomainActivity(t.Context(), "agent-1", "tasks", now.Add(-24*time.Hour))
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, driftPageSize+50, activity.MessageCount)
|
||||
}
|
||||
Loading…
Add table
Reference in a new issue