diff --git a/provider/foundryprovider/memory.go b/provider/foundryprovider/memory.go index a05a32aa..3fa39242 100644 --- a/provider/foundryprovider/memory.go +++ b/provider/foundryprovider/memory.go @@ -23,6 +23,11 @@ const ( defaultMaxMemories = 5 defaultSourceID = "foundrymemory" + // Session-state keys for the static (item-less) memory cache, fetched once + // per session and reused on later turns. + memoryStaticInitStateKey = "foundrymemory.staticInitialized" + memoryStaticMemoriesStateKey = "foundrymemory.staticMemories" + // memorySearchIDStateKey stores the last search id on the session so that // successive turns issue incremental memory searches. memorySearchIDStateKey = "foundrymemory.previousSearchID" @@ -203,34 +208,44 @@ func (p *MemoryProvider) EnsureStoredMemoriesDeleted(ctx context.Context, sessio func (p *MemoryProvider) provide(ctx context.Context, invoking agent.InvokingContext) ([]*message.Message, []agent.Option, error) { telemetry.MarkUsed(telemetry.FeatureFoundryMemory) - items := searchMemoryItems(invoking.Messages) - if len(items) == 0 { - return nil, nil, nil - } session, _ := agent.GetOption(invoking.Options, agent.WithSession) scope := p.scope(session) - searchOptions := &azaiprojects.MemoryStoresClientSearchMemoriesOptions{ - Items: items, - Options: &azaiprojects.MemorySearchResultOptions{MaxMemories: p.config.MaxMemories}, - } - // Thread the previous search id so successive turns perform incremental - // searches, matching the Python provider. - var previousSearchID string - if ok, _ := session.Get(memorySearchIDStateKey, &previousSearchID); ok && previousSearchID != "" { - searchOptions.PreviousSearchID = &previousSearchID - } - result, err := p.client.SearchMemories(ctx, p.memoryStoreName, scope, searchOptions) - if err != nil { - if isCallerCancellation(ctx, err) { - return nil, nil, ctx.Err() + + // Static (e.g. user-profile) memories are fetched once per session with an + // item-less search and prepended to the contextual memories on every turn, + // matching the Python provider. + staticMemories := p.staticMemories(ctx, session, scope) + + var contextual []string + if items := searchMemoryItems(invoking.Messages); len(items) > 0 { + searchOptions := &azaiprojects.MemoryStoresClientSearchMemoriesOptions{ + Items: items, + Options: &azaiprojects.MemorySearchResultOptions{MaxMemories: p.config.MaxMemories}, + } + // Thread the previous search id so successive turns perform incremental + // searches, matching the Python provider. + var previousSearchID string + if ok, _ := session.Get(memorySearchIDStateKey, &previousSearchID); ok && previousSearchID != "" { + searchOptions.PreviousSearchID = &previousSearchID + } + result, err := p.client.SearchMemories(ctx, p.memoryStoreName, scope, searchOptions) + if err != nil { + if isCallerCancellation(ctx, err) { + return nil, nil, ctx.Err() + } + p.log(ctx, slog.LevelError, "foundrymemory: failed to search memories", "memory_store", p.memoryStoreName, "error", err) + } else { + if result.SearchID != nil && *result.SearchID != "" { + session.Set(memorySearchIDStateKey, *result.SearchID) + } + contextual = memoryContents(result.Memories) } - p.log(ctx, slog.LevelError, "foundrymemory: failed to search memories", "memory_store", p.memoryStoreName, "error", err) - return nil, nil, nil - } - if result.SearchID != nil && *result.SearchID != "" { - session.Set(memorySearchIDStateKey, *result.SearchID) } - memories := memoryContents(result.Memories) + + // Copy into a fresh slice so appending never mutates the cached static list. + memories := make([]string, 0, len(staticMemories)+len(contextual)) + memories = append(memories, staticMemories...) + memories = append(memories, contextual...) p.log(ctx, slog.LevelInfo, "foundrymemory: retrieved memories", "memory_store", p.memoryStoreName, "count", len(memories)) if len(memories) == 0 { return nil, nil, nil @@ -239,6 +254,45 @@ func (p *MemoryProvider) provide(ctx context.Context, invoking agent.InvokingCon return []*message.Message{contextMessage}, nil, nil } +// staticMemories returns the scope's static (item-less) memories, fetching them +// once per session and caching them on the session for reuse on later turns. +// Without a session there is nowhere to cache, so it returns nothing rather +// than issuing an item-less search on every turn. +func (p *MemoryProvider) staticMemories(ctx context.Context, session *agent.Session, scope string) []string { + if session == nil { + return nil + } + // Key the cache by store and scope so a session shared across providers, or + // a scope callback whose result changes between turns, does not reuse the + // wrong scope's static memories. + initKey := p.staticStateKey(memoryStaticInitStateKey, scope) + memKey := p.staticStateKey(memoryStaticMemoriesStateKey, scope) + var initialized bool + if ok, _ := session.Get(initKey, &initialized); ok && initialized { + var cached []string + session.Get(memKey, &cached) + return cached + } + result, err := p.client.SearchMemories(ctx, p.memoryStoreName, scope, &azaiprojects.MemoryStoresClientSearchMemoriesOptions{ + Options: &azaiprojects.MemorySearchResultOptions{MaxMemories: p.config.MaxMemories}, + }) + if err != nil { + // Leave the session uninitialized so the next turn retries. + p.log(ctx, slog.LevelError, "foundrymemory: failed to search static memories", "memory_store", p.memoryStoreName, "error", err) + return nil + } + static := memoryContents(result.Memories) + session.Set(memKey, static) + session.Set(initKey, true) + return static +} + +// staticStateKey qualifies a static-memory session-state key with the store and +// scope so cached state is not shared across providers or scopes on one session. +func (p *MemoryProvider) staticStateKey(prefix, scope string) string { + return prefix + ":" + p.memoryStoreName + ":" + scope +} + func (p *MemoryProvider) store(ctx context.Context, invoked agent.InvokedContext) error { telemetry.MarkUsed(telemetry.FeatureFoundryMemory) items := updateMemoryItems(invoked.RequestMessages) diff --git a/provider/foundryprovider/memory_test.go b/provider/foundryprovider/memory_test.go index 0602397f..63b945f1 100644 --- a/provider/foundryprovider/memory_test.go +++ b/provider/foundryprovider/memory_test.go @@ -171,6 +171,72 @@ func TestNewMemoryProviderStoreFiltersDefaultToExternalOnlyRequestAndPassThrough } } +// Static (user-profile) memories are fetched once per session with an item-less +// search and prepended to the contextual memories on every turn, matching the +// Python provider. +func TestMemoryProviderInjectsStaticMemories(t *testing.T) { + transport := &recordingTransport{} + transport.handle = func(req *http.Request, body string) (*http.Response, error) { + // A populated items array marks the contextual search; the item-less + // static search sends no items. + if strings.Contains(body, `"items":[{`) { + return jsonResponse(req, http.StatusOK, `{"memories":[{"memory_item":{"content":"contextual memory","kind":"procedural","memory_id":"mem_c","scope":"user-456","updated_at":1}}]}`), nil + } + return jsonResponse(req, http.StatusOK, `{"memories":[{"memory_item":{"content":"static profile","kind":"user_profile","memory_id":"mem_s","scope":"user-456","updated_at":1}}]}`), nil + } + provider := foundryprovider.NewMemoryProvider(validEndpoint, validCredential, "memory", validScope, foundryprovider.MemoryProviderConfig{ + ClientOptions: azcore.ClientOptions{Transport: transport}, + ContextPrompt: new("Memories:"), + }) + session := &agent.Session{} + + // First turn: item-less static search + contextual search; both surfaced. + messages, _, err := provider.Invoking(t.Context(), agent.InvokingContext{ + Messages: []*message.Message{message.NewText("what do you remember?")}, + Options: []agent.Option{agent.WithSession(session)}, + }) + if err != nil { + t.Fatalf("Invoking error = %v", err) + } + // Static memories are prepended before the contextual ones, so assert order, + // not just presence. + if ctx := messagesString(messages); !strings.Contains(ctx, "static profile\ncontextual memory") { + t.Fatalf("first-turn context = %q, want static memory prepended before contextual", ctx) + } + + // Second turn: static is cached, so no second item-less search is issued, + // but the static memory is still prepended. + messages, _, err = provider.Invoking(t.Context(), agent.InvokingContext{ + Messages: []*message.Message{message.NewText("again")}, + Options: []agent.Option{agent.WithSession(session)}, + }) + if err != nil { + t.Fatalf("Invoking error = %v", err) + } + if ctx := messagesString(messages); !strings.Contains(ctx, "static profile") { + t.Fatalf("second-turn context = %q, want static memory prepended", ctx) + } + + // One item-less static search total (turn 1) + one contextual search per turn. + itemless := 0 + for _, r := range transport.Requests() { + if !strings.Contains(r.Body, `"items":[{`) { + itemless++ + } + } + if itemless != 1 { + t.Fatalf("item-less static searches = %d, want 1 (cached after first turn)", itemless) + } +} + +func messagesString(messages []*message.Message) string { + var b strings.Builder + for _, m := range messages { + b.WriteString(m.String()) + } + return b.String() +} + func TestMemoryProviderPanicsWhenScopeIsEmptyOnUse(t *testing.T) { provider := foundryprovider.NewMemoryProvider(validEndpoint, validCredential, "memory", func(*agent.Session) string { return " " }, foundryprovider.MemoryProviderConfig{}) assertPanics(t, func() { @@ -270,15 +336,23 @@ func TestMemoryProviderThreadsPreviousSearchID(t *testing.T) { } } - requests := transport.Requests() - if len(requests) != 2 { - t.Fatalf("request count = %d, want 2", len(requests)) + // Filter to the contextual (item-bearing) searches; the static + // (user-profile) search issued once per session is item-less and does not + // participate in incremental search-id threading. + var contextual []string + for _, r := range transport.Requests() { + if strings.Contains(r.Body, `"items":[{`) { + contextual = append(contextual, r.Body) + } + } + if len(contextual) != 2 { + t.Fatalf("contextual request count = %d, want 2", len(contextual)) } - if got := jsonMap(t, requests[0].Body)["previous_search_id"]; got != nil { - t.Fatalf("first request previous_search_id = %#v, want absent", got) + if got := jsonMap(t, contextual[0])["previous_search_id"]; got != nil { + t.Fatalf("first contextual request previous_search_id = %#v, want absent", got) } - if got := jsonMap(t, requests[1].Body)["previous_search_id"]; got != "search_1" { - t.Fatalf("second request previous_search_id = %#v, want search_1", got) + if got := jsonMap(t, contextual[1])["previous_search_id"]; got != "search_1" { + t.Fatalf("second contextual request previous_search_id = %#v, want search_1", got) } }