feat: support user-specific memory overlays in prompt context (#995)
This commit is contained in:
parent
4a7605ee14
commit
4bea9128d2
2 changed files with 81 additions and 0 deletions
|
|
@ -435,6 +435,47 @@ func (cb *ContextBuilder) buildDynamicContext(channel, chatID string) string {
|
|||
return sb.String()
|
||||
}
|
||||
|
||||
|
||||
func safePathSegment(v string) (string, bool) {
|
||||
v = strings.TrimSpace(v)
|
||||
if v == "" || v == "." || v == ".." {
|
||||
return "", false
|
||||
}
|
||||
if strings.ContainsAny(v, `/\`) {
|
||||
return "", false
|
||||
}
|
||||
return v, true
|
||||
}
|
||||
|
||||
// loadUserMemoryContext loads optional per-user memory content.
|
||||
// Supported lookup order:
|
||||
// 1) <workspace>/users/<channel>/<chatID>/MEMORY.md
|
||||
// 2) <workspace>/users/<chatID>/MEMORY.md
|
||||
func (cb *ContextBuilder) loadUserMemoryContext(channel, chatID string) string {
|
||||
channelSeg, okChannel := safePathSegment(channel)
|
||||
chatSeg, okChat := safePathSegment(chatID)
|
||||
if !okChat {
|
||||
return ""
|
||||
}
|
||||
|
||||
candidates := []string{}
|
||||
if okChannel {
|
||||
candidates = append(candidates, filepath.Join(cb.workspace, "users", channelSeg, chatSeg, "MEMORY.md"))
|
||||
}
|
||||
candidates = append(candidates, filepath.Join(cb.workspace, "users", chatSeg, "MEMORY.md"))
|
||||
|
||||
for _, p := range candidates {
|
||||
if data, err := os.ReadFile(p); err == nil {
|
||||
text := strings.TrimSpace(string(data))
|
||||
if text != "" {
|
||||
return "## User-specific Memory\n\n" + text
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
return ""
|
||||
}
|
||||
|
||||
func (cb *ContextBuilder) BuildMessages(
|
||||
history []providers.Message,
|
||||
summary string,
|
||||
|
|
@ -457,6 +498,7 @@ func (cb *ContextBuilder) BuildMessages(
|
|||
|
||||
// Build short dynamic context (time, runtime, session) — changes per request
|
||||
dynamicCtx := cb.buildDynamicContext(channel, chatID)
|
||||
userMemoryCtx := cb.loadUserMemoryContext(channel, chatID)
|
||||
|
||||
// Compose a single system message: static (cached) + dynamic + optional summary.
|
||||
// Keeping all system content in one message ensures every provider adapter can
|
||||
|
|
@ -474,6 +516,11 @@ func (cb *ContextBuilder) BuildMessages(
|
|||
{Type: "text", Text: dynamicCtx},
|
||||
}
|
||||
|
||||
if userMemoryCtx != "" {
|
||||
stringParts = append(stringParts, userMemoryCtx)
|
||||
contentBlocks = append(contentBlocks, providers.ContentBlock{Type: "text", Text: userMemoryCtx})
|
||||
}
|
||||
|
||||
if summary != "" {
|
||||
summaryText := fmt.Sprintf(
|
||||
"CONTEXT_SUMMARY: The following is an approximate summary of prior conversation "+
|
||||
|
|
|
|||
|
|
@ -667,3 +667,37 @@ func BenchmarkBuildMessagesWithCache(b *testing.B) {
|
|||
_ = cb.BuildMessages(history, "summary", "new message", nil, "cli", "test")
|
||||
}
|
||||
}
|
||||
|
||||
func TestBuildMessages_IncludesUserSpecificMemory(t *testing.T) {
|
||||
tmpDir := setupWorkspace(t, map[string]string{
|
||||
"IDENTITY.md": "# Identity\nTest agent.",
|
||||
"users/webchat/user123/MEMORY.md": "favorite color: blue",
|
||||
})
|
||||
defer os.RemoveAll(tmpDir)
|
||||
|
||||
cb := NewContextBuilder(tmpDir)
|
||||
msgs := cb.BuildMessages(nil, "", "hello", nil, "webchat", "user123")
|
||||
if len(msgs) == 0 || msgs[0].Role != "system" {
|
||||
t.Fatalf("expected first system message")
|
||||
}
|
||||
if !strings.Contains(msgs[0].Content, "User-specific Memory") {
|
||||
t.Fatalf("expected user memory section in system prompt")
|
||||
}
|
||||
if !strings.Contains(msgs[0].Content, "favorite color: blue") {
|
||||
t.Fatalf("expected user memory content in system prompt")
|
||||
}
|
||||
}
|
||||
|
||||
func TestBuildMessages_UserSpecificMemoryIgnoredForInvalidChatID(t *testing.T) {
|
||||
tmpDir := setupWorkspace(t, map[string]string{
|
||||
"IDENTITY.md": "# Identity\nTest agent.",
|
||||
"users/webchat/u1/MEMORY.md": "secret",
|
||||
})
|
||||
defer os.RemoveAll(tmpDir)
|
||||
|
||||
cb := NewContextBuilder(tmpDir)
|
||||
msgs := cb.BuildMessages(nil, "", "hello", nil, "webchat", "../u1")
|
||||
if strings.Contains(msgs[0].Content, "User-specific Memory") {
|
||||
t.Fatalf("expected no user memory section for invalid chat id")
|
||||
}
|
||||
}
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue