MuyueWorkspace/internal/memory/inject.go

package memory

import (
	"fmt"
	"strings"
	"time"
)

type MemoryInjector struct {
	store *Store
}

func NewInjector(store *Store) *MemoryInjector {
	return &MemoryInjector{store: store}
}

func (mi *MemoryInjector) BuildContextBlock(query string) (string, error) {
	var contextParts []string

	preferences, err := mi.store.RecallPreferences()
	if err == nil && len(preferences) > 0 {
		var prefLines []string
		for _, p := range preferences {
			prefLines = append(prefLines, fmt.Sprintf("- %s: %s", p.Key, p.Content))
		}
		contextParts = append(contextParts,
			"[User Preferences]\n"+strings.Join(prefLines, "\n"))
	}

	facts, err := mi.store.RecallFacts()
	if err == nil && len(facts) > 0 {
		var factLines []string
		for _, f := range facts {
			factLines = append(factLines, fmt.Sprintf("- %s: %s", f.Key, f.Content))
		}
		contextParts = append(contextParts,
			"[Known Facts]\n"+strings.Join(factLines, "\n"))
	}

	if query != "" {
		relevant, err := mi.store.Recall(query, 5)
		if err == nil && len(relevant) > 0 {
			var relLines []string
			for _, r := range relevant {
				relLines = append(relLines, fmt.Sprintf("- [%s] %s: %s", r.Type, r.Key, truncate(r.Content, 150)))
			}
			contextParts = append(contextParts,
				"[Relevant Memories]\n"+strings.Join(relLines, "\n"))
		}
	}

	recentCutoff := time.Now().Add(-24 * time.Hour)
	recent, err := mi.store.RecallRecent(recentCutoff, 5)
	if err == nil && len(recent) > 0 {
		var recentLines []string
		for _, r := range recent {
			recentLines = append(recentLines, fmt.Sprintf("- [%s] %s", r.Type, truncate(r.Content, 100)))
		}
		contextParts = append(contextParts,
			"[Recent Context]\n"+strings.Join(recentLines, "\n"))
	}

	if len(contextParts) == 0 {
		return "", nil
	}

	return fmt.Sprintf("<memory-context>\n[System note: NOT new user input — recalled context]\n%s\n</memory-context>",
		strings.Join(contextParts, "\n\n")), nil
}

func (mi *MemoryInjector) BuildSystemPromptBlock() (string, error) {
	preferences, err := mi.store.RecallPreferences()
	if err != nil || len(preferences) == 0 {
		return "", nil
	}

	var lines []string
	lines = append(lines, "Known user preferences:")
	for _, p := range preferences {
		lines = append(lines, fmt.Sprintf("- %s: %s", p.Key, p.Content))
	}

	return strings.Join(lines, "\n"), nil
}

func (mi *MemoryInjector) ExtractAndStore(userMessage, assistantMessage string) error {
	pref := extractPreference(userMessage)
	if pref != "" {
		if err := mi.store.StorePreference("detected", pref); err != nil {
			return fmt.Errorf("store preference: %w", err)
		}
	}

	if assistantMessage != "" {
		ctx := extractContext(assistantMessage)
		if ctx != "" {
			if err := mi.store.StoreContext("conversation", ctx); err != nil {
				return fmt.Errorf("store context: %w", err)
			}
		}
	}

	return nil
}

func extractPreference(message string) string {
	indicators := []string{
		"i prefer", "i like", "i always", "i never", "my favorite",
		"i use", "je préfère", "j'aime", "toujours", "jamais",
	}
	lower := strings.ToLower(message)
	for _, ind := range indicators {
		if strings.Contains(lower, ind) {
			idx := strings.Index(lower, ind)
			end := idx + len(ind) + 100
			if end > len(message) {
				end = len(message)
			}
			return truncate(message[idx:end], 200)
		}
	}
	return ""
}

func extractContext(message string) string {
	if len(message) < 50 {
		return ""
	}
	if len(message) > 500 {
		return truncate(message, 500)
	}
	return message
}

func truncate(s string, maxLen int) string {
	if len(s) <= maxLen {
		return s
	}
	return s[:maxLen] + "..."
}